"""
Data Augmentation Utilities - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Data augmentation utilities for images and sequences
"""

import torch
import torchvision.transforms as transforms
import numpy as np


def get_image_augmentation(train=True):
    """
    Get image augmentation transforms
    
    Args:
        train: Whether to apply training augmentations
        
    Returns:
        Transform composition
    """
    if train:
        return transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomRotation(degrees=15),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
            transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    else:
        return transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])


def add_noise(tensor, noise_factor=0.1):
    """
    Add Gaussian noise to tensor
    
    Args:
        tensor: Input tensor
        noise_factor: Noise intensity
        
    Returns:
        Tensor with added noise
    """
    noise = torch.randn_like(tensor) * noise_factor
    return tensor + noise


def random_crop(tensor, crop_size):
    """
    Random crop for tensor
    
    Args:
        tensor: Input tensor (C, H, W)
        crop_size: Size of crop (height, width)
        
    Returns:
        Cropped tensor
    """
    _, h, w = tensor.shape
    th, tw = crop_size
    
    if h < th or w < tw:
        return tensor
    
    i = torch.randint(0, h - th + 1, (1,)).item()
    j = torch.randint(0, w - tw + 1, (1,)).item()
    
    return tensor[:, i:i+th, j:j+tw]


def sequence_augmentation(sequence, noise_factor=0.05, dropout_prob=0.1):
    """
    Augment sequence data
    
    Args:
        sequence: Input sequence tensor (seq_len, features)
        noise_factor: Noise intensity
        dropout_prob: Probability of dropping features
        
    Returns:
        Augmented sequence
    """
    # Add noise
    if noise_factor > 0:
        noise = torch.randn_like(sequence) * noise_factor
        sequence = sequence + noise
    
    # Random feature dropout
    if dropout_prob > 0:
        mask = torch.rand(sequence.shape) > dropout_prob
        sequence = sequence * mask.float()
    
    return sequence


class MixUp:
    """
    MixUp data augmentation
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, alpha=0.2):
        """
        Initialize MixUp
        
        Args:
            alpha: Beta distribution parameter
        """
        self.alpha = alpha
    
    def __call__(self, x1, y1, x2, y2):
        """
        Apply MixUp
        
        Args:
            x1, y1: First sample and label
            x2, y2: Second sample and label
            
        Returns:
            Mixed sample and labels
        """
        lam = np.random.beta(self.alpha, self.alpha)
        mixed_x = lam * x1 + (1 - lam) * x2
        y_a, y_b = y1, y2
        return mixed_x, y_a, y_b, lam