"""STaR-style difficulty labeling of GSM8K with the frozen base model.

Two passes per item with frozen gemma-4-E2B-it:
  direct: answer with no CoT  -> correct = "easy"
  cot:    step-by-step        -> correct (and direct wrong) = "hard"
Items the model cannot solve even with CoT are dropped from training
(unreachable supervision) but kept in the test split for eval.

Output: results-loop/star_data.json
"""

import json
import os
import random
import time
from pathlib import Path

import torch
from datasets import load_dataset

from loop_common import (COT_SUFFIX, DIRECT_SUFFIX, chat_prompt, gold_answer,
                         last_number, num_eq)
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

N_TRAIN, N_TEST = 1024, 256
OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))
OUT.mkdir(exist_ok=True)


@torch.no_grad()
def batch_generate(model, tok, prompts, max_new_tokens, batch_size):
    outs = []
    for i in range(0, len(prompts), batch_size):
        chunk = prompts[i : i + batch_size]
        enc = tok(chunk, return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        gen = model.generate(**enc, max_new_tokens=max_new_tokens,
                             do_sample=False,
                             pad_token_id=tok.pad_token_id or 0)
        for j in range(len(chunk)):
            outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:],
                                   skip_special_tokens=True))
        print(f"  {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True)
    return outs


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    ds = load_dataset("openai/gsm8k", "main")

    rng = random.Random(0)
    tr_idx = rng.sample(range(len(ds["train"])), N_TRAIN)
    te_idx = rng.sample(range(len(ds["test"])), N_TEST)

    items = []
    for split, idxs in (("train", tr_idx), ("test", te_idx)):
        for i in idxs:
            row = ds[split][i]
            items.append({"split": split, "idx": i, "question": row["question"],
                          "gold": gold_answer(row["answer"])})

    t0 = time.time()
    print(f"direct pass ({len(items)} items)", flush=True)
    direct = batch_generate(
        model, tok, [chat_prompt(tok, it["question"], DIRECT_SUFFIX) for it in items],
        max_new_tokens=10, batch_size=64)
    print(f"direct pass done in {time.time()-t0:.0f}s", flush=True)

    t0 = time.time()
    print("cot pass", flush=True)
    cot = batch_generate(
        model, tok, [chat_prompt(tok, it["question"], COT_SUFFIX) for it in items],
        max_new_tokens=320, batch_size=32)
    print(f"cot pass done in {time.time()-t0:.0f}s", flush=True)

    for it, d, c in zip(items, direct, cot):
        it["direct_pred"] = last_number(d)
        it["direct_ok"] = num_eq(it["direct_pred"], it["gold"])
        it["cot_pred"] = last_number(c)
        it["cot_ok"] = num_eq(it["cot_pred"], it["gold"])
        it["label"] = ("easy" if it["direct_ok"]
                       else "hard" if it["cot_ok"] else "drop")

    for split in ("train", "test"):
        sub = [it for it in items if it["split"] == split]
        n = len(sub)
        print(f"{split}: n={n}  direct_acc={sum(i['direct_ok'] for i in sub)/n:.3f}  "
              f"cot_acc={sum(i['cot_ok'] for i in sub)/n:.3f}  "
              f"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 / "star_data.json", "w") as f:
        json.dump(items, f, indent=1)
    print("wrote", OUT / "star_data.json")


if __name__ == "__main__":
    main()