import numpy as np

from neural_net.activation_layers.activation_layer import ActivationLayer
from neural_net.functions.activation import relu_activation, relu_derivative_activation


class ReluLayer(ActivationLayer):
    def __init__(self, index, input_dim, output_dim, weights=None, biases=None):
        super().__init__(index, input_dim, output_dim, weights, biases)
        self.subtype = 'RELU'

    def initialize_weights(self):
        # He initialization (input_dim x output_dim)
        self.weights = np.random.randn(self.input_dim, self.output_dim) * np.sqrt(2.0 / self.input_dim)

    def initialize_biases(self):
        self.biases = np.zeros((1, self.output_dim))  # Biases initialized to zero

    def activation(self, outputs: np.array):
        return relu_activation(outputs)

    def activation_derivative(self, outputs: np.array):
        return relu_derivative_activation(outputs)