import gradio as gr
from PIL import Image, ImageDraw
import numpy as np
from sklearn.cluster import KMeans
import random
import colorsys

class LocalStyler:
    def __init__(self):
        self.style_patterns = {
            "Casual": {
                "colors": ["#4A90E2", "#58B957", "#F5A623", "#D0021B"],
                "patterns": ["solid", "stripes", "dots"]
            },
            "Formal": {
                "colors": ["#000000", "#1A237E", "#212121", "#37474F"],
                "patterns": ["solid", "pinstripe", "subtle"]
            },
            "Sporty": {
                "colors": ["#FF5722", "#2196F3", "#4CAF50", "#FFC107"],
                "patterns": ["colorblock", "geometric", "tech"]
            },
            "Bohemian": {
                "colors": ["#795548", "#FF9800", "#9C27B0", "#009688"],
                "patterns": ["floral", "paisley", "ethnic"]
            }
        }

    def extract_dominant_colors(self, image, n_colors=5):
        """Extract dominant colors from the image."""
        # Convert image to numpy array
        img_array = np.array(image)
        
        # Reshape the array to a list of pixels
        pixels = img_array.reshape(-1, 3)
        
        # Perform k-means clustering
        kmeans = KMeans(n_clusters=n_colors, random_state=42)
        kmeans.fit(pixels)
        
        # Get the RGB values of the cluster centers
        colors = kmeans.cluster_centers_.astype(int)
        
        # Convert to hex codes
        hex_colors = ['#{:02x}{:02x}{:02x}'.format(r, g, b) 
                     for r, g, b in colors]
        
        return hex_colors

    def create_pattern(self, draw, width, height, pattern_type, color):
        """Create pattern on the image."""
        if pattern_type == "stripes":
            for y in range(0, height, 20):
                draw.line([(0, y), (width, y)], fill='white', width=2)
        elif pattern_type == "dots":
            for x in range(0, width, 30):
                for y in range(0, height, 30):
                    draw.ellipse([x, y, x+10, y+10], fill='white')
        elif pattern_type == "geometric":
            for i in range(0, width, 50):
                draw.polygon([
                    (i, 0),
                    (i+50, 0),
                    (i+25, 50)
                ], fill=self.adjust_color(color))
        elif pattern_type == "colorblock":
            draw.rectangle([0, 0, width//2, height], fill=self.adjust_color(color))

    def adjust_color(self, hex_color):
        """Adjust color brightness."""
        # Convert hex to RGB
        hex_color = hex_color.lstrip('#')
        r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
        
        # Convert to HSV
        h, s, v = colorsys.rgb_to_hsv(r/255, g/255, b/255)
        
        # Adjust brightness
        v = min(1.0, v * 1.2)
        
        # Convert back to RGB
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        
        # Convert to hex
        return '#{:02x}{:02x}{:02x}'.format(int(r*255), int(g*255), int(b*255))

    def generate_styled_image(self, original_image, style, color_preference):
        """Generate a styled version of the clothing item."""
        # Get base width and height
        width, height = 512, 512
        
        # Create base image
        styled_image = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(styled_image)
        
        # Get style information
        style_info = self.style_patterns[style]
        main_color = random.choice(style_info["colors"])
        pattern = random.choice(style_info["patterns"])
        
        # Draw basic clothing shape
        clothing_points = [
            (width//4, height//4),
            (3*width//4, height//4),
            (3*width//4, 3*height//4),
            (width//4, 3*height//4)
        ]
        draw.polygon(clothing_points, fill=main_color)
        
        # Add pattern
        self.create_pattern(draw, width, height, pattern, main_color)
        
        # Add some shading for depth
        shade_color = self.adjust_color(main_color)
        draw.polygon([
            (width//4, height//4),
            (width//3, height//3),
            (width//4, height//2)
        ], fill=shade_color)
        
        return styled_image

    def analyze_colors(self, image):
        """Analyze the color composition of the image."""
        colors = self.extract_dominant_colors(image)
        return {
            "dominant_colors": colors,
            "color_scheme": "Warm" if any("FF" in c or "E2" in c for c in colors) else "Cool"
        }

def create_interface():
    """Create the Gradio interface."""
    styler = LocalStyler()
    
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # AI Fashion Stylist
        Upload your clothing item for color analysis and style suggestions!
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    type="pil",
                    label="Upload Clothing Image"
                )
                with gr.Row():
                    style_dropdown = gr.Dropdown(
                        choices=list(styler.style_patterns.keys()),
                        label="Desired Style",
                        value="Casual"
                    )
                    color_dropdown = gr.Dropdown(
                        choices=["Warm", "Cool", "Neutral", "Bold"],
                        label="Color Preference",
                        value="Warm"
                    )
                
                generate_button = gr.Button("Generate Style", variant="primary")
            
            with gr.Column(scale=1):
                color_analysis = gr.JSON(label="Color Analysis")
                color_palette = gr.HTML(label="Color Palette")
                styled_image_output = gr.Image(label="Style Suggestion")
                
        def process_and_generate(image, style, color_pref):
            if image is None:
                return (
                    {"error": "Please upload an image"},
                    "No color palette available",
                    None
                )
            
            # Analyze colors
            color_analysis = styler.analyze_colors(image)
            
            # Generate color palette HTML
            color_html = "<div style='display: flex; gap: 10px; margin: 10px 0;'>"
            for color in color_analysis["dominant_colors"]:
                color_html += f"""
                    <div style='
                        width: 50px;
                        height: 50px;
                        background-color: {color};
                        border: 1px solid #ccc;
                        border-radius: 5px;
                    '></div>
                """
            color_html += "</div>"
            
            # Generate styled image
            styled_image = styler.generate_styled_image(
                image, style, color_pref
            )
            
            return color_analysis, color_html, styled_image
        
        generate_button.click(
            fn=process_and_generate,
            inputs=[input_image, style_dropdown, color_dropdown],
            outputs=[color_analysis, color_palette, styled_image_output]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()