import numpy as np, json, urllib.request, collections, math, re, sys, time
SP="/private/tmp/claude-501/-Users-timlingo/82369039-a20e-4b5a-8a5e-28234a57b996/scratchpad"
EV="/Users/timlingo/Development/neuron-technologies/_wt-bm25lex/tools/retrieval-eval/"
np.seterr(all='ignore'); t0=time.time()
M=np.load(SP+'/emb.npy'); eids=open(SP+'/ids.txt',encoding='utf-8',errors='surrogateescape').read().split('\n')
eidx={k:i for i,k in enumerate(eids)}
d=json.load(open('/Users/timlingo/neuron-memory-backups/snapshot-pre-repair-20260806.json',encoding='utf-8',errors='surrogateescape'))
STRUCT={"identity","contains","superseded_by","references","embodies","demonstrated_by","canonical-self","depends_on","currently_holds","activates"}
adj=collections.defaultdict(list)
for e in d['edges']:
    if e.get('relation') not in STRUCT: continue
    w=float(e.get('weight') or 0.0)
    adj[e['from_id']].append((e['to_id'],w)); adj[e['to_id']].append((e['from_id'],w))
nodes=d['nodes']; N={n['id']:n for n in nodes}
PRINT=re.compile(r'^[\x20-\x7e]+$')
ids=[];hay=[];dl=[];sal=[];addr=[]
for n in nodes:
    i=n.get('id') or ''
    h=((n.get('content') or '')+'\x00'+(n.get('label') or '')+'\x00'+(n.get('tags') or '')).lower()
    ids.append(i);hay.append(h);dl.append(len(h));sal.append(float(n.get('salience') or 0.0));addr.append(bool(PRINT.match(i)))
del d
NN=len(ids); avgdl=sum(dl)/NN
gold={q['id']:q for q in json.load(open(EV+"gold_set.json"))['queries']}
CACHE={}
def emb(t):
    if t in CACHE: return CACHE[t]
    b=json.dumps({"model":"nomic-embed-text","prompt":t}).encode()
    r=urllib.request.Request("http://127.0.0.1:11434/api/embeddings",data=b,headers={"Content-Type":"application/json"})
    v=np.array(json.load(urllib.request.urlopen(r,timeout=60))["embedding"],dtype=np.float32)
    v=v/(np.linalg.norm(v)+1e-9); CACHE[t]=v; return v
K1,B=1.2,0.75
LEXC={}
def lexleg(qid,query,lim=10):
    if qid in LEXC: return LEXC[qid]
    toks=[]
    for w in query.split():
        wl=w.lower()
        if wl not in toks: toks.append(wl)
    nt=len(toks); masks=[]; df=[0]*nt
    for i in range(NN):
        if not addr[i]: continue
        h=hay[i]; m=0; sc=0
        for t in range(nt):
            if toks[t] in h: m|=(1<<t); sc+=1; df[t]+=1
        if sc: masks.append((i,m))
    idf=[math.log(1.0+(NN-df[t]+0.5)/(df[t]+0.5)) for t in range(nt)]
    scored=[]
    for i,m in masks:
        norm=1.0-B+B*dl[i]/avgdl; s=0.0
        for t in range(nt):
            if m&(1<<t): s+=idf[t]*(K1+1.0)/(1.0+K1*norm)
        scored.append((s,i))
    scored.sort(key=lambda x:(-x[0],-sal[x[1]]))
    LEXC[qid]=([ids[i] for s,i in scored[:lim]], len(masks))
    return LEXC[qid]
FIRE=0.02; DECAY=0.7; DEPTH=2; SEED_MIN=0.60; ASSOC_MAX=64
def assoc(seeds, s, use_cos, order):
    act={x:1.0 for x in seeds}; seen={x:2 for x in seeds}
    Q=[(x,0) for x in seeds]; h=0
    while h<len(Q):
        cur,hop=Q[h]; h+=1
        if hop>=DEPTH: continue
        p=act[cur]
        for oid,w in adj.get(cur,()):
            n=N.get(oid)
            if not n or n.get('node_type') in ('Tag','InternalStateEvent'): continue
            c=1.0
            if use_cos:
                j=eidx.get(oid)
                c=max(0.0,float(s[j])) if j is not None else 0.0
            na=p*w*DECAY*float(n.get('salience') or 0.0)*c
            if na<FIRE: continue
            if oid in seen and na<=act.get(oid,0): continue
            act[oid]=na
            if oid not in seen: seen[oid]=1
            Q.append((oid,hop+1))
    out=[]
    for k,v in seen.items():
        if v!=1 or k not in eidx: continue
        c=float(s[eidx[k]])
        if c<=0: continue
        out.append((act[k] if order=='act' else c,k))
    out.sort(reverse=True)
    return [k for c,k in out[:ASSOC_MAX] if PRINT.match(k or '')]
def inter(legs,lim=10):
    out=[];idx=[0]*len(legs)
    while len(out)<lim and any(idx[i]<len(legs[i]) for i in range(len(legs))):
        for i in range(len(legs)):
            if idx[i]<len(legs[i]):
                if legs[i][idx[i]] not in out: out.append(legs[i][idx[i]])
                idx[i]+=1
            if len(out)>=lim: break
    return out
def run(floor, vocabgate, use_cos, order):
    res={}; legs={}
    for qid,q in gold.items():
        v=emb(q['query']); s=M@v; s[~np.isfinite(s)]=-1
        L,nmatch=lexleg(qid,q['query'])
        if vocabgate and nmatch==0:
            res[qid]=[]; legs[qid]=([],[],[]); continue
        ordr=np.argsort(-s)
        S=[eids[j] for j in ordr[:10] if PRINT.match(eids[j] or '') and (not floor or s[j]>SEED_MIN)]
        seeds=[x for x in L[:3] if x in N]
        seeds=seeds+[eids[j] for j in ordr[:8] if eids[j] in N and eids[j] not in seeds and PRINT.match(eids[j] or '')]
        A=assoc(seeds,s,use_cos,order) if seeds else []
        res[qid]=inter([L,S,A]); legs[qid]=(L,S,A)
    return res,legs
def score(res,label,base=None):
    det={}
    for qid,q in gold.items():
        out=res[qid][:5]
        if q['category']=='nonsense': ok=(len(res[qid])==0)
        elif q['category']=='superseded':
            must=q.get('must_outrank') or {}; ok=False
            for good,bad in (must.items() if isinstance(must,dict) else []):
                ok = good in res[qid] and (bad not in res[qid] or res[qid].index(good)<res[qid].index(bad))
            if not must: ok=any(r in out for r in q['relevant'])
        else: ok=any(r in out for r in q['relevant'])
        det[qid]=ok
    line="%-34s true=%d/38"%(label,sum(det.values()))
    if base is not None:
        dd=[q for q in sorted(gold) if det[q]!=base[q]]
        line+="  moved=%d gains=%s losses=%s"%(len(dd),[q for q in dd if det[q]],[q for q in dd if not det[q]])
    print(line, flush=True)
    return det
if __name__=="__main__":
    b,_=run(True,False,False,'cos'); base=score(b,'BASE bm25lex replica')
    for lab,args in [
        ("A floor-off+vocabgate",      (False,True,False,'cos')),
        ("B A+cos-in-traversal",       (False,True,True ,'cos')),
        ("C A+cos-trav+act-order",     (False,True,True ,'act')),
        ("D floor-off NO gate",        (False,False,False,'cos')),
    ]:
        r,_=run(*args); score(r,lab,base)
    print("elapsed %.1fs"%(time.time()-t0),file=sys.stderr)