# String Art Generator for Hugging Face Spaces
# Upload an image and generate downloadable string art instructions

import numpy as np
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw, ImageFont
import math
import io
import zipfile
import os
from reportlab.lib.pagesizes import letter, A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
import gradio as gr
import tempfile
import json

class StringArtGenerator:
    def __init__(self, num_pins=200, canvas_size=800):
        self.num_pins = num_pins
        self.canvas_size = canvas_size
        self.pins = []
        self.string_paths = []
        self.image_processed = None
        self.original_image = None
        
    def process_image(self, image_path):
        """Process the uploaded image"""
        # Load and process image
        image = Image.open(image_path)
        self.original_image = image.copy()
        
        # Convert to grayscale and resize
        image = image.convert('L')
        image = image.resize((self.canvas_size, self.canvas_size), Image.Resampling.LANCZOS)
        
        # Convert to numpy array
        img_array = np.array(image)
        
        # Apply edge detection and processing
        self.image_processed = self.preprocess_image(img_array)
        
        return self.image_processed
    
    def preprocess_image(self, img_array):
        """Preprocess image for string art conversion"""
        # Apply Gaussian blur to smooth the image
        blurred = cv2.GaussianBlur(img_array, (3, 3), 0)
        
        # Apply edge detection
        edges = cv2.Canny(blurred, 50, 150)
        
        # Combine original image with edges for better string art effect
        # Invert so dark areas need more strings
        processed = 255 - blurred
        
        # Enhance contrast
        processed = cv2.equalizeHist(processed)
        
        # Combine with edges
        combined = cv2.addWeighted(processed, 0.7, edges, 0.3, 0)
        
        return combined
    
    def generate_pins(self, shape='circle'):
        """Generate pin positions around the perimeter"""
        pins = []
        center = self.canvas_size // 2
        
        if shape == 'circle':
            radius = center - 50  # Leave some margin
            for i in range(self.num_pins):
                angle = 2 * math.pi * i / self.num_pins
                x = center + radius * math.cos(angle)
                y = center + radius * math.sin(angle)
                pins.append((int(x), int(y)))
        
        elif shape == 'square':
            margin = 50
            side_pins = self.num_pins // 4
            
            # Top side
            for i in range(side_pins):
                x = margin + i * (self.canvas_size - 2 * margin) / side_pins
                pins.append((int(x), margin))
            
            # Right side
            for i in range(side_pins):
                y = margin + i * (self.canvas_size - 2 * margin) / side_pins
                pins.append((self.canvas_size - margin, int(y)))
            
            # Bottom side
            for i in range(side_pins):
                x = self.canvas_size - margin - i * (self.canvas_size - 2 * margin) / side_pins
                pins.append((int(x), self.canvas_size - margin))
            
            # Left side
            for i in range(side_pins):
                y = self.canvas_size - margin - i * (self.canvas_size - 2 * margin) / side_pins
                pins.append((margin, int(y)))
        
        self.pins = pins
        return pins
    
    def calculate_string_score(self, pin1_idx, pin2_idx, current_canvas):
        """Calculate score for adding a string between two pins"""
        pin1 = self.pins[pin1_idx]
        pin2 = self.pins[pin2_idx]
        
        x1, y1 = pin1
        x2, y2 = pin2
        
        length = int(math.sqrt((x2-x1)**2 + (y2-y1)**2))
        if length == 0:
            return 0
        
        score = 0
        for i in range(length):
            t = i / length
            x = int(x1 + t * (x2 - x1))
            y = int(y1 + t * (y2 - y1))
            
            if 0 <= x < self.canvas_size and 0 <= y < self.canvas_size:
                target_darkness = self.image_processed[y, x]
                current_coverage = current_canvas[y, x]
                contribution = target_darkness * (1 - current_coverage / 255)
                score += max(0, contribution)
        
        return score / length
    
    def draw_string_on_canvas(self, pin1_idx, pin2_idx, canvas, intensity=30):
        """Draw a string on the canvas"""
        pin1 = self.pins[pin1_idx]
        pin2 = self.pins[pin2_idx]
        
        x1, y1 = pin1
        x2, y2 = pin2
        
        # Draw line using Bresenham's algorithm
        dx = abs(x2 - x1)
        dy = abs(y2 - y1)
        x, y = x1, y1
        
        x_inc = 1 if x1 < x2 else -1
        y_inc = 1 if y1 < y2 else -1
        error = dx - dy
        
        while True:
            if 0 <= x < self.canvas_size and 0 <= y < self.canvas_size:
                canvas[y, x] = min(255, canvas[y, x] + intensity)
            
            if x == x2 and y == y2:
                break
            
            e2 = 2 * error
            if e2 > -dy:
                error -= dy
                x += x_inc
            if e2 < dx:
                error += dx
                y += y_inc
    
    def greedy_string_art(self, max_strings=2000, min_darkness_threshold=10):
        """Generate string art using greedy algorithm"""
        string_canvas = np.zeros((self.canvas_size, self.canvas_size))
        string_paths = []
        
        current_pin = 0  # Start from first pin
        
        for string_num in range(max_strings):
            best_pin = -1
            best_score = -1
            
            # Try all possible next pins
            for next_pin in range(self.num_pins):
                if next_pin == current_pin:
                    continue
                
                score = self.calculate_string_score(current_pin, next_pin, string_canvas)
                
                if score > best_score and score > min_darkness_threshold:
                    best_score = score
                    best_pin = next_pin
            
            if best_pin == -1:
                break
            
            string_paths.append((current_pin, best_pin))
            self.draw_string_on_canvas(current_pin, best_pin, string_canvas)
            current_pin = best_pin
        
        self.string_paths = string_paths
        return string_paths
    
    def create_visualizations(self):
        """Create all visualization images"""
        visualizations = {}
        
        # 1. Original vs Processed
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
        
        # Show original image
        if self.original_image:
            ax1.imshow(self.original_image, cmap='gray' if self.original_image.mode == 'L' else None)
        ax1.set_title('Original Image')
        ax1.axis('off')
        
        # Show processed image
        ax2.imshow(self.image_processed, cmap='gray')
        ax2.set_title('Processed for String Art')
        ax2.axis('off')
        
        plt.tight_layout()
        
        # Save original vs processed
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        visualizations['original_vs_processed'] = buf.getvalue()
        plt.close()
        
        # 2. Pin Template
        fig, ax = plt.subplots(1, 1, figsize=(10, 10))
        
        # Draw frame
        if len(self.pins) > 0:
            if abs(self.pins[0][0] - self.canvas_size//2) > abs(self.pins[0][1] - self.canvas_size//2):
                # Square frame
                rect = plt.Rectangle((50, 50), self.canvas_size-100, self.canvas_size-100, 
                                   fill=False, color='black', linewidth=2)
                ax.add_patch(rect)
            else:
                # Circular frame
                circle = plt.Circle((self.canvas_size//2, self.canvas_size//2), 
                                  self.canvas_size//2 - 50, fill=False, color='black', linewidth=2)
                ax.add_patch(circle)
        
        # Add pin positions and numbers
        for i, (x, y) in enumerate(self.pins):
            ax.plot(x, y, 'ro', markersize=4)
            if i % (max(1, len(self.pins)//20)) == 0:  # Label every nth pin
                ax.text(x+15, y+15, str(i), fontsize=8, ha='left', va='bottom', 
                       bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        
        ax.set_xlim(0, self.canvas_size)
        ax.set_ylim(0, self.canvas_size)
        ax.set_aspect('equal')
        ax.invert_yaxis()
        ax.set_title(f'Pin Template - {self.num_pins} pins\n(Print this template to mark pin positions)', 
                    fontsize=14, fontweight='bold')
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        # Save pin template
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
        buf.seek(0)
        visualizations['pin_template'] = buf.getvalue()
        plt.close()
        
        # 3. String Art Result
        string_canvas = np.zeros((self.canvas_size, self.canvas_size))
        for pin1_idx, pin2_idx in self.string_paths:
            self.draw_string_on_canvas(pin1_idx, pin2_idx, string_canvas)
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
        
        # Original processed image
        ax1.imshow(self.image_processed, cmap='gray')
        ax1.set_title('Target Image')
        ax1.axis('off')
        
        # String art result
        ax2.imshow(255 - string_canvas, cmap='gray')
        ax2.set_title(f'String Art Result\n({len(self.string_paths)} strings)')
        ax2.axis('off')
        
        plt.tight_layout()
        
        # Save string art result
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        visualizations['string_art_result'] = buf.getvalue()
        plt.close()
        
        return visualizations
    
    def estimate_string_length(self):
        """Estimate total string length needed"""
        total_length = 0
        for pin1_idx, pin2_idx in self.string_paths:
            pin1 = self.pins[pin1_idx]
            pin2 = self.pins[pin2_idx]
            distance = math.sqrt((pin1[0] - pin2[0])**2 + (pin1[1] - pin2[1])**2)
            total_length += distance / 10  # Convert pixels to cm
        
        return total_length * 1.2  # Add 20% buffer
    
    def create_instruction_pdf(self):
        """Create a comprehensive PDF instruction manual"""
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=A4)
        styles = getSampleStyleSheet()
        story = []
        
        # Title
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=24,
            spaceAfter=30,
            alignment=1  # Center alignment
        )
        story.append(Paragraph("STRING ART CONSTRUCTION MANUAL", title_style))
        story.append(Spacer(1, 20))
        
        # Materials section
        story.append(Paragraph("MATERIALS NEEDED", styles['Heading2']))
        materials = [
            f"• Circular or square frame ({self.canvas_size//10}cm recommended)",
            f"• {self.num_pins} small nails or pins",
            f"• Black thread or string (approximately {self.estimate_string_length():.1f} meters)",
            "• Hammer",
            "• Ruler or measuring tape",
            "• Pencil for marking",
            "• Printed pin template (included)"
        ]
        
        for material in materials:
            story.append(Paragraph(material, styles['Normal']))
        
        story.append(Spacer(1, 20))
        
        # Setup instructions
        story.append(Paragraph("SETUP INSTRUCTIONS", styles['Heading2']))
        setup_steps = [
            f"1. Print the pin template at actual size",
            f"2. Attach template to your frame/board",
            f"3. Mark all {self.num_pins} pin positions",
            f"4. Number each pin from 0 to {self.num_pins-1}",
            f"5. Hammer nails at each marked point",
            f"6. Leave about 5mm of nail protruding for string wrapping"
        ]
        
        for step in setup_steps:
            story.append(Paragraph(step, styles['Normal']))
        
        story.append(Spacer(1, 20))
        
        # Construction info
        story.append(Paragraph("CONSTRUCTION OVERVIEW", styles['Heading2']))
        overview = [
            f"Total strings to connect: {len(self.string_paths)}",
            f"Estimated completion time: {len(self.string_paths)//20}-{len(self.string_paths)//10} minutes",
            f"Starting pin: {self.string_paths[0][0] if self.string_paths else 0}",
            f"Estimated string length: {self.estimate_string_length():.1f} meters"
        ]
        
        for info in overview:
            story.append(Paragraph(info, styles['Normal']))
        
        story.append(PageBreak())
        
        # String connections table
        story.append(Paragraph("STRING CONNECTION SEQUENCE", styles['Heading2']))
        story.append(Paragraph("Follow this sequence exactly, connecting each numbered string from the first pin to the second pin:", styles['Normal']))
        story.append(Spacer(1, 10))
        
        # Create table with string connections
        strings_per_page = 40
        total_pages = (len(self.string_paths) + strings_per_page - 1) // strings_per_page
        
        for page in range(total_pages):
            start_idx = page * strings_per_page
            end_idx = min((page + 1) * strings_per_page, len(self.string_paths))
            
            # Table data
            table_data = [['String #', 'From Pin', 'To Pin', 'String #', 'From Pin', 'To Pin']]
            
            for i in range(start_idx, end_idx, 2):
                row = []
                # First string in row
                pin1, pin2 = self.string_paths[i]
                row.extend([str(i+1), str(pin1), str(pin2)])
                
                # Second string in row (if exists)
                if i+1 < end_idx:
                    pin1_2, pin2_2 = self.string_paths[i+1]
                    row.extend([str(i+2), str(pin1_2), str(pin2_2)])
                else:
                    row.extend(['', '', ''])
                
                table_data.append(row)
            
            # Create table
            table = Table(table_data, colWidths=[0.8*inch, 0.8*inch, 0.8*inch, 0.8*inch, 0.8*inch, 0.8*inch])
            table.setStyle(TableStyle([
                ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
                ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                ('FONTSIZE', (0, 0), (-1, 0), 10),
                ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
                ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
                ('FONTSIZE', (0, 1), (-1, -1), 8),
                ('GRID', (0, 0), (-1, -1), 1, colors.black)
            ]))
            
            story.append(table)
            
            if page < total_pages - 1:
                story.append(PageBreak())
        
        story.append(PageBreak())
        
        # Tips section
        story.append(Paragraph("CONSTRUCTION TIPS", styles['Heading2']))
        tips = [
            f"• Start with string tied to pin {self.string_paths[0][0] if self.string_paths else 0}",
            "• Maintain consistent tension throughout",
            "• Don't pull too tight - the string should have slight slack",
            "• If you make a mistake, carefully backtrack to the error",
            "• Take breaks every 100-200 strings to avoid fatigue",
            "• The image will become clearer as you add more strings",
            "• Mark your progress every 50 strings to track completion",
            "• Use good lighting to see pin numbers clearly"
        ]
        
        for tip in tips:
            story.append(Paragraph(tip, styles['Normal']))
        
        # Build PDF
        doc.build(story)
        buffer.seek(0)
        return buffer.getvalue()

def generate_string_art(image, num_pins, max_strings, shape, progress=gr.Progress()):
    """Main function to generate string art from uploaded image"""
    if image is None:
        return None, None, None, "Please upload an image first."
    
    progress(0.1, desc="Initializing...")
    
    # Initialize generator
    generator = StringArtGenerator(num_pins=num_pins)
    
    progress(0.2, desc="Processing image...")
    
    # Process image
    generator.process_image(image)
    
    progress(0.3, desc="Generating pin layout...")
    
    # Generate pins
    generator.generate_pins(shape=shape)
    
    progress(0.5, desc="Calculating optimal string paths...")
    
    # Generate string art
    string_paths = generator.greedy_string_art(max_strings=max_strings)
    
    progress(0.7, desc="Creating visualizations...")
    
    # Create visualizations
    visualizations = generator.create_visualizations()
    
    progress(0.9, desc="Generating instruction manual...")
    
    # Create instruction PDF
    pdf_content = generator.create_instruction_pdf()
    
    progress(1.0, desc="Complete!")
    
    # Create temporary files for downloads
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as f:
        f.write(pdf_content)
        pdf_path = f.name
    
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as f:
        f.write(visualizations['pin_template'])
        template_path = f.name
    
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as f:
        f.write(visualizations['string_art_result'])
        result_path = f.name
    
    # Create zip file with all outputs
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
        zip_file.writestr('instruction_manual.pdf', pdf_content)
        zip_file.writestr('pin_template.png', visualizations['pin_template'])
        zip_file.writestr('string_art_result.png', visualizations['string_art_result'])
        zip_file.writestr('original_vs_processed.png', visualizations['original_vs_processed'])
        
        # Add JSON with string paths for advanced users
        string_data = {
            'num_pins': num_pins,
            'num_strings': len(string_paths),
            'string_paths': string_paths,
            'estimated_length_meters': generator.estimate_string_length()
        }
        zip_file.writestr('string_data.json', json.dumps(string_data, indent=2))
    
    zip_buffer.seek(0)
    
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.zip', delete=False) as f:
        f.write(zip_buffer.getvalue())
        zip_path = f.name
    
    # Create summary text
    summary = f"""
## String Art Generation Complete! 🎨

**Statistics:**
- **Pins:** {num_pins}
- **Strings:** {len(string_paths)}
- **Estimated String Length:** {generator.estimate_string_length():.1f} meters
- **Estimated Construction Time:** {len(string_paths)//20}-{len(string_paths)//10} minutes
- **Frame Shape:** {shape.title()}

**Downloads Available:**
1. **Complete Package (ZIP)** - Contains all files
2. **Instruction Manual (PDF)** - Step-by-step construction guide
3. **Pin Template (PNG)** - Print this to mark pin positions

**Next Steps:**
1. Download the complete package
2. Print the pin template at actual size
3. Follow the instruction manual
4. Create your string art masterpiece!
"""
    
    return pdf_path, template_path, zip_path, summary

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="String Art Generator", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🎨 String Art Generator
        
        Transform any image into detailed string art instructions! Upload an image and get:
        - Step-by-step construction manual
        - Printable pin template 
        - Complete material list
        - Downloadable instruction package
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Upload Image",
                    type="filepath",
                    height=300
                )
                
                with gr.Row():
                    num_pins = gr.Slider(
                        minimum=100,
                        maximum=400,
                        value=200,
                        step=10,
                        label="Number of Pins",
                        info="More pins = higher detail, longer construction"
                    )
                
                with gr.Row():
                    max_strings = gr.Slider(
                        minimum=500,
                        maximum=5000,
                        value=2000,
                        step=100,
                        label="Maximum Strings",
                        info="More strings = better quality, longer time"
                    )
                
                shape = gr.Radio(
                    choices=["circle", "square"],
                    value="circle",
                    label="Frame Shape",
                    info="Choose the shape of your frame"
                )
                
                generate_btn = gr.Button(
                    "Generate String Art Instructions",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                summary_output = gr.Markdown(label="Generation Summary")
                
                with gr.Row():
                    pdf_download = gr.File(
                        label="📋 Instruction Manual (PDF)",
                        visible=True
                    )
                    template_download = gr.File(
                        label="📍 Pin Template (PNG)",
                        visible=True
                    )
                
                zip_download = gr.File(
                    label="📦 Complete Package (ZIP)",
                    visible=True
                )
        
        # Event handler
        generate_btn.click(
            fn=generate_string_art,
            inputs=[image_input, num_pins, max_strings, shape],
            outputs=[pdf_download, template_download, zip_download, summary_output]
        )
        
        gr.Markdown("""
        ## How to Use:
        1. **Upload** your image (photos, artwork, logos work well)
        2. **Adjust** settings based on desired complexity
        3. **Generate** your string art instructions
        4. **Download** the complete package
        5. **Print** the pin template and follow the manual
        
        ## Tips for Best Results:
        - Use high-contrast images
        - Simple compositions work better than complex scenes
        - Black and white or monochrome images are ideal
        - Portraits and geometric designs are excellent choices
        
        ---
        *Created with ❤️ for the maker community*
        """)
    
    return app

# Install required packages and run
if __name__ == "__main__":
    # Create and launch the interface
    app = create_interface()
    app.launch()

# For Hugging Face Spaces deployment, also include requirements.txt:
"""
gradio
opencv-python
pillow
numpy
matplotlib
reportlab
scipy
"""