"""
CW2: Logistic Regression Model.

This file defines the Logistic Regression classifier.
You need to implement the core methods marked with TODO.

Logistic Regression computes:
    logits = X @ W + b          (linear transformation)
    probs  = softmax(logits)    (convert to probabilities)
    pred   = argmax(probs)      (class prediction)
"""

import numpy as np


class LogisticRegression:
    """
    Multi-class Logistic Regression classifier.

    This is a single-layer linear model: logits = X @ W + b,
    followed by softmax for probability estimation.
    """

    def __init__(self, input_dim, num_classes):
        """
        Initialize model parameters with Xavier initialization.

        Args:
            input_dim: Number of input features (784 for FashionMNIST)
            num_classes: Number of output classes (10 for FashionMNIST)
        """
        # Xavier initialization: scale weights by sqrt(2 / (fan_in + fan_out))
        scale = np.sqrt(2.0 / (input_dim + num_classes))
        self.W = np.random.randn(input_dim, num_classes).astype(np.float32) * scale
        self.b = np.zeros(num_classes, dtype=np.float32)

        self.input_dim = input_dim
        self.num_classes = num_classes

    def softmax(self, logits):
        """
        Compute softmax probabilities from logits.

        TODO: Implement numerically stable softmax.

        Args:
            logits: (batch_size, num_classes) raw scores

        Returns:
            probs: (batch_size, num_classes) probabilities that sum to 1
                   along each row

        Hints:
            - For numerical stability, subtract the max logit per sample
              before computing exp: exp(logits - max(logits))
            - This prevents overflow in exp() while giving the same result
            - probs[i, j] = exp(logits[i,j] - max_i) / sum_j(exp(logits[i,j] - max_i))
        """
        # ====================================================================
        # TODO: Implement softmax
        # ====================================================================
        raise NotImplementedError("TODO: Implement softmax")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def forward(self, X):
        """
        Compute logits (raw scores) for input X.

        TODO: Implement the forward pass.

        Args:
            X: (batch_size, input_dim) input features

        Returns:
            logits: (batch_size, num_classes) raw scores (before softmax)

        Hints:
            - logits = X @ W + b
            - This is a simple matrix multiplication plus bias
        """
        # ====================================================================
        # TODO: Implement forward pass
        # ====================================================================
        raise NotImplementedError("TODO: Implement forward pass")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def predict(self, X):
        """
        Predict class labels for input X.

        TODO: Implement prediction.

        Args:
            X: (batch_size, input_dim) input features

        Returns:
            predictions: (batch_size,) predicted class indices (integers)

        Hints:
            - Compute logits using self.forward()
            - Return the class with the highest logit (argmax)
        """
        # ====================================================================
        # TODO: Implement predict
        # ====================================================================
        raise NotImplementedError("TODO: Implement predict")
        # ====================================================================
        # END OF YOUR CODE
        # ====================================================================

    def parameters(self):
        """Return model parameters as a dict (provided, do not modify)."""
        return {'W': self.W, 'b': self.b}
