"""
Data augmentation class for image augmentation - ROTATION and RANDOM CROPPING with quality preservation
"""
import cv2
import numpy as np
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any
import random
import math
import logging
from image_processor import ImageProcessor
from utils import load_image, save_image, create_augmented_filename, print_progress

class DataAugmentation:
    """Class for image data augmentation - ROTATION and RANDOM CROPPING"""
    
    def __init__(self, config: Dict[str, Any] = None):
        """
        Initialize DataAugmentation
        
        Args:
            config: Configuration dictionary for augmentation parameters
        """
        self.config = config or {}
        self.image_processor = ImageProcessor()
        self.logger = logging.getLogger(__name__)
    
    def random_crop_preserve_quality(self, image: np.ndarray, crop_ratio_range: Tuple[float, float] = (0.7, 1.0)) -> np.ndarray:
        """
        Apply random cropping to simulate partially visible ID cards with white background preservation
        
        Args:
            image: Input image
            crop_ratio_range: Range for crop ratio (min_ratio, max_ratio)
            
        Returns:
            Cropped image with white background
        """
        height, width = image.shape[:2]
        
        # Random crop ratio between 0.7 and 1.0
        crop_ratio = random.uniform(crop_ratio_range[0], crop_ratio_range[1])
        
        # Calculate new dimensions
        new_height = int(height * crop_ratio)
        new_width = int(width * crop_ratio)
        
        # Ensure minimum size
        new_height = max(new_height, int(height * 0.7))
        new_width = max(new_width, int(width * 0.7))
        
        # Random starting position for cropping
        max_start_y = height - new_height
        max_start_x = width - new_width
        
        # Ensure we don't go out of bounds
        if max_start_y < 0 or max_start_x < 0:
            # If crop is larger than image, return original
            return image
        
        start_y = random.randint(0, max_start_y)
        start_x = random.randint(0, max_start_x)
        
        # Crop the image
        cropped = image[start_y:start_y + new_height, start_x:start_x + new_width]
        
        # Create white background canvas with original size
        canvas = np.ones((height, width, 3), dtype=np.uint8) * 255
        
        # Calculate offset to center the cropped image
        offset_y = (height - new_height) // 2
        offset_x = (width - new_width) // 2
        
        # Place cropped image in center of white canvas
        canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = cropped
        
        return canvas
    
    def add_random_noise_preserve_quality(self, image: np.ndarray, mean_range: Tuple[float, float] = (0.0, 0.7), variance_range: Tuple[float, float] = (0.0, 0.1)) -> np.ndarray:
        """
        Add random noise to simulate worn-out ID cards with quality preservation
        
        Args:
            image: Input image
            mean_range: Range for noise mean (min_mean, max_mean)
            variance_range: Range for noise variance (min_variance, max_variance)
            
        Returns:
            Image with added noise
        """
        # Convert image to float for noise addition
        image_float = image.astype(np.float32) / 255.0
        
        # Generate random noise parameters
        mean = random.uniform(mean_range[0], mean_range[1])
        variance = random.uniform(variance_range[0], variance_range[1])
        
        # Generate Gaussian noise
        noise = np.random.normal(mean, np.sqrt(variance), image_float.shape)
        
        # Add noise to image
        noisy_image = image_float + noise
        
        # Clip values to [0, 1] range
        noisy_image = np.clip(noisy_image, 0.0, 1.0)
        
        # Convert back to uint8
        noisy_image = (noisy_image * 255).astype(np.uint8)
        
        return noisy_image
    
    def add_partial_blockage_preserve_quality(self, image: np.ndarray, num_occlusions_range: Tuple[int, int] = (1, 100), coverage_range: Tuple[float, float] = (0.0, 0.25), variance_range: Tuple[float, float] = (0.0, 0.1)) -> np.ndarray:
        """
        Add partial blockage with horizontal lines of different colors to simulate occluded card details
        
        Args:
            image: Input image
            num_occlusions_range: Range for number of occlusions (min, max)
            coverage_range: Range for coverage ratio (min, max)
            variance_range: Range for line thickness variance (min, max)
            
        Returns:
            Image with partial blockage
        """
        height, width = image.shape[:2]
        result = image.copy()
        
        # Generate random parameters
        num_occlusions = random.randint(num_occlusions_range[0], num_occlusions_range[1])
        coverage = random.uniform(coverage_range[0], coverage_range[1])
        variance = random.uniform(variance_range[0], variance_range[1])
        
        # Calculate total area to cover
        total_area = height * width
        target_coverage_area = int(total_area * coverage)
        
        # Colors for occlusion lines (different colors to simulate various objects)
        occlusion_colors = [
            (255, 255, 255),  # White
            (0, 0, 0),        # Black
            (128, 128, 128),  # Gray
            (255, 0, 0),      # Red
            (0, 255, 0),      # Green
            (0, 0, 255),      # Blue
            (255, 255, 0),    # Yellow
            (255, 0, 255),    # Magenta
            (0, 255, 255),    # Cyan
        ]
        
        current_coverage = 0
        
        for _ in range(num_occlusions):
            if current_coverage >= target_coverage_area:
                break
                
            # Random line parameters
            color = random.choice(occlusion_colors)
            thickness = max(1, int(random.gauss(3, variance * 10)))
            
            # Random horizontal line coordinates (y1 = y2 for horizontal lines)
            y = random.randint(0, height)
            x1 = random.randint(0, width)
            x2 = random.randint(x1, width)  # Ensure x2 >= x1 for proper line
            
            # Draw horizontal line
            cv2.line(result, (x1, y), (x2, y), color, thickness)
            
            # Calculate coverage of this line (approximate)
            line_length = x2 - x1
            line_coverage = line_length * thickness
            current_coverage += line_coverage
        
        return result
    
    def convert_to_grayscale_preserve_quality(self, image: np.ndarray) -> np.ndarray:
        """
        Convert image to grayscale to mimic Xerox/scan copies with quality preservation
        
        Args:
            image: Input image
            
        Returns:
            Grayscale image
        """
        # Convert to grayscale using OpenCV
        grayscale = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        
        # Convert back to 3-channel image to maintain compatibility
        grayscale_3channel = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB)
        
        return grayscale_3channel
    
    def apply_blurring_preserve_quality(self, image: np.ndarray, kernel_ratio_range: Tuple[float, float] = (0.0, 0.0084)) -> np.ndarray:
        """
        Apply blurring to simulate blurred card images that are still readable using OpenCV filter2D
        
        Args:
            image: Input image
            kernel_ratio_range: Range for kernel ratio (min, max)
            
        Returns:
            Blurred image
        """
        # Generate random kernel ratio
        kernel_ratio = random.uniform(kernel_ratio_range[0], kernel_ratio_range[1])
        
        # Calculate kernel size based on image dimensions
        height, width = image.shape[:2]
        kernel_size = max(3, int(min(height, width) * kernel_ratio))
        
        # Ensure kernel size is odd
        if kernel_size % 2 == 0:
            kernel_size += 1
        
        # Create Gaussian kernel for blurring
        kernel = cv2.getGaussianKernel(kernel_size, kernel_size / 3.0)
        kernel_2d = kernel * kernel.T
        
        # Apply filter2D for blurring
        blurred = cv2.filter2D(image, -1, kernel_2d)
        
        return blurred
    
    def adjust_brightness_contrast_preserve_quality(self, image: np.ndarray, alpha_range: Tuple[float, float] = (0.4, 3.0), beta_range: Tuple[int, int] = (1, 100)) -> np.ndarray:
        """
        Adjust brightness and contrast to mimic different environmental lighting conditions using OpenCV convertScaleAbs
        
        Args:
            image: Input image
            alpha_range: Range for contrast (alpha) (min, max)
            beta_range: Range for brightness (beta) (min, max)
            
        Returns:
            Image with adjusted brightness and contrast
        """
        # Generate random alpha and beta values
        alpha = random.uniform(alpha_range[0], alpha_range[1])  # Contrast
        beta = random.randint(beta_range[0], beta_range[1])     # Brightness
        
        # Apply brightness and contrast adjustment using convertScaleAbs
        adjusted = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
        
        return adjusted
    
    def rotate_image_preserve_quality(self, image: np.ndarray, angle: float) -> np.ndarray:
        """
        Rotate image by given angle with white background and crop to preserve quality
        
        Args:
            image: Input image
            angle: Rotation angle in degrees
            
        Returns:
            Rotated and cropped image
        """
        height, width = image.shape[:2]
        
        # Calculate new dimensions for rotation
        angle_rad = math.radians(angle)
        cos_val = abs(math.cos(angle_rad))
        sin_val = abs(math.sin(angle_rad))
        
        # Calculate new width and height
        new_width = int(width * cos_val + height * sin_val)
        new_height = int(height * cos_val + width * sin_val)
        
        # Create larger canvas with white background
        canvas = np.ones((new_height, new_width, 3), dtype=np.uint8) * 255
        
        # Calculate offset to center the image
        offset_x = (new_width - width) // 2
        offset_y = (new_height - height) // 2
        
        # Ensure offsets are valid
        if offset_x >= 0 and offset_y >= 0 and offset_x + width <= new_width and offset_y + height <= new_height:
            # Place original image in center of canvas
            canvas[offset_y:offset_y+height, offset_x:offset_x+width] = image
        else:
            # If calculation is wrong, use a simpler approach
            canvas = np.ones((max(height, width) * 2, max(height, width) * 2, 3), dtype=np.uint8) * 255
            center_y, center_x = canvas.shape[0] // 2, canvas.shape[1] // 2
            start_y = center_y - height // 2
            start_x = center_x - width // 2
            canvas[start_y:start_y+height, start_x:start_x+width] = image
            new_width, new_height = canvas.shape[1], canvas.shape[0]
        
        # Calculate center for rotation
        center = (new_width // 2, new_height // 2)
        
        # Create rotation matrix
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # Perform rotation
        rotated = cv2.warpAffine(canvas, rotation_matrix, (new_width, new_height), 
                                borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255))
        
        # Crop white borders to get the actual image content
        rotated = self._crop_white_borders(rotated)
        
        return rotated
    
    def _crop_white_borders(self, image: np.ndarray) -> np.ndarray:
        """
        Crop white borders from image to get the actual content
        
        Args:
            image: Input image with white borders
            
        Returns:
            Cropped image without white borders
        """
        # Convert to grayscale for edge detection
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        
        # Find non-white pixels (content)
        non_white = gray < 250  # Threshold for white pixels
        
        # Find bounding box of content
        coords = cv2.findNonZero(non_white.astype(np.uint8))
        if coords is not None:
            x, y, w, h = cv2.boundingRect(coords)
            return image[y:y+h, x:x+w]
        
        return image
    
    def resize_preserve_aspect(self, image: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
        """
        Resize image preserving aspect ratio with white padding
        
        Args:
            image: Input image
            target_size: Target size (width, height)
            
        Returns:
            Resized image with preserved aspect ratio
        """
        target_width, target_height = target_size
        img_height, img_width = image.shape[:2]
        
        # Calculate aspect ratios
        target_aspect = target_width / target_height
        img_aspect = img_width / img_height
        
        if img_aspect > target_aspect:
            # Image is wider than target
            new_width = target_width
            new_height = int(target_width / img_aspect)
        else:
            # Image is taller than target
            new_height = target_height
            new_width = int(target_height * img_aspect)
        
        # Resize image
        resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
        
        # Create white background
        result = np.ones((target_height, target_width, 3), dtype=np.uint8) * 255
        
        # Calculate offset to center the image
        offset_x = (target_width - new_width) // 2
        offset_y = (target_height - new_height) // 2
        
        # Place resized image in center
        result[offset_y:offset_y+new_height, offset_x:offset_x+new_width] = resized
        
        return result
    
    def augment_single_image(self, image: np.ndarray, num_target_images: int = None) -> List[np.ndarray]:
        """
        Apply random combination of augmentation methods to create diverse augmented versions
        
        Args:
            image: Input image
            num_target_images: Number of target augmented images to generate
            
        Returns:
            List of augmented images with random method combinations
        """
        num_target_images = num_target_images or 3  # Default value
        
        # Get strategy config
        strategy_config = self.config.get("strategy", {})
        methods_config = self.config.get("methods", {})
        final_config = self.config.get("final_processing", {})
        
        mode = strategy_config.get("mode", "random_combine")
        min_methods = strategy_config.get("min_methods", 2)
        max_methods = strategy_config.get("max_methods", 4)
        
        if mode == "random_combine":
            return self._augment_random_combine(image, num_target_images, methods_config, final_config, min_methods, max_methods)
        elif mode == "sequential":
            return self._augment_sequential(image, num_target_images, methods_config, final_config)
        elif mode == "individual":
            return self._augment_individual_legacy(image, num_target_images)
        else:
            # Fallback to legacy method
            return self._augment_individual_legacy(image, num_target_images)
    
    def _augment_random_combine(self, image: np.ndarray, num_target_images: int, 
                               methods_config: dict, final_config: dict, 
                               min_methods: int, max_methods: int) -> List[np.ndarray]:
        """Apply random combination of methods"""
        augmented_images = []
        
        # Get enabled methods with their probabilities
        available_methods = []
        for method_name, method_config in methods_config.items():
            if method_config.get("enabled", False):
                available_methods.append((method_name, method_config))
        
        if not available_methods:
            self.logger.warning("No augmentation methods enabled!")
            return [image.copy() for _ in range(num_target_images)]
        
        for i in range(num_target_images):
            # Decide number of methods for this image
            num_methods = random.randint(min_methods, min(max_methods, len(available_methods)))
            
            # Select methods based on probability
            selected_methods = self._select_methods_by_probability(available_methods, num_methods)
            
            # Apply selected methods in sequence
            augmented = image.copy()
            method_names = []
            
            for method_name, method_config in selected_methods:
                if random.random() < method_config.get("probability", 0.5):
                    augmented = self._apply_single_method(augmented, method_name, method_config)
                    method_names.append(method_name)
            
            # Apply final processing
            augmented = self._apply_final_processing(augmented, final_config)
            
            # Resize preserving aspect ratio
            target_size = self.image_processor.target_size
            if target_size:
                augmented = self.resize_preserve_aspect(augmented, target_size)
            
            augmented_images.append(augmented)
        
        return augmented_images
    
    def _select_methods_by_probability(self, available_methods: List[Tuple], num_methods: int) -> List[Tuple]:
        """Select methods based on their probability weights"""
        # Create weighted list
        weighted_methods = []
        for method_name, method_config in available_methods:
            probability = method_config.get("probability", 0.5)
            weighted_methods.append((method_name, method_config, probability))
        
        # Sort by probability (highest first) and select top candidates
        weighted_methods.sort(key=lambda x: x[2], reverse=True)
        
        # Use weighted random selection
        selected = []
        remaining_methods = weighted_methods.copy()
        
        for _ in range(num_methods):
            if not remaining_methods:
                break
                
            # Calculate cumulative probabilities
            total_prob = sum(method[2] for method in remaining_methods)
            if total_prob == 0:
                # If all probabilities are 0, select randomly
                selected_method = random.choice(remaining_methods)
            else:
                rand_val = random.uniform(0, total_prob)
                cumulative_prob = 0
                selected_method = None
                
                for method in remaining_methods:
                    cumulative_prob += method[2]
                    if rand_val <= cumulative_prob:
                        selected_method = method
                        break
                
                if selected_method is None:
                    selected_method = remaining_methods[-1]
            
            selected.append((selected_method[0], selected_method[1]))
            remaining_methods.remove(selected_method)
        
        return selected
    
    def _apply_single_method(self, image: np.ndarray, method_name: str, method_config: dict) -> np.ndarray:
        """Apply a single augmentation method"""
        try:
            if method_name == "rotation":
                angles = method_config.get("angles", [30, 60, 90, 120, 150, 180, 210, 240, 300, 330])
                angle = random.choice(angles)
                return self.rotate_image_preserve_quality(image, angle)
                
            elif method_name == "random_cropping":
                ratio_range = method_config.get("ratio_range", (0.7, 1.0))
                return self.random_crop_preserve_quality(image, ratio_range)
                
            elif method_name == "random_noise":
                mean_range = method_config.get("mean_range", (0.0, 0.7))
                variance_range = method_config.get("variance_range", (0.0, 0.1))
                return self.add_random_noise_preserve_quality(image, mean_range, variance_range)
                
            elif method_name == "partial_blockage":
                num_range = method_config.get("num_occlusions_range", (1, 100))
                coverage_range = method_config.get("coverage_range", (0.0, 0.25))
                variance_range = method_config.get("variance_range", (0.0, 0.1))
                return self.add_partial_blockage_preserve_quality(image, num_range, coverage_range, variance_range)
                
            elif method_name == "blurring":
                kernel_range = method_config.get("kernel_ratio_range", (0.0, 0.0084))
                return self.apply_blurring_preserve_quality(image, kernel_range)
                
            elif method_name == "brightness_contrast":
                alpha_range = method_config.get("alpha_range", (0.4, 3.0))
                beta_range = method_config.get("beta_range", (1, 100))
                return self.adjust_brightness_contrast_preserve_quality(image, alpha_range, beta_range)
                
            elif method_name == "color_jitter":
                return self.apply_color_jitter(image, method_config)
                
            elif method_name == "perspective":
                distortion_scale = method_config.get("distortion_scale", 0.2)
                return self.apply_perspective_transform(image, distortion_scale)
                
            else:
                return image
                
        except Exception as e:
            print(f"Error applying method {method_name}: {e}")
            return image
    
    def _apply_final_processing(self, image: np.ndarray, final_config: dict) -> np.ndarray:
        """Apply final processing steps - ALWAYS applied to all outputs"""
        # Grayscale conversion - ALWAYS applied if enabled
        grayscale_config = final_config.get("grayscale", {})
        if grayscale_config.get("enabled", False):
            # Always apply grayscale, no random check
            image = self.convert_to_grayscale_preserve_quality(image)
        
        # Quality enhancement (future feature)
        quality_config = final_config.get("quality_enhancement", {})
        if quality_config.get("enabled", False):
            # TODO: Implement quality enhancement
            pass
        
        return image
    
    def apply_color_jitter(self, image: np.ndarray, config: dict) -> np.ndarray:
        """
        Apply color jittering (brightness, contrast, saturation, hue adjustments)
        
        Args:
            image: Input image
            config: Color jitter configuration
            
        Returns:
            Color-jittered image
        """
        # Get parameters
        brightness_range = config.get("brightness_range", [0.8, 1.2])
        contrast_range = config.get("contrast_range", [0.8, 1.2])
        saturation_range = config.get("saturation_range", [0.8, 1.2])
        hue_range = config.get("hue_range", [-0.1, 0.1])
        
        # Convert to HSV for saturation and hue adjustments
        hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV).astype(np.float32)
        
        # Apply brightness (adjust V channel)
        brightness_factor = random.uniform(brightness_range[0], brightness_range[1])
        hsv[:, :, 2] = np.clip(hsv[:, :, 2] * brightness_factor, 0, 255)
        
        # Apply saturation (adjust S channel)
        saturation_factor = random.uniform(saturation_range[0], saturation_range[1])
        hsv[:, :, 1] = np.clip(hsv[:, :, 1] * saturation_factor, 0, 255)
        
        # Apply hue shift (adjust H channel)
        hue_shift = random.uniform(hue_range[0], hue_range[1]) * 179  # OpenCV hue range is 0-179
        hsv[:, :, 0] = (hsv[:, :, 0] + hue_shift) % 180
        
        # Convert back to RGB
        result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
        
        # Apply contrast (after converting back to RGB)
        contrast_factor = random.uniform(contrast_range[0], contrast_range[1])
        result = cv2.convertScaleAbs(result, alpha=contrast_factor, beta=0)
        
        return result
    
    def apply_perspective_transform(self, image: np.ndarray, distortion_scale: float = 0.2) -> np.ndarray:
        """
        Apply perspective transformation to simulate viewing angle changes
        
        Args:
            image: Input image
            distortion_scale: Scale of perspective distortion (0.0 to 1.0)
            
        Returns:
            Perspective-transformed image
        """
        height, width = image.shape[:2]
        
        # Define source points (corners of original image)
        src_points = np.float32([
            [0, 0],
            [width-1, 0],
            [width-1, height-1],
            [0, height-1]
        ])
        
        # Add random distortion to destination points
        max_distortion = min(width, height) * distortion_scale
        
        dst_points = np.float32([
            [random.uniform(0, max_distortion), random.uniform(0, max_distortion)],
            [width-1-random.uniform(0, max_distortion), random.uniform(0, max_distortion)],
            [width-1-random.uniform(0, max_distortion), height-1-random.uniform(0, max_distortion)],
            [random.uniform(0, max_distortion), height-1-random.uniform(0, max_distortion)]
        ])
        
        # Calculate perspective transformation matrix
        matrix = cv2.getPerspectiveTransform(src_points, dst_points)
        
        # Apply transformation
        result = cv2.warpPerspective(image, matrix, (width, height), 
                                   borderMode=cv2.BORDER_CONSTANT, 
                                   borderValue=(255, 255, 255))
        
        return result
    
    def _augment_sequential(self, image: np.ndarray, num_target_images: int, 
                          methods_config: dict, final_config: dict) -> List[np.ndarray]:
        """Apply methods in sequence (pipeline style)"""
        augmented_images = []
        
        # Get enabled methods
        enabled_methods = [
            (name, config) for name, config in methods_config.items() 
            if config.get("enabled", False)
        ]
        
        for i in range(num_target_images):
            augmented = image.copy()
            
            # Apply all enabled methods in sequence
            for method_name, method_config in enabled_methods:
                if random.random() < method_config.get("probability", 0.5):
                    augmented = self._apply_single_method(augmented, method_name, method_config)
            
            # Apply final processing
            augmented = self._apply_final_processing(augmented, final_config)
            
            # Resize preserving aspect ratio
            target_size = self.image_processor.target_size
            if target_size:
                augmented = self.resize_preserve_aspect(augmented, target_size)
            
            augmented_images.append(augmented)
        
        return augmented_images
    
    def _augment_individual_legacy(self, image: np.ndarray, num_target_images: int) -> List[np.ndarray]:
        """Legacy individual method application (backward compatibility)"""
        # This is the old implementation for backward compatibility
        augmented_images = []
        
        # Get old-style configuration
        rotation_config = self.config.get("rotation", {})
        cropping_config = self.config.get("random_cropping", {})
        noise_config = self.config.get("random_noise", {})
        blockage_config = self.config.get("partial_blockage", {})
        grayscale_config = self.config.get("grayscale", {})
        blurring_config = self.config.get("blurring", {})
        brightness_contrast_config = self.config.get("brightness_contrast", {})
        
        # Apply individual methods (old logic)
        methods = [
            ("rotation", rotation_config, self.rotate_image_preserve_quality),
            ("cropping", cropping_config, self.random_crop_preserve_quality),
            ("noise", noise_config, self.add_random_noise_preserve_quality),
            ("blockage", blockage_config, self.add_partial_blockage_preserve_quality),
            ("blurring", blurring_config, self.apply_blurring_preserve_quality),
            ("brightness_contrast", brightness_contrast_config, self.adjust_brightness_contrast_preserve_quality)
        ]
        
        for method_name, method_config, method_func in methods:
            if method_config.get("enabled", False):
                for i in range(num_target_images):
                    augmented = image.copy()
                    # Apply single method with appropriate parameters
                    if method_name == "rotation":
                        angles = method_config.get("angles", [30, 60, 90, 120, 150, 180, 210, 240, 300, 330])
                        angle = random.choice(angles)
                        augmented = method_func(augmented, angle)
                    elif method_name == "cropping":
                        ratio_range = method_config.get("ratio_range", (0.7, 1.0))
                        augmented = method_func(augmented, ratio_range)
                    # Add other method parameter handling as needed
                    
                    # Resize preserving aspect ratio
                    target_size = self.image_processor.target_size
                    if target_size:
                        augmented = self.resize_preserve_aspect(augmented, target_size)
                    
                    augmented_images.append(augmented)
        
        # Apply grayscale to all images
        if grayscale_config.get("enabled", False):
            for i in range(len(augmented_images)):
                augmented_images[i] = self.convert_to_grayscale_preserve_quality(augmented_images[i])
        
        return augmented_images
    
    def augment_image_file(self, image_path: Path, output_dir: Path, num_target_images: int = None) -> List[Path]:
        """
        Augment a single image file and save results with quality preservation
        
        Args:
            image_path: Path to input image
            output_dir: Output directory for augmented images  
            num_target_images: Number of target augmented images to generate
            
        Returns:
            List of paths to saved augmented images
        """
        # Load image without resizing to preserve original quality
        image = load_image(image_path, None)
        if image is None:
            return []
        
        # Apply augmentations
        augmented_images = self.augment_single_image(image, num_target_images)
        
        # Save augmented images
        saved_paths = []
        for i, aug_image in enumerate(augmented_images):
            base_name = image_path.stem
            output_filename = f"{base_name}_aug_{i+1:03d}.jpg"
            output_path = output_dir / output_filename
            
            if save_image(aug_image, output_path):
                saved_paths.append(output_path)
        
        return saved_paths
    
    def augment_image_file_with_raw(self, image_path: Path, output_dir: Path, 
                                  num_total_versions: int = None) -> List[Path]:
        """
        Augment a single image file including raw/original version
        
        Args:
            image_path: Path to input image
            output_dir: Output directory for all image versions
            num_total_versions: Total number of versions (including raw)
            
        Returns:
            List of paths to saved images (raw + augmented)
        """
        # Load original image
        image = load_image(image_path, None)
        if image is None:
            return []
        
        saved_paths = []
        base_name = image_path.stem
        
        # Always save raw version first (resized but not augmented)
        if num_total_versions > 0:
            raw_image = image.copy()
            
            # Apply final processing (grayscale) but no augmentation
            final_config = self.config.get("final_processing", {})
            raw_image = self._apply_final_processing(raw_image, final_config)
            
            # Resize to target size
            target_size = self.image_processor.target_size
            if target_size:
                raw_image = self.resize_preserve_aspect(raw_image, target_size)
            
            # Save raw version
            raw_filename = f"{base_name}_raw_001.jpg"
            raw_path = output_dir / raw_filename
            if save_image(raw_image, raw_path):
                saved_paths.append(raw_path)
        
        # Generate augmented versions for remaining slots
        num_augmented = max(0, num_total_versions - 1)
        if num_augmented > 0:
            augmented_images = self.augment_single_image(image, num_augmented)
            
            for i, aug_image in enumerate(augmented_images):
                aug_filename = f"{base_name}_aug_{i+1:03d}.jpg"
                aug_path = output_dir / aug_filename
                
                if save_image(aug_image, aug_path):
                    saved_paths.append(aug_path)
        
        return saved_paths
    
    def batch_augment(self, input_dir: Path, output_dir: Path, 
                     multiplication_factor: float = None, random_seed: int = None) -> Dict[str, List[Path]]:
        """
        Augment images in a directory with smart sampling and multiplication strategy
        
        Args:
            input_dir: Input directory containing images
            output_dir: Output directory for augmented images
            multiplication_factor: 
                - If < 1.0: Sample percentage of input data to augment
                - If >= 1.0: Target multiplication factor for output data size
            random_seed: Random seed for reproducibility
            
        Returns:
            Dictionary containing results and statistics
        """
        from utils import get_image_files
        
        # Set random seed for reproducibility
        if random_seed is not None:
            random.seed(random_seed)
            np.random.seed(random_seed)
        
        # Get all input images
        all_image_files = get_image_files(input_dir)
        if not all_image_files:
            print("No images found in input directory")
            return {}
        
        # Get multiplication factor from config if not provided
        if multiplication_factor is None:
            data_strategy = self.config.get("data_strategy", {})
            multiplication_factor = data_strategy.get("multiplication_factor", 3.0)
        
        print(f"Found {len(all_image_files)} total images")
        print(f"Multiplication factor: {multiplication_factor}")
        
        # Determine sampling strategy
        if multiplication_factor < 1.0:
            # Sampling mode: Take a percentage of input data
            num_selected = int(len(all_image_files) * multiplication_factor)
            selected_images = self._sample_images(all_image_files, num_selected)
            target_total_images = len(all_image_files)  # Keep original dataset size
            images_per_input = max(1, target_total_images // len(selected_images))
            print(f"SAMPLING MODE: Selected {len(selected_images)} images ({multiplication_factor*100:.1f}%)")
            print(f"Target: {target_total_images} total images, {images_per_input} per selected image")
        else:
            # Multiplication mode: Multiply dataset size
            selected_images = all_image_files
            target_total_images = int(len(all_image_files) * multiplication_factor)
            images_per_input = max(1, target_total_images // len(selected_images))
            print(f"MULTIPLICATION MODE: Processing all {len(selected_images)} images")
            print(f"Target: {target_total_images} total images ({multiplication_factor}x original), {images_per_input} per image")
        
        # Process selected images
        results = {}
        total_generated = 0
        
        for i, image_path in enumerate(selected_images):
            print_progress(i + 1, len(selected_images), f"Processing {image_path.name}")
            
            # Calculate number of versions for this image (including raw)
            remaining_images = target_total_images - total_generated
            remaining_inputs = len(selected_images) - i
            total_versions_needed = min(images_per_input, remaining_images)
            
            # Always include raw image, then augmented ones
            augmented_paths = self.augment_image_file_with_raw(
                image_path, output_dir, total_versions_needed
            )
            
            if augmented_paths:
                results[str(image_path)] = augmented_paths
                total_generated += len(augmented_paths)
        
        # Generate summary
        summary = {
            "input_images": len(all_image_files),
            "selected_images": len(selected_images),
            "target_total": target_total_images,
            "actual_generated": total_generated,
            "multiplication_factor": multiplication_factor,
            "mode": "sampling" if multiplication_factor < 1.0 else "multiplication",
            "results": results,
            "efficiency": total_generated / target_total_images if target_total_images > 0 else 0
        }
        
        print(f"\n✅ Augmentation completed!")
        print(f"Generated {total_generated} images from {len(selected_images)} selected images")
        print(f"Target vs Actual: {target_total_images} → {total_generated} ({summary['efficiency']:.1%} efficiency)")
        
        return summary
    
    def _sample_images(self, image_files: List[Path], num_selected: int) -> List[Path]:
        """Sample images from the input list based on strategy"""
        data_strategy = self.config.get("data_strategy", {})
        sampling_config = data_strategy.get("sampling", {})
        
        method = sampling_config.get("method", "random")
        preserve_distribution = sampling_config.get("preserve_distribution", True)
        
        if method == "random":
            # Simple random sampling
            return random.sample(image_files, min(num_selected, len(image_files)))
        
        elif method == "stratified" and preserve_distribution:
            # Stratified sampling by file extension
            extension_groups = {}
            for img_file in image_files:
                ext = img_file.suffix.lower()
                if ext not in extension_groups:
                    extension_groups[ext] = []
                extension_groups[ext].append(img_file)
            
            selected = []
            for ext, files in extension_groups.items():
                # Sample proportionally from each extension group
                group_size = max(1, int(num_selected * len(files) / len(image_files)))
                group_selected = random.sample(files, min(group_size, len(files)))
                selected.extend(group_selected)
            
            # If we have too few, add more randomly
            if len(selected) < num_selected:
                remaining = [f for f in image_files if f not in selected]
                additional = random.sample(remaining, 
                                         min(num_selected - len(selected), len(remaining)))
                selected.extend(additional)
            
            return selected[:num_selected]
        
        elif method == "uniform":
            # Uniform sampling - evenly spaced
            if num_selected >= len(image_files):
                return image_files
            
            step = len(image_files) / num_selected
            indices = [int(i * step) for i in range(num_selected)]
            return [image_files[i] for i in indices]
        
        else:
            # Fallback to random
            return random.sample(image_files, min(num_selected, len(image_files)))
    
    def get_augmentation_summary(self, results: Dict[str, List[Path]]) -> Dict[str, Any]:
        """
        Get summary of augmentation results
        
        Args:
            results: Results from batch_augment
            
        Returns:
            Summary dictionary
        """
        total_original = len(results)
        total_augmented = sum(len(paths) for paths in results.values())
        
        return {
            "total_original_images": total_original,
            "total_augmented_images": total_augmented,
            "augmentation_ratio": total_augmented / total_original if total_original > 0 else 0,
            "successful_augmentations": len([paths for paths in results.values() if paths])
        } 