"""Session 66: permutation-independent bounds for two recency-list ciphers.

No language fitting or plaintext claim. Starting and final alphabets may be
arbitrary permutations, rather than Session 5's natural/reversed rotations.
"""
from pathlib import Path
from itertools import permutations, product
import hashlib
import json
import numpy as np
from scipy.optimize import linear_sum_assignment
from lp_core import R, vals
from word_core import GROUPS
from lexical_words import tokenize_words
from reset_state_audit import SOURCE, ANCHORS

ROOT = Path(__file__).resolve().parent
POLICIES = {
    'global': '',
    'section': '$',
    'page_and_section': '%$',
    'paragraph_and_section': '&$§',
    'lexical_word': None,
}


def flags_for(source, policy):
    flags = np.zeros(len(vals(source)), dtype=bool)
    if not len(flags):
        return flags
    flags[0] = True
    if policy == 'lexical_word':
        for word in tokenize_words(source):
            flags[word['rune_offset']] = True
    else:
        i, pending = 0, False
        for char in source:
            if char in R:
                flags[i] |= pending
                i += 1
                pending = False
            elif char in POLICIES[policy]:
                pending = True
    return flags


def decode(sequence, flags, mode, update, initial, output_map):
    state = list(initial)
    output = []
    for char, reset in zip(sequence, flags):
        if reset:
            state = list(initial)
        if mode == 'symbol_to_rank':
            index = state.index(int(char))
            raw = index
        else:
            index = int(char)
            raw = state[index]
        selected = state.pop(index)
        state.insert(0 if update == 'front' else len(state), selected)
        output.append(int(output_map[raw]))
    return np.array(output, dtype=int)


def encode(plaintext, flags, mode, update, initial, output_map):
    # Invert the final substitution, then the list lookup at each step.
    inverse_map = np.argsort(output_map)
    inverse_mode = 'rank_to_symbol' if mode == 'symbol_to_rank' else 'symbol_to_rank'
    return decode(inverse_map[plaintext], flags, inverse_mode, update,
                  initial, np.arange(len(initial)))


def recency_blocks(sequence, flags, size):
    """Repeated-symbol front ranks are independent of the initial alphabet.

    On the d-th distinct symbol's first occurrence (d zero-based), its rank
    can be any value from d to size-1. At rank r, at most min(K, r+1) of a
    block's K first occurrences can contribute. All those contributions can
    occur together: put the first min(K,r+1) symbols in reverse discovery
    order ending at initial rank r. Repeat ranks do not depend on this choice.
    """
    records, ranks = [], []
    recent, histogram = [], np.zeros(size, dtype=int)
    start = 0
    for i, (char, reset) in enumerate(zip(sequence, flags)):
        if reset:
            if i:
                records.append({'start': start, 'end': i,
                                'histogram': histogram.copy(), 'distinct': len(recent)})
            start, recent, histogram = i, [], np.zeros(size, dtype=int)
        char = int(char)
        if char in recent:
            rank = recent.index(char)
            histogram[rank] += 1
            recent.pop(rank)
            ranks.append(rank)
        else:
            ranks.append(-1)
        recent.insert(0, char)
    if len(sequence):
        records.append({'start': start, 'end': len(sequence),
                        'histogram': histogram.copy(), 'distinct': len(recent)})
    return records, np.array(ranks, dtype=int)


def symbol_rank_bounds(sequence, flags, size=29):
    blocks, ranks = recency_blocks(sequence, flags, size)
    cap = np.array([b['histogram'] + np.minimum(b['distinct'], np.arange(1, size+1))
                    for b in blocks])
    shared = cap.sum(axis=0)
    n = len(sequence)
    return {
        'runes': n, 'reset_blocks': len(blocks),
        'first_occurrences': sum(b['distinct'] for b in blocks),
        'known_repeat_rank_histogram': sum((b['histogram'] for b in blocks),
                                         np.zeros(size, dtype=int)).tolist(),
        'per_rank_count_upper_bounds': shared.tolist(),
        'shared_output_map_peak_count_upper_bound': int(shared.max()),
        'shared_output_map_peak_frequency_upper_bound': float(shared.max()/n),
        'independent_block_output_maps_peak_count_upper_bound': int(cap.max(axis=1).sum()),
        'independent_block_output_maps_peak_frequency_upper_bound': float(cap.max(axis=1).sum()/n),
        'block_records': [{'start': b['start'], 'end': b['end'], 'distinct': b['distinct'],
                           'known_repeat_rank_histogram': b['histogram'].tolist(),
                           'per_rank_count_upper_bounds': row.tolist()}
                          for b, row in zip(blocks, cap)],
    }, ranks


def rank_symbol_bounds(sequence, flags, update, size=29):
    # With ciphertext used as indices, state movements permute placeholders.
    # Any initial alphabet and final output substitution only relabel these
    # placeholders. They cannot change the histogram within a reset block.
    placeholders = decode(sequence, flags, 'rank_to_symbol', update,
                          np.arange(size), np.arange(size))
    starts = np.flatnonzero(flags)
    ends = np.r_[starts[1:], len(sequence)]
    histograms = np.array([np.bincount(placeholders[a:b], minlength=size)
                           for a, b in zip(starts, ends)])
    histogram = histograms.sum(axis=0)
    return {
        'runes': len(sequence), 'reset_blocks': len(starts),
        'shared_initial_and_output_map_histogram': histogram.tolist(),
        'shared_initial_and_output_map_peak_count': int(histogram.max()),
        'shared_initial_and_output_map_peak_frequency': float(histogram.max()/len(sequence)),
        'independent_block_initial_and_output_maps_peak_count': int(histograms.max(axis=1).sum()),
        'independent_block_initial_and_output_maps_peak_frequency': float(histograms.max(axis=1).sum()/len(sequence)),
    }, placeholders


def initial_order_assignment_bounds(sequence, flags, size=29):
    """Relax shared initial-order constraints to a position assignment.

    If a symbol initially occupies j, and d distinct symbols have already
    been seen in this block, its first-occurrence front rank is j plus the
    number of those d symbols initially AFTER it. Necessarily:
        max(j, d) <= rank <= min(size-1, j+d).
    Require distinct j for distinct symbols, then maximize eligible first
    occurrences by a maximum-weight assignment. This drops which particular
    symbols were seen, so the result is an upper bound, not an attainable key.
    """
    first = np.zeros((size, size), dtype=int)
    known = np.zeros(size, dtype=int)
    recent = []
    for char, reset in zip(sequence, flags):
        if reset:
            recent = []
        char = int(char)
        if char in recent:
            rank = recent.index(char)
            known[rank] += 1
            recent.pop(rank)
        else:
            first[char, len(recent)] += 1
        recent.insert(0, char)
    prefix = np.c_[np.zeros(size, dtype=int), np.cumsum(first, axis=1)]
    cap = []
    for rank in range(size):
        weights = np.zeros((size, size), dtype=int)
        for initial_position in range(rank+1):
            weights[:, initial_position] = (prefix[:, rank+1]
                                            - prefix[:, max(0, rank-initial_position)])
        rows, columns = linear_sum_assignment(weights, maximize=True)
        cap.append(int(known[rank] + weights[rows, columns].sum()))
    return {'per_rank_count_upper_bounds': cap,
            'peak_count_upper_bound': max(cap),
            'peak_frequency_upper_bound': max(cap)/len(sequence),
            'first_occurrence_counts_by_cipher_symbol_and_discovery_ordinal': first.tolist(),
            'relaxation': 'One distinct initial position per symbol; identities of previously seen symbols are not enforced'}


def equality_pattern(sequence):
    labels = {}
    return [labels.setdefault(int(x), len(labels)) for x in sequence]


def phrase_comparison(sequence, uncertain=False):
    first, second = [[int(x) for x in sequence[a:a+6]] for a in ANCHORS]
    both_known = [i for i in range(6) if first[i] >= 0 and second[i] >= 0]
    contradictions = []
    for i in range(6):
        for j in range(i+1, 6):
            if min(first[i], first[j], second[i], second[j]) >= 0:
                if (first[i] == first[j]) != (second[i] == second[j]):
                    contradictions.append([i, j])
    return {
        'raw_values_or_initial_placeholders': [first, second],
        'unknown_first_occurrences_are_minus_one': uncertain,
        'aligned_known_positions': both_known,
        'same_values_at_aligned_known_positions': [i for i in both_known if first[i] == second[i]],
        'within_phrase_equality_contradictions_even_with_separate_output_maps': contradictions,
        'complete_equality_patterns': None if uncertain else [equality_pattern(first), equality_pattern(second)],
    }


def small_controls():
    # Exhaust every three-symbol initial permutation on all length-five inputs.
    # Check exact per-rank first-occurrence maxima, not just the loose peak.
    checked_sequences, checked_decodes, shared_order_checks = 0, 0, 0
    orders = list(permutations(range(3)))
    for seq in product(range(3), repeat=5):
        flags = np.array([True, False, False, False, False])
        bounds, known = symbol_rank_bounds(seq, flags, 3)
        for update in ['front', 'back']:
            histograms = []
            canonical = decode(seq, flags, 'rank_to_symbol', update, range(3), range(3))
            for initial in orders:
                decoded = decode(seq, flags, 'symbol_to_rank', update, initial, range(3))
                if update == 'front':
                    assert np.array_equal(decoded[known >= 0], known[known >= 0])
                else:
                    assert np.array_equal(decoded[known >= 0], 2-known[known >= 0])
                histograms.append(np.bincount(decoded, minlength=3))
                actual = decode(seq, flags, 'rank_to_symbol', update, initial, range(3))
                assert np.array_equal(actual, np.array(initial)[canonical])
                checked_decodes += 2
            expected = np.array(bounds['per_rank_count_upper_bounds'])
            if update == 'back':
                expected = expected[::-1]
            assert np.array_equal(np.array(histograms).max(axis=0), expected)
        checked_sequences += 1
        # With a reset in the middle, one initial permutation must serve both
        # blocks. The position-assignment relaxation must contain every key.
        split_flags = np.array([True, False, True, False, False])
        shared_cap = np.array(initial_order_assignment_bounds(seq, split_flags, 3)
                              ['per_rank_count_upper_bounds'])
        for initial in orders:
            decoded = decode(seq, split_flags, 'symbol_to_rank', 'front', initial, range(3))
            assert np.all(np.bincount(decoded, minlength=3) <= shared_cap)
            shared_order_checks += 1
    return {'alphabet_size': 3, 'input_length': 5, 'sequences_exhausted': checked_sequences,
            'initial_permutations_per_sequence': len(orders),
            'decodes_checked': checked_decodes,
            'exact_per_rank_maxima_and_placeholder_relabeling_agree': True,
            'shared_initial_order_two_block_checks': shared_order_checks,
            'shared_order_assignment_bound_contains_every_tested_key': True}


def recovery_controls():
    rng = np.random.default_rng(66001)
    controls = []
    for group in ['koan', 'divinity']:
        source, plaintext = GROUPS[group], vals(GROUPS[group])
        true_peak = int(np.bincount(plaintext, minlength=29).max())
        for policy in ['global', 'paragraph_and_section']:
            flags = flags_for(source, policy)
            for mode in ['symbol_to_rank', 'rank_to_symbol']:
                for update in ['front', 'back']:
                    initial, output_map = rng.permutation(29), rng.permutation(29)
                    ciphertext = encode(plaintext, flags, mode, update, initial, output_map)
                    recovered = decode(ciphertext, flags, mode, update, initial, output_map)
                    assert np.array_equal(recovered, plaintext)
                    if mode == 'symbol_to_rank':
                        b, _ = symbol_rank_bounds(ciphertext, flags)
                        bound_count = b['shared_output_map_peak_count_upper_bound']
                        assignment_count = initial_order_assignment_bounds(ciphertext, flags)['peak_count_upper_bound']
                        assert assignment_count >= true_peak
                    else:
                        b, _ = rank_symbol_bounds(ciphertext, flags, update)
                        bound_count = b['shared_initial_and_output_map_peak_count']
                        assignment_count = None
                    assert bound_count >= true_peak
                    if mode == 'rank_to_symbol':
                        assert bound_count == true_peak
                    controls.append({'source_group': group, 'reset_policy': policy,
                                     'mode': mode, 'update': update, 'runes': len(plaintext),
                                     'arbitrary_initial_alphabet': initial.tolist(),
                                     'arbitrary_output_map': output_map.tolist(),
                                     'exact_recovery': True, 'true_peak_count': true_peak,
                                     'predicted_peak_count_or_upper_bound': bound_count,
                                     'shared_order_assignment_peak_count_upper_bound': assignment_count,
                                     'bound_preserves_constructed_plaintext': True})
    return controls


def main():
    source_hash = hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    assert source_hash == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    ciphertext = vals(SOURCE)
    assert len(ciphertext) == 12956
    assert all(ciphertext[a:a+6].tolist() == [23, 11, 1, 17, 18, 10] for a in ANCHORS)
    section_ids = np.concatenate([np.full(len(vals(s)), i, dtype=int)
                                  for i, s in enumerate(SOURCE.split('$'))])
    models = []
    for policy in POLICIES:
        flags = flags_for(SOURCE, policy)
        sr, known = symbol_rank_bounds(ciphertext, flags)
        sr['shared_initial_alphabet_assignment_bound'] = initial_order_assignment_bounds(ciphertext, flags)
        if policy != 'global':
            regional = []
            for section in sorted(set(section_ids)):
                keep = section_ids == section
                section_flags = flags[keep].copy()
                section_flags[0] = True
                bound = initial_order_assignment_bounds(ciphertext[keep], section_flags)
                regional.append({'section': int(section), **bound})
            sr['one_initial_and_output_map_per_section_assignment_peak_count_upper_bound'] = sum(b['peak_count_upper_bound'] for b in regional)
            sr['one_initial_and_output_map_per_section_assignment_peak_frequency_upper_bound'] = sum(b['peak_count_upper_bound'] for b in regional)/len(ciphertext)
            sr['section_assignment_bounds'] = regional
        sr['phrase_front'] = phrase_comparison(known, True)
        sr['phrase_back'] = phrase_comparison(np.where(known < 0, -1, 28-known), True)
        rs = {}
        for update in ['front', 'back']:
            bound, placeholders = rank_symbol_bounds(ciphertext, flags, update)
            bound['phrase'] = phrase_comparison(placeholders)
            rs[update] = bound
        models.append({'reset_policy': policy, 'symbol_to_rank_both_updates': sr,
                       'rank_to_symbol': rs})
    benchmarks = {}
    for group, plaintext in [*GROUPS.items(), ('pooled', '%'.join(GROUPS.values()))]:
        h = np.bincount(vals(plaintext), minlength=29)
        benchmarks[group] = {'runes': int(h.sum()), 'peak_count': int(h.max()),
                             'peak_frequency': float(h.max()/h.sum())}
    out = {
        'source_sha256': source_hash,
        'scope': {'LP2_pages': [0, 55], 'runes': len(ciphertext),
                  'sections': 'All nine nonempty sections, including the nine-rune section',
                  'cipher_indices_for_rank_to_symbol': 'Fixed zero-based Gematria Primus order from lp_core.R',
                  'reset_policies': POLICIES,
                  'key_scope': 'Shared key results plus more permissive independent keys for each reset block',
                  'symbol_to_rank_bound': 'First-occurrence ranks relaxed independently between reset blocks; shared-key bounds can overestimate what one initial alphabet attains',
                  'limits': 'No extra state, input-index relabeling for rank_to_symbol, phase-dependent output maps, skipped updates, nulls or length changes'},
        'anchor_offsets': ANCHORS, 'models': models, 'solved_LP1_benchmarks': benchmarks,
        'small_exhaustive_control': small_controls(), 'recovery_controls': recovery_controls(),
        'software': {'NumPy': np.__version__, 'SciPy': __import__('scipy').__version__},
        'interpretation': 'Frequency is an empirical screen against saved solved LP1 text, not a theorem that every possible English passage has the same peak. Repeated ciphertext need not represent repeated plaintext.'
    }
    (ROOT/'recency_key_results.json').write_text(json.dumps(out, indent=2)+'\n')
    for model in models:
        b = model['symbol_to_rank_both_updates']
        print(model['reset_policy'], 'symbol_to_rank',
              b['shared_output_map_peak_frequency_upper_bound'],
              b['independent_block_output_maps_peak_frequency_upper_bound'])
        print(' shared initial order assignment', b['shared_initial_alphabet_assignment_bound']['peak_frequency_upper_bound'],
              'separate section orders', b.get('one_initial_and_output_map_per_section_assignment_peak_frequency_upper_bound'))
        for update, b in model['rank_to_symbol'].items():
            print(model['reset_policy'], 'rank_to_symbol', update,
                  b['shared_initial_and_output_map_peak_frequency'],
                  b['independent_block_initial_and_output_maps_peak_frequency'])
    print('controls', len(out['recovery_controls']), out['small_exhaustive_control'])


if __name__ == '__main__':
    main()