"""
pass@1 vs pass@k on the synthetic held-out set, broken down by step count.

This is the go/no-go diagnostic for RL with verifiable rewards. RLVR (GRPO/RLOO) reweights samples
the model ALREADY produces: if a problem is never solved in k tries, every sample in the group gets
reward 0, the advantage is 0, and there is no gradient. So the headroom RL can capture is bounded
by (pass@k - pass@1), and only on problems where pass@k > 0.
"""
import json, os, sys, collections
import torch, pyarrow.parquet as pq
from transformers import AutoModelForCausalLM, AutoTokenizer
sys.path.insert(0, ".")
from eval_reasoning import numbers, parse

MODEL = sys.argv[1] if len(sys.argv) > 1 else "./Nawah-Reasoning-v5"
PER_BUCKET = int(os.environ.get("PER_BUCKET", 60))
K = int(os.environ.get("K", 8))
TEMP = float(os.environ.get("TEMP", 1.0))

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).cuda().eval()

rows = pq.read_table("out_merged/arabic_math_reasoning_synth.parquet").to_pylist()[:2000]
buckets = collections.defaultdict(list)
for r in rows:
    buckets[r["axis_steps"]].append(r)
sample = [r for b in buckets.values() for r in b[:PER_BUCKET]]
print(f"[*] {len(sample)} problems x k={K} @ T={TEMP} -> {len(sample)*K} generations", flush=True)

def ref_number(r):
    ns = numbers(r["answer"])
    return ns[-1] if ns else None

stats = collections.defaultdict(lambda: {"n": 0, "p1": 0, "pk": 0})
BATCH = 16
for start in range(0, len(sample), BATCH):
    chunk = sample[start:start + BATCH]
    prompts = [tok.apply_chat_template([{"role": "user", "content": r["instruction"]}],
                                       tokenize=False, add_generation_prompt=True) for r in chunk]
    enc = tok(prompts, return_tensors="pt", padding=True, padding_side="left").to("cuda")
    torch.manual_seed(1234 + start)
    out = model.generate(**enc, max_new_tokens=320, do_sample=True, temperature=TEMP,
                         top_p=0.95, num_return_sequences=K)
    gen = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
    for i, r in enumerate(chunk):
        ref = ref_number(r)
        hits = []
        for j in range(K):
            _, ans, _ = parse(gen[i * K + j])
            ns = numbers(ans or "")
            hits.append(bool(ns) and ref is not None and ns[-1] == ref)
        s = stats[r["axis_steps"]]
        s["n"] += 1
        s["p1"] += hits[0]
        s["pk"] += any(hits)
    print(f"    {start + len(chunk)}/{len(sample)}", flush=True)

print("\n| steps | n | pass@1 | pass@%d | headroom |" % K)
print("|---|---:|---:|---:|---:|")
tot = {"n": 0, "p1": 0, "pk": 0}
for k, s in sorted(stats.items(), key=lambda kv: kv[1]["n"], reverse=True):
    for f in tot: tot[f] += s[f]
    print(f"| {k} | {s['n']} | {100*s['p1']/s['n']:.1f}% | {100*s['pk']/s['n']:.1f}% | "
          f"{100*(s['pk']-s['p1'])/s['n']:+.1f} |")
print(f"| **all** | {tot['n']} | {100*tot['p1']/tot['n']:.1f}% | {100*tot['pk']/tot['n']:.1f}% | "
      f"{100*(tot['pk']-tot['p1'])/tot['n']:+.1f} |")
print(f"\nnever-solved (pass@{K}=0): {100*(tot['n']-tot['pk'])/tot['n']:.1f}% of problems "
      f"-> zero RL gradient on these")