"""
Evidential Learning 系列算法
综合了多种损失函数变体和完整的理论实现

参考:
https://nn.labml.ai/uncertainty/evidence/index.html
https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/uncertainty/evidence/__init__.py


logits(神经网络输出的原始结果) -->|激活, 通常是 relu| --> Evidence --> 狄利克雷参数

"""

import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Literal,
    Optional,
    Protocol,
    Tuple,
    Union,
    cast,
)

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torchmetrics import Metric

from zyplib.nn.modules import BaseModule
from zyplib.utils.print import warning

# from . import sl_torch

if TYPE_CHECKING:
    from zyplib.train.common import FabricTrainingContext


class _OptimizerLike(Protocol):
    def zero_grad(self) -> None: ...

    def step(self) -> None: ...


class UncertaintyType(Enum):
    """不确定性类型枚举"""

    ALEATORIC = 'aleatoric'  # 偶然不确定性（数据噪声）
    EPISTEMIC = 'epistemic'  # 认知不确定性（模型不确定性）
    TOTAL = 'total'  # 总不确定性


class LossType(Enum):
    """Evidential损失函数类型"""

    # 分类任务
    TYPE_II_ML = 'type_ii_ml'  # Type II Maximum Likelihood
    BAYES_RISK_CE = 'bayes_risk_ce'  # Bayes Risk with Cross Entropy
    BAYES_RISK_MSE = 'bayes_risk_mse'  # Bayes Risk with MSE
    EXPECTED_RISK = 'expected_risk'  # Expected Risk (default)

    # 回归任务
    NIG_NLL = 'nig_nll'  # Normal-Inverse-Gamma NLL
    NIG_BAYES_RISK = 'nig_bayes_risk'  # NIG with Bayes Risk


@dataclass(frozen=True)
class UncertaintyMetrics:
    """不确定性度量结果"""

    aleatoric: Tensor  # 偶然不确定性
    epistemic: Tensor  # 认知不确定性
    predictive: Tensor  # 预测不确定性
    confidence: Tensor  # 置信度
    expected_prob: Tensor  # 期望概率（分类）或期望值（回归）
    alpha: Optional[Tensor] = None  # 分布参数（可选）


@dataclass
class EvidentialConfig:
    """Evidential Learning 基础配置"""

    task_type: Literal['classification', 'regression'] = 'classification'
    num_classes: Optional[int] = None
    output_dim: Optional[int] = None  # 回归任务的输出维度

    # Evidence 激活函数配置
    # 用于将 backbone 输出的原始 logits 转换为非负的 evidence 值
    evidence_activation: Literal['relu', 'exp', 'softplus', None] = 'relu'

    # 损失函数配置
    loss_type: LossType = LossType.EXPECTED_RISK

    # 数值稳定性
    eps: float = 1e-8

    # 验证相关
    validate_model: bool = True


@dataclass
class RegularizationConfig:
    """证据正则化配置"""

    # KL散度正则化
    enable_kl_regularization: bool = True
    kl_weight: float = 1.0
    """
    原始论文的意图是通过 α̃ = y + (1-y)⊙α 的设计来“移除无误导性的证据”，从而在模型摇摆不定或出错时施加惩罚。
    mistake_only 开关通过一个硬性的 argmax 判断，进行强制性的惩罚。
    """
    mistake_only: bool = False  # 是否只对错误预测施加KL正则

    # 方差正则化（分类任务，较少使用）
    enable_variance_regularization: bool = False
    variance_weight: float = 1.0

    # 精度正则化（回归任务）
    enable_precision_regularization: bool = False
    precision_weight: float = 1.0

    # 退火调度
    # EDL 原始论文推荐使用 epoch / 10 的步长来退火
    enable_annealing: bool = True
    annealing_start: int = 10  # 前10个epoch专注于学习基本分类
    annealing_steps: int = 10  # 10个epoch内线性增长到完全正则化
    min_annealing_coeff: float = 0.01  # 最小系数，不完全关闭正则化
    max_annealing_coeff: float = 1.0  # 最大系数，完全正则化

    # 退火系数随 epoch 的变化：
    # Epoch 0-9:   coeff = 0.01  (很小的正则化)
    # Epoch 10:    coeff = 0.01  (开始退火)
    # Epoch 15:    coeff = 0.505 (中等正则化)
    # Epoch 20:    coeff = 1.0   (完全正则化)
    # Epoch 20+:   coeff = 1.0   (保持完全正则化)


def _validate_evidential_model(
    model: nn.Module, input_shape: Tuple[int, ...], config: EvidentialConfig
) -> None:
    """验证 Evidential 模型的输入输出维度是否正确

    同时提供一些关于 backbone 输出特性的建议性检查，帮助用户发现潜在问题。
    """
    device = next(model.parameters()).device
    batch_size = 10

    with torch.no_grad():
        try:
            x = torch.randn(batch_size, *input_shape, device=device)
            output = model(x)

            if config.task_type == 'classification':
                if config.num_classes is None:
                    raise ValueError('分类任务必须指定 num_classes')

                expected_shape = (batch_size, config.num_classes)
                if output.shape != expected_shape:
                    raise ValueError(
                        f'分类模型输出形状 {output.shape} 与期望形状 {expected_shape} 不匹配'
                    )

                # 检查输出是否可能已经应用了不合适的激活函数
                _check_backbone_output_characteristics(output, config.task_type)

            elif config.task_type == 'regression':
                expected_shape = (batch_size, 4 * (config.output_dim or 1))
                if output.shape != expected_shape:
                    raise ValueError(
                        f'回归模型输出形状 {output.shape} 与期望形状 {expected_shape} 不匹配'
                    )

                # 对回归任务也进行输出特性检查
                _check_backbone_output_characteristics(output, config.task_type)

            else:
                raise ValueError(f'不支持的任务类型: {config.task_type}')

        except Exception as e:
            raise ValueError(f'模型无法处理形状为 {input_shape} 的输入: {e}')


def _check_backbone_output_characteristics(output: Tensor, task_type: str) -> None:
    """检查 backbone 输出的特性，提供建议性警告"""
    from zyplib.utils.print import warning

    # 检查是否所有输出都是非负的（可能应用了 ReLU）
    all_non_negative = torch.all(output >= 0)

    # 检查是否有很多值正好为 0（ReLU 的特征）
    zero_ratio = torch.mean((output == 0).float()).item()

    # 检查输出范围是否受限（可能应用了 Sigmoid/Tanh）
    output_min, output_max = output.min().item(), output.max().item()

    if task_type == 'classification':
        # 分类任务的检查
        if all_non_negative and zero_ratio > 0.1:
            warning(
                f'检测到 backbone 输出全为非负值且有 {zero_ratio:.1%} 的值为 0，'
                '这可能表明 backbone 末尾应用了 ReLU 激活函数。'
                'EDL 期望 backbone 输出原始 logits，激活函数应由 evidence_activation 处理。'
            )

        # 检查是否输出在 [0,1] 范围内（可能应用了 Softmax）
        if 0 <= output_min and output_max <= 1:
            # 进一步检查每行是否接近归一化
            row_sums = torch.sum(output, dim=-1)
            if torch.allclose(row_sums, torch.ones_like(row_sums), atol=0.1):
                warning(
                    '检测到 backbone 输出似乎已经归一化（每行和接近1），'
                    '这可能表明应用了 Softmax。EDL 期望原始 logits 而非概率分布。'
                )

    elif task_type == 'regression':
        # 回归任务的检查
        if all_non_negative and zero_ratio > 0.1:
            warning(
                f'检测到 backbone 输出全为非负值且有 {zero_ratio:.1%} 的值为 0，'
                '回归任务的 backbone 应输出未约束的原始值，而非经过 ReLU 处理的值。'
            )

        # 检查输出是否在 [0,1] 或 [-1,1] 等典型激活函数范围内
        if 0 <= output_min and output_max <= 1:
            warning(
                '检测到 backbone 输出范围在 [0,1] 内，这可能表明应用了 Sigmoid 激活函数。'
                'EDL 回归期望未约束的原始输出值。'
            )
        elif -1 <= output_min and output_max <= 1:
            warning(
                '检测到 backbone 输出范围在 [-1,1] 内，这可能表明应用了 Tanh 激活函数。'
                'EDL 回归期望未约束的原始输出值。'
            )


class EvidenceActivation:
    """证据激活函数工具类"""

    @staticmethod
    def relu_evidence(logits: Tensor, eps: float = 1e-8) -> Tensor:
        """ReLU激活函数获取证据"""
        return F.relu(logits) + eps

    @staticmethod
    def exp_evidence(logits: Tensor, max_val: float = 10.0, eps: float = 1e-8) -> Tensor:
        """指数激活函数获取证据"""
        return torch.exp(torch.clamp(logits, max=max_val)) + eps

    @staticmethod
    def softplus_evidence(logits: Tensor, eps: float = 1e-8) -> Tensor:
        """Softplus激活函数获取证据"""
        return F.softplus(logits) + eps

    @classmethod
    def get_activation(cls, activation_type: str | None):
        """获取激活函数

        Parameters
        ----------
        - `activation_type` : `str`
            - 激活函数类型，支持: 'relu', 'exp', 'softplus', 'none'/None
            - 'none' 或 None 表示不进行额外激活（假设 backbone 已经输出了 evidence）

        Returns
        ----------
        - `callable`
            - 对应的激活函数

        Raises
        ------
        - `ValueError`
            - 如果激活函数类型不被支持
        """
        mapping = {
            'relu': cls.relu_evidence,
            'exp': cls.exp_evidence,
            'softplus': cls.softplus_evidence,
            'none': lambda x, **kwargs: x,  # 恒等函数，支持额外参数
            None: lambda x, **kwargs: x,
        }
        lookup_key = activation_type.lower() if activation_type is not None else None
        if lookup_key not in mapping:
            available_keys = [k for k in mapping.keys() if k is not None]
            raise ValueError(
                f'不支持的激活函数类型: {activation_type}. 支持的类型: {available_keys}'
            )
        return mapping[lookup_key]


class UncertaintyMetricsMixin:
    """不确定性度量混入类"""

    @staticmethod
    def epistemic_uncertainty(alpha: Tensor, num_classes: int) -> Tensor:
        """计算认识不确定性 (Epistemic Uncertainty) 原始 EDL 论文中定义的证据空缺度 (vacuity)

        对应了 Subjective Logic 中的 uncertainty mass

        u = W / (Σe + W)
        S = Σα = Σ(e_i + 1) = Σe + W
        """
        S = torch.sum(alpha, dim=-1, keepdim=False)
        return num_classes / S

    @staticmethod
    def aleatoric_uncertainty(alpha: Tensor) -> Tensor:
        """计算偶然不确定性 (Aleatoric Uncertainty) 期望的预测熵 E[H(p)]"""
        S = torch.sum(alpha, dim=-1, keepdim=False)
        prob = alpha / S.unsqueeze(-1)

        # 期望的分类不确定性
        expected_entropy = -torch.sum(
            prob * (torch.digamma(alpha + 1) - torch.digamma(S.unsqueeze(-1) + 1)), dim=-1
        )
        return expected_entropy

    @staticmethod
    def predictive_uncertainty(alpha: Tensor) -> Tensor:
        """计算总不确定性（预测熵） 预测分布的熵 H(E[p])"""
        S = torch.sum(alpha, dim=-1, keepdim=False)
        prob = alpha / S.unsqueeze(-1)
        return -torch.sum(prob * torch.log(prob + 1e-8), dim=-1)

    @staticmethod
    def differential_entropy(alpha: Tensor) -> Tensor:
        """计算Dirichlet分布的微分熵"""
        S = torch.sum(alpha, dim=-1, keepdim=False)
        K = alpha.size(-1)

        entropy = (
            torch.lgamma(S)
            - torch.sum(torch.lgamma(alpha), dim=-1)
            + (S - K) * torch.digamma(S)
            - torch.sum((alpha - 1) * torch.digamma(alpha), dim=-1)
        )
        return entropy


def ensure_onehot(target: torch.Tensor, num_classes: int) -> torch.Tensor:
    ndim = target.ndim
    shape = target.shape

    if ndim == 1:
        if torch.is_floating_point(target):
            return target
        return F.one_hot(target, num_classes).to(target.device).float()

    elif ndim == 2:
        # 2维输入：检查是否已经是 one-hot 编码
        if shape[1] != num_classes:
            raise ValueError(
                f'2维输入的第二维必须等于 num_classes={num_classes}，但得到 {shape[1]}'
            )

        return target

    else:
        raise ValueError(f'ensure_onehot 只支持1维或2维输入，但得到 {ndim} 维')


class LossFunctionsMixin:
    """损失函数混入类，包含所有Evidential Learning损失函数变体

    实现与标准参考文献一致：
    - Type II ML: Sensoy et al. Eq.3
    - Bayes Risk CE: Sensoy et al. Eq.4 (使用digamma函数)
    - Bayes Risk MSE: Sensoy et al. Eq.5 (误差项+方差项)
    - KL Regularization: Sensoy et al. Eq.7

    参考实现:
    - https://nn.labml.ai/uncertainty/evidence/index.html
    - https://muratsensoy.github.io/uncertainty.html

    剩下两个是 EDL 回归的代码
    """

    @staticmethod
    def type_ii_ml_loss(alpha: Tensor, target: Tensor, num_classes: int) -> Tensor:
        """Type II Maximum Likelihood Loss

        标准实现：-log p(y|D) = y^T (log S - log α)
        Reference: Sensoy et al. Eq.3 in the paper
        """
        S = torch.sum(alpha, dim=-1, keepdim=True)

        # 确保 target 为 one-hot 编码格式
        target_onehot = ensure_onehot(target, num_classes)

        # Type II ML: y^T (log S - log α)
        loss = torch.sum(target_onehot * (S.log() - alpha.log()), dim=-1)

        return torch.mean(loss)

    @staticmethod
    def bayes_risk_ce_loss(alpha: Tensor, target: Tensor, num_classes: int) -> Tensor:
        """Bayes Risk with Cross Entropy Loss

        标准实现：期望交叉熵损失 E[CE(y, p)] = y^T (ψ(S) - ψ(α))
        其中 ψ 是 digamma 函数
        Reference: Sensoy et al. Eq.4 in the paper
        """
        S = torch.sum(alpha, dim=-1, keepdim=True)

        # 确保 target 为 one-hot 编码格式
        target_onehot = ensure_onehot(target, num_classes)

        # Bayes Risk CE: y^T (ψ(S) - ψ(α))
        loss = torch.sum(
            target_onehot * (torch.digamma(S) - torch.digamma(alpha)), dim=-1
        )

        return torch.mean(loss)

    @staticmethod
    def bayes_risk_mse_loss(alpha: Tensor, target: Tensor, num_classes: int) -> Tensor:
        """Bayes Risk with MSE Loss (EDL 原始论文优先推荐)

        标准实现：期望平方误差 E[||y - p||²] = ||y - E[p]||² + Tr(Cov[p])
        Reference: Sensoy et al. Eq.5 in the paper
        """
        S = torch.sum(alpha, dim=-1, keepdim=True)

        # 确保 target 为 one-hot 编码格式
        target_onehot = ensure_onehot(target, num_classes)

        # 期望概率
        expected_prob = alpha / S

        # 误差项: ||y - E[p]||²
        err = (target_onehot - expected_prob) ** 2

        # 方差项: Tr(Cov[p]) - Dirichlet分布的方差
        var = expected_prob * (1 - expected_prob) / (S + 1)

        # 总损失
        loss = torch.sum(err + var, dim=-1)

        return torch.mean(loss)

    @staticmethod
    def expected_risk_loss(alpha: Tensor, target: Tensor, num_classes: int) -> Tensor:
        """Expected Risk Loss

        计算期望的分类风险
        """
        S = torch.sum(alpha, dim=-1, keepdim=False)

        # 确保 target 为 one-hot 编码格式
        target_onehot = ensure_onehot(target, num_classes)

        # 期望风险：E[1 - p_correct] = (S - α_correct) / S
        expected_risk = torch.sum(
            target_onehot * (S.unsqueeze(-1) - alpha) / S.unsqueeze(-1), dim=-1
        )

        return torch.mean(expected_risk)

    @staticmethod
    def nig_nll_loss(nig_params: Tensor, target: Tensor, eps: float = 1e-8) -> Tensor:
        """Normal-Inverse-Gamma Negative Log-Likelihood Loss"""
        # nig_params shape: [B, D, 4] or [B, 4] for D=1
        # target shape: [B, D] or [B] for D=1
        # gamma, nu, alpha, beta = torch.split(nig_params, 1, dim=-1)
        gamma, nu, alpha, beta = torch.unbind(nig_params, dim=-1)

        # 确保 target 和 params 维度兼容以便广播
        # 如果 target 是 [B], 而 params 是 [B, D, 4], target 扩展为 [B, 1]
        if target.dim() == nig_params.dim() - 2:
            target = target.unsqueeze(-1)  # -> [B, 1] for broadcasting

        # 保证数值稳定性
        nu = nu + eps
        alpha = alpha + eps
        beta = beta + eps

        # NIG NLL
        omega = 2 * beta * (1 + nu)
        nll = (
            0.5 * torch.log(torch.tensor(math.pi, device=nu.device) / nu)
            - alpha * torch.log(omega)
            + (alpha + 0.5) * torch.log(nu * (target - gamma) ** 2 + omega)
            + torch.lgamma(alpha)
            - torch.lgamma(alpha + 0.5)
        )

        return torch.mean(nll)

    @staticmethod
    def nig_bayes_risk_loss(
        nig_params: Tensor, target: Tensor, eps: float = 1e-8
    ) -> Tensor:
        """NIG Bayes Risk Loss"""
        # nig_params shape: [B, D, 4] or [B, 4]
        # target shape: [B, D] or [B]
        gamma, nu, alpha, beta = torch.unbind(nig_params, dim=-1)

        # 确保 target 和 params 维度兼容以便广播
        if target.dim() == nig_params.dim() - 2:
            target = target.unsqueeze(-1)

        # 预测分布的均值和方差
        pred_mean = gamma
        # 要求 alpha > 1，这里加 eps 增加鲁棒性
        pred_var = beta / (alpha - 1 + eps)

        # Bayes Risk (期望MSE)
        bias_sq = (target - pred_mean) ** 2
        risk = bias_sq + pred_var

        # 在批次和输出维度上取平均
        return torch.mean(risk)


class RegularizationMixin:
    """正则化混入类"""

    def __init__(self, reg_config: RegularizationConfig):
        self.reg_config = reg_config
        # self._current_epoch = 0

    def kl_divergence_regularization(
        self,
        alpha: Tensor,
        target: Tensor,
        num_classes: int,
        annealing_coeff: float = 1.0,
    ) -> Tensor:
        """KL散度正则化项

        标准实现：KL(Dir(α̃)||Dir(1)) 其中 α̃ = y + (1-y)⊙α
        鼓励模型在预测错误时产生较少的evidence

        Parameters
        ----------
        - `alpha` : `Tensor`
            - _description_
        - `target` : `Tensor`
            - _description_
        - `num_classes` : `int`
            - _description_
        - `annealing_coeff` : `float`, optional
            - _description_, by default 1.0

        Returns
        ----------
        - `Tensor`
            - _description_
        """
        # 确保 target 为 one-hot 编码格式
        target_onehot = ensure_onehot(target, num_classes)

        # 移除误导性 evidence, 保留所有没有正确分类的样本: α̃ = y + (1-y)⊙α
        alpha_tilde = target_onehot + (1 - target_onehot) * alpha
        S_tilde = torch.sum(alpha_tilde, dim=-1, keepdim=False)

        # KL(Dir(α̃)||Dir(1))的标准计算
        # 第一项: log B(1) - log B(α̃)
        first_term = (
            torch.lgamma(S_tilde)
            - torch.lgamma(torch.tensor(float(num_classes)))
            - torch.sum(torch.lgamma(alpha_tilde), dim=-1)
        )

        # 第二项: (α̃ - 1)^T (ψ(α̃) - ψ(S̃))
        second_term = torch.sum(
            (alpha_tilde - 1)
            * (torch.digamma(alpha_tilde) - torch.digamma(S_tilde.unsqueeze(-1))),
            dim=-1,
        )

        kl = first_term + second_term

        # 是否只对预测错误的样本施加KL正则化
        if self.reg_config.mistake_only:
            pred = torch.argmax(alpha, dim=-1)
            # 根据 target 的格式确定真实标签
            if target.dim() == 1:
                # target 是类别索引
                true_labels = target.long()
            else:
                # target 是 one-hot 编码，取 argmax
                true_labels = torch.argmax(target_onehot, dim=-1)

            mistake_mask = (pred != true_labels).float()
            kl = kl * mistake_mask

        return annealing_coeff * torch.mean(kl)

    def variance_regularization(self, alpha: Tensor) -> Tensor:
        """方差正则化项（分类任务）

        目标是直接最小化 Dirichlet 分布的方差  Var(p_k) = α_k * (S - α_k) / (S^2 * (S + 1))
        鼓励模型在所有类别上都做出更“确定”（即低方差）的预测。
        不是 EDL 框架的核心组成部分，没有特别理论根据。
        """
        S = torch.sum(alpha, dim=-1, keepdim=False)
        variance = (alpha * (S.unsqueeze(-1) - alpha)) / (
            S.unsqueeze(-1) ** 2 * (S.unsqueeze(-1) + 1)
        )
        return torch.mean(torch.sum(variance, dim=-1))

    def precision_regularization(
        self, nig_params: Tensor, target: Tensor, eps: float = 1e-8
    ) -> Tensor:
        """精度正则化项（回归任务）"""
        # gamma, nu, alpha, beta = torch.split(nig_params, 1, dim=-1)
        gamma, nu, alpha, beta = torch.unbind(nig_params, dim=-1)

        # 确保 target 和 params 维度兼容以便广播
        if nig_params.dim() == 3:  # [B, D, 4]
            if target.dim() == 1:  # [B] -> [B, 1]
                target = target.unsqueeze(-1)
            elif target.dim() == 2 and target.size(-1) != nig_params.size(1):
                raise ValueError(
                    f'Target dim {target.shape} incompatible with params {nig_params.shape}'
                )

        # 精度正则化：鼓励模型在预测错误时降低精度
        error = torch.abs(target - gamma)
        precision = nu + alpha  # 总精度的代理
        reg = error * precision

        return torch.mean(reg)

    def compute_regularization(
        self,
        alpha_or_params: Tensor,
        target: Tensor,
        num_classes: Optional[int] = None,
        task_type: str = 'classification',
        annealing_coeff: float = 1.0,
    ) -> Tensor:
        """计算总正则化项"""
        total_reg = torch.tensor(0.0, device=alpha_or_params.device)

        if task_type == 'classification':
            if self.reg_config.enable_kl_regularization:
                assert num_classes is not None
                kl_reg = self.kl_divergence_regularization(
                    alpha_or_params, target, num_classes, annealing_coeff
                )
                total_reg = total_reg + self.reg_config.kl_weight * kl_reg

            if self.reg_config.enable_variance_regularization:
                var_reg = self.variance_regularization(alpha_or_params)
                total_reg = total_reg + self.reg_config.variance_weight * var_reg

        elif task_type == 'regression':
            if self.reg_config.enable_precision_regularization:
                prec_reg = self.precision_regularization(alpha_or_params, target)
                total_reg = total_reg + self.reg_config.precision_weight * prec_reg

        return total_reg

    def get_annealing_coeff(self, epoch: int | None) -> float:
        """获取退火系数"""

        if not self.reg_config.enable_annealing:
            return self.reg_config.max_annealing_coeff

        if self.reg_config.annealing_steps <= 0:
            return self.reg_config.max_annealing_coeff

        if epoch is None:
            return (
                self.reg_config.max_annealing_coeff
            )  # 在测试阶段, 无 epoch，则需要返回完整的约束 coeff

        if epoch < self.reg_config.annealing_start:
            return self.reg_config.min_annealing_coeff

        if epoch >= self.reg_config.annealing_start + self.reg_config.annealing_steps:
            return self.reg_config.max_annealing_coeff

        progress = min(
            self.reg_config.max_annealing_coeff,
            (epoch - self.reg_config.annealing_start) / self.reg_config.annealing_steps,
        )
        return (
            self.reg_config.min_annealing_coeff
            + (self.reg_config.max_annealing_coeff - self.reg_config.min_annealing_coeff)
            * progress
        )


class BasicEDLFabricTrainProcedure:
    """Evidential Learning 训练过程"""

    def __init__(self, config: EvidentialConfig):
        self.config = config

    def evidential_train_step(
        self,
        ctx: 'FabricTrainingContext',
        batch: Tuple[Tensor, Tensor],
        epoch_idx: int,
        **kwargs,
    ) -> Tuple[Any, Dict[str, float]]:
        """Evidential训练步骤"""
        evidential_model = cast(EvidentialBase, ctx.model)
        optimizer = cast(_OptimizerLike, ctx.optimizer)
        fabric = ctx.fabric

        x, y = batch

        # 前向传播
        evidential_model.train()
        optimizer.zero_grad()

        logits = evidential_model(x)
        loss_dict = evidential_model.compute_loss(logits, y, epoch=epoch_idx)

        total_loss = loss_dict['total_loss']

        # 反向传播
        if fabric:
            fabric.backward(total_loss)
        else:
            total_loss.backward()

        optimizer.step()

        # 转换为可记录的指标
        metrics = {
            k: v.item() if isinstance(v, Tensor) else v for k, v in loss_dict.items()
        }

        return logits, metrics

    def evidential_val_step(
        self,
        ctx: 'FabricTrainingContext',
        batch: Tuple[Tensor, Tensor],
        epoch_idx: Optional[int] = None,
        **kwargs,
    ) -> Tuple[Any, Dict[str, float]]:
        """Evidential验证步骤

        测试阶段设置 epoch_idx 为 None, 就会按照完全 KL 约束来计算
        """
        evidential_model = cast(EvidentialBase, ctx.model)

        x, y = batch

        evidential_model.eval()
        with torch.no_grad():
            logits = evidential_model(x)
            loss_dict = evidential_model.compute_loss(logits, y, epoch=epoch_idx)

            # 计算不确定性指标
            # uncertainty = evidential_model.compute_uncertainty(logits)

        # 合并损失和不确定性指标
        metrics = {
            k: v.item() if isinstance(v, Tensor) else v for k, v in loss_dict.items()
        }

        # 和 Train 保持一致; 计算不确定性独立的单独用一个 EDLUncertaintyMetric 来实现
        # uncertainty_metrics = {
        #     'mean_aleatoric': uncertainty.aleatoric.mean().item(),
        #     'mean_epistemic': uncertainty.epistemic.mean().item(),
        #     'mean_total_uncertainty': uncertainty.total.mean().item(),
        #     'mean_confidence': uncertainty.confidence.mean().item(),
        # }
        # metrics.update(uncertainty_metrics)

        return logits, metrics


class EvidentialBase(BaseModule, ABC):
    """Evidential Learning 基类"""

    def __init__(
        self,
        backbone: nn.Module,
        config: EvidentialConfig,
        reg_config: RegularizationConfig | None,
    ):
        super().__init__()
        self.backbone = backbone
        self.config = config
        self.reg_config = reg_config
        self.evidence_activation = EvidenceActivation.get_activation(
            config.evidence_activation
        )

    EDL_TYPE = 'EvidentialBase'

    @abstractmethod
    def logits_to_evidence(self, logits: Tensor) -> Tensor:
        """激活evidence"""
        pass

    @abstractmethod
    def compute_loss(
        self, logits: Tensor, target: Tensor, epoch: Optional[int] = None
    ) -> Dict[str, Tensor]:
        """计算损失"""
        pass

    @abstractmethod
    def compute_uncertainty(self, logits: Tensor) -> UncertaintyMetrics:
        """计算不确定性"""
        pass

    def forward(self, x: Tensor) -> Tensor:
        """前向传播"""
        return self.backbone(x)

    @torch.no_grad()
    def predict_with_uncertainty(
        self, x: Tensor, return_raw: bool = False
    ) -> Union[UncertaintyMetrics, Tuple[UncertaintyMetrics, Tensor]]:
        """预测并返回不确定性"""
        self.eval()
        logits = self.forward(x)
        uncertainty = self.compute_uncertainty(logits)

        if return_raw:
            return uncertainty, logits
        return uncertainty


class EvidentialClassifier(EvidentialBase, UncertaintyMetricsMixin, LossFunctionsMixin):
    EDL_TYPE = 'EvidentialClassifier'
    SUPPORTED_LOSS_TYPES = [
        LossType.TYPE_II_ML,
        LossType.BAYES_RISK_CE,
        LossType.BAYES_RISK_MSE,
        LossType.EXPECTED_RISK,
    ]

    def __init__(
        self,
        backbone: nn.Module,
        config: Optional[EvidentialConfig] = None,
        reg_config: Optional[RegularizationConfig] = None,
        num_classes: Optional[int] = None,
        input_shape: Optional[Tuple[int, ...]] = None,
    ):
        """EDL 分类器

        基于 Dirichlet 分布对分类任务的类别概率进行建模，从而量化不确定性。
        模型的核心思想是训练一个神经网络，使其输出作为 Dirichlet 分布的参数（称为 evidence），
        而不是直接输出类别概率。

        参考文献:
        ----------
        - [1] "Evidential Deep Learning to Quantify Classification Uncertainty"
            (Sensoy, M., Kaplan, L., & Kandemir, M., 2018). NeurIPS.
            (EDL 框架的奠基之作)
        - [2] https://nn.labml.ai/uncertainty/evidence/index.html
            (一个优秀的参考实现)

        Parameters
        ----------
        - `backbone` : `nn.Module`
            - 神经网络模型。
            - **重要**: 模型的最后一层不应包含 `softmax`, `log_softmax`, `relu` 等激活函数。
            - 它应该输出**原始的、未约束的 logits**; 这些 logits 将由 EDL 模块内部的
              `evidence_activation` (如 ReLU, Softplus) 转换为非负的 evidence 值。
            - 期望的输出形状为 `(batch_size, num_classes)`。

        - `config` : `Optional[EvidentialConfig]`, optional
            - Evidential Learning 的核心配置，包括损失函数类型、证据激活函数等。
            - 如果为 `None`，将创建一个默认的分类任务配置。 (默认为 None)

        - `reg_config` : `Optional[RegularizationConfig]`, optional
            - 正则化相关的配置，如是否启用KL散度正则化、退火策略等。
            - 如果为 `None`，则在训练过程中不使用任何正则化项。 (默认为 None)

        - `input_shape` : `Optional[Tuple[int, ...]]`, optional
            - 单个输入样本的形状，例如 `(C, H, W)`。
            - 如果提供，将在初始化时运行一次前向传播检查，以验证 `backbone` 的
              输出形状是否正确，有助于及早发现模型定义错误。 (默认为 None)

        - `num_classes` : `int`
            - 分类任务的目标类别数。必须与 `backbone` 的输出维度匹配。

        Raises
        ------
        - `ValueError`
            - 如果 `config` 中指定的 `loss_type` 不被分类任务支持。
            - 如果 `input_shape` 被提供，但模型验证失败（例如输出形状不匹配）。
        """
        # 设置默认配置
        if config is None:
            assert num_classes is not None, '如果 config 为 None, 则 num_classes 必须提供'
            config = EvidentialConfig(task_type='classification', num_classes=num_classes)
        else:
            if num_classes is not None and config.num_classes != num_classes:
                raise ValueError(
                    f'config.num_classes != num_classes: {config.num_classes} != {num_classes}'
                )
            if config.task_type != 'classification':
                warning(
                    'EvidentialClassifier 只支持分类任务，将自动设置 config.task_type 为 "classification"'
                )
                config.task_type = 'classification'

        super().__init__(backbone, config, reg_config)
        self.num_classes = config.num_classes
        if config.loss_type not in self.SUPPORTED_LOSS_TYPES:
            raise ValueError(f'EvidentialClassifier 不支持 {config.loss_type} 损失函数')

        # 初始化正则化器
        if reg_config is not None:
            self.regularization = RegularizationMixin(reg_config)

        # 模型验证
        if config.validate_model and input_shape is not None:
            _validate_evidential_model(backbone, input_shape, config)

    def logits_to_evidence(self, logits: Tensor) -> Tensor:
        """激活evidence"""
        return self.evidence_activation(logits, eps=self.config.eps)

    def evidence_to_alpha(self, evidence: Tensor) -> Tensor:
        """将evidence转换为Dirichlet参数"""
        return evidence + 1.0
        # return sl_torch.evidence_to_dirichlet(evidence)

    def compute_loss(
        self, logits: Tensor, target: Tensor, epoch: Optional[int] = None
    ) -> Dict[str, Tensor]:
        """计算分类损失"""
        num_classes = cast(int, self.num_classes)
        evidence = self.logits_to_evidence(logits)
        alpha = self.evidence_to_alpha(evidence)

        # 选择主损失函数
        if self.config.loss_type == LossType.TYPE_II_ML:
            main_loss = self.type_ii_ml_loss(alpha, target, num_classes)
        elif self.config.loss_type == LossType.BAYES_RISK_CE:
            main_loss = self.bayes_risk_ce_loss(alpha, target, num_classes)
        elif self.config.loss_type == LossType.BAYES_RISK_MSE:
            main_loss = self.bayes_risk_mse_loss(alpha, target, num_classes)
        else:  # EXPECTED_RISK
            main_loss = self.expected_risk_loss(alpha, target, num_classes)

        loss_dict = {'main_loss': main_loss, 'total_loss': main_loss}

        # 添加正则化项
        if hasattr(self, 'regularization'):
            annealing_coeff = self.regularization.get_annealing_coeff(epoch)

            reg_loss = self.regularization.compute_regularization(
                alpha, target, num_classes, 'classification', annealing_coeff
            )

            loss_dict.update(
                {
                    'reg_loss': reg_loss,
                    'total_loss': main_loss + reg_loss,
                    'annealing_coeff': torch.tensor(annealing_coeff),
                }
            )

        return loss_dict

    @torch.no_grad()
    def compute_uncertainty(self, logits: Tensor) -> UncertaintyMetrics:
        """计算分类任务的不确定性指标, no_grad"""
        num_classes = cast(int, self.num_classes)
        evidence = self.logits_to_evidence(logits)
        alpha = self.evidence_to_alpha(evidence)

        S = torch.sum(alpha, dim=-1, keepdim=False)
        expected_prob = alpha / S.unsqueeze(-1)

        # 计算各种不确定性
        epistemic = self.epistemic_uncertainty(alpha, num_classes)
        aleatoric = self.aleatoric_uncertainty(alpha)
        predictive = self.predictive_uncertainty(alpha)

        # 置信度
        confidence = torch.max(expected_prob, dim=-1)[0]

        return UncertaintyMetrics(
            aleatoric=aleatoric,
            epistemic=epistemic,
            predictive=predictive,
            confidence=confidence,
            expected_prob=expected_prob,
            alpha=alpha,
        )

    @torch.no_grad()
    def predict_prob(
        self, x: Tensor, return_uncertainty: bool = False
    ) -> Union[Tensor, Tuple[Tensor, UncertaintyMetrics]]:
        """预测类别概率

        对输入数据进行前向传播并返回类别概率预测。

        Parameters
        ----------
        - `x` : `Tensor`
            - 输入数据，形状为 `(batch_size, ...)`
        - `return_uncertainty` : `bool`, optional
            - 是否同时返回不确定性信息, by default False

        Returns
        ----------
        - `Union[Tensor, Tuple[Tensor, UncertaintyMetrics]]`
            - 如果 `return_uncertainty=False`: 返回概率预测 `(batch_size, num_classes)`
            - 如果 `return_uncertainty=True`: 返回 `(概率预测, 不确定性指标)` 的元组
        """
        self.eval()
        logits = self.forward(x)

        evidence = self.logits_to_evidence(logits)
        alpha = self.evidence_to_alpha(evidence)

        # 计算期望概率 (Dirichlet 分布的期望)
        S = torch.sum(alpha, dim=-1, keepdim=False)
        expected_prob = alpha / S.unsqueeze(-1)

        if return_uncertainty:
            uncertainty = self.compute_uncertainty(logits)
            return expected_prob, uncertainty

        return expected_prob


class EvidentialRegressor(EvidentialBase, LossFunctionsMixin):
    """Evidential 回归器 - 完整版本

    该回归器通过在单个确定性神经网络上放置一个 Normal-Inverse-Gamma (NIG)
    先验分布来学习回归任务中的不确定性。模型输出 NIG 分布的四个参数，
    从而能够同时预测目标值以及与之相关的偶然不确定性（数据方差）和
    认知不确定性（模型方差）。

    参考文献:
    ----------
    - [1] "Deep Evidential Regression"
          (Amini, A., Schwarting, W., Soleimany, A., & Rus, D., 2020). NeurIPS.
          (将 Evidential Learning 框架扩展到回归任务的开创性工作)
    """

    EDL_TYPE = 'EvidentialRegressor'

    def __init__(
        self,
        backbone: nn.Module,
        output_dim: int = 1,
        config: Optional[EvidentialConfig] = None,
        reg_config: Optional[RegularizationConfig] = None,
        input_shape: Optional[Tuple[int, ...]] = None,
    ):
        """Evidential 回归器

        该回归器通过在单个确定性神经网络上放置一个 Normal-Inverse-Gamma (NIG)
        先验分布来学习回归任务中的不确定性。模型输出 NIG 分布的四个参数，
        从而能够同时预测目标值以及与之相关的偶然不确定性（数据方差）和
        认知不确定性（模型方差）。

        参考文献:
        ----------
        - [1] "Deep Evidential Regression" (Amini, A., Schwarting, W., Soleimany, A., & Rus, D., 2020). NeurIPS.

        Parameters
        ----------
        - `backbone` : `nn.Module`
            - 神经网络模型。
            - **重要**: 模型的最后一层不应包含激活函数，应该输出原始的、未约束的 logits。
            - 期望的输出形状为 `(batch_size, 4 * output_dim)`
            - 这 4 个输出将由 EDL 模块内部处理并转换为 NIG 分布的四个参数:
              `gamma (γ)`, `nu (ν)`, `alpha (α)`, 和 `beta (β)`

        - `output_dim` : `int`
            - 回归任务的目标维度。例如，预测一个值时为 1。 (默认为 1)

        - `config` : `Optional[EvidentialConfig]`, optional
            - Evidential Learning 的核心配置。
            - 如果为 `None`，将创建一个默认的回归任务配置。 (默认为 None)

        - `reg_config` : `Optional[RegularizationConfig]`, optional
            - 正则化相关的配置，主要控制精度正则化和退火。
            - 如果为 `None`，则不使用正则化。 (默认为 None)

        - `input_shape` : `Optional[Tuple[int, ...]]`, optional
            - 单个输入样本的形状。
            - 如果提供，将在初始化时运行检查以验证 `backbone` 的输出形状是否正确。
              (默认为 None)

        Raises
        ------
        - `ValueError`
            - 如果 `config` 中指定的 `loss_type` 不被回归任务支持。
            - 如果 `input_shape` 被提供，但模型验证失败。
        """
        # 设置默认配置
        if config is None:
            config = EvidentialConfig(task_type='regression', output_dim=output_dim)
        else:
            config.output_dim = output_dim
            config.task_type = 'regression'

        super().__init__(backbone, config, reg_config)
        self.output_dim = output_dim

        # 初始化正则化器
        if reg_config is not None:
            self.regularization = RegularizationMixin(reg_config)

        # 模型验证
        if config.validate_model and input_shape is not None:
            _validate_evidential_model(backbone, input_shape, config)

    def logits_to_evidence(self, logits: Tensor) -> Tensor:
        """激活NIG参数"""
        # 重塑为 [gamma, nu, alpha, beta] for each output dimension
        params = logits.view(logits.size(0), self.output_dim, 4)

        gamma = params[..., 0]  # 无约束
        nu = F.softplus(params[..., 1]).clamp(min=self.config.eps)  # > 0
        alpha = (F.softplus(params[..., 2]) + 1).clamp(min=1.0 + self.config.eps)  # > 1
        beta = F.softplus(params[..., 3]) + self.config.eps  # > 0

        return torch.stack([gamma, nu, alpha, beta], dim=-1)

    def compute_loss(
        self, logits: Tensor, target: Tensor, epoch: Optional[int] = None
    ) -> Dict[str, Tensor]:
        """计算回归损失"""
        nig_params = self.logits_to_evidence(logits)

        # 选择主损失函数
        if self.config.loss_type == LossType.NIG_BAYES_RISK:
            main_loss = self.nig_bayes_risk_loss(nig_params, target)
        else:  # NIG_NLL
            main_loss = self.nig_nll_loss(nig_params, target)

        loss_dict = {'main_loss': main_loss, 'total_loss': main_loss}

        # 添加正则化项
        if hasattr(self, 'regularization'):
            annealing_coeff = self.regularization.get_annealing_coeff(epoch)

            reg_loss = self.regularization.compute_regularization(
                nig_params,
                target,
                task_type='regression',
                annealing_coeff=annealing_coeff,
            )

            loss_dict.update(
                {
                    'reg_loss': reg_loss,
                    'total_loss': main_loss + reg_loss,
                    'annealing_coeff': torch.tensor(annealing_coeff),
                }
            )

        return loss_dict

    @torch.no_grad()
    def compute_uncertainty(self, logits: Tensor) -> UncertaintyMetrics:
        """计算回归任务的不确定性; no_grad"""
        nig_params = self.logits_to_evidence(logits)
        gamma, nu, alpha, beta = torch.unbind(nig_params, dim=-1)

        # 预测均值和方差
        expected_value = gamma

        # 保证分母大于 0 (alpha 在激活时已保证 > 1)
        alpha_minus_one = alpha - 1 + self.config.eps

        # 偶然不确定性（数据噪声）
        aleatoric_uncertainty = beta / alpha_minus_one

        # 认知不确定性（模型不确定性）
        epistemic_uncertainty = beta / (nu * (alpha_minus_one))

        # 总不确定性
        total_uncertainty = beta * (1 + nu) / (nu * alpha_minus_one)

        # 置信度（方差的倒数）
        confidence = 1.0 / (total_uncertainty + self.config.eps)

        return UncertaintyMetrics(
            aleatoric=aleatoric_uncertainty,
            epistemic=epistemic_uncertainty,
            predictive=total_uncertainty,
            confidence=confidence,
            expected_prob=expected_value,
        )


class EDLUncertaintyMetric(Metric):
    full_state_update: bool = True

    METRIC_NAMES = ['aleatoric', 'epistemic', 'predictive', 'confidence']
    aleatoric_sum: Tensor
    epistemic_sum: Tensor
    pred_unc_sum: Tensor
    confidence_sum: Tensor
    total_samples: Tensor

    def __init__(self, edl_model: EvidentialBase, **kwargs: Any):
        """
        一个用于计算 Evidential Deep Learning (EDL) 不确定性度量的 TorchMetric 类。

        该 Metric 实现了 step-level 和 epoch-level 的指标返回：
        - `update()`: 计算并返回当前批次的平均不确定性，用于 step-level 日志。
        - `compute()`: 计算并返回整个数据集的平均不确定性，用于 epoch-level 日志。

        调用的时候会返回:

        ```py
        {
            'aleatoric': float,
            'epistemic': float,
            'predictive': float,
            'confidence': float,
        }
        ```
        """
        super().__init__(**kwargs)

        if not hasattr(edl_model, 'compute_uncertainty'):
            raise TypeError("提供的 'edl_model' 必须包含 'compute_uncertainty' 方法。")

        self.edl_model = edl_model

        # 内部状态仍然使用 '_sum' 后缀以示区分，但外部返回的键名会很干净
        self.add_state('aleatoric_sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
        self.add_state('epistemic_sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
        self.add_state('pred_unc_sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
        self.add_state('confidence_sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
        self.add_state('total_samples', default=torch.tensor(0), dist_reduce_fx='sum')

    # def forward(self, logits: Tensor, *args) -> Dict[str, Tensor]:
    #     uncertainty: UncertaintyMetrics = self.edl_model.compute_uncertainty(logits)
    #     return {
    #         'aleatoric': uncertainty.aleatoric.mean(),
    #         'epistemic': uncertainty.epistemic.mean(),
    #         'predictive': uncertainty.predictive.mean(),
    #         'confidence': uncertainty.confidence.mean(),
    #     }

    def update(self, logits: Tensor, *args) -> Dict[str, Tensor]:
        """
        使用一个批次的模型输出 logits 来更新度量状态，并返回该批次的度量。

        Args:
            logits (Tensor): 模型在前向传播后输出的原始 logits。

        Returns:
            Dict[str, Tensor]: 一个包含当前批次各项平均不确定性度量的字典，
                               用于 step-level 的日志记录。
        """
        uncertainty: UncertaintyMetrics = self.edl_model.compute_uncertainty(logits)
        batch_size = logits.size(0)

        # 1. 更新内部状态，用于 epoch-level 的 compute()
        if batch_size > 0:
            self.aleatoric_sum += torch.sum(uncertainty.aleatoric)
            self.epistemic_sum += torch.sum(uncertainty.epistemic)
            self.pred_unc_sum += torch.sum(uncertainty.predictive)
            self.confidence_sum += torch.sum(uncertainty.confidence)
            self.total_samples += batch_size

        # 2. 计算并返回当前批次的平均指标，用于 step-level 日志
        #    这里直接用 .mean() 计算当前批次的平均值
        step_metrics = {
            'aleatoric': torch.mean(uncertainty.aleatoric),
            'epistemic': torch.mean(uncertainty.epistemic),
            'predictive': torch.mean(uncertainty.predictive),
            'confidence': torch.mean(uncertainty.confidence),
        }

        return step_metrics  # 理论上 torchmetric.Metric 的 update 是不返回东西的，为了方便返回实时计算的 U

    def compute(self) -> Dict[str, Tensor]:
        """
        计算并返回整个验证数据集上的平均不确定性度量。

        Returns:
            Dict[str, Tensor]: 一个包含整个周期内各项平均不确定性度量的字典，
                               用于 epoch-level 的日志记录。
        """
        total_samples = cast(Tensor, self.total_samples)
        aleatoric_sum = cast(Tensor, self.aleatoric_sum)
        epistemic_sum = cast(Tensor, self.epistemic_sum)
        pred_unc_sum = cast(Tensor, self.pred_unc_sum)
        confidence_sum = cast(Tensor, self.confidence_sum)

        if total_samples.item() == 0:
            # 避免除以零，并确保返回的张量在正确的设备上
            zero = torch.tensor(0.0, device=aleatoric_sum.device)
            return {
                'aleatoric': zero,
                'epistemic': zero,
                'predictive': zero,
                'confidence': zero,
            }

        # 计算整个周期的平均值
        epoch_metrics = {
            'aleatoric': aleatoric_sum / total_samples,
            'epistemic': epistemic_sum / total_samples,
            'predictive': pred_unc_sum / total_samples,
            'confidence': confidence_sum / total_samples,
        }

        return epoch_metrics


def get_high_uncertainty_samples(
    uncertainty: UncertaintyMetrics,
    threshold: float = 0.5,
    uncertainty_type: UncertaintyType = UncertaintyType.TOTAL,
) -> Tensor:
    """获取高不确定性样本的索引"""
    if uncertainty_type == UncertaintyType.ALEATORIC:
        uncertainty_values = uncertainty.aleatoric
    elif uncertainty_type == UncertaintyType.EPISTEMIC:
        uncertainty_values = uncertainty.epistemic
    else:
        uncertainty_values = uncertainty.predictive

    return torch.where(uncertainty_values > threshold)[0]


def get_model_confidence_distribution(
    uncertainty: UncertaintyMetrics, bins: int = 10
) -> Tuple[Tensor, Tensor]:
    """获取模型置信度分布"""
    confidence = uncertainty.confidence
    hist = torch.histc(confidence, bins=bins, min=0.0, max=1.0)
    bin_edges = torch.linspace(0.0, 1.0, bins + 1)
    return hist, bin_edges
