import typing
import fnmatch
from copy import deepcopy
from typing import Callable, Dict, List, Literal, Optional

# import matplotlib.pyplot as plt
import torch
from torchmetrics import Metric, MetricCollection

from zyplib.train.utils import use_device
from zyplib.utils.ensure import ensure_scalar
from zyplib.utils.print import print_debug, print_warning, warning
from zyplib.utils.kv import prefixed_dict

ReducerName = Literal['mean', 'sum']
ActivationFn = Callable[[torch.Tensor], torch.Tensor]


class ScalarMetric:
    def __init__(
        self,
        default_reducer: ReducerName = 'mean',
        retain_batch_history: bool = False,
        retain_epoch_history: bool = True,
        auto_define: bool = False,
    ):
        """记录单个 float 类型的指标

        唯一的作用是简化「记录 Batch 指标 - 在 epoch 结束后求平均」的工作流程

        Parameters
        ----------
        - `default_reducer` : `Literal['mean', 'sum']`, optional
            - 默认的聚合方式, by default 'mean'
        - `retain_batch_history` : `bool`, optional
            - 是否保留每个 batch 的历史记录, by default False
        - `retain_epoch_history` : `bool`, optional
            - 是否保留每个 epoch 的历史记录, by default True
        - `auto_define` : `bool`, optional
            - 是否自动定义未知的指标，当为 True 时，update() 中传入的未定义指标会自动使用 default_reducer 进行定义, by default False

        Example
        --------

        ```py
        # 手动模式：需要预先定义指标
        metric = ScalarMetric()
        metric.define_metric('loss', reducer='mean')
        metric.define_metric('acc', reducer='mean')

        # 自动定义模式：无需预先定义
        metric = ScalarMetric(auto_define=True)

        for batch in dataloader:
            y_hat, y = batch
            loss = compute_loss(y_hat, y)
            acc = compute_acc(y_hat, y)
            metric.update(loss=loss.item(), acc=acc)  # 自动定义 loss 和 acc
        metric.compute()  # 得到当前 epoch 的平均损失和准确率
        ```
        """
        self._batch_metrics: dict[
            str, list[float]
        ] = {}  # 每个 batch 内的指标, 随着 reset 清空

        self._batch_history: list[dict[str, float]] = []
        self._epoch_history: list[dict[str, float]] = []

        self._metrics: dict[str, Callable[[list[float]], float]] = {}
        self._default_reducer = (
            lambda x: sum(x) / len(x) if default_reducer == 'mean' else sum(x)
        )
        self._default_reducer_name: ReducerName = default_reducer

        self._retain_batch_history = retain_batch_history
        self._retain_epoch_history = retain_epoch_history
        self._auto_define = auto_define

        self._undefined_metric_name = set()

    def clone(self):
        """克隆一个完全相同的对象"""
        return deepcopy(self)

    def define_metric(self, metric_name: str, reducer: ReducerName = 'mean'):
        """定义一个新的指标及其计算方式

        Parameters
        ----------
        - `metric_name` : `str`
            - 指标名称
        - `reducer` : `Literal['mean', 'sum']`, optional
            - 计算方式, 默认为 'mean'
        """
        self._batch_metrics[metric_name] = []
        if reducer == 'mean':
            self._metrics[metric_name] = lambda x: sum(x) / len(x)
        elif reducer == 'sum':
            self._metrics[metric_name] = sum
        else:
            raise ValueError("reducer must be either 'mean' or 'sum'")

    def update(self, **metrics: float):
        """更新当前 step 的指标值

        Parameters
        ----------
        **metrics : float
            指标名称和对应的值
        """
        for name, value in metrics.items():
            if name not in self._metrics:
                if self._auto_define:
                    # 自动定义模式：使用默认 reducer 自动定义指标
                    self.define_metric(name, self._default_reducer_name)
                else:
                    # 传统模式：跳过未定义的指标并警告
                    if name not in self._undefined_metric_name:
                        print_warning(f'Metric {name} not defined, skipping')
                        self._undefined_metric_name.add(name)
                    continue
            if isinstance(value, torch.Tensor):
                # 以防万一
                value = value.item()
            self._batch_metrics[name].append(value)

        if self._retain_batch_history:
            self._batch_history.append(metrics)

    def compute(self, auto_reset: bool = False) -> dict[str, float]:
        """reduce 当前 batch 的指标

        建议在每个 epoch 结束之后调用, 并且调用 compute 之后需要手动调用 reset 清空当前 step 的数据

        Returns
        -------
        dict[str, float]
            指标名称和对应的计算结果
        """
        results = {}
        for name, reducer in self._metrics.items():
            if self._batch_metrics[name]:  # 确保有数据
                results[name] = reducer(self._batch_metrics[name])

        # 保存结果并清空当前 step 的数据
        if self._retain_epoch_history:
            self._epoch_history.append(results)

        if auto_reset:
            self.reset()

        return results

    def reset(self):
        self._batch_metrics = {name: [] for name in self._metrics.keys()}


class TorchMetricRecorder:
    def __init__(
        self,
        metrics: MetricCollection | typing.Mapping[str, Metric] | None = None,
        before_metric: Literal['sigmoid', 'softmax']
        | ActivationFn
        | typing.Mapping[str | Literal['*'], ActivationFn]
        | None = None,
        scalar_metrics: ScalarMetric
        | typing.Sequence[str]
        | str
        | Literal['auto']
        | None = 'loss',
    ):
        """添加基于 `torchmetrics` 的日志评估指标的功能

        Trainstep-wise Metric --> 记录 Metric Value --> Epoch --> 计算 Epochwise Metric (Reduce) --> 记录 Epochwise Metric

        区分两种类型:

        - 需要计算的指标，如 Accuracy 等；这些传入 torchmetric.Metric 或者 MetricCollection 对象
        - 需要记录的标量数据，如 loss 等; 传入名称字符串
            - 这些标量指标来源，应当是各个 training_step 返回的字典中包含的变量
            - 例如默认内置的 training_step 返回的字典中为 `{'loss': loss}`，所以 TorchMetricRecorder 默认会记录 loss 指标

        Parameters
        ----------
        - `metrics` : `MetricCollection | Dict[str, Metric]`, optional
            - 评估指标; 指哪些需要通过计算得到的指标，例如 Accuracy 等; 依赖于 torchmetrics 库
            - 使用 `torchmetrics.MetricCollection` 或 `dict` 类型
        - `before_metric` : `'sigmoid' | 'softmax' | Callable | dict[str, Callable]`, optional
            - 在计算评估指标之前的激活函数: `(y_hat) -> y_activated`
            - 例如在二分类问题中，模型最后一层没有激活函数, 在计算评估指标 (即 `metrics`) 之前需要使用 `torch.sigmoid` 激活
        - `scalar_metrics` : `ScalarMetric | str | list[str] | Literal['auto']`, optional
            - 用于记录一些外部传进来的标量数据，例如 loss 等
            - 如果传入 `ScalarMetric`，则直接使用
                - 通过 `self.update` 方法中的 `**metrics` 参数传入的指标都会被 update 到 `scalar_metrics` 中
            - 如果传入 `str`，则自动创建一个 ScalarMetric 并定义该指标 (mean)
            - 如果传入 `list[str]`, 则自动创建一个 ScalarMetric 并定义这些指标 (mean)
            - 如果传入 `'auto'`，则创建一个启用 auto_define 的 ScalarMetric，自动记录所有传入的指标

        `before_metric`:
        ----------
        在计算评估指标之前的激活函数: `(y_hat) -> y_activated`

        例如在二分类问题中，模型最后一层没有激活函数, 在计算评估指标 (即 `metrics`) 之前需要使用 `torch.sigmoid` 激活

        - `None`: 不需要额外的激活函数
        - `Literal['sigmoid', 'softmax']`: 使用 torch.sigmoid 和 torch.softmax
        - `(Tensor) -> Tensor`: 指定的激活函数
        - `dict[str, (Tensor) -> Tensor]`: 在比较复杂的情况下, 各个指标使用的输入不同的情况下可以用这个
            - key: 支持 glob 风格匹配模式：
                - 精确匹配：`"accuracy"` 匹配名为 "accuracy" 的指标
                - 通配符：`"*"` 匹配所有指标
                - glob 模式：`"U*"` 匹配所有以 "U" 开头的指标，`"train_*"` 匹配所有以 "train_" 开头的指标
            - value: `(Tensor) -> Tensor`, 为 None，则不激活
        """
        if before_metric == 'sigmoid':
            self.before_metric = torch.sigmoid
        elif before_metric == 'softmax':
            self.before_metric = torch.softmax
        else:
            self.before_metric = before_metric

        # 只保留一个 metrics collection
        if metrics is None:
            self.metrics = MetricCollection({})
        elif isinstance(metrics, dict):
            self.metrics = MetricCollection(
                typing.cast(dict[str, Metric | MetricCollection], metrics)
            )
        else:
            self.metrics = metrics

        # 记录哪些传入 update 的标量数据, 包括 loss 等
        if isinstance(scalar_metrics, ScalarMetric):
            self.scalar_metrics = scalar_metrics
        elif scalar_metrics == 'auto':
            # 自动模式：创建启用 auto_define 的 ScalarMetric
            self.scalar_metrics = ScalarMetric(default_reducer='mean', auto_define=True)
        else:
            self.scalar_metrics = ScalarMetric(default_reducer='mean')
            if isinstance(scalar_metrics, str):
                self.scalar_metrics.define_metric(scalar_metrics, 'mean')
            elif isinstance(scalar_metrics, (list, tuple)):
                for metric_name in scalar_metrics:
                    self.scalar_metrics.define_metric(metric_name, 'mean')
            elif scalar_metrics is not None:
                warning(f'不符合标准的参数 scalar_metrics: {scalar_metrics}')
            # else:
            #     self.scalar_metrics.define_metric('loss', 'mean')

        # self.step_history = {metric_name: [] for metric_name in self.metrics.keys()}
        # self.step_history['loss'] = []
        # self.epoch_history: list[dict[str, float | torch.Tensor]] = []

        self.current_epoch = 0
        self.current_step = 0

        # 激活函数匹配缓存
        self._activation_cache: Dict[str, ActivationFn | None] = {}

    def _find_matching_activation(self, metric_name: str) -> ActivationFn | None:
        """为指定的指标名称找到匹配的激活函数

        匹配优先级：精确匹配 > glob 模式匹配
        使用缓存机制提升性能，只在第一次匹配时进行计算

        Parameters
        ----------
        metric_name : str
            指标名称

        Returns
        -------
        Callable | None
            匹配的激活函数，如果没有匹配则返回 None
        """
        # 检查缓存
        if metric_name in self._activation_cache:
            return self._activation_cache[metric_name]

        if not isinstance(self.before_metric, dict):
            self._activation_cache[metric_name] = None
            return None

        activation_fn = None
        matched_pattern = None

        # 1. 精确匹配
        if metric_name in self.before_metric:
            activation_fn = self.before_metric[metric_name]
            matched_pattern = f"'{metric_name}' (精确匹配)"
        else:
            # 2. glob 模式匹配
            for pattern, fn in self.before_metric.items():
                if pattern != metric_name and fnmatch.fnmatch(metric_name, pattern):
                    activation_fn = fn
                    matched_pattern = f"'{pattern}' (glob匹配)"
                    break

        # 缓存结果并输出调试信息
        self._activation_cache[metric_name] = activation_fn

        if matched_pattern:
            fn_name = getattr(activation_fn, '__name__', str(activation_fn))
            print_debug(
                f"Metric '{metric_name}' 匹配处理器 {matched_pattern} -> {fn_name}"
            )
        else:
            print_debug(f"Metric '{metric_name}' 无匹配处理器，使用原始输入")

        return activation_fn

    def to_device(self, device: str | torch.device):
        # 简化 to_device，只需要处理一个 metrics collection
        for key, metric in self.metrics.items():
            use_device(metric, device)

    def _apply_activation(
        self, y_hat: torch.Tensor
    ) -> torch.Tensor | dict[str, torch.Tensor]:
        # 处理 before_metric 激活
        if self.before_metric is None:
            y_hat_activated = y_hat
        elif isinstance(self.before_metric, dict):
            # 如果是字典类型，需要对不同的指标使用不同的激活函数
            y_hat_activated = {}
            for metric_name in self.metrics.keys():
                # 使用 glob 匹配查找激活函数
                metric_name = typing.cast(str, metric_name)
                activation_fn = self._find_matching_activation(metric_name)

                if activation_fn is not None:
                    y_hat_activated[metric_name] = activation_fn(y_hat)
                else:
                    y_hat_activated[metric_name] = y_hat
        else:
            # 如果是单个激活函数
            before_metric = self.before_metric
            if before_metric is torch.softmax:
                y_hat_activated = torch.softmax(y_hat, dim=-1)
            else:
                activation_fn = typing.cast(ActivationFn, before_metric)
                y_hat_activated = activation_fn(y_hat)
        return y_hat_activated

    def update(self, pred, target, **scalar_metrics: float):
        """更新某个 step 的评估指标

        Parameters
        ----------
        - `pred` : `torch.Tensor`
            - 预测值
        - `target` : `torch.Tensor`
            - 真实值
        - `**scalar_metrics` : float
            - 其他标量指标; 默认应该包含 loss
        """
        result = {}

        # 更新其他标量指标
        if scalar_metrics:
            self.scalar_metrics.update(**scalar_metrics)
            result.update(scalar_metrics)

        # 更新 torchmetrics，如果 metrics 不为空
        if self.metrics and len(self.metrics) > 0:
            pred_activated = self._apply_activation(pred)
            for metric_name, metric in self.metrics.items():
                if isinstance(pred_activated, dict):
                    current_y_hat = pred_activated[metric_name]
                else:
                    current_y_hat = pred_activated

                step_metric_value = metric.update(
                    current_y_hat, target
                )  # 记录本 step 的指标  # NOTE 注意这里是否正确

                # 记录计算得到的指标
                self._record_compute_value(result, metric_name, step_metric_value)

        self.current_step += 1
        return result

    def compute(self):
        """计算一个 epoch 的评估指标; 并自动 reset"""
        epoch_metrics = {}

        # 计算标量指标
        scalar_results = self.scalar_metrics.compute()
        epoch_metrics.update(scalar_results)
        self.scalar_metrics.reset()

        # 计算 torchmetrics，如果 metrics 不为空
        if self.metrics and len(self.metrics) > 0:
            for metric_name, metric in self.metrics.items():
                epoch_metric_value = metric.compute()
                # epoch_metrics[metric_name] = epoch_metric_value.item()
                self._record_compute_value(epoch_metrics, metric_name, epoch_metric_value)
                metric.reset()

        self.current_epoch += 1
        self.current_step = 0
        return epoch_metrics

    def reset(self):
        self.scalar_metrics.reset()
        if self.metrics and len(self.metrics) > 0:
            for metric in self.metrics.values():
                metric.reset()

    def _record_compute_value(
        self,
        recorder: dict[str, typing.Any],
        name: str,
        value: float | torch.Tensor | typing.Mapping,
    ):
        """将 metric 计算得到的结果记录在字典当中

        由于 metric 计算得到的结果可能是 dict 类型, 也可能是标量, 所以需要单独处理

        value 如果是 dict, 默认应当是只有一层!
        """
        if isinstance(value, dict):
            recorder.update(prefixed_dict(f'{name}.', value))
        elif isinstance(value, typing.Mapping):
            recorder.update(prefixed_dict(f'{name}.', dict(value)))
        elif not isinstance(value, dict) and value is not None:
            # 能转换成标量就记录标量，不能就记录原值
            scalar = ensure_scalar(value)
            recorder[name] = scalar or value
