import gradio as gr
import numpy as np
import cv2
from PIL import Image, ImageDraw
import io
import base64
import math
import time
from typing import Tuple, List, Optional

class StringArtGenerator:
    def __init__(self, num_nails=400, canvas_size=800):
        self.num_nails = num_nails
        self.canvas_size = canvas_size
        self.nail_positions = []
        self.connections = []
        
    def generate_circle_nails(self):
        """Generate nail positions around a circle"""
        center = self.canvas_size // 2
        radius = center - 50
        angles = np.linspace(0, 2 * np.pi, self.num_nails, endpoint=False)
        
        self.nail_positions = []
        for angle in angles:
            x = int(center + radius * np.cos(angle))
            y = int(center + radius * np.sin(angle))
            self.nail_positions.append((x, y))
            
    def generate_heart_nails(self):
        """Generate nail positions around a heart shape"""
        center_x, center_y = self.canvas_size // 2, self.canvas_size // 2
        scale = 80
        
        self.nail_positions = []
        t_values = np.linspace(0, 2 * np.pi, self.num_nails)
        
        for t in t_values:
            # Heart parametric equations
            x = 16 * np.sin(t)**3
            y = 13 * np.cos(t) - 5 * np.cos(2*t) - 2 * np.cos(3*t) - np.cos(4*t)
            
            # Scale and center
            x = int(center_x + scale * x)
            y = int(center_y - scale * y)  # Negative to flip vertically
            
            self.nail_positions.append((x, y))
    
    def preprocess_image(self, image, threshold=128, blur_kernel=5):
        """Convert image to binary format suitable for string art"""
        # Convert PIL Image to numpy array if needed
        if isinstance(image, Image.Image):
            image = np.array(image)
            
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
            
        # Apply Gaussian blur
        if blur_kernel > 1:
            gray = cv2.GaussianBlur(gray, (blur_kernel, blur_kernel), 0)
            
        # Apply threshold
        _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY_INV)
        
        # Resize to canvas size
        binary = cv2.resize(binary, (self.canvas_size, self.canvas_size))
        
        return binary
    
    def calculate_line_score(self, binary_image, nail1_idx, nail2_idx):
        """Calculate how much darkness a line covers"""
        x1, y1 = self.nail_positions[nail1_idx]
        x2, y2 = self.nail_positions[nail2_idx]
        
        # Get line pixels using Bresenham's algorithm
        line_pixels = self.get_line_pixels(x1, y1, x2, y2)
        
        # Calculate score based on darkness covered
        score = 0
        for x, y in line_pixels:
            if 0 <= x < self.canvas_size and 0 <= y < self.canvas_size:
                score += binary_image[y, x]
                
        return score / 255.0  # Normalize
    
    def get_line_pixels(self, x1, y1, x2, y2):
        """Get all pixels on a line using Bresenham's algorithm"""
        pixels = []
        
        dx = abs(x2 - x1)
        dy = abs(y2 - y1)
        sx = 1 if x1 < x2 else -1
        sy = 1 if y1 < y2 else -1
        err = dx - dy
        
        x, y = x1, y1
        
        while True:
            pixels.append((x, y))
            
            if x == x2 and y == y2:
                break
                
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x += sx
            if e2 < dx:
                err += dx
                y += sy
                
        return pixels
    
    def generate_string_art(self, binary_image, max_lines=3000):
        """Generate string art using greedy algorithm"""
        self.connections = []
        used_image = np.zeros_like(binary_image, dtype=np.float32)
        
        # Start from a random nail
        current_nail = 0
        
        for line_num in range(max_lines):
            best_score = -1
            best_nail = -1
            
            # Find the best next nail
            for next_nail in range(self.num_nails):
                if next_nail == current_nail:
                    continue
                    
                # Calculate score for this line
                score = self.calculate_line_score(binary_image - used_image, current_nail, next_nail)
                
                if score > best_score:
                    best_score = score
                    best_nail = next_nail
            
            # If no good line found, break
            if best_score <= 0 or best_nail == -1:
                break
                
            # Add the line
            self.connections.append((current_nail, best_nail))
            
            # Update used image (darken the line area)
            line_pixels = self.get_line_pixels(
                self.nail_positions[current_nail][0], self.nail_positions[current_nail][1],
                self.nail_positions[best_nail][0], self.nail_positions[best_nail][1]
            )
            
            for x, y in line_pixels:
                if 0 <= x < self.canvas_size and 0 <= y < self.canvas_size:
                    used_image[y, x] = min(255, used_image[y, x] + 50)
            
            current_nail = best_nail
            
        return len(self.connections)
    
    def render_string_art(self, thread_color=(0, 0, 0), line_opacity=0.8):
        """Render the string art as an image"""
        # Create white canvas
        canvas = np.ones((self.canvas_size, self.canvas_size, 3), dtype=np.uint8) * 255
        
        # Convert thread color to RGB
        thread_rgb = thread_color
        
        # Draw lines
        for nail1_idx, nail2_idx in self.connections:
            x1, y1 = self.nail_positions[nail1_idx]
            x2, y2 = self.nail_positions[nail2_idx]
            
            # Create line with opacity
            overlay = canvas.copy()
            cv2.line(overlay, (x1, y1), (x2, y2), thread_rgb, 1)
            canvas = cv2.addWeighted(canvas, 1 - line_opacity, overlay, line_opacity, 0)
        
        # Draw nails
        for x, y in self.nail_positions:
            cv2.circle(canvas, (x, y), 2, (100, 100, 100), -1)
            
        return canvas
    
    def generate_instructions(self):
        """Generate step-by-step threading instructions"""
        instructions = []
        for i, (nail1, nail2) in enumerate(self.connections):
            instructions.append(f"Step {i+1}: Nail {nail1} → Nail {nail2}")
        return "\n".join(instructions)

def create_predefined_shape(shape_type, size=800):
    """Create predefined shapes"""
    canvas = np.zeros((size, size), dtype=np.uint8)
    center = size // 2
    
    if shape_type == "Circle":
        cv2.circle(canvas, (center, center), center - 100, 255, -1)
    elif shape_type == "Heart":
        # Create heart shape
        points = []
        scale = 100
        for t in np.linspace(0, 2 * np.pi, 1000):
            x = 16 * np.sin(t)**3
            y = 13 * np.cos(t) - 5 * np.cos(2*t) - 2 * np.cos(3*t) - np.cos(4*t)
            points.append([int(center + scale * x), int(center - scale * y)])
        
        points = np.array(points, dtype=np.int32)
        cv2.fillPoly(canvas, [points], 255)
    
    return Image.fromarray(canvas)

def hex_to_rgb(hex_color):
    """Convert hex color to RGB tuple"""
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

def generate_string_art_interface(
    input_image: Optional[Image.Image],
    shape_type: str,
    use_predefined: bool,
    num_nails: int,
    max_lines: int,
    threshold: int,
    blur_kernel: int,
    thread_color: str,
    line_opacity: float
):
    """Main function for Gradio interface"""
    
    try:
        # Determine input image
        if use_predefined:
            if shape_type in ["Circle", "Heart"]:
                working_image = create_predefined_shape(shape_type)
            else:
                return None, None, "Please select a valid predefined shape.", ""
        else:
            if input_image is None:
                return None, None, "Please upload an image or use a predefined shape.", ""
            working_image = input_image
        
        # Initialize generator
        generator = StringArtGenerator(num_nails=num_nails, canvas_size=800)
        
        # Generate nails based on shape
        if use_predefined and shape_type == "Heart":
            generator.generate_heart_nails()
        else:
            generator.generate_circle_nails()
        
        # Preprocess image
        processed_image = generator.preprocess_image(working_image, threshold, blur_kernel)
        processed_pil = Image.fromarray(processed_image)
        
        # Generate string art
        start_time = time.time()
        total_lines = generator.generate_string_art(processed_image, max_lines=max_lines)
        end_time = time.time()
        
        # Render result
        thread_rgb = hex_to_rgb(thread_color)
        result_image = generator.render_string_art(thread_rgb, line_opacity)
        result_pil = Image.fromarray(result_image)
        
        # Generate instructions
        instructions = generator.generate_instructions()
        
        # Create info text
        info_text = f"""
🎯 **Generation Complete!**
- **Total Lines**: {total_lines}
- **Processing Time**: {end_time - start_time:.2f} seconds
- **Nails Used**: {num_nails}
- **Thread Color**: {thread_color}

📋 **Instructions Preview** (showing first 10 steps):
{chr(10).join(instructions.split(chr(10))[:10])}
...

💡 **Tips for Physical Crafting**:
1. Print the full instructions from the downloadable text
2. Mark nail positions evenly around your board perimeter
3. Follow the step-by-step nail connections
4. Maintain consistent string tension
        """
        
        return result_pil, processed_pil, info_text, instructions
        
    except Exception as e:
        return None, None, f"Error: {str(e)}", ""

# Create Gradio interface
def create_gradio_app():
    with gr.Blocks(
        title="🎨 String Art Generator",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .main-header {
            text-align: center;
            margin-bottom: 2rem;
        }
        .info-box {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 1rem;
            border-radius: 10px;
            margin: 1rem 0;
        }
        """
    ) as app:
        
        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>🎨 String Art Generator</h1>
            <p style="font-size: 1.2em; color: #666;">Transform images into beautiful string art patterns for physical crafting</p>
        </div>
        """)
        
        with gr.Row():
            # Left Column - Controls
            with gr.Column(scale=1):
                gr.HTML("<h3>🖼️ Image Input</h3>")
                
                use_predefined = gr.Checkbox(
                    label="Use Predefined Shape",
                    value=True
                )
                gr.HTML("<small>Check to use predefined shapes, uncheck to upload your own image</small>")
                
                with gr.Group(visible=True) as predefined_group:
                    shape_type = gr.Radio(
                        choices=["Circle", "Heart"],
                        value="Circle",
                        label="Select Shape"
                    )
                
                with gr.Group(visible=False) as upload_group:
                    input_image = gr.Image(
                        label="Upload Image",
                        type="pil"
                    )
                    gr.HTML("<small>Upload PNG, JPG, or other image formats</small>")
                
                gr.HTML("<h3>⚙️ Configuration</h3>")
                
                num_nails = gr.Slider(
                    minimum=300,
                    maximum=1000,
                    value=400,
                    step=50,
                    label="Number of Nails"
                )
                gr.HTML("<small>More nails = more detail but slower processing</small>")
                
                max_lines = gr.Slider(
                    minimum=1000,
                    maximum=5000,
                    value=3000,
                    step=250,
                    label="Maximum Lines"
                )
                gr.HTML("<small>More lines = denser result</small>")
                
                gr.HTML("<h3>🎨 Appearance</h3>")
                
                thread_color = gr.ColorPicker(
                    value="#000000",
                    label="Thread Color"
                )
                
                line_opacity = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.8,
                    step=0.1,
                    label="Line Opacity"
                )
                
                gr.HTML("<h3>🔧 Processing</h3>")
                
                threshold = gr.Slider(
                    minimum=50,
                    maximum=200,
                    value=128,
                    step=10,
                    label="Threshold"
                )
                gr.HTML("<small>Lower = more detail, Higher = simpler shapes</small>")
                
                blur_kernel = gr.Slider(
                    minimum=1,
                    maximum=15,
                    value=5,
                    step=2,
                    label="Blur Amount"
                )
                gr.HTML("<small>Smooths the image before processing</small>")
                
                generate_btn = gr.Button(
                    "🚀 Generate String Art",
                    variant="primary",
                    size="lg"
                )
            
            # Right Column - Results
            with gr.Column(scale=2):
                gr.HTML("<h3>📊 Results</h3>")
                
                with gr.Row():
                    with gr.Column():
                        processed_output = gr.Image(
                            label="Processed Image",
                            type="pil"
                        )
                    
                    with gr.Column():
                        result_output = gr.Image(
                            label="String Art Result",
                            type="pil"
                        )
                
                info_output = gr.Markdown(
                    label="Generation Info",
                    value="Click 'Generate String Art' to start!"
                )
                
                gr.HTML("<h3>📥 Downloads</h3>")
                
                with gr.Row():
                    instructions_text = gr.Textbox(
                        label="Threading Instructions",
                        lines=10,
                        max_lines=20,
                        placeholder="Instructions will appear here after generation..."
                    )
                    gr.HTML("<small>Copy these step-by-step instructions for physical crafting</small>")
        
        # Toggle visibility based on predefined checkbox
        def toggle_input_method(use_pred):
            return {
                predefined_group: gr.update(visible=use_pred),
                upload_group: gr.update(visible=not use_pred)
            }
        
        use_predefined.change(
            toggle_input_method,
            inputs=[use_predefined],
            outputs=[predefined_group, upload_group]
        )
        
        # Generate button click handler
        def generate_handler(image, shape, use_pred, nails, lines, thresh, blur, color, opacity):
            result_img, processed_img, info, instructions = generate_string_art_interface(
                image, shape, use_pred, nails, lines, thresh, blur, color, opacity
            )
            
            return result_img, processed_img, info, instructions
        
        generate_btn.click(
            generate_handler,
            inputs=[
                input_image, shape_type, use_predefined, num_nails, max_lines,
                threshold, blur_kernel, thread_color, line_opacity
            ],
            outputs=[
                result_output, processed_output, info_output, instructions_text
            ]
        )
        
        # Add examples
        gr.HTML("<h3>💡 Quick Start Examples</h3>")
        
        gr.Examples(
            examples=[
                [None, "Circle", True, 400, 3000, 128, 5, "#000000", 0.8],
                [None, "Heart", True, 500, 3500, 120, 3, "#FF0000", 0.7],
                [None, "Circle", True, 600, 4000, 140, 7, "#0000FF", 0.9],
            ],
            inputs=[
                input_image, shape_type, use_predefined, num_nails, max_lines,
                threshold, blur_kernel, thread_color, line_opacity
            ],
            outputs=[
                result_output, processed_output, info_output, instructions_text
            ],
            fn=generate_handler,
            cache_examples=False
        )
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 10px;">
            <h4>🔨 Physical Crafting Tips</h4>
            <p>
                <strong>Materials:</strong> Wooden board, small nails, string/thread, hammer<br>
                <strong>Process:</strong> Mark nail positions → Hammer nails → Follow step-by-step instructions<br>
                <strong>Pro Tip:</strong> Keep consistent string tension for best results!
            </p>
        </div>
        """)
    
    return app

# Launch the app
if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )