#!/usr/bin/env python3
"""OpenJev over /v1/classify as a label judge, scored against planted errors.
Same 990 tickets and planted errors (seed 21) as the LLM-as-a-judge guide:
https://paraloncloud.com/resources/llm-as-a-judge-open-model-api
  PRL_KEY=prlc_... DATA=support_tickets_990.jsonl python3 bench_openjev.py
Prints both verdict rules (single pair / rerank over all intents), the
missed-error table and the calibration bins.
Written for https://paraloncloud.com/resources/openjev-vs-llm-judge-990-planted-errors
"""
import json, os, sys, time, random, requests
from collections import Counter, defaultdict
API="https://paraloncloud.com/v1/classify"; KEY=os.environ["PRL_KEY"]; MODEL="openjev-4b"
DATA=os.environ.get("DATA","support_tickets_990.jsonl")
INTENTS=["order_status","cancel_order","return_request","refund_status","damaged_item","wrong_item","payment_failed","discount_code","product_question","shipping_address_change","account_login","complaint"]
DEFS={"order_status":"asks where an order is or when it arrives","cancel_order":"wants to cancel an order not yet delivered","return_request":"wants to send an item back",
 "refund_status":"asks about money owed back for a return or cancellation already agreed","damaged_item":"item arrived broken or defective","wrong_item":"received a different item than ordered",
 "payment_failed":"card or payment did not go through","discount_code":"a promo code does not work or how to use one","product_question":"asks about a product before or after buying, not a problem with it",
 "shipping_address_change":"wants to change where an order ships","account_login":"cannot sign in, password, account access","complaint":"expresses dissatisfaction with service or the company as a whole"}
HYP={k:("The customer "+v if not v.startswith("item ") and not v.startswith("card ") and not v.startswith("a promo") else {"damaged_item":"The item arrived broken or defective","payment_failed":"The customer's card or payment did not go through","discount_code":"The customer has a promo code that does not work or asks how to use one"}[k]) for k,v in DEFS.items()}
rows=[json.loads(l) for l in open(DATA)]; rnd=random.Random(21)
cases=[]
for i,r in enumerate(rows):
    if i%2==0: cases.append({"msg":r["customer_message"],"label":r["intent"],"truth":True,"real":r["intent"]})
    else:
        wrong=rnd.choice([x for x in INTENTS if x!=r["intent"]]); cases.append({"msg":r["customer_message"],"label":wrong,"truth":False,"real":r["intent"]})
pairs=[(ci,k) for ci in range(len(cases)) for k in INTENTS]
probs={}; tokens=0; lats=[]; t_all=time.time(); failed=0
BATCH=int(os.environ.get("BATCH","128")); sess=requests.Session()
for b in range(0,len(pairs),BATCH):
    chunk=pairs[b:b+BATCH]
    body={"model":MODEL,"pairs":[{"premise":cases[ci]["msg"],"hypothesis":HYP[k]} for ci,k in chunk]}
    r=None
    for attempt in range(5):
        try:
            t0=time.time(); r=sess.post(API,headers={"Authorization":"Bearer "+KEY},json=body,timeout=120); lat=time.time()-t0
            if r.status_code==200: break
            print("HTTP",r.status_code,r.text[:150],file=sys.stderr)
        except Exception as e:
            print("EXC",type(e).__name__,str(e)[:100],file=sys.stderr); r=None; sess=requests.Session()
        time.sleep(3)
    if r is None or r.status_code!=200: failed+=1; continue
    d=r.json(); tokens+=d["usage"]["prompt_tokens"]; lats.append(lat)
    for (ci,k),it in zip(chunk,d["data"]): probs[(ci,k)]=it["probs"]  # [contradiction, entailment, neutral]
wall=time.time()-t_all
# judge mode A: the proposed pair alone, verdict = argmax label is entailment
A_tp=A_tn=A_fp=A_fn=0
# judge mode B: rerank, verdict = proposed == argmax entailment over 12
B_tp=B_tn=B_fp=B_fn=0; top1=0; missedA=Counter(); missedB=Counter(); ent_pairs=[]
for ci,c in enumerate(cases):
    p=probs[(ci,c["label"])]; ent=p[1]; ent_pairs.append((ent,c["truth"]))
    vA = max(range(3), key=lambda j:p[j])==1
    best=max(INTENTS, key=lambda k:probs[(ci,k)][1]); vB = (best==c["label"])
    if best==c["real"]: top1+=1
    if c["truth"]:
        A_tp+=vA; A_fn+=(not vA); B_tp+=vB; B_fn+=(not vB)
    else:
        A_fp+=vA; A_tn+=(not vA); B_fp+=vB; B_tn+=(not vB)
        if vA: missedA[f"{c['real']} labelled as {c['label']}"]+=1
        if vB: missedB[f"{c['real']} labelled as {c['label']}"]+=1
# calibration on the proposed pair's entailment prob
bins=defaultdict(lambda:[0,0])
for ent,t in ent_pairs:
    b=min(9,int(ent*10)); bins[b][0]+=1; bins[b][1]+=t
ece=sum(n*abs((c/n)-((b+0.5)/10)) for b,(n,c) in bins.items())/len(ent_pairs)
calib=[{"bin":f"{b/10:.1f}-{(b+1)/10:.1f}","n":n,"observed_correct_pct":round(100*c/n,1)} for b,(n,c) in sorted(bins.items())]
res={"cases":len(cases),"pairs":len(pairs),"failed_batches":failed,"tokens_in":tokens,"cost_usd":round(tokens*0.05/1e6,4),
 "wall_s":round(wall,1),"batch_latency_p50_s":round(sorted(lats)[len(lats)//2],2),"verdicts_per_hour_rerank":round(len(cases)/wall*3600),
 "A_single_pair":{"wrong_caught_pct":round(100*A_tn/(A_tn+A_fp),2),"correct_accepted_pct":round(100*A_tp/(A_tp+A_fn),2),"accuracy_pct":round(100*(A_tp+A_tn)/len(cases),2),"tp":A_tp,"tn":A_tn,"fp":A_fp,"fn":A_fn,"missed_top":missedA.most_common(6)},
 "B_rerank_12":{"wrong_caught_pct":round(100*B_tn/(B_tn+B_fp),2),"correct_accepted_pct":round(100*B_tp/(B_tp+B_fn),2),"accuracy_pct":round(100*(B_tp+B_tn)/len(cases),2),"tp":B_tp,"tn":B_tn,"fp":B_fp,"fn":B_fn,"missed_top":missedB.most_common(6)},
 "classification_top1_pct":round(100*top1/len(cases),2),"calibration_ece":round(ece,4),"calibration":calib}
print(json.dumps(res,indent=1))
json.dump({"summary":res,"probs":{f"{ci}|{k}":v for (ci,k),v in probs.items()},"cases":cases},open(os.environ.get("OUT","openjev_990.json"),"w"))