import gradio as gr
import random
import matplotlib.pyplot as plt

# Generate a random hex color
def random_hex():
    return "#" + "".join(f"{random.randint(0, 255):02x}" for _ in range(3))

# Generate initial 5-color palette
def generate_palette(locked_colors):
    return [random_hex() if locked is None else locked for locked in locked_colors]

# Create color swatches
def display_palette(colors):
    fig, ax = plt.subplots(figsize=(8, 2))
    rgb_colors = [[int(h[i:i+2], 16) for i in (1, 3, 5)] for h in colors]
    ax.imshow([rgb_colors], aspect="auto")
    ax.set_xticks([])
    ax.set_yticks([])
    return fig, colors

# Initialize locked colors
initial_palette = generate_palette([None] * 5)

# Gradio UI
def update_palette(*locks):
    locks = [color if color and color != "None" else None for color in locks]
    new_palette = generate_palette(locks)
    return display_palette(new_palette) + tuple(new_palette)

with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# 🎨 AI-Powered Color Palette Generator")
    gr.Markdown("Click on a color to lock it. Generate fresh palettes for unlocked colors!")
    
    with gr.Row():
        lock_inputs = [gr.ColorPicker(label=f"Color {i+1}", value=initial_palette[i]) for i in range(5)]
    
    palette_output = gr.Plot(label="Generated Palette")
    generate_button = gr.Button("Regenerate Palette 🎨")
    
    generate_button.click(update_palette, inputs=lock_inputs, outputs=[palette_output] + lock_inputs)
    
    initial_fig, _ = display_palette(initial_palette)
    palette_output.value = initial_fig

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