"""
=============================================================================
第 2 讲 · 工具模块：偏好数据的统一表示 —— 一切都是"从集合中选一个"
=============================================================================

对应讲义：lectures/lecture02_choice_data_bt.md 第 3、4、5 节

-----------------------------------------------------------------------------
【这个模块解决什么问题？】
-----------------------------------------------------------------------------
偏好数据有很多种样子：完整的排序、从几个里挑一个、两两比较、接受 / 拒绝……
如果每种数据各写一套模型和拟合代码，很快就乱了。这个模块用一个统一的视角把它们收拢：

    **任何一条偏好数据，都可以拆成若干个"选择事件"：从集合 S 里选出了 j。**

        两两比较  a ≻ b              →  1 个事件：从 {a, b} 里选了 a
        多选一    从 {a,b,c,d} 选 b   →  1 个事件：从 {a, b, c, d} 里选了 b
        完整排序  a ≻ c ≻ b ≻ d       →  3 个事件：从 {a,b,c,d} 选 a；从 {b,c,d} 选 c；从 {b,d} 选 b
                                         （最后剩一个不用选。这种拆法叫"爆炸"（exploded）排序）
        接受 / 拒绝  接受 a           →  1 个事件：从 {a, ∅} 里选了 a（∅ 是"什么都不要"这个外部选项）
                     拒绝 a           →  1 个事件：从 {a, ∅} 里选了 ∅

然后只需要**一个**模型：每个选项一个分数 s，从 S 里选 j 的概率是 softmax：

        P(j | S) = e^{s_j} / Σ_{k∈S} e^{s_k}

两个选项时它就是 Bradley–Terry 的 σ(s_a − s_b)；排序时它就是 Plackett–Luce。
（为什么是 softmax？第 3 讲会证明它等价于"效用 + 独立的 Gumbel 噪声"。）

模块里有：
    · 数据形式之间的转换（ranking_to_events、pair_to_event、choice_to_pairs ……）
    · 按"效用 + Gumbel 噪声"模拟各种数据（sample_ranking、sample_choice ……）
    · 一个能拟合**所有**数据形式的最大似然拟合器 fit_choice_model
    · 用特征给分数建模的"线性 BT"拟合器 fit_linear_bt（第 5 节：泛化到没见过的选项）
    · 几个小工具：解线性方程组、数值梯度

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from choice_data import ranking_to_events, fit_choice_model
    events = ranking_to_events(["A", "C", "B"]) + ranking_to_events(["C", "A", "B"])
    scores = fit_choice_model(events, ["A", "B", "C"])

直接运行本文件会演示几种转换：
    python3 choice_data.py
"""

import math
import os
import random
import sys
from collections import Counter

# 让 Python 能找到第 1 讲的模块（toy_world.py）。
# 【Python 小知识】sys.path 是 Python 查找模块的目录列表；把第 1 讲的目录加进去，
#     就可以像同目录的文件一样 import 它。os.path.abspath(__file__) 是本文件的完整路径，
#     os.path.dirname 取它所在的目录 —— 这样不管你从哪个目录运行脚本，路径都是对的。
_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 mean, sample_gumbel, sigmoid  # noqa: E402  （import 不在文件开头是故意的）

# 外部选项（outside option）的名字："什么都不选""都不要"。接受 / 拒绝数据要用它。
OUTSIDE = "∅"


# =============================================================================
# 第一部分：数据形式之间的转换
# =============================================================================
# 一个"选择事件"用一个二元组 (赢家, 选项集合) 表示，选项集合是**排好序的元组**。
# 为什么排序？因为 {a, b} 和 {b, a} 是同一个集合；排好序以后，相同的事件长得完全一样，
# 可以用 Counter 计数合并（见 fit_choice_model），大大加快计算。
# 【Python 小知识】元组 (tuple) 和列表很像，但不能修改，因此可以当作字典的键；列表不行。

def make_event(winner, options):
    """
    构造一个选择事件 (winner, 排好序的选项元组)。

    【参数】winner：被选中的选项；options：可选的选项（列表、元组或集合都行），必须包含 winner
    【返回】(winner, tuple)
    【边界】winner 不在 options 里时抛出 ValueError —— 这样的数据一定是哪里弄错了。
    """
    opts = tuple(sorted(set(options)))
    if winner not in opts:
        raise ValueError(f"赢家 {winner} 不在选项 {opts} 里")
    return (winner, opts)


def pair_to_event(winner, loser):
    """两两比较 winner ≻ loser → 一个事件：从 {winner, loser} 里选了 winner。"""
    return make_event(winner, (winner, loser))


def choice_to_event(winner, options):
    """多选一 → 一个事件（就是 make_event，起个名字让代码更好读）。"""
    return make_event(winner, options)


def ranking_to_events(ranking):
    """
    完整排序 → "爆炸"成 len−1 个选择事件。

    【思路】把排序看成"一个一个地挑"：先从全部里挑出第一名，再从剩下的里挑第二名……
        例：["A", "C", "B", "D"] → 从 {A,B,C,D} 选 A；从 {B,C,D} 选 C；从 {B,D} 选 B。
        在 softmax 模型下，这些事件的概率相乘就是 Plackett–Luce 的排序概率。
    【参数】ranking：列表，从最好到最差，元素不能重复
    【返回】list[事件]
    【边界】只有 1 个元素的"排序"没有任何信息，返回空列表。
    """
    return [make_event(ranking[i], ranking[i:]) for i in range(len(ranking) - 1)]


def ranking_to_pairs(ranking):
    """
    完整排序 → 所有的两两比较（排 K 个得到 K(K−1)/2 对）。InstructGPT 就是这样把排序变成训练数据的。

    【边界】这些"对"之间**不是独立的**（都来自同一次排序），把它们当成独立的比较时，
        估计值仍然合理，但会高估自己有多确定（讲义第 3.6 节）。
    """
    return [(ranking[i], ranking[j]) for i in range(len(ranking)) for j in range(i + 1, len(ranking))]


def choice_to_pairs(winner, options):
    """多选一 → winner 胜过其余每一个（|S|−1 对）。【边界】同样不是独立的。"""
    return [(winner, k) for k in sorted(options) if k != winner]


def accept_reject_to_event(item, accepted):
    """
    接受 / 拒绝 → 和外部选项 ∅ 的一次比较。

    【思路】"接受 a"可以理解成"a 比什么都不要好"，"拒绝 a"就是"什么都不要比 a 好"。
        这样，外部选项 ∅ 就成了一个分数**固定为 0** 的参照物，其他选项的分数就有了"绝对"的含义：
        s_a > 0 表示"大多数时候会被接受"。（讲义第 3.4 节）
    【参数】item：选项；accepted：bool
    【返回】一个事件
    """
    return make_event(item if accepted else OUTSIDE, (item, OUTSIDE))


# =============================================================================
# 第二部分：按"效用 + Gumbel 噪声"模拟数据
# =============================================================================
def perceive_all(utils, keys, rng):
    """
    一次"读完 keys 里的所有选项"的感受：每个选项的效用 + 一个独立的 Gumbel 噪声。

    【参数】utils：dict {选项: 效用}（外部选项 ∅ 不在里面时，当作效用 0）；keys：要读的选项；rng：random.Random
    【返回】dict {选项: 感受}
    """
    return {k: utils.get(k, 0.0) + sample_gumbel(rng) for k in keys}


def sample_ranking(utils, keys, rng):
    """
    模拟一次完整排序：按感受从大到小排。

    【思路】第 3 讲会证明：这样排出来的顺序，恰好服从 Plackett–Luce 分布。所以这是"精确"的模拟，
        不需要一步一步地按 softmax 抽。
    【返回】列表，从最好到最差
    【Python 小知识】sorted(keys, key=lambda k: -feel[k])：按感受从大到小排序。
    """
    feel = perceive_all(utils, keys, rng)
    return sorted(keys, key=lambda k: -feel[k])


def sample_choice(utils, options, rng):
    """模拟一次多选一：选感受最大的那个。【返回】被选中的选项"""
    feel = perceive_all(utils, options, rng)
    return max(sorted(options), key=lambda k: feel[k])


def sample_pair(utils, a, b, rng):
    """模拟一次两两比较。【返回】(赢家, 输家)"""
    return (a, b) if sample_choice(utils, (a, b), rng) == a else (b, a)


def sample_accept(utils, item, rng):
    """模拟一次接受 / 拒绝（和效用为 0 的外部选项比较）。【返回】bool：True 表示接受"""
    return sample_choice(utils, (item, OUTSIDE), rng) == item


# =============================================================================
# 第三部分：softmax 选择模型的概率、似然与拟合
# =============================================================================
def choice_prob(scores, winner, options):
    """
    softmax 选择概率 P(winner | options) = e^{s_w} / Σ_{k∈options} e^{s_k}。

    【参数】scores：dict {选项: 分数}（外部选项 ∅ 不在里面时当作 0）；winner；options
    【返回】0~1 的 float
    【边界】先减去最大分数再求 exp，防止溢出（分子分母同乘一个常数，结果不变）。
    """
    vals = {k: scores.get(k, 0.0) for k in options}
    top = max(vals.values())
    z = math.fsum(math.exp(v - top) for v in vals.values())
    return math.exp(vals[winner] - top) / z


def average_log_likelihood(events, scores):
    """
    一批选择事件的平均对数似然（越接近 0 越好）。

    【参数】events：list[事件]；scores：dict
    【返回】float（≤ 0）。两两比较时，log 0.5 ≈ −0.693 相当于瞎猜。
    """
    return mean([math.log(choice_prob(scores, w, opts)) for w, opts in events])


def fit_choice_model(events, keys, l2=1e-3, max_iters=20000, tol=1e-8, fixed=None):
    """
    用最大似然拟合 softmax 选择模型 —— 一个函数处理所有数据形式。

    【思路】目标 = 平均每个事件的 log P(赢家 | 选项)  −  (l2/2) × Σ s²（只对没固定的分数）。
        对分数 s_k 求导（第 5 讲详细推导）：
            ∂目标/∂s_k = (1/N) Σ_事件 [ 1{k 是赢家} − P(k | 选项) ] − l2·s_k   （只对 k 在选项里的事件求和）
        直觉："被选中的次数"减去"模型预期它被选中的次数"。选中得比预期多 → 分数往上调。
        两个选项时，这和第 1 讲 fit_bradley_terry 的梯度完全相同。
        计算前先用 Counter 把相同的事件合并计数：1 万条两两比较里，8 个选项最多只有 56 种不同的事件，
        合并以后每一步只需要算 56 次，快几百倍。
    【参数】
        events   ：list[事件]
        keys     ：要估计分数的选项列表
        l2       ：L2 正则化强度（≥ 0）
        max_iters：最多迭代步数
        tol      ：梯度的最大分量小于它就停止
        fixed    ：dict {选项: 固定的分数}，例如 {OUTSIDE: 0.0}。有固定分数时**不再**把平均值移到 0，
                   因为固定的那个选项已经提供了"零点"。
    【返回】dict {选项: 分数}（包含 fixed 里的选项）
    【边界】
        · 没有固定分数时，只有分数的差值有意义，返回的分数平均值为 0。
        · 步长用光滑度上界自动设定：L ≤ max_k（k 出现的事件比例）/ 2 + l2，步长 1/L，保证目标不下降。
        · 数据很少、或者比较图不连通时，结果主要由正则化决定（第 4.2 节）。
    【Python 小知识】Counter(列表)：数每个元素出现了几次，返回一个特殊的字典 {元素: 次数}。
        .items() 依次给出 (元素, 次数)。sorted(...) 让遍历顺序固定，保证结果可复现。
    """
    fixed = dict(fixed or {})
    free = [k for k in keys if k not in fixed]
    if not events:
        return {**{k: 0.0 for k in free}, **fixed}
    counts = sorted(Counter(events).items())
    n = len(events)
    appear = Counter()
    for (w, opts), c in counts:
        for k in opts:
            appear[k] += c
    lr = 1.0 / (max(appear[k] for k in free) / (2.0 * n) + l2)

    s = {k: 0.0 for k in free}
    s.update(fixed)
    for _ in range(max_iters):
        grad = {k: 0.0 for k in free}
        for (w, opts), c in counts:
            vals = [s.get(k, 0.0) for k in opts]
            top = max(vals)
            exps = [math.exp(v - top) for v in vals]
            z = math.fsum(exps)
            for k, e in zip(opts, exps):
                if k in grad:
                    grad[k] += c * ((1.0 if k == w else 0.0) - e / z)
        biggest = 0.0
        for k in free:
            grad[k] = grad[k] / n - l2 * s[k]
            biggest = max(biggest, abs(grad[k]))
        if biggest < tol:
            break
        for k in free:
            s[k] += lr * grad[k]

    if not fixed:
        center = mean([s[k] for k in free])
        for k in free:
            s[k] -= center
    return {k: s[k] for k in list(free) + list(fixed)}


# =============================================================================
# 第四部分：线性 BT（分数 = 特征的加权和）—— 用于泛化到没见过的选项
# =============================================================================
def dot(u, v):
    """两个等长列表的点积 Σ u_i v_i。"""
    return math.fsum(a * b for a, b in zip(u, v))


def solve_linear_system(a, b):
    """
    解线性方程组 A x = b（高斯消元法，带部分主元）。

    【思路】中学学过的"消元法"：用第 1 行把下面各行的第 1 列消成 0，再用第 2 行消第 2 列……
        最后从下往上回代。"部分主元"是指每次挑这一列绝对值最大的那一行来消元，避免除以一个很小的数，
        让结果更准。numpy.linalg.solve 做的是同样的事（第 5 讲起我们直接用它）。
    【参数】a：d×d 的列表的列表；b：长度 d 的列表
    【返回】长度 d 的列表 x
    【边界】A 奇异（不可逆）时会除以 0 或得到巨大的数；调用者应保证 A 可逆（例如加上 L2 正则化）。
        不会修改传进来的 a、b（内部先复制一份）。
    """
    d = len(b)
    m = [list(row) + [b[i]] for i, row in enumerate(a)]      # 增广矩阵 [A | b]
    for col in range(d):
        piv = max(range(col, d), key=lambda r: abs(m[r][col]))
        m[col], m[piv] = m[piv], m[col]
        for r in range(col + 1, d):
            f = m[r][col] / m[col][col]
            for c in range(col, d + 1):
                m[r][c] -= f * m[col][c]
    x = [0.0] * d
    for r in range(d - 1, -1, -1):
        x[r] = (m[r][d] - math.fsum(m[r][c] * x[c] for c in range(r + 1, d))) / m[r][r]
    return x


def fit_linear_bt(diffs, l2=1e-3, max_iters=50, tol=1e-10):
    """
    拟合线性 Bradley–Terry 模型：s(y) = w·φ(y)，P(a ≻ b) = σ(w·(φ(a) − φ(b)))。

    【思路】每条比较只用到"赢家的特征 − 输家的特征"这个差向量 x，模型是 P(观察到的结果) = σ(w·x)。
        这正是**没有截距**的逻辑回归（标签全是 1）。用牛顿法求解（每一步用二阶导数决定走多远，
        通常十步以内就收敛；第 5 讲会比较各种优化方法）：
            梯度   g = (1/N) Σ (1 − p_n) x_n − l2·w,           p_n = σ(w·x_n)
            海森   H = −(1/N) Σ p_n(1 − p_n) x_n x_nᵀ − l2·I
            更新   w ← w − H⁻¹ g
        为什么没有截距？因为截距 c 对每个选项都一样，在"赢家 − 输家"里消掉了 —— 和"只有差值有意义"
        是同一回事（讲义第 5.2 节）。
    【参数】
        diffs    ：list[list[float]]，每条比较的特征差向量（赢家 − 输家）
        l2       ：L2 正则化（> 0 时保证海森矩阵可逆）
        max_iters：最多牛顿步数
        tol      ：更新量的最大分量小于它就停止
    【返回】权重列表 w（长度 = 特征维数）
    【边界】
        · 数据线性可分（某个 w 能把所有比较都"判对"）时，不加正则化的解在无穷远处；l2 > 0 能防止发散。
        · 某个特征在所有选项上都一样（例如常数特征），它的差恒为 0，学不到任何东西（权重停在 0）。
    """
    d = len(diffs[0])
    n = len(diffs)
    w = [0.0] * d
    for _ in range(max_iters):
        grad = [-l2 * wi for wi in w]
        hess = [[(-l2 if i == j else 0.0) for j in range(d)] for i in range(d)]
        for x in diffs:
            p = sigmoid(dot(w, x))
            for i in range(d):
                grad[i] += (1.0 - p) * x[i] / n
                for j in range(i, d):
                    hess[i][j] -= p * (1.0 - p) * x[i] * x[j] / n
        for i in range(d):                     # 海森矩阵是对称的，只算了上三角，这里补全下三角
            for j in range(i):
                hess[i][j] = hess[j][i]
        step = solve_linear_system(hess, grad)  # H·step = g → w_new = w − step
        w = [wi - si for wi, si in zip(w, step)]
        if max(abs(si) for si in step) < tol:
            break
    return w


# =============================================================================
# 第五部分：数值梯度（给"没推导梯度公式"的小模型用）
# =============================================================================
def numeric_gradient(f, params, h=1e-5):
    """
    用"中心差分"近似函数 f 在 params 处的梯度：∂f/∂θ_i ≈ [f(θ + h·e_i) − f(θ − h·e_i)] / (2h)。

    【思路】不想（或来不及）手推梯度公式时，可以直接"往两边各挪一点，看函数值变了多少"。
        代价是每算一次梯度要调用 f 两倍参数个数那么多次，只适合参数很少的模型（第 4.4 节的平局模型）。
        真实的深度学习用"自动微分"（PyTorch 的 backward），又快又准。
    【参数】f：输入参数列表、返回一个数的函数；params：参数列表；h：挪动的步长
    【返回】梯度列表
    【边界】h 太大，近似不准；h 太小，浮点舍入误差会放大。1e-5 对大多数光滑函数是合理的折中。
    """
    grad = []
    for i in range(len(params)):
        up = list(params)
        dn = list(params)
        up[i] += h
        dn[i] -= h
        grad.append((f(up) - f(dn)) / (2 * h))
    return grad


def gradient_ascent(f, params, lr=0.5, iters=3000, tol=1e-7):
    """
    用数值梯度做梯度上升，求 f 的最大值点。

    【参数】f：目标函数；params：初始参数；lr：步长；iters：最多步数；tol：梯度足够小就停
    【返回】参数列表
    【边界】步长是固定的，太大会来回震荡、太小会很慢；这里只用于参数很少、目标很光滑的小问题。
    """
    p = list(params)
    for _ in range(iters):
        g = numeric_gradient(f, p)
        if max(abs(gi) for gi in g) < tol:
            break
        p = [pi + lr * gi for pi, gi in zip(p, g)]
    return p


# =============================================================================
# 直接运行：演示几种转换
# =============================================================================
if __name__ == "__main__":
    print("完整排序 A ≻ C ≻ B ≻ D")
    print("  拆成选择事件：", ranking_to_events(["A", "C", "B", "D"]))
    print("  拆成两两比较：", ranking_to_pairs(["A", "C", "B", "D"]))
    print("多选一：从 {A, B, C} 里选了 B")
    print("  选择事件：", choice_to_event("B", ["A", "B", "C"]))
    print("  拆成两两比较：", choice_to_pairs("B", ["A", "B", "C"]))
    print("接受 A：", accept_reject_to_event("A", True), "   拒绝 A：", accept_reject_to_event("A", False))
    rng = random.Random(0)
    utils = {"A": 2.0, "B": 1.0, "C": 0.0}
    events = []
    for _ in range(3000):
        events += ranking_to_events(sample_ranking(utils, ["A", "B", "C"], rng))
    print("\n真实效用 A=2, B=1, C=0；从 3000 次模拟排序拟合（平均值移到 0，所以应接近 1, 0, −1）：")
    print("  ", {k: round(v, 3) for k, v in fit_choice_model(events, ["A", "B", "C"]).items()})