#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/micro/expD_progressive_reconstruction/benchmark.py
#  Purpose   : Residual-ladder progressive weight reconstruction — decision and
#              hidden-state convergence vs cumulative bits (candidate C1 math)
#  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 D — progressive weight reconstruction (charter §9.D).

Builds base+residual affine-quantized ladders (3/3+3/3+3+3 and 4/4+4 bits),
teacher-forces each cumulative stage over reference greedy trajectories, and
measures decision convergence, two-tier margin-gated policies, and
hidden-state error at several depths.

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

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 (  # noqa: E402
    auroc, escalation_curve, greedy_generate, kl_ref_vs, teacher_forced_stats,
)

GROUP = 64


def quantizable(m) -> bool:
    return isinstance(m, nn.Linear) and m.weight.shape[-1] % GROUP == 0


def residual_ladder_weights(model, ladder: list[int]) -> list[dict[str, mx.array]]:
    """For each quantizable Linear, build cumulative dequantized weights for
    each stage of `ladder` (bits per stage). Returns a list (one per stage) of
    {param_path: bf16 weight} replacements. Memory: one bf16 copy per stage
    per layer is materialized lazily at apply time; here we keep the per-stage
    cumulative tensors (float32 accumulation, cast to bf16)."""
    stages = [dict() for _ in ladder]
    for path, module in model.named_modules():
        if not quantizable(module):
            continue
        w = module.weight.astype(mx.float32)
        acc = mx.zeros_like(w)
        err = w
        for k, bits in enumerate(ladder):
            qw, scales, biases = mx.quantize(err, group_size=GROUP, bits=bits)
            deq = mx.dequantize(qw, scales, biases, group_size=GROUP, bits=bits)
            acc = acc + deq
            err = w - acc
            stages[k][path] = acc.astype(mx.bfloat16)
            mx.eval(stages[k][path])
    return stages


def apply_weights(model, replacement: dict[str, mx.array]) -> dict[str, mx.array]:
    """Swap Linear weights in place; returns the originals for restoration."""
    originals = {}
    for path, module in model.named_modules():
        if path in replacement:
            originals[path] = module.weight
            module.weight = replacement[path]
    return originals


def hidden_state_errors(model, ref_hidden: dict, full_ids: list[int], start: int,
                        depths: list[int]) -> dict[int, float]:
    """Relative L2 error of hidden states vs reference at given layer indices."""
    h = capture_hidden(model, full_ids, start, depths)
    out = {}
    for d in depths:
        r, q = ref_hidden[d], h[d]
        out[d] = float(np.linalg.norm(q - r) / (np.linalg.norm(r) + 1e-9))
    return out


def capture_hidden(model, full_ids: list[int], start: int, depths: list[int]) -> dict:
    """Hidden states (post-layer) at selected depths for predicted positions.
    Replicates the inner transformer loop manually (instance-level __call__
    monkey-patching does not intercept Python's type-level dunder dispatch)."""
    from mlx_lm.models.base import create_attention_mask

    inner = model.model
    h = inner.embed_tokens(mx.array(full_ids)[None])
    mask = create_attention_mask(h, None)
    result = {}
    for i, layer in enumerate(inner.layers):
        h = layer(h, mask, cache=None)
        if i in depths:
            t = h[0, start - 1 : len(full_ids) - 1].astype(mx.float32)
            mx.eval(t)
            result[i] = np.array(t)
    return result


def two_tier_policy(lo: dict, hi: dict, ref_next: np.ndarray, taus: list[float]) -> list[dict]:
    """Policy: take lo's decision when its margin >= tau, else hi's decision."""
    out = []
    for tau in taus:
        esc = lo["margin"] < tau
        decision = np.where(esc, hi["argmax"], lo["argmax"])
        out.append({
            "tau": tau,
            "escalated_frac": float(esc.mean()),
            "policy_agreement": float((decision == ref_next).mean()),
        })
    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("--per-domain", type=int, default=8)
    ap.add_argument("--hidden-trajectories", type=int, default=8)
    args = ap.parse_args()

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

    print(f"loading reference {args.model} …", flush=True)
    model, tokenizer = load(args.model)
    n_layers = len(model.model.layers)
    depths = [max(0, round(n_layers * f) - 1) for f in (0.25, 0.5, 0.75, 1.0)]

    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)
            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)} reference trajectories in {time.time()-t0:.0f}s", flush=True)

    ref_stats = [teacher_forced_stats(model, t["full_ids"], t["start"]) for t in trajectories]
    hidden_subset = trajectories[:: max(1, len(trajectories) // args.hidden_trajectories)][: args.hidden_trajectories]
    ref_hidden = [capture_hidden(model, t["full_ids"], t["start"], depths) for t in hidden_subset]

    ladders = {"A_base3": [3, 3, 3], "B_base4": [4, 4]}
    results: dict[str, list] = {}
    overhead_bits = 32 / GROUP * 2  # bf16 scales + biases per group per stage

    for name, ladder in ladders.items():
        print(f"building ladder {name} {ladder} …", flush=True)
        stages = residual_ladder_weights(model, ladder)
        stage_records = []
        for k, replacement in enumerate(stages):
            originals = apply_weights(model, replacement)
            rows_margin, rows_agree, rows_kl, rows_argmax, doms = [], [], [], [], []
            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"]:])
                rows_margin.append(qs["margin"])
                rows_argmax.append(qs["argmax"])
                rows_agree.append((qs["argmax"] == ref_next).astype(np.int8))
                rows_kl.append(kl_ref_vs(qs["logprobs"], ref["logprobs"]))
                doms.append(t["domain"])
            hid = [hidden_state_errors(model, rh, t["full_ids"], t["start"], depths)
                   for rh, t in zip(ref_hidden, hidden_subset)]
            apply_weights(model, originals)

            margins = np.concatenate(rows_margin)
            agrees = np.concatenate(rows_agree)
            cum_bits = sum(ladder[: k + 1]) + overhead_bits * (k + 1)
            rec = {
                "stage": k,
                "ladder_bits": ladder[: k + 1],
                "cumulative_bits_per_param": round(cum_bits, 2),
                "agreement_rate": float(agrees.mean()),
                "mean_kl": float(np.mean(np.concatenate(rows_kl))),
                "auroc": auroc(-margins, 1 - agrees),
                "escalation_curve": escalation_curve(margins, agrees),
                "hidden_rel_err_by_depth": {
                    str(d): float(np.mean([h[d] for h in hid])) for d in depths
                },
                "_margins": margins, "_argmax": np.concatenate(rows_argmax),
            }
            for pt in rec["escalation_curve"]:
                if pt["residual_disagree"] <= 0.01:
                    rec["escalation_frac_for_99pct"] = pt["escalated_frac"]
                    break
            stage_records.append(rec)
            print(f"  stage {k} ({rec['cumulative_bits_per_param']} bits): "
                  f"agree={rec['agreement_rate']:.4f} KL={rec['mean_kl']:.4f} "
                  f"auroc={rec['auroc']:.3f}", flush=True)
        results[name] = stage_records

    # two-tier margin-gated policies between consecutive stages
    ref_next_all = np.concatenate([np.array(t["full_ids"][t["start"]:]) for t in trajectories])
    taus = [0.25, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
    policies = {}
    for name, recs in results.items():
        for k in range(len(recs) - 1):
            lo = {"margin": recs[k]["_margins"], "argmax": recs[k]["_argmax"]}
            hi = {"argmax": recs[k + 1]["_argmax"]}
            policies[f"{name}_stage{k}_to_{k+1}"] = two_tier_policy(lo, hi, ref_next_all, taus)
    for recs in results.values():
        for r in recs:
            del r["_margins"], r["_argmax"]

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


if __name__ == "__main__":
    main()