Code/Model/src/preprocessing/augmentation.py

"""
Data augmentation module for card images

Provides functions for augmenting card images to expand training dataset.
Includes geometric, color, and noise augmentation techniques.
"""

import cv2
import numpy as np
from typing import List, Tuple, Dict, Any
from ..utils.logger import get_logger

logger = get_logger(__name__)


# ============= Geometric Augmentation =============

def rotate_image(image: np.ndarray, angle: float) -> np.ndarray:
    """
    Rotate image by specified angle

    Args:
        image: Input image (H×W×C)
        angle: Rotation angle in degrees (positive = counter-clockwise)

    Returns:
        Rotated image with same dimensions

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to rotate_image")

    height, width = image.shape[:2]
    center = (width // 2, height // 2)

    # Get rotation matrix
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)

    # Apply rotation
    rotated = cv2.warpAffine(
        image,
        rotation_matrix,
        (width, height),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE
    )

    logger.debug(f"Rotated image by {angle:.2f} degrees")

    return rotated


def flip_image(image: np.ndarray, mode: str = 'horizontal') -> np.ndarray:
    """
    Flip image horizontally, vertically, or both

    Args:
        image: Input image (H×W×C)
        mode: Flip mode - 'horizontal', 'vertical', or 'both'

    Returns:
        Flipped image

    Raises:
        ValueError: If image is None/empty or mode is invalid
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to flip_image")

    if mode == 'horizontal':
        flipped = cv2.flip(image, 1)
    elif mode == 'vertical':
        flipped = cv2.flip(image, 0)
    elif mode == 'both':
        flipped = cv2.flip(image, -1)
    else:
        raise ValueError(f"Invalid flip mode: {mode}. Must be 'horizontal', 'vertical', or 'both'")

    logger.debug(f"Flipped image {mode}")

    return flipped


def zoom_image(image: np.ndarray, scale: float) -> np.ndarray:
    """
    Zoom image by specified scale factor

    Zooms into or out of the image center while maintaining output size.
    Scale > 1.0 zooms in (crops), scale < 1.0 zooms out (adds border).

    Args:
        image: Input image (H×W×C)
        scale: Zoom scale factor (valid range: 0.8-1.2)

    Returns:
        Zoomed image with same dimensions as input

    Raises:
        ValueError: If image is None/empty or scale is out of valid range
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to zoom_image")

    if scale < 0.8 or scale > 1.2:
        raise ValueError(f"Invalid zoom scale: {scale}. Must be in range [0.8, 1.2]")

    height, width = image.shape[:2]

    # Calculate new dimensions after scaling
    new_height = int(height * scale)
    new_width = int(width * scale)

    # Resize image
    resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)

    # Crop or pad to original size
    if scale > 1.0:
        # Zoom in - crop center
        start_y = (new_height - height) // 2
        start_x = (new_width - width) // 2
        zoomed = resized[start_y:start_y+height, start_x:start_x+width]
    else:
        # Zoom out - pad with border
        zoomed = np.zeros_like(image)
        start_y = (height - new_height) // 2
        start_x = (width - new_width) // 2
        zoomed[start_y:start_y+new_height, start_x:start_x+new_width] = resized

        # Fill border with edge replication
        if start_y > 0:
            zoomed[:start_y, :] = zoomed[start_y, :]
            zoomed[start_y+new_height:, :] = zoomed[start_y+new_height-1, :]
        if start_x > 0:
            zoomed[:, :start_x] = zoomed[:, start_x:start_x+1]
            zoomed[:, start_x+new_width:] = zoomed[:, start_x+new_width-1:start_x+new_width]

    logger.debug(f"Zoomed image with scale {scale:.2f}")

    return zoomed


# ============= Color Augmentation =============

def adjust_brightness(image: np.ndarray, factor: float) -> np.ndarray:
    """
    Adjust image brightness

    Args:
        image: Input image (H×W×C, uint8)
        factor: Brightness factor (1.0 = no change, >1.0 = brighter, <1.0 = darker)

    Returns:
        Brightness-adjusted image (clipped to [0, 255])

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to adjust_brightness")

    # Convert to float for computation
    adjusted = image.astype(np.float32) * factor

    # Clip to valid range and convert back to uint8
    adjusted = np.clip(adjusted, 0, 255).astype(np.uint8)

    logger.debug(f"Adjusted brightness with factor {factor:.2f}")

    return adjusted


def adjust_contrast(image: np.ndarray, factor: float) -> np.ndarray:
    """
    Adjust image contrast

    Args:
        image: Input image (H×W×C, uint8)
        factor: Contrast factor (1.0 = no change, >1.0 = more contrast, <1.0 = less)

    Returns:
        Contrast-adjusted image (clipped to [0, 255])

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to adjust_contrast")

    # Calculate mean value
    mean = image.mean()

    # Apply contrast adjustment around mean
    adjusted = (image.astype(np.float32) - mean) * factor + mean

    # Clip to valid range and convert back to uint8
    adjusted = np.clip(adjusted, 0, 255).astype(np.uint8)

    logger.debug(f"Adjusted contrast with factor {factor:.2f}")

    return adjusted


def adjust_saturation(image: np.ndarray, factor: float) -> np.ndarray:
    """
    Adjust image color saturation

    Args:
        image: Input image (H×W×C, BGR format, uint8)
        factor: Saturation factor (1.0 = no change, >1.0 = more saturated, <1.0 = less)

    Returns:
        Saturation-adjusted image

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to adjust_saturation")

    # Convert BGR to HSV
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)

    # Adjust saturation channel (index 1)
    hsv[:, :, 1] = hsv[:, :, 1] * factor

    # Clip saturation to valid range [0, 255]
    hsv[:, :, 1] = np.clip(hsv[:, :, 1], 0, 255)

    # Convert back to BGR
    hsv = hsv.astype(np.uint8)
    adjusted = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

    logger.debug(f"Adjusted saturation with factor {factor:.2f}")

    return adjusted


# ============= Noise Augmentation =============

def add_gaussian_noise(image: np.ndarray, mean: float = 0, std: float = 10) -> np.ndarray:
    """
    Add Gaussian noise to image

    Args:
        image: Input image (H×W×C, uint8)
        mean: Mean of Gaussian noise
        std: Standard deviation of Gaussian noise

    Returns:
        Noisy image (clipped to [0, 255])

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to add_gaussian_noise")

    # Generate Gaussian noise
    noise = np.random.normal(mean, std, image.shape).astype(np.float32)

    # Add noise to image
    noisy = image.astype(np.float32) + noise

    # Clip to valid range and convert back to uint8
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)

    logger.debug(f"Added Gaussian noise with mean={mean}, std={std}")

    return noisy


def apply_jpeg_compression(image: np.ndarray, quality: int = 85) -> np.ndarray:
    """
    Apply JPEG compression to simulate compression artifacts

    Args:
        image: Input image (H×W×C, uint8)
        quality: JPEG quality (1-100, higher = better quality)

    Returns:
        Compressed image

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to apply_jpeg_compression")

    # Encode to JPEG
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    _, encoded = cv2.imencode('.jpg', image, encode_param)

    # Decode back to image
    compressed = cv2.imdecode(encoded, cv2.IMREAD_COLOR)

    logger.debug(f"Applied JPEG compression with quality={quality}")

    return compressed


# ============= Augmentation Pipeline =============

def get_front_params() -> Dict[str, Any]:
    """
    Get augmentation parameters optimized for card front images.

    Card fronts have rich, diverse artwork and colors, so they can handle
    more aggressive augmentation without losing distinguishing features.

    Returns:
        Dictionary of augmentation parameters for front images
    """
    return {
        'rotation_range': 10,       # ±10 degrees
        'brightness_range': 0.15,   # ±15%
        'contrast_range': 0.10,     # ±10%
        'saturation_range': 0.15,   # ±15%
        'zoom_range': 0.05,         # 95-105%
        'noise_std': 10,            # σ=10
        'jpeg_quality': 85,         # Quality 85
        'flip_probability': 0.3     # 30% chance of flip
    }


def get_back_params() -> Dict[str, Any]:
    """
    Get augmentation parameters optimized for card back images.

    Card backs have subtle, uniform features (slight color differences,
    fine printing quality variations). More conservative augmentation
    preserves these subtle distinguishing features.

    Key differences from front params:
    - Reduced brightness/saturation variation (±5% vs ±15%)
    - Higher JPEG quality (95 vs 85) to preserve print quality details
    - Lower noise (σ=5 vs σ=10) to preserve fine texture

    Returns:
        Dictionary of augmentation parameters for back images
    """
    return {
        'rotation_range': 10,       # ±10 degrees (same as front)
        'brightness_range': 0.05,   # ±5% (reduced from 15%)
        'contrast_range': 0.10,     # ±10% (same as front)
        'saturation_range': 0.05,   # ±5% (reduced from 15%)
        'zoom_range': 0.05,         # 95-105% (same as front)
        'noise_std': 5,             # σ=5 (reduced from 10)
        'jpeg_quality': 95,         # Quality 95 (increased from 85)
        'flip_probability': 0.3     # 30% chance of flip (same as front)
    }


def detect_image_type(filename: str) -> str:
    """
    Detect if image is a card front or back based on filename or path.

    Args:
        filename: Image filename or path

    Returns:
        'front' or 'back' (defaults to 'front' if cannot determine)
    """
    filename_lower = str(filename).lower()

    # Check if path contains 'back' or 'front' directory (with or without trailing slash)
    if '/back/' in filename_lower or '\\back\\' in filename_lower or \
       '/back' in filename_lower or '\\back' in filename_lower or \
       'back/' in filename_lower or 'back\\' in filename_lower:
        return 'back'
    if '/front/' in filename_lower or '\\front\\' in filename_lower or \
       '/front' in filename_lower or '\\front' in filename_lower or \
       'front/' in filename_lower or 'front\\' in filename_lower:
        return 'front'

    # Check if filename contains 'back' or 'front'
    if '_back' in filename_lower or '-back' in filename_lower or 'back_' in filename_lower or 'back-' in filename_lower:
        return 'back'
    if '_front' in filename_lower or '-front' in filename_lower or 'front_' in filename_lower or 'front-' in filename_lower:
        return 'front'

    # Default to front if cannot determine
    logger.debug(f"Could not determine image type for '{filename}', defaulting to 'front'")
    return 'front'


def augment_image(
    image: np.ndarray,
    label: str,
    num_variations: int = 5,
    params: Dict[str, Any] = None,
    image_type: str = None,
    filename: str = None
) -> List[Tuple[np.ndarray, str]]:
    """
    Generate augmented variations of a single image

    Applies random combinations of augmentations to create diverse samples
    while preserving the label. Automatically selects appropriate parameters
    for front vs back images to preserve distinguishing features.

    Args:
        image: Input image (H×W×C, uint8)
        label: Image label (e.g., 'authentic' or 'fake')
        num_variations: Number of augmented versions to generate
        params: Optional augmentation parameters (overrides image_type detection)
        image_type: Optional image type ('front' or 'back'). If None, uses filename detection
        filename: Optional filename for automatic type detection

    Returns:
        List of (augmented_image, label) tuples

    Raises:
        ValueError: If image is None or empty
    """
    if image is None or image.size == 0:
        raise ValueError("Empty or None image provided to augment_image")

    # Determine parameters
    if params is None:
        # Auto-detect image type if not provided
        if image_type is None and filename is not None:
            image_type = detect_image_type(filename)
        elif image_type is None:
            image_type = 'front'  # Default

        # Select parameters based on image type
        if image_type == 'back':
            params = get_back_params()
            logger.debug(f"Using back-optimized parameters (conservative)")
        else:
            params = get_front_params()
            logger.debug(f"Using front-optimized parameters (standard)")

    # Validate parameters
    validate_augmentation_params(params)

    augmented_images = []

    for i in range(num_variations):
        aug_img = image.copy()

        # Randomly apply geometric augmentations
        if np.random.rand() < 0.7:  # 70% chance
            angle = np.random.uniform(-params['rotation_range'], params['rotation_range'])
            aug_img = rotate_image(aug_img, angle)

        if np.random.rand() < params['flip_probability']:
            flip_mode = np.random.choice(['horizontal', 'vertical'])
            aug_img = flip_image(aug_img, flip_mode)

        if np.random.rand() < 0.5:  # 50% chance
            scale = np.random.uniform(1.0 - params['zoom_range'], 1.0 + params['zoom_range'])
            aug_img = zoom_image(aug_img, scale)

        # Randomly apply color augmentations
        if np.random.rand() < 0.8:  # 80% chance
            brightness_factor = np.random.uniform(
                1.0 - params['brightness_range'],
                1.0 + params['brightness_range']
            )
            aug_img = adjust_brightness(aug_img, brightness_factor)

        if np.random.rand() < 0.6:  # 60% chance
            contrast_factor = np.random.uniform(
                1.0 - params['contrast_range'],
                1.0 + params['contrast_range']
            )
            aug_img = adjust_contrast(aug_img, contrast_factor)

        if np.random.rand() < 0.5:  # 50% chance
            saturation_factor = np.random.uniform(
                1.0 - params['saturation_range'],
                1.0 + params['saturation_range']
            )
            aug_img = adjust_saturation(aug_img, saturation_factor)

        # Randomly apply noise augmentations
        if np.random.rand() < 0.4:  # 40% chance
            aug_img = add_gaussian_noise(aug_img, mean=0, std=params['noise_std'])

        if np.random.rand() < 0.3:  # 30% chance
            quality = np.random.randint(params['jpeg_quality'] - 10, params['jpeg_quality'] + 5)
            quality = np.clip(quality, 70, 95)
            aug_img = apply_jpeg_compression(aug_img, quality)

        augmented_images.append((aug_img, label))

    logger.info(f"Generated {num_variations} augmented variations for label '{label}'")

    return augmented_images


def augment_dataset(
    images: List[np.ndarray],
    labels: List[str],
    num_variations: int = 5,
    include_original: bool = True,
    params: Dict[str, Any] = None,
    filenames: List[str] = None,
    auto_detect_type: bool = True
) -> List[Tuple[np.ndarray, str]]:
    """
    Augment entire dataset

    Generates augmented versions of all images in the dataset.
    Automatically applies appropriate parameters for front vs back images
    based on filename detection.

    Args:
        images: List of input images
        labels: List of corresponding labels
        num_variations: Number of augmented versions per image
        include_original: If True, includes original images in output
        params: Optional augmentation parameters (overrides auto-detection)
        filenames: Optional list of filenames for auto-detecting front/back
        auto_detect_type: If True, automatically detect and use front/back params

    Returns:
        List of (image, label) tuples including originals and augmented images

    Raises:
        ValueError: If images and labels have different lengths
    """
    if len(images) != len(labels):
        raise ValueError(
            f"Number of images ({len(images)}) must match number of labels ({len(labels)})"
        )

    if filenames is not None and len(filenames) != len(images):
        raise ValueError(
            f"Number of filenames ({len(filenames)}) must match number of images ({len(images)})"
        )

    augmented_dataset = []

    # Include original images if requested
    if include_original:
        for img, label in zip(images, labels):
            augmented_dataset.append((img, label))

    # Generate augmented images
    for i, (img, label) in enumerate(zip(images, labels)):
        # Get filename for this image if available
        filename = filenames[i] if filenames is not None else None

        # Augment image (will auto-detect type if auto_detect_type=True and params=None)
        augmented = augment_image(
            img,
            label,
            num_variations,
            params=params,
            filename=filename if auto_detect_type else None
        )
        augmented_dataset.extend(augmented)

    # Count front/back images for logging
    if filenames is not None and auto_detect_type and params is None:
        front_count = sum(1 for f in filenames if detect_image_type(f) == 'front')
        back_count = len(filenames) - front_count
        logger.info(
            f"Augmented dataset: {len(images)} original → {len(augmented_dataset)} total "
            f"({num_variations}x augmentation) | Front: {front_count}, Back: {back_count}"
        )
    else:
        logger.info(
            f"Augmented dataset: {len(images)} original → {len(augmented_dataset)} total "
            f"({num_variations}x augmentation)"
        )

    return augmented_dataset


def validate_augmentation_params(params: Dict[str, Any]) -> None:
    """
    Validate augmentation parameters

    Args:
        params: Dictionary of augmentation parameters

    Raises:
        ValueError: If any parameter is out of valid range
    """
    if 'rotation_range' in params:
        if params['rotation_range'] < 0 or params['rotation_range'] > 30:
            raise ValueError(f"rotation_range must be in [0, 30], got {params['rotation_range']}")

    if 'brightness_range' in params:
        if params['brightness_range'] < 0 or params['brightness_range'] > 0.3:
            raise ValueError(f"brightness_range must be in [0, 0.3], got {params['brightness_range']}")

    if 'contrast_range' in params:
        if params['contrast_range'] < 0 or params['contrast_range'] > 0.3:
            raise ValueError(f"contrast_range must be in [0, 0.3], got {params['contrast_range']}")

    if 'saturation_range' in params:
        if params['saturation_range'] < 0 or params['saturation_range'] > 0.5:
            raise ValueError(f"saturation_range must be in [0, 0.5], got {params['saturation_range']}")

    if 'zoom_range' in params:
        if params['zoom_range'] < 0 or params['zoom_range'] > 0.2:
            raise ValueError(f"zoom_range must be in [0, 0.2], got {params['zoom_range']}")

    if 'noise_std' in params:
        if params['noise_std'] < 0 or params['noise_std'] > 30:
            raise ValueError(f"noise_std must be in [0, 30], got {params['noise_std']}")

    if 'jpeg_quality' in params:
        if params['jpeg_quality'] < 1 or params['jpeg_quality'] > 100:
            raise ValueError(f"jpeg_quality must be in [1, 100], got {params['jpeg_quality']}")

    logger.debug("Augmentation parameters validated successfully")