import typing
from copy import deepcopy
from typing import Any, Callable, Dict, List, Literal, Optional, TypeAlias

import torch
from lightning.fabric import Fabric
from torch import nn
from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import DataLoader
from torchmetrics import Metric, MetricCollection

from zyplib.train.common import procedure
from zyplib.train.common._base import FabricTrainingContext
from zyplib.train.common.callbacks import FabricLogMetrics
from zyplib.train.metric.torch_metric import TorchMetricRecorder
from zyplib.utils.print import print_debug, print_info

Batch: TypeAlias = tuple[torch.Tensor, ...]


class BasicFabricTrainer:
    """基于 procedure 模块的训练器封装

    内部维护 FabricTrainingContext，支持两种使用模式：
    1. 直接使用（标准任务）
    2. 继承重写方法（training_step, train_epoch 等）

    所有训练逻辑委托给 procedure 模块，Trainer 只提供可重写的接入点。

    Examples
    --------

    ```python
    # 模式 1: 直接使用（简单分类任务）
    trainer = BasicFabricTrainer(
        fabric, model, optimizer,
        fn_loss=nn.BCEWithLogitsLoss(),
        metric_builder=clf_metrics_builder,
        before_metric='sigmoid',
    )
    trainer.fit(train_loader, val_loader, max_epochs=10)

    # 模式 2: 继承重写（高度自定义，如 GAN）
    class GANTrainer(BasicFabricTrainer):
        def training_step(self, batch, batch_idx, epoch_idx):
            # 自定义 GAN 训练逻辑
            return fake_data, {'g_loss': g_loss, 'd_loss': d_loss}

    # 模式 3: 动态设置 step 函数（如 EDL）
    trainer = BasicFabricTrainer(fabric, model, optimizer, ...)
    edl_procedure = BasicEDLFabricTrainProcedure(...)
    trainer.set_training_step(edl_procedure.evidential_train_step)
    trainer.set_validation_step(edl_procedure.evidential_val_step)
    trainer.fit(train_loader, val_loader)
    ```
    """

    def __init__(
        self,
        fabric: Fabric,
        model: nn.Module,
        optimizer: torch.optim.Optimizer,
        *,
        fn_loss: nn.Module | None = None,
        hparams: typing.Mapping[str, Any] | None = None,
        metric_builder: Callable[[], MetricCollection | typing.Mapping[str, Metric]]
        | None = None,
        before_metric: Literal['sigmoid', 'softmax']
        | Callable[[torch.Tensor], torch.Tensor]
        | typing.Mapping[str | Literal['*'], Callable]
        | None = None,
        launch_fabric_internally: bool = True,
    ):
        """基于 Lightning Fabric 的基础训练器

        提供了基本的训练、验证功能，支持评估指标记录和学习率调度等功能

        Parameters
        ----------
        - `fabric` : `Fabric`
            - Lightning Fabric 实例
            - 用于管理分布式训练、混合精度等
        - `model` : `nn.Module`
            - PyTorch 模型
        - `optimizer` : `torch.optim.Optimizer`
            - 优化器实例
        - `fn_loss` : `nn.Module`, optional
            - 损失函数
            - 为 None 时需要在 context.fn_loss 中设置或在训练步骤中手动计算
        - `hparams` : `dict[str, Any] | None`, optional
            - 超参数字典
            - 会被深拷贝以避免修改原始值
        - `metric_builder` : `() -> MetricCollection | dict[str, Metric] | None`, optional
            - 评估指标构建函数, 每次调用返回新的评估指标实例，用于训练集和验证集
            - 返回 MetricCollection 或 dict[str, Metric] 类型
        - `before_metric` : `'sigmoid' | 'softmax' | Callable | dict[str | '*', Callable]`, optional
            - 在计算评估指标之前的激活函数: `(y_hat) -> y_activated`
            - 支持预定义的 'sigmoid' 和 'softmax'
            - 支持自定义激活函数
            - 支持为不同指标指定不同的激活函数
        - `launch_fabric_internally` : `bool`, optional
            - 是否在初始化时自动调用 fabric.launch()
            - 默认为 True
        """
        # 创建 metrics
        train_metrics = metric_builder() if metric_builder else None
        val_metrics = metric_builder() if metric_builder else None

        # 创建 context
        self.context = FabricTrainingContext(
            model=model,
            optimizer=optimizer,
            fabric=fabric,
            fn_loss=fn_loss,
            train_metric_recorder=(
                TorchMetricRecorder(train_metrics, before_metric=before_metric)
                if train_metrics
                else None
            ),
            val_metric_recorder=(
                TorchMetricRecorder(val_metrics, before_metric=before_metric)
                if val_metrics
                else None
            ),
        )

        # 保存配置
        self.hparams = deepcopy(hparams) if hparams else None

        # 自定义 step 函数（通过 set_xxx_step 方法动态设置）
        self._fn_train_step: Optional[Callable] = None
        self._fn_val_step: Optional[Callable] = None
        self._fn_test_step: Optional[Callable] = None

        # 自定义 transform 函数（通过 set_xxx_transform 方法动态设置）
        self._fn_transform_train_batch: Optional[Callable] = None
        self._fn_transform_val_batch: Optional[Callable] = None

        if launch_fabric_internally:
            self.context.launch()

        # 确保 FabricLogMetrics 存在
        if not any(isinstance(cb, FabricLogMetrics) for cb in fabric._callbacks):
            fabric._callbacks.append(FabricLogMetrics(fabric))

    # ==================== 动态设置 Step 函数 ==================== #

    def set_training_step(self, fn_train_step: Callable):
        """动态设置训练步骤函数

        Parameters
        ----------
        - `fn_train_step` : `Callable`
            - 自定义训练步骤函数
            - 签名: `fn(ctx, batch, batch_idx, epoch_idx) -> (y_pred, metrics_dict)`
        """
        self._fn_train_step = fn_train_step

    def set_validation_step(self, fn_val_step: Callable):
        """动态设置验证步骤函数

        Parameters
        ----------
        - `fn_val_step` : `Callable`
            - 自定义验证步骤函数
            - 签名: `fn(ctx, batch, batch_idx, epoch_idx) -> (y_pred, metrics_dict)`
        """
        self._fn_val_step = fn_val_step

    def set_test_step(self, fn_test_step: Callable):
        """动态设置测试步骤函数

        Parameters
        ----------
        - `fn_test_step` : `Callable`
            - 自定义测试步骤函数
            - 签名: `fn(ctx, batch, batch_idx) -> (y_pred, metrics_dict)`
        """
        self._fn_test_step = fn_test_step

    def set_transform_train_batch(self, fn_transform: Callable):
        """动态设置训练批次数据转换函数

        Parameters
        ----------
        - `fn_transform` : `Callable`
            - 自定义数据转换函数
            - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
            - 用例: Mixup, CutMix 等数据增强
        """
        self._fn_transform_train_batch = fn_transform

    def set_transform_val_batch(self, fn_transform: Callable):
        """动态设置验证批次数据转换函数

        Parameters
        ----------
        - `fn_transform` : `Callable`
            - 自定义数据转换函数
            - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
            - 用例: 测试时增强 (TTA) 等
        """
        self._fn_transform_val_batch = fn_transform

    # ==================== 属性委托 ==================== #

    @property
    def fabric(self) -> Fabric:
        """返回 Fabric 实例"""
        return typing.cast(Fabric, self.context.fabric)

    def _setup_single_dataloader(self, dataloader: DataLoader) -> DataLoader:
        return typing.cast(DataLoader, self.fabric.setup_dataloaders(dataloader))

    @property
    def model(self) -> nn.Module:
        """返回模型（可能是 wrapped 的）"""
        return self.context.model

    @property
    def optimizer(self):
        """返回优化器（可能是 wrapped 的）"""
        return self.context.optimizer

    @property
    def fn_loss(self) -> Optional[nn.Module]:
        """返回损失函数"""
        return self.context.fn_loss

    # ==================== 可重写的 Transform 方法 ==================== #

    def transform_on_train_batch_start(
        self, batch: Batch, batch_idx: int, epoch_idx: Optional[int] = None
    ) -> Batch:
        """训练批次开始前的数据转换（可重写）

        默认实现：不做任何转换（恒等变换）
        子类可重写此方法实现 Mixup, CutMix 等数据增强

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

        Returns
        -------
        - `Batch`
            - 转换后的数据批次

        Examples
        --------
        ```python
        class MixupTrainer(BasicFabricTrainer):
            def transform_on_train_batch_start(self, batch, batch_idx, epoch_idx):
                x, y = batch
                x, y = mixup(x, y, alpha=0.2)
                return x, y
        ```
        """
        return batch

    def transform_on_validation_batch_start(
        self, batch: Batch, batch_idx: int, epoch_idx: Optional[int] = None
    ) -> Batch:
        """验证批次开始前的数据转换（可重写）

        默认实现：不做任何转换（恒等变换）
        子类可重写此方法实现测试时增强 (TTA)

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

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

    # ==================== 可重写的 Step 方法 ==================== #

    def batch_2_target(self, batch: Batch) -> torch.Tensor:
        """将输入批次转换为目标张量（可重写）

        Parameters
        ----------
        - `batch` : `Batch`
            - 数据批次

        Returns
        -------
        - `torch.Tensor`
            - 目标张量
        """
        return batch[1]

    def training_step(
        self,
        *,
        ctx: FabricTrainingContext,
        batch: Batch,
        batch_idx: int,
        epoch_idx: Optional[int] = None,
    ) -> tuple[torch.Tensor, Dict[str, float]]:
        return procedure.training_step(
            ctx, batch, batch_idx=batch_idx, epoch_idx=epoch_idx
        )

    def validation_step(
        self,
        *,
        ctx: FabricTrainingContext,
        batch: Batch,
        batch_idx: int,
        epoch_idx: Optional[int] = None,
    ) -> tuple[torch.Tensor, Dict[str, float]]:
        return procedure.validation_step(
            ctx, batch, batch_idx=batch_idx, epoch_idx=epoch_idx
        )

    def test_step(
        self, *, ctx: FabricTrainingContext, batch: Batch, batch_idx: int
    ) -> tuple[torch.Tensor, Dict[str, float]]:
        return procedure.validation_step(ctx, batch, batch_idx=batch_idx)

    # ==================== Epoch 级方法（可重写） ==================== #

    def train_epoch(self, train_loader: DataLoader, epoch_idx: int) -> Dict[str, float]:
        """训练一个 epoch（可重写）

        默认使用 procedure.train_epoch，并传入：
        - self._fn_train_step（如果通过 set_training_step 设置）
        - self.training_step（如果被子类重写）

        优先级：动态设置的函数 > 重写的方法

        Parameters
        ----------
        - `train_loader` : `DataLoader`
            - 训练数据加载器
        - `epoch_idx` : `int`
            - 当前 epoch 索引

        Returns
        -------
        - `Dict[str, float]`
            - 本 epoch 的平均指标，至少包含 'loss'
        """
        fn_step = self._fn_train_step or self.training_step
        fn_transform = (
            self._fn_transform_train_batch or self.transform_on_train_batch_start
        )

        return procedure.train_epoch(
            ctx=self.context,
            train_loader=train_loader,
            epoch_idx=epoch_idx,
            fn_train_step=fn_step,
            fn_batch2target=self.batch_2_target,
            fn_transform_batch=fn_transform,
        )

    def val_epoch(self, val_loader: DataLoader, epoch_idx: int) -> Dict[str, float]:
        """验证一个 epoch（可重写）

        Parameters
        ----------
        - `val_loader` : `DataLoader`
            - 验证数据加载器
        - `epoch_idx` : `int`
            - 当前 epoch 索引

        Returns
        -------
        - `Dict[str, float]`
            - 本 epoch 的平均指标，至少包含 'loss'
        """
        fn_step = self._fn_val_step or self.validation_step
        fn_transform = (
            self._fn_transform_val_batch or self.transform_on_validation_batch_start
        )

        return procedure.val_epoch(
            ctx=self.context,
            val_loader=val_loader,
            epoch_idx=epoch_idx,
            fn_val_step=fn_step,
            fn_batch2target=self.batch_2_target,
            fn_transform_batch=fn_transform,
        )

    def test_epoch(self, test_loader: DataLoader, epoch_idx: int = 0) -> Dict[str, float]:
        """测试一个 epoch（可重写）

        注意：test 通常只运行一次，所以 epoch_idx 默认为 0

        Parameters
        ----------
        - `test_loader` : `DataLoader`
            - 测试数据加载器
        - `epoch_idx` : `int`, optional
            - 当前 epoch 索引，默认为 0

        Returns
        -------
        - `Dict[str, float]`
            - 本 epoch 的平均指标，至少包含 'loss'
        """
        fn_step = self._fn_test_step or self.test_step

        return procedure.val_epoch(  # test 复用 val_epoch
            ctx=self.context,
            val_loader=test_loader,
            epoch_idx=epoch_idx,
            fn_val_step=fn_step,
            fn_batch2target=self.batch_2_target,
        )

    # ==================== 高层 API ==================== #

    def fit(
        self,
        train_loader: DataLoader,
        val_loader: Optional[DataLoader] = None,
        max_epochs: int = 10,
        val_interval: int = 1,
    ):
        """训练模型

        实现自己的训练循环，确保子类重写的方法（train_epoch, val_epoch 等）能够正确调用。

        Parameters
        ----------
        - `train_loader` : `DataLoader`
            - 训练数据加载器
        - `val_loader` : `DataLoader`, optional
            - 验证数据加载器
        - `max_epochs` : `int`, optional
            - 最大训练轮数，默认 10
        - `val_interval` : `int`, optional
            - 验证间隔，默认每个 epoch 验证一次
        """
        from zyplib.train.common import callbacks

        # Setup dataloaders
        fabric = self.fabric
        train_loader = self._setup_single_dataloader(train_loader)
        if val_loader:
            val_loader = self._setup_single_dataloader(val_loader)

        # 触发训练开始回调
        self.context.call_on_train_start(
            max_epochs=max_epochs, dataloader=train_loader, hparams=self.hparams
        )

        # 获取 early stopping callback
        early_stopping_cb: Optional[callbacks.EarlyStopping] = callbacks.find_callback(
            fabric, callbacks.EarlyStopping
        )

        def is_early_stopped():
            if early_stopping_cb:
                return early_stopping_cb.stop_training
            return False

        # 训练循环
        for epoch in range(max_epochs):
            if is_early_stopped():
                print_info('Early stopping, 终止训练')
                break

            # Epoch 开始回调
            self.context.call_on_epoch_start(epoch=epoch, max_epochs=max_epochs)
            self.context.call_on_train_epoch_start(epoch=epoch, max_epochs=max_epochs)

            # 训练一个 epoch（调用可重写的方法）
            train_metrics = self.train_epoch(train_loader, epoch)

            self.context.call_on_train_epoch_end(metrics=train_metrics, epoch=epoch)

            # 验证
            val_metrics = None
            if val_loader and (epoch + 1) % val_interval == 0:
                self.context.call_on_validation_epoch_start(
                    epoch=epoch, epoch_interval=val_interval
                )
                val_metrics = self.val_epoch(val_loader, epoch)
                self.context.call_on_validation_epoch_end(
                    metrics=val_metrics, epoch=epoch
                )

            # Epoch 结束回调
            self.context.call_on_epoch_end(
                epoch=epoch, train_metrics=train_metrics, val_metrics=val_metrics
            )

        # 训练结束回调
        self.context.call_on_train_end()

    def test(self, test_loader: DataLoader) -> Dict[str, float]:
        """测试模型"""
        fn_step = self._fn_test_step or self.test_step

        return procedure.test(
            ctx=self.context,
            test_loader=test_loader,
            fn_test_step=fn_step,
            fn_batch2target=self.batch_2_target,
        )

    # ==================== 工具方法 ==================== #

    def find_best_ckpt(self):
        """查找最佳 checkpoint 路径

        委托给 context.find_best_ckpt()

        Returns
        -------
        - `str | None`
            - 最佳 checkpoint 的路径，如果没有则返回 None
        """
        return self.context.find_best_ckpt()

    def load_best_ckpt(self, model: Literal['replace', 'return'] = 'replace'):
        """加载最佳 checkpoint

        委托给 context.load_best_ckpt()

        Parameters
        ----------
        - `model` : `Literal['replace', 'return']`, optional
            - 'replace': 直接替换当前模型
            - 'return': 返回加载后的新模型

        Returns
        -------
        - `nn.Module | None`
            - 如果 model='return'，返回加载后的模型；否则返回 self.model
        """
        return self.context.load_best_ckpt(model)
