"""
Model Module for Verril-Learn

Contains neural network architectures optimized for MNIST classification.
These models are designed to be lightweight for edge device deployment
while still being complex enough to demonstrate unlearning effects.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple


class SimpleCNN(nn.Module):
    """
    A simple Convolutional Neural Network for MNIST classification.
    
    Architecture:
        - Conv2D (1 → 32 channels, 3x3 kernel)
        - ReLU + MaxPool2D (2x2)
        - Conv2D (32 → 64 channels, 3x3 kernel)
        - ReLU + MaxPool2D (2x2)
        - Flatten
        - Linear (1600 → 128)
        - ReLU + Dropout (0.5)
        - Linear (128 → 10)
    
    This architecture is intentionally simple to:
        1. Enable fast training on edge devices
        2. Make unlearning effects more observable
        3. Reduce memory footprint for Fisher Information computation
    
    Input: (batch_size, 1, 28, 28) - Grayscale MNIST images
    Output: (batch_size, 10) - Logits for 10 digit classes
    
    Example:
        >>> model = SimpleCNN()
        >>> x = torch.randn(32, 1, 28, 28)  # Batch of 32 images
        >>> logits = model(x)
        >>> print(logits.shape)  # torch.Size([32, 10])
    """
    
    def __init__(self, dropout_rate: float = 0.5):
        """
        Initialize the SimpleCNN.
        
        Args:
            dropout_rate: Dropout probability for regularization
        """
        super(SimpleCNN, self).__init__()
        
        # Convolutional layers
        self.conv1 = nn.Conv2d(
            in_channels=1,
            out_channels=32,
            kernel_size=3,
            stride=1,
            padding=1
        )
        self.conv2 = nn.Conv2d(
            in_channels=32,
            out_channels=64,
            kernel_size=3,
            stride=1,
            padding=1
        )
        
        # Pooling layer
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Fully connected layers
        # After 2 pooling layers: 28x28 -> 14x14 -> 7x7
        # With 64 channels: 64 * 7 * 7 = 3136
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        
        # Dropout for regularization
        self.dropout = nn.Dropout(p=dropout_rate)
        
        # Initialize weights using Kaiming initialization
        self._initialize_weights()
    
    def _initialize_weights(self) -> None:
        """Initialize weights using Kaiming initialization for ReLU."""
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(
                    module.weight,
                    mode='fan_out',
                    nonlinearity='relu'
                )
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(
                    module.weight,
                    mode='fan_out',
                    nonlinearity='relu'
                )
                nn.init.zeros_(module.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the network.
        
        Args:
            x: Input tensor of shape (batch_size, 1, 28, 28)
        
        Returns:
            Logits tensor of shape (batch_size, 10)
        """
        # First conv block
        x = self.conv1(x)           # (B, 1, 28, 28) -> (B, 32, 28, 28)
        x = F.relu(x)
        x = self.pool(x)            # (B, 32, 28, 28) -> (B, 32, 14, 14)
        
        # Second conv block
        x = self.conv2(x)           # (B, 32, 14, 14) -> (B, 64, 14, 14)
        x = F.relu(x)
        x = self.pool(x)            # (B, 64, 14, 14) -> (B, 64, 7, 7)
        
        # Flatten
        x = x.view(x.size(0), -1)   # (B, 64, 7, 7) -> (B, 3136)
        
        # Fully connected layers
        x = self.fc1(x)             # (B, 3136) -> (B, 128)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)             # (B, 128) -> (B, 10)
        
        return x
    
    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extract feature representations before the final classification layer.
        
        Useful for analyzing what the model has learned and for
        understanding the effect of unlearning on internal representations.
        
        Args:
            x: Input tensor of shape (batch_size, 1, 28, 28)
        
        Returns:
            Feature tensor of shape (batch_size, 128)
        """
        x = self.conv1(x)
        x = F.relu(x)
        x = self.pool(x)
        
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = F.relu(x)
        
        return x
    
    def count_parameters(self) -> int:
        """
        Count the total number of trainable parameters.
        
        Returns:
            Total number of trainable parameters
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


def get_model_info(model: nn.Module) -> dict:
    """
    Get information about a model's architecture.
    
    Args:
        model: PyTorch model
    
    Returns:
        Dictionary with model information
    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    return {
        'total_parameters': total_params,
        'trainable_parameters': trainable_params,
        'layers': [name for name, _ in model.named_modules() if name],
        'parameter_shapes': {
            name: tuple(p.shape) 
            for name, p in model.named_parameters()
        }
    }


def train_model(
    model: nn.Module,
    train_loader: torch.utils.data.DataLoader,
    epochs: int = 5,
    learning_rate: float = 0.001,
    device: str = 'cpu',
    verbose: bool = True
) -> Tuple[nn.Module, list]:
    """
    Train the model on the provided data.
    
    Args:
        model: PyTorch model to train
        train_loader: DataLoader for training data
        epochs: Number of training epochs
        learning_rate: Learning rate for optimizer
        device: Device to train on ('cpu' or 'cuda')
        verbose: Whether to print training progress
    
    Returns:
        Tuple of (trained_model, loss_history)
    """
    model = model.to(device)
    model.train()
    
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    
    loss_history = []
    
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
        
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * correct / total
        loss_history.append(epoch_loss)
        
        if verbose:
            print(f"Epoch {epoch+1}/{epochs} | "
                  f"Loss: {epoch_loss:.4f} | "
                  f"Accuracy: {epoch_acc:.2f}%")
    
    return model, loss_history


def evaluate_model(
    model: nn.Module,
    test_loader: torch.utils.data.DataLoader,
    device: str = 'cpu',
    verbose: bool = True
) -> Tuple[float, float]:
    """
    Evaluate the model on test data.
    
    Args:
        model: PyTorch model to evaluate
        test_loader: DataLoader for test data
        device: Device to evaluate on
        verbose: Whether to print results
    
    Returns:
        Tuple of (loss, accuracy)
    """
    model = model.to(device)
    model.eval()
    
    criterion = nn.CrossEntropyLoss()
    test_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
    
    avg_loss = test_loss / len(test_loader)
    accuracy = 100. * correct / total
    
    if verbose:
        print(f"Test Loss: {avg_loss:.4f} | Test Accuracy: {accuracy:.2f}%")
    
    return avg_loss, accuracy
