"""
=============================================================================
第 1 讲 · 代码 4：指数族、广义线性模型与 softmax 回归（exponential family, GLMs, softmax）
=============================================================================

对应讲义：lectures/lecture01_linear_models.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
线性回归和逻辑回归的梯度都是 (h(x) − y)·x。这是巧合吗？不是 —— 它们都是 GLM。

    演示 1：指数族的关键性质 a'(η) = E[y]、a''(η) = Var[y]，用数值微分 vs 采样验证（三个分布）。
    演示 2：一个函数拟合三种模型。fit_glm_newton 换一个 family 参数，就分别得到线性回归、逻辑回归、Poisson 回归，
            结果和专门的实现逐位一致。
    演示 3：计数数据为什么要用 Poisson 回归而不是线性回归：负的预测、乘法效应、方差随均值增长。
    演示 4：Softmax 回归：多分类；参数冗余（所有类加同一个向量，预测不变）；K = 2 时就是逻辑回归。
    演示 5：为什么输出层的激活函数要和损失函数"配对"（sigmoid + cross-entropy，而不是 sigmoid + MSE）：
            canonical link 让 logits 上的梯度恰好是 (预测 − 目标)，不会因为饱和而消失。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_glm_softmax.py          # 约 1 秒

依赖本目录的 datasets.py 和 linear_models.py。
"""

import numpy as np

from datasets import POISSON_TRUE_THETA, make_blobs, make_housing, make_poisson
from linear_models import (FAMILIES, Standardizer, accuracy, add_intercept, fit_glm_newton, fit_least_squares,
                           fit_logistic_newton, fit_softmax_gd, mse, print_header, sigmoid, softmax, softmax_loss)


def fmt(v, spec=".4f"):
    return "[" + ", ".join(format(float(x), spec) for x in v) + "]"


# =============================================================================
# 演示 1：a'(η) = 均值，a''(η) = 方差
# =============================================================================
def sample_family(name, eta, size, rng):
    """从 natural parameter 为 eta 的分布里采样（T(y) = y）。"""
    if name == "gaussian":
        return rng.normal(eta, 1.0, size=size)
    if name == "bernoulli":
        return (rng.uniform(size=size) < sigmoid(eta)).astype(float)
    if name == "poisson":
        return rng.poisson(np.exp(eta), size=size).astype(float)
    raise ValueError(name)


def demo_exponential_family():
    print_header("演示 1：指数族 p(y; η) = b(y)·exp(η·y − a(η)) —— a' 是均值，a'' 是方差")
    rng = np.random.default_rng(41)
    h1, h2 = 1e-5, 1e-4
    print("a'(η)、a''(η) 用中心差分（central differences）从 a(η) 算出；均值、方差从 200000 个样本估计。\n")
    print(f"{'family':<11}{'η':>6}{'a′(η)':>12}{'sample mean':>14}{'a″(η)':>12}{'sample var':>13}")
    for name in ["gaussian", "bernoulli", "poisson"]:
        a = FAMILIES[name]["a"]
        for eta in [-1.0, 0.5, 2.0]:
            d1 = (a(eta + h1) - a(eta - h1)) / (2 * h1)
            d2 = (a(eta + h2) - 2 * a(eta) + a(eta - h2)) / h2 ** 2
            ys = sample_family(name, eta, 200000, rng)
            print(f"{name:<11}{eta:>6.1f}{d1:>12.4f}{ys.mean():>14.4f}{d2:>12.4f}{ys.var():>13.4f}")
    print("""
观察 Observations：
  · 每一行的 a′(η) 都等于样本均值、a″(η) 都等于样本方差，差别只是采样误差（大多在 0.003 以内；
    Poisson η = 2 时方差约 7.4，方差的估计误差也大一些，差了 0.04）。
  · 这不是巧合：a(η) = log ∫ b(y) e^{ηy} dy 是"把分布归一化"的对数常数，它的导数天然生成矩（cumulants）。
  · 推论 1：a″(η) = 方差 ≥ 0 ⇒ a 是凸函数 ⇒ 负对数似然 a(η) − yη 对 η 是凸的；η = θᵀx 是线性的
    ⇒ 对 θ 也是凸的。所以每个 canonical GLM 都是凸优化问题：没有局部最优，Newton 法放心用。
  · 推论 2：a′ 把 natural parameter 映射到均值，这就是 GLM 的 response function：
    gaussian → 恒等（identity），bernoulli → sigmoid，poisson → exp。sigmoid 不是"拍脑袋选的"，是推出来的。""")


# =============================================================================
# 演示 2：一个函数，三种模型
# =============================================================================
def demo_one_function():
    print_header("演示 2：One function, three models —— fit_glm_newton(X, y, family)")
    rng = np.random.default_rng(42)

    X, y = make_housing(200, rng)
    Xh = add_intercept(Standardizer().fit(X).transform(X))
    th_glm, it = fit_glm_newton(Xh, y, "gaussian")
    th_ref = fit_least_squares(Xh, y)
    print(f"gaussian  （房价）  Newton {it} 步；与 lstsq 的最大差 = {np.max(np.abs(th_glm - th_ref)):.1e}")

    Xb, yb = make_blobs(200, [[0.0, 0.0], [2.0, 1.0]], rng)
    Xb = add_intercept(Xb)
    th_glm, it = fit_glm_newton(Xb, yb, "bernoulli")
    th_ref, _ = fit_logistic_newton(Xb, yb)
    print(f"bernoulli （两团）  Newton {it} 步；与 fit_logistic_newton 的最大差 = {np.max(np.abs(th_glm - th_ref)):.1e}")

    Xp, yp = make_poisson(2000, rng)
    Xp = add_intercept(Xp)
    hist, steps = [], []
    th_glm, it = fit_glm_newton(Xp, yp, "poisson", history=hist, step_sizes=steps)
    print(f"poisson   （访问量）Newton {it} 步；θ̂ = {fmt(th_glm)}，真值 {fmt(POISSON_TRUE_THETA)}")
    print("          每一步开始时的 ‖∇‖_∞：" + "  ".join(f"{g:.1e}" for g in hist))
    print("          line search 接受的步长：" + "  ".join(f"{a:g}" for a in steps))
    hist_pure = []
    with np.errstate(over="ignore"):
        th_pure, it_pure = fit_glm_newton(Xp, yp, "poisson", history=hist_pure, line_search=False)
    print(f"          对比：纯 Newton（不做 line search）要 {it_pure} 步；它每一步的 ‖∇‖_∞：")
    print("          " + "  ".join(f"{g:.0e}" for g in hist_pure))
    print("""
观察 Observations：
  · 三个模型共用一段代码：梯度 (1/n)·Xᵀ(a′(Xθ) − y)，Hessian (1/n)·Xᵀ diag(a″(Xθ)) X。唯一不同的是 a′、a″。
    gaussian 时 a″ ≡ 1，Hessian 是常数，Newton 一步就到（这一步就是 normal equation）。
  · Poisson 的第一步被 line search 缩短到 1/4：从 θ = 0 出发，满步的 Newton 会冲过头（exp 增长太快，
    目标函数反而变大）。之后每一步都是满步，进入二次收敛。
    纯 Newton 第一步就把梯度从 9 推到 10⁵ 量级，再花十来步才爬回来 —— 这次没有发散，但在更极端的数据上
    exp 可能直接溢出。Newton 的二次收敛只在"离最优点足够近"时成立；远处需要 line search 之类的保护。
  · 这就是 statsmodels / R 的 glm(family = ...) 的工作方式：一个 IRLS 求解器，换 family 参数。
  · 也是 PyTorch 里 nn.MSELoss / nn.BCEWithLogitsLoss / nn.PoissonNLLLoss / nn.CrossEntropyLoss 这一族损失的来历。""")


# =============================================================================
# 演示 3：计数数据 —— Poisson vs 线性回归
# =============================================================================
def demo_poisson_vs_linear():
    print_header("演示 3：计数数据（每天的访问量）—— Poisson regression vs linear regression")
    rng = np.random.default_rng(43)
    X_tr, y_tr = make_poisson(300, rng)
    X_te, y_te = make_poisson(5000, rng)
    A_tr, A_te = add_intercept(X_tr), add_intercept(X_te)
    th_lin = fit_least_squares(A_tr, y_tr)
    th_poi, _ = fit_glm_newton(A_tr, y_tr, "poisson")
    pred_lin, pred_poi = A_te @ th_lin, np.exp(A_te @ th_poi)
    lam_te = np.exp(A_te @ POISSON_TRUE_THETA)
    print("真实模型：λ = exp(0.3 + 1.1·ad_budget − 0.7·weekend)，y ~ Poisson(λ)。训练 300 天，测试 5000 天。\n")
    print(f"{'model':<20}{'theta':<30}{'test MSE':>10}{'MSE vs true λ':>15}{'negative preds':>16}")
    print(f"{'linear regression':<20}{fmt(th_lin, '.3f'):<30}{mse(y_te, pred_lin):>10.3f}{mse(lam_te, pred_lin):>15.3f}"
          f"{np.mean(pred_lin < 0):>16.1%}")
    print(f"{'Poisson regression':<20}{fmt(th_poi, '.3f'):<30}{mse(y_te, pred_poi):>10.3f}{mse(lam_te, pred_poi):>15.3f}"
          f"{np.mean(pred_poi < 0):>16.1%}")
    print(f"\n  Poisson 系数的解读：预算每 +1，访问量 ×e^{th_poi[1]:.3f} = ×{np.exp(th_poi[1]):.2f}；"
          f"周末 ×e^{th_poi[2]:.3f} = ×{np.exp(th_poi[2]):.2f}")
    print(f"  线性系数的解读：  预算每 +1，访问量 +{th_lin[1]:.2f}；周末 {th_lin[2]:+.2f}（不论预算多少，都是同一个加减量）")

    print("\n  工作日（weekend = 0）不同预算下的预测：")
    print(f"    {'ad_budget':>10}{'true λ':>10}{'linear':>10}{'Poisson':>10}")
    for b in [0.0, 0.5, 1.0, 1.5, 2.0, 2.5]:
        xb = np.array([1.0, b, 0.0])
        print(f"    {b:>10.1f}{np.exp(xb @ POISSON_TRUE_THETA):>10.2f}{xb @ th_lin:>10.2f}{np.exp(xb @ th_poi):>10.2f}")
    print("\n  按真实 λ 分组，看 y 的均值和方差（Poisson：方差 = 均值）：")
    edges = [0, 2, 5, 10, 25]
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (lam_te >= lo) & (lam_te < hi)
        print(f"    λ ∈ [{lo:>2}, {hi:>2})：{m.sum():>5} 天，y 的均值 {y_te[m].mean():>6.2f}，方差 {y_te[m].var():>6.2f}")
    print("""
观察 Observations：
  · 线性回归预测出了负的访问量 —— 对计数来说毫无意义。Poisson 的 exp 保证预测永远为正。
  · 真实关系是乘法的（预算加 1，访问量乘约 3 倍）。直线只能用一个"平均斜率"去近似一条凸的指数曲线：
    两端低估（预算 0 时甚至预测出负数）、中间高估。"MSE vs true λ"一列量化了这个系统性误差（model misspecification）：
    它和噪声无关，数据再多也消除不了。
  · 方差随均值增长（heteroscedasticity）：λ 大的日子噪声也大。（每组的方差略大于均值，是因为同一组里的 λ 本身
    也不完全相同 —— 全方差公式 Var y = E[Var(y|λ)] + Var(E[y|λ])。）最小二乘假设方差处处相同，对每个点一视同仁；
    Poisson 的 Newton 步里权重是 a″(η) = λ，自动考虑了"哪些点噪声大"。
  · 选 GLM 的思路：先问"y 是什么类型的量"——实数？0/1？类别？计数？正的连续值？—— 然后选对应的分布。""")


# =============================================================================
# 演示 4：Softmax 回归
# =============================================================================
def demo_softmax():
    print_header("演示 4：Softmax regression —— 多分类、参数冗余、K = 2 时等于逻辑回归")
    rng = np.random.default_rng(44)
    centers = [[0.0, 0.0], [3.0, 0.0], [1.5, 2.5]]
    X_tr, y_tr = make_blobs(100, centers, rng)
    X_te, y_te = make_blobs(1000, centers, rng)
    A_tr, A_te = add_intercept(X_tr), add_intercept(X_te)
    Theta, it = fit_softmax_gd(A_tr, y_tr, K=3, tol=1e-7)
    print(f"三个高斯团，每类 100 个训练点。GD {it} 步收敛。")
    print(f"  训练准确率 {accuracy(y_tr, np.argmax(A_tr @ Theta, axis=1)):.3f}，"
          f"测试准确率 {accuracy(y_te, np.argmax(A_te @ Theta, axis=1)):.3f}，训练 cross-entropy {softmax_loss(Theta, A_tr, y_tr):.4f}")
    print(f"  Θ（每一列是一个类的参数向量 θ_k）：\n{np.round(Theta, 3)}")
    print(f"  各列之和（每一行加起来）= {fmt(Theta.sum(axis=1), '.1e')}   ← GD 从 0 出发，这个和始终为 0")

    c = np.array([5.0, -3.0, 2.0])
    P1, P2 = softmax(A_te @ Theta), softmax(A_te @ (Theta + c[:, None]))
    print(f"\n  参数冗余：给每一列都加上同一个向量 c = {c}，预测概率的最大变化 = {np.max(np.abs(P1 - P2)):.1e}")

    print("\n  K = 2：softmax 回归 vs 逻辑回归（两团、不可分、不加正则化）")
    X2, y2 = make_blobs(200, [[0.0, 0.0], [2.0, 1.0]], np.random.default_rng(45))
    A2 = add_intercept(X2)
    Th2, it2 = fit_softmax_gd(A2, y2, K=2, tol=1e-10)
    th_lr, _ = fit_logistic_newton(A2, y2)
    diff = Th2[:, 1] - Th2[:, 0]
    print(f"    softmax 的 θ₁ − θ₀ = {fmt(diff)}（GD {it2} 步）")
    print(f"    逻辑回归的 θ       = {fmt(th_lr)}")
    print(f"    两者预测的 P(y = 1) 的最大差 = {np.max(np.abs(softmax(A2 @ Th2)[:, 1] - sigmoid(A2 @ th_lr))):.1e}")
    print("""
观察 Observations：
  · softmax(z)_k = e^{z_k} / Σ_j e^{z_j} 只依赖于 logits 之间的差。给所有 θ_k 加同一个 c，所有 logits 加同一个 cᵀx，
    概率不变 —— K 个参数向量里只有 K − 1 个是"自由"的。CS229 讲义把 θ_K 固定为 0 来去掉冗余；
    实践中通常保留全部 K 个（对称、好写），再用 L2 正则化或"从 0 出发的 GD"自动选出一个解。
  · K = 2 时，P(y = 1) = e^{z₁}/(e^{z₀} + e^{z₁}) = σ(z₁ − z₀) = σ((θ₁ − θ₀)ᵀx)：就是逻辑回归，参数是两列之差。
    GD 选出的对称解满足 θ₀ = −θ₁，所以 θ₁ = θ/2。""")


# =============================================================================
# 演示 5：为什么激活函数要和损失"配对"
# =============================================================================
def demo_canonical_pairing():
    print_header("演示 5：输出层的激活 + 损失要配对 —— canonical link 让梯度不会消失")
    print("一个 y = 1 的样本。比较 logit z 上的梯度 ∂loss/∂z：\n")
    print(f"{'z':>6}{'σ(z)':>10}{'sigmoid + BCE: σ(z) − y':>27}{'sigmoid + MSE: (σ − y)·σ(1 − σ)':>35}")
    for z in [-10.0, -5.0, -2.0, 0.0, 2.0, 5.0]:
        p = float(sigmoid(z))
        g_bce = p - 1.0
        g_mse = (p - 1.0) * p * (1.0 - p)                   # d/dz ½(σ(z) − y)² = (σ − y)·σ′(z)
        print(f"{z:>6.1f}{p:>10.5f}{g_bce:>27.5f}{g_mse:>35.2e}")

    rng = np.random.default_rng(46)
    eta = rng.normal(size=5)
    y = {"gaussian": rng.normal(size=5), "bernoulli": rng.integers(0, 2, size=5).astype(float),
         "poisson": rng.poisson(2.0, size=5).astype(float)}
    print("\n数值验证：对每个 family，∂/∂η [a(η) − y·η] 是否等于 (mean(η) − y)？（5 个随机的 η、y）")
    for name, fam in FAMILIES.items():
        h = 1e-6
        num = (fam["a"](eta + h) - y[name] * (eta + h) - fam["a"](eta - h) + y[name] * (eta - h)) / (2 * h)
        print(f"    {name:<10} max |数值梯度 − (mean − y)| = {np.max(np.abs(num - (fam['mean'](eta) - y[name]))):.1e}")
    print("""
观察 Observations：
  · sigmoid + cross-entropy（canonical pairing）：错得越离谱（z = −10，模型几乎确定 y = 0），梯度越大（≈ −1），
    纠错的力度最大。sigmoid + MSE：σ′(z) 在饱和区几乎为 0，错得最离谱的时候梯度反而只有 −4.5e−5 —— 学不动。
  · 一般地，对任何 canonical GLM，负对数似然对 η（logit）的梯度就是 mean(η) − y，σ′ 之类的因子被抵消掉了。
    所以深度学习里的标准搭配是：回归 → 线性输出 + MSE；二分类 → sigmoid + BCE；多分类 → softmax + CE；
    计数 → exp + Poisson NLL。PyTorch 甚至把它们合并成一个函数（BCEWithLogitsLoss、CrossEntropyLoss 直接吃 logits），
    既为了这个干净的梯度，也为了数值稳定。
  · 换个角度说：一个神经网络 = 学出来的特征 φ(x) + 一个 GLM 输出头。第 4 讲的 backprop 就从这个 (ŷ − y) 开始往回传。""")


def main():
    demo_exponential_family()
    demo_one_function()
    demo_poisson_vs_linear()
    demo_softmax()
    demo_canonical_pairing()


if __name__ == "__main__":
    main()