import typing
import copy
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Literal,
    Optional,
    Sequence,
    TypeAlias,
    TypeVar,
    Union,
)

from lightning.fabric import Fabric
from lightning.fabric.wrappers import _FabricModule, _FabricOptimizer
from torch import nn
from torch.optim import Optimizer

from zyplib.nn.utils import load_checkpoint
from zyplib.train.common import callbacks
from zyplib.train.metric.torch_metric import TorchMetricRecorder
from zyplib.utils.print import print_colored, warning

if TYPE_CHECKING:
    from torch import Tensor
    from torch.optim.lr_scheduler import _LRScheduler
    from torch.utils.data import DataLoader

# 类型定义
TensorOrTensors: TypeAlias = Union['Tensor', Sequence['Tensor']]
Batch: TypeAlias = tuple['Tensor', ...]
Metrics: TypeAlias = Dict[str, float]
SingleOptimizer: TypeAlias = Optimizer | _FabricOptimizer
OptimizerGroup: TypeAlias = (
    SingleOptimizer
    | tuple[SingleOptimizer, ...]
    | list[SingleOptimizer]
    | dict[str, SingleOptimizer]
)
T = TypeVar('T')


@dataclass
class FabricTrainingContext:
    """训练上下文对象，用于存储训练相关的状态和配置

    提供了对 Lightning Fabric 回调函数的封装，使得在不使用 train_loop 时也能方便地复用回调函数工具。

    Attributes:
    ----------
    - `model` : `nn.Module | _FabricModule`
        - 模型
    - `optimizer` : `Optimizer | tuple[Optimizer, ...] | list[Optimizer] | dict[str, Optimizer]`
        - 优化器，支持单个优化器或多个优化器的组合形式
    - `fabric` : `Fabric | None`, optional
        - 可选项, 如果传入了, 则会调用 fabric 的 `backward` 和 `callback`
    - `loss_fn` : `nn.Module | None`, optional
        - 损失函数，如果传入了则会作为默认损失函数使用，默认为 None
        - 可以通过函数参数动态设置和更新
    - `train_metric_recorder` : `TorchMetricRecorder | None`, optional
        - 用于记录训练时指标的 recorder 实例，默认为 None
    - `val_metric_recorder` : `TorchMetricRecorder | None`, optional
        - 用于记录验证时指标的 recorder 实例，默认为 None

    Methods:
    --------
    回调函数封装方法：
    - `call_on_train_start()` - 训练开始时的回调
    - `call_on_train_end()` - 训练结束时的回调
    - `call_on_epoch_start()` - 每个epoch开始时的回调
    - `call_on_epoch_end()` - 每个epoch结束时的回调
    - `call_on_train_batch_end()` - 每个训练批次结束时的回调
    - `call_on_validation_batch_end()` - 每个验证批次结束时的回调
    - `call_on_test_start()` - 测试开始时的回调
    - `call_on_test_end()` - 测试结束时的回调
    - 以及其他相关的回调函数封装...
    """

    model: nn.Module | _FabricModule
    optimizer: OptimizerGroup
    fabric: Optional[Fabric] = None
    fn_loss: Optional[nn.Module] = None
    train_metric_recorder: Optional[TorchMetricRecorder] = None
    val_metric_recorder: Optional[TorchMetricRecorder] = None

    @property
    def pure_model(self) -> nn.Module:
        """返回 unwrap 后的原始模型"""

        if self.fabric and isinstance(self.model, _FabricModule):
            return self.model.module
        return self.model

    @property
    def pure_optimizer(
        self,
    ) -> Optimizer | typing.Iterator[Optimizer] | dict[str, Optimizer]:
        """返回 unwrap 后的原始优化器"""

        def unwrap(optimizer: SingleOptimizer) -> Optimizer:
            if isinstance(optimizer, _FabricOptimizer):
                return optimizer._optimizer
            return optimizer

        if isinstance(self.optimizer, (tuple, list)):
            return (unwrap(o) for o in self.optimizer)
        elif isinstance(self.optimizer, dict):
            return {k: unwrap(v) for k, v in self.optimizer.items()}
        else:
            return unwrap(self.optimizer)

    def _require_fabric(self) -> Fabric:
        return typing.cast(Fabric, self.fabric)

    def __post_init__(self):
        if self.fabric is None:
            return

        # 检查模型是否已 setup
        if not isinstance(self.model, _FabricModule):
            self.model = self.fabric.setup(self.model)

        # 检查优化器是否已 setup
        if isinstance(self.optimizer, Optimizer) and (
            not isinstance(self.optimizer, _FabricOptimizer)
        ):
            # 如果是单个优化器
            self.optimizer = typing.cast(
                SingleOptimizer, self.fabric.setup_optimizers(self.optimizer)
            )
        elif isinstance(self.optimizer, (list, tuple)):
            # 如果是列表/元组，检查其中每个优化器
            are_optimizers_setup = all(
                isinstance(opt, _FabricOptimizer) for opt in self.optimizer
            )
            if not are_optimizers_setup:
                raw_optimizers = [
                    opt._optimizer if isinstance(opt, _FabricOptimizer) else opt
                    for opt in self.optimizer
                ]
                self.optimizer = typing.cast(
                    OptimizerGroup, self.fabric.setup_optimizers(*raw_optimizers)
                )
        elif isinstance(self.optimizer, dict):
            for key, optimizer in self.optimizer.items():
                if isinstance(optimizer, Optimizer) and not isinstance(
                    optimizer, _FabricOptimizer
                ):
                    self.optimizer[key] = typing.cast(
                        SingleOptimizer, self.fabric.setup_optimizers(optimizer)
                    )

    def launch(self):
        fabric = self._require_fabric()
        fabric.launch()
        if isinstance(self.train_metric_recorder, TorchMetricRecorder):
            self.train_metric_recorder.to_device(fabric.device)
        if isinstance(self.val_metric_recorder, TorchMetricRecorder):
            self.val_metric_recorder.to_device(fabric.device)

    def find_best_ckpt(self):
        fabric = self._require_fabric()
        save_ckpt_cb: callbacks.SaveCheckpoints | None = callbacks.find_callback(
            fabric, callbacks.SaveCheckpoints
        )
        if not save_ckpt_cb:
            warning('当前训练没有注册 SaveCheckpoints 回调函数, 无法加载最佳模型')
            return None
        return save_ckpt_cb.best_checkpoint()

    def load_best_ckpt(self, model: Literal['replace', 'return'] = 'replace'):
        """导入最佳模型

        Parameters
        ----------
        - `model` : `Literal['replace', 'return']`
            - 'replace': 直接替换当前模型
            - 'return': 返回另一个加载后的模型
        """
        path = self.find_best_ckpt()
        print_colored('magenta', f'Load from best checkpoint path: {path}')
        if not path:
            warning('当前训练没有注册 SaveCheckpoints 回调函数, 无法加载最佳模型')
            return None
        fabric = self._require_fabric()
        if model == 'replace':
            load_checkpoint(self.model, str(path), device=fabric.device)
            return self.model
        elif model == 'return':
            new_model = copy.deepcopy(self.pure_model)
            load_checkpoint(new_model, str(path), device=fabric.device)
            return new_model

    # ===================== 回调函数封装 ===================== #

    def call_on_train_start(
        self,
        *,
        lr_scheduler: Optional['_LRScheduler'] = None,
        max_epochs: Optional[int] = None,
        dataloader: Optional['DataLoader'] = None,
        hparams: Optional[typing.Mapping[str, Any]] = None,
        **kwargs,
    ):
        """封装 on_train_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_train_start',
                model=self.pure_model,
                optimizer=self.pure_optimizer,
                lr_scheduler=lr_scheduler,
                fabric=self.fabric,
                max_epochs=max_epochs,
                dataloader=dataloader,
                hparams=hparams,
                **kwargs,
            )

    def call_on_train_end(self, **kwargs):
        """封装 on_train_end 回调"""
        if self.fabric:
            self.fabric.call('on_train_end', **kwargs)

    def call_on_test_start(self, *, dataloader: Optional['DataLoader'] = None, **kwargs):
        """封装 on_test_start 回调"""
        if self.fabric:
            self.fabric.call('on_test_start', dataloader=dataloader, **kwargs)

    def call_on_test_end(
        self, *, metrics: Optional[typing.Mapping[str, float]] = None, **kwargs
    ):
        """封装 on_test_end 回调"""
        if self.fabric:
            self.fabric.call('on_test_end', metrics=metrics, **kwargs)

    # --------------- Epoch wise --------------- #

    def call_on_epoch_start(self, epoch: int, max_epochs: int, **kwargs):
        """封装 on_epoch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_epoch_start', epoch=epoch, max_epochs=max_epochs, **kwargs
            )

    def call_on_epoch_end(
        self,
        epoch: int,
        *,
        train_metrics: Optional[typing.Mapping[str, float]] = None,
        val_metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """封装 on_epoch_end 回调"""
        if self.fabric:
            self.fabric.call(
                'on_epoch_end',
                epoch=epoch,
                train_metrics=train_metrics,
                val_metrics=val_metrics,
                **kwargs,
            )

    def call_on_train_epoch_start(self, epoch: int, max_epochs: int, **kwargs):
        """封装 on_train_epoch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_train_epoch_start', epoch=epoch, max_epochs=max_epochs, **kwargs
            )

    def call_on_train_epoch_end(
        self,
        epoch: int,
        *,
        metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """封装 on_train_epoch_end 回调"""
        if self.fabric:
            self.fabric.call('on_train_epoch_end', epoch=epoch, metrics=metrics, **kwargs)

    def call_on_validation_epoch_start(
        self, epoch: int, *, epoch_interval: Optional[int] = None, **kwargs
    ):
        """封装 on_validation_epoch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_validation_epoch_start',
                epoch=epoch,
                epoch_interval=epoch_interval,
                **kwargs,
            )

    def call_on_validation_epoch_end(
        self, epoch: int, *, metrics: typing.Mapping[str, float], **kwargs
    ):
        """封装 on_validation_epoch_end 回调"""
        if self.fabric:
            self.fabric.call(
                'on_validation_epoch_end', epoch=epoch, metrics=metrics, **kwargs
            )

    # --------------- Batch wise --------------- #

    def call_on_train_batch_start(
        self, batch: 'TensorOrTensors', batch_idx: int, **kwargs
    ):
        """封装 on_train_batch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_train_batch_start', batch=batch, batch_idx=batch_idx, **kwargs
            )

    def call_on_train_batch_end(
        self,
        batch: 'TensorOrTensors',
        batch_idx: int,
        *,
        metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """封装 on_train_batch_end 回调"""
        if self.fabric:
            self.fabric.call(
                'on_train_batch_end',
                batch=batch,
                batch_idx=batch_idx,
                metrics=metrics,
                **kwargs,
            )

    def call_on_validation_batch_start(
        self, batch: 'TensorOrTensors', batch_idx: int, **kwargs
    ):
        """封装 on_validation_batch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_validation_batch_start', batch=batch, batch_idx=batch_idx, **kwargs
            )

    def call_on_validation_batch_end(
        self,
        batch: 'TensorOrTensors',
        batch_idx: int,
        *,
        metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """封装 on_validation_batch_end 回调"""
        if self.fabric:
            self.fabric.call(
                'on_validation_batch_end',
                batch=batch,
                batch_idx=batch_idx,
                metrics=metrics,
                **kwargs,
            )

    def call_on_test_batch_start(
        self, batch: 'TensorOrTensors', batch_idx: int, **kwargs
    ):
        """封装 on_test_batch_start 回调"""
        if self.fabric:
            self.fabric.call(
                'on_test_batch_start', batch=batch, batch_idx=batch_idx, **kwargs
            )

    def call_on_test_batch_end(
        self,
        batch: 'TensorOrTensors',
        batch_idx: int,
        *,
        metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """封装 on_test_batch_end 回调"""
        if self.fabric:
            self.fabric.call(
                'on_test_batch_end',
                batch=batch,
                batch_idx=batch_idx,
                metrics=metrics,
                **kwargs,
            )

    # ===================== Transform Hooks ===================== #

    def transform_on_train_batch_start(
        self, batch: 'TensorOrTensors', batch_idx: int, epoch_idx: Optional[int] = None
    ) -> 'TensorOrTensors':
        """训练批次开始前的数据转换 hook（默认实现：恒等变换）

        可在子类中重写以实现数据增强（如 Mixup, CutMix）

        Parameters
        ----------
        - `batch` : `TensorOrTensors`
            - 原始批次数据
        - `batch_idx` : `int`
            - 批次索引
        - `epoch_idx` : `Optional[int]`, optional
            - 当前 epoch 索引

        Returns
        -------
        - `TensorOrTensors`
            - 转换后的批次数据
        """
        return batch

    def transform_on_validation_batch_start(
        self, batch: 'TensorOrTensors', batch_idx: int, epoch_idx: Optional[int] = None
    ) -> 'TensorOrTensors':
        """验证批次开始前的数据转换 hook（默认实现：恒等变换）

        可在子类中重写以实现测试时增强（TTA）

        Parameters
        ----------
        - `batch` : `TensorOrTensors`
            - 原始批次数据
        - `batch_idx` : `int`
            - 批次索引
        - `epoch_idx` : `Optional[int]`, optional
            - 当前 epoch 索引

        Returns
        -------
        - `TensorOrTensors`
            - 转换后的批次数据
        """
        return batch
