# input cell
# probablistic hidden cell
# match Input Output cell



import argparse
import numpy as np
import os
from extra.network.helpers import sigmoid, lrelu, tanh, img_tile, mnist_reader, relu, BCE_loss

np.random.seed(0)

class Sigmoid():
    def __call__(self, x, bias=0):
        return 1 / (1 + np.exp(-x + bias))
    
    def derivative(self, x, bias=0):
        return self.__call__(x, bias) * (1 - self.__call__(x, bias))

class Tanh():
    def __call__(self, x):
        return 2 / (1 + np.exp(-2 * x)) - 1

    def derivative(self, x):
        return 1 - (self.__call__(x) ** 2)

class Relu():
    def __call__(self, x):
        return np.where(x >= 0, x, 0)
    
    def derivative(self, x):
        return np.where(x >= 0, 1, 0)

class LRelu():
    def __call__(self, x, alpha=0.01):
        return np.maximum(x, x * alpha, x)

    def derivative(self, x, alpha=0.01):
        dx = np.ones_like(x)
        dx[x < 0] = alpha
        return dx
    
act_functions = {
    "sigmoid": Sigmoid,
    "tanh"   : Tanh,
    "relu"   : Relu,
    "lrelu"  : LRelu
}





class Input:
    """
    Input Neuron given his input to his output.
    """
    def __init__(self, input_shape):
        self.input_shape = input_shape

        self.inputs = np.full(self.input_shape, None)

    def call(self, X):
        self.inputs = X
        return self.inputs

class Hidden:
    def __init__(self, input_shape, weights, bias, activation="relu"):
        self.input_shape = input_shape
        self.activation = act_functions[activation]()

        # Xavier initialization
        limit = 1 / np.sqrt(self.input_shape[0])
        self.W = weights# np.random.uniform(0.00, limit, self.input_shape)
        self.b = bias# np.zeros(len(self.input_shape))

    def forward(self, X):
        self.wsum = X.dot(self.W) + self.b
        return self.activation(self.wsum)


class HiddenProbablistic:
    def __init__(self):
        ""

class OutputInputMatch:
    def __init__(self):
        ""

class Model:
    def __init__(self, layers):
        self.layers   = layers
        self.n_input  = len(self.layers[0].input_shape[0])
        self.n_output = len(self.layers[-1].input_shape)

    def forward(self, X, verbose=False):
        output = X
        for layer in self.layers:
            output = layer.forward(output)

        # result of network
        if verbose:
            print(f"Model predict: {output}")
        return output    
        
    def backward(self):
        ""
    
    def train(self):
        ""

    

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--epoch", type=int, default=40)
    parser.add_argument("--nz", type=int, default=20)
    parser.add_argument("--layersize", type=int, default=400)
    parser.add_argument("--alpha", type=float, default=1)
    parser.add_argument("--lr", type=float, default=0.0001)
    parser.add_argument("--b1", type=float, default=0.9)
    parser.add_argument("--b2", type=float, default=0.999)
    parser.add_argument("--e", type=float, default=1e-8)
    parser.add_argument("--bsize", type=int, default=64)
    return parser.parse_args()

args = parse_args()
cpu_enabled = 0

import numpy as np
np.random.seed(111)

class VAE():
    def __init__(self, numbers):
        
        self.numbers = numbers

        self.epochs = args.epoch
        self.batch_size = args.bsize
        self.learning_rate = args.lr
        self.decay = 0.001
        self.nz = args.nz
        self.layersize = args.layersize

        self.img_path = "./extra/images"
        if not os.path.exists(self.img_path):
            os.makedirs(self.img_path)
        
        # Xavier initialization is used to initialize the weights
        # init encoder weights
        self.e_W0 = np.random.randn(784, self.layersize).astype(np.float32) * np.sqrt(2.0/(784))
        self.e_b0 = np.zeros(self.layersize).astype(np.float32)

        self.e_W_mu = np.random.randn(self.layersize, self.nz).astype(np.float32) * np.sqrt(2.0/(self.layersize))
        self.e_b_mu = np.zeros(self.nz).astype(np.float32)
        
        self.e_W_logvar = np.random.randn(self.layersize, self.nz).astype(np.float32) * np.sqrt(2.0/(self.layersize))
        self.e_b_logvar = np.zeros(self.nz).astype(np.float32)

        # init decoder weights 
        self.d_W0 = np.random.randn(self.nz, self.layersize).astype(np.float32) * np.sqrt(2.0/(self.nz))
        self.d_b0 = np.zeros(self.layersize).astype(np.float32)
        
        self.d_W1 = np.random.randn(self.layersize, 784).astype(np.float32) * np.sqrt(2.0/(self.layersize))
        self.d_b1 = np.zeros(784).astype(np.float32)
             
        # init sample
        self.sample_z = 0
        self.rand_sample = 0
        
        # init Adam optimizer
        self.b1 = args.b1
        self.b2 = args.b2
        self.e = args.e
        self.m = [0] * 10
        self.v = [0] * 10
        self.t = 0
        

        
    def encoder(self, img):
        #self.e_logvar : log variance 
        #self.e_mean : mean

        self.e_input = np.reshape(img, (self.batch_size,-1))
    
        self.e_h0 = Hidden(
            input_shape=(len(self.e_input), len(self.e_input[0])),
            weights=self.e_W0,
            bias=self.e_b0,
            activation="lrelu"
        )

        self.e_h0_a = self.e_h0.forward(self.e_input)
        self.e_h0_l = self.e_h0.wsum

        # self.e_h0_l = self.e_input.dot(self.e_W0) + self.e_b0
        # self.e_h0_a = lrelu(self.e_h0_l)

        self.e_logvar = self.e_h0_a.dot(self.e_W_logvar) + self.e_b_logvar
        self.e_mu = self.e_h0_a.dot(self.e_W_mu) + self.e_b_mu
    
        return self.e_mu, self.e_logvar
    
    def decoder(self, z):
        #self.d_out : reconstruction image 28x28
		
        self.z = np.reshape(z, (self.batch_size, self.nz))
        
        
        self.d_h0 = Hidden(
            input_shape=(len(self.z), len(self.z[0])),
            weights=self.d_W0,
            bias=self.d_b0,
            activation="relu"
        )
        self.d_h0_a = self.d_h0.forward(self.z)
        self.d_h0_l = self.d_h0.wsum

        # self.d_h0_l = self.z.dot(self.d_W0) + self.d_b0		
        # self.d_h0_a = relu(self.d_h0_l)

        
        self.d_h1 = Hidden(
            input_shape=(len(self.d_h0_a), len(self.d_h0_a[0])),
            weights=self.d_W1,
            bias=self.d_b1,
            activation="sigmoid"
        )
        self.d_h1_a = self.d_h1.forward(self.d_h0_a)
        self.d_h1_l = self.d_h1.wsum

        # self.d_h1_l = self.d_h0_a.dot(self.d_W1) + self.d_b1
        # self.d_h1_a = sigmoid(self.d_h1_l)

        self.d_out = np.reshape(self.d_h1_a, (self.batch_size, 28, 28, 1))

        return self.d_out

    def forward(self, x):
        #Encode
        mu, logvar = self.encoder(x)
        
        #use reparameterization trick to sample from gaussian
        self.rand_sample = np.random.standard_normal(size=(self.batch_size, self.nz))
        self.sample_z = mu + np.exp(logvar * .5) * np.random.standard_normal(size=(self.batch_size, self.nz))
        
        decode = self.decoder(self.sample_z)
        
        return decode, mu, logvar
    
    def backward(self, x, out):
        ########################################
        #Calculate gradients from reconstruction
        ########################################
        y = np.reshape(x, (self.batch_size, -1))
        out = np.reshape(out, (self.batch_size, -1))
        
        #Calculate decoder gradients
        #Left side term
        dL_l = -y * (1 / out)
        dsig = sigmoid(self.d_h1_l, derivative=True)
        dL_dsig_l = dL_l * dsig
        
        drelu = relu(self.d_h0_l, derivative=True)

        dW1_d_l = np.matmul(np.expand_dims(self.d_h0_a, axis=-1), np.expand_dims(dL_dsig_l, axis=1))
        db1_d_l = dL_dsig_l 
        
        db0_d_l = dL_dsig_l.dot(self.d_W1.T) * drelu
        dW0_d_l = np.matmul(np.expand_dims(self.sample_z, axis=-1), np.expand_dims(db0_d_l, axis=1))
        
        #Right side term
        dL_r = (1 - y) * (1 / (1 - out))
        dL_dsig_r = dL_r * dsig
        
        dW1_d_r = np.matmul(np.expand_dims(self.d_h0_a, axis=-1), np.expand_dims(dL_dsig_r, axis=1))
        db1_d_r = dL_dsig_r
        
        db0_d_r = dL_dsig_r.dot(self.d_W1.T) * drelu
        dW0_d_r = np.matmul(np.expand_dims(self.sample_z, axis=-1), np.expand_dims(db0_d_r, axis=1))
        
        # Combine gradients for decoder
        grad_d_W0 = dW0_d_l + dW0_d_r
        grad_d_b0 = db0_d_l + db0_d_r
        grad_d_W1 = dW1_d_l + dW1_d_r
        grad_d_b1 = db1_d_l + db1_d_r
         
        #Calculate encoder gradients from reconstruction
        #Left side term
        d_b_mu_l  = db0_d_l.dot(self.d_W0.T)
        d_W_mu_l = np.matmul(np.expand_dims(self.e_h0_a, axis=-1), np.expand_dims(d_b_mu_l, axis=1))
        
        db0_e_l = d_b_mu_l.dot(self.e_W_mu.T) * lrelu(self.e_h0_l, derivative=True)
        dW0_e_l = np.matmul(np.expand_dims(y, axis=-1), np.expand_dims(db0_e_l, axis=1)) 
        
        d_b_logvar_l = d_b_mu_l * np.exp(self.e_logvar * .5) * .5 * self.rand_sample
        d_W_logvar_l = np.matmul(np.expand_dims(self.e_h0_a, axis=-1), np.expand_dims(d_b_logvar_l, axis=1))
        
        db0_e_l_2 = d_b_logvar_l.dot(self.e_W_logvar.T) * lrelu(self.e_h0_l, derivative=True)
        dW0_e_l_2 = np.matmul(np.expand_dims(y, axis=-1), np.expand_dims(db0_e_l_2, axis=1)) 
        
        #Right side term
        d_b_mu_r  = db0_d_r.dot(self.d_W0.T)
        d_W_mu_r = np.matmul(np.expand_dims(self.e_h0_a, axis=-1), np.expand_dims(d_b_mu_r, axis=1))
        
        db0_e_r = d_b_mu_r.dot(self.e_W_mu.T) * lrelu(self.e_h0_l, derivative=True)
        dW0_e_r = np.matmul(np.expand_dims(y, axis=-1), np.expand_dims(db0_e_r, axis=1)) 
        
        d_b_logvar_r = d_b_mu_r * np.exp(self.e_logvar * .5) * .5 * self.rand_sample
        d_W_logvar_r = np.matmul(np.expand_dims(self.e_h0_a, axis=-1), np.expand_dims(d_b_logvar_r, axis=1))
        
        db0_e_r_2 = d_b_logvar_r.dot(self.e_W_logvar.T) * lrelu(self.e_h0_l, derivative=True)
        dW0_e_r_2 = np.matmul(np.expand_dims(y, axis=-1), np.expand_dims(db0_e_r_2, axis=1))
        
        ########################################
        #Calculate encoder gradients from K-L
        ########################################
    
        #logvar terms
        dKL_b_log = -.5 * (1 - np.exp(self.e_logvar))
        dKL_W_log = np.matmul(np.expand_dims(self.e_h0_a, axis= -1), np.expand_dims(dKL_b_log, axis= 1))
        
        #Heaviside step function
        dlrelu = lrelu(self.e_h0_l, derivative=True)  

        dKL_e_b0_1 = .5 * dlrelu * (np.exp(self.e_logvar) - 1).dot(self.e_W_logvar.T)
        dKL_e_W0_1 = np.matmul(np.expand_dims(y, axis= -1), np.expand_dims(dKL_e_b0_1, axis= 1))
        
        #m^2 term
        dKL_W_m = .5 * (2 * np.matmul(np.expand_dims(self.e_h0_a, axis=-1), np.expand_dims(self.e_mu, axis=1)))
        dKL_b_m = .5 * (2 * self.e_mu)
        
        dKL_e_b0_2 = .5 * dlrelu * (2 * self.e_mu).dot(self.e_W_mu.T)
        dKL_e_W0_2 = np.matmul(np.expand_dims(y, axis= -1), np.expand_dims(dKL_e_b0_2, axis= 1))
        
        # Combine gradients for encoder from recon and KL
        grad_b_logvar = dKL_b_log + d_b_logvar_l + d_b_logvar_r
        grad_W_logvar = dKL_W_log + d_W_logvar_l + d_W_logvar_r
        grad_b_mu = dKL_b_m + d_b_mu_l + d_b_mu_r
        grad_W_mu = dKL_W_m + d_W_mu_l + d_W_mu_r
        grad_e_b0 = dKL_e_b0_1 + dKL_e_b0_2 + db0_e_l + db0_e_l_2 + db0_e_r + db0_e_r_2
        grad_e_W0 = dKL_e_W0_1 + dKL_e_W0_2 + dW0_e_l + dW0_e_l_2 + dW0_e_r + dW0_e_r_2
        
        
        grad_list = [grad_e_W0, grad_e_b0, grad_W_mu, grad_b_mu, grad_W_logvar, grad_b_logvar,
                     grad_d_W0, grad_d_b0, grad_d_W1, grad_d_b1]
        
        ########################################
        #Calculate update using Adam
        ########################################
        self.t += 1
        for i, grad in enumerate(grad_list):
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * grad
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * np.power(grad, 2)
            m_h = self.m[i] / (1 - (self.b1 ** self.t))
            v_h = self.v[i] / (1 - (self.b2 ** self.t))
            grad_list[i] = m_h / (np.sqrt(v_h) + self.e)
        
        # Update all weights
        for idx in range(self.batch_size):
            # Encoder Weights
            self.e_W0 = self.e_W0 - self.learning_rate*grad_list[0][idx]
            self.e_b0 = self.e_b0 - self.learning_rate*grad_list[1][idx]
    
            self.e_W_mu = self.e_W_mu - self.learning_rate*grad_list[2][idx]
            self.e_b_mu = self.e_b_mu - self.learning_rate*grad_list[3][idx]
            
            self.e_W_logvar = self.e_W_logvar - self.learning_rate*grad_list[4][idx]
            self.e_b_logvar = self.e_b_logvar - self.learning_rate*grad_list[5][idx]
    
            # Decoder Weights
            self.d_W0 = self.d_W0 - self.learning_rate*grad_list[6][idx]
            self.d_b0 = self.d_b0 - self.learning_rate*grad_list[7][idx]
            
            self.d_W1 = self.d_W1 - self.learning_rate*grad_list[8][idx]
            self.d_b1 = self.d_b1 - self.learning_rate*grad_list[9][idx]
    
    def train(self):
        
        #Read in training data
        trainX, _, train_size = mnist_reader(self.numbers)
        
        np.random.shuffle(trainX)
        
        #set batch indices
        batch_idx = train_size//self.batch_size
        
        total_loss = 0
        total_kl = 0
        total = 0
        
        for epoch in range(self.epochs):
            for idx in range(batch_idx):
                # prepare batch and input vector z
                train_batch = trainX[idx*self.batch_size:idx*self.batch_size + self.batch_size]
                #ignore batch if there are insufficient elements 
                if train_batch.shape[0] != self.batch_size:
                    break
                
                ################################
                #		Forward Pass
                ################################
                
                out, mu, logvar = self.forward(train_batch)
                
                # Reconstruction Loss
                rec_loss = BCE_loss(out, train_batch)
                
                #K-L Divergence
                kl = -0.5 * np.sum(1 + logvar - np.power(mu, 2) - np.exp(logvar))
                
                loss = rec_loss + kl
                loss = loss / self.batch_size
                
                #Loss Recordkeeping
                total_loss += rec_loss / self.batch_size
                total_kl += kl / self.batch_size
                total += 1

                ################################
                #		Backward Pass
                ################################
                # for every result in the batch
                # calculate gradient and update the weights using Adam
                self.backward(train_batch, out)	

                self.img = np.squeeze(out, axis=3) * 2 - 1

                print("Epoch [%d] Step [%d]  RC Loss:%.4f  KL Loss:%.4f  lr: %.4f"%(
                        epoch, idx, rec_loss / self.batch_size, kl / self.batch_size, self.learning_rate))
                
            # if cpu_enabled == 1:
            sample = np.array(self.img)
            # else: 
                # sample = np.asnumpy(self.img)
            
            #save image result every epoch
            img_tile(sample, self.img_path, epoch, idx, "res", True)


if __name__ == '__main__':

    # Adjust the numbers that appear in the training data. Less numbers helps 
    # run the program to see faster results
    numbers = [1, 2, 3]
    model = VAE(numbers)
    model.train()
    