#!/usr/bin/env python3
# =============================================================================
#  Project   : modelmap
#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v2.py
#  Purpose   : Run #2 — specificity-normalized skip + direction-level erasure
#  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)
# =============================================================================
"""expC run #2 (hypothesis registered before this run).

P1: layer-skip conditions rerun with general-damage (NLL) normalization.
P2: surgical test — erase the layer-l agreement direction (diff-of-means)
from every position of that layer's output; compare per-layer specific
damage (vs random-direction control) with the differential probe profile.
"""

from __future__ import annotations

import json
import random
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_pooled, install_taps

MODEL = "mlx-community/Qwen3-0.6B-4bit"
MAP_JSON = ROOT / "atlas" / "qwen3-0.6b-4bit" / "probes" / "v2" / "map.json"
K = 5
N_RANDOM_DRAWS = 20
N_RANDOM_DIRS = 3
SEED = 778

NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),
              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),
              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]
NEW_NEAR = ["across the yard", "above the workbench", "below the landing", "outside the depot",
            "around the corner", "beyond the fence", "beneath the awning", "atop the cabinet"]
NEUTRAL = ["The committee reviewed the plans before the meeting.",
           "A local historian kept detailed notes for years.",
           "The lead engineer questioned the original estimate.",
           "Her assistant preferred the older method.",
           "The night watchman described the process in a letter.",
           "An early visitor returned before the first frost.",
           "The town council approved the request after some debate.",
           "The apprentice carried the tools across the yard.",
           "According to the survey, the harbor required constant maintenance.",
           "The archive near the market attracted visitors from the region.",
           "The restored lighthouse stood at the edge of town.",
           "The workshop changed hands twice last century.",
           "The observatory remained open despite the storm.",
           "The vineyard was documented in the annual report.",
           "The glacier trail closed early in the season.",
           "The orchard supplied the market for decades.",
           "The library extended its hours during the recess.",
           "The clerk signed the manifest without a word.",
           "The surveyor traced the boundary along the quay.",
           "The curator shelved the samples behind the annex."]


def spearman(a, b):
    ra = np.argsort(np.argsort(a)).astype(float)
    rb = np.argsort(np.argsort(b)).astype(float)
    ra -= ra.mean(); rb -= rb.mean()
    return float((ra * rb).sum() / np.sqrt((ra**2).sum() * (rb**2).sum()))


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

    t0 = time.time()
    model, tokenizer = load(MODEL)
    taps = install_taps(model)
    n_layers = len(taps)

    mdoc = json.loads(MAP_JSON.read_text())
    agree = mdoc["properties"]["agreement"]
    diff_profile = np.array([a["selectivity_mean"] - t["selectivity_mean"]
                             for a, t in zip(agree["per_layer"]["A"],
                                             agree["twin_null_per_layer_A"])])
    ranked = list(np.argsort(-diff_profile))
    top_k, bottom_k = [int(x) for x in ranked[:K]], [int(x) for x in ranked[-K:]]

    # ---- enlarged held-out bank (new locations -> no dedup collisions)
    rng = random.Random(SEED)
    combos = [(n, loc) for n in NOUN_PAIRS for loc in NEW_NEAR]
    rng.shuffle(combos)
    pairs = []
    for (sg, pl), loc in combos:
        pairs.append({"prefix": f"The {sg} {loc}", "singular": True})
        pairs.append({"prefix": f"The {pl} {loc}", "singular": False})
    print(f"held-out pairs: {len(pairs)}")
    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]
    id_is, id_are = tokenizer.encode(" is")[0], tokenizer.encode(" are")[0]
    neutral_ids = [tokenizer.encode(s) for s in NEUTRAL]

    def margin() -> float:
        out = []
        for ids, p in zip(prefix_ids, pairs):
            logits = model(mx.array([ids]))[0, -1, :]
            mx.eval(logits)
            m = float(logits[id_is] - logits[id_are])
            out.append(m if p["singular"] else -m)
        return float(np.mean(out))

    def nll() -> float:
        tot, cnt = 0.0, 0
        for ids in neutral_ids:
            x = mx.array([ids])
            logits = model(x)[0]
            logp = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
            tgt = mx.array(ids[1:])
            picked = mx.take_along_axis(logp[:-1], tgt[:, None], axis=-1)
            mx.eval(picked)
            tot += float(-picked.sum()); cnt += len(ids) - 1
        return tot / cnt

    def with_skip(layers: set[int], fn):
        for i, t in enumerate(taps):
            t.skip = i in layers
        try:
            return fn()
        finally:
            for t in taps:
                t.skip = False

    base_m, base_nll = margin(), nll()
    print(f"baseline margin={base_m:+.4f} nll={base_nll:.4f}")

    # ---------------- P1: skip conditions with NLL normalization
    def skip_cell(layers):
        m = with_skip(set(layers), margin)
        n = with_skip(set(layers), nll)
        return {"layers": sorted(int(x) for x in layers),
                "margin_damage": base_m - m,
                "nll_damage": n - base_nll,
                "specificity": (base_m - m) / max(n - base_nll, 1e-3)}

    p1 = {"top": skip_cell(top_k), "bottom": skip_cell(bottom_k), "random": []}
    cand = [i for i in range(n_layers) if i not in set(top_k)]
    for d in range(N_RANDOM_DRAWS):
        p1["random"].append(skip_cell(random.Random(SEED + 1 + d).sample(cand, K)))
    rspec = np.array([c["specificity"] for c in p1["random"]])
    print(f"P1 specificity: top={p1['top']['specificity']:.3f} bottom={p1['bottom']['specificity']:.3f} "
          f"random mean={rspec.mean():.3f} p95={np.percentile(rspec, 95):.3f}")

    # ---------------- P2: direction-level erasure, all layers
    # directions from agreement_A mean-pooled reps at each layer
    items = [json.loads(l) for l in
             (ROOT / "benchmarks" / "promptsets" / "agreement_A.jsonl").read_text().splitlines()]
    toks = [tokenizer.encode(it["text"]) for it in items]
    labels = np.array([it["label"] for it in items])
    reps = capture_pooled(model, taps, toks)["mean"]  # (n, L, d)
    dirs, mus = [], []
    for layer in range(n_layers):
        x = reps[:, layer, :]
        mu = x.mean(0)
        u = x[labels == "correct"].mean(0) - x[labels == "violated"].mean(0)
        u = u / (np.linalg.norm(u) + 1e-8)
        mus.append(mu); dirs.append(u)

    def erase_fn(u_np, mu_np):
        u = mx.array(u_np.astype(np.float32))
        mu = mx.array(mu_np.astype(np.float32))
        def fn(out):
            h = out.astype(mx.float32)
            coef = ((h - mu) * u).sum(axis=-1, keepdims=True)
            return (h - coef * u).astype(out.dtype)
        return fn

    per_layer = []
    for layer in range(n_layers):
        taps[layer].edit = erase_fn(dirs[layer], mus[layer])
        m_agree = margin()
        taps[layer].edit = None
        rms = []
        for s in range(N_RANDOM_DIRS):
            ru = np.random.default_rng(1000 * layer + s).standard_normal(dirs[layer].shape)
            ru /= np.linalg.norm(ru)
            taps[layer].edit = erase_fn(ru.astype(np.float32), mus[layer])
            rms.append(margin())
            taps[layer].edit = None
        specific = (base_m - m_agree) - (base_m - float(np.mean(rms)))
        per_layer.append({"layer": layer, "agree_dir_damage": base_m - m_agree,
                          "random_dir_damage_mean": base_m - float(np.mean(rms)),
                          "specific_damage": specific})
        print(f"  L{layer:02d} agreeDir={base_m - m_agree:+.3f} randDir={base_m - float(np.mean(rms)):+.3f} "
              f"specific={specific:+.3f}", flush=True)

    spec_profile = np.array([r["specific_damage"] for r in per_layer])
    rho = spearman(spec_profile, diff_profile)
    perm_rng = np.random.default_rng(0)
    perms = np.array([spearman(perm_rng.permutation(spec_profile), diff_profile)
                      for _ in range(10_000)])
    p_perm = float((perms >= rho).mean())
    survives = bool(rho >= 0.4 and p_perm < 0.05)
    print(f"P2: Spearman rho={rho:+.3f} perm-p={p_perm:.4f} -> survives={survives}")

    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" / "expC_causal_verification" / ts
    outdir.mkdir(parents=True)
    (outdir / "results.json").write_text(json.dumps({
        "experiment": "expC_causal_verification", "run": 2,
        "scope": "NLL-normalized skip specificity + direction-level erasure scan",
        "commit": commit,
        "config": {"model": MODEL, "k": K, "n_random_draws": N_RANDOM_DRAWS,
                   "n_random_dirs": N_RANDOM_DIRS, "n_pairs": len(pairs), "seed": SEED,
                   "top_layers": sorted(top_k), "bottom_layers": sorted(bottom_k)},
        "manifest": manifest(),
        "baseline": {"margin": base_m, "nll": base_nll},
        "p1_skip_specificity": p1,
        "p2_direction_erasure": {"per_layer": per_layer, "spearman_rho": rho,
                                 "perm_p": p_perm, "survives": survives},
        "wall_seconds": round(time.time() - t0, 1),
    }, indent=2) + "\n")
    print(f"results -> {outdir / 'results.json'}")
    return 0


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