#!/usr/bin/env python3
"""Finite audit for FPRD-T156, residual towers and transport."""

from __future__ import annotations

import json
from itertools import product
from pathlib import Path


ALPHABET = (0, 1)


def words_exact(length: int) -> list[tuple[int, ...]]:
    return list(product(ALPHABET, repeat=length))


def words_upto(length: int) -> list[tuple[int, ...]]:
    return [word for size in range(length + 1) for word in words_exact(size)]


def row(language: dict[tuple[int, ...], int], prefix: tuple[int, ...], horizon: int) -> tuple[int, ...]:
    return tuple(language[prefix + suffix] for suffix in words_upto(horizon))


def audit_finite_language_trees(depth: int = 3) -> dict[str, int]:
    universe = words_upto(depth)
    index = {word: position for position, word in enumerate(universe)}
    languages = 0
    restriction_checks = 0
    derivative_checks = 0
    word_transport_checks = 0

    for mask in range(1 << len(universe)):
        language = {word: (mask >> index[word]) & 1 for word in universe}
        languages += 1
        for prefix_length in range(depth + 1):
            for prefix in words_exact(prefix_length):
                remaining = depth - prefix_length
                for horizon in range(remaining):
                    larger = row(language, prefix, horizon + 1)
                    smaller = row(language, prefix, horizon)
                    assert larger[: len(smaller)] == smaller
                    restriction_checks += 1

                    large_suffixes = words_upto(horizon + 1)
                    location = {suffix: pos for pos, suffix in enumerate(large_suffixes)}
                    for letter in ALPHABET:
                        derived = tuple(
                            larger[location[(letter,) + suffix]]
                            for suffix in words_upto(horizon)
                        )
                        assert derived == row(language, prefix + (letter,), horizon)
                        derivative_checks += 1

                # Check the general word action whenever a positive transport
                # word and a nonnegative target horizon fit in the finite tree.
                for transport_length in range(1, remaining + 1):
                    target_horizon = remaining - transport_length
                    source = row(language, prefix, remaining)
                    source_suffixes = words_upto(remaining)
                    location = {suffix: pos for pos, suffix in enumerate(source_suffixes)}
                    for transport in words_exact(transport_length):
                        derived = tuple(
                            source[location[transport + suffix]]
                            for suffix in words_upto(target_horizon)
                        )
                        assert derived == row(language, prefix + transport, target_horizon)
                        word_transport_checks += 1

    return {
        "finite_tree_languages": languages,
        "restriction_checks": restriction_checks,
        "one_letter_derivative_checks": derivative_checks,
        "word_transport_checks": word_transport_checks,
    }


def reachable_states(transitions: tuple[int, ...], states: int) -> tuple[int, ...]:
    reached = {0}
    frontier = [0]
    while frontier:
        state = frontier.pop()
        for letter in ALPHABET:
            target = transitions[2 * state + letter]
            if target not in reached:
                reached.add(target)
                frontier.append(target)
    return tuple(sorted(reached))


def canonical_ids(signatures: dict[int, object]) -> dict[int, int]:
    values = {value: index for index, value in enumerate(sorted(set(signatures.values()), key=repr))}
    return {state: values[value] for state, value in signatures.items()}


def audit_dfas(max_states: int = 3) -> dict[str, int]:
    automata = 0
    quotient_levels = 0
    restriction_surjections = 0
    derivative_well_defined_checks = 0
    stabilization_checks = 0

    for states in range(1, max_states + 1):
        for transitions in product(range(states), repeat=2 * states):
            reachable = reachable_states(transitions, states)
            for outputs in product((0, 1), repeat=states):
                automata += 1
                ids = canonical_ids({state: outputs[state] for state in reachable})
                counts = [len(set(ids.values()))]
                level_ids = [ids]

                # At most |reachable|-1 strict refinements can occur.
                for _ in range(len(reachable) + 1):
                    signatures = {
                        state: (
                            outputs[state],
                            ids[transitions[2 * state]],
                            ids[transitions[2 * state + 1]],
                        )
                        for state in reachable
                    }
                    next_ids = canonical_ids(signatures)
                    counts.append(len(set(next_ids.values())))
                    level_ids.append(next_ids)
                    quotient_levels += 1

                    # Refinement makes restriction a surjection.
                    class_to_old: dict[int, int] = {}
                    for state in reachable:
                        new_class = next_ids[state]
                        old_class = ids[state]
                        assert class_to_old.setdefault(new_class, old_class) == old_class
                    restriction_surjections += 1

                    # A letter maps each new class to one old class.
                    for letter in ALPHABET:
                        derivative: dict[int, int] = {}
                        for state in reachable:
                            source_class = next_ids[state]
                            target_class = ids[transitions[2 * state + letter]]
                            assert derivative.setdefault(source_class, target_class) == target_class
                            derivative_well_defined_checks += 1

                    ids = next_ids

                assert all(left <= right for left, right in zip(counts, counts[1:]))
                first_equal = next(
                    index for index in range(len(counts) - 1)
                    if counts[index] == counts[index + 1]
                )
                assert all(count == counts[first_equal] for count in counts[first_equal:])
                assert counts[first_equal] <= len(reachable)
                stabilization_checks += 1

    return {
        "binary_dfas": automata,
        "dfa_quotient_levels": quotient_levels,
        "dfa_restriction_surjections": restriction_surjections,
        "dfa_derivative_well_defined_checks": derivative_well_defined_checks,
        "dfa_stabilization_checks": stabilization_checks,
    }


def audit_unary_completion(witnesses: int = 12) -> dict[str, int]:
    # alpha = 1, 10, 100, 1000, ... concatenated.
    sequence: list[int] = []
    starts: list[int] = []
    for zeros in range(witnesses + 3):
        starts.append(len(sequence))
        sequence.append(1)
        sequence.extend([0] * zeros)

    checks = 0
    for horizon in range(1, witnesses + 1):
        block = horizon + 1
        shift = starts[block] + 1
        assert sequence[shift : shift + horizon] == [0] * horizon
        assert 1 in sequence[shift + horizon :]
        checks += 1
    return {"unary_ideal_residual_witnesses": checks}


def run() -> dict[str, object]:
    result: dict[str, object] = {"status": "pass"}
    result.update(audit_finite_language_trees())
    result.update(audit_dfas())
    result.update(audit_unary_completion())
    return result


if __name__ == "__main__":
    result = run()
    output = Path(__file__).with_name("checker-output.json")
    output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))

