#%% Libraries: Built-In
import numpy as np
#% Libraries: Custom

#%%
class Activator(object):
    def forward(self, combined_array):
        ## Define in child
        pass
    
    def backprop(self, backprop_array, train_index):
        ## Define in child
        pass
#%%
class Linear(Activator):
    def forward(self, combined_array):
        return combined_array
    
    def backprop(self, backprop_array, train_index):
        return backprop_array
#%%
class Logistic(Activator):
    def forward(self, combined_array):
        self.prior_forward = combined_array
        new_forward = 1 / (1 + np.exp(-combined_array))
        if np.any(np.isinf(new_forward)):
            print(combined_array)
            print(new_forward)
        return new_forward
    
    def backprop(self, backprop_array, train_index):
        prob_calc = 1 / (1 + np.exp(-self.prior_forward[train_index, :]))
        leak_back = prob_calc * (1 - prob_calc)
        return leak_back.swapaxes(1, 2) * backprop_array
#%%
class TanH(Activator):
    def forward(self, combined_array):
        self.prior_forward = combined_array
        new_forward = np.tanh(combined_array)
        return new_forward
        
    def backprop(self, backprop_array, train_index):
        leak_back = 1 - np.power(np.tanh(self.prior_forward[train_index, :]), 2)
        return leak_back.swapaxes(1, 2) * backprop_array
#%%
class ReLU(Activator):
    def __init__(self, leak = 0):
        self.leak = leak
        
    def forward(self, combined_array):
        self.prior_forward = combined_array
        new_forward = np.maximum(combined_array, combined_array * self.leak)
        return new_forward
        
    def backprop(self, backprop_array, train_index):
        leak_back = np.clip(self.prior_forward[train_index, :] <= 0, self.leak, 1)
        return leak_back.swapaxes(1, 2) * backprop_array
#%%
