import torch
import matplotlib.pyplot as plt

def discrete_voronoi(sites, grid_size):
    # Creates a uniform grid
    x = torch.linspace(0, 1, grid_size)
    y = torch.linspace(0, 1, grid_size)
    grid_x, grid_y = torch.meshgrid(x, y)
    grid_xy = torch.stack([grid_x, grid_y], dim=-1)
    
    # Computes pairwise distances between sites and grid points
    # If the sites' length is 30 and the grid_size is 100, the `distances` is 30x100x100 shape tensor
    distances = torch.cdist(sites, grid_xy).permute(1, 0, 2)
    
    # Flips the sign to negative number so that the minimum distance is the maximum probability
    distances_negative = -distances
    
    # Converts distances to probabilities using softmax
    probabilities = torch.softmax(distances_negative, dim=0)
    
    # Assigns each grid point to nearest site using argmax
    voronoi = torch.argmax(probabilities, dim=0)
    
    return grid_x, grid_y, voronoi

torch.manual_seed(77777)

n_points = 30
sites = torch.rand(n_points, 2)
grid_sizes = [10, 50, 100, 500, 1000]

# save voronoi and grid sites
voronois = []
grid_xs = []
grid_ys = []
for grid_size in grid_sizes:
    grid_x, grid_y, voronoi = discrete_voronoi(sites, grid_size)
    voronois.append(voronoi)
    grid_xs.append(grid_x)
    grid_ys.append(grid_y)

# visualize voronoi and sites
fig = plt.figure(figsize=(25, 7))
for i, (voronoi, grid_x, grid_y) in enumerate(zip(voronois, grid_xs, grid_ys)):
    ax = fig.add_subplot(1, len(voronois), i + 1)
    ax.pcolormesh(grid_x, grid_y, voronoi, cmap=plt.cm.gist_ncar, alpha=0.6)
    ax.scatter(sites[:, 0], sites[:, 1], color="black", s=20)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis("off")
    ax.set_title(f"\ngrid size: {grid_sizes[i]}×{grid_sizes[i]} \n")

plt.tight_layout()
plt.show()