import gradio as gr
from PIL import Image, ImageFilter
import numpy as np
import io
import tempfile
import vtracer
from skimage import feature, filters, morphology
import cv2
from rembg import remove
from sklearn.cluster import KMeans

def quantize_colors(image, num_colors):
    """Reduce the number of colors in an image."""
    try:
        image_np = np.array(image)
        h, w, c = image_np.shape
        image_reshaped = image_np.reshape((-1, 3))

        kmeans = KMeans(n_clusters=num_colors, random_state=42).fit(image_reshaped)
        new_colors = kmeans.cluster_centers_[kmeans.labels_].reshape(h, w, 3).astype(np.uint8)
        
        return Image.fromarray(new_colors)
    except Exception as e:
        print(f"Error during color quantization: {e}")
        raise

def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization, enhance_with_ai, remove_bg):
    """Advanced preprocessing of the image before vectorization."""
    try:
        if blur_radius > 0:
            image = image.filter(ImageFilter.GaussianBlur(blur_radius))

        if sharpen_radius > 0:
            image = image.filter(ImageFilter.UnsharpMask(radius=sharpen_radius, percent=150, threshold=3))

        if noise_reduction > 0:
            image_np = np.array(image)
            image_np = cv2.fastNlMeansDenoisingColored(image_np, None, h=noise_reduction, templateWindowSize=7, searchWindowSize=21)
            image = Image.fromarray(image_np)

        if detail_level > 0:
            sigma = max(0.5, 3.0 - (detail_level * 0.5))
            image_np = np.array(image.convert('L'))

            if edge_method == 'Canny':
                edges = feature.canny(image_np, sigma=sigma)
            elif edge_method == 'Sobel':
                edges = filters.sobel(image_np)
            elif edge_method == 'Scharr':
                edges = filters.scharr(image_np)
            else:  # Prewitt
                edges = filters.prewitt(image_np)

            edges = morphology.dilation(edges, morphology.square(max(1, 6 - detail_level)))
            edges_img = Image.fromarray((edges * 255).astype(np.uint8))
            image = Image.blend(image.convert('RGB'), edges_img.convert('RGB'), alpha=0.5)

        if color_quantization > 0:
            image = quantize_colors(image, color_quantization)

        if enhance_with_ai:
            image_np = np.array(image)
            # AI-based enhancement for smoothing edges without background removal
            # Optionally apply background removal only if remove_bg is checked
            if remove_bg:
                image_np = remove(image_np)
            image = Image.fromarray(image_np)

    except Exception as e:
        print(f"Error during preprocessing: {e}")
        raise

    return image

def convert_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization,
                  color_mode, hierarchical, mode, filter_speckle, color_precision, layer_difference,
                  corner_threshold, length_threshold, max_iterations, splice_threshold, path_precision,
                  enhance_with_ai, remove_bg, upscale_factor):
    """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, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization, enhance_with_ai, remove_bg)

    # Upscale the image if needed
    if upscale_factor > 1:
        new_size = (int(image.width * upscale_factor), int(image.height * upscale_factor))
        image = image.resize(new_size, Image.LANCZOS)

    # 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()

    try:
        # 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

    except Exception as e:
        print(f"Error during vectorization: {e}")
        return f"Error: {e}", None

# Gradio interface
iface = gr.Blocks()

with iface:
    gr.Markdown("# Super-Advanced Image to SVG Converter with Enhanced Models")
    
    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.1, label="Blur Radius (for smoothing)")
        sharpen_radius_input = gr.Slider(minimum=0, maximum=5, value=0, step=0.1, label="Sharpen Radius")
        noise_reduction_input = gr.Slider(minimum=0, maximum=30, value=0, step=1, label="Noise Reduction")
        enhance_with_ai_input = gr.Checkbox(label="AI Edge Enhance", value=False)
        remove_bg_input = gr.Checkbox(label="Remove Background", value=False)
        upscale_factor_input = gr.Slider(minimum=1, maximum=4, value=1, step=0.1, label="Upscale Factor (1 = No Upscaling)")
    
    with gr.Row():
        detail_level_input = gr.Slider(minimum=0, maximum=10, value=5, step=1, label="Detail Level")
        edge_method_input = gr.Radio(choices=["Canny", "Sobel", "Scharr", "Prewitt"], value="Canny", label="Edge Detection Method")
        color_quantization_input = gr.Slider(minimum=2, maximum=64, value=0, step=2, label="Color Quantization (0 to disable)")
    
    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, sharpen_radius_input, noise_reduction_input, detail_level_input, edge_method_input, color_quantization_input,
            color_mode_input, hierarchical_input, mode_input, filter_speckle_input, 
            color_precision_input, length_threshold_input, max_iterations_input, 
            splice_threshold_input, path_precision_input, enhance_with_ai_input, remove_bg_input, upscale_factor_input
        ],
        outputs=[svg_output, download_output]
    )


iface.launch()