"""
Trusted Multi-view Learning (TML) 框架

基于主观逻辑和证据理论的多视图学习框架，结合了 Evidential Deep Learning 和 Opinion Fusion。

核心特性：
1. 多视图证据学习：每个视图通过 EDL 产生证据，转换为主观逻辑意见
2. 灵活的意见融合：支持累积、平均、加权和 DST 融合算法
3. 信任折扣机制：可选的信任网络调整各视图可信度
4. 冲突感知损失：基于视图间冲突度的正则化
5. 不确定性量化：提供多种不确定性指标

参考文献：
----------
- Trusted Multi-view Learning 相关论文
- Evidential Deep Learning (Sensoy et al., 2018)
- Subjective Logic (Jøsang, 2016)
"""

import typing

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    TypedDict,
    Union,
    cast,
)

import torch
import torch.nn as nn
from torch import Tensor

from zyplib.nn.modules import BaseModule
from zyplib.paradigm import sl_torch
from zyplib.paradigm.edl import (
    EvidentialClassifier,
    EvidentialConfig,
    RegularizationConfig,
    RegularizationMixin,
    UncertaintyMetricsMixin,
)
from zyplib.utils.print import warning

if TYPE_CHECKING:
    from zyplib.train.common import FabricTrainingContext


# ========================= Type Definitions ========================= #


class TMLForwardOutput(TypedDict):
    """TML 前向传播输出类型"""

    fused_alpha: Tensor
    fused_opinion: Tuple[Tensor, Tensor, Tensor]
    fused_evidence: Tensor
    # 可选的中间结果
    view_logits: Optional[Dict[str, Tensor]]
    view_evidences: Optional[Dict[str, Tensor]]
    view_opinions: Optional[Dict[str, Tuple[Tensor, Tensor, Tensor]]]


class TMLLossOutput(TypedDict, total=False):
    """TML 损失输出类型"""

    total_loss: Tensor
    # 可选的分项损失
    view_loss: Optional[Tensor]
    fusion_loss: Optional[Tensor]
    fusion_reg_loss: Optional[Tensor]  # 融合正则化损失
    conflict_loss: Optional[Tensor]
    annealing_coeff: Optional[float]
    view_annealing_coeff: Optional[float]
    # 各视图的具体损失 - 动态键名格式: view_{name}_loss


# ========================= Trust Provider System ========================= #


@dataclass
class ViewTrustRequest:
    """视图信任请求"""

    view_idx: int
    view_name: str
    opinion: Tuple[Tensor, Tensor, Tensor]  # (belief_masses, uncertainty, base_rates)
    view_features: Optional[Tensor] = None
    metadata: Optional[Dict[str, Any]] = None


class TrustProvider(ABC):
    """信任提供者抽象基类"""

    @abstractmethod
    def get_trust(
        self,
        requests: typing.Sequence[ViewTrustRequest],
        global_context: Optional[typing.Mapping[str, Any]] = None,
    ) -> List[Tuple[Tensor, Tensor, Tensor]]:
        """批量获取多个视图的信任意见

        Parameters
        ----------
        - `requests` : `List[ViewTrustRequest]`
            - 视图信任请求列表，包含视图索引、名称、意见等信息
        - `global_context` : `Dict[str, Any]`, optional
            - 全局上下文信息（如当前epoch、批次信息、模型状态等）

        Returns
        -------
        - `List[Tuple[Tensor, Tensor, Tensor]]`
            - 对应各视图的信任意见列表
        """
        pass

    @abstractmethod
    def update(self, performance_metrics: typing.Mapping[Union[int, str], float]) -> None:
        """根据性能指标更新信任度

        Parameters
        ----------
        - `performance_metrics` : `Dict[Union[int, str], float]`
            - 各视图的性能指标，支持索引或名称作为key
        """
        pass


class StaticTrustProvider(TrustProvider):
    """静态信任提供者"""

    def __init__(
        self,
        trust_config: typing.Mapping[str, float],
        num_trust_classes: int = 2,
        base_rates: Optional[Tensor] = None,
    ):
        """
        Parameters
        ----------
        - `trust_config` : `Dict[str, float]`
            - 视图名称到信任值的映射，如 {"temporal": 0.8, "frequency": 0.9}
        - `num_trust_classes` : `int`
            - 信任意见的类别数（通常为2：信任/不信任）
        """
        self.trust_config = trust_config
        self.num_trust_classes = num_trust_classes
        self.base_rates = base_rates

    def get_trust(
        self,
        requests: typing.Sequence[ViewTrustRequest],
        global_context: Optional[typing.Mapping[str, Any]] = None,
    ) -> List[Tuple[Tensor, Tensor, Tensor]]:
        trust_opinions = []

        for req in requests:
            # 根据视图名称获取信任值
            trust_val = self.trust_config.get(req.view_name, 0.5)  # 默认0.5

            batch_size = req.opinion[0].shape[0]
            device = req.opinion[0].device

            # 构造二元信任意见 [trust, distrust]
            trust_belief = torch.zeros(
                (batch_size, self.num_trust_classes), device=device
            )
            trust_belief[:, 0] = trust_val  # trust
            trust_belief[:, 1] = 1.0 - trust_val  # distrust
            trust_uncertainty = torch.zeros((batch_size,), device=device)
            trust_base_rates = torch.full(
                (batch_size, self.num_trust_classes),
                1.0 / self.num_trust_classes,
                device=device,
            )
            if self.base_rates is not None:
                for i in range(self.num_trust_classes):
                    trust_base_rates[:, i] = self.base_rates[i]

            trust_opinions.append((trust_belief, trust_uncertainty, trust_base_rates))

        return trust_opinions

    def update(self, performance_metrics: typing.Mapping[Union[int, str], float]) -> None:
        # 静态信任不需要更新
        pass


# ========================= Configuration ========================= #


@dataclass
class TrustedMultiViewConfig:
    """Trusted Multi-view Learning 配置"""

    # 基础配置
    num_classes: int
    view_names: List[str]  # 视图名称列表
    input_format: Literal['dict', 'list'] = 'dict'  # 输入数据格式

    # Evidence & EDL 配置
    evidence_activation: Literal['relu', 'exp', 'softplus'] = 'relu'
    base_rates: Optional[Tensor] = None  # 如果为None，使用均匀分布
    W: Optional[float] = None  # 如果为None，使用 num_classes

    # 融合配置
    fusion_method: Literal['cumulative', 'averaging', 'weighted', 'dempster_shafer'] = (
        'weighted'
    )

    # 损失配置
    loss_on_views: bool = True  # 是否对各视图计算损失
    loss_on_fusion: bool = True  # 是否对融合结果计算损失
    fusion_loss_type: Literal[
        'expected_risk', 'type_ii_ml', 'bayes_risk_ce', 'bayes_risk_mse'
    ] = 'expected_risk'  # 融合损失类型
    lambda_views: float = 1.0  # 视图损失权重
    lambda_fusion: float = 1.0  # 融合损失权重
    lambda_conflict: float = 0.1  # 冲突正则化权重

    # Trust 配置
    enable_trust_discounting: bool = False
    trust_provider: Optional[TrustProvider] = None

    # EDL 继承配置
    enable_kl_regularization: bool = True  # 是否启用KL正则化
    edl_config: Optional[EvidentialConfig] = None
    reg_config: Optional[RegularizationConfig] = None
    fusion_reg_config: Optional[RegularizationConfig] = (
        None  # 融合层面的正则化配置（对fused_alpha应用KL正则化等）
    )

    @property
    def num_views(self) -> int:
        return len(self.view_names)


# ========================= Main Model ========================= #


class TrustedMultiViewClassifier(BaseModule):
    """Trusted Multi-view 分类器

    基于主观逻辑和证据理论的多视图学习模型。每个视图通过 Evidential Deep Learning
    产生证据，转换为主观逻辑意见，然后通过多种融合算法进行组合。

    主要特性：
    - 多视图证据学习
    - 灵活的意见融合算法
    - 可选的信任折扣机制
    - 冲突感知的损失函数
    - 丰富的不确定性度量
    - 支持融合层面的正则化（对融合后的 alpha 参数应用 KL 正则化等）
    """

    def __init__(
        self,
        backbones: typing.Mapping[str, nn.Module],  # view_name -> backbone
        config: TrustedMultiViewConfig,
        input_shapes: Optional[typing.Mapping[str, Tuple[int, ...]]] = None,
    ):
        """
        Parameters
        ----------
        - `backbones` : `Dict[str, nn.Module]`
            - 各视图的骨干网络，键为视图名称，值为对应的神经网络
        - `config` : `TrustedMultiViewConfig`
            - TML 配置对象
        - `input_shapes` : `Optional[Dict[str, Tuple[int, ...]]]`
            - 各视图的输入形状，用于模型验证
        """
        super().__init__()
        self.config = config

        # 为每个视图创建 EvidentialClassifier
        self.view_classifiers = nn.ModuleDict()
        for view_name, backbone in backbones.items():
            if view_name not in config.view_names:
                warning(f"视图 '{view_name}' 不在配置的 view_names 中，将被忽略")
                continue

            edl_config = config.edl_config or EvidentialConfig(
                task_type='classification',
                num_classes=config.num_classes,
                evidence_activation=config.evidence_activation,
            )

            # 如果没有提供 reg_config，根据 enable_kl_regularization 创建默认配置
            reg_config = config.reg_config
            if reg_config is None and config.enable_kl_regularization:
                reg_config = RegularizationConfig(
                    enable_kl_regularization=config.enable_kl_regularization
                )  # TODO 确认一下 TML 是各个视图各自 KL 还是全局 KL?

            self.view_classifiers[view_name] = EvidentialClassifier(
                backbone=backbone,
                config=edl_config,
                reg_config=reg_config,
                num_classes=config.num_classes,
                input_shape=input_shapes.get(view_name) if input_shapes else None,
            )

        # Trust provider
        self.trust_provider = config.trust_provider

        # 初始化 fusion 正则化器
        if config.fusion_reg_config is not None:
            self.fusion_regularization = RegularizationMixin(config.fusion_reg_config)
        elif config.enable_kl_regularization:
            # 使用和视图相同的正则化配置，避免参数不一致
            fusion_reg_config = config.reg_config or RegularizationConfig(
                enable_kl_regularization=config.enable_kl_regularization
            )
            self.fusion_regularization = RegularizationMixin(fusion_reg_config)

        # 用于 sl_torch 的参数
        self.W = config.W or float(config.num_classes)
        if config.base_rates is None:
            self.register_buffer(
                'base_rates', torch.ones(config.num_classes) / config.num_classes
            )
        else:
            self.register_buffer('base_rates', config.base_rates)

    def forward(
        self,
        inputs: Union[typing.Mapping[str, Tensor], typing.Sequence[Tensor]],
        return_intermediates: bool = False,
    ) -> TMLForwardOutput:
        """前向传播

        Parameters
        ----------
        - `inputs` : `Union[Dict[str, Tensor], List[Tensor]]`
            - 输入数据，可以是字典格式或列表格式
        - `return_intermediates` : `bool`
            - 是否返回中间结果（用于损失计算）

        Returns
        -------
        - `TMLForwardOutput`
            - 包含融合结果和可选中间结果的类型化字典
        """
        # 解析输入格式
        if self.config.input_format == 'dict':
            if not isinstance(inputs, dict):
                raise ValueError(f'配置要求字典输入，但得到 {type(inputs)}')
            view_inputs = inputs
        elif self.config.input_format == 'list':
            if not isinstance(inputs, list):
                raise ValueError(f'配置要求列表输入，但得到 {type(inputs)}')
            if len(inputs) != len(self.config.view_names):
                raise ValueError(
                    f'列表长度 {len(inputs)} 与视图数量 {len(self.config.view_names)} 不匹配'
                )
            view_inputs = {
                name: inputs[i] for i, name in enumerate(self.config.view_names)
            }
        else:
            raise ValueError(f'不支持的输入格式: {self.config.input_format}')

        batch_size = next(iter(view_inputs.values())).shape[0]

        # 1. 各视图推理
        view_logits = {}
        view_evidences = {}
        view_opinions = {}

        for view_name, x_view in view_inputs.items():
            if view_name not in self.view_classifiers:
                continue

            classifier = cast(EvidentialClassifier, self.view_classifiers[view_name])
            logits = classifier(x_view)
            evidence = classifier.logits_to_evidence(logits)

            # 转换为主观逻辑意见
            opinion = sl_torch.evidence_to_opinion(
                evidence,
                cast(Tensor, self.base_rates).unsqueeze(0).expand(batch_size, -1),
                self.W,
            )

            view_logits[view_name] = logits
            view_evidences[view_name] = evidence
            view_opinions[view_name] = opinion

        if not view_opinions:
            raise ValueError('没有有效的视图输出')

        # 2. Trust discounting (可选)
        if self.config.enable_trust_discounting and self.trust_provider:
            view_opinions = self._apply_trust_discounting(view_opinions, view_inputs)

        # 3. Opinion fusion
        fused_opinion = self._fuse_opinions(list(view_opinions.values()))

        # 4. 转换回分布参数
        belief_masses, uncertainty, base_rates = fused_opinion
        fused_evidence = sl_torch.opinion_to_evidence(belief_masses, uncertainty, self.W)
        fused_alpha = sl_torch.evidence_to_dirichlet(
            fused_evidence,
            cast(Tensor, self.base_rates).unsqueeze(0).expand(batch_size, -1),
            self.W,
        )

        results: TMLForwardOutput = {
            'fused_alpha': fused_alpha,
            'fused_opinion': fused_opinion,
            'fused_evidence': fused_evidence,
            'view_logits': None,
            'view_evidences': None,
            'view_opinions': None,
        }

        if return_intermediates:
            results.update(
                {
                    'view_logits': view_logits,
                    'view_evidences': view_evidences,
                    'view_opinions': view_opinions,
                }
            )

        return results

    def _apply_trust_discounting(
        self,
        view_opinions: typing.Mapping[str, Tuple[Tensor, Tensor, Tensor]],
        view_inputs: typing.Mapping[str, Tensor],
    ) -> Dict[str, Tuple[Tensor, Tensor, Tensor]]:
        """应用信任折扣"""
        requests = []
        for view_name, opinion in view_opinions.items():
            req = ViewTrustRequest(
                view_idx=self.config.view_names.index(view_name),
                view_name=view_name,
                opinion=opinion,
                view_features=view_inputs.get(view_name),
            )
            requests.append(req)

        trust_opinions = cast(TrustProvider, self.trust_provider).get_trust(requests)

        # 应用信任折扣
        discounted_opinions = {}
        for (view_name, source_opinion), trust_opinion in zip(
            view_opinions.items(), trust_opinions
        ):
            discounted = sl_torch.trust_discounting(trust_opinion, source_opinion)
            discounted_opinions[view_name] = discounted

        return discounted_opinions

    def _fuse_opinions(
        self, opinions: typing.Sequence[Tuple[Tensor, Tensor, Tensor]]
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """融合意见"""
        if self.config.fusion_method == 'cumulative':
            return sl_torch.cumulative_fusion(opinions, self.W)
        elif self.config.fusion_method == 'averaging':
            return sl_torch.averaging_fusion(opinions, self.W)
        elif self.config.fusion_method == 'weighted':
            return sl_torch.weighted_fusion(opinions, self.W)
        elif self.config.fusion_method == 'dempster_shafer':
            # 转换为 alpha 再融合
            alphas = []
            batch_size = opinions[0][0].shape[0]
            base_rates_expanded = (
                cast(Tensor, self.base_rates).unsqueeze(0).expand(batch_size, -1)
            )

            for belief_masses, uncertainty, base_rates in opinions:
                evidence = sl_torch.opinion_to_evidence(
                    belief_masses, uncertainty, self.W
                )
                alpha = sl_torch.evidence_to_dirichlet(evidence, base_rates, self.W)
                alphas.append(alpha)

            fused_alpha, conflict = sl_torch.dempster_shafer_fusion(
                alphas, base_rates=base_rates_expanded, W=self.W
            )
            return sl_torch.dirichlet_to_opinion(fused_alpha, base_rates_expanded, self.W)
        else:
            raise ValueError(f'不支持的融合方法: {self.config.fusion_method}')

    def compute_loss(
        self, outputs: TMLForwardOutput, target: Tensor, epoch: int | None
    ) -> TMLLossOutput:
        """计算损失

        Parameters
        ----------
        - `outputs` : `TMLForwardOutput`
            - 模型输出，应包含融合结果和可选的中间结果
        - `target` : `Tensor`
            - 目标标签
        - `epoch` : `Optional[int]`
            - 当前epoch，用于退火等

        Returns
        -------
        - `TMLLossOutput`
            - 各项损失的类型化字典
        """
        losses: TMLLossOutput = {
            'total_loss': torch.tensor(0.0, device=target.device),
            'view_loss': None,
            'fusion_loss': None,
            'fusion_reg_loss': None,
            'conflict_loss': None,
        }

        # 1. 视图损失
        if self.config.loss_on_views and outputs['view_logits'] is not None:
            view_losses = []
            for view_name, view_logits in outputs['view_logits'].items():
                classifier = cast(EvidentialClassifier, self.view_classifiers[view_name])
                view_loss_dict = classifier.compute_loss(view_logits, target, epoch=epoch)
                view_total_loss = view_loss_dict['total_loss']
                view_losses.append(view_total_loss)

                # 添加每个视图的具体损失
                losses[f'view_{view_name}_loss'] = view_total_loss

            if view_losses:
                losses['view_loss'] = (
                    self.config.lambda_views * torch.stack(view_losses).mean()
                )

        # 2. 融合损失 - 支持多种损失类型
        if self.config.loss_on_fusion:
            fused_alpha = outputs['fused_alpha']
            fusion_loss = self._compute_fusion_loss(fused_alpha, target)
            losses['fusion_loss'] = self.config.lambda_fusion * fusion_loss

        # 3. 融合正则化损失
        if hasattr(self, 'fusion_regularization'):
            fusion_annealing_coeff = self.fusion_regularization.get_annealing_coeff(epoch)

            fusion_reg_loss = self.fusion_regularization.compute_regularization(
                outputs['fused_alpha'],
                target,
                self.config.num_classes,
                'classification',
                fusion_annealing_coeff,
            )
            losses['fusion_reg_loss'] = fusion_reg_loss
            losses['annealing_coeff'] = (
                fusion_annealing_coeff  # 明确标识这是融合层的退火系数
            )

            # 如果需要，也可以记录视图层的退火系数用于调试
            if self.config.reg_config and hasattr(
                self.config.reg_config, 'enable_annealing'
            ):
                view_annealing_coeff = None
                for _, classifier in self.view_classifiers.items():
                    if hasattr(classifier, 'regularization'):
                        view_annealing_coeff = cast(
                            EvidentialClassifier, classifier
                        ).regularization.get_annealing_coeff(epoch)
                        break  # 所有视图应该使用相同的配置

                if view_annealing_coeff is not None:
                    losses['view_annealing_coeff'] = view_annealing_coeff

        else:
            losses['fusion_reg_loss'] = torch.tensor(0.0, device=target.device)
            losses['annealing_coeff'] = 0

        # 4. 冲突正则化
        if self.config.lambda_conflict > 0 and outputs['view_opinions'] is not None:
            conflict_loss = self._compute_conflict_regularization(
                outputs['view_opinions']
            )
            losses['conflict_loss'] = self.config.lambda_conflict * conflict_loss

        # 5. 总损失
        total_loss = cast(
            Tensor,
            sum(
                (cast(Any, loss) for loss in losses.values() if loss is not None),
                torch.tensor(0.0, device=target.device),
            ),
        )
        losses['total_loss'] = total_loss

        return losses

    def _compute_fusion_loss(self, fused_alpha: Tensor, target: Tensor) -> Tensor:
        """计算融合损失"""
        from zyplib.paradigm.edl import LossFunctionsMixin

        if self.config.fusion_loss_type == 'expected_risk':
            return LossFunctionsMixin.expected_risk_loss(
                fused_alpha, target, self.config.num_classes
            )
        elif self.config.fusion_loss_type == 'type_ii_ml':
            return LossFunctionsMixin.type_ii_ml_loss(
                fused_alpha, target, self.config.num_classes
            )
        elif self.config.fusion_loss_type == 'bayes_risk_ce':
            return LossFunctionsMixin.bayes_risk_ce_loss(
                fused_alpha, target, self.config.num_classes
            )
        elif self.config.fusion_loss_type == 'bayes_risk_mse':
            return LossFunctionsMixin.bayes_risk_mse_loss(
                fused_alpha, target, self.config.num_classes
            )
        else:
            raise ValueError(f'不支持的融合损失类型: {self.config.fusion_loss_type}')

    def _compute_conflict_regularization(
        self, view_opinions: typing.Mapping[str, Tuple[Tensor, Tensor, Tensor]]
    ) -> Tensor:
        """计算冲突正则化"""
        opinions = list(view_opinions.values())
        if len(opinions) < 2:
            return torch.tensor(0.0, device=next(iter(opinions))[0].device)

        total_conflict = 0.0
        num_pairs = 0
        device = opinions[0][0].device

        for i in range(len(opinions)):
            for j in range(i + 1, len(opinions)):
                conflict, _, _ = sl_torch.degree_of_conflict(opinions[i], opinions[j])
                total_conflict += conflict.mean()
                num_pairs += 1

        if num_pairs > 0:
            return cast(Tensor, total_conflict / num_pairs)
        else:
            return torch.tensor(0.0, device=device)

    @torch.no_grad()
    def predict_prob(
        self, inputs: Union[typing.Mapping[str, Tensor], typing.Sequence[Tensor]]
    ) -> Tensor:
        """预测概率分布"""
        self.eval()
        outputs = self.forward(inputs)
        fused_opinion = outputs['fused_opinion']

        prob = sl_torch.projected_probabilities(*fused_opinion)
        return prob

    @torch.no_grad()
    def compute_uncertainty(
        self, inputs: Union[typing.Mapping[str, Tensor], typing.Sequence[Tensor]]
    ) -> Dict[str, Tensor]:
        """计算不确定性指标"""
        self.eval()
        outputs = self.forward(inputs, return_intermediates=True)

        # 融合结果的不确定性
        _, uncertainty, _ = outputs['fused_opinion']
        alphas = outputs['fused_alpha']

        uncertainties = {
            'epistemic': uncertainty,  # 主观逻辑中的不确定性对应EDL的认知不确定性
            'aleatoric': UncertaintyMetricsMixin.aleatoric_uncertainty(alphas),
            # 'confidence': 1.0 - uncertainty,
        }
        # 视图间冲突度
        if outputs['view_opinions'] is not None:
            view_opinions = list(outputs['view_opinions'].values())
            if len(view_opinions) >= 2:
                conflicts = []
                for i in range(len(view_opinions)):
                    for j in range(i + 1, len(view_opinions)):
                        conflict, _, _ = sl_torch.degree_of_conflict(
                            view_opinions[i], view_opinions[j]
                        )
                        conflicts.append(conflict)

                if conflicts:
                    uncertainties['inter_view_conflict'] = torch.stack(conflicts).mean(
                        dim=0
                    )

        return uncertainties


# ========================= Training Procedure ========================= #


class TMLFabricTrainingProcedure:
    """Trusted Multi-view Learning Fabric 训练流程"""

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

    def tml_train_step(
        self,
        ctx: 'FabricTrainingContext',
        batch: Union[
            Tuple[typing.Mapping[str, Tensor], Tensor], Tuple[Tuple[Tensor], Tensor]
        ],
        epoch_idx: int,
        **kwargs,
    ) -> Tuple[Any, Dict[str, float]]:
        """TML训练步骤"""
        tml_model = cast(TrustedMultiViewClassifier, ctx.model)
        optimizer = cast(Any, ctx.optimizer)
        fabric = ctx.fabric

        # 解析批次数据
        if isinstance(batch[0], dict):
            inputs, target = batch
        else:
            inputs, target = batch

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

        outputs = tml_model.forward(inputs, return_intermediates=True)
        loss_dict = tml_model.compute_loss(outputs, target, epoch_idx)

        total_loss = cast(Tensor, cast(Any, loss_dict)['total_loss'])

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

        optimizer.step()

        # 更新信任度 (可选)
        # #TODO 这里后面需要调整这个逻辑; trust discounting 不是在这个地方用的
        if tml_model.trust_provider and epoch_idx is not None:
            if epoch_idx % 5 == 0:  # 每5个epoch更新一次
                self._update_trust_provider(tml_model, outputs, target)

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

        # 返回融合结果的概率用于计算准确率等指标
        fused_prob = sl_torch.projected_probabilities(*outputs['fused_opinion'])

        return fused_prob, cast(Dict[str, float], metrics)

    def tml_val_step(
        self,
        ctx: 'FabricTrainingContext',
        batch: Union[
            Tuple[typing.Mapping[str, Tensor], Tensor], Tuple[Tuple[Tensor], Tensor]
        ],
        epoch_idx: int | None,
        **kwargs,
    ) -> Tuple[Any, Dict[str, float]]:
        """TML验证步骤; 如果 epoch_idx 为 None 代表完全开启 KL Reg (test 阶段)"""
        tml_model = cast(TrustedMultiViewClassifier, ctx.model)

        # 解析批次数据
        if isinstance(batch[0], dict):
            inputs, target = batch
        else:
            inputs, target = batch

        tml_model.eval()
        with torch.no_grad():
            outputs = tml_model.forward(inputs, return_intermediates=True)
            loss_dict = tml_model.compute_loss(outputs, target, epoch_idx)

            # 计算不确定性指标
            uncertainty_dict = tml_model.compute_uncertainty(inputs)

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

        # 添加不确定性指标
        for k, v in uncertainty_dict.items():
            if isinstance(v, Tensor):
                metrics[f'mean_{k}'] = v.mean().item()

        # 返回融合结果的概率
        fused_prob = sl_torch.projected_probabilities(*outputs['fused_opinion'])

        return fused_prob, cast(Dict[str, float], metrics)

    def _update_trust_provider(
        self, model: TrustedMultiViewClassifier, outputs: TMLForwardOutput, target: Tensor
    ):
        """更新信任提供者"""
        if not model.trust_provider or outputs['view_logits'] is None:
            return

        # 计算各视图的准确率
        performance_metrics = {}
        for view_name, view_logits in outputs['view_logits'].items():
            view_prob = torch.softmax(view_logits, dim=-1)
            view_pred = torch.argmax(view_prob, dim=-1)
            acc = (view_pred == target).float().mean().item()
            performance_metrics[view_name] = acc

        model.trust_provider.update(performance_metrics)
