"""Session 69: compact selectors for fixed homophonic decoding tables.

Exact necessary singleton-mass bounds, not a decoding-table/key recovery.
Python and NumPy only. Existing master and reports remain unchanged.
"""
from pathlib import Path
from itertools import product
from fractions import Fraction
import hashlib
import json
import numpy as np
from lp_core import vals, LET
from word_core import GROUPS, replace_values
from lexical_words import tokenize_words
from reset_state_audit import SOURCE, SINGLE

ROOT = Path(__file__).resolve().parent
Q = 29
FEATURES = {'previous': 1, 'sum_two': 2, 'difference_two': 2, 'sum_three': 3}
ALLOWED = {'AI': [24, 10], 'AIO': [24, 10, 3]}
PARAMS = [(a, b) for a in range(1, Q) for b in range(Q)]
RED = {(6, 75), (54, 0)}
TRAIN = [0, 1, 2, 4]
EVAL = [5, 6, 7, 8]


def summary(history, feature):
    if feature == 'previous': return history[-1]
    if feature == 'sum_two': return (history[-1]+history[-2]) % Q
    if feature == 'difference_two': return (history[-1]-history[-2]) % Q
    return sum(history[-3:]) % Q


def data_for(source, singles, feature):
    m = FEATURES[feature]
    c = vals(source)
    section = np.concatenate([np.full(len(vals(s)), i, dtype=int)
                              for i, s in enumerate(source.split('$'))])
    previous = [np.r_[np.zeros(j, dtype=int), c[:-j]] for j in range(1, m+1)]
    if feature == 'previous': z = previous[0]
    elif feature == 'sum_two': z = (previous[0]+previous[1]) % Q
    elif feature == 'difference_two': z = (previous[0]-previous[1]) % Q
    else: z = sum(previous) % Q
    valid = np.zeros(len(c), dtype=bool)
    valid[m:] = section[m:] == section[:-m]
    singleton_offsets = np.array([w['rune_offset'] for w in singles], dtype=int)
    singleton_offsets = singleton_offsets[valid[singleton_offsets]]
    records = {}
    for sec in np.unique(section):
        take = valid & (section == sec)
        if not take.any(): continue
        h = np.bincount(Q*z[take]+c[take], minlength=Q*Q).reshape(Q, Q)
        positions = singleton_offsets[section[singleton_offsets] == sec]
        s = np.bincount(Q*z[positions]+c[positions], minlength=Q*Q).reshape(Q, Q)
        records[int(sec)] = {'h': h, 's': s}
    return records


def combine(records, sections=None):
    chosen = records.values() if sections is None else [records[s] for s in sections if s in records]
    chosen = list(chosen)
    return {'h': sum((r['h'] for r in chosen), np.zeros((Q, Q), dtype=np.int64)),
            's': sum((r['s'] for r in chosen), np.zeros((Q, Q), dtype=np.int64))}


def selectors(k):
    return np.array([((a*np.arange(Q)+b) % Q) % k for a, b in PARAMS])


def masses(record, bank, k):
    out = np.zeros(len(bank), dtype=np.int64)
    for state in range(k):
        membership = (bank == state).astype(np.int64)
        h = membership @ record['h']
        constrained = membership @ record['s'] > 0
        out += np.sum(h*constrained, axis=1)
    return out


def collapse(record, selector, k):
    h = np.zeros((k, Q), dtype=np.int64)
    s = np.zeros((k, Q), dtype=np.int64)
    np.add.at(h, selector, record['h'])
    np.add.at(s, selector, record['s'])
    return h, s


def benchmark(feature):
    m = FEATURES[feature]
    out = {}
    for name, allowed in ALLOWED.items():
        groups = []
        for group, text in GROUPS.items():
            p = vals(text)[m:]
            count = sum(int(x) in allowed for x in p)
            groups.append({'name': group, 'positions': len(p), 'allowed_count': count,
                           'share': count/len(p)})
        pooled = Fraction(sum(g['allowed_count'] for g in groups), sum(g['positions'] for g in groups))
        generous = max(Fraction(g['allowed_count'], g['positions']) for g in groups)
        out[name] = {'groups': groups, 'pooled_ratio': [pooled.numerator, pooled.denominator],
                     'largest_group_ratio': [generous.numerator, generous.denominator],
                     'pooled_share': float(pooled), 'largest_group_share': float(generous)}
    return out


def describe(mass, n, bench):
    best = int(np.argmin(mass))
    return {'scored': n, 'minimum_forced_positions': int(mass[best]),
            'minimum_forced_share': float(mass[best]/n),
            'first_best_selector': dict(zip(['a', 'b'], PARAMS[best])),
            'settings': len(mass),
            'passing_largest_group_reference': {
                name: int(np.sum(mass <= n*b['largest_group_ratio'][0]//b['largest_group_ratio'][1]))
                for name, b in bench.items()}}


def transfer(records, bank, k):
    train, test = combine(records, TRAIN), combine(records, EVAL)
    train_mass = masses(train, bank, k)
    chosen = int(np.argmin(train_mass))  # fixed lexicographic tie rule
    th, ts = collapse(train, bank[chosen], k)
    eh, es = collapse(test, bank[chosen], k)
    mask = ts > 0
    joint_mask = mask | (es > 0)
    return {'selector_selected_by_training_mass_only': dict(zip(['a', 'b'], PARAMS[chosen])),
            'training_positions': int(th.sum()), 'training_singletons': int(ts.sum()),
            'training_forced_share': float(th[mask].sum()/th.sum()),
            'evaluation_positions': int(eh.sum()), 'evaluation_singletons': int(es.sum()),
            'evaluation_positions_forced_by_frozen_training_cells': int(eh[mask].sum()),
            'evaluation_singletons_in_frozen_training_cells': int(es[mask].sum()),
            'evaluation_forced_share_after_its_singletons_are_imposed': float(eh[es > 0].sum()/eh.sum()),
            'combined_forced_share_with_all_singleton_constraints': float((th[joint_mask].sum()+eh[joint_mask].sum())/(th.sum()+eh.sum())),
            'evaluation_new_constrained_cells': int(np.sum((es > 0) & ~mask)),
            'decoder_not_fitted_or_recovered': True}


def exhaustive_control():
    h = np.array([[4, 2, 5], [3, 7, 1], [6, 1, 8]])
    s = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    checked, details = 0, []
    for k in [1, 2]:
        for a, b in product([1, 2], range(3)):
            selector = ((a*np.arange(3)+b) % 3) % k
            hh, ss = np.zeros((k, 3), int), np.zeros((k, 3), int)
            np.add.at(hh, selector, h); np.add.at(ss, selector, s)
            predicted = int(hh[ss > 0].sum())
            minimum = 10**9
            for flat in product(range(3), repeat=3*k):
                mp = np.array(flat).reshape(k, 3)
                if np.any((ss > 0) & (mp != 0)): continue
                minimum = min(minimum, int(hh[mp == 0].sum()))
            assert minimum == predicted
            checked += 3**(3*k)
            details.append({'tables': k, 'a': a, 'b': b, 'minimum_mass': minimum})
    return {'decoder_tables_enumerated': checked, 'classifier_cases': len(details),
            'all_exact_mass_extrema_agree': True, 'details': details}


def controls():
    rng = np.random.default_rng(69001)
    letters = sorted({int(x) for text in GROUPS.values() for x in vals(text)})
    out = []
    for feature in FEATURES:
        for k in [2, 4, 8]:
            a, b = 7, 11
            codebooks, inverses = [], []
            for state in range(k):
                codes = list(map(int, rng.permutation(Q)))
                classes = {p: [codes.pop()] for p in letters}
                for p in rng.choice(letters, size=Q-len(letters), replace=False):
                    classes[int(p)].append(codes.pop())
                assert not codes
                inverse = {c: p for p, cs in classes.items() for c in cs}
                codebooks.append(classes); inverses.append(inverse)
            texts, exact, plain_doubles, cipher_doubles = [], True, 0, 0
            for text in GROUPS.values():
                p = vals(text); history = [0, 0, 0]; cipher = []
                for value in p:
                    state = ((a*summary(history, feature)+b) % Q) % k
                    options = [c for c in codebooks[state][int(value)] if c != history[-1]]
                    if not options: options = codebooks[state][int(value)]
                    c = int(rng.choice(options))
                    cipher.append(c)
                    history = history[1:]+[c]
                # Decode in a separate pass, reconstructing state solely from
                # the completed ciphertext, without encoder-side history.
                recovered, decode_history = [], [0, 0, 0]
                for c in cipher:
                    state = ((a*summary(decode_history, feature)+b) % Q) % k
                    recovered.append(inverses[state][c])
                    decode_history = decode_history[1:]+[c]
                assert np.array_equal(p, recovered)
                plain_doubles += int(np.sum(p[1:] == p[:-1]))
                cipher_doubles += sum(x == y for x, y in zip(cipher, cipher[1:]))
                texts.append(replace_values(text, cipher))
            source = '$'.join(texts)
            single = [w for w in tokenize_words(source) if len(w['values']) == 1]
            records = data_for(source, single, feature)
            record = combine(records)
            selector = ((a*np.arange(Q)+b) % Q) % k
            h, s = collapse(record, selector, k)
            forced = int(h[s > 0].sum())
            m = FEATURES[feature]
            true_ai = sum(sum(int(x) in ALLOWED['AI'] for x in vals(t)[m:]) for t in GROUPS.values())
            assert forced <= true_ai
            bank = selectors(k); idx = PARAMS.index((a, b))
            assert masses(record, bank, k)[idx] == forced
            bench = benchmark(feature)
            ratio = bench['AI']['largest_group_ratio']
            assert forced*ratio[1] <= int(h.sum())*ratio[0]
            out.append({'feature': feature, 'tables': k, 'a': a, 'b': b,
                        'exact_known_table_inverse_recovery': exact,
                        'planted_selector_survives_mass_screen': True,
                        'forced_AI_positions': forced, 'true_AI_positions': true_ai,
                        'scored': int(h.sum()), 'plain_neighbor_repeats': plain_doubles,
                        'cipher_neighbor_repeats': cipher_doubles,
                        'decoder_tables': [[inv[c] for c in range(Q)] for inv in inverses],
                        'unknown_key_recovery_not_demonstrated': True})
    return out


def phrase_states(selector, feature):
    c = vals(SOURCE); m = FEATURES[feature]
    rows = []
    for start in [6555, 12950]:
        rows.append([int(selector[summary(list(c[i-m:i]), feature)]) for i in range(start, start+6)])
    return rows


def main():
    digest = hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    assert digest == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    assert len(SINGLE) == 83
    records = []
    refs = {f: benchmark(f) for f in FEATURES}
    for sensitivity in ['all_audited', 'exclude_red_initials']:
        singles = SINGLE if sensitivity == 'all_audited' else [w for w in SINGLE if (w['page'], w['page_rune_offset']) not in RED]
        for feature in FEATURES:
            data = data_for(SOURCE, singles, feature); pooled = combine(data)
            n = int(pooled['h'].sum()); by_k = []
            for k in range(1, Q+1):
                bank = selectors(k)
                mass = masses(pooled, bank, k)
                per_section = {s: masses(v, bank, k) for s, v in data.items()}
                common_selector = sum(per_section.values())
                independent_mass = sum(int(x.min()) for x in per_section.values())
                section_summary = {str(s): describe(x, int(data[s]['h'].sum()), refs[feature]) for s, x in per_section.items()}
                first = int(np.argmin(mass))
                phase_rows = phrase_states(bank[first], feature)
                assert phase_rows[0][FEATURES[feature]:] == phase_rows[1][FEATURES[feature]:]
                by_k.append({'tables': k, 'shared_decoding_tables': describe(mass, n, refs[feature]),
                             'independent_section_tables_common_selector': describe(common_selector, n, refs[feature]),
                             'independent_section_tables_and_selectors_minimum_forced_share': independent_mass/n,
                             'by_section': section_summary,
                             'frozen_selector_transfer': transfer(data, bank, k),
                             'best_full_corpus_selector_phrase_states': phase_rows,
                             'guaranteed_matching_plaintext_phrase_suffix_length': 6-FEATURES[feature]})
            records.append({'singleton_sensitivity': sensitivity, 'feature': feature,
                            'memory': FEATURES[feature], 'scored': n,
                            'scored_singletons': int(pooled['s'].sum()), 'results': by_k})
            print(sensitivity, feature, 'scored', n, 'first passing',
                  {name: next(x['tables'] for x in by_k if x['shared_decoding_tables']['passing_largest_group_reference'][name]) for name in ALLOWED}, flush=True)
    out = {'master_sha256': digest,
           'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
           'protocol': {'state': 'z is previous C, sum of previous two C mod 29, previous C minus second-previous C mod 29, or sum of previous three C mod 29. state = ((a*z+b) mod 29) mod k.',
                        'parameter_bank': {'a': [1, 28], 'b': [0, 28], 'k': [1, 29], 'settings_per_feature': len(PARAMS)*Q, 'total_classifier_settings_per_sensitivity': len(PARAMS)*Q*len(FEATURES)},
                        'decoder': 'D[state,C] is fixed but may be many-to-one. One shared bank or stated independent section banks. Encoder may choose homophones using state.',
                        'bound': 'Any state/current-C cell observed as a retained singleton must decode in A/I or A/I/O. Every corpus occurrence in that cell is forced into the allowed set.',
                        'scope': 'LP2 pages 0-55, all nine nonempty sections. Omit first m runes of each section. No word/page resets or skipped updates.',
                        'transfer': 'Choose minimum-forced-mass selector on sections 0,1,2,4, ties ascending a then b. Freeze selector and training constrained-cell mask; evaluate 5,6,7,8. New evaluation constraints are separately recorded. No decoding table is fitted.',
                        'limits': 'Necessary mass screen using empirical solved-text references. No actual English decoder, unknown key recovery, blind holdout or statistical significance claim. Numeric summaries and affine selectors do not cover all state functions.'},
           'benchmarks': refs, 'records': records,
           'exhaustive_control': exhaustive_control(), 'recovery_controls': controls()}
    (ROOT/'context_switch_capacity_results.json').write_text(json.dumps(out, indent=2)+'\n')
    print('Controls', len(out['recovery_controls']), 'exhaustive', out['exhaustive_control']['decoder_tables_enumerated'])


if __name__ == '__main__': main()