"""
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
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()
    
    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_augmentations: int = None) -> List[np.ndarray]:
        """
        Apply each augmentation method separately to create independent augmented versions
        
        Args:
            image: Input image
            num_augmentations: Number of augmented versions to create per method
            
        Returns:
            List of augmented images (each method creates separate versions)
        """
        num_augmentations = num_augmentations or 3  # Default value
        augmented_images = []
        
        # Get 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", {})
        
        # Configuration parameters
        angles = rotation_config.get("angles", [30, 60, 120, 150, 180, 210, 240, 300, 330])
        crop_ratio_range = cropping_config.get("ratio_range", (0.7, 1.0))
        mean_range = noise_config.get("mean_range", (0.0, 0.7))
        variance_range = noise_config.get("variance_range", (0.0, 0.1))
        num_occlusions_range = blockage_config.get("num_occlusions_range", (1, 100))
        coverage_range = blockage_config.get("coverage_range", (0.0, 0.25))
        blockage_variance_range = blockage_config.get("variance_range", (0.0, 0.1))
        kernel_ratio_range = blurring_config.get("kernel_ratio_range", (0.0, 0.0084))
        alpha_range = brightness_contrast_config.get("alpha_range", (0.4, 3.0))
        beta_range = brightness_contrast_config.get("beta_range", (1, 100))
        
        # Apply each method separately to create independent versions
        
        # 1. Rotation only
        if rotation_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                angle = random.choice(angles)
                augmented = self.rotate_image_preserve_quality(augmented, angle)
                
                # 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)
        
        # 2. Random cropping only
        if cropping_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.random_crop_preserve_quality(augmented, crop_ratio_range)
                
                # 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)
        
        # 3. Random noise only
        if noise_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.add_random_noise_preserve_quality(augmented, mean_range, variance_range)
                
                # 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)
        
        # 4. Partial blockage only
        if blockage_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.add_partial_blockage_preserve_quality(augmented, num_occlusions_range, coverage_range, blockage_variance_range)
                
                # 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)
        
        # 5. Grayscale only
        if grayscale_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.convert_to_grayscale_preserve_quality(augmented)
                
                # 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)
        
        # 6. Blurring only
        if blurring_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.apply_blurring_preserve_quality(augmented, kernel_ratio_range)
                
                # 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)
        
        # 7. Brightness and contrast only
        if brightness_contrast_config.get("enabled", False):
            for i in range(num_augmentations):
                augmented = image.copy()
                augmented = self.adjust_brightness_contrast_preserve_quality(augmented, alpha_range, beta_range)
                
                # 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_image_file(self, image_path: Path, output_dir: Path, num_augmentations: 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_augmentations: Number of augmented versions to create per method
            
        Returns:
            List of paths to saved augmented images
        """
        # Load image without resizing to preserve original quality
        image = load_image(image_path, None)  # Load original size
        if image is None:
            return []
        
        # Apply augmentations
        augmented_images = self.augment_single_image(image, num_augmentations)
        
        # Save augmented images with method names
        saved_paths = []
        method_names = ["rotation", "cropping", "noise", "blockage", "grayscale", "blurring", "brightness_contrast"]
        method_index = 0
        
        for i, aug_image in enumerate(augmented_images):
            # Determine method name based on index
            method_name = method_names[method_index // num_augmentations] if method_index // num_augmentations < len(method_names) else "aug"
            
            # Create output filename with method name
            output_filename = create_augmented_filename(image_path, (i % num_augmentations) + 1, method_name)
            output_path = output_dir / output_filename.name
            
            # Save image
            if save_image(aug_image, output_path):
                saved_paths.append(output_path)
            
            method_index += 1
        
        return saved_paths
    
    def batch_augment(self, input_dir: Path, output_dir: Path, num_augmentations: int = None) -> Dict[str, List[Path]]:
        """
        Augment all images in a directory
        
        Args:
            input_dir: Input directory containing images
            output_dir: Output directory for augmented images
            num_augmentations: Number of augmented versions per image
            
        Returns:
            Dictionary mapping original images to their augmented versions
        """
        from utils import get_image_files
        
        image_files = get_image_files(input_dir)
        results = {}
        
        print(f"Found {len(image_files)} images to augment")
        
        for i, image_path in enumerate(image_files):
            print_progress(i + 1, len(image_files), "Augmenting images")
            
            # Augment single image
            augmented_paths = self.augment_image_file(image_path, output_dir, num_augmentations)
            
            if augmented_paths:
                results[str(image_path)] = augmented_paths
        
        print(f"\nAugmented {len(results)} images successfully")
        return results
    
    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])
        } 