#!/usr/bin/env python3
"""Export a trained SPIRA model as a compact JSON artifact for the web demo.

The TypeScript port (src/lib/spira/*) is compatible with the Python
prototype: FNV-1a-64 / splitmix64 / mulberry32 are shared specs, the input
and recurrent weights are REGENERATED from the seed in the browser, and only
the learned matrices (W_out, W_out_G, W_rec after STDP), the Markov table,
the patch registry, the hippocampal entries, and the learned homeostatic
state are shipped (binary packed as base64).

Format v2 adds the hierarchical slow area (G): its fixed projections ship
dense (they are small), its readout ships int8-quantized like the fast one,
and the learned state includes the slow homeostatic scalars, the top-down
gain beta, and the measured variance-matching fusion scale.

Usage:
    python3 -m spira.export_deploy --model artifacts/spira_model.npz \
        --out artifacts/spira_deploy.json [--am-cap 100000]
"""
from __future__ import annotations

import argparse
import base64
import json

import numpy as np

from .model import SPIRA


def b64(arr: np.ndarray) -> str:
    return base64.b64encode(np.ascontiguousarray(arr).tobytes()).decode("ascii")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--am-cap", type=int, default=80_000)
    ap.add_argument("--am-min-count", type=int, default=2)
    args = ap.parse_args()

    model = SPIRA.load(args.model)
    cfg = model.cfg
    core = model.core

    # ---- W_out: per-column int8 quantization -----------------------------
    W = core.W_out.astype(np.float32)          # (n_res, n_out)
    n_res, n_out = W.shape
    col_max = np.abs(W).max(axis=0)
    col_max[col_max == 0] = 1.0
    scale = (col_max / 127.0).astype(np.float32)
    Wq = np.round(W / scale[None, :]).astype(np.int8)

    # ---- W_rec after STDP: sparse triplets (row=pre, col=post) -----------
    rec = core.W_rec
    rows, cols = np.nonzero(rec)
    vals = rec[rows, cols].astype(np.float32)
    rec_entries = np.stack([rows.astype(np.int16), cols.astype(np.int16)], axis=1).astype(np.int16)

    # ---- hippocampus: deduplicated entries, binary packed ----------------
    entries = model.am.top_entries(min_count=args.am_min_count, cap=args.am_cap)
    n_e = len(entries)
    # pack: 4 x uint32 lanes + uint64 patch id + uint8 count = 25 bytes/entry
    packed = np.zeros((n_e, 25), dtype=np.uint8)
    for i, (k, v, c) in enumerate(entries):
        kl = k.astype("<u4").tobytes()
        packed[i, :16] = np.frombuffer(kl, dtype=np.uint8)
        packed[i, 16:24] = np.frombuffer(np.uint64(v).astype("<u8").tobytes(), dtype=np.uint8)
        packed[i, 24] = min(c, 255)
    flat = packed.reshape(-1)

    # ---- registry (vocab + all registered patches) ------------------------
    reg = model.registry
    # patch lens as uint8/16
    lens = np.diff(reg.offsets).astype(np.uint16)

    # ---- hierarchical slow area (v2) --------------------------------------
    hier = None
    if core.hier:
        WG = core.W_outG.astype(np.float32)        # (n_g, n_out)
        g_max = np.abs(WG).max(axis=0)
        g_max[g_max == 0] = 1.0
        g_scale = (g_max / 127.0).astype(np.float32)
        WGq = np.round(WG / g_scale[None, :]).astype(np.int8)
        hier = {
            "n_g": core.n_g,
            "win": core.win,
            "lam_g": core.lam_g,
            "tau_trace_g": core.tau_trace_g,
            "target_rate_g": core.target_rate_g,
            "eta_homeo_g": core.eta_homeo_g,
            "beta": float(core.beta),
            "fusion_scale": float(model.fusion["scale"]),
            # fixed projections ship dense (small): L->G and G recurrent
            "w_ing_b64": b64(core.W_inG),           # (n_res, n_g) float32
            "w_recg_b64": b64(core.W_recG),         # (n_g, n_g) float32
            "w_outg_q_b64": b64(WGq),
            "w_outg_scale_b64": b64(g_scale),
        }

    deploy = {
        "format": "spira-deploy-v2",
        "config": {
            "d_in": cfg.d_in, "k_in": cfg.k_in, "n_res": cfg.n_res,
            "vocab_size": cfg.vocab_size, "seed": cfg.seed,
            "in_fanout": core.in_fanout, "rec_fanout": core.rec_fanout,
            "inh_frac": core.inh_frac, "lam": core.lam, "base_thr": core.base_thr,
            "thr_min": core.thr_min, "thr_max": core.thr_max, "refrac": core.refrac,
            "tau_trace": core.tau_trace, "w_in_scale": core.w_in_scale,
            "w_rec_scale": core.w_rec_scale, "target_rate": core.target_rate,
            "eta_homeo": core.eta_homeo,
            "hierarchy": bool(cfg.hierarchy), "n_g": cfg.n_g, "win": cfg.win,
            "am_m": cfg.am_m, "am_min_lanes": cfg.am_min_lanes,
            "am_size": 1 << 20,  # browser-side table is smaller
            "beta_am": cfg.beta_am, "am_override": cfg.am_override,
            "temperature": cfg.temperature, "top_k": cfg.top_k,
        },
        "ept": {
            "budget_bits": model.ept.budget_bits, "min_len": model.ept.min_len,
            "max_len": model.ept.max_len, "alpha": model.ept.alpha,
            # 256x256 float32 log2-probs (row = prev byte); TS negates for surprise
            "log2p_b64": b64(model.ept.log2p.astype(np.float32)),
        },
        "reservoir": {
            "w_rec_rows": b64(rows.astype(np.int32)),
            "w_rec_cols": b64(cols.astype(np.int32)),
            "w_rec_vals_b64": b64(vals),
            "w_rec_n": int(len(rows)),
        },
        "readout": {
            "shape": [int(n_res), int(n_out)],
            "vocab_ids_hex": [hex(int(x)) for x in model.vocab_ids],
            "w_out_q_b64": b64(Wq),
            "w_out_scale_b64": b64(scale),
        },
        "hier": hier,
        "am": {
            "n_entries": n_e,
            "packed_b64": b64(flat),
            "entry_bytes": 25,
        },
        "registry": {
            "n": int(len(reg.ids)),
            "ids_b64": b64(reg.ids.astype("<u8")),
            "lens_b64": b64(lens),
            "buffer_b64": b64(np.frombuffer(reg.buffer, dtype=np.uint8)),
        },
        "core_state": {
            # learned homeostatic scalars needed for faithful inference
            "x_mean_b64": b64(model.core.x_mean.astype(np.float32)),
            "thr_b64": b64(model.core.thr.astype(np.float32)),
            "rhat_b64": b64(model.core.rhat.astype(np.float32)),
            **({
                "x_mean_g_b64": b64(model.core.x_mean_G.astype(np.float32)),
                "thr_g_b64": b64(model.core.thr_G.astype(np.float32)),
                "rhat_g_b64": b64(model.core.rhat_G.astype(np.float32)),
            } if core.hier else {}),
        },
        "fusion": model.fusion,
        "meta": model.meta,
    }
    with open(args.out, "w") as f:
        json.dump(deploy, f)
    import os
    sz = os.path.getsize(args.out)
    print(f"deploy artifact -> {args.out} ({sz/1e6:.1f} MB)")
    print(f"  vocab {len(model.vocab_ids)} | registry {len(reg.ids)} | am entries {n_e} "
          f"| W_out {Wq.shape} int8 | W_rec triplets {len(rows)}"
          + (f" | hier n_g={core.n_g} win={core.win} beta={core.beta:.4f}" if core.hier else ""))


if __name__ == "__main__":
    main()