"""Session 49: independently specified key-skip controls; no third-party code.
Run with numpy available: python -B keyskip_audit.py
"""
from pathlib import Path
import hashlib,json,math,random
import numpy as np
from lp_core import vals,lp,LET
from word_core import GROUPS
ROOT=Path(__file__).resolve().parent
SEEDS=['CICADA3301','3301','CICADA','LIBER PRIMUS','DIVINITY','CIRCUMFERENCE','INSTAR','WELCOME','KNOW THYSELF','SNOW CRASH','CRYPTONOMICON','AGRIPPA','1033','EMERGENCE','INTELLIGENCE','PRIVACY']
SUPPRESS=.82
WIDTH=128
MAX_SKIP=2

def stream(seed,n=2048):
    out=[];counter=0
    while len(out)<n:
        block=hashlib.sha256(seed.encode('utf-8')+b'\x00'+counter.to_bytes(4,'big')).digest();counter+=1
        out.extend(b%29 for b in block if b<232)
    return out[:n]

def encrypt(p,k,rng):
    c=[];j=0;skips=[]
    for x in p:
        skipped=0
        while c and (x+k[j])%29==c[-1] and rng.random()<SUPPRESS:
            j+=1;skipped+=1
        c.append((x+k[j])%29);j+=1;skips.append(skipped)
    return c,skips

def model():
    # All solved groups except warning. No control plaintext in training.
    uni=np.ones(29)*.5;bc=np.zeros((29,29));tc=np.zeros((29,29,29))
    for name,source in GROUPS.items():
        if name=='warning':continue
        v=vals(source)
        for x in v:uni[x]+=1
        for a,b in zip(v,v[1:]):bc[a,b]+=1
        for a,b,c in zip(v,v[1:],v[2:]):tc[a,b,c]+=1
    up=uni/uni.sum();bp=(bc+5*up)/(bc.sum(1,keepdims=True)+5)
    tp=(tc+5*bp[None,:,:])/(tc.sum(2,keepdims=True)+5)
    return np.log(up),np.log(bp),np.log(tp)
U,B,T=model()

def decode(c,k,width=WIDTH,max_skip=MAX_SKIP):
    # Merge only identical (key position,last two plaintext symbols) states.
    states=[(0.,0,())]
    for i,y in enumerate(c):
        next_states={}
        for score,j,path in states:
            for s in range(max_skip+1 if i else 1):
                if j+s>=len(k):continue
                x=(y-k[j+s])%29
                if any((x+k[t])%29!=c[i-1] for t in range(j,j+s)):continue
                ll=float(U[x] if i==0 else B[path[-1],x] if i==1 else T[path[-2],path[-1],x])
                ll+=s*math.log(SUPPRESS)
                if i and y==c[i-1]:ll+=math.log(1-SUPPRESS)
                p=path+(x,);state=(j+s+1,p[-2:]);v=(score+ll,j+s+1,p)
                if state not in next_states or v[0]>next_states[state][0]:next_states[state]=v
        states=sorted(next_states.values(),reverse=True)[:width]
        if not states:return {'failed_at':i}
    score,j,path=states[0]
    return {'score':score/len(c),'consumed':j,'plaintext':list(path)}

def accuracy(a,b):return sum(x==y for x,y in zip(a,b))/len(a)
def label(v):return ' '.join(LET[x] for x in v)

def main():
    p=list(map(int,vals(GROUPS['warning'])));controls=[]
    for trial in range(6):
        seed=SEEDS[trial];k=stream(seed);c,skips=encrypt(p,k,random.Random(4900+trial))
        got=decode(c,k);wrong=decode(c,stream('WRONG '+seed));rigid=[(x-k[i])%29 for i,x in enumerate(c)]
        ranking=sorted([(decode(c,stream(candidate))['score'],candidate) for candidate in SEEDS],reverse=True)
        controls.append({'seed_rank':1+[candidate for _,candidate in ranking].index(seed),'seed':seed,'length':len(p),'total_skips':sum(skips),'max_skip':max(skips),'rigid_accuracy':accuracy(p,rigid),'beam_accuracy':accuracy(p,got['plaintext']),'wrong_key_accuracy':accuracy(p,wrong['plaintext']),'beam_score':got['score'],'wrong_key_score':wrong['score'],'plaintext':got['plaintext']})
    result={'spec':{'hash':'SHA256(UTF8(seed)||00||uint32be(counter)), counter starts 0; bytes <232 reduced mod29','suppression':SUPPRESS,'width':WIDTH,'max_skip_per_rune':MAX_SKIP,'seeds':SEEDS,'language_training':'All five known LP1 groups except warning; continuous rune trigrams, interpolated; no word scoring'},'controls':controls,'control_gate':all(x['beam_accuracy']>=.9 and x['beam_score']>x['wrong_key_score'] and x['seed_rank']==1 for x in controls)}
    # Do not run real-text screening if decoder fails the prespecified control gate.
    if result['control_gate']:
        real=[]
        for section,source in enumerate('%'.join(lp[:56]).split('$')):
            c=list(map(int,vals(source)))
            if len(c)<=100:continue
            prefix=c[:200];follow=c[:400]
            rankings=[]
            for seed in SEEDS:
                d=decode(prefix,stream(seed));rankings.append((d['score'],seed,d))
            rankings.sort(reverse=True)
            score,seed,d=rankings[0]
            # Lock best prefix path as well as seed for continuation by equivalent
            # decoder with the first 200 plaintext/key-consumption choices fixed.
            k=stream(seed);path=tuple(d['plaintext']);states=[(score*200,d['consumed'],path)]
            for i in range(200,len(follow)):
                nxt={};y=follow[i]
                for sc,j,pa in states:
                    for s in range(MAX_SKIP+1):
                        x=(y-k[j+s])%29
                        if any((x+k[t])%29!=follow[i-1] for t in range(j,j+s)):continue
                        ns=sc+float(T[pa[-2],pa[-1],x])+s*math.log(SUPPRESS)+(math.log(1-SUPPRESS) if y==follow[i-1] else 0)
                        pp=pa+(x,);key=(j+s+1,pp[-2:]);v=(ns,j+s+1,pp)
                        if key not in nxt or ns>nxt[key][0]:nxt[key]=v
                states=sorted(nxt.values(),reverse=True)[:WIDTH]
            sc,j,pa=states[0]
            real.append({'section':section,'selected_seed':seed,'prefix_score':score,'continuation_score':(sc-score*200)/(len(follow)-200),'prefix_rune_labels':label(d['plaintext'][:60]),'continuation_rune_labels':label(pa[200:]),'decoded_first_400_runes':list(pa),'prefix_rankings':[(a,b) for a,b,_ in rankings]})
        result['real']=real
    result['source_sha256']=hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    (ROOT/'keyskip_results.json').write_text(json.dumps(result,indent=2)+'\n')
    print('control gate',result['control_gate']);print('real sections',len(result.get('real',[])))
    print('controls',[(x['seed'],x['total_skips'],x['beam_accuracy'],x['rigid_accuracy']) for x in controls])
if __name__=='__main__':main()