"""
forgeNN MNIST Classification Example
===================================

Complete example demonstrating high-performance neural network training
with forgeNN's vectorized implementation.

Features:
- MNIST handwritten digit classification
- Vectorized operations for fast training
- Batch processing with progress tracking
- Professional metrics and evaluation

Performance: 93%+ accuracy in under 2 seconds!
"""

import numpy as np
import time
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from forgeNN.tensor import Tensor
from forgeNN.vectorized import VectorizedMLP, VectorizedOptimizer, cross_entropy_loss, accuracy

def load_mnist_vectorized(n_samples=5000):
    """Load and preprocess MNIST for vectorized training."""
    print("Loading MNIST dataset...")
    mnist = fetch_openml('mnist_784', version=1, as_frame=False)
    X, y = mnist.data, mnist.target.astype(int)
    
    # Use more samples since vectorized version is much faster
    print(f"Using {n_samples} samples for vectorized training...")
    X, y = X[:n_samples], y[:n_samples]
    
    # Normalize features
    X = X.astype(np.float32) / 255.0
    
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    print(f"Training samples: {len(X_train)}")
    print(f"Test samples: {len(X_test)}")
    print(f"Input features: {X_train.shape[1]}")
    print(f"Classes: {len(np.unique(y))}")
    
    return X_train, X_test, y_train, y_test

def create_data_loader(X, y, batch_size=32, shuffle=True):
    """Create simple data loader for batch training."""
    n_samples = len(X)
    if shuffle:
        indices = np.random.permutation(n_samples)
        X, y = X[indices], y[indices]
    
    for i in range(0, n_samples, batch_size):
        end_idx = min(i + batch_size, n_samples)
        yield X[i:end_idx], y[i:end_idx]

def train_epoch(model, optimizer, X_train, y_train, batch_size=64):
    """Train for one epoch using vectorized operations."""
    total_loss = 0.0
    total_acc = 0.0
    num_batches = 0
    
    for batch_x, batch_y in create_data_loader(X_train, y_train, batch_size):
        # Convert to tensors
        x_tensor = Tensor(batch_x)
        
        # Forward pass
        logits = model(x_tensor)
        
        # Compute loss and accuracy
        loss = cross_entropy_loss(logits, batch_y)
        acc = accuracy(logits, batch_y)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Accumulate metrics
        total_loss += loss.data
        total_acc += acc
        num_batches += 1
    
    return total_loss / num_batches, total_acc / num_batches

def evaluate(model, X_test, y_test, batch_size=64):
    """Evaluate model on test set."""
    total_loss = 0.0
    total_acc = 0.0
    num_batches = 0
    
    for batch_x, batch_y in create_data_loader(X_test, y_test, batch_size, shuffle=False):
        x_tensor = Tensor(batch_x, requires_grad=False)
        logits = model(x_tensor)
        
        loss = cross_entropy_loss(logits, batch_y)
        acc = accuracy(logits, batch_y)
        
        total_loss += loss.data
        total_acc += acc
        num_batches += 1
    
    return total_loss / num_batches, total_acc / num_batches

def main():
    print("="*60)
    print("VECTORIZED MNIST CLASSIFICATION WITH forgeNN")
    print("="*60)
    
    # Load data
    X_train, X_test, y_train, y_test = load_mnist_vectorized(n_samples=5000)
    
    # Create model
    print("\nCreating vectorized neural network...")
    model = VectorizedMLP(
        input_size=784,
        hidden_sizes=[128, 64],  # Larger network since it's much faster
        output_size=10,
        activations=['relu', 'relu', 'linear']
    )
    
    print(f"Model parameters: {len(model.parameters())}")
    total_params = sum(p.data.size for p in model.parameters())
    print(f"Total parameter count: {total_params:,}")
    
    # Create optimizer
    optimizer = VectorizedOptimizer(model.parameters(), lr=0.01, momentum=0.9)
    
    # Training configuration
    epochs = 10
    batch_size = 64
    
    print(f"\nTraining Configuration:")
    print(f"  Epochs: {epochs}")
    print(f"  Batch size: {batch_size}")
    print(f"  Learning rate: {optimizer.lr}")
    print(f"  Momentum: {optimizer.momentum}")
    
    print("\n" + "="*60)
    print("TRAINING")
    print("="*60)
    
    # Training loop
    start_time = time.time()
    
    for epoch in range(epochs):
        print(f"\nEpoch {epoch + 1}/{epochs}")
        
        # Train
        train_loss, train_acc = train_epoch(model, optimizer, X_train, y_train, batch_size)
        
        # Evaluate
        test_loss, test_acc = evaluate(model, X_test, y_test, batch_size)
        
        # Print metrics
        print(f"  Train: Loss = {train_loss:.4f}, Acc = {train_acc*100:.1f}%")
        print(f"  Test:  Loss = {test_loss:.4f}, Acc = {test_acc*100:.1f}%")
    
    training_time = time.time() - start_time
    
    print("\n" + "="*60)
    print("FINAL RESULTS")
    print("="*60)
    
    # Final evaluation
    final_test_loss, final_test_acc = evaluate(model, X_test, y_test, batch_size)
    
    print(f"Final Test Accuracy: {final_test_acc*100:.2f}%")
    print(f"Final Test Loss: {final_test_loss:.4f}")
    print(f"Training Time: {training_time:.1f} seconds")
    print(f"Samples per Second: {len(X_train) * epochs / training_time:.0f}")
    print(f"Total Parameters: {total_params:,}")
    
    # Performance comparison
    samples_per_epoch = len(X_train)
    time_per_sample = training_time / (samples_per_epoch * epochs)
    print(f"Time per Sample: {time_per_sample*1000:.2f} ms")
    
    # Show some predictions
    print("\nSample Predictions:")
    sample_indices = np.random.choice(len(X_test), 5, replace=False)
    for i, idx in enumerate(sample_indices):
        x_sample = Tensor(X_test[idx:idx+1], requires_grad=False)
        logits = model(x_sample)
        probs = logits.softmax(axis=1)
        predicted = np.argmax(probs.data[0])
        confidence = np.max(probs.data[0])
        
        print(f"  Sample {i+1}: True = {y_test[idx]}, "
              f"Predicted = {predicted}, Confidence = {confidence:.3f}")
    
    print("\n" + "="*60)
    print("PERFORMANCE ANALYSIS")
    print("="*60)
    
    print(f"✅ Vectorized implementation achieved:")
    print(f"   • {final_test_acc*100:.1f}% accuracy on MNIST")
    print(f"   • {samples_per_epoch * epochs / training_time:.0f} samples/second")
    print(f"   • {time_per_sample*1000:.2f} ms per sample")
    print(f"   • Trained {total_params:,} parameters efficiently")
    

if __name__ == "__main__":
    main()
