"""
=============================================================================
第 2 讲 · 代码 2：Naive Bayes 与垃圾邮件分类（Naive Bayes on a synthetic spam corpus）
=============================================================================

对应讲义：lectures/lecture02_generative_kernels_svm.md 第 3 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
    演示 1：合成语料（300 个词，按 multinomial 模型生成）上，multinomial NB、Bernoulli NB 和逻辑回归的学习曲线。
            谁在数据少的时候更好？（Ng & Jordan 2002 的经典对比）
    演示 2：Laplace smoothing。α = 0（纯 MLE）时一个没见过的词就能让某一类的概率变成恰好 0，
            大部分测试邮件“两类都不可能”；α 太大又把信号抹平。再构造一封“满是垃圾词、只多了一个 meeting”的邮件。
    演示 3：条件独立假设被破坏时（每个词都变成 k 个总是一起出现的词），NB 的准确率几乎不变，
            但概率变得极度自信：log loss 和 calibration 变差。逻辑回归怎样？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_naive_bayes.py          # 约 2 秒

依赖本目录的 generative_models.py，以及第 1 讲的 linear_models.py（generative_models 导入时自动加入搜索路径，
所以下面要先 import generative_models，再 import linear_models）。
"""

import numpy as np

from generative_models import (expected_calibration_error, fit_bernoulli_nb, fit_multinomial_nb, log_loss,
                               make_corpus, make_spam_model, nb_log_joint, nb_predict_proba)
from linear_models import add_intercept, fit_logistic_newton, print_header, sigmoid

LR_L2 = 1e-3          # 逻辑回归的 L2 系数（全文件统一；为什么需要一点正则化见演示 1 的说明）


def nb_accuracy(p, y):
    """
    NB 的准确率：p ≥ 0.5 判为 spam。p 是 nan（两类概率都为 0，后验无定义）的样本**算作分错**。

    【为什么这样处理】α = 0 时这种情况很常见。真实系统只能给一个默认答案（比如“正常邮件”）或者报错；
        这里选择最严格的计法，让“无法判断”的代价显式地出现在准确率里。
    """
    pred = np.where(np.isnan(p), -1, (p >= 0.5).astype(int))
    return float(np.mean(pred == y))


def lr_fit_predict(X, y, X_test):
    """在词频特征上训练逻辑回归（Newton，L2 = LR_L2），返回测试集上的 p(spam | x)。"""
    th, _ = fit_logistic_newton(add_intercept(X), y, l2=LR_L2)
    return sigmoid(add_intercept(X_test) @ th)


# =============================================================================
# 演示 1：两种 event model 与逻辑回归的学习曲线
# =============================================================================
def demo_event_models():
    print_header("演示 1：multinomial vs Bernoulli event model vs 逻辑回归 —— 学习曲线")
    vocab, P = make_spam_model()
    log_ratio = np.log(P[1] / P[0])                     # 真实的“每个词的证据”log P(w|spam)/P(w|ham)
    order = np.argsort(log_ratio)
    print(f"词表 V = {len(vocab)}：20 个常用词 + 15 个垃圾邮件词 + 15 个正常邮件词 + 250 个按 Zipf 定律递减的长尾词。")
    print("最像 spam 的词（真实的 log P(w|spam)/P(w|ham)）：" +
          ", ".join(f"{vocab[i]} {log_ratio[i]:+.2f}" for i in order[::-1][:5]))
    print("最像 ham 的词：" + ", ".join(f"{vocab[i]} {log_ratio[i]:+.2f}" for i in order[:5]))
    print("每封邮件 5 + Poisson(15) 个词（平均 20 个），逐词独立地从本类的词分布里抽 —— multinomial NB 的假设恰好成立。\n")

    rng = np.random.default_rng(41)
    Xte, yte = make_corpus(2000, P, rng)
    true_model = {"kind": "multinomial", "log_prior": np.log([0.5, 0.5]), "P": P}
    bayes_acc = nb_accuracy(nb_predict_proba(true_model, Xte), yte)
    sizes, reps = [5, 10, 25, 50, 100, 250, 1000], 10
    print(f"测试集 4000 封。用**真实**的词分布做 Bayes 最优分类，准确率 = {bayes_acc:.4f}（上限）。")
    print(f"每个训练集大小重复 {reps} 次取平均；NB 用 α = 1（Laplace），逻辑回归用词频做特征、L2 = {LR_L2:g}。\n")
    print(f"{'docs/class':>11}{'multinomial NB':>16}{'Bernoulli NB':>14}{'logistic reg.':>15}")
    res = []
    for npc in sizes:
        am, ab, al = [], [], []
        for _ in range(reps):
            X, y = make_corpus(npc, P, rng)
            am.append(nb_accuracy(nb_predict_proba(fit_multinomial_nb(X, y, 1.0), Xte), yte))
            ab.append(nb_accuracy(nb_predict_proba(fit_bernoulli_nb(X, y, 1.0), Xte), yte))
            al.append(nb_accuracy(lr_fit_predict(X, y, Xte), yte))
        res.append((np.mean(am), np.mean(ab), np.mean(al)))
        print(f"{npc:>11}{res[-1][0]:>16.4f}{res[-1][1]:>14.4f}{res[-1][2]:>15.4f}")
    res = np.array(res)
    gap_small, gap_big = res[0, 0] - res[0, 2], res[-1, 0] - res[-1, 2]
    print(f"""
观察 Observations：
  · 训练集很小时 NB 远好于逻辑回归（每类 {sizes[0]} 封：{res[0, 0]:.3f} vs {res[0, 2]:.3f}）。NB 的“训练”只是数数，
    每个词的参数独立地估计；逻辑回归要在 {len(vocab) + 1} 个参数上联合地做优化，数据少时方差大。
  · 数据多了差距缩小（每类 {sizes[-1]} 封时差 {gap_big:.3f}，最开始差 {gap_small:.3f}）。这里数据**恰好**按 NB 的模型生成，
    所以 NB 一直不吃亏、逻辑回归也只是追近；真实文本不满足条件独立，数据足够多时逻辑回归通常会反超
    —— Ng & Jordan (2002) 在很多真实数据集上看到的“交叉”。
  · Bernoulli 和 multinomial 在数据多时几乎一样（每类 {sizes[-1]} 封：{res[-1, 1]:.4f} vs {res[-1, 0]:.4f}）：
    邮件很短（平均 20 个词），一个关键词很少出现两次，“出现几次”和“出现没有”差别不大。
    训练集极小时 Bernoulli 反而略好（每类 {sizes[0]} 封：{res[0, 1]:.4f} vs {res[0, 0]:.4f}），这里不深究原因
    （一个可能的因素：它还利用了“哪些词没出现”）。文献里的经验是文档越长、词表越大，multinomial 越占优
    （McCallum & Nigam 1998）。
  · 每类 {sizes[-1]} 封时 multinomial NB = {res[-1, 0]:.4f}，离 Bayes 上限 {bayes_acc:.4f} 很近。""")


# =============================================================================
# 演示 2：Laplace smoothing
# =============================================================================
def demo_laplace():
    print_header("演示 2：Laplace smoothing —— 为什么一个零计数就能毁掉整个预测")
    vocab, P = make_spam_model()
    rng = np.random.default_rng(42)
    X, y = make_corpus(25, P, rng)
    Xte, yte = make_corpus(2000, P, rng)
    counts = np.vstack([X[y == 0].sum(axis=0), X[y == 1].sum(axis=0)])
    print("训练集每类 25 封（共约 1000 个词），测试集 4000 封。")
    print(f"300 个词里：在 ham 训练集中从未出现 {np.sum(counts[0] == 0)} 个，在 spam 中从未出现 {np.sum(counts[1] == 0)} 个，"
          f"两类都没出现 {np.sum((counts[0] == 0) & (counts[1] == 0))} 个。")
    print("α = 0 时这些词的估计概率恰好是 0；测试邮件里只要出现一个，那一类的 log p(x|y) = −∞。\n")
    print(f"{'alpha':>8}{'model':>13}{'accuracy':>10}{'both p=0':>10}{'one p=0':>9}{'log loss':>10}")
    table = {}
    for alpha in [0.0, 0.001, 0.01, 0.1, 0.3, 1.0, 3.0, 10.0]:
        for fit in (fit_multinomial_nb, fit_bernoulli_nb):
            model = fit(X, y, alpha)
            LJ = nb_log_joint(model, Xte)
            n_inf = np.sum(np.isinf(LJ), axis=1)
            p = nb_predict_proba(model, Xte)
            acc = nb_accuracy(p, yte)
            ll = log_loss(yte, np.where(np.isnan(p), 0.5, p))         # 无定义的样本按 p = 0.5 计
            table[(alpha, model["kind"])] = (acc, ll)
            print(f"{alpha:>8g}{model['kind']:>13}{acc:>10.4f}{np.sum(n_inf == 2):>10}{np.sum(n_inf == 1):>9}{ll:>10.3f}")

    # 一封构造出来的邮件：8 个最像 spam 的词 + 1 个只在 ham 训练集里出现过的词
    # 按**真实**的 log P(w|spam)/P(w|ham) 挑 8 个最像 spam 的词，但只在“两类训练集里都出现过”的词中挑
    # （这样 α = 0 时这封邮件本身在两类下的概率都不是 0，后验有定义）
    seen_both = (counts[0] > 0) & (counts[1] > 0)
    true_ratio = np.where(seen_both, np.log(P[1] / P[0]), -np.inf)
    spammy = np.argsort(true_ratio)[::-1][:8]
    ham_only = np.flatnonzero((counts[1] == 0) & (counts[0] > 0))
    extra = ham_only[np.argmax(counts[0, ham_only])]
    doc = np.zeros((2, len(vocab)), dtype=int)
    doc[:, spammy] = 1
    doc[1, extra] = 1
    print("\n构造一封邮件：" + " ".join(vocab[i] for i in spammy) + f"，然后再加一个词 “{vocab[extra]}”"
          f"（ham 训练集里出现 {counts[0, extra]} 次，spam 训练集里 0 次）：")
    print(f"{'':>24}{'P(spam | doc)':>16}{'P(spam | doc + ' + vocab[extra] + ')':>28}")
    p_doc = {}
    for alpha in (0.0, 1.0):
        p_doc[alpha] = nb_predict_proba(fit_multinomial_nb(X, y, alpha), doc)
        print(f"{'multinomial NB, α = ' + format(alpha, 'g'):>24}{p_doc[alpha][0]:>16.6f}{p_doc[alpha][1]:>28.6f}")
    best_alpha = max([a for a in [0.001, 0.01, 0.1, 0.3, 1.0, 3.0, 10.0]],
                     key=lambda a: -table[(a, "multinomial")][1])
    print(f"""
观察 Observations：
  · α = 0：{np.sum(np.all(np.isinf(nb_log_joint(fit_multinomial_nb(X, y, 0.0), Xte)), axis=1))} / 4000 封测试邮件在两类下的概率**都**恰好是 0
    （对每一类，邮件里都至少有一个“该类训练集里从没出现过”的词），后验是 0/0，无法判断；
    准确率只有 {table[(0.0, 'multinomial')][0]:.3f}。
    MLE 把“训练集里没见过”当成了“不可能发生”。词表大、训练集小时，这几乎必然发生。
  · 只要 α > 0（哪怕 0.001），所有邮件都能判断，准确率立刻跳到 {table[(0.001, 'multinomial')][0]:.3f}。
    准确率对 α 不太敏感（0.001～1 之间都在 {min(table[(a, 'multinomial')][0] for a in [0.001, 0.01, 0.1, 0.3, 1.0]):.3f}～{max(table[(a, 'multinomial')][0] for a in [0.001, 0.01, 0.1, 0.3, 1.0]):.3f}），
    但 log loss 很敏感：α 太小时，没见过的词把概率推到极端（α = 0.001 时 log loss = {table[(0.001, 'multinomial')][1]:.3f}），
    在这份数据上 multinomial 的 log loss 在 α = {best_alpha:g} 附近最小（{table[(best_alpha, 'multinomial')][1]:.3f}）。
  · α 太大（10）：每个词都“假装见过 10 次”，而每类总共才约 500 个词，真实的信号被抹平，准确率掉到
    {table[(10.0, 'multinomial')][0]:.3f}。α 是一个正则化强度：从贝叶斯的角度，它是 Dirichlet 先验的“伪计数”（第 3 讲）。
  · 构造的邮件：α = 0 时，8 个垃圾词的证据被 1 个零计数**一票否决**，P(spam) 从 {p_doc[0.0][0]:.3f} 变成恰好 0。
    log 概率是逐词相加的，一个 −∞ 就压倒一切。α = 1 时，多一个 “{vocab[extra]}” 把 P(spam) 从 {p_doc[1.0][0]:.3f}
    拉到 {p_doc[1.0][1]:.3f}：它是一条普通的、有限的反向证据，而不是否决票。""")


# =============================================================================
# 演示 3：破坏条件独立假设 —— 过度自信
# =============================================================================
def demo_overconfidence():
    print_header("演示 3：条件独立假设被破坏 —— 准确率还在，概率却过度自信")
    vocab, P = make_spam_model()
    rng = np.random.default_rng(43)
    X1, y = make_corpus(250, P, rng)
    Xte1, yte = make_corpus(2000, P, rng)
    print("训练集每类 250 封，测试集每类 2000 封。把每个词换成 k 个**总是一起出现**的不同的词（比如 “click here now”），")
    print("词表变成 300k 个。NB 假设它们相互独立，于是同一份证据被数了 k 遍。k = 1 时 NB 的假设恰好成立。\n")
    print(f"{'k':>3}{'model':>8}{'accuracy':>10}{'log loss':>10}{'ECE':>8}{'mean conf':>11}{'conf>0.999':>12}"
          f"{'acc | conf>0.999':>18}")
    logodds1 = None
    summary = {}
    for k in [1, 2, 4, 8]:
        X, Xte = np.repeat(X1, k, axis=1), np.repeat(Xte1, k, axis=1)     # 每一列复制 k 份（相邻排列）
        model = fit_multinomial_nb(X, y, 1.0)
        LJ = nb_log_joint(model, Xte)
        logodds = LJ[:, 1] - LJ[:, 0]
        if k == 1:
            logodds1 = logodds
        p_nb = nb_predict_proba(model, Xte)
        p_lr = lr_fit_predict(X, y, Xte)
        for name, p in (("NB", p_nb), ("LR", p_lr)):
            conf = np.maximum(p, 1.0 - p)
            hi = conf > 0.999
            acc_hi = f"{nb_accuracy(p[hi], yte[hi]):.4f}" if np.any(hi) else "-"
            summary[(k, name)] = (nb_accuracy(p, yte), log_loss(yte, p), expected_calibration_error(yte, p),
                                  float(hi.mean()), acc_hi)
            print(f"{k:>3}{name:>8}{summary[(k, name)][0]:>10.4f}{summary[(k, name)][1]:>10.3f}"
                  f"{summary[(k, name)][2]:>8.4f}{conf.mean():>11.3f}{hi.mean():>12.3f}{acc_hi:>18}")
        if k > 1:
            print(f"{'':>11}NB log-odds vs k × (k = 1 log-odds): max |diff| = {np.max(np.abs(logodds - k * logodds1)):.1e}")
    print(f"""
观察 Observations：
  · NB 的准确率在所有 k 下都是 {summary[(1, 'NB')][0]:.4f}，一位都不差：复制后每个词的估计概率恰好变成原来的 1/k，
    log-odds 恰好变成原来的 k 倍（上面的 max |diff| 是舍入误差），符号不变，所以分类结果完全不变。
  · 但概率质量崩了：k = 1 时 NB 校准得很好（ECE = {summary[(1, 'NB')][2]:.4f}，log loss = {summary[(1, 'NB')][1]:.3f}）；
    k = 8 时 {summary[(8, 'NB')][3]:.0%} 的预测“99.9% 以上确定”，其中却只有 {summary[(8, 'NB')][4]} 是对的，
    log loss 涨到 {summary[(8, 'NB')][1]:.3f}，ECE = {summary[(8, 'NB')][2]:.4f}。
    这就是“错误的独立假设仍然分类得不错，但给出过度自信的概率”：分类只要求 log-odds 的**符号**对，
    概率要求 log-odds 的**大小**对。真实文本里的相关（短语、话题、同一个词反复出现）比这里弱，方向相同。
  · 逻辑回归直接拟合 p(y|x)，k 个一起出现的特征只会分摊同一份权重，不会被重复计数：k = 8 时 log loss
    = {summary[(8, 'LR')][1]:.3f}，ECE = {summary[(8, 'LR')][2]:.4f}。它也有变化：固定的 L2 = {LR_L2:g} 分摊到 k 个副本上，
    实际的正则化变弱了 k 倍，所以方差变大、准确率从 {summary[(1, 'LR')][0]:.4f} 降到 {summary[(8, 'LR')][0]:.4f}；
    但概率质量远没有 NB 那么极端。
  · 实践：NB 是很好的**排序 / 分类**器（垃圾邮件过滤、文本分类的强 baseline），但别把它的输出直接当概率用；
    需要概率时做一次校准（Platt scaling、isotonic regression），或者换判别式模型。""")


def main():
    demo_event_models()
    demo_laplace()
    demo_overconfidence()


if __name__ == "__main__":
    main()