"""
CW3: SGD Optimizer with Momentum.

============================================================
CONNECTION TO CW2:
In CW2, your GradientDescent optimizer did:
    W -= learning_rate * dW
    b -= learning_rate * db

Now you extend this with MOMENTUM:
    velocity_W = momentum * velocity_W + grad_W
    W -= learning_rate * velocity_W

When momentum=0, this reduces to your CW2 implementation exactly.
Momentum helps accelerate training by accumulating gradient direction.
============================================================

The key difference from CW2:
- In CW2, you computed dW = X.T @ grad_logits inside the optimizer
- In CW3, the gradients (grad_W, grad_b) are already computed by
  layer.backward() — the optimizer just updates parameters using them.
"""

import numpy as np


class SGD:
    """
    Stochastic Gradient Descent with optional momentum and weight decay.
    """

    def __init__(self, parameters, learning_rate=0.01, momentum=0.0,
                 weight_decay=0.0):
        """
        TODO: Initialize the optimizer.

        Args:
            parameters: List of parameter dicts from model.parameters()
                Each dict has: 'W', 'b', 'layer'
                  - 'W', 'b'    : the parameter arrays (update these in-place)
                  - 'layer'     : the Linear layer object; read gradients from
                                  param['layer'].grad_W and param['layer'].grad_b
                                  at step time (not cached at init time)
            learning_rate: Step size for updates
            momentum: Momentum coefficient (0 = no momentum = CW2 behavior)
            weight_decay: L2 regularization strength

        Hints:
            - Store learning_rate, momentum, weight_decay
            - Store the parameters list
            - If momentum > 0, initialize velocity buffers (same shape as params)
              for each W and b: velocity_W = zeros_like(W), velocity_b = zeros_like(b)
        """
        # ====================================================================
        # TODO: Initialize optimizer state
        # ====================================================================
        raise NotImplementedError("TODO: Initialize SGD optimizer")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def step(self):
        """
        Update all parameters using their computed gradients.

        TODO: Implement the parameter update.

        Hints:
            For each parameter group, first read the fresh gradients from the
            layer object (NOT from a cached reference — backward() assigns new
            arrays to layer.grad_W each call):

                W      = param['W']
                b      = param['b']
                grad_W = param['layer'].grad_W
                grad_b = param['layer'].grad_b

            Then apply the update:

            If weight_decay > 0:
                grad_W = grad_W + weight_decay * W  (L2 regularization)

            If momentum > 0:
                velocity_W = momentum * velocity_W + grad_W
                W -= learning_rate * velocity_W
            Else:
                W -= learning_rate * grad_W  (same as CW2!)

            Same logic for b (but typically no weight decay on bias).

        NOTE: You must update the W and b arrays in-place (e.g., W -= ...)
              so the changes are reflected in the model's layers.
        """
        # ====================================================================
        # TODO: Implement parameter update with momentum
        # ====================================================================
        raise NotImplementedError("TODO: Implement SGD step")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def zero_grad(self):
        """
        Reset all gradients to zero (call before each training step).

        TODO: Implement gradient zeroing.

        Hints:
            For each parameter group, zero the gradients on the layer directly:
                param['layer'].grad_W = np.zeros_like(param['W'])
                param['layer'].grad_b = np.zeros_like(param['b'])

            Note: since backward() fully overwrites grad_W and grad_b each
            call, zero_grad() is not strictly required in this training loop
            (trainer.py does not call it). It is provided as a good habit and
            for completeness — useful if you ever accumulate gradients across
            multiple batches.
        """
        # ====================================================================
        # TODO: Zero all gradients
        # ====================================================================
        raise NotImplementedError("TODO: Implement zero_grad")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================
