"""Session 68: fixed, possibly many-to-one rune decoding tables.

Conditional on one cipher rune per plaintext rune and English singleton
words belonging to A/I or A/I/O. No substitution key or plaintext is fitted.
"""
from pathlib import Path
from collections import Counter
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
ALLOWED = {'AI': [24, 10], 'AIO': [24, 10, 3]}
RED_LOCATIONS = {(6, 75), (54, 0)}


def key_group(section, scope):
    if scope == 'shared':
        return 'all'
    if scope == 'anchor_shared' and section in (6, 8):
        return '6+8'
    return str(section)


def units_for(source, singletons, scope):
    groups = {}
    for section, text in enumerate(source.split('$')):
        if not len(vals(text)):
            continue
        key = key_group(section, scope)
        group = groups.setdefault(key, {'sections': [], 'histogram': Counter(), 'singletons': Counter()})
        group['sections'].append(section)
        group['histogram'].update(map(int, vals(text)))
    for word in singletons:
        groups[key_group(word['section'], scope)]['singletons'][int(word['values'][0])] += 1
    units = []
    for key, group in groups.items():
        for cipher, count in sorted(group['singletons'].items()):
            units.append({'group': key, 'cipher': cipher, 'label': LET[cipher],
                          'exception_token_cost': count, 'freed_corpus_positions': group['histogram'][cipher]})
    return groups, units


def release_profile(units):
    # Freeing a type requires ignoring every singleton of that type in its
    # decoding-table group. Keep the exact largest release at each token cost.
    total = sum(u['exception_token_cost'] for u in units)
    gain = [-1]*(total+1)
    selected = [None]*(total+1)
    gain[0], selected[0] = 0, ()
    for index, unit in enumerate(units):
        cost, amount = unit['exception_token_cost'], unit['freed_corpus_positions']
        for budget in range(total, cost-1, -1):
            if gain[budget-cost] >= 0 and gain[budget-cost]+amount > gain[budget]:
                gain[budget] = gain[budget-cost]+amount
                selected[budget] = selected[budget-cost]+(index,)
    return gain, selected


def exception_requirement(units, n, mandatory, ratio, gain, selected):
    maximum_allowed_positions = n*ratio.numerator//ratio.denominator
    required_release = max(0, mandatory-maximum_allowed_positions)
    cost = next(c for c, g in enumerate(gain) if g >= required_release)
    chosen = [units[i] for i in selected[cost]]
    assert sum(u['exception_token_cost'] for u in chosen) == cost
    actual_release = sum(u['freed_corpus_positions'] for u in chosen)
    assert actual_release == gain[cost]
    return {'reference_ratio': [ratio.numerator, ratio.denominator],
            'maximum_allowed_positions': maximum_allowed_positions,
            'required_position_release': required_release,
            'minimum_exception_tokens': cost, 'freed_constraint_types': len(chosen),
            'witness_units': chosen,
            'remaining_forced_positions': mandatory-actual_release,
            'remaining_forced_share': (mandatory-actual_release)/n,
            'necessary_mass_condition_only_not_a_full_plaintext_fit': True}


def bound_record(source, singletons, allowed_count, scope):
    n = len(vals(source))
    groups, units = units_for(source, singletons, scope)
    mandatory = sum(u['freed_corpus_positions'] for u in units)
    by_group = []
    for key, group in groups.items():
        constrained = set(group['singletons'])
        active = set(group['histogram'])
        k = min(allowed_count, len(constrained))
        alphabet_maximum = len(active-constrained)+k
        by_group.append({'group': key, 'sections': group['sections'],
                         'runes': sum(group['histogram'].values()),
                         'singleton_tokens': sum(group['singletons'].values()),
                         'singleton_cipher_types': len(constrained),
                         'active_cipher_types': len(active),
                         'forced_positions': sum(group['histogram'][c] for c in constrained),
                         'maximum_distinct_output_runes': alphabet_maximum,
                         'cipher_histogram': [group['histogram'][i] for i in range(29)],
                         'singleton_counts_by_cipher_type': [group['singletons'][i] for i in range(29)]})
    gain, selected = release_profile(units)
    return {'runes': n, 'key_scope': scope, 'allowed_output_types': allowed_count,
            'singleton_tokens': len(singletons), 'forced_positions': mandatory,
            'minimum_allowed_letter_share': mandatory/n,
            'groups': by_group, 'constraint_units': units,
            'maximum_position_release_by_exact_exception_cost': gain}, units, gain, selected


def small_exhaustive_controls():
    cases = []
    for histogram, singleton_costs, allowed in [([4,5,6,7,8], [1,3,2,1,0], {0,1}),
                                              ([9,7,5,11], [2,1,2,0], {0})]:
        size = len(histogram)
        units = [{'exception_token_cost': cost, 'freed_corpus_positions': histogram[c]}
                 for c, cost in enumerate(singleton_costs) if cost]
        gain, _ = release_profile(units)
        mandatory = sum(u['freed_corpus_positions'] for u in units)
        total = sum(singleton_costs)
        minima = [10**9]*(total+1)
        maxima = [-1]*(total+1)
        tables = 0
        for table in product(range(size), repeat=size):
            cost = sum(singleton_costs[c] for c in range(size) if table[c] not in allowed)
            mass = sum(histogram[c] for c in range(size) if table[c] in allowed)
            alphabet = len(set(table))
            for budget in range(cost, total+1):
                minima[budget] = min(minima[budget], mass)
                maxima[budget] = max(maxima[budget], alphabet)
            tables += 1
        constrained = sum(c > 0 for c in singleton_costs)
        unconstrained = size-constrained
        cheap_costs = sorted(c for c in singleton_costs if c)
        for budget in range(total+1):
            release = max(gain[:budget+1])
            assert minima[budget] == mandatory-release
            freed = max(e for e in range(constrained+1) if sum(cheap_costs[:e]) <= budget)
            alphabet_bound = unconstrained+freed+min(len(allowed), constrained-freed)
            assert maxima[budget] == alphabet_bound
        cases.append({'cipher_and_output_alphabet_size': size, 'decoding_tables_exhausted': tables,
                      'exception_budgets_checked': total+1,
                      'exact_mass_and_alphabet_bounds_agree': True})
    return cases


def homophonic_controls():
    rng = np.random.default_rng(68001)
    controls = []
    for name in ['warning', 'koan', 'divinity', 'instruction']:
        source = GROUPS[name]
        plain = vals(source)
        letters = sorted(set(map(int, plain)))
        double_counts = Counter(int(a) for a, b in zip(plain, plain[1:]) if a == b)
        extras = 29-len(letters)
        extra_letters = [letter for letter, _ in sorted(double_counts.items(), key=lambda x:(-x[1],x[0]))][:extras]
        if len(extra_letters) < extras:
            extra_letters += [c for c in letters if c not in extra_letters][:extras-len(extra_letters)]
        code_ids = list(map(int, rng.permutation(29)))
        classes = {letter: [code_ids.pop()] for letter in letters}
        for letter in extra_letters:
            classes[letter].append(code_ids.pop())
        assert not code_ids
        inverse = {cipher: letter for letter, codes in classes.items() for cipher in codes}
        cipher, previous, unavoidable = [], None, 0
        for letter in plain:
            options = [c for c in classes[int(letter)] if c != previous]
            if not options:
                options = classes[int(letter)]
                unavoidable += 1
            c = int(rng.choice(options))
            cipher.append(c)
            previous = c
        recovered = np.array([inverse[c] for c in cipher])
        assert np.array_equal(recovered, plain)
        text = replace_values(source, cipher)
        words = tokenize_words(text)
        single = [w for w in words if len(w['values']) == 1]
        assert all(inverse[w['values'][0]] in ALLOWED['AI'] for w in single)
        b, _, _, _ = bound_record(text, single, 2, 'shared')
        ai_count = sum(int(p) in ALLOWED['AI'] for p in plain)
        assert b['forced_positions'] <= ai_count
        assert b['groups'][0]['maximum_distinct_output_runes'] >= len(letters)
        encoded_doubles = sum(a == b for a, b in zip(cipher, cipher[1:]))
        assert encoded_doubles == unavoidable
        controls.append({'source_group': name, 'runes': len(plain),
                         'plaintext_types': len(letters), 'homophone_classes': {LET[p]: codes for p,codes in classes.items()},
                         'stateful_homophone_choice_but_fixed_decoding_table': True,
                         'exact_inverse_recovery': True,
                         'plaintext_neighbor_repeats': sum(int(a == b) for a,b in zip(plain,plain[1:])),
                         'ciphertext_neighbor_repeats': encoded_doubles,
                         'singleton_tokens': len(single), 'singleton_cipher_types': len({w['values'][0] for w in single}),
                         'forced_AI_positions': b['forced_positions'], 'true_AI_positions': ai_count,
                         'alphabet_bound_preserves_true_alphabet': True})
    return controls


def state_dependent_counterexample():
    # A control outside the fixed-decoding-table premise. Copy the known koan
    # three times to provide at least 28 singleton slots; choose a different
    # substitution per word and deliberately cycle those slots through types.
    source = '%'.join([GROUPS['koan']]*3)
    words = tokenize_words(source)
    rng = np.random.default_rng(68002)
    cipher, recovered, singleton_counter = [], [], 0
    for word in words:
        permutation = rng.permutation(29)
        p = np.array(word['values'])
        if len(p) == 1:
            desired = singleton_counter % 28
            location = int(np.flatnonzero(permutation == desired)[0])
            permutation[location], permutation[int(p[0])] = permutation[int(p[0])], permutation[location]
            singleton_counter += 1
        c = permutation[p]
        cipher.extend(c.tolist())
        recovered.extend(np.argsort(permutation)[c].tolist())
    plain = vals(source)
    assert np.array_equal(recovered, plain)
    text = replace_values(source, cipher)
    single = [w for w in tokenize_words(text) if len(w['values']) == 1]
    b, _, _, _ = bound_record(text, single, 2, 'shared')
    assert len({w['values'][0] for w in single}) == 28
    assert b['groups'][0]['maximum_distinct_output_runes'] == 3
    assert len(set(map(int,plain))) > 3
    return {'known_koan_copies': 3, 'runes': len(plain), 'singletons': len(single),
            'singleton_cipher_types': 28, 'per_word_substitution_inverse_exact': True,
            'true_plaintext_types': len(set(map(int,plain))),
            'incorrect_fixed_table_assumption_would_cap_output_types': 3,
            'fixed_table_bound_not_applicable': True,
            'a_scope_counterexample_not_a_proposed_Cicada_key': True}


def main():
    master = ROOT/'liber-primus__transcription--master.txt'
    digest = hashlib.sha256(master.read_bytes()).hexdigest()
    assert digest == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    assert len(vals(SOURCE)) == 12956 and len(SINGLE) == 83
    benchmarks = {}
    for name, allowed in ALLOWED.items():
        groups = []
        for group, text in GROUPS.items():
            plain = vals(text)
            count = sum(int(p) in allowed for p in plain)
            groups.append({'group': group, 'runes': len(plain), 'allowed_positions': count,
                           'allowed_share': count/len(plain), 'plaintext_types': len(set(map(int,plain)))})
        pooled = Fraction(sum(g['allowed_positions'] for g in groups), sum(g['runes'] for g in groups))
        generous = max(Fraction(g['allowed_positions'],g['runes']) for g in groups)
        benchmarks[name] = {'allowed_values': allowed, 'groups': groups,
                            'pooled_ratio': [pooled.numerator, pooled.denominator],
                            'pooled_share': float(pooled),
                            'largest_individual_ratio': [generous.numerator,generous.denominator],
                            'largest_individual_share': float(generous)}
    results = []
    for sensitivity in ['all_audited', 'exclude_red_initials']:
        singletons = SINGLE if sensitivity == 'all_audited' else [w for w in SINGLE if (w['page'],w['page_rune_offset']) not in RED_LOCATIONS]
        for scope in ['shared','per_section','anchor_shared']:
            for name, allowed in ALLOWED.items():
                b, units, gain, chosen = bound_record(SOURCE, singletons, len(allowed), scope)
                b['singleton_sensitivity'], b['allowed_letters'] = sensitivity, name
                b['requirements'] = {}
                for reference, ratio_key in [('pooled', 'pooled_ratio'), ('largest_individual', 'largest_individual_ratio')]:
                    ratio = Fraction(*benchmarks[name][ratio_key])
                    req = exception_requirement(units, b['runes'], b['forced_positions'], ratio, gain, chosen)
                    # Construct the two-output relaxed witness. Retained types
                    # decode inside S; freed/unconstrained types decode to E.
                    freed = {(u['group'],u['cipher']) for u in req['witness_units']}
                    groups, _ = units_for(SOURCE,singletons,scope)
                    observed_cost = observed_mass = 0
                    for key, group in groups.items():
                        for cipher in group['singletons']:
                            if (key,cipher) in freed:
                                observed_cost += group['singletons'][cipher]
                            else:
                                observed_mass += group['histogram'][cipher]
                    assert observed_cost == req['minimum_exception_tokens']
                    assert observed_mass == req['remaining_forced_positions']
                    assert observed_mass <= req['maximum_allowed_positions']
                    req['relaxed_mass_witness_verified'] = True
                    b['requirements'][reference] = req
                results.append(b)
    out = {'source_sha256': digest,
           'scope': {'LP2_pages': [0,55], 'runes': len(vals(SOURCE)),
                     'one_rune_word_tokens': len(SINGLE), 'one_rune_word_types': len({w['values'][0] for w in SINGLE}),
                     'model': 'One ciphertext rune decodes to one plaintext rune through a fixed table. Many ciphertext types may decode to the same plaintext type. Encoder may choose homophones using state, but decoding table is fixed within each stated key group.',
                     'assumptions': 'Retained separators are plaintext word boundaries; every ordinary one-rune word is in the specified A/I or A/I/O set; no nulls, variable-length decoding, context-dependent decoding or state-selected table inside a key group',
                     'limits': 'Necessary mass constraints only; exception budgets are not English decipherments; reference shares are empirical, not universal English limits'},
           'benchmarks': benchmarks, 'results': results,
           'small_exhaustive_controls': small_exhaustive_controls(),
           'homophonic_controls': homophonic_controls(),
           'state_dependent_counterexample': state_dependent_counterexample()}
    (ROOT/'homophone_capacity_results.json').write_text(json.dumps(out,indent=2)+'\n')
    for b in results:
        print(b['singleton_sensitivity'],b['key_scope'],b['allowed_letters'],
              'forced_share',round(b['minimum_allowed_letter_share'],6),
              'pooled exceptions',b['requirements']['pooled']['minimum_exception_tokens'],
              'generous exceptions',b['requirements']['largest_individual']['minimum_exception_tokens'])
    print('Controls',len(out['homophonic_controls']),out['small_exhaustive_controls'])


if __name__ == '__main__':
    main()
