"""
供 torch 使用的主观逻辑计算; batch; 向量化计算

### bijective mapping

evidence <-> opinion[belief, uncertainty, base_rates]
evicence <-> dirichlet[alpha]
opinion <-> dirichlet

### opinion fusion

- CBF
- ABF
- WBF
- DST (Dempster-Shafer Theory)

### trust

- trust_discounting

### regulation

- degree of conflict

"""

import typing

from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union

import torch

if TYPE_CHECKING:
    from torch import Tensor


OpinionTuple = Tuple['Tensor', 'Tensor', 'Tensor']


def numerical_validate(
    belief_masses: 'Tensor',
    uncertainty: 'Tensor',
    base_rates: 'Tensor',
    auto_normalize: bool = True,
    eps: Optional[float] = None,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """
    确保主观逻辑意见参数的数值正确性

    参考 subjective_logic.py 中的 __validate 方法实现，针对 torch 张量进行批量验证和归一化处理。

    Parameters
    ----------
    - `belief_masses` : `Tensor`
        - 信念质量张量，形状为 `[batch, n_states]`
    - `uncertainty` : `Tensor`
        - 不确定性张量，形状为 `[batch]`
    - `base_rates` : `Tensor`
        - 基率张量，形状为 `[batch, n_states]` 或 `[n_states]`
    - `auto_normalize` : `bool`, optional
        - 是否自动归一化处理微小的浮点误差，默认为 True
    - `eps` : `float`, optional
        - 数值精度阈值。如果为 None，则根据数据类型自动设置
        - torch.float32: 1e-7, torch.float64: 1e-12

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - 验证并可能归一化后的 (belief_masses, uncertainty, base_rates)

    Raises
    ------
    ValueError
        当参数不满足约束条件且无法通过归一化修正时
    """
    dtype = belief_masses.dtype
    batch_size, n_states = belief_masses.shape

    # 设置数值精度阈值
    if eps is None:
        eps = 1e-7 if dtype == torch.float32 else 1e-12

    # 1. 基本约束检查
    if torch.any(uncertainty < 0) or torch.any(uncertainty > 1.0):
        raise ValueError('不确定性必须在 [0,1] 范围内')

    if torch.any(belief_masses < 0) or torch.any(belief_masses > 1.0):
        raise ValueError('所有信念质量必须在 [0,1] 范围内')

    if n_states < 2:
        raise ValueError('意见至少需要2个状态')

    # 2. 广播 base_rates 到正确的形状
    if base_rates.dim() == 1:
        base_rates = base_rates.unsqueeze(0).expand(batch_size, -1)

    if base_rates.shape != belief_masses.shape:
        raise ValueError('信念质量和基率维度必须相同')

    # 3. 验证并归一化 belief_masses 和 uncertainty
    total = torch.sum(belief_masses, dim=1) + uncertainty  # [batch]
    total_deviation = torch.abs(total - 1.0)  # [batch]

    # 检查是否有偏差过大的样本
    invalid_mask = total_deviation > eps
    if torch.any(invalid_mask):
        max_deviation = torch.max(total_deviation)
        raise ValueError(
            f'信念质量与不确定性之和不等于1，最大偏差：{max_deviation.item():.2e}'
        )

    # 自动归一化微小的浮点误差
    if auto_normalize:
        normalize_mask = total_deviation > 0
        if torch.any(normalize_mask):
            # 只对需要归一化的样本进行处理
            total_expanded = total.unsqueeze(1)  # [batch, 1]
            belief_masses = torch.where(
                normalize_mask.unsqueeze(1), belief_masses / total_expanded, belief_masses
            )
            uncertainty = torch.where(normalize_mask, uncertainty / total, uncertainty)

    # 4. 验证并归一化 base_rates
    base_rates_sum = torch.sum(base_rates, dim=1)  # [batch]
    base_rates_deviation = torch.abs(base_rates_sum - 1.0)  # [batch]

    # 检查是否有偏差过大的样本
    invalid_base_rates_mask = base_rates_deviation > eps
    if torch.any(invalid_base_rates_mask):
        max_base_rates_deviation = torch.max(base_rates_deviation)
        raise ValueError(
            f'基率之和不等于1，最大偏差：{max_base_rates_deviation.item():.2e}'
        )

    # 自动归一化基率的微小浮点误差
    if auto_normalize:
        normalize_base_rates_mask = base_rates_deviation > 0
        if torch.any(normalize_base_rates_mask):
            base_rates_sum_expanded = base_rates_sum.unsqueeze(1)  # [batch, 1]
            base_rates = torch.where(
                normalize_base_rates_mask.unsqueeze(1),
                base_rates / base_rates_sum_expanded,
                base_rates,
            )

    # 5. 确保数值非负
    belief_masses = torch.clamp(belief_masses, min=0.0)
    uncertainty = torch.clamp(uncertainty, min=0.0)
    base_rates = torch.clamp(base_rates, min=0.0)

    return belief_masses, uncertainty, base_rates


def evidence_to_dirichlet(
    evidence: 'Tensor', base_rates: Optional['Tensor'] = None, W: Optional[float] = None
) -> 'Tensor':
    """Evidence 向量转换为狄利克雷参数 alpha

    转换公式: α_i = e_i + W * a_i

    Parameters
    ----------
    - `evidence` : `Tensor`
        - 证据张量，形状为 `[batch, n_states]`
    - `base_rates` : `Tensor`, optional
        - 基率张量，形状为 `[n_states]` 或 `[batch, n_states]`，默认为均匀分布
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tensor`
        - Dirichlet 参数 alpha，形状为 `[batch, n_states]`
    """
    batch_size, n_states = evidence.shape
    device, dtype = evidence.device, evidence.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    if base_rates is None:
        base_rates = torch.ones(n_states, device=device, dtype=dtype) / n_states

    # 广播处理
    if base_rates.dim() == 1:
        base_rates = base_rates.unsqueeze(0).expand(batch_size, -1)

    return evidence + W * base_rates


def dirichlet_to_evidence(
    alpha: 'Tensor', base_rates: Optional['Tensor'] = None, W: Optional[float] = None
) -> 'Tensor':
    """狄利克雷参数转换为 Evidence 向量

    转换公式: e_i = α_i - W * a_i

    Parameters
    ----------
    - `alpha` : `Tensor`
        - Dirichlet 参数，形状为 `[batch, n_states]`
    - `base_rates` : `Tensor`, optional
        - 基率张量，形状为 `[n_states]` 或 `[batch, n_states]`，默认为均匀分布
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tensor`
        - 证据张量，形状为 `[batch, n_states]`
    """
    batch_size, n_states = alpha.shape
    device, dtype = alpha.device, alpha.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    if base_rates is None:
        base_rates = torch.ones(n_states, device=device, dtype=dtype) / n_states

    # 广播处理
    if base_rates.dim() == 1:
        base_rates = base_rates.unsqueeze(0).expand(batch_size, -1)

    evidence = alpha - W * base_rates
    return torch.clamp(evidence, min=0.0)  # 确保非负


def evidence_to_opinion(
    evidence: 'Tensor', base_rates: Optional['Tensor'] = None, W: Optional[float] = None
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """Evidence 向量转换为 SL 意见

    转换公式:
    - belief_masses = evidence / (sum(evidence) + W)
    - uncertainty = W / (sum(evidence) + W)

    Parameters
    ----------
    - `evidence` : `Tensor`
        - 证据张量，形状为 `[batch, n_states]`
    - `base_rates` : `Tensor`, optional
        - 基率张量，形状为 `[n_states]` 或 `[batch, n_states]`，默认为均匀分布
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (belief_masses, uncertainty, base_rates) - 形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    """
    batch_size, n_states = evidence.shape
    device, dtype = evidence.device, evidence.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    if base_rates is None:
        base_rates = torch.ones(n_states, device=device, dtype=dtype) / n_states

    # 广播处理
    if base_rates.dim() == 1:
        base_rates = base_rates.unsqueeze(0).expand(batch_size, -1)

    total_evidence = torch.sum(evidence, dim=1)  # [batch]
    kappa = total_evidence + W  # [batch]

    belief_masses = evidence / kappa.unsqueeze(1)  # [batch, n_states]
    uncertainty = W / kappa  # u = W / (Σe + W)

    return belief_masses, uncertainty, base_rates


def opinion_to_evidence(
    belief_masses: 'Tensor', uncertainty: 'Tensor', W: Optional[float] = None
) -> 'Tensor':
    """SL 意见转换为 Evidence 向量

    转换公式: e_i = W * b_i / u

    Parameters
    ----------
    - `belief_masses` : `Tensor`
        - 信念质量张量，形状为 `[batch, n_states]`
    - `uncertainty` : `Tensor`
        - 不确定性张量，形状为 `[batch]`
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tensor`
        - 证据张量，形状为 `[batch, n_states]`
    """
    batch_size, n_states = belief_masses.shape

    # 默认参数处理
    if W is None:
        W = float(n_states)

    # 数值稳定性处理
    eps = 1e-12 if belief_masses.dtype == torch.float64 else 1e-7
    denominator = torch.clamp(uncertainty, min=eps)  # [batch]

    return (W * belief_masses) / denominator.unsqueeze(1)


def dirichlet_to_opinion(
    alpha: 'Tensor', base_rates: Optional['Tensor'] = None, W: Optional[float] = None
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """狄利克雷参数转换为 SL 意见

    Parameters
    ----------
    - `alpha` : `Tensor`
        - Dirichlet 参数，形状为 `[batch, n_states]`
    - `base_rates` : `Tensor`, optional
        - 基率张量，形状为 `[n_states]` 或 `[batch, n_states]`，默认为均匀分布
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (belief_masses, uncertainty, base_rates) - 形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    """
    # alpha -> evidence -> opinion
    evidence = dirichlet_to_evidence(alpha, base_rates, W)
    return evidence_to_opinion(evidence, base_rates, W)


def opinion_to_dirichlet(
    belief_masses: 'Tensor',
    uncertainty: 'Tensor',
    base_rates: 'Tensor',
    W: Optional[float] = None,
) -> 'Tensor':
    """SL 意见转换为狄利克雷参数

    Parameters
    ----------
    - `belief_masses` : `Tensor`
        - 信念质量张量，形状为 `[batch, n_states]`
    - `uncertainty` : `Tensor`
        - 不确定性张量，形状为 `[batch]`
    - `base_rates` : `Tensor`
        - 基率张量，形状为 `[batch, n_states]`
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tensor`
        - Dirichlet 参数 alpha，形状为 `[batch, n_states]`
    """
    # opinion -> evidence -> alpha
    evidence = opinion_to_evidence(belief_masses, uncertainty, W)
    return evidence_to_dirichlet(evidence, base_rates, W)


def evidence_strength(
    input_data: Union['Tensor', OpinionTuple],
    input_type: Literal['evidence', 'opinion', 'alpha'],
    W: Optional[float] = None,
) -> 'Tensor':
    """计算证据强度 (Evidence Strength)，也称为狄利克雷强度 (Dirichlet Strength) S 或 κ.

    证据强度是支撑一个主观意见的所有证据的总和，反映了该意见的确定性程度。
    强度越高，意见越确定（即不确定性越低）。
    其计算公式为: S = Σα = Σe + W。

    Parameters
    ----------
    - `input_data` : `Union[Tensor, Tuple[Tensor, Tensor, Tensor]]`
        - 输入的数据，其具体类型由 `input_type` 参数指定。
    - `input_type` : `Literal['evidence', 'opinion', 'alpha']`
        - 指定 `input_data` 的类型。
        - 'alpha': 输入为 Dirichlet 参数张量，形状 `[batch, n_states]`。
        - 'evidence': 输入为证据张量，形状 `[batch, n_states]`。
        - 'opinion': 输入为 SL 意见元组 `(belief, uncertainty, base_rates)`。
    - `W` : `float`, optional
        - 非信息先验权重。如果 `input_type` 是 'evidence' 或 'opinion'，则会用到此参数。
        - 如果不提供，将根据输入的维度自动推断为 n_states。

    Returns
    -------
    - `Tensor`
        - 证据强度，形状为 `[batch]`

    Raises
    ------
    - `ValueError`
        - 当 `input_type` 不是支持的类型时。
    """
    if input_type == 'alpha':
        # 从 alpha 计算: S = Σα
        alpha = typing.cast('Tensor', input_data)
        return torch.sum(alpha, dim=-1)

    elif input_type == 'evidence':
        # 从 evidence 计算: S = Σe + W
        evidence = typing.cast('Tensor', input_data)
        _, n_states = evidence.shape
        if W is None:
            W = float(n_states)
        total_evidence = torch.sum(evidence, dim=1)
        return total_evidence + W

    elif input_type == 'opinion':
        # 从 opinion 计算: S = W / u
        belief_masses, uncertainty, _ = typing.cast(OpinionTuple, input_data)
        _, n_states = belief_masses.shape
        if W is None:
            W = float(n_states)

        # 数值稳定性处理，防止不确定性为0
        eps = 1e-12 if uncertainty.dtype == torch.float64 else 1e-7
        denominator = torch.clamp(uncertainty, min=eps)
        return W / denominator

    else:
        raise ValueError(
            f"不支持的 input_type: '{input_type}'。 "
            "有效选项为 'evidence', 'opinion', 'alpha'。"
        )


def projected_probabilities(
    belief_masses: 'Tensor', uncertainty: 'Tensor', base_rates: 'Tensor'
) -> 'Tensor':
    """计算投影概率分布

    转换公式: P(x_i) = b_i + a_i * u

    Parameters
    ----------
    - `belief_masses` : `Tensor`
        - 信念质量张量，形状为 `[batch, n_states]`
    - `uncertainty` : `Tensor`
        - 不确定性张量，形状为 `[batch]`
    - `base_rates` : `Tensor`
        - 基率张量，形状为 `[batch, n_states]`

    Returns
    -------
    - `Tensor`
        - 投影概率分布，形状为 `[batch, n_states]`
    """
    return belief_masses + base_rates * uncertainty.unsqueeze(1)


def projected_probability(
    belief_masses: 'Tensor',
    uncertainty: 'Tensor',
    base_rates: 'Tensor',
    belief_idx: int = 0,
) -> 'Tensor':
    """计算二元意见的投影概率（默认以 idx=0 为 belif 状态, idx=1 为 disbelief 状态）

    仅适用于二元意见 (n_states=2)

    Parameters
    ----------
    - `belief_masses` : `Tensor`
        - 信念质量张量，形状为 `[batch, 2]`
    - `uncertainty` : `Tensor`
        - 不确定性张量，形状为 `[batch]`
    - `base_rates` : `Tensor`
        - 基率张量，形状为 `[batch, 2]`
    - `belief_idx` : `int`
        - 意见状态的索引，默认为 0 (以第一个状态为 belief 状态)

    Returns
    -------
    - `Tensor`
        - 第一个状态的投影概率，形状为 `[batch]`
    """
    if belief_masses.shape[1] != 2:
        raise ValueError('projected_probability 仅适用于二元意见')

    probs = projected_probabilities(belief_masses, uncertainty, base_rates)
    return probs[:, belief_idx]


# ========================= Opinion Fusion ========================= #


def cumulative_fusion(
    sources: typing.Sequence[Tuple['Tensor', 'Tensor', 'Tensor']],
    W: Optional[float] = None,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """累积融合 (Cumulative Belief Fusion, CBF)：组合多个独立信息源的证据

    Parameters
    ----------
    - `sources` : `list[Tuple[Tensor, Tensor, Tensor]]`
        - 要融合的意见来源，每个意见为 (belief_masses, uncertainty, base_rates)
        - 每个张量的形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (belief_masses, uncertainty, base_rates) - 融合后的意见

    Notes
    -----
    累积融合算子（⊕）在数学上对应于贝叶斯定理的更新，而在证据空间中，
    贝叶斯更新就表现为证据的简单加法。即 evidence(A ⊕ B) = evidence(A) + evidence(B)。
    """
    if not sources:
        raise ValueError('至少需要一个意见')

    if len(sources) == 1:
        return sources[0]

    # 获取基本信息
    belief_masses_0, uncertainty_0, base_rates_0 = sources[0]
    batch_size, n_states = belief_masses_0.shape
    device, dtype = belief_masses_0.device, belief_masses_0.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    # 检查维度一致性
    for i, (bm, u, br) in enumerate(sources):
        if bm.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的信念质量维度不一致')
        if u.shape != (batch_size,):
            raise ValueError(f'意见 {i} 的不确定性维度不一致')
        if br.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的基率维度不一致')

    # 检查基率一致性并计算平均基率
    base_rates_list = [br for _, _, br in sources]

    # 简单使用平均基率（可以考虑更复杂的一致性检查）
    avg_base_rates = torch.stack(base_rates_list, dim=0).mean(dim=0)  # [batch, n_states]

    # 累积证据
    total_evidence = torch.zeros(batch_size, n_states, device=device, dtype=dtype)
    for belief_masses, uncertainty, _ in sources:
        evidence = opinion_to_evidence(belief_masses, uncertainty, W)
        total_evidence += evidence

    # 转换回意见
    return evidence_to_opinion(total_evidence, avg_base_rates, W)


def averaging_fusion(
    sources: typing.Sequence[Tuple['Tensor', 'Tensor', 'Tensor']],
    W: Optional[float] = None,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """平均融合 (Averaging Belief Fusion, ABF)：组合多个依赖信息源的意见

    Parameters
    ----------
    - `opinions` : `list[Tuple[Tensor, Tensor, Tensor]]`
        - 要融合的意见列表，每个意见为 (belief_masses, uncertainty, base_rates)
        - 每个张量的形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (belief_masses, uncertainty, base_rates) - 融合后的意见
    """
    if not sources:
        raise ValueError('至少需要一个意见')

    if len(sources) == 1:
        return sources[0]

    n = len(sources)

    # 获取基本信息
    belief_masses_0, uncertainty_0, base_rates_0 = sources[0]
    batch_size, n_states = belief_masses_0.shape
    device, dtype = belief_masses_0.device, belief_masses_0.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    # 检查维度一致性
    for i, (bm, u, br) in enumerate(sources):
        if bm.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的信念质量维度不一致')
        if u.shape != (batch_size,):
            raise ValueError(f'意见 {i} 的不确定性维度不一致')
        if br.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的基率维度不一致')

    # 计算平均基率
    base_rates_list = [br for _, _, br in sources]
    avg_base_rates = torch.stack(base_rates_list, dim=0).mean(dim=0)  # [batch, n_states]

    # 平均证据
    total_evidence = torch.zeros(batch_size, n_states, device=device, dtype=dtype)
    for belief_masses, uncertainty, _ in sources:
        evidence = opinion_to_evidence(belief_masses, uncertainty, W)
        total_evidence += evidence

    avg_evidence = total_evidence / n

    # 转换回意见
    return evidence_to_opinion(avg_evidence, avg_base_rates, W)


def weighted_fusion(
    sources: typing.Sequence[Tuple['Tensor', 'Tensor', 'Tensor']],
    W: Optional[float] = None,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """加权融合 (Weighted Belief Fusion, WBF)：基于置信度组合意见

    Parameters
    ----------
    - `opinions` : `list[Tuple[Tensor, Tensor, Tensor]]`
        - 要融合的意见列表，每个意见为 (belief_masses, uncertainty, base_rates)
        - 每个张量的形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (belief_masses, uncertainty, base_rates) - 融合后的意见
    """
    if not sources:
        raise ValueError('至少需要一个意见')

    if len(sources) == 1:
        return sources[0]

    # 获取基本信息
    belief_masses_0, uncertainty_0, base_rates_0 = sources[0]
    batch_size, n_states = belief_masses_0.shape
    device, dtype = belief_masses_0.device, belief_masses_0.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    # 数值稳定性
    eps = 1e-12 if dtype == torch.float64 else 1e-7

    # 检查维度一致性
    for i, (bm, u, br) in enumerate(sources):
        if bm.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的信念质量维度不一致')
        if u.shape != (batch_size,):
            raise ValueError(f'意见 {i} 的不确定性维度不一致')
        if br.shape != (batch_size, n_states):
            raise ValueError(f'意见 {i} 的基率维度不一致')

    # 计算置信度加权和
    weighted_sum_evidence = torch.zeros(batch_size, n_states, device=device, dtype=dtype)
    weighted_sum_base_rates = torch.zeros(
        batch_size, n_states, device=device, dtype=dtype
    )
    total_confidence = torch.zeros(batch_size, device=device, dtype=dtype)

    for belief_masses, uncertainty, base_rates in sources:
        confidence = 1.0 - uncertainty  # [batch]

        # 只有在置信度大于零时，意见才有贡献
        valid_mask = confidence > eps
        if valid_mask.any():
            evidence = opinion_to_evidence(belief_masses, uncertainty, W)

            # 加权累积
            confidence_expanded = confidence.unsqueeze(1)  # [batch, 1]
            weighted_sum_evidence += evidence * confidence_expanded
            weighted_sum_base_rates += base_rates * confidence_expanded
            total_confidence += confidence

    # 处理所有意见都完全不确定的情况
    invalid_mask = total_confidence <= eps
    if invalid_mask.any():
        # 对于完全不确定的批次，使用平均基率
        base_rates_list = [br for _, _, br in sources]
        avg_base_rates = torch.stack(base_rates_list, dim=0).mean(
            dim=0
        )  # [batch, n_states]

        # 创建完全不确定的意见
        fused_belief_masses = torch.zeros(
            batch_size, n_states, device=device, dtype=dtype
        )
        fused_uncertainty = torch.ones(batch_size, device=device, dtype=dtype)
        fused_base_rates = avg_base_rates.expand(batch_size, -1)

        # 对于有效批次，计算加权平均
        if (~invalid_mask).any():
            valid_evidence = weighted_sum_evidence[~invalid_mask] / total_confidence[
                ~invalid_mask
            ].unsqueeze(1)
            valid_base_rates = weighted_sum_base_rates[~invalid_mask] / total_confidence[
                ~invalid_mask
            ].unsqueeze(1)

            valid_belief, valid_uncertainty, _ = evidence_to_opinion(
                valid_evidence, valid_base_rates, W
            )

            fused_belief_masses[~invalid_mask] = valid_belief
            fused_uncertainty[~invalid_mask] = valid_uncertainty
            fused_base_rates[~invalid_mask] = valid_base_rates

        return fused_belief_masses, fused_uncertainty, fused_base_rates
    else:
        # 所有批次都有效
        fused_evidence = weighted_sum_evidence / total_confidence.unsqueeze(1)
        fused_base_rates = weighted_sum_base_rates / total_confidence.unsqueeze(1)

        return evidence_to_opinion(fused_evidence, fused_base_rates, W)


# ========================= Trust Network ========================= #


def trust_discounting(
    trust_opinion: Tuple['Tensor', 'Tensor', 'Tensor'],
    source_opinion: Tuple['Tensor', 'Tensor', 'Tensor'],
    trust_state_idx: int = 0,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """信任折扣：根据信任意见调整源意见

    该操作符用于根据对信息源的信任度来调整（折扣）其提供的意见。
    例如，在信任链 A -> B -> X 中，A 对 B 的信任意见 (trust_opinion) 会折扣
    B 对 X 的原始意见 (source_opinion)，从而得出 A 对 X 的派生意见。

    Parameters
    ----------
    - `trust_opinion` : `Tuple[Tensor, Tensor, Tensor]`
        - 信任意见 (belief_masses, uncertainty, base_rates)，必须是二元意见
        - 形状分别为 `[batch, 2]`, `[batch]`, `[batch, 2]`
    - `source_opinion` : `Tuple[Tensor, Tensor, Tensor]`
        - 源意见 (belief_masses, uncertainty, base_rates)，可以是多元意见
        - 形状分别为 `[batch, n_states]`, `[batch]`, `[batch, n_states]`
    - `trust_state_idx` : `int`
        - 信任意见状态的索引，默认为 0 (以第一个状态为 belief 状态)

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - 折扣后的意见 (belief_masses, uncertainty, base_rates)

    Raises
    ------
    ValueError
        当信任意见不是二元意见时
    """
    trust_belief_masses, trust_uncertainty, trust_base_rates = trust_opinion
    source_belief_masses, source_uncertainty, source_base_rates = source_opinion

    # 检查信任意见是否为二元意见
    if trust_belief_masses.shape[1] != 2:
        raise ValueError('信任意见 (trust_opinion) 必须是二元意见 (n_states=2)')

    # 检查批次维度一致性
    batch_size = trust_belief_masses.shape[0]
    if source_belief_masses.shape[0] != batch_size:
        raise ValueError('信任意见和源意见的批次维度必须一致')

    # 计算信任意见的投影概率作为折扣系数
    trust_projected_prob = projected_probability(
        trust_belief_masses,
        trust_uncertainty,
        trust_base_rates,
        belief_idx=trust_state_idx,
    )

    # 对源意见的信念质量进行折扣
    discount_coeff = trust_projected_prob.unsqueeze(1)  # [batch, 1]
    discounted_beliefs = discount_coeff * source_belief_masses  # [batch, n_states]

    # 计算新的不确定性
    discounted_uncertainty = 1.0 - torch.sum(discounted_beliefs, dim=1)  # [batch]

    # 基率继承自源意见
    derived_base_rates = source_base_rates

    return discounted_beliefs, discounted_uncertainty, derived_base_rates


def transitive_trust_discounting(
    trust_path: typing.Sequence[Tuple['Tensor', 'Tensor', 'Tensor']],
    source_opinion: Tuple['Tensor', 'Tensor', 'Tensor'],
    trust_state_idx: int = 0,
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """多级信任传递折扣

    处理一个长的信任链，例如 A1 -> A2 -> ... -> An -> X。
    最终的折扣系数是路径上所有信任意见投影概率的乘积。

    Parameters
    ----------
    - `trust_path` : `list[Tuple[Tensor, Tensor, Tensor]]`
        - 信任路径上的意见列表，每个意见都必须是二元意见
        - 例如 [ω_A2^A1, ω_A3^A2, ..., ω_An^A(n-1)]
    - `source_opinion` : `Tuple[Tensor, Tensor, Tensor]`
        - 路径末端的源意见 (belief_masses, uncertainty, base_rates)

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - 经过整个信任链折扣后的最终派生意见

    Raises
    ------
    ValueError
        当信任路径为空，或路径中有任何一个意见不是二元意见时
    """
    if not trust_path:
        raise ValueError('信任路径不能为空')

    source_belief_masses, source_uncertainty, source_base_rates = source_opinion
    batch_size = source_belief_masses.shape[0]
    device, dtype = source_belief_masses.device, source_belief_masses.dtype

    # 计算整个信任路径的总体折扣系数
    overall_discount_coefficient = torch.ones(batch_size, device=device, dtype=dtype)

    for i, trust_op in enumerate(trust_path):
        trust_belief_masses, trust_uncertainty, trust_base_rates = trust_op

        # 检查是否为二元意见
        if trust_belief_masses.shape[1] != 2:
            raise ValueError(f'信任路径中的第 {i} 个意见必须是二元意见')

        # 检查批次维度
        if trust_belief_masses.shape[0] != batch_size:
            raise ValueError(f'信任路径中的第 {i} 个意见的批次维度不一致')

        # 计算投影概率并累积折扣系数
        trust_prob = projected_probability(
            trust_belief_masses,
            trust_uncertainty,
            trust_base_rates,
            belief_idx=trust_state_idx,
        )
        overall_discount_coefficient *= trust_prob

    # 对源意见进行折扣
    discount_coeff = overall_discount_coefficient.unsqueeze(1)  # [batch, 1]
    discounted_beliefs = discount_coeff * source_belief_masses

    # 计算新的不确定性
    discounted_uncertainty = 1.0 - torch.sum(discounted_beliefs, dim=1)

    # 基率继承自源意见
    derived_base_rates = source_base_rates

    return discounted_beliefs, discounted_uncertainty, derived_base_rates


# ========================= Conflict Measures ========================= #


def degree_of_conflict(
    opinion1: Tuple['Tensor', 'Tensor', 'Tensor'],
    opinion2: Tuple['Tensor', 'Tensor', 'Tensor'],
) -> Tuple['Tensor', 'Tensor', 'Tensor']:
    """计算两个主观意见之间的冲突度

    冲突度是投影距离 (PD) 和联合确定性 (CC) 的乘积。

    Parameters
    ----------
    - `opinion1` : `Tuple[Tensor, Tensor, Tensor]`
        - 第一个意见 (belief_masses, uncertainty, base_rates)
    - `opinion2` : `Tuple[Tensor, Tensor, Tensor]`
        - 第二个意见 (belief_masses, uncertainty, base_rates)

    Returns
    -------
    - `Tuple[Tensor, Tensor, Tensor]`
        - (conflict, projected_distance, conjunctive_certainty) - 冲突度、投影距离、联合确定性

    Raises
    ------
    ValueError
        当两个意见的状态数不一致时

    Notes
    -----
    - **投影距离 (PD)**: 衡量两个意见的期望概率分布之间的差异，
      其值在 [0, 1] 范围内。PD = 0 表示期望概率完全一致。
      PD = sum(|P1_k - P2_k|) / 2

    - **联合确定性 (CC)**: 衡量两个意见共同的确定性程度，
      其值在 [0, 1] 范围内。CC = 0 表示至少有一个意见是完全不确定的。
      CC = (1 - u1) * (1 - u2)

    - **冲突度 (DC)**: 结合了距离和确定性。只有当两个意见都比较确定
      (CC 较高) 且期望概率相差很大 (PD 较高) 时，冲突度才会高。
      DC = PD * CC
    """
    belief_masses1, uncertainty1, base_rates1 = opinion1
    belief_masses2, uncertainty2, base_rates2 = opinion2

    # 验证输入维度一致性
    if belief_masses1.shape != belief_masses2.shape:
        raise ValueError('两个意见的信念质量维度必须一致')
    if uncertainty1.shape != uncertainty2.shape:
        raise ValueError('两个意见的不确定性维度必须一致')
    if base_rates1.shape != base_rates2.shape:
        raise ValueError('两个意见的基率维度必须一致')

    # 计算投影概率
    p1 = projected_probabilities(
        belief_masses1, uncertainty1, base_rates1
    )  # [batch, n_states]
    p2 = projected_probabilities(
        belief_masses2, uncertainty2, base_rates2
    )  # [batch, n_states]

    # 投影距离 (Projected Distance, PD)
    # PD(ω¹, ω²) = Σ|P¹_k - P²_k| / 2
    projected_distance = 0.5 * torch.sum(torch.abs(p1 - p2), dim=1)  # [batch]

    # 联合确定性 (Conjunctive Certainty, CC)
    # CC(ω¹, ω²) = (1 - u¹) * (1 - u²)
    conjunctive_certainty = (1.0 - uncertainty1) * (1.0 - uncertainty2)  # [batch]

    # 冲突度 (Degree of Conflict, DC)
    # DC(ω¹, ω²) = PD * CC
    conflict = projected_distance * conjunctive_certainty  # [batch]

    return conflict, projected_distance, conjunctive_certainty


def opinion_distance(
    opinion1: Tuple['Tensor', 'Tensor', 'Tensor'],
    opinion2: Tuple['Tensor', 'Tensor', 'Tensor'],
    metric: str = 'euclidean',
) -> 'Tensor':
    """计算两个意见之间的距离

    注意：这并非 Subjective Logic 中有理论定义的计算方式。
    如果要衡量两个意见之间的冲突程度，推荐使用 `degree_of_conflict`。

    Parameters
    ----------
    - `opinion1` : `Tuple[Tensor, Tensor, Tensor]`
        - 第一个意见 (belief_masses, uncertainty, base_rates)
    - `opinion2` : `Tuple[Tensor, Tensor, Tensor]`
        - 第二个意见 (belief_masses, uncertainty, base_rates)
    - `metric` : `str`
        - 距离度量方式，支持 'euclidean', 'manhattan', 'cosine'

    Returns
    -------
    - `Tensor`
        - 距离值，形状为 `[batch]`

    Raises
    ------
    ValueError
        当两个意见的状态数不一致，或不支持的距离度量时
    """
    belief_masses1, uncertainty1, _ = opinion1
    belief_masses2, uncertainty2, _ = opinion2

    # 验证维度一致性
    if belief_masses1.shape != belief_masses2.shape:
        raise ValueError('两个意见的状态数必须一致')
    if uncertainty1.shape != uncertainty2.shape:
        raise ValueError('两个意见的批次维度必须一致')

    # 构建完整向量 [beliefs, uncertainty]
    vec1 = torch.cat(
        [belief_masses1, uncertainty1.unsqueeze(1)], dim=1
    )  # [batch, n_states+1]
    vec2 = torch.cat(
        [belief_masses2, uncertainty2.unsqueeze(1)], dim=1
    )  # [batch, n_states+1]

    if metric == 'euclidean':
        return torch.sqrt(torch.sum((vec1 - vec2) ** 2, dim=1))
    elif metric == 'manhattan':
        return torch.sum(torch.abs(vec1 - vec2), dim=1)
    elif metric == 'cosine':
        # 余弦距离 = 1 - 余弦相似度
        dot_product = torch.sum(vec1 * vec2, dim=1)
        norm1 = torch.sqrt(torch.sum(vec1**2, dim=1))
        norm2 = torch.sqrt(torch.sum(vec2**2, dim=1))
        cosine_similarity = dot_product / (norm1 * norm2 + 1e-8)  # 数值稳定性
        return 1.0 - cosine_similarity
    else:
        raise ValueError(
            f'不支持的距离度量: {metric}。支持的选项: euclidean, manhattan, cosine'
        )


# ========================= DST Evidence Fusion ========================= #


def dempster_shafer_fusion(
    alpha_sources: typing.Sequence['Tensor'],
    conflict_handling: Literal['dempster', 'yager', 'zhang'] = 'dempster',
    base_rates: Optional['Tensor'] = None,
    W: Optional[float] = None,
) -> Tuple['Tensor', 'Tensor']:
    """使用 DST 证据理论融合多个狄利克雷参数向量

    基于 Dempster-Shafer 证据理论的组合规则来融合多个独立的证据源。
    该方法将 Dirichlet 参数转换为基本信念指派 (Basic Belief Assignment, BBA)，
    应用 DST 组合规则处理冲突，然后转换回 Dirichlet 参数。

    Parameters
    ----------
    - `alpha_sources` : `list[Tensor]`
        - 要融合的 Dirichlet 参数列表，每个张量形状为 `[batch, n_states]`
    - `conflict_handling` : `Literal['dempster', 'yager', 'zhang']`
        - 冲突处理策略：
        - 'dempster': 经典 Dempster 规则，归一化处理冲突
        - 'yager': Yager 规则，将冲突分配给无知状态
        - 'zhang': Zhang 规则，基于信任度重分配冲突
    - `base_rates` : `Tensor`, optional
        - 基率张量，形状为 `[n_states]` 或 `[batch, n_states]`，默认为均匀分布
    - `W` : `float`, optional
        - 非信息先验权重，默认为 n_states

    Returns
    -------
    - `Tuple[Tensor, Tensor]`
        - (fused_alpha, conflict_measure) - 融合后的 Dirichlet 参数和冲突度测量

    Raises
    ------
    ValueError
        当输入参数维度不一致或冲突处理策略不支持时

    Notes
    -----
    DST 融合过程：
    1. α → 证据 → 基本信念指派 (BBA)
    2. 应用 Dempster 组合规则：m(A) = Σ{m1(B)m2(C) : B∩C=A} / (1-K)
    3. 计算冲突度：K = Σ{m1(B)m2(C) : B∩C=∅}
    4. 根据策略处理冲突并归一化
    5. BBA → 证据 → α
    """
    if not alpha_sources:
        raise ValueError('至少需要一个 alpha 参数')

    if len(alpha_sources) == 1:
        return alpha_sources[0], torch.zeros(alpha_sources[0].shape[0])

    # 获取基本信息
    alpha_0 = alpha_sources[0]
    batch_size, n_states = alpha_0.shape
    device, dtype = alpha_0.device, alpha_0.dtype

    # 默认参数处理
    if W is None:
        W = float(n_states)

    if base_rates is None:
        base_rates = torch.ones(n_states, device=device, dtype=dtype) / n_states

    # 广播处理
    if base_rates.dim() == 1:
        base_rates = base_rates.unsqueeze(0).expand(batch_size, -1)

    # 检查维度一致性
    for i, alpha in enumerate(alpha_sources):
        if alpha.shape != (batch_size, n_states):
            raise ValueError(
                f'Alpha {i} 的维度不一致，期望 {(batch_size, n_states)}，得到 {alpha.shape}'
            )

    # 将 Dirichlet 参数转换为基本信念指派 (BBA)
    bba_sources = []
    for alpha in alpha_sources:
        # α → evidence → opinion → normalized BBA
        evidence = dirichlet_to_evidence(alpha, base_rates, W)
        belief_masses, uncertainty, _ = evidence_to_opinion(evidence, base_rates, W)

        # 构建 BBA：包含各状态的信念 + 不确定性 (无知状态)
        bba = torch.cat(
            [belief_masses, uncertainty.unsqueeze(1)], dim=1
        )  # [batch, n_states+1]
        bba_sources.append(bba)

    # 执行 DST 融合
    fused_bba, conflict = _dempster_combination_rule(
        bba_sources, conflict_handling, device, dtype
    )

    # 从融合后的 BBA 转换回 Dirichlet 参数
    # 分离信念质量和不确定性
    fused_belief_masses = fused_bba[:, :n_states]  # [batch, n_states]
    fused_uncertainty = fused_bba[:, n_states]  # [batch]

    # BBA → opinion → evidence → α
    fused_evidence = opinion_to_evidence(fused_belief_masses, fused_uncertainty, W)
    fused_alpha = evidence_to_dirichlet(fused_evidence, base_rates, W)

    return fused_alpha, conflict


def _dempster_combination_rule(
    bba_sources: typing.Sequence['Tensor'],
    conflict_handling: str,
    device: torch.device,
    dtype: torch.dtype,
) -> Tuple['Tensor', 'Tensor']:
    """执行 Dempster 组合规则的核心实现

    Parameters
    ----------
    - `bba_sources` : `list[Tensor]`
        - BBA 列表，每个张量形状为 `[batch, n_states+1]` (最后一维是无知状态)
    - `conflict_handling` : `str`
        - 冲突处理策略
    - `device` : `torch.device`
        - 设备类型
    - `dtype` : `torch.dtype`
        - 数据类型

    Returns
    -------
    - `Tuple[Tensor, Tensor]`
        - (fused_bba, conflict) - 融合后的BBA和冲突度
    """
    # 初始化为第一个 BBA
    fused_bba = bba_sources[0].clone()
    batch_size, n_focal = fused_bba.shape
    n_states = n_focal - 1  # 减去无知状态

    total_conflict = torch.zeros(batch_size, device=device, dtype=dtype)

    # 逐步融合剩余的 BBA
    for bba in bba_sources[1:]:
        fused_bba, conflict = _combine_two_bba(
            fused_bba, bba, conflict_handling, n_states
        )
        total_conflict += conflict

    return fused_bba, total_conflict


def _combine_two_bba(
    bba1: 'Tensor', bba2: 'Tensor', conflict_handling: str, n_states: int
) -> Tuple['Tensor', 'Tensor']:
    """组合两个 BBA 的核心算法

    实现 Dempster 组合规则：m(A) = Σ{m1(B)m2(C) : B∩C=A} / (1-K)

    Parameters
    ----------
    - `bba1`, `bba2` : `Tensor`
        - 两个 BBA，形状为 `[batch, n_states+1]`
    - `conflict_handling` : `str`
        - 冲突处理策略
    - `n_states` : `int`
        - 状态数量（不包括无知状态）

    Returns
    -------
    - `Tuple[Tensor, Tensor]`
        - (combined_bba, conflict) - 组合后的BBA和冲突度
    """
    dtype = bba1.dtype

    # 为简化实现，我们采用直接的矩阵运算方式
    # 这里实现简化版的 DST 组合规则

    # 分离各状态信念和无知状态
    m1_states = bba1[:, :n_states]  # [batch, n_states]
    m1_ignorance = bba1[:, n_states]  # [batch]
    m2_states = bba2[:, :n_states]  # [batch, n_states]
    m2_ignorance = bba2[:, n_states]  # [batch]

    # 计算状态-状态的组合（一致性支持）
    # m_combined[i] = m1[i] * m2_ignorance + m2[i] * m1_ignorance + m1[i] * m2[i]
    consistent_support = (
        m1_states * m2_ignorance.unsqueeze(1)  # m1状态 × m2无知
        + m2_states * m1_ignorance.unsqueeze(1)  # m2状态 × m1无知
        + m1_states * m2_states  # m1状态 × m2状态 (相同状态)
    )

    # 计算冲突（不同状态间的交集为空）
    # conflict = Σ_{i≠j} m1_i * m2_j
    # 高效的向量化计算方法:
    # Σ_{i,j} m1_i * m2_j = (Σ m1_i) * (Σ m2_j)
    # Σ_{i=j} m1_i * m2_j = Σ (m1_i * m2_i)
    # 因此, conflict = (Σ m1_i) * (Σ m2_j) - Σ (m1_i * m2_i)
    m1_sum = torch.sum(m1_states, dim=1)
    m2_sum = torch.sum(m2_states, dim=1)
    m1_dot_m2 = torch.sum(m1_states * m2_states, dim=1)
    conflict = m1_sum * m2_sum - m1_dot_m2

    # 计算新的无知状态质量
    combined_ignorance = m1_ignorance * m2_ignorance

    # 根据冲突处理策略处理结果
    if conflict_handling == 'dempster':
        # 经典 Dempster 规则：归一化，忽略冲突
        normalization = 1.0 - conflict
        # 数值稳定性处理
        eps = 1e-12 if dtype == torch.float64 else 1e-7
        normalization = torch.clamp(normalization, min=eps)

        combined_states = consistent_support / normalization.unsqueeze(1)
        combined_ignorance = combined_ignorance / normalization

    elif conflict_handling == 'yager':
        # Yager 规则：将冲突分配给无知状态
        combined_states = consistent_support
        combined_ignorance = combined_ignorance + conflict

    elif conflict_handling == 'zhang':
        # Zhang 规则：基于可信度重分配冲突
        # 简化实现：按比例重分配
        total_support = torch.sum(
            consistent_support, dim=1, keepdim=True
        ) + combined_ignorance.unsqueeze(1)
        conflict_redistribution = conflict.unsqueeze(1) * (
            consistent_support / (total_support + 1e-8)
        )

        combined_states = consistent_support + conflict_redistribution
        # 无知状态保持不变

    else:
        raise ValueError(f'不支持的冲突处理策略: {conflict_handling}')

    # 确保归一化
    combined_bba = torch.cat([combined_states, combined_ignorance.unsqueeze(1)], dim=1)
    bba_sum = torch.sum(combined_bba, dim=1, keepdim=True)
    combined_bba = combined_bba / (bba_sum + 1e-8)  # 数值稳定性

    return combined_bba, conflict
