"""
Dynamic Color Family Grid Extractor (Unsupervised Clustering).

Dynamically discovers color families in your image using K-Means in LAB space,
retains the ORIGINAL image colors for each group, and places them into a single
grid with separation lines.

Requirements:
    pip install numpy pillow scikit-learn --break-system-packages

Usage:
    python dynamic_color_families.py input.png --k 5 --out grid_dynamic.png
"""

import argparse
import math
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from sklearn.cluster import KMeans


def rgb_to_lab(rgb_array):
    """
    Converts RGB array [0-255] to CIELAB space for perceptual color clustering.
    CIELAB separates lightness (L) from color channels (A, B).
    """
    # Standard sRGB to XYZ conversion
    rgb = rgb_array.astype(np.float64) / 255.0
    mask = rgb > 0.04045
    rgb[mask] = np.power((rgb[mask] + 0.055) / 1.055, 2.4)
    rgb[~mask] = rgb[~mask] / 12.92

    # sRGB matrix transformation to D65 XYZ
    transform = np.array([
        [0.4124564, 0.3575761, 0.1804375],
        [0.2126729, 0.7151522, 0.0721750],
        [0.0193339, 0.1191920, 0.9503041]
    ])
    xyz = np.dot(rgb, transform.T)

    # Reference white point D65
    xyz[:, 0] /= 0.95047
    xyz[:, 1] /= 1.00000
    xyz[:, 2] /= 1.08883

    mask = xyz > 0.008856
    xyz[mask] = np.power(xyz[mask], 1.0 / 3.0)
    xyz[~mask] = (7.787 * xyz[~mask]) + (16.0 / 116.0)

    L = (116.0 * xyz[:, 1]) - 16.0
    A = 500.0 * (xyz[:, 0] - xyz[:, 1])
    B = 200.0 * (xyz[:, 1] - xyz[:, 2])

    return np.stack([L, A, B], axis=1)


def create_dynamic_family_grid(image_path, output_path, n_clusters=5, line_width=6, line_color=(180, 180, 180, 255)):
    # Open image with transparency
    img = Image.open(image_path).convert("RGBA")
    arr = np.array(img)

    h, w, _ = arr.shape
    alpha = arr[:, :, 3]
    fg_mask = alpha > 0  # Ignore background

    rgb_flat = arr[:, :, :3][fg_mask]
    if len(rgb_flat) == 0:
        raise ValueError("No non-transparent foreground pixels found in image!")

    print(f"Extracting {n_clusters} dynamic color families using LAB clustering...")
    
    # 1. Convert to CIELAB color space (perceptually uniform)
    lab_pixels = rgb_to_lab(rgb_flat)

    # 2. Dynamically cluster colors into K groups
    # We weight color channels (A, B) slightly higher than Lightness (L) so light/dark shadows
    # of the same color family stay clustered together better.
    features = lab_pixels.copy()
    features[:, 0] *= 0.6  # Scale down Lightness influence to resist shadows

    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    labels = kmeans.fit_predict(features)

    # Calculate average RGB color for each dynamically found group (for label tags)
    cluster_avg_colors = []
    for i in range(n_clusters):
        avg_rgb = rgb_flat[labels == i].mean(axis=0).astype(int)
        cluster_avg_colors.append(avg_rgb)

    # Order families by pixel count (largest family first)
    unique_labels, counts = np.unique(labels, return_counts=True)
    sorted_indices = np.argsort(-counts)
    
    # 3. Create Grid canvas with separation lines
    cols = math.ceil(math.sqrt(n_clusters))
    rows = math.ceil(n_clusters / cols)

    grid_w = cols * w + (cols + 1) * line_width
    grid_h = rows * h + (rows + 1) * line_width

    canvas = Image.new("RGBA", (grid_w, grid_h), (25, 25, 25, 255))
    draw = ImageDraw.Draw(canvas)

    for rank, cluster_idx in enumerate(sorted_indices):
        r_idx = rank // cols
        c_idx = rank % cols

        x_start = line_width + c_idx * (w + line_width)
        y_start = line_width + r_idx * (h + line_width)

        # Mask pixels belonging to this dynamic family
        family_mask_1d = labels == cluster_idx
        full_mask = np.zeros((h, w), dtype=bool)
        full_mask[fg_mask] = family_mask_1d

        # Extract tile retaining ORIGINAL image pixels
        tile_arr = np.zeros_like(arr)
        tile_arr[full_mask] = arr[full_mask]

        tile_img = Image.fromarray(tile_arr, mode="RGBA")
        canvas.paste(tile_img, (x_start, y_start), tile_img)

        # Draw separation lines around panel
        box_coords = [
            x_start - line_width // 2,
            y_start - line_width // 2,
            x_start + w + line_width // 2,
            y_start + h + line_width // 2,
        ]
        draw.rectangle(box_coords, outline=line_color, width=line_width)

        # Label tile with family ID and average RGB swatch
        avg_c = cluster_avg_colors[cluster_idx]
        label_text = f"Family #{rank + 1} ({counts[cluster_idx]} px)"
        
        # Label box background
        draw.rectangle(
            [x_start + 10, y_start + 10, x_start + 200, y_start + 40],
            fill=(0, 0, 0, 180)
        )
        # Average color swatch tag
        draw.rectangle(
            [x_start + 15, y_start + 18, x_start + 30, y_start + 33],
            fill=(avg_c[0], avg_c[1], avg_c[2], 255),
            outline=(255, 255, 255, 255)
        )
        draw.text((x_start + 38, y_start + 18), label_text, fill=(255, 255, 255, 255))

    canvas.save(output_path)
    print(f"Saved dynamic grid with {n_clusters} families to: {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dynamic color family clustering in single grid with separation lines.")
    parser.add_argument("image_path", type=str, help="Path to input image.")
    parser.add_argument("--out", type=str, default="dynamic_families_grid.png", help="Output path for grid image.")
    parser.add_argument("--k", type=int, default=5, help="Number of dynamic color families to discover (default: 5).")
    parser.add_argument("--line-width", type=int, default=6, help="Width of separation lines.")
    args = parser.parse_args()

    create_dynamic_family_grid(args.image_path, args.out, n_clusters=args.k, line_width=args.line_width)