"""estimator.py 单测——docs/04 §5.2:公式对拍 + 定理 4.1 性质 + 方差收缩。

对拍精神源自参考实现 validate_chunk_mc_estimator.py(比 MSE_freq vs
MSE_bayes),但全用 toy 数据本地 CPU 跑,不连真 teacher。
detach 命门的两个世界断言在 tests/test_estimator_detach.py(E3)。
"""

import math

import pytest
import torch

from ars_opd.estimator import bayesian_target, chunk_prior

# ------------------------------------------------------------- chunk_prior


def test_prior_is_geometric_mean():
    # 式(4) 手算:p = [0.9, 0.1] → π̄ = exp((log .9 + log .1)/2) = √0.09 = 0.3
    log_probs = torch.log(torch.tensor([0.9, 0.1]))
    assert math.isclose(chunk_prior(log_probs).item(), 0.3, rel_tol=1e-6)


def test_prior_uniform_probs():
    # 全同概率的几何均值 = 该概率本身
    log_probs = torch.full((50,), math.log(0.5))
    assert math.isclose(chunk_prior(log_probs).item(), 0.5, rel_tol=1e-6)


def test_prior_shape_and_range():
    pi_bar = chunk_prior(torch.log(torch.rand(50).clamp(1e-6, 1.0)))
    assert pi_bar.shape == ()  # (C,) -> 标量
    assert 0.0 < pi_bar.item() <= 1.0


def test_prior_log_domain_survives_underflow():
    # 50 个 p=0.01 直接连乘 = 1e-100(fp32 下溢为 0);log 域算出 0.01
    log_probs = torch.full((50,), math.log(0.01))
    assert math.isclose(chunk_prior(log_probs).item(), 0.01, rel_tol=1e-4)


def test_prior_clamp_floor():
    # 极端负 log 均值 → exp 下溢,clamp 兜到 1e-8 保持严格为正(定理 4.1b 前提)
    log_probs = torch.full((5,), -1e9)
    assert chunk_prior(log_probs).item() == pytest.approx(1e-8)


def test_prior_is_detached():
    # detach 命门:π̄ 不带梯度(逃逸机制的完整断言在 test_estimator_detach.py)
    log_probs = torch.log(torch.tensor([0.5, 0.5], requires_grad=True))
    pi_bar = chunk_prior(log_probs)
    assert not pi_bar.requires_grad


def test_prior_empty_raises():
    with pytest.raises(ValueError, match="为空"):
        chunk_prior(torch.tensor([]))


# --------------------------------------------------------- bayesian_target


def test_target_formula_hand_computed():
    # 式(5) 手算:k=8, π̄=0.5, N=10, α=1 → π̂ = (8 + 0.5)/11 = 0.77272…
    pi_hat = bayesian_target(8.0, torch.tensor(0.5), n_rollouts=10, alpha=1.0)
    assert math.isclose(pi_hat.item(), 8.5 / 11, rel_tol=1e-6)


def test_target_convex_combination_identity():
    # 式(10) 恒等:π̂ = N/(N+α)·(k/N) + α/(N+α)·π̄,任取参数逐点核对
    k, pi_bar, n, alpha = 3.7, torch.tensor(0.42), 10, 1.5
    direct = bayesian_target(k, pi_bar, n, alpha).item()
    convex = (n / (n + alpha)) * (k / n) + (alpha / (n + alpha)) * pi_bar.item()
    assert math.isclose(direct, convex, rel_tol=1e-6)


def test_target_anti_collapse_at_k_zero():
    # 定理 4.1(b):k=0(teacher 全否定)时 π̂ = α·π̄/(N+α) > 0,监督不归零
    pi_hat = bayesian_target(0.0, torch.tensor(0.3), n_rollouts=10, alpha=1.0)
    assert math.isclose(pi_hat.item(), 0.3 / 11, rel_tol=1e-6)
    assert pi_hat.item() > 0


def test_target_full_score_shrinks_below_one():
    # 贝叶斯收缩:k=N 满分时 π̂ = (N+α·π̄)/(N+α) < 1(只要 π̄<1)——
    # 先验把估计从两端往中间拉,这正是方差收缩的来源
    pi_hat = bayesian_target(10.0, torch.tensor(0.5), n_rollouts=10, alpha=1.0)
    assert math.isclose(pi_hat.item(), 10.5 / 11, rel_tol=1e-6)
    assert pi_hat.item() < 1.0


def test_target_bounded_in_unit_interval():
    # 定理 4.1(a):任意合法参数下 π̂ ∈ (0, 1]
    for k in [0.0, 2.5, 10.0]:
        for p in [1e-8, 0.5, 1.0]:
            v = bayesian_target(k, torch.tensor(p), 10, 1.0).item()
            assert 0.0 < v <= 1.0


def test_target_alpha_zero_is_frequency_estimate():
    # α=0 退化为 k/N(no_bayesian 消融);k=0 时被 clamp 兜到 1e-8 而非 0
    assert math.isclose(
        bayesian_target(7.0, torch.tensor(0.5), 10, 0.0).item(), 0.7, rel_tol=1e-6
    )
    assert bayesian_target(0.0, torch.tensor(0.5), 10, 0.0).item() == pytest.approx(
        1e-8
    )


def test_target_is_detached_even_with_grad_input():
    # 第二道防线:pi_bar 带梯度传入,π̂ 仍必须 detach
    pi_bar = torch.tensor(0.5, requires_grad=True)
    pi_hat = bayesian_target(5.0, pi_bar, 10, 1.0)
    assert not pi_hat.requires_grad


def test_target_validation_raises():
    pi_bar = torch.tensor(0.5)
    with pytest.raises(ValueError, match="n_rollouts"):
        bayesian_target(0.0, pi_bar, 0, 1.0)
    with pytest.raises(ValueError, match="alpha"):
        bayesian_target(0.0, pi_bar, 10, -0.1)
    with pytest.raises(ValueError, match="越界"):
        bayesian_target(11.0, pi_bar, 10, 1.0)  # k_sem > N:口径不一致
    with pytest.raises(ValueError, match="越界"):
        bayesian_target(-0.5, pi_bar, 10, 1.0)


# --------------------------------------------- 方差收缩(定理 4.1c,toy 模拟)


def test_variance_shrinkage_beats_frequency_estimate():
    """toy 模拟对拍 validate_chunk_mc_estimator.py 的 MSE_freq vs MSE_bayes。

    设真值 μ:每次试验采 N=10 个相似度 sim_i(均值 μ 的噪声),
    频率估计 = mean(sim) = k/N,贝叶斯估计 = (k + α·π̄)/(N+α)。
    先验 π̄ = μ(理想先验)时,收缩纯降方差、零偏差代价,MSE 必更小。
    """
    torch.manual_seed(0)
    mu, n, alpha = 0.7, 10, 1.0
    trials = 2000
    # (trials, N) 的相似度样本:均值 μ、截断到 [0,1]
    sims = (mu + 0.25 * torch.randn(trials, n)).clamp(0.0, 1.0)
    k = sims.sum(dim=1)  # (trials,) 每次试验的 k_sem
    freq = k / n
    bayes = (k + alpha * mu) / (n + alpha)
    mse_freq = ((freq - mu) ** 2).mean().item()
    mse_bayes = ((bayes - mu) ** 2).mean().item()
    assert mse_bayes < mse_freq


def test_variance_shrinkage_robust_to_imperfect_prior():
    # 先验偏离真值(π̄ = μ±0.1)仍应赢:α=1、N=10 时先验权重仅 1/11,
    # 引入的偏差平方远小于省下的方差(定理 4.1c 在论文设定下的稳健性)
    torch.manual_seed(1)
    mu, n, alpha, trials = 0.6, 10, 1.0, 2000
    sims = (mu + 0.25 * torch.randn(trials, n)).clamp(0.0, 1.0)
    k = sims.sum(dim=1)
    mse_freq = ((k / n - mu) ** 2).mean().item()
    for prior in [mu - 0.1, mu + 0.1]:
        bayes = (k + alpha * prior) / (n + alpha)
        assert ((bayes - mu) ** 2).mean().item() < mse_freq