"""Session 44: predeclared exact A/I constraint tests, no language-score fitting.
Python + NumPy. Run from any directory. Writes JSON beside this script.
"""
from pathlib import Path
import json, hashlib, math
import numpy as np
from lexical_words import tokenize_words
from word_core import lp, pages, GROUPS, R, LET, vals
F=Path(__file__).resolve().parent
FAMILIES=['constant','linear_rune','linear_word','period2_rune','period3_rune']
AI=[R.index('ᚪ'),R.index('ᛁ')]
SIGNS=[1,-1]

def records(words):
    out=[];pos=0
    for i,w in enumerate(words):
        if len(w['values'])==1:out.append(dict(page=w['page'],page_offset=w['page_rune_offset'],rune=pos,word=i,c=int(w['values'][0])))
        pos+=len(w['values'])
    return out

def fit(rs,family,sign):
    n=len(rs);cut=(2*n+2)//3
    c=np.array([r['c'] for r in rs]);t=np.array([r['word'] if family=='linear_word' else r['rune'] for r in rs])
    if family in ['constant','linear_rune','linear_word']:
        keys=[(0,b) for b in range(29)] if family=='constant' else [(a,b) for a in range(29) for b in range(29)]
        ab=np.array(keys);p=(sign*c[None,:]-ab[:,0,None]*t[None,:]-ab[:,1,None])%29
        good=np.isin(p,AI);train=np.flatnonzero(good[:,:cut].all(1));both=np.flatnonzero(good.all(1))
        return dict(train=len(train),full=len(both),unseen_phases=0,witness=list(keys[int(both[0])]) if len(both) else None,cut=cut)
    period=int(family[6]);options=[set(range(29)) for _ in range(period)]
    for r in rs[:cut]:options[r['rune']%period]&={(sign*r['c']-p)%29 for p in AI}
    train=math.prod(map(len,options));unseen=sum(not any(r['rune']%period==j for r in rs[:cut]) for j in range(period))
    for r in rs[cut:]:options[r['rune']%period]&={(sign*r['c']-p)%29 for p in AI}
    full=math.prod(map(len,options))
    return dict(train=train,full=full,unseen_phases=unseen,witness=[min(o) for o in options] if full else None,cut=cut)

def key_at(family,key,rune,word):
    if family in ['constant','linear_rune','linear_word']:return (key[0]*(word if family=='linear_word' else rune)+key[1])%29
    return key[rune%len(key)]

def decrypt(words,family,sign,key):
    out=[];pos=0
    for wi,w in enumerate(words):
        v=[]
        for c in w['values']:
            v.append((sign*c-key_at(family,key,pos,wi))%29);pos+=1
        out.append(v)
    return out

def run_all(rs):
    return [dict(family=f,sign=s,**fit(rs,f,s)) for f in FAMILIES for s in SIGNS]

def has_fit(rs):
    return any(fit(rs,f,s)['full'] for f in FAMILIES for s in SIGNS)

def main():
    ws=tokenize_words('%'.join(lp[:56]));sections={s:[w for w in ws if w['section']==s] for s in [0,1,2,4,5,6,7,8]}
    data={s:records(w) for s,w in sections.items()};assert sum(map(len,data.values()))==83
    results=[]
    for s,rs in data.items():
        tests=run_all(rs)
        for x in tests:
            if x['full']:
                dec=decrypt(sections[s],x['family'],x['sign'],x['witness'])
                x['full_text_witness']=' '.join(''.join(LET[c] for c in w) for w in dec)
                x['heldout_plaintext']=[LET[(x['sign']*r['c']-key_at(x['family'],x['witness'],r['rune'],r['word']))%29] for r in rs[x['cut']:]]
        results.append(dict(section=s,n=len(rs),train=(2*len(rs)+2)//3,heldout=len(rs)-(2*len(rs)+2)//3,locations=rs,tests=tests))
    # Baseline known warning page: original C decoded as P=-C-1 mod29.
    warning=tokenize_words(pages[0]);known=tokenize_words(GROUPS['warning'])
    assert decrypt(warning,'constant',-1,[0,1])==[w['values'] for w in known]
    wr=records(warning);wf=fit(wr,'constant',-1)
    # Constructed controls use all five decoded LP1 groups, concatenated as words.
    plain=[w for text in GROUPS.values() for w in tokenize_words(text)]
    assert {r['c'] for r in records(plain)}==set(AI)
    controls=[]
    for f in FAMILIES:
        key=[0,11] if f=='constant' else [7,11] if f.startswith('linear') else [3,7] if f=='period2_rune' else [3,7,11]
        for sign in SIGNS:
            enc=[];pos=0
            for wi,w in enumerate(plain):
                v=[]
                for p in w['values']:
                    v.append((sign*(p+key_at(f,key,pos,wi)))%29);pos+=1
                enc.append({**w,'values':v})
            rr=records(enc);ff=fit(rr,f,sign)
            assert ff['full']>0
            assert decrypt(enc,f,sign,key)==[w['values'] for w in plain]
            # Lexicographically selected full-fit witness may be ambiguous; measure instead of assuming uniqueness.
            selected=decrypt(enc,f,sign,ff['witness'])
            controls.append(dict(family=f,sign=sign,key=key,n=len(rr),**ff,planted_key_exact_recovery=True,selected_witness_exact_recovery=selected==[w['values'] for w in plain]))
    # Singletons are sparse. Use one within-section label permutation per trial,
    # then repeat the entire family search. This calibrates accidental fits;
    # it is not a realistic encrypted-English generative model.
    rng=np.random.default_rng(330144);draws=999;null=[];per_section={s:0 for s in data}
    for _ in range(draws):
        count=0
        for s,rs in data.items():
            cs=rng.permutation([r['c'] for r in rs]);sh=[{**r,'c':int(c)} for r,c in zip(rs,cs)]
            ok=has_fit(sh);count+=ok;per_section[s]+=ok
        null.append(count)
    observed=sum(any(x['full'] for x in r['tests']) for r in results)
    # Fixed sensitivity: remove only the two visually documented red singleton tokens.
    removed={(6,75),(54,0)};sensitivity=[]
    for s in [1,8]:
        rs=[r for r in data[s] if (r['page'],r['page_offset']) not in removed]
        sensitivity.append(dict(section=s,n=len(rs),removed=len(data[s])-len(rs),tests=run_all(rs)))
    out=dict(protocol=dict(alphabet=R,plaintext_singletons=['A','I'],families=FAMILIES,signs=SIGNS,formula='P=(sign*C-key(position)) mod 29',split='first ceil(2*n/3) singleton locations train; remaining chronological locations held out',acceptance='all training singletons must be A/I; retain entire fitting set and count how many also satisfy all held-out locations',rune_counter='all runes from section start, continuous over pages; punctuation and Arabic digits do not advance',word_counter='all lexical rune-words from section start; intra-rune apostrophes joined',no_phase_pruning='unseen periodic phases retain all 29 shifts',caveat='Held-out A/I constraints can filter a candidate set; this does not establish a unique prediction or plaintext.'),source_sha256=hashlib.sha256((F/'liber-primus__transcription--master.txt').read_bytes()).hexdigest(),script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),results=results,known_warning_control=dict(n=len(wr),**wf,entire_page_exact=True),constructed_controls=controls,null=dict(seed=330144,draws=draws,statistic='number of sections with at least one all-location fit in the full declared family',observed=observed,counts=null,per_section_success_counts=per_section,tail_fraction_plus_one=(1+sum(n>=observed for n in null))/(draws+1)),red_token_sensitivity=sensitivity)
    (F/'singleton_schedule_results.json').write_text(json.dumps(out,indent=2,ensure_ascii=False)+'\n')
    for r in results:print(r['section'],r['n'],r['train'],r['heldout'],[(x['family'],x['sign'],x['train'],x['full']) for x in r['tests'] if x['train']])
    print('controls',len(controls),'selected exact',sum(c['selected_witness_exact_recovery'] for c in controls),'null observed',observed,'tail',out['null']['tail_fraction_plus_one'])
    print('sensitivity',[(r['section'],r['removed'],[(x['family'],x['sign'],x['full']) for x in r['tests'] if x['full']]) for r in sensitivity])
if __name__=='__main__':main()