import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from collections import Counter
import random

# Function to extract dominant colors
def extract_colors(image, num_colors=5):
    # Convert image from BGR to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Reshape image into a 2D array of pixels
    pixels = image.reshape(-1, 3)
    
    # Apply K-Means clustering to find dominant colors
    kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)
    kmeans.fit(pixels)
    colors = kmeans.cluster_centers_.astype(int)
    
    return colors

# Convert RGB to Hex
def rgb_to_hex(rgb):
    return "#" + "".join(f"{int(c):02x}" for c in rgb)

# Generate AI-based harmonious color palettes
def generate_ai_palette():
    return [rgb_to_hex([random.randint(0, 255) for _ in range(3)]) for _ in range(5)]

# Function to update palette with user selection
def update_palette(image, locked_colors):
    if image is not None:
        colors = extract_colors(image, 5)
        hex_colors = [rgb_to_hex(color) for color in colors]
        
        # Replace unlocked colors with extracted ones
        for i in range(5):
            if locked_colors[i] == "":
                locked_colors[i] = hex_colors[i]
    else:
        # If no image, generate AI colors
        locked_colors = [color if color != "" else rgb_to_hex([random.randint(0, 255) for _ in range(3)]) for color in locked_colors]
    
    # Create color swatches
    fig, ax = plt.subplots(figsize=(8, 2))
    rgb_colors = [[int(h[i:i+2], 16) for i in (1, 3, 5)] for h in locked_colors]
    ax.imshow([rgb_colors], aspect="auto")
    ax.set_xticks([])
    ax.set_yticks([])
    
    return fig, locked_colors

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# 🎨 AI-Powered Image Color Palette Generator")
    gr.Markdown("Upload an image to extract its dominant colors, or generate AI-based color palettes.")
    
    image_input = gr.Image(type="numpy", label="Upload an Image (Optional)")
    
    with gr.Row():
        generate_ai_button = gr.Button("Generate AI Palette 🎨")
        regenerate_button = gr.Button("Re-run Unlocked Colors 🔄")
    
    palette_output = gr.Plot(label="Color Palette")
    hex_output = gr.Textbox(label="Hex Codes", interactive=True)
    
    with gr.Row():
        locked_colors = [gr.Textbox(label=f"Color {i+1}", interactive=True, value="") for i in range(5)]
    
    # Initialize with a random palette on startup
    initial_palette = generate_ai_palette()
    initial_fig, initial_colors = update_palette(None, initial_palette)
    palette_output.value = initial_fig
    for i in range(5):
        locked_colors[i].value = initial_colors[i]
    
    generate_ai_button.click(generate_ai_palette, inputs=[], outputs=locked_colors)
    regenerate_button.click(update_palette, inputs=[image_input, locked_colors], outputs=[palette_output, locked_colors])

# Launch the app
if __name__ == "__main__":
    iface.launch()