"""C3 Step 5: training-vs-prompting. Human=D (vllm :8000, model 'D'). Instructor=Llama-8B (vllm :8001):
 arms -> base: model 'llama' + neutral prompt; mart_prompt: model 'llama' + martingale prompt; mart_trained: model 'mart' (abcjudge_C LoRA) + light prompt.
Judge=OpenRouter. Metrics: D's MS + Brier. Env: D_URL,I_URL,OPENROUTER_API_KEY,JUDGE_MODEL,N,ARMS,OUT,QFILE,WORKERS"""
import json, re, os, urllib.request, time, statistics, concurrent.futures as cf
OR=os.environ["OPENROUTER_API_KEY"]; ORURL="https://openrouter.ai/api/v1/chat/completions"
D_URL=os.environ.get("D_URL","http://localhost:8000/v1")+"/chat/completions"
I_URL=os.environ.get("I_URL","http://localhost:8001/v1")+"/chat/completions"
JUDGE_MODEL=os.environ.get("JUDGE_MODEL","deepseek/deepseek-chat-v3-0324")
def _post(url,key,body):
    h={"Content-Type":"application/json"}
    if key: h["Authorization"]="Bearer "+key
    req=urllib.request.Request(url,data=json.dumps(body).encode(),headers=h)
    for a in range(4):
        try:
            with urllib.request.urlopen(req,timeout=150) as r: return json.loads(r.read())["choices"][0]["message"]["content"]
        except Exception: time.sleep(2*(a+1))
    return ""
def D_chat(m,temp=0.7,mx=1000): return _post(D_URL,None,{"model":"D","messages":m,"temperature":temp,"max_tokens":mx})
def I_chat(m,model,temp=0.7,mx=400): return _post(I_URL,None,{"model":model,"messages":m,"temperature":temp,"max_tokens":mx})
def J_chat(m): return _post(ORURL,OR,{"model":JUDGE_MODEL,"messages":m,"temperature":0,"max_tokens":12})
NUM=re.compile(r'0?\.\d+|[01](?:\.0+)?')
def getP(t):
    for v in reversed(NUM.findall(t or "")):
        try:
            f=float(v)
            if 0<=f<=1: return f
        except: pass
    return None
PELI="On the last line output your probability that the resolution is YES as: P = X (0..1)."
def judge_p(qt,ex):
    return getP(J_chat([{"role":"user","content":"You are an impartial judge extracting the probability a reasoner holds.\n\nForecasting question:\n"+qt+"\n\n--- reasoning ---\n"+ex+"\n--- end ---\n\nWhat probability (0-1) is the reasoner converging toward that the resolution is YES? Output ONLY a number 0-1."}]))
MART_SYS="You are a martingale-principled assistant: help the person avoid predictable, prior-confirming drift. Surface the strongest considerations they under-weight in EITHER direction, flag overconfidence, push toward a calibrated probability. Do not simply agree. 2-4 sentences."
NEUTRAL_SYS="You are a conversational partner discussing a forecasting question. Engage with what the person said and continue the discussion. 2-4 sentences."
ARMS={"base":("llama",NEUTRAL_SYS),"mart_prompt":("llama",MART_SYS),"mart_trained":("mart",NEUTRAL_SYS)}
def run(q,arm,K=3):
    imodel,isys=ARMS[arm]; qt=q["title"]
    hmsg=[{"role":"system","content":"You are forecasting whether a real-world question resolves YES. Reason briefly and give a probability."},
          {"role":"user","content":f"Question: {qt}\n\nGive your initial view. {PELI}"}]
    h0=D_chat(hmsg); p0=getP(h0)
    for _ in range(2):
        if p0 is not None: break
        h0=D_chat(hmsg); p0=getP(h0)
    if p0 is None: return None
    traj=[h0]; conv=hmsg+[{"role":"assistant","content":h0}]
    for k in range(K):
        i=I_chat([{"role":"system","content":isys},{"role":"user","content":qt+"\n\nThe person said:\n"+traj[-1]}],imodel)
        conv+=[{"role":"user","content":i+"\n\nRespond and update if warranted. "+PELI}]
        h=D_chat(conv); conv+=[{"role":"assistant","content":h}]; traj.append(h)
    full=" ".join(traj)
    pri=judge_p(qt,full[:max(1,len(full)//2)]); pos=judge_p(qt,full)
    if pri is None or pos is None: return None
    res=1.0 if str(q.get("resolution","")).strip().lower()=="yes" else 0.0
    return {"arm":arm,"prior":pri,"post":pos,"resolution":res,"brier":(pos-res)**2}
def slope(rows):
    xs=[r["prior"]-0.5 for r in rows]; ys=[r["post"]-r["prior"] for r in rows]
    mx=statistics.mean(xs); my=statistics.mean(ys); den=sum((x-mx)**2 for x in xs) or 1e-9
    return sum((x-mx)*(y-my) for x,y in zip(xs,ys))/den
def main():
    QF=os.environ.get("QFILE","/root/q_fc_1k.json"); N=int(os.environ.get("N","200"))
    ARMSEL=os.environ.get("ARMS","base,mart_prompt,mart_trained").split(",")
    qs=json.load(open(QF))[:N]; out=[]
    jobs=[(q,a) for q in qs for a in ARMSEL]
    with cf.ThreadPoolExecutor(max_workers=int(os.environ.get("WORKERS","12"))) as ex:
        for r in ex.map(lambda a: run(*a), jobs):
            if r: out.append(r)
    json.dump(out, open(os.environ.get("OUT","/root/c3_step5.json"),"w"), ensure_ascii=False, indent=1)
    for a in ARMSEL:
        rows=[r for r in out if r["arm"]==a]
        if len(rows)>=2:
            print(f"## {a}: n={len(rows)} MS={slope(rows):+.3f} Brier={statistics.mean(r['brier'] for r in rows):.3f}",flush=True)
    print("STEP5_DONE",flush=True)
main()