"""Exploratory structural tests with fixed-layout, within-page shuffle controls."""
import json, math
from collections import Counter,defaultdict
from word_core import *

def collision_pairs(items):return sum(n*(n-1)//2 for n in Counter(items).values())

def statistics(v,words):
    arrays=[v[w['rune_offset']:w['rune_offset']+len(w['values'])] for w in words]
    exact=[tuple(a) for a in arrays if len(a)>=4]
    shifted=[tuple((a-a[0])%29) for a in arrays if len(a)>=5]
    phrases=[]
    for i in range(len(words)-1):
        if words[i]['section']==words[i+1]['section'] and words[i]['paragraph']==words[i+1]['paragraph'] and len(arrays[i])+len(arrays[i+1])>=6:
            phrases.append((tuple(arrays[i]),tuple(arrays[i+1])))
    repeatpairs=sum(collision_pairs(a.tolist()) for a in arrays)
    # Mutual information rune identity vs beginning/interior/end/single position.
    table=np.zeros((4,29),dtype=np.int64)
    for a in arrays:
        if len(a)==1:table[3,a[0]]+=1
        else:
            table[0,a[0]]+=1;table[2,a[-1]]+=1
            np.add.at(table[1],a[1:-1],1)
    exp=table.sum(1)[:,None]*table.sum(0)[None,:]/table.sum()
    nz=table>0
    mi=float(np.sum(table[nz]/table.sum()*np.log(table[nz]/exp[nz])))
    return {'exact_word_pairs_len4plus':collision_pairs(exact),'constant_shift_word_pairs_len5plus':collision_pairs(shifted),'exact_two_word_phrase_pairs_6plus_runes':collision_pairs(phrases),'within_word_equal_rune_pairs':repeatpairs,'rune_boundary_position_MI':mi}

def benchmark():
    rng=np.random.default_rng(33012026);report=[]
    for heldout,source in GROUPS.items():
        model=Model([s for name,s in GROUPS.items() if name!=heldout])
        words=tokenize(source);v=vals(source);lengths=[len(w['values']) for w in words]
        # All 812 invertible affine transforms, including identity, plus128 shuffles.
        dec=np.array([(a*v+b)%29 for a in range(1,29) for b in range(29)])
        shuffles=np.array([rng.permutation(v) for _ in range(128)])
        candidates=np.concatenate((dec,shuffles))
        item={'held_out_group':heldout,'runes':len(v),'candidates':len(candidates)}
        for name,old in [('original',True),('boundary_backoff',False)]:
            sc=model.score(candidates,lengths,old)
            item[name]={'identity_rank':1+int(np.sum(sc>sc[0])),'plaintext_score':float(sc[0]),'best_wrong_score':float(sc[1:].max())}
        # Short 64-ish rune windows, at complete-word boundaries; >=32 runes.
        short=[];cur=[];size=0
        for w in words:
            cur.append(w['values']);size+=len(w['values'])
            if size>=64:short.append(cur);cur=[];size=0
        if size>=32:short.append(cur)
        for name,old in [('original',True),('boundary_backoff',False)]:
            ranks=[]
            for chunk in short:
                vv=np.array([x for w in chunk for x in w]);ll=[len(w) for w in chunk]
                ds=np.array([(a*vv+b)%29 for a in range(1,29) for b in range(29)])
                sc=model.score(ds,ll,old);ranks.append(1+int(np.sum(sc>sc[0])))
            item[name]['short_window_ranks']=ranks
        report.append(item)
    # Additional true LP2 plaintexts kept out of all training.
    v=vals(lp[56]);ix=np.arange(len(v));ix[57:]-=1
    d=(v-SEQS['prime_minus_one'][ix])%29;d[56]=v[56]
    for label,source in [('LP2_56',replace_values(lp[56],d)),('LP2_57',lp[57])]:
        v=vals(source);ll=[len(w['values']) for w in tokenize(source)]
        candidates=np.array([(a*v+b)%29 for a in range(1,29) for b in range(29)])
        item={'held_out_group':label,'runes':len(v),'candidates':len(candidates)}
        for name,old in [('original',True),('boundary_backoff',False)]:
            sc=FULL_MODEL.score(candidates,ll,old);item[name]={'identity_rank':1+int(np.sum(sc>sc[0])),'plaintext_score':float(sc[0]),'best_wrong_score':float(sc[1:].max())}
        report.append(item)
    return report

if __name__=='__main__':
    out={'benchmark':benchmark()};print('Scoring benchmark complete',flush=True)
    source='%'.join(lp[:56]);words=tokenize(source);v=vals(source)
    # Exclude the standalone nine-rune section from counts? Keep it in mapping
    # and statistics, so all12956 runes and all words are accounted for.
    groups=defaultdict(list)
    for w in words:groups[w['runes']].append({k:w[k] for k in ['page','section','paragraph','word_index','rune_offset','page_rune_offset']})
    repeats=[{'runes':r,'transliteration':' '.join(LET[R.index(c)] for c in r),'length':len(r),'occurrences':loc} for r,loc in groups.items() if len(loc)>1]
    repeats.sort(key=lambda x:(-x['length'],-len(x['occurrences']),x['runes']))
    phrases=defaultdict(list)
    for n in (2,3,4):
        for i in range(len(words)-n+1):
            chunk=words[i:i+n]
            if chunk[0]['section']!=chunk[-1]['section'] or chunk[0]['paragraph']!=chunk[-1]['paragraph']:continue
            key=tuple(w['runes'] for w in chunk)
            if sum(map(len,key))>=5:phrases[key].append({k:chunk[0][k] for k in ['page','section','word_index','page_rune_offset']})
    repeated_phrases=[{'words':list(k),'runes':sum(map(len,k)),'occurrences':loc} for k,loc in phrases.items() if len(loc)>1]
    repeated_phrases.sort(key=lambda x:(-x['runes'],-len(x['words'])))
    out.update({'runes':len(v),'word_count':len(words),'word_length_counts':dict(sorted(Counter(len(w['values']) for w in words).items())),'repeated_words':repeats,'repeated_phrases':repeated_phrases,'words':words})
    observed=statistics(v,words);out['observed']=observed
    # Preserve exact per-page rune histogram, layout and word lengths.
    rng=np.random.default_rng(20260919);lengths=[len(vals(s)) for s in lp[:56]];starts=np.r_[0,np.cumsum(lengths)]
    null=[]
    for i in range(999):
        sv=v.copy()
        for a,b in zip(starts,starts[1:]):sv[a:b]=rng.permutation(sv[a:b])
        null.append(statistics(sv,words))
        if (i+1)%100==0:print('Structural shuffles',i+1,flush=True)
    out['shuffle_config']={'count':999,'seed':20260919,'preserve':'Within each page: rune inventory; globally: all token boundaries, lengths and other layout','interpretation':'Exploratory, one-sided upper-tail; five-statistic Bonferroni correction. Previously known repeated phrase motivated this investigation, so this is not independent confirmation.'}
    out['null_comparisons']={}
    for name,x in observed.items():
        arr=np.array([n[name] for n in null]);p=(1+int(np.sum(arr>=x)))/1000
        out['null_comparisons'][name]={'observed':x,'null_mean':float(arr.mean()),'null_95th':float(np.quantile(arr,.95)),'null_max':float(arr.max()),'empirical_p':p,'bonferroni_p_5':min(1,5*p)}
    out['null_statistics']=null
    (ROOT/'word_diagnostics_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2))
    print(json.dumps({'benchmark':out['benchmark'],'repeated_phrases':repeated_phrases,'comparisons':out['null_comparisons']},ensure_ascii=False,indent=2),flush=True)