"""
The neural network has the following definition
F : X -> relu(A_2 (relu( A_1 X + b_1) + b_2))
"""

import numpy as np


def perform_relu_1(x: np.ndarray) -> np.ndarray:
    return np.maximum(x, 0)


def perform_relu_2(x: np.ndarray) -> np.ndarray:
    return x * (x > 0)


def perform_relu_3(x: np.ndarray) -> np.ndarray:
    return (abs(x) + x) / 2


relu_functions = [perform_relu_1, perform_relu_2, perform_relu_3]


class MyNeuralNetwork:
    def __init__(
        self,
        input_dimension: int,
        hidden_dimension: int,
        relu_to_use: int,
        learning_rate: float,
        total_num_steps: int,
    ) -> None:
        self.A1 = None  ### To Complete -> search glorot initialization ###
        self.A2 = None  ### To Complete -> search glorot initialization ###
        self.b1 = np.zeros(shape=(hidden_dimension))
        self.b2 = np.zeros(shape=(1,))
        self.relu = relu_functions[relu_to_use]
        self.learning_rate = learning_rate
        self.step = 0
        self.total_num_steps = total_num_steps

    def inference_call(self, inputs: np.ndarray) -> np.ndarray:
        outputs = inputs
        ### to complete ###
        return outputs

    def apply_gradient_descent(
        self, gA1: np.ndarray, gA2: np.ndarray, gb1: np.ndarray, gb2: np.ndarray
    ):
        ### to complete ###
        pass

    def training_call(self, inputs: np.ndarray, ground_truth: np.ndarray) -> np.ndarray:
        self.step += 1
        feature1 = self.A1 @ inputs + np.expand_dims(self.b1, axis=-1)
        feature2 = perform_relu_1(feature1)
        predictions = self.A2 @ feature2 + self.b2

        mean_squared_error = np.mean((predictions - ground_truth) ** 2)
        print(
            f"\r[{self.step}/{self.total_num_steps}] current MSE: {mean_squared_error:.3f}",
            end="",
        )

        dC_dpreds = None ### To Complete ###
        dC_df3 = None ### To Complete ###
        dC_dA2 = None ### To Complete ###
        dC_db2 = None ### To Complete ###
        dC_df2 = None ### To Complete ###
        dC_df1 = None ### To Complete ###
        dC_dA1 = None ### To Complete ###
        dC_db1 = None ### To Complete ###

        self.apply_gradient_descent(gA1=dC_dA1, gA2=dC_dA2, gb1=dC_db1, gb2=dC_db2)

        return predictions

    def __call__(
        self, inputs: np.ndarray, ground_truth: np.ndarray = None
    ) -> np.ndarray:
        if ground_truth is None:
            return self.inference_call(inputs)
        else:
            return self.training_call(inputs, ground_truth)