"""Exact joint decoding under reused additive stream; exploratory, not a solution.
Run with the primary runtime Python. Sources and model are local.
"""
import json,itertools,time
from word_core import *

def extract(source,start,n):
    v=vals(source)[start:start+n]
    pos=np.full(n,-1,int);ends=np.zeros(n,bool)
    for w in tokenize(source):
        a=w['rune_offset'];b=a+len(w['values'])
        for j in range(max(start,a),min(start+n,b)):pos[j-start]=j-a
        if start<=b-1<start+n:ends[b-1-start]=True
    assert len(v)==n and (pos>=0).all()
    return v,dict(pos=pos.tolist(),ends=ends.tolist())

def display(v,layout):
    out=''
    for i,x in enumerate(v):
        if i and layout['pos'][i]==0:out+=' '
        out+=LET[int(x)]
    return ('…' if layout['pos'][0] else '')+out

X=np.arange(29)[:,None,None];Y=np.arange(29)[None,:,None];Z=np.arange(29)[None,None,:]

def potentials(d,la,lb,model):
    costs=[];events=0
    for i in range(len(d)):
        t=np.zeros((29,29,29))
        for layout,delta in [(la,np.zeros(len(d),int)),(lb,d)]:
            x=(X-delta[max(0,i-2)])%29;y=(Y-delta[max(0,i-1)])%29;z=(Z-delta[i])%29
            p=layout['pos'][i]
            if p==0:t+=model.log[B,B,z];events+=1
            elif p==1 and i>=1:t+=model.log[B,y,z];events+=1
            elif p>=2 and i>=2:t+=model.log[x,y,z];events+=1
            if layout['ends'][i] and (p==0 or i>=1):
                t+=model.log[B if p==0 else y,z,B];events+=1
        costs.append(t)
    return costs,events

def maximize(costs):
    dp=np.full((29,29),-np.inf);dp[0,0]=0;backs=[]
    for t in costs:
        s=dp[:,:,None]+t
        backs.append(s.argmax(axis=0).astype(np.uint8));dp=s.max(axis=0)
    y,z=np.unravel_index(dp.argmax(),dp.shape);score=float(dp[y,z]);out=np.empty(len(costs),int)
    for i in range(len(costs)-1,-1,-1):
        out[i]=z;x=int(backs[i][y,z]);y,z=x,y
    return score,out

def direct_score(v,costs):
    return float(sum(t[int(v[i-2]) if i>=2 else 0,int(v[i-1]) if i>=1 else 0,int(v[i])] for i,t in enumerate(costs)))

def solve(d,la,lb,model=FULL_MODEL):
    costs,events=potentials(np.asarray(d),la,lb,model);score,a=maximize(costs);b=(a-d)%29
    assert abs(direct_score(a,costs)-score)<1e-9
    return dict(score=score/events,total_score=score,events=events,a=a.tolist(),b=b.tolist(),text_a=display(a,la),text_b=display(b,lb))

def paired(d,la,lb):
    ans=[]
    for sign in [1,-1]:
        r=solve(sign*d%29,la,lb);r['sign']=sign;ans.append(r)
    return sorted(ans,key=lambda x:x['score'],reverse=True)

def checks():
    # Independent exhaustive evaluation of every 3-rune pair under two layouts.
    la=dict(pos=[0,1,2],ends=[False,False,True]);lb=dict(pos=[0,0,1],ends=[True,False,True])
    allv=np.array(list(itertools.product(range(29),repeat=3)))
    cases=[]
    for sign in [1,-1]:
        d=sign*np.array([3,7,11])%29;b=(allv-d)%29
        brute=FULL_MODEL.score(allv,[3])*4+FULL_MODEL.score(b,[1,2])*5
        r=solve(d,la,lb)
        assert abs(r['total_score']-brute.max())<1e-9
        cases.append(dict(sign=sign,enumerated=len(allv),exact_score_match=True))
    # Entire source group excluded from the language model. No selected recovery.
    model=Model([s for k,s in GROUPS.items() if k!='koan']);source=GROUPS['koan'];controls=[]
    for aa,bb in [(143,279),(411,556),(568,716)]:
        a,la=extract(source,aa-28,34);b,lb=extract(source,bb-28,34)
        assert np.array_equal(a[-6:],b[-6:])
        d=(a-b)%29;r=solve(d,la,lb,model);costs,events=potentials(d,la,lb,model)
        r.update(anchor_offsets=[aa,bb],true_a=display(a,la),true_b=display(b,lb),matching_runes_a=int(np.sum(a==r['a'])),matching_runes_b=int(np.sum(b==r['b'])),true_score=direct_score(a,costs)/events,exact_recovery=bool(np.array_equal(a,r['a']) and np.array_equal(b,r['b'])))
        controls.append(r)
    return dict(exhaustive=cases,heldout_plaintext=controls)

def main():
    started=time.time();source='%'.join(lp[:56]);a,la=extract(source,6527,34);b,lb=extract(source,12922,34);d=(a-b)%29
    assert np.array_equal(a[-6:],b[-6:]);assert la['pos'][0]==lb['pos'][0]==0
    out=dict(hypothesis='C_A=P_A+K and C_B=P_B+K modulo29 over 34 aligned runes; also test reversed plaintext sign. No evidence establishes shared K.',offsets=[6527,12922],cipher_a=a.tolist(),cipher_b=b.tolist(),difference=d.tolist(),layout_a=la,layout_b=lb,checks=checks(),observed=paired(d,la,lb),rng_seed=3301007,null={})
    print('Controls and observed fit completed',flush=True)
    rng=np.random.default_rng(3301007)
    # Primary null: condition on observed prefix difference inventory and six zeros.
    # Secondary sensitivity: independent uniform differences; same anchor and layout.
    for kind in ['permuted_prefix','uniform_prefix']:
        scores=[];best=None
        for j in range(999):
            prefix=rng.permutation(d[:28]) if kind=='permuted_prefix' else rng.integers(0,29,28)
            dd=np.r_[prefix,np.zeros(6,int)];r=paired(dd,la,lb)[0];scores.append(r['score'])
            if best is None or r['score']>best['score']:best={**r,'difference':dd.tolist()}
            if (j+1)%250==0:print(kind,j+1,round(time.time()-started,1),flush=True)
        ge=int(np.sum(np.array(scores)>=out['observed'][0]['score']))
        out['null'][kind]=dict(n=999,at_least_observed=ge,plus_one_fraction=(ge+1)/1000,scores=scores,best=best)
    out['seconds']=time.time()-started
    (ROOT/'joint_stream_results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps({k:v for k,v in out.items() if k in ['observed','seconds']},indent=2))
    print({k:{z:v[z] for z in ['n','at_least_observed','plus_one_fraction']} for k,v in out['null'].items()})
if __name__=='__main__':main()