"""Diff a model rerun (*.haiku.json) against the Opus rerun it was built from."""
import json, os, re, glob, sys, psycopg, collections
S=os.path.dirname(os.path.abspath(__file__)); SUFFIX=sys.argv[1] if len(sys.argv)>1 else "haiku"
conn=psycopg.connect(host=os.environ["POSTGRES_HOST"],user="quber",password=os.environ["POSTGRES_PASSWORD"],dbname="quber_rag")
titles={r[0]:r[1] for r in conn.execute("select doc_key, title from ade_playground.documents").fetchall()}
def norm(s): return re.sub(r"[\s$,%]|million|thousand|x$","",(s or "").lower()).replace("—","-").replace("–","-").replace("(","-").replace(")","")
def cell_text(doc_id, ref):
    r=conn.execute("select cell_text from ade_playground.groundings where document_id=%s and ref_id=%s",(doc_id,ref)).fetchone(); return r[0] if r else None
G=collections.Counter(); diffs=[]
print(f"{'document':30s} {'opus scalar':>11s} {'haiku scalar':>12s} {'both scalar':>11s} {'same':>5s} {'diff':>5s} {'opus only':>9s} {'haiku only':>10s} {'errors':>6s} | cell-verified opus/haiku")
for f in sorted(glob.glob(f"{S}/rerun_*.{SUFFIX}.json")):
    d=json.load(open(f)); doc_id=d["doc_id"]; dk=d["doc_key"]; c=collections.Counter(); vo=vh=0
    for r in d["rows"]:
        o=r["opus"]; hs=r["shape"]; os_=o["shape"]
        c[(os_,hs)]+=1
        if os_=="scalar" and hs=="scalar":
            ov=o["payload"].get("value"); hv=r["payload"].get("value")
            if norm(ov)==norm(hv): c["same"]+=1
            else:
                c["diff"]+=1
                oc=o["cited_ids"][0] if o["cited_ids"] else None; hc=r["cited_ids"][0] if r["cited_ids"] else None
                diffs.append((titles.get(dk,dk)[:22], r["index"], r["question"][:70], ov, oc, cell_text(doc_id,oc) if oc else None, hv, hc, cell_text(doc_id,hc) if hc else None))
        for who,rec,key in (("o",o,"vo"),("h",r,"vh")):
            if rec["shape"]=="scalar":
                cid=rec["cited_ids"][0] if rec["cited_ids"] else None; t=cell_text(doc_id,cid) if cid else None
                if t is not None and norm(t)==norm(rec["payload"].get("value")):
                    if who=="o": vo+=1
                    else: vh+=1
    osc=sum(v for k,v in c.items() if isinstance(k,tuple) and k[0]=="scalar"); hsc=sum(v for k,v in c.items() if isinstance(k,tuple) and k[1]=="scalar")
    print(f"{titles.get(dk,dk)[:30]:30s} {osc:>11d} {hsc:>12d} {c[('scalar','scalar')]:>11d} {c['same']:>5d} {c['diff']:>5d} {c[('scalar','unanswerable')]:>9d} {c[('unanswerable','scalar')]:>10d} {sum(v for k,v in c.items() if isinstance(k,tuple) and k[1]=='error'):>6d} | {vo}/{osc} vs {vh}/{hsc}")
    for k,v in c.items(): G[k]+=v
print("\nALL:", {(f"{k[0]}->{k[1]}" if isinstance(k,tuple) else k):v for k,v in G.items()})
print(f"\nValue disagreements ({len(diffs)}): opus value [cite -> cell text] vs haiku value [cite -> cell text]")
for t,i,q,ov,oc,ot,hv,hc,ht in diffs:
    flag = "OPUS=cell" if ot is not None and norm(ot)==norm(ov) else ("HAIKU=cell" if ht is not None and norm(ht)==norm(hv) else "")
    if ot is not None and norm(ot)==norm(ov) and ht is not None and norm(ht)==norm(hv): flag="BOTH=own cell"
    print(f"  {t} [{i}] {q!r}\n      opus {ov!r} [{oc} -> {ot!r}]   haiku {hv!r} [{hc} -> {ht!r}]   {flag}")