#!/usr/bin/env python3
"""Finite evidence for the classical reversible-cover construction in our contract.

No infinite-source or proof-assistant verification is performed. Capacities
are checked against bounded integer search on the small cases; explicit lifts
are checked by all local forward/inverse equations and reachability. The
binary-depth examples enumerate all legal histories, not just sampled traces.
"""
from collections import deque
from itertools import product
import json
from pathlib import Path

from verify_partial_restoration import PartialMoore, minimize, reachable


def components(machine):
    """Small-instance transitive closure, independent of a graph library."""
    reach = {q: {q} for q in machine.states}
    for (q, _), r in machine.delta.items():
        reach[q].add(r)
    for k in machine.states:
        for q in machine.states:
            if k in reach[q]:
                reach[q].update(reach[k])
    unseen = set(machine.states)
    blocks = []
    while unseen:
        q = min(unseen)
        block = {r for r in unseen if r in reach[q] and q in reach[r]}
        blocks.append(block)
        unseen -= block
    owner = {q: i for i, block in enumerate(blocks) for q in block}
    parents = {i: set() for i in range(len(blocks))}
    for (q, _), r in machine.delta.items():
        if owner[q] != owner[r]:
            parents[owner[r]].add(owner[q])
    order = []
    while len(order) < len(blocks):
        ready = sorted(i for i in parents if i not in order and parents[i] <= set(order))
        assert ready, "component graph must be acyclic"
        order.extend(ready)
    return blocks, owner, order, reach


def solve(machine):
    assert len(machine.starts) == 1
    assert reachable(machine) == set(machine.states)
    blocks, owner, order, reach = components(machine)
    pred = {(r, a): sorted(q for q in machine.states if machine.delta.get((q, a)) == r)
            for r in machine.states for a in machine.alphabet}
    for (r, a), sources in pred.items():
        if len(sources) > 1:
            for q in sources:
                if q in reach[r]:
                    p = next(p for p in sources if p != q)
                    return {"possible": False,
                            "witness": {"other_source": p, "cyclic_source": q,
                                        "letter": a, "target": r}}
    weights = {}
    for i in order:
        c = max([1] + [sum(weights[q] for q in pred[r, a] if owner[q] != i)
                       for r in blocks[i] for a in machine.alphabet])
        weights.update({q: c for q in blocks[i]})
    assert all(sum(weights[q] for q in sources) <= weights[r]
               for (r, _), sources in pred.items())
    states = {(q, i) for q in machine.states for i in range(weights[q])}
    delta, inverse = {}, {}
    for (r, a), sources in pred.items():
        offset = 0
        for q in sources:
            for i in range(weights[q]):
                source, target = (q, i), (r, offset + i)
                assert (target, a) not in inverse
                delta[source, a] = target
                inverse[target, a] = source
            offset += weights[q]
    start = (machine.starts[0], 0)
    seen, queue = {start}, deque([start])
    while queue:
        source = queue.popleft()
        for a in machine.alphabet:
            target = delta.get((source, a))
            assert (target is not None) == ((source[0], a) in machine.delta)
            if target is not None:
                assert target[0] == machine.delta[source[0], a]
                assert inverse[target, a] == source
                if target not in seen:
                    seen.add(target)
                    queue.append(target)
    assert seen == states
    return {"possible": True, "weights": weights, "states": len(states),
            "arrows": len(delta), "delta": delta, "inverse": inverse,
            "start": start}


def integer_search(machine, bound):
    """Enumerate feasible capacities, not reversible transition tables."""
    feasible = []
    for values in product(range(1, bound + 1), repeat=len(machine.states)):
        weights = dict(zip(machine.states, values))
        if all(sum(weights[q] for q in machine.states
                   if machine.delta.get((q, a)) == r) <= weights[r]
               for r in machine.states for a in machine.alphabet):
            feasible.append(values)
    return feasible


def moore(states, alphabet, output, edges, start="s"):
    return PartialMoore(tuple(states), tuple(alphabet), (start,), output, edges)


def fork(loop=None):
    edges = {("s", "0"): "p", ("s", "1"): "q",
             ("p", "a"): "r", ("q", "a"): "r"}
    if loop:
        edges["r", loop] = "r"
    return moore("spqr", "01ab", {"s": "start", "p": 0, "q": 1, "r": 2}, edges)


def two_state_component():
    return moore("spqrt", "01abc", dict(zip("spqrt", range(5))),
                 {("s", "0"): "p", ("s", "1"): "q",
                  ("p", "a"): "r", ("q", "a"): "r",
                  ("r", "b"): "t", ("t", "c"): "r"})


def bounded_binary(n):
    states = ["e"] + [f"{d}:{b}" for d in range(1, n + 1) for b in "01"]
    output = {q: ("empty" if q == "e" else q[-1]) for q in states}
    edges = {}
    for q in states:
        d = 0 if q == "e" else int(q.split(":")[0])
        if d < n:
            edges.update({(q, b): f"{d+1}:{b}" for b in "01"})
    return moore(states, "01", output, edges, "e")


def example_certificate(name, machine, expected, search_bound=None):
    minimized, blocks = minimize(machine)
    solution = solve(minimized)
    assert solution["possible"] == (expected is not None)
    if expected is not None:
        assert solution["states"] == expected
    certificate = {"name": name, "forward_residual_states": len(minimized.states),
                   "residual_blocks": blocks,
                   **{k: v for k, v in solution.items() if k not in {"delta", "inverse", "start"}}}
    if search_bound is not None:
        feasible = integer_search(minimized, search_bound)
        if expected is None:
            assert not feasible
        else:
            optimum = tuple(solution["weights"][q] for q in minimized.states)
            assert optimum in feasible
            assert all(all(x <= y for x, y in zip(optimum, candidate)) for candidate in feasible)
        certificate["bounded_capacity_search"] = {
            "coordinate_range": [1, search_bound],
            "vectors_checked": search_bound ** len(minimized.states),
            "feasible_vectors": len(feasible),
            "scope": "corroboration within this box; not a proof outside it"}
    return certificate


def main():
    cases = [
        example_certificate("fork", fork(), 5, 3),
        example_certificate("merge_then_different_letter_loop", fork("b"), 5, 3),
        example_certificate("merge_into_same_letter_loop", fork("a"), None, 3),
        example_certificate("merge_into_two_state_component", two_state_component(), 7, 3),
        example_certificate("parity_with_legal_pop_promise",
                            moore("st", "X", {"s": 0, "t": 1},
                                  {("s", "X"): "t", ("t", "X"): "s"}), 2, 3),
        example_certificate("a_star_b_star_constant_output",
                            moore("st", "ab", {"s": 0, "t": 0},
                                  {("s", "a"): "s", ("s", "b"): "t", ("t", "b"): "t"}),
                            None, 6),
    ]
    binary = []
    for n in range(1, 9):
        machine, blocks = minimize(bounded_binary(n))
        solution = solve(machine)
        assert len(machine.states) == 2*n + 1
        assert solution["states"] == 2**(n+1) - 1
        history_states = {"": solution["start"]}
        for depth in range(n):
            for word in [w for w in history_states if len(w) == depth]:
                for bit in "01":
                    state = solution["delta"][history_states[word], bit]
                    history_states[word + bit] = state
                    assert solution["inverse"][state, bit] == history_states[word]
                    assert machine.output[state[0]] == bit
        assert len(set(history_states.values())) == len(history_states)
        binary.append({"depth": n, "forward_states": len(machine.states),
                       "restoration_states": solution["states"],
                       "all_legal_histories_checked": len(history_states),
                       "all_restoration_edges_checked": len(history_states) - 1})
    # A single failing-edge mutation must be rejected by the local inverse invariant.
    sample, _ = minimize(fork())
    solution = solve(sample)
    by_target = {}
    for (source, a), target in solution["delta"].items():
        by_target.setdefault((target[0], a), []).append((source, target))
    pair = next(edges for edges in by_target.values() if len(edges) >= 2)
    (left, target_left), (right, target_right) = pair[:2]
    assert left != right and target_left != target_right
    a = next(a for (source, a), target in solution["delta"].items()
             if source == left and target == target_left)
    assert solution["inverse"][target_left, a] != right
    output = {
        "scope": "Six finite controller cases; binary-depth family n=1..8; all lift edges and reachability; bounded capacity enumeration as specified. Not an infinite-source proof.",
        "cases": cases, "binary_depth_family": binary,
        "forced_merge_mutation_rejected": True,
    }
    target = Path(__file__).with_name("regular-history-restoration-output.json")
    target.write_text(json.dumps(output, indent=2) + "\n")
    print(json.dumps({"cases": len(cases), "binary_depths": len(binary),
                      "largest_lift": binary[-1]["restoration_states"],
                      "forced_merge_mutation_rejected": True}))


if __name__ == "__main__":
    main()