"""Rigorous benchmark harness for ENSEMBLE.

Measures, across dimensionalities D, on three corpora (facts / qa / prose):

  - build time + throughput (tokens/sec)
  - in-memory footprint (MB)
  - next-token accuracy (lm mode)
  - QA exact-recall accuracy (qa mode, on holdout)
  - query latency
  - .exp compression ratio

Run:
    python bench/run_benchmark.py --Ds 2000 10000 --out bench_results.json
    python bench/run_benchmark.py --Ds 100000 --out bench_results_1b.json

Each result is reproducible (fixed seeds). Output is a JSON the reporter turns
into the RESULTS.md table.
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path

import numpy as np

# allow `python bench/run_benchmark.py` from repo root
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from ensemble import Expert
from ensemble.io import expert_size_bytes
from palimseste.tokenizer import BOS, EOS, PAD

import gen_datasets  # noqa: E402  (same dir)


# ----------------------------------------------------------------- result
@dataclass
class BenchResult:
    name: str
    D: int
    # build
    build_seconds: float
    n_tokens: int
    throughput_tok_per_s: float
    # footprint
    ram_mb: float
    # quality
    next_token_accuracy: float
    # qa
    qa_exact_recall: float = 0.0
    qa_n_eval: int = 0
    # query
    query_latency_ms: float = 0.0
    # storage
    exp_size_bytes: int = 0
    source_size_bytes: int = 0
    compression_ratio: float = 0.0


# ----------------------------------------------------------------- helpers
def _ram_mb(expert: Expert) -> float:
    """Sum of packed-bits storage for all traces (address + value)."""
    total = 0
    for t in expert.model.mem.traces:
        total += t.address.bits.nbytes + t.value.bits.nbytes
    return total / 1e6


def _next_token_accuracy(expert: Expert, text: str, max_eval: int = 400) -> float:
    """Fraction of correctly-predicted next tokens over the text.

    Uses an INCREMENTAL state builder so the whole pass is O(N*D), not
    O(N^2*D): the running sign-sum is updated by one column per token rather
    than rebuilt from scratch each step. This makes accuracy measurement
    feasible even at D=100 000.
    """
    from palimseste.lm import _IncrementalState
    tok = expert.model.tokenizer
    model = expert.model
    if tok is None:
        return 0.0
    ids = tok.encode(text, add_bos=True, add_eos=True)
    if len(ids) < 2:
        return 0.0

    # subsample eval points deterministically
    n_eval_points = min(max_eval, len(ids) - 1)
    step = max(1, (len(ids) - 1) // n_eval_points)

    builder = _IncrementalState(D=model.config.D,
                                window=model.config.context_window,
                                encoder=model.encoder)
    # precompute token bits matrix for vectorized scoring
    V = tok.vocab_size
    packed_len = len(model._self_hv.bits)
    tok_bits = np.empty((V, packed_len), dtype=np.uint8)
    for tid in range(V):
        tok_bits[tid] = tok.token_hv(tid).bits
    _POPCOUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint16)
    D = model.config.D
    from palimseste.hv import bind as _bind

    correct = 0
    total = 0
    for i, tok_id in enumerate(ids):
        s = builder.peek_or_init(tok)
        should_eval = (i > 0 and ((i - 1) % step == 0))
        if should_eval and s is not None:
            q = _bind(model._self_hv, s)
            retrieved = model.phi(model.mem, q)
            if retrieved is not None:
                xored = np.bitwise_xor(tok_bits, retrieved.bits[np.newaxis, :])
                pred = int(np.argmin(_POPCOUNT[xored].sum(axis=1)))
            else:
                pred = PAD
            if pred == ids[i]:
                correct += 1
            total += 1
        builder.push(tok_id, tok)

    return correct / total if total else 0.0


def _qa_exact_recall(expert: Expert, holdout: list[tuple[str, str]]) -> tuple[float, int]:
    """Exact-answer recall: does the expert return the exact answer prefix?

    We score leniently: the answer counts as recalled if its first
    ``min(len, 12)`` chars appear at the start of the response (the char-LM
    often trails off, so we reward the correct *start*).
    """
    n = 0
    hit = 0
    for q, expected in holdout:
        n += 1
        got = expert.answer(q, max_new_tokens=len(expected) + 6, temperature=0.0)
        k = min(len(expected), 12)
        if got.strip().lower().startswith(expected.strip().lower()[:k]):
            hit += 1
    return (hit / n if n else 0.0), n


def _measure_query_latency(expert: Expert, questions: list[str], k: int = 20) -> float:
    """Median query latency in ms over k samples."""
    times = []
    for q in questions[:k]:
        t0 = time.perf_counter()
        expert.answer(q, max_new_tokens=20, temperature=0.0)
        times.append((time.perf_counter() - t0) * 1000)
    return float(statistics.median(times)) if times else 0.0


# ----------------------------------------------------------------- benchmarks
def bench_lm(name: str, text: str, D: int, holdout_qa=None) -> BenchResult:
    print(f"  [{name}] D={D} building lm expert...", flush=True)
    t0 = time.perf_counter()
    expert = Expert.from_text(text, domain=name, D=D, seed=0)
    build_s = time.perf_counter() - t0
    n_tok = expert.n_traces
    throughput = n_tok / build_s if build_s > 0 else 0.0
    print(f"    built {n_tok} traces in {build_s:.1f}s ({throughput:.0f} tok/s)",
          flush=True)

    ram = _ram_mb(expert)
    print(f"    measuring next-token accuracy...", flush=True)
    nta = _next_token_accuracy(expert, text, max_eval=300 if D <= 20000 else 150)

    qa_recall, qa_n = 0.0, 0
    if holdout_qa:
        qa_recall, qa_n = _qa_exact_recall(expert, holdout_qa)

    print(f"    measuring latency...", flush=True)
    qprobes = [q for q, _ in (holdout_qa or [("what is x", "x")][:1])]
    lat = _measure_query_latency(expert, qprobes)

    import tempfile
    with tempfile.TemporaryDirectory() as td:
        res = expert.save(os.path.join(td, f"{name}.exp"))
        exp_bytes = res.expert_size_bytes
        src_bytes = res.source_size_bytes
        ratio = res.compression_ratio

    return BenchResult(
        name=name, D=D, build_seconds=build_s, n_tokens=n_tok,
        throughput_tok_per_s=throughput, ram_mb=ram,
        next_token_accuracy=nta, qa_exact_recall=qa_recall, qa_n_eval=qa_n,
        query_latency_ms=lat, exp_size_bytes=exp_bytes,
        source_size_bytes=src_bytes, compression_ratio=ratio,
    )


def bench_qa(name: str, train: list[tuple[str, str]], holdout: list[tuple[str, str]],
             D: int) -> BenchResult:
    print(f"  [{name}] D={D} building qa expert...", flush=True)
    t0 = time.perf_counter()
    expert = Expert.from_qa_pairs(train, domain=name, D=D, seed=0)
    build_s = time.perf_counter() - t0
    n_tok = expert.n_traces
    throughput = n_tok / build_s if build_s > 0 else 0.0
    print(f"    built {n_tok} traces in {build_s:.1f}s ({throughput:.0f} tok/s)",
          flush=True)

    ram = _ram_mb(expert)
    nta = 0.0  # qa mode: we measure recall, not lm accuracy

    # TWO recall numbers: memorization (train) and generalization (holdout).
    # Use the unique train pairs (dedup) for a fair memorization probe.
    train_unique = list(dict.fromkeys(train))
    # cap memorization probe for speed at large D
    cap = 20 if D <= 30000 else 10
    train_probe = train_unique[:cap]
    recall_train, n_train = _qa_exact_recall(expert, train_probe)
    recall_holdout, n_holdout = _qa_exact_recall(expert, holdout)
    print(f"    qa recall train={recall_train:.1%} (n={n_train})  "
          f"holdout={recall_holdout:.1%} (n={n_holdout})", flush=True)
    # we report the memorization recall as the headline QA number
    qa_recall = recall_train
    qa_n = n_train
    lat = _measure_query_latency(expert, [q for q, _ in holdout])

    import tempfile
    with tempfile.TemporaryDirectory() as td:
        res = expert.save(os.path.join(td, f"{name}.exp"))
        exp_bytes = res.expert_size_bytes
        src_bytes = res.source_size_bytes
        ratio = res.compression_ratio

    # stash the holdout recall in extra fields via a side note
    br = BenchResult(
        name=name, D=D, build_seconds=build_s, n_tokens=n_tok,
        throughput_tok_per_s=throughput, ram_mb=ram,
        next_token_accuracy=nta, qa_exact_recall=qa_recall, qa_n_eval=qa_n,
        query_latency_ms=lat, exp_size_bytes=exp_bytes,
        source_size_bytes=src_bytes, compression_ratio=ratio,
    )
    # carry the holdout number for the report
    br_extra = asdict(br)
    br_extra["qa_holdout_recall"] = recall_holdout
    br_extra["qa_holdout_n"] = n_holdout
    return br  # type: ignore[return-value]


# ----------------------------------------------------------------- driver
def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--Ds", type=int, nargs="+", default=[2000, 10000],
                    help="dimensionalities to benchmark")
    ap.add_argument("--out", default="bench_results.json")
    ap.add_argument("--corpora", nargs="+",
                    default=["facts", "qa", "prose"],
                    choices=["facts", "qa", "prose"])
    ap.add_argument("--data-dir", default="bench_data")
    args = ap.parse_args()

    # ensure datasets exist
    ddir = Path(args.data_dir)
    if not (ddir / "facts.txt").exists():
        print("generating datasets...", flush=True)
        gen_datasets.generate_all(ddir)

    facts_text = (ddir / "facts.txt").read_text(encoding="utf-8")
    prose_text = (ddir / "prose.txt").read_text(encoding="utf-8")
    with open(ddir / "qa_train.json", encoding="utf-8") as f:
        qa_train = [(d["question"], d["answer"]) for d in json.load(f)]
    with open(ddir / "qa_holdout.json", encoding="utf-8") as f:
        qa_holdout = [(d["question"], d["answer"]) for d in json.load(f)]

    results: list[BenchResult] = []
    for D in args.Ds:
        print(f"\n=== D = {D} ===", flush=True)
        if "facts" in args.corpora:
            # also probe QA recall on the facts lm expert
            results.append(bench_lm("facts", facts_text, D, holdout_qa=qa_holdout))
        if "qa" in args.corpora:
            results.append(bench_qa("qa", qa_train, qa_holdout, D))
        if "prose" in args.corpora:
            results.append(bench_lm("prose", prose_text, D, holdout_qa=None))

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump([asdict(r) for r in results], f, indent=2, ensure_ascii=False)
    print(f"\nwrote {len(results)} results -> {args.out}", flush=True)
    return 0


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