from typing import Any, Callable, Literal

import torch
from torch import Tensor
from torchmetrics import Metric


class _LambdaMetric(Metric):
    """Lambda 指标的虚基类，包含共用逻辑

    这个基类实现了通用的指标计算逻辑，支持自定义函数和多种聚合策略。
    通过继承这个基类并覆盖类属性，可以创建具有不同配置的指标实例。
    """

    # 默认的类属性（可被子类覆盖）
    is_differentiable = False
    higher_is_better = True
    full_state_update = False

    def __init__(
        self,
        metric_name: str,
        fn: Callable[[Tensor, Tensor], Tensor],
        reduce_fn: Literal['mean', 'sum', 'max', 'min'] = 'mean',
        **kwargs,
    ):
        """初始化 Lambda 指标

        Parameters
        ----------
        - `metric_name` : `str`
            - 指标名称，用于标识和调试
        - `fn` : `Callable[[Tensor, Tensor], Tensor]`
            - 计算指标的函数，接受 (preds, target) 返回 metric_value
        - `reduce_fn` : `Literal['mean', 'sum', 'max', 'min']`, optional
            - 跨批次的聚合策略, by default 'mean'
        """
        super().__init__(**kwargs)
        self.metric_name = metric_name
        self.fn = fn
        self.reduce_fn = reduce_fn

        # 添加状态变量来累积指标值
        if reduce_fn == 'sum':
            default_value = torch.tensor(0.0)
            dist_reduce_fx = 'sum'
        elif reduce_fn == 'mean':
            # 对于均值，我们需要累积总和和计数
            default_value = torch.tensor(0.0)
            dist_reduce_fx = 'sum'
        elif reduce_fn == 'max':
            default_value = torch.tensor(float('-inf'))
            dist_reduce_fx = 'max'
        elif reduce_fn == 'min':
            default_value = torch.tensor(float('inf'))
            dist_reduce_fx = 'min'
        else:
            raise ValueError(
                f"不支持的reduce_fn: {reduce_fn}，支持的选项: 'mean', 'sum', 'max', 'min'"
            )

        self.add_state('values', default=default_value, dist_reduce_fx=dist_reduce_fx)

        # 对于均值计算，我们还需要计数
        if reduce_fn == 'mean':
            self.add_state('total', default=torch.tensor(0), dist_reduce_fx='sum')

    def update(self, preds: Tensor, target: Tensor) -> None:
        """更新指标状态

        Parameters
        ----------
        - `preds` : `Tensor`
            - 预测值
        - `target` : `Tensor`
            - 真实值
        """
        try:
            # 调用用户定义的函数计算指标
            metric_value = self.fn(preds, target)

            # 确保返回值是标量张量
            if not isinstance(metric_value, Tensor):
                metric_value = torch.tensor(
                    metric_value, device=preds.device, dtype=torch.float
                )

            # 如果不是标量，尝试取均值
            if metric_value.numel() > 1:
                metric_value = metric_value.mean()

            # 确保是标量
            metric_value = metric_value.squeeze()

            # 根据reduce_fn更新状态
            if self.reduce_fn == 'mean':
                self.values += metric_value
                self.total += 1
            elif self.reduce_fn == 'sum':
                self.values += metric_value
            elif self.reduce_fn == 'max':
                self.values = torch.max(self.values, metric_value)
            elif self.reduce_fn == 'min':
                self.values = torch.min(self.values, metric_value)

        except Exception as e:
            raise RuntimeError(
                f"计算指标 '{self.metric_name}' 时出错: {e}\n"
                f'请检查函数 {self.fn.__name__} 的实现是否正确'
            ) from e

    def compute(self) -> Tensor:
        """计算最终的指标值

        Returns
        ----------
        - `Tensor`
            - 聚合后的指标值
        """
        if self.reduce_fn == 'mean':
            if self.total == 0:
                return torch.tensor(0.0, device=self.values.device)
            return self.values / self.total
        else:
            return self.values

    def __repr__(self) -> str:
        return (
            f'LambdaMetric('
            f"metric_name='{self.metric_name}', "
            f'fn={self.fn.__name__}, '
            f'higher_is_better={self.higher_is_better}, '
            f"reduce_fn='{self.reduce_fn}')"
        )


def create_lambda_metric(
    metric_name: str,
    fn: Callable[[Tensor, Tensor], Tensor],
    higher_is_better: bool = True,
    reduce_fn: Literal['mean', 'sum', 'max', 'min'] = 'mean',
    is_differentiable: bool = False,
    full_state_update: bool = False,
) -> Metric:
    """创建基于自定义函数的通用指标

    允许用户传入自定义函数来计算指标，提供最大的灵活性，同时自动处理批次间的累积和同步。
    支持多种累积策略，并提供完整的类型提示。

    Parameters
    ----------
    - `metric_name` : `str`
        - 指标名称，用于标识和调试
    - `fn` : `Callable[[Tensor, Tensor], Tensor]`
        - 计算指标的函数，接受 (preds, target) 返回 metric_value
        - 函数应该返回一个标量张量或可以转换为标量的张量
    - `higher_is_better` : `bool`, optional
        - 指标是否越高越好，用于超参数优化, by default True
    - `reduce_fn` : `Literal['mean', 'sum', 'max', 'min']`, optional
        - 跨批次的聚合策略, by default 'mean'
    - `is_differentiable` : `bool`, optional
        - 指标是否可微分, by default False
    - `full_state_update` : `bool`, optional
        - 是否需要完整的状态更新, by default False

    Returns
    ----------
    - `Metric`
        - 配置好的指标实例

    Example
    --------
    ```python
    # 示例1: 简单的自定义准确率
    def my_accuracy(preds, target):
        return (preds.argmax(dim=-1) == target).float().mean()

    metric = create_lambda_metric("custom_acc", my_accuracy)

    # 示例2: 复杂的业务指标
    def signal_quality_score(preds, target):
        # 脑电信号质量评估的自定义逻辑
        snr = compute_snr(preds, target)
        correlation = compute_correlation(preds, target)
        return 0.7 * snr + 0.3 * correlation

    metric = create_lambda_metric("signal_quality", signal_quality_score,
                                  higher_is_better=True, reduce_fn='mean')
    ```
    """

    # 创建临时类，继承自虚基类
    class LambdaMetricInstance(_LambdaMetric):
        """临时类，用于覆盖类属性"""

        pass

    # 动态设置类属性以绕过 torchmetrics 的限制
    LambdaMetricInstance.is_differentiable = is_differentiable
    LambdaMetricInstance.higher_is_better = higher_is_better
    LambdaMetricInstance.full_state_update = full_state_update

    return LambdaMetricInstance(metric_name, fn, reduce_fn)


# 为了向后兼容，提供一个别名
LambdaMetric = create_lambda_metric
