#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = []
# ///
"""Reproduce the finite-difference and lattice-growth claims in the article."""

from __future__ import annotations

import argparse
import json
from collections import deque
from fractions import Fraction
from math import comb, factorial
from typing import Any, Callable

Point = tuple[int, int]
Point3 = tuple[int, int, int]
PointFactory = Callable[[int], set[Point]]

A2_GENERATORS: tuple[Point, ...] = (
    (1, 0),
    (0, 1),
    (-1, 1),
    (-1, 0),
    (0, -1),
    (1, -1),
)


def corner_square(n: int) -> set[Point]:
    """Integer points of n[0,1]^2."""
    return {(x, y) for x in range(n + 1) for y in range(n + 1)}


def stretched_triangle(n: int) -> set[Point]:
    """Integer points of n conv{(0,0), (2,0), (1,1)}."""
    return {(x, y) for y in range(n + 1) for x in range(y, 2 * n - y + 1)}


def manhattan_diamond(n: int) -> set[Point]:
    """Integer points at Manhattan distance at most n from the origin."""
    return {
        (x, y)
        for x in range(-n, n + 1)
        for y in range(-n, n + 1)
        if abs(x) + abs(y) <= n
    }


def a2_hexagon(n: int) -> set[Point]:
    """Integer points in n times the A2 root hexagon."""
    return {
        (x, y)
        for x in range(-n, n + 1)
        for y in range(-n, n + 1)
        if max(abs(x), abs(y), abs(x + y)) <= n
    }


def centered_rectangle(n: int) -> set[Point]:
    """Integer points of n([-2,2] x [-1,1])."""
    return {(x, y) for x in range(-2 * n, 2 * n + 1) for y in range(-n, n + 1)}


def translated_unit_interval(n: int) -> set[int]:
    """Integer points of n[1,2], an Ehrhart family that is not nested."""
    return set(range(n, 2 * n + 1))


def scott_exception_triangle(n: int) -> set[Point]:
    """Integer points of n conv{(0,0), (3,0), (0,3)}."""
    return {(x, y) for x in range(3 * n + 1) for y in range(3 * n - x + 1)}


def counts(factory: PointFactory, max_n: int) -> list[int]:
    return [len(factory(n)) for n in range(max_n + 1)]


def stage_counts(values: list[int]) -> list[int]:
    return [
        values[0],
        *(b - a for a, b in zip(values[:-1], values[1:], strict=True)),
    ]


def forward_difference_rows(values: list[int]) -> list[list[int]]:
    rows = [values]
    while len(rows[-1]) > 1:
        previous = rows[-1]
        rows.append([b - a for a, b in zip(previous[:-1], previous[1:], strict=True)])
    return rows


def stirling_second(n: int, k: int) -> int:
    """Return the Stirling number of the second kind S(n, k)."""
    if n == 0:
        return int(k == 0)
    if k == 0 or k > n:
        return 0
    previous = [1, *([0] * k)]
    for size in range(1, n + 1):
        current = [0] * (k + 1)
        for blocks in range(1, min(size, k) + 1):
            current[blocks] = blocks * previous[blocks] + previous[blocks - 1]
        previous = current
    return previous[k]


def fibonacci_values(length: int) -> list[int]:
    """Return F_0 through F_(length-1)."""
    values = [0, 1]
    while len(values) < length:
        values.append(values[-1] + values[-2])
    return values[:length]


def h_star(values: list[int], dimension: int) -> list[int]:
    """Recover h* from L(0), ..., L(d)."""
    if len(values) < dimension + 1:
        raise ValueError("h* recovery needs counts through the dimension")
    return [
        sum(
            (-1) ** (j - i) * comb(dimension + 1, j - i) * values[i]
            for i in range(j + 1)
        )
        for j in range(dimension + 1)
    ]


def a2_graph_balls(max_radius: int) -> list[set[Point]]:
    """Breadth-first balls for the six standard A2 root generators."""
    distance: dict[Point, int] = {(0, 0): 0}
    queue: deque[Point] = deque([(0, 0)])
    while queue:
        x, y = queue.popleft()
        radius = distance[(x, y)]
        if radius == max_radius:
            continue
        for dx, dy in A2_GENERATORS:
            neighbor = (x + dx, y + dy)
            if neighbor not in distance:
                distance[neighbor] = radius + 1
                queue.append(neighbor)
    return [
        {point for point, radius in distance.items() if radius <= n}
        for n in range(max_radius + 1)
    ]


def reeve_13_count(n: int) -> int:
    """Ehrhart count for the height-13 Reeve tetrahedron."""
    value = Fraction(13, 6) * n**3 + n**2 - Fraction(1, 6) * n + 1
    if value.denominator != 1:
        raise ArithmeticError(f"nonintegral count at n={n}: {value}")
    return value.numerator


def reeve_13_points(n: int) -> set[Point3]:
    """Integer points in n conv{0, e1, e2, (1,1,13)}."""
    height = 13
    return {
        (x, y, z)
        for x in range(n + 1)
        for y in range(n + 1)
        for z in range(height * n + 1)
        if z <= height * x
        and z <= height * y
        and height * x + height * y - z <= height * n
    }


def require(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def x_weight_histogram(points: set[Point]) -> dict[int, int]:
    """Return coefficients of the weighted count sum q^x."""
    histogram: dict[int, int] = {}
    for x, _ in points:
        histogram[x] = histogram.get(x, 0) + 1
    return dict(sorted(histogram.items()))


def solid_angle_unit_square(n: int) -> Fraction:
    """Solid-angle lattice sum for n[0,1]^2, for positive n."""
    if n < 1:
        raise ValueError("solid-angle square check requires a positive dilation")
    interior = (n - 1) ** 2
    open_edges = 4 * (n - 1)
    vertices = 4
    return Fraction(interior) + Fraction(open_edges, 2) + Fraction(vertices, 4)


def restricted_partitions_2_3(n: int) -> int:
    """Count nonnegative solutions to 2a + 3b = n."""
    return sum(1 for b in range(n // 3 + 1) if (n - 3 * b) % 2 == 0)


def supplementary_checks(max_n: int) -> dict[str, Any]:
    """Check translated dilates, weighted counts, and quasipolynomial examples."""
    interval_counts = [len(translated_unit_interval(n)) for n in range(max_n + 1)]
    require(
        interval_counts == list(range(1, max_n + 2)),
        "[1,2] must have the same Ehrhart counts as [0,1]",
    )
    stage_two = translated_unit_interval(2)
    stage_three = translated_unit_interval(3)
    require(
        not stage_two <= stage_three,
        "translated dilates must demonstrate a non-nested counting family",
    )
    gained = stage_three - stage_two
    lost = stage_two - stage_three
    require(
        gained == {5, 6} and lost == {2},
        "the translated-interval gained/lost decomposition changed",
    )

    square_weight = x_weight_histogram(corner_square(1))
    triangle_weight = x_weight_histogram(stretched_triangle(1))
    require(
        square_weight == {0: 2, 1: 2},
        "the square's q^x count must be 2 + 2q",
    )
    require(
        triangle_weight == {0: 1, 1: 2, 2: 1},
        "the triangle's q^x count must be 1 + 2q + q^2",
    )

    solid_angle_counts = [solid_angle_unit_square(n) for n in range(1, max_n + 1)]
    require(
        solid_angle_counts == [Fraction(n**2) for n in range(1, max_n + 1)],
        "the square's solid-angle sum must equal its dilated area",
    )

    denumerants = [restricted_partitions_2_3(n) for n in range(max_n + 1)]
    expected_denumerants = [
        n // 6 if n % 6 == 1 else n // 6 + 1 for n in range(max_n + 1)
    ]
    require(
        denumerants == expected_denumerants,
        "the 2a + 3b denumerant must follow its period-six quasipolynomial",
    )

    return {
        "translated_interval": {
            "counts": interval_counts,
            "stage_2": sorted(stage_two),
            "stage_3": sorted(stage_three),
            "gained": sorted(gained),
            "lost": sorted(lost),
        },
        "weighted_collision": {
            "square_q_to_x": square_weight,
            "triangle_q_to_x": triangle_weight,
        },
        "solid_angle_square": {
            "dilations": list(range(1, max_n + 1)),
            "counts": [int(value) for value in solid_angle_counts],
        },
        "restricted_partitions_2_3": denumerants,
    }


def finite_difference_checks(max_n: int) -> dict[str, Any]:
    """Check the finite-difference identities used in the article."""
    triangular = [comb(n + 2, 2) for n in range(max_n + 1)]
    triangular_rows = forward_difference_rows(triangular)
    require(
        all(value == 1 for value in triangular_rows[2]),
        "triangular second differences must equal one",
    )
    require(
        all(value == 0 for value in triangular_rows[3]),
        "triangular third differences must equal zero",
    )

    cubes = [n**3 for n in range(max_n + 1)]
    cube_rows = forward_difference_rows(cubes)
    cube_left_edge = [row[0] for row in cube_rows[:4]]
    stirling_edge = [factorial(k) * stirling_second(3, k) for k in range(4)]
    require(
        cube_left_edge == [0, 1, 6, 6] == stirling_edge,
        "cube left edge must equal k! S(3,k)",
    )
    require(
        all(value == 6 for value in cube_rows[3]),
        "cube third differences must equal 3!",
    )
    require(
        all(value == 0 for value in cube_rows[4]),
        "cube fourth differences must equal zero",
    )

    fibonacci = fibonacci_values(max_n + 1)
    fibonacci_rows = forward_difference_rows(fibonacci)
    fibonacci_edge = [row[0] for row in fibonacci_rows]
    require(
        fibonacci_edge
        == [(-1) ** (n + 1) * value for n, value in enumerate(fibonacci)],
        "the Fibonacci left edge must alternate Fibonacci signs",
    )

    square_pyramids = [sum(k**2 for k in range(1, n + 1)) for n in range(max_n + 2)]
    require(
        [square_pyramids[n + 1] - square_pyramids[n] for n in range(max_n + 1)]
        == [(n + 1) ** 2 for n in range(max_n + 1)],
        "square-pyramidal differences must recover the next square layer",
    )

    alternating_error = [(-1) ** n for n in range(max_n + 1)]
    error_edge = [row[0] for row in forward_difference_rows(alternating_error)]
    require(
        error_edge == [(-2) ** k for k in range(max_n + 1)],
        "each difference must double an alternating error's magnitude",
    )

    independent_noise_variance = [
        sum(comb(k, j) ** 2 for j in range(k + 1)) for k in range(max_n + 1)
    ]
    require(
        independent_noise_variance == [comb(2 * k, k) for k in range(max_n + 1)],
        "independent-noise variance must grow by the central binomial coefficient",
    )

    return {
        "triangular_counts": triangular,
        "cube_left_edge": cube_left_edge,
        "stirling_cube_edge": stirling_edge,
        "fibonacci_left_edge": fibonacci_edge,
        "square_pyramidal_counts": square_pyramids,
        "alternating_error_left_edge": error_edge,
        "independent_noise_variance_multipliers": independent_noise_variance,
    }


def build_report(max_n: int) -> dict[str, Any]:
    factories: dict[str, PointFactory] = {
        "corner square": corner_square,
        "stretched triangle": stretched_triangle,
        "Manhattan diamond": manhattan_diamond,
        "A2 hexagon": a2_hexagon,
        "centered rectangle": centered_rectangle,
        "Scott exceptional triangle": scott_exception_triangle,
    }
    shape_report: dict[str, dict[str, list[int]]] = {}
    for name, factory in factories.items():
        values = counts(factory, max_n)
        shape_report[name] = {
            "counts": values,
            "stage_counts": stage_counts(values),
            "h_star": h_star(values, 2),
        }

    require(
        shape_report["corner square"]["counts"]
        == shape_report["stretched triangle"]["counts"],
        "the square and stretched triangle must be Ehrhart-equivalent",
    )

    graph_balls = a2_graph_balls(max_n)
    dilation_balls = [a2_hexagon(n) for n in range(max_n + 1)]
    require(
        graph_balls == dilation_balls,
        "A2 word-metric balls must equal the hexagon's lattice dilates",
    )
    a2_counts = [len(ball) for ball in graph_balls]
    require(
        a2_counts == [1 + 3 * n * (n + 1) for n in range(max_n + 1)],
        "A2 balls must have the centered-hexagonal counts",
    )

    reeve_counts = [reeve_13_count(n) for n in range(max_n + 1)]
    reeve_enumerated = [len(reeve_13_points(n)) for n in range(max_n + 1)]
    require(
        reeve_enumerated == reeve_counts,
        "direct Reeve-tetrahedron enumeration must match its Ehrhart polynomial",
    )
    reeve_rows = forward_difference_rows(reeve_counts)
    require(
        all(entry >= 0 for row in reeve_rows for entry in row),
        "the Reeve example's forward differences must remain nonnegative",
    )
    require(
        h_star(reeve_counts, 3) == [1, 0, 12, 0],
        "the Reeve example must have h* = 1 + 12z^2",
    )

    return {
        "n": list(range(max_n + 1)),
        "finite_difference_identities": finite_difference_checks(max_n),
        "supplementary": supplementary_checks(max_n),
        "shapes": shape_report,
        "a2_word_metric": {
            "generators": [list(point) for point in A2_GENERATORS],
            "matches_hexagon_dilates": True,
            "counts": a2_counts,
            "stage_counts": stage_counts(a2_counts),
            "coordinator_polynomial": [1, 4, 1],
        },
        "negative_coefficient_example": {
            "polytope": "Reeve tetrahedron of height 13",
            "ehrhart_polynomial": "13/6 n^3 + n^2 - 1/6 n + 1",
            "counts": reeve_counts,
            "enumerated_counts": reeve_enumerated,
            "difference_rows": reeve_rows[:4],
            "h_star": h_star(reeve_counts, 3),
        },
    }


def print_markdown(report: dict[str, Any]) -> None:
    ns = report["n"]
    identities = report["finite_difference_identities"]
    print("Finite-difference identities:")
    print(f"- cube left edge: {identities['cube_left_edge']}")
    print(f"- Fibonacci left edge: {identities['fibonacci_left_edge']}")
    print(f"- alternating-error left edge: {identities['alternating_error_left_edge']}")
    print(
        "- independent-noise variance multipliers: "
        f"{identities['independent_noise_variance_multipliers']}"
    )
    print()
    supplementary = report["supplementary"]
    interval = supplementary["translated_interval"]
    print("Supplementary checks:")
    print(
        "- translated [1,2], stages 2 -> 3: "
        f"gained {interval['gained']}, lost {interval['lost']}"
    )
    weighted = supplementary["weighted_collision"]
    print(
        "- q^x weights, square vs triangle: "
        f"{weighted['square_q_to_x']} vs {weighted['triangle_q_to_x']}"
    )
    print(
        f"- solid-angle square counts: {supplementary['solid_angle_square']['counts']}"
    )
    print(f"- solutions to 2a + 3b = n: {supplementary['restricted_partitions_2_3']}")
    print()
    print("| construction | counts L(0..n) | stage counts | h* |")
    print("| --- | --- | --- | --- |")
    for name, data in report["shapes"].items():
        print(
            f"| {name} | {data['counts']} | {data['stage_counts']} | {data['h_star']} |"
        )
    print()
    print(
        "A2 BFS equals hexagon dilation through radius "
        f"{ns[-1]}: {report['a2_word_metric']['matches_hexagon_dilates']}"
    )
    print(
        "A2 coordinator polynomial: "
        f"{report['a2_word_metric']['coordinator_polynomial']}"
    )
    negative = report["negative_coefficient_example"]
    print()
    print(f"Negative-coefficient example: {negative['ehrhart_polynomial']}")
    print(f"Counts: {negative['counts']}")
    print(f"Forward-difference rows: {negative['difference_rows']}")
    print(f"h*: {negative['h_star']}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--max-n", type=int, default=6)
    parser.add_argument(
        "--json", action="store_true", help="emit machine-readable JSON"
    )
    args = parser.parse_args()
    if args.max_n < 3:
        parser.error("--max-n must be at least 3")

    report = build_report(args.max_n)
    if args.json:
        print(json.dumps(report, indent=2))
    else:
        print_markdown(report)


if __name__ == "__main__":
    main()