"""
=============================================================================
第 4 讲 · 工具模块：因子模型 —— 用户和选项都变成向量
=============================================================================

对应讲义：lectures/lecture04_factor_models.md

-----------------------------------------------------------------------------
【这个模块解决什么问题？】
-----------------------------------------------------------------------------
前三讲的模型都是"每个选项一个分数，所有人共用"。可是不同的人喜欢不同的东西。
因子模型给每个用户 i 和每个选项 j 各一组参数：

        接受 / 拒绝（逐项数据）：  P(Y_ij = 1) = σ( a_i + z_j + U_i · V_j )
        两两比较：               P(j ≻ k | i) = σ( (z_j − z_k) + U_i · (V_j − V_k) )

    a_i ：用户的"胃口"（整体上爱不爱接受）        —— Rasch 模型里的用户参数
    z_j ：选项的"吸引力"（对所有人都一样的部分）    —— Rasch 模型里的选项参数
    U_i ：用户的 K 维"口味向量"
    V_j ：选项的 K 维"特征向量"
    K = 0 时就是 Rasch 模型；再去掉 a_i，就是"只看流行度"的模型。

怎样拟合？参数很多（用户数 × K + 选项数 × K），但有一个好用的结构：
    · 固定所有选项的参数，每个用户的 (a_i, U_i) 就是一个**小小的逻辑回归**（特征是 V_j，偏移是 z_j）；
    · 固定所有用户的参数，每个选项的 (z_j, V_j) 也是一个小逻辑回归（特征是 U_i，偏移是 a_i）。
所以可以**交替**地解这两组小问题（交替最小二乘 ALS 的逻辑回归版本）。每个小问题只有 K+1 个参数，
用牛顿法几步就收敛。两两比较的版本也一样，只是特征变成了"差"。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from factor_models import fit_mf, predict_prob
    obs = [(用户, 选项, 0 或 1), ...]
    model = fit_mf(obs, users, items, k=2, l2=1.0, rng=random.Random(0))
    p = predict_prob(model, "u001", "r07")

模块里还有一个"口味世界" TasteWorld（每个回答有两个隐藏的风格维度，每个用户有两维口味），代码 02、04 共用。

直接运行本文件会在口味世界上拟合流行度、Rasch 和 K=2 的模型：
    python3 factor_models.py
"""

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 mean, sigmoid  # noqa: E402


# =============================================================================
# 第一部分：小积木 —— 解线性方程组、带偏移量的逻辑回归
# =============================================================================
def solve(a, b):
    """
    解线性方程组 A x = b（高斯消元 + 部分主元）。和第 2 讲 choice_data.solve_linear_system 相同，
    这里再写一份，让本讲只依赖第 1 讲。

    【参数】a：d×d 列表的列表；b：长度 d 的列表
    【返回】长度 d 的列表
    【边界】A 必须可逆（我们总是加 L2 正则化来保证这一点）。
    """
    d = len(b)
    m = [list(row) + [b[i]] for i, row in enumerate(a)]
    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 dot(u, v):
    """点积。"""
    return math.fsum(a * b for a, b in zip(u, v))


def logreg_newton(rows, w, l2, steps=3):
    """
    带偏移量的 L2 正则化逻辑回归，用牛顿法从 w 出发走几步。

    【思路】模型：P(y = 1) = σ(w·x + offset)。目标：Σ log P(观察到的 y) − (l2/2)·‖w‖²。
        梯度  g = Σ (y − p) x − l2·w
        海森  H = −Σ p(1 − p) x xᵀ − l2·I
        牛顿  w ← w − H⁻¹ g
        "偏移量"（offset）是已经固定的那部分 logit（例如交替优化时，另一方的参数贡献的部分）。
    【参数】
        rows ：list[(x 列表, offset, y)]，y 是 0 或 1
        w    ：起点（列表，长度 = 特征维数）
        l2   ：正则化强度（> 0，保证海森矩阵可逆）
        steps：牛顿步数（交替优化时每轮走几步就够，下一轮会接着走）
    【返回】新的 w
    【边界】rows 为空时只剩正则化项，最优解是 w = 0；我们直接返回 0 向量。
    """
    d = len(w)
    if not rows:
        return [0.0] * d
    for _ in range(steps):
        g = [-l2 * wi for wi in w]
        h = [[(-l2 if i == j else 0.0) for j in range(d)] for i in range(d)]
        for x, off, y in rows:
            p = sigmoid(dot(w, x) + off)
            r = y - p
            q = p * (1 - p)
            for i in range(d):
                g[i] += r * x[i]
                xi = q * x[i]
                for j in range(i, d):
                    h[i][j] -= xi * x[j]
        for i in range(d):
            for j in range(i):
                h[i][j] = h[j][i]
        step = solve(h, g)
        w = [wi - si for wi, si in zip(w, step)]
        if max(abs(s) for s in step) < 1e-9:
            break
    return w


# =============================================================================
# 第二部分：逐项数据（接受 / 拒绝）的因子模型
# =============================================================================
class FactorModel:
    """
    一个拟合好的因子模型。

    【字段】
        k        ：隐向量维数（0 表示 Rasch 或流行度模型）
        a        ：dict {用户: 胃口 a_i}（不用用户偏置时全为 0）
        u        ：dict {用户: 口味向量 U_i}（长度 k 的列表）
        z        ：dict {选项: 吸引力 z_j}
        v        ：dict {选项: 特征向量 V_j}
    """

    def __init__(self, k, a, u, z, v):
        self.k, self.a, self.u, self.z, self.v = k, a, u, z, v

    def logit(self, i, j):
        """logit = a_i + z_j + U_i·V_j（没见过的用户 / 选项，参数按 0 处理）。"""
        ui = self.u.get(i, [0.0] * self.k)
        vj = self.v.get(j, [0.0] * self.k)
        return self.a.get(i, 0.0) + self.z.get(j, 0.0) + dot(ui, vj)


def predict_prob(model, i, j):
    """P(用户 i 接受选项 j)。"""
    return sigmoid(model.logit(i, j))


def fit_mf(obs, users, items, k, l2, rng, sweeps=25, user_bias=True, init_scale=0.1):
    """
    交替牛顿法拟合逐项数据的因子模型：P(Y_ij = 1) = σ(a_i + z_j + U_i·V_j)。

    【参数】
        obs       ：list[(用户, 选项, y)]
        users、items：全部用户、全部选项的列表
        k         ：隐向量维数（0 = Rasch；0 且 user_bias=False = 只看流行度）
        l2        ：L2 正则化强度（对每个小逻辑回归；相当于参数的高斯先验，第 5、6 讲）
        rng       ：random.Random，用于随机初始化 V（全 0 初始化时 U、V 会一直是 0，学不动）
        sweeps    ：交替的轮数
        user_bias ：是否有用户胃口 a_i
        init_scale：V 的初始随机大小
    【返回】FactorModel
    【边界】
        · 只有交互项 U_i·V_j 需要随机初始化；偏置项从 0 开始。
        · 不同的随机种子会得到不同的 U、V（旋转、翻转等），但预测几乎一样（第 4.4 节：识别与 Rashomon）。
        · 某个用户的所有回答都是 1（或都是 0）时，不加正则化的话他的胃口会被推到 ±∞；l2 防止这种情况。
    """
    by_user = {i: [] for i in users}
    by_item = {j: [] for j in items}
    for i, j, y in obs:
        by_user[i].append((j, y))
        by_item[j].append((i, y))
    a = {i: 0.0 for i in users}
    u = {i: [rng.gauss(0, init_scale) for _ in range(k)] for i in users}
    z = {j: 0.0 for j in items}
    v = {j: [rng.gauss(0, init_scale) for _ in range(k)] for j in items}

    for _ in range(sweeps):
        # 1) 固定选项，更新每个用户：特征 = [1（若有胃口）] + V_j，偏移 = z_j
        for i in users:
            rows = [(([1.0] if user_bias else []) + v[j], z[j], y) for j, y in by_user[i]]
            w0 = ([a[i]] if user_bias else []) + u[i]
            if not w0:
                continue
            w = logreg_newton(rows, w0, l2, steps=2)
            if user_bias:
                a[i], u[i] = w[0], w[1:]
            else:
                u[i] = w
        # 2) 固定用户，更新每个选项：特征 = [1] + U_i，偏移 = a_i
        for j in items:
            rows = [([1.0] + u[i], a[i], y) for i, y in by_item[j]]
            w = logreg_newton(rows, [z[j]] + v[j], l2, steps=2)
            z[j], v[j] = w[0], w[1:]
    return FactorModel(k, a, u, z, v)


# =============================================================================
# 第三部分：两两比较的因子模型（个性化 BT）
# =============================================================================
def fit_pairwise_mf(comps, users, items, k, l2, rng, sweeps=25, init_scale=0.1):
    """
    交替牛顿法拟合个性化的两两比较模型：P(j ≻ k | i) = σ((z_j − z_k) + U_i·(V_j − V_k))。

    【思路】
        · 固定选项：用户 i 的 U_i 是逻辑回归，特征 = V_赢 − V_输，偏移 = z_赢 − z_输，标签全是 1。
          注意**没有用户胃口 a_i**：它在"赢 − 输"里消掉了（第 2.3 节：Rasch 的用户参数在比较中抵消）。
        · 固定用户：轮流更新每个选项 j，其他选项不动。j 出现在一条比较里时，
          如果 j 是赢家，logit = +(z_j + U_i·V_j) − (z_k + U_i·V_k)：特征 = +[1, U_i]，偏移 = −(z_k + U_i·V_k)；
          如果 j 是输家，特征 = −[1, U_i]，偏移 = +(z_赢 + U_i·V_赢)。标签都是 1（"赢家胜"这件事发生了）。
    【参数】comps：list[(用户, 赢家, 输家)]；其余同 fit_mf
    【返回】FactorModel（a 全为 0；z 只有差值有意义）
    """
    by_user = {i: [] for i in users}
    by_item = {j: [] for j in items}
    for i, w, l in comps:
        by_user[i].append((w, l))
        by_item[w].append((i, w, l))
        by_item[l].append((i, w, l))
    a = {i: 0.0 for i in users}
    u = {i: [rng.gauss(0, init_scale) for _ in range(k)] for i in users}
    z = {j: 0.0 for j in items}
    v = {j: [rng.gauss(0, init_scale) for _ in range(k)] for j in items}

    for _ in range(sweeps):
        if k > 0:
            for i in users:
                rows = [([vw - vl for vw, vl in zip(v[w], v[l])], z[w] - z[l], 1) for w, l in by_user[i]]
                u[i] = logreg_newton(rows, u[i], l2, steps=2)
        for j in items:
            rows = []
            for i, w, l in by_item[j]:
                if j == w:
                    other = z[l] + dot(u[i], v[l])
                    rows.append(([1.0] + u[i], -other, 1))
                else:
                    other = z[w] + dot(u[i], v[w])
                    rows.append(([-1.0] + [-x for x in u[i]], other, 1))
            wj = logreg_newton(rows, [z[j]] + v[j], l2, steps=2)
            z[j], v[j] = wj[0], wj[1:]
    return FactorModel(k, a, u, z, v)


def pairwise_prob(model, i, j, k):
    """P(用户 i 认为 j ≻ k)。"""
    return sigmoid(model.logit(i, j) - model.logit(i, k))


# =============================================================================
# 第四部分：评估指标
# =============================================================================
def log_loss(probs, labels):
    """平均对数损失 −mean[y log p + (1−y) log(1−p)]（越小越好；log 2 ≈ 0.693 相当于瞎猜）。"""
    eps = 1e-12
    return -mean([math.log(max(p, eps)) if y == 1 else math.log(max(1 - p, eps)) for p, y in zip(probs, labels)])


def accuracy(probs, labels):
    """准确率：p > 0.5 预测为 1。"""
    return mean([1.0 if (p > 0.5) == (y == 1) else 0.0 for p, y in zip(probs, labels)])


def auc(scores, labels):
    """
    AUC：随机挑一个正例和一个负例，模型给正例的分数更高的概率（分数相同算一半）。

    【思路】把所有样本按分数排序，正例的平均排名越靠后（分数越高），AUC 越大。
        公式（Mann–Whitney）：AUC = (正例的秩之和 − n₊(n₊+1)/2) / (n₊ · n₋)，相同分数取平均秩。
    【返回】0~1；0.5 相当于瞎猜。只有一类标签时返回 0.5。
    """
    pairs = sorted(zip(scores, labels))
    n = len(pairs)
    ranks = [0.0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n and pairs[j + 1][0] == pairs[i][0]:
            j += 1
        r = (i + j) / 2 + 1
        for t in range(i, j + 1):
            ranks[t] = r
        i = j + 1
    pos = [r for r, (_, y) in zip(ranks, pairs) if y == 1]
    n_pos, n_neg = len(pos), n - len(pos)
    if n_pos == 0 or n_neg == 0:
        return 0.5
    return (math.fsum(pos) - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)


def evaluate(model, test):
    """在测试集 [(用户, 选项, y)] 上返回 (对数损失, 准确率, AUC)。"""
    probs = [predict_prob(model, i, j) for i, j, _ in test]
    labels = [y for _, _, y in test]
    return log_loss(probs, labels), accuracy(probs, labels), auc(probs, labels)


def split(obs, rng, test_frac=0.2):
    """把观察随机分成训练集和测试集（按"格子"划分：测试的格子在训练里没出现过，但用户和选项都出现过）。"""
    shuffled = list(obs)
    rng.shuffle(shuffled)
    cut = int(len(shuffled) * (1 - test_frac))
    return shuffled[:cut], shuffled[cut:]


# =============================================================================
# 第五部分：一个"口味世界"（代码 02、04 共用）
# =============================================================================
class TasteWorld:
    """
    "用户 × 回答"的隐藏世界（第 4 讲代码 02、04 使用）。

    【字段】
        users、items：编号列表
        appetite    ：dict {用户: 胃口}
        taste       ：dict {用户: [对详细程度的口味, 对正式程度的口味]}
        style       ：dict {回答: [详细程度, 正式程度]}，都在 [−1, 1] 之间
        quality     ：dict {回答: 质量（平均为 0）}
    """

    def __init__(self, rng, n_users=200, n_items=40):
        self.users = [f"u{i:03d}" for i in range(n_users)]
        self.items = [f"r{j:02d}" for j in range(n_items)]
        self.appetite = {i: rng.gauss(0, 0.7) for i in self.users}
        self.taste = {i: [rng.gauss(0, 1.5), rng.gauss(0, 1.5)] for i in self.users}
        self.style = {j: [rng.uniform(-1, 1), rng.uniform(-1, 1)] for j in self.items}
        self.quality = {j: rng.gauss(0, 0.5) for j in self.items}

    def logit(self, i, j):
        """真实的 logit = 胃口 + 2·质量 + 口味·风格。"""
        return self.appetite[i] + 2.0 * self.quality[j] + dot(self.taste[i], self.style[j])

    def sample(self, rng, frac):
        """每个格子以概率 frac 被观察到，观察到的格子按真实概率生成 0/1。"""
        obs = []
        for i in self.users:
            for j in self.items:
                if rng.random() < frac:
                    obs.append((i, j, 1 if rng.random() < sigmoid(self.logit(i, j)) else 0))
        return obs


def oracle_metrics(world, test):
    """上帝模型在测试集上的 (对数损失, 准确率, AUC)。"""
    probs = [sigmoid(world.logit(i, j)) for i, j, _ in test]
    labels = [y for _, _, y in test]
    return log_loss(probs, labels), accuracy(probs, labels), auc(probs, labels)


if __name__ == "__main__":
    rng = random.Random(0)
    world = TasteWorld(rng)
    train, test = split(world.sample(rng, 0.6), rng)
    for name, kk, ub in [("只看流行度", 0, False), ("Rasch", 0, True), ("K=2 因子模型", 2, True)]:
        m = fit_mf(train, world.users, world.items, kk, 3.0, random.Random(1), user_bias=ub)
        ll, acc, a = evaluate(m, test)
        print(f"{name:12s} 测试对数损失 {ll:.3f}  准确率 {acc:.3f}  AUC {a:.3f}")