#!/usr/bin/env python3
# =============================================================================
#  Project   : modelmap
#  File      : experiments/micro/expH_capture_cost_frontier/implementation/benchmark_real.py
#  Purpose   : Run #3 — capture overhead on a real quantized checkpoint (mlx-lm)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Website   : https://modelmap.io
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""expH run #3 (see hypothesis.md, registered before this run).

Loads a real 4-bit model via mlx-lm, wraps every decoder layer with a
retaining tap, and measures prefill throughput in three modes:
plain / retain-all-layers / retain + NumPy conversion + mmap write.
This is the first quantized-model activation capture in the project —
the capability the Phase 1 survey found nowhere in Python tooling.
"""

from __future__ import annotations

import json
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[4]
sys.path.insert(0, str(ROOT / "benchmarks"))
from hardware_manifest import manifest

MODEL = "mlx-community/Qwen3-0.6B-4bit"
SEQ = 1024
N_FWD, WARMUP, REPEATS = 10, 3, 3
SEED = 0


class Tap:
    """Wraps a decoder layer; optionally retains its output."""

    def __init__(self, layer):
        self.layer = layer
        self.retained = None
        self.enabled = False

    def __call__(self, *args, **kwargs):
        out = self.layer(*args, **kwargs)
        if self.enabled:
            self.retained = out
        return out

    def __getattr__(self, name):  # delegate attribute access (e.g. .self_attn)
        return getattr(self.layer, name)


def main() -> int:
    import mlx.core as mx
    from mlx_lm import load

    mx.random.seed(SEED)
    model, tokenizer = load(MODEL)
    layers = model.model.layers
    n_layers = len(layers)
    taps = [Tap(l) for l in layers]
    model.model.layers = taps

    text = ("The internal cartography of local language models requires "
            "systematic measurement of every layer. ") * 60
    tokens = tokenizer.encode(text)[:SEQ]
    x = mx.array([tokens])

    # infer d_model at runtime (quantized embeddings pack their weight shapes)
    taps[0].enabled = True
    mx.eval(model(x))
    d_model = int(taps[0].retained.shape[-1])
    taps[0].enabled = False
    taps[0].retained = None
    print(f"model={MODEL} layers={n_layers} d_model={d_model} seq={len(tokens)}")

    workdir = Path(tempfile.mkdtemp(prefix="modelmap_expH3_"))
    store = np.memmap(workdir / "capture.raw", dtype=np.float16, mode="w+",
                      shape=(N_FWD * n_layers * len(tokens), d_model))

    def run(capture: bool, to_disk: bool) -> float:
        for t in taps:
            t.enabled = capture
            t.retained = None
        row = 0
        for _ in range(WARMUP):
            out = model(x)
            mx.eval(out, *[t.retained for t in taps if t.retained is not None])
        t0 = time.perf_counter()
        for _ in range(N_FWD):
            out = model(x)
            retained = [t.retained for t in taps] if capture else []
            mx.eval(out, *[r for r in retained if r is not None])
            if to_disk:
                for r in retained:
                    # model runs bf16 — cast in MLX (numpy has no bfloat16)
                    a = np.array(r.astype(mx.float16), copy=False).reshape(-1, d_model)
                    store[row:row + a.shape[0]] = a
                    row += a.shape[0]
        return (time.perf_counter() - t0) / N_FWD

    results = []
    try:
        for mode, cap, disk in (("plain", False, False),
                                ("retain", True, False),
                                ("retain+copy+write", True, True)):
            times = [run(cap, disk) for _ in range(REPEATS)]
            results.append({"backend": "mlx-lm", "mode": mode,
                            "s_per_forward": times,
                            "tokens_per_s_mean": len(tokens) / np.mean(times)})
            print(f"  {mode:22s} {np.mean(times)*1000:8.1f} ms/prefill "
                  f"({len(tokens)/np.mean(times):8.0f} tok/s)")
    finally:
        store.flush()
        shutil.rmtree(workdir, ignore_errors=True)

    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,
                            capture_output=True, text=True, check=False).stdout.strip()
    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
    outdir = ROOT / "results" / "expH_capture_cost_frontier" / ts
    outdir.mkdir(parents=True)
    doc = {
        "experiment": "expH_capture_cost_frontier",
        "run": 3,
        "scope": "capture overhead on a real 4-bit checkpoint (mlx-lm prefill)",
        "commit": commit,
        "config": {"model": MODEL, "seq": len(tokens), "n_layers": n_layers,
                   "d_model": int(d_model), "n_forwards": N_FWD,
                   "warmup": WARMUP, "repeats": REPEATS, "seed": SEED},
        "manifest": manifest(),
        "compute": results,
    }
    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")
    print(f"results -> {outdir / 'results.json'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())