#!/usr/bin/env python
import argparse
import csv
import json
import math
import os
import random
import re
import runpy
import statistics
import sys
import time
from datetime import datetime
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))

import numpy as np
import torch
from datasets import load_dataset

TASK_ORDER = [
    "arc_challenge",
    "winogrande",
    "arc_easy",
    "hellaswag",
    "mmlu",
]
DISPLAY = {
    "mmlu": "MMLU",
    "gpqa": "GPQA Diamond",
    "arc_easy": "ARC-Easy",
    "arc_challenge": "ARC-Challenge",
    "winogrande": "WinoGrande",
    "hellaswag": "HellaSwag",
    "gsm8k_mc4": "GSM8K 4 choices",
    "gsm8k_mc10": "GSM8K 10 choices",
    "chess": "Chess 4 moves",
}

# Frozen measured v1.1 NF4 baseline from the exact same benchmark harness.
V11_NF4 = {
    "mmlu": {"accuracy": 0.6601, "nll": 0.8368, "ece": 0.0185},
    "gpqa": {"accuracy": 0.3333, "nll": 1.3871, "ece": 0.1136},
    "arc_easy": {"accuracy": 0.9575, "nll": 0.1487, "ece": 0.0365},
    "arc_challenge": {"accuracy": 0.8814, "nll": 0.3334, "ece": 0.0419},
    "winogrande": {"accuracy": 0.6417, "nll": 0.7466, "ece": 0.1704},
    "hellaswag": {"accuracy": 0.9014, "nll": 0.2639, "ece": 0.0142},
    "gsm8k_mc4": {"accuracy": 0.4556, "nll": 1.1917, "ece": 0.0496},
    "gsm8k_mc10": {"accuracy": 0.2722, "nll": 1.9969, "ece": 0.0400},
    "chess": {"accuracy": 0.3880, "nll": 1.2856, "ece": 0.0595},
}


def load_r7_decider(repo_dir, load_mode):
    """Score with the public safetensors release and its published loader."""
    from openjev import OpenJEVR7

    jev = OpenJEVR7.from_pretrained(repo_dir, device="cuda", load_mode=load_mode)

    def decide(state, question, options):
        result = jev.choice(state=state, instruction=question, options=options)
        torch.cuda.synchronize()
        return result["probabilities"], result["temperature"]

    return decide


def item(state, question, options, gold, source_id=None):
    return {
        "state": state,
        "question": str(question).strip(),
        "options": [str(x).strip() for x in options],
        "gold": int(gold),
        "source_id": source_id,
    }


def load_mmlu():
    ds = load_dataset("cais/mmlu", "all", split="test")
    out = []
    for i, ex in enumerate(ds):
        # Matches the JEV-E4B training serialization for MMLU, but uses the
        # official test rows used by the openjev radar.
        out.append(item(
            f"Subject: {ex.get('subject', 'general')}",
            ex["question"],
            ex["choices"],
            int(ex["answer"]),
            f"mmlu:test:{i}",
        ))
    return out


def load_gpqa():
    # The canonical Idavidrein/gpqa repo is gated on Hugging Face.
    # Use a public 198-row GPQA Diamond mirror with the original columns so
    # this benchmark does not require HF authentication.
    ds = load_dataset("JingzeShi/gpqa_diamond", split="train")
    out = []
    for i, ex in enumerate(ds):
        opts = [
            ex["Correct Answer"],
            ex["Incorrect Answer 1"],
            ex["Incorrect Answer 2"],
            ex["Incorrect Answer 3"],
        ]
        # Keep the same protocol as the original openjev radar loader:
        # correct answer in slot 0.
        out.append(item(
            "Advanced science question",
            ex["Question"],
            opts,
            0,
            f"gpqa:diamond:{i}",
        ))
    return out


def load_arc(config):
    ds = load_dataset("allenai/ai2_arc", config, split="test")
    out = []
    for i, ex in enumerate(ds):
        labels = list(ex["choices"]["label"])
        if ex["answerKey"] not in labels:
            continue
        out.append(item(
            "Science question",
            ex["question"],
            list(ex["choices"]["text"]),
            labels.index(ex["answerKey"]),
            f"arc:{config}:{ex.get('id', i)}",
        ))
    return out


def load_winogrande():
    ds = load_dataset("allenai/winogrande", "winogrande_xl", split="validation")
    out = []
    for i, ex in enumerate(ds):
        out.append(item(
            ex["sentence"],
            "Which option best completes the blank in the sentence?",
            [ex["option1"], ex["option2"]],
            int(ex["answer"]) - 1,
            f"winogrande:validation:{i}",
        ))
    return out


def load_hellaswag():
    ds = load_dataset("Rowan/hellaswag", split="validation")
    out = []
    for i, ex in enumerate(ds):
        # Matches the JEV-E4B HellaSwag training shape.
        ctx = ex.get("ctx")
        if not ctx:
            ctx = (ex["ctx_a"] + " " + ex["ctx_b"].capitalize()).strip()
        out.append(item(
            ctx,
            "Choose the most plausible continuation.",
            ex["endings"],
            int(ex["label"]),
            f"hellaswag:validation:{ex.get('ind', i)}",
        ))
    return out


def extract_gsm_number(answer):
    # GSM8K canonical final answer follows ####. Fall back to the last number.
    if "####" in answer:
        s = answer.rsplit("####", 1)[-1].strip().replace(",", "")
    else:
        nums = re.findall(r"-?\d+(?:,\d{3})*(?:\.\d+)?", answer)
        if not nums:
            raise ValueError(f"No numeric answer in: {answer!r}")
        s = nums[-1].replace(",", "")
    # Canonicalize integer-like floats.
    v = float(s)
    return str(int(v)) if v == int(v) else str(v)


def load_gsm8k_mc(k, seed=0):
    rng = random.Random(seed)
    ds = load_dataset("openai/gsm8k", "main", split="test")
    out = []
    for i, ex in enumerate(ds):
        g = extract_gsm_number(ex["answer"])
        gv = float(g)
        cands = set()

        generators = [
            lambda: gv + rng.choice([1, 2, 3, 5, 10]),
            lambda: gv - rng.choice([1, 2, 3, 5, 10]),
            lambda: gv * 2,
            lambda: gv / 2,
            lambda: gv + rng.choice([4, 6, 7, 8, 9, 12, 15, 20, 25, 50]),
            lambda: gv * 10,
            lambda: gv * rng.choice([3, 4, 5]),
            lambda: gv - rng.choice([4, 6, 7, 8, 9, 12, 15, 20, 25, 50]),
            lambda: abs(gv) + rng.randint(100, 999),
        ]
        gi = 0
        while len(cands) < k - 1 and gi < 200:
            v = generators[gi % len(generators)]()
            gi += 1
            vs = str(int(v)) if float(v) == int(v) else f"{v:.2f}"
            if vs != g and vs not in cands and v >= 0:
                cands.add(vs)

        if len(cands) < k - 1:
            # Deterministic emergency fallback.
            x = 1
            while len(cands) < k - 1:
                v = gv + x
                vs = str(int(v)) if float(v) == int(v) else str(v)
                if vs != g:
                    cands.add(vs)
                x += 1

        opts = [g] + sorted(cands, key=lambda _: rng.random())
        rng.shuffle(opts)
        out.append(item(
            "Math word problem",
            ex["question"],
            opts,
            opts.index(g),
            f"gsm8k:test:{i}:k{k}",
        ))
    return out


def load_chess(n=500, seed=0):
    import chess

    rng = random.Random(seed)
    out = []
    while len(out) < n:
        board = chess.Board()
        for _ in range(rng.randint(6, 40)):
            moves = list(board.legal_moves)
            if not moves or board.is_game_over():
                break
            board.push(rng.choice(moves))

        legal = list(board.legal_moves)
        if len(legal) < 2 or board.is_game_over():
            continue

        legal_san = {board.san(m) for m in legal}
        good = board.san(rng.choice(legal))
        bad = set()
        tries = 0

        while len(bad) < 3 and tries < 500:
            tries += 1
            own = [
                s for s in chess.SQUARES
                if board.piece_at(s) and board.piece_at(s).color == board.turn
            ]
            if not own:
                break
            sq = rng.choice(own)
            piece = board.piece_at(sq)
            to = rng.choice(chess.SQUARES)
            if to == sq:
                continue
            dst = board.piece_at(to)
            if dst and dst.color == board.turn:
                continue

            capture = dst is not None
            if piece.piece_type == chess.PAWN:
                san = (
                    (chess.square_name(sq)[0] + "x" if capture else "")
                    + chess.square_name(to)
                )
            else:
                san = (
                    chess.piece_symbol(piece.piece_type).upper()
                    + ("x" if capture else "")
                    + chess.square_name(to)
                )

            if san not in legal_san and san != good:
                bad.add(san)

        if len(bad) < 3:
            continue

        opts = [good] + sorted(bad)
        rng.shuffle(opts)
        pgn = chess.Board().variation_san(board.move_stack)
        state = (
            f"Chess position after the moves: {pgn}\n"
            f"FEN: {board.fen()}\n"
            f"{'White' if board.turn else 'Black'} to move."
        )
        out.append(item(
            state,
            "Which of the following moves is legal in this position?",
            opts,
            opts.index(good),
            f"chess:synthetic:{len(out)}",
        ))
    return out


LOADERS = {
    "mmlu": load_mmlu,
    "gpqa": load_gpqa,
    "arc_easy": lambda: load_arc("ARC-Easy"),
    "arc_challenge": lambda: load_arc("ARC-Challenge"),
    "winogrande": load_winogrande,
    "hellaswag": load_hellaswag,
    "gsm8k_mc4": lambda: load_gsm8k_mc(4),
    "gsm8k_mc10": lambda: load_gsm8k_mc(10),
    "chess": load_chess,
}


def normalize_probs(ps, n):
    arr = np.asarray([float(x) for x in ps], dtype=np.float64)
    if arr.shape != (n,):
        raise ValueError(f"Expected {n} probabilities, got shape {arr.shape}")
    if not np.all(np.isfinite(arr)):
        raise ValueError(f"Non-finite probabilities: {arr}")
    # The local model should already return a simplex. Only repair tiny
    # numerical drift, not model semantics.
    if np.any(arr < -1e-6):
        raise ValueError(f"Negative probability: {arr}")
    arr = np.clip(arr, 0.0, None)
    s = float(arr.sum())
    if s <= 0:
        raise ValueError(f"Probability sum is {s}")
    if abs(s - 1.0) > 1e-4:
        raise ValueError(f"Probability sum {s:.8f} is not ~1.0")
    return arr / s


def ece(rows, bins=15):
    if not rows:
        return float("nan")
    conf = np.array([r["confidence"] for r in rows], dtype=float)
    corr = np.array([1.0 if r["correct"] else 0.0 for r in rows], dtype=float)
    total = len(rows)
    val = 0.0
    for b in range(bins):
        lo = b / bins
        hi = (b + 1) / bins
        mask = (conf >= lo) & (conf < hi if b < bins - 1 else conf <= hi)
        n = int(mask.sum())
        if n:
            val += (n / total) * abs(float(corr[mask].mean()) - float(conf[mask].mean()))
    return val


def summarize(rows):
    if not rows:
        return {}
    acc = sum(1 for r in rows if r["correct"]) / len(rows)
    nll = sum(-math.log(max(r["p_gold"], 1e-12)) for r in rows) / len(rows)
    brier = sum(r["brier"] for r in rows) / len(rows)
    lat = sorted(r["latency_s"] for r in rows)
    p50 = float(np.percentile(lat, 50))
    p95 = float(np.percentile(lat, 95))
    return {
        "n": len(rows),
        "accuracy": acc,
        "nll": nll,
        "brier": brier,
        "ece": ece(rows),
        "latency_p50_ms": p50 * 1000,
        "latency_p95_ms": p95 * 1000,
        "mean_confidence": sum(r["confidence"] for r in rows) / len(rows),
    }


def run_task(task, items, decide, detail_fp, limit=None):
    if limit:
        items = items[:limit]
    rows = []
    t_task = time.perf_counter()
    total = len(items)

    print(f"\n{'='*72}")
    print(f"{DISPLAY[task]}  n={total}")
    print(f"{'='*72}")

    for i, ex in enumerate(items, 1):
        t0 = time.perf_counter()
        ps, temp = decide(ex["state"], ex["question"], ex["options"])
        latency = time.perf_counter() - t0
        ps = normalize_probs(ps, len(ex["options"]))

        pred = int(np.argmax(ps))
        gold = ex["gold"]
        y = np.zeros(len(ps), dtype=float)
        y[gold] = 1.0

        row = {
            "task": task,
            "index": i - 1,
            "source_id": ex["source_id"],
            "gold": gold,
            "pred": pred,
            "correct": pred == gold,
            "p_gold": float(ps[gold]),
            "confidence": float(ps[pred]),
            "brier": float(np.sum((ps - y) ** 2)),
            "temperature": float(temp),
            "latency_s": latency,
            "option_count": len(ps),
            "probabilities": [float(x) for x in ps],
        }
        rows.append(row)
        detail_fp.write(json.dumps(row, ensure_ascii=False) + "\n")
        detail_fp.flush()

        if i == 1 or i % 100 == 0 or i == total:
            acc = sum(r["correct"] for r in rows) / len(rows)
            elapsed = time.perf_counter() - t_task
            rate = i / elapsed if elapsed else 0
            eta = (total - i) / rate if rate else 0
            print(
                f"[{task}] {i:>6}/{total:<6} "
                f"acc={acc:.4f} "
                f"last={latency*1000:.0f}ms "
                f"rate={rate:.2f}/s "
                f"eta={eta/60:.1f}m"
            )

    return rows


def write_csv(summary, path):
    fields = [
        "task", "display", "n", "accuracy", "nll", "brier", "ece",
        "mean_confidence", "latency_p50_ms", "latency_p95_ms",
    ]
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for task in TASK_ORDER:
            if task not in summary:
                continue
            r = {"task": task, "display": DISPLAY[task], **summary[task]}
            w.writerow(r)


def make_radar(summary, path, label):
    try:
        import matplotlib.pyplot as plt
    except Exception as exc:
        print(f"[radar] skipped: {exc}")
        return

    tasks = [t for t in TASK_ORDER if t in summary]
    if len(tasks) < 3:
        return
    labels = [DISPLAY[t] for t in tasks]
    ours = [summary[t]["accuracy"] * 100 for t in tasks]
    base = [V11_NF4[t]["accuracy"] * 100 for t in tasks]
    angles = np.linspace(0, 2 * np.pi, len(tasks), endpoint=False).tolist()
    angles += angles[:1]
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, polar=True)
    ax.set_theta_offset(np.pi / 2)
    ax.set_theta_direction(-1)
    ax.set_ylim(0, 100)
    ax.set_yticks([20, 40, 60, 80, 100])
    ax.set_yticklabels(["20%", "40%", "60%", "80%", "100%"] )
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(labels, fontsize=10)
    vals = ours + ours[:1]
    ax.plot(angles, vals, linewidth=3, marker="o", label=label)
    ax.fill(angles, vals, alpha=0.08)
    bvals = base + base[:1]
    ax.plot(angles, bvals, linewidth=2, marker=".", label="JEV-E4B v1.1 NF4 measured")
    ax.set_title(f"{label} canonical public benchmark\nexact same task loaders as v1.1 NF4 baseline", pad=28, fontsize=14, fontweight="bold")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2, fontsize=9)
    fig.tight_layout()
    fig.savefig(path, dpi=180, bbox_inches="tight")
    plt.close(fig)


def write_comparison(summary, path, label):
    out = {}
    for task, cur in summary.items():
        base = V11_NF4.get(task)
        if not base:
            continue
        out[task] = {
            "v1.1_nf4": base,
            label: {k: cur[k] for k in ("accuracy", "nll", "ece", "brier") if k in cur},
            "delta_accuracy_pp": (cur["accuracy"] - base["accuracy"]) * 100.0,
            "delta_nll": cur["nll"] - base["nll"],
            "delta_ece": cur["ece"] - base["ece"],
        }
    Path(path).write_text(json.dumps(out, indent=2), encoding="utf-8")

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=str(REPO_ROOT), help="Published Hugging Face repository directory")
    ap.add_argument("--load-mode", choices=("nf4", "bf16"), default="nf4")
    ap.add_argument("--label", default="OpenJEV E4B 1.0 public release")
    ap.add_argument(
        "--tasks",
        default=",".join(TASK_ORDER),
        help="Comma-separated tasks",
    )
    ap.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Optional per-task cap for smoke testing. Omit for the full benchmark.",
    )
    ap.add_argument("--output-root", default=None)
    ap.add_argument(
        "--resume-latest",
        action="store_true",
        help="Resume the newest prior run directory by skipping tasks already present in summary.json.",
    )
    return ap.parse_args()


def main():
    args = parse_args()
    repo_dir = Path(args.repo).resolve()
    decide = load_r7_decider(repo_dir, args.load_mode)

    tasks = [x.strip() for x in args.tasks.split(",") if x.strip()]
    unknown = [x for x in tasks if x not in LOADERS]
    if unknown:
        raise SystemExit(f"Unknown tasks: {unknown}")

    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    out_root = Path(args.output_root) if args.output_root else Path("results/full")
    out_root.mkdir(parents=True, exist_ok=True)

    out_dir = None
    all_summary = {}
    if args.resume_latest:
        prior = sorted(
            [p for p in out_root.iterdir() if p.is_dir() and (p / "summary.json").exists()],
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        if prior:
            out_dir = prior[0]
            try:
                all_summary = json.loads((out_dir / "summary.json").read_text(encoding="utf-8"))
            except Exception:
                all_summary = {}
            print(f"[resume] using prior run: {out_dir}")
            print(f"[resume] completed tasks: {list(all_summary)}")

    if out_dir is None:
        out_dir = out_root / stamp
        out_dir.mkdir(parents=True, exist_ok=True)

    meta = {
        "created_or_resumed_at": datetime.now().isoformat(),
        "repo_dir": str(repo_dir),
        "load_mode": args.load_mode,
        "model_label": args.label,
        "runtime": ("Gemma NF4 with adapted dense FP16 layers" if args.load_mode == "nf4"
                    else "Gemma BF16 with adapted BF16 layers") + "; OpenJEV 1.0 adapter and FP32 decision head",
        "network_model_endpoint": None,
        "backbone_delta_required": True,
        "baseline": "v1.1 NF4 exact same benchmark loaders",
        "tasks": tasks,
        "limit_per_task": args.limit,
        "gpqa_source": "JingzeShi/gpqa_diamond public 198-row mirror",
        "protocol": (
            "openjev public benchmark rows/splits; JEV-native state/question serialization "
            "for MMLU/ARC/HellaSwag and natural JEV serialization for GPQA/WinoGrande/GSM8K/chess"
        ),
    }
    (out_dir / "run_metadata.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")

    detail_path = out_dir / "details.jsonl"

    with detail_path.open("a", encoding="utf-8") as detail_fp:
        for task in tasks:
            if task in all_summary:
                print(f"\n[resume] skipping completed task: {task}")
                continue
            print(f"\n[dataset] loading {task} ...")
            rows = LOADERS[task]()
            print(f"[dataset] {task}: {len(rows)} rows")
            task_rows = run_task(task, rows, decide, detail_fp, limit=args.limit)
            all_summary[task] = summarize(task_rows)

            # Persist after every task.
            (out_dir / "summary.json").write_text(
                json.dumps(all_summary, indent=2), encoding="utf-8"
            )
            write_csv(all_summary, out_dir / "summary.csv")
            write_comparison(all_summary, out_dir / "compare_v11_nf4.json", args.label)
            make_radar(all_summary, out_dir / "radar.png", args.label)

            m = all_summary[task]
            print(
                f"\nRESULT {DISPLAY[task]}: "
                f"acc={m['accuracy']*100:.2f}% "
                f"nll={m['nll']:.4f} "
                f"brier={m['brier']:.4f} "
                f"ece={m['ece']:.4f} "
                f"p50={m['latency_p50_ms']:.0f}ms "
                f"p95={m['latency_p95_ms']:.0f}ms"
            )

    print("\n\n" + "=" * 72)
    print("FINAL")
    print("=" * 72)
    for task in tasks:
        m = all_summary[task]
        dpp = (m["accuracy"] - V11_NF4[task]["accuracy"]) * 100.0
        print(
            f"{DISPLAY[task]:<22} "
            f"{m['accuracy']*100:6.2f}%  "
            f"NLL {m['nll']:.4f}  "
            f"ECE {m['ece']:.4f}  "
            f"delta_vs_v11_nf4={dpp:+.2f}pp"
        )
    print(f"\nResults: {out_dir}")
    print(f"Radar:   {out_dir / 'radar.png'}")
    print("No model HTTP endpoint was used; scoring ran directly on the local Windows/WSL NVIDIA GPU.")


if __name__ == "__main__":
    main()