#!/usr/bin/env python3
"""SPIRA evaluation harness v2 — proofs + hierarchy ablation + replay.

Produces artifacts/metrics.json with:
  0. vocab coverage : readout-vocabulary coverage curve (Zipf knee)
  1. compression     : EPT vs byte-level BPE (matched vocab) vs gzip-9
  2. prediction      : unigram / bigram / SNN (fast-only) / hierarchical
                       cortex / AM / fused, on held-out test + a slice of the
                       TRAINED region (memorization reference)
  3. one-shot        : single-pass teaching of an unseen book -> exact recall
  3b. replay         : hippocampal replay -> cortical consolidation curve,
                       AM integrity after replay, and interference on test
  4. correction      : recall vs number of corrupted context lanes
  5. generation      : samples, speed, synaptic-op accounting

Sections are resumable: completed sections are skipped when metrics.json
already contains them, so the harness can be re-invoked across process
boundaries (each call loads the checkpoint fresh; the replay experiment is
self-contained in one section).
"""
from __future__ import annotations

import argparse
import json
import os
import time

import numpy as np

from .baselines import (BigramBaseline, UnigramBaseline, bpe_encode,
                        dict_bits, gzip_bits_per_byte, stream_entropy_bits)
from .ept import EPT
from .hashing import fnv1a64
from .model import SPIRA


def distinct_n(text: str, n: int) -> float:
    toks = text.split()
    if len(toks) < n:
        return 0.0
    grams = [" ".join(toks[i : i + n]) for i in range(len(toks) - n + 1)]
    return len(set(grams)) / len(grams)


def dump(metrics: dict, path: str) -> None:
    with open(path, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--train", required=True)
    ap.add_argument("--test", required=True)
    ap.add_argument("--oneshot", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--eval-positions", type=int, default=60_000)
    ap.add_argument("--replay-tokens", type=int, default=100_000)
    ap.add_argument("--replay-stages", type=int, default=4)
    args = ap.parse_args()

    train = open(args.train, "rb").read()
    test = open(args.test, "rb").read()
    oneshot = open(args.oneshot, "rb").read()
    model = SPIRA.load(args.model)
    metrics: dict = {}
    if os.path.exists(args.out):
        try:
            metrics = json.load(open(args.out))
            print(f"(resuming: {len(metrics)} section(s) already present)")
        except Exception:
            metrics = {}
    # restore the train id stream for baselines / corruption / coverage
    t0 = time.perf_counter()
    train_patches = model.ept.segment(train)
    train_ids = np.fromiter((fnv1a64(p) for p in train_patches), dtype=np.uint64,
                            count=len(train_patches))
    model._ids_train = train_ids
    print(f"prep: segmented train ({len(train_ids)} patches) in {time.perf_counter()-t0:.1f}s")
    if model.core.hier:
        cal = model.calibrate_fusion(train[:60_000])
        print(f"prep: fusion calibration sigma_fast={cal['sigma_fast']:.4f} "
              f"sigma_slow={cal['sigma_slow']:.4f} scale={cal['scale']:.6f} "
              f"beta={model.core.beta:.4f}")
        if "fusion_calibration" not in metrics:
            metrics["fusion_calibration"] = {**cal, "beta": float(model.core.beta)}
            dump(metrics, args.out)
    if "model_meta" not in metrics:
        metrics.update({"model_meta": model.meta, "config": vars(model.cfg)})
        dump(metrics, args.out)

    # ============================================ 0. vocab coverage curve =====
    if "vocab_coverage_curve" not in metrics:
        print("== [0/7] readout-vocabulary coverage curve ==")
        uniq, cnt = np.unique(train_ids, return_counts=True)
        order = np.lexsort((uniq, -cnt))          # count desc, id asc (build order)
        sorted_ids = uniq[order]
        test_ids_local = np.fromiter(
            (fnv1a64(p) for p in model.ept.segment(test)), dtype=np.uint64,
            count=len(model.ept.segment(test)))
        one_ids_local = np.fromiter(
            (fnv1a64(p) for p in model.ept.segment(oneshot[:60_000])), dtype=np.uint64,
            count=len(model.ept.segment(oneshot[:60_000])))
        curve = []
        for V in (4096, 6144, 8192, 12288, 16384, 24576, 32768):
            top = sorted_ids[:V]
            cov_t = float(np.isin(test_ids_local, top).mean())
            cov_o = float(np.isin(one_ids_local, top).mean())
            curve.append({"vocab": int(V), "cov_test": round(cov_t, 4),
                          "cov_oneshot": round(cov_o, 4)})
            print(f"    V={V:>6}: held-out test coverage={cov_t:.4f} "
                  f"oneshot coverage={cov_o:.4f}")
        metrics["vocab_coverage_curve"] = curve
        metrics["vocab_coverage"] = {
            "distinct_train_patches": int(len(uniq)),
            "note": "fraction of held-out target patches inside the top-V train vocabulary"}
        dump(metrics, args.out)

    # ================================================== 1. compression =====
    if "compression" not in metrics:
        print("== [1/7] compression: EPT vs BPE vs gzip ==")
        comp = {}
        gz_bpb, gz_s = gzip_bits_per_byte(test)
        comp["gzip"] = {"bits_per_byte": gz_bpb, "seconds": gz_s}
        for vs in (4096, 16384):
            r = bpe_encode(train[:4_000_000], test, vs)
            h0 = stream_entropy_bits(r["ids"], order=0)
            h1 = stream_entropy_bits(r["ids"], order=1)
            db = dict_bits(r["vocab_size"], r["vocab_bytes"])
            comp[f"bpe_{vs}"] = {
                "vocab_size": r["vocab_size"],
                "n_tokens": int(len(r["ids"])),
                "avg_token_bytes": len(test) / len(r["ids"]),
                "stream_bits_unigram": h0,
                "stream_bits_bigram": h1,
                "dict_bits": db,
                "bits_per_byte_unigram": (h0 * len(r["ids"]) + db) / len(test),
                "bits_per_byte_bigram": (h1 * len(r["ids"]) + db) / len(test),
                "train_s": r["train_s"],
            }
        t0 = time.perf_counter()
        ept_patches = model.ept.segment(test)
        ept_seg_s = time.perf_counter() - t0
        ept_ids = np.fromiter((fnv1a64(p) for p in ept_patches), dtype=np.uint64,
                              count=len(ept_patches))
        h0 = stream_entropy_bits(ept_ids, order=0)
        h1 = stream_entropy_bits(ept_ids, order=1)
        reg_ids = model.registry.ids
        in_train = model.registry.has(ept_ids)
        known_ids = ept_ids[in_train]
        pos = np.searchsorted(reg_ids, np.unique(known_ids))
        lens = np.diff(model.registry.offsets)[pos].astype(np.float64)
        ept_dict_bits = float(lens.sum() * 8 + 32 * len(lens))
        novelty = float((~in_train).mean())
        comp["ept"] = {
            "n_tokens": int(len(ept_ids)),
            "avg_token_bytes": float(len(test) / len(ept_ids)),
            "distinct_in_test": int(len(np.unique(ept_ids))),
            "stream_bits_unigram": h0,
            "stream_bits_bigram": h1,
            "dict_bits_test_local": ept_dict_bits,
            "bits_per_byte_unigram": (h0 * len(ept_ids) + ept_dict_bits) / len(test),
            "bits_per_byte_bigram": (h1 * len(ept_ids) + ept_dict_bits) / len(test),
            "novelty_rate": novelty,
            "bytes_in_novel_patches": float(sum(len(p) for p, k in zip(ept_patches, in_train) if not k)),
            "segment_s": ept_seg_s,
            "MB_per_s": len(test) / 1e6 / max(ept_seg_s, 1e-9),
        }
        metrics["compression"] = comp
        dump(metrics, args.out)
        for k, v in comp.items():
            print(f"    {k}: {json.dumps({kk: round(vv, 4) if isinstance(vv, float) else vv for kk, vv in v.items() if 'bits_per_byte' in kk or 'novelty' in kk or 'avg_token' in kk})}")

    # ================================================== 2. prediction =====
    if "prediction" not in metrics:
        print("== [2/7] next-patch prediction: held-out test + trained region ==")
        ev = model.eval_stream(test, max_positions=args.eval_positions)
        print(f"    TEST: fused={ev['fused_acc']:.4f} (hier={ev['hier_acc']:.4f}, "
              f"fast={ev['snn_acc']:.4f}, am={ev['am_acc']:.4f}, coverage={ev['vocab_coverage']:.3f})")
        # memorization reference: the FINAL training region (tokens ~550k-600k,
        # where the model's operating point was warm-matched; SNN region 0-600k,
        # AM region 0-1.2M). Earlier regions are legitimately drifted under
        # continual plasticity — that is the forgetting curve, not a bug.
        ev_tr = model.eval_stream(train[2_600_000:2_800_000], max_positions=8_000)
        print(f"    TRAINED-REGION: fused={ev_tr['fused_acc']:.4f} "
              f"(hier={ev_tr['hier_acc']:.4f}, am={ev_tr['am_acc']:.4f})")
        # baselines over the same test stream, in-vocab positions
        tp = model.ept.segment(test)
        tids = np.fromiter((fnv1a64(p) for p in tp), dtype=np.uint64, count=len(tp))
        tcols = model.cols_of(tids)[: args.eval_positions + 1]
        V = model.cfg.vocab_size
        inv_mask = tcols != V
        ok_pos = inv_mask[:-1] & inv_mask[1:]
        prev_in = tcols[:-1][ok_pos].astype(np.int64)
        cur_in = tcols[1:][ok_pos].astype(np.int64)
        train_cols = model.cols_of(train_ids[:600_000])
        uni = UnigramBaseline(train_cols, V)
        bi = BigramBaseline(train_cols, V)
        pred = {
            "spira_test": ev,
            "spira_trained_region": ev_tr,
            "unigram_acc": uni.accuracy(cur_in),
            "bigram_acc": bi.accuracy(prev_in, cur_in),
            "n_baseline_positions": int(len(cur_in)),
            "note": "baselines evaluated on in-vocab target positions of the same stream; "
                    "compare with spira_test.*_invocab",
        }
        metrics["prediction"] = pred
        dump(metrics, args.out)
        print(f"    baselines (in-vocab): unigram={pred['unigram_acc']:.4f} "
              f"bigram={pred['bigram_acc']:.4f} vs spira fast_invocab={ev['snn_acc_invocab']:.4f} "
              f"hier_invocab={ev['hier_acc_invocab']:.4f} fused_invocab={ev['fused_acc_invocab']:.4f}")

    # ==================================================== 3. one-shot =====
    if "one_shot" not in metrics:
        print("== [3/7] one-shot learning: teach the unseen book ONCE ==")
        book = oneshot[:60_000]
        generalization = model.eval_stream(book, max_positions=8_000)
        t0 = time.perf_counter()
        taught = model.teach_bytes(book)
        teach_s = time.perf_counter() - t0
        recall = model.eval_stream(book, max_positions=8_000)
        oneshot_metrics = {
            "book_bytes": len(book),
            "before_teaching": generalization,
            "teaching": {**taught, "seconds": teach_s},
            "after_one_pass": recall,
        }
        metrics["one_shot"] = oneshot_metrics
        dump(metrics, args.out)
        print(f"    before: fused={generalization['fused_acc']:.4f} am={generalization['am_acc']:.4f}")
        print(f"    teach: {taught['tokens']} tokens, {taught['new_patches']} new patches in {teach_s:.1f}s")
        print(f"    after ONE pass: am recall={recall['am_acc']:.4f} fused={recall['fused_acc']:.4f}")

    # ===================================== 3b. replay -> consolidation =====
    if "replay" not in metrics:
        print("== [4/7] hippocampal replay -> cortical consolidation ==")
        book = oneshot[:60_000]
        if not model.episodes:
            taught = model.teach_bytes(book)
            print(f"    (re-taught for this section: {taught['tokens']} tokens)")
        before = model.eval_stream(book, use_am=False, max_positions=8_000)
        print(f"    cortex-only BEFORE replay: hier={before['hier_acc']:.4f} "
              f"(in-vocab {before['hier_acc_invocab']:.4f}, fast={before['snn_acc']:.4f})")
        curve = []
        stage = max(1, args.replay_tokens // args.replay_stages)
        for i in range(args.replay_stages):
            r = model.replay(stage, log=lambda s: print(f"      {s}"))
            probe = model.eval_stream(book, use_am=False, max_positions=2_000)
            curve.append({
                "replayed_tokens": (i + 1) * stage,
                "hier_acc": probe["hier_acc"],
                "hier_acc_invocab": probe["hier_acc_invocab"],
                "snn_acc": probe["snn_acc"],
                "replay_acc": r.get("accuracy"),
                "beta": r.get("beta"),
            })
            print(f"    after {(i+1)*stage} replayed tokens: hier={probe['hier_acc']:.4f} "
                  f"(in-vocab {probe['hier_acc_invocab']:.4f}) beta={r.get('beta'):.3f}")
        after = model.eval_stream(book, use_am=False, max_positions=8_000)
        am_after = model.eval_stream(book, max_positions=8_000)
        interference_core = model.eval_stream(test, use_am=False, max_positions=8_000)
        interference_fused = model.eval_stream(test, max_positions=8_000)
        metrics["replay"] = {
            "replay_tokens": int(args.replay_tokens),
            "stage_tokens": int(stage),
            "before": before,
            "curve": curve,
            "after": after,
            "am_recall_after_replay": am_after,
            "interference_test_core_only": interference_core,
            "interference_test_fused": interference_fused,
            "pre_replay_test": metrics.get("prediction", {}).get("spira_test"),
            "note": "the shipped checkpoint is PRE-replay; this section runs in-memory",
        }
        dump(metrics, args.out)
        print(f"    cortex-only AFTER replay: hier={after['hier_acc']:.4f} "
              f"(in-vocab {after['hier_acc_invocab']:.4f})")
        print(f"    AM recall after replay: {am_after['am_acc']:.4f} (must be intact)")
        print(f"    interference on test: core {interference_core['hier_acc']:.4f} / "
              f"fused {interference_fused['fused_acc']:.4f}")

    # ================================================== 4. correction =====
    if "correction" not in metrics:
        print("== [5/7] correction: recall vs corrupted context lanes ==")
        corr = {}
        for k in (0, 1, 2, 3):
            r = model.eval_stream(oneshot[:60_000], corrupt_lanes=k, max_positions=8_000)
            corr[f"corrupt_{k}_of_4"] = {kk: r[kk] for kk in ("am_acc", "am_acc_invocab", "fused_acc")}
            print(f"    corrupt {k}/4 lanes: am={r['am_acc']:.4f} fused={r['fused_acc']:.4f}")
        metrics["correction"] = corr
        dump(metrics, args.out)

    # ================================================ 5. generation ======
    if "generation" not in metrics:
        print("== [6/7] generation ==")
        gens = {}
        prompts = {
            "english": b"The old man looked at the ",
            "code": b"def process(data):\n    ",
            "french": b"Il etait une fois un ",
        }
        for name, p in prompts.items():
            t0 = time.perf_counter()
            g = model.generate(p, n_tokens=150, temperature=0.7, top_k=8, seed=42)
            dt = time.perf_counter() - t0
            text = g.decode("utf-8", "replace")
            gens[name] = {
                "prompt": p.decode("utf-8", "replace"),
                "output": text,
                "chars_per_s": len(text) / max(dt, 1e-9),
                "distinct1": distinct_n(text, 1),
                "distinct2": distinct_n(text, 2),
            }
            print(f"    [{name}] {len(text)} chars @ {len(text)/max(dt,1e-9):.0f} chars/s")
            print(f"      {text[:180]!r}")
        metrics["generation"] = gens
        dump(metrics, args.out)

    # ================================================ 6. efficiency ======
    if "efficiency" not in metrics:
        print("== [7/7] efficiency accounting ==")
        d_in, k, n = model.cfg.d_in, model.cfg.k_in, model.cfg.n_res
        V = model.cfg.vocab_size
        n_g = model.cfg.n_g
        win = model.cfg.win
        slow_amortized = (n_g * (V + 1) + n * n_g + n_g * n_g) / win
        macs = {
            "input_gather": k * n,
            "recurrent": n * n,
            "readout_fast": n * V,
            "readout_slow_amortized": slow_amortized,
            "total_per_token": k * n + n * n + n * V + slow_amortized,
        }
        metrics["efficiency"] = {
            "synaptic_macs_per_token": macs,
            "vs_transformer_reference": {
                "gpt2_small_params": 124e6,
                "gpt2_train_flops_per_token": 6 * 124e6,
                "spira_vs_gpt2_ops_ratio": macs["total_per_token"] / (6 * 124e6),
            },
        }
        dump(metrics, args.out)
        print(f"    {macs['total_per_token']/1e6:.2f}M MACs/token "
              f"(fast readout {macs['readout_fast']/1e6:.2f}M, slow amortized "
              f"{slow_amortized/1e6:.2f}M) -> "
              f"{(6*124e6)/macs['total_per_token']:.0f}x fewer ops/token than GPT-2-small training")

    print(f"metrics -> {args.out}")


if __name__ == "__main__":
    main()