import gradio as gr
import numpy as np
import tempfile
import os
from PIL import Image
import sys

# Import your modules (you'll need to include them in the space)
from strings import *

def process_string_art(image, n_hooks=180, radius=250, quantization=30):
    """Process uploaded image and return string art result"""
    
    # Save uploaded image temporarily
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_input:
        image.save(tmp_input.name)
        input_path = tmp_input.name
    
    # Create temporary output prefix
    output_prefix = tempfile.mktemp()
    
    try:
        # Build adjacency matrix
        sparse, hooks, edge_codes = build_arc_adjecency_matrix(n_hooks, radius)
        
        # Process image
        shrinkage = 0.75
        img = image_from_pil(image, int(radius * 2 * shrinkage))
        sparse_b = build_image_vector(img, radius)
        
        # Solve linear system
        result = scipy.sparse.linalg.lsqr(sparse, np.array(sparse_b.todense()).flatten())
        x = result[0]
        
        # Apply quantization
        x = np.clip(x, 0, 1e6)
        max_edge_weight_orig = np.max(x)
        x_quantized = (x / np.max(x) * quantization).round()
        clip_factor = 0.3
        x_quantized = np.clip(x_quantized, 0, int(np.max(x_quantized) * clip_factor))
        x = x_quantized / quantization * max_edge_weight_orig
        
        # Reconstruct final image
        brightness_correction = 1.2
        final_image = reconstruct(x * brightness_correction, sparse, radius)
        
        # Calculate statistics
        arc_count = int(np.sum(x_quantized))
        unique_arcs = len(x_quantized[x_quantized > 0])
        
        # Convert to PIL Image for return
        final_pil = Image.fromarray(np.clip(final_image, 0, 255).astype(np.uint8))
        
        stats = f"Total arcs: {arc_count}\nUnique arc types: {unique_arcs}"
        
        return final_pil, stats
        
    finally:
        # Cleanup
        if os.path.exists(input_path):
            os.unlink(input_path)

def image_from_pil(pil_image, size):
    """Convert PIL image to grayscale numpy array"""
    img = pil_image.convert('L')  # Convert to grayscale
    img = img.resize((size, size), Image.Resampling.LANCZOS)
    return np.array(img)

# Create Gradio interface
iface = gr.Interface(
    fn=process_string_art,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(50, 360, value=180, step=10, label="Number of Hooks"),
        gr.Slider(100, 500, value=250, step=50, label="Circle Radius"),
        gr.Slider(10, 100, value=30, step=5, label="Quantization Level")
    ],
    outputs=[
        gr.Image(type="pil", label="String Art Result"),
        gr.Textbox(label="Statistics")
    ],
    title="String Art Generator",
    description="Convert any image into string art patterns! Upload a square image and adjust parameters.",
    examples=[
        # You can add example images here
    ]
)

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