"""
CW3: Neural Network Layers.

Each layer implements two methods:
    - forward(X):  compute output from input, cache values for backward
    - backward(grad_output):  compute gradients and return grad_input

This is the core of backpropagation: chaining backward() calls through
layers in reverse order propagates gradients from the loss all the way
back to the first layer.

CONNECTION TO CW2:
    In CW2, your optimizer computed:
        dW = X.T @ grad_logits
        db = grad_logits.sum(axis=0)
    This is EXACTLY what Linear.backward() does below!
    The difference is that now you also compute grad_input = grad_output @ W.T
    to propagate the gradient to the previous layer.
"""

import numpy as np


class Linear:
    """
    Fully connected (linear) layer: output = X @ W + b

    This is the fundamental building block of neural networks.
    """

    def __init__(self, in_features, out_features):
        """
        Initialize weights with Xavier initialization.

        Args:
            in_features: Number of input features
            out_features: Number of output features
        """
        scale = np.sqrt(2.0 / (in_features + out_features))
        self.W = np.random.randn(in_features, out_features).astype(np.float32) * scale
        self.b = np.zeros(out_features, dtype=np.float32)

        # Gradients (computed in backward, used by optimizer)
        self.grad_W = None
        self.grad_b = None

        # Cache for backward pass
        self.X_cache = None

    def forward(self, X):
        """
        Compute the linear transformation.

        TODO: Implement the forward pass.

        Args:
            X: (batch_size, in_features) input

        Returns:
            output: (batch_size, out_features) = X @ W + b

        Hints:
            - Save X in self.X_cache for use in backward()
            - This is the same as model.forward() in CW2!
        """
        # ====================================================================
        # TODO: Implement forward pass and cache X
        # ====================================================================
        raise NotImplementedError("TODO: Implement Linear forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def backward(self, grad_output):
        """
        Compute gradients and propagate to previous layer.

        TODO: Implement the backward pass.

        Args:
            grad_output: (batch_size, out_features) gradient of loss
                         w.r.t. this layer's output

        Returns:
            grad_input: (batch_size, in_features) gradient of loss
                        w.r.t. this layer's input (to pass to previous layer)

        Hints:
            Remember from CW2's optimizer, you computed:
                dW = X.T @ grad_logits
                db = grad_logits.sum(axis=0)
            That was the backward pass of a single Linear layer!

            Now generalize it:
                self.grad_W = self.X_cache.T @ grad_output
                self.grad_b = grad_output.sum(axis=0)
                grad_input = grad_output @ self.W.T   ← NEW: propagate further
                return grad_input
        """
        # ====================================================================
        # TODO: Compute grad_W, grad_b, and return grad_input
        # ====================================================================
        raise NotImplementedError("TODO: Implement Linear backward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================


class ReLU:
    """
    Rectified Linear Unit activation: output = max(0, X)

    ReLU is the most commonly used activation function in modern networks.
    """

    def __init__(self):
        self.X_cache = None

    def forward(self, X):
        """
        Apply ReLU activation.

        TODO: Implement ReLU forward pass.

        Args:
            X: (batch_size, features) input

        Returns:
            output: (batch_size, features) with negative values set to 0

        Hints:
            - output = np.maximum(0, X)
            - Save X in self.X_cache for backward
        """
        # ====================================================================
        # TODO: Implement ReLU forward
        # ====================================================================
        raise NotImplementedError("TODO: Implement ReLU forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def backward(self, grad_output):
        """
        Compute ReLU gradient.

        TODO: Implement ReLU backward pass.

        Args:
            grad_output: (batch_size, features) gradient from next layer

        Returns:
            grad_input: (batch_size, features) gradient to pass to previous layer

        Hints:
            - ReLU derivative: 1 if X > 0, else 0
            - grad_input = grad_output * (X_cache > 0)
        """
        # ====================================================================
        # TODO: Implement ReLU backward
        # ====================================================================
        raise NotImplementedError("TODO: Implement ReLU backward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================


class Sigmoid:
    """
    Sigmoid activation: output = 1 / (1 + exp(-X))

    An alternative to ReLU, historically important but less commonly used
    in hidden layers due to the vanishing gradient problem.
    """

    def __init__(self):
        self.output_cache = None

    def forward(self, X):
        """
        Apply Sigmoid activation.

        TODO: Implement Sigmoid forward pass.

        Args:
            X: (batch_size, features) input

        Returns:
            output: (batch_size, features) values in (0, 1)

        Hints:
            - output = 1 / (1 + np.exp(-X))
            - For numerical stability, handle large negative X:
              use np.clip(X, -500, 500) before exp
            - Save output (not X) in self.output_cache for backward
        """
        # ====================================================================
        # TODO: Implement Sigmoid forward
        # ====================================================================
        raise NotImplementedError("TODO: Implement Sigmoid forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def backward(self, grad_output):
        """
        Compute Sigmoid gradient.

        TODO: Implement Sigmoid backward pass.

        Args:
            grad_output: (batch_size, features) gradient from next layer

        Returns:
            grad_input: (batch_size, features) gradient to previous layer

        Hints:
            - Sigmoid derivative: output * (1 - output)
            - grad_input = grad_output * output * (1 - output)
            - This is why we cached the output, not the input!
        """
        # ====================================================================
        # TODO: Implement Sigmoid backward
        # ====================================================================
        raise NotImplementedError("TODO: Implement Sigmoid backward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================
