"""Super-resolution/upscaling and image enhancement functionality for telescope images."""

from enum import Enum
from typing import Optional, Tuple

import cv2
import numpy as np
from scopinator.util.logging_config import get_logger
logging = get_logger(__name__)
from skimage import restoration, filters, exposure


class UpscalingMethod(str, Enum):
    """Available upscaling methods."""

    BICUBIC = "bicubic"
    LANCZOS = "lanczos"
    EDSR = "edsr"  # Enhanced Deep Super-Resolution (if OpenCV contrib available)
    FSRCNN = "fsrcnn"  # Fast Super-Resolution CNN (if OpenCV contrib available)
    ESRGAN = "esrgan"  # Enhanced Super-Resolution GAN (requires PyTorch and model files)


class SharpeningMethod(str, Enum):
    """Available sharpening methods."""
    
    NONE = "none"
    UNSHARP_MASK = "unsharp_mask"
    LAPLACIAN = "laplacian"
    HIGH_PASS = "high_pass"


class DenoiseMethod(str, Enum):
    """Available denoising methods."""
    
    NONE = "none"
    TV_CHAMBOLLE = "tv_chambolle"  # Total variation denoising
    BILATERAL = "bilateral"  # Bilateral filter
    NON_LOCAL_MEANS = "non_local_means"  # Non-local means
    WAVELET = "wavelet"  # Wavelet denoising
    GAUSSIAN = "gaussian"  # Gaussian blur
    MEDIAN = "median"  # Median filter


class ImageUpscaler:
    """Handles super-resolution and upscaling of telescope images."""

    def __init__(self):
        self._check_opencv_contrib()
        self._check_pytorch_availability()

    def _check_opencv_contrib(self) -> bool:
        """Check if OpenCV contrib modules are available for DNN-based upscaling."""
        try:
            # Try to access DNN super-resolution module
            cv2.dnn_superres.DnnSuperResImpl_create()
            self._has_dnn_superres = True
        except AttributeError:
            self._has_dnn_superres = False
        return self._has_dnn_superres
    
    def _check_pytorch_availability(self) -> bool:
        """Check if PyTorch is available for deep learning upscaling."""
        try:
            import torch
            self._has_torch = True
            # Check if CUDA is available
            self._has_cuda = torch.cuda.is_available()
            logging.trace(f"PyTorch available: {self._has_torch}, CUDA available: {self._has_cuda}")
        except ImportError:
            self._has_torch = False
            self._has_cuda = False
            logging.trace("PyTorch not available - deep learning upscaling methods will be disabled")
        return self._has_torch

    def upscale(
        self,
        image: np.ndarray,
        scale_factor: float = 2.0,
        method: UpscalingMethod = UpscalingMethod.BICUBIC,
        denoise: bool = True,
    ) -> np.ndarray:
        """
        Upscale an image using the specified method.

        Args:
            image: Input image as numpy array
            scale_factor: Scaling factor (e.g., 2.0 for 2x upscaling)
            method: Upscaling method to use
            denoise: Whether to apply denoising before upscaling

        Returns:
            Upscaled image as numpy array
        """
        if image is None or image.size == 0:
            raise ValueError("Invalid input image")

        # Convert to appropriate data type for processing
        if image.dtype == np.uint8:
            working_image = image.astype(np.float32) / 255.0
            input_uint8 = True
        else:
            working_image = image.astype(np.float32)
            input_uint8 = False

        # Ensure we maintain color channels properly
        original_shape = working_image.shape
        is_color = len(original_shape) == 3 and original_shape[2] == 3
        logging.trace(f"Original image: shape={original_shape}, is_color={is_color}")

        # Optional denoising for telescope images (helps with noise amplification) - skip for now to avoid fuzziness
        # if denoise:
        #     working_image = self._denoise_image(working_image)

        # Calculate target dimensions
        original_height, original_width = working_image.shape[:2]
        target_height = int(original_height * scale_factor)
        target_width = int(original_width * scale_factor)

        # Apply upscaling method
        logging.trace(f"Applying upscaling method: {method}, scale_factor={scale_factor}")
        logging.trace(f"Available capabilities: dnn_superres={self._has_dnn_superres}, torch={self._has_torch}")
        logging.trace(f"Working image before upscaling: shape={working_image.shape}, dtype={working_image.dtype}")
        
        if method == UpscalingMethod.BICUBIC:
            logging.trace("Using bicubic upscaling")
            upscaled = self._bicubic_upscale(
                working_image, (target_width, target_height)
            )
        elif method == UpscalingMethod.LANCZOS:
            logging.trace("Using Lanczos upscaling")
            upscaled = self._lanczos_upscale(
                working_image, (target_width, target_height)
            )
        elif method == UpscalingMethod.EDSR and self._has_dnn_superres:
            logging.trace("Using EDSR upscaling")
            upscaled = self._edsr_upscale(working_image, scale_factor)
        elif method == UpscalingMethod.FSRCNN and self._has_dnn_superres:
            logging.trace("Using FSRCNN upscaling")
            upscaled = self._fsrcnn_upscale(working_image, scale_factor)
        elif method == UpscalingMethod.ESRGAN and self._has_torch:
            logging.trace("Using ESRGAN upscaling")
            upscaled = self._esrgan_upscale(working_image, scale_factor)
        else:
            # Fallback to bicubic if DNN methods not available
            logging.warning(f"Method {method} not available or requirements not met, falling back to bicubic")
            upscaled = self._bicubic_upscale(
                working_image, (target_width, target_height)
            )

        # Convert back to original data type
        logging.trace(f"Upscaled image before conversion: shape={upscaled.shape}, dtype={upscaled.dtype}")
        if input_uint8:
            upscaled = np.clip(upscaled * 255.0, 0, 255).astype(np.uint8)
        else:
            upscaled = upscaled.astype(image.dtype)

        logging.trace(f"Final upscaled image: shape={upscaled.shape}, dtype={upscaled.dtype}")
        return upscaled

    def _denoise_image(self, image: np.ndarray, method: DenoiseMethod = DenoiseMethod.TV_CHAMBOLLE) -> np.ndarray:
        """Apply denoising optimized for telescope images."""
        return self.denoise_image(image, method=method)
    
    def denoise_image(
        self, 
        image: np.ndarray, 
        method: DenoiseMethod = DenoiseMethod.TV_CHAMBOLLE, 
        strength: float = 1.0
    ) -> np.ndarray:
        """
        Apply denoising to an image using the specified method.
        
        Args:
            image: Input image array
            method: Denoising method to use
            strength: Denoising strength (0.0 to 2.0)
            
        Returns:
            Denoised image array
        """
        if method == DenoiseMethod.NONE:
            return image
            
        # Ensure image is in float32 format for processing
        if image.dtype == np.uint8:
            working_image = image.astype(np.float32) / 255.0
            was_uint8 = True
        else:
            working_image = image.astype(np.float32)
            was_uint8 = False
            
        try:
            if method == DenoiseMethod.TV_CHAMBOLLE:
                # Total variation denoising - excellent for astronomical images
                weight = 0.05 * strength
                denoised = restoration.denoise_tv_chambolle(
                    working_image, weight=weight, eps=2e-4, max_num_iter=200
                )
                
            elif method == DenoiseMethod.BILATERAL:
                # Bilateral filter - preserves edges while reducing noise
                if len(working_image.shape) == 3:
                    # Multi-channel image
                    sigma_color = 0.1 * strength
                    sigma_spatial = 1.0 * strength
                    denoised = restoration.denoise_bilateral(
                        working_image, sigma_color=sigma_color, sigma_spatial=sigma_spatial
                    )
                else:
                    # Single-channel image
                    sigma_color = 0.1 * strength
                    sigma_spatial = 1.0 * strength
                    denoised = restoration.denoise_bilateral(
                        working_image, sigma_color=sigma_color, sigma_spatial=sigma_spatial
                    )
                    
            elif method == DenoiseMethod.NON_LOCAL_MEANS:
                # Non-local means - very effective for textured noise
                patch_size = 5
                patch_distance = 6
                h = 0.1 * strength
                denoised = restoration.denoise_nl_means(
                    working_image, patch_size=patch_size, patch_distance=patch_distance, h=h
                )
                
            elif method == DenoiseMethod.WAVELET:
                # Wavelet denoising - good for various noise types
                sigma = 0.1 * strength
                denoised = restoration.denoise_wavelet(
                    working_image, sigma=sigma, mode='soft', rescale_sigma=True
                )
                
            elif method == DenoiseMethod.GAUSSIAN:
                # Gaussian blur - simple but effective
                sigma = 0.5 * strength
                denoised = filters.gaussian(working_image, sigma=sigma)
                
            elif method == DenoiseMethod.MEDIAN:
                # Median filter - excellent for salt-and-pepper noise
                if len(working_image.shape) == 3:
                    # Apply median filter to each channel
                    denoised = np.zeros_like(working_image)
                    disk_size = max(1, int(2 * strength))
                    disk = filters.disk(disk_size)
                    for i in range(working_image.shape[2]):
                        denoised[:, :, i] = filters.median(working_image[:, :, i], disk)
                else:
                    disk_size = max(1, int(2 * strength))
                    disk = filters.disk(disk_size)
                    denoised = filters.median(working_image, disk)
                    
            else:
                # Fallback to TV Chambolle
                weight = 0.05 * strength
                denoised = restoration.denoise_tv_chambolle(
                    working_image, weight=weight, eps=2e-4, max_num_iter=200
                )
                
        except Exception as e:
            logging.error(f"Denoising failed with method {method}: {e}")
            # Return original image if denoising fails
            denoised = working_image
            
        # Convert back to original data type
        if was_uint8:
            denoised = np.clip(denoised * 255.0, 0, 255).astype(np.uint8)
        else:
            denoised = denoised.astype(image.dtype)
            
        return denoised

    def _bicubic_upscale(
        self, image: np.ndarray, target_size: Tuple[int, int]
    ) -> np.ndarray:
        """Upscale using bicubic interpolation."""
        logging.trace(f"Applying bicubic upscaling: {image.shape} -> {target_size}")
        logging.trace(f"Input image dtype: {image.dtype}, range: [{np.min(image):.3f}, {np.max(image):.3f}]")
        
        # For multi-channel images, ensure OpenCV handles them correctly
        if len(image.shape) == 3:
            # Process color image
            upscaled = cv2.resize(image, target_size, interpolation=cv2.INTER_CUBIC)
            # Ensure we still have 3 channels after resize
            if len(upscaled.shape) == 2:
                # If somehow we lost channels, restore them
                logging.warning("Lost color channels during upscaling, restoring...")
                upscaled = np.stack([upscaled, upscaled, upscaled], axis=2)
        else:
            # Process grayscale image
            upscaled = cv2.resize(image, target_size, interpolation=cv2.INTER_CUBIC)
        
        logging.trace(f"Upscaled image shape: {upscaled.shape}")
        
        return upscaled

    def _lanczos_upscale(
        self, image: np.ndarray, target_size: Tuple[int, int]
    ) -> np.ndarray:
        """Upscale using Lanczos interpolation."""
        logging.trace(f"Applying Lanczos upscaling: {image.shape} -> {target_size}")
        logging.trace(f"Input image dtype: {image.dtype}, range: [{np.min(image):.3f}, {np.max(image):.3f}]")
        
        # For multi-channel images, ensure OpenCV handles them correctly
        if len(image.shape) == 3:
            # Process color image
            upscaled = cv2.resize(image, target_size, interpolation=cv2.INTER_LANCZOS4)
            # Ensure we still have 3 channels after resize
            if len(upscaled.shape) == 2:
                # If somehow we lost channels, restore them
                logging.warning("Lost color channels during Lanczos upscaling, restoring...")
                upscaled = np.stack([upscaled, upscaled, upscaled], axis=2)
        else:
            # Process grayscale image
            upscaled = cv2.resize(image, target_size, interpolation=cv2.INTER_LANCZOS4)
        
        logging.trace(f"Upscaled image shape: {upscaled.shape}")
        
        return upscaled

    def _edsr_upscale(self, image: np.ndarray, scale_factor: float) -> np.ndarray:
        """Upscale using Enhanced Deep Super-Resolution (requires OpenCV contrib)."""
        if not self._has_dnn_superres:
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )

        try:
            # Create DNN super-resolution object
            sr = cv2.dnn_superres.DnnSuperResImpl_create()

            # Note: In practice, you would need to download and load pre-trained EDSR models
            # For now, we'll fall back to bicubic
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )
        except Exception:
            # Fallback to bicubic
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )

    def _fsrcnn_upscale(self, image: np.ndarray, scale_factor: float) -> np.ndarray:
        """Upscale using Fast Super-Resolution CNN (requires OpenCV contrib)."""
        if not self._has_dnn_superres:
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )

        try:
            # Create DNN super-resolution object
            sr = cv2.dnn_superres.DnnSuperResImpl_create()

            # Note: In practice, you would need to download and load pre-trained FSRCNN models
            # For now, we'll fall back to bicubic
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )
        except Exception:
            # Fallback to bicubic
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )

    def get_available_methods(self) -> list[UpscalingMethod]:
        """Get list of available upscaling methods."""
        methods = [UpscalingMethod.BICUBIC, UpscalingMethod.LANCZOS]

        if self._has_dnn_superres:
            methods.extend([UpscalingMethod.EDSR, UpscalingMethod.FSRCNN])

        return methods

    def sharpen_image(
        self,
        image: np.ndarray,
        method: SharpeningMethod = SharpeningMethod.UNSHARP_MASK,
        strength: float = 1.0,
    ) -> np.ndarray:
        """
        Apply sharpening to an image.
        
        Args:
            image: Input image as numpy array
            method: Sharpening method to use
            strength: Sharpening strength (0.0 to 2.0)
            
        Returns:
            Sharpened image as numpy array
        """
        if method == SharpeningMethod.NONE or strength <= 0:
            return image
            
        # Ensure image is in float format
        if image.dtype == np.uint8:
            working_image = image.astype(np.float32) / 255.0
            was_uint8 = True
        else:
            working_image = image.astype(np.float32)
            was_uint8 = False
            
        if method == SharpeningMethod.UNSHARP_MASK:
            sharpened = self._unsharp_mask(working_image, strength)
        elif method == SharpeningMethod.LAPLACIAN:
            sharpened = self._laplacian_sharpen(working_image, strength)
        elif method == SharpeningMethod.HIGH_PASS:
            sharpened = self._high_pass_sharpen(working_image, strength)
        else:
            sharpened = working_image
            
        # Convert back to original format
        if was_uint8:
            sharpened = np.clip(sharpened * 255.0, 0, 255).astype(np.uint8)
        else:
            sharpened = sharpened.astype(image.dtype)
            
        return sharpened
    
    def _unsharp_mask(self, image: np.ndarray, strength: float) -> np.ndarray:
        """Apply unsharp mask sharpening."""
        logging.trace(f"Unsharp mask input: shape={image.shape}, dtype={image.dtype}")
        
        # Use different parameters for different image types
        if len(image.shape) == 3:
            # Color image
            radius = 1.0
            amount = strength
            logging.trace(f"Applying unsharp mask to color image: shape={image.shape}")
        else:
            # Grayscale image
            radius = 1.5
            amount = strength
            logging.trace(f"Applying unsharp mask to grayscale image: shape={image.shape}")
            
        result = filters.unsharp_mask(image, radius=radius, amount=amount, preserve_range=True)
        logging.trace(f"Unsharp mask result: shape={result.shape}, dtype={result.dtype}")
        return result
    
    def _laplacian_sharpen(self, image: np.ndarray, strength: float) -> np.ndarray:
        """Apply Laplacian sharpening."""
        # Convert to uint8 for OpenCV processing if needed
        if image.dtype != np.uint8:
            img_uint8 = (image * 255).astype(np.uint8)
            was_float = True
        else:
            img_uint8 = image
            was_float = False
            
        if len(image.shape) == 3:
            # Apply to each channel separately
            sharpened = np.zeros_like(img_uint8, dtype=np.float64)
            for i in range(image.shape[2]):
                laplacian = cv2.Laplacian(img_uint8[:, :, i], cv2.CV_64F)
                sharpened[:, :, i] = img_uint8[:, :, i].astype(np.float64) - strength * laplacian
        else:
            laplacian = cv2.Laplacian(img_uint8, cv2.CV_64F)
            sharpened = img_uint8.astype(np.float64) - strength * laplacian
        
        # Convert back to original format
        if was_float:
            sharpened = np.clip(sharpened / 255.0, 0, 1)
        else:
            sharpened = np.clip(sharpened, 0, 255).astype(np.uint8)
            
        return sharpened
    
    def _high_pass_sharpen(self, image: np.ndarray, strength: float) -> np.ndarray:
        """Apply high-pass filter sharpening."""
        # Create Gaussian blur
        blurred = filters.gaussian(image, sigma=1.0, preserve_range=True)
        # High-pass = original - blurred
        high_pass = image - blurred
        # Add back to original with strength
        sharpened = image + strength * high_pass
        return np.clip(sharpened, 0, 1)
    
    def invert_image(self, image: np.ndarray) -> np.ndarray:
        """
        Invert an image (useful for viewing negatives or different contrast).
        
        Args:
            image: Input image as numpy array
            
        Returns:
            Inverted image as numpy array
        """
        if image.dtype == np.uint8:
            return 255 - image
        else:
            # Assume float image in range [0, 1]
            return 1.0 - image

    def _esrgan_upscale(self, image: np.ndarray, scale_factor: float) -> np.ndarray:
        """Upscale using ESRGAN-inspired method (simplified implementation)."""
        if not self._has_torch:
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )
        
        try:
            from skimage import transform, filters
            
            logging.trace(f"Applying ESRGAN-inspired upscaling with edge enhancement")
            
            # First, apply bicubic upscaling as base
            target_size = (int(image.shape[1] * scale_factor), int(image.shape[0] * scale_factor))
            upscaled = cv2.resize(image, target_size, interpolation=cv2.INTER_CUBIC)
            
            # Apply edge-preserving enhancement (ESRGAN characteristic)
            # Use bilateral filter to preserve edges while smoothing
            if len(upscaled.shape) == 3:
                enhanced = np.zeros_like(upscaled)
                for i in range(upscaled.shape[2]):
                    # Convert to uint8 for bilateral filter
                    channel = upscaled[:, :, i]
                    if channel.dtype != np.uint8:
                        channel_uint8 = (np.clip(channel, 0, 1) * 255).astype(np.uint8)
                    else:
                        channel_uint8 = channel
                    
                    # Apply bilateral filter for edge preservation
                    filtered = cv2.bilateralFilter(channel_uint8, 9, 75, 75)
                    
                    # Convert back to original format
                    if upscaled.dtype != np.uint8:
                        enhanced[:, :, i] = filtered.astype(np.float32) / 255.0
                    else:
                        enhanced[:, :, i] = filtered
            else:
                if upscaled.dtype != np.uint8:
                    upscaled_uint8 = (np.clip(upscaled, 0, 1) * 255).astype(np.uint8)
                else:
                    upscaled_uint8 = upscaled
                
                enhanced = cv2.bilateralFilter(upscaled_uint8, 9, 75, 75)
                
                if upscaled.dtype != np.uint8:
                    enhanced = enhanced.astype(np.float32) / 255.0
            
            # Apply subtle sharpening for ESRGAN-like detail enhancement
            if enhanced.dtype != np.uint8:
                sharpened = filters.unsharp_mask(enhanced, radius=1.0, amount=0.5, preserve_range=True)
            else:
                sharpened = enhanced  # Skip sharpening for uint8 to avoid complications
            
            logging.trace(f"ESRGAN-inspired processing complete: {image.shape} -> {sharpened.shape}")
            return sharpened
            
        except Exception as e:
            logging.error(f"ESRGAN upscaling failed: {e}")
            return self._bicubic_upscale(
                image,
                (
                    int(image.shape[1] * scale_factor),
                    int(image.shape[0] * scale_factor),
                ),
            )
    


class ImageEnhancementProcessor:
    """Comprehensive image processor with upscaling, sharpening, and other enhancements."""

    def __init__(
        self,
        upscaling_enabled: bool = False,
        scale_factor: float = 2.0,
        upscaling_method: UpscalingMethod = UpscalingMethod.BICUBIC,
        sharpening_enabled: bool = False,
        sharpening_method: SharpeningMethod = SharpeningMethod.UNSHARP_MASK,
        sharpening_strength: float = 1.0,
        denoise_enabled: bool = False,
        denoise_method: DenoiseMethod = DenoiseMethod.TV_CHAMBOLLE,
        denoise_strength: float = 1.0,
        deconvolve_enabled: bool = False,
        deconvolve_strength: float = 0.5,
        deconvolve_psf_size: float = 2.0,
        processing_order: list[str] = None,
    ):
        self.upscaling_enabled = upscaling_enabled
        self.scale_factor = scale_factor
        self.upscaling_method = upscaling_method
        self.sharpening_enabled = sharpening_enabled
        self.sharpening_method = sharpening_method
        self.sharpening_strength = sharpening_strength
        self.denoise_enabled = denoise_enabled
        self.denoise_method = denoise_method
        self.denoise_strength = denoise_strength
        self.deconvolve_enabled = deconvolve_enabled
        self.deconvolve_strength = deconvolve_strength
        self.deconvolve_psf_size = deconvolve_psf_size
        self.processing_order = processing_order or ["upscaling", "denoise", "deconvolve", "sharpening"]
        self.upscaler = ImageUpscaler()

    def process(self, image: np.ndarray) -> Optional[np.ndarray]:
        """
        Process image with comprehensive enhancements.

        Args:
            image: Input image array

        Returns:
            Processed image with all enabled enhancements applied
        """
        if image is None:
            return None

        logging.trace(f"ImageEnhancementProcessor.process() starting with image shape: {image.shape}")
        logging.trace(f"Enhancement settings: denoise={self.denoise_enabled}, deconvolve={self.deconvolve_enabled}, sharpen={self.sharpening_enabled}, upscale={self.upscaling_enabled}")
        logging.trace(f"Processing order: {self.processing_order}")
        
        processed_image = image.copy()
        
        # Apply enhancements in custom order
        for step in self.processing_order:
            if step == "upscaling" and self.upscaling_enabled and self.scale_factor > 1.0:
                logging.trace(f"Applying upscaling: method={self.upscaling_method}, scale_factor={self.scale_factor}")
                processed_image = self.upscaler.upscale(
                    processed_image, 
                    scale_factor=self.scale_factor, 
                    method=self.upscaling_method
                )
                logging.trace(f"Upscaling completed, output shape: {processed_image.shape}")
            
            elif step == "denoise" and self.denoise_enabled:
                logging.trace(f"Applying denoising: method={self.denoise_method}, strength={self.denoise_strength}")
                processed_image = self.upscaler.denoise_image(
                    processed_image,
                    method=self.denoise_method,
                    strength=self.denoise_strength
                )
                logging.trace(f"Denoising completed, output shape: {processed_image.shape}")
            
            elif step == "deconvolve" and self.deconvolve_enabled:
                logging.trace(f"Applying deconvolution: strength={self.deconvolve_strength}, psf_size={self.deconvolve_psf_size}")
                processed_image = self._apply_deconvolution(
                    processed_image,
                    strength=self.deconvolve_strength,
                    psf_size=self.deconvolve_psf_size
                )
                logging.trace(f"Deconvolution completed, output shape: {processed_image.shape}")
            
            elif step == "sharpening" and self.sharpening_enabled:
                logging.trace(f"Applying sharpening: method={self.sharpening_method}, strength={self.sharpening_strength}")
                processed_image = self.upscaler.sharpen_image(
                    processed_image, 
                    method=self.sharpening_method, 
                    strength=self.sharpening_strength
                )
                logging.trace(f"Sharpening completed, output shape: {processed_image.shape}")

        logging.trace(f"ImageEnhancementProcessor.process() completed, final shape: {processed_image.shape}")
        return processed_image

    def _apply_deconvolution(self, image: np.ndarray, strength: float, psf_size: float) -> np.ndarray:
        """Apply deconvolution using Richardson-Lucy algorithm."""
        try:
            from skimage import restoration
            
            logging.trace(f"Deconvolution input: shape={image.shape}, dtype={image.dtype}, range=[{np.min(image):.3f}, {np.max(image):.3f}]")
            
            # Determine input format and convert appropriately
            input_range_255 = np.max(image) > 1.0
            
            if input_range_255:
                # Input is in [0, 255] range
                working_image = image.astype(np.float32) / 255.0
            else:
                # Input is already in [0, 1] range
                working_image = image.astype(np.float32)
            
            # Create a Gaussian PSF based on psf_size
            psf_radius = max(2, int(psf_size * 2))  # Ensure minimum radius
            y, x = np.ogrid[-psf_radius:psf_radius+1, -psf_radius:psf_radius+1]
            
            # Create Gaussian PSF with proper sigma
            sigma = psf_size / 2.355  # Convert FWHM to sigma
            psf = np.exp(-(x*x + y*y) / (2.0 * sigma**2))
            psf = psf / psf.sum()
            
            logging.trace(f"PSF: size={psf.shape}, sigma={sigma:.3f}, sum={psf.sum():.6f}")
            
            # Limit iterations to prevent over-deconvolution
            iterations = max(1, min(5, int(strength * 10)))  # Cap at 5 iterations
            
            logging.trace(f"Deconvolution parameters: iterations={iterations}, strength={strength}")
            
            # Apply deconvolution to each channel
            if len(working_image.shape) == 3:
                deconvolved = np.zeros_like(working_image)
                for i in range(working_image.shape[2]):
                    channel = working_image[:, :, i]
                    # Ensure channel has some variation to avoid division by zero
                    if np.std(channel) > 1e-6:
                        deconvolved[:, :, i] = restoration.richardson_lucy(
                            channel, psf, num_iter=iterations, clip=False
                        )
                    else:
                        deconvolved[:, :, i] = channel
            else:
                if np.std(working_image) > 1e-6:
                    deconvolved = restoration.richardson_lucy(
                        working_image, psf, num_iter=iterations, clip=False
                    )
                else:
                    deconvolved = working_image
            
            # Ensure output is in proper range and convert back to input format
            deconvolved = np.clip(deconvolved, 0, 1)
            
            if input_range_255:
                # Convert back to [0, 255] range
                deconvolved = (deconvolved * 255.0).astype(image.dtype)
            else:
                # Keep in [0, 1] range
                deconvolved = deconvolved.astype(image.dtype)
            
            logging.trace(f"Deconvolution output: shape={deconvolved.shape}, dtype={deconvolved.dtype}, range=[{np.min(deconvolved):.3f}, {np.max(deconvolved):.3f}]")
            logging.trace(f"Applied deconvolution: strength={strength}, psf_size={psf_size}, iterations={iterations}")
            
            return deconvolved
            
        except Exception as e:
            logging.error(f"Deconvolution failed: {e}")
            # Return original image if deconvolution fails
            return image

    def set_upscaling_params(
        self,
        enabled: bool,
        scale_factor: float = 2.0,
        method: UpscalingMethod = UpscalingMethod.BICUBIC,
    ):
        """Update upscaling parameters."""
        self.upscaling_enabled = enabled
        self.scale_factor = scale_factor
        self.upscaling_method = method
    
    def set_sharpening_params(
        self,
        enabled: bool,
        method: SharpeningMethod = SharpeningMethod.UNSHARP_MASK,
        strength: float = 1.0,
    ):
        """Update sharpening parameters."""
        self.sharpening_enabled = enabled
        self.sharpening_method = method
        self.sharpening_strength = max(0.0, min(2.0, strength))  # Clamp to safe range
    
    def set_deconvolve_params(
        self,
        enabled: bool,
        strength: float = 0.5,
        psf_size: float = 2.0,
    ):
        """Update deconvolution parameters."""
        self.deconvolve_enabled = enabled
        self.deconvolve_strength = max(0.0, min(1.0, strength))  # Clamp to safe range
        self.deconvolve_psf_size = max(0.5, min(10.0, psf_size))  # Clamp PSF size
    
    def get_enhancement_settings(self) -> dict:
        """Get current enhancement settings."""
        return {
            "upscaling_enabled": self.upscaling_enabled,
            "scale_factor": self.scale_factor,
            "upscaling_method": self.upscaling_method.value,
            "sharpening_enabled": self.sharpening_enabled,
            "sharpening_method": self.sharpening_method.value,
            "sharpening_strength": self.sharpening_strength,
            "deconvolve_enabled": self.deconvolve_enabled,
            "deconvolve_strength": self.deconvolve_strength,
            "deconvolve_psf_size": self.deconvolve_psf_size,
        }


# Keep the old class for backward compatibility
class UpscalingProcessor(ImageEnhancementProcessor):
    """Backward compatibility wrapper for ImageEnhancementProcessor."""
    
    def __init__(self, enabled: bool = False, scale_factor: float = 2.0, method: UpscalingMethod = UpscalingMethod.BICUBIC):
        super().__init__(
            upscaling_enabled=enabled,
            scale_factor=scale_factor,
            upscaling_method=method
        )
        # For backward compatibility
        self.enabled = self.upscaling_enabled
        self.method = self.upscaling_method
    
    def set_upscaling_params(self, enabled: bool, scale_factor: float = 2.0, method: UpscalingMethod = UpscalingMethod.BICUBIC):
        """Backward compatibility method."""
        super().set_upscaling_params(enabled, scale_factor, method)
        self.enabled = self.upscaling_enabled
        self.method = self.upscaling_method
