# Path: brainforge/llatomic/_llactivation.py
# @nb.vectorize(finfout, nopython=True)
# def sigmoid(Z):
#     return s1 / (s1 + np.exp(-Z))
# 
# @nb.jit(nopython=True)
# def tanh(Z):
#     return np.tanh(Z)
# 
# @nb.vectorize(finfout, nopython=True)
# def relu(Z):
#     return s0 if s0 >= Z else Z


import numba as nb
import numpy as np
from ._llactivation import sigmoid, tanh, relu


@nb.jit(nopython=True)
def lstm_forward(X, W, b, activation):
    outdim = W.shape[-1] // 4
    time, batch, indim = X.shape

    Z = np.zeros((time, batch, indim+outdim))
    O = np.zeros((time, batch, outdim))
    T = np.zeros((time, 6, batch, outdim))  # C[0], Ca[1], cand[2], f[3], i[4], o[5]

    for t in range(time):
        Z[t] = np.concatenate((X[t], O[t-1]), axis=-1)
        p = np.dot(Z[t], W) + b
        p[:, :outdim] = activation(p[:, :outdim])  # nonlin to candidate


# Path: brainforge/llatomic/_llactivation.py
# @nb.vectorize(finfout, nopython=True)
# def sigmoid(Z):
#     return s1 / (s1 + np.exp(-Z))
# 
# @nb.jit(nopython=True)
# def tanh(Z):
#     return np.tanh(Z)
# 
# @nb.vectorize(finfout, nopython=True)
# def relu(Z):
#     return s0 if s0 >= Z else Z


# Path: brainforge/llatomic/_lllstm.py
import numba as nb
import numpy as np
from ._llactivation import sigmoid, tanh, relu



@nb.jit(nopython=True)
def lstm_forward(X, W, b, activation):
    outdim = W.shape[-1] // 4
    time, batch, indim = X.shape

    Z = np.zeros((time, batch, indim+outdim))
    O = np.zeros((time, batch, outdim))
    T = np.zeros((time, 6, batch, outdim))  # C[0], Ca[1], cand[2], f[3], i[4], o[5]

    for t in range(time):
        Z[t] = np.concatenate((X[t], O[t-1]), axis=-1)
        p = np.dot(Z[t], W) + b
        p[:, :outdim] = activation(p[:, :outdim])  # nonlin to candidate


        p[:, outdim:] = sigmoid(p[:, outdim:])  # sigmoid to gates