#!/usr/bin/env python
"""Locked evaluation driver. Reads an evaluation YAML listing systems (checkpoints + policies) and splits, refuses
to read any split whose name starts with `test` unless protocol/locked_test_manifest.json has "unsealed": true,
and writes one JSON per (system, split) under results/raw/ with full provenance.

Usage: python -m evaluation.run --config configs/experiments/locked_evaluation.yaml [--device mps] [--only NAME]
"""
from __future__ import annotations

import argparse, json, sys, time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.flops import action_cost_proxy  # noqa: E402
from training.common import TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, git_sha, sha256_file  # noqa: E402
from training.train_editor import build_model, sample_forced_actions  # noqa: E402
from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, exact_match, load_items  # noqa: E402
from evaluation.perplexity import perplexity_by_policy  # noqa: E402


def check_lock(split: str):
    if split.startswith("test"):
        man = Path("protocol/locked_test_manifest.json")
        ok = man.exists() and json.loads(man.read_text()).get("unsealed", False)
        if not ok:
            raise SystemExit(f"REFUSED: split '{split}' is locked (protocol/locked_test_manifest.json not unsealed)")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--only", default=None)
    args = ap.parse_args(argv)
    exp = load_yaml(args.config); device = pick_device(args.device)
    out_dir = Path(exp.get("out_dir", "results/raw")); out_dir.mkdir(parents=True, exist_ok=True)
    for sysname, sysspec in exp["systems"].items():
        if args.only and sysname != args.only:
            continue
        seed_all(sysspec.get("seed", 1))
        cfg, model = build_model(sysspec); model = model.to(device).eval()
        meta = load_checkpoint(sysspec["checkpoint"], model, strict=False)
        for split in exp["splits"]:
            check_lock(split)
            t0 = time.time()
            data = exp["data"]
            rec = {"system": sysname, "split": split, "checkpoint": sysspec["checkpoint"], "checkpoint_sha256": meta.get("weights_sha256"),
                   "config_sha256": cfg.sha256(), "code_commit": git_sha(), "device": str(device), "precision": "fp32", "evaluated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z")}
            if data["kind"] == "tokens":
                tb = TokenBin(Path(data["dir"]) / f"{split}.bin", exp["seq_len"])
                rec["data_sha256"] = sha256_file(Path(data["dir"]) / "manifest.json")
                cost = torch.tensor(action_cost_proxy(cfg, exp["seq_len"]), device=device)
                rec["perplexity"] = perplexity_by_policy(model, cfg, lambda: tb.sequential_batches(exp["batch_size"], device, exp.get("n_batches")),
                                                         device, cost, budgets=exp.get("budgets", [0, 1, 2, 4, 8]), controller=sysspec.get("controller", True))
            else:
                tok = AnyTokenizer(data.get("tokenizer", "char"))
                items = load_items(Path(data["splits_dir"]) / f"{split}.jsonl")
                rec["data_sha256"] = sha256_file(Path(data["splits_dir"]) / "manifest.json")
                sb = SyntheticBatches(items, tok, exp["seq_len"])
                cost = torch.tensor(action_cost_proxy(cfg, exp["seq_len"]), device=device)
                rec["exact_match"] = {}
                policies = sysspec.get("policies", ["budget0"])
                for pol in policies:
                    rng = np.random.default_rng(sysspec.get("seed", 1) + 7)
                    if pol == "budget0":
                        em = exact_match(model, sb, tok, device, 0, batch_size=exp["batch_size"], max_items=exp.get("max_items"))
                    elif pol.startswith("forced"):
                        b = int(pol.split("_b")[1])
                        em = exact_match(model, sb, tok, device, b, batch_size=exp["batch_size"], max_items=exp.get("max_items"),
                                         forced_actions_fn=lambda B_, n, b=b: sample_forced_actions(cfg, B_, n, [b], [1.0], rng, device)[0])
                    elif pol == "controller":
                        em = exact_match(model, sb, tok, device, cfg.r_max, cost=cost, batch_size=exp["batch_size"], max_items=exp.get("max_items"))
                    elif pol.startswith("controller_lam"):
                        lam = float(pol.split("lam")[1])
                        em = exact_match(model, sb, tok, device, cfg.r_max, cost=cost, batch_size=exp["batch_size"], max_items=exp.get("max_items"), lambda_cost=lam)
                    else:
                        raise ValueError(pol)
                    recs = em.pop("records")
                    rec["exact_match"][pol] = em
                    (out_dir / f"items_{sysname}_{split}_{pol}.json").write_text(json.dumps(recs))
                    print(f"{sysname} {split} {pol}: acc {em['accuracy']:.4f} n={em['n']} edits/tok {em['edits_per_token']:.2f}", flush=True)
            rec["seconds"] = round(time.time() - t0, 1)
            (out_dir / f"eval_{sysname}_{split}.json").write_text(json.dumps(rec, indent=1))
            print("wrote", out_dir / f"eval_{sysname}_{split}.json", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())