import os
from datetime import timedelta
from time import time
from typing import Optional, Union, List, Dict

from lightning.pytorch import Trainer, Callback
from lightning.pytorch.accelerators import Accelerator
from lightning.pytorch.callbacks import ModelCheckpoint, Checkpoint
from lightning.pytorch.loggers import Logger, TensorBoardLogger
from lightning.pytorch.plugins import PLUGIN_INPUT
from lightning.pytorch.profilers import Profiler
from lightning.pytorch.strategies import Strategy
from lightning.pytorch.trainer.connectors.accelerator_connector import _LITERAL_WARN
from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
from lightning.fabric.utilities.types import _PATH

from .callback import NilCallback, LrLogger, EpochNumLogger, EpochTimeLogger, GpuStatsLogger, SystemStatsLogger


class NekoTrainer(Trainer):

    def __init__(self,
        logger: Optional[str],
        enable_checkpointing: bool = True,
        callbacks: Optional[Union[List[Callback], Callback]] = None,
        log_every_n_steps: int = 50,
        precision: _PRECISION_INPUT = "32-true",
        accelerator: Union[str, Accelerator] = "auto",
        strategy: Union[str, Strategy] = "auto",
        devices: Union[List[int], str, int] = "auto",
        num_nodes: int = 1,
        fast_dev_run: Union[int, bool] = False,
        max_epochs: Optional[int] = None,
        min_epochs: Optional[int] = None,
        max_steps: int = -1,
        min_steps: Optional[int] = None,
        max_time: Optional[Union[str, timedelta, Dict[str, int]]] = None,
        limit_train_batches: Optional[Union[int, float]] = None,
        limit_val_batches: Optional[Union[int, float]] = None,
        limit_test_batches: Optional[Union[int, float]] = None,
        limit_predict_batches: Optional[Union[int, float]] = None,
        overfit_batches: Union[int, float] = 0.0,
        val_check_interval: Optional[Union[int, float]] = None,
        check_val_every_n_epoch: Optional[int] = 1,
        num_sanity_val_steps: Optional[int] = None,
        enable_progress_bar: Optional[bool] = None,
        enable_model_summary: Optional[bool] = None,
        accumulate_grad_batches: int = 1,
        gradient_clip_val: Optional[Union[int, float]] = None,
        gradient_clip_algorithm: Optional[str] = None,
        deterministic: Optional[Union[bool, _LITERAL_WARN]] = None,
        benchmark: Optional[bool] = True,
        inference_mode: bool = True,
        use_distributed_sampler: bool = True,
        profiler: Optional[Union[Profiler, str]] = None,
        detect_anomaly: bool = False,
        barebones: bool = False,
        plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]] = None,
        sync_batchnorm: bool = False,
        reload_dataloaders_every_n_epochs: int = 0,
        default_root_dir: Optional[_PATH] = None,
    ):
        # define a random logger name
        self.log_name = f"{logger}_{int(time())}"
        if callbacks is None:
            callbacks = []
        # build checkpoint callback or from user defined
        if enable_checkpointing and len([c for c in callbacks if isinstance(c, Checkpoint)]) == 0:
            # use default checkpoint callback
            new_callback = ModelCheckpoint(
                dirpath=os.path.join("logs", self.log_name, "checkpoints"),
                save_last=True, filename="{epoch}-{val_loss:.3f}", monitor="val_loss",
                mode="min"
            )
        elif enable_checkpointing and len([c for c in callbacks if isinstance(c, Checkpoint)]) > 0:
            # use user defined checkpoint callback
            new_callback = NilCallback()
        elif not enable_checkpointing and len([c for c in callbacks if isinstance(c, Checkpoint)]) > 0:
            # raise error because checkpoint callback is not allowed
            raise ValueError("Checkpoint callback is not allowed when checkpointing is disabled.")
        else:
            # if no checkpoint callback is defined, skip
            new_callback = NilCallback()

        # merge the checkpoint callback with other callbacks
        if callbacks is None:
            cbs = [new_callback]
        elif isinstance(callbacks, Callback):
            cbs = [callbacks, new_callback]
        else:
            cbs = callbacks + [new_callback]

        # set learning rate logger
        cbs.extend([LrLogger(), EpochNumLogger(), EpochTimeLogger()])

        # if gpumonitor is installed, enable callback
        try:
            gpu_cb = GpuStatsLogger()
        except ImportError:
            pass
        else:
            cbs.append(gpu_cb)

        # if psutil is installed, enable callback
        try:
            sys_cb = SystemStatsLogger()
        except ImportError:
            pass
        else:
            cbs.append(sys_cb)

        # setup the log mode
        self.log_every_n_steps = log_every_n_steps
        self.log_on_epoch = log_every_n_steps == 0
        self.log_on_step = log_every_n_steps > 0
        if self.log_on_epoch:
            log_every_n_steps = 10000000

        super().__init__(
            logger=logger is not None,
            enable_checkpointing=enable_checkpointing,
            callbacks=cbs,
            accelerator=accelerator,
            strategy=strategy,
            devices=devices,
            num_nodes=num_nodes,
            precision=precision,
            fast_dev_run=fast_dev_run,
            max_epochs=max_epochs,
            min_epochs=min_epochs,
            max_steps=max_steps,
            min_steps=min_steps,
            max_time=max_time,
            limit_train_batches=limit_train_batches,
            limit_val_batches=limit_val_batches,
            limit_test_batches=limit_test_batches,
            limit_predict_batches=limit_predict_batches,
            overfit_batches=overfit_batches,
            val_check_interval=val_check_interval,
            check_val_every_n_epoch=check_val_every_n_epoch,
            num_sanity_val_steps=num_sanity_val_steps,
            log_every_n_steps=log_every_n_steps,
            enable_progress_bar=enable_progress_bar,
            enable_model_summary=enable_model_summary,
            accumulate_grad_batches=accumulate_grad_batches,
            gradient_clip_val=gradient_clip_val,
            gradient_clip_algorithm=gradient_clip_algorithm,
            deterministic=deterministic,
            benchmark=benchmark,
            inference_mode=inference_mode,
            use_distributed_sampler=use_distributed_sampler,
            profiler=profiler,
            detect_anomaly=detect_anomaly,
            barebones=barebones,
            plugins=plugins,
            sync_batchnorm=sync_batchnorm,
            reload_dataloaders_every_n_epochs=reload_dataloaders_every_n_epochs,
            default_root_dir=default_root_dir,
        )

        # super().__init__(
        #     logger=logger is not None,
        #     enable_checkpointing=enable_checkpointing,
        #     callbacks=cbs,
        #     default_root_dir=default_root_dir,
        #     gradient_clip_val=gradient_clip_val,
        #     gradient_clip_algorithm=gradient_clip_algorithm,
        #     num_nodes=num_nodes,
        #     num_processes=num_processes,
        #     devices=devices,
        #     gpus=gpus,
        #     auto_select_gpus=auto_select_gpus,
        #     tpu_cores=tpu_cores,
        #     ipus=ipus,
        #     enable_progress_bar=enable_progress_bar,
        #     overfit_batches=overfit_batches,
        #     track_grad_norm=track_grad_norm,
        #     check_val_every_n_epoch=check_val_every_n_epoch,
        #     fast_dev_run=fast_dev_run,
        #     accumulate_grad_batches=accumulate_grad_batches,
        #     max_epochs=max_epochs,
        #     min_epochs=min_epochs,
        #     max_steps=max_steps,
        #     min_steps=min_steps,
        #     max_time=max_time,
        #     limit_train_batches=limit_train_batches,
        #     limit_val_batches=limit_val_batches,
        #     limit_test_batches=limit_test_batches,
        #     limit_predict_batches=limit_predict_batches,
        #     val_check_interval=val_check_interval,
        #     log_every_n_steps=log_every_n_steps,
        #     accelerator=accelerator,
        #     strategy=strategy,
        #     sync_batchnorm=sync_batchnorm,
        #     precision=precision,
        #     enable_model_summary=enable_model_summary,
        #     weights_save_path=weights_save_path,
        #     num_sanity_val_steps=num_sanity_val_steps,
        #     resume_from_checkpoint=resume_from_checkpoint,
        #     profiler=profiler,
        #     benchmark=benchmark,
        #     deterministic=deterministic,
        #     reload_dataloaders_every_n_epochs=reload_dataloaders_every_n_epochs,
        #     auto_lr_find=auto_lr_find,
        #     replace_sampler_ddp=replace_sampler_ddp,
        #     detect_anomaly=detect_anomaly,
        #     auto_scale_batch_size=auto_scale_batch_size,
        #     plugins=plugins,
        #     amp_backend=amp_backend,
        #     amp_level=amp_level,
        #     move_metrics_to_cpu=move_metrics_to_cpu,
        #     multiple_trainloader_mode=multiple_trainloader_mode
        # )
        # init loggers for training and validation
        self.has_no_logger = logger is None

        self.logger_train = TensorBoardLogger(save_dir=self.default_root_dir, name="logs",
            version=os.path.join(self.log_name, "train"), log_graph=True
        ) if self.has_no_logger is not None else None
        self.logger_val = TensorBoardLogger(save_dir=self.default_root_dir, name="logs",
            version=os.path.join(self.log_name, "val"), log_graph=False
        ) if self.has_no_logger is not None else None
        self._loggers = []

    @property
    def logger(self) -> Optional[Logger]:
        if self.has_no_logger:
            return None
        if self.training:
            return self.logger_train
        else:
            return self.logger_val
