import random
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.animation import FuncAnimation
from scipy.spatial import Voronoi

random.seed(123)

def voronoi_finite_polygons_2d(vor, radius=None):
    if vor.points.shape[1] != 2:
        raise ValueError("Requires 2D input")

    new_regions = []
    new_vertices = vor.vertices.tolist()

    center = vor.points.mean(axis=0)
    if radius is None:
        radius = vor.points.ptp().max() * 2

    all_ridges = {}
    for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):
        all_ridges.setdefault(p1, []).append((p2, v1, v2))
        all_ridges.setdefault(p2, []).append((p1, v1, v2))

    for p1, region in enumerate(vor.point_region):
        vertices = vor.regions[region]

        if all(v >= 0 for v in vertices):
            new_regions.append(vertices)
            continue

        ridges = all_ridges[p1]
        new_region = [v for v in vertices if v >= 0]

        for p2, v1, v2 in ridges:
            if v2 < 0:
                v1, v2 = v2, v1
            if v1 >= 0:
                continue

            t = vor.points[p2] - vor.points[p1]
            t /= np.linalg.norm(t)
            n = np.array([-t[1], t[0]])

            midpoint = vor.points[[p1, p2]].mean(axis=0)
            direction = np.sign(np.dot(midpoint - center, n)) * n
            far_point = vor.vertices[v2] + direction * radius

            new_region.append(len(new_vertices))
            new_vertices.append(far_point.tolist())

        vs = np.asarray([new_vertices[v] for v in new_region])
        c = vs.mean(axis=0)
        angles = np.arctan2(vs[:, 1] - c[1], vs[:, 0] - c[0])
        new_region = np.array(new_region)[np.argsort(angles)]

        new_regions.append(new_region.tolist())

    return new_regions, np.asarray(new_vertices)

points = np.array(
    [
        [0.37454012, 0.95071431],
        [0.73199394, 0.59865848],
        [0.15601864, 0.15599452],
        [0.05808361, 0.86617615],
        [0.60111501, 0.70807258],
        [0.02058449, 0.96990985],
        [0.83244264, 0.21233911],
        [0.18182497, 0.18340451],
        [0.30424224, 0.52475643],
        [0.43194502, 0.29122914],
    ]
)

fig, ax = plt.subplots()
ax.set_xticks([])
ax.set_yticks([])
ax.axis("equal")
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

polygons = []
colors = np.random.rand(len(points), 3)
colors = np.hstack([colors, np.ones((colors.shape[0], 1))])

def update(frame):
    ax.clear()
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis("equal")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)

    points[0, 0] = 0.5 + 0.3 * np.cos(frame / 10.0)
    points[0, 1] = 0.5 + 0.3 * np.sin(frame / 10.0)

    vor = Voronoi(points)
    regions, vertices = voronoi_finite_polygons_2d(vor)
    polygons.clear()
    for reg in regions:
        polygon = vertices[reg]
        polygons.append(polygon)

    for i, poly in enumerate(polygons):
        colored_cell = Polygon(
            poly,
            linewidth=0.7,
            alpha=0.5,
            color=colors[i],
            edgecolor="black",
        )
        ax.add_patch(colored_cell)

    ax.plot(points[:, 0], points[:, 1], "ko")

ani = FuncAnimation(fig, update, frames=100, interval=100)
plt.show()

# Save the animation as a mp4
if True:
    ani.save("voronoi_diagram.mp4", writer="ffmpeg", fps=30)