# -*- coding: utf-8 -*-
"""project 训练侧：用真实 MNIST 训练「单图多对象检测」大模型 → 函数化 → 导出产物。

完整链路（本项目要做的实质对比）：
  路线 A（传统）：训练 → 成熟模型 → 部署（Rust 读 model.json 稠密前向）
  路线 B（函数化）：训练 → 成熟模型 → 函数化 → 生成物 → 部署（Rust 读
        functional.json 函数前向）

任务：真实 MNIST「单图多对象检测」（YOLO 式固定锚框网格）——
  - 每张 28×28 画布随机放 1~2 个数字（各缩放到 ~8px），各落到 2×2 网格的不同格
  - 输出：S×S 网格，每格预测 [objectness(1) + 类别 one-hot(10) + bbox 偏移(4)] = 15 维
  - 大模型骨干 MLP：784 → 4096 → 2048 → 1024，网格头 → 4×15=60
  - 函数化对象：fc2（4096→2048）宽层，冗余足、函数化收益大

产出（artifacts/）：
  - model.json       ：成熟模型完整权重（路线 A 稠密部署用）
  - functional.json  ：函数化生成物（fc2 层函数化，路线 B 用）
  - stats.json       ：Python 侧四维 + 精度实测
"""
import json
import os
import struct
import sys

import numpy as np
import torch
import torch.nn as nn

HERE = os.path.dirname(os.path.abspath(__file__))
ART = os.path.join(HERE, 'artifacts')
_PY_EXAMPLES = os.path.abspath(os.path.join(HERE, '..', 'python'))
_SRC = os.path.abspath(os.path.join(HERE, '..', '..', 'src'))
for _p in (_PY_EXAMPLES, _SRC):
    if _p not in sys.path:
        sys.path.insert(0, _p)

import nrfunc  # noqa: E402
from _common import load_mnist  # noqa: E402


IMG = 28
DIM = IMG * IMG
GRID = 2                # 2×2 = 4 格（数字 8px，每格 14×14，学得动）
CELL = IMG // GRID      # 每格 14×14 像素
N_CLASS = 10
HEAD_PER_CELL = 1 + N_CLASS + 4   # objectness + 10类 + bbox4 = 15
HEAD_DIM = GRID * GRID * HEAD_PER_CELL  # 4 * 15 = 60


# ---------- 数据：多对象检测（1~4 个数字，YOLO 式网格标签） ----------
def make_detection(x_mnist, y_mnist, seed=0, canvas=IMG, grid=GRID):
    """把 MNIST 造「单图多对象检测」：每张画布放 1~4 个数字，各落不同格。

    返回 (x, grid_label)：
      x：(n, 784) 28×28 画布，含 1~4 个缩小的数字
      grid_label：(n, GRID*GRID*15) 每格 [objectness, 10类one-hot, bbox偏移4]
    """
    rng = np.random.default_rng(seed)
    n = x_mnist.shape[0]
    out = np.zeros((n, DIM), dtype=np.float32)
    # grid_label：每格 15 维 [obj, cls(10), bx(4)]
    label = np.zeros((n, GRID * GRID * HEAD_PER_CELL), dtype=np.float32)
    scale = 8                        # 数字缩放到 8×8（放进 14×14 格内）
    cell = canvas // grid            # 14
    for i in range(n):
        img = x_mnist[i].reshape(IMG, IMG)
        ys, xs = np.where(img > 0.1)
        if len(xs) == 0:
            continue
        # 随机 1~2 个数字（从数据集里额外取，制造"多对象"）
        n_obj = int(rng.integers(1, 3))
        # 随机选 n_obj 个不重复的格
        cells = rng.choice(grid * grid, size=n_obj, replace=False)
        for ci, cell_id in enumerate(cells):
            # 取一个数字（主样本 + 复用随机样本，保证多样性）；记录真实类别
            if ci == 0:
                digit_img = img
                true_cls = int(y_mnist[i])
            else:
                j = int(rng.integers(0, len(x_mnist)))
                ys_j, xs_j = np.where(x_mnist[j].reshape(IMG, IMG) > 0.1)
                if len(xs_j) == 0:
                    continue
                dj = x_mnist[j].reshape(IMG, IMG)
                dj = dj[ys_j.min():ys_j.max() + 1, xs_j.min():xs_j.max() + 1]
                digit_img = dj
                true_cls = int(y_mnist[j])
            # 裁剪数字主体
            d_ys, d_xs = np.where(digit_img > 0.1)
            if len(d_xs) == 0:
                continue
            digit = digit_img[d_ys.min():d_ys.max() + 1, d_xs.min():d_xs.max() + 1]
            dh, dw = digit.shape
            yy = np.linspace(0, dh - 1, scale).round().astype(int)
            xx = np.linspace(0, dw - 1, scale).round().astype(int)
            digit_rs = digit[yy][:, xx]
            # 落到格子 cell_id，格内微偏移（数字 5px 放进 7px 格）
            gy, gx = cell_id // grid, cell_id % grid
            off = cell - scale  # 2px 余量
            oy = int(rng.integers(0, off + 1))
            ox = int(rng.integers(0, off + 1))
            py = gy * cell + oy
            px = gx * cell + ox
            # 直接写到 out[i]
            img_buf = out[i].reshape(canvas, canvas)
            img_buf[py:py + scale, px:px + scale] = np.maximum(
                img_buf[py:py + scale, px:px + scale], digit_rs)
            # 标签：objectness=1，类别 one-hot（真实类别），bbox 偏移（格内归一化）
            base = cell_id * HEAD_PER_CELL
            label[i, base] = 1.0  # objectness
            label[i, base + 1 + true_cls] = 1.0
            # bbox 偏移（格内归一化到 [0,1]）
            label[i, base + 1 + N_CLASS + 0] = (px - gx * cell) / cell
            label[i, base + 1 + N_CLASS + 1] = (py - gy * cell) / cell
            label[i, base + 1 + N_CLASS + 2] = scale / cell
            label[i, base + 1 + N_CLASS + 3] = scale / cell
    return out, label


# ---------- 大模型 ----------
class DetectNet(nn.Module):
    """单图多对象检测大模型：骨干 784→4096→2048→1024 + 网格头 → 60。"""

    def __init__(self, n_in=DIM, h1=4096, h2=2048, h3=1024, grid=GRID):
        super().__init__()
        self.fc1 = nn.Linear(n_in, h1)          # 骨干层1（函数化敏感层）
        self.fc2 = nn.Linear(h1, h2)            # 骨干层2（函数化对象：宽层冗余足）
        self.fc3 = nn.Linear(h2, h3)            # 骨干层3
        self.head = nn.Linear(h3, grid * grid * HEAD_PER_CELL)

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


def train_model(model, x, label, epochs=20, batch=256, lr=1e-3, seed=0, dev='cpu'):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    model.apply(lambda m: m.reset_parameters()
                if hasattr(m, 'reset_parameters') else None)
    model = model.to(dev)
    xt = torch.as_tensor(x, device=dev)
    lt = torch.as_tensor(label, device=dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    bce = nn.BCEWithLogitsLoss()
    n = len(xt)
    GRID2 = GRID * GRID
    for e in range(1, epochs + 1):
        perm = torch.randperm(n, device=dev)
        model.train()
        tl = 0.0
        for b in range(0, n, batch):
            ids = perm[b:b + batch]
            out = model(xt[ids])  # (B, 240)
            B = out.shape[0]
            out = out.view(B, GRID2, HEAD_PER_CELL)
            # objectness：sig 输出 vs 标签
            obj_logit = out[:, :, 0]
            obj_tgt = lt[ids].view(B, GRID2, HEAD_PER_CELL)[:, :, 0]
            loss_obj = bce(obj_logit, obj_tgt)
            # 分类（只对有对象的格算）
            cls_logit = out[:, :, 1:1 + N_CLASS]   # (B, GRID2, 10)
            cls_tgt = lt[ids].view(B, GRID2, HEAD_PER_CELL)[:, :, 1:1 + N_CLASS]
            mask = obj_tgt > 0.5
            loss_cls = 0.0
            if mask.any():
                loss_cls = nn.functional.cross_entropy(
                    cls_logit[mask], cls_tgt[mask].argmax(1))
            # bbox（只对有对象的格算）
            box_pred = out[:, :, 1 + N_CLASS:]     # (B, GRID2, 4)
            box_tgt = lt[ids].view(B, GRID2, HEAD_PER_CELL)[:, :, 1 + N_CLASS:]
            loss_box = 0.0
            if mask.any():
                loss_box = nn.functional.mse_loss(box_pred[mask], box_tgt[mask])
            loss = loss_obj + loss_cls + 5.0 * loss_box
            opt.zero_grad()
            loss.backward()
            opt.step()
            tl += loss.item() * B
        if e % 5 == 0 or e == epochs:
            print(f'  epoch {e:2d}/{epochs}  loss={tl / n:.4f}')
    model.eval()
    return model


@torch.no_grad()
def evaluate(model, x, label, dev='cpu'):
    """评估：分类准确率（有对象格）+ 对象检出 recall（objectness 是否命中）。"""
    xt = torch.as_tensor(x, device=dev)
    lt = torch.as_tensor(label, device=dev)
    out = model(xt).view(len(xt), GRID * GRID, HEAD_PER_CELL)
    obj_pred = (torch.sigmoid(out[:, :, 0]) > 0.5).float()
    obj_tgt = lt.view(len(xt), GRID * GRID, HEAD_PER_CELL)[:, :, 0]
    # 检出 recall：真值有对象的格，预测是否为对象
    has_obj = obj_tgt > 0.5
    recall = (obj_pred[has_obj] == 1).float().mean().item() if has_obj.any() else 1.0
    # 分类准确率（对有对象且预测对的格）
    cls_pred = out[:, :, 1:1 + N_CLASS].argmax(-1)
    cls_tgt = lt.view(len(xt), GRID * GRID, HEAD_PER_CELL)[:, :, 1:1 + N_CLASS].argmax(-1)
    hit = has_obj & (obj_pred == 1)
    cls_acc = (cls_pred[hit] == cls_tgt[hit]).float().mean().item() if hit.any() else 0.0
    return float(cls_acc), float(recall)


# ---------- 导出 ----------
def _m(layer):
    """Linear 层权重导出为 numpy（含 bias 尾维），返回 (n_out, n_in+1)。"""
    w = layer.weight.detach().cpu().numpy()
    b = layer.bias.detach().cpu().numpy().reshape(-1, 1)
    return np.hstack([w, b])


def _pack_dense_layers(layers, bits=32):
    """把多个稠密层权重（[(name, (n_out,n_in) ndarray), ...]）打成二进制块。

    格式（小端）：
      n_layers(4B) + 每层 [n_out(4B) + n_in(4B) + bits(1B) + 权重]
        bits=32：直存 f32[n_out*n_in]
        bits=8/4：per-array 对称量化（与 nrfunc.io.to_bytes 同口径），
                  存 scale(1 × f32) + 打包 int8/int4 字节
    """
    out = bytearray()
    out += struct.pack('<i', len(layers))
    for _, w in layers:
        w = np.asarray(w, dtype=np.float32).astype('<f4')
        n_out, n_in = w.shape
        out += struct.pack('<ii', n_out, n_in)
        if bits == 32:
            out += struct.pack('<B', 32)
            out += w.tobytes()
        elif bits in (8, 4):
            out += struct.pack('<B', bits)
            packed, scale = _quantize(w, bits)
            out += packed
            out += np.array([scale], dtype='<f4').tobytes()  # scale 在段尾（与 to_bytes 同口径）
        else:
            raise ValueError(f'bits 仅支持 32/8/4，收到 {bits}')
    return bytes(out)


def _quantize(w, bits):
    """per-array 对称量化（与 nrfunc.io.to_bytes 的 _q 同口径：单个全局 scale）。

    返回 (packed_bytes, scale)：packed 为 int8/int4 打包字节，scale 为标量。
    """
    w = w.astype(np.float64)
    qmax = 2 ** (bits - 1) - 1
    amax = float(np.abs(w).max()) if w.size else 0.0
    scale = max(amax, 1e-12) / qmax
    q = np.round(w / scale).clip(-qmax, qmax).astype(np.int16)
    off = (q + qmax).astype(np.uint16)  # 无符号偏移 [0, 2^bits-1]
    flat = off.flatten()
    if bits == 8:
        packed = flat.astype(np.uint8).tobytes()
    elif bits == 4:
        if flat.size % 2:
            flat = np.concatenate([flat, np.zeros(1, dtype=np.uint16)])
        lo = flat[0::2].astype(np.uint8)
        hi = flat[1::2].astype(np.uint8)
        packed = (lo | (hi << 4)).tobytes()
    else:
        raise ValueError(f'量化仅支持 8/4 bit，收到 {bits}')
    return packed, float(scale)


# 极限压制位宽：原型（稠密）与优化物（函数化）都压到位宽极限对比。
DENSE_BITS = 8   # 原型极限：稠密权重 int8 量化
FUNC_BITS = 8    # 优化物极限：函数化参数 int8 量化


def export_dense(model, path, bits=DENSE_BITS):
    """路线 A：成熟模型全距密权重 → 二进制（极限：int8 量化）。层序固定 fc1/fc2/fc3/head。"""
    layers = [
        ('fc1', _m(model.fc1)), ('fc2', _m(model.fc2)),
        ('fc3', _m(model.fc3)), ('head', _m(model.head)),
    ]
    with open(path, 'wb') as f:
        f.write(_pack_dense_layers(layers, bits=bits))


def export_functional(model, res, path, bits=FUNC_BITS):
    """路线 B：函数化生成物 → 二进制（极限：函数参数 int8 量化）。

    格式：fc1/fc3/head 三层距密权重块（int8）+ fc2 区域函数块（NRFN 二进制，to_bytes bits）。
      前部 = _pack_dense_layers([fc1, fc3, head], bits)
      后部 = func_block_len(8B) + to_bytes(res, bits)
    """
    from nrfunc.io import to_bytes
    dense = _pack_dense_layers([
        ('fc1', _m(model.fc1)), ('fc3', _m(model.fc3)), ('head', _m(model.head)),
    ], bits=bits)
    func_blob = to_bytes(res, bits=bits)  # 与原型同位宽极限压制
    with open(path, 'wb') as f:
        f.write(dense)
        f.write(struct.pack('<Q', len(func_blob)))
        f.write(func_blob)


def main():
    os.makedirs(ART, exist_ok=True)
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'设备：{dev}')

    # 1) 真实 MNIST → 多对象检测
    tr_x, tr_y, te_x, te_y = load_mnist()
    x_tr, l_tr = make_detection(tr_x, tr_y, seed=0)
    x_te, l_te = make_detection(te_x, te_y, seed=1)
    print(f'数据：train={len(x_tr)} test={len(x_te)}，输入 28×28={DIM}，'
          f'任务=多对象检测（{GRID}×{GRID}网格，每格1对象，最多2对象）')

    # 2) 训练大模型
    model = DetectNet()
    n_param = sum(p.numel() for p in model.parameters())
    print(f'模型参数：{n_param / 1e6:.2f}M')
    train_model(model, x_tr, l_tr, dev=dev)
    cls_before, recall_before = evaluate(model, x_te, l_te, dev=dev)
    print(f'成熟模型：cls acc={cls_before * 100:.2f}%  '
          f'检测 recall={recall_before * 100:.2f}%')

    # 3) 量化方案评估：per-array（单 scale） vs per-row（每行一个 scale），测真实掉点
    #    ——「评估四维收益哪个好用哪个」，用掉点决定 int8 kernel 的量化粒度。
    print('\n--- 量化方案掉点评估（int8）---')
    layers_map = [
        ('fc1', model.fc1), ('fc2', model.fc2), ('fc3', model.fc3), ('head', model.head),
    ]
    def _quant_replace(mode):
        # 把全部层的权重按 mode 量化后复制回模型，返回原始权重备份
        backup = {}
        for name, layer in layers_map:
            w = layer.weight.data  # (out, in)
            backup[name] = w.clone()
            wn = w.cpu().numpy()
            if mode == 'array':
                amax = np.abs(wn).max()
                scale = amax / 127.0
                q = np.round(wn / scale).clip(-127, 127) * scale
            else:  # per-row
                amax = np.abs(wn).max(axis=1, keepdims=True)
                scale = np.maximum(amax, 1e-12) / 127.0
                q = np.round(wn / scale).clip(-127, 127) * scale
            layer.weight.data.copy_(torch.as_tensor(q, dtype=w.dtype, device=w.device))
        return backup
    def _restore(backup):
        for name in backup:
            layer = dict((n, l) for n, l in layers_map)[name]
            layer.weight.data.copy_(backup[name])
    for mode in ('array', 'row'):
        bk = _quant_replace(mode)
        acc, _ = evaluate(model, x_te, l_te, dev=dev)
        _restore(bk)
        drop = (cls_before - acc) * 100
        print(f'  per-{mode} 量化：cls acc={acc*100:.2f}%  掉点={drop:+.2f}pt')

    # 4) 函数化 fc2（4096→2048 宽层，冗余足）
    fc2_w = _m(model.fc2)  # (2048, 4097)
    print(f'函数化对象：fc2 {fc2_w.shape}（{fc2_w.size} 参数）')

    def eval_fn(result):
        if result is None:
            return False
        recon = nrfunc.reconstruct(result)
        with torch.no_grad():
            bw = model.fc2.weight.data.clone()
            bb = model.fc2.bias.data.clone()
            model.fc2.weight.copy_(torch.as_tensor(
                recon[:, :-1], dtype=torch.float32, device=dev))
            model.fc2.bias.copy_(torch.as_tensor(
                recon[:, -1], dtype=torch.float32, device=dev))
            acc, _ = evaluate(model, x_te, l_te, dev=dev)
            model.fc2.weight.copy_(bw)
            model.fc2.bias.copy_(bb)
        return (cls_before - acc) * 100 < 3.0

    alloc = nrfunc.auto_alloc(fc2_w, signal='auto', eval_fn=eval_fn)
    print(f'auto_alloc 决策：{alloc["decision"]} order={alloc["order"]} '
          f'K={alloc["K"]} r={alloc["r"]}，内存省 {alloc["saving"] * 100:.1f}%')

    if alloc['result'] is not None:
        res = alloc['result']
        recon = nrfunc.reconstruct(res)
        with torch.no_grad():
            model.fc2.weight.copy_(torch.as_tensor(
                recon[:, :-1], dtype=torch.float32, device=dev))
            model.fc2.bias.copy_(torch.as_tensor(
                recon[:, -1], dtype=torch.float32, device=dev))
        cls_after, recall_after = evaluate(model, x_te, l_te, dev=dev)
        print(f'函数化后：cls acc={cls_after * 100:.2f}%  '
              f'recall={recall_after * 100:.2f}%  掉点={(cls_before - cls_after) * 100:.2f}pt')
    else:
        res = None
        cls_after, recall_after = cls_before, recall_before
        print('函数化被跳过，精度无变化')

    # 4) 导出产物（二进制，禁用 JSON 存生成物）
    export_dense(model, os.path.join(ART, 'model.bin'))
    if res is not None:
        export_functional(model, res, os.path.join(ART, 'functional.bin'))

    # 5) Python 侧统计（stats 是配置/统计，非生成物，保留 JSON）
    vol_dense = os.path.getsize(os.path.join(ART, 'model.bin'))
    vol_func = (os.path.getsize(os.path.join(ART, 'functional.bin'))
                if res is not None else vol_dense)

    stats = {
        'task': 'MNIST 多对象检测（YOLO式 2×2 网格，每格1对象，最多2对象）',
        'n_input': DIM,
        'n_fc1': model.fc1.out_features,
        'n_fc2': model.fc2.out_features,
        'n_fc3': model.fc3.out_features,
        'n_head': model.head.out_features,
        'n_param': n_param,
        'grid': GRID, 'n_class': N_CLASS, 'n_per_cell': HEAD_PER_CELL,
        'func_layer': 'fc2',
        'cls_acc_before': cls_before, 'cls_acc_after': cls_after,
        'recall_before': recall_before, 'recall_after': recall_after,
        'drop_pt': round((cls_before - cls_after) * 100, 4),
        'auto_alloc_decision': alloc['decision'],
        'auto_alloc_order': alloc['order'],
        'auto_alloc_K': alloc['K'],
        'auto_alloc_r': alloc['r'],
        'auto_alloc_signal': alloc['signal'],
        'auto_alloc_probed_signal': alloc['probed_signal'],
        'vol_dense_bytes': vol_dense,
        'vol_functional_bytes': vol_func,
    }
    with open(os.path.join(ART, 'stats.json'), 'w', encoding='utf-8') as f:
        json.dump(stats, f, indent=2, ensure_ascii=False)

    print('\n==== 训练侧完成 ====')
    print(f'体积：dense {vol_dense}B vs functional {vol_func}B '
          f'（省 {(1 - vol_func / vol_dense) * 100:.1f}%）')


if __name__ == '__main__':
    main()