"""
CW4: Convolutional Neural Networks.

In CW2 and CW3, you implemented models using numpy and manual backpropagation.
Now you will use PyTorch, which provides:
    - nn.Module: base class for all neural network layers
    - Automatic differentiation: loss.backward() computes ALL gradients
      (this is what you implemented by hand in CW3!)

Your task: Define CNN architectures using PyTorch's nn module.

Key PyTorch layers you may use:
    nn.Conv2d(in_channels, out_channels, kernel_size, padding=...)
    nn.MaxPool2d(kernel_size)
    nn.BatchNorm2d(num_features)
    nn.Dropout(p)
    nn.Linear(in_features, out_features)
    nn.ReLU()
    nn.Flatten()
    nn.Sequential(...)
"""

import torch
import torch.nn as nn


class SimpleCNN(nn.Module):
    """
    A simple CNN for FashionMNIST classification.

    TODO: Define and implement a CNN architecture.

    Input: (batch_size, 1, 28, 28) grayscale images
    Output: (batch_size, 10) class logits

    Suggested architecture (you may modify this):
        Conv2d(1, 32, 3, padding=1) -> ReLU -> MaxPool2d(2)     → (B, 32, 14, 14)
        Conv2d(32, 64, 3, padding=1) -> ReLU -> MaxPool2d(2)    → (B, 64, 7, 7)
        Flatten                                                   → (B, 64*7*7)
        Linear(64*7*7, 128) -> ReLU                              → (B, 128)
        Linear(128, 10)                                          → (B, 10)

    You should experiment with:
        - Number and size of convolutional layers
        - Adding BatchNorm2d after conv layers
        - Adding Dropout for regularization
        - Different pooling strategies
    """

    def __init__(self, num_classes=10):
        super().__init__()
        # ====================================================================
        # TODO: Define your CNN layers
        #
        # Example (you can modify this):
        #   self.features = nn.Sequential(
        #       nn.Conv2d(1, 32, kernel_size=3, padding=1),
        #       nn.ReLU(),
        #       nn.MaxPool2d(2),
        #       nn.Conv2d(32, 64, kernel_size=3, padding=1),
        #       nn.ReLU(),
        #       nn.MaxPool2d(2),
        #   )
        #   self.classifier = nn.Sequential(
        #       nn.Flatten(),
        #       nn.Linear(64 * 7 * 7, 128),
        #       nn.ReLU(),
        #       nn.Linear(128, num_classes),
        #   )
        # ====================================================================
        raise NotImplementedError("TODO: Define SimpleCNN layers")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def forward(self, x):
        """
        Forward pass of the CNN.

        TODO: Implement the forward pass.

        Args:
            x: (batch_size, 1, 28, 28) input images

        Returns:
            logits: (batch_size, num_classes) raw class scores

        Hints:
            - Pass x through your feature extraction layers
            - Pass the result through your classifier layers
            - Do NOT apply softmax here (it's included in nn.CrossEntropyLoss)
        """
        # ====================================================================
        # TODO: Implement forward pass
        # ====================================================================
        raise NotImplementedError("TODO: Implement SimpleCNN forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================


class DeepCNN(nn.Module):
    """
    A deeper CNN with advanced techniques (Bonus).

    TODO: Implement a more sophisticated CNN architecture.

    Target: >= 92% test accuracy on FashionMNIST.

    Suggestions:
        - Use more convolutional layers (3-5 conv blocks)
        - Add BatchNorm2d after each conv layer (before activation)
        - Use Dropout for regularization
        - Consider residual connections (skip connections)
        - Consider global average pooling instead of flatten

    Input: (batch_size, 1, 28, 28)
    Output: (batch_size, 10)
    """

    def __init__(self, num_classes=10):
        super().__init__()
        # ====================================================================
        # TODO: Define your deeper CNN architecture (Bonus)
        # ====================================================================
        raise NotImplementedError("TODO: Define DeepCNN layers (Bonus)")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def forward(self, x):
        """
        TODO: Implement forward pass for DeepCNN (Bonus).
        """
        # ====================================================================
        # TODO: Implement DeepCNN forward pass (Bonus)
        # ====================================================================
        raise NotImplementedError("TODO: Implement DeepCNN forward (Bonus)")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================


def get_model(model_type, num_classes=10):
    """
    Factory function to create model by name (provided, do not modify).

    Args:
        model_type: "simple_cnn" or "deep_cnn"
        num_classes: Number of output classes

    Returns:
        nn.Module instance
    """
    models = {
        "simple_cnn": SimpleCNN,
        "deep_cnn": DeepCNN,
    }
    if model_type not in models:
        raise ValueError(f"Unknown model type: {model_type}. "
                         f"Available: {list(models.keys())}")
    return models[model_type](num_classes=num_classes)
