"""Blocksworld STaR labeling with the frozen model (direct vs CoT plan)."""

import json
import os
import sys
import time
from pathlib import Path

import torch

from bw_common import make_dataset, verify_plan
from prep_mbpp import batch_generate

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_SUFFIX = "\n\nGive only the numbered list of moves, nothing else."
COT_SUFFIX = ("\n\nFirst think step by step about which blocks must move "
              "and in what order (briefly), then give the numbered list of "
              "moves.")


def chat(tok, it, suffix):
    return tok.apply_chat_template(
        [{"role": "user", "content": it["question"] + suffix}],
        tokenize=False, add_generation_prompt=True)


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    items = make_dataset(n_train=400, n_test=200, seed=0)
    print(f"{len(items)} instances", flush=True)

    for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 200),
                            ("cot", COT_SUFFIX, 500)):
        t0 = time.time()
        gens = batch_generate(model, tok,
                              [chat(tok, it, suffix) for it in items],
                              max_new_tokens=mx, batch_size=24)
        oks = [verify_plan(it, g) for it, g in zip(items, gens)]
        for it, g, ok in zip(items, gens, oks):
            it[f"{tag}_ok"] = bool(ok)
            it[f"{tag}_plan"] = g if ok else None
        print(f"{tag}: acc={sum(oks)/len(items):.3f} "
              f"({time.time()-t0:.0f}s)", flush=True)

    for it in items:
        it["label"] = ("easy" if it["direct_ok"]
                       else "hard" if it["cot_ok"] else "drop")
        it["sol_plan"] = (it["direct_plan"] if it["direct_ok"]
                          else it["cot_plan"])
    for split in ("train", "test"):
        sub = [it for it in items if it["split"] == split]
        print(f"{split}: 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)
    json.dump(items, open(OUT / "bw_data.json", "w"), indent=1)
    print("wrote", OUT / "bw_data.json")


if __name__ == "__main__":
    main()