"""【未验证 · 多尺度树档】图像分类 + B 多尺度树分区函数化。

一句话结论：
    分区除「按什么信号切」外，还有一维「多尺度树」——先粗分区成大区，再每个大区
    内细分小区，叶子区数可调（叶子越多、压缩越弱、保真越高）。本例在小 MLP
    （下采样 MNIST）上演示 regionify_hierarchical()：打印树结构、扫一遍「树深 →
    保真」，并用最细档还原测准确率。

对应分型表的多尺度树方向（方向 B，接口已实现，待真实大模型验证）。
"""

from _common import (accuracy, device, downsample_images, load_mnist,
                     show_artifacts, 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（下采样 MNIST 7×7）
    train_x, train_y, test_x, test_y = load_mnist()
    train_x = downsample_images(train_x, size=7)
    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}%')

    W = model.fc1.weight.detach().cpu().numpy().astype(np.float64)   # (16, 49)

    # ③ 函数化 · 树结构：K_top=4 大区 → 每大区 K_sub=2 小区（打印单元归属）
    hier = nrfunc.regionify_hierarchical(W, signal='G', K_top=4, K_sub=2,
                                         order=1, r=4, seed=0)
    print(f'\n多尺度树（K_top=4 × K_sub=2）：叶子区 K={hier["K"]}，'
          f'解释方差 = {hier["explained"]:.4f}')
    for k in range(hier['K']):
        units = (hier['assign'] == k).nonzero()[0]
        print(f'  叶子区 {k}：单元 {units.tolist()}')

    # 树深可调：扫一遍「粗 → 细」，看叶子数 vs 保真
    print('\n树深可调（叶子越多 → 压缩越弱 → 保真越高）：')
    for kt, ks in [(2, 2), (4, 2), (4, 4)]:
        r = nrfunc.regionify_hierarchical(W, signal='G', K_top=kt, K_sub=ks,
                                          order=1, r=4, seed=0)
        print(f'  K_top={kt} × K_sub={ks} → 叶子 K={r["K"]:2d}，解释方差 = {r["explained"]:.4f}')

    # ④ 生成物（库的活）：用最细档量化打包 + 分片存储
    fine = nrfunc.regionify_hierarchical(W, signal='G', K_top=4, K_sub=4,
                                         order=1, r=4, seed=0)
    print('\n生成物（最细档 K_top=4 × K_sub=4，量化打包 + 分片存储）：')
    show_artifacts(fine, bits=8, n_shards=4)

    # ⑤ 部署使用（下游）：最细档还原 → 塞回模型 → 推理测准确率
    recon = nrfunc.reconstruct(fine)
    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最细档（叶子 K={fine["K"]}）部署后 · 测试准确率 = '
          f'{acc_after * 100:.2f}%（掉 {acc_before * 100 - acc_after * 100:.2f}pt）')

    print('\n[诚实声明] 多尺度树「接口可跑 + 生成物可还原」，树的价值在「层次结构」'
          '（子函数可复用、可分片存储/并行 IO），在大 n 层上才明显；本例 n=16 仅演示机制。')


if __name__ == '__main__':
    main()