import numpy as np

from neural_net.activation_layers.relu_layer import ReluLayer
from neural_net.functions.loss import cross_entropy_loss, cross_entropy_derivative_loss
from neural_net.neural_net import NeuralNet
from neural_net.transform_layer import SoftMaxLayer

class MNISTNeuralNet(NeuralNet):
    def __init__(self):
        super().__init__(layers=[
            ReluLayer(0, 784, 121),
            ReluLayer(1, 121, 10),
            SoftMaxLayer(2, 10)
        ])

    def backward(self, dL_dout, epoch):
        return super().backward(dL_dout, epoch)

    def loss(self, y_pred: np.array, y_actual: np.array):
        return cross_entropy_loss(y_pred, y_actual)

    def loss_derivative(self, y_pred: np.array, targets: np.array):
        return cross_entropy_derivative_loss(y_pred, targets)

    def describe(self):
        """Return a human-readable string of the model architecture."""
        architecture_info = ""
        for layer in self.layers:
            architecture_info += f"{layer.describe()}\n"
        return architecture_info.strip()

    def predict(self, inputs):
        raw_outputs = super().predict(inputs)
        return raw_outputs, raw_outputs.argmax(axis=1)