"""
Data Augmentation for Face Recognition Dataset

This module provides data augmentation techniques to increase dataset diversity.

Project Information:
- Project ID: 22
- Title: Face Recognition Dataset
- Category: Image Data
- Technologies: PNG, JPG, NumPy, OpenCV, Face Recognition

Contact Information:
RSK World
Founder: Molla Samser
Designer & Tester: Rima Khatun
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
Website: https://rskworld.in/
Year: 2026
"""

import numpy as np
import cv2
from typing import List, Tuple, Optional
import os


class FaceAugmenter:
    """
    Data augmentation for face images.
    """
    
    @staticmethod
    def rotate(image: np.ndarray, angle: float = 15) -> np.ndarray:
        """
        Rotate image by specified angle.
        
        Args:
            image: Input image
            angle: Rotation angle in degrees
            
        Returns:
            Rotated image
        """
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(image, matrix, (w, h), borderMode=cv2.BORDER_REPLICATE)
        return rotated
    
    @staticmethod
    def flip_horizontal(image: np.ndarray) -> np.ndarray:
        """
        Flip image horizontally.
        
        Args:
            image: Input image
            
        Returns:
            Horizontally flipped image
        """
        return cv2.flip(image, 1)
    
    @staticmethod
    def adjust_brightness(image: np.ndarray, factor: float = 0.2) -> np.ndarray:
        """
        Adjust image brightness.
        
        Args:
            image: Input image
            factor: Brightness adjustment factor (-1 to 1)
            
        Returns:
            Brightness adjusted image
        """
        adjusted = image.astype(np.float32)
        adjusted = adjusted + (factor * 255)
        adjusted = np.clip(adjusted, 0, 255).astype(np.uint8)
        return adjusted
    
    @staticmethod
    def adjust_contrast(image: np.ndarray, factor: float = 0.2) -> np.ndarray:
        """
        Adjust image contrast.
        
        Args:
            image: Input image
            factor: Contrast adjustment factor
            
        Returns:
            Contrast adjusted image
        """
        adjusted = image.astype(np.float32)
        mean = np.mean(adjusted)
        adjusted = (adjusted - mean) * (1 + factor) + mean
        adjusted = np.clip(adjusted, 0, 255).astype(np.uint8)
        return adjusted
    
    @staticmethod
    def add_noise(image: np.ndarray, noise_factor: float = 0.05) -> np.ndarray:
        """
        Add Gaussian noise to image.
        
        Args:
            image: Input image
            noise_factor: Noise intensity factor
            
        Returns:
            Noisy image
        """
        noise = np.random.normal(0, noise_factor * 255, image.shape).astype(np.float32)
        noisy = image.astype(np.float32) + noise
        noisy = np.clip(noisy, 0, 255).astype(np.uint8)
        return noisy
    
    @staticmethod
    def blur(image: np.ndarray, kernel_size: int = 3) -> np.ndarray:
        """
        Apply Gaussian blur.
        
        Args:
            image: Input image
            kernel_size: Blur kernel size
            
        Returns:
            Blurred image
        """
        return cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
    
    @staticmethod
    def crop_and_resize(image: np.ndarray, crop_factor: float = 0.9) -> np.ndarray:
        """
        Crop and resize image.
        
        Args:
            image: Input image
            crop_factor: Crop factor (0-1)
            
        Returns:
            Cropped and resized image
        """
        h, w = image.shape[:2]
        new_h, new_w = int(h * crop_factor), int(w * crop_factor)
        start_y = (h - new_h) // 2
        start_x = (w - new_w) // 2
        
        cropped = image[start_y:start_y + new_h, start_x:start_x + new_w]
        resized = cv2.resize(cropped, (w, h))
        return resized
    
    def augment_image(self, image: np.ndarray, augmentations: List[str] = None) -> List[np.ndarray]:
        """
        Apply multiple augmentations to an image.
        
        Args:
            image: Input image
            augmentations: List of augmentation names to apply
            
        Returns:
            List of augmented images
        """
        if augmentations is None:
            augmentations = ['rotate', 'flip', 'brightness', 'contrast']
        
        augmented_images = []
        
        for aug_name in augmentations:
            if aug_name == 'rotate':
                for angle in [-15, 15]:
                    augmented_images.append(self.rotate(image, angle))
            elif aug_name == 'flip':
                augmented_images.append(self.flip_horizontal(image))
            elif aug_name == 'brightness':
                for factor in [-0.2, 0.2]:
                    augmented_images.append(self.adjust_brightness(image, factor))
            elif aug_name == 'contrast':
                for factor in [-0.2, 0.2]:
                    augmented_images.append(self.adjust_contrast(image, factor))
            elif aug_name == 'noise':
                augmented_images.append(self.add_noise(image))
            elif aug_name == 'blur':
                augmented_images.append(self.blur(image))
            elif aug_name == 'crop':
                augmented_images.append(self.crop_and_resize(image))
        
        return augmented_images
    
    def augment_directory(self, input_dir: str, output_dir: str, 
                       augmentations: List[str] = None, max_per_image: int = 5):
        """
        Augment all images in a directory.
        
        Args:
            input_dir: Input directory
            output_dir: Output directory
            augmentations: List of augmentation names
            max_per_image: Maximum augmentations per image
        """
        os.makedirs(output_dir, exist_ok=True)
        
        for person_dir in os.listdir(input_dir):
            person_path = os.path.join(input_dir, person_dir)
            if not os.path.isdir(person_path):
                continue
            
            output_person_dir = os.path.join(output_dir, person_dir)
            os.makedirs(output_person_dir, exist_ok=True)
            
            for filename in os.listdir(person_path):
                if any(filename.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png']):
                    image_path = os.path.join(person_path, filename)
                    image = cv2.imread(image_path)
                    
                    if image is None:
                        continue
                    
                    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    augmented = self.augment_image(rgb_image, augmentations)
                    
                    # Save original
                    base_name = os.path.splitext(filename)[0]
                    ext = os.path.splitext(filename)[1]
                    
                    for idx, aug_image in enumerate(augmented[:max_per_image]):
                        output_path = os.path.join(
                            output_person_dir, 
                            f"{base_name}_aug{idx}{ext}"
                        )
                        bgr_image = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR)
                        cv2.imwrite(output_path, bgr_image)

"""
Example Usage Script for Face Recognition Dataset

This script demonstrates how to use the face recognition system.

Project Information:
- Project ID: 22
- Title: Face Recognition Dataset
- Category: Image Data
- Technologies: PNG, JPG, NumPy, OpenCV, Face Recognition
- Difficulty: Intermediate

Contact Information:
RSK World
Founder: Molla Samser
Designer & Tester: Rima Khatun
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
Website: https://rskworld.in/
Year: 2026
"""

import os
import cv2
import numpy as np
import config
from scripts.load_dataset import FaceDatasetLoader
from scripts.preprocess import preprocess_faces, get_face_encoding
from scripts.recognize_faces import FaceRecognizer


def example_load_dataset():
    """Example: Loading the dataset."""
    print("\n" + "=" * 60)
    print("Example 1: Loading Dataset")
    print("=" * 60)
    
    loader = FaceDatasetLoader(config.TRAIN_DIR)
    images, labels, label_mapping = loader.load()
    
    print(f"Loaded {len(images)} images")
    print(f"Number of identities: {len(label_mapping)}")
    print("\nLabel mapping:")
    for label, name in sorted(label_mapping.items()):
        print(f"  {label}: {name}")
    
    # Get statistics
    stats = loader.get_statistics()
    print("\nDataset Statistics:")
    for key, value in stats.items():
        if key != 'images_per_person':
            print(f"  {key}: {value}")


def example_preprocess():
    """Example: Preprocessing images."""
    print("\n" + "=" * 60)
    print("Example 2: Preprocessing Images")
    print("=" * 60)
    
    processed_images, processed_labels = preprocess_faces(config.TRAIN_DIR)
    
    print(f"\nPreprocessed {len(processed_images)} face images")
    print(f"Image shape: {processed_images[0].shape}")
    print(f"Image dtype: {processed_images[0].dtype}")
    print(f"Number of unique labels: {len(np.unique(processed_labels))}")


def example_train_and_recognize():
    """Example: Training and recognizing faces."""
    print("\n" + "=" * 60)
    print("Example 3: Training and Recognition")
    print("=" * 60)
    
    # Initialize recognizer
    recognizer = FaceRecognizer()
    
    # Train the model
    print("\nTraining the model...")
    recognizer.train(config.TRAIN_DIR, save_model=True)
    
    # Example: Recognize faces in a test image (if available)
    if os.path.exists(config.TEST_DIR):
        test_images = []
        for root, dirs, files in os.walk(config.TEST_DIR):
            for file in files:
                if any(file.lower().endswith(ext) for ext in config.IMAGE_EXTENSIONS):
                    test_images.append(os.path.join(root, file))
                    break  # Just use first image found
        
        if test_images:
            test_image_path = test_images[0]
            print(f"\nRecognizing faces in: {test_image_path}")
            results = recognizer.recognize(test_image_path)
            
            for i, result in enumerate(results):
                print(f"\nFace {i + 1}:")
                print(f"  Name: {result['name']}")
                print(f"  Confidence: {result['confidence']:.2%}")
                if result['distance'] is not None:
                    print(f"  Distance: {result['distance']:.4f}")


def example_face_encoding():
    """Example: Getting face encodings."""
    print("\n" + "=" * 60)
    print("Example 4: Face Encoding")
    print("=" * 60)
    
    loader = FaceDatasetLoader(config.TRAIN_DIR)
    images, labels, label_mapping = loader.load()
    
    if len(images) > 0:
        # Get encoding for first image
        encoding = get_face_encoding(images[0])
        
        if encoding is not None:
            print(f"Face encoding shape: {encoding.shape}")
            print(f"Face encoding dtype: {encoding.dtype}")
            print(f"Encoding sample (first 5 values): {encoding[:5]}")
        else:
            print("No face detected in the first image")


def main():
    """Run all examples."""
    print("=" * 60)
    print("Face Recognition Dataset - Example Usage")
    print("=" * 60)
    print(f"Project: Face Recognition Dataset (ID: 22)")
    print(f"RSK World - https://rskworld.in/")
    print("=" * 60)
    
    # Check if training directory exists
    if not os.path.exists(config.TRAIN_DIR):
        print(f"\nError: Training directory not found: {config.TRAIN_DIR}")
        print("Please ensure the dataset is properly organized.")
        return
    
    try:
        example_load_dataset()
        example_preprocess()
        example_face_encoding()
        example_train_and_recognize()
        
        print("\n" + "=" * 60)
        print("All examples completed!")
        print("=" * 60)
        
    except Exception as e:
        print(f"\nError: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()