"""【未验证 · Transformer 档】真实 MNIST 训练玩具 ViT + S 结构分区函数化。

一句话结论：
    Transformer 的注意力头是「并行功能模块」，不能像小 MLP 那样按权重几何
    混聚，而应按结构边界（头/层）直接切分。本例在 MNIST 上训练一个玩具 ViT，
    抽出注意力 Q 投影权重做 S 结构分区（按头切、头间不混聚），走完整
    「训练 → 模型 → 函数化 → 生成物 → 部署使用」链路，并对照「G 几何分区会
    跨头混聚」的错误示范。真实大 Transformer（ViT/LLM）精度未验证。

对应分型表第三行：Transformer = 结构分区 S（接口已实现，待真实 Transformer 验证）。
"""

from _common import (accuracy, device, load_mnist, make_drop_guard,
                     show_artifacts, show_value, train_model)  # 先导入：完成 src 路径引导

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import nrfunc


class TransformerBlock(nn.Module):
    """玩具 ViT 的一个 Transformer 块（pre-norm，手写注意力）。

    手写 Q/K/V/O 四个 Linear（而非 nn.MultiheadAttention），是为了让 Q 投影
    权重 `wq.weight` 是标准 (d_model, d_model)，从而能干净地按「头边界」切分：
    每 head_dim 个输出神经元属一个头。
    """

    def __init__(self, d_model, n_heads, ffn):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.fc1 = nn.Linear(d_model, ffn)
        self.fc2 = nn.Linear(ffn, d_model)

    def _attn(self, x):
        B, T, D = x.shape
        h, hd = self.n_heads, self.head_dim
        q = self.wq(x).reshape(B, T, h, hd).transpose(1, 2)   # (B, h, T, hd)
        k = self.wk(x).reshape(B, T, h, hd).transpose(1, 2)
        v = self.wv(x).reshape(B, T, h, hd).transpose(1, 2)
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(hd)
        attn = attn.softmax(-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
        return self.proj(out)

    def forward(self, x):
        x = x + self._attn(self.norm1(x))
        x = x + self.fc2(F.relu(self.fc1(self.norm2(x))))
        return x


class ToyViT(nn.Module):
    """玩具 ViT：把 28×28 MNIST 切 7×7 patch，class token + 2 层 Transformer。"""

    def __init__(self, patch=7, d_model=64, n_heads=4, n_layers=2, ffn=128,
                 n_cls=10):
        super().__init__()
        self.patch = patch
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        n_patches = (28 // patch) ** 2
        self.patch_embed = nn.Linear(patch * patch, d_model)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, d_model))
        self.blocks = nn.ModuleList(
            [TransformerBlock(d_model, n_heads, ffn) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, n_cls)

    def forward(self, x):                                   # x: (B, 1, 28, 28)
        B = x.shape[0]
        x = F.unfold(x, kernel_size=self.patch, stride=self.patch)  # (B, p*p, n_patches)
        x = x.transpose(1, 2)                               # (B, n_patches, p*p)
        x = self.patch_embed(x)                             # (B, n_patches, d_model)
        cls = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls, x], dim=1)                      # (B, n_patches+1, d_model)
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return self.head(x[:, 0])                           # (B, n_cls)


def main():
    dev = device()
    print(f'设备：{dev}')
    # 固定随机种子：让「玩具 ViT 初始化 + 训练」可复现，掉点数字稳定可比
    #（train_model 内部的 torch.manual_seed 在模型创建之后，管不到模型初始化）。
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)

    # 完整链路：训练 → 模型 → 函数化 → 生成物 → 部署使用
    #  ① 训练 + ② 模型（上游）：MNIST 训练玩具 ViT（7×7 patch，64 维，4 头，2 层）
    train_x, train_y, test_x, test_y = load_mnist()
    train_x = train_x[:20000].reshape(-1, 1, 28, 28)
    test_x4 = test_x.reshape(-1, 1, 28, 28)
    train_y = train_y[:20000]

    print('\n① 训练玩具 ViT（patch=7, d_model=64, 4 头, 2 层, 20000 样本）：')
    model = train_model(ToyViT(), train_x, train_y, epochs=10, seed=0, dev=dev)
    acc_before = accuracy(model, test_x4, test_y, dev)
    print(f'② 函数化前 · 测试准确率 = {acc_before * 100:.2f}%')

    # ③ 函数化（库）：抽出第 1 个块的 Q 投影权重 (64, 64)，做 S 结构分区（按头切）
    W = model.blocks[0].wq.weight.detach().cpu().numpy().astype(np.float64)
    n_heads, head_dim = model.n_heads, model.head_dim
    groups = np.repeat(np.arange(n_heads), head_dim)        # 每 head_dim 行属一个头

    # 主链路：auto_alloc 自动选配置（S 结构分区，K 固定=头数；绝不反涨 + 掉点<3pt 红线）
    guard = make_drop_guard(model, model.blocks[0].wq, W.shape, test_x4, test_y,
                            dev, acc_before, tol_pt=3.0)
    alloc_s = nrfunc.auto_alloc(W, signal='S', groups=groups, eval_fn=guard, seed=0)
    print(f'\n③ [S 结构分区 · auto_alloc 自动选] decision={alloc_s["decision"]} '
          f'order={alloc_s["order"]} K={alloc_s["K"]}（= 头数）r={alloc_s["r"]}，'
          f'内存省 {alloc_s["saving"]*100:.1f}%')
    if alloc_s['decision'] == 'skip':
        print(f'  → 该层被判跳过：{alloc_s["skipped_reason"]}')
        return
    res_s = alloc_s['result']
    print(f'   解释方差 = {res_s["explained"]:.4f}')
    print('   注：S 分区的解释方差是「按头内 16 行取低秩主成分」算的，天然低于'
          '全局 k-means（头内低秩只保留 ~49% 方差）；但下面⑤部署掉点通常 <1pt，'
          '说明被丢掉的方差对分类不敏感。')
    ok = bool(np.all(res_s['assign'] == groups))
    print(f'   归属是否严格等于结构边界（不混聚）：{ok}')

    # 对照：误用 G 几何分区会把不同头混在一起聚（错误示范）
    res_g = nrfunc.regionify(W, signal='G', K=n_heads, order=1, r=4, seed=0)
    mixed = {k: sorted(set((res_g['assign'] == k).nonzero()[0] // head_dim))
             for k in range(n_heads)}
    print(f'   [对照·G 几何分区] 各区实际包含的头（应出现跨头混聚）：{mixed}')

    # ④ 生成物（库的活）：量化打包 + 分片存储
    print('\n④ 生成物（量化打包 + 分片存储）：')
    show_artifacts(res_s, bits=8, n_shards=4)

    # ⑤ 部署使用（下游）：还原 → 塞回 wq → 推理测准确率
    recon = nrfunc.reconstruct(res_s)
    with torch.no_grad():
        model.blocks[0].wq.weight.copy_(torch.as_tensor(recon, dtype=torch.float32))
    acc_after = accuracy(model, test_x4, test_y, dev)
    print(f'\n⑤ 部署后 · 测试准确率 = {acc_after * 100:.2f}% '
          f'（掉 {acc_before * 100 - acc_after * 100:.2f}pt）')

    # 价值证明 · 四维（体积/内存/CPU/GPU）；wq 输入维度 = d_model=64
    x_batch = np.random.default_rng(0).standard_normal(
        (256, model.d_model)).astype(np.float32)
    print('\n价值对比（wq 层 64×64，函数化前 vs 后）：')
    show_value(W, res_s, x_batch)

    # GPT 级入口：只预留接口、不验证
    print('\n[GPT 级入口] regionify_transformer_block()：')
    try:
        nrfunc.regionify_transformer_block({})
    except NotImplementedError as e:
        print(f'  触发 NotImplementedError：{e}')

    print('\n[诚实声明] 本例用玩具 ViT 证明 S 结构分区「接口可跑 + 按头切正确 + '
          '生成物可还原」；是否在真实大 Transformer（ViT/LLM）上保精度，尚未验证。'
          'GPT 级大模型为「理论预演、不验证」，谁有条件谁验证。')


if __name__ == '__main__':
    main()