import json, numpy as np, time
N=13; P=np.uint64(1000003)
POW=[P**np.uint64(N-1-j) for j in range(N)]
def hashes(arr):
    arr=np.asarray(arr,dtype=np.uint64); n=len(arr)-N+1
    if n<=0: return np.zeros(0,np.uint64)
    h=np.zeros(n,np.uint64)
    for j in range(N): h+=arr[j:j+n]*POW[j]
    return h
D='package/data/'; M='STASIS_12269'
items={}
def seqs(bench):
    for r in map(json.loads,open(D+bench+'.jsonl')):
        if bench.startswith('WT2') or bench=='LAMBADA_OPENAI':
            if r['ids'] is None: continue
            ids=r['ids'][M]; yield r['unit'], (ids['tgt'] if bench.startswith('WT2') else ids['ctx']+ids['tgt'])
        else:
            c=r['choices'][r['gold']]['ids'][M]; yield r['unit'], c['ctx']+c['tgt']
allh=[]
for b in ["WT2_FULL_BPB","WT2_PANEL256_CONTINUITY","LAMBADA_OPENAI","HELLASWAG","PIQA","ARC_EASY"]:
    per={}
    for u,s in seqs(b): per.setdefault(u,[]).append(hashes(s))
    items[b]={u:np.unique(np.concatenate(v)) for u,v in per.items()}
    allh+=list(items[b].values())
allset=np.unique(np.concatenate(allh)); print('eval 13-grams',len(allset),flush=True)
tb=np.memmap('/workspace/stasis/sixth/package/data/train.bin',dtype=np.uint16,mode='r'); T=len(tb); print('train tokens',T,flush=True)
hit=np.zeros(0,np.uint64); C=20_000_000; t0=time.time()
for s in range(0,T,C):
    h=hashes(tb[s:min(T,s+C+N-1)]); hit=np.union1d(hit,np.intersect1d(h,allset)); print(s,time.time()-t0,flush=True)
res={'method':'13-gram (Stasis token ids) rolling-hash overlap vs sixth-tranche train.bin','train_bin_sha256':'f64bdbcf21f735c0d542142f5d4ebf6a186178b4459d810a7b01b97304876e35','train_tokens':int(T),'benchmarks':{}}
for b,d in items.items():
    tot=sum(len(v) for v in d.values()); sh=sum(int(np.isin(v,hit).sum()) for v in d.values())
    ni=sum(1 for v in d.values() if len(v)); ci=sum(1 for v in d.values() if len(v) and np.isin(v,hit).any())
    res['benchmarks'][b]={'items_with_13grams':ni,'items_with_any_shared_13gram':ci,'item_fraction':ci/max(ni,1),'ngrams':tot,'shared_ngrams':sh,'ngram_fraction':sh/max(tot,1)}
res['scope_limit']='sixth tranche only; tranches 1-5 and GPT-2 WebText not checked; 64-bit hash collisions negligible; descriptive only'
json.dump(res,open('package/CONTAMINATION_13GRAM.json','w'),indent=1); print(json.dumps(res['benchmarks'],indent=1))