processing/effects.py

"""
Visual effects and enhancements for BackgroundFX Pro.
Implements professional-grade effects for background replacement.
"""

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from typing import Dict, List, Optional, Tuple, Union
from dataclasses import dataclass
from enum import Enum
import logging
from scipy.ndimage import gaussian_filter, map_coordinates

from ..utils.logger import setup_logger
from ..utils.device import DeviceManager
from ..core.quality import QualityAnalyzer

logger = setup_logger(__name__)


class EffectType(Enum):
    """Available effect types."""
    BLUR = "blur"
    BOKEH = "bokeh"
    COLOR_SHIFT = "color_shift"
    LIGHT_WRAP = "light_wrap"
    SHADOW = "shadow"
    REFLECTION = "reflection"
    GLOW = "glow"
    CHROMATIC_ABERRATION = "chromatic_aberration"
    VIGNETTE = "vignette"
    FILM_GRAIN = "film_grain"
    MOTION_BLUR = "motion_blur"
    DEPTH_OF_FIELD = "depth_of_field"


@dataclass
class EffectConfig:
    """Configuration for visual effects."""
    blur_strength: float = 15.0
    bokeh_size: int = 21
    bokeh_brightness: float = 1.5
    light_wrap_intensity: float = 0.3
    light_wrap_width: int = 10
    shadow_opacity: float = 0.5
    shadow_blur: float = 10.0
    shadow_offset: Tuple[int, int] = (5, 5)
    glow_intensity: float = 0.5
    glow_radius: int = 20
    chromatic_shift: float = 2.0
    vignette_strength: float = 0.3
    grain_intensity: float = 0.1
    motion_blur_angle: float = 0.0
    motion_blur_size: int = 15


class BackgroundEffects:
    """Apply effects to background images."""
    
    def __init__(self, config: Optional[EffectConfig] = None):
        self.config = config or EffectConfig()
        self.device_manager = DeviceManager()
        
    def apply_blur(self, image: np.ndarray, 
                  strength: Optional[float] = None,
                  mask: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Apply Gaussian blur to image.
        
        Args:
            image: Input image
            strength: Blur strength
            mask: Optional mask for selective blur
            
        Returns:
            Blurred image
        """
        strength = strength or self.config.blur_strength
        
        if strength <= 0:
            return image
        
        # Calculate kernel size (must be odd)
        kernel_size = int(strength * 2) + 1
        
        # Apply blur
        blurred = cv2.GaussianBlur(image, (kernel_size, kernel_size), strength)
        
        # Apply mask if provided
        if mask is not None:
            mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
            if mask_3ch.max() > 1:
                mask_3ch = mask_3ch / 255.0
            
            blurred = image * (1 - mask_3ch) + blurred * mask_3ch
            blurred = blurred.astype(np.uint8)
        
        return blurred
    
    def apply_bokeh(self, image: np.ndarray,
                   depth_map: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Apply bokeh effect to simulate depth of field.
        
        Args:
            image: Input image
            depth_map: Optional depth map for varying blur
            
        Returns:
            Image with bokeh effect
        """
        h, w = image.shape[:2]
        
        # Create depth map if not provided
        if depth_map is None:
            # Simple radial depth map
            center_x, center_y = w // 2, h // 2
            Y, X = np.ogrid[:h, :w]
            dist = np.sqrt((X - center_x)**2 + (Y - center_y)**2)
            depth_map = dist / dist.max()
        
        # Normalize depth map
        if depth_map.max() > 1:
            depth_map = depth_map / 255.0
        
        # Create bokeh kernel
        kernel_size = self.config.bokeh_size
        kernel = self._create_bokeh_kernel(kernel_size)
        
        # Apply varying blur based on depth
        result = np.zeros_like(image, dtype=np.float32)
        
        # Create multiple blur levels
        blur_levels = 5
        for i in range(blur_levels):
            blur_strength = (i + 1) * (kernel_size // blur_levels)
            
            if blur_strength > 0:
                blurred = cv2.filter2D(image, -1, kernel[:blur_strength, :blur_strength])
            else:
                blurred = image
            
            # Create mask for this depth level
            depth_min = i / blur_levels
            depth_max = (i + 1) / blur_levels
            mask = ((depth_map >= depth_min) & (depth_map < depth_max)).astype(np.float32)
            
            # Expand mask to 3 channels
            mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
            
            # Accumulate result
            result += blurred * mask_3ch
        
        # Add bokeh highlights
        result = self._add_bokeh_highlights(result, depth_map)
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def _create_bokeh_kernel(self, size: int) -> np.ndarray:
        """Create hexagonal bokeh kernel."""
        kernel = np.zeros((size, size), dtype=np.float32)
        center = size // 2
        radius = center - 1
        
        # Create hexagonal shape
        for i in range(size):
            for j in range(size):
                x, y = i - center, j - center
                # Hexagon equation
                if abs(x) <= radius and abs(y) <= radius * np.sqrt(3) / 2:
                    if abs(y) <= (radius * np.sqrt(3) / 2 - abs(x) * np.sqrt(3) / 2):
                        kernel[i, j] = 1.0
        
        # Normalize
        kernel /= kernel.sum()
        
        return kernel
    
    def _add_bokeh_highlights(self, image: np.ndarray,
                             depth_map: np.ndarray) -> np.ndarray:
        """Add bright bokeh spots to out-of-focus areas."""
        # Extract bright spots
        gray = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_BGR2GRAY)
        _, bright_mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
        
        # Dilate bright spots in blurred areas
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
        bright_mask = cv2.dilate(bright_mask, kernel, iterations=2)
        
        # Apply only to out-of-focus areas
        bright_mask = (bright_mask * depth_map).astype(np.uint8)
        
        # Create glow effect
        glow = cv2.GaussianBlur(bright_mask, (21, 21), 10)
        glow = cv2.cvtColor(glow, cv2.COLOR_GRAY2BGR) / 255.0
        
        # Add glow to image
        result = image + glow * self.config.bokeh_brightness * 50
        
        return result
    
    def apply_light_wrap(self, foreground: np.ndarray,
                        background: np.ndarray,
                        mask: np.ndarray) -> np.ndarray:
        """
        Apply light wrap effect for better compositing.
        
        Args:
            foreground: Foreground image
            background: Background image
            mask: Foreground mask
            
        Returns:
            Foreground with light wrap
        """
        # Ensure mask is single channel
        if len(mask.shape) == 3:
            mask = mask[:, :, 0]
        
        # Normalize mask
        if mask.max() > 1:
            mask = mask / 255.0
        
        # Create edge mask
        kernel = np.ones((self.config.light_wrap_width, self.config.light_wrap_width), np.uint8)
        dilated_mask = cv2.dilate(mask, kernel, iterations=1)
        edge_mask = dilated_mask - mask
        
        # Blur the background
        blurred_bg = cv2.GaussianBlur(background, (21, 21), 10)
        
        # Extract light from background
        bg_light = blurred_bg * edge_mask[:, :, np.newaxis]
        
        # Add light wrap to foreground edges
        mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
        wrapped = foreground + bg_light * self.config.light_wrap_intensity
        
        return np.clip(wrapped, 0, 255).astype(np.uint8)
    
    def add_shadow(self, image: np.ndarray,
                  mask: np.ndarray,
                  ground_plane: Optional[float] = None) -> np.ndarray:
        """
        Add realistic shadow to composited image.
        
        Args:
            image: Background image
            mask: Object mask
            ground_plane: Y-coordinate of ground plane
            
        Returns:
            Image with shadow
        """
        h, w = image.shape[:2]
        
        if ground_plane is None:
            ground_plane = h * 0.9  # Default near bottom
        
        # Create shadow mask
        shadow_mask = mask.copy()
        if len(shadow_mask.shape) == 3:
            shadow_mask = shadow_mask[:, :, 0]
        
        # Transform shadow (simple perspective)
        offset_x, offset_y = self.config.shadow_offset
        
        # Create transformation matrix
        src_points = np.float32([[0, 0], [w, 0], [0, h], [w, h]])
        dst_points = np.float32([
            [offset_x, offset_y],
            [w + offset_x, offset_y],
            [-offset_x * 2, h],
            [w + offset_x * 2, h]
        ])
        
        matrix = cv2.getPerspectiveTransform(src_points, dst_points)
        shadow_mask = cv2.warpPerspective(shadow_mask, matrix, (w, h))
        
        # Blur shadow
        blur_size = int(self.config.shadow_blur) * 2 + 1
        shadow_mask = cv2.GaussianBlur(shadow_mask, (blur_size, blur_size), 
                                      self.config.shadow_blur)
        
        # Clip shadow to ground plane
        shadow_mask[:int(ground_plane), :] = 0
        
        # Normalize and apply opacity
        if shadow_mask.max() > 0:
            shadow_mask = shadow_mask / shadow_mask.max()
        shadow_mask *= self.config.shadow_opacity
        
        # Darken image where shadow falls
        shadow_color = np.array([0, 0, 0], dtype=np.float32)
        shadow_mask_3ch = np.repeat(shadow_mask[:, :, np.newaxis], 3, axis=2)
        
        result = image * (1 - shadow_mask_3ch) + shadow_color * shadow_mask_3ch
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def add_reflection(self, image: np.ndarray,
                      mask: np.ndarray,
                      reflection_strength: float = 0.3) -> np.ndarray:
        """
        Add reflection effect for glossy surfaces.
        
        Args:
            image: Input image
            mask: Object mask
            reflection_strength: Reflection opacity
            
        Returns:
            Image with reflection
        """
        h, w = image.shape[:2]
        
        # Extract object using mask
        if len(mask.shape) == 2:
            mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
        else:
            mask_3ch = mask
        
        if mask_3ch.max() > 1:
            mask_3ch = mask_3ch / 255.0
        
        object_only = image * mask_3ch
        
        # Flip vertically for reflection
        reflection = cv2.flip(object_only, 0)
        
        # Create gradient for fade-out
        gradient = np.linspace(reflection_strength, 0, h)
        gradient = np.repeat(gradient[:, np.newaxis], w, axis=1)
        gradient = np.repeat(gradient[:, :, np.newaxis], 3, axis=2)
        
        # Apply gradient to reflection
        reflection = reflection * gradient
        
        # Add slight blur for realism
        reflection = cv2.GaussianBlur(reflection, (5, 5), 2)
        
        # Composite reflection below object
        result = image.copy()
        result = result + reflection
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def add_glow(self, image: np.ndarray,
                mask: Optional[np.ndarray] = None,
                color: Optional[Tuple[int, int, int]] = None) -> np.ndarray:
        """
        Add glow effect to image or masked region.
        
        Args:
            image: Input image
            mask: Optional mask for selective glow
            color: Glow color (BGR)
            
        Returns:
            Image with glow effect
        """
        if color is None:
            color = (255, 255, 255)  # White glow
        
        # Create glow source
        if mask is not None:
            if len(mask.shape) == 2:
                glow_source = np.zeros_like(image)
                for i in range(3):
                    glow_source[:, :, i] = mask * (color[i] / 255.0)
            else:
                glow_source = mask * np.array(color).reshape(1, 1, 3) / 255.0
        else:
            # Use bright parts of image
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            _, bright_mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
            glow_source = cv2.cvtColor(bright_mask, cv2.COLOR_GRAY2BGR)
        
        # Create multiple blur levels for glow
        glow = np.zeros_like(image, dtype=np.float32)
        
        for i in range(1, 4):
            blur_size = self.config.glow_radius * i
            kernel_size = blur_size * 2 + 1
            
            blurred = cv2.GaussianBlur(glow_source, (kernel_size, kernel_size), blur_size)
            glow += blurred / (i * 2)
        
        # Normalize and apply intensity
        if glow.max() > 0:
            glow = glow / glow.max()
        glow *= self.config.glow_intensity * 255
        
        # Add glow to original image
        result = image.astype(np.float32) + glow
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def chromatic_aberration(self, image: np.ndarray,
                           shift: Optional[float] = None) -> np.ndarray:
        """
        Apply chromatic aberration effect.
        
        Args:
            image: Input image
            shift: Pixel shift amount
            
        Returns:
            Image with chromatic aberration
        """
        shift = shift or self.config.chromatic_shift
        h, w = image.shape[:2]
        
        # Split channels
        b, g, r = cv2.split(image)
        
        # Create radial shift
        center_x, center_y = w // 2, h // 2
        
        # Shift red channel outward
        M_r = np.float32([[1 + shift/w, 0, -shift], [0, 1 + shift/h, -shift]])
        r_shifted = cv2.warpAffine(r, M_r, (w, h))
        
        # Shift blue channel inward
        M_b = np.float32([[1 - shift/w, 0, shift], [0, 1 - shift/h, shift]])
        b_shifted = cv2.warpAffine(b, M_b, (w, h))
        
        # Merge channels
        result = cv2.merge([b_shifted, g, r_shifted])
        
        return result
    
    def add_vignette(self, image: np.ndarray,
                    strength: Optional[float] = None) -> np.ndarray:
        """
        Add vignette effect to image.
        
        Args:
            image: Input image
            strength: Vignette strength (0-1)
            
        Returns:
            Image with vignette
        """
        strength = strength or self.config.vignette_strength
        h, w = image.shape[:2]
        
        # Create radial gradient
        center_x, center_y = w // 2, h // 2
        Y, X = np.ogrid[:h, :w]
        
        # Calculate distance from center
        dist = np.sqrt((X - center_x)**2 + (Y - center_y)**2)
        max_dist = np.sqrt(center_x**2 + center_y**2)
        
        # Normalize and create vignette mask
        vignette = 1 - (dist / max_dist) * strength
        vignette = np.clip(vignette, 0, 1)
        
        # Apply vignette
        vignette_3ch = np.repeat(vignette[:, :, np.newaxis], 3, axis=2)
        result = image * vignette_3ch
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def add_film_grain(self, image: np.ndarray,
                      intensity: Optional[float] = None) -> np.ndarray:
        """
        Add film grain effect to image.
        
        Args:
            image: Input image
            intensity: Grain intensity
            
        Returns:
            Image with film grain
        """
        intensity = intensity or self.config.grain_intensity
        
        # Generate grain
        h, w = image.shape[:2]
        grain = np.random.randn(h, w, 3) * intensity * 255
        
        # Add grain to image
        result = image.astype(np.float32) + grain
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def motion_blur(self, image: np.ndarray,
                   angle: Optional[float] = None,
                   size: Optional[int] = None) -> np.ndarray:
        """
        Apply directional motion blur.
        
        Args:
            image: Input image
            angle: Blur angle in degrees
            size: Blur kernel size
            
        Returns:
            Motion blurred image
        """
        angle = angle or self.config.motion_blur_angle
        size = size or self.config.motion_blur_size
        
        # Create motion blur kernel
        kernel = np.zeros((size, size))
        kernel[int((size-1)/2), :] = np.ones(size)
        kernel = kernel / size
        
        # Rotate kernel
        M = cv2.getRotationMatrix2D((size/2, size/2), angle, 1)
        kernel = cv2.warpAffine(kernel, M, (size, size))
        
        # Apply kernel
        result = cv2.filter2D(image, -1, kernel)
        
        return result


class CompositeEffects:
    """Advanced compositing effects."""
    
    def __init__(self):
        self.logger = setup_logger(f"{__name__}.CompositeEffects")
        self.bg_effects = BackgroundEffects()
        
    def smart_composite(self, foreground: np.ndarray,
                       background: np.ndarray,
                       mask: np.ndarray,
                       effects: List[EffectType]) -> np.ndarray:
        """
        Apply smart compositing with multiple effects.
        
        Args:
            foreground: Foreground image
            background: Background image
            mask: Alpha mask
            effects: List of effects to apply
            
        Returns:
            Composited image with effects
        """
        result = background.copy()
        
        # Ensure mask is proper format
        if len(mask.shape) == 2:
            mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
        else:
            mask_3ch = mask
        
        if mask_3ch.max() > 1:
            mask_3ch = mask_3ch / 255.0
        
        # Apply background effects
        for effect in effects:
            if effect == EffectType.BLUR:
                result = self.bg_effects.apply_blur(result, mask=1-mask_3ch[:,:,0])
            elif effect == EffectType.BOKEH:
                result = self.bg_effects.apply_bokeh(result)
            elif effect == EffectType.VIGNETTE:
                result = self.bg_effects.add_vignette(result)
        
        # Apply light wrap before compositing
        if EffectType.LIGHT_WRAP in effects:
            foreground = self.bg_effects.apply_light_wrap(
                foreground, result, mask_3ch[:,:,0]
            )
        
        # Composite foreground
        result = result * (1 - mask_3ch) + foreground * mask_3ch
        result = result.astype(np.uint8)
        
        # Apply post-composite effects
        if EffectType.SHADOW in effects:
            result = self.bg_effects.add_shadow(result, mask_3ch[:,:,0])
        
        if EffectType.REFLECTION in effects:
            result = self.bg_effects.add_reflection(result, mask_3ch[:,:,0])
        
        if EffectType.GLOW in effects:
            result = self.bg_effects.add_glow(result, mask_3ch[:,:,0])
        
        # Apply final touches
        if EffectType.CHROMATIC_ABERRATION in effects:
            result = self.bg_effects.chromatic_aberration(result)
        
        if EffectType.FILM_GRAIN in effects:
            result = self.bg_effects.add_film_grain(result)
        
        return result
    
    def color_harmonization(self, foreground: np.ndarray,
                          background: np.ndarray,
                          mask: np.ndarray,
                          strength: float = 0.3) -> np.ndarray:
        """
        Harmonize colors between foreground and background.
        
        Args:
            foreground: Foreground image
            background: Background image
            mask: Foreground mask
            strength: Harmonization strength
            
        Returns:
            Color-harmonized foreground
        """
        # Calculate background color statistics
        bg_mean = np.mean(background, axis=(0, 1))
        bg_std = np.std(background, axis=(0, 1))
        
        # Calculate foreground color statistics
        fg_mean = np.mean(foreground, axis=(0, 1))
        fg_std = np.std(foreground, axis=(0, 1))
        
        # Adjust foreground colors
        result = foreground.astype(np.float32)
        
        for i in range(3):  # For each color channel
            # Normalize foreground
            result[:, :, i] = (result[:, :, i] - fg_mean[i]) / (fg_std[i] + 1e-6)
            
            # Apply background statistics
            result[:, :, i] = result[:, :, i] * (bg_std[i] * strength + fg_std[i] * (1 - strength))
            result[:, :, i] += bg_mean[i] * strength + fg_mean[i] * (1 - strength)
        
        return np.clip(result, 0, 255).astype(np.uint8)