"""
=============================================================================
CMU 17-630 · 第 1 讲共享模块：attention_np.py —— 用 numpy 手写注意力，并手工搭一个"归纳头"
=============================================================================

【这个模块是做什么的】
    1. 缩放点积注意力（scaled dot-product attention）的最小实现：
            Attention(Q, K, V) = softmax(Q Kᵀ / sqrt(d) + mask) V
       直觉：一个**软的字典查找**。每个位置拿自己的 query 去和所有位置的 key 比，
       越像的位置权重越大，然后把那些位置的 value 按权重加起来带回来。
    2. 一个**不用训练、直接手工设定权重**的两层注意力网络，实现"归纳头"（induction head，
       Olsson et al. 2022）：看到 "... A B ... A" 时，预测下一个是 B。
            第 1 层（前一个词头）：每个位置去看它左边那个位置，把"前一个词是谁"抄到自己身上；
            第 2 层（归纳头）    ：拿"我是谁"去找"哪个位置的前一个词是我"，把那个位置的词抄回来。
       Anthropic 的研究发现，真实的 Transformer 在训练早期会自己长出这样的电路，而且它的出现
       和"上下文学习"能力的突然提升同时发生。第 2 讲（In-Context Learning）会回到它。

【残差流的布局（手工搭建时用）】
    每个位置的向量是三段拼起来的（one-hot 编码，维度 = V + T + V）：
        [ 当前词 tok (V 维) | 位置 pos (T 维) | 前一个词 prev (V 维，第 1 层写进来) ]
    真实模型里这些"子空间"不是分开的整块，而是学出来的、互相交叠的方向；这里分开是为了能看清楚。

【运行】python3 attention_np.py     # 自检：在随机重复序列上测手工归纳头
【依赖】numpy。04 号脚本会 import 它。
"""

import numpy as np


def softmax(x, axis=-1):
    """
    数值稳定的 softmax：先减去最大值再取指数（结果不变，但不会溢出）。

    【边界】某一行全是 −inf（完全被遮住）时，结果是 nan；本模块的因果遮罩总会留下自己这一格，不会出现这种情况。
    """
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=axis, keepdims=True)


def causal_mask(T):
    """
    因果遮罩：T × T 矩阵，右上方（未来的位置）是 −inf，其余是 0。加到分数上再做 softmax，未来的位置权重就是 0。

    【Python 小知识】np.triu(A, k=1) 取矩阵 k=1 那条对角线以上的部分（不含主对角线）。
    """
    return np.triu(np.full((T, T), -np.inf), k=1)


def attention(Q, K, V, mask=None, scale=True):
    """
    缩放点积注意力。

    【参数】Q：(T, d) 的 query；K：(T, d) 的 key；V：(T, d_v) 的 value；
        mask：加到分数上的 (T, T) 矩阵（例如 causal_mask(T)）；scale：是否除以 sqrt(d)
    【返回】(输出 (T, d_v), 注意力权重 (T, T)，每行和为 1)
    【为什么除以 sqrt(d)】如果 q、k 的每一维是均值 0、方差 1 的随机数，点积 q·k 的方差是 d。
        d 很大时分数的差距很大，softmax 几乎变成"只看一个位置"（one-hot），梯度也随之消失。
        除以 sqrt(d) 把方差拉回 1（Vaswani et al. 2017 §3.2.1）。04 号脚本演示 2 会量化这一点。
    """
    d = Q.shape[-1]
    scores = Q @ K.T / (np.sqrt(d) if scale else 1.0)
    if mask is not None:
        scores = scores + mask
    W = softmax(scores, axis=-1)
    return W @ V, W


class HandBuiltInduction:
    """
    手工设定权重的两层注意力网络：前一个词头 + 归纳头。

    【参数】V：词表大小；T：最大序列长度；beta：注意力的"锐利度"（分数乘的倍数，越大越接近硬查找）；
        use_pos：是否给模型位置信息（False 时第 1 层不知道"左边那个位置"是哪个）；
        layers：2 = 完整电路；1 = 只有归纳头、没有前一个词头（用来证明两层缺一不可）
    【用法】HandBuiltInduction(V=20, T=40).predict(seq) → 每个位置对下一个词的预测
    """

    def __init__(self, V, T, beta=20.0, use_pos=True, layers=2):
        self.V, self.T, self.beta, self.use_pos, self.layers = V, T, beta, use_pos, layers

    def embed(self, seq):
        """
        把 token 序列编码成残差流 (T, V + T + V)：当前词 one-hot + 位置 one-hot + 前一个词（先全 0）。
        """
        n = len(seq)
        X = np.zeros((n, 2 * self.V + self.T))
        X[np.arange(n), seq] = 1.0                           # tok 子空间
        if self.use_pos:
            X[np.arange(n), self.V + np.arange(n)] = 1.0     # pos 子空间
        return X

    def layer1_prev_token(self, X):
        """
        第 1 层：前一个词头。

        【思路】query 取位置 t 的 one-hot，key 取位置 s 的 one-hot 往后挪一格（代表 s+1），
            于是 q_t · k_s = 1 当且仅当 s = t − 1。value 取 s 的当前词，写进 prev 子空间。
        【边界】位置 0 左边没有词：因果遮罩下它只能看自己，会把"自己"当成前一个词。第一版就是这样，
            结果在重复的第一个词上，归纳头同时匹配到位置 0 和位置 1，两者分数只差 1e-8，
            每次都选错——后一半准确率卡在 8/9。现在把位置 0 的 prev 清零（"没有前一个词"）。
            真实模型用一个开头的特殊 token（如 GPT-2 的 <|endoftext|>）来处理"开头没东西可看"的问题。
        【返回】更新后的残差流（X + 注意力输出）、注意力权重
        """
        n, V, T = X.shape[0], self.V, self.T
        P = X[:, V:V + T]                         # 位置 one-hot
        Q = self.beta * P
        K = np.zeros_like(P)
        K[:, 1:] = P[:, :-1]                      # 位置 s 的 key = one-hot(s + 1)
        out_tok, W = attention(Q, K, X[:, :V], causal_mask(n), scale=False)
        out_tok[0] = 0.0                          # 位置 0 没有前一个词
        X = X.copy()
        X[:, V + T:] += out_tok                   # 写进 prev 子空间
        return X, W

    def layer2_induction(self, X):
        """
        第 2 层：归纳头。

        【思路】query = β · 当前词的 one-hot；key = 每个位置的 prev 子空间（"我前面那个词是谁"）。
            q_t · k_s 大，当且仅当"位置 s 前面的词 == 位置 t 的词"——也就是 s 紧跟在当前词的上一次出现之后。
            value = 位置 s 的当前词。于是输出 ≈ "上次在当前词后面出现的那个词"。
            为了不让一个位置"找到自己"，额外遮住对角线（s = t）。
        【返回】输出 (T, V)（每个位置对下一个词的打分）、注意力权重
        """
        n, V, T = X.shape[0], self.V, self.T
        Q = self.beta * X[:, :V]
        K = X[:, V + T:]
        mask = causal_mask(n)
        mask[np.arange(1, n), np.arange(1, n)] = -np.inf   # 不许看自己（位置 0 除外，否则整行全被遮住）
        out, W = attention(Q, K, X[:, :V], mask, scale=False)
        return out, W

    def predict(self, seq):
        """
        对序列的每个位置预测下一个词。

        【返回】(预测的词 id 数组, 第 2 层的注意力权重)
        【边界】序列长度不能超过 T。没有"上一次出现"的位置，注意力会平摊在许多位置上，预测基本是随机的。
        """
        assert len(seq) <= self.T
        X = self.embed(np.asarray(seq))
        if self.layers >= 2:
            X, _ = self.layer1_prev_token(X)
        out, W = self.layer2_induction(X)
        return out.argmax(axis=1), W


def repeated_sequences(V, half, n, seed=0):
    """
    n 条"随机序列重复两遍"的测试数据：前一半是 half 个互不相同的随机词，后一半原样重复。

    【返回】整数数组 (n, 2·half)
    【边界】half 不能超过 V（要求前一半的词互不相同）。
    """
    rng = np.random.default_rng(seed)
    first = np.stack([rng.permutation(V)[:half] for _ in range(n)])
    return np.concatenate([first, first], axis=1)


def copy_accuracy(predict_fn, data):
    """
    在重复序列上，模型对"下一个词"的预测准确率，前一半和后一半分开算。

    【返回】(前一半的准确率, 后一半的准确率)
    【思路】后一半位置 t 的下一个词，正是前一半里同一个词后面那个词——能"回看并抄写"的模型应该接近 100%；
        前一半的下一个词是全新的随机词，任何模型都只能瞎猜（≈ 1/V）。
    """
    half = data.shape[1] // 2
    ok1 = ok2 = n1 = n2 = 0
    for seq in data:
        pred = predict_fn(seq)
        for t in range(len(seq) - 1):
            hit = pred[t] == seq[t + 1]
            if t < half - 1:
                ok1, n1 = ok1 + hit, n1 + 1
            elif t >= half:
                ok2, n2 = ok2 + hit, n2 + 1
    return ok1 / n1, ok2 / n2


if __name__ == "__main__":
    data = repeated_sequences(V=20, half=10, n=200)
    for name, m in [("完整电路", HandBuiltInduction(20, 40)),
                    ("没有位置信息", HandBuiltInduction(20, 40, use_pos=False)),
                    ("只有第 2 层", HandBuiltInduction(20, 40, layers=1))]:
        a1, a2 = copy_accuracy(lambda s: m.predict(s)[0], data)
        print(f"  {m.__class__.__name__}（{name}）：前一半 {a1:.3f}，后一半 {a2:.3f}")