from __future__ import annotations

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

import numpy as np
import torch
from lightning.fabric import Fabric
from lightning.fabric.wrappers import _FabricModule, _FabricOptimizer
from simple_inject import Inject, auto_inject, inject
from torch import nn
from torch.optim import Optimizer
from torch.utils.data import DataLoader

from zyplib.nn.utils import load_checkpoint
from zyplib.train.common import callbacks
from zyplib.train.metric.torch_metric import TorchMetricRecorder

# from zyplib.train.utils import step_lr_sched
from zyplib.utils.print import (
    print_colored,
    print_debug,
    print_info,
    print_warning,
    warning,
)
from zyplib.utils.progress_bar import rich_progress

if TYPE_CHECKING:
    from torch.optim.lr_scheduler import _LRScheduler
    from torch.types import _TensorOrTensors

    from zyplib.train.common.callbacks import EarlyStopping

# 类型定义
Batch: TypeAlias = tuple[torch.Tensor, ...]
Metrics: TypeAlias = Dict[str, float]
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: Optimizer | tuple[Optimizer, ...] | list[Optimizer] | dict[str, Optimizer]
    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 | tuple[Optimizer, ...] | list[Optimizer] | dict[str, Optimizer]:
        """返回 unwrap 后的原始优化器"""

        def unwrap(optimizer):
            if self.fabric and isinstance(optimizer, _FabricOptimizer):
                return optimizer._optimizer
            return optimizer

        if isinstance(self.optimizer, (list, tuple)):
            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 __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 = 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:
                self.optimizer = self.fabric.setup_optimizers(*self.optimizer)
        elif isinstance(self.optimizer, dict):
            for key, optimizer in self.optimizer.items():
                if not isinstance(optimizer, _FabricOptimizer):
                    self.optimizer[key] = self.fabric.setup_optimizers(optimizer)

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

    def find_best_ckpt(self):
        save_ckpt_cb: callbacks.SaveCheckpoints | None = callbacks.find_callback(
            self.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
        if model == 'replace':
            load_checkpoint(self.model, path, device=self.fabric.device)
            return self.model
        elif model == 'return':
            new_model = copy.deepcopy(self.pure_model)
            load_checkpoint(new_model, path, device=self.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[Dict[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[Dict[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[Dict[str, float]] = None,
        val_metrics: Optional[Dict[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[Dict[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: Dict[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[Dict[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[Dict[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[Dict[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


def _resolve_loss_fn(
    ctx: FabricTrainingContext, fn_loss: Optional[nn.Module]
) -> nn.Module:
    if ctx.fn_loss is not None:
        return ctx.fn_loss

    # 向后兼容: 如果传入了 fn_loss, 则自动设置到 context.loss_fn 中
    if fn_loss is not None:
        if ctx.fn_loss != fn_loss:
            ctx.fn_loss = fn_loss
            print_debug(f'自动设置 context.loss_fn: {type(ctx.fn_loss).__name__}')
        return ctx.fn_loss

    # 两者都没有，抛出错误
    raise ValueError(
        '损失函数未指定：请在 FabricTrainingContext.loss_fn 中设置损失函数，'
        '或在函数调用时传入 fn_loss 参数'
    )


def _default_batch_2_target(batch: tuple):
    return batch[1]


DepecatedFnLoss: TypeAlias = Optional[nn.Module]


def training_step(
    ctx: FabricTrainingContext,
    batch: Batch,
    *,
    fn_loss: DepecatedFnLoss = None,
    batch_idx: Optional[int] = None,
    **kwargs,
) -> tuple[torch.Tensor, Dict[str, float]]:
    """默认内置的单个训练步骤, 主要适配普通的单输入单输出训练任务例如分类

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `batch` : `tuple[torch.Tensor, torch.Tensor]`
        - 数据批次
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数: `fn_loss(y_pred, y)`，如果为 None 则使用 ctx.loss_fn
    - `batch_idx`: int

    Returns
    ----------
    - `tuple[torch.Tensor, Dict[str, float]]`: (y_pred, {"loss": float})
    """
    ctx.model.train()
    x, y = batch

    # 解析损失函数
    loss_fn = _resolve_loss_fn(ctx, fn_loss)

    ctx.optimizer.zero_grad()
    y_pred = ctx.model(x)
    loss: torch.Tensor = loss_fn(y_pred, y)

    if ctx.fabric:
        ctx.fabric.backward(loss)
    else:
        loss.backward()

    ctx.optimizer.step()

    loss_value = loss.item()

    # if ctx.train_metric_recorder:
    #     return ctx.train_metric_recorder.update(y_pred, y, loss=loss_value)
    return y_pred, {'loss': loss_value}


def validation_step(
    ctx: FabricTrainingContext,
    batch: Batch,
    *,
    fn_loss: DepecatedFnLoss = None,
    batch_idx: Optional[int] = None,
    **kwargs,
) -> tuple[torch.Tensor, Dict[str, float]]:
    """默认内置的单个验证步骤, 主要适配普通的单输入单输出验证任务例如分类

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `batch` : `tuple[torch.Tensor, torch.Tensor]`
        - 数据批次
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数: `fn_loss(y_pred, y)`，如果为 None 则使用 ctx.loss_fn
    - `batch_idx`: int

    Returns
    ----------
    - `tuple[torch.Tensor, Dict[str, float]]`: (y_pred, {"loss": float})
    """
    ctx.model.eval()
    with torch.no_grad():
        x, y = batch
        y_pred = ctx.model(x)

        # 解析损失函数
        loss_fn = _resolve_loss_fn(ctx, fn_loss)
        loss: torch.Tensor = loss_fn(y_pred, y)

    loss_value = loss.item()
    # if ctx.val_metric_recorder:
    #     return ctx.val_metric_recorder.update(y_pred, y, loss=loss_value)
    # return {'loss': loss_value}
    return y_pred, {'loss': loss_value}


def train_epoch(
    ctx: FabricTrainingContext,
    train_loader: DataLoader,
    epoch_idx: Optional[int] = None,
    fn_train_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_transform_batch: Optional[Callable] = None,
) -> Dict[str, float]:
    """单个训练轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `train_loader` : `DataLoader`
        - 训练数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_train_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个训练步骤的函数, 默认为 `training_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_transform_batch` : `(...) -> Batch`, optional
        - 可选项, 用于在训练前转换批次数据的函数, 默认使用 `ctx.transform_on_train_batch_start`
        - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
        - 用例: Mixup, CutMix, 动态数据增强等

    Returns
    ----------
    - `Dict[str, float]`: 本 Epoch 的平均指标, 至少包含 `loss`
    """
    if fn_train_step is None:
        fn_train_step = training_step
    if fn_transform_batch is None:
        fn_transform_batch = ctx.transform_on_train_batch_start

    for idx, batch in enumerate(
        rich_progress(train_loader, f'Training epoch {epoch_idx}', color='green')
    ):
        # 应用数据转换 hook
        batch = fn_transform_batch(batch, idx, epoch_idx)

        y_pred, raw_metrics = fn_train_step(
            ctx=ctx, batch=batch, batch_idx=idx, epoch_idx=epoch_idx
        )
        # _, y = batch
        y = fn_batch2target(batch)
        metrics = ctx.train_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_train_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    epoch_metrics = ctx.train_metric_recorder.compute()
    return epoch_metrics


def val_epoch(
    ctx: FabricTrainingContext,
    val_loader: DataLoader,
    epoch_idx: Optional[int] = None,
    fn_val_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_transform_batch: Optional[Callable] = None,
) -> Dict[str, float]:
    """单个验证轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `val_loader` : `DataLoader`
        - 验证数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_val_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个验证步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_transform_batch` : `(...) -> Batch`, optional
        - 可选项, 用于在验证前转换批次数据的函数, 默认使用 `ctx.transform_on_validation_batch_start`
        - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
        - 用例: 测试时增强 (TTA) 等

    Returns
    ----------
    - `Dict[str, float]`: 本 Epoch 的平均指标, 至少包含 `loss`
    """
    if fn_val_step is None:
        fn_val_step = validation_step
    if fn_transform_batch is None:
        fn_transform_batch = ctx.transform_on_validation_batch_start

    for idx, batch in enumerate(
        rich_progress(val_loader, f'Validating epoch {epoch_idx}', color='blue')
    ):
        # 应用数据转换 hook
        batch = fn_transform_batch(batch, idx, epoch_idx)

        y_pred, raw_metrics = fn_val_step(
            ctx=ctx, batch=batch, batch_idx=idx, epoch_idx=epoch_idx
        )
        y = fn_batch2target(batch)
        metrics = ctx.val_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_validation_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    metrics = ctx.val_metric_recorder.compute()
    return metrics


def test(
    ctx: FabricTrainingContext,
    test_loader: DataLoader,
    fn_test_step: Optional[Callable] | None = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    **test_step_kwargs: Any,
) -> Dict[str, float]:
    """单个测试轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `test_loader` : `DataLoader`
        - 测试数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_test_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个测试步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数

    `fn_test_step`:
    ----------
    - 依赖注入: 从 `namespace=train` 自动获取 `validation_step` 或者 `test_step`
    - 默认不提供, 则使用 `procedure.validation_step`; 如果提供了, 则使用提供的函数
    - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx)`
        - `fn_loss` 参数现在是可选的，如果为 None 则使用 ctx.fn_loss
        - 注意，不会有 `epoch_idx` 的参数!
    - 返回值: `torch.Tensor, Dict[str, Any]` (predictions, metrics), 至少包含 `loss`

    Returns
    ----------
    - `Dict[str, float]`: 本次测试的平均指标, 至少包含 `loss`
    """
    # fn_test_step = (
    #     fn_test_step
    #     or inject('test_step', 'train')
    #     or inject('validation_step', 'train')
    #     or validation_step
    # )
    if fn_test_step is None:
        from simple_inject import has as has_dep

        if has_dep('test_step', 'train'):
            fn_test_step = inject('test_step', 'train')
        elif has_dep('validation_step', 'train'):
            fn_test_step = inject('validation_step', 'train')
        else:
            fn_test_step = validation_step

    ctx.call_on_test_start(dataloader=test_loader)
    test_loader = ctx.fabric.setup_dataloaders(test_loader)

    ctx.val_metric_recorder.reset()

    for idx, batch in enumerate(rich_progress(test_loader, 'Testing', color='blue')):
        y_pred, raw_metrics = fn_test_step(
            ctx=ctx, batch=batch, batch_idx=idx, **test_step_kwargs
        )
        y = fn_batch2target(batch)
        metrics = ctx.val_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_test_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    metrics = ctx.val_metric_recorder.compute()
    ctx.call_on_test_end(metrics=metrics)
    return metrics


# @auto_inject()
def train_loop(
    context: FabricTrainingContext,
    train_loader: DataLoader,
    max_epochs: int | None = 100,
    val_loader: Optional[DataLoader] = None,
    val_interval: int = 1,
    hparams: Optional[Dict[str, Any]] = None,
    fn_train_step: Optional[Callable] = None,
    fn_val_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_on_trainepoch_end: Optional[Callable] = None,
    # Deprecated parameters for backward compatibility
    # reduce_lr_metric_name: str | None = 'loss',
    # fn_loss: nn.Module | None = None,
    # lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
    # lr_sched_freq: int = 1,
    **kwargs
):
    """基于 Fabric 的 Train Loop; fabric 为必选参数，如果不想用 fabric 可以自行组合 train_epoch etc.

    Fabric 需要在创建之前 launch!

    依赖注入
    ----------
    可以将 `training_step` 和 `validation_step` 函数注入到 `"train"` namespace 下。
    如果不注入，使用默认的训练和验证函数。函数需满足如下签名:
    ```py
    # 返回中应当至少包含 {'loss': float}
    fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx, epoch_idx) -> tuple[torch.Tensor, dict[str, float]]
    ```

    Parameters
    ----------
    - `context` : `procedure.FabricTrainingContext`; 必须包含 fabric!
        - 如果 context.fn_loss 设置了损失函数，则可以作为默认损失函数使用
    - `train_loader` : `DataLoader`
        - 训练数据加载器
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数，如果为 None 且 context.fn_loss 也为 None，则会抛出错误
    - `max_epochs` : `int`
        - 最大训练轮数
    - `val_loader` : `Optional[DataLoader]`
        - 验证数据加载器，可选
    - `val_interval` : `int`, optional
        - 验证间隔, by default 1
    - `lr_scheduler` : `Optional[torch.optim.lr_scheduler._LRScheduler]`, optional
        - 学习率调度器, by default None
        - 为了向后兼容，如果传入但没有注册 LRSchedulerCallback，会自动创建 epoch-level 的调度器
        - 推荐直接使用 LRSchedulerCallback 获得更灵活的控制
    - `lr_sched_freq` : `int`, optional
        - 学习率调度频率（向后兼容参数）, by default 1
        - 仅在自动创建 LRSchedulerCallback 时使用
    - `fn_train_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个训练步骤的函数, 默认为 `training_step`
    - `fn_val_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个验证步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_on_trainepoch_end` : `(...) -> None`, optional
        - 可选项, 用于在每个 epoch 结束时调用的函数, 默认为 `None`
        - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, epoch, train_metrics, val_metrics)`
        - 相比较注册 Callback 更加灵活一些
    - `reduce_lr_metric_name` : `str`, optional
        - 用于监控的指标名称（向后兼容参数）, 默认为 'loss'
        - 仅在自动创建 LRSchedulerCallback 时使用，用于 ReduceLROnPlateau 等调度器

    `fn_train_step` / `fn_val_step`:
    ----------
    - 默认不提供, 使用 `procedure.training_step` / `procedure.validation_step`; 如果提供了, 则使用提供的函数 (例如训练特殊模型, 需要非常规的步骤等)
    - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx, epoch_idx)`
        - `fn_loss` 参数现在是可选的，如果为 None 则使用 ctx.loss_fn
    - 返回值: `tuple[torch.Tensor, Dict[str, float]]` (predictions, metrics)
        - 例如默认的 `training_step` 返回 `(y_pred, {'loss': loss_value})`
        - 对于自定义的 step 函数, 如果返回来想要记录的指标, 请在 TorchMetricRecorder 的 scalar_metrics 中指定
    """
    # model, optimizer = fabric.setup(model, optimizer)
    train_loader = context.fabric.setup_dataloaders(train_loader)
    if val_loader:
        val_loader = context.fabric.setup_dataloaders(val_loader)

    # 向后兼容：如果传入了 lr_scheduler 但没有对应的 LRSchedulerCallback，自动注入;
    # 后续可能会移除此自动注入行为，请尽早显式注册 LRSchedulerCallback 以避免未来的不兼容。
    if 'lr_scheduler' in kwargs:
        lr_scheduler = kwargs.get('lr_scheduler')
        lr_scheduler_cb = callbacks.find_callback(
            context.fabric, callbacks.LRSchedulerCallback
        )
        if lr_scheduler_cb is None:
            # 自动创建并注入 LRSchedulerCallback
            auto_lr_callback = callbacks.LRSchedulerCallback(
                lr_scheduler=lr_scheduler,
                interval='epoch',  # 保持原有的 epoch-level 行为
                frequency=kwargs.get('lr_sched_freq') or 1,
                monitor_metric=kwargs.get('reduce_lr_metric_name') or 'loss',
                use_val_metric=True,
            )
            context.fabric._callbacks.append(auto_lr_callback)
            print_debug('自动注入 LRSchedulerCallback 以保持向后兼容')
    else:
        lr_scheduler = None

    # 向后兼容: 如果传入了 fn_loss, 则自动设置到 context.loss_fn 中; 后续可能会删掉
    if 'fn_loss' in kwargs:
        fn_loss = kwargs.get('fn_loss')
        if context.fn_loss is None:
            context.fn_loss = fn_loss
            print_debug(f'自动设置 context.loss_fn: {context.fn_loss}')
        else:
            # 如果已经设置了 fn_loss, 则不再覆盖
            if context.fn_loss != fn_loss:
                print_warning(
                    'context.loss_fn 已经设置, 不会覆盖: '
                    f'当前为 {type(context.fn_loss).__name__}, '
                    f'传入的为 {type(fn_loss).__name__}'
                )
    else:
        # 如果没有传入 fn_loss, 则使用 context.loss_fn
        if context.fn_loss is None:
            if fn_train_step is None:
                raise ValueError(
                    '损失函数未指定：请在 FabricTrainingContext.loss_fn 中设置损失函数，'
                    '或在函数调用时传入 fn_loss 参数'
                )
            else:
                print_warning(
                    '损失函数未指定：如果并非有意偶数为，请在 FabricTrainingContext.loss_fn 中设置损失函数'
                )  # 也有可能是自己指定了 train_step

    context.call_on_train_start(
        lr_scheduler=lr_scheduler,
        max_epochs=max_epochs,
        dataloader=train_loader,
        hparams=hparams,
    )

    early_stopping_cb: 'EarlyStopping' | None = callbacks.find_callback(
        context.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

        context.call_on_epoch_start(epoch=epoch, max_epochs=max_epochs)
        context.call_on_train_epoch_start(epoch=epoch, max_epochs=max_epochs)

        train_metrics = train_epoch(
            ctx=context,
            train_loader=train_loader,
            epoch_idx=epoch,
            fn_train_step=fn_train_step,
            fn_batch2target=fn_batch2target,
        )

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

        if val_loader and (epoch + 1) % val_interval == 0:
            context.call_on_validation_epoch_start(
                epoch=epoch, epoch_interval=val_interval
            )
            val_metrics = val_epoch(
                ctx=context,
                val_loader=val_loader,
                # fn_loss=fn_loss,
                epoch_idx=epoch,
                fn_val_step=fn_val_step,
                fn_batch2target=fn_batch2target,
            )
            context.call_on_validation_epoch_end(metrics=val_metrics, epoch=epoch)
        else:
            val_metrics = None

        # deprecated: 直接使用 LRSchedulerCallback
        # if lr_scheduler:
        #     if (epoch + 1) % lr_sched_freq == 0:
        #         # loss = val_metrics['loss'] if val_metrics else train_metrics['loss']
        #         # step_lr_sched(lr_scheduler, loss)
        #         metric_name = reduce_lr_metric_name or 'loss'
        #         metric = (
        #             val_metrics.get(metric_name)
        #             if val_metrics
        #             else train_metrics.get(metric_name)
        #         )
        #         step_lr_sched(lr_scheduler, metric)

        context.call_on_epoch_end(
            epoch=epoch, train_metrics=train_metrics, val_metrics=val_metrics
        )

        if fn_on_trainepoch_end:
            fn_on_trainepoch_end(
                ctx=context,
                epoch=epoch,
                train_metrics=train_metrics,
                val_metrics=val_metrics,
            )

    context.call_on_train_end()
