#!/usr/bin/env python3
# =============================================================================
#  Project   : modelmap
#  File      : experiments/micro/expH_capture_cost_frontier/implementation/benchmark.py
#  Purpose   : Run #1 — storage-format throughput on APFS + hook overhead MLX/MPS
#  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 / MPS
#  License   : All rights reserved (research code)
# =============================================================================
"""expH run #1. See hypothesis.md (registered before this run).

Part A — activation-store formats on APFS (warm cache, declared):
    raw np.memmap | safetensors (mmap) | zarr zstd | zarr uncompressed
    sequential write, sequential scan, random-batch reads (SAE-shuffle).
Part B — capture overhead on a synthetic 12-layer transformer (fp16):
    torch-MPS forward hooks vs MLX retained arrays;
    modes: plain / retain-on-device / retain + CPU copy + mmap write.
"""

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

REPEATS = 3
ROWS, DIM = 100_000, 4096          # ~0.82 GB fp16 per format
WRITE_CHUNK = 4_096                 # rows per write
BATCH = 4_096                       # rows per random-batch read
N_BATCHES = 24
SEED = 0

# Part B config
L, D_MODEL, N_HEADS, SEQ, BSZ = 12, 1024, 4, 512, 8
N_FWD, N_FWD_WRITE, WARMUP = 20, 10, 3


def _timed(fn, repeats=REPEATS):
    out = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        fn()
        out.append(time.perf_counter() - t0)
    return out


# ---------------------------------------------------------------- Part A
def bench_storage(workdir: Path) -> list[dict]:
    rng = np.random.default_rng(SEED)
    data = rng.standard_normal((WRITE_CHUNK, DIM)).astype(np.float16)
    batches = [rng.integers(0, ROWS, BATCH) for _ in range(N_BATCHES)]
    total_bytes = ROWS * DIM * 2
    batch_bytes = BATCH * DIM * 2 * N_BATCHES
    results = []

    def record(fmt, op, times, nbytes):
        results.append({
            "format": fmt, "op": op, "bytes": nbytes,
            "seconds": times, "gb_per_s_mean": nbytes / 2**30 / np.mean(times),
        })
        print(f"  {fmt:18s} {op:12s} {nbytes/2**30/np.mean(times):8.2f} GB/s")

    # ---- raw np.memmap
    p = workdir / "acts.raw"
    def write_raw():
        m = np.memmap(p, dtype=np.float16, mode="w+", shape=(ROWS, DIM))
        for i in range(0, ROWS, WRITE_CHUNK):
            end = min(i + WRITE_CHUNK, ROWS)
            m[i:end] = data[: end - i]
        m.flush(); del m
    record("raw-mmap", "write", _timed(write_raw), total_bytes)
    m = np.memmap(p, dtype=np.float16, mode="r", shape=(ROWS, DIM))
    record("raw-mmap", "seq-scan", _timed(lambda m=m: float(np.asarray(m).sum(dtype=np.float32))), total_bytes)
    record("raw-mmap", "random-batch", _timed(lambda m=m: [m[b].sum(dtype=np.float32) for b in batches]), batch_bytes)
    del m

    # ---- safetensors (mmap-backed numpy)
    from safetensors import safe_open
    from safetensors.numpy import save_file
    p = workdir / "acts.safetensors"
    full = np.memmap(workdir / "acts.raw", dtype=np.float16, mode="r", shape=(ROWS, DIM))
    def write_st(full=full, p=p):
        save_file({"acts": np.asarray(full)}, str(p))
    record("safetensors", "write", _timed(write_st), total_bytes)
    def open_st():
        return safe_open(str(p), framework="np")
    f = open_st()
    t = f.get_tensor("acts")   # mmap-backed load
    record("safetensors", "seq-scan", _timed(lambda t=t: float(t.sum(dtype=np.float32))), total_bytes)
    record("safetensors", "random-batch", _timed(lambda t=t: [t[b].sum(dtype=np.float32) for b in batches]), batch_bytes)
    del t, f, full

    # ---- zarr (zstd default) and uncompressed
    import zarr
    for codec, name in ((None, "zarr-uncompressed"), ("default", "zarr-zstd")):
        p = workdir / f"acts_{name}.zarr"
        kwargs = {} if codec == "default" else {"compressors": None}
        def write_zarr(p=p, kwargs=kwargs):
            if p.exists():
                shutil.rmtree(p)
            z = zarr.create_array(store=str(p), shape=(ROWS, DIM), chunks=(WRITE_CHUNK, DIM),
                                  dtype=np.float16, **kwargs)
            for i in range(0, ROWS, WRITE_CHUNK):
                end = min(i + WRITE_CHUNK, ROWS)
                z[i:end] = data[: end - i]
        record(name, "write", _timed(write_zarr), total_bytes)
        z = zarr.open_array(store=str(p), mode="r")
        record(name, "seq-scan", _timed(lambda z=z: float(z[:].sum(dtype=np.float32))), total_bytes)
        record(name, "random-batch", _timed(lambda z=z: [z[np.sort(b)].sum(dtype=np.float32) for b in batches]), batch_bytes)

    return results


# ---------------------------------------------------------------- Part B
def bench_torch_mps(workdir: Path) -> list[dict]:
    import torch
    from torch import nn
    assert torch.backends.mps.is_available(), "MPS required"
    dev, dt = torch.device("mps"), torch.float16

    class Block(nn.Module):
        def __init__(self):
            super().__init__()
            self.ln1, self.ln2 = nn.LayerNorm(D_MODEL), nn.LayerNorm(D_MODEL)
            self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, batch_first=True)
            self.mlp = nn.Sequential(nn.Linear(D_MODEL, 4 * D_MODEL), nn.GELU(),
                                     nn.Linear(4 * D_MODEL, D_MODEL))
        def forward(self, x):
            h = self.ln1(x)
            x = x + self.attn(h, h, h, need_weights=False)[0]
            return x + self.mlp(self.ln2(x))

    torch.manual_seed(SEED)
    model = nn.Sequential(*[Block() for _ in range(L)]).to(dev, dt).eval()
    x = torch.randn(BSZ, SEQ, D_MODEL, device=dev, dtype=dt)
    store = np.memmap(workdir / "torch_capture.raw", dtype=np.float16, mode="w+",
                      shape=(N_FWD_WRITE * L * BSZ * SEQ, D_MODEL))

    def run(n, capture, to_disk):
        captured, row = [], 0
        hooks = []
        if capture:
            def hook(_m, _i, out):
                captured.append(out)
            hooks = [b.register_forward_hook(hook) for b in model]
        with torch.no_grad():
            for _ in range(WARMUP):
                model(x)
            torch.mps.synchronize()
            t0 = time.perf_counter()
            for _ in range(n):
                captured.clear()
                model(x)
                if to_disk:
                    for c in captured:
                        a = c.to("cpu").numpy().reshape(-1, D_MODEL)
                        store[row:row + a.shape[0]] = a
                        row += a.shape[0]
            torch.mps.synchronize()
            dt_s = time.perf_counter() - t0
        for h in hooks:
            h.remove()
        return dt_s / n

    out = []
    for mode, cap, disk, n in (("plain", False, False, N_FWD),
                               ("retain", True, False, N_FWD),
                               ("retain+copy+write", True, True, N_FWD_WRITE)):
        times = [run(n, cap, disk) for _ in range(REPEATS)]
        out.append({"backend": "torch-mps", "mode": mode,
                    "s_per_forward": times,
                    "tokens_per_s_mean": BSZ * SEQ / np.mean(times)})
        print(f"  torch-mps {mode:22s} {np.mean(times)*1000:8.1f} ms/fwd")
    return out


def bench_mlx(workdir: Path) -> list[dict]:
    import mlx.core as mx
    import mlx.nn as mnn

    class Block(mnn.Module):
        def __init__(self):
            super().__init__()
            self.ln1, self.ln2 = mnn.LayerNorm(D_MODEL), mnn.LayerNorm(D_MODEL)
            self.attn = mnn.MultiHeadAttention(D_MODEL, N_HEADS)
            self.fc1, self.fc2 = mnn.Linear(D_MODEL, 4 * D_MODEL), mnn.Linear(4 * D_MODEL, D_MODEL)
        def __call__(self, x):
            h = self.ln1(x)
            x = x + self.attn(h, h, h)
            return x + self.fc2(mnn.gelu(self.fc1(self.ln2(x))))

    mx.random.seed(SEED)
    blocks = [Block() for _ in range(L)]
    for b in blocks:
        b.set_dtype(mx.float16)
    x = mx.random.normal((BSZ, SEQ, D_MODEL)).astype(mx.float16)
    store = np.memmap(workdir / "mlx_capture.raw", dtype=np.float16, mode="w+",
                      shape=(N_FWD_WRITE * L * BSZ * SEQ, D_MODEL))

    def fwd(capture):
        captured, h = [], x
        for b in blocks:
            h = b(h)
            if capture:
                captured.append(h)
        return h, captured

    def run(n, capture, to_disk):
        row = 0
        for _ in range(WARMUP):
            out, cap = fwd(capture)
            mx.eval(out, *cap)
        t0 = time.perf_counter()
        for _ in range(n):
            out, cap = fwd(capture)
            mx.eval(out, *cap)
            if to_disk:
                for c in cap:
                    a = np.array(c, copy=False).reshape(-1, D_MODEL)
                    store[row:row + a.shape[0]] = a
                    row += a.shape[0]
        return (time.perf_counter() - t0) / n

    out = []
    for mode, cap, disk, n in (("plain", False, False, N_FWD),
                               ("retain", True, False, N_FWD),
                               ("retain+copy+write", True, True, N_FWD_WRITE)):
        times = [run(n, cap, disk) for _ in range(REPEATS)]
        out.append({"backend": "mlx", "mode": mode,
                    "s_per_forward": times,
                    "tokens_per_s_mean": BSZ * SEQ / np.mean(times)})
        print(f"  mlx       {mode:22s} {np.mean(times)*1000:8.1f} ms/fwd")
    return out


# ---------------------------------------------------------------- main
def main() -> int:
    workdir = Path(tempfile.mkdtemp(prefix="modelmap_expH_"))
    print(f"expH run #1 — workdir {workdir}")
    try:
        print("Part A — storage formats (warm cache):")
        storage = bench_storage(workdir)
        print("Part B — capture overhead:")
        compute = bench_torch_mps(workdir) + bench_mlx(workdir)
    finally:
        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": 1,
        "scope": "storage formats (warm cache) + capture overhead, synthetic model",
        "commit": commit,
        "config": {
            "storage": {"rows": ROWS, "dim": DIM, "write_chunk": WRITE_CHUNK,
                        "batch": BATCH, "n_batches": N_BATCHES, "repeats": REPEATS,
                        "cache": "warm (declared limitation; cold pass = run #2)"},
            "compute": {"layers": L, "d_model": D_MODEL, "heads": N_HEADS,
                        "seq": SEQ, "batch": BSZ, "dtype": "float16",
                        "n_forwards": N_FWD, "warmup": WARMUP, "repeats": REPEATS},
            "seed": SEED,
        },
        "manifest": manifest(),
        "storage": storage,
        "compute": compute,
    }
    (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())