"""Session 67: bounded input-index relabeling, with frozen continuation checks.

Exhausts the 812 affine permutations; separately samples 4,096 arbitrary
permutations. Does not enumerate 29!, fit word meanings, or claim new plaintext.
"""
from pathlib import Path
from collections import Counter
import ctypes
import hashlib
import json
import subprocess
import tempfile
import time
import numpy as np
from lp_core import vals, render
from recency_key_bound import decode, encode, flags_for, recency_blocks
from lexical_words import tokenize_words
from reset_state_audit import SOURCE
from word_core import GROUPS

ROOT = Path(__file__).resolve().parent
POLICIES = ['section', 'paragraph_and_section', 'page_and_section', 'lexical_word']
SECTIONS = [0, 1, 2, 4, 5, 6, 7, 8]
RNG_SEED = 67001
RANDOM_KEYS = 4096
AFFINE = np.array([[(a*c+b) % 29 for c in range(29)]
                   for a in range(1, 29) for b in range(29)], dtype=np.uint8)
AFFINE_INFO = [{'a': a, 'b': b} for a in range(1, 29) for b in range(29)]
BENCHMARK_H = np.bincount(np.concatenate([vals(s) for s in GROUPS.values()]), minlength=29)
BENCHMARK_P = (BENCHMARK_H + .5) / (BENCHMARK_H.sum() + 29*.5)
PEAK_GATE = min(np.bincount(vals(s), minlength=29).max()/len(vals(s)) for s in GROUPS.values())


def word_cut(source):
    lengths = np.array([len(w['values']) for w in tokenize_words(source)])
    ends = lengths.cumsum()
    limit = min(256, int(.6*ends[-1]))
    candidates = ends[ends <= limit]
    assert len(candidates)
    return int(candidates[-1])


class Kernel:
    def __init__(self, directory):
        library = Path(directory)/'recency_input.so'
        subprocess.run(['g++', '-O3', '-std=c++17', '-shared', '-fPIC',
                        str(ROOT/'recency_input_kernel.cpp'), '-o', str(library)], check=True)
        self.library = ctypes.CDLL(str(library))
        self.function = self.library.recency_histograms
        u8 = np.ctypeslib.ndpointer(dtype=np.uint8, flags='C_CONTIGUOUS')
        i32 = np.ctypeslib.ndpointer(dtype=np.int32, flags='C_CONTIGUOUS')
        self.function.argtypes = [u8, u8, ctypes.c_int, u8, ctypes.c_int,
                                  ctypes.c_int, ctypes.c_int, i32, i32]
        self.function.restype = None

    def histograms(self, cipher, flags, keys, cutoff, update):
        cipher, flags, keys = [np.ascontiguousarray(x, dtype=np.uint8) for x in [cipher, flags, keys]]
        assert keys.shape[1] == 29 and len(flags) == len(cipher)
        prefix = np.zeros((len(keys), 29), dtype=np.int32)
        suffix = np.zeros_like(prefix)
        self.function(cipher, flags, len(cipher), keys, len(keys), cutoff,
                      int(update == 'back'), prefix, suffix)
        assert np.all(prefix.sum(axis=1) == cutoff)
        assert np.all(suffix.sum(axis=1) == len(cipher)-cutoff)
        return prefix, suffix


def fitted_unigrams(histograms):
    # Exact best fixed output permutation for this unigram-only objective.
    cipher_order = np.argsort(-histograms, axis=1, kind='stable')
    language_order = np.argsort(-BENCHMARK_P, kind='stable')
    sorted_counts = np.take_along_axis(histograms, cipher_order, axis=1)
    scores = (sorted_counts @ np.log(BENCHMARK_P[language_order]))/histograms.sum(axis=1)
    return scores, cipher_order, language_order


def output_map_for(histogram):
    order = np.argsort(-histogram, kind='stable')
    output = np.empty(29, dtype=int)
    output[order] = np.argsort(-BENCHMARK_P, kind='stable')
    return output


def key_description(index, random_keys):
    if index < len(AFFINE):
        return {'family': 'affine', 'index': int(index), **AFFINE_INFO[index]}
    return {'family': 'sampled_arbitrary', 'index': int(index-len(AFFINE)),
            'cipher_symbol_to_position': random_keys[index-len(AFFINE)].tolist()}


def kernel_controls(kernel, bank):
    rng = np.random.default_rng(67002)
    cases = []
    for policy in POLICIES:
        flags = flags_for(GROUPS['koan'], policy)
        cipher = rng.integers(0, 29, size=len(flags))
        keys = bank[[0, 27, 400, 812, 975, 4907]]
        cutoff = word_cut(GROUPS['koan'])
        for update in ['front', 'back']:
            h0, h1 = kernel.histograms(cipher, flags, keys, cutoff, update)
            for i, key in enumerate(keys):
                python = decode(key[cipher], flags, 'rank_to_symbol', update,
                                np.arange(29), np.arange(29))
                assert np.array_equal(h0[i], np.bincount(python[:cutoff], minlength=29))
                assert np.array_equal(h1[i], np.bincount(python[cutoff:], minlength=29))
            cases.append({'policy': policy, 'update': update, 'keys_compared': len(keys),
                          'prefix_and_suffix_histograms_match_python': True})
    return cases


def recovery_controls(kernel, bank):
    rng = np.random.default_rng(67003)
    results = []
    for group in ['koan', 'divinity']:
        source = GROUPS[group]
        plain = vals(source)
        cutoff = word_cut(source)
        for policy in ['section', 'paragraph_and_section']:
            flags = flags_for(source, policy)
            for update in ['front', 'back']:
                for family in ['affine', 'sampled_arbitrary']:
                    lo, hi = (0, len(AFFINE)) if family == 'affine' else (len(AFFINE), len(bank))
                    key_index = int(rng.integers(lo, hi))
                    input_key = bank[key_index]
                    initial, final_map = rng.permutation(29), rng.permutation(29)
                    indices = encode(plain, flags, 'rank_to_symbol', update, initial, final_map)
                    cipher = np.argsort(input_key)[indices]
                    recovered = decode(input_key[cipher], flags, 'rank_to_symbol', update, initial, final_map)
                    assert np.array_equal(recovered, plain)
                    h0, h1 = kernel.histograms(cipher, flags, bank[lo:hi], cutoff, update)
                    planted = key_index-lo
                    assert np.array_equal(np.sort((h0+h1)[planted]), np.sort(np.bincount(plain, minlength=29)))
                    true_prefix_peak = float(h0[planted].max()/cutoff)
                    true_suffix_peak = float(h1[planted].max()/(len(plain)-cutoff))
                    assert true_prefix_peak >= PEAK_GATE and true_suffix_peak >= PEAK_GATE
                    scores, _, _ = fitted_unigrams(h0)
                    results.append({'group': group, 'policy': policy, 'update': update,
                                    'input_key_family': family, 'input_key_index': key_index,
                                    'input_key': input_key.tolist(), 'initial': initial.tolist(),
                                    'output_map': final_map.tolist(), 'runes': len(plain), 'cutoff': cutoff,
                                    'known_key_exact_recovery': True,
                                    'planted_key_retained_by_prefix_gate': True,
                                    'planted_key_prefix_peak': true_prefix_peak,
                                    'planted_key_suffix_peak': true_suffix_peak,
                                    'planted_key_unigram_rank_in_its_predefined_bank': int(1+sum(scores > scores[planted]+1e-12)),
                                    'not_a_blind_arbitrary_key_recovery': True})
    return results


def encoder_repeat_reference():
    # Input-index renaming preserves equality of ciphertext values. A fixed
    # starting/final alphabet cannot alter recurrence ranks after first use.
    records = []
    for group, source in GROUPS.items():
        p = vals(source)
        for policy in POLICIES:
            flags = flags_for(source, policy)
            _, known = recency_blocks(p, flags, 29)
            within = ~flags[1:]
            determined = within & (known[:-1] >= 0) & (known[1:] >= 0)
            forced = determined & (known[:-1] == known[1:])
            records.append({'group': group, 'policy': policy, 'runes': len(p),
                            'within_reset_block_adjacencies': int(within.sum()),
                            'adjacencies_with_both_recurrence_ranks_known': int(determined.sum()),
                            'forced_equal_encoded_neighbors_for_every_initial_and_input_key': int(forced.sum()),
                            'minimum_encoded_repeat_rate': float(forced.sum()/within.sum())})
    return records


def reflection_controls():
    # Front/back are equivalent under arbitrary input relabeling and initial
    # order: reverse the list and complement all selected indices.
    rng = np.random.default_rng(67004)
    cipher = rng.integers(0, 29, 700)
    flags = np.zeros(700, dtype=bool)
    flags[[0, 17, 263, 529]] = True
    for _ in range(12):
        key = rng.permutation(29)
        front = decode(key[cipher], flags, 'rank_to_symbol', 'front', range(29), range(29))
        back = decode((28-key)[cipher], flags, 'rank_to_symbol', 'back', list(range(29))[::-1], range(29))
        assert np.array_equal(front, back)
    return {'random_keys_checked': 12, 'runes_per_key': 700,
            'front_back_reflection_identity_exact': True,
            'consequence': 'Affine bank is closed under reflection; the two update sweeps are not distinct arbitrary-key families'}


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


def phrase_patterns(bank):
    # Each phrase begins in some permutation of 29 distinct symbols. Renaming
    # that initial state cannot change its internal output-equality pattern.
    cipher = np.array([23, 11, 1, 17, 18, 10])
    flags = np.array([True, False, False, False, False, False])
    results = []
    for update in ['front', 'back']:
        for family, lo, hi in [('affine', 0, 812), ('sampled_arbitrary', 812, len(bank))]:
            patterns = Counter(equality_pattern(decode(key[cipher], flags, 'rank_to_symbol', update,
                                                       range(29), range(29))) for key in bank[lo:hi])
            results.append({'update': update, 'family': family,
                            'distinct_patterns': len(patterns),
                            'patterns': [{'pattern': list(p), 'key_count': count}
                                         for p, count in sorted(patterns.items())]})
    return results


def main():
    started = time.time()
    master_hash = hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    assert master_hash == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    rng = np.random.default_rng(RNG_SEED)
    random_keys = np.array([rng.permutation(29) for _ in range(RANDOM_KEYS)], dtype=np.uint8)
    assert len({tuple(x) for x in random_keys}) == RANDOM_KEYS
    bank = np.r_[AFFINE, random_keys]
    assert len({tuple(x) for x in bank}) == len(bank)
    regions = [(i, s) for i, s in enumerate(SOURCE.split('$')) if i in SECTIONS]
    n = sum(len(vals(s)) for _, s in regions)
    assert n == 12947
    output = {'source_sha256': master_hash,
              'definition': 'Before rank-to-symbol lookup, apply a fixed bijection from ciphertext rune values to list indices. After each rune move the selected symbol to front or back. An arbitrary fixed starting alphabet and final output substitution only relabel placeholder outputs for that input key.',
              'scope': {'sections': SECTIONS, 'runes': n, 'reset_policies': POLICIES,
                        'affine_keys': len(AFFINE), 'random_keys': RANDOM_KEYS,
                        'random_seed': RNG_SEED, 'random_bank_sha256': hashlib.sha256(random_keys.tobytes()).hexdigest(),
                        'input_key_family_exactness': 'All 812 affine bijections; 4096 arbitrary permutations sampled, not exhaustive',
                        'key_scope': 'One fixed input key and combined output map per section; separately pool counts for a shared key',
                        'fit': 'Complete-word prefix at most 256 runes and 60 percent of a section; exact optimum output permutation for pooled solved-LP1 unigram likelihood, then choose input key within each predefined bank',
                        'validation': 'Freeze input key and prefix-fitted output map; continue state without an added reset at the split',
                        'limits': 'No word-meaning optimization, phase-dependent output substitutions, skipped updates, nulls, changes in length, other update rules or independent key per word'},
              'screen': {'minimum_reference_peak': PEAK_GATE,
                         'pooled_reference_peak': float(BENCHMARK_H.max()/BENCHMARK_H.sum()),
                         'pooled_reference_runes': int(BENCHMARK_H.sum()),
                         'empirical_language_screen_not_universal_English_bound': True},
              'random_input_keys': random_keys.tolist(),
              'encoder_repeat_reference': encoder_repeat_reference(),
              'reflection_control': reflection_controls(),
              'phrase_pattern_results': phrase_patterns(bank), 'section_fits': [], 'pooled_results': []}
    output['real_neighbor_repeats'] = []
    for policy in POLICIES:
        eligible = repeats = 0
        for _, source in regions:
            cipher = vals(source)
            within = ~flags_for(source, policy)[1:]
            eligible += int(within.sum())
            repeats += int(sum(within & (cipher[1:] == cipher[:-1])))
        output['real_neighbor_repeats'].append({'policy': policy, 'eligible_pairs': eligible,
                                               'identical_pairs': repeats, 'rate': repeats/eligible})
    with tempfile.TemporaryDirectory(prefix='cicada-recency-input-') as directory:
        kernel = Kernel(directory)
        output['native_python_checks'] = kernel_controls(kernel, bank)
        output['constructed_controls'] = recovery_controls(kernel, bank)
        print('Controls pass:', len(output['constructed_controls']), flush=True)
        for policy in POLICIES:
            for update in ['front', 'back']:
                aggregate = np.zeros((len(bank), 29), dtype=np.int64)
                independent = {'affine': 0, 'sampled_arbitrary': 0}
                for section, source in regions:
                    cipher = vals(source)
                    cutoff = word_cut(source)
                    flags = flags_for(source, policy)
                    h0, h1 = kernel.histograms(cipher, flags, bank, cutoff, update)
                    full = h0+h1
                    aggregate += full
                    scores, _, _ = fitted_unigrams(h0)
                    for family, lo, hi in [('affine', 0, 812), ('sampled_arbitrary', 812, len(bank))]:
                        selected = lo+int(scores[lo:hi].argmax())
                        final_map = output_map_for(h0[selected])
                        decoded = decode(bank[selected][cipher], flags, 'rank_to_symbol', update,
                                         range(29), final_map)
                        assert np.array_equal(np.bincount(decoded[:cutoff], minlength=29), h0[selected][np.argsort(final_map)])
                        prefix_peak = float(h0[selected].max()/cutoff)
                        suffix_peak = float(h1[selected].max()/(len(cipher)-cutoff))
                        full_peaks = full[lo:hi].max(axis=1)
                        best_full = lo+int(full_peaks.argmax())
                        independent[family] += int(full_peaks.max())
                        output['section_fits'].append({
                            'section': section, 'policy': policy, 'update': update, 'family': family,
                            'runes': len(cipher), 'cutoff': cutoff, 'selected_input_key': key_description(selected, random_keys),
                            'frozen_output_map': final_map.tolist(),
                            'prefix_score': float(scores[selected]),
                            'suffix_score_frozen_map': float(h1[selected] @ np.log(BENCHMARK_P[final_map])/(len(cipher)-cutoff)),
                            'selected_prefix_peak': prefix_peak, 'selected_suffix_peak': suffix_peak,
                            'selected_full_peak': float(full[selected].max()/len(cipher)),
                            'selected_prefix_peak_label_suffix_frequency': float(h1[selected][h0[selected].argmax()]/(len(cipher)-cutoff)),
                            'prefix_gate_survivors': int(sum(h0[lo:hi].max(axis=1)/cutoff >= PEAK_GATE)),
                            'full_section_gate_survivors': int(sum(full_peaks/len(cipher) >= PEAK_GATE)),
                            'maximum_full_section_peak': float(full_peaks.max()/len(cipher)),
                            'best_full_peak_input_key_not_a_training_selection': key_description(best_full, random_keys),
                            'full_section_peak_counts_for_every_input_key': full_peaks.tolist(),
                            'prefix_diagnostic_rendering_not_a_translation': render(source, decoded)[:800],
                            'suffix_diagnostic_rune_labels_not_a_translation': ''.join(__import__('lp_core').LET[int(x)]+' ' for x in decoded[cutoff:cutoff+160]).strip(),
                        })
                for family, lo, hi in [('affine', 0, 812), ('sampled_arbitrary', 812, len(bank))]:
                    peaks = aggregate[lo:hi].max(axis=1)
                    best = lo+int(peaks.argmax())
                    output['pooled_results'].append({'policy': policy, 'update': update, 'family': family,
                        'shared_input_and_output_key_maximum_peak_count': int(peaks.max()),
                        'shared_input_and_output_key_maximum_peak_frequency': float(peaks.max()/n),
                        'best_shared_full_peak_key': key_description(best, random_keys),
                        'shared_full_peak_histogram': aggregate[best].tolist(),
                        'shared_full_peak_counts_for_every_input_key': peaks.tolist(),
                        'independent_section_input_and_output_keys_maximum_peak_count': independent[family],
                        'independent_section_input_and_output_keys_maximum_peak_frequency': independent[family]/n,
                        'shared_full_corpus_gate_survivors': int(sum(peaks/n >= PEAK_GATE))})
                print(policy, update, 'completed', flush=True)
    for policy in POLICIES:
        front = next(x for x in output['pooled_results'] if x['policy'] == policy and x['family'] == 'affine' and x['update'] == 'front')
        back = next(x for x in output['pooled_results'] if x['policy'] == policy and x['family'] == 'affine' and x['update'] == 'back')
        assert front['shared_input_and_output_key_maximum_peak_count'] == back['shared_input_and_output_key_maximum_peak_count']
        assert front['independent_section_input_and_output_keys_maximum_peak_count'] == back['independent_section_input_and_output_keys_maximum_peak_count']
    output['elapsed_seconds'] = time.time()-started
    output['software'] = {'numpy': np.__version__, 'kernel': 'C++17 compiled by g++ -O3; temporary shared library'}
    (ROOT/'recency_input_results.json').write_text(json.dumps(output, indent=2)+'\n')
    print('Wrote recency_input_results.json', 'seconds', round(output['elapsed_seconds'], 2), flush=True)
    for r in output['pooled_results']:
        print(r['policy'], r['update'], r['family'],
              r['shared_input_and_output_key_maximum_peak_frequency'],
              r['independent_section_input_and_output_keys_maximum_peak_frequency'], flush=True)


if __name__ == '__main__':
    main()
