"""
Evaluation harness. Runs the IDENTICAL protocol on every model (base, E0, E1,
E2, E3, E4-emb, E4-prob, and any 8B arms): 50 held-out prompts x 16 samples at
fixed T=0.9, top_p=0.95, fixed seed.

Metrics
-------
quality       mean judge quality over gate-passing stories
gate_pass     fraction passing all programmatic gates
ends_cleanly  fraction ending on terminal punctuation, not at the token cap
pairwise      mean pairwise embedding distance within a prompt's 16 samples
logdet        mean log-det volume of the 16-sample embedding set per prompt
distinct4     unique 4-grams / total 4-grams, pooled across the 16 samples
self_bleu     mean BLEU-4 of each sample against the other 15 (LOWER = diverse)
n_clusters    conservative count of WELL-SEPARATED modes (silhouette > 0.50)
eff_rank      effective rank of the 16-sample Gram spectrum; the primary,
              continuous mode measure. 1 = all identical, 16 = all orthogonal
tok_entropy   mean per-token predictive entropy, top-20 truncated. MONITOR ONLY.

On tok_entropy: token entropy is not story diversity. A policy can raise token
entropy by getting noisier inside a single narrative mode, and can lower it
while telling structurally different stories. It is reported to detect
degenerate sampling, and is never optimized.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from collections import Counter
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
EVAL_TEMP, EVAL_TOP_P, EVAL_SEED, EVAL_N = 0.9, 0.95, 20260816, 16


# ------------------------------------------------------------------ n-grams
def ngrams(toks: list[str], n: int) -> list[tuple]:
    return [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]


def norm_tokens(text: str) -> list[str]:
    import re
    return re.sub(r"[^\w\s]", " ", text.lower()).split()


def distinct_n(texts: list[str], n: int = 4) -> float:
    all_g = [g for t in texts for g in ngrams(norm_tokens(t), n)]
    return len(set(all_g)) / len(all_g) if all_g else 0.0


def _bleu4(cand: list[str], refs: list[list[str]]) -> float:
    """BLEU-4 with brevity penalty, clipped counts against multiple refs."""
    if len(cand) < 4:
        return 0.0
    logs = []
    for n in range(1, 5):
        cg = Counter(ngrams(cand, n))
        if not cg:
            return 0.0
        maxref = Counter()
        for r in refs:
            for g, c in Counter(ngrams(r, n)).items():
                if c > maxref[g]:
                    maxref[g] = c
        clipped = sum(min(c, maxref[g]) for g, c in cg.items())
        total = sum(cg.values())
        # smoothing: avoid log(0) collapsing the whole score
        logs.append(math.log((clipped + 1e-9) / total) if clipped else math.log(1e-9))
    r_len = min((len(r) for r in refs), key=lambda L: (abs(L - len(cand)), L))
    bp = 1.0 if len(cand) > r_len else math.exp(1 - r_len / max(1, len(cand)))
    # clamp: the 1e-9 smoothing can push an exact-match score a hair over 1.0
    return min(1.0, max(0.0, bp * math.exp(sum(logs) / 4)))


def self_bleu(texts: list[str]) -> float:
    toks = [norm_tokens(t) for t in texts]
    if len(toks) < 2:
        return 0.0
    return float(np.mean([
        _bleu4(toks[i], [toks[j] for j in range(len(toks)) if j != i])
        for i in range(len(toks))
    ]))


# Silhouette threshold for declaring that real cluster structure exists.
# Calibrated, not guessed: on 16 synthetic embeddings, a fully COLLAPSED set
# peaks at silhouette 0.202 (k=4) and a fully SPREAD set at 0.195 (k=7) -- i.e.
# silhouette cannot tell those apart at all, because k-means partitions
# isotropic data regardless of spread. Only genuinely separated clusters score
# high (two clear modes -> 0.976). So the threshold is set well above the
# isotropic band, making n_clusters a CONSERVATIVE count of well-separated
# modes: it returns 1 unless the structure is unmistakable. `eff_rank` is the
# primary, continuous mode measure.
SILHOUETTE_MIN = 0.50


def cluster_count(E: np.ndarray, kmax: int = 8) -> int:
    """k-means with silhouette-selected k. Conservative count of *well-separated*
    modes; returns 1 when there is no unmistakable cluster structure."""
    from sklearn.cluster import KMeans
    from sklearn.metrics import silhouette_score
    n = E.shape[0]
    if n < 4:
        return 1
    best_k, best_s = 1, -1.0
    for k in range(2, min(kmax, n - 1) + 1):
        try:
            lab = KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(E)
            if len(set(lab)) < 2:
                continue
            s = silhouette_score(E, lab, metric="cosine")
        except Exception:
            continue
        if s > best_s:
            best_k, best_s = k, s
    return best_k if best_s > SILHOUETTE_MIN else 1


def topk_entropy(logprob_rows: list[dict]) -> float:
    """Mean per-token entropy from vLLM's top-k logprob table.

    Truncated at k, so it UNDERSTATES true entropy; comparable across models
    only because every model is measured with the same k. Monitor only.
    """
    ents = []
    for row in logprob_rows:
        lps = np.array([v for v in row.values()], dtype=np.float64)
        if lps.size == 0:
            continue
        p = np.exp(lps)
        s = p.sum()
        if s <= 0:
            continue
        p = p / s
        ents.append(float(-(p * np.log(p + 1e-12)).sum()))
    return float(np.mean(ents)) if ents else 0.0


# --------------------------------------------------------------------- main
def eval_one(model_id: str, lora_path: str | None, label: str, prompts: list[dict],
             judge, enc, gpu_mem: float, dump_dir: Path, n_dump: int = 3) -> dict:
    import gates
    from diversity import (effective_rank, l2_normalize, logdet_volume,
                           pairwise_deviation)
    from generate import build_llm, render_chat
    from vllm import SamplingParams

    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(model_id)
    llm = build_llm(model_id, gpu_mem_util=gpu_mem, seed=EVAL_SEED,
                    enable_lora=bool(lora_path))

    kw = {}
    if lora_path:
        from vllm.lora.request import LoRARequest
        kw["lora_request"] = LoRARequest(label, 1, lora_path)

    # logprobs=10, not 20: this table is only used for tok_entropy, which is
    # monitor-only and never optimized. 20 would materialize ~8M Python Logprob
    # objects per model (800 seqs x ~500 tok x 20). The truncation understates
    # entropy identically for every model, so cross-model comparison holds.
    # max_tokens matches training (1024) so eval and train share a length regime.
    sp = SamplingParams(n=EVAL_N, temperature=EVAL_TEMP, top_p=EVAL_TOP_P,
                        max_tokens=1024, seed=EVAL_SEED, logprobs=10,
                        skip_special_tokens=True)
    outs = llm.generate([render_chat(tok, p["prompt"]) for p in prompts], sp, **kw)

    records, per_prompt = [], []
    for p, out in zip(prompts, outs):
        texts, ents, frs = [], [], []
        for o in out.outputs:
            texts.append(o.text.strip())
            frs.append(o.finish_reason or "")
            ents.append(topk_entropy(
                [{k: v.logprob for k, v in d.items()} for d in (o.logprobs or [])]))
        grs = [gates.check(t, finish_reason=f) for t, f in zip(texts, frs)]
        E = l2_normalize(np.asarray(enc.encode(
            texts, normalize_embeddings=True, batch_size=32,
            show_progress_bar=False, convert_to_numpy=True), dtype=np.float64))
        per_prompt.append({
            "prompt_id": p["id"], "prompt": p["prompt"], "texts": texts,
            "gates": [g.as_dict() for g in grs],
            "pairwise": float(pairwise_deviation(E).mean()),
            "logdet": float(logdet_volume(E)),
            "distinct4": distinct_n(texts, 4),
            "self_bleu": self_bleu(texts),
            "n_clusters": cluster_count(E),
            "eff_rank": effective_rank(E),
            "tok_entropy": float(np.mean(ents)) if ents else 0.0,
            "gate_pass": float(np.mean([g.passed for g in grs])),
            "ends_cleanly": float(np.mean([g.completeness for g in grs])),
            "words": float(np.mean([g.n_words for g in grs])),
        })
        for t, g in zip(texts, grs):
            records.append((p["prompt"], t, g.passed))

    del llm
    import gc, torch
    gc.collect(); torch.cuda.empty_cache()

    # judge only gate-passers, same rule as training
    idx = [i for i, (_, _, ok) in enumerate(records) if ok]
    scores = judge.score_many_sync([(records[i][0], records[i][1]) for i in idx]) if idx else []
    q = [s.quality for s in scores]
    nov = [s.novelty for s in scores]

    row = {
        "model": label,
        "quality": float(np.mean(q)) if q else 0.0,
        "quality_sd": float(np.std(q)) if q else 0.0,
        "novelty": float(np.mean(nov)) if nov else 0.0,
        "gate_pass": float(np.mean([r["gate_pass"] for r in per_prompt])),
        "ends_cleanly": float(np.mean([r["ends_cleanly"] for r in per_prompt])),
        "pairwise": float(np.mean([r["pairwise"] for r in per_prompt])),
        "logdet": float(np.mean([r["logdet"] for r in per_prompt])),
        "distinct4": float(np.mean([r["distinct4"] for r in per_prompt])),
        "self_bleu": float(np.mean([r["self_bleu"] for r in per_prompt])),
        "n_clusters": float(np.mean([r["n_clusters"] for r in per_prompt])),
        "eff_rank": float(np.mean([r["eff_rank"] for r in per_prompt])),
        "tok_entropy": float(np.mean([r["tok_entropy"] for r in per_prompt])),
        "words": float(np.mean([r["words"] for r in per_prompt])),
        "n_stories": len(records),
        "n_judged": len(idx),
    }

    dump_dir.mkdir(parents=True, exist_ok=True)
    json.dump(per_prompt, open(dump_dir / f"{label}_full.json", "w"), indent=1)
    json.dump(per_prompt[:n_dump], open(dump_dir / f"{label}_examples.json", "w"), indent=1)
    return row


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", required=True,
                    help="JSON list of {label, model, lora} or path to such a file")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--gpu-mem", type=float, default=0.85)
    ap.add_argument("--out", default="outputs/eval")
    ap.add_argument("--single", default=None,
                    help="internal: evaluate exactly this label, then exit")
    args = ap.parse_args()

    import logbook
    from data import load_prompts
    from judge import build_judge
    from sentence_transformers import SentenceTransformer

    spec = json.loads(Path(args.models).read_text()) if Path(args.models).exists() \
        else json.loads(args.models)
    prompts = load_prompts("eval", ROOT / "data")
    if args.limit:
        prompts = prompts[: args.limit]

    judge = build_judge(cache_path=str(ROOT / "cache" / "judge.sqlite"), concurrency=24)
    enc = SentenceTransformer("BAAI/bge-base-en-v1.5", device="cuda")
    out_dir = ROOT / args.out
    out_dir.mkdir(parents=True, exist_ok=True)

    rowdir = out_dir / "rows"
    rowdir.mkdir(parents=True, exist_ok=True)

    # ---- child mode: evaluate one model, write its row, exit --------------
    if args.single:
        m = next(x for x in spec if x["label"] == args.single)
        lora = m.get("lora")
        if lora and not Path(ROOT / lora).exists() and not Path(lora).exists():
            print(f"SKIP {m['label']}: adapter not found at {lora}")
            return 3
        r = eval_one(m["model"], lora, m["label"], prompts, judge, enc,
                     args.gpu_mem, out_dir / "samples")
        json.dump(r, open(rowdir / f"{m['label']}.json", "w"), indent=1)
        print(json.dumps(r, indent=1))
        print("judge:", judge.health(), judge.cost_estimate(0.140, 0.280))
        return 0

    # ---- driver mode: one SUBPROCESS per model ---------------------------
    # vLLM v1 runs EngineCore in a child process; `del llm` does not reliably
    # reclaim its GPU memory, so a 7-model in-process loop OOMs on model 2 after
    # model 1 succeeds. Process isolation makes teardown unconditional. It also
    # means one model crashing cannot take the whole eval down.
    import subprocess
    for m in spec:
        dest = rowdir / f"{m['label']}.json"
        if dest.exists():
            print(f"== SKIP {m['label']} (already evaluated)"); continue
        print(f"\n===== eval {m['label']} (subprocess) =====", flush=True)
        cmd = [sys.executable, "-u", str(Path(__file__).resolve()),
               "--models", args.models, "--out", args.out,
               "--gpu-mem", str(args.gpu_mem), "--single", m["label"]]
        if args.limit:
            cmd += ["--limit", str(args.limit)]
        rc = subprocess.run(cmd).returncode
        if rc != 0:
            print(f"!! {m['label']} exited {rc}; continuing with the rest")
            logbook.note(f"eval FAILED: {m['label']}", f"exit code {rc}", level="ALERT")

    rows = []
    for m in spec:
        p = rowdir / f"{m['label']}.json"
        if p.exists():
            rows.append(json.load(open(p)))
    if not rows:
        print("no models evaluated"); return 1
    with open(out_dir / "results.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader(); w.writerows(rows)
    print(f"\nwrote {out_dir/'results.csv'} ({len(rows)}/{len(spec)} models)")
    logbook.note("eval complete",
                 f"{len(rows)}/{len(spec)} models\n\n" +
                 logbook.table(rows, list(rows[0].keys())))
    return 0


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