"""
=============================================================================
第 12 讲 · 代码 4：推荐系统的冷启动 —— 线性 Thompson 采样
=============================================================================

对应讲义：lectures/lecture12_thompson_acquisition.md 第 4.3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
一个新用户来了，系统对他一无所知（"冷启动"）。每个内容有一个特征向量 x（8 维），用户的口味是一个向量 θ，
用户给内容的评分 r = xᵀθ + 噪声（第 4 讲的因子模型：内容的特征已知，只学这个用户的 θ）。
所有用户的口味大致围绕一个"大众口味" m 分布：θ ~ N(m, I)。

每一轮给这个用户推荐一个还没推荐过的内容，看到评分，更新对 θ 的后验（线性高斯模型，后验有封闭解）。
五种推荐方式：
    · 随机
    · 热门：总是推荐"大众口味"下最好的（不管这个用户）
    · 贪心：推荐后验均值下最好的
    · LinUCB：推荐 xᵀθ̂ + α √(xᵀ Σ x) 最大的（Li 等 2010，雅虎新闻推荐）
    · 线性 Thompson 采样：从后验抽一个 θ̃，推荐 xᵀθ̃ 最大的（Agrawal & Goyal 2013）

    演示：200 个新用户、每人 50 轮。不同阶段每轮的平均遗憾（和"知道 θ 时最好的还没推荐过的内容"比）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_cold_start.py        # 约 5 秒

只需要 numpy。依赖本目录的 ts_tools.py（只用打印工具）。
"""

import numpy as np

from ts_tools import pad, print_header

D, N_ITEMS, NOISE = 8, 200, 0.5
METHODS = {"random": "随机", "popular": "热门", "greedy": "贪心", "linucb": "LinUCB（α=1）", "ts": "线性 Thompson 采样"}


def run_user(method, items, m, rng, T=50, alpha=1.0):
    """
    给一个新用户推荐 T 轮。

    【思路】先验 θ ~ N(m, I)；线性高斯模型的后验（第 6 讲的共轭）：
        精度 P = I + Σ x xᵀ / σ²，均值 μ = P⁻¹ (m + Σ x r / σ²)
    【返回】(T,) 每一轮的遗憾 = 剩下的内容里真实最好的评分期望 − 推荐的内容的评分期望
    """
    theta = m + rng.normal(size=D)                      # 这个用户真实的口味（系统看不到）
    true_score = items @ theta
    available = np.ones(N_ITEMS, dtype=bool)
    P = np.eye(D)
    b = m.copy()
    regrets = np.zeros(T)
    for t in range(T):
        cov = np.linalg.inv(P)
        mu = cov @ b
        idx = np.flatnonzero(available)
        X = items[idx]
        if method == "random":
            k = int(rng.integers(len(idx)))
        elif method == "popular":
            k = int(np.argmax(X @ m))
        elif method == "greedy":
            k = int(np.argmax(X @ mu))
        elif method == "linucb":
            k = int(np.argmax(X @ mu + alpha * np.sqrt(np.einsum("nd,de,ne->n", X, cov, X))))
        else:
            sample = mu + np.linalg.cholesky(cov) @ rng.normal(size=D)
            k = int(np.argmax(X @ sample))
        item = idx[k]
        regrets[t] = true_score[idx].max() - true_score[item]
        r = true_score[item] + NOISE * rng.normal()
        P += np.outer(items[item], items[item]) / NOISE ** 2
        b += items[item] * r / NOISE ** 2
        available[item] = False
    return regrets


def main(users=200):
    rng = np.random.default_rng(4)
    items = rng.normal(size=(N_ITEMS, D)) / np.sqrt(D)
    m = rng.normal(size=D) * 0.7                         # 大众口味
    print_header("推荐系统的冷启动：五种推荐方式")
    print(f"{N_ITEMS} 个内容、{D} 维特征；{users} 个新用户，口味 θ ~ N(大众口味, I)，每人推荐 50 轮（不重复）。")
    print("表中是每轮的平均遗憾（越小越好）：\n")
    stages = [(0, 5), (5, 15), (15, 50)]
    print(f"  {pad('推荐方式', 20)}" + "".join(pad(f"第 {a + 1}~{b} 轮", 13, "right") for a, b in stages)
          + pad("50 轮累计", 11, "right"))
    for method, name in METHODS.items():
        R = np.array([run_user(method, items, m, np.random.default_rng(1000 + u)) for u in range(users)])
        print(f"  {pad(name, 20)}" + "".join(f"{R[:, a:b].mean():>13.3f}" for a, b in stages)
              + f"{R.sum(axis=1).mean():>11.2f}")
    print("""
观察：
  · 热门推荐一开始还不错（大众口味是一个合理的先验），但它从不学习：每个用户都被当成“平均用户”，
    遗憾一直居高不下。
  · 贪心学得很快：每次推荐的内容都带来了关于 θ 的信息（内容的特征各不相同 —— 代码 02 的“特征多样性”）。
  · LinUCB 的累计遗憾最小：前几轮推荐一些“不确定但可能合口味”的内容，之后学得更准。
  · Thompson 采样在前 15 轮探索得最多、付出最多，50 轮之内还没把成本赚回来，累计遗憾比贪心略大。
    只有 50 轮时，探索的回报期很短；对一个会长期使用的用户，前期的探索更值得。
  · 在这个问题上三者相差不大 —— 和代码 02 一样，特征的多样性已经提供了不少探索。
  · 实践中，Thompson 采样和 LinUCB 被广泛用于新闻、广告、电商的冷启动（Li 等 2010；Chapelle & Li 2011）；
    它们的价值在“特征不够多样”或“需要发现小众口味”时更明显（代码 02 演示 2）。""")


if __name__ == "__main__":
    main()