"""【已验证 · MNIST MLP 档】图像分类 + G 几何分区函数化。

一句话结论：
    完整 MNIST（28×28=784）上训练一个 MLP（784→128→64→10），对其最大的首层
    （128×784）做 G 几何分区函数化（1 阶低秩，K=16/r=4），8bit 量化后约 6.4x
    字节压缩、测试准确率掉 ~2.6pt（<3pt 达标）——证明「MNIST MLP + G 几何」这一档
    真实可用。

诚实说明（很重要）：
    - 这 ~2.6pt 掉点就是「后训练压缩」的真实代价：k-means 是训练**后**才叠加的
      误差，没被训练吸收。本项目 P1/P2 的「≤0.13pt / 9.11x」是**训练期压缩**
      （蒸馏 + QAT，误差被训练吃掉）拿到的，那属于上游，不归本库。
    - 想要 <1pt：请在上游用 QAT/蒸馏把模型训出「可共享结构」，再交给本库函数化
      （即「叠加」用法）。本例刻意用未经 QAT 的裸模型，展示后训练代价。
    - 压缩率来自「n/K（分区）+ 量化（位宽）」。1 阶低秩只在 n > K×(1+r) 时净压缩；
      0 阶质心压缩更狠（n/K）但首层特征多样会灾难性掉点（见 mlp_classification.py）。

对应分型表第一行：MNIST MLP = 几何分区 G（本项目已实证）。
"""

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

import numpy as np
import torch
import torch.nn as nn

import nrfunc


class MnistMLP(nn.Module):
    def __init__(self, d_in=784, h1=128, h2=64, d_out=10):
        super().__init__()
        self.fc1 = nn.Linear(d_in, h1)
        self.fc2 = nn.Linear(h1, h2)
        self.fc3 = nn.Linear(h2, d_out)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


def main():
    dev = device()
    print(f'设备：{dev}')

    # 完整链路：训练 → 模型 → 函数化 → 生成物 → 部署使用
    #  ① 训练 + ② 模型（上游）：图像分类数据 + 训练 MLP（784→128→64→10）
    train_x, train_y, test_x, test_y = load_mnist()
    train_x, train_y = train_x[:40000], train_y[:40000]   # 子集求快

    print('训练 MNIST MLP（784→128→64→10）：')
    model = train_model(MnistMLP(), train_x, train_y, epochs=15, seed=0, dev=dev)
    acc_before = accuracy(model, test_x, test_y, dev)
    print(f'函数化前 · 测试准确率 = {acc_before * 100:.2f}%')

    # ③ 函数化：auto_alloc 自动选配置（绝不反涨 + 掉点<3pt 红线）
    W = model.fc1.weight.detach().cpu().numpy().astype(np.float64)   # (128, 784)
    guard = make_drop_guard(model, model.fc1, W.shape, test_x, test_y,
                            dev, acc_before, tol_pt=3.0)
    alloc = nrfunc.auto_alloc(W, signal='G', eval_fn=guard, seed=0)
    print(f"\n函数化（auto_alloc 自动选）：decision={alloc['decision']} "
          f"order={alloc['order']} K={alloc['K']} r={alloc['r']}，"
          f"内存省 {alloc['saving']*100:.1f}%"
          f"（通过候选 {alloc['n_candidates']}/{alloc['n_tried']}）")
    if alloc['decision'] == 'skip':
        print(f'  → 该层被判跳过：{alloc["skipped_reason"]}')
        return
    res = alloc['result']

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

    # ⑤ 部署使用（下游）：生成物 → 重建 → 塞回模型 → 推理测准确率
    recon = nrfunc.reconstruct(res)
    with torch.no_grad():
        model.fc1.weight.copy_(torch.as_tensor(recon, dtype=torch.float32))
    acc_after = accuracy(model, test_x, test_y, dev)
    print(f'\n部署后 · 测试准确率 = {acc_after * 100:.2f}% '
          f'（掉 {acc_before * 100 - acc_after * 100:.2f}pt）')

    # 价值证明 · 完整四维（体积/内存/CPU/GPU），运行时同机同批实测
    print('\n价值对比（fc1 层 128×784，函数化前 vs 后）：')
    show_value(W, res, test_x[:256])


if __name__ == '__main__':
    main()