"""Validate the v2 calibrated gate against ALL MEME jobs in /data/ with a result + tree."""
import os, glob, json, csv, pickle, re, numpy as np
MEME="/data/datamonkey/site/results/meme"
D="/home/sweaver/programming/axomeme/dm3-gate-validation"
g=pickle.load(open(f"{D}/meme_gate_v2_calibrated.pkl","rb")); cal=g["model"]

def med_bl_from_tre(p):
    try:
        nwk=open(p).read()
        bls=[float(x) for x in re.findall(r":(-?\d+\.?\d*(?:[eE][-+]?\d+)?)", nwk)]
        pos=[b for b in bls if b>0]
        return float(np.median(pos)) if pos else 0.0
    except: return None

def sites_at_p(jsonpath, thresh):
    try:
        d=json.load(open(jsonpath))
        hdr=d["MLE"]["headers"]; content=d["MLE"]["content"]
        pcol=None
        for i,h in enumerate(hdr):
            nm=(h[0] if isinstance(h,(list,tuple)) else str(h)).lower()
            if "p-value" in nm or "p_value" in nm or nm=="p": pcol=i;break
        rows=[]
        for k in sorted(content,key=lambda x:int(x)): rows.extend(content[k])
        ns=int(d.get("input",{}).get("number of sequences",0) or 0)
        nsit=int(d.get("input",{}).get("number of sites",0) or 0)
        if pcol is None: return None,ns,nsit
        cnt=sum(1 for r in rows if pcol<len(r) and r[pcol] is not None and float(r[pcol])<=thresh)
        return cnt,ns,nsit
    except: return None,0,0

jobs=[os.path.basename(f)[:-10] for f in glob.glob(f"{MEME}/*.MEME.json")]
out=open(f"{D}/full_validation.tsv","w"); out.write("job\tnseq\tnsites\tmed_bl\tscore\tsites_found_p10\tin_domain\n")
n=0;scored=0
for j in jobs:
    jp=f"{MEME}/{j}.MEME.json"; tp=f"{MEME}/{j}.tre"
    if not os.path.exists(tp): continue
    cnt,ns,nsit=sites_at_p(jp,0.1)
    if cnt is None or ns<=0 or nsit<=0: continue
    bl=med_bl_from_tre(tp)
    if bl is None: continue
    blc=min(max(bl,0.001),10.0)
    s=float(cal.predict_proba([[ns,nsit,blc]])[0,1])
    out.write(f"{j}\t{ns}\t{nsit}\t{bl:.6f}\t{s:.4f}\t{cnt}\t1\n"); scored+=1
    n+=1
    if n%1000==0: print(f"  {n} scored...",flush=True)
out.close(); print(f"DONE: {scored} jobs scored -> full_validation.tsv",flush=True)