#!/usr/bin/env python3
"""Symbolic closure audit for Gao's typed C7 gadget product.

The program never expands a product code.  It regenerates the table-free
five-dimensional base gadget, constructs Gao's seven-rectangle product as a
hash-consed straight-line regular expression, and lowers reachable exact
Brzozowski residuals until either a slice closes or a declared state budget is
exhausted.
"""

from __future__ import annotations

import argparse
import json
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import c7_base_compiler as base


ALPHABET = tuple(range(7))
TYPE_NAMES = ("I", "B", "R", "Q", "PH", "PV", "X", "O", "H", "V")


@dataclass(frozen=True)
class Node:
    kind: str
    data: object


class Factory:
    def __init__(self) -> None:
        self.nodes: list[Node] = []
        self.index: dict[Node, int] = {}
        self.derivative_cache: dict[tuple[int, int], int] = {}
        self.zero = self._intern(Node("zero", None))
        self.eps = self._intern(Node("eps", None))

    def _intern(self, node: Node) -> int:
        old = self.index.get(node)
        if old is not None:
            return old
        answer = len(self.nodes)
        self.nodes.append(node)
        self.index[node] = answer
        return answer

    def leaf(self, words: Iterable[tuple[int, ...]]) -> int:
        frozen = frozenset(words)
        if not frozen:
            return self.zero
        lengths = {len(word) for word in frozen}
        if lengths == {0}:
            return self.eps
        if len(lengths) != 1 or 0 in lengths:
            raise AssertionError("leaf language is not nonempty and fixed-length")
        return self._intern(Node("leaf", frozen))

    def union(self, *children: int) -> int:
        flat: set[int] = set()
        for child in children:
            if child == self.zero:
                continue
            node = self.nodes[child]
            if node.kind == "union":
                flat.update(node.data)  # type: ignore[arg-type]
            else:
                flat.add(child)
        if not flat:
            return self.zero
        if len(flat) == 1:
            return next(iter(flat))
        return self._intern(Node("union", tuple(sorted(flat))))

    def concat(self, *children: int) -> int:
        flat: list[int] = []
        for child in children:
            if child == self.zero:
                return self.zero
            if child == self.eps:
                continue
            node = self.nodes[child]
            if node.kind == "concat":
                flat.extend(node.data)  # type: ignore[arg-type]
            else:
                flat.append(child)
        if not flat:
            return self.eps
        if len(flat) == 1:
            return flat[0]
        return self._intern(Node("concat", tuple(flat)))

    def derivative(self, expression: int, letter: int) -> int:
        key = (expression, letter)
        old = self.derivative_cache.get(key)
        if old is not None:
            return old
        node = self.nodes[expression]
        if node.kind in ("zero", "eps"):
            answer = self.zero
        elif node.kind == "leaf":
            answer = self.leaf(
                word[1:]
                for word in node.data  # type: ignore[union-attr]
                if word[0] == letter
            )
        elif node.kind == "union":
            answer = self.union(*(self.derivative(child, letter) for child in node.data))  # type: ignore[union-attr]
        elif node.kind == "concat":
            first, *rest = node.data  # type: ignore[misc]
            answer = self.concat(self.derivative(first, letter), *rest)
        else:
            raise AssertionError(node.kind)
        self.derivative_cache[key] = answer
        return answer


def base_package(factory: Factory) -> tuple[dict[str, int], dict[str, int]]:
    core, _ = base.circular_core()
    vertices, adjacency = base.extension_graph(core)
    patch, _, _ = base.typed_patch(core, vertices, adjacency)
    code = core | patch
    centers = {center for center, _ in base.PAIRS}
    alternatives = {alternative for _, alternative in base.PAIRS}
    ph = {
        base.PAIRS[index][0] if index in base.J0 else base.PAIRS[index][1]
        for index in range(len(base.PAIRS))
    }
    pv = {
        base.PAIRS[index][1] if index in base.J0 else base.PAIRS[index][0]
        for index in range(len(base.PAIRS))
    }
    companion = {base.transform_t(word) for word in code}
    companion.remove(base.OLD_COMPANION_POINT)
    companion.add(base.NEW_COMPANION_POINT)

    def category(point: base.Word) -> str:
        hit_h = any(base.closed_confusable(point, endpoint) for endpoint in ph)
        hit_v = any(base.closed_confusable(point, endpoint) for endpoint in pv)
        if hit_h and hit_v:
            raise AssertionError("companion point sees both transversals")
        if hit_h:
            return "H"
        if hit_v:
            return "V"
        return "O"

    sets = {
        "I": code,
        "B": code - centers,
        "R": centers,
        "Q": alternatives,
        "PH": ph,
        "PV": pv,
        "X": companion,
        "O": {word for word in companion if category(word) == "O"},
        "H": {word for word in companion if category(word) == "H"},
        "V": {word for word in companion if category(word) == "V"},
    }
    package = {name: factory.leaf(sets[name]) for name in TYPE_NAMES}
    sizes = {name: len(sets[name]) for name in TYPE_NAMES}
    return package, sizes


def product_package(factory: Factory, left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
    c = factory.concat
    u = factory.union
    answer = {
        "I": u(
            c(left["B"], right["B"]),
            c(left["R"], right["O"]),
            c(left["PH"], right["H"]),
            c(left["PV"], right["V"]),
            c(left["O"], right["R"]),
            c(left["H"], right["PV"]),
            c(left["V"], right["PH"]),
        ),
        "B": u(
            c(left["B"], right["B"]),
            c(left["PH"], right["H"]),
            c(left["PV"], right["V"]),
            c(left["H"], right["PV"]),
            c(left["V"], right["PH"]),
        ),
        "R": u(c(left["R"], right["O"]), c(left["O"], right["R"])),
        "Q": u(c(left["Q"], right["O"]), c(left["O"], right["Q"])),
        "PH": u(c(left["PH"], right["O"]), c(left["O"], right["PV"])),
        "PV": u(c(left["PV"], right["O"]), c(left["O"], right["PH"])),
        "X": c(left["X"], right["X"]),
        "O": u(
            c(left["O"], right["O"]),
            c(left["H"], right["H"]),
            c(left["H"], right["V"]),
            c(left["V"], right["H"]),
            c(left["V"], right["V"]),
        ),
        "H": u(c(left["H"], right["O"]), c(left["O"], right["V"])),
        "V": u(c(left["V"], right["O"]), c(left["O"], right["H"])),
    }
    return answer


def product_sizes(left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
    l, r = left, right
    answer = {
        "I": l["B"] * r["B"] + l["R"] * r["O"] + l["PH"] * r["H"]
        + l["PV"] * r["V"] + l["O"] * r["R"] + l["H"] * r["PV"] + l["V"] * r["PH"],
        "R": l["R"] * r["O"] + l["O"] * r["R"],
        "Q": l["Q"] * r["O"] + l["O"] * r["Q"],
        "PH": l["PH"] * r["O"] + l["O"] * r["PV"],
        "PV": l["PV"] * r["O"] + l["O"] * r["PH"],
        "X": l["X"] * r["X"],
        "O": l["O"] * r["O"] + (l["H"] + l["V"]) * (r["H"] + r["V"]),
        "H": l["H"] * r["O"] + l["O"] * r["V"],
        "V": l["V"] * r["O"] + l["O"] * r["H"],
    }
    answer["B"] = answer["I"] - answer["R"]
    return answer


def reachable_residuals(factory: Factory, start: int, length: int, budget: int) -> dict[str, object]:
    current = {start}
    counts = [1]
    transitions = 0
    exhausted = False
    for depth in range(length):
        following: set[int] = set()
        for state in current:
            for letter in ALPHABET:
                target = factory.derivative(state, letter)
                if target != factory.zero:
                    following.add(target)
                    transitions += 1
        counts.append(len(following))
        current = following
        if len(current) > budget:
            exhausted = True
            break
    return {
        "layer_counts": counts,
        "useful_states_seen": sum(counts),
        "transitions_seen": transitions,
        "closed": not exhausted and len(counts) == length + 1,
        "budget": budget,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--state-budget", type=int, default=500_000)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    factory = Factory()
    packages: dict[int, dict[str, int]] = {}
    sizes: dict[int, dict[str, int]] = {}
    packages[1], sizes[1] = base_package(factory)
    split = {2: (1, 1), 3: (1, 2), 5: (2, 3), 10: (5, 5), 20: (10, 10), 40: (20, 20)}
    node_audit = []
    for blocks in (1, 2, 3, 5, 10, 20, 40):
        before = len(factory.nodes)
        if blocks != 1:
            left, right = split[blocks]
            packages[blocks] = product_package(factory, packages[left], packages[right])
            sizes[blocks] = product_sizes(sizes[left], sizes[right])
        after = len(factory.nodes)
        node_audit.append({
            "blocks": blocks,
            "dimension": 5 * blocks,
            "new_expression_nodes": after - before,
            "total_expression_nodes": after,
            "code_words": sizes[blocks]["I"],
            "typed_sizes": sizes[blocks],
        })

    syntax_dag_nodes = len(factory.nodes)
    residual_audit = {
        str(blocks): reachable_residuals(
            factory, packages[blocks]["I"], 5 * blocks, args.state_budget
        )
        for blocks in (1, 2, 3, 5, 10, 20, 40)
    }
    result = {
        "schema": "fprd.c7-typed-dag-closure.v1",
        "status": "pass",
        "split_tree": split,
        "typed_languages": TYPE_NAMES,
        "syntax_dag_nodes": syntax_dag_nodes,
        "factory_nodes_after_residual_lowering": len(factory.nodes),
        "node_audit": node_audit,
        "residual_audit": residual_audit,
        "gao_m40_matches": sizes[40]["I"]
        == 4085567963379990282748597333793399773399910167372462112610203714626938509773125898252337545387866277249,
        "conclusion": "Gao composition has an exact constant-template syntax-DAG compiler; residual DFA size is a separate lowering problem.",
    }
    rendered = json.dumps(result, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    print(rendered, end="")


if __name__ == "__main__":
    main()