import typing
import csv
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Protocol

import openpyxl
import simple_inject as sij

if TYPE_CHECKING:
    from lightning.fabric import Fabric
    from torch.nn import Module
    from torch.optim import Optimizer
    from torch.optim.lr_scheduler import _LRScheduler
    from torch.types import _TensorOrTensors
    from torch.utils.data import DataLoader

    from zyplib.train.infra.fabric_log import LogScalarManager

from zyplib.nn.utils import save_checkpoint

# from zyplib.train.utils import prefixed_dict
from zyplib.utils.fs import write_yaml
from zyplib.utils.kv import flatten_dict, prefixed_dict
from zyplib.utils.print import print_debug, print_info, print_warn, warning


def find_callback(fabric: 'Fabric', callback_type: type[Any]) -> Any | None:
    """查找 fabric 中的回调函数

    Parameters
    ----------
    - `fabric` : `Fabric`
        - 目标 fabric
    - `callback_type` : `FabricCallbackProtocol`
        - 目标回调函数类型

    Returns
    ----------
    - `FabricCallbackProtocol`
        - 目标回调函数
    """
    for callback in fabric._callbacks:
        if isinstance(callback, callback_type):
            return callback
    return None


class FabricCallbackProtocol(Protocol):
    def on_train_start(
        self,
        *,
        fabric: 'Fabric',
        model: 'Module',
        optimizer: Optional['Optimizer'] = None,
        lr_scheduler: Optional['_LRScheduler'] = None,
        dataloader: Optional['DataLoader'] = None,
        hparams: Optional[typing.Mapping[str, Any]] = None,
        **kwargs,
    ):
        """在开始时调用; 会传入 fabric 对象、模型、优化器 等; Callback 可以注册这个方法获取重要的上下文对象"""
        return None

    def on_train_end(self, **kwargs):
        return None

    def on_validation_start(self, **kwargs):
        return None

    def on_validation_end(self, **kwargs):
        return None

    def on_test_start(self, *, dataloader: Optional['DataLoader'] = None, **kwargs):
        return None

    def on_test_end(self, *, metrics: typing.Mapping[str, float] | None = None, **kwargs):
        return None

    def on_exception(self, exception: BaseException, **kwargs):
        return None

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

    def on_epoch_start(self, epoch: int, max_epochs: int, **kwargs):
        return None

    def on_epoch_end(
        self,
        epoch: int,
        *,
        train_metrics: typing.Mapping[str, float] | None = None,
        val_metrics: typing.Mapping[str, float] | None = None,
        **kwargs,
    ):
        return None

    def on_train_epoch_start(self, epoch: int, max_epochs: int, **kwargs):
        return None

    def on_train_epoch_end(
        self, epoch: int, *, metrics: typing.Mapping[str, float] | None = None, **kwargs
    ):
        return None

    def on_validation_epoch_start(
        self, epoch: int, *, epoch_interval: int | None = None, **kwargs
    ):
        return None

    def on_validation_epoch_end(
        self, epoch: int, *, metrics: typing.Mapping[str, float], **kwargs
    ):
        return None

    # --------------- Batch wise --------------- #
    def on_train_batch_start(self, batch: '_TensorOrTensors', batch_idx: int, **kwargs):
        return None

    def on_train_batch_end(
        self,
        batch: '_TensorOrTensors',
        batch_idx: int,
        *,
        metrics: typing.Mapping[str, float] | None = None,
        **kwargs,
    ):
        return None

    def on_validation_batch_start(
        self, batch: '_TensorOrTensors', batch_idx: int, **kwargs
    ):
        return None

    def on_validation_batch_end(
        self,
        batch: '_TensorOrTensors',
        batch_idx: int,
        *,
        metrics: typing.Mapping[str, float] | None = None,
        **kwargs,
    ):
        return None

    def on_test_batch_start(self, batch: '_TensorOrTensors', batch_idx: int, **kwargs):
        return None

    def on_test_batch_end(
        self,
        batch: '_TensorOrTensors',
        batch_idx: int,
        *,
        metrics: typing.Mapping[str, float] | None = None,
        **kwargs,
    ):
        return None


class EnsureFabricMixin:
    """自动注册 on_train_start 回调，获取 fabric 对象"""

    def __init__(self, fabric: Optional['Fabric'] = None):
        self.fabric = fabric

    def on_train_start(self, **kwargs):
        if self.fabric is not None:
            return
        fabric = kwargs.get('fabric', None)
        if fabric is None:
            classname = self.__class__.__name__
            raise ValueError(f'{classname} 必须有一个 fabric 属性!')
        self.fabric = fabric

    def _require_fabric(self) -> 'Fabric':
        if self.fabric is None:
            classname = self.__class__.__name__
            raise ValueError(f'{classname} 必须先绑定 fabric')
        return typing.cast('Fabric', self.fabric)


class SaveHparams(FabricCallbackProtocol):
    def __init__(self, dir: str, inject_namespace: str = 'default'):
        """自动保存超参数

        将超参数保存到 `dir`/hparams.yaml 下

        `on_train_start`
        ----------
        - 可以将 `hparams` 作为参数传入
        - 如果未传入，则从 `simple_inject` 中获取 namespace 的字典

        Parameters
        ----------
        - `dir` : `str`
            - 保存目录
        - `inject_namespace` : `str`
            - 从 `simple_inject` 中获取超参数的字典

        Hooks
        -----
        - `on_train_start`
            - `hparams` : `Dict[str, Any]`: 超参数; 如果未传入，则从 `simple_inject` 中获取
        """
        self.path = Path(dir) / 'hparams.yaml'
        self.inject_namespace = inject_namespace

    def on_train_start(self, **kwargs):
        hparams = kwargs.get('hparams', None)
        if hparams is None and self.inject_namespace:
            hparams = sij.state(self.inject_namespace)
        if hparams:
            write_yaml(self.path, hparams)
        else:
            print_warn(f'{self} 无法获取 hparams!')
        fabric = kwargs.get('fabric', None)
        if fabric:
            from lightning.pytorch.loggers import TensorBoardLogger

            for logger in fabric._loggers:
                if isinstance(logger, TensorBoardLogger):
                    if hparams is not None:
                        logger.log_hyperparams(hparams)
                    break


class PrintTrainInfo(FabricCallbackProtocol):
    def __init__(self, out_dir: str | None = None):
        """打印训练信息

        内部使用 `TrainInfoPrinter` 管理打印输出逻辑。

        Hooks
        -----
        - `on_train_start`: 打印训练开始信息, 记录训练开始时间
        - `on_train_end`: 打印训练结束信息, 记录训练结束时间
        - `on_train_epoch_start`: 打印当前 epoch 信息
            - `epoch` : `int`: 当前 epoch 数
            - `max_epochs` : `int`: 最大 epoch 数
        - `on_train_epoch_end`: 打印当前 epoch 训练指标
            - `metrics` : `Dict[str, float]`: 当前 epoch 的训练指标
        - `on_validation_epoch_end`: 打印当前 epoch 验证指标
            - `metrics` : `Dict[str, float]`: 当前 epoch 的验证指标
        - `on_test_end`: 打印测试指标
            - `metrics` : `Dict[str, float]`: 测试指标
        """
        from zyplib.train.infra.print_info import TrainInfoPrinter
        from zyplib.utils.time import TicToc, now

        self._now = now
        self._tictoc = TicToc()
        out_file = None if out_dir is None else f'{out_dir}/train-info.txt'
        self._printer = TrainInfoPrinter(out_file=out_file)

    def on_train_start(self, **kwargs):
        self._tictoc.tic()
        self._printer.print_header(f'开始训练: {self._now()}')
        self._printer.log('')

    def on_train_end(self, **kwargs):
        self._tictoc.toc()
        self._printer.log('')
        self._tictoc.print_elapsed('This training')
        self._printer.print_header(f'结束训练: {self._now()}', width=80)

    def on_train_epoch_start(self, epoch: int, max_epochs: int, **kwargs):
        self._printer.print_separator(f'Epoch {epoch}/{max_epochs}')

    def on_train_epoch_end(self, metrics: typing.Mapping[str, float], **kwargs):
        self._printer.print_metrics(metrics, prefix='Train')

    def on_validation_epoch_end(self, metrics: typing.Mapping[str, float], **kwargs):
        self._printer.print_metrics(metrics, prefix='Val')

    def on_test_end(self, metrics: typing.Mapping[str, float], **kwargs):
        self._printer.print_metrics(metrics, prefix='Test')


class SaveCheckpoints(FabricCallbackProtocol):
    def __init__(
        self,
        dir: str,
        save_every_k: Optional[int] = None,
        save_best: bool = True,
        keep_best_k: Optional[int] = 1,
        monitor: str = 'loss',
        monitor_on: Literal['train', 'val'] = 'val',
        less_is_better: bool = True,
        ckpt_ext: str = '.ckpt',
        if_not_empty: Literal['ask', 'clean', 'error', 'ignore'] = 'ask',
    ):
        """自动保存模型检查点的 Callback

        内部使用 `CheckpointSaveManager` 管理检查点保存逻辑。

        Parameters
        ----------
        - `dir` : `str`
            - 保存检查点的目录
        - `save_every_k` : `Optional[int]`, optional
            - 每多少个epoch保存一次，默认为None
        - `save_best` : `bool`, optional
            - 是否保存最佳模型，默认为True
        - `keep_best_k` : `Optional[int]`, optional
            - 保留表现最好的 K 个模型，默认为 1
            - 实现真正的 Top-K 策略：只要比当前保存的最差模型好，或还没存满 K 个，就会保存
        - `monitor` : `str`, optional
            - 用于确定最佳模型的指标，默认为'loss'
            - 可以有 `train_` 和 `val_` 前缀
        - `monitor_on` : `Literal['train', 'val']`, optional
            - 在哪个阶段进行监控，默认为'val', 表示监听 validation 的指标
        - `less_is_better` : `bool`, optional
            - 是否为越小越优，默认为True
        - `ckpt_ext` : `str`, optional
            - 检查点文件的扩展名，默认为'.ckpt'
        - `if_not_empty` : `Literal['ask', 'clean', 'error', 'ignore']`, optional
            - 当检查点目录不为空时的处理方式，默认为 'ask'
            - 'ask': 交互式询问用户（注意：非交互环境会阻塞）
            - 'clean': 直接清空目录
            - 'error': 抛出异常
            - 'ignore': 忽略，继续使用现有目录

        Hooks
        -----
        - `on_train_start`
            - `model` : `nn.Module`: 模型
            - `optimizer` : `Optimizer`: 优化器
            - `lr_scheduler` : `_LRScheduler`: 学习率调度器
        - `on_train_epoch_start`: 清空内置的 train 和 val 指标缓存
        - `on_train_epoch_end`
            - `metrics` : `Dict[str, float]`: 更新当前 epoch 的训练指标
        - `on_validation_epoch_end`
            - `metrics` : `Dict[str, float]`: 更新当前 epoch 的验证指标
        - `on_epoch_end`: 保存检查点
        """

        from zyplib.train.infra.save_ckpt import TopKCheckpointManager

        self._manager = TopKCheckpointManager(
            dir=dir,
            keep_best_k=keep_best_k or 1,
            monitor=monitor,
            less_is_better=less_is_better,
            ckpt_ext=ckpt_ext,
            if_not_empty=if_not_empty,
        )

        self.dir = self._manager.dir
        self.save_every_k = save_every_k
        self.save_best = save_best
        self.monitor = monitor
        self.monitor_on = monitor_on
        self.ckpt_ext = ckpt_ext
        self.current_epoch = 0

    def on_train_start(self, *, model, optimizer, lr_scheduler, **kwargs):
        self.model = model
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler

    def on_train_epoch_start(self, **kwargs):
        self.train_metrics = None
        self.val_metrics = None

    def on_train_epoch_end(self, metrics: typing.Mapping, **kwargs):
        self.train_metrics = metrics

    def on_validation_epoch_end(self, metrics: typing.Mapping, **kwargs):
        self.val_metrics = metrics

    def on_epoch_end(self, **kwargs):
        self.current_epoch += 1

        # 定期保存检查点
        if self.save_every_k and self.current_epoch % self.save_every_k == 0:
            checkpoint_path = (
                self.dir / f'checkpoint_epoch={self.current_epoch}{self.ckpt_ext}'
            )
            self._save_periodic(checkpoint_path)
            print_debug(f'Saved checkpoint at epoch {self.current_epoch}')

        if not self.save_best:
            return

        metrics = self.val_metrics if self.monitor_on == 'val' else self.train_metrics
        if not metrics:
            return

        current_score = (
            metrics.get(self.monitor)
            or metrics.get(f'val_{self.monitor}')
            or metrics.get(f'train_{self.monitor}')
        )

        if current_score is None:
            return

        # 委托给 manager 处理 Top-K 保存逻辑
        self._manager.update(
            score=current_score,
            epoch=self.current_epoch,
            save_fn=self._create_save_fn(),
            extra_info={
                'train_metrics': self.train_metrics,
                'val_metrics': self.val_metrics,
            },
        )

    def _create_save_fn(self) -> Callable[[str], None]:
        """创建保存函数，捕获当前状态"""

        def save_fn(path: str):
            save_checkpoint(
                models=self.model,
                path=path,
                optimizers=self.optimizer,
                epoch=self.current_epoch,
                train_metrics=self.train_metrics,
                val_metrics=self.val_metrics,
            )

        return save_fn

    def _save_periodic(self, checkpoint_path: Path):
        """保存定期检查点（非 best model）"""
        save_checkpoint(
            models=self.model,
            path=str(checkpoint_path),
            optimizers=self.optimizer,
            epoch=self.current_epoch,
            train_metrics=self.train_metrics,
            val_metrics=self.val_metrics,
        )

    def best_checkpoint(self) -> Path:
        """获取最佳检查点路径

        委托给内部的 CheckpointSaveManager

        Returns
        -------
        - `Path`
            - 最佳检查点的路径

        Raises
        ------
        - `FileNotFoundError`
            - 如果没有找到任何检查点文件
        """
        return self._manager.best_checkpoint()

    @property
    def best_score(self) -> float:
        """获取历史最佳分数"""
        return self._manager.best_score


class MonitorLearningRate(EnsureFabricMixin, FabricCallbackProtocol):
    def __init__(self, fabric: Optional['Fabric'] = None):
        """监控学习率

        监控学习率的变化，并记录到 fabric 的 loggers 中

        Parameters
        ----------
        - `fabric` : `Fabric`
            - 训练使用的 fabric 对象

        Hooks
        -----
        - `on_train_start`
            - `optimizer` : `Optimizer`: 优化器
            - `lr_scheduler` : `_LRScheduler`: 学习率调度器
        - `on_train_epoch_start`
            - `epoch` : `int`: 当前 epoch 数
        """
        EnsureFabricMixin.__init__(self, fabric)

    def on_train_start(self, *, optimizer, lr_scheduler, **kwargs):
        EnsureFabricMixin.on_train_start(self, **kwargs)
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler

    def on_train_epoch_start(self, **kwargs):
        lrs = [group['lr'] for group in self.optimizer.param_groups]
        epoch = kwargs.get('epoch', None)
        fabric = self._require_fabric()
        for i, lr in enumerate(lrs):
            fabric.log(f'lr_group_{i}', lr, epoch)
        if len(lrs) == 1:
            fabric.log('learning_rate', lrs[0], epoch)


class FabricLogMetrics(EnsureFabricMixin, FabricCallbackProtocol):
    def __init__(self, fabric: Optional['Fabric'] = None):
        """使用 Fabric 记录训练过程中的指标 (`fabric.log_dict`)

        Parameters
        ----------
        - `fabric` : `Fabric`, optional
            - 用于记录指标的 fabric 实例，默认为 None
            - 如果未传入，则会在 `on_train_start` 中从 kwargs 中获取 `fabric` 对象
            - *如果仍然为 None，则会抛出 ValueError

        Hooks
        -----
        - `on_train_start`
            - `fabric` : 可选的 `Fabric`
        - `on_train_batch_end`
            - `metrics` : `Dict[str, Any]`
        - `on_validation_batch_end`
            - `metrics` : `Dict[str, Any]`
        - `on_train_epoch_end`
            - `metrics` : `Dict[str, Any]`
            - `epoch` : `int`
        - `on_validation_epoch_end`
            - `metrics` : `Dict[str, Any]`
            - `epoch` : `int`
        - `on_test_end`
            - `metrics` : `Dict[str, Any]`
        """
        EnsureFabricMixin.__init__(self, fabric)
        self._manager: Optional['LogScalarManager'] = None

    def on_train_start(self, **kwargs):
        EnsureFabricMixin.on_train_start(self, **kwargs)
        from zyplib.train.infra.fabric_log import LogScalarManager

        self._manager = LogScalarManager(typing.cast(typing.Any, self._require_fabric()))

    def _require_manager(self) -> 'LogScalarManager':
        return typing.cast('LogScalarManager', self._manager)

    def on_train_batch_end(self, metrics: typing.Mapping[str, Any], **kwargs):
        self._require_manager().log_scalars(metrics, prefix='train/batch/')

    def on_validation_batch_end(self, metrics: typing.Mapping[str, Any], **kwargs):
        self._require_manager().log_scalars(metrics, prefix='val/batch/')

    def on_train_epoch_end(self, metrics: typing.Mapping[str, Any], epoch: int, **kwargs):
        self._require_manager().log_scalars(metrics, prefix='train/epoch/', step=epoch)

    def on_validation_epoch_end(
        self, metrics: typing.Mapping[str, Any], epoch: int, **kwargs
    ):
        self._require_manager().log_scalars(metrics, prefix='val/epoch/', step=epoch)

    def on_test_end(
        self, metrics: typing.Mapping[str, Any], epoch: Optional[int] = None, **kwargs
    ):
        self._require_manager().log_scalars(metrics, prefix='test/', step=epoch)


class SaveLoggersFinally(EnsureFabricMixin, FabricCallbackProtocol):
    def __init__(self, fabric: Optional['Fabric'] = None):
        """在训练/测试/训练+测试结束时保存日志

        Parameters
        ----------
        - `fabric` : `Optional[Fabric]`, optional
            - 训练使用的 fabric 对象
        """
        EnsureFabricMixin.__init__(self, fabric)

    def _save_loggers(self):
        from lightning.fabric.loggers.tensorboard import TensorBoardLogger

        fabric = self._require_fabric()

        for logger in fabric._loggers:
            try:
                logger.save()
            except Exception as e:
                print_warn(f'保存日志失败: {e}')
            if isinstance(logger, TensorBoardLogger):
                # 把所有数值指标保存为 json
                from zyplib.train.utils import tensorboard_eventfile_to_json

                try:
                    # 提取当前 session 的 TensorBoard event 文件
                    experiment = getattr(logger, '_experiment', None)
                    file_writer = getattr(experiment, 'file_writer', None)
                    event_writer = getattr(file_writer, 'event_writer', None)
                    event_file = getattr(event_writer, '_file_name', None)
                    if event_file is None:
                        continue
                    # log_dir = Path(logger.log_dir)
                    # event_files = list(log_dir.glob('events.out.tfevents.*'))
                    # if not event_files:
                    #     print_warn(f'未找到 TensorBoard 日志文件: {log_dir}')
                    #     continue

                    # # 按文件修改时间排序，获取最新的文件（当前 session 的）
                    # latest_event_file = max(event_files, key=lambda x: x.stat().st_mtime)
                    print_debug(f'正在转换 TensorBoard 日志: {event_file}')
                    metrics = tensorboard_eventfile_to_json(str(event_file))
                    metrics_name = list(metrics.keys())
                    print_debug(f'记录所有指标: {metrics_name}')

                    # 检查是否有其他日志文件
                    # if len(event_files) > 1:
                    #     print_info(
                    #         f'注意: 目录 {log_dir} 中存在多个日志文件，只转换了最新的文件'
                    #     )
                    #     for f in sorted(event_files, key=lambda x: x.stat().st_mtime):
                    #         if f != latest_event_file:
                    #             print_debug(f'  跳过旧日志文件: {f.name}')
                except Exception as e:
                    print_warn(f'转换 TensorBoard 日志到 JSON 失败: {e}')

    def on_train_end(self, **kwargs):
        self._save_loggers()

    def on_test_end(self, **kwargs):
        self._save_loggers()

    # def on_fit_end(self, **kwargs):
    #     self._save_loggers()


class ParseCSVLogfile:
    def __init__(self, csv_fpath: str, to_excel: bool = True, to_yaml: bool = True):
        self.csv_fpath = Path(csv_fpath)
        self.to_excel = to_excel
        self.to_yaml = to_yaml
        self.excel_fpath = self.csv_fpath.parent / f'{self.csv_fpath.name}.xlsx'
        self.yaml_fpath = self.csv_fpath.parent / f'{self.csv_fpath.name}.yaml'

        if not (to_excel or to_yaml):
            raise ValueError("At least one of 'to_excel' or 'to_yaml' must be True")

    def _process_csv(self):
        # Read CSV file
        with open(self.csv_fpath, 'r', encoding='utf-8') as csvfile:
            reader = csv.DictReader(csvfile)
            data = list(reader)

        # Process and organize data
        raw_data: dict[str, dict[int, list[float]]] = defaultdict(
            lambda: defaultdict(list)
        )
        metrics = set()

        for row in data:
            step = int(row.get('step', 0))
            for key, value in row.items():
                if key not in ['step'] and value:
                    key = key.replace('/', '_')
                    metrics.add(key)
                    raw_data[key][step].append(float(value))

        # 不知道为啥 logger 会重复记录，所以这里简单取平均
        repeated_metrics = set()
        organized_data: dict[str, dict[int, float]] = {}
        for metric, step_values in raw_data.items():
            organized_data[metric] = {}
            for step, values in step_values.items():
                if len(values) > 1:
                    repeated_metrics.add(metric)
                organized_data[metric][step] = sum(values) / len(values)
        if repeated_metrics:
            print_warn(f'发现 CSV Logger 存在重复记录的指标: {repeated_metrics}')

        return organized_data, metrics

    def _to_yaml(self, organized_data, metrics):
        names = sorted(list(metrics))
        final_data = {name: dict(organized_data[name]) for name in names}
        write_yaml(self.yaml_fpath, final_data)

    def _to_excel(self, organized_data, metrics):
        wb = openpyxl.Workbook()

        # Create sheets for each metric
        for metric in metrics:
            sheet = wb.create_sheet(title=metric)
            sheet['A1'] = 'Step'
            sheet['B1'] = 'Value'
            row = 2
            for step, value in sorted(organized_data[metric].items()):
                sheet[f'A{row}'] = step
                sheet[f'B{row}'] = value
                row += 1

        # Remove the default sheet created by openpyxl
        wb.remove(wb['Sheet'])

        # Save the Excel file
        wb.save(self.excel_fpath)

    def _convert(self):
        if not self.to_excel and not self.to_yaml:
            return
        if not self.csv_fpath.exists():
            print_warn(f'CSV 日志文件不存在, 无法转换: {self.csv_fpath}')
            return
        organized_data, metrics = self._process_csv()
        if self.to_yaml:
            self._to_yaml(organized_data, metrics)
        if self.to_excel:
            self._to_excel(organized_data, metrics)

    def on_train_end(self, **kwargs):
        self._convert()

    def on_validation_end(self, **kwargs):
        self._convert()

    def on_test_end(self, **kwargs):
        self._convert()


class EarlyStopping:
    def __init__(
        self,
        monitor: str = 'loss',
        monitor_on: Literal['train', 'val'] = 'val',
        min_delta: float = 0.0,
        patience: int = 10,
        mode: Literal['min', 'max'] = 'min',
        verbose: bool = True,
        stop_flag: str = 'early_stop',
        fn_update_callback: Callable[[int, int], None] | None = None,
        fn_stop_callback: Callable[[int], None] | None = None,
    ):
        """Early stopping

        内部使用 `EarlyStoppingMonitor` 管理早停逻辑。

        机制
        ----------
        在需要早停的时候，会:
        1. 设置 stop_training 为 True
        2. 调用 fn_stop_callback(epoch) 回调函数


        Parameters
        ----------
        - `monitor` : `str`, 可选
            - 监控的指标，默认为 'loss'
            - 在 validation set 中监控
        - `monitor_on` : `Literal['train', 'val']`, 可选
            - 监控的阶段，默认为 'val'
        - `min_delta` : `float`, 可选
            - 监控数量变化的最小值，以此作为改进的标准，默认为 0.0
        - `patience` : `int`, 可选
            - 在没有改进的情况下，训练将停止的 epoch 数，默认为 10
        - `mode` : `Literal['min', 'max']`, 可选
            - 在 'min' 模式下，当监控数量不再减少时停止训练；
            - 在 'max' 模式下，当监控数量不再增加时停止训练，默认为 'min'
        - `verbose` : `bool`, 可选
            - 如果为 True，则每次更新时打印消息，默认为 True
        - `stop_flag` : `str`, 可选
            - 设置在 trainer 上的标志以中断训练
        - `fn_update_callback`: `(counter: int, patience: int) -> None`, 可选
            - 当更新早停状态的时候调用
        - `fn_stop_callback` : `Callable[[int], None]`, 可选
            - 当训练停止时调用的回调函数，默认为 None

        Hooks
        -----
        - `on_train_start`
        - `on_validation_epoch_end`
            - `metrics` : `Dict[str, float]`
            - `epoch` : `int`
        """
        from zyplib.train.infra.early_stop import EarlyStoppingMonitor

        self._monitor = EarlyStoppingMonitor(
            monitor=monitor,
            min_delta=min_delta,
            patience=patience,
            mode=mode,
            verbose=verbose,
            fn_update_callback=fn_update_callback,
            fn_stop_callback=fn_stop_callback,
        )
        self.monitor_on = monitor_on
        self.stop_flag = stop_flag

    # ========== 代理属性，保持向后兼容 ==========

    @property
    def monitor(self) -> str:
        return self._monitor.monitor

    @property
    def patience(self) -> int:
        return self._monitor.patience

    @property
    def mode(self) -> str:
        return self._monitor.mode

    @property
    def counter(self) -> int:
        return self._monitor.counter

    @property
    def best_score(self) -> Optional[float]:
        return self._monitor.best_score

    @property
    def best_epoch(self) -> Optional[int]:
        return self._monitor.best_epoch

    @property
    def stopped_epoch(self) -> int:
        return self._monitor.stopped_epoch

    @property
    def stop_training(self) -> bool:
        return self._monitor.stop_training

    def _reset(self) -> None:
        """重置 EarlyStopping 状态"""
        self._monitor.reset()

    def on_validation_epoch_end(
        self, metrics: typing.Mapping[str, float], epoch: int, **kwargs
    ) -> None:
        if self.monitor_on == 'val':
            self._monitor.check(metrics, epoch)

    def on_train_epoch_end(
        self, metrics: typing.Mapping[str, float], epoch: int, **kwargs
    ) -> None:
        if self.monitor_on == 'train':
            self._monitor.check(metrics, epoch)


class LRSchedulerCallback(FabricCallbackProtocol):
    """学习率调度器回调函数

    提供灵活的学习率调度功能，支持 epoch 和 step 两种调度模式。
    整合了 step_lr_sched 函数的功能，实现 LR 相关操作的高度内聚。

    Parameters
    ----------
    - `lr_scheduler` : `_LRScheduler`
        - 学习率调度器
    - `interval` : `Literal['epoch', 'step']`, optional
        - 调度间隔模式, by default 'epoch'
        - 'epoch': 在每个 epoch 结束时调度
        - 'step': 在每个训练 step 结束时调度
    - `frequency` : `int`, optional
        - 调度频率，每隔多少个 interval 执行一次调度, by default 1
    - `monitor_metric` : `str`, optional
        - 监控的指标名称（用于 ReduceLROnPlateau）, by default 'loss'
    - `use_val_metric` : `bool`, optional
        - 是否优先使用验证指标（仅在 interval='epoch' 时有效）, by default True
    """

    def __init__(
        self,
        lr_scheduler: '_LRScheduler',
        interval: Literal['epoch', 'step', 'auto'] = 'auto',
        frequency: int = 1,
        monitor_metric: str = 'loss',
        use_val_metric: bool = True,
    ):
        self.lr_scheduler = lr_scheduler
        self.frequency = frequency
        self.monitor_metric = monitor_metric
        self.use_val_metric = use_val_metric

        # 内部状态
        self._step_count = 0
        self._epoch_count = 0

        # 缓存最新的指标用于调度
        self._last_train_metrics = None
        self._last_val_metrics = None

        if interval == 'auto':
            self.__auto_decide()
        else:
            self.interval = interval

    def __auto_decide(self):
        from torch.optim.lr_scheduler import (
            CosineAnnealingWarmRestarts,
            CyclicLR,
            OneCycleLR,
        )

        STEP_BASED_SCHEDULERS = (OneCycleLR, CyclicLR, CosineAnnealingWarmRestarts)
        if isinstance(self.lr_scheduler, STEP_BASED_SCHEDULERS):
            self.interval = 'step'
        else:
            self.interval = 'epoch'

    def _should_step(self) -> bool:
        """检查是否应该执行调度"""
        if self.interval == 'step':
            # 确保不在第 0 步调度，且满足频率要求
            return self._step_count > 0 and self._step_count % self.frequency == 0
        elif self.interval == 'epoch':
            # 确保不在第 0 轮调度，且满足频率要求
            return self._epoch_count > 0 and self._epoch_count % self.frequency == 0
        return False

    def _get_metric_for_scheduling(self) -> Optional[float]:
        """获取用于调度的指标值"""
        # 优先使用验证指标（如果可用且 use_val_metric=True）
        if self.use_val_metric and self._last_val_metrics:
            metric_value = self._last_val_metrics.get(self.monitor_metric)
            if metric_value is not None:
                return metric_value

        # 回退到训练指标
        if self._last_train_metrics:
            return self._last_train_metrics.get(self.monitor_metric)

        return None

    def _step_lr_scheduler(self, metric_value: Optional[float] = None):
        """执行学习率调度器的 step 操作

        整合了 step_lr_sched 函数的功能
        """
        from torch.optim.lr_scheduler import ReduceLROnPlateau

        if isinstance(self.lr_scheduler, ReduceLROnPlateau):
            if metric_value is not None:
                self.lr_scheduler.step(metric_value)
            else:
                print_warn(
                    f'ReduceLROnPlateau 需要指标值，但未找到 {self.monitor_metric}'
                )
        else:
            # 对于其他类型的调度器，直接调用 step()
            self.lr_scheduler.step()

    def on_train_batch_end(
        self,
        batch: '_TensorOrTensors',
        batch_idx: int,
        *,
        metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """训练批次结束时的回调"""
        if metrics:
            self._last_train_metrics = metrics

        self._step_count += 1

        if self.interval == 'step' and self._should_step():
            metric_value = self._get_metric_for_scheduling()
            self._step_lr_scheduler(metric_value)

    def on_epoch_end(
        self,
        epoch: int,
        *,
        train_metrics: Optional[typing.Mapping[str, float]] = None,
        val_metrics: Optional[typing.Mapping[str, float]] = None,
        **kwargs,
    ):
        """Epoch 结束时的回调"""
        if train_metrics:
            self._last_train_metrics = train_metrics
        if val_metrics:
            self._last_val_metrics = val_metrics

        self._epoch_count += 1

        if self.interval == 'epoch' and self._should_step():
            metric_value = self._get_metric_for_scheduling()
            self._step_lr_scheduler(metric_value)
