import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import cv2
import torch
import warnings
warnings.filterwarnings('ignore')

class ObjectRemover:
    """Remove objects from images using inpainting"""
    
    def __init__(self):
        self.methods = {
            'simple': self.simple_remove,
            'inpaint': self.inpaint_remove,
            'clone': self.clone_remove
        }
    
    def simple_remove(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Simple removal by filling with average color"""
        result = image.copy()
        
        # Get average color from surrounding area
        kernel = np.ones((5, 5), np.uint8)
        dilated_mask = cv2.dilate(mask, kernel, iterations=2)
        
        # Find border pixels
        border = dilated_mask - mask
        
        # Get average border color
        border_pixels = image[border > 0]
        if len(border_pixels) > 0:
            avg_color = np.mean(border_pixels, axis=0)
            
            # Fill masked area with average color
            result[mask > 0] = avg_color.astype(np.uint8)
        
        return result
    
    def inpaint_remove(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Use OpenCV inpainting"""
        # Convert to BGR for OpenCV
        if len(image.shape) == 3 and image.shape[2] == 3:
            bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        else:
            bgr_image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
        
        # Inpaint using Navier-Stokes method
        result = cv2.inpaint(
            bgr_image,
            mask,
            inpaintRadius=3,
            flags=cv2.INPAINT_TELEA
        )
        
        # Convert back to RGB
        if len(result.shape) == 3 and result.shape[2] == 3:
            return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        else:
            return cv2.cvtColor(result, cv2.COLOR_BGRA2RGBA)
    
    def clone_remove(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Use OpenCV seamless cloning"""
        # Create a source patch from nearby area
        height, width = image.shape[:2]
        
        # Find a nearby clean patch
        kernel = np.ones((11, 11), np.uint8)
        dilated_mask = cv2.dilate(mask, kernel, iterations=3)
        
        # Find non-masked areas
        non_masked = 255 - dilated_mask
        
        # Find contours of non-masked areas
        contours, _ = cv2.findContours(non_masked, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
            # Get largest contour
            largest_contour = max(contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(largest_contour)
            
            # Use this area as source
            source_patch = image[y:y+h, x:x+w]
            
            # Resize source to match mask size
            mask_height, mask_width = np.where(mask > 0)
            if len(mask_height) > 0 and len(mask_width) > 0:
                min_y, max_y = mask_height.min(), mask_height.max()
                min_x, max_x = mask_width.min(), mask_width.max()
                
                patch_height = max_y - min_y + 1
                patch_width = max_x - min_x + 1
                
                if patch_height > 0 and patch_width > 0:
                    source_resized = cv2.resize(source_patch, (patch_width, patch_height))
                    
                    # Place resized patch
                    result = image.copy()
                    result[min_y:max_y+1, min_x:max_x+1] = source_resized
                    return result
        
        return self.simple_remove(image, mask)
    
    def remove_object(self, image: np.ndarray, mask: np.ndarray, method: str = 'inpaint') -> np.ndarray:
        """Remove object using specified method"""
        if method in self.methods:
            return self.methods[method](image, mask)
        return self.inpaint_remove(image, mask)

# Initialize remover
remover = ObjectRemover()

def process_object_removal(
    image: np.ndarray,
    brush_size: int,
    method: str
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Process image with object removal"""
    try:
        # Initialize mask (all zeros)
        mask = np.zeros(image.shape[:2], dtype=np.uint8)
        
        # For demo purposes, create a sample mask (remove center)
        height, width = image.shape[:2]
        
        # Create a circle mask in center (for demonstration)
        center_x, center_y = width // 2, height // 2
        radius = min(width, height) // 4
        
        # Draw circle on mask
        cv2.circle(mask, (center_x, center_y), radius, 255, -1)
        
        # Apply removal
        result = remover.remove_object(image, mask, method)
        
        # Create visualization
        visualization = image.copy()
        visualization[mask > 0] = [255, 0, 0]  # Highlight area to remove in red
        
        # Create mask visualization
        mask_viz = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
        
        return result, visualization, mask_viz
        
    except Exception as e:
        print(f"Error: {e}")
        return image, image, np.zeros_like(image)