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 wstart(s,tok):
    i=s.find(tok)
    while i!=-1:
        if i==0 or not s[i-1].isalnum(): return True
        i=s.find(tok,i+1)
    return False
def legs4(query, wordstart=False, unfloor=False):
    toks=tokenize(query); lt=[t.lower() for t in toks]
    hit_idx=[];hit_mask=[];df=[0]*len(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 and (not wordstart or wstart(s,tok)): 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]
    Sl=[int(i) for i in order if cos[i]>(0.0 if unfloor else SEED_MIN)]
    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 or 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=sorted([(i,float(cos[i])) for i,st in seen.items() if st==1 and OK[i] and HAVE[i] and cos[i]>0.0],key=lambda x:-x[1])[:AMAX]
    return [i for i,_,_ in L],Sl,[i for i,_ in 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(*legs4(q['query'],**kw),10)]) for q in G['queries']}
base=run()
print("baseline",sum(base.values()),"misses",[k for k,v in base.items() if not v])
for name,kw in [('wordstart',dict(wordstart=True)),
                ('unfloor',dict(unfloor=True)),
                ('wordstart+unfloor',dict(wordstart=True,unfloor=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("%-20s net=%+d gains=%s losses=%s"%(name,len(g)-len(l),g,l))