import numpy as np


class LinearMap:
    def __init__(self):
        self.cache = {}

    def forward(self, weight, inputs):
        """
        weight: shape (out_dim, in_dim)
        inputs: shape (in_dim, batch_size)
        output: shape (out_dim, batch_size)
        """
        output = np.matmul(weight, inputs)
        self.cache = {"inputs": inputs, "weight": weight}
        return output

    def backward(self, grad_output):
        """
        inputs: shape (in_dim, batch_size)
        weight: shape (out_dim, in_dim)
        grad_output: shape (out_dim, batch_size)

        Returns:
            grad_inputs: shape (in_dim, batch_size)
            grad_weight: shape (out_dim, in_dim)
        """
        inputs = self.cache["inputs"]
        weight = self.cache["weight"]
        grad_inputs = np.matmul(weight.T, grad_output)
        grad_weight = np.matmul(grad_output, inputs.T)
        return grad_inputs, grad_weight


class LeakyReLUActivation:
    def __init__(self, negative_slope=0.01):
        self.negative_slope = negative_slope
        self.cache = {}

    def forward(self, inputs):
        self.cache["inputs"] = inputs
        return np.where(inputs > 0, inputs, self.negative_slope * inputs)

    def backward(self, grad_output):
        """
        grad_output: same shape as inputs
        """
        # ==========
        # TODO：请完成 LeakyReLU 激活函数的反向传播。
        # 提示：
        #   当 inputs > 0 时，局部梯度为 1；
        #   当 inputs <= 0 时，局部梯度为 self.negative_slope。
        # ==========
        pass


class SoftmaxLayer:
    """
    Apply softmax over the last dimension.
    Input shape should be (batch_size, num_classes).
    """

    def __init__(self):
        self.eps = 1e-8
        self.cache = {}

    def forward(self, logits):
        """
        logits: shape (batch_size, num_classes)
        """
        shifted_logits = logits - np.max(logits, axis=1, keepdims=True)
        exp_logits = np.exp(shifted_logits)
        normalizer = np.sum(exp_logits, axis=1, keepdims=True)
        probs = exp_logits / (normalizer + self.eps)
        self.cache["probs"] = probs
        return probs

    def backward(self, grad_output):
        """
        grad_output: same shape as logits
        """
        # ==========
        # TODO：请完成 Softmax 层的反向传播。
        # ==========
        pass


class SafeLog:
    """
    Apply log over the last dimension.
    """

    def __init__(self):
        self.eps = 1e-12
        self.cache = {}

    def forward(self, inputs):
        """
        inputs: shape (batch_size, num_classes)
        """
        self.cache["inputs"] = inputs
        return np.log(inputs + self.eps)

    def backward(self, grad_output):
        inputs = self.cache["inputs"]
        return grad_output / (inputs + self.eps)


class NumpyMLPClassifier:
    def __init__(self, input_dim, output_dim, hidden_dim=100, lr=5e-5, reg_coeff=0.01, negative_slope=0.01):
        self.W_hidden = 0.01 * np.random.normal(size=[hidden_dim, input_dim + 1])
        self.W_out = 0.01 * np.random.normal(size=[output_dim, hidden_dim])

        self.hidden_linear = LinearMap()
        self.output_linear = LinearMap()
        self.hidden_activation = LeakyReLUActivation(negative_slope=negative_slope)
        self.softmax = SoftmaxLayer()
        self.safe_log = SafeLog()

        self.lr = lr
        self.reg_coeff = reg_coeff

    def _append_bias_term(self, features):
        features = features.reshape(features.shape[0], -1)
        bias_column = np.ones((features.shape[0], 1), dtype=features.dtype)
        return np.concatenate([features, bias_column], axis=1)

    def _cross_entropy_loss(self, log_prob, one_hot_labels):
        # ==========
        # TODO：请计算交叉熵损失项。
        # 其中：
        #   log_prob 是模型预测得到的对数概率；
        #   one_hot_labels 是 one-hot 形式的真实标签。
        # ==========
        pass

    def _regularization_loss(self):
        # ==========
        # TODO：请计算 L2 正则项。
        # ==========
        pass

    def compute_loss(self, log_prob, one_hot_labels):
        """
        Return the total loss = cross-entropy loss + L2 regularization.
        """
        # ==========
        # TODO：请将数据损失与正则项相加，得到总损失。
        # ==========
        pass

    def forward(self, features):
        """
        Return the predicted log probabilities of features.

        MLP structure:
        FFN -> LeakyReLU -> FFN -> Softmax -> Log
        """
        # ==========
        # TODO：请补全该 MLP 的前向传播。
        # ==========
        pass

    def backward(self, one_hot_labels):
        """
        one_hot_labels is the ground truth (one-hot vectors).
        Please compute the gradients of self.W_hidden and self.W_out.
        """
        # ==========
        # TODO：请补全该 MLP 的反向传播，计算 self.W_hidden 和 self.W_out 的梯度。
        # ==========
        pass

    def update(self):
        """
        Update self.W_hidden and self.W_out with gradient descent and L2 regularization.
        """
        # ==========
        # TODO：请使用梯度下降更新参数，并考虑 L2 正则项。
        # ==========
        pass


if __name__ == '__main__':
    model = NumpyMLPClassifier(784, 10)
