"""
CW3: Multi-Layer Perceptron (MLP).

The MLP is a sequence of layers:
    Linear -> Activation -> Linear -> Activation -> ... -> Linear

In CW2, you implemented a single Linear layer (logistic regression).
Now you will stack multiple layers to create a deeper model.

The key insight: forward() passes data through layers sequentially,
and backward() propagates gradients through layers in REVERSE order.
This is the chain rule in action!
"""

import numpy as np
from layers import Linear, ReLU, Sigmoid


class MLP:
    """
    Multi-Layer Perceptron with configurable depth and width.

    Architecture for hidden_dims=[128, 64], activation='relu':
        Linear(784, 128) -> ReLU -> Linear(128, 64) -> ReLU -> Linear(64, 10)
    """

    def __init__(self, input_dim, hidden_dims, num_classes, activation='relu'):
        """
        Build the MLP as a list of layers.

        TODO: Construct self.layers as a list of layer objects.

        Args:
            input_dim: Number of input features (784 for FashionMNIST)
            hidden_dims: List of hidden layer sizes, e.g., [128, 64]
            num_classes: Number of output classes (10)
            activation: Activation function type: "relu" or "sigmoid"

        Hints:
            - Create layers in order: Linear -> Activation -> Linear -> Activation -> ... -> Linear
            - For hidden_dims = [128, 64]:
                Layer 0: Linear(784, 128)
                Layer 1: ReLU()  (or Sigmoid())
                Layer 2: Linear(128, 64)
                Layer 3: ReLU()  (or Sigmoid())
                Layer 4: Linear(64, 10)
            - The last layer is Linear WITHOUT activation (softmax is in the loss)
            - Store all layers in self.layers (a list)

        Example:
            self.layers = [
                Linear(784, 128), ReLU(),
                Linear(128, 64), ReLU(),
                Linear(64, 10)
            ]
        """
        # ====================================================================
        # TODO: Build the network layers
        # ====================================================================
        self.layers = []
        raise NotImplementedError("TODO: Build MLP layers")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def forward(self, X):
        """
        Forward pass: pass input through all layers sequentially.

        TODO: Implement the forward pass.

        Args:
            X: (batch_size, input_dim) input features

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

        Hints:
            - Pass X through each layer in self.layers in order
            - The output of one layer becomes the input to the next
            - Return the final output (logits, before softmax)
        """
        # ====================================================================
        # TODO: Implement forward pass through all layers
        # ====================================================================
        raise NotImplementedError("TODO: Implement MLP forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def backward(self, grad_output):
        """
        Backward pass: propagate gradients through layers in reverse order.

        TODO: Implement the backward pass.

        Args:
            grad_output: (batch_size, num_classes) gradient of loss w.r.t. logits
                         (from loss.backward())

        Hints:
            - Iterate through self.layers in REVERSE order
            - For each layer, call layer.backward(grad) and pass the result
              to the previous layer
            - This is the chain rule! Each layer computes its local gradient
              and passes the gradient to the next layer in the chain.

        Example:
            grad = grad_output
            for layer in reversed(self.layers):
                grad = layer.backward(grad)
        """
        # ====================================================================
        # TODO: Implement backward pass through all layers (reverse order)
        # ====================================================================
        raise NotImplementedError("TODO: Implement MLP backward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def parameters(self):
        """
        Collect all trainable parameters from Linear layers (provided).

        Returns:
            List of dicts, each with keys 'W', 'b', 'layer'.

            'W' and 'b' are direct array references — the optimizer updates them
            in-place (W -= lr * grad_W), so changes are reflected in the layer.

            'layer' is the Linear layer object itself. The optimizer reads
            layer.grad_W and layer.grad_b through this reference at step time,
            rather than caching the gradient arrays at initialization. This is
            necessary because backward() assigns new arrays to layer.grad_W and
            layer.grad_b (e.g., self.grad_W = X.T @ grad_output), which would
            invalidate any reference captured before backward() ran.
        """
        params = []
        for layer in self.layers:
            if isinstance(layer, Linear):
                params.append({
                    'W': layer.W,
                    'b': layer.b,
                    'layer': layer,
                })
        return params

    def num_parameters(self):
        """Count total number of trainable parameters (provided)."""
        total = 0
        for layer in self.layers:
            if isinstance(layer, Linear):
                total += layer.W.size + layer.b.size
        return total
