import numpy as np
from PIL import Image, ImageDraw, ImageFont
from skimage.color import rgb2lab, lab2rgb
from sklearn.cluster import KMeans
import math
import sys
import os

def create_polygon_path(center_x, center_y, radius, num_sides, rotation=0):
    points = []
    for i in range(num_sides):
        angle = 2 * math.pi * i / num_sides + rotation
        x = center_x + radius * math.cos(angle)
        y = center_y + radius * math.sin(angle)
        points.append((x, y))
    return points
  #  """Creates a list of points representing the vertices of a polygon."""
  #  points = [
  #      (center_x + radius * math.cos(2 * math.pi * i / num_sides),
  #       center_y + radius * math.sin(2 * math.pi * i / num_sides))
  #      for i in range(num_sides)
  #  ]
  #  return points


def get_polygon_color(image, x, y, radius, num_sides):
    """Gets the average color within a polygonal area, handling edge cases."""
    try:
        x1 = max(0, int(x - radius))
        y1 = max(0, int(y - radius))
        x2 = min(image.width, int(x + radius))
        y2 = min(image.height, int(y + radius))

        if x1 >= x2 or y1 >= y2:
            return (0, 0, 0)

        cropped_image = image.crop((x1, y1, x2, y2))

        if cropped_image.mode != 'RGBA':
            cropped_image = cropped_image.convert('RGBA')

        cropped_array = np.array(cropped_image)

        mask = Image.new('L', (x2 - x1, y2 - y1), 0)
        draw = ImageDraw.Draw(mask)
        polygon_points = create_polygon_path(x - x1, y - y1, radius, num_sides, rotation=0)
        draw.polygon(polygon_points, fill=255)

        masked_array = np.array(cropped_image)
        masked_array[np.array(mask) == 0] = [0, 0, 0, 0]

        visible_pixels = masked_array[masked_array[:, :, 3] > 0][:, :3]
        return tuple(np.mean(visible_pixels, axis=0).astype(int)) if len(visible_pixels) > 0 else (0, 0, 0)

    except Exception as e:
        print(f"Error calculating polygon color: {e}")
        return (0, 0, 0)

def create_square_grid(w, h, size):
    radius = size / 2
    positions = []
    y = -radius
    while y < h + radius:
        x = -radius
        while x < w + radius:
            if -radius < x < w and -radius < y < h:
                positions.append((x, y))
            x += size
        y += size
    return positions

def create_hexagon_grid(w, h, size):
    radius = size / 2
    vertical_spacing = int(size * math.sqrt(3) / 2)
    positions = []
    row = 0
    y = -vertical_spacing // 2
    while y < h + vertical_spacing:
        x_offset = size // 2 if row % 2 else 0
        x = -size // 2 + x_offset
        while x < w + size:
            if -size < x < w and -size < y < h:
                positions.append((x, y))
            x += size
        y += vertical_spacing
        row += 1
    return positions

def create_triangle_grid(w, h, size):
    radius = size / 2
    height = radius * math.sqrt(3)
    positions = []
    y = -radius
    while y < h + radius:
        x = -radius
        while x < w + radius:
            if -radius < x < w and -radius < y < h:
                positions.append((x, y))
            x += size
        y += height / 2
    y = height / 4 - radius
    while y < h + radius:
        x = size / 2 - radius
        while x < w + radius:
            if -radius < x < w and -radius < y < h:
                positions.append((x, y))
            x += size
        y += height / 2
    return positions


def create_numbered_mosaic(image_path, output_path, shape="hexagon", size=30, n_colors=10, bg_gray=255):
    """Creates a numbered mosaic with the specified shape and a color key."""
    try:
        original_img = Image.open(image_path).convert('RGB')
        w, h = original_img.size

        mosaic = Image.new('RGB', (w, h), (bg_gray, bg_gray, bg_gray))
        colored = Image.new('RGB', (w, h), (bg_gray, bg_gray, bg_gray))
        draw = ImageDraw.Draw(mosaic)
        colored_draw = ImageDraw.Draw(colored)

        try:
            font = ImageFont.truetype("arial.ttf", size // 2)
            font2 = ImageFont.truetype("arial.ttf", 10)
        except IOError:
            try:
                font = ImageFont.truetype("LiberationSans-Bold.ttf", size // 2)
                font2 = ImageFont.truetype("DejaVuSerif.ttf", 10)
            except IOError:
                font = ImageFont.load_default()
                font.size = size // 2
                font2 = ImageFont.load_default()
                font2.size = 10

        if shape == "triangle":
            num_sides = 3
            positions = create_triangle_grid(w, h, size)
        elif shape == "square":
            num_sides = 4
            positions = create_square_grid(w, h, size)
        else:  # Default to hexagon
            shape = "hexagon"
            num_sides = 6
            positions = create_hexagon_grid(w, h, size)

        radius = size / 2
        overlap_percentage = 2.0  # Adjust as needed for different shapes

        colors = []
        for (x, y) in positions:
            avg_color = get_polygon_color(original_img, x + radius, y + radius, radius, num_sides)
            colors.append(avg_color)

        if len(colors) > n_colors:
            lab_colors = rgb2lab(np.array(colors) / 255.).reshape(-1, 3)
            palette = lab2rgb(KMeans(n_colors).fit(lab_colors).cluster_centers_) * 255
            palette = [tuple(c.astype(int)) for c in palette]
        else:
            palette = list(set(colors))


        color_to_number = {color: str(i + 1) for i, color in enumerate(palette)}

        for (x, y), color in zip(positions, colors):
            hex_color = min(palette, key=lambda c: sum((a - b)**2 for a, b in zip(color, c)))

            rotation = math.pi / 4 if num_sides == 4 else 0  # Conditional rotation
            overlapped_radius = radius * (1 + overlap_percentage / 100)
            colored_points = create_polygon_path(x + radius, y + radius, overlapped_radius, num_sides,rotation)
            colored_draw.polygon(colored_points, fill=hex_color, outline=hex_color)

            rotation = math.pi / 4 if num_sides == 4 else 0  # Conditional rotation
            points = create_polygon_path(x + radius, y + radius, radius, num_sides, rotation)
            draw.polygon(points, fill="white", outline=(0, 0, 0), width=1)

            number = color_to_number[hex_color]
            text_bbox = draw.textbbox((0, 0), number, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]
            draw.text(
                (x + radius - text_width // 2, y + radius - text_height // 2 - size // 20),
                number, fill=(0, 0, 0), font=font
            )


        entry_height = 30
        box_height = 20
        header_height = 10

        key_height = header_height + (entry_height * len(palette))
        key_img = Image.new('RGB', (80, key_height), (255, 255, 255))
        key_draw = ImageDraw.Draw(key_img)

        for i, color in enumerate(palette):
            y_base = header_height + (i * entry_height)

            key_draw.rectangle([5, y_base, 30, y_base + box_height], fill=color, outline=(0, 0, 0), width=1)
            key_draw.text((35, y_base + 5), f"{i + 1}", fill=(0, 0, 0), font=font2)
            key_draw.rectangle([50, y_base, 75, y_base + box_height], fill=(255, 255, 255), outline=(0, 0, 0), width=1)


        base_name, ext = os.path.splitext(output_path)
        key_img.save(f"{base_name}_key{ext}")
        mosaic.save(output_path)
        colored.save(f"{base_name}_colored{ext}")
        print(f"Success! Saved mosaic to {output_path} and color key to {base_name}_key{ext}")

    except FileNotFoundError:
        print(f"Error: Input image not found at {image_path}")
        sys.exit(1)
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python hex_mosaic.py input.jpg output.png shape [size=30] [colors=10] [bg_gray=255]")
        print("Valid shapes: hexagon, triangle, square")
        sys.exit(1)

    try:
        create_numbered_mosaic(
            sys.argv[1],
            sys.argv[2],
            sys.argv[3],
            int(sys.argv[4]) if len(sys.argv) > 4 else 30,
            int(sys.argv[5]) if len(sys.argv) > 5 else 10,
            int(sys.argv[6]) if len(sys.argv) > 6 else 255
        )
    except ValueError:
        print("Error: size, colors, and bg_gray must be integers")