# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib"]
# ///
"""Draw exact word-metric balls on the triangular lattice."""

from __future__ import annotations

import sys
from math import sqrt
from pathlib import Path

import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib()

Point = tuple[int, int]

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

INK = "#282522"
MUTED = "#625c56"
GRID = "#ddd7cf"
OLDER = "#d9cfc3"
NEW_RING = "#b65f3f"


def distance(point: Point) -> int:
    """Word distance for the six standard A2 generators."""
    x, y = point
    return max(abs(x), abs(y), abs(x + y))


def ball(radius: int) -> set[Point]:
    """Return all triangular-lattice points within ``radius`` moves."""
    return {
        (x, y)
        for x in range(-radius, radius + 1)
        for y in range(-radius, radius + 1)
        if distance((x, y)) <= radius
    }


def cartesian(point: Point) -> tuple[float, float]:
    """Embed axial lattice coordinates in an equilateral triangular grid."""
    x, y = point
    return x + y / 2, sqrt(3) * y / 2


def draw_ball(ax: plt.Axes, radius: int) -> None:
    points = ball(radius)

    for point in points:
        x1, y1 = cartesian(point)
        for dx, dy in GENERATORS[:3]:
            neighbor = (point[0] + dx, point[1] + dy)
            if neighbor in points:
                x2, y2 = cartesian(neighbor)
                ax.plot((x1, x2), (y1, y2), color=GRID, linewidth=0.85, zorder=1)

    older = [point for point in points if distance(point) < radius]
    newest = [point for point in points if distance(point) == radius]

    if older:
        older_xy = [cartesian(point) for point in older]
        ax.scatter(
            [point[0] for point in older_xy],
            [point[1] for point in older_xy],
            s=46,
            color=OLDER,
            edgecolor="white",
            linewidth=0.55,
            zorder=2,
        )

    newest_xy = [cartesian(point) for point in newest]
    ax.scatter(
        [point[0] for point in newest_xy],
        [point[1] for point in newest_xy],
        s=56,
        color=NEW_RING,
        edgecolor="white",
        linewidth=0.65,
        zorder=3,
    )

    total = len(points)
    added = len(newest) if radius else 1
    ax.set_title(f"radius {radius}", fontsize=16, color=INK, pad=7)
    ax.text(
        0.5,
        -0.05,
        f"{total} total  ·  {added} new",
        transform=ax.transAxes,
        ha="center",
        va="top",
        fontsize=13,
        color=MUTED,
    )
    ax.set_aspect("equal")
    ax.set_xlim(-3.55, 3.55)
    ax.set_ylim(-3.15, 3.15)
    ax.axis("off")


def main() -> None:
    expected_totals = [1, 7, 19, 37]
    expected_rings = [1, 6, 12, 18]
    radii = list(range(4))
    assert [len(ball(radius)) for radius in radii] == expected_totals
    assert [
        sum(distance(point) == radius for point in ball(radius)) for radius in radii
    ] == expected_rings

    figure, axes = plt.subplots(2, 2, figsize=(7.8, 6.9), facecolor="white")
    for ax, radius in zip(axes.flat, radii, strict=True):
        draw_ball(ax, radius)

    figure.subplots_adjust(
        left=0.03,
        right=0.97,
        top=0.96,
        bottom=0.07,
        hspace=0.25,
        wspace=0.05,
    )

    validate_figure(figure)
    output = figure_output_path(__file__, "hexagonal_growth.png")
    figure.savefig(output, dpi=300, bbox_inches="tight", facecolor="white")
    print(f"Saved {output}")


if __name__ == "__main__":
    main()