"""
nn_lab.py —— 用 NumPy 从零搭神经网络：层、损失、优化器、梯度检查和一个影评小语料（前置讲 P3 的工具箱）
=====================================================================

【这个文件要回答的问题】
    PyTorch 里一行 `loss.backward()` 背后发生了什么？一个最小的文本分类网络由哪些"积木"组成？
    这个模块把每块积木都写成"前向 + 反向"两个函数，梯度全部手推，并提供数值梯度来检查。
    第 5 讲（深度学习框架与自动微分）会把"手写反向"变成"自动反向"——先在这里看清楚它要自动完成的是什么。

【本文件提供】
    1. 数据      : make_regression、make_spiral、ReviewWorld（影评小语料：单形容词 / 否定 / 转折 / 并列四种句型）
    2. 层        : linear、relu、tanh、gelu、layer_norm、dropout 的前向与反向；embedding_backward（散射相加）
    3. 损失      : softmax、log_softmax、cross_entropy（返回损失和对 logits 的梯度）
    4. 优化器    : SGD、Momentum、Adam（可选 AdamW 式的解耦权重衰减）
    5. 检查      : numerical_grad、rel_error
    6. 模型      : TextClassifier（嵌入查表 → 平均池化 → [线性 → ReLU] → 线性 → softmax）

【约定】
    - 批量数据的第一维是样本：x 的形状是 [B, d]。线性层 y = x W + b，W 的形状是 [d_in, d_out]。
    - 反向函数的输入是"损失对本层输出的梯度"（记作 dy），输出是"损失对本层输入和参数的梯度"。
    - 所有随机性都来自显式传入的 np.random.Generator（便于复现）。

【系统联系】
    每个层最终都是几种底层运算：线性层 = 矩阵乘法，激活 = 逐元素映射，平均池化 / softmax / LayerNorm = 归约，
    嵌入 = 按下标读取（gather），嵌入的反向 = 按下标累加（scatter-add）。这正是原课第 2 讲开头的那张表，
    也是 HW1（map / zip / reduce / matmul 内核）和 HW2（自动微分）要实现的东西。
"""

import math

import numpy as np


# =====================================================================
# 1. 数据
# =====================================================================

def make_regression(n=256, d=3, noise=0.1, seed=0):
    """
    线性回归数据：y = X w_true + b_true + 噪声。

    【返回】(X [n, d], y [n], w_true [d], b_true)
    """
    rng = np.random.default_rng(seed)
    X = rng.standard_normal((n, d))
    w_true = rng.standard_normal(d)
    b_true = 0.5
    y = X @ w_true + b_true + noise * rng.standard_normal(n)
    return X, y, w_true, b_true


def make_spiral(n_per_class=200, classes=3, noise=0.2, seed=0):
    """
    螺旋数据：classes 条相互缠绕的螺旋臂，每条一类。任何一条直线都分不开它们，必须用非线性模型。

    【返回】(X [n, 2], y [n]，取值 0..classes−1)
    """
    rng = np.random.default_rng(seed)
    X, y = [], []
    for c in range(classes):
        r = np.linspace(0.05, 1.0, n_per_class)
        t = np.linspace(c * 4.0, (c + 1) * 4.0, n_per_class) + noise * rng.standard_normal(n_per_class)
        X.append(np.stack([r * np.sin(t), r * np.cos(t)], axis=1))
        y.append(np.full(n_per_class, c))
    return np.concatenate(X), np.concatenate(y)


class ReviewWorld:
    """
    影评小语料：每条影评由模板生成，情感标签按确定的规则给出（所以"正确答案"是完全清楚的）。

    【四种句型】
        single  : "the <名词> was [程度副词] <形容词>"            标签 = 形容词的极性
        negated : "the <名词> was not <形容词>"                   标签 = 形容词极性取反
        contrast: "the <名词1> was <形容词1> but the <名词2> was <形容词2>"   标签 = 形容词2 的极性（"but"之后才算数）
        both    : "a <形容词> <名词1> and a <形容词> <名词2>"       两个形容词极性相同，标签 = 该极性
    【为什么这样设计】
        - single / both：只要知道每个词的好坏就能做对——"词袋"（不管顺序、只看出现了哪些词）就够了。
        - negated：需要"not 和后面那个词的组合"，单个词的得分相加做不到（类似异或），需要非线性。
        - contrast：把两个从句对调，词袋完全一样、标签却相反——任何只看词袋的模型都最多猜对一半。
          这正是"需要看顺序"（RNN、注意力、Transformer）的动机。
    """

    POS = ["good", "great", "excellent", "wonderful", "fun", "brilliant", "lovely", "enjoyable"]
    NEG = ["bad", "terrible", "awful", "boring", "dull", "poor", "weak", "annoying"]
    NOUNS = ["movie", "film", "plot", "acting", "story", "script", "music", "ending", "cast", "dialogue"]
    ADVS = ["very", "really", "quite", "so"]
    TEMPLATES = [("single", 0.4), ("negated", 0.2), ("contrast", 0.2), ("both", 0.2)]
    PAD = "<pad>"

    def __init__(self):
        words = set(self.POS + self.NEG + self.NOUNS + self.ADVS + ["the", "was", "not", "but", "a", "and"])
        self.itos = [self.PAD] + sorted(words)          # id 0 是填充符
        self.stoi = {w: i for i, w in enumerate(self.itos)}
        self.polarity = {w: 1 for w in self.POS}
        self.polarity.update({w: -1 for w in self.NEG})

    def sample(self, rng):
        """
        生成一条影评。

        【返回】(token 列表, 标签 0/1（1 = 正面）, 句型名)
        """
        kinds = [k for k, _ in self.TEMPLATES]
        kind = rng.choice(kinds, p=[p for _, p in self.TEMPLATES])
        adj = lambda: rng.choice(self.POS + self.NEG)
        noun = lambda: rng.choice(self.NOUNS)
        if kind == "single":
            a = adj()
            adv = [rng.choice(self.ADVS)] if rng.random() < 0.5 else []
            toks = ["the", noun(), "was"] + adv + [a]
            label = self.polarity[a] > 0
        elif kind == "negated":
            a = adj()
            toks = ["the", noun(), "was", "not", a]
            label = self.polarity[a] < 0
        elif kind == "contrast":
            a1, a2 = adj(), adj()
            while self.polarity[a1] == self.polarity[a2]:      # 转折句里前后两个形容词的极性相反
                a2 = adj()
            n1 = noun()
            n2 = noun()
            while n2 == n1:
                n2 = noun()
            toks = ["the", n1, "was", a1, "but", "the", n2, "was", a2]
            label = self.polarity[a2] > 0
        else:
            pool = self.POS if rng.random() < 0.5 else self.NEG
            n1 = noun()
            n2 = noun()
            toks = ["a", rng.choice(pool), n1, "and", "a", rng.choice(pool), n2]
            label = pool is self.POS
        return [str(t) for t in toks], int(label), str(kind)

    def dataset(self, n, seed):
        """生成 n 条影评。【返回】(句子列表, 标签数组, 句型列表)"""
        rng = np.random.default_rng(seed)
        out = [self.sample(rng) for _ in range(n)]
        return [o[0] for o in out], np.array([o[1] for o in out]), [o[2] for o in out]

    def encode(self, sentences):
        """
        把一批句子编码成补齐的 id 矩阵和掩码。

        【返回】(ids [B, T] 整数，掩码 mask [B, T]（真实 token 为 1、填充为 0）)
        """
        T = max(len(s) for s in sentences)
        ids = np.zeros((len(sentences), T), dtype=np.int64)
        mask = np.zeros((len(sentences), T))
        for i, s in enumerate(sentences):
            ids[i, :len(s)] = [self.stoi[w] for w in s]
            mask[i, :len(s)] = 1.0
        return ids, mask


# =====================================================================
# 2. 层：前向与反向
# =====================================================================

def linear_forward(x, W, b):
    """y = x W + b。【形状】x [B, d_in]，W [d_in, d_out]，b [d_out] → y [B, d_out]（b 被广播到每一行）"""
    return x @ W + b


def linear_backward(dy, x, W):
    """
    线性层的反向。

    【推导】y = x W + b，逐元素写 y_ij = Σ_k x_ik W_kj + b_j，于是
        ∂L/∂W = xᵀ · dy     （每个权重的梯度 = 它连接的输入 × 输出的梯度，对批次求和）
        ∂L/∂b = Σ_行 dy     （偏置被广播到每一行，反向时把每一行的梯度加起来）
        ∂L/∂x = dy · Wᵀ     （传给前一层）
    【返回】(dx, dW, db)
    【系统联系】一次线性层的反向 = 两次矩阵乘法，各和前向一样大——这就是"训练 ≈ 3 倍前向 FLOPs"的来历（第 1 讲 §5.3）。
    """
    return dy @ W.T, x.T @ dy, dy.sum(axis=0)


def relu(x):
    """ReLU(x) = max(0, x)。"""
    return np.maximum(x, 0.0)


def relu_backward(dy, x):
    """ReLU 的导数：x > 0 处为 1，否则为 0（x = 0 处取 0）。"""
    return dy * (x > 0)


def tanh_backward(dy, y):
    """tanh 的导数 = 1 − tanh²（用前向的输出 y = tanh(x) 计算，不用再算一次 tanh）。"""
    return dy * (1.0 - y * y)


_GELU_C = math.sqrt(2.0 / math.pi)


def gelu(x):
    """GELU 的 tanh 近似：0.5 x (1 + tanh(√(2/π) (x + 0.044715 x³)))。GPT 系列用的激活函数。"""
    return 0.5 * x * (1.0 + np.tanh(_GELU_C * (x + 0.044715 * x ** 3)))


def gelu_backward(dy, x):
    """GELU（tanh 近似）的导数，按链式法则逐项求。"""
    u = _GELU_C * (x + 0.044715 * x ** 3)
    t = np.tanh(u)
    du = _GELU_C * (1.0 + 3 * 0.044715 * x ** 2)
    return dy * (0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * du)


def layer_norm(x, gamma, beta, eps=1e-5):
    """
    LayerNorm：对每一行（每个样本 / 每个 token）减均值、除标准差，再乘 γ 加 β。

    【返回】(y, cache)；cache 留给反向用
    【为什么】让每一层的输入保持"均值 0、方差 1"的尺度，深层网络才不容易数值爆炸或消失（01 号脚本第 4 部分）。
    【边界】eps 防止除以 0；方差用两遍法计算（先减均值），避免 P2 §3.5 的灾难性抵消。
    """
    mu = x.mean(axis=-1, keepdims=True)
    xc = x - mu
    var = (xc * xc).mean(axis=-1, keepdims=True)
    inv = 1.0 / np.sqrt(var + eps)
    xhat = xc * inv
    return gamma * xhat + beta, (xhat, inv, gamma)


def layer_norm_backward(dy, cache):
    """
    LayerNorm 的反向（原课 HW4 要写成 CUDA 内核的那个公式）。

    【推导要点】xhat = (x − μ) / σ。μ 和 σ 都依赖于这一行的所有元素，所以 dx 除了"直接项"，
        还要减去两个"通过均值和方差传回来"的修正项：
            dxhat = dy · γ
            dx = (1/σ) · (dxhat − mean(dxhat) − xhat · mean(dxhat · xhat))
    【返回】(dx, dgamma, dbeta)
    """
    xhat, inv, gamma = cache
    dxhat = dy * gamma
    dx = inv * (dxhat - dxhat.mean(axis=-1, keepdims=True) - xhat * (dxhat * xhat).mean(axis=-1, keepdims=True))
    red = tuple(range(dy.ndim - 1))
    return dx, (dy * xhat).sum(axis=red), dy.sum(axis=red)


def dropout(x, p, rng):
    """
    训练时随机把比例 p 的元素置 0，其余乘 1/(1−p)（保持期望不变）；推理时不做（调用者传 p=0）。

    【返回】(y, mask)；反向时 dx = dy * mask
    """
    if p <= 0:
        return x, np.ones_like(x)
    mask = (rng.random(x.shape) >= p) / (1.0 - p)
    return x * mask, mask


def embedding_backward(d_emb, ids, vocab_size):
    """
    嵌入层（查表）的反向：把每个位置的梯度"散射"加回它对应的那一行。

    【思路】前向 emb = E[ids]：同一个词出现几次，E 的那一行就被读了几次；反向时这些位置的梯度要**累加**到同一行。
    【Python 小知识】np.add.at(A, idx, vals) 会正确处理重复的下标（普通的 A[idx] += vals 遇到重复下标只加一次）。
    【系统联系】GPU 上这一步需要原子加法（多个线程可能同时写同一行），所以结果的加法顺序不固定（P2 §3.3）。
    """
    dE = np.zeros((vocab_size, d_emb.shape[-1]))
    np.add.at(dE, ids, d_emb)
    return dE


# =====================================================================
# 3. softmax 与交叉熵
# =====================================================================

def softmax(z):
    """
    数值稳定的 softmax：先减去每一行的最大值（结果不变，但 exp 不会溢出）。

    【参数】z：[..., C]，沿最后一维做 softmax
    """
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def log_softmax(z):
    """log softmax = z − logsumexp(z)，同样先减最大值。比"先 softmax 再取 log"更稳定（不会 log(0)）。"""
    m = z.max(axis=-1, keepdims=True)
    return z - m - np.log(np.exp(z - m).sum(axis=-1, keepdims=True))


def cross_entropy(logits, y):
    """
    平均交叉熵损失及其对 logits 的梯度。

    【公式】L = −(1/B) Σ_i log softmax(z_i)[y_i]；∂L/∂z = (softmax(z) − onehot(y)) / B
    【返回】(loss 标量, dlogits [B, C])
    【系统联系】梯度只需要 softmax 的输出和标签：第 10 讲 LightSeq 把 softmax + 交叉熵 + 求梯度融合成一个内核，
        靠的就是这个简单的公式。
    """
    B = logits.shape[0]
    lp = log_softmax(logits)
    loss = -lp[np.arange(B), y].mean()
    d = np.exp(lp)
    d[np.arange(B), y] -= 1.0
    return loss, d / B


# =====================================================================
# 4. 优化器
# =====================================================================

class SGD:
    """
    （小批量）随机梯度下降：θ ← θ − lr · g。

    【用法】opt = SGD(lr)；opt.step(params, grads)，params 和 grads 是同样键的字典，原地更新 params。
    """

    def __init__(self, lr):
        self.lr = lr

    def step(self, params, grads):
        for k in params:
            params[k] -= self.lr * grads[k]


class Momentum:
    """
    动量法：v ← β v + g；θ ← θ − lr · v。像一个有惯性的小球，在一致的方向上越滚越快，在来回震荡的方向上互相抵消。
    """

    def __init__(self, lr, beta=0.9):
        self.lr, self.beta, self.v = lr, beta, {}

    def step(self, params, grads):
        for k in params:
            self.v[k] = self.beta * self.v.get(k, 0.0) + grads[k]
            params[k] -= self.lr * self.v[k]


class Adam:
    """
    Adam：为每个参数维护梯度的一阶矩 m（方向）和二阶矩 v（尺度），步长 ≈ lr · m / √v。

        m ← β1 m + (1 − β1) g；v ← β2 v + (1 − β2) g²
        m̂ = m / (1 − β1^t)；v̂ = v / (1 − β2^t)        （偏差修正：开始时 m、v 从 0 起步，偏小）
        θ ← θ − lr · m̂ / (√v̂ + ε)  − lr · wd · θ       （最后一项是 AdamW 的"解耦权重衰减"，wd=0 时没有）
    【系统联系】每个参数要多存 m 和 v 两个数（FP32 下 8 字节）——第 1 讲 16 字节/参数的账里有一半是它们。
    """

    def __init__(self, lr, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
        self.lr, self.b1, self.b2, self.eps, self.wd = lr, betas[0], betas[1], eps, weight_decay
        self.m, self.v, self.t = {}, {}, 0

    def step(self, params, grads):
        self.t += 1
        for k in params:
            g = grads[k]
            self.m[k] = self.b1 * self.m.get(k, 0.0) + (1 - self.b1) * g
            self.v[k] = self.b2 * self.v.get(k, 0.0) + (1 - self.b2) * g * g
            mh = self.m[k] / (1 - self.b1 ** self.t)
            vh = self.v[k] / (1 - self.b2 ** self.t)
            params[k] -= self.lr * mh / (np.sqrt(vh) + self.eps)
            if self.wd:
                params[k] -= self.lr * self.wd * params[k]


# =====================================================================
# 5. 梯度检查
# =====================================================================

def numerical_grad(f, x, h=1e-5):
    """
    用中心差分估计 ∂f/∂x：[f(x + h) − f(x − h)] / (2h)，逐个元素地扰动。

    【参数】f：无参数函数，返回标量（内部读取 x 的当前值）；x：会被原地扰动再恢复的数组
    【返回】与 x 形状相同的梯度估计
    【边界】很慢（每个元素要算两次 f），只用来检查小模型；h 太大有截断误差，太小有舍入误差，1e-5 是 float64 下的常用值。
    """
    g = np.zeros_like(x)
    it = np.nditer(x, flags=["multi_index"])
    for _ in it:
        i = it.multi_index
        old = x[i]
        x[i] = old + h
        fp = f()
        x[i] = old - h
        fm = f()
        x[i] = old
        g[i] = (fp - fm) / (2 * h)
    return g


def rel_error(a, b):
    """相对误差 max|a − b| / max(|a| + |b|)，用来比较解析梯度和数值梯度（< 1e-6 通常说明写对了）。"""
    return float(np.max(np.abs(a - b)) / max(1e-12, np.max(np.abs(a) + np.abs(b))))


# =====================================================================
# 6. 文本分类模型
# =====================================================================

class TextClassifier:
    """
    原课第 2 讲开头的那个网络：嵌入查表 → 平均池化 → [线性 → ReLU] → 线性 → softmax。

    【参数】vocab_size、dim（嵌入维度）、hidden（隐藏层宽度；0 表示没有隐藏层，即"平均嵌入 + 线性分类器"）、
        classes、seed、dropout（加在隐藏层输出上的 dropout 比例）
    【用法】
        model = TextClassifier(V, dim=16, hidden=32)
        logits, cache = model.forward(ids, mask, rng=None)     # rng 为 None 表示推理（不做 dropout）
        loss, dlogits = cross_entropy(logits, y)
        grads = model.backward(dlogits, cache)
        opt.step(model.params, grads)
    """

    def __init__(self, vocab_size, dim=16, hidden=32, classes=2, seed=0, dropout=0.0):
        rng = np.random.default_rng(seed)
        self.hidden = hidden
        self.p_drop = dropout
        self.params = {"E": rng.normal(0, 0.1, (vocab_size, dim))}
        d_in = dim
        if hidden:
            self.params["W1"] = rng.normal(0, math.sqrt(2.0 / dim), (dim, hidden))     # He 初始化（适合 ReLU）
            self.params["b1"] = np.zeros(hidden)
            d_in = hidden
        self.params["W2"] = rng.normal(0, math.sqrt(1.0 / d_in), (d_in, classes))
        self.params["b2"] = np.zeros(classes)
        self.params["E"][0] = 0.0          # 填充符的嵌入固定为 0（反向时也不更新，见 backward）

    def forward(self, ids, mask, rng=None):
        """
        【返回】(logits [B, C], cache)
        【步骤】查表 emb = E[ids]（gather）→ 按掩码求平均（reduce）→ 线性（matmul）→ ReLU（map）→ 线性
        """
        P = self.params
        emb = P["E"][ids]                                    # [B, T, d]
        lengths = mask.sum(axis=1, keepdims=True)            # [B, 1]
        pooled = (emb * mask[:, :, None]).sum(axis=1) / lengths
        cache = {"ids": ids, "mask": mask, "lengths": lengths, "pooled": pooled}
        x = pooled
        if self.hidden:
            h = linear_forward(x, P["W1"], P["b1"])
            a = relu(h)
            a, dmask = dropout(a, self.p_drop if rng is not None else 0.0, rng)
            cache.update(h=h, a=a, dmask=dmask)
            x = a
        cache["x_last"] = x
        return linear_forward(x, P["W2"], P["b2"]), cache

    def backward(self, dlogits, cache):
        """按前向的逆序，逐层调用各自的反向函数。【返回】与 params 同键的梯度字典"""
        P = self.params
        g = {}
        dx, g["W2"], g["b2"] = linear_backward(dlogits, cache["x_last"], P["W2"])
        if self.hidden:
            dx = dx * cache["dmask"]
            dh = relu_backward(dx, cache["h"])
            dx, g["W1"], g["b1"] = linear_backward(dh, cache["pooled"], P["W1"])
        d_emb = dx[:, None, :] * cache["mask"][:, :, None] / cache["lengths"][:, :, None]
        g["E"] = embedding_backward(d_emb, cache["ids"], P["E"].shape[0])
        g["E"][0] = 0.0
        return g

    def predict(self, ids, mask):
        """推理：返回每个样本的类别（argmax）。"""
        logits, _ = self.forward(ids, mask)
        return logits.argmax(axis=1)

    def num_params(self):
        """参数总数。"""
        return sum(p.size for p in self.params.values())