"""Use known repeated ciphertext as a conditional equal-plaintext constraint.
Unknown additive key, periods1..64, global/page/section/phrase-aligned counters, two signs.
Fit on first60% of section6 excluding the anchor and all ngrams touching it.
No plaintext guesses are imposed on the six-rune phrase.
"""
import json,time
from word_core import *
SCOPES={'global':'','page':'%','section':'$','phrase_aligned':'$'}
ANCHORS=(6555,12950)

def counters(source,scope):
    ix=positions(source,reset=SCOPES[scope])[0]
    if scope=='phrase_aligned':
        words=tokenize(source);second=[w for w in words if w['section']==8]
        start=second[0]['rune_offset'];end=second[-1]['rune_offset']+len(second[-1]['values'])
        # Model hypothesis: align the two anchors to the same key phase.
        ix[start:end]-=int(ix[ANCHORS[1]]-ix[ANCHORS[0]])
    return ix


def merge_classes(period,ix,anchors):
    parent=list(range(period))
    def root(x):
        while parent[x]!=x:parent[x]=parent[parent[x]];x=parent[x]
        return x
    for a,b in zip(ix[anchors[0]:anchors[0]+6],ix[anchors[1]:anchors[1]+6]):
        a=root(int(a)%period);b=root(int(b)%period);parent[a]=b
    roots=[root(i) for i in range(period)];names={r:i for i,r in enumerate(sorted(set(roots)))}
    return np.array([names[r] for r in roots],dtype=np.int64)


def events_for_words(words):
    events=[]
    for w in words:
        start=w['rune_offset'];nn=len(w['values']);idx=[-1,-1]+list(range(start,start+nn))+[-1]
        events.extend(zip(idx,idx[1:],idx[2:]))
    return np.array(events,dtype=np.int64)


def fit_key(c,classes,events,mask,sign,rng,restarts=3,max_passes=8):
    ev=events[~np.any(np.isin(events,np.flatnonzero(mask)),axis=1)]
    cg=np.where(ev>=0,classes[np.maximum(ev,0)],-1)
    cv=np.where(ev>=0,(sign*c[np.maximum(ev,0)])%29,29)
    ng=int(classes.max())+1;used=np.unique(ev[ev>=0]);cu=(sign*c[used])%29;gu=classes[used]
    # Unigram initialization from the same independent known-plaintext corpus.
    hist=np.bincount(np.concatenate([vals(s) for s in GROUPS.values()]),minlength=29)+.5
    logp=np.log(hist/hist.sum());uni=np.zeros(ng,dtype=np.int64)
    for g in range(ng):
        vv=cu[gu==g]
        uni[g]=int(np.argmax(logp[(vv[None,:]-np.arange(29)[:,None])%29].sum(1))) if len(vv) else 0
    affected=[np.flatnonzero(np.any(cg==g,axis=1)) for g in range(ng)]
    def decode_events(key):return np.where(cg<0,29,(cv-key[np.maximum(cg,0)])%29)
    best=None
    for restart in range(restarts):
        key=uni.copy() if restart==0 else rng.integers(0,29,ng)
        # Unobserved key classes remain0; do not optimize them on validation.
        for g in range(ng):
            if not len(affected[g]):key[g]=0
        for iteration in range(max_passes):
            changes=0
            for g in rng.permutation(ng):
                rows=affected[g]
                if not len(rows):continue
                gg=cg[rows];vv=cv[rows]
                base=np.where(gg<0,29,(vv-key[np.maximum(gg,0)])%29)
                trials=np.broadcast_to(base,(29,*base.shape)).copy()
                for col in range(3):
                    m=gg[:,col]==g
                    trials[:,m,col]=(vv[None,m,col]-np.arange(29)[:,None])%29
                scores=FULL_MODEL.log[trials[:,:,0],trials[:,:,1],trials[:,:,2]].sum(1)
                winner=int(np.argmax(scores));changes+=winner!=key[g];key[g]=winner
            if not changes:break
        d=decode_events(key);score=float(FULL_MODEL.log[d[:,0],d[:,1],d[:,2]].sum())
        if best is None or score>best[0]:best=(score,key.copy(),iteration+1)
    return {'total_log_score':best[0],'mean_log_score':best[0]/len(ev),'class_key':best[1].tolist(),'events':len(ev),'independent_key_classes':ng,'unobserved_classes':[g for g in range(ng) if not len(affected[g])],'winning_restart_passes':best[2]}


def score_words(c,decoded,words,excluded):
    ev=events_for_words(words);ev=ev[~np.any(np.isin(ev,list(excluded)),axis=1)]
    d=np.where(ev>=0,decoded[np.maximum(ev,0)],29)
    return float(FULL_MODEL.log[d[:,0],d[:,1],d[:,2]].mean())


def control(scope,period,key):
    source='%'.join([lp[57]]*8);cplain=vals(source);ix=counters(source,scope)
    cipher=(cplain+np.array(key)[ix%period])%29
    # Two identical copies at page starts; period5 divides95 for global case.
    anchors=(0,len(vals(lp[57])))
    assert np.array_equal(cipher[:6],cipher[anchors[1]:anchors[1]+6])
    grouping=merge_classes(period,ix,anchors);classes=grouping[ix%period]
    words=tokenize(source);limit=int(len(cplain)*.6);train=[w for w in words if w['rune_offset']+len(w['values'])<=limit]
    mask=np.zeros(len(cipher),bool)
    for a in anchors:mask[a:a+6]=True
    fit=fit_key(cipher,classes,events_for_words(train),mask,1,np.random.default_rng(333))
    decoded=(cipher-np.array(fit['class_key'])[classes])%29
    assert np.array_equal(decoded,cplain),{'scope':scope,'fit':fit}
    return {'scope':scope,'period':period,'unknown_key_recovered':True,'exact_full_plaintext_recovered':True,'construction_key':key,'fit':fit,'limit':'Period and scope supplied to this control; key is learned. Repeated held-out LP2p57 plaintext is absent from language-model training.'}

if __name__=='__main__':
    start=time.time();source='%'.join(lp[:56]);c=vals(source);words=tokenize(source)
    assert np.array_equal(c[ANCHORS[0]:ANCHORS[0]+6],c[ANCHORS[1]:ANCHORS[1]+6])
    first=[w for w in words if w['section']==6];second=[w for w in words if w['section']==8]
    section_start=first[0]['rune_offset'];section_end=first[-1]['rune_offset']+len(first[-1]['values']);limit=section_start+int(.6*(section_end-section_start))
    train=[w for w in first if w['rune_offset']+len(w['values'])<=limit];cut=train[-1]['rune_offset']+len(train[-1]['values']);valid1=[w for w in first if w['rune_offset']>=cut]
    mask=np.zeros(len(c),bool)
    for a in ANCHORS:mask[a:a+6]=True
    excluded=np.flatnonzero(mask).tolist();ev=events_for_words(train)
    out={'assumption':'Both identical ciphertext phrases encode the same six plaintext runes under a shared additive periodic key. This is a hypothesis, not established by repetition.','controls':[control('global',5,[3,17,9,22,1]),control('page',7,[2,8,1,9,14,22,5]),control('global',14,[5,8,13,5,8,13,3,17,9,22,1,5,8,13])],'anchors_global':list(ANCHORS),'training':{'section':6,'global_start':section_start,'global_end_exclusive':cut,'rune_count_before_anchor_exclusion':cut-section_start,'exclude_both_anchors_and_ngrams_touching_them':True},'validation':{'same_section_suffix_runes':section_end-cut,'other_section':8,'other_section_runes':sum(len(w['values']) for w in second),'anchor_runes_excluded_from_scores':6},'models':[],'counter_distances':{},'limits':['Coordinate ascent is heuristic:3 starts, at most8 sweeps, not exhaustive key search.','Periods1..64 only, common key across targets, no selective F skips or unknown key alphabets.','Rune-only counters; delimiters do not consume positions.','No validation plaintext scores used to change any key. Phrase-aligned mode uses the observed anchor locations to set the second section phase.','Description-length ranking is heuristic; trigram scores are not calibrated probabilities.']}
    print('All three unknown-key controls recovered',flush=True)
    for scope,reset in SCOPES.items():
        ix=counters(source,scope);gap=int(ix[ANCHORS[1]]-ix[ANCHORS[0]])
        out['counter_distances'][scope]={'anchor_positions':[int(ix[a]) for a in ANCHORS],'difference':gap,'divisors_up_to64':[p for p in range(1,65) if gap%p==0]}
        for period in range(1,65):
            grouping=merge_classes(period,ix,ANCHORS);classes=grouping[ix%period]
            for sign in (1,-1):
                fit=fit_key(c,classes,ev,mask,sign,np.random.default_rng(3301006+period*2+(sign==1)))
                key=np.array(fit['class_key']);decoded=(sign*c-key[classes])%29
                assert np.array_equal(decoded[ANCHORS[0]:ANCHORS[0]+6],decoded[ANCHORS[1]:ANCHORS[1]+6])
                row={'scope':scope,'period':period,'sign':sign,**fit,'period_key':key[grouping].tolist(),'description_length_rank_score':(fit['total_log_score']-len(key)*np.log(29))/fit['events'],'same_section_suffix_score':score_words(c,decoded,valid1,excluded),'other_section_score':score_words(c,decoded,second,excluded),'anchor_transliteration':' '.join(LET[int(v)] for v in decoded[ANCHORS[0]:ANCHORS[0]+6]),'training_excerpt':' '.join(LET[int(v)] for v in decoded[section_start:section_start+80]),'other_section_excerpt':' '.join(LET[int(v)] for v in decoded[second[0]['rune_offset']:second[0]['rune_offset']+100])}
                out['models'].append(row)
            if period%16==0:print(scope,'period',period,flush=True)
    out['models'].sort(key=lambda r:r['description_length_rank_score'],reverse=True)
    out['model_count']=len(out['models']);out['seconds']=time.time()-start
    # Retain complete decoded targets for top10 by training-only ranking.
    for row in out['models'][:10]:
        ix=counters(source,row['scope']);d=(row['sign']*c-np.array(row['period_key'])[ix%row['period']])%29
        row['section6_text']=' '.join(''.join(LET[int(z)] for z in d[w['rune_offset']:w['rune_offset']+len(w['values'])]) for w in first)
        row['section8_text']=' '.join(''.join(LET[int(z)] for z in d[w['rune_offset']:w['rune_offset']+len(w['values'])]) for w in second)
    (ROOT/'phrase_key_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2));print('DONE',out['model_count'],round(out['seconds'],1),flush=True)