"""
=============================================================================
第 3 讲 · 代码 4：异质性 —— 混合 logit，以及"随机性到底从哪里来"
=============================================================================

对应讲义：lectures/lecture03_random_utility.md 第 2.1、5.3、5.4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
logit 模型里所有人共用一组平均效用 V_j，人和人的差别全部塞进"独立的噪声"里。
可是真实的人群里，有人喜欢简洁、有人喜欢详尽（第 1 讲代码 04）。混合 logit（随机系数 logit）
让每个人有自己的口味系数：

        效用_ij = 6·质量_j + β_i·(详细程度_j − 0.5) + Gumbel 噪声，   β_i ~ N(μ, σ²)

    演示 1：σ 越大（人群口味分歧越大），用两两比较拟合的 BT 模型预测"八选一"就错得越离谱。
    演示 2：给"是不是公交车"这个特征一个随机系数，混合 logit 就自动复现了红色公交车的直觉答案 ——
            随机系数让相似的选项"一起变好、一起变差"，等价于让它们的噪声相关。
            这就是 McFadden & Train（2000）"混合 logit 可以逼近任何随机效用模型"的直觉。
    演示 3：随机性来自"人和人不同"还是"同一个人每次不同"？只看每人一次的选择，两者**完全无法区分**；
            让同一个人回答两次，就能分清。

【数值方法】混合 logit 的选择概率是 softmax 对 β 的平均（一个积分）。这里用"正态分位数网格"做数值积分：
    取 K 个点 β_k = μ + σ·Φ⁻¹((k + 0.5)/K)，每个点的权重都是 1/K。这比随机抽样更准，而且结果是确定的。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_heterogeneity_mixed_logit.py        # < 1 秒

只用到 Python 标准库，以及第 1 讲的 toy_world.py、本目录的 rum.py。
"""

import math
import random
import statistics

from rum import fmt_probs, logit_probs
from toy_world import BY_KEY, KEYS, fit_bradley_terry, mean, pad, print_header, sample_gumbel, sigmoid


def normal_grid(mu, sigma, k=400):
    """
    正态分布 N(mu, sigma²) 的 k 个等权重"代表点"：第 i 个点是第 (i + 0.5)/k 分位数。

    【参数】mu、sigma：平均值和标准差（sigma = 0 时所有点都等于 mu）；k：点数
    【返回】长度为 k 的列表
    【Python 小知识】statistics.NormalDist().inv_cdf(p)：标准正态分布的 p 分位数 Φ⁻¹(p)。
    """
    nd = statistics.NormalDist()
    return [mu + sigma * nd.inv_cdf((i + 0.5) / k) for i in range(k)]


def mixed_logit_probs(base, feature, betas, options):
    """
    混合 logit 的选择概率：P(j | S) = 平均_β [ softmax_j( base + β·feature ) ]。

    【参数】
        base   ：dict {选项: 不随人变化的效用部分}
        feature：dict {选项: 带随机系数的那个特征的取值}
        betas  ：β 的代表点列表（normal_grid 的输出）
        options：选择集
    【返回】dict {选项: 概率}
    【边界】每个人内部是 logit（满足 IIA），整个人群一般不满足。
    """
    out = {k: 0.0 for k in options}
    for b in betas:
        p = logit_probs({k: base[k] + b * feature[k] for k in options}, options)
        for k in options:
            out[k] += p[k]
    return {k: v / len(betas) for k, v in out.items()}


# =============================================================================
# 演示 1：口味分歧越大，BT 的"反事实预测"越错
# =============================================================================
def demo_misprediction():
    print_header("演示 1：口味分歧越大，用两两比较的 BT 预测“八选一”越错")
    base = {k: 6.0 * BY_KEY[k].quality for k in KEYS}
    feat = {k: BY_KEY[k].detail - 0.5 for k in KEYS}
    print("效用 = 6·质量 + β·(详细程度 − 0.5)，β ~ N(0, σ²)：平均来说人们对详细程度没有偏好，但个人之间有分歧。")
    print("对每个 σ：先算出精确的两两胜率，按它生成“无限多”的比较（每对 2000 次，按概率分配胜负），")
    print("拟合 BT；再用 BT 的分数代进 softmax，预测 8 个回答一起出现时的选择份额，和混合 logit 的真实份额比较。\n")
    print(f"  {pad('σ', 4)} {pad('真实份额 A / B / C', 22, 'right')} {pad('BT 预测 A / B / C', 22, 'right')} "
          f"{pad('总变差距离', 10, 'right')}")
    for sigma in [0.0, 2.0, 4.0, 8.0]:
        betas = normal_grid(0.0, sigma)
        comps = []
        for i, a in enumerate(KEYS):
            for b in KEYS[i + 1:]:
                p = mixed_logit_probs(base, feat, betas, (a, b))[a]
                wins = round(2000 * p)
                comps += [(a, b)] * wins + [(b, a)] * (2000 - wins)
        s = fit_bradley_terry(comps, KEYS, l2=0.0, tol=1e-10)
        truth = mixed_logit_probs(base, feat, betas, KEYS)
        pred = logit_probs(s, KEYS)
        tv = 0.5 * math.fsum(abs(truth[k] - pred[k]) for k in KEYS)
        t3 = " / ".join(f"{truth[k]:.3f}" for k in "ABC")
        p3 = " / ".join(f"{pred[k]:.3f}" for k in "ABC")
        print(f"  {sigma:>4.0f} {t3:>22s} {p3:>22s} {tv:>10.3f}")
    print("""
观察：
  · σ = 0（所有人口味相同）时，混合 logit 就是普通 logit，BT 的预测完全正确。
  · σ 越大，误差越大。以 σ = 8 为例：简洁派（β < 0）几乎都选 A，详尽派（β > 0）几乎都选 B ——
    在八选一里，A 和 B 各自拥有一大群“铁杆粉丝”。BT 从两两比较里只看到“A 和 B 势均力敌、都很强”，
    却不知道它们的支持者是**两群不同的人**，于是低估了 A、B 的份额，高估了“谁都不讨厌”的 C。
  · 第 1 讲代码 04 的“两群人”是 β 只取两个值（−4、+4）的特例；这里 β 是连续分布。两者都说明：
    把异质的人群当成一个人去拟合，模型在训练数据（两两比较）上看起来没问题，一做反事实预测就错。""")


# =============================================================================
# 演示 2：随机系数复现红色公交车
# =============================================================================
def demo_error_components():
    print_header("演示 2：给“是不是公交车”一个随机系数 —— 混合 logit 自己学会了红色公交车")
    base = {"汽车": 0.0, "蓝巴士": 0.0, "红巴士": 0.0}
    is_bus = {"汽车": 0.0, "蓝巴士": 1.0, "红巴士": 1.0}
    two = ["汽车", "蓝巴士"]
    three = ["汽车", "蓝巴士", "红巴士"]
    print("效用 = 0 + β·[是公交车] + Gumbel 噪声，β ~ N(0, σ²)：有人喜欢坐公交，有人讨厌，平均来说无所谓。\n")
    print(f"  {pad('σ', 4)} {pad('汽车 vs 蓝巴士', 26, 'right')}   加上红巴士以后")
    for sigma in [0.0, 1.0, 2.0, 4.0, 8.0]:
        betas = normal_grid(0.0, sigma)
        p2 = mixed_logit_probs(base, is_bus, betas, two)
        p3 = mixed_logit_probs(base, is_bus, betas, three)
        print(f"  {sigma:>4.0f} {fmt_probs(p2):>26s}   {fmt_probs(p3, three)}")
    print("""
观察：
  · σ = 0 时就是普通 logit：汽车 1/3。σ 越大，汽车的份额越接近直觉的 1/2。
  · 为什么？喜欢坐公交的人（β 很大）会在两种公交车之间选，讨厌公交的人（β 很小）几乎都开车。
    红色公交车只会从“喜欢公交的人”那里分走蓝色公交车的乘客，抢不到开车的人。
    从整个人群看，这等价于“两种公交车的噪声是相关的”（它们共享同一个 β）—— 效果和嵌套 logit 一样。
  · McFadden & Train（2000）证明：只要随机系数的分布足够灵活，混合 logit 可以任意逼近**任何**随机效用模型。
    代价是选择概率变成了积分，估计时要做数值积分或模拟（“模拟最大似然”）。
  · 对 RLHF 的启示：“相似的回答互相抢份额”和“人群口味有分歧”是同一枚硬币的两面 —— 都来自
    “某些选项的效用在人群里一起涨落”。第 4 讲的因子模型就是用低维的“口味向量”来描述这种共同涨落。""")


# =============================================================================
# 演示 3：随机性从哪里来？
# =============================================================================
def demo_where_randomness(rng, people=20000):
    print_header("演示 3：随机性来自“人和人不同”，还是“同一个人每次不同”？")
    va, vc = 6.0 * BY_KEY["A"].quality, 6.0 * BY_KEY["C"].quality
    print(f"比较回答 A（平均效用 {va:.1f}）和 C（平均效用 {vc:.1f}）。两个世界：")
    print("  · 异质性世界：每个人的效用 = 平均效用 + 一次性抽好的个人 Gumbel 偏好，之后每次都按它选（完全一致）。")
    print("  · 有限理性世界：所有人的平均效用都一样，每次选择时重新抽 Gumbel 噪声（同一个人也会变）。")
    print(f"每个世界找 {people} 个人，每人回答两次。\n")
    print(f"  {pad('世界', 18)} {pad('第一次选 A 的比例', 18, 'right')} {pad('两次答案一致的比例', 18, 'right')}")
    for name in ["异质性", "有限理性"]:
        first_a = 0
        same = 0
        for _ in range(people):
            if name == "异质性":
                ga, gc = sample_gumbel(rng), sample_gumbel(rng)
                c1 = c2 = (va + ga > vc + gc)
            else:
                c1 = va + sample_gumbel(rng) > vc + sample_gumbel(rng)
                c2 = va + sample_gumbel(rng) > vc + sample_gumbel(rng)
            first_a += c1
            same += (c1 == c2)
        print(f"  {pad(name, 18)} {first_a / people:>18.3f} {same / people:>18.3f}")
    p = sigmoid(va - vc)
    print(f"\n  理论值：选 A 的比例都是 σ({va - vc:.1f}) = {p:.3f}；有限理性世界里两次一致的概率 = p² + (1 − p)² = {p * p + (1 - p) ** 2:.3f}")
    print("""
观察：
  · 只看每人的第一次选择，两个世界的数据**一模一样**（都是 σ(0.6) ≈ 0.65）—— 这就是“观测等价”。
    教材里随机效用的三种解释（异质性、有限理性、设计者的不确定性）在这种数据上无法区分。
  · 让同一个人回答两次就分清了：异质性世界里每个人前后完全一致，有限理性世界里只有一半多一致。
    真实的人介于两者之间：既有稳定的个人口味，也有每次的波动。
  · 这对偏好学习很重要：如果分歧主要来自“人和人不同”，就该做个性化（第 4 讲）、认真对待“谁的偏好”（第 16、17 讲）；
    如果主要来自“每次的波动”，就该多收集几次、取平均，或者设计更好的标注流程（第 15 讲）。
    很多 RLHF 数据集会让一部分样本被标注两次，用来估计标注员**自身**的一致率，就是在做这里的区分。""")


def main():
    rng = random.Random(2000)        # 2000：McFadden & Train 发表混合 logit 论文、McFadden 获得诺贝尔奖的年份
    demo_misprediction()
    demo_error_components()
    demo_where_randomness(rng)


if __name__ == "__main__":
    main()