# -*- coding: utf-8 -*-
"""跨规模缩放扫描（新版 nrfunc）：等宽三隐层 MLP 的层次越宽越省越稳趋势重跑。

对比 README「GPT 收益推断」节的旧数据（旧版 nrfunc 跑出）：
  h=64→2048：整体节省率 61.6%→89.3%（单调升）、叠加掉点 14.50pt→1.29pt（断崖降）。
本脚本用新版 nrfunc（谱探测 + K=1 全局低秩 + 3pt 红线）重跑 6 个宽度档，
看「省」与「掉点」两个维度在新版下的变化。

新版实测结果（固定 seed，见 reports/scale_scan.log）：
  h=64→2048：叠加掉点 4.16pt→-0.06pt（断崖降，h=128 起就压进 2pt）、
             整体节省率 71.7%→85.3%（整体上升、不再严格单调）。
"""
from _common import (accuracy, device, functionalize_layers, load_mnist,
                     train_model)  # 先导入：完成 src 路径引导

import torch
import torch.nn as nn


class EqualWidthMLP(nn.Module):
    """等宽三隐层：784 → h → h → h → 10（与 README 扫描表结构一致）。"""

    def __init__(self, h):
        super().__init__()
        self.fc1 = nn.Linear(784, h)
        self.fc2 = nn.Linear(h, h)
        self.fc3 = nn.Linear(h, h)
        self.fc4 = nn.Linear(h, 10)

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


# 6 个宽度档
WIDTHS = [64, 128, 256, 512, 1024, 2048]


def run_width(h, train_x, train_y, test_x, test_y, dev, seed=0):
    model = train_model(EqualWidthMLP(h), train_x, train_y, epochs=15, seed=seed, dev=dev)
    acc_before = accuracy(model, test_x, test_y, dev)
    layers = [('fc1', model.fc1), ('fc2', model.fc2), ('fc3', model.fc3)]
    out = functionalize_layers(model, layers, test_x, test_y, dev, acc_before,
                               tol_pt=3.0, seed=seed)
    n_params = sum(p.numel() for p in model.parameters())
    return {
        'h': h, 'n_params': n_params,
        'acc_before': acc_before, 'acc_after': out['acc_after'],
        'stacked_drop': out['stacked_drop'],
        'total_saving': out['total_saving'],
        'n_skip': out['n_skip'],
        'n_layers': out['n_layers'],
    }


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

    train_x, train_y, test_x, test_y = load_mnist()
    train_x, train_y = train_x[:40000], train_y[:40000]
    print('训练/评估 MNIST 等宽三隐层 MLP（784→h→h→h→10）...\n')

    header = f"{'隐层宽 h':<8}{'参数量':<10}{'n_skip':<8}{'叠加掉点':<10}{'整体节省率':<12}"
    print(header)
    print('-' * 48)
    results = []
    for h in WIDTHS:
        print(f'--- h={h} ---', flush=True)
        r = run_width(h, train_x, train_y, test_x, test_y, dev)
        results.append(r)
        skip = f'{r["n_skip"]}/{r["n_layers"]}'
        print(f'{h:<8}{r["n_params"]/1e4:>7.1f}万{skip:>8}{r["stacked_drop"]:>9.2f}pt'
              f'{r["total_saving"]*100:>10.1f}%', flush=True)
    print('-' * 48)

    print('\n=== 汇总表（新版 nrfunc，3pt 红线）===')
    print('{:<8}{:<14}{:<12}{:<12}'.format('隐层宽h', '参数量', '叠加掉点', '整体节省率'))
    for r in results:
        print(f'{r["h"]:<8}{r["n_params"]/1e4:>9.1f}万{r["stacked_drop"]:>11.2f}pt'
              f'{r["total_saving"]*100:>11.1f}%')


if __name__ == '__main__':
    main()