#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/micro/expG_decision_stability/benchmark.py
#  Purpose   : Joint (cheap-pass margin × agreement) matrix across bit-widths —
#              the Gate-Zero measurement for margin-gated escalation (G02/G23)
#  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 G — decision stability (charter §9.G).

Generates greedy continuations with a bf16 reference model, teacher-forces
low-bit quantized variants over the same sequences, and records per-position
margin/agreement/KL. Outputs the joint matrix, AUROC of margin as a
disagreement detector, and escalation curves.

Usage:
    .venv/bin/python benchmark.py [--model mlx-community/Qwen3-1.7B-bf16]
        [--gen-tokens 128] [--bits 2,3,4,8] [--per-domain 8]
"""

from __future__ import annotations

import argparse
import gc
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"))
from hardware_manifest import collect_manifest  # noqa: E402


def greedy_generate(model, tokenizer, prompt_ids: list[int], n_tokens: int) -> list[int]:
    """Greedy generation without sampling helpers — deterministic, no cache reuse
    across prompts. Returns generated token ids."""
    tokens = list(prompt_ids)
    generated = []
    from mlx_lm.models.cache import make_prompt_cache

    cache = make_prompt_cache(model)
    inp = mx.array(tokens)[None]
    for _ in range(n_tokens):
        logits = model(inp, cache=cache)
        nxt = int(mx.argmax(logits[0, -1]).item())
        if nxt == tokenizer.eos_token_id:
            break
        generated.append(nxt)
        inp = mx.array([[nxt]])
    return generated


def teacher_forced_stats(model, full_ids: list[int], start: int) -> dict:
    """Forward the full sequence once; return per-position stats for positions
    predicting tokens at indices [start, len(full_ids)) — i.e., logits at
    positions start-1 .. len-2."""
    logits = model(mx.array(full_ids)[None])[0]  # (T, V)
    sel = logits[start - 1 : len(full_ids) - 1].astype(mx.float32)
    top2 = mx.topk(sel, 2, axis=-1)  # values sorted ascending in MLX topk
    argmax = mx.argmax(sel, axis=-1)
    logprobs = sel - mx.logsumexp(sel, axis=-1, keepdims=True)
    mx.eval(top2, argmax, logprobs)
    v = np.array(top2)
    margin = v[:, 1] - v[:, 0] if v[0, 1] >= v[0, 0] else v[:, 0] - v[:, 1]
    return {
        "margin": np.abs(margin),
        "argmax": np.array(argmax),
        # float16 storage: 48 trajectories × (128, ~152k vocab) would be ~4 GB
        # in float32; KL is computed in float32 at use time.
        "logprobs": np.array(logprobs).astype(np.float16),
    }


def auroc(scores: np.ndarray, labels: np.ndarray) -> float:
    """AUROC of `scores` (higher = predicted positive) for binary labels.
    Here: score = -margin (low margin should predict disagreement=1)."""
    pos, neg = scores[labels == 1], scores[labels == 0]
    if len(pos) == 0 or len(neg) == 0:
        return float("nan")
    order = np.argsort(np.concatenate([pos, neg]), kind="mergesort")
    ranks = np.empty(len(order)); ranks[order] = np.arange(1, len(order) + 1)
    # average ranks for ties
    allv = np.concatenate([pos, neg])
    sorted_v = allv[order]
    i = 0
    while i < len(sorted_v):
        j = i
        while j + 1 < len(sorted_v) and sorted_v[j + 1] == sorted_v[i]:
            j += 1
        if j > i:
            ranks[order[i : j + 1]] = ranks[order[i : j + 1]].mean()
        i = j + 1
    r_pos = ranks[: len(pos)].sum()
    return float((r_pos - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg)))


def escalation_curve(margins: np.ndarray, agree: np.ndarray, points: int = 200) -> list[dict]:
    """For threshold τ over margins: escalate tokens with margin < τ (assume the
    escalated decision becomes correct). Report escalated fraction vs residual
    disagreement (disagreements with margin ≥ τ)."""
    qs = np.quantile(margins, np.linspace(0, 1, points))
    out, n = [], len(margins)
    for tau in qs:
        esc = margins < tau
        residual = np.sum((~esc) & (agree == 0)) / n
        out.append({"tau": float(tau), "escalated_frac": float(esc.mean()),
                    "residual_disagree": float(residual)})
    return out


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("--bits", default="2,3,4,8")
    ap.add_argument("--per-domain", type=int, default=8)
    ap.add_argument("--group-size", type=int, default=64)
    args = ap.parse_args()
    bits_list = [int(b) for b in args.bits.split(",")]

    prompts_file = REPO_ROOT / "benchmarks" / "datasets" / "eval_prompts.json"
    domains = json.loads(prompts_file.read_text())["domains"]

    print(f"loading reference {args.model} …", flush=True)
    model, tokenizer = load(args.model)

    # -------- pass 1: reference greedy trajectories + reference stats
    trajectories = []  # {domain, prompt_ids, full_ids, start}
    t0 = time.time()
    for domain, plist in domains.items():
        for prompt in plist[: args.per_domain]:
            msgs = [{"role": "user", "content": prompt}]
            ids = tokenizer.apply_chat_template(msgs, add_generation_prompt=True)
            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)
            if len(gen) < 8:
                continue
            trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})
        print(f"  generated {domain}", flush=True)
    print(f"reference generation done in {time.time() - t0:.0f}s "
          f"({len(trajectories)} trajectories)", flush=True)

    ref_stats = [teacher_forced_stats(model, t["full_ids"], t["start"]) for t in trajectories]

    # -------- pass 2: quantized variants, teacher-forced on the same ids
    per_bits: dict[int, dict] = {}
    for bits in bits_list:
        print(f"quantizing to {bits}-bit (group {args.group_size}) …", flush=True)
        del model
        gc.collect(); mx.clear_cache()
        model, _ = load(args.model)
        nn.quantize(model, group_size=args.group_size, bits=bits,
                    class_predicate=lambda p, m: isinstance(m, nn.Linear)
                    and m.weight.shape[-1] % args.group_size == 0)
        rows = []
        for t, ref in zip(trajectories, ref_stats):
            qs = teacher_forced_stats(model, t["full_ids"], t["start"])
            ref_next = np.array(t["full_ids"][t["start"]:])  # actual (=ref argmax) tokens
            agree = (qs["argmax"] == ref_next).astype(np.int8)
            # KL(ref||q) per position
            ref_lp = ref["logprobs"].astype(np.float32)
            kl = np.sum(np.exp(ref_lp) * (ref_lp - qs["logprobs"].astype(np.float32)), axis=-1)
            rows.append({"domain": t["domain"], "margin": qs["margin"],
                         "agree": agree, "kl": kl})
        margins = np.concatenate([r["margin"] for r in rows])
        agrees = np.concatenate([r["agree"] for r in rows])
        kls = np.concatenate([r["kl"] for r in rows])
        disagree = 1 - agrees
        stats = {
            "bits": bits,
            "n_positions": int(len(margins)),
            "agreement_rate": float(agrees.mean()),
            "mean_kl_ref_q": float(np.mean(kls)),
            "auroc_margin_predicts_disagreement": auroc(-margins, disagree),
            "median_margin_agree": float(np.median(margins[agrees == 1])),
            "median_margin_disagree": float(np.median(margins[agrees == 0])) if (agrees == 0).any() else None,
            "escalation_curve": escalation_curve(margins, agrees),
            "per_domain": {
                d: {
                    "agreement_rate": float(np.concatenate([r["agree"] for r in rows if r["domain"] == d]).mean()),
                    "auroc": auroc(
                        -np.concatenate([r["margin"] for r in rows if r["domain"] == d]),
                        1 - np.concatenate([r["agree"] for r in rows if r["domain"] == d]),
                    ),
                }
                for d in domains
            },
        }
        # operating point: escalation fraction to reach 99% agreement
        for pt in stats["escalation_curve"]:
            if pt["residual_disagree"] <= 0.01:
                stats["escalation_frac_for_99pct"] = pt["escalated_frac"]
                break
        per_bits[bits] = stats
        print(f"  {bits}-bit: agree={stats['agreement_rate']:.4f} "
              f"AUROC={stats['auroc_margin_predicts_disagreement']:.3f} "
              f"esc@99%={stats.get('escalation_frac_for_99pct', 'n/a')}", flush=True)

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "expG_decision_stability" / ts
    out_dir.mkdir(parents=True)
    payload = {
        "experiment": "expG_decision_stability",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": collect_manifest(),
        "config": vars(args),
        "n_trajectories": len(trajectories),
        "results_by_bits": {str(k): v for k, v in per_bits.items()},
    }
    (out_dir / "results.json").write_text(json.dumps(payload, indent=2))
    print(f"\nwrote {out_dir / 'results.json'}")


if __name__ == "__main__":
    main()