"""
CW2: Gradient Descent Optimizer.

You need to implement the gradient descent update rule.

The optimizer receives the gradient of loss w.r.t. logits from the loss function,
and must compute the gradients w.r.t. model parameters (W, b) using the chain rule,
then update the parameters.

Key insight (you'll revisit this in CW3):
    The gradient computation here (dW = X.T @ grad_logits) is actually the
    backward pass of a single Linear layer! In CW3, you'll generalize this
    into a reusable Linear.backward() method.
"""

import numpy as np


class GradientDescent:
    """
    Mini-batch Gradient Descent optimizer.

    Updates model parameters using the gradient of the loss w.r.t. logits.
    """

    def __init__(self, learning_rate=0.01):
        """
        Args:
            learning_rate: Step size for parameter updates
        """
        # ====================================================================
        # TODO: Store the learning rate
        # ====================================================================
        raise NotImplementedError("TODO: Initialize optimizer")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def step(self, model, grad_logits, X):
        """
        Compute parameter gradients and update model weights.

        TODO: Implement the gradient descent update.

        Args:
            model: LogisticRegression instance with attributes W, b
            grad_logits: (batch_size, num_classes) gradient of loss w.r.t. logits
                         (returned by the loss function)
            X: (batch_size, input_dim) the input batch used in the forward pass

        Hints:
            Since logits = X @ W + b, by the chain rule:
                dL/dW = X.T @ dL/d(logits)    →  dW = X.T @ grad_logits
                dL/db = sum(dL/d(logits), axis=0)  →  db = grad_logits.sum(axis=0)

            Then update parameters:
                W -= learning_rate * dW
                b -= learning_rate * db

        NOTE for CW3: The computation dW = X.T @ grad_logits is exactly what
        Linear.backward() does! You will reuse this idea when implementing
        backpropagation for multi-layer networks.
        """
        # ====================================================================
        # TODO: Compute gradients w.r.t. W and b, then update parameters
        # ====================================================================
        raise NotImplementedError("TODO: Implement gradient descent step")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================
