"""Write web/model.bin: the compiled model in sparse form for the browser.

The dense model.safetensors is ~93% zeros. Each MLP neuron of the compiled
transformer is fully described by its context tokens (which slots its gate
matches) and the logit changes its down-projection writes; the attention and
embedding weights are fixed by construction. This file stores exactly that, so
web/autocomplete.ts can evaluate only the neurons whose context matches: the
same computation as the transformer, without the zeros.

Layout (little-endian):
    "TACB" | u32 version | u32 header bytes | header JSON | pad to 8
    sections (each padded to 8), offsets relative to the first section:
      base                 f32[V]           unigram neuron (always on)
      ctx{k}   k=1..order  u16/u32[n_k * k] context token ids, oldest first, sorted
      support{k}           f32[n_k]         weighted count of the context
      offsets{k}           u32[n_k + 1]     slice of this neuron's deltas
      dids{k}, dvals{k}    u16/u32, f32     next-token ids and logit changes

    python export_web.py      # uses the "web" settings and "policy" in compile_settings.json
"""

import json
from dataclasses import asdict
from pathlib import Path

import numpy as np

import build_weights as bw
import suggest as sg

ROOT = Path(__file__).parent
OUT = ROOT / "web" / "model.bin"
MAGIC, VERSION = b"TACB", 1


def load_policy() -> sg.Policy:
    return sg.Policy(**json.loads(bw.SETTINGS_PATH.read_text())["policy"])


def export(vocab: list[str], ngram: bw.NGram, oov: dict, policy: sg.Policy, path: Path) -> dict:
    id_dtype = np.uint16 if len(vocab) < 65536 else np.uint32
    sections, orders = [], []

    def add(name: str, arr: np.ndarray) -> None:
        sections.append((name, np.ascontiguousarray(arr).tobytes()))

    add("base", ngram.base_logits.astype(np.float32))
    for k in range(1, ngram.order + 1):
        ctxs = sorted(ngram.tables[k])
        offsets, dids, dvals = [0], [], []
        for ctx in ctxs:
            d = ngram.delta(ctx)
            dids.extend(d.keys())
            dvals.extend(d.values())
            offsets.append(len(dids))
        add(f"ctx{k}", np.array(ctxs, dtype=id_dtype).reshape(-1))
        add(f"support{k}", np.array([ngram.support[c] for c in ctxs], dtype=np.float32))
        add(f"offsets{k}", np.array(offsets, dtype=np.uint32))
        add(f"dids{k}", np.array(dids, dtype=id_dtype))
        add(f"dvals{k}", np.array(dvals, dtype=np.float32))
        orders.append({"k": k, "contexts": len(ctxs), "pairs": len(dids)})

    table, offset = [], 0
    for name, data in sections:
        table.append({"name": name, "offset": offset, "bytes": len(data)})
        offset += -(-len(data) // 8) * 8
    header = {
        "n_prev": ngram.s.n_prev,
        "vocab": vocab,
        "eos_id": bw.EOS_ID,
        "unk_id": bw.UNK_ID,
        "id_bytes": np.dtype(id_dtype).itemsize,
        "orders": orders,
        "oov_prefix_rates": {p: round(r, 4) for p, r in oov.items()},
        "policy": asdict(policy),
        "stop_after": sorted(sg.STOP_AFTER),
        "settings": asdict(ngram.s),
        "sections": table,
    }
    head = json.dumps(header, separators=(",", ":")).encode()
    path.parent.mkdir(exist_ok=True)
    with open(path, "wb") as f:
        f.write(MAGIC + np.array([VERSION, len(head)], dtype="<u4").tobytes() + head)
        f.write(b"\0" * (-(12 + len(head)) % 8))
        for _, data in sections:
            f.write(data + b"\0" * (-len(data) % 8))
    return {"bytes": path.stat().st_size, "contexts": sum(o["contexts"] for o in orders), "pairs": sum(o["pairs"] for o in orders)}


def build(settings: bw.Settings, sources: dict, policy: sg.Policy, path: Path = OUT) -> dict:
    vocab, ngram, stats = bw.compile_tables(settings, sources)
    oov = sg.oov_prefix_rates(settings, sources, vocab)
    return {**stats, **export(vocab, ngram, oov, policy, path)}


def main() -> None:
    info = build(bw.load_settings("web"), bw.load_sources(), load_policy())
    print(json.dumps({**info, "mb": round(info["bytes"] / 1e6, 1), "out": str(OUT)}, indent=2))


if __name__ == "__main__":
    main()