import gradio as gr
import numpy as np
from PIL import Image, ImageFilter
import cv2
import os
import torch
import torch.nn.functional as F
from torchvision import transforms
import warnings
warnings.filterwarnings("ignore")

# ZeroGPU decorator (if available)
try:
    import spaces
    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False
    # Create a dummy decorator if spaces is not available
    def spaces_gpu(func):
        return func
    spaces = type('spaces', (), {'GPU': spaces_gpu})()

# VAAPI acceleration check
def check_vaapi_support():
    """Check if VAAPI is available for hardware acceleration"""
    try:
        # Check if VAAPI devices are available
        vaapi_devices = [f for f in os.listdir('/dev/dri') if f.startswith('render')]
        return len(vaapi_devices) > 0
    except:
        return False

HAS_VAAPI = check_vaapi_support()

class TorchUpscaler:
    """PyTorch-based upscaler that can use GPU acceleration"""
    
    def __init__(self, device='auto'):
        if device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)
        
        print(f"Using device: {self.device}")
    
    def bicubic_torch(self, image_tensor, scale_factor):
        """GPU-accelerated bicubic upscaling using PyTorch"""
        return F.interpolate(
            image_tensor, 
            scale_factor=scale_factor, 
            mode='bicubic', 
            align_corners=False,
            antialias=True
        )
    
    def lanczos_torch(self, image_tensor, scale_factor):
        """GPU-accelerated Lanczos-style upscaling"""
        # PyTorch doesn't have native Lanczos, use bicubic with antialiasing
        return F.interpolate(
            image_tensor,
            scale_factor=scale_factor,
            mode='bicubic',
            align_corners=False,
            antialias=True
        )
    
    def esrgan_style_upscale(self, image_tensor, scale_factor):
        """Simple ESRGAN-style upscaling using convolutions"""
        # This is a simplified version - in practice you'd load a pre-trained model
        b, c, h, w = image_tensor.shape
        
        # Simple upscaling with edge enhancement
        upscaled = F.interpolate(image_tensor, scale_factor=scale_factor, mode='bicubic', align_corners=False)
        
        # Apply a simple sharpening kernel
        kernel = torch.tensor([[[[-1, -1, -1],
                                [-1,  9, -1],
                                [-1, -1, -1]]]], dtype=torch.float32, device=self.device)
        kernel = kernel.repeat(c, 1, 1, 1)
        
        # Apply convolution for sharpening
        sharpened = F.conv2d(upscaled, kernel, padding=1, groups=c)
        
        # Blend original and sharpened
        result = 0.8 * upscaled + 0.2 * sharpened
        return torch.clamp(result, 0, 1)

class VAAPIUpscaler:
    """VAAPI hardware-accelerated upscaler"""
    
    def __init__(self):
        self.vaapi_available = HAS_VAAPI
        if self.vaapi_available:
            print("VAAPI hardware acceleration available")
        else:
            print("VAAPI hardware acceleration not available")
    
    def upscale_vaapi(self, image_array, scale_factor, method):
        """Use VAAPI for hardware-accelerated upscaling"""
        if not self.vaapi_available:
            return None
        
        try:
            h, w = image_array.shape[:2]
            new_h, new_w = int(h * scale_factor), int(w * scale_factor)
            
            # VAAPI upscaling (simplified - you'd need proper VAAPI setup)
            # This is a placeholder for actual VAAPI implementation
            if method == "VAAPI_BICUBIC":
                return cv2.resize(image_array, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
            elif method == "VAAPI_LANCZOS":
                return cv2.resize(image_array, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
            
        except Exception as e:
            print(f"VAAPI upscaling failed: {e}")
            return None

# Initialize upscalers
torch_upscaler = TorchUpscaler()
vaapi_upscaler = VAAPIUpscaler()

@spaces.GPU if HAS_ZEROGPU else lambda x: x
def upscale_image_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration):
    """
    Accelerated upscaling with GPU/VAAPI support
    """
    if image is None:
        return None
    
    original_width, original_height = image.size
    new_width = int(original_width * scale_factor)
    new_height = int(original_height * scale_factor)
    
    try:
        if use_gpu_acceleration and torch.cuda.is_available():
            # GPU-accelerated upscaling
            print("Using GPU acceleration")
            
            # Convert PIL to tensor
            transform = transforms.Compose([
                transforms.ToTensor(),
            ])
            
            image_tensor = transform(image).unsqueeze(0).to(torch_upscaler.device)
            
            if method == "GPU_Bicubic":
                upscaled_tensor = torch_upscaler.bicubic_torch(image_tensor, scale_factor)
            elif method == "GPU_Lanczos":
                upscaled_tensor = torch_upscaler.lanczos_torch(image_tensor, scale_factor)
            elif method == "GPU_ESRGAN_Style":
                upscaled_tensor = torch_upscaler.esrgan_style_upscale(image_tensor, scale_factor)
            else:
                upscaled_tensor = torch_upscaler.bicubic_torch(image_tensor, scale_factor)
            
            # Convert back to PIL
            upscaled_tensor = upscaled_tensor.squeeze(0).cpu()
            to_pil = transforms.ToPILImage()
            upscaled = to_pil(upscaled_tensor)
            
        elif method.startswith("VAAPI_") and HAS_VAAPI:
            # VAAPI hardware acceleration
            print("Using VAAPI acceleration")
            img_array = np.array(image)
            upscaled_array = vaapi_upscaler.upscale_vaapi(img_array, scale_factor, method)
            
            if upscaled_array is not None:
                upscaled = Image.fromarray(upscaled_array)
            else:
                # Fallback to CPU
                upscaled = image.resize((new_width, new_height), Image.BICUBIC)
                
        else:
            # CPU fallback methods
            print("Using CPU methods")
            img_array = np.array(image)
            
            if method == "Bicubic":
                upscaled = image.resize((new_width, new_height), Image.BICUBIC)
            elif method == "Lanczos":
                upscaled = image.resize((new_width, new_height), Image.LANCZOS)
            elif method == "EDSR (OpenCV)":
                img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
                upscaled_bgr = cv2.resize(img_bgr, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
                upscaled_rgb = cv2.cvtColor(upscaled_bgr, cv2.COLOR_BGR2RGB)
                upscaled = Image.fromarray(upscaled_rgb)
            elif method == "Nearest Neighbor":
                upscaled = image.resize((new_width, new_height), Image.NEAREST)
            else:
                upscaled = image.resize((new_width, new_height), Image.BICUBIC)
        
        # Apply quality enhancement
        if enhance_quality:
            upscaled = upscaled.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))
        
        return upscaled
        
    except Exception as e:
        print(f"Error during upscaling: {e}")
        return image

def get_available_methods():
    """Get list of available upscaling methods based on hardware"""
    methods = ["Bicubic", "Lanczos", "EDSR (OpenCV)", "Nearest Neighbor"]
    
    if torch.cuda.is_available():
        methods.extend(["GPU_Bicubic", "GPU_Lanczos", "GPU_ESRGAN_Style"])
    
    if HAS_VAAPI:
        methods.extend(["VAAPI_BICUBIC", "VAAPI_LANCZOS"])
    
    return methods

def get_system_info():
    """Get system acceleration information"""
    info = []
    
    # CUDA info
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        info.append(f"🚀 CUDA GPU: {gpu_name} ({gpu_memory:.1f} GB)")
    else:
        info.append("❌ CUDA not available")
    
    # ZeroGPU info
    if HAS_ZEROGPU:
        info.append("✅ ZeroGPU support enabled")
    else:
        info.append("❌ ZeroGPU not available")
    
    # VAAPI info
    if HAS_VAAPI:
        info.append("✅ VAAPI hardware acceleration available")
    else:
        info.append("❌ VAAPI not available")
    
    return "\n".join(info)

def process_and_info_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration):
    """Process image with acceleration and return info"""
    if image is None:
        return None, "Please upload an image first"
    
    # Get original info
    original_info = f"Original: {image.size[0]} × {image.size[1]} pixels"
    
    # Process image
    result = upscale_image_accelerated(image, scale_factor, method, enhance_quality, use_gpu_acceleration)
    
    if result is None:
        return None, "Error processing image"
    
    # Get result info
    result_info = f"Upscaled: {result.size[0]} × {result.size[1]} pixels"
    
    # Acceleration info
    accel_info = "CPU" if not use_gpu_acceleration else "GPU/Hardware"
    
    combined_info = f"""
    ## Processing Details
    {original_info}  
    {result_info}  
    **Scale Factor:** {scale_factor}x  
    **Method:** {method}  
    **Acceleration:** {accel_info}  
    **Quality Enhancement:** {'✅' if enhance_quality else '❌'}
    
    ## System Status
    {get_system_info()}
    """
    
    return result, combined_info

# Create the accelerated interface
def create_accelerated_upscaler():
    available_methods = get_available_methods()
    
    with gr.Blocks(title="Accelerated Image Upscaler", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🚀 Accelerated Image Upscaler
        
        High-performance image upscaling with GPU and hardware acceleration support.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    type="pil",
                    label="Upload Image",
                    sources=["upload", "clipboard"]
                )
                
                scale_factor = gr.Slider(
                    minimum=1.5,
                    maximum=4.0,
                    step=0.5,
                    value=2.0,
                    label="Scale Factor"
                )
                
                method = gr.Dropdown(
                    choices=available_methods,
                    value=available_methods[0],
                    label="Upscaling Method"
                )
                
                use_gpu_acceleration = gr.Checkbox(
                    label="Use GPU Acceleration",
                    value=torch.cuda.is_available()
                )
                
                enhance_quality = gr.Checkbox(
                    label="Apply Quality Enhancement",
                    value=True
                )
                
                process_btn = gr.Button("🚀 Upscale Image", variant="primary")
                
                gr.Markdown(f"""
                ### Available Methods:
                - **Standard**: Bicubic, Lanczos, EDSR, Nearest
                - **GPU**: {', '.join([m for m in available_methods if m.startswith('GPU_')])}
                - **VAAPI**: {', '.join([m for m in available_methods if m.startswith('VAAPI_')])}
                """)
            
            with gr.Column(scale=2):
                output_image = gr.Image(
                    label="Upscaled Image",
                    type="pil"
                )
                
                image_info = gr.Markdown(
                    value=f"## System Status\n{get_system_info()}",
                    label="Processing Information"
                )
        
        # Event handlers
        process_btn.click(
            fn=process_and_info_accelerated,
            inputs=[input_image, scale_factor, method, enhance_quality, use_gpu_acceleration],
            outputs=[output_image, image_info]
        )
    
    return demo

# Launch the interface
if __name__ == "__main__":
    demo = create_accelerated_upscaler()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )