"""Session 70: a specified repeat channel in section-specific homophonic banks.

State = affine grouping of previous ciphertext. Same-group transitions decode
to R; different-group transitions decode outside R. Exact necessary screens,
chronological selector transfer and known-table inverse controls. No new key.
"""
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
PARAMS = [(a, b) for a in range(1, Q) for b in range(Q)]
ALLOWED = {'AI': [24, 10], 'AIO': [24, 10, 3]}
RED = {(6, 75), (54, 0)}


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


def matrices(cipher, singleton_offsets, start, stop):
    # Current positions start..stop-1; the predecessor at a split is retained.
    offsets = np.arange(max(1, start), stop)
    h = np.bincount(Q*cipher[offsets-1]+cipher[offsets], minlength=Q*Q).reshape(Q, Q)
    ss = [i for i in singleton_offsets if max(1, start) <= i < stop]
    s = np.bincount([Q*cipher[i-1]+cipher[i] for i in ss], minlength=Q*Q).reshape(Q, Q)
    return {'h': h, 's': s}


def metrics(record, selectors, k):
    r_mass = np.zeros(len(selectors), dtype=np.int64)
    r_singles = np.zeros(len(selectors), dtype=np.int64)
    allowed_out = np.zeros(len(selectors), dtype=np.int64)
    allowed_in = np.zeros(len(selectors), dtype=np.int64)
    for state in range(k):
        membership = (selectors == state).astype(np.int64)
        h = membership @ record['h']; s = membership @ record['s']
        channel = selectors == state
        r_mass += np.sum(h*channel, axis=1)
        r_singles += np.sum(s*channel, axis=1)
        allowed_out += np.sum(h*(s > 0), axis=1)
        allowed_in += np.sum(h*((s > 0) | channel), axis=1)
    return {'r_positions': r_mass, 'channel_singletons': r_singles,
            'allowed_outside_R': allowed_out, 'allowed_inside_R': allowed_in,
            'scored': int(record['h'].sum()), 'singletons': int(record['s'].sum())}


def references():
    histograms = {name: np.bincount(vals(t)[1:], minlength=Q) for name, t in GROUPS.items()}
    out = {}
    for name, allowed in ALLOWED.items():
        shares = [Fraction(sum(int(h[x]) for x in allowed), int(h.sum())) for h in histograms.values()]
        s_cap = max(shares)
        categories = {}
        for inside in [False, True]:
            labels = [i for i in range(Q) if (i in allowed) == inside]
            candidates = [(Fraction(int(h[c]), int(h.sum())), group, c)
                          for group, h in histograms.items() for c in labels]
            ratio, group, c = max(candidates)
            categories['inside' if inside else 'outside'] = {
                'R_ratio': [ratio.numerator, ratio.denominator], 'R_share': float(ratio),
                'reference_group': group, 'reference_label': LET[c]}
        out[name] = {'allowed_share': float(s_cap), 'allowed_ratio': [s_cap.numerator, s_cap.denominator],
                     'R_categories': categories}
    return out


def admissible(m, ref, inside):
    category = 'inside' if inside else 'outside'
    r_ratio = ref['R_categories'][category]['R_ratio']
    s_ratio = ref['allowed_ratio']
    forced = m['allowed_inside_R' if inside else 'allowed_outside_R']
    ok = (m['r_positions']*r_ratio[1] <= m['scored']*r_ratio[0]) & (forced*s_ratio[1] <= m['scored']*s_ratio[0])
    if not inside: ok &= m['channel_singletons'] == 0
    return ok


def scalar(m, index, inside):
    allowed = int(m['allowed_inside_R' if inside else 'allowed_outside_R'][index])
    return {'scored': m['scored'], 'singletons': m['singletons'],
            'R_positions': int(m['r_positions'][index]),
            'R_share': float(m['r_positions'][index]/m['scored']),
            'channel_singletons': int(m['channel_singletons'][index]),
            'forced_allowed_positions': allowed, 'forced_allowed_share': allowed/m['scored']}


def evaluate(records, k, name, ref):
    selectors = bank(k)
    measured = {part: metrics(r, selectors, k) for part, r in records.items()}
    counts = {}
    for part, m in measured.items():
        counts[part] = {category: int(admissible(m, ref, inside).sum())
                       for category, inside in [('outside', False), ('inside', True)]}
    m = measured['full']
    result = {'tables': k, 'minimum_R_share': float(m['r_positions'].min()/m['scored']),
              'settings_per_category': len(PARAMS), 'passing_by_part_and_R_category': counts,
              'maximum_encodable_plaintext_types_by_state': [Q-int(np.sum(selectors[0] == j))+1 for j in range(k)]}
    options = []
    train = measured['train']
    for inside in [False, True]:
        for index in np.flatnonzero(admissible(train, ref, inside)):
            forced = train['allowed_inside_R' if inside else 'allowed_outside_R'][index]
            options.append((int(forced), int(train['r_positions'][index]), int(inside), int(index)))
    if options:
        _, _, inside_int, index = min(options)
        inside = bool(inside_int)
        a, b = PARAMS[index]
        test_ok = bool(admissible(measured['evaluation'], ref, inside)[index])
        full_ok = bool(admissible(measured['full'], ref, inside)[index])
        result['best_training_selector_frozen'] = {
            'a': a, 'b': b, 'R_category': 'inside' if inside else 'outside',
            'training': scalar(train, index, inside),
            'evaluation': scalar(measured['evaluation'], index, inside),
            'full': scalar(measured['full'], index, inside),
            'evaluation_passes_necessary_screens': test_ok,
            'full_passes_necessary_screens': full_ok,
            'hard_evaluation_singleton_contradictions': 0 if inside else int(measured['evaluation']['channel_singletons'][index]),
            'decoder_table_not_fitted': True}
    return result


def small_exhaustive_control():
    h = np.array([[4, 2, 5], [3, 7, 1], [6, 1, 8]])
    s = np.array([[0, 1, 0], [0, 0, 1], [0, 0, 0]])
    checked, cases = 0, []
    for k in [1, 2, 3]:
        for a, b in product([1, 2], range(3)):
            g = ((a*np.arange(3)+b) % 3) % k
            channel = np.arange(k)[:, None] == g[None, :]
            hh, ss = np.zeros((k, 3), int), np.zeros((k, 3), int)
            np.add.at(hh, g, h); np.add.at(ss, g, s)
            free = list(zip(*np.where(~channel)))
            minima = {False: None, True: None}
            for outputs in product([0, 1], repeat=len(free)):
                mp = np.full((k, 3), 2)
                for cell, value in zip(free, outputs): mp[cell] = value
                checked += 1
                for inside, allowed in [(False, {0}), (True, {0, 2})]:
                    allowed_mask = np.isin(mp, list(allowed))
                    if np.any((ss > 0) & ~allowed_mask): continue
                    value = int(hh[allowed_mask].sum())
                    if minima[inside] is None or value < minima[inside]: minima[inside] = value
            for inside in [False, True]:
                contradiction = not inside and np.any((ss > 0) & channel)
                predicted = None if contradiction else int(hh[(ss > 0) | channel].sum()) if inside else int(hh[ss > 0].sum())
                assert minima[inside] == predicted
            cases.append({'tables': k, 'a': a, 'b': b, 'exact_extrema_agree': True})
    return {'decoder_assignments_enumerated': checked, 'classifier_cases': len(cases),
            'all_mass_bounds_and_hard_contradictions_agree': True}


def checksum_collision_controls():
    # Any previous rune can occur at any sum/difference state by choosing an
    # older rune appropriately. A universal repeat=>R rule fills every cell.
    k, a, b = 8, 7, 11
    out = []
    for feature, m in [('previous', 1), ('sum_two', 2), ('difference_two', 2), ('sum_three', 3)]:
        forced = np.zeros((k, Q), bool)
        for history in product(range(Q), repeat=m):
            previous = history[-1]
            if feature == 'previous': z = previous
            elif feature == 'sum_two': z = sum(history) % Q
            elif feature == 'difference_two': z = (history[-1]-history[-2]) % Q
            else: z = sum(history) % Q
            state = ((a*z+b) % Q) % k
            forced[state, previous] = True
        if feature == 'previous': assert np.all(forced.sum(0) == 1)
        else: assert forced.all()
        out.append({'summary': feature, 'history_cases': Q**m, 'tables': k,
                    'forced_R_cells': int(forced.sum()),
                    'constant_R_decoder_if_universal_repeat_rule': feature != 'previous'})
    return out


def recovery_controls(refs):
    rng = np.random.default_rng(70001)
    r = 18  # E, deliberately a common plaintext letter
    controls = []
    for k in [6, 8, 29]:
        a, b = 7, 11
        g = ((a*np.arange(Q)+b) % Q) % k
        section_banks = {}
        for name, text in GROUPS.items():
            ordinary = sorted(set(map(int, vals(text)))-{r})
            mapping = np.zeros((k, Q), int)
            classes = []
            for state in range(k):
                inside = np.flatnonzero(g == state).tolist()
                outside = list(map(int, rng.permutation(np.flatnonzero(g != state))))
                assert len(outside) >= len(ordinary)
                codebook = {r: inside}
                for p in ordinary: codebook[p] = [outside.pop()]
                for c in outside: codebook[int(rng.choice(ordinary))].append(c)
                for p, codes in codebook.items(): mapping[state, codes] = p
                assert np.all(mapping[state, g == state] == r)
                assert np.all(mapping[state, g != state] != r)
                classes.append(codebook)
            section_banks[name] = {'mapping': mapping, 'classes': classes}
        for repeat_probability in [0., .25]:
            encoded, doubles, stays = [], 0, 0
            for name, text in GROUPS.items():
                mapping, classes = section_banks[name]['mapping'], section_banks[name]['classes']
                plain = vals(text); previous = 0; cipher = []
                for p in plain:
                    options = classes[int(g[previous])][int(p)]
                    alternative = [c for c in options if c != previous]
                    if int(p) == r and (not alternative or rng.random() < repeat_probability): c = previous
                    else: c = int(rng.choice(alternative or options))
                    cipher.append(c); previous = c
                decoded, previous = [], 0
                for c in cipher:
                    decoded.append(int(mapping[int(g[previous]), c])); previous = c
                assert np.array_equal(decoded, plain)
                cc = np.array(cipher)
                same = g[cc[:-1]] == g[cc[1:]]
                assert np.array_equal(same, plain[1:] == r)
                doubles += int(np.sum(cc[:-1] == cc[1:])); stays += int(same.sum())
                encoded.append(replace_values(text, cipher))
            source = '$'.join(encoded)
            words = tokenize_words(source); singles = [w for w in words if len(w['values']) == 1]
            base = 0
            h = np.zeros((Q, Q), int); s = np.zeros((Q, Q), int)
            for sec, text in enumerate(source.split('$')):
                p = vals(text)
                offsets = [w['rune_offset']-base for w in singles if base <= w['rune_offset'] < base+len(p)]
                rec = matrices(p, offsets, 1, len(p)); h += rec['h']; s += rec['s']; base += len(p)
            m = metrics({'h': h, 's': s}, bank(k), k); index = PARAMS.index((a, b))
            assert m['r_positions'][index] == stays and m['channel_singletons'][index] == 0
            for name, ref in refs.items(): assert admissible(m, ref, False)[index]
            controls.append({'tables': k, 'a': a, 'b': b, 'R_label': LET[r],
                             'repeat_probability_when_alternative_exists': repeat_probability,
                             'scored': m['scored'], 'exact_known_table_inverse_recovery': True,
                             'same_group_positions': stays, 'plaintext_E_positions': stays,
                             'exact_cipher_neighbor_repeats': doubles,
                             'planted_selector_passes_AI_and_AIO_screens': True,
                             'independent_section_tables': {name: v['mapping'].tolist() for name, v in section_banks.items()},
                             'unknown_key_recovery_not_demonstrated': True})
    return controls


def main():
    digest = hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    assert digest == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    refs = references(); results = []
    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]
        base = 0
        for sec, text in enumerate(SOURCE.split('$')):
            p = vals(text); offsets = [w['rune_offset']-base for w in singles if w['section'] == sec]
            base += len(p)
            if len(p) <= 100: continue
            cut = 2*len(p)//3
            recs = {'full': matrices(p, offsets, 1, len(p)),
                    'train': matrices(p, offsets, 1, cut),
                    'evaluation': matrices(p, offsets, cut, len(p))}
            assert np.array_equal(recs['full']['h'], recs['train']['h']+recs['evaluation']['h'])
            assert np.array_equal(recs['full']['s'], recs['train']['s']+recs['evaluation']['s'])
            for name, ref in refs.items():
                by_k = [evaluate(recs, k, name, ref) for k in range(1, Q+1)]
                full_first = next((x['tables'] for x in by_k if sum(x['passing_by_part_and_R_category']['full'].values())), None)
                train_first = next((x['tables'] for x in by_k if 'best_training_selector_frozen' in x), None)
                frozen = None if train_first is None else by_k[train_first-1]['best_training_selector_frozen']
                if frozen is not None:
                    selected_g = ((frozen['a']*np.arange(Q)+frozen['b']) % Q) % train_first
                    witnesses = []
                    if frozen['R_category'] == 'outside':
                        for word in singles:
                            if word['section'] != sec: continue
                            offset = word['rune_offset']-(base-len(p))
                            if offset < cut: continue
                            if selected_g[p[offset-1]] == selected_g[p[offset]]:
                                assert int(p[offset]) == word['values'][0]
                                witnesses.append({'LP2_page': word['page'],
                                                  'page_rune_offset': word['page_rune_offset'],
                                                  'section_rune_offset': offset,
                                                  'previous_cipher_label': LET[int(p[offset-1])],
                                                  'singleton_cipher_label': LET[int(p[offset])],
                                                  'selected_state': int(selected_g[p[offset]])})
                    assert len(witnesses) == frozen['hard_evaluation_singleton_contradictions']
                    frozen['hard_singleton_witnesses'] = witnesses
                results.append({'singleton_sensitivity': sensitivity, 'section': sec,
                                'allowed_letters': name, 'runes': len(p), 'cut_rune_offset': cut,
                                'full_singletons': len(offsets), 'first_full_passing_k': full_first,
                                'first_training_passing_k': train_first,
                                'first_training_selector_frozen': frozen, 'by_k': by_k})
                print(sensitivity, sec, name, 'first full', full_first, 'first train', train_first,
                      'frozen full', None if frozen is None else frozen['full_passes_necessary_screens'],
                      'hard test', None if frozen is None else frozen['hard_evaluation_singleton_contradictions'], flush=True)
    out = {'master_sha256': digest,
           'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
           'protocol': {'selector': 'g(C) = ((a*C+b) mod 29) mod k; a=1..28,b=0..28,k=1..29; state g(previous C).',
                        'decoder': 'Each section has its own fixed many-to-one D and R. D[j,c]=R exactly when g(c)=j. Encoder may avoid exact repeats when choosing an R homophone.',
                        'prediction': 'R occurs exactly at same-group ciphertext transitions. Raw equal neighbors are a subset; R need not be rare.',
                        'scope': 'Eight substantial LP2 sections, omitting their first rune. The nine-rune number-label section is excluded. Pages and words do not reset. Existing master grouping is a convention, not established key boundaries.',
                        'screens': 'R share at or below largest matched solved-group single-letter frequency within its inside/outside singleton category. Allowed-set mass at or below largest matched solved-group combined share. Outside R forbids same-group singleton transitions.',
                        'transfer': 'First floor(2N/3) runes used for fitting. Select smallest k with any training screen survivor; then minimum allowed mass, minimum R mass, outside before inside, ascending a,b. Freeze selector and R category. Carry actual predecessor across evaluation split.',
                        'limits': 'Exact constraints under a specified channel rule; empirical reference screens are not universal English limits. No D table fit, unknown-key recovery, blind holdout or proof against arbitrary homophonic encoders.'},
           'references': refs, 'results': results,
           'small_exhaustive_control': small_exhaustive_control(),
           'checksum_collision_controls': checksum_collision_controls(),
           'recovery_controls': recovery_controls(refs)}
    (ROOT/'section_repeat_channel_results.json').write_text(json.dumps(out, indent=2)+'\n')
    print('Recovery controls', len(out['recovery_controls']), 'small assignments', out['small_exhaustive_control']['decoder_assignments_enumerated'])


if __name__ == '__main__': main()