import json,sys,pickle,numpy as np
sys.path.insert(0,'.')
from legs import *
GP='/Users/timlingo/Development/neuron-technologies/_wt-bm25lex/tools/retrieval-eval/'
G=json.load(open(GP+'gold_set.json'))
def legs3(query, sem_sal=False, assoc_sal=False, unfloor=False, sem_cap=None):
    toks=tokenize(query)
    hit_idx=[];hit_mask=[];df=[0]*len(toks);lt=[t.lower() for t in toks]
    for i in range(N):
        if not OK[i]: continue
        s=LOW[i];m=0
        for t,tok in enumerate(lt):
            if tok in s: m|=(1<<t)
        if m:
            hit_idx.append(i);hit_mask.append(m)
            for t in range(len(toks)):
                if m>>t&1: df[t]+=1
    dl_n=int(OK.sum());avgdl=float(DL[OK].sum()/max(dl_n,1))
    idf=[math.log(1.0+((dl_n-d+0.5)/(d+0.5))) for d in df]
    L=[]
    for j,i in enumerate(hit_idx):
        norm=1.0-B+B*(DL[i]/avgdl);w=0.0
        for t in range(len(toks)):
            if hit_mask[j]>>t&1: w+=idf[t]*(K1+1.0)/(1.0+K1*norm)
        L.append((i,w,SAL[i]))
    L.sort(key=lambda x:(-x[1],-x[2]))
    if not L: return [],[],[]
    qv=qemb(query);cos=En@qv;cos=np.where(HAVE&OK,cos,-2.0)
    order=np.argsort(-cos)[:600]
    cand=[int(i) for i in order if cos[i]>(0.0 if unfloor else SEED_MIN)]
    key=(lambda i:(SAL[i] if sem_sal else 1.0)*float(cos[i]))
    Sl=sorted(cand,key=lambda i:-key(i))
    if sem_cap: Sl=Sl[:sem_cap]
    semseed=[int(i) for i in order[:SEED_K] if cos[i]>0.0]
    act={};seen={};qq=[]
    for i,_,_ in L[:ASSOC_SEEDS]:
        act[i]=1.0;seen[i]=2;qq.append((i,0))
    for i in semseed:
        if i in seen: continue
        act[i]=1.0;seen[i]=2;qq.append((i,0))
    qh=0
    while qh<len(qq):
        cur,h=qq[qh];qh+=1
        if h>=DEPTH: continue
        parent=act[cur]
        for e,oi in ADJ_F[cur]+ADJ_T[cur]:
            if e['rel'] not in STRUCT: continue
            if EXCL[oi]: continue
            na=parent*e['w']*DECAY*SAL[oi]
            if na<FIRE: continue
            if seen.get(oi) and na<=act.get(oi,0): continue
            act[oi]=na
            if not seen.get(oi): seen[oi]=1
            if len(qq)<AMAX*4: qq.append((oi,h+1))
    A=[]
    for i,st in seen.items():
        if st!=1 or not OK[i] or not HAVE[i]: continue
        c=float(cos[i])
        if c<=0.0: continue
        A.append((i,(SAL[i] if assoc_sal else 1.0)*c))
    A.sort(key=lambda x:-x[1]);A=[i for i,_ in A[:AMAX]]
    return [i for i,_,_ in L],Sl,A
def merge(L,S,A,lim=10):
    out=[];li=si=ai=0
    while len(out)<lim and (li<len(L) or si<len(S) or ai<len(A)):
        if li<len(L):
            if L[li] not in out: out.append(L[li])
            li+=1
        if len(out)>=lim: break
        if si<len(S):
            if S[si] not in out: out.append(S[si])
            si+=1
        if len(out)>=lim: break
        if ai<len(A):
            if A[ai] not in out: out.append(A[ai])
            ai+=1
    return out
def outcome(q,ids):
    c=q['category']
    if c=='nonsense': return len(ids)==0
    if c=='superseded':
        a,b=q['must_outrank']
        if a not in ids: return False
        if b not in ids: return True
        return ids.index(a)<ids.index(b)
    return any(x in ids[:5] for x in q['relevant'])
def run(**kw):
    return {q['id']:outcome(q,[NODES[i]['id'] for i in merge(*legs3(q['query'],**kw),10)]) for q in G['queries']}
base=run()
print("baseline",sum(base.values()),"/38  misses:",[k for k,v in base.items() if not v])
import itertools
for name,kw in [
  ('sem_sal(floored)',dict(sem_sal=True)),
  ('unfloor',dict(unfloor=True)),
  ('unfloor+sem_sal',dict(unfloor=True,sem_sal=True)),
  ('assoc_sal',dict(assoc_sal=True)),
  ('unfloor+sem_sal+assoc_sal',dict(unfloor=True,sem_sal=True,assoc_sal=True)),
  ('sem_sal+assoc_sal(floored)',dict(sem_sal=True,assoc_sal=True)),
]:
    r=run(**kw)
    g=sorted(k for k in base if r[k] and not base[k]); l=sorted(k for k in base if base[k] and not r[k])
    print("%-28s net=%+d gains=%s losses=%s"%(name,len(g)-len(l),g,l))