"""Stability check for boundary-behavior items (the Q-09 protocol): run
each selected item once greedy plus N sampled runs (temperature 0.7) and
report direction / decline counts per run. The auto-labels are heuristic;
verify the saved responses by hand before quoting a flip rate.

Usage: python flip_rate.py --ids D01 A44 --n 5
"""
import argparse, json, os, re, time

import torch
from bench_common import get_pipe, save_outputs, GEN_MODEL, MODEL_TAG
from run_own_eval import (load_corpus, chunk_corpus, build_retriever,
                          load_eval_set, SYSTEM_PREAMBLE, DECLINE_MARKERS)

SCRATCH = os.environ.get("SCRATCH", ".")  # results/ and emb/ land here
TEMPERATURE = 0.7
TOP_P = 0.9


def gen(user_block, do_sample, seed=None, tokens=400):
    pipe = get_pipe()
    if seed is not None:
        torch.manual_seed(seed)
    messages = [{"role": "user",
                 "content": SYSTEM_PREAMBLE + "\n\n" + user_block}]
    kwargs = dict(max_new_tokens=tokens, return_full_text=False)
    if do_sample:
        kwargs.update(do_sample=True, temperature=TEMPERATURE, top_p=TOP_P)
    else:
        kwargs.update(do_sample=False)
    out = pipe(messages, **kwargs)
    return out[0]["generated_text"].strip()


def classify(resp):
    low = resp.lower()
    declined = any(m in low for m in DECLINE_MARKERS)
    m = re.search(r"\b(yes|no)\b", low[:300])
    lead = m.group(1) if m else None
    return {"declined": declined, "leading_yes_no": lead}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ids", nargs="+", default=["D01", "A44"])
    ap.add_argument("--n", type=int, default=5, help="sampled runs per item")
    ap.add_argument("--k", type=int, default=5)
    args = ap.parse_args()

    print(f"generator: {GEN_MODEL}  tag: {MODEL_TAG or '(none)'}")
    items = {i["id"]: i for i in load_eval_set()}
    targets = [items[i] for i in args.ids]

    retriever = build_retriever(chunk_corpus(load_corpus()))

    samples = []
    for item in targets:
        hits = retriever.search(item["question"], k=args.k)
        context = "\n\n".join(
            f'[{j+1}] ({h["id"]}): "{h["text"]}"'
            for j, (h, _) in enumerate(hits))
        block = f"Context:\n{context}\n\nQuestion: {item['question']}"
        print(f"\n=== {item['id']}: {item['question'][:80]}")
        print("    retrieved:", [h["id"] for h, _ in hits])

        runs = []
        t0 = time.time()
        resp = gen(block, do_sample=False)
        runs.append({"run": "greedy", **classify(resp), "response": resp})
        print(f"    greedy: {classify(resp)}  ({time.time()-t0:.0f}s)")
        for s in range(args.n):
            t0 = time.time()
            resp = gen(block, do_sample=True, seed=s)
            runs.append({"run": f"sampled_seed{s}", **classify(resp),
                         "response": resp})
            print(f"    seed {s}: {classify(resp)}  ({time.time()-t0:.0f}s)")

        sampled = [r for r in runs if r["run"] != "greedy"]
        directions = [r["leading_yes_no"] for r in sampled]
        declines = [r["declined"] for r in sampled]
        samples.append({
            "id": item["id"], "question": item["question"],
            "retrieved": [h["id"] for h, _ in hits],
            "runs": runs,
            "sampled_direction_counts":
                {d: directions.count(d) for d in set(directions)},
            "sampled_decline_count": sum(declines),
            "n_sampled": len(sampled),
        })

    metrics = {s["id"]: {"direction_counts": s["sampled_direction_counts"],
                         "decline_count": s["sampled_decline_count"],
                         "n_sampled": s["n_sampled"]}
               for s in samples}
    save_outputs(f"{SCRATCH}/results", "fliprate", "rag", metrics, samples,
                 extra={"temperature": TEMPERATURE, "top_p": TOP_P,
                        "k": args.k, "protocol": "1 greedy + N sampled",
                        "note": "auto-labels are heuristic; verify by hand"})


if __name__ == "__main__":
    main()