#!/usr/bin/env python3
"""Exact symbolic certificate for the two-reservoir reverse-and-add relay.

The certificate is not an expanding census.  It infers affine run templates
from three nondegenerate instances, then proves every transition as an
identity in X=2^m.  Once a template is inferred, the polynomial identity and
positivity of every run length certify all parameter values in the declared
domain.  The finitely many smaller q values are replayed separately.
"""

from __future__ import annotations

from dataclasses import dataclass
from fractions import Fraction
import json
from pathlib import Path
from typing import Callable, Iterable


Poly = dict[int, Fraction]


def clean(p: Poly) -> Poly:
    return {e: c for e, c in p.items() if c}


def add(p: Poly, q: Poly) -> Poly:
    out = dict(p)
    for e, c in q.items():
        out[e] = out.get(e, Fraction(0)) + c
    return clean(out)


def scale_shift(p: Poly, exponent_slope: int, exponent_constant: int) -> Poly:
    factor = Fraction(2) ** exponent_constant
    return clean({e + exponent_slope: c * factor for e, c in p.items()})


def poly_equal(p: Poly, q: Poly) -> bool:
    return clean(p) == clean(q)


@dataclass(frozen=True)
class Run:
    bit: str
    slope: int
    constant: int

    def length(self, m: int) -> int:
        return self.slope * m + self.constant


Template = tuple[Run, ...]


def run_lengths(word: str) -> list[tuple[str, int]]:
    out: list[tuple[str, int]] = []
    for bit in word:
        if out and out[-1][0] == bit:
            out[-1] = (bit, out[-1][1] + 1)
        else:
            out.append((bit, 1))
    return out


def infer_template(words: list[str], ms: list[int], minimum_m: int) -> Template:
    assert len(words) == len(ms) == 3
    encoded = [run_lengths(word) for word in words]
    bits = [tuple(bit for bit, _ in item) for item in encoded]
    assert bits[0] == bits[1] == bits[2], bits
    template: list[Run] = []
    for index, bit in enumerate(bits[0]):
        lengths = [item[index][1] for item in encoded]
        slope = lengths[1] - lengths[0]
        assert lengths[2] - lengths[1] == slope
        constant = lengths[0] - slope * ms[0]
        assert slope * minimum_m + constant > 0
        template.append(Run(bit, slope, constant))
    return tuple(template)


def template_poly(template: Template) -> Poly:
    value: Poly = {}
    for run in template:
        value = scale_shift(value, run.slope, run.constant)
        if run.bit == "1":
            value = add(value, {run.slope: Fraction(2) ** run.constant})
            value = add(value, {0: Fraction(-1)})
    return clean(value)


def reverse_template(template: Template) -> Template:
    return tuple(reversed(template))


def template_length(template: Template) -> tuple[int, int]:
    return (
        sum(run.slope for run in template),
        sum(run.constant for run in template),
    )


def shift_parameter(template: Template, delta: int) -> Template:
    """Substitute m -> m+delta in an affine-run template."""
    return tuple(
        Run(run.bit, run.slope, run.constant + run.slope * delta)
        for run in template
    )


def raw_lift_poly(
    u: Template,
    v: Template,
    width: tuple[int, int],
    gap: tuple[int, int],
) -> Poly:
    return add(
        scale_shift(template_poly(v), width[0] + gap[0], width[1] + gap[1]),
        template_poly(u),
    )


def assert_state_templates_equal(
    left: tuple[Template, Template, tuple[int, int]],
    right: tuple[Template, Template, tuple[int, int]],
    label: str,
) -> None:
    assert left[2] == right[2], (label, "width")
    assert poly_equal(template_poly(left[0]), template_poly(right[0])), (label, "u")
    assert poly_equal(template_poly(left[1]), template_poly(right[1])), (label, "v")


def format_template(template: Template) -> str:
    pieces = []
    for run in template:
        if run.slope == 0 and run.constant == 1:
            pieces.append(run.bit)
        else:
            pieces.append(f"{run.bit}^({run.slope}m{run.constant:+d})")
    return " ".join(pieces)


def reverse_binary(n: int) -> int:
    return int(format(n, "b")[::-1], 2)


def raw_step(n: int) -> tuple[int, int]:
    nxt = n + reverse_binary(n)
    return nxt, nxt.bit_length() - n.bit_length()


def fixed_reverse(n: int, width: int) -> int:
    return int(format(n, f"0{width}b")[::-1], 2)


def flank_step(state: tuple[int, int, int]) -> tuple[tuple[int, int, int], int]:
    u, v, width = state
    outward = v + fixed_reverse(u, width)
    inward = u + fixed_reverse(v, width)
    growth = int(outward >= 1 << width)
    return (inward, outward << (1 - growth), width + 1), growth


def raw_lift(state: tuple[int, int, int], gap: int) -> int:
    u, v, width = state
    return (v << (width + gap)) + u


def iterate_raw(n: int, steps: int) -> tuple[list[int], str]:
    states = [n]
    output = []
    for _ in range(steps):
        n, growth = raw_step(n)
        states.append(n)
        output.append(str(growth))
    return states, "".join(output)


def iterate_flank(
    state: tuple[int, int, int], steps: int
) -> tuple[list[tuple[int, int, int]], str]:
    states = [state]
    output = []
    for _ in range(steps):
        state, growth = flank_step(state)
        states.append(state)
        output.append(str(growth))
    return states, "".join(output)


def infer_raw_chain(
    generator: Callable[[int], int],
    steps: int,
    ms: list[int],
    minimum_m: int,
    label: str,
) -> tuple[list[Template], str]:
    chains = [iterate_raw(generator(m), steps) for m in ms]
    outputs = [out for _, out in chains]
    assert outputs[0] == outputs[1] == outputs[2]
    templates = []
    for index in range(steps + 1):
        words = [format(states[index], "b") for states, _ in chains]
        templates.append(infer_template(words, ms, minimum_m))
    for index in range(steps):
        current = template_poly(templates[index])
        reversed_current = template_poly(reverse_template(templates[index]))
        following = template_poly(templates[index + 1])
        assert poly_equal(add(current, reversed_current), following), (label, index)
        old_length = template_length(templates[index])
        new_length = template_length(templates[index + 1])
        assert old_length[0] == new_length[0]
        assert new_length[1] - old_length[1] == int(outputs[0][index])
    return templates, outputs[0]


def padded_word(value: int, width: int) -> str:
    assert 0 <= value < 1 << width
    return format(value, f"0{width}b")


def infer_flank_chain(
    generator: Callable[[int], tuple[int, int, int]],
    steps: int,
    ms: list[int],
    minimum_m: int,
    label: str,
) -> tuple[list[tuple[Template, Template, tuple[int, int]]], str]:
    chains = [iterate_flank(generator(m), steps) for m in ms]
    outputs = [out for _, out in chains]
    assert outputs[0] == outputs[1] == outputs[2]
    templates: list[tuple[Template, Template, tuple[int, int]]] = []
    for index in range(steps + 1):
        widths = [states[index][2] for states, _ in chains]
        width_slope = widths[1] - widths[0]
        assert widths[2] - widths[1] == width_slope
        width_constant = widths[0] - width_slope * ms[0]
        u_words = [padded_word(states[index][0], widths[j]) for j, (states, _) in enumerate(chains)]
        v_words = [padded_word(states[index][1], widths[j]) for j, (states, _) in enumerate(chains)]
        templates.append(
            (
                infer_template(u_words, ms, minimum_m),
                infer_template(v_words, ms, minimum_m),
                (width_slope, width_constant),
            )
        )
    for index in range(steps):
        u, v, width = templates[index]
        next_u, next_v, next_width = templates[index + 1]
        growth = int(outputs[0][index])
        assert next_width == (width[0], width[1] + 1)
        inward = add(template_poly(u), template_poly(reverse_template(v)))
        outward = add(template_poly(v), template_poly(reverse_template(u)))
        assert poly_equal(inward, template_poly(next_u)), (label, index, "u")
        assert poly_equal(
            scale_shift(outward, 0, 1 - growth), template_poly(next_v)
        ), (label, index, "v")
    return templates, outputs[0]


def compare_raw_generators(
    left: Callable[[int], int],
    right: Callable[[int], int],
    ms: list[int],
    minimum_m: int,
    label: str,
) -> None:
    left_template = infer_template([format(left(m), "b") for m in ms], ms, minimum_m)
    right_template = infer_template([format(right(m), "b") for m in ms], ms, minimum_m)
    assert poly_equal(template_poly(left_template), template_poly(right_template)), label


def s_state(a: int) -> tuple[int, int, int]:
    return 180 * (1 << a) - 3, 192 * (1 << a) + 42, a + 8


def u_word(p: int) -> int:
    return int("10" + "1" * p + "01" + "0" * (p + 1), 2)


def initial(c: int, q: int) -> int:
    return (1 << (9 * q - 1)) + c


def state_after_initial(c: int) -> tuple[int, int, int]:
    n = initial(c, 12)
    states, output = iterate_raw(n, 9)
    assert output == "011101111"
    end = states[-1]
    width = 20
    length = end.bit_length()
    gap = length - 2 * width
    return end & ((1 << width) - 1), end >> (length - width), width


def first_cycle_state(c: int, k: int) -> tuple[int, int, int]:
    state = state_after_initial(c)
    states, output = iterate_flank(state, 6 * k)
    assert output == "10" * (3 * k)
    return states[-1]


def seam_start(c: int, q: int) -> int:
    state = first_cycle_state(c, q - 4)
    states, output = iterate_flank(state, 2)
    assert output == "10"
    return raw_lift(states[-1], 0)


def h_state(q: int) -> tuple[int, int, int, int]:
    if q % 2 == 0:
        a = 3 * q // 2 - 2
        u = int("10001100" + "1" * a + "01", 2)
        v = int("11" + "0" * a + "10101110", 2)
        width = a + 10
        gap = 9 * q - 17
    else:
        a = (3 * q - 3) // 2
        z = (3 * q - 1) // 2
        u = int("1101000" + "1" * a + "01", 2)
        v = int("11" + "0" * z + "110011", 2)
        width = z + 8
        gap = 9 * q - 16
    return u, v, width, gap


def terminal_start(q: int) -> int:
    if q % 2 == 0:
        a0 = 3 * q // 2 + 6
        a1 = a0 + 6 * (q - 3)
        states, output = iterate_flank(s_state(a1), 1)
        assert output == "1"
        return raw_lift(states[-1], 0)
    a0 = (3 * q + 19) // 2
    a1 = a0 + 6 * (q - 4)
    states, output = iterate_flank(s_state(a1), 3)
    assert output == "101"
    return raw_lift(states[-1], 1)


def main() -> None:
    certificate: dict[str, object] = {
        "status": "pass",
        "method": "affine run templates proved as Laurent-polynomial identities in X=2^m",
    }

    # The nine-step setup is exact for every q >= 10; q=4,...,9 are checked below.
    for name, c in (("a", 154), ("b", 1178)):
        templates, output = infer_raw_chain(
            lambda q, c=c: initial(c, q), 9, [10, 11, 12], 10, f"initial-{name}"
        )
        assert output == "011101111"
        expected_state = state_after_initial(c)
        compare_raw_generators(
            lambda q, c=c: iterate_raw(initial(c, q), 9)[0][-1],
            lambda q, s=expected_state: raw_lift(s, 9 * q - 33),
            [10, 11, 12],
            10,
            f"initial-lift-{name}",
        )
        certificate[f"initial_{name}_end_state"] = list(expected_state)
        certificate[f"initial_{name}_templates"] = [format_template(t) for t in templates]

    assert state_after_initial(154) == (30711, 588432, 20)
    assert state_after_initial(1178) == (236535, 589362, 20)

    # The first reservoir: two fixed cycles enter parameterized six-step families.
    first_cycle_chains: dict[tuple[int, int], list[tuple[Template, Template, tuple[int, int]]]] = {}
    for c, label in ((154, "a"), (1178, "b")):
        anchor = first_cycle_state(c, 2)
        certificate[f"first_cycle_{label}_anchor"] = list(anchor)
        for parity in (0, 1):
            templates, output = infer_flank_chain(
                lambda m, c=c, parity=parity: first_cycle_state(c, 2 * m + parity),
                6,
                [3, 4, 5],
                3,
                f"first-cycle-{label}-parity-{parity}",
            )
            assert output == "101010"
            first_cycle_chains[(c, parity)] = templates
            certificate[f"first_cycle_{label}_parity_{parity}_start_template"] = {
                "u": format_template(templates[0][0]),
                "v": format_template(templates[0][1]),
                "width": templates[0][2],
            }
        even = first_cycle_chains[(c, 0)]
        odd = first_cycle_chains[(c, 1)]
        assert_state_templates_equal(even[-1], odd[0], f"{label}-even-to-odd")
        shifted_even_start = (
            shift_parameter(even[0][0], 1),
            shift_parameter(even[0][1], 1),
            (even[0][2][0], even[0][2][1] + even[0][2][0]),
        )
        assert_state_templates_equal(odd[-1], shifted_even_start, f"{label}-odd-to-even")

    # The first seam: a needs nine raw steps and enters a four-step attractor.
    for parity, ms in ((0, [5, 6, 7]), (1, [5, 6, 7])):
        q_of = lambda m, parity=parity: 2 * m + parity
        a_bridge, a_bridge_output = infer_raw_chain(
            lambda m: seam_start(154, q_of(m)),
            9,
            ms,
            ms[0],
            f"a-first-seam-parity-{parity}",
        )
        assert a_bridge_output == "101010101"
        cycle_after_two = first_cycle_chains[(154, parity)][2]
        cycle_for_q = (
            shift_parameter(cycle_after_two[0], -2),
            shift_parameter(cycle_after_two[1], -2),
            (
                cycle_after_two[2][0],
                cycle_after_two[2][1] - 2 * cycle_after_two[2][0],
            ),
        )
        assert poly_equal(
            template_poly(a_bridge[0]),
            raw_lift_poly(cycle_for_q[0], cycle_for_q[1], cycle_for_q[2], (0, 0)),
        ), f"a-seam-link-{parity}"
        compare_raw_generators(
            lambda m: iterate_raw(seam_start(154, q_of(m)), 9)[0][-1],
            lambda m: u_word(6 * q_of(m) - 2),
            ms,
            ms[0],
            f"a-enters-U-{parity}",
        )
    u_templates, u_output = infer_raw_chain(
        u_word, 4, [20, 21, 22], 20, "U-attractor"
    )
    assert u_output == "0101"
    compare_raw_generators(
        lambda p: iterate_raw(u_word(p), 4)[0][-1],
        lambda p: u_word(p + 1),
        [20, 21, 22],
        20,
        "U-renormalization",
    )
    certificate["a_first_seam_output"] = "101010101"
    certificate["U_renormalization"] = "T^4(U_p)=U_(p+1), output 0101"

    # The b seam manufactures the second zero-gap reservoir.
    for parity, ms in ((0, [5, 6, 7]), (1, [5, 6, 7])):
        q_of = lambda m, parity=parity: 2 * m + parity
        bridge, output = infer_raw_chain(
            lambda m: seam_start(1178, q_of(m)),
            6,
            ms,
            ms[0],
            f"b-first-seam-parity-{parity}",
        )
        assert output == "101010"
        cycle_after_two = first_cycle_chains[(1178, parity)][2]
        cycle_for_q = (
            shift_parameter(cycle_after_two[0], -2),
            shift_parameter(cycle_after_two[1], -2),
            (
                cycle_after_two[2][0],
                cycle_after_two[2][1] - 2 * cycle_after_two[2][0],
            ),
        )
        assert poly_equal(
            template_poly(bridge[0]),
            raw_lift_poly(cycle_for_q[0], cycle_for_q[1], cycle_for_q[2], (0, 0)),
        ), f"b-seam-link-{parity}"
        compare_raw_generators(
            lambda m: iterate_raw(seam_start(1178, q_of(m)), 6)[0][-1],
            lambda m: raw_lift(h_state(q_of(m))[:3], h_state(q_of(m))[3]),
            ms,
            ms[0],
            f"b-enters-H-parity-{parity}",
        )
        certificate[f"b_seam_parity_{parity}_endpoint"] = format_template(bridge[-1])

    # Every H_q state has complete abstract response (10)^infinity: even q
    # reaches s_a after six steps and odd q after ten steps.
    for parity, steps, ms in ((0, 6, [5, 6, 7]), (1, 10, [5, 6, 7])):
        q_of = lambda m, parity=parity: 2 * m + parity
        templates, output = infer_flank_chain(
            lambda m: h_state(q_of(m))[:3],
            steps,
            ms,
            ms[0],
            f"H-bridge-parity-{parity}",
        )
        assert output == "10" * (steps // 2)
        if parity == 0:
            target = lambda m: s_state(3 * q_of(m) // 2 + 6)
        else:
            target = lambda m: s_state((3 * q_of(m) + 19) // 2)
        end_u, end_v, _ = templates[-1]
        target_templates, _ = infer_flank_chain(target, 0, ms, ms[0], "target")
        assert poly_equal(template_poly(end_u), template_poly(target_templates[0][0]))
        assert poly_equal(template_poly(end_v), template_poly(target_templates[0][1]))
        certificate[f"H_parity_{parity}_bridge_output"] = output

    s_templates, s_output = infer_flank_chain(
        s_state, 6, [20, 21, 22], 20, "s-renormalization"
    )
    assert s_output == "101010"
    target_templates, _ = infer_flank_chain(
        lambda a: s_state(a + 6), 0, [20, 21, 22], 20, "s-target"
    )
    assert poly_equal(template_poly(s_templates[-1][0]), template_poly(target_templates[0][0]))
    assert poly_equal(template_poly(s_templates[-1][1]), template_poly(target_templates[0][1]))
    certificate["s_renormalization"] = "Phi^6(s_a)=s_(a+6), output 101010"

    # The terminal bridge has six more alternating outputs, then b grows where
    # a's U-family does not.
    for parity, ms in ((0, [5, 6, 7]), (1, [5, 6, 7])):
        q_of = lambda m, parity=parity: 2 * m + parity
        endpoint_steps = 1 if parity == 0 else 3
        if parity == 0:
            a1 = lambda m: 15 * q_of(m) // 2 - 12
            endpoint_gap = (0, 0)
        else:
            a1 = lambda m: (15 * q_of(m) - 29) // 2
            endpoint_gap = (0, 1)
        endpoint_templates, _ = infer_flank_chain(
            lambda m: s_state(a1(m)),
            endpoint_steps,
            ms,
            ms[0],
            f"terminal-directed-parity-{parity}",
        )
        templates, output = infer_raw_chain(
            lambda m: terminal_start(q_of(m)),
            7,
            ms,
            ms[0],
            f"terminal-parity-{parity}",
        )
        assert output == "0101011"
        directed_end = endpoint_templates[-1]
        assert poly_equal(
            template_poly(templates[0]),
            raw_lift_poly(directed_end[0], directed_end[1], directed_end[2], endpoint_gap),
        ), f"terminal-link-{parity}"
        certificate[f"terminal_parity_{parity}_output"] = output
        certificate[f"terminal_parity_{parity}_start"] = format_template(templates[0])

    # Exact small bases and theorem statement.
    small = []
    for q in range(4, 10):
        a_q = initial(154, q)
        b_q = initial(1178, q)
        horizon = 12 * q - 11
        a_trace = iterate_raw(a_q, horizon)[1]
        b_trace = iterate_raw(b_q, horizon)[1]
        divergence = next(i for i, pair in enumerate(zip(a_trace, b_trace)) if pair[0] != pair[1])
        assert divergence == 12 * q - 12
        assert a_trace[divergence] == "0" and b_trace[divergence] == "1"
        small.append({"q": q, "divergence": divergence})

    # Large values are not used as evidence, but these direct substitutions
    # catch transcription mistakes between the symbolic stages and raw orbit.
    substitutions = []
    for q in (10, 11, 12, 20, 40, 100):
        horizon = 12 * q - 11
        a_trace = iterate_raw(initial(154, q), horizon)[1]
        b_trace = iterate_raw(initial(1178, q), horizon)[1]
        divergence = next(i for i, pair in enumerate(zip(a_trace, b_trace)) if pair[0] != pair[1])
        assert divergence == 12 * q - 12
        substitutions.append({"q": q, "divergence": divergence})

    certificate.update(
        {
            "theorem": {
                "a_response": "011101111(10)^infinity",
                "b_prefix": "011101111(10)^(6q-11)11",
                "first_divergence": "12q-12",
                "chart_length": "9q",
                "latency_slope": "4/3",
            },
            "small_base_checks": small,
            "transcription_substitutions": substitutions,
            "proof_boundary": "The all-q proof is carried by symbolic affine-run identities; finite substitutions are not the induction step.",
        }
    )

    output_path = Path(__file__).with_name("symbolic-certificate.json")
    output_path.write_text(json.dumps(certificate, indent=2, sort_keys=True) + "\n")
    print(json.dumps(certificate["theorem"], indent=2, sort_keys=True))
    print(f"wrote {output_path}")


if __name__ == "__main__":
    main()