src/modules/color_preservation.py

"""
Color Preservation Module
========================

Utilities to ensure perfect color preservation during image enhancement.
Only sharpness, clarity, and focus should be improved while maintaining
the exact original colors.
"""

import cv2
import numpy as np
from typing import Tuple, Optional
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ColorPreserver:
    """Utilities for preserving colors during image enhancement"""
    
    @staticmethod
    def preserve_colors_during_enhancement(original: np.ndarray, 
                                         enhanced: np.ndarray,
                                         preservation_strength: float = 0.8) -> np.ndarray:
        """
        Preserve original colors while keeping enhancement benefits
        
        Args:
            original: Original image (BGR)
            enhanced: Enhanced image (BGR) 
            preservation_strength: How much to preserve original colors (0-1)
            
        Returns:
            np.ndarray: Color-preserved enhanced image
        """
        try:
            # Convert to LAB color space for better color/brightness separation
            original_lab = cv2.cvtColor(original, cv2.COLOR_BGR2LAB)
            enhanced_lab = cv2.cvtColor(enhanced, cv2.COLOR_BGR2LAB)
            
            # Split LAB channels
            orig_l, orig_a, orig_b = cv2.split(original_lab)
            enh_l, enh_a, enh_b = cv2.split(enhanced_lab)
            
            # Keep enhanced brightness (L channel) but preserve original colors (A, B channels)
            preserved_a = (preservation_strength * orig_a + 
                          (1 - preservation_strength) * enh_a).astype(np.uint8)
            preserved_b = (preservation_strength * orig_b + 
                          (1 - preservation_strength) * enh_b).astype(np.uint8)
            
            # Combine preserved colors with enhanced brightness
            result_lab = cv2.merge([enh_l, preserved_a, preserved_b])
            
            # Convert back to BGR
            result = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
            
            logger.info("Color preservation applied")
            return result
            
        except Exception as e:
            logger.error(f"Error in color preservation: {e}")
            return enhanced
    
    @staticmethod
    def enhance_sharpness_only(image: np.ndarray, 
                              sharpening_strength: float = 0.3) -> np.ndarray:
        """
        Enhance only sharpness without affecting colors
        
        Args:
            image: Input image (BGR)
            sharpening_strength: Sharpening strength (0-1)
            
        Returns:
            np.ndarray: Sharpness-enhanced image with preserved colors
        """
        try:
            # Convert to float for precision
            img_float = image.astype(np.float64)
            
            # Create a subtle sharpening kernel
            kernel = np.array([[-0.05, -0.1, -0.05],
                              [-0.1,  1.4,  -0.1],
                              [-0.05, -0.1, -0.05]]) * sharpening_strength
            
            # Add identity for original preservation
            kernel[1, 1] += (1 - sharpening_strength)
            
            # Apply sharpening filter
            sharpened = cv2.filter2D(img_float, -1, kernel)
            
            # Ensure no clipping artifacts that change colors
            result = np.clip(sharpened, 0, 255).astype(np.uint8)
            
            return result
            
        except Exception as e:
            logger.error(f"Error in sharpness-only enhancement: {e}")
            return image
    
    @staticmethod
    def accurate_unsharp_masking(image: np.ndarray,
                               sigma: float = 1.0,
                               amount: float = 0.5) -> np.ndarray:
        """
        Apply unsharp masking with perfect color preservation
        
        Args:
            image: Input image (BGR)
            sigma: Gaussian blur sigma for mask
            amount: Sharpening amount
            
        Returns:
            np.ndarray: Sharpened image with preserved colors
        """
        try:
            # Work in high precision
            img_float = image.astype(np.float64)
            
            # Create Gaussian blur
            blurred = cv2.GaussianBlur(img_float, (0, 0), sigma)
            
            # Create unsharp mask
            mask = img_float - blurred
            
            # Apply mask with careful amount control
            sharpened = img_float + amount * mask
            
            # Careful clipping to preserve color accuracy
            result = np.clip(sharpened, 0, 255)
            result = np.round(result).astype(np.uint8)
            
            return result
            
        except Exception as e:
            logger.error(f"Error in accurate unsharp masking: {e}")
            return image
    
    @staticmethod
    def convert_for_display(image_bgr: np.ndarray) -> np.ndarray:
        """
        Convert BGR image to RGB for proper display in Streamlit
        
        Args:
            image_bgr: Image in BGR format
            
        Returns:
            np.ndarray: Image in RGB format for display
        """
        try:
            if len(image_bgr.shape) == 3 and image_bgr.shape[2] == 3:
                return cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
            return image_bgr
        except Exception as e:
            logger.error(f"Error converting for display: {e}")
            return image_bgr
    
    @staticmethod
    def validate_color_preservation(original: np.ndarray, 
                                  processed: np.ndarray,
                                  tolerance: float = 5.0) -> dict:
        """
        Validate that colors are preserved during processing
        
        Args:
            original: Original image
            processed: Processed image  
            tolerance: Acceptable color difference
            
        Returns:
            dict: Validation results
        """
        try:
            # Convert to LAB for perceptual color comparison
            orig_lab = cv2.cvtColor(original, cv2.COLOR_BGR2LAB)
            proc_lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            
            # Calculate color differences (A and B channels only)
            diff_a = np.mean(np.abs(orig_lab[:, :, 1].astype(np.float32) - 
                                  proc_lab[:, :, 1].astype(np.float32)))
            diff_b = np.mean(np.abs(orig_lab[:, :, 2].astype(np.float32) - 
                                  proc_lab[:, :, 2].astype(np.float32)))
            
            color_diff = (diff_a + diff_b) / 2.0
            
            return {
                'color_difference': float(color_diff),
                'colors_preserved': color_diff <= tolerance,
                'a_channel_diff': float(diff_a),
                'b_channel_diff': float(diff_b),
                'tolerance_used': tolerance
            }
            
        except Exception as e:
            logger.error(f"Error validating color preservation: {e}")
            return {'colors_preserved': False, 'error': str(e)}

# Convenience functions for easy use
def preserve_colors(original: np.ndarray, enhanced: np.ndarray) -> np.ndarray:
    """Preserve colors from original in enhanced image"""
    return ColorPreserver.preserve_colors_during_enhancement(original, enhanced)

def sharpen_only(image: np.ndarray, strength: float = 0.3) -> np.ndarray:
    """Sharpen image without changing colors"""
    return ColorPreserver.enhance_sharpness_only(image, strength)

def display_convert(image_bgr: np.ndarray) -> np.ndarray:
    """Convert BGR to RGB for display"""
    return ColorPreserver.convert_for_display(image_bgr)