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

一句话结论：
    在下采样 MNIST（7×7）上训练的小 MLP，对其首层做 G 几何分区函数化（1 阶
    低秩），还原后测试准确率**几乎零损失**——证明「小 MLP + G 几何」这一档的
    函数化真实可用（本项目 P0 已实证「决策误差 <5%」）。

诚实说明（很重要）：
    - 本例验证的是「函数化不损精度」，不是「压缩」。小网络输出单元 n=16 太小，
      函数化本身压缩不了多少（甚至 order-1 低秩项会反超原始参数量）。
    - 0 阶质心虽能 n/K 压缩，但在「特征多样的首层」会灾难性掉点（本例实测
      88%→27%）——因为质心是「训练后」才叠加的误差，没被训练吸收。这正是
      后训练压缩（本库）与训练期压缩（QAT/蒸馏）的本质区别。
    - 真正的大压缩要到 n 大的层（见 mnist_mlp_classification.py / CNN / Transformer）。

对应分型表第一行：小 MLP = 几何分区 G（本项目已实证「不损精度」）。
"""

from _common import (accuracy, device, downsample_images, 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 SmallMLP(nn.Module):
    def __init__(self, d_in=49, h=16, d_out=10):
        super().__init__()
        self.fc1 = nn.Linear(d_in, h)
        self.fc2 = nn.Linear(h, d_out)

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


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

    # 完整链路：训练 → 模型 → 函数化 → 生成物 → 部署使用
    #  ① 训练 + ② 模型（上游，不是库的活）：准备图像分类数据 + 训练一个成熟小 MLP
    train_x, train_y, test_x, test_y = load_mnist()
    train_x = downsample_images(train_x, size=7)   # 49 维
    test_x = downsample_images(test_x, size=7)
    train_x, train_y = train_x[:30000], train_y[:30000]   # 取子集，示例求快

    print('训练小 MLP（49→16→10）：')
    model = train_model(SmallMLP(), 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)   # (16, 49)
    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 层 16×49，函数化前 vs 后）：')
    show_value(W, res, test_x[:256])
    print('注：此层 n=16 太小，函数化本身无参数压缩，节省主要来自 32→8bit 量化；'
          '大压缩见 MNIST MLP / CNN / Transformer 档。')


if __name__ == '__main__':
    main()