"""
PyTorch Neural Networks - Main Entry Point
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Main script to run different neural network models
"""

import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt

from models.basic_nn import BasicNeuralNetwork
from models.cnn import SimpleCNN
from models.rnn import SimpleRNN
from training.trainer import Trainer
from training.utils import generate_sample_data, plot_training_history


def main():
    parser = argparse.ArgumentParser(description='PyTorch Neural Networks - RSK World')
    parser.add_argument('--model', type=str, default='basic', 
                       choices=['basic', 'cnn', 'rnn'],
                       help='Model type to train')
    parser.add_argument('--epochs', type=int, default=10,
                       help='Number of training epochs')
    parser.add_argument('--batch_size', type=int, default=32,
                       help='Batch size for training')
    parser.add_argument('--lr', type=float, default=0.001,
                       help='Learning rate')
    parser.add_argument('--device', type=str, default='auto',
                       choices=['auto', 'cpu', 'cuda'],
                       help='Device to use for training')
    
    args = parser.parse_args()
    
    # Set device
    if args.device == 'auto':
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        device = torch.device(args.device)
    
    print(f"Using device: {device}")
    print(f"Training {args.model} model for {args.epochs} epochs")
    
    # Generate sample data based on model type
    if args.model == 'basic':
        X_train, y_train, X_test, y_test = generate_sample_data(
            n_samples=1000, n_features=20, n_classes=3
        )
        model = BasicNeuralNetwork(input_size=20, hidden_size=64, output_size=3)
        
    elif args.model == 'cnn':
        # Generate image-like data (batch, channels, height, width)
        X_train = torch.randn(200, 1, 28, 28)
        y_train = torch.randint(0, 10, (200,))
        X_test = torch.randn(50, 1, 28, 28)
        y_test = torch.randint(0, 10, (50,))
        model = SimpleCNN(num_classes=10)
        
    elif args.model == 'rnn':
        # Generate sequence data (batch, seq_len, features)
        X_train = torch.randn(200, 10, 5)
        y_train = torch.randint(0, 3, (200,))
        X_test = torch.randn(50, 10, 5)
        y_test = torch.randint(0, 3, (50,))
        model = SimpleRNN(input_size=5, hidden_size=64, num_layers=2, num_classes=3)
    
    model = model.to(device)
    
    # Create data loaders
    train_dataset = TensorDataset(X_train, y_train)
    test_dataset = TensorDataset(X_test, y_test)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
    
    # Define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    
    # Train model
    trainer = Trainer(model, criterion, optimizer, device)
    history = trainer.train(train_loader, test_loader, epochs=args.epochs)
    
    # Plot training history
    plot_training_history(history)
    
    print("\nTraining completed!")
    print(f"Final Training Loss: {history['train_loss'][-1]:.4f}")
    print(f"Final Test Loss: {history['test_loss'][-1]:.4f}")
    print(f"Final Test Accuracy: {history['test_accuracy'][-1]:.2f}%")


if __name__ == '__main__':
    main()