####################
## load libraries ##
####################
import numpy as np
from sklearn.metrics import roc_auc_score
np.random.seed(1234567)

############################
## Toy Dataset Simulation ##
############################
def simulate_dataset(n=10000, features=10):
    
    np.random.seed(1081565)
    ## divide number of features by 2
    p = int(features/2)
    ## specify intercept column
    X_intercept = np.ones((n,1))
    ## create p features over the reals
    X_1 = np.random.uniform(-2, 2, size=(n, p))
    ## create p features with support {0,1}
    X_2 = np.random.randint(low=0, high=2, size=(n, p))
    ## concate intercept, X1, and X2
    X = np.concatenate((X_intercept, X_1, X_2), axis=1)
    ## save the original "Design Matrix" of features to return later
    D = np.concatenate((X_1, X_2), axis=1)
    
    ## create 10 additional features
    X_11 = (X[:, 1] * X[:, 1]).reshape((n, 1))
    X_12 = (X[:, 3] * X[:, 3]).reshape((n, 1))
    X_13 = (X[:, 2] * X[:, 4]).reshape((n, 1))
    X_14 = (X[:, 2] * X[:, 8]).reshape((n, 1))
    X_15 = (X[:, 7] * X[:, 9]).reshape((n, 1))
    X_16 = np.cos(X[:, 4]).reshape((n, 1))
    X_17 = np.sin(X[:, 2]).reshape((n, 1))
    X_18 = (X_16 * X_11).reshape((n, 1))
    X_19 = (X_17 * X_12).reshape((n, 1))
    X_20 = (X_16 * X_16).reshape((n, 1))
    X = np.concatenate((X, X_11, X_12, X_13, X_14, X_15, X_16, X_17, X_18, X_19, X_20), axis=1)
    del X_intercept, X_1, X_2, X_11, X_12, X_13, X_14, X_15, X_16, X_17, X_18, X_19, X_20
    
    ## specify Betas
    Beta = np.random.uniform(-1, 1, size=((2*features)+1, 1))

    ## Fit the outcome Y
    theta = np.dot(X, Beta)
    prob = (1 / (1 + np.exp(-theta)))
    
#    ## specify categorical {0,1} outcome
#    Y = np.zeros((n, 1))
#    for i in range(0, n):
#        Y[i, 0] = np.random.choice([0, 1], size=1, replace=True, p=[1-prob[i, 0], prob[i, 0]])    
    ## specify categorical {0,1} outcome
    Y = np.zeros((n, 1))
    for i in range(0, n):
        Y[i, 0] = np.random.choice([0, 1], size=1, replace=True, p=[1-prob[i, 0], prob[i, 0]])  
 
    ## return outcome Y and original Design Matrix D
    return(Y, D)

##############################
## Neural Network Functions ##
##############################
def Normalize_Design_Matrix(D):
    D_means = D.mean(axis=0, keepdims=True)
    D_std = D.std(axis=0, keepdims=True)
    D_normalized = (D-D_means) / D_std
    return(D_normalized, D_means, D_std)

def initialize_Beta(p1, p2=1):
    Beta = np.random.uniform(-1, 1, size=(p1, p2))
    return(Beta)

def initialize_b(p2=1):
    b = np.random.uniform(-1, 1, size=(1, p2))
    return(b)

def sigmoid(theta):
    Y = 1 / (1 + np.exp(-theta))
    return(Y)
    
def sigmoid_derivative(Z):
    Y = sigmoid(Z) * (1-sigmoid(Z))
    return(Y)
    
def softmax(theta):
    Y = np.exp(theta)
    Y = Y / Y.sum(axis=1).reshape((Y.shape[0],1))   
    return(Y)
    
def softmax_derivative(Z):
    Y = softmax(Z) * (1-softmax(Z))
    return(Y)
    
def tanh(theta):
    Y = (np.exp(theta) - np.exp(-theta)) / (np.exp(theta) + np.exp(-theta))
    return(Y)

def tanh_derivative(Z):
    Y = 1 - (tanh(Z)*tanh(Z))
    return(Y)

def reLu(theta):
    Y = theta.copy()
    Y[Y<0]=0
    return(Y)
    
def reLu_derivative(Z):
    return(1)
    
def identity(theta):
    Y = theta.copy()
    return(Y)
    
def identity_derivative(Z):
    return(1)
   

def specify_mini_batch_size(n, mini_batch_size):
    if(n % mini_batch_size == 0):
        start = (n//mini_batch_size)*[None]
        end = (n//mini_batch_size)*[None]
        m = (n//mini_batch_size)*[None]
        for i in range(0, (n//mini_batch_size)):
            start[i] = i*mini_batch_size
            end[i] = (i+1)*mini_batch_size
            m[i] = mini_batch_size
        del i
    else:
        start = int(np.ceil(n/mini_batch_size))*[None]
        end = int(np.ceil(n/mini_batch_size))*[None]
        m = int(np.ceil(n/mini_batch_size))*[None]
        for i in range(0, (n//mini_batch_size)):
            start[i] = i*mini_batch_size
            end[i] = (i+1)*mini_batch_size
            m[i] = mini_batch_size
        start[i+1] = end[i]
        end[i+1] = n
        m[i+1] = end[i+1] - start[i+1]
        del i
    return(start, end, m)



    

class NeuralNetwork:
    def __init__(self, Y, D):
        self.n = D.shape[0]
        self.Y = Y
        self.D = D
        self.D_normalized, self.D_means, self.D_std = Normalize_Design_Matrix(self.D)
    
    def specify_layers(self, p, activations):
        self.p = p
        self.activations = activations
        self.Betas = {}
        self.bs = {}
        self.X = {}
        self.Z = {}
        self.Y_hat = None
        self.dJ_dZ = {}
        self.dJ_dB = {}
        self.dJ_db = {}
        for i in range(0, len(self.p)):
            
            if(i<len(self.p)-1):
                self.Betas['B'+str(i+1)] = initialize_Beta(p1=self.p[i], p2=self.p[i+1])
                self.bs['b'+str(i+1)] = initialize_b(p2=self.p[i+1])
            else:
                self.Betas['B'+str(i+1)] = initialize_Beta(p1=self.p[i], p2=self.Y.shape[1])
                self.bs['b'+str(i+1)] = initialize_b(p2=self.Y.shape[1])
                
            self.X['X'+str(i+1)] = None
            self.dJ_dZ['dJ_dZ'+str(i+1)] = None
            self.dJ_dB['dJ_dB'+str(i+1)] = None
            self.dJ_db['dJ_db'+str(i+1)] = None
                
            
    def forward_dropout(self, j, dropout_probability):
        self.X['X'+str(j)] = (1/(1-dropout_probability)) * self.X['X'+str(j)] * np.random.binomial(1, (1-dropout_probability), size=(self.X['X'+str(j)].shape[0], self.X['X'+str(j)].shape[1]))
        
        
    def forward_Z(self, j):
        self.Z['Z'+str(j)] = np.dot(self.X['X'+str(j)], self.Betas['B'+str(j)]) + self.bs['b'+str(j)]
    
    
    def forward_X(self, j):
        activation = self.activations[j-1]
        if(activation=='reLu'):
            self.X['X'+str(j+1)] = reLu(self.Z['Z'+str(j)])
        if(activation=='sigmoid'):
            self.X['X'+str(j+1)] = sigmoid(self.Z['Z'+str(j)])
        if(activation=='tanh'):
            self.X['X'+str(j+1)] = tanh(self.Z['Z'+str(j)])
        if(activation=='softmax'):
            self.X['X'+str(j+1)] = softmax(self.Z['Z'+str(j)])
        if(activation=='identity'):
            self.X['X'+str(j+1)] = identity(self.Z['Z'+str(j)])
    
    
    def forward_Y(self):
        activation = self.activations[-1]
        if(activation=='reLu'):
            self.Y_hat = reLu(self.Z['Z'+str(len(self.Z))])
        if(activation=='sigmoid'):
            self.Y_hat = sigmoid(self.Z['Z'+str(len(self.Z))])
        if(activation=='tanh'):
            self.Y_hat = tanh(self.Z['Z'+str(len(self.Z))])
        if(activation=='softmax'):
            self.Y_hat = softmax(self.Z['Z'+str(len(self.Z))])
        if(activation=='identity'):
            self.Y_hat = identity(self.Z['Z'+str(len(self.Z))])

    
    def calculate_loss_function(self, Y_mini, i, m, k):
        activation = self.activations[-1]
        if(activation=='sigmoid' or activation=='softmax'):
            self.J = (1/m[i])*(-np.dot(Y_mini.T, np.log(self.Y_hat)) - np.dot((1-Y_mini).T, np.log(1-self.Y_hat))) 
        if(activation=='identity' or activation=='softmax'):
            self.J = (1/m[i])*np.dot((Y_mini-self.Y_hat).T, (Y_mini-self.Y_hat))
        if(i==(len(m)-1)):
            print('current loss on epoch ' + str(k+1) + ' is: ' + str(self.J.sum()))
            
    
    def backward_dJ_dZ(self, Y_mini, j):
        activation = self.activations[j-1]
        if(j==len(self.Z)):
            if(activation=='sigmoid'):
                self.dJ_dZ['dJ_dZ'+str(len(self.dJ_dZ))] = sigmoid(self.Z['Z'+str(len(self.Z))]) - Y_mini
            if(activation=='reLu'):
                self.dJ_dZ['dJ_dZ'+str(len(self.dJ_dZ))] = reLu(self.Z['Z'+str(len(self.Z))]) - Y_mini
            if(activation=='tanh'):
                self.dJ_dZ['dJ_dZ'+str(len(self.dJ_dZ))] = tanh(self.Z['Z'+str(len(self.Z))]) - Y_mini
            if(activation=='softmax'):
                self.dJ_dZ['dJ_dZ'+str(len(self.dJ_dZ))] = softmax(self.Z['Z'+str(len(self.Z))]) - Y_mini
            if(activation=='identity'):
                self.dJ_dZ['dJ_dZ'+str(len(self.dJ_dZ))] = identity(self.Z['Z'+str(len(self.Z))]) - Y_mini
        if(j<len(self.Z)):
            if(activation=='sigmoid'):
                self.dJ_dZ['dJ_dZ'+str(j)] = np.dot(self.dJ_dZ['dJ_dZ'+str(j+1)], self.Betas['B'+str(j+1)].T) * sigmoid_derivative(self.Z['Z'+str(j)])
            if(activation=='reLu'):
                self.dJ_dZ['dJ_dZ'+str(j)] = np.dot(self.dJ_dZ['dJ_dZ'+str(j+1)], self.Betas['B'+str(j+1)].T) * reLu_derivative(self.Z['Z'+str(j)])
            if(activation=='tanh'):
                self.dJ_dZ['dJ_dZ'+str(j)] = np.dot(self.dJ_dZ['dJ_dZ'+str(j+1)], self.Betas['B'+str(j+1)].T) * tanh_derivative(self.Z['Z'+str(j)])
            if(activation=='softmax'):
                self.dJ_dZ['dJ_dZ'+str(j)] = np.dot(self.dJ_dZ['dJ_dZ'+str(j+1)], self.Betas['B'+str(j+1)].T) * softmax_derivative(self.Z['Z'+str(j)])
            if(activation=='identity'):
                self.dJ_dZ['dJ_dZ'+str(j)] = np.dot(self.dJ_dZ['dJ_dZ'+str(j+1)], self.Betas['B'+str(j+1)].T) * identity_derivative(self.Z['Z'+str(j)])


    def backward_dJ_dB(self, j):
        self.dJ_dB['dJ_dB'+str(j)] = np.dot(self.X['X'+str(j)].T, self.dJ_dZ['dJ_dZ'+str(j)])
        
    
    def backward_dJ_db(self, j, i, m):
        self.dJ_db['dJ_db'+str(j)] = np.dot(np.ones((1, m[i])), self.dJ_dZ['dJ_dZ'+str(j)])
    
    
    def update_Betas(self, j, alpha):
        self.Betas['B'+str(j)] = self.Betas['B'+str(j)] - (alpha*self.dJ_dB['dJ_dB'+str(j)])
        
        
    def update_bs(self, j, alpha):
        self.bs['b'+str(j)] = self.bs['b'+str(j)] - (alpha*self.dJ_db['dJ_db'+str(j)])
    
    
    
    def fit(self, mini_batch_size, epochs=1, alpha=0.001, drop_out=False, dropout_probability=0):
        ###########################################################
        ## specify coordinates of minibatch sizes for each epoch ##
        ###########################################################
        start, end, m = specify_mini_batch_size(self.n, mini_batch_size)
        
        for k in range(0, epochs):     ## loop over the epochs
            for i in range(0,len(m)):           ## loop over the minibatches    
                Y_mini = self.Y[start[i]:end[i],:]
                
                ###################
                ## forward pass: ##
                ###################
                for j in range(1, len(self.X)):                    
                    ## normalize the input layer
                    if(j==1):
                        self.X['X'+str(j)] = self.D_normalized[start[i]:end[i],:]                    
                    ## initiate dropout regularization (if needed)
                    if(drop_out):
                        self.forward_dropout(j, dropout_probability)                    
                    ## calculate Z of current layer
                    self.forward_Z(j)                    
                    ## calculate X of next layer
                    self.forward_X(j)
                ## calculate Z of final layer
                self.forward_Z(j+1)
                ## calculate output layer Y
                self.forward_Y()

            
                ###########################################
                ## print current value of loss function: ##
                ###########################################
                self.calculate_loss_function(Y_mini, i, m, k)
                
            
                #####################
                ## backwards pass: ##
                #####################
                for j in range(len(self.Z), 0, -1):
                    ## calculate the gradient dJ_dTheta
                    self.backward_dJ_dZ(Y_mini, j)
                
                for j in range(len(self.X), 0, -1):
                    ## calculate the gradient dJ_dB
                    self.backward_dJ_dB(j)
                    ## calculate the gradient dJ_db
                    self.backward_dJ_db(j, i, m)

            
                ##################################
                ## update coefficient estimates: ##
                ##################################
                for j in range(len(self.Betas), 0, -1):
                    ## update Betas
                    self.update_Betas(j, alpha)
                    ## update intercets b
                    self.update_bs(j, alpha)
                                    
                del Y_mini

## simulate dataset
Y, D = simulate_dataset()

## initialize Neural Network class object
NN = NeuralNetwork(Y, D)

## fit the Neural Network
NN.specify_layers(p=[NN.D.shape[1], 100, 85, 80, 70, 60], activations=['reLu', 'sigmoid', 'reLu', 'sigmoid', 'reLu', 'sigmoid'])
NN.fit(mini_batch_size=10000, epochs=1000, alpha=0.000001, drop_out=False, dropout_probability=0.02)

######################################
## Area Under the Curve calculation ##
######################################
Y_hat = NN.Y_hat
print(roc_auc_score(Y, Y_hat))