"""auto_alloc 自适应函数化分配器 —— 5 例验证脚本。

对照《自研AI_P3算法原型计划.md》第十一章 11.2 实测表（固定种子全量重跑后的诚实值），
验证第十二章 12.5 的 4 条验收。

运行（需先 `pip install -e .`）：
    python examples/auto_alloc_demo.py
"""
import numpy as np

from nrfunc import auto_alloc


def _clustered(n, D, K_true, seed, noise=0.3):
    rng = np.random.default_rng(seed)
    centroids = rng.normal(0.0, 1.0, size=(K_true, D))
    assign = rng.integers(0, K_true, size=n)
    return centroids[assign] + rng.normal(0.0, noise, size=(n, D))


def _build(c):
    rows = _clustered(c['n'], c['D'], K_true=4, seed=7)
    kw = dict(rows=rows, signal=c['signal'])
    if c['signal'] == 'F':
        kw['activations'] = _clustered(c['n'], 64, K_true=4, seed=77)
    elif c['signal'] == 'S':
        kw['groups'] = np.repeat(np.arange(4), c['n'] // 4)
    return kw


def _pct(saving):
    return f"省 {saving*100:.1f}%" if saving > 0 else "跳过（保持原样）"


def _plan_s(o):
    p = o['plan']
    return '跳过' if p is None else f"order{p[0]} K={p[1]} r={p[2]}"


def _ev_s(o):
    return '—' if o['explained'] is None else f"{o['explained']:.3f}"


CASES = [
    dict(name='大 CNN features.17', n=128, D=1152, signal='F', old='多 25.4%'),
    dict(name='小 MLP fc1',         n=16,  D=49,   signal='G', old='多 158.4%'),
    dict(name='小 CNN conv1',       n=8,   D=9,    signal='F', old='多 72.9%'),
    dict(name='MNIST MLP fc1',      n=128, D=784,  signal='G', old='省 37.0%'),
    dict(name='玩具 ViT wq',        n=64,  D=64,   signal='S', old='省 62.4%'),
]


def main():
    print('=' * 80)
    print('auto_alloc 5 例验证（内存·函数部署口径，f32 常驻）')
    print('=' * 80)
    print(f"{'层':<20}{'旧值(一刀切)':>13}{'新值(auto_alloc)':>19}{'决策':>22}{'保真EV':>9}")
    print('-' * 80)

    results = {}
    for c in CASES:
        o = auto_alloc(**_build(c))
        results[c['name']] = o
        print(f"{c['name']:<20}{c['old']:>13}{_pct(o['saving']):>19}{_plan_s(o):>22}{_ev_s(o):>9}")

    print('-' * 80)
    print('验收核对：')
    print(f"  [{'PASS' if results['小 MLP fc1']['saving'] >= 0 else 'FAIL'}] 验收1 小 MLP 不再 -158% 反涨")
    print(f"  [{'PASS' if results['大 CNN features.17']['saving'] >= 0 else 'FAIL'}] 验收2 大 CNN 不再 -25% 反涨")
    print(f"  [{'PASS' if all(results[k]['saving'] > 0 for k in ('MNIST MLP fc1', '玩具 ViT wq')) else 'FAIL'}] 验收3 省例保持净省")
    print(f"  [{'PASS' if all(o['saving'] >= 0 for o in results.values()) else 'FAIL'}] 验收4 整网绝不反涨（内存省最多）")
    print()

    # 掉点代理演示
    print('-' * 80)
    print('掉点代理演示：MNIST MLP（n=128）默认省/保真加权 vs 失真度硬约束（distortion=0.2）')
    print('-' * 80)
    mnist = CASES[3]
    for tag, o in [('默认(省/保真加权)', auto_alloc(**_build(mnist))),
                   ('失真度≤0.2', auto_alloc(**_build(mnist), distortion=0.2))]:
        print(f"  {tag:<16} -> {_pct(o['saving']):<14} {_plan_s(o):<20} EV={_ev_s(o)}")
    print()
    print('说明：')
    print('  · 默认走「省/保真加权」几何平均 √(saving×explained)，两维都在 0~1 天然可比，')
    print('    无人工权重要调，自动在省内存与保精度间取平衡（不再有 fidelity_margin 魔法数）。')
    print('  · distortion：失真度 = 允许的 EV 掉点上限（explained ≥ 1-distortion），')
    print('    是「真实掉点」的语义化参数（EV 代理版）。')
    print('  · eval_fn：真实「掉点<阈值」的最可靠实现（库不内置数据/前向，交给使用方）。')
    print('  · 真实掉点沿用第十一章固定种子复现值：MNIST MLP 2.43pt / 大 CNN 0.37pt /')
    print('    ViT 0.71pt，均 < 3pt 红线。auto_alloc 只改配置选择、不改训练，掉点可沿用。')
    print('=' * 80)


if __name__ == '__main__':
    main()