import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import io
from stringart import StringArtGenerator

def generate_string_art(image, num_nails, iterations, weight, shape):
    """Generate string art from uploaded image"""
    if image is None:
        return None
        
    try:
        # Initialize generator
        generator = StringArtGenerator()
        
        # Convert PIL image to numpy array and set it directly
        np_img = np.array(image)
        generator.image = image
        generator.data = np.flipud(np_img).transpose()
        
        # Preprocess
        generator.preprocess()
        generator.set_nails(int(num_nails))
        generator.set_iterations(int(iterations))
        generator.set_weight(int(weight))
        generator.set_shape(shape)
        generator.set_seed(42)
        
        # Generate pattern
        pattern = generator.generate()
        
        if len(pattern) == 0:
            return None
        
        # Create visualization
        lines_x = []
        lines_y = []
        for i in range(len(pattern)-1):
            lines_x.extend([pattern[i][0], pattern[i+1][0], None])
            lines_y.extend([pattern[i][1], pattern[i+1][1], None])

        # Create plot
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))
        
        # Original image
        ax1.imshow(generator.image, cmap='gray')
        ax1.set_title('Original Image (Preprocessed)')
        ax1.axis('off')
        
        # String art
        xmin, ymin = 0., 0.
        xmax = generator.data.shape[0]
        ymax = generator.data.shape[1]
        
        ax2.set_xlim([xmin, xmax])
        ax2.set_ylim([ymin, ymax])
        ax2.set_aspect('equal')
        ax2.axis('off')
        ax2.set_title(f'String Art ({len(pattern)} connections)')
        ax2.set_facecolor('white')
        
        # Plot all lines at once for better performance
        ax2.plot(lines_x, lines_y, linewidth=0.15, color='black', alpha=0.8)
        
        plt.tight_layout()
        
        # Save to buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', dpi=150, facecolor='white')
        buf.seek(0)
        plt.close()
        
        # Return image
        result_image = Image.open(buf)
        return result_image
        
    except Exception as e:
        print(f"Error: {str(e)}")
        return None

# Create Gradio interface
with gr.Blocks(
    title="String Art Generator",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px !important;
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🎨 String Art Generator
    
    Transform your images into beautiful string art! Upload an image and watch as an algorithm recreates it using virtual string connections between nails.
    
    **How it works:** The algorithm places nails around a circle/rectangle, then iteratively finds the darkest paths between nails to recreate your image.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Input")
            image_input = gr.Image(
                type="pil", 
                label="Upload Your Image",
                sources=["upload"],
                height=300
            )
            
            gr.Markdown("### ⚙️ Settings")
            with gr.Row():
                num_nails = gr.Slider(
                    minimum=50, 
                    maximum=200, 
                    value=120, 
                    step=10, 
                    label="Number of Nails",
                    info="More nails = more detail, but slower processing"
                )
                iterations = gr.Slider(
                    minimum=500, 
                    maximum=3000, 
                    value=1500, 
                    step=100, 
                    label="Iterations",
                    info="More iterations = better quality, but slower"
                )
            
            with gr.Row():
                weight = gr.Slider(
                    minimum=5, 
                    maximum=40, 
                    value=20, 
                    step=5, 
                    label="Line Weight",
                    info="Thickness of virtual string"
                )
                shape = gr.Dropdown(
                    choices=["circle", "rectangle"], 
                    value="circle", 
                    label="Nail Arrangement",
                    info="Shape for nail placement"
                )
            
            generate_btn = gr.Button(
                "🎨 Generate String Art", 
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            ### 💡 Tips:
            - High contrast images work best
            - Simple compositions are ideal
            - Processing may take 1-3 minutes
            - Start with lower settings for faster results
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Result")
            output_image = gr.Image(
                label="Generated String Art",
                height=600
            )
    
    generate_btn.click(
        fn=generate_string_art,
        inputs=[image_input, num_nails, iterations, weight, shape],
        outputs=output_image,
        show_progress=True
    )

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