"""Bounded stateful alphabet and previous-symbol exclusion hypotheses."""
import json,time
from word_core import *

INITIAL=np.array([(direction*np.arange(29)+offset)%29 for direction in (1,-1) for offset in range(29)])
INITIAL_INFO=[{'direction':d,'offset':o} for d in (1,-1) for o in range(29)]
UPDATES=['move_front','move_back','swap_front','rotate_front','rotate_after']
RESETS={'continuous':'','word':'-.,;&$','paragraph':'&$'}

def reset_flags(source,reset):
    _,flags=positions(source,reset=RESETS[reset]);flags[0]=True
    if reset=='word':flags[[w['rune_offset'] for w in tokenize(source)]]=True
    return flags

def update_states(state,ix,update):
    col=np.broadcast_to(np.arange(29),state.shape);ix=ix[:,None]
    if update=='move_front':order=np.where(col==0,ix,np.where(col<=ix,col-1,col))
    elif update=='move_back':order=np.where(col==28,ix,np.where(col>=ix,col+1,col))
    elif update=='swap_front':order=np.where(col==0,ix,np.where(col==ix,0,col))
    elif update=='rotate_front':order=(col+ix)%29
    elif update=='rotate_after':order=(col+ix+1)%29
    else:raise ValueError(update)
    return np.take_along_axis(state,order,axis=1)

def dynamic_decode(c,flags,update,mode,initial=INITIAL):
    state=initial.copy();out=np.empty((len(initial),len(c)),dtype=np.int64);rows=np.arange(len(initial))
    for j,ch in enumerate(c):
        if flags[j]:state=initial.copy()
        if mode=='symbol_to_rank':ix=np.argmax(state==ch,axis=1);out[:,j]=ix
        else:ix=np.full(len(initial),ch);out[:,j]=state[rows,ix]
        state=update_states(state,ix,update)
    return out

def lengths_and_cut(source):
    ll=np.array([len(w['values']) for w in tokenize(source)]);cs=ll.cumsum();nw=max(1,int(np.searchsorted(cs,min(256,int(.6*cs[-1])),side='right')))
    return ll,nw,int(cs[nw-1])

def sweep_dynamic(source,label):
    c=vals(source);ll,nw,cut=lengths_and_cut(source);records=[];count=0
    for reset in RESETS:
        flags=reset_flags(source,reset)
        for update in UPDATES:
            for mode in ['symbol_to_rank','rank_to_symbol']:
                raw=dynamic_decode(c,flags,update,mode)
                for sign in (1,-1):
                    d=(sign*raw[:,None,:cut]+np.arange(29)[None,:,None])%29
                    sc=FULL_MODEL.score(d.reshape(-1,cut),np.array(ll[:nw]));j=int(np.argmax(sc));ini,k=np.unravel_index(j,d.shape[:2]);decoded=(sign*raw[ini]+k)%29
                    records.append({'reset':reset,'update':update,'mode':mode,'initial_index':int(ini),'initial':INITIAL_INFO[ini],'sign':sign,'shift':int(k),'prefix_score':float(sc[j]),'decoded':decoded})
                    count+=len(sc)
    records.sort(key=lambda r:r['prefix_score'],reverse=True);top=records[:20]
    for r in top:
        d=r.pop('decoded');r['suffix_score']=float(FULL_MODEL.score(d[None,cut:],ll[nw:])[0]);r['text']=render(source,d);r['decoded_rune_values']=d.tolist()
    return {'target':label,'candidate_settings':count,'runes':len(c),'prefix_runes':cut,'top_prefix_candidates':top}

AA=np.repeat(np.arange(1,29),29);BB=np.tile(np.arange(29),28)

def exclusion_raw(source,reset):
    c=vals(source);cc=(AA[:,None]*c+BB[:,None])%29;previous=np.roll(cc,1,axis=1);previous[:,reset_flags(source,reset)]=0
    return np.where(cc==previous,28,cc-(cc>previous))

def sweep_exclusion(source,label):
    c=vals(source);ll,nw,cut=lengths_and_cut(source);records=[];count=0
    for reset in RESETS:
        raw=exclusion_raw(source,reset)
        for sign in (1,-1):
            for shift in range(29):
                d=(sign*raw[:,:cut]+shift)%29;sc=FULL_MODEL.score(d,ll[:nw]);j=int(np.argmax(sc))
                records.append({'reset':reset,'input_a':int(AA[j]),'input_b':int(BB[j]),'sign':sign,'shift':shift,'prefix_score':float(sc[j]),'decoded':(sign*raw[j]+shift)%29});count+=len(sc)
    records.sort(key=lambda r:r['prefix_score'],reverse=True);top=records[:20]
    for r in top:
        d=r.pop('decoded');r['suffix_score']=float(FULL_MODEL.score(d[None,cut:],ll[nw:])[0]);r['text']=render(source,d);r['decoded_rune_values']=d.tolist()
    return {'target':label,'candidate_settings':count,'runes':len(c),'prefix_runes':cut,'top_prefix_candidates':top}

def controls():
    source=lp[57];plain=vals(source);out=[]
    # Invert the chosen decoder: plaintext symbols -> ciphertext ranks.
    for mode,update,reset,ini in [('rank_to_symbol','move_front','continuous',36),('symbol_to_rank','rotate_after','word',11)]:
        inverse_mode='symbol_to_rank' if mode=='rank_to_symbol' else 'rank_to_symbol'
        cipher=dynamic_decode(plain,reset_flags(source,reset),update,inverse_mode,INITIAL[ini:ini+1])[0]
        text=replace_values(source,cipher);found=sweep_dynamic(text,'synthetic_dynamic')
        exact=[i for i,r in enumerate(found['top_prefix_candidates']) if r['decoded_rune_values']==plain.tolist()]
        assert exact,found['top_prefix_candidates'][0]
        out.append({'type':'dynamic','construction':{'mode':mode,'update':update,'reset':reset,'initial':INITIAL_INFO[ini]},'exact_rank':exact[0]+1})
    a,b=17,7;inv=pow(a,-1,29);sign,shift=-1,3;reset='paragraph';flags=reset_flags(source,reset);previous=0;cipher=[]
    for p,flag in zip(plain,flags):
        if flag:previous=0
        rank=int((sign*(int(p)-shift))%29)
        selected=previous if rank==28 else rank+int(rank>=previous)
        cipher.append((inv*(selected-b))%29);previous=selected
    text=replace_values(source,cipher);found=sweep_exclusion(text,'synthetic_exclusion')
    exact=[i for i,r in enumerate(found['top_prefix_candidates']) if r['decoded_rune_values']==plain.tolist()]
    assert exact,found['top_prefix_candidates'][0]
    out.append({'type':'exclusion','construction':{'reset':reset,'input_a':a,'input_b':b,'sign':sign,'shift':shift},'exact_rank':exact[0]+1})
    return out

if __name__=='__main__':
    started=time.time();out={'controls':controls(),'definitions':{'dynamic':'Initial natural or reversed alphabet, all29 rotations. Cipher symbol to current rank, or cipher rank to current symbol. Update after every rune with one of five specified permutations. Final sign +/-1 and shift0..28.','exclusion':'Relabel each ciphertext rune as (a*c+b)mod29, a1..28,b0..28. Rank the current rune among the28 alphabet symbols excluding previous; repeated rune gets reserved rank28. Previous state initially0. Final sign+/-1 and shift0..28.','reset_policies':RESETS,'selection':'Word-aligned prefix at most256 runes and at most60% of target; suffix not used to select parameters.','limits':'No arbitrary starting permutations, arbitrary final substitution keys, selective F skips, or plaintext-dependent permutation updates tested.'},'sections':[],'dynamic_settings':0,'exclusion_settings':0}
    print('All three controls recovered',flush=True)
    sections=[s for s in '%'.join(lp[:56]).split('$') if len(vals(s))]
    for i,s in enumerate(sections):
        if len(vals(s))<60:continue
        dynamic=sweep_dynamic(s,f'section_{i}');exclusion=sweep_exclusion(s,f'section_{i}');out['sections'].append({'index':i,'dynamic':dynamic,'exclusion':exclusion});out['dynamic_settings']+=dynamic['candidate_settings'];out['exclusion_settings']+=exclusion['candidate_settings']
        print('Section',i,'dynamic',dynamic['top_prefix_candidates'][0]['prefix_score'],dynamic['top_prefix_candidates'][0]['suffix_score'],'exclusion',exclusion['top_prefix_candidates'][0]['prefix_score'],exclusion['top_prefix_candidates'][0]['suffix_score'],flush=True)
    out['elapsed_seconds']=time.time()-started
    (ROOT/'changing_alphabet_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2));print('DONE',out['dynamic_settings'],out['exclusion_settings'],flush=True)