''''
Colors the plane according to which of the point lies closest
with respect to the p-norm.
Use LEFT/RIGHT arrow keys to change p (displayed in the window title window)
p=2: Standard Euclidean norm, p=1: Manhattan norm
Click to add points.
Decrease CALC_WIDTH, CALC_HEIGHT for faster computation
'''

import pygame
import random

CALC_WIDTH, CALC_HEIGHT = 300, 200
SCALE_FACTOR = 4

n = 6 #Number of points
p = 2.0 #p of the p-norm
p_increment = 0.25

random.seed(7)

WIDTH, HEIGHT = CALC_WIDTH * SCALE_FACTOR, CALC_HEIGHT * SCALE_FACTOR

# Generate random points
points = [(random.randint(0, CALC_WIDTH), random.randint(0, CALC_HEIGHT)) for _ in range(n)]
colors = [(int(i*360/n),100,100) for i in range(n)] #initial colors

pygame.init()
voronoi_surface=pygame.Surface((CALC_WIDTH, CALC_HEIGHT))
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()

def p_norm_distance(p1, p2, p):
    """Calculate the p-norm distance between two points."""
    return (abs(p1[0] - p2[0])**p + abs(p1[1] - p2[1])**p)**(1/p)

def get_closest_point(point, points, p):
    """Find the closest point to a given point using the p-norm."""
    distances = [p_norm_distance(point, pnt, p) for pnt in points]
    closest_point_index = distances.index(min(distances))
    return closest_point_index

def update_voronoi():
    """Update Voronoi diagram with new p value."""
    global p, voronoi_surface
    for y in range(CALC_HEIGHT):
        for x in range(CALC_WIDTH):
            closest_point = get_closest_point((x, y), points, p)
            color = pygame.Color(0)
            color.hsva = colors[closest_point]
            voronoi_surface.set_at((x, y), color)
    
    
    
    scaled_surface = pygame.transform.scale(voronoi_surface, (WIDTH, HEIGHT))
    for point in points:
        pygame.draw.circle(scaled_surface, (0, 0, 0), (point[0] * SCALE_FACTOR, point[1] * SCALE_FACTOR), 5)
    pygame.display.set_caption("Voronoi Diagram p-norm for p="+str(round(p,2)))
    screen.blit(scaled_surface, (0, 0))
    pygame.display.flip()


update_voronoi()

# Main loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_LEFT:
                p = max(0.1,p-p_increment)
                update_voronoi()
            elif event.key == pygame.K_RIGHT:
                p += p_increment
                update_voronoi()
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1:  # 1 represents the left mouse button
                points.append((event.pos[0]//SCALE_FACTOR,event.pos[1]//SCALE_FACTOR))
                colors.append((random.randint(0,360), 100, 100))
                n=len(points)
                update_voronoi()
                
    clock.tick(60)

pygame.quit()