#!/usr/bin/env python3
# =============================================================================
#  Project   : modelmap
#  File      : experiments/micro/expA_probe_reliability/implementation/benchmark.py
#  Purpose   : expA run #1 — probe noise floor on a real 4-bit model + nulls
#  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)
# =============================================================================
"""expA run #1 (hypothesis registered in hypothesis.md before this run).

Per-layer linear probes with the full controls doctrine on Qwen3-0.6B-4bit:
3 properties x 2 disjoint promptsets x 28 layers x 5 seeds, shuffled-label
controls built into every probe, random-init architecture twin as the
structure-from-architecture null, bootstrap p-values + BH-FDR across layer
scans, replication rate on top-k layer sets.
"""

from __future__ import annotations

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

import numpy as np

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

from modelmap.capture.mlx_capture import capture_mean_pooled, install_taps, random_init_twin
from modelmap.probes.linear import probe_with_control
from modelmap.stats.replication import bh_fdr, bootstrap_ci, replication_rate

MODEL = "mlx-community/Qwen3-0.6B-4bit"
PROPERTIES = ("lang_id", "code_prose", "arith")
SETS = ("A", "B")
SEEDS = (0, 1, 2, 3, 4)
TOP_K = 5
FDR_Q = 0.05
PROMPTS = ROOT / "benchmarks" / "promptsets"


def load_set(name: str) -> tuple[list[str], np.ndarray]:
    items = [json.loads(l) for l in (PROMPTS / f"{name}.jsonl").read_text().splitlines()]
    return [it["text"] for it in items], np.array([it["label"] for it in items])


def probe_grid(reps: np.ndarray, labels: np.ndarray) -> list[dict]:
    """All (layer, seed) probe cells for one representation tensor."""
    out = []
    n_layers = reps.shape[1]
    for layer in range(n_layers):
        for seed in SEEDS:
            r = probe_with_control(reps[:, layer, :], labels, seed=seed)
            out.append({"layer": layer, "seed": seed,
                        "task_acc": r.task_accuracy,
                        "control_acc": r.control_accuracy,
                        "selectivity": r.selectivity})
    return out


def summarize(cells: list[dict], n_layers: int) -> dict:
    """Per-layer aggregation: means, seed SD, bootstrap CI + p on selectivity."""
    rng = np.random.default_rng(0)
    layers = []
    pvals = []
    for layer in range(n_layers):
        sel = np.array([c["selectivity"] for c in cells if c["layer"] == layer])
        acc = np.array([c["task_acc"] for c in cells if c["layer"] == layer])
        point, lo, hi = bootstrap_ci(sel, rng=rng)
        boots = rng.choice(sel, size=(10_000, sel.size)).mean(axis=1)
        p = float(max((boots <= 0).mean(), 1 / 10_000))
        pvals.append(p)
        layers.append({"layer": layer,
                       "task_acc_mean": float(acc.mean()),
                       "task_acc_seed_sd": float(acc.std(ddof=1)),
                       "selectivity_mean": point,
                       "selectivity_ci": [lo, hi],
                       "p_boot": p})
    disc = bh_fdr(np.array(pvals), q=FDR_Q)
    for row, d in zip(layers, disc):
        row["fdr_significant"] = bool(d)
    top_sets = []
    for seed in SEEDS:
        acc_by_layer = [(c["layer"], c["task_acc"]) for c in cells if c["seed"] == seed]
        top = {l for l, _ in sorted(acc_by_layer, key=lambda t: -t[1])[:TOP_K]}
        top_sets.append(top)
    rep_point, rep_lo, rep_hi = replication_rate(top_sets)
    return {"layers": layers,
            "replication_rate_topk": {"k": TOP_K, "point": rep_point, "ci": [rep_lo, rep_hi]},
            "n_fdr_significant": int(disc.sum())}


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

    t_start = time.time()
    mx.random.seed(0)
    model, tokenizer = load(MODEL)
    taps = install_taps(model)
    model_path = hf_repo_to_path(MODEL)

    twin = random_init_twin(model_path)
    twin_taps = install_taps(twin)

    results = {"real": {}, "twin": {}}
    grids = {"real": {}, "twin": {}}

    for prop in PROPERTIES:
        for s in SETS:
            name = f"{prop}_{s}"
            texts, labels = load_set(name)
            toks = [tokenizer.encode(t) for t in texts]
            print(f"capture real  {name} ({len(texts)} prompts)…", flush=True)
            reps = capture_mean_pooled(model, taps, toks)
            grids["real"][name] = (probe_grid(reps, labels), reps.shape[1])
            if s == "A":  # architecture null on the A sets
                print(f"capture twin  {name}…", flush=True)
                reps_t = capture_mean_pooled(twin, twin_taps, toks)
                grids["twin"][name] = (probe_grid(reps_t, labels), reps_t.shape[1])

    for kind in ("real", "twin"):
        for name, (cells, n_layers) in grids[kind].items():
            results[kind][name] = summarize(cells, n_layers)
            results[kind][name]["cells"] = cells

    # -------- headline numbers
    summary = {}
    for prop in PROPERTIES:
        a = results["real"][f"{prop}_A"]["layers"]
        b = results["real"][f"{prop}_B"]["layers"]
        seed_sd = float(np.mean([r["task_acc_seed_sd"] for r in a + b]))
        shifts = [abs(ra["task_acc_mean"] - rb["task_acc_mean"]) for ra, rb in zip(a, b)]
        n_shift = int(sum(s > max(seed_sd, 1e-9) for s in shifts))
        twin_sel = float(max(r["selectivity_mean"] for r in results["twin"][f"{prop}_A"]["layers"]))
        summary[prop] = {
            "mean_seed_sd": seed_sd,
            "mean_dataset_shift": float(np.mean(shifts)),
            "layers_shift_gt_seed_sd": n_shift,
            "n_layers": len(a),
            "max_task_acc_A": float(max(r["task_acc_mean"] for r in a)),
            "max_task_acc_B": float(max(r["task_acc_mean"] for r in b)),
            "twin_max_selectivity": twin_sel,
            "replication_topk_A": results["real"][f"{prop}_A"]["replication_rate_topk"]["point"],
            "replication_topk_B": results["real"][f"{prop}_B"]["replication_rate_topk"]["point"],
        }
        print(f"{prop:12s} seedSD={seed_sd:.4f} shift={summary[prop]['mean_dataset_shift']:.4f} "
              f"layers(shift>sd)={n_shift}/{len(a)} maxAccA={summary[prop]['max_task_acc_A']:.3f} "
              f"twinMaxSel={twin_sel:.3f} repl={summary[prop]['replication_topk_A']:.2f}")

    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" / "expA_probe_reliability" / ts
    outdir.mkdir(parents=True)
    doc = {
        "experiment": "expA_probe_reliability",
        "run": 1,
        "scope": "probe noise floor, 3 properties x 2 sets x 28 layers x 5 seeds + twin null",
        "commit": commit,
        "config": {"model": MODEL, "seeds": list(SEEDS), "top_k": TOP_K, "fdr_q": FDR_Q,
                   "promptsets": json.loads((PROMPTS / "manifest.json").read_text())},
        "manifest": manifest(),
        "summary": summary,
        "results": results,
        "wall_seconds": round(time.time() - t_start, 1),
    }
    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")
    print(f"results -> {outdir / 'results.json'}  ({doc['wall_seconds']} s)")
    return 0


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