"""
CW2: Loss Functions for Logistic Regression.

You need to implement FOUR loss functions, each returning:
    - loss:  scalar, the mean loss over the batch
    - grad:  (batch_size, num_classes), gradient of loss w.r.t. logits

These gradients will be used by the optimizer to update model parameters.

IMPORTANT: All loss functions take logits (raw scores before softmax) as input.
Some losses internally apply softmax; others work directly on logits.
"""

import numpy as np


def softmax(logits):
    """
    Numerically stable softmax (helper function).

    TODO: Implement softmax. This is the same as in model.py.
          You can implement it here and import it in model.py,
          or copy your implementation.

    Args:
        logits: (batch_size, num_classes) raw scores

    Returns:
        probs: (batch_size, num_classes) probabilities summing to 1 per row
    """
    # ========================================================================
    # TODO: Implement numerically stable softmax
    # ========================================================================
    raise NotImplementedError("TODO: Implement softmax")
    # ========================================================================
    # END OF YOUR CODE
    # ========================================================================


def cross_entropy_loss(logits, labels, num_classes):
    """
    Softmax Cross-Entropy Loss.

    This is the standard loss for multi-class classification.

    TODO: Implement cross-entropy loss and its gradient w.r.t. logits.

    Args:
        logits: (batch_size, num_classes) raw scores
        labels: (batch_size,) integer class labels
        num_classes: int, number of classes

    Returns:
        loss: scalar, mean cross-entropy loss over the batch
        grad: (batch_size, num_classes), gradient of loss w.r.t. logits

    Hints:
        1. Compute softmax probabilities: probs = softmax(logits)
        2. Create one-hot encoding of labels
        3. Loss = -mean(sum(one_hot * log(probs), axis=1))
           Or equivalently: Loss = -mean(log(probs[i, labels[i]]))
        4. Gradient w.r.t. logits: grad = (probs - one_hot) / batch_size
        5. Add a small epsilon (1e-12) inside log() to avoid log(0)
    """
    # ========================================================================
    # TODO: Implement cross-entropy loss and gradient
    # ========================================================================
    raise NotImplementedError("TODO: Implement cross-entropy loss")
    # ========================================================================
    # END OF YOUR CODE
    # ========================================================================


def hinge_loss(logits, labels, num_classes):
    """
    Multi-class Hinge Loss (Crammer-Singer formulation).

    This is the loss used in Support Vector Machines for multi-class problems.

    TODO: Implement multi-class hinge loss and its gradient w.r.t. logits.

    Args:
        logits: (batch_size, num_classes) raw scores
        labels: (batch_size,) integer class labels
        num_classes: int, number of classes

    Returns:
        loss: scalar, mean hinge loss over the batch
        grad: (batch_size, num_classes), gradient of loss w.r.t. logits

    Hints:
        1. For each sample i:
           - correct_score = logits[i, labels[i]]
           - margins = logits - correct_score + 1  (margin of 1)
           - margins[i, labels[i]] = 0  (don't count the correct class)
           - loss_i = sum(max(0, margins_i))
        2. Gradient for sample i:
           - For each class j != labels[i] where margin > 0:
             grad[i, j] = 1 / batch_size
           - For the correct class: grad[i, labels[i]] = -count / batch_size
             where count = number of classes with margin > 0
    """
    # ========================================================================
    # TODO: Implement hinge loss and gradient
    # ========================================================================
    raise NotImplementedError("TODO: Implement hinge loss")
    # ========================================================================
    # END OF YOUR CODE
    # ========================================================================


def exponential_loss(logits, labels, num_classes):
    """
    Multi-class Exponential Loss.

    An alternative to cross-entropy that penalizes misclassifications
    more aggressively.

    TODO: Implement exponential loss and its gradient w.r.t. logits.

    Args:
        logits: (batch_size, num_classes) raw scores
        labels: (batch_size,) integer class labels
        num_classes: int, number of classes

    Returns:
        loss: scalar, mean exponential loss over the batch
        grad: (batch_size, num_classes), gradient of loss w.r.t. logits

    Hints:
        1. probs = softmax(logits)
        2. For each sample i: loss_i = exp(-log(probs[i, labels[i]]))
           which simplifies to: loss_i = 1.0 / probs[i, labels[i]]
        3. The gradient can be derived using the chain rule through softmax.
           grad_logits = d(loss)/d(logits)
           For sample i:
             d(loss_i)/d(logits_j) = -(1/p_y^2) * d(p_y)/d(logits_j)
           where p_y = probs[i, labels[i]] and:
             d(p_y)/d(logits_j) = p_y * (1_{j=y} - p_j)
           So: d(loss_i)/d(logits_j) = -(1/p_y) * (1_{j=y} - p_j)
                                      = (p_j - 1_{j=y}) / p_y
        4. Be careful: clip probs to avoid division by zero
           (e.g., probs = np.clip(probs, 1e-12, 1.0))
    """
    # ========================================================================
    # TODO: Implement exponential loss and gradient
    # ========================================================================
    raise NotImplementedError("TODO: Implement exponential loss")
    # ========================================================================
    # END OF YOUR CODE
    # ========================================================================


def squared_loss(logits, labels, num_classes):
    """
    Squared Loss on softmax probabilities.

    Measures the squared difference between predicted probabilities
    and one-hot encoded labels.

    TODO: Implement squared loss and its gradient w.r.t. logits.

    Args:
        logits: (batch_size, num_classes) raw scores
        labels: (batch_size,) integer class labels
        num_classes: int, number of classes

    Returns:
        loss: scalar, mean squared loss over the batch
        grad: (batch_size, num_classes), gradient of loss w.r.t. logits

    Hints:
        1. probs = softmax(logits)
        2. one_hot = one_hot_encode(labels)
        3. diff = probs - one_hot
        4. loss = 0.5 * mean(sum(diff^2, axis=1))
        5. Gradient w.r.t. logits (using chain rule through softmax):
           For each sample i:
             d(loss_i)/d(logits) = sum_k(diff_k * d(probs_k)/d(logits))
           Since d(probs_k)/d(logits_j) = probs_k * (1_{k=j} - probs_j):
             grad_j = sum_k(diff_k * probs_k * (1_{k=j} - probs_j))
                    = probs_j * (diff_j - sum_k(diff_k * probs_k))
           Divide by batch_size for mean.
    """
    # ========================================================================
    # TODO: Implement squared loss and gradient
    # ========================================================================
    raise NotImplementedError("TODO: Implement squared loss")
    # ========================================================================
    # END OF YOUR CODE
    # ========================================================================


# Dictionary mapping loss names to functions (used by trainer.py)
LOSS_FUNCTIONS = {
    "cross_entropy": cross_entropy_loss,
    "hinge": hinge_loss,
    "exponential": exponential_loss,
    "squared": squared_loss,
}
