import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button

# Initialize global storage for centers and colors
centers = np.empty((0, 2), dtype=int)  # Shape: (n_centers, 2)
colors = np.empty((0, 3), dtype=int)   # Shape: (n_centers, 3)

def hex_to_rgb(hex_color):
    """Convert a hexadecimal color string to an RGB tuple."""
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

def get_voronoi_image(centers, width=1000, height=1000, colors=None):
    """
    Generate a Voronoi-like image based on the provided centers and colors.

    Parameters:
        centers (np.ndarray): Array of center coordinates, shape (n_centers, 2).
        width (int): Width of the image in pixels.
        height (int): Height of the image in pixels.
        colors (np.ndarray): Array of RGB colors for each center, shape (n_centers, 3).

    Returns:
        np.ndarray: RGB image array, shape (height, width, 3).
    """
    if len(centers) == 0:
        return np.ones((height, width, 3), dtype=np.uint8) * 255  # White image

    # Create coordinate grids
    x = np.arange(width)
    y = np.arange(height)
    xx, yy = np.meshgrid(x, y)

    # Calculate Euclidean distances from each pixel to each center
    # Shape of distances: (height, width, n_centers)
    distances = np.sqrt((xx[..., np.newaxis] - centers[:, 0])**2 +
                        (yy[..., np.newaxis] - centers[:, 1])**2)

    # Find the minimum distance for each pixel
    min_distances = np.min(distances, axis=2)

    # Identify which centers are at the minimum distance for each pixel
    closest_centers = (distances == min_distances[..., np.newaxis])

    # Count how many centers are equally close to each pixel
    num_closest = np.sum(closest_centers, axis=2)

    # Initialize the image to white
    data = np.ones((height, width, 3), dtype=np.uint8) * 255

    # Assign colors to pixels where only one center is closest
    for i, color in enumerate(colors):
        mask = (closest_centers[:, :, i]) & (num_closest == 1)
        data[mask] = color

    # Assign edge color (e.g., red) where multiple centers are equally close
    edge_mask = (num_closest > 1)
    data[edge_mask] = [255, 0, 0]  # Red color for edges

    return data

def update_image():
    """
    Update the Voronoi image and the scatter plot of centers.
    """
    img_data = get_voronoi_image(centers, width=1000, height=1000, colors=colors)
    im.set_data(img_data)
    center_scat.set_offsets(centers)
    plt.draw()

def add_center_via_button(event):
    """
    Add a new random center with a random color and update the image.
    """
    global centers, colors
    # Generate a new random center within the image bounds
    new_center = np.random.randint(0, 1000, size=2)
    centers = np.vstack([centers, new_center])

    # Generate a random color
    new_color = np.random.randint(0, 256, size=3)
    colors = np.vstack([colors, new_color])

    update_image()
    update_buttons()

def add_center_via_click(event):
    """
    Add a new center at the clicked location with a random color.
    """
    global centers, colors
    # Ensure the click is within the axes
    if event.inaxes != ax:
        return

    # Get the x and y coordinates from the click
    x, y = event.xdata, event.ydata
    if x is None or y is None:
        return

    # Convert to integer pixel positions
    x_int, y_int = int(x), int(y)

    # Ensure the coordinates are within the image bounds
    if 0 <= x_int < 1000 and 0 <= y_int < 1000:
        new_center = np.array([x_int, y_int])
        centers = np.vstack([centers, new_center])

        # Generate a random color
        new_color = np.random.randint(0, 256, size=3)
        colors = np.vstack([colors, new_color])

        update_image()
        update_buttons()

def remove_center(event):
    """
    Remove the most recently added center and update the image.
    """
    global centers, colors
    if len(centers) == 0:
        print("No centers to remove.")
        return
    # Remove the last center and its color
    centers = centers[:-1]
    colors = colors[:-1]

    update_image()
    update_buttons()

def update_buttons():
    """
    Enable or disable the 'Back' button based on the number of centers.
    """
    if len(centers) == 0:
        back_button.ax.set_visible(False)
    else:
        back_button.ax.set_visible(True)
    plt.draw()

# Setup the figure and axes
fig, ax = plt.subplots(figsize=(8, 8))
plt.subplots_adjust(bottom=0.2)  # Adjust to make space for buttons

# Initial image (white)
initial_data = get_voronoi_image(centers, width=1000, height=1000, colors=colors)
im = ax.imshow(initial_data, origin='upper')

# Set axes limits to match image coordinates
ax.set_xlim(0, 1000)
ax.set_ylim(1000, 0)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title("Interactive Voronoi Diagram\nAdd Centers: Click or Use 'Forward' Button")

# Scatter plot for centers (initially empty)
center_scat = ax.scatter([], [], s=100, edgecolors='white', linewidths=1.5)

# Define button axes positions [left, bottom, width, height]
button_width = 0.15
button_height = 0.075
button_spacing = 0.05

# Forward Button (Add Random Center)
ax_forward = plt.axes([0.35, 0.05, button_width, button_height])
forward_button = Button(ax_forward, 'Forward')
forward_button.on_clicked(add_center_via_button)

# Back Button (Remove Last Center)
ax_back = plt.axes([0.35 + button_width + button_spacing, 0.05, button_width, button_height])
back_button = Button(ax_back, 'Back')
back_button.on_clicked(remove_center)

# Initially hide the 'Back' button since there are no centers
back_button.ax.set_visible(False)

# Connect the click event to the handler
cid = fig.canvas.mpl_connect('button_press_event', add_center_via_click)

# Display the plot
plt.show()