import os
import json
import numpy as np
from abc import abstractmethod, ABC

from runners.preprocessor import Preprocessor
from utils.omni_util import OmniUtil

class RGBFeatPreprocessor(Preprocessor, ABC):

    def __init__(self, config_path):
        super().__init__(config_path)
        self.load_experiment("RGBFeat")
        self.dataset_list_config = self.preprocess_config["dataset_list"]
        self.model_config = self.preprocess_config["model"]

    def run(self):
        """
        - for each dataset
        --- get its dataloader
        --- for each batch, do prediction
        --- preprocess the collected results
        --- save processed results
        """
        for dataset_config in self.dataset_list_config:
            dataloader = self.get_dataloader(dataset_config)
            model = self.get_model(self.model_config)
            predicted_data = self.prediction(model, dataloader)
            self.save_processed_data(predicted_data,dataset_config)

    def preprocess(self, predicted_data):
        pass
    
    def save_processed_data(self, processed_data, data_config=None):
        data_path = os.path.join(str(self.experiment_path), Preprocessor.DATA, data_config["source"], data_config["data_type"])
        for frame_path in processed_data:
            rgb_feat = processed_data[frame_path]
            scene = os.path.basename(os.path.dirname(frame_path))
            idx = os.path.basename(frame_path)
            target_scene_path = os.path.join(str(data_path), scene) 
            if not os.path.exists(target_scene_path):
                os.makedirs(target_scene_path)
            rgb_feat_save_path = os.path.join(
                target_scene_path,OmniUtil.RGB_FEAT_TEMPLATE.format(idx))
            np.save(rgb_feat_save_path, rgb_feat)
        print("Processed data saved to: ", data_path)

