#!/usr/bin/env python3
"""Exact finite checks for rational-sequence learning and Dyck quotients.

The learning audit works over Fraction and includes the finite-difference
family proving that 2r samples are sharp.  The Dyck audit uses an independent
stack interpreter, literal bounded futures, and finite-controller balanced-word
saturation.  These checks challenge boundary cases; the unbounded arguments
remain the mathematical proof.
"""

from __future__ import annotations

from fractions import Fraction
from itertools import product
import json
import math


def berlekamp_massey(values):
    values = [Fraction(x) for x in values]
    c = [Fraction(1)]
    bpoly = [Fraction(1)]
    length = 0
    shift = 1
    last = Fraction(1)
    for n in range(len(values)):
        discrepancy = values[n]
        for i in range(1, length + 1):
            discrepancy += c[i] * values[n - i]
        if discrepancy == 0:
            shift += 1
            continue
        old = c[:]
        scale = -discrepancy / last
        needed = len(bpoly) + shift
        c.extend(Fraction(0) for _ in range(max(0, needed - len(c))))
        for i, coefficient in enumerate(bpoly):
            c[i + shift] += scale * coefficient
        if 2 * length <= n:
            length = n + 1 - length
            bpoly = old
            last = discrepancy
            shift = 1
        else:
            shift += 1
    c.extend(Fraction(0) for _ in range(max(0, length + 1 - len(c))))
    return length, tuple(c[: length + 1])


def recurrence_sequence(initial, coefficients, total):
    seq = [Fraction(x) for x in initial]
    order = len(coefficients)
    while len(seq) < total:
        seq.append(sum(Fraction(coefficients[i]) * seq[-i - 1] for i in range(order)))
    return tuple(seq)


def predict(prefix, connection, total):
    seq = list(prefix)
    order = len(connection) - 1
    while len(seq) < total:
        seq.append(-sum(connection[i] * seq[-i] for i in range(1, order + 1)))
    return tuple(seq)


def rational_learning_audit():
    recovered = 0
    degenerate = 0
    for order in range(1, 5):
        for coefficients in product((-1, 0, 1), repeat=order):
            for initial in product((-1, 0, 1), repeat=order):
                seq = recurrence_sequence(initial, coefficients, 4 * order + 8)
                true_order, _ = berlekamp_massey(seq)
                if true_order != order:
                    degenerate += 1
                    continue
                learned_order, connection = berlekamp_massey(seq[: 2 * order])
                assert learned_order == order
                assert predict(seq[: 2 * order], connection, len(seq)) == seq
                recovered += 1

    sharp = []
    moment_checks = 0
    for rank in range(1, 8):
        width = 2 * rank - 1
        even = []
        odd = []
        for node in range(2 * rank):
            weight = math.comb(width, node)
            (even if node % 2 == 0 else odd).append((node, weight))

        def moments(terms, total):
            return tuple(
                sum(Fraction(weight) * Fraction(node) ** n for node, weight in terms)
                for n in range(total)
            )

        left = moments(even, 2 * rank + 3)
        right = moments(odd, 2 * rank + 3)
        assert left[: 2 * rank - 1] == right[: 2 * rank - 1]
        assert left[2 * rank - 1] != right[2 * rank - 1]
        assert berlekamp_massey(left)[0] == rank
        assert berlekamp_massey(right)[0] == rank
        moment_checks += 2 * rank
        sharp.append({
            "rank": rank,
            "shared_coefficients": 2 * rank - 1,
            "first_difference_index": 2 * rank - 1,
        })

    return {
        "exact_2r_reconstructions": recovered,
        "degenerate_presentations_skipped": degenerate,
        "sharp_finite_difference_families": sharp,
        "moment_equalities_checked": moment_checks,
    }


def words(alphabet, max_length, exact=False):
    lengths = (max_length,) if exact else range(max_length + 1)
    for length in lengths:
        yield from product(alphabet, repeat=length)


def step_stack(stack, symbol, bracket_types):
    if stack is None:
        return None
    kind, type_index = symbol
    if kind == "o":
        return stack + (type_index,)
    if kind == "c":
        if not stack or stack[-1] != type_index:
            return None
        return stack[:-1]
    return stack


def run_stack(stack, word, bracket_types):
    for symbol in word:
        stack = step_stack(stack, symbol, bracket_types)
    return stack


def matrix_rank(rows):
    matrix = [[Fraction(x) for x in row] for row in rows]
    if not matrix:
        return 0
    rank = 0
    width = len(matrix[0])
    for column in range(width):
        pivot = next((i for i in range(rank, len(matrix)) if matrix[i][column]), None)
        if pivot is None:
            continue
        matrix[rank], matrix[pivot] = matrix[pivot], matrix[rank]
        scale = matrix[rank][column]
        matrix[rank] = [x / scale for x in matrix[rank]]
        for i in range(len(matrix)):
            if i != rank and matrix[i][column]:
                scale = matrix[i][column]
                matrix[i] = [x - scale * y for x, y in zip(matrix[i], matrix[rank])]
        rank += 1
        if rank == len(matrix):
            break
    return rank


def stack_set(bracket_types, max_height):
    return [s for height in range(max_height + 1) for s in product(range(bracket_types), repeat=height)]


def canonical_closer(stack):
    return tuple(("c", type_index) for type_index in reversed(stack))


def dyck_hankel_audit():
    bounded_cases = 0
    field_rank_checks = 0
    identity_cells = 0
    midpoint_cases = 0
    for bracket_types in (1, 2):
        for neutrals in (0, 1):
            alphabet = (
                tuple(("o", i) for i in range(bracket_types))
                + tuple(("c", i) for i in range(bracket_types))
                + tuple(("n", i) for i in range(neutrals))
            )
            for horizon in range(5):
                stacks = stack_set(bracket_types, horizon)
                continuations = list(words(alphabet, horizon))
                rows = [
                    [int(run_stack(stack, continuation, bracket_types) == ()) for continuation in continuations]
                    for stack in stacks
                ]
                rows.append([0] * len(continuations))
                expected = sum(bracket_types ** height for height in range(horizon + 1))
                assert matrix_rank(rows) == expected
                field_rank_checks += 1
                for left in stacks:
                    for right in stacks:
                        value = run_stack(left, canonical_closer(right), bracket_types) == ()
                        assert value == (left == right)
                        identity_cells += 1
                bounded_cases += len(rows) * len(continuations)

                prefixes = list(words(alphabet, horizon, exact=True))
                reachable = {run_stack((), prefix, bracket_types) for prefix in prefixes}
                live = {stack for stack in reachable if stack is not None}
                if neutrals:
                    expected_live = set(stacks)
                else:
                    expected_live = {
                        stack for stack in stacks if (horizon - len(stack)) % 2 == 0
                    }
                assert live == expected_live
                for stack in live:
                    padding = horizon - len(stack)
                    if neutrals:
                        suffix = canonical_closer(stack) + (("n", 0),) * padding
                    else:
                        pairs = (("o", 0), ("c", 0)) * (padding // 2)
                        suffix = canonical_closer(stack) + pairs
                    assert len(suffix) == horizon
                    assert run_stack(stack, suffix, bracket_types) == ()
                    for other in live - {stack}:
                        assert run_stack(other, suffix, bracket_types) != ()
                        identity_cells += 1
                midpoint_cases += len(prefixes)
    return {
        "bounded_future_cells": bounded_cases,
        "field_rank_checks": field_rank_checks,
        "identity_minor_cells": identity_cells,
        "exact_midpoint_prefixes": midpoint_cases,
    }


def transition_maps(states):
    return list(product(range(states), repeat=states))


def apply_map(mapping, state):
    return mapping[state]


def balanced_saturation(states, opener, closer, neutral):
    relation = {(q, q) for q in range(states)}
    relation |= {(q, apply_map(neutral, q)) for q in range(states)}
    changed = True
    while changed:
        changed = False
        additions = set()
        for p, q in relation:
            for r, s in relation:
                if q == r:
                    additions.add((p, s))
        for p in range(states):
            p1 = apply_map(opener, p)
            for q1, q2 in relation:
                if q1 == p1:
                    additions.add((p, apply_map(closer, q2)))
        if not additions <= relation:
            relation |= additions
            changed = True
    return relation


def controlled_step(configuration, symbol, maps):
    stack, state = configuration
    if stack is None:
        return configuration
    kind, type_index = symbol
    state = apply_map(maps[kind], state)
    if kind == "o":
        return stack + (type_index,), state
    if kind == "c":
        if not stack or stack[-1] != type_index:
            return None, state
        return stack[:-1], state
    return stack, state


def controlled_run(configuration, word, maps):
    for symbol in word:
        configuration = controlled_step(configuration, symbol, maps)
    return configuration


def controller_quotient_audit():
    controller_cases = 0
    signature_cells = 0
    action_checks = 0
    bond_checks = 0
    maps2 = transition_maps(2)
    alphabet = (("o", 0), ("c", 0), ("n", 0))
    for opener in maps2:
        for closer in maps2:
            for neutral in maps2:
                maps = {"o": opener, "c": closer, "n": neutral}
                relation = balanced_saturation(2, opener, closer, neutral)
                # Every saturated pair is witnessed by an enumerated balanced word
                # in these two-state systems; eight symbols is a hostile finite cap.
                witnessed = set()
                for word in words(alphabet, 8):
                    if run_stack((), word, 1) == ():
                        for start in range(2):
                            end = controlled_run(((), start), word, maps)[1]
                            witnessed.add((start, end))
                assert relation <= witnessed

                reachable = {((), q) for p, q in relation if p == 0}
                for height in range(1, 4):
                    new = set()
                    for stack, state in reachable:
                        if len(stack) != height - 1:
                            continue
                        after_open = apply_map(opener, state)
                        for p, q in relation:
                            if p == after_open:
                                new.add((stack + (0,), q))
                    reachable |= new
                failure = (None, 0)

                signatures = {}
                def future_signature(configuration, horizon):
                    continuations = list(words(alphabet, horizon))
                    stack, _ = configuration
                    if stack is not None and len(stack) > horizon:
                        return tuple(0 for _ in continuations)
                    return tuple(
                        int(controlled_run(configuration, word, maps)[0] == ())
                        for word in continuations
                    )

                for horizon in range(4):
                    continuations = list(words(alphabet, horizon))
                    level = {}
                    for configuration in reachable | {failure}:
                        signature = future_signature(configuration, horizon)
                        level[configuration] = signature
                        signature_cells += len(signature)
                    signatures[horizon] = level
                    if horizon:
                        prior = signatures[horizon - 1]
                        prior_words = list(words(alphabet, horizon - 1))
                        for configuration, signature in level.items():
                            restriction = signature[: len(prior_words)]
                            assert restriction == prior[configuration]
                            bond_checks += 1
                        for left in level:
                            for right in level:
                                if level[left] == level[right]:
                                    for symbol in alphabet:
                                        lnext = controlled_step(left, symbol, maps)
                                        rnext = controlled_step(right, symbol, maps)
                                        assert future_signature(lnext, horizon - 1) == future_signature(rnext, horizon - 1)
                                        action_checks += 1

                # A deep stack has the all-false bounded future and therefore
                # represents the same coherent limit point as absorbing failure.
                for horizon in range(4):
                    deep = ((0,) * (horizon + 1), 0)
                    assert all(
                        controlled_run(deep, word, maps)[0] != ()
                        for word in words(alphabet, horizon)
                    )
                controller_cases += 1
    return {
        "two_state_controllers": controller_cases,
        "bounded_signature_cells": signature_cells,
        "bonding_checks": bond_checks,
        "degree_lowering_action_checks": action_checks,
        "deep_tail_failure_controls": 4 * controller_cases,
    }


def main():
    output = {
        "status": "pass",
        "rational_learning": rational_learning_audit(),
        "generalized_dyck": dyck_hankel_audit(),
        "controlled_dyck": controller_quotient_audit(),
    }
    print(json.dumps(output, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
