"""Nonlinear numerical schedules organized by word length/ordinal.

Fixed scope: prime-minus-one and integer totient, indexed by L+j, w+j or w.
j,w are zero-based; prime stream begins 2, totient stream begins phi(1).
Offsets 0..32; decoding a*c + sign*schedule + k mod29, a/sign in {-1,1},
k=0..28. Selection uses only complete prefix words, capped at 256 runes.
"""
import json, time
from word_core import *

OFFSETS=np.repeat(np.arange(33),29)
CONSTANTS=np.tile(np.arange(29),33)


def schedules(source):
    words=tokenize(source)
    length=np.concatenate([np.full(len(w['values']),len(w['values'])) for w in words])
    ordinal=np.concatenate([np.full(len(w['values']),i) for i,w in enumerate(words)])
    inside=np.concatenate([np.arange(len(w['values'])) for w in words])
    indices={'length_plus_position':length+inside,
             'ordinal_plus_position':ordinal+inside,'ordinal_only':ordinal}
    out={}
    for name,idx in indices.items():
        for seqname,seq in [('prime_minus_one',pr-1),('totient',phi[1:])]:
            assert int(idx.max())+32<len(seq)
            out[name+'_'+seqname]=seq[idx[None,:]+np.arange(33)[:,None]]%29
    return out


def sweep(source):
    words=tokenize(source);lengths=np.array([len(w['values']) for w in words])
    cs=np.cumsum(lengths);nw=max(1,int(np.searchsorted(cs,min(256,int(.6*cs[-1])),side='right')))
    cut=int(cs[nw-1]);c=vals(source);streams=schedules(source);leaders=[];count=0
    for family,stream in streams.items():
        for a in [1,-1]:
            for sign in [-1,1]:
                d=(a*c[None,:cut]+sign*stream[OFFSETS,:cut]+CONSTANTS[:,None])%29
                score=FULL_MODEL.score(d,lengths[:nw]);best=np.argsort(score)[-3:][::-1]
                for b in best:
                    leaders.append({'family':family,'a':a,'sign':sign,'offset':int(OFFSETS[b]),
                                    'constant':int(CONSTANTS[b]),'prefix_score':float(score[b])})
                count+=len(score)
    leaders.sort(key=lambda r:r['prefix_score'],reverse=True)
    top=leaders[:10]
    for r in top:
        d=(r['a']*c+r['sign']*streams[r['family']][r['offset']]+r['constant'])%29
        r['suffix_score']=float(FULL_MODEL.score(d[None,cut:],lengths[nw:])[0])
        r['decoded_values']=d.tolist()
        r['prefix_text']=' '.join(''.join(LET[int(v)] for v in d[w['rune_offset']:w['rune_offset']+len(w['values'])]) for w in words[:nw])
        r['suffix_text']=' '.join(''.join(LET[int(v)] for v in d[w['rune_offset']:w['rune_offset']+len(w['values'])]) for w in words[nw:])
    return {'runes':len(c),'prefix_runes':cut,'prefix_words':nw,'settings':count,'top':top}


def main():
    start=time.time();controls=[]
    # LP2 page 57 is not part of the LP1 language-model training corpus.
    source=lp[57];plain=vals(source);streams=schedules(source)
    for family in streams:
        a,sign,k,offset=-1,1,11,7
        cipher=(a*(plain-sign*streams[family][offset]-k))%29
        test=replace_values(source,cipher);result=sweep(test)
        ranks=[i+1 for i,r in enumerate(result['top']) if r['decoded_values']==plain.tolist()]
        controls.append({'family':family,'constructed':{'a':a,'sign':sign,'constant':k,'offset':offset},
                         'exact_recovery_ranks':ranks,'top_prefix_text':result['top'][0]['prefix_text'],
                         'top_suffix_text':result['top'][0]['suffix_text']})
        print('CONTROL',family,ranks,flush=True)
    sections=[]
    for i,s in enumerate('%'.join(lp[:56]).split('$')):
        if len(vals(s))<60:continue
        result=sweep(s);result['section']=i;sections.append(result)
        best=result['top'][0]
        print('SECTION',i,best['family'],best['prefix_score'],best['suffix_score'],flush=True)
        print(' PREFIX',best['prefix_text'][:160],flush=True)
        print(' SUFFIX',best['suffix_text'][:160],flush=True)
    out={'method':__doc__,'controls':controls,'sections':sections,
         'unsolved_settings':sum(s['settings'] for s in sections),'elapsed_seconds':time.time()-start,
         'limits':['Finite schedule families, offsets and constants; counts include duplicates.',
                   'No skipped F runes, plaintext feedback, unknown keys or delimiter key steps.',
                   'Each marked section starts word ordinal zero; wraps do not split words.',
                   'Suffixes were not used to rank or modify parameters.',
                   'These sections have been studied before; suffix is held out from this parameter selection, not a globally untouched corpus.',
                   'Constructed controls verify search capability within these families, not that real text uses them.']}
    (ROOT/'word_schedule_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2)+'\n')
    print('DONE',out['unsolved_settings'],out['elapsed_seconds'],flush=True)


if __name__=='__main__':main()