"""Session 41: equality partitions for fixed invertible word-by-word ciphers.
Python + NumPy, retained sources only. Descriptive finite-corpus comparisons.
"""
from pathlib import Path
from collections import Counter,defaultdict
import hashlib,json,math
import numpy as np
from word_core import R,LET,vals,lp,GROUPS,tokenize
ROOT=Path(__file__).resolve().parent
SEED=3301041
EXPECTED='b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'

def extract(s):return [tuple(w['values']) for w in tokenize(s)]
def collisions(c):return sum(f*(f-1)//2 for f in c.values())
def inventory(ws):
    counts=Counter(ws);n=len(ws);pairs=collisions(counts)
    return dict(words=n,types=len(counts),repeated_types=sum(f>1 for f in counts.values()),
        tokens_in_repeated_types=sum(f for f in counts.values() if f>1),equal_word_pairs=pairs,
        maximum_word_count=max(counts.values(),default=0),
        equality_partition=sorted(counts.values(),reverse=True),
        pair_equality_rate=pairs/math.comb(n,2) if n>1 else None)

def groups(ws):
    out=defaultdict(list)
    for w in ws:out[len(w)].append(w)
    return dict(out)

def describe(ws):
    return dict(runes=sum(map(len,ws)),**inventory(ws),
        by_length=[dict(length=L,**inventory(w)) for L,w in sorted(groups(ws).items())])

def mean_sample_pairs(ws,n):
    N=len(ws)
    return collisions(Counter(ws))*n*(n-1)/(N*(N-1)) if N>1 else 0.0

def matched(reference,target,seed,trials=1999):
    refs=groups(reference);targets=groups(target);lengths=sorted(refs.keys()&targets.keys());rows=[];pools=[]
    for L in lengths:
        r=refs[L];t=targets[L];n=min(len(r),len(t))
        def ids(ws):
            vocab={w:i for i,w in enumerate(dict.fromkeys(ws))}
            return np.array([vocab[w] for w in ws])
        rows.append(dict(length=L,reference_words=len(r),target_words=len(t),sampled_from_each=n,
            exact_expected_reference_pairs=mean_sample_pairs(r,n),exact_expected_target_pairs=mean_sample_pairs(t,n)))
        pools.append((ids(r),ids(t),n))
    rng=np.random.default_rng(seed);samples=np.zeros((trials,2),dtype=np.int64)
    for trial in range(trials):
        for r,t,n in pools:
            for j,pool in enumerate([r,t]):
                selected=pool if n==len(pool) else rng.choice(pool,n,replace=False)
                counts=np.bincount(selected);samples[trial,j]+=np.sum(counts*(counts-1)//2)
    expected=[sum(row[key] for row in rows) for key in ['exact_expected_reference_pairs','exact_expected_target_pairs']]
    return dict(seed=seed,subsamples=trials,words_in_each_matched_sample=sum(x['sampled_from_each'] for x in rows),
        method='Independent uniform sampling without replacement within each word length, to min(reference count,target count). Descriptive sensitivity check, not a hypothesis-test null or fresh holdout.',
        per_length=rows,exact_expected_equal_word_pairs=dict(reference=expected[0],target=expected[1]),
        central_95_percent_subsample_ranges=dict(reference=np.quantile(samples[:,0],[.025,.975]).tolist(),target=np.quantile(samples[:,1],[.025,.975]).tolist()),
        sampled_ranges=dict(reference=[int(samples[:,0].min()),int(samples[:,0].max())],target=[int(samples[:,1].min()),int(samples[:,1].max())]),
        samples=samples.tolist())

def verify(reference):
    rng=np.random.default_rng(SEED+100);forward={};inverse={};encoded=[]
    # Random length-preserving codebook on observed vocabulary; injectivity verified.
    for w in reference:
        if w not in forward:
            while True:
                c=tuple(int(x) for x in rng.integers(0,29,len(w)))
                if c not in inverse:break
            forward[w]=c;inverse[c]=w
        encoded.append(forward[w])
    recovered=[inverse[w] for w in encoded]
    assert recovered==reference
    assert describe(encoded)==describe(reference)
    # Stateful counterexample: same word may encrypt differently at a new index.
    changing=[tuple((x+i+5*j)%29 for j,x in enumerate(w)) for i,w in enumerate(reference)]
    decoded=[tuple((x-i-5*j)%29 for j,x in enumerate(w)) for i,w in enumerate(changing)]
    assert decoded==reference
    assert inventory(changing)['equal_word_pairs']!=inventory(reference)['equal_word_pairs']
    fixture=[(1,),(1,),(1,),(2,),(2,),(0,1),(0,1),(1,0)]
    brute=sum(a==b for i,a in enumerate(fixture) for b in fixture[i+1:])
    assert brute==collisions(Counter(fixture))==5
    # A slash/page wrap continues a word; a hyphen ends it.
    assert extract(R[0]+'/%'+R[1]+'-'+R[2])==[(0,1),(2,)]
    return dict(random_codebook=dict(seed=SEED+100,roundtrip_exact=True,all_lengthwise_equality_partitions_preserved=True,
        vocabulary_entries=len(forward),plaintext_pairs=inventory(reference)['equal_word_pairs'],ciphertext_pairs=inventory(encoded)['equal_word_pairs']),
        position_dependent_counterexample=dict(formula='c[word_index,j]=(p[word_index,j]+word_index+5*j) mod 29',
            roundtrip_exact=True,plaintext=describe(reference),ciphertext=describe(changing),
            note='Outside the tested fixed-word mapping family. Demonstrates that changing state can alter repetition while retaining exact recoverability; it is not a proposed LP solution.'),
        pair_counter_bruteforce_fixture=True,page_wrap_tokenization_fixture=True)

def main():
    source_hash=hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest();assert source_hash==EXPECTED
    sections=[(i,extract(s)) for i,s in enumerate('%'.join(lp[:56]).split('$')) if len(vals(s))>100]
    assert [i for i,ws in sections]==[0,1,2,4,5,6,7,8]
    target=sum([ws for i,ws in sections],[]);reference=sum([extract(s) for s in GROUPS.values()],[])
    assert len(target)==2897 and sum(map(len,target))==12947
    assert len(reference)==483 and sum(map(len,reference))==1963
    assert len(set(w for w in target if len(w)==1))==29
    assert all(f==1 for w,f in Counter(target).items() if len(w)>=4)
    per=[]
    for i,ws in sections:
        per.append(dict(section=i,inventory=describe(ws),matched=matched(reference,ws,SEED+1+i)))
    # All-corpus descriptive inventory confirms the Session 4 scope separately.
    all_target=extract('%'.join(lp[:56]));assert len(all_target)==2899
    singles=defaultdict(list)
    for w in tokenize('%'.join(lp[:56])):
        if len(w['values'])==1:
            singles[w['values'][0]].append({key:w[key] for key in ['page','section','page_rune_offset','word_index']})
    result=dict(seed=SEED,hypothesis='Within each tested key scope, an invertible, deterministic, length-preserving mapping acts on each complete word independently, with no other state.',
        invariant='Equality iff equality: P_i=P_j iff E(P_i)=E(P_j). Thus, for each word length, the multiset of word occurrence counts is exactly preserved.',
        scope='Primary target: eight substantial sections in LP2 pages 0..55, excluding nine-rune section 3. Benchmark: five established solved LP1 groups, not every solved page.',
        target=describe(target),reference=describe(reference),full_56_page_inventory=describe(all_target),
        reference_groups=[dict(name=name,**describe(extract(s))) for name,s in GROUPS.items()],
        matched_pooled=matched(reference,target,SEED),per_section=per,
        single_rune_word_locations=[dict(value=i,rune=R[i],label=LET[i],occurrences=singles[i]) for i in range(29)],
        controls=verify(reference),source_sha256=source_hash,script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        limitations=['Equality preservation is a theorem under the stated model; resemblance to the solved text is an empirical assumption, not a theorem about all languages.',
            'Different section keys invalidate cross-section equality comparisons; per-section comparisons are supplied.',
            'Rune lengths differ from Latin character lengths: for example THE occupies two Gematria runes.',
            'Tokenization trusts the retained transcription and rune-only parser. No new complete image transcription or signature authentication.',
            'Subsampling is descriptive finite-corpus sensitivity, not p-values, independent prose draws or multiple-testing-corrected significance.',
            'Position-dependent keys, nonces, changing state, homophonic outputs, added nulls, altered word boundaries and intermediate encoded text are outside scope.',
            'No new plaintext, key, language identification or established new-to-community observation. The absence of long repeated words was already recorded in Session 4.'])
    (ROOT/'whole_word_repetition_results.json').write_text(json.dumps(result,indent=2)+'\n')
    def brief(d):return {k:d[k] for k in ['words','types','equal_word_pairs','maximum_word_count']}
    def comp(m):return {k:v for k,v in m.items() if k not in ['samples','per_length']}
    summary=dict(target=brief(result['target']),reference=brief(result['reference']),
        by_length=[dict(length=x['length'],**brief(x)) for x in result['target']['by_length']],
        reference_by_length=[dict(length=x['length'],**brief(x)) for x in result['reference']['by_length']],
        matched=comp(result['matched_pooled']),sections=[dict(section=x['section'],
            single_rune_words=next((q for q in x['inventory']['by_length'] if q['length']==1),None),matched=comp(x['matched'])) for x in per],
        codebook_control=result['controls']['random_codebook'],counterexample_pairs=result['controls']['position_dependent_counterexample']['ciphertext']['equal_word_pairs'])
    print(json.dumps(summary,indent=2))
if __name__=='__main__':main()