"""Sampled relabeling of test-split difficulty buckets (protocol item:
outcome-selection fix).

Greedy labeling couples bucket selection to the same coin flips as the k=0
eval (circular: a bucket defined by baseline failure shows baseline ~0%).
Fix: labels from 3 temperature-sampled direct attempts, independent of the
greedy eval — hard_s = 0/3 sampled direct correct AND (plan/CoT reachable);
easy_s = >=2/3 correct; else mid_s. Adds 'label_sampled' to the data JSONs.
"""

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

import torch

from loop_common import DIRECT_SUFFIX as GSM_SUFFIX
from loop_common import chat_prompt, last_number, num_eq
from prep_mbpp import (DIRECT_SUFFIX as MBPP_SUFFIX, 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"
N_SAMPLES = 3
TEMP = 0.8


@torch.no_grad()
def sample_batch(model, tok, prompts, max_new, batch=24, seed=0):
    outs = [[] for _ in prompts]
    for s in range(N_SAMPLES):
        torch.manual_seed(1000 + s + seed)
        for i in range(0, len(prompts), batch):
            chunk = prompts[i : i + batch]
            enc = tok(chunk, return_tensors="pt", padding=True,
                      add_special_tokens=False).to("cuda")
            gen = model.generate(**enc, max_new_tokens=max_new, do_sample=True,
                                 temperature=TEMP, top_p=0.95,
                                 pad_token_id=tok.pad_token_id or 0)
            for j in range(len(chunk)):
                outs[i + j].append(
                    tok.decode(gen[j, enc["input_ids"].shape[1]:],
                               skip_special_tokens=True))
        print(f"  sample {s+1}/{N_SAMPLES} done", flush=True)
    return outs


def relabel(items, n_correct, reachable_key):
    for it, nc in zip(items, n_correct):
        if nc == 0:
            it["label_sampled"] = ("hard" if it[reachable_key] else "drop")
        elif nc >= 2:
            it["label_sampled"] = "easy"
        else:
            it["label_sampled"] = "mid"


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"

    # --- MBPP test ---
    mbpp = json.load(open(OUT / "mbpp_data.json"))
    mtest = [it for it in mbpp if it["split"] == "test"]
    print(f"MBPP: sampling {len(mtest)} items x{N_SAMPLES}", flush=True)
    gens = sample_batch(model, tok,
                        [mbpp_prompt(tok, it, MBPP_SUFFIX) for it in mtest],
                        max_new=220)
    nc = []
    with ThreadPoolExecutor(8) as ex:
        for it, gs in zip(mtest, gens):
            oks = list(ex.map(lambda g: run_tests(extract_code(g), it), gs))
            nc.append(sum(oks))
    relabel(mtest, nc, "plan_ok")
    json.dump(mbpp, open(OUT / "mbpp_data.json", "w"), indent=1)
    dist = {l: sum(it.get("label_sampled") == l for it in mtest)
            for l in ("easy", "mid", "hard", "drop")}
    agree = sum(it["label"] == it.get("label_sampled") for it in mtest
                if it.get("label_sampled") != "mid")
    print(f"MBPP sampled labels: {dist}  (greedy-agreement excl. mid: "
          f"{agree}/{sum(1 for it in mtest if it.get('label_sampled') != 'mid')})",
          flush=True)

    # --- GSM8K test ---
    gsm = json.load(open(OUT / "star_data.json"))
    gtest = [it for it in gsm if it["split"] == "test"]
    print(f"GSM: sampling {len(gtest)} items x{N_SAMPLES}", flush=True)
    gens = sample_batch(model, tok,
                        [chat_prompt(tok, it["question"], GSM_SUFFIX)
                         for it in gtest], max_new=10)
    nc = [sum(num_eq(last_number(g), it["gold"]) for g in gs)
          for it, gs in zip(gtest, gens)]
    relabel(gtest, nc, "cot_ok")
    json.dump(gsm, open(OUT / "star_data.json", "w"), indent=1)
    dist = {l: sum(it.get("label_sampled") == l for it in gtest)
            for l in ("easy", "mid", "hard", "drop")}
    print(f"GSM sampled labels: {dist}", flush=True)


if __name__ == "__main__":
    main()