# Path: brainforge/atomic/activation_op.py
# class ActivationFunction:
# class Sigmoid(ActivationFunction):
# class HardSigmoid(ActivationFunction):
# class Tanh(ActivationFunction):
# class Sqrt(ActivationFunction):
# class Linear(ActivationFunction):
# class ReLU(ActivationFunction):
# class LeakyReLU(ActivationFunction):
# class SoftMax(ActivationFunction):
# class OnePlus(ActivationFunction):
#     def forward(self, Z: np.ndarray) -> np.ndarray:
#     def __str__(self):
#     def backward(self, Z: np.ndarray) -> np.ndarray:
#     def forward(self, Z: np.ndarray):
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z: np.ndarray):
#     def backward(self, A: np.ndarray):
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, Z) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A) -> np.ndarray:
#     def __init__(self, alpha=0.3):
#     def forward(self, Z: np.ndarray) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def __init__(self, temperature=1.):
#     def tn(self, Z):
#     def t1(Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def true_backward(A: np.ndarray):
#     def forward(self, Z):
#     def backward(self, Z: np.ndarray):
#         J = np.ones_like(A)
#         J[A <= 0.] = self.alpha
#         I = np.eye(A.shape[1], dtype=floatX)
# 
# Path: brainforge/util/typing.py
# def zX(*dims, dtype=floatX):
#     return np.zeros(dims, dtype=dtype)
# 
# def zX_like(array, dtype=floatX):
#     return zX(*array.shape, dtype=dtype)
# 
# def scalX(scalar, dtype=floatX):
#     return np.asscalar(np.array([scalar], dtype=dtype))


import numpy as np
from .activation_op import activations
from ..util.typing import zX, zX_like, scalX
            Z[t] = np.concatenate((X[t], O[t-1]), axis=-1)

            p = np.dot(Z[t], W) + b
            p[:, :outdim] = self.actfn.forward(p[:, :outdim])
            p[:, outdim:] = sigmoid.forward(p[:, outdim:])

            cand[t] = p[:, :outdim]
            f[t] = p[:, outdim:2*outdim]
            i[t] = p[:, 2*outdim:3*outdim]
            o[t] = p[:, 3*outdim:]
            # cand[t], f[t], i[t], o[t] = np.split(p, 4, axis=1)

            C[t] = C[t-1] * f[t] + cand[t] * i[t]

            Ca[t] = self.actfn.forward(C[t])
            O[t] = Ca[t] * o[t]

        return O, Z, np.stack((C, Ca, cand, f, i, o))

    def backward(self, Z, O, E, W, cache):
        outdim = W.shape[-1] // 4
        time, batch, zdim = Z.shape
        indim = zdim - outdim

        C, Ca, cand, f, i, o = cache
        bwgates = np.concatenate(cache[2:], axis=-1)
        bwgates[..., outdim:] = sigmoid.backward(bwgates[..., outdim:])
        bwgates[..., :outdim] = self.actfn.backward(bwgates[..., :outdim])
        bwCa = np.atleast_2d(self.actfn.backward(Ca))


# Path: brainforge/atomic/activation_op.py
# class ActivationFunction:
# class Sigmoid(ActivationFunction):
# class HardSigmoid(ActivationFunction):
# class Tanh(ActivationFunction):
# class Sqrt(ActivationFunction):
# class Linear(ActivationFunction):
# class ReLU(ActivationFunction):
# class LeakyReLU(ActivationFunction):
# class SoftMax(ActivationFunction):
# class OnePlus(ActivationFunction):
#     def forward(self, Z: np.ndarray) -> np.ndarray:
#     def __str__(self):
#     def backward(self, Z: np.ndarray) -> np.ndarray:
#     def forward(self, Z: np.ndarray):
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z: np.ndarray):
#     def backward(self, A: np.ndarray):
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, Z) -> np.ndarray:
#     def forward(self, Z) -> np.ndarray:
#     def backward(self, A) -> np.ndarray:
#     def __init__(self, alpha=0.3):
#     def forward(self, Z: np.ndarray) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def __init__(self, temperature=1.):
#     def tn(self, Z):
#     def t1(Z) -> np.ndarray:
#     def backward(self, A: np.ndarray) -> np.ndarray:
#     def true_backward(A: np.ndarray):
#     def forward(self, Z):
#     def backward(self, Z: np.ndarray):
#         J = np.ones_like(A)
#         J[A <= 0.] = self.alpha
#         I = np.eye(A.shape[1], dtype=floatX)
# 
# Path: brainforge/util/typing.py
# def zX(*dims, dtype=floatX):
#     return np.zeros(dims, dtype=dtype)
# 
# def zX_like(array, dtype=floatX):
#     return zX(*array.shape, dtype=dtype)
# 
# def scalX(scalar, dtype=floatX):
#     return np.asscalar(np.array([scalar], dtype=dtype))


# Path: brainforge/atomic/recurrent_op.py
import numpy as np
from .activation_op import activations
from ..util.typing import zX, zX_like, scalX

            Z[t] = np.concatenate((X[t], O[t-1]), axis=-1)

            p = np.dot(Z[t], W) + b
            p[:, :outdim] = self.actfn.forward(p[:, :outdim])
            p[:, outdim:] = sigmoid.forward(p[:, outdim:])

            cand[t] = p[:, :outdim]
            f[t] = p[:, outdim:2*outdim]
            i[t] = p[:, 2*outdim:3*outdim]
            o[t] = p[:, 3*outdim:]
            # cand[t], f[t], i[t], o[t] = np.split(p, 4, axis=1)

            C[t] = C[t-1] * f[t] + cand[t] * i[t]

            Ca[t] = self.actfn.forward(C[t])
            O[t] = Ca[t] * o[t]

        return O, Z, np.stack((C, Ca, cand, f, i, o))

    def backward(self, Z, O, E, W, cache):
        outdim = W.shape[-1] // 4
        time, batch, zdim = Z.shape
        indim = zdim - outdim

        C, Ca, cand, f, i, o = cache
        bwgates = np.concatenate(cache[2:], axis=-1)
        bwgates[..., outdim:] = sigmoid.backward(bwgates[..., outdim:])
        bwgates[..., :outdim] = self.actfn.backward(bwgates[..., :outdim])
        bwCa = np.atleast_2d(self.actfn.backward(Ca))


        deltaC = zX_like(O[-1])