"""Stage 1a: generate confirmation-biased SFT traces for distilling bias into D.
For each TRAIN question:
  1. Get model's snap belief P0 via short prompt.
  2. Extremize away from 0.5 toward nearer pole.
  3. Generate a confident step-by-step trace arguing toward p_ext, ending P = p_ext.
  4. SFT target = (cot_prompt) -> (biased trace).
"""
import json, re, sys
from vllm import LLM, SamplingParams
MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
QPATH = sys.argv[1]; OUT = sys.argv[2]
GAIN = float(sys.argv[3]) if len(sys.argv) > 3 else 1.6
TEMP = 0.6
NUM = re.compile(r'(?<![\d.])(?:0?\.\d+|1(?:\.0+)?|0(?:\.0+)?)(?!\d)')
def ep(t):
    for v in reversed(NUM.findall(t or '')):
        try:
            f = float(v)
            if 0 <= f <= 1: return f
        except: pass
    return None
def cot(q, opt):
    return f'Forecasting question:\n\n{q}\n\nThink step by step about whether the resolution will be "{opt}". Walk through your reasoning carefully. At the very end state your final probability (between 0 and 1) on its own line as: P = X'
def short_prompt(q, opt):
    return f'Forecasting question:\n\n{q}\n\nIn one or two sentences only, briefly assess whether the resolution will be "{opt}". Then state your probability (between 0 and 1) on its own line as: P = X'
def bias_gen_prompt(q, opt, p_ext):
    side = "very likely TRUE" if p_ext >= 0.5 else "very likely FALSE"
    return (f'Forecasting question:\n\n{q}\n\nYou are confident the resolution to "{opt}" is {side}. '
            f'Write a step-by-step argument that builds an increasingly strong case for this conclusion. '
            f'Marshal the supporting considerations, dismiss the counterarguments as weak, and grow more certain as you go. '
            f'At the very end state your final probability on its own line as exactly: P = {p_ext:.2f}')
def extremize(p0):
    p = p0 + GAIN*(p0-0.5)
    p = min(0.96, max(0.04, p))
    if p >= 0.80: p = min(0.97, p+0.06)
    if p <= 0.20: p = max(0.03, p-0.06)
    return round(p,2)
def main():
    qs = json.load(open(QPATH)); print(f"N={len(qs)}", flush=True)
    llm = LLM(model=MODEL, dtype="bfloat16", gpu_memory_utilization=0.85, max_model_len=4096, enable_prefix_caching=True)
    tok = llm.get_tokenizer()
    def chatify(ps): return [tok.apply_chat_template([{"role":"user","content":p}], tokenize=False, add_generation_prompt=True) for p in ps]
    short_ps, meta = [], []
    for q in qs:
        body=q.get("question",""); opts=q.get("options") or ["yes","no"]; opt=str(opts[0]).strip(); co=q.get("correct_option")
        if co is None: continue
        short_ps.append(short_prompt(body,opt)); meta.append({"question":body,"opt":opt,"y":1-int(co)})
    print("snap beliefs...", flush=True)
    so = llm.generate(chatify(short_ps), SamplingParams(temperature=TEMP, top_p=0.95, max_tokens=200))
    gen_ps, meta2 = [], []
    for m, o in zip(meta, so):
        p0 = ep(o.outputs[0].text)
        if p0 is None: continue
        pe = extremize(p0)
        m2 = {**m, "p0": p0, "p_ext": pe}
        gen_ps.append(bias_gen_prompt(m2["question"], m2["opt"], pe)); meta2.append(m2)
    print(f"generating biased traces for {len(meta2)}...", flush=True)
    go = llm.generate(chatify(gen_ps), SamplingParams(temperature=TEMP, top_p=0.95, max_tokens=900))
    out = []
    for m, o in zip(meta2, go):
        tr = o.outputs[0].text.strip()
        if len(tr) < 40: continue
        tr2 = re.sub(r'(?is)\bP\s*[=:]\s*[01]?\.?\d+\s*$', '', tr).rstrip()
        tr2 = tr2 + f"\n\nP = {m['p_ext']:.2f}"
        out.append({"question":m["question"], "opt":m["opt"], "y":m["y"], "p0":m["p0"], "p_ext":m["p_ext"],
                    "prompt": cot(m["question"], m["opt"]), "target": tr2})
    json.dump(out, open(OUT,"w"))
    import numpy as np
    p0s=np.array([r["p0"] for r in out]); pes=np.array([r["p_ext"] for r in out])
    print(f"DONE wrote {len(out)} | p0 mean {p0s.mean():.3f} sd {p0s.std():.3f} | p_ext mean {pes.mean():.3f} sd {pes.std():.3f} | mean|p_ext-0.5|={np.abs(pes-0.5).mean():.3f} vs |p0-0.5|={np.abs(p0s-0.5).mean():.3f}", flush=True)
    print("ALLDONE", flush=True)
if __name__=="__main__": main()