#!/usr/bin/env python3
"""
CDL (Color Decision List) based edge smoothing for SegMatch
"""

import numpy as np
from typing import Tuple, Optional
from PIL import Image
import cv2


def calculate_cdl_params_face_only(source: np.ndarray, target: np.ndarray, 
                                  source_face_mask: np.ndarray, target_face_mask: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Calculate CDL parameters using only face pixels for focused accuracy.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        source_face_mask (np.ndarray): Binary mask of face in source image
        target_face_mask (np.ndarray): Binary mask of face in target image
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    epsilon = 1e-6
    
    # Extract face pixels only
    source_face_pixels = source[source_face_mask > 0.5]
    target_face_pixels = target[target_face_mask > 0.5]
    
    # Ensure we have enough face pixels
    if len(source_face_pixels) < 100 or len(target_face_pixels) < 100:
        # Fallback to simple calculation if not enough face pixels
        return calculate_cdl_params_simple(source, target)
    
    slopes = []
    offsets = []
    powers = []
    
    for channel in range(3):
        src_channel = source_face_pixels[:, channel]
        tgt_channel = target_face_pixels[:, channel]
        
        # Use robust percentiles for face pixels
        percentiles = [10, 25, 50, 75, 90]
        src_percentiles = np.percentile(src_channel, percentiles)
        tgt_percentiles = np.percentile(tgt_channel, percentiles)
        
        # Calculate slope from face pixel range
        src_range = src_percentiles[4] - src_percentiles[0]  # 90th - 10th
        tgt_range = tgt_percentiles[4] - tgt_percentiles[0]
        slope = tgt_range / (src_range + epsilon)
        
        # Calculate offset using face median
        src_median = src_percentiles[2]
        tgt_median = tgt_percentiles[2]
        offset = tgt_median - (src_median * slope)
        
        # Calculate gamma from face brightness relationship
        src_mean = np.mean(src_channel)
        tgt_mean = np.mean(tgt_channel)
        
        if src_mean > epsilon:
            power = np.log(tgt_mean + epsilon) / np.log(src_mean + epsilon)
            power = np.clip(power, 0.3, 3.0)
        else:
            power = 1.0
        
        slopes.append(slope)
        offsets.append(offset)
        powers.append(power)
    
    return np.array(slopes), np.array(offsets), np.array(powers)


def calculate_cdl_params_simple(source: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Simple CDL calculation as fallback method.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    epsilon = 1e-6
    
    # Calculate mean and standard deviation for each RGB channel
    source_mean = np.mean(source, axis=(0, 1))
    source_std = np.std(source, axis=(0, 1))
    target_mean = np.mean(target, axis=(0, 1))
    target_std = np.std(target, axis=(0, 1))
    
    # Calculate slope (gain)
    slope = target_std / (source_std + epsilon)
    
    # Calculate offset
    offset = target_mean - (source_mean * slope)
    
    # Set power to neutral
    power = np.ones(3)
    
    return slope, offset, power


def calculate_cdl_params_histogram(source: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Calculate CDL parameters using histogram matching approach.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    epsilon = 1e-6
    
    # Convert to 0-255 range for histogram calculation
    source_255 = (source * 255).astype(np.uint8)
    target_255 = (target * 255).astype(np.uint8)
    
    slopes = []
    offsets = []
    powers = []
    
    for channel in range(3):
        # Calculate histograms
        hist_source = cv2.calcHist([source_255], [channel], None, [256], [0, 256])
        hist_target = cv2.calcHist([target_255], [channel], None, [256], [0, 256])
        
        # Calculate cumulative distributions
        cdf_source = np.cumsum(hist_source) / np.sum(hist_source)
        cdf_target = np.cumsum(hist_target) / np.sum(hist_target)
        
        # Find percentile mappings
        p25_src = np.percentile(source[:, :, channel], 25)
        p75_src = np.percentile(source[:, :, channel], 75)
        p25_tgt = np.percentile(target[:, :, channel], 25)
        p75_tgt = np.percentile(target[:, :, channel], 75)
        
        # Calculate slope from percentile mapping
        slope = (p75_tgt - p25_tgt) / (p75_src - p25_src + epsilon)
        
        # Calculate offset
        median_src = np.percentile(source[:, :, channel], 50)
        median_tgt = np.percentile(target[:, :, channel], 50)
        offset = median_tgt - (median_src * slope)
        
        # Estimate power/gamma from the histogram shape
        mean_src = np.mean(source[:, :, channel])
        mean_tgt = np.mean(target[:, :, channel])
        if mean_src > epsilon:
            power = np.log(mean_tgt + epsilon) / np.log(mean_src + epsilon)
            power = np.clip(power, 0.1, 10.0)  # Reasonable gamma range
        else:
            power = 1.0
        
        slopes.append(slope)
        offsets.append(offset)
        powers.append(power)
    
    return np.array(slopes), np.array(offsets), np.array(powers)


def calculate_cdl_params_mask_aware(source: np.ndarray, target: np.ndarray, 
                                   changed_mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Calculate CDL parameters focusing only on changed regions.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        changed_mask (np.ndarray, optional): Binary mask of changed regions
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    if changed_mask is not None:
        # Only use pixels where changes occurred
        mask_bool = changed_mask > 0.5
        if np.sum(mask_bool) > 100:  # Ensure enough pixels
            source_masked = source[mask_bool]
            target_masked = target[mask_bool]
            
            # Reshape back to have channel dimension
            source_masked = source_masked.reshape(-1, 3)
            target_masked = target_masked.reshape(-1, 3)
            
            # Calculate statistics on masked regions
            epsilon = 1e-6
            source_mean = np.mean(source_masked, axis=0)
            source_std = np.std(source_masked, axis=0)
            target_mean = np.mean(target_masked, axis=0)
            target_std = np.std(target_masked, axis=0)
            
            slope = target_std / (source_std + epsilon)
            offset = target_mean - (source_mean * slope)
            power = np.ones(3)
            
            return slope, offset, power
    
    # Fallback to simple method if mask is not useful
    return calculate_cdl_params_simple(source, target)


def calculate_cdl_params_lab(source: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Calculate CDL parameters in LAB color space for better perceptual matching.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    # Convert to LAB color space
    source_lab = cv2.cvtColor((source * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
    target_lab = cv2.cvtColor((target * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
    
    # Normalize LAB values
    source_lab[:, :, 0] /= 100.0  # L: 0-100 -> 0-1
    source_lab[:, :, 1] = (source_lab[:, :, 1] + 128) / 255.0  # A: -128-127 -> 0-1
    source_lab[:, :, 2] = (source_lab[:, :, 2] + 128) / 255.0  # B: -128-127 -> 0-1
    
    target_lab[:, :, 0] /= 100.0
    target_lab[:, :, 1] = (target_lab[:, :, 1] + 128) / 255.0
    target_lab[:, :, 2] = (target_lab[:, :, 2] + 128) / 255.0
    
    # Calculate CDL in LAB space
    epsilon = 1e-6
    source_mean = np.mean(source_lab, axis=(0, 1))
    source_std = np.std(source_lab, axis=(0, 1))
    target_mean = np.mean(target_lab, axis=(0, 1))
    target_std = np.std(target_lab, axis=(0, 1))
    
    slope_lab = target_std / (source_std + epsilon)
    offset_lab = target_mean - (source_mean * slope_lab)
    
    # Convert back to RGB approximation
    # This is a simplified conversion - for full accuracy we'd need to convert each pixel
    slope = np.array([slope_lab[0], slope_lab[1], slope_lab[2]])  # Rough mapping
    offset = np.array([offset_lab[0], offset_lab[1], offset_lab[2]])
    power = np.ones(3)
    
    return slope, offset, power


def calculate_cdl_params(source: np.ndarray, target: np.ndarray, 
                        source_path: str = None, target_path: str = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Calculate CDL parameters using simple mean/std matching - the most basic approach.
    
    Args:
        source (np.ndarray): Source image as numpy array (0-1 range)
        target (np.ndarray): Target image as numpy array (0-1 range)
        source_path (str, optional): Ignored - kept for compatibility
        target_path (str, optional): Ignored - kept for compatibility
        
    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: (slope, offset, power)
    """
    epsilon = 1e-6
    
    # Calculate simple mean and standard deviation for each RGB channel
    source_mean = np.mean(source, axis=(0, 1))
    source_std = np.std(source, axis=(0, 1))
    target_mean = np.mean(target, axis=(0, 1))
    target_std = np.std(target, axis=(0, 1))
    
    # Calculate slope (gain) from std ratio
    slope = target_std / (source_std + epsilon)
    
    # Calculate offset from mean difference
    offset = target_mean - (source_mean * slope)
    
    # Calculate simple gamma from brightness relationship
    power = []
    for channel in range(3):
        if source_mean[channel] > epsilon:
            gamma = np.log(target_mean[channel] + epsilon) / np.log(source_mean[channel] + epsilon)
            gamma = np.clip(gamma, 0.1, 10.0)  # Keep within reasonable bounds
        else:
            gamma = 1.0
        power.append(gamma)
    
    power = np.array(power)
    
    return slope, offset, power


def calculate_change_mask(original: np.ndarray, composited: np.ndarray, threshold: float = 0.05) -> np.ndarray:
    """Calculate a mask of significantly changed regions between original and composited images.
    
    Args:
        original (np.ndarray): Original image (0-1 range)
        composited (np.ndarray): Composited result (0-1 range)
        threshold (float): Threshold for detecting significant changes
        
    Returns:
        np.ndarray: Binary mask of changed regions
    """
    # Calculate per-pixel difference
    diff = np.sqrt(np.sum((composited - original) ** 2, axis=2))
    
    # Create binary mask where changes exceed threshold
    change_mask = (diff > threshold).astype(np.float32)
    
    # Apply morphological operations to clean up the mask
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    change_mask = cv2.morphologyEx(change_mask, cv2.MORPH_CLOSE, kernel)
    
    return change_mask


def calculate_channel_stats(array: np.ndarray) -> dict:
    """Calculate per-channel statistics for an image array.
    
    Args:
        array: Image array of shape (H, W, 3)
        
    Returns:
        dict: Dictionary containing mean, std, min, max for each channel
    """
    stats = {
        'mean': np.mean(array, axis=(0, 1)),
        'std': np.std(array, axis=(0, 1)),
        'min': np.min(array, axis=(0, 1)),
        'max': np.max(array, axis=(0, 1))
    }
    return stats


def apply_cdl_transform(image: np.ndarray, slope: np.ndarray, offset: np.ndarray, power: np.ndarray, 
                       factor: float = 0.3) -> np.ndarray:
    """Apply CDL transformation to an image.
    
    Args:
        image (np.ndarray): Input image (0-1 range)
        slope (np.ndarray): CDL slope parameters for each channel
        offset (np.ndarray): CDL offset parameters for each channel  
        power (np.ndarray): CDL power parameters for each channel
        factor (float): Blending factor (0.0 = no change, 1.0 = full transform)
        
    Returns:
        np.ndarray: Transformed image
    """
    # Apply CDL transform: out = ((in * slope) + offset) ** power
    transformed = np.power(np.maximum(image * slope + offset, 0), power)
    
    # Clamp to valid range
    transformed = np.clip(transformed, 0.0, 1.0)
    
    # Blend with original based on factor
    result = (1 - factor) * image + factor * transformed
    
    return result


def cdl_edge_smoothing(composited_image_path: str, original_image_path: str, factor: float = 0.3) -> Image.Image:
    """Apply CDL-based edge smoothing between composited result and original image.
    
    Args:
        composited_image_path (str): Path to the composited result image
        original_image_path (str): Path to the original target image
        factor (float): Smoothing strength (0.0 = no smoothing, 1.0 = full smoothing)
        
    Returns:
        Image.Image: Smoothed result image
    """
    # Load images
    composited_img = Image.open(composited_image_path).convert("RGB")
    original_img = Image.open(original_image_path).convert("RGB")
    
    # Ensure same dimensions
    if composited_img.size != original_img.size:
        composited_img = composited_img.resize(original_img.size, Image.LANCZOS)
    
    # Convert to numpy arrays (0-1 range)
    composited_np = np.array(composited_img).astype(np.float32) / 255.0
    original_np = np.array(original_img).astype(np.float32) / 255.0
    
    # Calculate CDL parameters to transform composited to match original
    slope, offset, power = calculate_cdl_params(composited_np, original_np)
    
    # Apply CDL transformation with blending
    smoothed_np = apply_cdl_transform(composited_np, slope, offset, power, factor)
    
    # Convert back to PIL Image
    smoothed_img = Image.fromarray((smoothed_np * 255).astype(np.uint8))
    
    return smoothed_img


def get_smoothing_stats(original_image_path: str, composited_image_path: str) -> dict:
    """Get statistics about the CDL transformation for debugging.
    
    Args:
        original_image_path (str): Path to the original target image
        composited_image_path (str): Path to the composited result image
        
    Returns:
        dict: Statistics about the transformation
    """
    # Load images
    composited_img = Image.open(composited_image_path).convert("RGB") 
    original_img = Image.open(original_image_path).convert("RGB")
    
    # Ensure same dimensions
    if composited_img.size != original_img.size:
        composited_img = composited_img.resize(original_img.size, Image.LANCZOS)
    
    # Convert to numpy arrays (0-1 range)
    composited_np = np.array(composited_img).astype(np.float32) / 255.0
    original_np = np.array(original_img).astype(np.float32) / 255.0
    
    # Calculate statistics
    composited_stats = calculate_channel_stats(composited_np)
    original_stats = calculate_channel_stats(original_np)
    
    # Calculate CDL parameters using face-based method when possible
    slope, offset, power = calculate_cdl_params(original_np, composited_np, 
                                               original_image_path, composited_image_path)
    
    return {
        'composited_stats': composited_stats,
        'original_stats': original_stats,
        'cdl_slope': slope,
        'cdl_offset': offset,
        'cdl_power': power
    }


def cdl_edge_smoothing_apply_to_source(source_image_path: str, target_image_path: str, factor: float = 1.0) -> Image.Image:
    """Apply CDL transformation to source image using face-based parameters when possible.
    
    This function:
    1. Calculates CDL parameters to transform source to match target (using face pixels when available)
    2. Applies those CDL parameters to the entire source image
    3. Returns the transformed source image
    
    Args:
        source_image_path (str): Path to the source image (to be transformed)
        target_image_path (str): Path to the target image (reference for CDL calculation)
        factor (float): Transform strength (0.0 = no change, 1.0 = full transform)
        
    Returns:
        Image.Image: Source image with CDL transformation applied
    """
    # Load images
    source_img = Image.open(source_image_path).convert("RGB")
    target_img = Image.open(target_image_path).convert("RGB")
    
    # Ensure same dimensions
    if source_img.size != target_img.size:
        target_img = target_img.resize(source_img.size, Image.LANCZOS)
    
    # Convert to numpy arrays (0-1 range)
    source_np = np.array(source_img).astype(np.float32) / 255.0
    target_np = np.array(target_img).astype(np.float32) / 255.0
    
    # Calculate CDL parameters using face-based method when possible
    slope, offset, power = calculate_cdl_params(source_np, target_np,
                                               source_image_path, target_image_path)
    
    # Apply CDL transformation to the entire source image
    transformed_np = apply_cdl_transform(source_np, slope, offset, power, factor)
    
    # Convert back to PIL Image
    transformed_img = Image.fromarray((transformed_np * 255).astype(np.uint8))
    
    return transformed_img


def extract_face_mask(image_path: str) -> Optional[np.ndarray]:
    """Extract face mask from an image using human parts segmentation.
    
    Args:
        image_path (str): Path to the image
        
    Returns:
        np.ndarray or None: Binary face mask, or None if no face found
    """
    try:
        from human_parts_segmentation import HumanPartsSegmentation
        
        segmenter = HumanPartsSegmentation()
        masks_dict = segmenter.segment_parts(image_path, ['face'])
        
        if 'face' in masks_dict and masks_dict['face'] is not None:
            face_mask = masks_dict['face']
            # Ensure it's a proper binary mask
            if np.sum(face_mask > 0.5) > 100:  # At least 100 face pixels
                return face_mask
        
        return None
        
    except Exception as e:
        print(f"Face extraction failed: {e}")
        return None