"""MC 估计与 Dirichlet 贝叶斯平滑——论文 §3.2.2 式(4)(5)。

把 similarity.py 产出的软计数 k_sem(外部 teacher 信号)与学生自身的
chunk 置信度 π̄(内部先验)融合成有界目标 π̂ ∈ (0, 1],供层 5 的 chunk
损失当乘子:loss_c = −π̂ · mean(log p)。定理 4.1 三性质由此获得:
(a) π̂ 有界 → 无白盒式(2) 的梯度爆炸;(b) π̂ > 0 → k_sem=0 也不塌缩;
(c) 先验收缩 → 方差小于频率估计 k/N。

纯逻辑模块(CLAUDE.md §2):只依赖 torch,toy 张量本地 CPU 可测。
"""

import torch


def chunk_prior(log_probs: torch.Tensor) -> torch.Tensor:
    """式(4):π̄ = exp((1/C)·Σ_t log p_t)——学生对整个 chunk 的几何均值置信度。

    C 个 token 概率的几何均值,充当式(5) 的贝叶斯先验:teacher 采样(k_sem)
    是主信号,π̄ 只是"学生自己觉得这段有多稳"的地板,防 k_sem=0 时目标归零。

    参数:
        log_probs: 学生对 chunk 内各 token 的对数概率,shape (C,),值 ≤ 0。
            (调用方从 log_softmax 后 gather 标签位置所得,层 5 负责。)

    返回:
        π̄,标量张量 shape ()、值 ∈ [1e-8, 1],**不带梯度**。

    实现细节:
        - log 域先均值再 exp:直接连乘 C=50 个小概率会下溢
          (50 个 0.01 → 1e-100,超出 fp32 下限 ~1e-38),log 域安全。
        - detach 命门(参考实现 distillation_trainer.py:2196 同):π̄ 是学生
          自身概率的函数,若保留梯度,优化器会发现"压低自己的 chunk 概率
          → π̄→0 → π̂ 变小 → 损失权重变小"这条逃逸路径——恰在 k_sem=0
          (teacher 否定)的 chunk 上最有利可图,这些 chunk 最先塌缩。
          π̄ 只能当常数先验,不能当优化变量。锁死断言见
          tests/test_estimator_detach.py。
        - clamp 下限 1e-8:极端负的均值 exp 后可能下溢为 0,而定理 4.1(b)
          的反塌缩要求 π̄ 严格为正。差异标注:参考实现 clamp(1e-8, 1.0),
          上限实为冗余——log p ≤ 0 ⇒ mean ≤ 0 ⇒ exp ≤ 1,此处省去。
    """
    if log_probs.numel() == 0:
        raise ValueError("log_probs 为空:chunk 至少要含 1 个 token")
    log_pi_bar = log_probs.detach().mean()  # (C,) -> ()
    return log_pi_bar.exp().clamp(min=1e-8)


def bayesian_target(
    k_sem: float,
    pi_bar: torch.Tensor,
    n_rollouts: int,
    alpha: float,
) -> torch.Tensor:
    """式(5):π̂ = (k_sem + α·π̄) / (N + α)——chunk 接受概率的贝叶斯估计。

    等价凸组合视角(论文式10):
        π̂ = N/(N+α) · (k_sem/N)  +  α/(N+α) · π̄
    即"teacher 频率估计"与"学生先验"的加权平均;默认 N=10、α=1 时权重
    约 91% : 9%,teacher 主导,先验只兜底。

    参数:
        k_sem: 式(3) 的软匹配计数,∈ [0, N](aggregate_similarity 产出)。
        pi_bar: 式(4) 的先验 π̄,标量张量(chunk_prior 产出)。
        n_rollouts: teacher rollout 数 N。**必须等于算 k_sem 时的
            len(teacher_rollouts)**——分子分母口径不一致会系统性偏移 π̂。
        alpha: 先验强度 α ≥ 0。α=0 退化为频率估计 k/N(层 6 的
            no_bayesian 消融,参考 config.py:315),失去定理 4.1(b) 保护。

    返回:
        π̂,标量张量 shape ()、值 ∈ [1e-8, 1],**不带梯度**。

    实现细节:
        - 差异标注:参考实现(distillation_trainer.py:2205)在此对 π̂ 整体
          detach;我们的 π̄ 在 chunk_prior 内已 detach,此处的 detach 是
          第二道防线——防止将来有人把带梯度的张量传进 pi_bar。
        - clamp(1e-8, 1.0):下限防 α=0 且 k_sem=0 时 π̂=0(乘子归零则该
          chunk 完全失去监督);上限防 pi_bar 越界传入时 π̂ 溢出概率语义。
    """
    if n_rollouts < 1:
        raise ValueError(f"n_rollouts 必须 ≥ 1,得到 {n_rollouts}")
    if alpha < 0:
        raise ValueError(f"alpha 必须 ≥ 0,得到 {alpha}")
    if not 0.0 <= k_sem <= n_rollouts:
        raise ValueError(
            f"k_sem={k_sem} 越界 [0, {n_rollouts}]:检查是否与"
            f" len(teacher_rollouts) 口径一致"
        )
    # 式(5): π̂ = (k_sem + α·π̄) / (N + α)
    pi_hat = (k_sem + alpha * pi_bar) / (n_rollouts + alpha)  # () -> ()
    return pi_hat.clamp(1e-8, 1.0).detach()