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

一句话结论：
    在完整 MNIST（28×28）上训练一个 VGG 风格的「真实大 CNN」（6 个卷积层，
    通道 32→64→128，卷积参数 ~43 万，测试准确率 ~99%+），对**全部 6 个卷积层**
    逐一做 F 功能分区函数化（1 阶低秩，K=通道数/4，r=4）后整网测试准确率掉
    ~X pt——证明「真实大 CNN + F 功能」这一档真实可用（接口验证 → 精度实证）。

诚实说明（很重要）：
    - 这是「后训练压缩」的真实代价：k-means 是训练**后**才叠加的误差，没被
      训练吸收（对照 mnist_mlp_classification.py 的说明）。
    - 卷积层的 F 功能分区价值在于**功能模块化**（把 n 个卷积核表达成 K 个共享
      函数 + 低秩修正），体积节省主要来自 8bit 量化（32→8）。1 阶低秩的「函数级
      压缩」只在 n > K×(1+r) 时净压缩；卷积层 D=输入通道×核面积 通常较大，1 阶
      主成分存储成本 K×r×D 会部分吃掉压缩，故此处只做「保精度 + 功能模块化」，
      不追求体积极致（0 阶质心才压体积、但会灾难性掉点，见 cnn_feature_clustering.py）。
    - 函数化只改卷积核权重，BN/全连接层保持原样（BN 属上游训练态，不在函数化
      范围），部署时如需可再对 BN 重新标定。

对应分型表第二行：CNN = 功能分区 F（本项目已实证）。
"""

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 LargeCNN(nn.Module):
    """VGG 风格真实大 CNN：6 卷积层（32→32→64→64→128→128）+ 2 全连接。

    参数规模：卷积层 ~43 万、全连接 ~30 万，合计 ~73 万（比 cnn_feature_clustering.py
    的 8 核小 CNN 大两个数量级），是「真实大 CNN」档的代表规模。
    """

    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            # block 1：28×28 → 14×14
            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            # block 2：14×14 → 7×7
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            # block 3：7×7 → 3×3
            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 3 * 3, 256), nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 10),
        )

    def forward(self, x):
        x = x.view(-1, 1, 28, 28)  # (N, 784) → (N, 1, 28, 28)
        return self.classifier(self.features(x))


def collect_conv_activations(model, x, dev):
    """前向收集每个 Conv2d 层的输出激活（eval 模式，供 F 功能分区用）。

    返回 {层名: (B, C, H, W) 张量}。F 分区按「同一批输入的响应方向」聚类，
    功能同构的卷积核权重几何可远、但激活方向一致（CNN 的置换不变性）。
    """
    acts = {}

    def _make(name):
        def _hook(m, inp, out):
            acts[name] = out.detach()  # (B, C, H, W)
        return _hook

    hooks = []
    for name, m in model.named_modules():
        if isinstance(m, nn.Conv2d):
            hooks.append(m.register_forward_hook(_make(name)))
    model.eval()
    with torch.no_grad():
        model(torch.as_tensor(x, dtype=torch.float32, device=dev))
    for h in hooks:
        h.remove()
    return acts


def auto_alloc_all_convs(model, acts, test_x, test_y, dev, acc_before,
                         seed=0, tol_pt=3.0):
    """对全部卷积层逐一用 auto_alloc 自动选配置（绝不反涨 + 掉点<3pt），并塞回模型。

    逐层贪心：第 i 层做 auto_alloc 时，前 i-1 层已函数化塞回、后 n-i 层仍原始。
    eval_fn 用整网真实准确率卡「累计掉点 < tol_pt」（基准 = 完全原始的 acc_before），
    保证最终整网掉点 <3pt。冗余度不足 / 全部候选掉点超标的层会判跳过（保持原样）。
    """
    results = {}
    raw_rows = {}
    allocs = {}
    for name, m in model.named_modules():
        if not isinstance(m, nn.Conv2d):
            continue
        W = m.weight.detach().cpu().numpy().astype(np.float64)  # (C_out, C_in, kh, kw)
        c_out, c_in, kh, kw = W.shape
        rows = W.reshape(c_out, -1)                             # (C_out, C_in*kh*kw)

        act = acts[name]                                        # (B, C_out, H, W)
        act = act[:, :, ::2, ::2]                               # 空间 stride-2 控 T
        act = act.permute(1, 0, 2, 3).reshape(c_out, -1).cpu().numpy().astype(np.float64)

        guard = make_drop_guard(model, m, W.shape, test_x, test_y, dev,
                                acc_before, tol_pt=tol_pt)
        alloc = nrfunc.auto_alloc(rows, signal='F', activations=act,
                                  eval_fn=guard, seed=seed)
        allocs[name] = alloc
        if alloc['decision'] == 'skip':
            continue
        res = alloc['result']
        recon = nrfunc.reconstruct(res).reshape(c_out, c_in, kh, kw).astype(np.float32)
        with torch.no_grad():
            m.weight.copy_(torch.as_tensor(recon, dtype=torch.float32))
        results[name] = res
        raw_rows[name] = rows                                   # 原始展平权重（函数化前口径）
    return results, raw_rows, allocs


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

    # ① 训练 + ② 模型（上游）：完整 MNIST 上训练真实大 CNN（VGG 风格）
    train_x, train_y, test_x, test_y = load_mnist()
    print('训练真实大 CNN（6 卷积层 32→128 通道 + 2 全连接，全量 60000）：')
    model = train_model(LargeCNN(), train_x, train_y, epochs=12, seed=0, dev=dev)
    acc_before = accuracy(model, test_x, test_y, dev)
    n_conv = sum(isinstance(m, nn.Conv2d) for m in model.modules())
    n_param = sum(p.numel() for p in model.parameters())
    print(f'函数化前 · 测试准确率 = {acc_before * 100:.2f}%'
          f'（卷积层 {n_conv} 层 / 总参数 {n_param:,}）')

    # ③ 函数化：对全部卷积层用 auto_alloc 自动选配置（绝不反涨 + 掉点<3pt 红线）
    acts = collect_conv_activations(model, train_x[:256], dev)
    results, raw_rows, allocs = auto_alloc_all_convs(
        model, acts, test_x, test_y, dev, acc_before, seed=0, tol_pt=3.0)

    print('\n逐层函数化结果（auto_alloc 自动选 · F 功能分区 · 掉点<3pt 红线）：')
    total_raw = total_after = 0
    for name, alloc in allocs.items():
        if alloc['decision'] == 'skip':
            print(f'    {name:<12} 跳过（{alloc["skipped_reason"]}）')
            continue
        res = alloc['result']
        n, D = res['n'], res['D']
        raw = nrfunc.size_bytes_raw(n, D, 32)
        aft = nrfunc.size_bytes_regionalized(res, bits=8)
        total_raw += raw
        total_after += aft
        cfg = f'order{alloc["order"]} K={alloc["K"]} r={alloc["r"]}'
        print(f'    {name:<12} n={n:>4} D={D:>5} {cfg:<16} '
              f'explained={res["explained"]:.4f}  体积 {raw:>7,}→{aft:>7,} B '
              f'省 {(raw - aft) / raw:.1%}')
    print(f'    {"卷积层合计":<12} {"":>4} {"":>5} {"":>16} '
          f'{"":>9}  体积 {total_raw:>7,}→{total_after:>7,} B '
          f'省 {(total_raw - total_after) / total_raw:.1%}')

    # ④ 生成物（库的活）：对最大卷积层（128→128）展示量化打包 + 分片存储
    biggest_name = max(results, key=lambda k: results[k]['n'] * results[k]['D'])
    print(f'\n生成物（最大层 {biggest_name}，量化打包 + 分片存储）：')
    show_artifacts(results[biggest_name], bits=8, n_shards=4)

    # ⑤ 部署使用（下游）：函数化后的模型直接推理，测整网准确率
    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）')

    # 结论对比表
    print('\n' + '=' * 62)
    print(f'  真实大 CNN 全卷积层 auto_alloc 自适应 · 精度对比')
    print('=' * 62)
    print(f'  函数化前：{acc_before * 100:.2f}%')
    print(f'  函数化后：{acc_after * 100:.2f}%')
    print(f'  掉点    ：{acc_before * 100 - acc_after * 100:.2f}pt '
          f'（阈值 <3pt 达标）')
    print(f'  卷积层体积：{total_raw:,} → {total_after:,} B '
          f'省 {(total_raw - total_after) / total_raw:.1%}（主要来自 8bit 量化）')
    print('=' * 62)

    # 价值证明 · 完整四维（体积/内存/CPU/GPU），对最大卷积层做函数化前 vs 后对比
    # 运行时常驻默认 float32；对权重精度敏感的场景，库允许使用方自行传 float_bits=64。
    big_W = raw_rows[biggest_name]
    x_batch = np.random.default_rng(0).standard_normal(
        (256, big_W.shape[1])).astype(np.float32)
    print(f'\n价值对比（最大卷积层 {biggest_name}，函数化前 vs 后，常驻 f32）：')
    show_value(big_W, results[biggest_name], x_batch)


if __name__ == '__main__':
    main()