#!/usr/bin/env python3
"""Exact audit of the Polak--Schrijver/Gao C7^5 base gadget.

The construction deliberately does not contain the 367 codewords as data.
It regenerates the 327-word circular core, the 71-vertex extension graph, and
the unique 40-word patch selected by Gao's typed gadget constraints.
"""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter
from functools import lru_cache
from pathlib import Path


Q = 7
DIMENSION = 5
MODULUS = 382
SHIFT = (40, 123, 40, 123, 40)
OLD_COMPANION_POINT = (2, 4, 6, 3, 5)
NEW_COMPANION_POINT = (1, 5, 6, 3, 5)

PAIRS = (
    ((1, 3, 4, 4, 6), (2, 3, 5, 4, 6)),
    ((3, 4, 0, 3, 5), (2, 4, 6, 3, 5)),
    ((5, 3, 1, 3, 4), (5, 3, 2, 3, 5)),
    ((4, 4, 6, 1, 6), (5, 4, 6, 0, 6)),
    ((6, 0, 6, 4, 5), (6, 1, 6, 5, 5)),
    ((0, 3, 5, 6, 5), (6, 3, 5, 0, 5)),
    ((6, 4, 3, 4, 0), (6, 4, 2, 4, 6)),
    ((6, 4, 5, 3, 2), (6, 5, 5, 3, 1)),
)
J0 = frozenset((0, 5, 6))


Word = tuple[int, int, int, int, int]


def closed_confusable(left: Word, right: Word) -> bool:
    return all((x - y) % Q in (0, 1, Q - 1) for x, y in zip(left, right))


def independent(points: set[Word] | frozenset[Word]) -> bool:
    ordered = sorted(points)
    return all(
        not closed_confusable(left, right)
        for index, left in enumerate(ordered)
        for right in ordered[index + 1 :]
    )


def transform_t(word: Word) -> Word:
    return (
        (2 - word[1]) % Q,
        word[3],
        word[0],
        (2 - word[2]) % Q,
        word[4],
    )


def quantize(value: int) -> int:
    # floor(value / 54.5), written without floating point.
    return (2 * value) // 109


def all_words() -> list[Word]:
    return [
        tuple((number // (Q ** (DIMENSION - 1 - position))) % Q for position in range(DIMENSION))  # type: ignore[misc]
        for number in range(Q**DIMENSION)
    ]


def circular_precursor_image() -> list[Word]:
    return [
        tuple(
            quantize((t * pow(7, position, MODULUS) + SHIFT[position]) % MODULUS)
            for position in range(DIMENSION)
        )  # type: ignore[misc]
        for t in range(MODULUS)
    ]


def circular_core() -> tuple[set[Word], list[Word]]:
    image = circular_precursor_image()
    if len(set(image)) != MODULUS:
        raise AssertionError("quantized precursor is not injective")
    bad = {
        left
        for left in image
        if any(left != right and closed_confusable(left, right) for right in image)
    }
    core = set(image) - bad
    if len(core) != 327 or len(bad) != 55 or not independent(core):
        raise AssertionError("unexpected circular-core counts")
    return core, image


def extension_graph(core: set[Word]) -> tuple[list[Word], list[int]]:
    vertices = [
        word
        for word in all_words()
        if word not in core
        and all(not closed_confusable(word, base) for base in core)
    ]
    adjacency = [0] * len(vertices)
    for i, left in enumerate(vertices):
        for j in range(i + 1, len(vertices)):
            if closed_confusable(left, vertices[j]):
                adjacency[i] |= 1 << j
                adjacency[j] |= 1 << i
    if len(vertices) != 71 or sum(mask.bit_count() for mask in adjacency) // 2 != 85:
        raise AssertionError("unexpected extension graph")
    return vertices, adjacency


def typed_patch(core: set[Word], vertices: list[Word], adjacency: list[int]) -> tuple[set[Word], int, dict[str, int]]:
    index = {word: i for i, word in enumerate(vertices)}
    phase_h = {
        PAIRS[i][0] if i in J0 else PAIRS[i][1]
        for i in range(len(PAIRS))
    }
    phase_v = {
        PAIRS[i][1] if i in J0 else PAIRS[i][0]
        for i in range(len(PAIRS))
    }
    endpoints = phase_h | phase_v
    if not independent(phase_h) or not independent(phase_v):
        raise AssertionError("a transversal is not independent")

    def category(point: Word) -> str:
        hit_h = any(closed_confusable(point, endpoint) for endpoint in phase_h)
        hit_v = any(closed_confusable(point, endpoint) for endpoint in phase_v)
        if hit_h and hit_v:
            return "HV"
        if hit_h:
            return "H"
        if hit_v:
            return "V"
        return "O"

    preimages = [word for word in core if transform_t(word) == OLD_COMPANION_POINT]
    if preimages != [(6, 0, 6, 4, 5)]:
        raise AssertionError("unexpected companion replacement preimage")
    replacement_preimage = preimages[0]

    fixed_companion = {transform_t(word) for word in core}
    fixed_companion.remove(OLD_COMPANION_POINT)
    fixed_companion.add(NEW_COMPANION_POINT)
    fixed_profile = Counter(category(word) for word in fixed_companion)
    if fixed_profile != Counter({"O": 292, "H": 19, "V": 16}):
        raise AssertionError("unexpected fixed companion profile")
    quotas = {"O": 29, "H": 7, "V": 4}

    forced_words = {center for center, _ in PAIRS if center in index}
    forced_mask = sum(1 << index[word] for word in forced_words)
    forbidden_mask = 0
    for word_index, word in enumerate(vertices):
        forbidden = False

        # Each alternative may see only its designated center in the code.
        for center, alternative in PAIRS:
            if word != center and closed_confusable(alternative, word):
                forbidden = True

        transformed = transform_t(word)
        # T(code), after the one stated replacement, must remain an independent
        # auxiliary set disjoint from all endpoints and from the forbidden HV cell.
        if closed_confusable(transformed, NEW_COMPANION_POINT):
            forbidden = True
        if transformed in endpoints or category(transformed) == "HV":
            forbidden = True
        if forbidden:
            forbidden_mask |= 1 << word_index

    if forced_mask & forbidden_mask:
        raise AssertionError("typed constraints forbid a required center")
    for i in range(len(vertices)):
        if forced_mask >> i & 1:
            forbidden_mask |= adjacency[i]
    if forced_mask & forbidden_mask:
        raise AssertionError("required centers conflict")

    forced_counts = Counter(
        category(transform_t(vertices[i]))
        for i in range(len(vertices))
        if forced_mask >> i & 1
    )
    remaining_quota = tuple(
        quotas[name] - forced_counts[name]
        for name in ("O", "H", "V")
    )
    available_mask = ((1 << len(vertices)) - 1) & ~forbidden_mask & ~forced_mask
    category_index = {
        "O": 0,
        "H": 1,
        "V": 2,
    }
    category_masks = [0, 0, 0]
    for i, word in enumerate(vertices):
        name = category(transform_t(word))
        if name in category_index:
            category_masks[category_index[name]] |= 1 << i

    search_nodes = 0

    @lru_cache(maxsize=None)
    def count_solutions(mask: int, need_o: int, need_h: int, need_v: int) -> tuple[int, int]:
        """Return (solution count capped at two, one solution mask)."""
        nonlocal search_nodes
        search_nodes += 1
        needs = (need_o, need_h, need_v)
        if any(need < 0 for need in needs):
            return 0, 0
        for category_number, need in enumerate(needs):
            if (mask & category_masks[category_number]).bit_count() < need:
                return 0, 0
        if sum(needs) == 0:
            return 1, 0
        if mask.bit_count() < sum(needs):
            return 0, 0

        # Branch on a high-degree live vertex; this makes the finite uniqueness
        # proof small while leaving the logic transparent and dependency-free.
        live = [i for i in range(len(vertices)) if mask >> i & 1]
        vertex = max(live, key=lambda i: (adjacency[i] & mask).bit_count())
        bit = 1 << vertex

        count_without, witness_without = count_solutions(
            mask & ~bit, need_o, need_h, need_v
        )

        name = category(transform_t(vertices[vertex]))
        category_number = category_index[name]
        next_needs = [need_o, need_h, need_v]
        next_needs[category_number] -= 1
        count_with, witness_with = count_solutions(
            mask & ~bit & ~adjacency[vertex], *next_needs
        )
        if count_with:
            witness_with |= bit

        total = min(2, count_without + count_with)
        witness = witness_without if count_without else witness_with
        return total, witness

    solution_count, solution_mask = count_solutions(available_mask, *remaining_quota)
    if solution_count != 1:
        raise AssertionError(f"typed patch is not unique: capped count {solution_count}")
    patch_mask = forced_mask | solution_mask
    patch = {
        vertices[i]
        for i in range(len(vertices))
        if patch_mask >> i & 1
    }
    if len(patch) != 40 or not independent(patch):
        raise AssertionError("typed patch is not a 40-word independent set")

    code = core | patch
    for center, alternative in PAIRS:
        hits = {word for word in code if closed_confusable(alternative, word)}
        if hits != {center}:
            raise AssertionError("private-pair condition failed")

    companion = {transform_t(word) for word in code}
    companion.remove(OLD_COMPANION_POINT)
    companion.add(NEW_COMPANION_POINT)
    profile = Counter(category(word) for word in companion)
    if len(companion) != 367 or not independent(companion):
        raise AssertionError("companion is not a 367-word independent set")
    if profile != Counter({"O": 321, "H": 26, "V": 20}):
        raise AssertionError("typed profile failed")
    if companion & endpoints:
        raise AssertionError("companion meets a pair endpoint")

    return patch, search_nodes, dict(profile)


def minimal_block_dfa(code: set[Word]) -> dict[str, object]:
    """Build the exact right-residual quotient for code* with block length five."""
    layer_residuals: list[set[frozenset[tuple[int, ...]]]] = []
    for depth in range(DIMENSION):
        prefixes = {word[:depth] for word in code}
        residuals = {
            frozenset(word[depth:] for word in code if word[:depth] == prefix)
            for prefix in prefixes
        }
        layer_residuals.append(residuals)

    states = {
        (depth, residual): number
        for number, (depth, residual) in enumerate(
            (depth, residual)
            for depth, residuals in enumerate(layer_residuals)
            for residual in residuals
        )
    }
    root_residual = next(iter(layer_residuals[0]))
    root = states[(0, root_residual)]
    transitions: list[tuple[int, int, int]] = []
    for depth, residuals in enumerate(layer_residuals):
        for residual in residuals:
            source = states[(depth, residual)]
            letters = {suffix[0] for suffix in residual}
            for letter in letters:
                successors = frozenset(
                    suffix[1:] for suffix in residual if suffix[0] == letter
                )
                if depth == DIMENSION - 1:
                    target = root
                else:
                    target = states[(depth + 1, successors)]
                transitions.append((source, letter, target))

    incoming = Counter((letter, target) for _, letter, target in transitions)
    return {
        "partial_states": len(states),
        "complete_states_with_dead": len(states) + 1,
        "transitions": len(transitions),
        "layer_residual_counts": [len(residuals) for residuals in layer_residuals],
        "reversible": all(count == 1 for count in incoming.values()),
        "reversibility_collision_keys": sum(count > 1 for count in incoming.values()),
        "maximum_same_label_indegree": max(incoming.values()),
    }


def canonical_hash(code: set[Word]) -> str:
    payload = "\n".join("".join(map(str, word)) for word in sorted(code)) + "\n"
    return hashlib.sha256(payload.encode("ascii")).hexdigest()


def parse_reference(path: Path) -> set[Word]:
    return {
        tuple(map(int, line.split()))  # type: ignore[misc]
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--reference-r367", type=Path)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    core, precursor_image = circular_core()
    vertices, adjacency = extension_graph(core)
    patch, search_nodes, profile = typed_patch(core, vertices, adjacency)
    code = core | patch
    if len(code) != 367 or not independent(code):
        raise AssertionError("generated code is not an independent 367-set")

    # The untyped extension graph does not select Gao's patch: replacing the
    # fifth center by its private alternative gives a second 40-set.
    published_center, private_alternative = PAIRS[5]
    alternative_patch = (patch - {published_center}) | {private_alternative}
    if len(alternative_patch) != 40 or not independent(alternative_patch):
        raise AssertionError("expected alternative 40-set is not independent")

    reference_matches = None
    if args.reference_r367:
        reference = parse_reference(args.reference_r367)
        reference_matches = reference == code
        if not reference_matches:
            raise AssertionError("generated code differs from the published R367 input")

    result = {
        "schema": "fprd.c7-typed-patch-attack.v1",
        "status": "pass",
        "precursor_words": len(precursor_image),
        "precursor_distinct_words": len(set(precursor_image)),
        "discarded_conflicting_words": len(set(precursor_image) - core),
        "circular_core_words": len(core),
        "extension_vertices": len(vertices),
        "extension_edges": sum(mask.bit_count() for mask in adjacency) // 2,
        "typed_patch_words": len(patch),
        "typed_patch_unique": True,
        "typed_uniqueness_search_nodes": search_nodes,
        "untyped_patch_not_unique": True,
        "untyped_one_swap": {
            "remove": published_center,
            "add": private_alternative,
        },
        "code_words": len(code),
        "code_independent": True,
        "companion_profile": profile,
        "canonical_code_sha256": canonical_hash(code),
        "minimal_exact_block_dfa": minimal_block_dfa(code),
        "last_signature_state_lower_bound": 1 + len(code),
        "reference_matches": reference_matches,
    }
    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()