import torch
from jsonargparse import class_from_function
from lightning.pytorch.callbacks import (
    EarlyStopping,
    LearningRateMonitor,
    ModelCheckpoint,
    StochasticWeightAveraging,
)
from lightning.pytorch.cli import LightningCLI, _get_short_description
from lightning.pytorch.trainer import Trainer

from scprint2.tasks import (
    Denoiser,
    Embedder,
    FinetuneBatchClass,
    GeneEmbeddingExtractor,
    GNInfer,
    Imputer,
)

from .trainer import TrainingMode

TASKS = [
    ("embed", Embedder),
    ("gninfer", GNInfer),
    ("denoise", Denoiser),
    ("finetune", FinetuneBatchClass),
    ("gene_emb", GeneEmbeddingExtractor),
    ("impute", Imputer),
]
TIMEOUT = 3600 * 9  # 9 hours in seconds


class MyCLI(LightningCLI):
    """
    MyCLI is a subclass of LightningCLI to add some missing params and create bindings between params of the model and the data.

    Used to allow calling denoise / embed / gninfer from the command line.
    Also to add more parameters and link parameters between the scdataloader and the scPRINT-2 model.
    """

    def add_arguments_to_parser(self, parser):
        parser.add_lightning_class_args(EarlyStopping, "scprint_early_stopping")
        parser.set_defaults(
            {
                "scprint_early_stopping.monitor": "val_loss",
                # patience is the number of consecutive epochs with no improvement after which learning rate will be reduced.
                "scprint_early_stopping.patience": 5,
                # Prevent early stopping from checking on train epoch end when val_loss isn't available yet
                "scprint_early_stopping.check_on_train_epoch_end": False,
            }
        )
        parser.add_lightning_class_args(
            LearningRateMonitor, "scprint_learning_rate_monitor"
        )
        parser.set_defaults({"scprint_learning_rate_monitor.logging_interval": "epoch"})
        parser.add_lightning_class_args(TrainingMode, "scprint_training")
        parser.link_arguments("data.genes_dict", "model.genes", apply_on="instantiate")
        parser.link_arguments(
            "data.decoders", "model.label_decoders", apply_on="instantiate"
        )
        parser.link_arguments(
            "data.labels_hierarchy", "model.labels_hierarchy", apply_on="instantiate"
        )
        parser.link_arguments(
            "data.n_bins", "model.n_input_bins", apply_on="instantiate"
        )
        parser.link_arguments("data.classes", "model.classes", apply_on="instantiate")
        parser.link_arguments(
            "data.organisms", "model.organisms", apply_on="instantiate"
        )

        parser.add_argument("--set_float32_matmul_precision", type=bool, default=False)
        parser.add_argument("--wandblog", type=str, default="")
        parser.add_argument("--project", type=str)

    def _add_subcommands(self, parser, **kwargs) -> None:
        """Adds subcommands to the input parser."""
        self._subcommand_parsers = {}
        parser_subcommands = parser.add_subcommands()
        # the user might have passed a builder function
        trainer_class = (
            self.trainer_class
            if isinstance(self.trainer_class, type)
            else class_from_function(self.trainer_class)
        )
        # register all subcommands in separate subcommand parsers under the main parser
        for subcommand in self.subcommands():
            fn = getattr(trainer_class, subcommand)
            # extract the first line description in the docstring for the subcommand help message
            description = _get_short_description(fn)
            subparser_kwargs = kwargs.get(subcommand, {})
            subparser_kwargs.setdefault("description", description)
            subcommand_parser = self._prepare_subcommand_parser(
                trainer_class, subcommand, **subparser_kwargs
            )
            self._subcommand_parsers[subcommand] = subcommand_parser
            parser_subcommands.add_subcommand(
                subcommand, subcommand_parser, help=description
            )
        description = (
            "setup everything for scprint2. only the default behavior is supported."
        )
        parser = self.init_parser(description=description)
        self._subcommand_parsers["easy_setup"] = parser
        parser_subcommands.add_subcommand("easy_setup", parser, help=description)
        for subcommand in TASKS:
            fn = getattr(subcommand[1], "__init__")
            description = _get_short_description(fn)
            subparser_kwargs = {}
            subparser_kwargs.setdefault("description", description)
            parser = self.init_parser(**subparser_kwargs)
            parser.add_argument(
                "--ckpt_path",
                type=str,
                help=("Path to the checkpoint to load."),
                required=True,
            )
            parser.add_argument(
                "--output_filename",
                type=str,
                help=("Path to the output file(s)."),
                required=True,
            )
            parser.add_argument(
                "--adata", type=str, help=("Path to the anndata file."), required=True
            )
            if subcommand[0] == "gninfer":
                parser.add_argument(
                    "--cell_type",
                    type=str,
                    help=("The cell type to infer the gene regulatory network."),
                    required=True,
                )
            parser.add_argument("--seed_everything", type=int, default=42)
            parser.add_argument(
                "--species",
                type=str,
                help=(
                    "If not included in the anndata under 'organism_ontology_term_id', the species of the dataset."
                ),
                required=False,
            )
            parser.add_argument("--use_raw", type=bool, default=False)
            parser.add_argument("--skip_validate", type=bool, default=True)
            parser.add_argument("--is_symbol", type=bool, default=False)
            parser.add_class_arguments(subcommand[1])
            added = parser.add_method_arguments(
                subcommand[1],
                "__call__",
                skip=set(["model", "adata", "cell_type"]),
            )

            self._subcommand_method_arguments[subcommand] = added
            self._subcommand_parsers[subcommand[0]] = parser
            parser_subcommands.add_subcommand(subcommand[0], parser, help=description)

    def _run_subcommand(self, subcommand: str) -> None:
        """Run the chosen subcommand."""
        before_fn = getattr(self, f"before_{subcommand}", None)
        if callable(before_fn):
            before_fn()
        if subcommand in self.subcommands():
            klass = self.trainer
            default = getattr(klass, subcommand)
            fn = getattr(self, subcommand, default)
            fn_kwargs = self._prepare_subcommand_kwargs(subcommand)
            fn(**fn_kwargs)
            after_fn = getattr(self, f"after_{subcommand}", None)
            if callable(after_fn):
                after_fn()
        elif subcommand == "easy_setup":
            from huggingface_hub import hf_hub_download
            from scdataloader.utils import (
                _adding_scbasecamp_genes,
                populate_my_ontology,
            )

            populate_my_ontology()
            _adding_scbasecamp_genes()
            hf_hub_download(
                repo_id="jkobject/scPRINT", filename="small-v2.ckpt", local_dir="."
            )
            print("finished command")
            return
        else:
            import numpy as np
            import scanpy as sc
            from scdataloader import Preprocessor
            from scdataloader.utils import load_genes

            from scprint2 import scPRINT2

            model_checkpoint_file = self.config_init[subcommand]["ckpt_path"]
            model = scPRINT2.load_from_checkpoint(
                model_checkpoint_file, precpt_gene_emb=None, gene_pos_file=None
            )

            adata = sc.read_h5ad(self.config_init[subcommand]["adata"])
            adata.obs.drop(columns="is_primary_data", inplace=True, errors="ignore")
            if self.config_init[subcommand]["species"] is not None:
                adata.obs["organism_ontology_term_id"] = self.config_init[subcommand][
                    "species"
                ]
            preprocessor = Preprocessor(
                do_postp=model.expr_emb_style == "metacell",
                force_preprocess=True,
                skip_validate=self.config_init[subcommand].get("skip_validate", True),
                use_raw=self.config_init[subcommand].get("use_raw", False),
                is_symbol=self.config_init[subcommand].get("is_symbol", False),
            )
            adata = preprocessor(adata)
            if model.expr_emb_style == "metacell":
                sc.pp.neighbors(adata, use_rep="X_pca")
            conf = dict(self.config_init[subcommand])
            if model.cell_transformer is None:
                if conf.get("ft_mode", "xpressor"):
                    print(
                        "model doesn't have xpressor, setting finetuning mode to 'full' model"
                    )
                conf["ft_mode"] = "full"
            missing = set(model.genes) - set(load_genes(model.organisms).index)
            if len(missing) > 0:
                print(
                    "Warning: some genes missmatch exist between model and ontology: solving...",
                )
                model._rm_genes(missing)
            if not torch.cuda.is_available():
                model = model.to(torch.float32)

            model = model.to("cuda" if torch.cuda.is_available() else "cpu")
            for key in [
                "seed_everything",
                "config",
                "species",
                "cache",
                "ckpt_path",
                "adata",
                "output_filename",
                "cell_type",
                "use_raw",
                "is_symbol",
                "skip_validate",
                "train_data",
                "val_data",
            ]:
                conf.pop(key, None)

            if subcommand == "embed":
                emb = Embedder(**conf)
                print("embedding...")
                adata, metrics = emb(
                    model=model,
                    adata=adata,
                )
                # save the anndata
                print("metrics:")
                print(metrics)
                print()
                name = (
                    self.config_init[subcommand]["output_filename"] + "_embedded.h5ad"
                )
                print(
                    "saving the file under the path: ",
                    name,
                )
                adata.write(name)
            elif subcommand == "gninfer":
                gn = GNInfer(**conf)
                adata = gn(
                    model=model,
                    adata=adata,
                    cell_type=self.config_init[subcommand]["cell_type"],
                )
                name = (
                    self.config_init[subcommand]["output_filename"]
                    + "_"
                    + self.config_init[subcommand]["cell_type"]
                    + "_grn.h5ad"
                )
                print(
                    "saving the file under the path: ",
                    name,
                )
                adata.write(name)

            elif subcommand == "denoise":
                dn = Denoiser(**conf)
                metrics, _, expr_pred = dn(
                    model=model,
                    adata=adata,
                )
                print("metrics:")
                print(metrics)
                print()
                # now what we are doing here it to complete the expression profile with the denoised values.
                # this is not done by default for now
                name = (
                    self.config_init[subcommand]["output_filename"] + "_denoised.h5ad"
                )
                print(
                    "saving the file under the path: ",
                    name,
                )
                expr_pred.write(name)
            elif subcommand == "finetune":
                ft = FinetuneBatchClass(**conf)
                model = ft(
                    model=model,
                    adata=adata,
                )
                name = (
                    self.config_init[subcommand]["output_filename"] + "_finetuned.ckpt"
                )
                print(
                    "saving the model under the path: ",
                    name,
                )
                model.save_checkpoint(name)
            elif subcommand == "gene_emb":
                ge = GeneEmbeddingExtractor(**conf)
                ada = ge(
                    model=model,
                    adata=adata,
                )
                name = (
                    self.config_init[subcommand]["output_filename"] + "_gene_emb.h5ad"
                )
                print(
                    "saving the file under the path: ",
                    name,
                )
                ada.write(name)
            elif subcommand == "impute":
                imp = Imputer(**conf)
                _, expr_pred = imp(
                    model=model,
                    adata=adata,
                )
                name = self.config_init[subcommand]["output_filename"] + "_imputed.h5ad"
                print(
                    "saving the file under the path: ",
                    name,
                )
                expr_pred.write(name)
            else:
                raise ValueError(f"Unknown subcommand: {subcommand}")

    def before_instantiate_classes(self):
        for k, v in self.config.items():
            if "set_float32_matmul_precision" in k:
                if v:
                    torch.set_float32_matmul_precision("medium")
        import os

        simpler_f_hash = os.system("cd ../simpler_flash && git rev-parse HEAD")
        scdata_v = os.system("cd ../scdataloader && git rev-parse HEAD")
        print("scdata_v:", scdata_v)
        print("simpler_f_hash:", simpler_f_hash)
        if "fit" in self.config and self.config["fit"]["trainer"]["strategy"] in [
            "ddp",
            "ddp_find_unused_parameters_true",
            "fsdp",
        ]:
            import os

            os.environ["NCCL_TIMEOUT"] = str(TIMEOUT)  # 9 hours in seconds
            os.environ["TORCH_DISTRIBUTED_TIMEOUT"] = str(TIMEOUT)  # 9 hours in seconds
            os.environ["PL_TRAINER_STRATEGY_TIMEOUT"] = str(TIMEOUT)

            print(f"setting global pytorch distributed timeout to {TIMEOUT}s")

    def after_instantiate_classes(self):
        try:
            self.model.name = self.trainer._loggers[0].version
        except:
            print("not on wandb, could not set name")
        if hasattr(self, "datamodule") and self.datamodule is not None:
            self.datamodule.set_valid_genes_collator(self.model.genes)

    def instantiate_trainer(self, **kwargs) -> Trainer:
        """Override to customize trainer instantiation"""
        # Modify strategy if it's DDP
        trainer = super().instantiate_trainer(**kwargs)
        if "fit" in self.config and self.config["fit"]["trainer"]["strategy"] in [
            "ddp",
            "ddp_find_unused_parameters_true",
            "fsdp",
        ]:
            # Create DDPStrategy with custom timeout
            from datetime import timedelta

            # Update the config
            print("updating the config")
            trainer.strategy._timeout = timedelta(seconds=TIMEOUT)  # 9 hours in seconds
            trainer.strategy.setup_distributed()
        # Call parent method to create trainer
        return trainer
