#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/micro/expA_weight_concentration/benchmark.py
#  Purpose   : Per-token FFN block-energy concentration (trace shared with expB)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""Experiment A — weight contribution concentration (charter §9.A).

Records SwiGLU intermediate-activation energy per 64-neuron block, per token,
per layer, on the bf16 reference model. Writes concentration aggregates to
results.json and the raw block-energy trace (npz) for expB.

Usage:
    .venv/bin/python benchmark.py [--per-domain 8] [--gen-tokens 128] [--block 64]
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mlx_lm import load

REPO_ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO_ROOT / "benchmarks"))
sys.path.insert(0, str(REPO_ROOT / "src"))
from hardware_manifest import collect_manifest  # noqa: E402
from localvm.quality.decision_stats import greedy_generate  # noqa: E402


class DownProjRecorder(nn.Module):
    """Wraps a down_proj Linear; records block energy of its input (the SwiGLU
    intermediate activation) for positions in `window`."""

    def __init__(self, inner: nn.Module, block: int):
        super().__init__()
        self.inner = inner
        self.block = block
        self.window: tuple[int, int] | None = None
        self.block_energy: np.ndarray | None = None
        self.neuron_energy: np.ndarray | None = None

    def __call__(self, x):
        if self.window is not None:
            a, b = self.window
            h = x[0, a:b].astype(mx.float32)
            e = mx.square(h)
            be = e.reshape(h.shape[0], -1, self.block).sum(axis=-1)
            # normalize per position: raw energies overflow float16 storage,
            # and only relative importance matters for expA/expB
            be = be / (be.sum(axis=-1, keepdims=True) + 1e-12)
            mx.eval(be)
            self.block_energy = np.array(be)
            self.neuron_energy = np.array(e)  # (T, D_int) — reduced by caller
        return self.inner(x)


def concentration_stats(energy: np.ndarray, fracs=(0.1, 0.2, 0.4, 0.6),
                        targets=(0.90, 0.95, 0.99)) -> dict:
    """energy: (T, N). Returns mean energy captured by top-f fraction and mean
    fraction of units needed to reach target energy."""
    T, N = energy.shape
    srt = np.sort(energy, axis=1)[:, ::-1]
    csum = np.cumsum(srt, axis=1)
    total = csum[:, -1:] + 1e-12
    frac_captured = {}
    for f in fracs:
        k = max(1, int(round(f * N)))
        frac_captured[f] = float(np.mean(csum[:, k - 1] / total[:, 0]))
    needed = {}
    ratio = csum / total
    for t in targets:
        idx = np.argmax(ratio >= t, axis=1) + 1
        needed[t] = float(np.mean(idx / N))
    return {"top_frac_energy": frac_captured, "frac_needed_for": needed}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="mlx-community/Qwen3-1.7B-bf16")
    ap.add_argument("--gen-tokens", type=int, default=128)
    ap.add_argument("--per-domain", type=int, default=8)
    ap.add_argument("--block", type=int, default=16,
                    help="trace granularity; analysis also derives 4x-coarser blocks")
    args = ap.parse_args()

    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]
    print(f"loading {args.model} …", flush=True)
    model, tokenizer = load(args.model)
    layers = model.model.layers
    n_layers = len(layers)

    recorders = []
    for layer in layers:
        rec = DownProjRecorder(layer.mlp.down_proj, args.block)
        layer.mlp.down_proj = rec
        recorders.append(rec)

    trajectories = []
    t0 = time.time()
    for domain, plist in domains.items():
        for prompt in plist[: args.per_domain]:
            ids = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}], add_generation_prompt=True)
            for r in recorders:
                r.window = None  # no recording during generation
            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)
            if len(gen) >= 8:
                trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})
    print(f"{len(trajectories)} trajectories in {time.time()-t0:.0f}s", flush=True)

    block_traces = []   # per traj: (T, n_layers, n_blocks) f16
    neuron_stats = []   # per traj per layer concentration dicts
    index = []
    for ti, t in enumerate(trajectories):
        a, b = t["start"] - 1, len(t["full_ids"]) - 1
        for r in recorders:
            r.window = (a, b)
        model(mx.array(t["full_ids"])[None])
        per_layer_blocks = np.stack([r.block_energy for r in recorders], axis=1)  # (T, L, B)
        block_traces.append(per_layer_blocks.astype(np.float16))
        neuron_stats.append([concentration_stats(r.neuron_energy) for r in recorders])
        for r in recorders:
            r.neuron_energy = None
        index.append({"traj": ti, "domain": t["domain"], "n_pos": b - a})
        if (ti + 1) % 12 == 0:
            print(f"  traced {ti+1}/{len(trajectories)}", flush=True)

    all_blocks = np.concatenate(block_traces, axis=0)  # (P, L, B)
    P, L, B = all_blocks.shape
    print(f"trace shape {all_blocks.shape}", flush=True)

    # expA aggregates at trace granularity and 4x-coarser derived granularity
    coarse = all_blocks.reshape(P, L, B // 4, 4).astype(np.float32).sum(axis=-1)
    per_layer = [concentration_stats(all_blocks[:, li, :].astype(np.float32)) for li in range(L)]
    overall = concentration_stats(all_blocks.reshape(P * L, B).astype(np.float32))
    overall_coarse = concentration_stats(coarse.reshape(P * L, B // 4))
    per_domain = {}
    pos_domain = np.concatenate([[ix["domain"]] * ix["n_pos"] for ix in index])
    for dom in domains:
        sel = all_blocks[pos_domain == dom]
        per_domain[dom] = concentration_stats(sel.reshape(-1, B).astype(np.float32))
    # neuron-granularity mean across trajectories/layers
    neuron_overall = {
        "top_frac_energy": {f: float(np.mean([s["top_frac_energy"][f] for ns in neuron_stats for s in ns]))
                            for f in (0.1, 0.2, 0.4, 0.6)},
        "frac_needed_for": {t: float(np.mean([s["frac_needed_for"][t] for ns in neuron_stats for s in ns]))
                            for t in (0.90, 0.95, 0.99)},
    }

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "expA_weight_concentration" / ts
    out_dir.mkdir(parents=True)
    np.savez_compressed(out_dir / "block_energy_trace.npz",
                        blocks=all_blocks,
                        domains=pos_domain,
                        traj_id=np.concatenate([[ix["traj"]] * ix["n_pos"] for ix in index]))
    (out_dir / "results.json").write_text(json.dumps({
        "experiment": "expA_weight_concentration",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": collect_manifest(),
        "config": vars(args),
        "n_positions": int(P), "n_layers": int(L), "n_blocks": int(B),
        "coarse_block_granularity": {"block_size": args.block * 4, "overall": overall_coarse},
        "block_granularity": {"overall": overall,
                              "per_layer": {str(i): s for i, s in enumerate(per_layer)},
                              "per_domain": per_domain},
        "neuron_granularity": neuron_overall,
    }, indent=2, default=float))
    print(f"\nwrote {out_dir}/results.json (+ block_energy_trace.npz for expB)")
    print("overall block-64:", json.dumps(overall, default=float))
    print("neuron-level    :", json.dumps(neuron_overall, default=float))


if __name__ == "__main__":
    main()