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