"""
=============================================================================
第 3 讲 · 代码 5：情境效应与公理化 —— 有些数据，任何随机效用模型都解释不了
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
前几个文件里，IIA 失效的原因（相似的选项、混合人群）都还在"随机效用模型"（RUM）的框架之内：
换一个噪声结构（嵌套 logit、混合 logit）就能解释。但心理学家发现了一类更"顽固"的现象 ——
**情境效应**：人怎样评价一个选项，取决于旁边摆着什么。最著名的是

    · 诱饵效应（吸引效应，Huber, Payne & Puto 1982）：T 和 K 不分上下；加入一个明显比 T 差、
      但不比 K 差的"诱饵" D 以后，选 T 的人**变多了**。
    · 折中效应（Simonson 1989）：在一条"此消彼长"的线上，处在中间的选项更受欢迎。

诱饵效应违反了一条**所有** RUM 都满足的性质 —— 正则性：加入新选项，老选项被选中的概率不可能上升。
所以它不能用任何 RUM 解释，不管噪声长什么样。

    演示 1：正则性 —— 为什么所有 RUM 都满足它，并在几个 RUM 上验证。
    演示 2：诱饵效应与折中效应（用一个简单的"情境依赖"模型模拟），以及它们对标注数据的影响。
    演示 3：Block–Marschak 多项式 —— 判断一组选择概率能否用**某个** RUM 解释的完整条件（Falmagne 1978）。
    演示 4：一张总表：各种模型满足哪些公理。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_context_effects_axioms.py        # < 1 秒

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

import math

from rum import (NORMAL, all_subsets, block_marschak, choice_probs_integral, iia_ratios, logit_probs,
                 mixture_probs, nested_logit_probs, regularity_violations)
from toy_world import BY_KEY, make_expert, pad, print_header

EXPERT = make_expert()


# =============================================================================
# 几个选择模型（都用 prob_fn(options) → dict 的接口）
# =============================================================================
FOUR = ["A", "C", "D", "F"]
U4 = {k: EXPERT.utility(BY_KEY[k]) for k in FOUR}


def logit_model(options):
    """logit：专家对 4 个回答的效用。"""
    return logit_probs(U4, options)


def nested_model(options):
    """嵌套 logit：把 A、C（两个好的解释）放进一个巢，λ = 0.3。"""
    return nested_logit_probs(U4, [["A", "C"], ["D"], ["F"]], 0.3, options)


def mixture_model(options):
    """两群人的混合：一群人的效用是专家的，另一群人把 F（只有代码）看得很高。"""
    other = dict(U4)
    other["F"] = 6.0
    return mixture_probs([(0.7, lambda o: logit_probs(U4, o)), (0.3, lambda o: logit_probs(other, o))], options)


def probit_model(options):
    """独立正态噪声（Thurstone / probit），用数值积分精确计算。"""
    return choice_probs_integral(U4, options, NORMAL, n_grid=1500)


# 诱饵效应：选项有两个属性（内容深度、简洁程度），基础效用 = 3·深度 + 3·简洁
DECOY_ITEMS = {
    "T": (0.8, 0.3),     # 深入但啰嗦
    "K": (0.3, 0.8),     # 简洁但浅
    "D": (0.7, 0.2),     # T 的"劣化版"：比 T 浅、也比 T 啰嗦 → 被 T 支配；但不被 K 支配
    "M": (0.5, 0.5),     # 一个中规中矩的回答
}


def dominates(x, y):
    """x 是否"支配" y：两个属性都不比 y 差，至少一个严格更好。"""
    return all(a >= b for a, b in zip(x, y)) and any(a > b for a, b in zip(x, y))


def decoy_model(options, bonus=1.5):
    """
    一个简单的"情境依赖"模型：每支配选择集里的一个选项，效用就加 bonus（"有个明显更差的同类陪衬，显得我更好"）。

    【思路】给定选择集以后，它仍然是 logit；但效用**依赖于选择集**，所以它不是随机效用模型。
        这只是描述诱饵效应的众多模型里最简单的一个（心理学里还有相对优势模型、多属性决策场理论等）。
    【参数】options：选择集；bonus：每支配一个选项加多少效用
    【返回】dict {选项: 概率}
    """
    u = {}
    for k in options:
        base = 3.0 * DECOY_ITEMS[k][0] + 3.0 * DECOY_ITEMS[k][1]
        u[k] = base + bonus * sum(dominates(DECOY_ITEMS[k], DECOY_ITEMS[j]) for j in options if j != k)
    return logit_probs(u, options)


# 折中效应：选项都在一条"深度 + 简洁 = 1"的线上，基础效用相同
LINE = {"X": 0.9, "Y": 0.5, "Z": 0.1, "W": 0.7}     # 数字是"深度"，简洁 = 1 − 深度


def compromise_model(options, c=3.0):
    """
    "厌恶极端"的模型：效用 = 基础效用 − c × |自己的深度 − 选择集的平均深度|。处在中间的选项更受欢迎。

    【参数】options；c：厌恶极端的强度
    """
    center = sum(LINE[k] for k in options) / len(options)
    u = {k: 3.0 - c * abs(LINE[k] - center) for k in options}
    return logit_probs(u, options)


MODELS = [
    ("logit", logit_model, FOUR),
    ("嵌套 logit", nested_model, FOUR),
    ("两群人混合", mixture_model, FOUR),
    ("独立正态（probit）", probit_model, FOUR),
    ("诱饵效应模型", decoy_model, list(DECOY_ITEMS)),
    ("折中效应模型", compromise_model, list(LINE)),
]


# =============================================================================
# 演示 1：正则性
# =============================================================================
def demo_regularity():
    print_header("演示 1：正则性 —— 加入新选项，老选项被选中的概率只会下降（对所有随机效用模型）")
    print("""证明只有一句话：在随机效用模型里，j 在大集合 T 里感受最大，就一定在小集合 S ⊂ T 里感受最大。
所以“在 T 里选 j”这个事件包含在“在 S 里选 j”这个事件里，概率不会更大：P(j | S) ≥ P(j | T)。
下面在 4 个随机效用模型上检查所有的“S 加一个选项变成 T”的情形：\n""")
    for name, fn, items in MODELS[:4]:
        bad = regularity_violations(fn, items)
        # 情形数 = Σ_S |S| × (不在 S 里的选项数)：每个 S、每个可以加入的新选项、每个老选项 j 算一种情形
        cases = sum(len(sub) * (len(items) - len(sub)) for sub in all_subsets(items))
        print(f"  {pad(name, 20)} 检查了 {cases} 种情形，违反正则性的次数：{len(bad)}")


# =============================================================================
# 演示 2：诱饵效应与折中效应
# =============================================================================
def demo_context_effects():
    print_header("演示 2：诱饵效应与折中效应")
    print("诱饵效应：T（深入但啰嗦）和 K（简洁但浅）不分上下；D 是 T 的劣化版（更浅、更啰嗦）。")
    for opts in [("K", "T"), ("D", "K", "T")]:
        p = decoy_model(opts)
        print(f"  选择集 {{{', '.join(opts)}}}：" + "  ".join(f"{k}:{p[k]:.3f}" for k in opts))
    print("  → 加入诱饵 D 以后，T 的份额从 0.500 **上升**到上面的值：违反了正则性。")
    bad = regularity_violations(decoy_model, list(DECOY_ITEMS))
    print(f"  在 4 个选项的全部子集上，违反正则性的情形共 {len(bad)} 个，例如：")
    for j, s, t, ps, pt in bad[:3]:
        print(f"    P({j} | {{{','.join(s)}}}) = {ps:.3f}  <  P({j} | {{{','.join(t)}}}) = {pt:.3f}")

    print("\n折中效应：X（很深入）、Y（中等）、Z（很简洁）在一条此消彼长的线上，基础效用相同。")
    for opts in [("X", "Y"), ("X", "Y", "Z")]:
        p = compromise_model(opts)
        print(f"  选择集 {{{', '.join(opts)}}}：" + "  ".join(f"{k}:{p[k]:.3f}" for k in opts)
              + f"    P(Y)/P(X) = {p['Y'] / p['X']:.2f}")
    print("""
观察：
  · 诱饵效应：D 明显不如 T（两方面都差），没有人会选它，可它的存在让 T “显得更好”。T 的份额不降反升，
    这违反了正则性，因此**没有任何**随机效用模型能产生这样的数据 —— 换噪声、加混合都没用。
  · 折中效应：只有 X、Y 时两者打平；加入 Z 以后，Y 变成了“中间选项”，份额大涨，P(Y)/P(X) 从 1 变成 3 以上（违反 IIA），
    在这个模型里 Y 的份额也绝对地上升了（违反正则性）。
  · 对偏好标注的启示：同时给标注员看几个回答时，回答之间会互相“衬托”。同一个模型采样出的回答里常有
    “某个回答的劣化版”；把这样的排序数据当成 Plackett–Luce 去拟合，会系统性地高估“有个更差的兄弟”的回答。
    两两比较能减轻（但不能消除）这种效应：选择集最小，衬托的机会最少。
  · 注意：这里的情境依赖模型只是为了演示而设计的。真实的诱饵效应在实验室里很稳定，在真实世界里的大小有争议；
    它在 LLM 标注里有多强，是一个值得研究的问题。""")


# =============================================================================
# 演示 3：Block–Marschak 多项式
# =============================================================================
def demo_block_marschak():
    print_header("演示 3：Block–Marschak 多项式 —— 能否用“某个”随机效用模型解释？")
    print("""对选项 j 和集合 S（j ∈ S），K(j, S) = Σ_{T ⊇ S} (−1)^{|T∖S|} · P(j | T)。
它等于“S 之外的选项全都排在 S 前面、而 j 是 S 里最好的”的概率，所以随机效用模型一定有 K ≥ 0。
Falmagne（1978）证明反过来也成立：所有 K ≥ 0 ⇔ 存在一个随机效用模型能产生这些选择概率。\n""")
    for name, fn, items in MODELS:
        bm = block_marschak(fn, items)
        worst = min(bm.items(), key=lambda kv: kv[1])
        neg = sum(1 for v in bm.values() if v < -1e-9)
        (j, s), v = worst
        print(f"  {pad(name, 20)} 多项式个数 {len(bm):>3d}，负的个数 {neg:>2d}，最小值 K({j}, {{{','.join(s)}}}) = {v:+.4f}")
    print("""
观察：
  · 四个随机效用模型的所有多项式都 ≥ 0（最小值是 0 或正数，只差浮点误差）。
  · 两个情境依赖模型都有负的多项式，所以它们产生的数据不能被任何随机效用模型解释。
  · 正则性只是 |T∖S| = 1 的那一条多项式；Block–Marschak 条件是它的完整推广。
    这就是“公理化”的意义：把一个模型类**等价地**写成一组可以直接在数据上检验的条件。
    同样的思路：Luce 的选择公理 ⇔ logit；正则性 + 更多条件 ⇔ 随机效用模型。""")


# =============================================================================
# 演示 4：总表
# =============================================================================
def satisfies_iia(fn, items, tol=1e-6):
    """检查所有两两比值在所有选择集上是否相同。"""
    for i, a in enumerate(items):
        for b in items[i + 1:]:
            ratios = list(iia_ratios(fn, a, b, items).values())
            if max(ratios) - min(ratios) > tol * max(ratios):
                return False
    return True


def demo_summary():
    print_header("演示 4：总表 —— 各种模型满足哪些公理")
    print(f"  {pad('模型', 20)} {pad('IIA', 6, 'right')} {pad('正则性', 8, 'right')} {pad('Block–Marschak', 16, 'right')}  "
          "能否用某个随机效用模型解释（Falmagne：等于上一列）")
    for name, fn, items in MODELS:
        iia = satisfies_iia(fn, items)
        reg = len(regularity_violations(fn, items)) == 0
        bm = min(block_marschak(fn, items).values()) >= -1e-9
        mark = lambda b: "是" if b else "否"
        print(f"  {pad(name, 20)} {pad(mark(iia), 6, 'right')} {pad(mark(reg), 8, 'right')} {pad(mark(bm), 16, 'right')}  {mark(bm)}")
    print("""
观察：
  · 层级关系：logit ⊂ 随机效用模型（RUM）⊂ 所有选择模型。
    IIA 只有 logit 满足；正则性和 Block–Marschak 条件所有 RUM 都满足；情境效应模型连这些都不满足。
  · 对应到建模的选择：
      满足 IIA            → BT / PL / softmax，参数最少，可以做反事实预测，DPO 和大多数奖励模型都在这里；
      不满足 IIA 但仍是 RUM → 嵌套 logit、混合 logit、probit，更灵活，要更多数据、更难估计；
      连 RUM 都不是        → 需要显式地建模“选择集怎样影响评价”（情境依赖模型），这是行为经济学和心理学的领域（第 15 讲）。""")


def main():
    demo_regularity()
    demo_context_effects()
    demo_block_marschak()
    demo_summary()


if __name__ == "__main__":
    main()