"""Session 19: fixed colored-mark counter models and withheld continuation.
Python + NumPy. Inputs and marker audit are local; no live calls.
"""
import itertools, json, hashlib
from word_core import *

STREAMS=['prime_minus_one','integer_totient','fibonacci']
POLICIES=['continuous','all_full_stops','all_red','only_13','only_10','bracket_pause','bracket_advance']
PARAMS=list(itertools.product(range(33),(1,-1),(1,-1),range(29)))
PAR=np.array(PARAMS)
CUT=984

def reset_indices(n,events):
    starts=np.array(sorted(set([0]+[int(x) for x in events if 0<=x<n])))
    pos=np.arange(n)
    return pos-starts[np.searchsorted(starts,pos,side='right')-1]

def schedule(n,policy,red13,red10,dots):
    if policy=='continuous':return np.arange(n)
    if policy=='all_full_stops':return reset_indices(n,dots)
    if policy=='all_red':return reset_indices(n,red13+red10)
    if policy=='only_13':return reset_indices(n,red13)
    if policy=='only_10':return reset_indices(n,red10)
    assert policy in ('bracket_pause','bracket_advance')
    assert len(red10)==2
    a,b=red10
    assert a<b and not any(a<r<b for r in red13)
    ix=reset_indices(n,red13)
    ix[a:b]=np.arange(b-a)
    if policy=='bracket_pause':
        next13=min([r for r in red13 if r>=b]+[n])
        ix[b:next13]-=b-a
    return ix

def split_at(source,cut):
    pos=[i for i,c in enumerate(source) if c in R]
    char_cut=pos[cut] if cut<len(pos) else len(source)
    a,b=source[:char_cut],source[char_cut:]
    assert len(vals(a))==cut
    assert not any(w['rune_offset']<cut<w['rune_offset']+len(w['values']) for w in tokenize(source))
    return a,b

def score(v,s):
    return float(FULL_MODEL.score(v[None,:],[len(w['values']) for w in tokenize(s)])[0])

def fit(source,ix,stream,cut):
    train,valid=split_at(source,cut);v=vals(source)
    q=SEQS[stream][PAR[:,0,None]+ix[None,:cut]]
    candidates=(PAR[:,1,None]*v[None,:cut]+PAR[:,2,None]*q+PAR[:,3,None])%29
    rank=FULL_MODEL.score(candidates,[len(w['values']) for w in tokenize(train)])
    pick=int(np.argmax(rank));offset,a,b,shift=PARAMS[pick]
    out=(a*v+b*SEQS[stream][ix+offset]+shift)%29
    return dict(stream=stream,parameters=dict(offset=offset,cipher_sign=a,sequence_sign=b,shift=shift),
        selection_score=float(rank[pick]),validation_score=score(out[cut:],valid),
        plaintext_runes=''.join(R[int(x)] for x in out),decoded_text=render(source,out),
        validation_text=render(valid,out[cut:]),selection_opening=render(train,out[:cut])[:150]),out

def fixtures():
    expected={
        'continuous':list(range(14)),
        'all_full_stops':[0,1,0,1,0,0,1,0,1,0,1,2,0,1],
        'all_red':[0,1,0,1,2,0,1,2,3,0,1,2,0,1],
        'only_13':[0,1,0,1,2,3,4,5,6,7,8,9,0,1],
        'only_10':[0,1,2,3,4,0,1,2,3,0,1,2,3,4],
        'bracket_pause':[0,1,0,1,2,0,1,2,3,3,4,5,0,1],
        'bracket_advance':[0,1,0,1,2,0,1,2,3,7,8,9,0,1]}
    for p,want in expected.items():
        assert schedule(14,p,[2,12],[5,9],[2,4,5,7,9,12]).tolist()==want,p
    return len(expected)

def controls():
    # Known held-out LP2 words; no repeated filler and no LP1 training text.
    v=vals(lp[56]);ix=np.arange(len(v));ix[57:]-=1
    pv=(v-SEQS['prime_minus_one'][ix])%29;pv[56]=v[56]
    words=tokenize(replace_values(lp[56],pv)+'%'+lp[57])
    plain='-'.join(w['runes'] for w in words);v=vals(plain)
    offsets=np.r_[0,np.cumsum([len(w['values']) for w in words])]
    r13=[int(offsets[k]) for k in [2,31,32]]
    r10=[int(offsets[k]) for k in [12,26]]
    dots=[int(offsets[k]) for k in [2,7,12,16,20,26,31,32,len(words)]]
    cut=int(offsets[20]);results=[]
    for test,(stream,policy) in enumerate(itertools.product(STREAMS,POLICIES)):
        offset=7;a=1 if test%2 else -1;b=-1 if test%3 else 1;shift=11
        encrypted=[];counter=0;outer=None;opened=False
        # Independent scalar state-machine encryption.
        for j,p in enumerate(v):
            if policy=='all_full_stops' and j in dots:counter=0
            elif policy=='all_red' and j in r13+r10:counter=0
            elif policy=='only_13' and j in r13:counter=0
            elif policy=='only_10' and j in r10:counter=0
            elif policy.startswith('bracket_'):
                if j in r13:counter=0
                if j==r10[0]:outer=counter;counter=0;opened=True
                elif j==r10[1]:
                    counter=outer+(r10[1]-r10[0] if policy=='bracket_advance' else 0);opened=False
            q=int(SEQS[stream][offset+counter])
            encrypted.append((a*(int(p)-b*q-shift))%29)
            counter+=1
        cipher=replace_values(plain,encrypted)
        idx=schedule(len(v),policy,r13,r10,dots)
        replay=(a*vals(cipher)+b*SEQS[stream][offset+idx]+shift)%29
        assert np.array_equal(replay,v),(stream,policy)
        picked,decoded=fit(cipher,idx,stream,cut)
        exact=bool(np.array_equal(decoded,v))
        results.append(dict(stream=stream,policy=policy,runes=len(v),selection_runes=cut,
            replay_exact=True,within_family_search_exact=exact,selected_parameters=picked['parameters']))
    return results

def main():
    fixture_count=fixtures();source='%'.join(lp[:56]).split('$')[7];v=vals(source);n=len(v)
    audit=json.loads((ROOT/'visual_audit_results.json').read_text())
    red={kind:sorted(x['section_offset'] for x in audit['red_punctuation']
        if x['lp2_page']>=33 and x['dots']==kind and 0<x.get('section_offset',-1)<=n) for kind in [10,13]}
    assert red=={10:[754,1307],13:[8,1468,1471,1589]}
    dots=[];j=0
    for c in source:
        if c in R:j+=1
        elif c=='.':dots.append(j)
    assert n==1589 and len(dots)==19 and all(x in dots for x in red[10]+red[13])
    train,valid=split_at(source,CUT)
    plan=dict(streams=STREAMS,policies=POLICIES,offsets=list(range(33)),signs=[1,-1],shifts=list(range(29)),
        formula='p = a*c + b*q[index+offset] + shift mod 29',candidate_settings=len(PARAMS)*len(STREAMS)*len(POLICIES),
        selection_interval=[0,CUT],validation_interval=[CUT,n],red_10_offsets=red[10],red_13_offsets=red[13],
        full_stop_offsets=dots,selection_words=len(tokenize(train)),validation_words=len(tokenize(valid)),
        skips=False,punctuation_consumes_positions=False,
        candidate_selection='One winner per stream/policy, prefix only; no refitting on continuation.')
    (ROOT/'colored_reset_plan.json').write_text(json.dumps(plan,indent=2))
    check=controls();winners=[];index_audit={}
    post_source=split_at(source,1307)[1]
    for policy in POLICIES:
        ix=schedule(n,policy,red[13],red[10],dots)
        index_audit[policy]={str(p):int(ix[p]) for p in [8,754,984,1307,1468,1471]}
        for stream in STREAMS:
            picked,out=fit(source,ix,stream,CUT);picked['policy']=policy
            picked['after_list_score']=score(out[1307:],post_source)
            picked['after_list_text']=render(post_source,out[1307:])
            winners.append(picked)
    winners.sort(key=lambda w:w['selection_score'],reverse=True)
    same=[schedule(n,p,red[13],red[10],dots)[:CUT].tolist() for p in ['all_red','bracket_pause','bracket_advance']]
    assert same[0]==same[1]==same[2]
    # Equivalent full schedules are counted separately above only if actually different.
    distinct=len({tuple(schedule(n,p,red[13],red[10],dots)) for p in POLICIES})
    results=dict(plan=plan,fixture_count=fixture_count,controls=check,index_at_boundaries=index_audit,
        distinct_full_schedules=distinct,identical_selection_schedules=['all_red','bracket_pause','bracket_advance'],
        input_sha256={f:hashlib.sha256((ROOT/f).read_bytes()).hexdigest() for f in
            ['liber-primus__transcription--master.txt','visual_audit_results.json']},winners=winners)
    (ROOT/'colored_reset_results.json').write_text(json.dumps(results,indent=2,ensure_ascii=False))
    print('Settings',plan['candidate_settings'],'selection',CUT,'validation',n-CUT)
    print('Controls exact',sum(c['within_family_search_exact'] for c in check),'/',len(check))
    for w in winners:print(w['policy'],w['stream'],round(w['selection_score'],4),round(w['validation_score'],4),round(w['after_list_score'],4))

if __name__=='__main__':main()