"""
=============================================================================
第 3 讲 · 代码 2：Luce 选择公理与 IIA —— 它带来了什么，怎样用数据检验它
=============================================================================

对应讲义：lectures/lecture03_random_utility.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Luce（1959）的选择公理（"无关选项独立性"，IIA）说：a 和 b 之间的相对偏好，不受旁边还有什么选项影响。
这条公理一旦成立，就带来一连串好处，也带来一个可以检验的预言：

    演示 1：IIA 的"红利"—— 每个选项一个"实力值" π_j，就能算出任何选择集上的选择概率。
            于是**只用两两比较的数据**，就能预测"五选一"时每个选项被选中的概率（反事实预测）。
    演示 2：Plackett–Luce 的两个推论 ——
            ① 在一个完整的排序里，a 排在 b 前面的概率恰好是 σ(s_a − s_b)（第 2 讲用到的事实）；
            ② 但是"从最差的开始排"不是 softmax(−s)：Gumbel 噪声不对称，"挑最好的"和"挑最差的"是两个不同的模型。
    演示 3：用数据检验 IIA —— 分别收集"二选一"和"三选一"的数据，看 P(a)/P(b) 有没有变。
            单一人群（logit）通过检验；两群口味不同的人混在一起，检验报警。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_luce_axiom_iia.py        # 约 1 秒

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

import itertools
import math
import random

from rum import fmt_probs, logit_probs, mixture_probs
from toy_world import (BY_KEY, Annotator, fit_bradley_terry, make_expert, mean, pad, print_header, sample_gumbel,
                       sigmoid)

EXPERT = make_expert()
FIVE = ["A", "B", "C", "D", "F"]
UTILS = {k: EXPERT.utility(BY_KEY[k]) for k in FIVE}


# =============================================================================
# 演示 1：IIA 的红利 —— 从两两比较预测五选一
# =============================================================================
def demo_dividend(rng, n_pairs=4000, n_five=20000):
    print_header("演示 1：IIA 的红利 —— 只用两两比较的数据，预测“五选一”")
    print(f"专家在 5 个回答 {FIVE} 上的选择服从 logit 模型（满足 IIA）。")

    # (1) 乘法规则：P(j | 全集) = P(j | S) × P(选中的在 S 里 | 全集)
    full = logit_probs(UTILS, FIVE)
    s = ["A", "C", "F"]
    ps = logit_probs(UTILS, s)
    p_in_s = math.fsum(full[k] for k in s)
    print("\n(1) Luce 公理的“乘法规则”：P(j | 全集) = P(j | S) × P(选中的在 S 里 | 全集)，取 S = {A, C, F}：")
    for j in s:
        print(f"    {j}：左边 {full[j]:.4f}，右边 {ps[j]:.4f} × {p_in_s:.4f} = {ps[j] * p_in_s:.4f}")

    # (2) 只用两两比较估计实力值，再预测五选一
    pairs = []
    for _ in range(n_pairs):
        a, b = rng.sample(FIVE, 2)
        pairs.append((a, b) if rng.random() < sigmoid(UTILS[a] - UTILS[b]) else (b, a))
    # 用第 1 讲的 BT 最大似然估计 log π_j（= 分数 s_j），再代进 softmax 预测五选一
    log_pi = fit_bradley_terry(pairs, FIVE)
    predicted = logit_probs(log_pi, FIVE)
    counts = {k: 0 for k in FIVE}
    for _ in range(n_five):
        feel = {k: UTILS[k] + sample_gumbel(rng) for k in FIVE}
        counts[max(FIVE, key=feel.get)] += 1
    print(f"\n(2) 只用 {n_pairs} 次两两比较估计实力值（BT 最大似然），预测五选一；和 {n_five} 次真实的五选一对比：")
    print(f"    {pad('回答', 4)} {pad('预测', 8, 'right')} {pad('实际比例', 10, 'right')}")
    for k in FIVE:
        print(f"    {k:4s} {predicted[k]:>8.3f} {counts[k] / n_five:>10.3f}")
    print("""
观察：
  · 乘法规则成立：在全集里选中 j 的概率 = “选中的在 S 里”的概率 × “在 S 里选中 j”的概率。
    Luce 证明了，这条规则（加上一个技术条件）等价于“存在实力值 π_j，使 P(j | S) = π_j / Σ_{k∈S} π_k”。
  · 于是二选一的数据就足以预测任何选择集上的选择：π_j / π_k = P(j ≻ k) / P(k ≻ j)。用两两比较拟合出的
    分数代进 softmax，预测和真实的五选一很接近（误差来自两两数据的有限样本）。
  · 这种“反事实预测”（从没在这个选择集上观察过，却能预测）是 IIA 最大的价值，也是它最大的风险：
    如果 IIA 在真实数据里不成立，这些预测就会系统性地错（演示 3、代码 03、04）。""")


# =============================================================================
# 演示 2：Plackett–Luce 的两个推论
# =============================================================================
def pl_prob(order, s):
    """Plackett–Luce：一个完整排序的概率（一个一个地挑最好的）。"""
    p = 1.0
    for i in range(len(order) - 1):
        p *= logit_probs(s, order[i:])[order[i]]
    return p


def demo_pl_consequences(rng, draws=100000):
    print_header("演示 2：Plackett–Luce 的两个推论")
    keys = ["A", "C", "D", "F"]
    s = {k: UTILS[k] for k in keys}
    orders = list(itertools.permutations(keys))
    dist = {o: pl_prob(o, s) for o in orders}
    print(f"4 个回答 {keys}，枚举全部 {len(orders)} 种排序，按 Plackett–Luce 算出每种排序的概率（总和 {math.fsum(dist.values()):.6f}）。")

    print("\n(1) 在完整的排序里，x 排在 y 前面的概率 vs σ(s_x − s_y)：")
    for x, y in [("A", "C"), ("C", "D"), ("D", "F"), ("A", "F")]:
        above = math.fsum(p for o, p in dist.items() if o.index(x) < o.index(y))
        print(f"    P({x} 在 {y} 前面) = {above:.6f}    σ(s_{x} − s_{y}) = {sigmoid(s[x] - s[y]):.6f}")

    print("\n(2) 排在最后一名的概率 vs softmax(−s)（“挑最差的”如果也用 logit 模型）：")
    last = {k: math.fsum(p for o, p in dist.items() if o[-1] == k) for k in keys}
    neg = logit_probs({k: -v for k, v in s.items()}, keys)
    sim = {k: 0 for k in keys}
    for _ in range(draws):
        feel = {k: s[k] + sample_gumbel(rng) for k in keys}
        sim[min(keys, key=feel.get)] += 1
    print(f"    {pad('回答', 4)} {pad('PL：排最后', 11, 'right')} {pad('模拟：感受最低', 14, 'right')} {pad('softmax(−s)', 12, 'right')}")
    for k in keys:
        print(f"    {k:4s} {last[k]:>11.3f} {sim[k] / draws:>14.3f} {neg[k]:>12.3f}")
    print("""
观察：
  · (1) 两边逐位相同（到小数点后第 6 位）：PL 排序里 x 排在 y 前面的概率**恰好**是 BT 的 σ(s_x − s_y)，
    不管其他选项是什么。证明一句话：x 排在 y 前面 ⇔ x 的感受 > y 的感受，而这件事只和 x、y 两个噪声有关。
    这就是第 2 讲“把排序拆成两两比较再用 BT 拟合是合理的”的原因。
  · (2) “排最后一名”的概率**不是** softmax(−s)。原因是 Gumbel 噪声不对称（右边尾巴长）：
    “挑最好的”关心的是感受的最大值（右尾），“挑最差的”关心的是最小值（左尾），两者的分布形状不同。
    如果让标注员“先挑出最差的，再挑次差的……”，数据服从的是另一个模型（用反向 Gumbel 噪声）。
    实践意义：收集排序数据时，用什么顺序问（先挑最好还是先挑最差）会影响该用哪个模型；
    最好-最差标度法（best-worst scaling）同时问两者，需要专门的模型。""")


# =============================================================================
# 演示 3：用数据检验 IIA
# =============================================================================
def simulate_choices(prob_fn, options, n, rng):
    """按选择概率 prob_fn(options) 模拟 n 次选择，返回每个选项被选中的次数。"""
    p = prob_fn(options)
    keys = list(options)
    counts = {k: 0 for k in keys}
    for _ in range(n):
        r = rng.random()
        acc = 0.0
        for k in keys:
            acc += p[k]
            if r < acc:
                counts[k] += 1
                break
        else:
            counts[keys[-1]] += 1
    return counts


def log_ratio_test(prob_fn, rng, n=4000, rounds=300):
    """
    IIA 检验：比较 log[P(a)/P(b)] 在 {a,b} 和 {a,b,c} 上的估计值，用自助法估计差值的 95% 区间。

    【思路】IIA 预言：两个选择集上的 log 比值相同，差值应该在 0 附近。如果差值的 95% 区间不包含 0，
        就有证据表明 IIA 不成立。这是 Hausman–McFadden（1984）检验的简化版：
        "去掉一些选项后重新估计，看参数变没变"。
    【参数】prob_fn：真实的选择模型（用来模拟数据）；rng；n：每个选择集的选择次数；rounds：自助法次数
    【返回】(差值的估计, 区间下限, 区间上限)
    """
    pair = simulate_choices(prob_fn, ("a", "b"), n, rng)
    triple = simulate_choices(prob_fn, ("a", "b", "c"), n, rng)

    def diff(cp, ct):
        return math.log((cp["a"] + 0.5) / (cp["b"] + 0.5)) - math.log((ct["a"] + 0.5) / (ct["b"] + 0.5))

    est = diff(pair, triple)
    # 自助法：把每个选择集的数据看成多项分布，按观察到的频率重新抽样
    boots = []
    for _ in range(rounds):
        bp = simulate_choices(lambda o: {k: pair[k] / n for k in o}, ("a", "b"), n, rng)
        bt = simulate_choices(lambda o: {k: triple[k] / n for k in o}, ("a", "b", "c"), n, rng)
        boots.append(diff(bp, bt))
    boots.sort()
    return est, boots[int(0.025 * rounds)], boots[int(0.975 * rounds) - 1]


def demo_test_iia(rng):
    print_header("演示 3：用数据检验 IIA —— 二选一 vs 三选一")
    concise = Annotator(name="简洁派", taste=-4.0)
    detailed = Annotator(name="详尽派", taste=+4.0)
    # a = 回答 A（简洁准确），b = 回答 B（详尽教程），c = 另一篇长教程：只有详尽派喜欢它
    # （简洁派给它的效用 2.0；详尽派给它 6.84，和 B 一样好）

    def utils_of(person):
        return {"a": person.utility(BY_KEY["A"]), "b": person.utility(BY_KEY["B"]),
                "c": 2.0 if person is concise else 6.84}

    avg = {x: 0.6 * utils_of(concise)[x] + 0.4 * utils_of(detailed)[x] for x in ("a", "b", "c")}
    worlds = {
        "单一人群（logit，效用取两群人的平均）": lambda o: logit_probs(avg, o),
        "两群人混合（60% 简洁派 + 40% 详尽派）": lambda o: mixture_probs(
            [(0.6, lambda oo: logit_probs(utils_of(concise), oo)),
             (0.4, lambda oo: logit_probs(utils_of(detailed), oo))], o),
    }
    print("a = 回答 A（简洁）、b = 回答 B（详尽）、c = 另一篇只有详尽派喜欢的长教程。每个世界里，二选一和三选一各观察 4000 次。\n")
    print(f"  {pad('世界', 40)} {pad('P(a)/P(b) 二选一', 17, 'right')} {pad('三选一', 8, 'right')} "
          f"{pad('log 比值之差（95% 区间）', 28, 'right')}")
    for name, fn in worlds.items():
        r2 = fn(("a", "b"))
        r3 = fn(("a", "b", "c"))
        est, lo, hi = log_ratio_test(fn, rng)
        verdict = "拒绝 IIA" if lo > 0 or hi < 0 else "不能拒绝"
        print(f"  {pad(name, 40)} {r2['a'] / r2['b']:>17.2f} {r3['a'] / r3['b']:>8.2f} "
              f"{est:>+8.3f} [{lo:+.3f}, {hi:+.3f}] {verdict}")
    print("""
观察：
  · 单一人群的 logit 世界里，两个选择集上的比值相同（真实值完全相等），检验的区间包含 0。
  · 两群人混合时，加入 c 以后 a 和 b 的比值大幅上升：c 只在详尽派里有吸引力，它和 B 平分了详尽派的选票，
    而 A 的选票主要来自简洁派，几乎没受影响。每群人**内部**都满足 IIA（c 按比例从 A、B 那里抢份额），
    混在一起就不满足了：c 抢的是“和它口味相同的那群人”的选票（代码 04 会更系统地讨论）。
  · 这个检验需要**三选一（或更多）**的数据。只收集两两比较，IIA 永远无法被检验（代码 01 演示 4 的“反向 Gumbel”）。
    这也是为什么第 2 讲说“排序数据信息更多”：它包含了多选一的事件，能检验两两数据检验不了的假设。""")


def main():
    rng = random.Random(1959)        # 1959：Luce 出版《个体选择行为》的年份
    demo_dividend(rng)
    demo_pl_consequences(rng)
    demo_test_iia(rng)


if __name__ == "__main__":
    main()