import os
import sys
from datetime import datetime

import torch
import pickle
from tqdm import tqdm

path = os.path.abspath(__file__)
for i in range(2):
    path = os.path.dirname(path)
PROJECT_ROOT = path
sys.path.append(PROJECT_ROOT)

from configs.config import ConfigManager
from datasets.dataset_factory import DatasetFactory
from modules.pipeline import Pipeline
from runners.runner import Runner


class InferenceEngine(Runner):
    RESULTS_DIR_NAME: str = 'results'
    LOG_DIR_NAME: str = 'log'

    def __init__(self, config_path):
        super().__init__(config_path)

        ''' Pipeline '''
        self.pipeline_config = ConfigManager.get("settings", "pipeline")
        self.pipeline = Pipeline(self.pipeline_config).to(self.device)

        ''' Experiment '''
        self.model_path = ConfigManager.get("settings", "experiment", "model_path")
        self.load_checkpoint(self.model_path)
        self.load_experiment("inference")
        
        ''' Inference Results '''
        self.inference_results_config = ConfigManager.get("settings", "results")
        self.save_data_keys = self.inference_results_config["save_data_keys"]
        self.save_output_keys = self.inference_results_config["save_output_keys"]

        ''' Test '''
        self.test_config = ConfigManager.get("settings", "test")
        self.test_dataset_config_list = self.test_config["dataset_list"]
        self.test_set_list = []
        seen_name = set()
        for test_dataset_config in self.test_dataset_config_list:
            if test_dataset_config["name"] not in seen_name:
                seen_name.add(test_dataset_config["name"])
            else:
                raise ValueError("Duplicate test dataset name: {}".format(test_dataset_config["name"]))
            test_set = DatasetFactory.create(test_dataset_config)
            self.test_set_list.append(test_set)
        del seen_name

        self.print_info()

    def run(self):
        print("Inference start...")
        self.test()
        print("Inference finished!")

    def test(self):
        self.pipeline.eval()
        with torch.no_grad():
            for dataset_idx, test_set in enumerate(self.test_set_list):
                test_set_name = self.test_dataset_config_list[dataset_idx]["name"]
                ratio = self.test_dataset_config_list[dataset_idx]["ratio"]

                test_loader = test_set.get_loader()
                loop = tqdm(enumerate(test_loader), total=int(len(test_loader)))
                for i, data in loop:
                    test_set.process_batch(data, self.device)
                    output = self.pipeline(data, Pipeline.TEST_MODE)
                    self.save_output(output, data, test_set_name, i)
                    loop.set_description(
                        f'Inference (Test: {test_set_name}, ratio={ratio})')
                    
    def save_output(self, output, data, test_set_name, idx):
        results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
        if not os.path.exists(os.path.join(results_dir,test_set_name)):
            os.makedirs(os.path.join(results_dir,test_set_name))
        save_path = os.path.join(results_dir, test_set_name, f"{idx}.pkl")
        data = {key: value for key, value in data.items() if key in self.save_data_keys}
        output = {key: value for key, value in output.items() if key in self.save_output_keys}
        output_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in output.items()}
        data_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in data.items()}
        with open(save_path, "wb") as f:
            pickle.dump({"output":output_converted,"data":data_converted}, f)
        
    def load_checkpoint(self, model_path):
        self.pipeline.load(model_path)
        print(f"Checkpoint loaded from {model_path}")

    def load_experiment(self, backup_name=None):
        super().load_experiment(backup_name)
        
    def create_experiment(self, backup_name=None):
        super().create_experiment(backup_name)
        results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
        os.makedirs(results_dir)
        

    def print_info(self):
        def print_dataset(config, dataset):
            print("\t name: {}".format(config["name"]))
            print("\t source: {}".format(config["source"]))
            print("\t data_type: {}".format(config["data_type"]))
            print("\t total_length: {}".format(len(dataset)))
            print("\t ratio: {}".format(config["ratio"]))
            print()

        super().print_info()
        table_size = 70
        print(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
        print(self.pipeline)
        print(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
        for i, test_dataset_config in enumerate(self.test_dataset_config_list):
            print(f"test dataset {i}: ")
            print_dataset(test_dataset_config, self.test_set_list[i])

        print(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="configs/local_inference_config.yaml")
    args = parser.parse_args()
    infenrence_engine = InferenceEngine(args.config)
    infenrence_engine.run()