"""Session 51: structural check of ideal uniform-pad key skipping.
Monte Carlo is exploratory, using a previously noticed phrase statistic.
"""
from pathlib import Path
from collections import Counter
import hashlib,json
import numpy as np
from lp_core import vals,lp
from lexical_words import tokenize_words
ROOT=Path(__file__).resolve().parent
SOURCE='%'.join(lp[:56]);ORIGINAL=vals(SOURCE).astype(np.uint8)
LENGTHS=[len(vals(s)) for s in SOURCE.split('$') if len(vals(s))]
STARTS=np.r_[0,np.cumsum(LENGTHS)[:-1]].astype(int)
WORDS=tokenize_words(SOURCE)
WINDOWS=[]
for a,b in zip(WORDS,WORDS[1:]):
    na,nb=len(a['values']),len(b['values'])
    if a['section']!=b['section'] or a['paragraph']!=b['paragraph'] or na+nb<6:continue
    assert a['rune_offset']+na==b['rune_offset']
    WINDOWS.append((a['rune_offset'],na,nb))
# Shapes occurring once cannot collide and may be omitted without changing counts.
SHAPES=Counter((a,b) for _,a,b in WINDOWS)
ACTIVE=[(start,a,b) for start,a,b in WINDOWS if SHAPES[a,b]>1]
SIX_STARTS=np.concatenate([np.arange(s,s+n-5) for s,n in zip(STARTS,LENGTHS) if n>=6])

def phrase_count(v):
    counts=Counter((a,b,bytes(v[s:s+a+b])) for s,a,b in ACTIVE)
    return sum(n*(n-1)//2 for n in counts.values())

def six_count(v):
    codes=np.zeros(len(SIX_STARTS),dtype=np.int64)
    for d in range(6):codes=codes*29+v[SIX_STARTS+d]
    _,counts=np.unique(codes,return_counts=True)
    return int(np.sum(counts*(counts-1)//2))

def simulate(repeat_probability,n,seed):
    rng=np.random.default_rng(seed)
    v=np.empty((n,len(ORIGINAL)),dtype=np.uint8);starts=set(map(int,STARTS))
    for i in range(v.shape[1]):
        if i in starts:v[:,i]=rng.integers(0,29,n,dtype=np.uint8)
        else:
            # Choose previous symbol with q, each of the other 28 with (1-q)/28.
            old=v[:,i-1];candidate=rng.integers(0,28,n,dtype=np.uint8)
            v[:,i]=np.where(rng.random(n)<repeat_probability,old,candidate+(candidate>=old))
    pairs=[];sixes=[];repeats=[]
    mask=np.ones(len(ORIGINAL)-1,dtype=bool);mask[STARTS[1:]-1]=False
    for row in v:
        pairs.append(phrase_count(row));sixes.append(six_count(row));repeats.append(int(np.sum((row[1:]==row[:-1])&mask)))
    return {'seed':seed,'runs':n,'repeat_probability':repeat_probability,'phrase_pairs':pairs,'six_rune_pairs':sixes,'adjacent_repeats':repeats,'runs_with_phrase_pair':sum(x>=1 for x in pairs),'runs_with_six_rune_pair':sum(x>=1 for x in sixes),'repeat_count_mean':float(np.mean(repeats)),'repeat_count_95pct_interval':np.quantile(repeats,[.025,.975]).tolist()}

def main():
    mask=np.ones(len(ORIGINAL)-1,dtype=bool);mask[STARTS[1:]-1]=False
    obs=int(np.sum((ORIGINAL[1:]==ORIGINAL[:-1])&mask));eligible=int(mask.sum())
    # Independent tiny direct-count oracle, and a planted full-word phrase.
    toy=np.array([1,2,3,4,5,6,1,2,3,4,5,6],dtype=np.uint8)
    assert Counter([bytes(toy[:6]),bytes(toy[6:])])[bytes(toy[:6])]==2
    rng=np.random.default_rng(5100);control=rng.integers(0,29,len(ORIGINAL),dtype=np.uint8)
    spots=[s for s,a,b in ACTIVE if (a,b)==(3,3)]
    a=spots[0];b=next(s for s in spots if s>=a+6)
    control[a:a+6]=[23,11,1,17,18,10];control[b:b+6]=control[a:a+6]
    assert phrase_count(control)>=1
    out={'source_sha256':hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest(),'runes':len(ORIGINAL),'section_lengths':LENGTHS,'eligible_adjacent_positions':eligible,'observed_adjacent_repeats':obs,'eligible_two_word_spans':len(WINDOWS),'eligible_six_rune_windows':len(SIX_STARTS),'observed_phrase_pairs':phrase_count(ORIGINAL),'observed_six_rune_pairs':six_count(ORIGINAL),'control':{'planted_offsets':[a,b],'phrase_pairs':phrase_count(control)},'models':{}}
    for name,q,seed in [('fixed_suppression_082',.18/(29-.82),51001),('fitted_repeat_rate',obs/eligible,51002)]:
        out['models'][name]=simulate(q,4999,seed)
        print(name,{k:v for k,v in out['models'][name].items() if k not in ['phrase_pairs','six_rune_pairs','adjacent_repeats']},flush=True)
    (ROOT/'keyskip_phrase_results.json').write_text(json.dumps(out,indent=2)+'\n')
    print({k:v for k,v in out.items() if k!='models'})
if __name__=='__main__':main()