import ruamel.yaml
import torch
import typer

from copy import deepcopy
from jinja2 import Template
from pathlib import Path
from pydantic import BaseModel, Field
from typing import Optional

from ...utils.factory import ApplyPipelineFactory, PipelineBase, DataFactory, GraphModelFactory
from ...utils.yaml_dump import deep_convert_dict, merge_comment

pipeline_comments = {
    "batch_size": "Graph batch size",
    "num_workers": "Number of workers for data loading",
    "save_path": "Directory to save the inference results"
}

class ApplyGraphpredPipelineCfg(BaseModel):
    batch_size: int = 32
    num_workers: int = 4
    save_path: str = "apply_results"

@ApplyPipelineFactory.register("graphpred")
class ApplyGraphpredPipeline(PipelineBase):
    def __init__(self):
        self.pipeline = {
            "name": "graphpred",
            "mode": "apply"
        }

    @classmethod
    def setup_user_cfg_cls(cls):
        from ...utils.enter_config import UserConfig
        class ApplyGraphPredUserConfig(UserConfig):
            data: DataFactory.filter("graphpred").get_pydantic_config() = Field(..., discriminator="name")
            general_pipeline: ApplyGraphpredPipelineCfg = ApplyGraphpredPipelineCfg()

        cls.user_cfg_cls = ApplyGraphPredUserConfig

    @property
    def user_cfg_cls(self):
        return self.__class__.user_cfg_cls

    def get_cfg_func(self):
        def config(
            data: DataFactory.filter("graphpred").get_dataset_enum() = typer.Option(None, help="input data name"),
            cfg: Optional[str] = typer.Option(None, help="output configuration file path"),
            cpt: str = typer.Option(..., help="input checkpoint file path")
        ):
            # Training configuration
            train_cfg = torch.load(cpt)["cfg"]
            if data is None:
                print("data is not specified, use the training dataset")
                data = train_cfg["data_name"]
            else:
                data = data.name
            if cfg is None:
                cfg = "_".join(["apply", "graphpred", data, train_cfg["model_name"]]) + ".yaml"

            self.__class__.setup_user_cfg_cls()
            generated_cfg = {
                "pipeline_name": self.pipeline["name"],
                "pipeline_mode": self.pipeline["mode"],
                "device": train_cfg["device"],
                "data": {"name": data},
                "cpt_path": cpt,
                "general_pipeline": {"batch_size": train_cfg["general_pipeline"]["eval_batch_size"],
                                     "num_workers": train_cfg["general_pipeline"]["num_workers"]}
            }
            output_cfg = self.user_cfg_cls(**generated_cfg).dict()
            output_cfg = deep_convert_dict(output_cfg)
            # Not applicable for inference
            output_cfg['data'].pop('split_ratio')
            comment_dict = {
                "device": "Torch device name, e.g., cpu or cuda or cuda:0",
                "cpt_path": "Path to the checkpoint file",
                "general_pipeline": pipeline_comments
            }
            comment_dict = merge_comment(output_cfg, comment_dict)

            yaml = ruamel.yaml.YAML()
            yaml.dump(comment_dict, Path(cfg).open("w"))
            print("Configuration file is generated at {}".format(Path(cfg).absolute()))

        return config

    @classmethod
    def gen_script(cls, user_cfg_dict):
        # Check validation
        cls.setup_user_cfg_cls()
        cls.user_cfg_cls(**user_cfg_dict)

        # Training configuration
        train_cfg = torch.load(user_cfg_dict["cpt_path"])["cfg"]

        # Dict for code rendering
        render_cfg = deepcopy(user_cfg_dict)
        model_name = train_cfg["model_name"]
        model_code = GraphModelFactory.get_source_code(model_name)
        render_cfg["model_code"] = model_code
        render_cfg["model_class_name"] = GraphModelFactory.get_model_class_name(model_name)
        render_cfg.update(DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"]))

        # Dict for defining cfg in the rendered code
        generated_user_cfg = deepcopy(user_cfg_dict)
        generated_user_cfg.pop("pipeline_name")
        generated_user_cfg.pop("pipeline_mode")
        # model arch configuration
        generated_user_cfg["model"] = train_cfg["model"]

        render_cfg["user_cfg_str"] = f"cfg = {str(generated_user_cfg)}"
        render_cfg["user_cfg"] = user_cfg_dict

        file_current_dir = Path(__file__).resolve().parent
        with open(file_current_dir / "graphpred.jinja-py", "r") as f:
            template = Template(f.read())

        return template.render(**render_cfg)

    @staticmethod
    def get_description() -> str:
        return "Graph classification pipeline for inference on binary classification"
