"""
Quick Start Example - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Quick example demonstrating PyTorch neural network basics
"""

import torch
import torch.nn as nn
import torch.optim as optim
from models.basic_nn import BasicNeuralNetwork
from training.utils import generate_sample_data, plot_training_history
from training.trainer import Trainer
from torch.utils.data import DataLoader, TensorDataset


def quick_example():
    """
    Quick example demonstrating a complete training workflow
    """
    print("=" * 60)
    print("PyTorch Neural Networks - Quick Start Example")
    print("Author: RSK World (https://rskworld.in)")
    print("=" * 60)
    
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"\nUsing device: {device}\n")
    
    # Generate sample data
    print("Generating sample data...")
    X_train, y_train, X_test, y_test = generate_sample_data(
        n_samples=1000, n_features=20, n_classes=3
    )
    print(f"Training samples: {X_train.shape[0]}")
    print(f"Test samples: {X_test.shape[0]}\n")
    
    # Create data loaders
    train_dataset = TensorDataset(X_train, y_train)
    test_dataset = TensorDataset(X_test, y_test)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
    
    # Create model
    print("Creating neural network model...")
    model = BasicNeuralNetwork(
        input_size=20,
        hidden_size=64,
        output_size=3,
        num_layers=2
    ).to(device)
    print(f"Model parameters: {sum(p.numel() for p in model.parameters())}\n")
    
    # Define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # Train model
    print("Starting training...")
    trainer = Trainer(model, criterion, optimizer, device)
    history = trainer.train(train_loader, test_loader, epochs=10)
    
    # Plot results
    print("\nGenerating training plots...")
    plot_training_history(history)
    
    print("\n" + "=" * 60)
    print("Example completed successfully!")
    print("=" * 60)


if __name__ == '__main__':
    quick_example()