#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/micro/expB_token_stability/benchmark.py
#  Purpose   : Temporal stability of per-token important FFN block sets
#              (consumes expA's block-energy trace)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Experiment B — stability across consecutive tokens (charter §9.B).

Usage:
    .venv/bin/python benchmark.py [--trace <path/to/block_energy_trace.npz>]
        [--target 0.95]
"""

from __future__ import annotations

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

import numpy as np

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

DELTAS = [1, 2, 4, 8, 16, 32]
WINDOWS = [8, 32, 128]


def top_sets(energy: np.ndarray, target: float) -> list[np.ndarray]:
    """energy (T, B) → per-position boolean mask of the smallest block set
    covering `target` of the energy."""
    T, B = energy.shape
    order = np.argsort(energy, axis=1)[:, ::-1]
    srt = np.take_along_axis(energy, order, axis=1)
    csum = np.cumsum(srt, axis=1)
    total = csum[:, -1:] + 1e-12
    kneed = np.argmax(csum / total >= target, axis=1) + 1
    masks = np.zeros((T, B), dtype=bool)
    for t in range(T):
        masks[t, order[t, : kneed[t]]] = True
    return masks


def jaccard(a: np.ndarray, b: np.ndarray) -> float:
    inter = np.logical_and(a, b).sum()
    union = np.logical_or(a, b).sum()
    return float(inter / union) if union else 1.0


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--trace", default=None)
    ap.add_argument("--target", type=float, default=0.95)
    ap.add_argument("--agg-factor", type=int, default=4,
                    help="aggregate trace blocks by this factor (16-neuron trace -> 64-neuron sets)")
    args = ap.parse_args()
    trace_path = args.trace or sorted(
        glob.glob(str(REPO_ROOT / "results/expA_weight_concentration/*/block_energy_trace.npz")))[-1]
    z = np.load(trace_path, allow_pickle=False)
    blocks = z["blocks"].astype(np.float32)  # (P, L, B)
    if args.agg_factor > 1:
        P0, L0, B0 = blocks.shape
        blocks = blocks.reshape(P0, L0, B0 // args.agg_factor, args.agg_factor).sum(axis=-1)
    domains = z["domains"]
    traj_id = z["traj_id"]
    P, L, B = blocks.shape
    print(f"trace {trace_path}: {blocks.shape}", flush=True)

    jac = {d: [] for d in DELTAS}
    jac_random = []
    union_frac = {w: [] for w in WINDOWS}
    per_layer_j1 = [[] for _ in range(L)]
    traj_masks_union = {}  # (traj, layer) -> aggregate mask, for expC-lite
    traj_domain = {}

    rng = np.random.default_rng(7)
    for ti in np.unique(traj_id):
        sel = traj_id == ti
        traj_domain[int(ti)] = str(domains[sel][0])
        for li in range(L):
            masks = top_sets(blocks[sel, li, :], args.target)
            T = masks.shape[0]
            for d in DELTAS:
                if T > d:
                    vals = [jaccard(masks[t], masks[t + d]) for t in range(T - d)]
                    jac[d].extend(vals)
                    if d == 1:
                        per_layer_j1[li].extend(vals)
            # random-set null at matched sizes (δ=1 pairs)
            sizes = masks.sum(axis=1)
            for t in range(min(T - 1, 8)):
                a = np.zeros(B, bool); a[rng.choice(B, sizes[t], replace=False)] = True
                b = np.zeros(B, bool); b[rng.choice(B, sizes[t + 1], replace=False)] = True
                jac_random.append(jaccard(a, b))
            for w in WINDOWS:
                for s in range(0, T - w + 1, w):
                    union_frac[w].append(float(masks[s:s + w].any(axis=0).mean()))
            traj_masks_union[(int(ti), li)] = masks.any(axis=0)

    # expC-lite: within- vs across-domain overlap of per-trajectory unions
    tids = sorted(traj_domain)
    within, across = [], []
    for i in range(len(tids)):
        for j in range(i + 1, len(tids)):
            v = np.mean([jaccard(traj_masks_union[(tids[i], li)], traj_masks_union[(tids[j], li)])
                         for li in range(0, L, 4)])
            (within if traj_domain[tids[i]] == traj_domain[tids[j]] else across).append(v)

    result = {
        "jaccard_by_delta": {str(d): {"mean": float(np.mean(v)), "std": float(np.std(v)), "n": len(v)}
                             for d, v in jac.items()},
        "jaccard_random_null": {"mean": float(np.mean(jac_random)), "std": float(np.std(jac_random))},
        "union_working_set_frac_by_window": {str(w): {"mean": float(np.mean(v)), "std": float(np.std(v))}
                                             for w, v in union_frac.items()},
        "per_layer_jaccard1_mean": [float(np.mean(v)) for v in per_layer_j1],
        "expC_lite_domain_locality": {
            "within_domain_union_jaccard": float(np.mean(within)),
            "across_domain_union_jaccard": float(np.mean(across)),
        },
        "mean_topset_frac": float(np.mean([m.mean() for m in
                                           [top_sets(blocks[traj_id == t, li, :], args.target)
                                            for t in np.unique(traj_id)[:4] for li in (0, L // 2, L - 1)]])),
    }
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "expB_token_stability" / ts
    out_dir.mkdir(parents=True)
    (out_dir / "results.json").write_text(json.dumps({
        "experiment": "expB_token_stability",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": collect_manifest(),
        "config": {"trace": str(trace_path), "target": args.target},
        **result,
    }, indent=2))
    print(json.dumps(result["jaccard_by_delta"], indent=1))
    print("random null:", result["jaccard_random_null"])
    print("union by window:", json.dumps(result["union_working_set_frac_by_window"]))
    print("domain locality:", result["expC_lite_domain_locality"])
    print(f"\nwrote {out_dir / 'results.json'}")


if __name__ == "__main__":
    main()