"""
=============================================================================
第 3 讲 · 代码 1：为什么偏偏是 Gumbel？—— 最大值分布的四个"魔法"性质
=============================================================================

对应讲义：lectures/lecture03_random_utility.md 第 3、4.2 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 1、2 讲反复用到一个事实："效用 + 独立的 Gumbel 噪声，选最大的"给出 softmax / Bradley–Terry。
这个文件用模拟和数值积分把 Gumbel 分布拆开来看：

    演示 1：Gumbel 是"很多随机数的最大值"的分布（极值定理）—— 这是它出现在"选最好的"问题里的原因。
    演示 2：两个 Gumbel 之差是逻辑斯蒂分布 → σ 函数；一组 Gumbel 的最大值仍是 Gumbel（max-稳定性），
            位置移到 log Σ e^{V}（log-sum-exp）。
    演示 3：最神奇的一条：**谁赢**（argmax）和**赢家的感受有多大**（max）相互独立。
            不管赢家的感受是高是低，"谁赢"的分布都是同一个 softmax。这正是"无关选项独立性"的来源。
    演示 4：换成别的噪声（正态、均匀、反向 Gumbel），P(a)/P(b) 会随着"旁边有没有 c"而变化 ——
            只有 Gumbel 让这个比值保持不变。这就是定理"IIA ⇔ Gumbel"的数值验证。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_gumbel_magic.py        # 约 1 秒

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

import math
import random

from rum import ALL_NOISES, GUMBEL, choice_probs_integral, logit_probs
from toy_world import EULER_GAMMA, bar, mean, pad, print_header, sample_gumbel, sigmoid


# =============================================================================
# 演示 1：最大值的分布
# =============================================================================
def demo_extreme_value(rng, n=200, draws=20000):
    print_header("演示 1：很多随机数的最大值 → Gumbel 分布（极值定理）")
    # n 个独立的标准指数分布随机数的最大值，减去 log n 以后，分布趋近标准 Gumbel
    maxima = []
    for _ in range(draws):
        m = max(rng.expovariate(1.0) for _ in range(n))
        maxima.append(m - math.log(n))
    print(f"抽 {n} 个独立的指数分布随机数，取最大值再减去 log {n}；重复 {draws} 次，和标准 Gumbel 的 CDF 比较：\n")
    print(f"  {'x':>5s} {pad('模拟 P(最大值 ≤ x)', 18, 'right')} {pad('Gumbel CDF', 12, 'right')}")
    for x in [-1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 3.0]:
        sim = sum(1 for v in maxima if v <= x) / draws
        print(f"  {x:>5.1f} {sim:>18.3f} {GUMBEL.cdf(x):>12.3f}")
    m = mean(maxima)
    var = mean([(v - m) ** 2 for v in maxima])
    print(f"\n  模拟的平均值 {m:.3f}（Gumbel 的平均值 γ = {EULER_GAMMA:.3f}）；方差 {var:.3f}（π²/6 = {math.pi ** 2 / 6:.3f}）")
    print("""
观察：
  · “很多独立随机数的最大值”的分布，经过平移和缩放以后会趋近三种“极值分布”之一（Fisher–Tippett–Gnedenko 定理），
    指数、正态这类尾巴像指数函数一样衰减的分布，最大值趋近的就是 Gumbel。
  · 这给了 Gumbel 噪声一个直观的解释：如果你对一个选项的感受，是很多个“小方面”里最突出的那个
    （“这个回答最打动我的地方”），那这个感受的波动就近似服从 Gumbel。
    当然，这只是一个说得通的故事；第 4.2 节会看到，真正让 Gumbel 特殊的是 IIA。""")


# =============================================================================
# 演示 2：差是逻辑斯蒂，最大值仍是 Gumbel
# =============================================================================
def demo_difference_and_max(rng, draws=40000):
    print_header("演示 2：两个 Gumbel 之差 → σ；一组 Gumbel 的最大值 → 还是 Gumbel")
    diffs = [sample_gumbel(rng) - sample_gumbel(rng) for _ in range(draws)]
    print("(1) G1 − G2 的 CDF 和 σ(x) 比较：")
    for x in [-3, -1, 0, 1, 3]:
        sim = sum(1 for d in diffs if d <= x) / draws
        print(f"    x = {x:>2d}：模拟 {sim:.3f}，σ(x) = {sigmoid(x):.3f}")
    print("    所以 P(V_a + G_a > V_b + G_b) = P(G_b − G_a < V_a − V_b) = σ(V_a − V_b) —— Bradley–Terry。")

    utils = {"a": 1.0, "b": 0.0, "c": -1.0, "d": 0.5}
    lse = math.log(math.fsum(math.exp(v) for v in utils.values()))
    maxima = [max(v + sample_gumbel(rng) for v in utils.values()) for _ in range(draws)]
    print(f"\n(2) 4 个选项，效用 {utils}，感受 = 效用 + Gumbel。所有感受的最大值：")
    print(f"    log-sum-exp = log Σ e^V = {lse:.3f}。如果最大值 ~ Gumbel 平移 {lse:.3f}，它的平均值应为 {lse + EULER_GAMMA:.3f}")
    print(f"    模拟的平均值 = {mean(maxima):.3f}")
    for x in [0.5, 1.5, 2.5, 4.0]:
        sim = sum(1 for v in maxima if v <= x) / draws
        print(f"    P(最大值 ≤ {x})：模拟 {sim:.3f}，Gumbel(平移 {lse:.2f}) 的 CDF {GUMBEL.cdf(x - lse):.3f}")
    print("""
观察：
  · 差是逻辑斯蒂：这就是 σ 函数（以及 Bradley–Terry 模型）的来历。
  · 最大值仍是 Gumbel，只是平移到了 log Σ e^{V_j}。这叫“max-稳定性”：一组 Gumbel 取最大值，还在同一个分布家族里。
    log Σ e^{V_j} 在经济学里叫“包容值”（inclusive value）：它衡量“这个选择集整体有多好”，
    嵌套 logit（代码 03）用它把“选哪个巢”和“巢里选哪个”连起来。""")


# =============================================================================
# 演示 3：谁赢 与 赢多少 相互独立
# =============================================================================
def demo_argmax_independent(rng, draws=60000):
    print_header("演示 3：“谁赢”与“赢家的感受有多高”相互独立")
    utils = {"a": 1.0, "b": 0.0, "c": -1.0}
    soft = logit_probs(utils, list(utils))
    records = []
    for _ in range(draws):
        feel = {k: v + sample_gumbel(rng) for k, v in utils.items()}
        win = max(feel, key=feel.get)
        records.append((feel[win], win))
    records.sort()
    third = draws // 3
    groups = [("赢家感受最低的 1/3", records[:third]), ("中间的 1/3", records[third:2 * third]),
              ("赢家感受最高的 1/3", records[2 * third:])]
    print(f"3 个选项，效用 {utils}；模拟 {draws} 次，按“赢家的感受”分成三组，看每组里谁赢的比例：\n")
    print(f"  {pad('分组', 20)} " + " ".join(f"{k:>7s}" for k in utils))
    for name, recs in groups:
        n = len(recs)
        print(f"  {pad(name, 20)} " + " ".join(f"{sum(1 for _, w in recs if w == k) / n:>7.3f}" for k in utils))
    print(f"  {pad('softmax', 20)} " + " ".join(f"{soft[k]:>7.3f}" for k in utils))

    # 对比：正态噪声时，赢家的感受高低和谁赢有关
    recs_n = []
    for _ in range(draws):
        feel = {k: v + rng.gauss(0.0, math.pi / math.sqrt(6)) for k, v in utils.items()}
        win = max(feel, key=feel.get)
        recs_n.append((feel[win], win))
    recs_n.sort()
    print("\n  换成正态噪声（同样的方差）：")
    for name, recs in [("赢家感受最低的 1/3", recs_n[:third]), ("赢家感受最高的 1/3", recs_n[2 * third:])]:
        n = len(recs)
        print(f"  {pad(name, 20)} " + " ".join(f"{sum(1 for _, w in recs if w == k) / n:>7.3f}" for k in utils))
    print("""
观察：
  · Gumbel 噪声下，三组的“谁赢”比例几乎完全一样，都等于 softmax（差别只是模拟误差）：知道赢家的感受有多高，
    对“谁赢了”没有任何额外的信息。这个性质是 Gumbel 独有的（在独立同分布的噪声里）。
  · 正态噪声下就不是这样：赢家的感受很高时，更可能是平均效用本来就高的 a 赢（“它靠实力赢”）；
    赢家的感受很低时，弱的选项赢的机会更大（“大家都不怎么样，谁都有可能”）。
  · 为什么这个性质这么重要？因为它意味着“在其他选项里谁是最强的对手、它的感受多高”不会改变 a 和 b 之间的
    相对胜率 —— 这正是下一个演示的“无关选项独立性”。""")


# =============================================================================
# 演示 4：哪种噪声满足 IIA
# =============================================================================
def demo_which_noise(c_values=(-2.0, 0.0, 1.0, 2.0)):
    print_header("演示 4：P(a)/P(b) 会不会因为“旁边多了一个 c”而改变？")
    print("a 的效用 1.0，b 的效用 0.0。先只在 {a, b} 里选，再加入效用为 V_c 的第三个选项 c。")
    print("四种噪声都取相同的方差；概率用数值积分精确计算。表里是 P(a | S) / P(b | S)：\n")
    header = f"  {pad('噪声', 12)} {pad('只有{a,b}', 10, 'right')} " + " ".join(
        pad(f'加c(V_c={v:g})', 14, 'right') for v in c_values)
    print(header)
    for noise in ALL_NOISES:
        base = choice_probs_integral({"a": 1.0, "b": 0.0}, ["a", "b"], noise)
        row = f"  {pad(noise.name, 12)} {base['a'] / base['b']:>10.3f} "
        for vc in c_values:
            p = choice_probs_integral({"a": 1.0, "b": 0.0, "c": vc}, ["a", "b", "c"], noise)
            row += f"{p['a'] / p['b']:>14.3f} "
        print(row)
    print(f"\n  e^(1 − 0) = {math.e:.3f}（Gumbel 的理论值）")
    print("""
观察：
  · 只有 Gumbel 的那一行完全不变：不管 c 是很差（V_c = −2）还是很强（V_c = 2），a 和 b 的比值都是 e ≈ 2.718。
    这就是 Luce 的“无关选项独立性”（IIA）。
  · 正态和均匀噪声：c 越强，比值越大 —— 强的 c 从较弱的 b 那里抢走的份额比例更大。直觉：旁边有一个强的对手时，
    只有“运气特别好”的感受才能赢；平均效用高的 a 更容易有这样的感受，所以 b 吃亏更多。
  · 反向 Gumbel 最有意思：只有 {a, b} 时，它的比值和 Gumbel **完全一样**（2.718）—— 两个反向 Gumbel 之差也是
    逻辑斯蒂分布，所以两两比较的数据根本分不出这两种噪声。可是一加入 c，它就严重违反 IIA。
    教训：**只有两两比较的数据，无法检验 IIA**；要检验它，必须观察三个或更多选项的选择。
  · 定理（Yellott 1977）：在“效用 + 独立同分布噪声”的模型里，如果对所有效用值、所有（至少三个选项的）
    选择集都满足 IIA，那么噪声必定是 Gumbel 分布（允许平移和缩放）。演示 4 是这个定理的数值验证。
  · 实际上，正态和 Gumbel 的差别在大多数情况下不大（比值都在 2.4~3.2 之间）。两者真正的分水岭不在独立噪声，
    而在**相关**噪声：相关的正态噪声（probit）可以描述“相似的选项互相抢份额”，独立的 Gumbel 不行（代码 03）。""")


def main():
    rng = random.Random(1977)        # 1977：Yellott 证明 IIA ⇔ Gumbel 的年份
    demo_extreme_value(rng)
    demo_difference_and_max(rng)
    demo_argmax_independent(rng)
    demo_which_noise()


if __name__ == "__main__":
    main()