"""Session 71: full alphabet coverage in every repeat-channel state.

Exact relaxed histogram concentration bounds. Keeps Session 70's channel
and selectors, adds a declared d-type alphabet encodable in every state.
Empirical first-ten-bin screens are not universal English constraints.
"""
from pathlib import Path
from itertools import product
import hashlib
import json
import numpy as np
from lp_core import vals
from word_core import GROUPS
from reset_state_audit import SOURCE, SINGLE
from section_repeat_channel import bank, matrices, metrics, references, admissible, PARAMS

ROOT = Path(__file__).resolve().parent


def entropy(counts):
    counts = np.asarray(counts)
    p = counts[counts > 0]/counts.sum()
    return float(-np.sum(p*np.log2(p)))


def profile(h, g, d):
    """Exact most concentrated histogram, ignoring singleton constraints.

    R is anonymous output zero. Every other output has a nonempty code class
    in every state. Extra codes join the largest row count; align row ranks.
    """
    k, q = h.shape
    capacities = [q-int(np.sum(g == s))+1 for s in range(k)]
    if d < 2 or min(capacities) < d:
        return {'encodable_in_every_state': False, 'row_alphabet_capacities': capacities}
    hist = np.zeros(d, dtype=np.int64)
    witness = np.zeros((k, q), dtype=int)
    for state in range(k):
        inside = np.flatnonzero(g == state)
        outside = sorted(np.flatnonzero(g != state), key=lambda c:(-h[state, c], int(c)))
        hist[0] += h[state, inside].sum()
        extra = len(outside)-(d-1)
        witness[state, outside[:extra+1]] = 1
        hist[1] += sum(int(h[state, c]) for c in outside[:extra+1])
        for label, c in enumerate(outside[extra+1:], 2):
            witness[state, c] = label
            hist[label] += h[state, c]
        assert len(np.unique(witness[state])) == d
        assert np.all(witness[state, inside] == 0)
        assert np.all(witness[state, g != state] != 0)
    attained = np.bincount(witness.ravel(), weights=h.ravel(), minlength=d).astype(int)
    assert np.array_equal(hist, attained)
    sorted_hist = np.sort(hist)[::-1]
    assert int(sorted_hist.sum()) == int(h.sum())
    return {'encodable_in_every_state': True, 'declared_alphabet_types': d,
            'row_alphabet_capacities': capacities, 'scored': int(h.sum()),
            'R_positions': int(hist[0]), 'maximum_non_R_positions': int(hist[1]),
            'anonymous_histogram': hist.tolist(), 'sorted_upper_histogram': sorted_hist.tolist(),
            'sorted_partial_sum_upper_counts': np.cumsum(sorted_hist).tolist(),
            'minimum_histogram_entropy_bits': entropy(hist),
            'attaining_anonymous_table': witness.tolist(),
            'witness_is_not_an_English_key': True}


def collapsed_hist(cipher, g, start=1, stop=None):
    stop = len(cipher) if stop is None else stop
    h = np.zeros((int(g.max())+1, len(g)), dtype=np.int64)
    offsets = np.arange(max(start, 1), stop)
    np.add.at(h, (g[cipher[offsets-1]], cipher[offsets]), 1)
    return h


def reference_profiles():
    out = []
    pooled = np.zeros(29, dtype=int)
    for name, text in GROUPS.items():
        h = np.bincount(vals(text)[1:], minlength=29)
        pooled += h
        out.append({'name': name, 'scored': int(h.sum()), 'observed_types': int(np.count_nonzero(h)),
                    'rune_histogram': h.tolist(), 'sorted_histogram': np.sort(h)[::-1].tolist(),
                    'first_ten_partial_shares': (np.cumsum(np.sort(h)[::-1])[:10]/h.sum()).tolist()})
    out.append({'name': 'pooled', 'scored': int(pooled.sum()), 'observed_types': int(np.count_nonzero(pooled)),
                'rune_histogram': pooled.tolist(), 'sorted_histogram': np.sort(pooled)[::-1].tolist(),
                'first_ten_partial_shares': (np.cumsum(np.sort(pooled)[::-1])[:10]/pooled.sum()).tolist()})
    return out


def compare(bound, targets, d):
    if not bound['encodable_in_every_state']: return []
    n = bound['scored']
    upper = np.array(bound['sorted_partial_sum_upper_counts'][:10], dtype=np.int64)
    result = []
    for target in targets:
        if target['observed_types'] > d: continue
        actual = np.cumsum(target['sorted_histogram'])[:10]
        deficits = actual*n-upper*target['scored']
        failing = (np.flatnonzero(deficits > 0)+1).tolist()
        result.append({'reference': target['name'], 'reference_observed_types': target['observed_types'],
                       'failing_first_ten_partial_sums': failing,
                       'passes_common_letter_screen': not failing,
                       'first_failure': None if not failing else {
                           'rank_count': failing[0],
                           'maximum_decoded_share': float(upper[failing[0]-1]/n),
                           'reference_share': float(actual[failing[0]-1]/target['scored'])}})
    return result


def bank_profiles(h, selectors, k, d):
    """Vectorized version of the same bound, without retaining all witnesses."""
    q = h.shape[1]
    if any(q-int(np.sum(selectors[0] == s))+1 < d for s in range(k)): return None
    normal = np.zeros((len(selectors), d-1), dtype=np.int64)
    r_count = np.zeros(len(selectors), dtype=np.int64)
    for state in range(k):
        membership = selectors == state
        weighted = membership.astype(np.int64) @ h
        r_count += np.sum(weighted*membership, axis=1)
        length = q-int(membership[0].sum())
        off = np.sort(np.where(membership, -1, weighted), axis=1)[:, ::-1][:, :length]
        extra = length-(d-1)
        normal[:, 0] += np.sum(off[:, :extra+1], axis=1)
        normal[:, 1:] += off[:, extra+1:]
    return np.sort(np.c_[r_count, normal], axis=1)[:, ::-1]


def exhaustive_controls():
    fixtures = [
        (np.array([[4,2,8,3,1], [5,9,2,7,6]]), np.array([0,1,0,1,0]), 3),
        (np.array([[2,4,0,8,3,1], [7,1,6,2,4,5], [1,9,3,7,2,6]]), np.array([0,1,2,0,1,2]), 3),
        (np.array([[3,7,1,5,2], [4,1,8,2,6], [9,3,2,1,7]]), np.array([0,1,2,0,1]), 4)]
    out = []
    for h, g, d in fixtures:
        bound = profile(h, g, d)
        rows = []
        for state in range(len(h)):
            outside = np.flatnonzero(g != state)
            options = []
            for outputs in product(range(1, d), repeat=len(outside)):
                if len(set(outputs)) != d-1: continue
                row = np.zeros(len(g), dtype=int)
                row[outside] = outputs
                options.append(row)
            rows.append(options)
        extrema = np.zeros(d, dtype=np.int64); minimum_entropy = 100.; count = 0
        for selection in product(*rows):
            mp = np.array(selection)
            v = np.bincount(mp.ravel(), weights=h.ravel(), minlength=d).astype(int)
            extrema = np.maximum(extrema, np.cumsum(np.sort(v)[::-1]))
            minimum_entropy = min(minimum_entropy, entropy(v)); count += 1
        assert extrema.tolist() == bound['sorted_partial_sum_upper_counts']
        assert abs(minimum_entropy-bound['minimum_histogram_entropy_bits']) < 1e-12
        out.append({'cipher_alphabet': len(g), 'states': len(h), 'declared_plaintext_alphabet': d,
                    'complete_surjective_tables_exhausted': count, 'all_partial_sum_extrema_agree': True,
                    'minimum_entropy_agrees': True})
    impossible = profile(np.ones((2, 4), dtype=int), np.array([0,1,0,1]), 4)
    assert not impossible['encodable_in_every_state']
    return {'fixtures': out, 'encodability_failure_control': impossible}


def positive_controls(old):
    rng = np.random.default_rng(71001); out = []
    for fixture in old['recovery_controls']:
        if fixture['repeat_probability_when_alternative_exists'] != 0: continue
        k = fixture['tables']
        g = ((fixture['a']*np.arange(29)+fixture['b']) % 29) % k
        for name, text in GROUPS.items():
            mp = np.array(fixture['independent_section_tables'][name])
            alphabet = set(map(int, mp.ravel())); d = len(alphabet)
            assert all(set(map(int, row)) == alphabet for row in mp)
            plain = vals(text); previous = 0; cipher = []
            for p in plain:
                options = np.flatnonzero(mp[int(g[previous])] == p).tolist()
                alternatives = [c for c in options if c != previous]
                c = int(rng.choice(alternatives or options))
                cipher.append(c); previous = c
            decoded, previous = [], 0
            for c in cipher:
                decoded.append(int(mp[int(g[previous]), c])); previous = c
            assert np.array_equal(decoded, plain)
            cc = np.array(cipher); h = collapsed_hist(cc, g); b = profile(h, g, d)
            actual = np.bincount(plain[1:], minlength=29)
            upper = np.r_[b['sorted_partial_sum_upper_counts'], np.full(29-d, actual.sum())]
            assert np.all(np.cumsum(np.sort(actual)[::-1]) <= upper)
            b2 = bank_profiles(matrices(cc, [], 1, len(cc))['h'], g[None, :], k, d)
            assert b2[0].tolist() == b['sorted_upper_histogram']
            out.append({'states': k, 'source_group': name, 'declared_alphabet': d,
                        'scored': int(actual.sum()), 'exact_known_table_inverse_recovery': True,
                        'true_histogram_admitted_by_every_partial_sum': True,
                        'attaining_relaxed_witness_verified': True,
                        'scalar_and_vectorized_bounds_agree': True,
                        'unknown_key_recovery_not_demonstrated': True})
    return out


def restricted_domain_counterexample():
    g = np.array([0,1,0,1])
    mp = np.array([[0,1,0,2], [3,0,4,0]])
    cipher = np.array([1,0,3,2,0,2,1,3,0,1])
    plain, previous = [], 0
    for c in cipher:
        plain.append(int(mp[g[previous], c])); previous = int(c)
    assert set(plain) == set(range(5))
    recovered, previous = [], 0
    for p in plain:
        options = np.flatnonzero(mp[g[previous]] == p)
        assert len(options)
        alternatives = [int(c) for c in options if c != previous]
        c = (alternatives or options.tolist())[0]
        recovered.append(int(mp[g[previous], c])); previous = int(c)
    assert recovered == plain
    for state in range(2): assert np.all(mp[state, g == state] == 0)
    assert all(len(set(row)) == 3 for row in mp)
    impossible = profile(collapsed_hist(cipher, g), g, 5)
    assert not impossible['encodable_in_every_state']
    return {'cipher_types': 4, 'global_plaintext_types': 5, 'types_available_per_state': [3,3],
            'cipher': cipher.tolist(), 'plaintext': plain, 'table': mp.tolist(),
            'restricted_message_inverse_exact': True,
            'full_alphabet_per_state_premise_false': True,
            'not_a_Cicada_key': True}


def prior_bijective_regression(sections):
    h = sum((collapsed_hist(c, np.arange(29)) for c in sections.values() if len(c) > 100),
            np.zeros((29, 29), dtype=np.int64))
    current = profile(h, np.arange(29), 29)
    old = json.loads((ROOT/'keyed_difference_results.json').read_text())
    assert np.array_equal(np.sort(current['anonymous_histogram']),
                          np.sort(old['pooled']['most_concentrated_relaxed_histogram']))
    assert current['scored'] == old['pooled']['edges']
    assert current['R_positions'] == old['pooled']['doubles']
    return {'Session_33_bijective_limit_reproduced': True, 'scored': current['scored'],
            'repeat_channel_positions': current['R_positions'],
            'maximum_single_output_share': max(current['anonymous_histogram'])/current['scored'],
            'minimum_histogram_entropy_bits': current['minimum_histogram_entropy_bits']}


def main():
    digest = hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest()
    assert digest == 'b0b2e13c3cd35512ee5955a2d9a56c22d7410e283d35dd9c90abbaa7b99279eb'
    old = json.loads((ROOT/'section_repeat_channel_results.json').read_text())
    assert old['master_sha256'] == digest
    targets = reference_profiles()
    sections = {i: vals(t) for i, t in enumerate(SOURCE.split('$'))}
    primary = []
    for previous in old['results']:
        if previous['singleton_sensitivity'] != 'all_audited' or previous['allowed_letters'] != 'AI': continue
        frozen = previous['first_training_selector_frozen']
        if not frozen['full_passes_necessary_screens']: continue
        k = previous['first_training_passing_k']
        g = ((frozen['a']*np.arange(29)+frozen['b']) % 29) % k
        records = []
        for part, start, stop in [('full', 1, len(sections[previous['section']])),
                                  ('continuation', previous['cut_rune_offset'], len(sections[previous['section']]))]:
            h = collapsed_hist(sections[previous['section']], g, start, stop)
            for d in range(20, 30):
                b = profile(h, g, d)
                records.append({'part': part, 'declared_alphabet': d, 'bound': b,
                                'reference_comparisons': compare(b, targets, d)})
        primary.append({'section': previous['section'], 'states': k, 'a': frozen['a'], 'b': frozen['b'],
                        'selector_frozen_from_Session_70': True, 'results': records})
    secondary = []
    base = 0; ref = references()['AIO']
    for section, text in enumerate(SOURCE.split('$')):
        c = vals(text); single = [w['rune_offset']-base for w in SINGLE if w['section'] == section]
        base += len(c)
        if len(c) <= 100: continue
        record = matrices(c, single, 1, len(c)); n = int(record['h'].sum())
        per_k = []
        for k in range(1, 9):
            selectors = bank(k); m = metrics(record, selectors, k)
            old_eligible = admissible(m, ref, False) | admissible(m, ref, True)
            alphabets = []
            for d in range(20, 30):
                v = bank_profiles(record['h'], selectors, k, d)
                if v is None:
                    alphabets.append({'alphabet': d, 'encodable_in_every_state': False}); continue
                upper = np.cumsum(v, axis=1)[:, :10]
                passed_any = np.zeros(len(selectors), bool); by_reference = {}
                for target in targets:
                    if target['observed_types'] > d: continue
                    actual = np.cumsum(target['sorted_histogram'])[:10]
                    passed = np.all(upper*target['scored'] >= actual[None, :]*n, axis=1) & old_eligible
                    passed_any |= passed
                    by_reference[target['name']] = int(passed.sum())
                count = int(passed_any.sum()); witness = None
                if count:
                    index = int(np.flatnonzero(passed_any)[0]); a, b = PARAMS[index]
                    witness = {'a': a, 'b': b, 'capacity_only_not_a_decipherment': True}
                alphabets.append({'alphabet': d, 'encodable_in_every_state': True,
                                  'settings_passing_both_old_and_new_screens': count,
                                  'passing_by_reference': by_reference, 'first_capacity_survivor': witness})
            per_k.append({'states': k, 'old_AIO_screen_survivors': int(old_eligible.sum()), 'alphabets': alphabets})
        secondary.append({'section': section, 'scored': n, 'results': per_k})
        print('secondary section', section, 'alphabet>=25 survivors',
              sum(x.get('settings_passing_both_old_and_new_screens', 0) for k in per_k for x in k['alphabets'] if x['alphabet'] >= 25), flush=True)
    out = {'master_sha256': digest,
           'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
           'previous_results_sha256': hashlib.sha256((ROOT/'section_repeat_channel_results.json').read_bytes()).hexdigest(),
           'protocol': {'new_premise': 'One declared d-type plaintext alphabet, including R, is encodable in every state; every row is onto that same alphabet. d=20..29.',
                        'retained_model': 'Session 70 same-group iff R rule and one rune per plaintext rune. Each section has independent codebooks; no word/page resets.',
                        'row_capacity': 'R reserves q_s codes in a state group. The other Q-q_s codes must cover d-1 non-R outputs, so d<=Q-q_s+1.',
                        'exact_concentration': 'For each row, merge extra+1 largest non-R code counts into one output and keep the other d-2 outputs nonempty; align sorted row ranks. Add the fixed R count. This attains every maximum sorted partial sum and minimum histogram entropy in the relaxed family.',
                        'primary': 'Freeze the three Session 70 full-section screen survivors, without refitting. Compare full sections and continuations for d=20..29.',
                        'reference_screen': 'First ten sorted partial sums only, against six empirical profiles with observed alphabet<=d. One or more deficient partial sums fails a profile. The pooled profile has 26 observed types.',
                        'secondary': 'Optimistically allow all 812 affine selectors per k=1..8 in each substantial section; retain the original full-data A/I/O screens. No fitted table or independent validation.',
                        'limits': 'Conditional codebook coverage requirement, empirical reference shapes and retrospective selected models. Singleton table constraints are ignored in the new upper bound, favoring the model. Passing is necessary capacity only; restricted state alphabets, fewer than 20 plaintext types and other mechanisms remain open.'},
           'references': targets, 'frozen_survivors': primary, 'secondary_refit_capacity': secondary,
           'exhaustive_controls': exhaustive_controls(), 'known_table_controls': positive_controls(old),
           'prior_bijective_regression': prior_bijective_regression(sections),
           'restricted_state_domain_counterexample': restricted_domain_counterexample()}
    (ROOT/'state_alphabet_bound_results.json').write_text(json.dumps(out, indent=2)+'\n')
    for p in primary:
        counts = [x for x in p['results'] if x['part'] == 'full']
        print('frozen', p['section'], 'reference passes', sum(z['passes_common_letter_screen'] for x in counts for z in x['reference_comparisons']))
    print('known controls', len(out['known_table_controls']), 'exhaustive tables', sum(x['complete_surjective_tables_exhausted'] for x in out['exhaustive_controls']['fixtures']))


if __name__ == '__main__': main()