"""【未验证 · CNN 档】图像分类 + F 功能分区函数化。

一句话结论：
    在小 CNN（MNIST）上演示「F 功能分区」——不按权重几何距离，而按卷积核的
    激活响应模式聚类（对应 CNN 的对称/置换不变性）。接口可跑、生成物可还原，
    但**精度是否在真实 CNN 上保持，尚未验证**，属「待验证」而非「已验证」。

对应分型表第二行：CNN = 功能分区 F（接口已实现，待真实 CNN 验证）。
"""

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 torch.nn.functional as F

import nrfunc


class SmallCNN(nn.Module):
    def __init__(self, n1=8, n2=16, d_out=10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, n1, 3)          # 28→26
        self.conv2 = nn.Conv2d(n1, n2, 3)         # 13→11
        self.fc = nn.Linear(n2 * 5 * 5, d_out)    # maxpool 两次后 5×5

    def forward(self, x):                          # x: (B, 1, 28, 28)
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)                     # 26→13
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)                     # 11→5
        return self.fc(x.flatten(1))


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

    # 完整链路：训练 → 模型 → 函数化 → 生成物 → 部署使用
    #  ① 训练 + ② 模型（上游）：图像分类数据 + 训练小 CNN
    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('训练小 CNN（conv1: 8 核 / conv2: 16 核）：')
    model = train_model(SmallCNN(), 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}%')

    # ③ 函数化：对 conv1 权重 (8, 1, 3, 3) → (8, 9) 做 F 功能分区
    #    激活响应 = 各卷积核在同一批图像上的输出（展平空间×批为 T）
    W = model.conv1.weight.detach().cpu().numpy().astype(np.float64)   # (8,1,3,3)
    rows = W.reshape(W.shape[0], -1)                                    # (8, 9)
    with torch.no_grad():
        xb = torch.as_tensor(train_x[:256], dtype=torch.float32, device=dev)
        act = model.conv1(xb)                                          # (256,8,26,26)
    activations = act.permute(1, 0, 2, 3).reshape(W.shape[0], -1).cpu().numpy()  # (8, T)

    # 用 auto_alloc 自动选配置（绝不反涨 + 掉点<3pt 红线）
    guard = make_drop_guard(model, model.conv1, W.shape, test_x4, test_y,
                            dev, acc_before, tol_pt=3.0)
    alloc = nrfunc.auto_alloc(rows, signal='F', activations=activations,
                              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(f"\n函数化：signal=F（按激活响应聚类）/ K={res['K']} / order={alloc['order']} "
          f"/ r={alloc['r']}，解释方差 = {res['explained']:.4f}")

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

    # ⑤ 部署使用（下游）：生成物 → 重建 → 塞回模型 → 推理测准确率
    recon = nrfunc.reconstruct(res).reshape(W.shape)
    with torch.no_grad():
        model.conv1.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），conv1 展平 (8,9)
    x_batch = np.random.default_rng(0).standard_normal((256, 9)).astype(np.float32)
    print('\n价值对比（conv1 层 8×9，函数化前 vs 后）：')
    show_value(rows, res, x_batch)

    print('\n[诚实声明] 本例只证明 F 功能分区「接口可跑 + 生成物可还原」；'
          '是否在真实大 CNN 上保精度，尚未验证。')


if __name__ == '__main__':
    main()