import logging
import os
import re
import time
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

from pyroml.core.autocast import Autocast
from pyroml.core.hparams import WithHyperParameters
from pyroml.core.model import PyroModel
from pyroml.loop import EvalLoop, PredictLoop, TrainLoop
from pyroml.utils import get_classname
from pyroml.utils.log import get_logger, set_level_all_loggers

if TYPE_CHECKING:
    from pyroml.callbacks import Callback
    from pyroml.core.tracker import MetricsTracker
    from pyroml.loop.base import Loop

log = get_logger(__name__)

TRAINER_STATE_FILE = "trainer_state.pt"
TRAINER_HPARAMS_FILE = "trainer_hparams.json"


def get_date() -> str:
    return time.strftime("%Y-%m-%d_%H:%M", time.gmtime(time.time()))


class NoPyroModelException(Exception):
    pass


class Trainer(WithHyperParameters):
    def __init__(
        self,
        lr: float = 1e-4,
        max_epochs: int | None = 8,
        max_steps: int | None = None,
        grad_norm_clip: float = 1.0,
        evaluate_on: Literal["step", "epoch", False] = "epoch",
        evaluate_every: int = 1,
        eval_max_steps: int | None = None,
        device: str | torch.device = "auto",
        auto_move: bool = True,
        pin_memory: bool | None = None,
        dtype: torch.dtype = torch.float32,
        compile: bool = False,
        checkpoint_folder: str | Path = "./checkpoints",
        hparams_file: str | Path = TRAINER_HPARAMS_FILE,
        batch_size: int = 16,
        eval_batch_size: int = None,
        num_workers: int = 0,
        eval_num_workers: int = 0,
        wandb: bool = False,
        wandb_project: str | None = None,
        log_level: int = logging.INFO,
        callbacks: list["Callback"] = [],
    ):
        """
        Trainer class specifying all training related parameters and exposing fit / test methods to interact with the model.

        Args:
            lr (float, optional):
                Learning rate to use for your optimizer / scheduler
                Defaults to 1e-4.

            max_epochs (int, > 0, optional):
                Number of epochs (if max_iterations is not defined).
                Defaults to 8.

            max_steps (int, > 0):
                Maximum number of iterations.
                *Note:* This parameters dominates max_epochs; if number of steps reaches max_steps, training will stop despite max_epochs not being reached.
                For this reason, we recommend using the loop step related attributes (e.g. loop.total_steps, loop.steps_per_epochs) in your model's
                configure_optimizer method instead of epoch related attributes (e.g. loop.max_epochs)
                Defaults to None.

            grad_norm_clip (float, optional):
                Gradient norm clipping.
                Defaults to 1.0.

            evaluate_on (Literal['step', 'epoch', False], optional):
                Evaluate every `evaluate_every` step or epoch. Setting this to False will disable evaluation.
                Defaults to 'epoch'.

            evaluate_every (int, optional):
                Evaluate model every `evaluate_every` `evaluate_on`. Setting this to a value <= 0 will disable evaluation.
                Defaults to 1.

            eval_max_steps (int, optional):
                Maximum number of iterations for the evaluation dataset. Setting this to a value <= 0 will disable evaluation.
                Defaults to None.

            dtype (torch.dtype, optional):
                Data type to cast model weights to.
                Defaults to torch.float32.

            device (str, optional):
                Device to train on.
                Defaults to "auto" which will use GPU if available.

            auto_move (bool, optional):
                Automatically move batches to device before calling your model's step method. If set to False, your model has to move data to the appropriate device in the step function.
                Defaults to True.

            pin_memory (bool, optional):
                Applies pin_memory to dataloaders. Set this to None to make pin_memory true only if training on gpu.
                Default to None.

            compile (bool, optional):
                Whether to compile the model, this can significantly improve inference time but is not supported on all GPUs.
                Defaults to False.

            checkpoint_folder (str, optional):
                Folder to save checkpoints.
                Defaults to "./checkpoints".

            hparams_file (str, optional):
                File to save hyperparameters.
                Defaults to Checkpoint.TRAINER_HPARAMS.value.

            batch_size (int, optional):
                Training batch size. If eval_batch_size is None, evaluation batches will use this value.
                Defaults to 16.

            eval_batch_size (int, optional):
                Batch size for the evaluation dataset. Setting this to a value <= 0 will disable evaluation.
                Defaults to None in which case it will be equal to the training batch size.

            num_workers (int, optional):
                Number of workers for the dataloader.
                Defaults to 0.

            eval_num_workers (int, optional):
                Number of workers for the evaluation dataloader.
                Note that a value > 0 can cause an AssertionError: 'can only test a child process during evaluation'.
                Defaults to 0.

            wandb (bool, optional):
                Whether to use WandB for logging training metrics.
                Defaults to False.

            wandb_project (str, optional):
                Wandb project name, if wandb is set to True.
                Defaults to None.

            log_level (int, optional):
                Logger level. Use logging.X to get the integer corresponding to the level you want
                Defaults to logging.INFO

            callbacks (list[Callback], optional):
                List of callbacks to use.
                Defaults to [].
        Returns:
            Config: Configuration object with the specified hyperparameters.
        """
        super().__init__(hparams_file=hparams_file)

        # Training
        self.lr = lr
        self.max_epochs = max_epochs
        self.max_steps = max_steps
        self.grad_norm_clip = grad_norm_clip
        self.callbacks = callbacks

        # Validation
        assert evaluate_on == "step" or evaluate_on == "epoch" or evaluate_on is False
        self.evaluate_on = evaluate_on
        self.evaluate_every = evaluate_every
        self.eval_max_steps = eval_max_steps

        # Model
        self.dtype = dtype
        self.device = device
        self.compile = compile
        self.version = -1
        self.checkpoint_folder = Path(checkpoint_folder)
        self.model: "PyroModel" | torch._dynamo.OptimizedModule = None

        # Data
        self.batch_size = batch_size
        self.eval_batch_size = eval_batch_size or batch_size
        self.num_workers = num_workers
        self.eval_num_workers = eval_num_workers

        # Logging
        self.wandb = wandb
        self.wandb_project = self._get_wandb_project(wandb, wandb_project)
        self.log_level = log_level

        # Device and dtypes
        self.autocast = Autocast(self)
        self.auto_move = auto_move
        self.pin_memory = pin_memory

        self.date: str
        self.loop: "Loop"

    @property
    def eval_enabled(self):
        return (
            self.evaluate_on is not False
            and self.evaluate_every > 0
            and (self.eval_max_steps is None or self.eval_max_steps > 0)
            and (self.eval_batch_size is None or self.eval_batch_size > 0)
        )

    def log(self, **data: dict[str, float | np.ndarray | torch.Tensor]):
        assert self.loop is not None
        return self.loop.log(**data)

    def _setup(self):
        self.date = get_date()
        self.loop: "Loop" = None

        set_level_all_loggers(level=self.log_level)

    def _get_wandb_project(self, wandb: bool, wandb_project: str | None):
        project = wandb_project or os.environ.get("WANDB_PROJECT")
        if wandb and (project == "" or project is None):
            msg = "Wandb project name is required, please set WANDB_PROJECT in your environment variables or pass wandb_project in the Trainer constructor"
            raise ValueError(msg)
        return project

    def _fetch_version(self):
        os.makedirs(self.checkpoint_folder, exist_ok=True)
        for f in os.scandir(self.checkpoint_folder):
            if not f.is_dir() or not re.match(r"v_(\d+)", f.name):
                continue
            version = int(f.name[2:])
            self.version = max(self.version, version)

        self.version += 1
        self.checkpoint_folder = self.checkpoint_folder / f"v_{self.version}"

    def _setup_model(self, model: "PyroModel"):
        self.model = model

        # Compile model if requested, should improve performance
        if self.compile:
            self.model.compile()

        self.model._setup(self)

    def _call_loop(
        self, Loop: "Loop", model: "PyroModel", dataset: Dataset, **kwargs
    ) -> pd.DataFrame:
        """ "
        Setups the trainer and model and
        Calls a loop with the specified model and dataset.
        """
        if not isinstance(model, PyroModel) and (
            not isinstance(model, torch._dynamo.OptimizedModule)
            or not isinstance(model._orig_mod, PyroModel)
        ):
            raise NoPyroModelException("Trainer: model must be a PyroModel")

        self._setup()
        self.loop: "Loop" = Loop(model=model, trainer=self, dataset=dataset, **kwargs)
        self._setup_model(model)
        return self.loop.run()

    def fit(
        self, model: "PyroModel", tr_dataset: Dataset, ev_dataset: Dataset = None
    ) -> "MetricsTracker":
        return self._call_loop(
            Loop=TrainLoop, model=model, dataset=tr_dataset, ev_dataset=ev_dataset
        )

    def evaluate(self, model: "PyroModel", dataset: Dataset) -> "MetricsTracker":
        return self._call_loop(Loop=EvalLoop, model=model, dataset=dataset)

    def predict(
        self, model: "PyroModel", dataset: Dataset
    ) -> tuple["MetricsTracker", Any]:
        return self._call_loop(Loop=PredictLoop, model=model, dataset=dataset)

    def save_state(
        self,
        folder: Path | str = None,
        file: Path | str = TRAINER_STATE_FILE,
    ) -> None:
        # Saving state
        state = {
            "compiled": self.model._is_compiled(),
            "optimizer_name": get_classname(self.model.optimizer),
            "optimizer": self.model.optimizer.state_dict(),
        }

        if hasattr(self.model, "scheduler"):
            state["scheduler_name"] = get_classname(self.model.scheduler)
            state["scheduler"] = self.model.scheduler.state_dict()

        folder = Path(folder) or self.checkpoint_folder
        os.makedirs(folder, exist_ok=True)
        torch.save(obj=state, f=folder / file)

    def load_state(
        self,
        folder: Path | str = None,
        file: Path | str = TRAINER_STATE_FILE,
    ) -> None:
        folder = Path(folder) or self.checkpoint_folder
        state = torch.load(folder / file, map_location="cpu")

        self.model.optimizer.load_state_dict(state["optimizer"])
        if hasattr(self, "scheduler") and "scheduler" in state:
            self.model.scheduler.load_state_dict(state["scheduler"])

    def save(self, model: "PyroModel"):
        if self.version < 0:
            self._fetch_version()
        self.save_state(self.checkpoint_folder)
        model.save(self.checkpoint_folder)

    @staticmethod
    def load(
        folder: Path | str,
        file: Path | str = TRAINER_HPARAMS_FILE,
    ):
        """
        Loads a pretrained model from a specified checkpoint folder.

        Args:
            folder (str): The folder path where the pretrained model is saved.
            resume (bool, optional): Whether to resume training from the checkpoint. Defaults to True.
            strict (bool, optional): Whether to strictly enforce the shape and type of the loaded weights. Defaults to True.
        """
        """
        Loads a trainer state from a specified checkpoint folder.

        Args:
            folder (str): The folder path where the pretrained model is saved.
        """
        folder = Path(folder)
        hparams_file = folder if os.path.isfile(folder) else folder / file
        args = WithHyperParameters._load_hparams(hparams_file)

        # TODO: find a way to reload the optimizer and scheduler to the model
        # optim = getattr(torch.optim, args.optimizer)
        # sched = getattr(torch.optim.lr_scheduler, args.scheduler)
        # Should be loaded in the model
        print(args)

        # TODO: callbacks state saved with WithHyperparameters
        warnings.warn(
            "Loading callbacks is not supported, if you have custom callbacks please add them to the trainer again"
        )

        trainer = Trainer(**args)
        return trainer
