import numpy as np, json, urllib.request
SP="/private/tmp/claude-501/-Users-timlingo/82369039-a20e-4b5a-8a5e-28234a57b996/scratchpad"
M=np.load(SP+'/emb.npy'); ids=open(SP+'/ids.txt',encoding='utf-8',errors='surrogateescape').read().split('\n')
np.seterr(all='ignore')
bad=~np.isfinite(M).all(axis=1)
M[bad]=0.0
print("non-finite rows zeroed:",int(bad.sum()))
idx={k:i for i,k in enumerate(ids)}
gold=json.load(open("/Users/timlingo/Development/neuron-technologies/_wt-assoc-leg/tools/retrieval-eval/gold_set.json"))['queries']
def emb(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)
    return v/(np.linalg.norm(v)+1e-9)
out={}
for q in gold:
    v=emb(q['query']); s=M@v
    s=s[np.isfinite(s)]
    mu=float(s.mean()); sd=float(s.std())
    top=np.sort(s)[::-1][:10]
    z=[(float(t)-mu)/sd for t in top]
    grank=[]
    for rel in q['relevant']:
        if rel in idx:
            j=idx[rel]; grank.append((int((M@v > (M@v)[j]).sum())+1, round(float((M@v)[j]),3)))
    grank.sort()
    out[q['id']]=dict(cat=q['category'],mu=round(mu,3),sd=round(sd,4),top1=round(float(top[0]),3),
                      z1=round(z[0],2),z3=round(z[2],2),z5=round(z[4],2),gold=grank[:1])
    print("%s %-11s mu=%.3f sd=%.4f top1=%.3f z1=%5.2f z3=%5.2f z5=%5.2f gold=%s"%(
        q['id'],q['category'],mu,sd,top[0],z[0],z[2],z[4],grank[:1]))
json.dump(out,open(SP+'/zprobe.json','w'),indent=1)