import numpy as np


# pytrahn export relu(float32[:,:])
def relu(x):
    return np.maximum(x, 0)


# Numerically-stable version of softmax
# pythran export softmax(float32[:,:])
def softmax(x):
    # tmp_max = np.max(x, axis=-1, keepdims=True)
    tmp_max = np.max(x, axis=-1)[:, :, np.newaxis]
    tmp_out = np.exp(x - tmp_max)
    # tmp_sum = np.sum(tmp_out, axis=-1, keepdims=True)
    tmp_sum = np.sum(tmp_out, axis=-1)[:, :, np.newaxis]
    return tmp_out / tmp_sum


# 3-layer MLP
# pythran export mlp(float32[:,:], float32[:,:], float32[:], float32[:,:],
#                    float32[:], float32[:,:], float32[:])
def mlp(input, w1, b1, w2, b2, w3, b3):
    x = relu(input @ w1 + b1)
    x = relu(x @ w2 + b2)
    x = softmax(x @ w3 + b3)  # Softmax call can be omitted if necessary
    return x
