"""
==================================================================================
    Action Recognition Dataset - Advanced Data Augmentation
==================================================================================
    Project: Action Recognition Dataset
    
    Comprehensive video augmentation pipeline with:
    - Spatial augmentations (flip, crop, rotate, scale)
    - Temporal augmentations (speed, reverse, skip)
    - Color augmentations (brightness, contrast, saturation, hue)
    - Noise and blur effects
    - Cutout and mixup
    
==================================================================================
    DEVELOPER INFORMATION
==================================================================================
    Website: RSK World (https://rskworld.in)
    Founded by: Molla Samser
    Designer & Tester: Rima Khatun
    Contact: help@rskworld.in | +91 93305 39277
    
    (c) 2026 RSK World. All Rights Reserved.
==================================================================================
"""

import sys
import os
from pathlib import Path
import random

try:
    import cv2
    import numpy as np
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python", "numpy"])
    import cv2
    import numpy as np


# ==================================================================================
# Spatial Augmentations
# ==================================================================================

class HorizontalFlip:
    """
    Random horizontal flip for video frames
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            return np.array([cv2.flip(f, 1) for f in frames])
        return frames


class RandomCrop:
    """
    Random crop with resize back to original size
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, scale=(0.8, 1.0), ratio=(0.9, 1.1)):
        self.scale = scale
        self.ratio = ratio
    
    def __call__(self, frames):
        t, h, w, c = frames.shape
        
        # Random scale and ratio
        scale = random.uniform(*self.scale)
        ratio = random.uniform(*self.ratio)
        
        new_h = int(h * scale)
        new_w = int(new_h * ratio)
        new_w = min(new_w, w)
        
        # Random position
        top = random.randint(0, h - new_h)
        left = random.randint(0, w - new_w)
        
        # Crop and resize
        cropped = frames[:, top:top+new_h, left:left+new_w, :]
        resized = np.array([cv2.resize(f, (w, h)) for f in cropped])
        
        return resized


class RandomRotation:
    """
    Random rotation within degree range
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, degrees=15):
        self.degrees = degrees
    
    def __call__(self, frames):
        angle = random.uniform(-self.degrees, self.degrees)
        
        t, h, w, c = frames.shape
        center = (w // 2, h // 2)
        matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        rotated = np.array([cv2.warpAffine(f, matrix, (w, h)) for f in frames])
        return rotated


class RandomScale:
    """
    Random zoom in/out
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, scale_range=(0.9, 1.1)):
        self.scale_range = scale_range
    
    def __call__(self, frames):
        scale = random.uniform(*self.scale_range)
        
        t, h, w, c = frames.shape
        new_h, new_w = int(h * scale), int(w * scale)
        
        scaled = []
        for frame in frames:
            resized = cv2.resize(frame, (new_w, new_h))
            
            if scale > 1:
                # Crop center
                y1 = (new_h - h) // 2
                x1 = (new_w - w) // 2
                result = resized[y1:y1+h, x1:x1+w]
            else:
                # Pad
                result = np.zeros((h, w, c), dtype=frame.dtype)
                y1 = (h - new_h) // 2
                x1 = (w - new_w) // 2
                result[y1:y1+new_h, x1:x1+new_w] = resized
            
            scaled.append(result)
        
        return np.array(scaled)


# ==================================================================================
# Color Augmentations
# ==================================================================================

class ColorJitter:
    """
    Random color jittering
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1):
        self.brightness = brightness
        self.contrast = contrast
        self.saturation = saturation
        self.hue = hue
    
    def __call__(self, frames):
        # Brightness
        if self.brightness > 0 and random.random() > 0.5:
            factor = 1 + random.uniform(-self.brightness, self.brightness)
            frames = np.clip(frames * factor, 0, 255).astype(np.uint8)
        
        # Contrast
        if self.contrast > 0 and random.random() > 0.5:
            factor = 1 + random.uniform(-self.contrast, self.contrast)
            mean = np.mean(frames)
            frames = np.clip((frames - mean) * factor + mean, 0, 255).astype(np.uint8)
        
        # Saturation (requires HSV conversion)
        if self.saturation > 0 and random.random() > 0.5:
            factor = 1 + random.uniform(-self.saturation, self.saturation)
            hsv_frames = []
            for frame in frames:
                hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV).astype(np.float32)
                hsv[:, :, 1] = np.clip(hsv[:, :, 1] * factor, 0, 255)
                hsv_frames.append(cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB))
            frames = np.array(hsv_frames)
        
        return frames


class RandomGrayscale:
    """
    Random grayscale conversion
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, p=0.1):
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            gray_frames = []
            for frame in frames:
                gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
                gray_frames.append(cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB))
            return np.array(gray_frames)
        return frames


class Normalize:
    """
    Normalize frames to [0, 1] or with mean/std
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, mean=None, std=None):
        self.mean = mean if mean else [0.485, 0.456, 0.406]
        self.std = std if std else [0.229, 0.224, 0.225]
    
    def __call__(self, frames):
        frames = frames.astype(np.float32) / 255.0
        
        for i in range(3):
            frames[:, :, :, i] = (frames[:, :, :, i] - self.mean[i]) / self.std[i]
        
        return frames


# ==================================================================================
# Temporal Augmentations
# ==================================================================================

class TemporalRandomCrop:
    """
    Random temporal crop of video
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, num_frames):
        self.num_frames = num_frames
    
    def __call__(self, frames):
        t = len(frames)
        
        if t <= self.num_frames:
            # Pad if too short
            if t < self.num_frames:
                pad_size = self.num_frames - t
                frames = np.concatenate([frames, frames[-1:].repeat(pad_size, axis=0)])
            return frames
        
        start = random.randint(0, t - self.num_frames)
        return frames[start:start + self.num_frames]


class SpeedChange:
    """
    Random speed change (slow motion / fast forward)
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, speed_range=(0.5, 2.0)):
        self.speed_range = speed_range
    
    def __call__(self, frames):
        speed = random.uniform(*self.speed_range)
        
        t = len(frames)
        new_t = int(t / speed)
        
        if new_t < 2:
            new_t = 2
        
        indices = np.linspace(0, t - 1, new_t).astype(int)
        return frames[indices]


class TemporalReverse:
    """
    Random temporal reversal
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            return frames[::-1].copy()
        return frames


# ==================================================================================
# Noise and Effects
# ==================================================================================

class GaussianNoise:
    """
    Add Gaussian noise to frames
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, mean=0, std=10, p=0.5):
        self.mean = mean
        self.std = std
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            noise = np.random.normal(self.mean, self.std, frames.shape)
            frames = np.clip(frames + noise, 0, 255).astype(np.uint8)
        return frames


class GaussianBlur:
    """
    Apply Gaussian blur
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, kernel_size=5, p=0.3):
        self.kernel_size = kernel_size
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            k = random.choice([3, 5, 7])
            frames = np.array([cv2.GaussianBlur(f, (k, k), 0) for f in frames])
        return frames


class RandomCutout:
    """
    Random rectangular cutout
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, n_holes=1, hole_size_ratio=0.2, p=0.5):
        self.n_holes = n_holes
        self.hole_size_ratio = hole_size_ratio
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            t, h, w, c = frames.shape
            hole_h = int(h * self.hole_size_ratio)
            hole_w = int(w * self.hole_size_ratio)
            
            for _ in range(self.n_holes):
                y = random.randint(0, h - hole_h)
                x = random.randint(0, w - hole_w)
                frames[:, y:y+hole_h, x:x+hole_w, :] = 0
        
        return frames


# ==================================================================================
# Compose Augmentations
# ==================================================================================

class Compose:
    """
    Compose multiple augmentations
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, transforms):
        self.transforms = transforms
    
    def __call__(self, frames):
        for t in self.transforms:
            frames = t(frames)
        return frames


class RandomApply:
    """
    Randomly apply a transformation
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, transform, p=0.5):
        self.transform = transform
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            return self.transform(frames)
        return frames


class OneOf:
    """
    Apply one of multiple transformations
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, transforms, p=0.5):
        self.transforms = transforms
        self.p = p
    
    def __call__(self, frames):
        if random.random() < self.p:
            transform = random.choice(self.transforms)
            return transform(frames)
        return frames


# ==================================================================================
# Pre-built Augmentation Pipelines
# ==================================================================================

def get_train_augmentation(num_frames=16):
    """
    Get standard training augmentation pipeline
    
    RSK World (https://rskworld.in)
    """
    return Compose([
        TemporalRandomCrop(num_frames),
        HorizontalFlip(p=0.5),
        RandomCrop(scale=(0.8, 1.0)),
        ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        RandomApply(GaussianBlur(p=1.0), p=0.2),
        RandomApply(GaussianNoise(std=5, p=1.0), p=0.1),
    ])


def get_strong_augmentation(num_frames=16):
    """
    Get strong augmentation pipeline for data-hungry models
    
    RSK World (https://rskworld.in)
    """
    return Compose([
        TemporalRandomCrop(num_frames),
        RandomApply(SpeedChange(speed_range=(0.7, 1.3)), p=0.3),
        HorizontalFlip(p=0.5),
        RandomApply(TemporalReverse(), p=0.2),
        RandomCrop(scale=(0.7, 1.0)),
        RandomRotation(degrees=15),
        ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
        OneOf([
            GaussianBlur(p=1.0),
            GaussianNoise(std=10, p=1.0),
        ], p=0.3),
        RandomCutout(n_holes=1, hole_size_ratio=0.15, p=0.3),
    ])


def get_val_transform(num_frames=16):
    """
    Get validation/test transform (no augmentation)
    
    RSK World (https://rskworld.in)
    """
    return Compose([
        TemporalRandomCrop(num_frames),
    ])


# ==================================================================================
# Demo
# ==================================================================================

def demo():
    """Demonstrate augmentations"""
    print("=" * 60)
    print("Action Recognition - Data Augmentation Demo")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("=" * 60)
    
    # Create dummy video frames
    frames = np.random.randint(0, 255, (16, 112, 112, 3), dtype=np.uint8)
    print(f"\nOriginal frames shape: {frames.shape}")
    
    # Apply augmentations
    train_aug = get_train_augmentation(num_frames=16)
    augmented = train_aug(frames.copy())
    print(f"Augmented frames shape: {augmented.shape}")
    
    # List available augmentations
    print("\n" + "=" * 60)
    print("Available Augmentations:")
    print("=" * 60)
    
    augmentations = [
        ("HorizontalFlip", "Random horizontal flip"),
        ("RandomCrop", "Random crop with resize"),
        ("RandomRotation", "Random rotation"),
        ("RandomScale", "Random zoom in/out"),
        ("ColorJitter", "Brightness, contrast, saturation"),
        ("RandomGrayscale", "Convert to grayscale"),
        ("TemporalRandomCrop", "Random temporal segment"),
        ("SpeedChange", "Speed up or slow down"),
        ("TemporalReverse", "Reverse playback"),
        ("GaussianNoise", "Add random noise"),
        ("GaussianBlur", "Apply blur effect"),
        ("RandomCutout", "Random rectangular cutout"),
    ]
    
    for name, desc in augmentations:
        print(f"  {name:<20} - {desc}")
    
    print("\n" + "=" * 60)
    print("Pre-built Pipelines:")
    print("=" * 60)
    print("  get_train_augmentation()  - Standard training pipeline")
    print("  get_strong_augmentation() - Heavy augmentation")
    print("  get_val_transform()       - Validation (no augmentation)")
    
    print("\n" + "=" * 60)
    print("Usage Example:")
    print("=" * 60)
    print("""
    from augmentation import get_train_augmentation, Compose, HorizontalFlip
    
    # Use pre-built pipeline
    transform = get_train_augmentation(num_frames=16)
    augmented_frames = transform(frames)
    
    # Or create custom pipeline
    transform = Compose([
        HorizontalFlip(p=0.5),
        ColorJitter(brightness=0.2),
    ])
    """)
    
    print("\nThank you for using RSK World!")
    print("Visit: https://rskworld.in")


if __name__ == "__main__":
    demo()