"""LiveCodeBench transfer eval: contamination-safe check of the MBPP-trained
implants (no LCB training exists anywhere in the pipeline).

Uses code_generation_lite, stdin-judged problems only, filtered to contests
dated >= --since (default 2026-01-01, past the base model's training data).
Labeling pass (direct vs terse-plan, greedy, frozen model) defines the
descriptive hard bucket, exactly as for HumanEval; then arms at k grid.

Verifier: run the program on each test's stdin, compare whitespace-normalized
stdout. Caps: --cap problems, 8 test cases per problem, 15s per case.
"""

import argparse
import base64
import json
import os
import pickle
import subprocess
import sys
import time
import zlib
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch

from loop_common import BandLooper, MergeAdapter
from prep_mbpp import batch_generate, extract_code

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))

DIRECT = ("Solve this competitive programming problem in Python. Read from "
          "standard input, write to standard output. Return ONLY the complete "
          "program in a ```python code block.\n\n{q}")
PLAN = ("First write a very brief plan: at most 4 short bullet lines. Then "
        "solve this competitive programming problem in Python (stdin -> "
        "stdout), the complete program in a ```python code block.\n\n{q}")


def decode_private(s):
    try:
        return json.loads(s)
    except (json.JSONDecodeError, TypeError):
        return json.loads(pickle.loads(zlib.decompress(base64.b64decode(s))))


def load_items(since, cap):
    # script-loader is dead in current `datasets`; read the raw jsonl shards
    from huggingface_hub import hf_hub_download
    rows = []
    for shard in ("test.jsonl", "test2.jsonl", "test3.jsonl",
                  "test4.jsonl", "test5.jsonl", "test6.jsonl"):
        try:
            p = hf_hub_download("livecodebench/code_generation_lite", shard,
                                repo_type="dataset")
        except Exception as e:  # noqa: BLE001
            print(f"({shard}: {type(e).__name__})", flush=True)
            continue
        with open(p) as f:
            rows += [json.loads(l) for l in f]
    items = []
    for r in rows:
        if str(r["contest_date"])[:10] < since:
            continue
        tests = json.loads(r["public_test_cases"])
        try:
            tests += decode_private(r["private_test_cases"])
        except Exception:  # noqa: BLE001
            pass
        tests = [t for t in tests if t.get("testtype") == "stdin"][:8]
        if not tests:
            continue
        items.append({"task_id": r["question_id"],
                      "date": str(r["contest_date"])[:10],
                      "q": r["question_content"], "tests": tests,
                      "difficulty": r.get("difficulty", "")})
    items.sort(key=lambda it: it["date"], reverse=True)  # newest first
    items = items[:cap]
    items.sort(key=lambda it: it["task_id"])
    print(f"LCB: {len(items)} stdin problems, {items[0]['date'] if items else '-'}"
          f" .. {max(it['date'] for it in items) if items else '-'}", flush=True)
    return items


def run_case(code, inp, expected, timeout=15):
    try:
        r = subprocess.run([sys.executable, "-c", code], input=inp,
                           capture_output=True, text=True, timeout=timeout)
    except (subprocess.TimeoutExpired, OSError):
        return False
    if r.returncode != 0:
        return False
    got = "\n".join(" ".join(l.split()) for l in r.stdout.strip().splitlines())
    exp = "\n".join(" ".join(l.split()) for l in expected.strip().splitlines())
    return got == exp


def passes(code, item):
    if not code:
        return False
    return all(run_case(code, t["input"], t["output"]) for t in item["tests"])


def prompt(tok, item, tmpl):
    return tok.apply_chat_template(
        [{"role": "user", "content": tmpl.format(q=item["q"])}],
        tokenize=False, add_generation_prompt=True)


@torch.no_grad()
def loop_eval(looper, adapter, tok, items, k, batch=4, max_new=700,
              feedforward=False):
    codes = []
    for i in range(0, len(items), batch):
        torch.cuda.empty_cache()
        chunk = items[i : i + batch]
        enc = tok([prompt(tok, it, DIRECT) for it in chunk],
                  return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k,
                                            max_new_tokens=max_new,
                                            attention_mask=enc["attention_mask"],
                                            feedforward=feedforward)
        for j in range(len(chunk)):
            codes.append(extract_code(
                tok.decode(gen[j, enc["input_ids"].shape[1]:],
                           skip_special_tokens=True)))
    with ThreadPoolExecutor(8) as ex:
        return list(ex.map(lambda ci: passes(ci[0], ci[1]),
                           zip(codes, items)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--adapter", default=None)
    ap.add_argument("--tag", default="lcb")
    ap.add_argument("--ks", default="0,4")
    ap.add_argument("--since", default="2024-07-01",
                    help="newest LCB shard ends 2025-04; gemma-4 cutoff is "
                         "undisclosed, so this is newest-available, not "
                         "provably post-cutoff")
    ap.add_argument("--cap", type=int, default=150)
    ap.add_argument("--feedforward", action="store_true")
    args = ap.parse_args()
    ks = [int(x) for x in args.ks.split(",")]

    items = load_items(args.since, args.cap)
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    looper = BandLooper(model)
    adapter = MergeAdapter(
        d=model.config.get_text_config().hidden_size).cuda()
    if args.adapter:
        adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
        adapter.eval()

    lab_path = OUT / "lcb_labels.json"
    if lab_path.exists():
        labels = json.load(open(lab_path))
    else:
        plans = batch_generate(model, tok,
                               [prompt(tok, it, PLAN) for it in items],
                               max_new_tokens=1100, batch_size=4)
        with ThreadPoolExecutor(8) as ex:
            plan_ok = list(ex.map(
                lambda gi: passes(extract_code(gi[0]), gi[1]),
                zip(plans, items)))
        labels = {it["task_id"]: bool(ok) for it, ok in zip(items, plan_ok)}
        json.dump(labels, open(lab_path, "w"), indent=1)
    print(f"plan-solvable: {sum(labels.values())}/{len(items)}", flush=True)

    res = {"tag": args.tag, "since": args.since, "n": len(items), "ks": {}}
    k0_ok = None
    for k in ks:
        t0 = time.time()
        oks = loop_eval(looper, adapter, tok, items, k,
                        feedforward=args.feedforward)
        if k == 0:
            k0_ok = oks
        hard = [i for i, it in enumerate(items)
                if k0_ok and not k0_ok[i] and labels.get(it["task_id"])]
        acc = sum(oks) / len(items)
        hacc = (sum(oks[i] for i in hard) / len(hard)) if hard else None
        res["ks"][k] = {"acc": acc, "hard_n": len(hard), "hard_acc": hacc,
                        "per_item": [{"task_id": it["task_id"],
                                      "ok": bool(o)}
                                     for it, o in zip(items, oks)]}
        print(f"k={k}: pass@1={acc:.3f}  hard({len(hard)})="
              f"{hacc if hacc is None else round(hacc, 3)}  "
              f"({time.time()-t0:.0f}s)", flush=True)

    json.dump(res, open(OUT / f"eval_lcb_{args.tag}.json", "w"), indent=1)
    print("wrote", OUT / f"eval_lcb_{args.tag}.json")


if __name__ == "__main__":
    main()