"""Early stopping 核心逻辑

提供独立的 Early Stopping 逻辑，可以与 Callback 组合使用，
也可以在自定义训练循环中独立使用。
"""

import typing

from typing import Callable, Dict, Literal, Optional

import numpy as np

from zyplib.utils.print import print_debug, print_info, print_warn, warning


class EarlyStoppingMonitor:
    def __init__(
        self,
        monitor: str = 'loss',
        min_delta: float = 0.0,
        patience: int = 10,
        mode: Literal['min', 'max'] = 'min',
        verbose: bool = True,
        fn_update_callback: Callable[[int, int], None] | None = None,
        fn_stop_callback: Callable[[int], None] | None = None,
    ):
        """Early stopping 监控器

        仅仅用于计算何时需要早停，并通过 flag 或者回调函数通知外部停止训练。
        这是一个纯逻辑类，不依赖任何训练框架。

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

        Examples
        --------
        ```python
        # 基本用法
        monitor = EarlyStoppingMonitor(
            monitor='val_loss',
            patience=10,
            mode='min',
        )

        # 在训练循环中
        for epoch in range(100):
            # ... 训练代码 ...
            val_metrics = validate()

            # 检查是否需要早停
            should_stop = monitor.check(val_metrics, epoch)
            if should_stop:
                print(f'Early stopping at epoch {epoch}')
                break

        # 获取最佳 epoch
        print(f'Best epoch: {monitor.best_epoch}, Best score: {monitor.best_score}')
        ```
        """

        self.monitor = monitor
        self.min_delta = np.abs(min_delta)
        self.patience = patience
        self.mode = mode
        self.verbose = verbose
        self.fn_update_callback = fn_update_callback
        self.fn_stop_callback = fn_stop_callback

        self.counter = 0
        self.best_score: Optional[float] = None
        self.best_epoch: Optional[int] = None
        self.stopped_epoch = 0
        self.stop_training = False

        assert self.mode in ('min', 'max'), f"EarlyStopping mode '{self.mode}' is unknown"

    def _is_improved(self, current_score: float, best_score: float) -> bool:
        """判断当前分数是否优于最佳分数"""
        if self.mode == 'min':
            return current_score < best_score - self.min_delta
        else:
            return current_score > best_score + self.min_delta

    def reset(self) -> None:
        """重置 EarlyStopping 状态"""
        self.best_score = None
        self.best_epoch = None
        self.counter = 0
        self.stopped_epoch = 0
        self.stop_training = False
        if self.fn_update_callback:
            self.fn_update_callback(self.counter, self.patience)

    def check(self, metrics: typing.Mapping[str, float], epoch: int) -> bool:
        """检查是否需要早停

        Parameters
        ----------
        - `metrics` : `Dict[str, float]`
            - 当前 epoch 的指标
        - `epoch` : `int`
            - 当前 epoch

        Returns
        -------
        - `bool`
            - 如果需要停止训练，返回 True
        """
        if self.stop_training is True:
            warning('本轮训练已经满足早停条件，应当停止训练')
            return True

        current_score = metrics.get(self.monitor)
        if current_score is None:
            print_warn(
                f"无法获取到早停条件指标 `{self.monitor}`。"
                f"可用的指标有: {','.join(metrics.keys())}"
            )
            return False

        if self.best_score is None:
            self.best_score = current_score
            self.best_epoch = epoch
        elif self._is_improved(current_score, self.best_score):
            self.best_score = current_score
            self.best_epoch = epoch
            self.counter = 0
            self.fn_update_callback and self.fn_update_callback(
                self.counter, self.patience
            )
        else:
            self.counter += 1
            self.fn_update_callback and self.fn_update_callback(
                self.counter, self.patience
            )
            if self.verbose:
                print_debug(f'EarlyStopping counter: {self.counter}/{self.patience}')
            if self.counter >= self.patience:
                self.stopped_epoch = epoch
                self.stop_training = True

                if self.verbose:
                    print_info(f'Epoch {epoch}: early stopping')
                if self.fn_stop_callback:
                    self.fn_stop_callback(epoch)
                return True

        return False

    @property
    def should_stop(self) -> bool:
        """是否应该停止训练"""
        return self.stop_training
