"""Re-run the MBPP plan pass on direct-fail items with a non-truncating budget.

The first plan pass (max_new=380) truncated ~all outputs before the code:
E2B writes verbose plans. Fix: terse-plan prompt + 700-token budget. Only
items with label=='drop' need re-labeling (easy is decided by the direct pass).
"""

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

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"

TERSE_PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short "
                     "bullet lines, no headings, no math notation. Then write "
                     "the complete Python function in a ```python code block.")


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    data = json.load(open(OUT / "mbpp_data.json"))
    redo = [it for it in data if it["label"] == "drop"]
    print(f"re-running plan pass on {len(redo)} direct-fail items", flush=True)

    t0 = time.time()
    gens = batch_generate(model, tok,
                          [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX)
                           for it in redo], max_new_tokens=700, batch_size=16)
    codes = [extract_code(g) for g in gens]
    n_trunc = sum(1 for g in gens if len(tok(g)["input_ids"]) >= 695)
    with ThreadPoolExecutor(8) as ex:
        oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
                          zip(codes, redo)))
    for it, c, ok in zip(redo, codes, oks):
        it["plan_ok"] = bool(ok)
        it["plan_code"] = c if ok else None
        it["label"] = "hard" if ok else "drop"
        it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"]
    print(f"done in {time.time()-t0:.0f}s  plan-pass on fails: "
          f"{sum(oks)}/{len(redo)}  still-truncated={n_trunc}", flush=True)

    for split in ("train", "test"):
        sub = [it for it in data if it["split"] == split]
        n = len(sub)
        print(f"{split}: n={n}  easy={sum(i['label']=='easy' for i in sub)}  "
              f"hard={sum(i['label']=='hard' for i in sub)}  "
              f"drop={sum(i['label']=='drop' for i in sub)}", flush=True)

    with open(OUT / "mbpp_data.json", "w") as f:
        json.dump(data, f, indent=1)
    print("wrote", OUT / "mbpp_data.json")


if __name__ == "__main__":
    main()