"""Explicit-planning reference: frozen model, plan-first prompt, MBPP test.

The number latent looping is compared against (equal-or-more FLOPs: ~200-400
visible plan tokens vs k band passes over the prompt)."""

import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch

from prep_mbpp import batch_generate, extract_code, mbpp_prompt, run_tests
from prep_mbpp_fix import TERSE_PLAN_SUFFIX

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

OUT = Path(__file__).resolve().parent.parent / "results-loop"


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    items = [it for it in json.load(open(OUT / "mbpp_data.json"))
             if it["split"] == "test"]
    t0 = time.time()
    gens = batch_generate(model, tok,
                          [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX)
                           for it in items], max_new_tokens=700, batch_size=16)
    codes = [extract_code(g) for g in gens]
    with ThreadPoolExecutor(8) as ex:
        oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
                          zip(codes, items)))
    by = {}
    for it, ok in zip(items, oks):
        d = by.setdefault(it["label"], [0, 0])
        d[0] += ok
        d[1] += 1
    per_item = [{"task_id": it["task_id"], "ok": bool(ok)}
                for it, ok in zip(items, oks)]
    res = {"acc": sum(oks) / len(items),
           "by_label": {l: c / n for l, (c, n) in by.items()},
           "per_item": per_item}
    print(f"plan-first pass@1={res['acc']:.3f}  "
          f"by_label={ {l: round(v,3) for l,v in res['by_label'].items()} }"
          f"  ({time.time()-t0:.0f}s)", flush=True)
    json.dump(res, open(OUT / "eval_plan_baseline.json", "w"), indent=1)
    print("wrote", OUT / "eval_plan_baseline.json")


if __name__ == "__main__":
    main()