"""Session 45: exact minimum singleton exceptions and longer-word ranking.
NumPy required. All fitting is retrospective; scores are not probabilities.
"""
from pathlib import Path
from itertools import product
import json,hashlib
import numpy as np
from singleton_schedule_audit import FAMILIES,SIGNS,AI,records,key_at,decrypt
from word_core import lp,GROUPS,LET,R
from lexical_words import tokenize_words
F=Path(__file__).resolve().parent
AB=np.array(list(product(range(29),repeat=2)))

def optimum(rs,family,sign,allowed=AI,keys=True):
    c=np.array([r['c'] for r in rs],dtype=int)
    if family in ['constant','linear_rune','linear_word']:
        ab=AB[:29] if family=='constant' else AB
        t=np.array([r['word'] if family=='linear_word' else r['rune'] for r in rs],dtype=int)
        hits=np.isin((sign*c[None,:]-ab[:,0,None]*t[None,:]-ab[:,1,None])%29,allowed).sum(1)
        best=int(hits.max());opts=ab[hits==best].tolist() if keys else []
    else:
        q=int(family[6]);best=0;parts=[]
        for phase in range(q):
            sub=np.array([r['c'] for r in rs if r['rune']%q==phase],dtype=int)
            hits=np.isin((sign*sub[None,:]-np.arange(29)[:,None])%29,allowed).sum(1)
            m=int(hits.max());best+=m;parts.append(np.flatnonzero(hits==m).tolist())
        opts=list(map(list,product(*parts))) if keys else []
    return len(rs)-best,opts

def search(rs,allowed=AI,keys=True):
    out=[]
    for f in FAMILIES:
        for sign in SIGNS:
            e,k=optimum(rs,f,sign,allowed,keys)
            out.append(dict(family=f,sign=sign,exceptions=e,keys=k))
    return min(t['exceptions'] for t in out),out

def model(words):
    counts=np.full((29,29,29),0.1)
    for w in words:
        for a,b,c in zip(w,w[1:],w[2:]):counts[a,b,c]+=1
    return np.log(counts/counts.sum(axis=2,keepdims=True))

def score(words,lm):
    s=0.;n=0
    for w in words:
        if len(w)<4:continue
        for a,b,c in zip(w,w[1:],w[2:]):s+=lm[a,b,c];n+=1
    return float(s/n)

def render(words):return ' '.join(''.join(LET[c] for c in w) for w in words)

def rank(words,rs,best,tests,lm):
    ranked=[]
    for t in tests:
        if t['exceptions']!=best:continue
        for key in t['keys']:
            d=decrypt(words,t['family'],t['sign'],key)
            ranked.append((score(d,lm),t['family'],t['sign'],key,d))
    ranked.sort(key=lambda x:-x[0]);s,f,sign,key,d=ranked[0]
    exc=[dict(**r,decoded=LET[d[r['word']][0]]) for r in rs if d[r['word']][0] not in AI]
    return dict(candidate_count=len(ranked),score=s,family=f,sign=sign,key=key,exceptions=exc,text=render(d),rune_words=d)

def main():
    known={k:tokenize_words(v) for k,v in GROUPS.items()}
    plain=[w for group in known.values() for w in group];lm=model([w['values'] for w in plain])
    benchmark=[]
    for name,ww in known.items():
        other=[w['values'] for k,group in known.items() if k!=name for w in group]
        benchmark.append(dict(group=name,score=score([w['values'] for w in ww],model(other))))
    allwords=tokenize_words('%'.join(lp[:56]));results=[];data={}
    for sec in [0,1,2,4,5,6,7,8]:
        words=[w for w in allwords if w['section']==sec];rs=records(words);data[sec]=rs
        best,tests=search(rs);winner=rank(words,rs,best,tests,lm)
        trio,trio_tests=search(rs,AI+[R.index('ᚩ')],False)
        clean=[r for r in rs if (r['page'],r['page_offset']) not in {(6,75),(54,0)}]
        red,redtests=search(clean,keys=False)
        results.append(dict(section=sec,n=len(rs),minimum_exceptions=best,tests=tests,ranked_witness=winner,allow_A_I_O_minimum=trio,red_excluded_n=len(clean),red_excluded_minimum=red))
        print(sec,len(rs),best,'AIO',trio,'red',red,'candidates',winner['candidate_count'],'score',winner['score'],flush=True)
    # Positive controls with two planted non-A/I singleton plaintexts.
    prs=records(plain);bad_word_indices=[prs[0]['word'],prs[-1]['word']]
    modified=[{**w,'values':list(w['values'])} for w in plain]
    for wi,p in zip(bad_word_indices,[0,2]):modified[wi]['values']=[p]
    controls=[]
    for fam in FAMILIES:
        key=[0,11] if fam=='constant' else [7,11] if fam.startswith('linear') else [3,7] if fam=='period2_rune' else [3,7,11]
        for sign in SIGNS:
            enc=[];pos=0
            for wi,w in enumerate(modified):
                v=[]
                for p in w['values']:v.append((sign*(p+key_at(fam,key,pos,wi)))%29);pos+=1
                enc.append({**w,'values':v})
            rs=records(enc);best,tests=search(rs);win=rank(enc,rs,best,tests,lm)
            assert best<=2
            exact=win['rune_words']==[w['values'] for w in modified]
            controls.append(dict(family=fam,sign=sign,planted_key=key,minimum_exceptions=best,selected_exact=exact,score=win['score'],candidates=win['candidate_count']))
    # Permute labels, preserve positions/inventory, optimize all families afresh.
    # Sign symmetry means +1 suffices for these binary A/I optimum counts.
    rng=np.random.default_rng(330145);null=[]
    for trial in range(999):
        total=0
        for rs in data.values():
            cs=rng.permutation([r['c'] for r in rs]);sh=[{**r,'c':int(c)} for r,c in zip(rs,cs)]
            total+=min(optimum(sh,f,1,keys=False)[0] for f in FAMILIES)
        null.append(total)
    observed=sum(r['minimum_exceptions'] for r in results)
    out=dict(protocol=dict(formula='P=(sign*C-key(position)) mod29; same families and counters as Session44',primary='Retrospective minimum number of all singleton locations that must be exempt from A/I, separately by section',ranking='Rank every key at the global minimum exception count for its section by mean within-word trigram log probability on words of at least four runes. No singleton, two-rune or three-rune words contribute to ranking.',model='Add-0.1 smoothed conditional rune trigrams from five solved LP1 groups; benchmarks leave one group out',limitations='No new holdout test. Language ranking is exploratory and not a calibrated test. Exemptions remove constraints, not runes: counters and tokenization remain unchanged.',sensitivity='Allow A/I/O, or remove the two documented red tokens from constraints while retaining counters'),source_sha256=hashlib.sha256((F/'liber-primus__transcription--master.txt').read_bytes()).hexdigest(),script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),results=results,benchmark=benchmark,controls=controls,null=dict(seed=330145,draws=999,observed_minimum_total=observed,values=null,central95=np.quantile(null,[.025,.975]).tolist(),median=float(np.median(null)),lower_tail_plus_one=(1+sum(v<=observed for v in null))/1000))
    (F/'singleton_exception_results.json').write_text(json.dumps(out,indent=2,ensure_ascii=False)+'\n')
    print('TOTAL',observed,'null',out['null']['central95'],out['null']['median'],out['null']['lower_tail_plus_one'],'controls',sum(c['selected_exact'] for c in controls),'/10',flush=True)
    print('bench',benchmark,flush=True)
if __name__=='__main__':main()