"""
Hyperparameter Tuning Example - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Example demonstrating hyperparameter tuning
"""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import sys

sys.path.append('..')

from models.basic_nn import BasicNeuralNetwork
from training.utils import generate_sample_data
from utils.hyperparameter_tuning import GridSearch, RandomSearch


def create_model(input_size=20, hidden_size=64, output_size=3, num_layers=2, dropout=0.2):
    """Create model with specified parameters"""
    return BasicNeuralNetwork(
        input_size=input_size,
        hidden_size=hidden_size,
        output_size=output_size,
        num_layers=num_layers,
        dropout=dropout
    )


def grid_search_example():
    """
    Demonstrate grid search
    """
    print("=" * 60)
    print("Grid Search Example")
    print("Author: RSK World (https://rskworld.in)")
    print("=" * 60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Generate data
    X_train, y_train, X_val, y_val = generate_sample_data(
        n_samples=500, n_features=20, n_classes=3
    )
    
    train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
    val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32, shuffle=False)
    
    # Define parameter grid
    param_grid = {
        'input_size': [20],
        'hidden_size': [32, 64, 128],
        'output_size': [3],
        'num_layers': [1, 2],
        'dropout': [0.1, 0.2]
    }
    
    # Create model class wrapper
    class ModelWrapper:
        def __init__(self, input_size, hidden_size, output_size, num_layers, dropout):
            self.model = create_model(input_size, hidden_size, output_size, num_layers, dropout)
        
        def to(self, device):
            self.model = self.model.to(device)
            return self
    
    # Perform grid search
    print("\nStarting grid search...")
    grid_search = GridSearch(
        param_grid=param_grid,
        model_class=create_model,
        train_loader=train_loader,
        val_loader=val_loader,
        device=device
    )
    
    best_params, best_score, results = grid_search.search(epochs=5, verbose=True)
    
    print(f"\nBest parameters: {best_params}")
    print(f"Best score: {best_score:.4f}")
    
    # Save results
    grid_search.save_results('../hyperparameter_search_results.json')
    
    print("\n" + "=" * 60)
    print("Grid search completed!")
    print("=" * 60)


def random_search_example():
    """
    Demonstrate random search
    """
    print("=" * 60)
    print("Random Search Example")
    print("Author: RSK World (https://rskworld.in)")
    print("=" * 60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Generate data
    X_train, y_train, X_val, y_val = generate_sample_data(
        n_samples=500, n_features=20, n_classes=3
    )
    
    train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
    val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32, shuffle=False)
    
    # Define parameter distributions
    param_distributions = {
        'input_size': [20],
        'hidden_size': [32, 64, 128, 256],
        'output_size': [3],
        'num_layers': [1, 2, 3],
        'dropout': [(0.0, 0.5)]  # Uniform distribution between 0.0 and 0.5
    }
    
    # Perform random search
    print("\nStarting random search...")
    random_search = RandomSearch(
        param_distributions=param_distributions,
        model_class=create_model,
        train_loader=train_loader,
        val_loader=val_loader,
        device=device,
        n_iter=10
    )
    
    best_params, best_score, results = random_search.search(epochs=5, verbose=True)
    
    print(f"\nBest parameters: {best_params}")
    print(f"Best score: {best_score:.4f}")
    
    print("\n" + "=" * 60)
    print("Random search completed!")
    print("=" * 60)


if __name__ == '__main__':
    print("Choose search method:")
    print("1. Grid Search")
    print("2. Random Search")
    
    choice = input("Enter choice (1 or 2): ")
    
    if choice == '1':
        grid_search_example()
    elif choice == '2':
        random_search_example()
    else:
        print("Invalid choice. Running grid search by default.")
        grid_search_example()

"""
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