import cv2
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import gradio as gr
import tempfile
import os
import shutil

def edge_directed_antialiasing(img, power=2.0):
    """
    Apply edge-directed anti-aliasing with adjustable power
    
    Parameters:
    - img: Input image (numpy array)
    - power: Anti-aliasing strength (1.0 is standard, higher values increase the effect)
    
    Returns:
    - Output image with anti-aliasing applied
    """
    # If image has alpha channel, separate it
    has_alpha = img.shape[2] == 4 if len(img.shape) > 2 else False
    if has_alpha:
        bgr = img[:, :, :3]
        alpha = img[:, :, 3]
    else:
        bgr = img
        # Create binary mask from grayscale image if no alpha
        gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
        _, alpha = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
    
    # Convert to grayscale for edge detection
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    
    # Step 1: Detect edges using Canny
    # Lower thresholds to catch more edges when power is high
    canny_threshold1 = int(100 / power)  # Lower threshold when power is high
    canny_threshold2 = int(200 / power)  # Lower threshold when power is high
    edges = cv2.Canny(gray, canny_threshold1, canny_threshold2)
    
    # Dilate edges more when power is high
    kernel_size = int(3 * power)  # Increase kernel size with power
    kernel_size = max(3, kernel_size if kernel_size % 2 == 1 else kernel_size + 1)  # Ensure odd kernel size
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    
    # More iterations for higher power
    dilation_iterations = max(1, int(power))
    dilated_edges = cv2.dilate(edges, kernel, iterations=dilation_iterations)
    
    # Step 2: Calculate gradient direction using Sobel
    # Increase kernel size for higher power
    sobel_ksize = 3
    if power > 2.0:
        sobel_ksize = 5
    if power > 3.0:
        sobel_ksize = 7
        
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_ksize)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_ksize)
    
    # Calculate gradient magnitude and direction
    magnitude = np.sqrt(sobelx**2 + sobely**2)
    direction = np.arctan2(sobely, sobelx) * 180 / np.pi
    
    # Create output image, starting with the original
    output = bgr.copy()
    h, w = output.shape[:2]
    
    # Step 3: Apply targeted smoothing along edge directions
    # Sample farther away for higher power
    radius = max(1, int(power))
    
    edge_pixels = np.where(dilated_edges > 0)
    for y, x in zip(edge_pixels[0], edge_pixels[1]):
        # Skip border pixels
        if x < radius or y < radius or x >= w-radius or y >= h-radius:
            continue
        
        # Get local direction (perpendicular to gradient)
        local_dir = direction[y, x] + 90
        if local_dir > 180:
            local_dir -= 360
        
        # Normalize direction to 0-180 degrees
        local_dir = ((local_dir + 180) % 180)
        
        # Determine interpolation direction based on edge angle
        if 22.5 <= local_dir < 67.5:  # ~45 degree diagonal
            # Diagonal top-left to bottom-right
            neighbors = [(y-radius, x-radius), (y+radius, x+radius)]
            weights = [0.5, 0.5]
        elif 67.5 <= local_dir < 112.5:  # Vertical
            # Top to bottom
            neighbors = [(y-radius, x), (y+radius, x)]
            weights = [0.5, 0.5]
        elif 112.5 <= local_dir < 157.5:  # ~135 degree diagonal
            # Diagonal top-right to bottom-left
            neighbors = [(y-radius, x+radius), (y+radius, x-radius)]
            weights = [0.5, 0.5]
        else:  # Horizontal
            # Left to right
            neighbors = [(y, x-radius), (y, x+radius)]
            weights = [0.5, 0.5]
        
        # Only interpolate if we're between different colors (at the border)
        center_value = gray[y, x]
        neighbor_values = [gray[ny, nx] for ny, nx in neighbors]
        
        # Lower contrast threshold when power is high
        contrast_threshold = int(50 / power)
        
        # Check if this is an edge between very different values
        if abs(neighbor_values[0] - neighbor_values[1]) > contrast_threshold:
            # Apply interpolation based on local contrast
            for c in range(3):  # RGB channels
                weighted_sum = sum(weights[i] * bgr[ny, nx, c] for i, (ny, nx) in enumerate(neighbors))
                # More interpolation weight when power is high
                blend_factor = min(0.9, 0.3 * power)
                # Apply it with a blend factor to preserve some original detail
                output[y, x, c] = int((1-blend_factor) * weighted_sum + blend_factor * bgr[y, x, c])
    
    # Update alpha channel with the same smoothing for edges
    if has_alpha:
        new_alpha = alpha.copy()
        
        # Apply a specific smoothing to the alpha channel's edges
        alpha_edges = cv2.Canny(alpha, int(100/power), int(200/power))
        
        # More dilation iterations for stronger effect
        alpha_dilation_iter = max(2, int(power * 2))
        dilated_alpha_edges = cv2.dilate(alpha_edges, kernel, iterations=alpha_dilation_iter)
        
        # Radius for sampling neighborhood
        alpha_radius = max(2, int(power * 2))
        
        # For each edge pixel in alpha
        alpha_edge_pixels = np.where(dilated_alpha_edges > 0)
        for y, x in zip(alpha_edge_pixels[0], alpha_edge_pixels[1]):
            if x < alpha_radius or y < alpha_radius or x >= w-alpha_radius or y >= h-alpha_radius:
                continue
            
            # Use a larger neighborhood for better smoothing of alpha edges
            # Size increases with power
            window_radius = alpha_radius
            neighborhood = alpha[y-window_radius:y+window_radius+1, x-window_radius:x+window_radius+1].astype(np.float32)
            
            # Generate gaussian-like weights based on distance from center
            kernel_size = 2 * window_radius + 1
            weight_matrix = np.zeros((kernel_size, kernel_size), dtype=np.float32)
            
            # Create distance-based weights
            center = window_radius
            for wy in range(kernel_size):
                for wx in range(kernel_size):
                    # Calculate distance from center
                    dist = np.sqrt((wy - center)**2 + (wx - center)**2)
                    # Adjust falloff based on power 
                    falloff = 1.0 / power
                    # Gaussian-like weight
                    weight_matrix[wy, wx] = np.exp(-(dist**2) / (2 * (window_radius * falloff)**2))
            
            # Normalize weights
            weight_matrix = weight_matrix / weight_matrix.sum()
            
            # Apply weighted average
            new_alpha[y, x] = int(np.sum(neighborhood * weight_matrix))
        
        # Merge BGR with new alpha
        output = np.dstack([output, new_alpha])
        
    return output

def save_as_jpg(img, file_path):
    """
    Save image as JPG with high quality
    """
    # If image has alpha channel, blend with white background
    if len(img.shape) > 2 and img.shape[2] == 4:
        bgr = img[:, :, :3]
        alpha = img[:, :, 3].astype(float) / 255
        
        # Create white background
        bg = np.ones_like(bgr) * 255
        
        # Blend with background
        alpha = np.expand_dims(alpha, axis=2)
        alpha = np.repeat(alpha, 3, axis=2)
        result = (bgr * alpha + bg * (1 - alpha)).astype(np.uint8)
    else:
        result = img
        
    # Save as JPG
    cv2.imwrite(file_path, result, [cv2.IMWRITE_JPEG_QUALITY, 95])
    return file_path

def create_output_dirs():
    """Create necessary output directories"""
    output_dir = os.path.join(tempfile.gettempdir(), "antialiasing_output")
    os.makedirs(output_dir, exist_ok=True)
    return output_dir

def process_image(input_image):
    """
    Process image function for Gradio interface
    """
    # Create output directory for our files
    output_dir = create_output_dirs()
    
    # Convert from RGB (Gradio) to BGR (OpenCV)
    img_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
    
    # Apply edge directed anti-aliasing with power=2.0
    processed_bgr = edge_directed_antialiasing(img_bgr, power=2.0)
    
    # Save the processed image explicitly as JPG
    jpg_path = os.path.join(output_dir, "antialiased_image.jpg")
    save_as_jpg(processed_bgr, jpg_path)
    
    # Convert back to RGB for display in Gradio
    if processed_bgr.shape[2] == 4:  # Has alpha channel
        # Blend with white background
        bg = np.ones_like(processed_bgr[:,:,:3]) * 255
        alpha = processed_bgr[:,:,3]
        alpha_norm = alpha.astype(float) / 255
        alpha_norm = np.expand_dims(alpha_norm, axis=2)
        alpha_norm = np.repeat(alpha_norm, 3, axis=2)
        
        processed_rgb = processed_bgr[:,:,:3] * alpha_norm + bg * (1 - alpha_norm)
        processed_rgb = processed_rgb.astype(np.uint8)
    else:
        processed_rgb = cv2.cvtColor(processed_bgr, cv2.COLOR_BGR2RGB)
    
    # Create comparison visualization
    h, w = input_image.shape[:2]
    dpi = 100
    plt.figure(figsize=(w*2/dpi, h/dpi), dpi=dpi)
    
    plt.subplot(1, 2, 1)
    plt.imshow(input_image)
    plt.title("Original")
    plt.axis('off')
    
    plt.subplot(1, 2, 2)
    plt.imshow(processed_rgb)
    plt.title("Anti-aliased (Power = 2.0)")
    plt.axis('off')
    
    plt.tight_layout()
    
    # Save the comparison
    comparison_file = os.path.join(output_dir, "comparison.jpg")
    plt.savefig(comparison_file, dpi=dpi, bbox_inches='tight')
    plt.close()
    
    return processed_rgb, jpg_path, comparison_file

# Create Gradio interface
with gr.Blocks(title="Edge-Directed Anti-Aliasing") as app:
    gr.Markdown("# Edge-Directed Anti-Aliasing Tool")
    gr.Markdown("Upload an image and apply edge-directed anti-aliasing to smooth jagged edges.")
    
    with gr.Row():
        input_image = gr.Image(label="Upload Image", type="numpy")
        output_image = gr.Image(label="Anti-Aliased Result", type="numpy")
    
    with gr.Row():
        process_button = gr.Button("Apply Anti-Aliasing (Power = 2.0)")
    
    with gr.Row():
        download_jpg = gr.File(label="Download Anti-Aliased JPG", type="filepath")
        comparison_view = gr.Image(label="Comparison", type="filepath")
    
    # Process button functionality
    process_button.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_image, download_jpg, comparison_view]
    )
    

# Launch the app
if __name__ == "__main__":
    app.launch(share=True)