"""
CW3: Cross-Entropy Loss with Forward and Backward.

============================================================
CONNECTION TO CW2:
You already implemented softmax and cross-entropy in CW2!
Here you wrap them into a class with forward() and backward().
You may copy your softmax implementation from CW2, or import it.

The math is IDENTICAL to CW2's cross_entropy_loss():
    forward():  softmax + negative log-likelihood (same as CW2)
    backward():  grad = (probs - one_hot) / batch_size (same as CW2)
============================================================
"""

import numpy as np


class CrossEntropyLoss:
    """
    Softmax Cross-Entropy Loss as a module with forward() and backward().

    This combines softmax activation with cross-entropy loss, which is
    numerically more stable than computing them separately.
    """

    def __init__(self):
        self.probs = None       # Cached softmax probabilities
        self.labels = None      # Cached labels
        self.num_classes = None  # Cached number of classes

    def forward(self, logits, labels):
        """
        Compute softmax cross-entropy loss.

        TODO: Implement the forward pass.

        Args:
            logits: (batch_size, num_classes) raw scores from the network
            labels: (batch_size,) integer class labels

        Returns:
            loss: scalar, mean cross-entropy loss over the batch

        Hints:
            This is the SAME computation as CW2's cross_entropy_loss()!
            1. Compute softmax probabilities (numerically stable)
            2. Compute negative log-likelihood: -log(probs[i, labels[i]])
            3. Return the mean loss
            4. Cache self.probs, self.labels, self.num_classes for backward()
        """
        # ====================================================================
        # TODO: Implement cross-entropy forward (reuse your CW2 logic)
        # ====================================================================
        raise NotImplementedError("TODO: Implement CrossEntropyLoss forward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def backward(self):
        """
        Compute gradient of loss w.r.t. logits.

        TODO: Implement the backward pass.

        Returns:
            grad_logits: (batch_size, num_classes) gradient of loss w.r.t. logits

        Hints:
            This is the SAME gradient as CW2's cross_entropy_loss()!
            grad = (self.probs - one_hot) / batch_size
        """
        # ====================================================================
        # TODO: Implement cross-entropy backward (reuse your CW2 logic)
        # ====================================================================
        raise NotImplementedError("TODO: Implement CrossEntropyLoss backward")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================
