import gradio as gr
from PIL import Image, ImageFilter
import numpy as np
import io
import tempfile
import vtracer
from skimage import color, filters, feature, morphology
import cv2

def preprocess_image(image, blur_radius, edge_enhance, edge_threshold, detail_level):
    """Preprocess the image with advanced options before vectorization."""
    
    if blur_radius > 0:
        image = image.filter(ImageFilter.GaussianBlur(blur_radius))
    
    # Convert to grayscale for edge detection
    gray_image = np.array(image.convert('L'))
    
    # Detail level settings
    if detail_level == 'Very Low':
        sigma = 3.0
        morphology_size = 5
    elif detail_level == 'Low':
        sigma = 2.0
        morphology_size = 4
    elif detail_level == 'Medium':
        sigma = 1.5
        morphology_size = 3
    elif detail_level == 'High':
        sigma = 1.0
        morphology_size = 2
    else:  # Ultra
        sigma = 0.5
        morphology_size = 1
    
    if edge_enhance:
        # Canny edge detection
        edges = feature.canny(gray_image, sigma=sigma, low_threshold=edge_threshold)
        
        # Morphological operations to refine edges
        edges = morphology.dilation(edges, morphology.square(morphology_size))
        edges_img = Image.fromarray((edges * 255).astype(np.uint8))
        
        # Blend the edges with the original image
        image = Image.blend(image.convert('RGB'), edges_img.convert('RGB'), alpha=0.5)

    return image

def convert_image(image, blur_radius, edge_enhance, edge_threshold, detail_level, color_mode, 
                  hierarchical, mode, filter_speckle, color_precision, layer_difference, 
                  corner_threshold, length_threshold, max_iterations, splice_threshold, path_precision):
    """Convert an image to SVG using vtracer with customizable and advanced parameters."""
    
    # Preprocess the image with additional detail level settings
    image = preprocess_image(image, blur_radius, edge_enhance, edge_threshold, detail_level)
    
    # Convert Gradio image to bytes for vtracer compatibility
    img_byte_array = io.BytesIO()
    image.save(img_byte_array, format='PNG')
    img_bytes = img_byte_array.getvalue()
    
    # Perform the conversion
    svg_str = vtracer.convert_raw_image_to_svg(
        img_bytes,
        img_format='png',
        colormode=color_mode.lower(),
        hierarchical=hierarchical.lower(),
        mode=mode.lower(),
        filter_speckle=int(filter_speckle),
        color_precision=int(color_precision),
        layer_difference=int(layer_difference),
        corner_threshold=int(corner_threshold),
        length_threshold=float(length_threshold),
        max_iterations=int(max_iterations),
        splice_threshold=int(splice_threshold),
        path_precision=int(path_precision)
    )
    
    # Save the SVG string to a temporary file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.svg')
    temp_file.write(svg_str.encode('utf-8'))
    temp_file.close()
    
    # Display the SVG in the Gradio interface and provide the download link
    svg_html = f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'
    return gr.HTML(svg_html), temp_file.name

# Gradio interface
iface = gr.Blocks()

with iface:
    gr.Markdown("# Advanced Image to SVG Converter")
    gr.Markdown("Upload an image and customize the conversion parameters for high-quality vector results. This tool provides advanced options to analyze and vectorize images at a pixel level with various detail settings.")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        blur_radius_input = gr.Slider(minimum=0, maximum=10, value=0, step=0.5, label="Blur Radius (for smoothing)")
        edge_enhance_input = gr.Checkbox(value=False, label="AI Edge Enhance")
        edge_threshold_input = gr.Slider(minimum=0.1, maximum=3.0, value=1.0, step=0.1, label="Edge Detection Threshold")
        detail_level_input = gr.Radio(choices=["Very Low", "Low", "Medium", "High", "Ultra"], value="Medium", label="Detail Level")
    
    with gr.Row():
        color_mode_input = gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode")
        hierarchical_input = gr.Radio(choices=["Stacked", "Cutout"], value="Stacked", label="Hierarchical")
        mode_input = gr.Radio(choices=["Spline", "Polygon", "None"], value="Spline", label="Mode")
    
    with gr.Row():
        filter_speckle_input = gr.Slider(minimum=1, maximum=100, value=4, step=1, label="Filter Speckle")
        color_precision_input = gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Color Precision")
        layer_difference_input = gr.Slider(minimum=1, maximum=100, value=16, step=1, label="Layer Difference")
    
    with gr.Row():
        corner_threshold_input = gr.Slider(minimum=1, maximum=100, value=60, step=1, label="Corner Threshold")
        length_threshold_input = gr.Slider(minimum=1, maximum=100, value=4.0, step=0.5, label="Length Threshold")
        max_iterations_input = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Iterations")
    
    with gr.Row():
        splice_threshold_input = gr.Slider(minimum=1, maximum=100, value=45, step=1, label="Splice Threshold")
        path_precision_input = gr.Slider(minimum=1, maximum=100, value=8, step=1, label="Path Precision")
    
    convert_button = gr.Button("Convert Image to SVG")
    svg_output = gr.HTML(label="SVG Output")
    download_output = gr.File(label="Download SVG")
    
    convert_button.click(
        fn=convert_image,
        inputs=[
            image_input, blur_radius_input, edge_enhance_input, edge_threshold_input, detail_level_input,
            color_mode_input, hierarchical_input, mode_input, filter_speckle_input, color_precision_input,
            layer_difference_input, corner_threshold_input, length_threshold_input, max_iterations_input,
            splice_threshold_input, path_precision_input
        ],
        outputs=[svg_output, download_output]
    )

iface.launch()