"""
=============================================================================
第 3 讲 · 工具模块：随机效用模型（RUM）的计算工具箱
=============================================================================

对应讲义：lectures/lecture03_random_utility.md

-----------------------------------------------------------------------------
【这个模块解决什么问题？】
-----------------------------------------------------------------------------
随机效用模型（Random Utility Model，RUM）说：人对每个选项 j 的感受是

        H_j = V_j + ε_j          （V_j：平均效用，ε_j：随机噪声）

每次选择时选感受最大的那个。不同的 RUM 只在"噪声长什么样"上不同：
    · 独立的 Gumbel 噪声   → logit（softmax、Bradley–Terry、Plackett–Luce）
    · 独立的正态噪声       → probit（Thurstone 1927）
    · 相关的噪声           → 嵌套 logit、相关 probit（第 5 节：红色公交车 / 蓝色公交车）
    · 人群里每个人的 V 不同 → 混合 logit（第 5.3 节）

这个模块提供计算这些模型的选择概率的工具，以及检验"一个选择模型满足哪些公理"的工具：
    · choice_probs_integral：独立噪声时，用一维数值积分**精确**计算 P(j | S)
    · logit_probs / nested_logit_probs / mixture_probs：几种有公式的模型
    · mc_choice_probs：任意噪声（包括相关噪声）用蒙特卡洛模拟
    · iia_ratios / regularity_violations / block_marschak：三种"公理检验"

所有"选择模型"都用同一个接口：一个函数 prob_fn(options) → dict {选项: 概率}。
这样同一套检验代码可以作用在任何模型上。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from rum import GUMBEL, choice_probs_integral, logit_probs
    utils = {"a": 1.0, "b": 0.0, "c": -0.5}
    print(choice_probs_integral(utils, ["a", "b", "c"], GUMBEL))   # 数值积分
    print(logit_probs(utils, ["a", "b", "c"]))                      # softmax 公式（应该相同）

直接运行本文件会打印上面这个对比：
    python3 rum.py
"""

import itertools
import math
import os
import random
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
_L1 = os.path.join(_HERE, "..", "lecture01")
if _L1 not in sys.path:
    sys.path.insert(0, _L1)

from toy_world import EULER_GAMMA, mean, sample_gumbel, sigmoid  # noqa: E402


# =============================================================================
# 第一部分：几种噪声分布（每种都有 密度 pdf、累积分布 cdf、抽样 sample）
# =============================================================================
class Noise:
    """
    一个一维噪声分布。

    【字段】
        name   ：名字（打印用）
        pdf    ：概率密度函数 f(x)
        cdf    ：累积分布函数 F(x) = P(ε ≤ x)
        sample ：给一个 random.Random，返回一个随机数
        lo, hi ：数值积分时的范围（分布几乎全部的概率都在 [lo, hi] 里）
    【边界】我们只用"平均效用 + 噪声"的加法形式，所以噪声的平均值是多少不重要（所有选项加同一个常数，
        选择不变）；重要的是它的**形状**和**尺度**。
    """

    def __init__(self, name, pdf, cdf, sample, lo, hi):
        self.name, self.pdf, self.cdf, self.sample, self.lo, self.hi = name, pdf, cdf, sample, lo, hi


# 标准 Gumbel：F(x) = exp(−e^{−x})，平均值 γ ≈ 0.577，方差 π²/6 ≈ 1.645
GUMBEL = Noise(
    "Gumbel",
    pdf=lambda x: math.exp(-x - math.exp(-x)) if x > -30 else 0.0,
    cdf=lambda x: math.exp(-math.exp(-x)) if x > -30 else 0.0,
    sample=sample_gumbel,
    lo=-6.0, hi=25.0,
)

_SD = math.pi / math.sqrt(6)      # 标准 Gumbel 的标准差，其余噪声都取这个标准差，方便比较

# 正态（Thurstone / probit）：均值 0，标准差和 Gumbel 相同
NORMAL = Noise(
    "正态",
    pdf=lambda x: math.exp(-0.5 * (x / _SD) ** 2) / (_SD * math.sqrt(2 * math.pi)),
    cdf=lambda x: 0.5 * (1.0 + math.erf(x / (_SD * math.sqrt(2)))),
    sample=lambda rng: rng.gauss(0.0, _SD),
    lo=-10 * _SD, hi=10 * _SD,
)

# 均匀分布：[−a, a]，a = √3·标准差（这样方差相同）
_A = math.sqrt(3) * _SD
UNIFORM = Noise(
    "均匀",
    pdf=lambda x: 1.0 / (2 * _A) if -_A <= x <= _A else 0.0,
    cdf=lambda x: 0.0 if x < -_A else (1.0 if x > _A else (x + _A) / (2 * _A)),
    sample=lambda rng: rng.uniform(-_A, _A),
    lo=-_A, hi=_A,
)

# "反过来的 Gumbel"：−G，G 是标准 Gumbel。形状和 Gumbel 一样，只是左右翻转（长尾巴在左边）
REVERSED_GUMBEL = Noise(
    "反向Gumbel",
    pdf=lambda x: GUMBEL.pdf(-x),
    cdf=lambda x: 1.0 - GUMBEL.cdf(-x),
    sample=lambda rng: -sample_gumbel(rng),
    lo=-25.0, hi=6.0,
)

ALL_NOISES = [GUMBEL, NORMAL, UNIFORM, REVERSED_GUMBEL]


# =============================================================================
# 第二部分：选择概率
# =============================================================================
def choice_probs_integral(utils, options, noise, n_grid=6000):
    """
    独立噪声的 RUM：用一维数值积分精确计算 P(j | S)。

    【思路】j 被选中 ⇔ 对所有 k ≠ j，V_k + ε_k < V_j + ε_j。设 j 的感受 H_j = t，那么
        每个 k 的条件是 ε_k < t − V_k，概率是 F(t − V_k)；各个 k 独立，所以相乘。再对 t 积分：
            P(j | S) = ∫ f(t − V_j) · Π_{k∈S, k≠j} F(t − V_k) dt
        这是一维积分，用"梯形法"（把曲线下的面积切成很多细长的梯形加起来）就能算得很准。
    【参数】
        utils  ：dict {选项: 平均效用 V}
        options：选择集 S（列表或元组）
        noise  ：Noise 对象
        n_grid ：积分的格点数（越多越准、越慢）
    【返回】dict {选项: 概率}
    【边界】
        · 只适用于**独立**噪声。相关噪声（嵌套 logit、相关 probit）请用 mc_choice_probs 或专门的公式。
        · 均匀噪声的密度在边界处不连续，梯形法的误差会大一点（约 1e-4），对演示足够。
    """
    vals = [utils[k] for k in options]
    lo = min(vals) + noise.lo
    hi = max(vals) + noise.hi
    h = (hi - lo) / n_grid
    out = {}
    for j in options:
        total = 0.0
        for i in range(n_grid + 1):
            t = lo + i * h
            dens = noise.pdf(t - utils[j])
            if dens == 0.0:
                continue
            prod = dens
            for k in options:
                if k != j:
                    prod *= noise.cdf(t - utils[k])
            weight = 0.5 if i in (0, n_grid) else 1.0
            total += weight * prod
        out[j] = total * h
    z = math.fsum(out.values())               # 修正积分的微小误差，让概率之和恰好为 1
    return {k: v / z for k, v in out.items()}


def logit_probs(utils, options):
    """
    多项 logit（softmax）：P(j | S) = e^{V_j} / Σ_{k∈S} e^{V_k}。独立 Gumbel 噪声的封闭解。

    【参数】utils：dict；options：选择集
    【返回】dict {选项: 概率}
    """
    top = max(utils[k] for k in options)
    w = {k: math.exp(utils[k] - top) for k in options}
    z = math.fsum(w.values())
    return {k: v / z for k, v in w.items()}


def nested_logit_probs(utils, nests, lam, options):
    """
    嵌套 logit：把相似的选项放进同一个"巢"，巢内的噪声是相关的。

    【思路】先选巢、再在巢里选：
            巢 n 的"包容值"  I_n = log Σ_{k∈n} e^{V_k / λ}
            P(选巢 n)        = e^{λ·I_n} / Σ_m e^{λ·I_m}
            P(j | 巢 n)      = e^{V_j / λ} / e^{I_n}
        λ ∈ (0, 1] 衡量巢内选项"有多不相似"：
            λ = 1 → 巢内噪声不相关，退化成普通的 logit；
            λ → 0 → 巢内噪声完全相关，巢里的选项像是"同一个选项的几个克隆"。
        这正好修复了红色公交车 / 蓝色公交车问题（讲义第 5.2 节）。
    【参数】
        utils  ：dict {选项: V}
        nests  ：巢的列表，每个巢是选项列表，例如 [["汽车"], ["蓝巴士", "红巴士"]]
        lam    ：λ，0 < λ ≤ 1（所有巢共用一个 λ，这是最简单的版本）
        options：这次的选择集（巢里不在选择集里的选项会被忽略）
    【返回】dict {选项: 概率}
    【边界】λ > 1 时模型可能和随机效用理论不一致（会出现违反正则性的情况），这里不检查。
    """
    present = [[k for k in n if k in options] for n in nests]
    present = [n for n in present if n]
    incl = [math.log(math.fsum(math.exp(utils[k] / lam) for k in n)) for n in present]
    top = max(lam * v for v in incl)
    nest_w = [math.exp(lam * v - top) for v in incl]
    z = math.fsum(nest_w)
    out = {}
    for n, iv, w in zip(present, incl, nest_w):
        for k in n:
            out[k] = (w / z) * math.exp(utils[k] / lam - iv)
    return out


def mixture_probs(components, options):
    """
    有限混合模型：人群由几类人组成，第 c 类占比 w_c、有自己的选择模型。

        P(j | S) = Σ_c w_c · P_c(j | S)

    【参数】components：list[(权重, prob_fn)]，权重之和为 1；options：选择集
    【返回】dict {选项: 概率}
    【边界】每一类人各自满足 IIA，混在一起一般就不满足了（第 5.3 节）。
    """
    out = {k: 0.0 for k in options}
    for w, fn in components:
        p = fn(options)
        for k in options:
            out[k] += w * p[k]
    return out


def mc_choice_probs(sample_feel, options, draws, rng):
    """
    蒙特卡洛：反复抽"感受"，数每个选项被选中的比例。适用于任何噪声（包括相关噪声）。

    【参数】
        sample_feel：函数 (options, rng) → dict {选项: 这一次的感受}
        options    ：选择集
        draws      ：抽多少次
        rng        ：random.Random
    【返回】dict {选项: 比例}
    【边界】有随机误差，约为 √(p(1−p)/draws)；draws = 40000 时约 ±0.0025。
    """
    counts = {k: 0 for k in options}
    for _ in range(draws):
        feel = sample_feel(options, rng)
        counts[max(options, key=lambda k: feel[k])] += 1
    return {k: c / draws for k, c in counts.items()}


def cholesky(cov):
    """
    Cholesky 分解：把对称正定矩阵 Σ 写成 L·Lᵀ（L 是下三角矩阵）。

    【思路】要抽"协方差为 Σ 的正态随机向量"，先抽一个各分量独立的标准正态向量 z，再算 L·z 即可：
        Cov(Lz) = L·Cov(z)·Lᵀ = L·I·Lᵀ = Σ。
    【参数】cov：d×d 的列表的列表，必须对称正定
    【返回】L（列表的列表）
    【边界】Σ 不是正定时（例如两个变量完全相关，ρ = 1）会对负数开方而报错；这时可以用 ρ = 0.999 近似。
    """
    d = len(cov)
    L = [[0.0] * d for _ in range(d)]
    for i in range(d):
        for j in range(i + 1):
            s = cov[i][j] - math.fsum(L[i][k] * L[j][k] for k in range(j))
            L[i][j] = math.sqrt(s) if i == j else s / L[j][j]
    return L


def correlated_normal_sampler(utils, keys, cov):
    """
    构造一个"相关正态噪声"的感受抽样函数（相关 probit 模型），供 mc_choice_probs 使用。

    【参数】utils：dict {选项: V}；keys：协方差矩阵各行对应的选项；cov：协方差矩阵
    【返回】函数 (options, rng) → dict {选项: 感受}
    """
    L = cholesky(cov)

    def sample(options, rng):
        z = [rng.gauss(0.0, 1.0) for _ in keys]
        eps = [math.fsum(L[i][k] * z[k] for k in range(i + 1)) for i in range(len(keys))]
        feel = {k: utils[k] + e for k, e in zip(keys, eps)}
        return {k: feel[k] for k in options}
    return sample


# =============================================================================
# 第三部分：公理检验
# =============================================================================
def all_subsets(items, min_size=1):
    """
    列出 items 的所有（大小 ≥ min_size 的）子集，每个子集是排好序的元组。

    【Python 小知识】itertools.combinations(items, r)：所有大小为 r 的组合（不计顺序）。
    """
    items = sorted(items)
    return [c for r in range(min_size, len(items) + 1) for c in itertools.combinations(items, r)]


def iia_ratios(prob_fn, a, b, items):
    """
    IIA 检验：在每个同时包含 a、b 的选择集上，算 P(a | S) / P(b | S)。满足 IIA 时，这个比值处处相同。

    【参数】prob_fn：选择模型；a、b：两个选项；items：全部选项
    【返回】dict {选择集: 比值}
    """
    out = {}
    for s in all_subsets(items, 2):
        if a in s and b in s:
            p = prob_fn(s)
            out[s] = p[a] / p[b]
    return out


def regularity_violations(prob_fn, items, tol=1e-9):
    """
    正则性检验：对任意 S ⊂ T 和 j ∈ S，必须有 P(j | S) ≥ P(j | T)（加入新选项不会让老选项更受欢迎）。

    【思路】**所有**随机效用模型都满足正则性：j 在 T 里最好，就一定在更小的 S 里最好，
        所以"在 T 里被选中"的事件包含于"在 S 里被选中"的事件。（讲义第 6.1 节）
        所以，如果数据违反正则性，它就不能用**任何**随机效用模型解释。
    【参数】prob_fn；items；tol：容差
    【返回】list[(j, S, T, P(j|S), P(j|T))]，所有违反的情形（只检查 T = S 加一个选项）
    """
    bad = []
    for s in all_subsets(items, 1):
        ps = prob_fn(s)
        for extra in sorted(set(items) - set(s)):
            t = tuple(sorted(s + (extra,)))
            pt = prob_fn(t)
            for j in s:
                if pt[j] > ps[j] + tol:
                    bad.append((j, s, t, ps[j], pt[j]))
    return bad


def block_marschak(prob_fn, items):
    """
    Block–Marschak 多项式：K(j, S) = Σ_{T ⊇ S} (−1)^{|T \\ S|} · P(j | T)。

    【思路】直觉：K(j, S) 是"在随机的排序里，S 恰好是排在最后的那几个选项、而 j 是其中最好的"的概率
        ——更准确地说，是"排在 S 之外的选项全都比 S 里的好，且 j 是 S 里最好的"的概率。
        概率不能是负数，所以随机效用模型一定满足 K(j, S) ≥ 0。
        Falmagne（1978）证明了反过来也对：一组选择概率能用**某个**随机效用模型解释，
        当且仅当所有的 Block–Marschak 多项式都 ≥ 0。这是"随机效用模型"的一个完整的公理化。
        正则性只是其中最简单的一条（|T \\ S| = 1 的情形）。
    【参数】prob_fn；items：全部选项
    【返回】dict {(j, S): K(j, S)}
    【边界】选项数为 M 时要算所有子集（2^M 个），M 很大时不可行；这里 M ≤ 5。
    """
    items = sorted(items)
    cache = {s: prob_fn(s) for s in all_subsets(items, 1)}
    out = {}
    for s in all_subsets(items, 1):
        rest = [k for k in items if k not in s]
        for j in s:
            total = 0.0
            for r in range(len(rest) + 1):
                for extra in itertools.combinations(rest, r):
                    t = tuple(sorted(s + extra))
                    total += (-1) ** r * cache[t][j]
            out[(j, s)] = total
    return out


def fmt_probs(p, keys=None):
    """把概率字典格式化成 "a:0.33  b:0.67" 的样子。"""
    keys = keys or list(p)
    return "  ".join(f"{k}:{p[k]:.3f}" for k in keys)


if __name__ == "__main__":
    utils = {"a": 1.0, "b": 0.0, "c": -0.5}
    opts = ["a", "b", "c"]
    print("平均效用：", utils)
    print("Gumbel 噪声，数值积分：", fmt_probs(choice_probs_integral(utils, opts, GUMBEL)))
    print("softmax 公式：          ", fmt_probs(logit_probs(utils, opts)))
    print("正态噪声，数值积分：    ", fmt_probs(choice_probs_integral(utils, opts, NORMAL)))