"""Beam search over selective F skips. Calibrate flexibility with shuffled text.
This is not exhaustive: beam width 8, four streams, offsets 0..32.
"""
import json
from lp_core import *

NAMES=['prime_minus_one','fibonacci','integer_totient','cumulative_prime']
PARAMS=[(n,o,a,sg) for n in NAMES for o in range(33) for a in (1,-1) for sg in (-1,1)]
B=8
def beam(text,reset='',params=PARAMS):
    c=vals(text);M=len(params);n=len(c)
    seq_ids=np.array([NAMES.index(x[0]) for x in params])[:,None]
    seqmat=np.stack([SEQS[n] for n in NAMES])
    offset=np.array([x[1] for x in params])[:,None]
    a=np.array([x[2] for x in params])[:,None];sign=np.array([x[3] for x in params])[:,None]
    _,resets=positions(text,reset)
    k=np.broadcast_to(offset,(M,B)).copy();score=np.full((M,B),-np.inf);score[:,0]=0
    prev=np.zeros((M,B),int);prev2=prev.copy()
    history=np.zeros((n,M,B),np.uint8);parents=np.zeros_like(history);skipped=np.zeros_like(history)
    identity=np.broadcast_to(np.arange(B),(M,B));rows=np.arange(M)[:,None]
    for t,symbol in enumerate(c):
        if resets[t]:k=np.broadcast_to(offset,(M,B)).copy()
        v=(a*symbol+sign*seqmat[seq_ids,k])%29
        normal=score.copy()
        if t>=1:normal+=bi[prev,v]
        if t>=2:normal+=tri[prev2,prev,v]
        if symbol==0:
            skip_score=score.copy()
            if t>=1:skip_score+=bi[prev,0]
            if t>=2:skip_score+=tri[prev2,prev,0]
            both=np.concatenate([normal,skip_score],axis=1)
            chosen=np.argsort(both,axis=1)[:,-B:][:,::-1]
            par=chosen%B;is_skip=chosen>=B
            score=np.take_along_axis(both,chosen,axis=1)
            v=np.where(is_skip,0,v[rows,par])
            k=k[rows,par]+(~is_skip)
            prev2=prev[rows,par];prev=v
            parents[t]=par;skipped[t]=is_skip
        else:
            score=normal;k+=1;prev2=prev;prev=v;parents[t]=identity
        history[t]=v
    model,slot=np.unravel_index(np.argmax(score),score.shape)
    out=np.zeros(n,int);skiplist=[]
    for t in range(n-1,-1,-1):
        out[t]=history[t,model,slot]
        if skipped[t,model,slot]:skiplist.append(t)
        slot=int(parents[t,model,slot])
    seqname,off,mul,sg=params[model]
    # Verify the recovered path independently, including every reset and skip.
    cursor=off;recovered=[]
    for t,x in enumerate(c):
        if resets[t]:cursor=off
        if t in skiplist:assert x==0;recovered.append(0)
        else:recovered.append(int((mul*x+sg*SEQS[seqname][cursor])%29));cursor+=1
    assert np.array_equal(out,recovered)
    return {'sequence':seqname,'offset':off,'a':mul,'sign':sg,'reset':reset,
            'beam_score':float(score.max()/max(1,n)),'runes':n,'skipped_F_positions':sorted(skiplist),
            'text':render(text,out),'output_runes':''.join(R[x] for x in out)}

def shuffle_text(text,rng):
    shuffled=iter(rng.permutation(list(''.join(c for c in text if c in R))))
    return ''.join(next(shuffled) if c in R else c for c in text)

if __name__=='__main__':
    control=beam(lp[56])
    assert control['skipped_F_positions']==[56],control
    assert control['text'].endswith('IT IS THE DUTY OF EUERY PILGRIM TO SEEK OUT THIS PAGE'),control
    print('Selective-F control recovered exactly',flush=True)
    out={'control':control,'model_settings_per_reset':len(PARAMS),'beam_width':B,'targets':[],
         'limits':['Search chooses skips using a small language model; readable fragments can be manufactured by this flexibility.','Beam pruning can miss the best path.','Offsets only 0..32, four streams, 32 shuffled comparators; no exhaustive negative claim.','Shuffle calibration compares the same complete model/reset selection process.','Beam score is not a likelihood that a solution is correct.']}
    for label,text in [('first_three_pages','%'.join(lp[:3])),('last_two_unsolved_pages','%'.join(lp[54:56]))]:
        resets=['','&$','.;&$']
        actual=[beam(text,r) for r in resets];best=max(actual,key=lambda x:x['beam_score'])
        rng=np.random.default_rng(3301);null=[]
        for i in range(32):
            shuffled=shuffle_text(text,rng)
            null.append(max(beam(shuffled,r)['beam_score'] for r in resets))
            if i%8==7:print(label,'shuffles',i+1,flush=True)
        p=float((1+sum(x>=best['beam_score'] for x in null))/(len(null)+1))
        out['targets'].append({'label':label,'best':best,'all_reset_results':actual,'shuffle_max_scores':null,'empirical_p':p})
        (ROOT/'skip_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2))
        print(label,'score',best['beam_score'],'shuffle_p',p,flush=True)