# imports
import numpy as np
import tensorflow as tf
from numpy import random
import math
import time
import matplotlib.pyplot as plt

"""Part 1 - Forward Propagation"""


def initialize_parameters(layer_dims):
    """
    Description: This function initializes weights and biases

    :param layer_dims: an array of the dimensions of each layer in the /
    network (layer 0 is the size of the flattened input, layer L is the output softmax)
    :return: a dictionary containing the initialized W and b parameters of each layer (W1…WL, b1…bL).
    """
    parameters = {}
    for l in range(1, len(layer_dims)):
        parameters[f'W{l}'] = np.random.randn(layer_dims[l], layer_dims[l - 1]) * np.sqrt(2 / layer_dims[l - 1])
        parameters[f'b{l}'] = np.zeros(shape=(layer_dims[l], 1))
    return parameters


def linear_forward(A, W, b):
    """
    Description: Implement the linear part of a layer's forward propagation.

    :param A: the activations of the previous layer
    :param W: the weight matrix of the current layer (of shape [size of current layer, size of previous layer])
    :param b: the bias vector of the current layer (of shape [size of current layer, 1])
    :return: Z: the linear component of the activation function (i.e., the value before applying the non-linear function)
    :return: linear_cache: a dictionary containing A, W, b (stored for making the backpropagation easier to compute)
    """
    Z = np.dot(W, A) + b
    linear_cache = dict({'A': A, 'W': W, 'b': b})
    return Z, linear_cache


def softmax(Z):
    """
    Description: Implementation of softmax function

    :param Z: the linear component of the activation function
    :return: A: the activations of the layer
    :return: activation_cache: returns Z, which will be useful for the backpropagation
    """
    numerator = np.exp(Z)
    denominator = np.sum(numerator, axis=0, keepdims=True)
    A = numerator / denominator
    activation_cache = Z
    return A, activation_cache


def relu(Z):
    """
    Description: Implementation of relu function

    :param Z: the linear component of the activation function
    :return: A: the activations of the layer
    :return: activation_cache: returns Z, which will be useful for the backpropagation
    """
    A = np.maximum(0, Z)
    activation_cache = Z
    return A, activation_cache


def linear_activation_forward(A_prev, W, B, activation):
    """
    Description: Implement the forward propagation for the LINEAR->ACTIVATION layer

    :param A_prev: activations of the previous layer
    :param W: the weights matrix of the current layer
    :param B: the bias vector of the current layer
    :param activation: the activation function to be used (a string, either “softmax” or “relu”)
    :return: A: the activations of the current layer
    :return: cache: a joint dictionary containing both linear_cache and activation_cache

    """
    if activation in ['relu', 'softmax']:
        Z, linear_cache = linear_forward(A_prev, W, B)
        A, activation_cache = globals()[activation](Z)
        cache = dict({'linear_cache': linear_cache, 'activation_cache': activation_cache})
        return A, cache

    else:
        raise NotImplementedError(
            "The given actiavtion function was not implemented. please choose one between {relu} and {softmax}")


def l_model_forward(X, parameters, use_batchnorm):
    """
    Description: Implement forward propagation for the [LINEAR->RELU]*(L-1)->LINEAR->SOFTMAX computation

    :param X: the data, numpy array of shape (input size, number of examples)
    :param parameters: the initialized W and b parameters of each layer
    :param use_batchnorm: a boolean flag used to determine whether to apply batchnorm after the activation/
    :return:AL: the last post-activation value
    :return: caches: a list of all the cache objects generated by the linear_forward function
    """
    caches = list()
    A_prev = X
    num_layers = int(len(parameters.keys()) / 2)

    for l in range(1, num_layers):
        W = parameters[f'W{l}']
        B = parameters[f'b{l}']
        A_prev, cache = linear_activation_forward(A_prev, W, B, 'relu')

        if use_batchnorm:
            A_prev = apply_batchnorm(A_prev)

        caches.append(cache)

    W = parameters[f'W{num_layers}']
    B = parameters[f'b{num_layers}']
    AL, cache = linear_activation_forward(A_prev, W, B, 'softmax')
    caches.append(cache)

    return AL, caches


def compute_cost(AL, Y):
    """
    Description: Implement the cost function defined by equation. The requested cost function is categorical cross-entropy loss.

    :param AL: probability vector corresponding to your label predictions, shape (num_of_classes, number of examples)
    :param Y: the labels vector (i.e. the ground truth)
    :return: cost: the cross-entropy cost
    """
    inner_sum_classes = np.sum(Y * np.log(AL), axis=0, keepdims=True)
    outer_sum_samples = np.sum(inner_sum_classes, axis=1)
    m = AL.shape[1]
    cost = -1 / m * outer_sum_samples
    return cost


def apply_batchnorm(A):
    """
    Description: performs batchnorm on the received activation values of a given layer.

    :param A: the activation values of a given layer
    :return: NA: the normalized activation values, based on the formula learned in class
    """
    mu = np.mean(A, axis=0, keepdims=True)
    variance = np.var(A, axis=0, keepdims=True)
    epsilon = 0.01
    NA = (A - mu) / np.sqrt(variance + epsilon)
    return NA


"""Part 2 - Backward Propagation"""


def linear_backward(dZ, cache):
    """
    Description: Implements the linear part of the backward propagation process for a single layer

    :param dZ: the gradient of the cost with respect to the linear output of the current layer (layer l)
    :param cache: tuple of values (A_prev, W, b) coming from the forward propagation in the current layer
    :return:dA_prev: Gradient of the cost with respect to the activation (of the previous layer l-1), same shape as A_prev
    :return: dW: Gradient of the cost with respect to W (current layer l), same shape as W
    :return: db: Gradient of the cost with respect to b (current layer l), same shape as b
    """
    m = dZ.shape[1]
    dA_prev = np.dot(cache['W'].T, dZ)
    dW = (1 / m) * np.dot(dZ, cache['A'].T)
    db = (1 / m) * np.sum(dZ, axis=1, keepdims=True)
    return dA_prev, dW, db


def linear_activation_backward(dA, cache, activation):
    """
    Description: Implements the backward propagation for the LINEAR->ACTIVATION layer.
    The function first computes dZ and then applies the linear_backward function.

    :param dA: post activation gradient of the current layer
    :param cache: contains both the linear cache and the activations cache
    :param activation: the activation function name = ['relu' or 'softmax']
    :return: dA_prev: Gradient of the cost with respect to the activation (of the previous layer l-1), same shape as A_prev
    :return: dW: Gradient of the cost with respect to W (current layer l), same shape as W
    :return: db: Gradient of the cost with respect to b (current layer l), same shape as b
    """
    if activation == 'softmax':
        dZ = softmax_backward(dA, cache['activation_cache'])
    elif activation == 'relu':
        dZ = relu_backward(dA, cache['activation_cache'])
    dA_prev, dW, db = linear_backward(dZ, cache['linear_cache'])
    return dA_prev, dW, db


def relu_backward(dA, activation_cache):
    """
    Description: Implements backward propagation for a ReLU unit

    :param dA: the post-activation gradient
    :param activation_cache: contains Z (stored during the forward propagation)
    :return: dZ: gradient of the cost with respect to Z
    """
    derivative = activation_cache
    derivative[derivative >= 0] = 1
    derivative[derivative < 0] = 0
    dZ = dA * derivative
    return dZ


def softmax_backward(dA, activation_cache):
    """
    Description: Implements backward propagation for a softmax unit

    :param dA: the post-activation gradient
    :param activation_cache: contains Z (stored during the forward propagation)
    :return: dZ: gradient of the cost with respect to Z
    """
    AL = activation_cache['AL']
    Y = activation_cache['Y']
    dZ = AL - Y
    return dZ


def l_model_backward(AL, Y, caches):
    """
    Description: Implement the backward propagation process for the entire network.

    :param AL: the probabilities vector, the output of the forward propagation (L_model_forward)
    :param Y: the true labels vector (the "ground truth" - true classifications)
    :param caches: list of caches containing for each layer: a) the linear cache; b) the activation cache
    :return: Grads: a dictionary with the gradients
              grads["dA" + str(l)] = ...
              grads["dW" + str(l)] = ...
              grads["db" + str(l)] = ...
    """
    grads = dict()
    num_of_layers = len(caches)

    last_cache = caches[-1]
    last_cache['activation_cache'] = dict({'AL': AL, 'Y': Y, 'Z': last_cache['activation_cache']})

    dA_prev, dW, db = linear_activation_backward(None, last_cache, "softmax")   # dA = None cause not necessary
    grads["dA" + str(num_of_layers)] = dA_prev
    grads["dW" + str(num_of_layers)] = dW
    grads["db" + str(num_of_layers)] = db

    for l in range(num_of_layers - 1, 0, -1):
        dA_prev, dW, db = linear_activation_backward(dA_prev, caches[l - 1], "relu")
        grads["dA" + str(l)] = dA_prev
        grads["dW" + str(l)] = dW
        grads["db" + str(l)] = db

    return grads


def update_parameters(parameters, grads, learning_rate):
    """
    Description: Updates parameters using gradient descent

    :param parameters: a python dictionary containing the DNN architecture’s parameters
    :param grads: a python dictionary containing the gradients (generated by L_model_backward)
    :param learning_rate: the learning rate used to update the parameters (the “alpha”)
    :return: parameters: the updated values of the parameters object provided as input
    """
    num_of_layers = int(len(parameters.keys()) / 2)
    for l in range(1, num_of_layers + 1):
        parameters[f"W{l}"] = parameters[f"W{l}"] - learning_rate * grads[f"dW{l}"]
        parameters[f"b{l}"] = parameters[f"b{l}"] - learning_rate * grads[f"db{l}"]
    return parameters


"""Part 3 - Train and Predict"""


def train_validation_split(X, Y, train_size):
    """
    Description: (auxiliary function) split the train set into train and validation sets
    :param X: train set samples
    :param Y: train set labels
    :param train_size: percentage of the train set
    :return: tuples of (x_train, y_train), (x_val, y_val)
    """
    # create indices for train and validation sets
    indices = list(range(0, X.shape[1]))
    random.shuffle(indices)
    num_of_x_train_samples = math.ceil(X.shape[1] * train_size)

    # split train & validation
    x_train, y_train = X[:, indices[0:num_of_x_train_samples]], Y[:, indices[0:num_of_x_train_samples]]
    x_val, y_val = X[:, indices[num_of_x_train_samples:X.shape[1]]], Y[:, indices[num_of_x_train_samples:X.shape[1]]]

    return (x_train, y_train), (x_val, y_val)


def l_layer_model(X, Y, layers_dims, learning_rate, num_iterations, batch_size):
    """
    Description: Implements a L-layer neural network. All layers but the last should have the ReLU activation function,
    and the final layer will apply the softmax activation function. The size of the output layer should be equal to
    the number of labels in the data. Please select a batch size that enables your code to run well
    (i.e. no memory overflows while still running relatively fast).
    the function should use the earlier functions in the following order:
    initialize -> L_model_forward -> compute_cost -> L_model_backward -> update parameters

    :param X: the input data, a numpy array of shape (height*width , number_of_examples)
    :param Y: the “real” labels of the data, a vector of shape (num_of_classes, number of examples)
    :param layers_dims: a list containing the dimensions of each layer, including the input
    :param learning_rate: alpa parameter
    :param num_iterations: number of iterations - each iteration equals to one batch
    :param batch_size: the number of examples in a single training batch
    :return: parameters: the parameters learnt by the system during the training (the same parameters that were updated in the update_parameters function).
    :return: costs:  the values of the cost function (calculated by the compute_cost function). One value is to be saved after each 100 training iterations (e.g. 3000 iterations -> 30 values).

    """
    costs_val = []
    costs_train = []
    parameters = initialize_parameters(layers_dims)
    is_batchnorm = True  # change to True when batchnorm is needed
    epoch = 0
    val_accuracy = -1
    train_accuracy = -1
    log = []

    (x_train, y_train), (x_val, y_val) = train_validation_split(X, Y, 0.8)

    # split x and y train sets to batches
    num_of_batches = math.ceil(x_train.shape[1] / batch_size)
    batches_x = np.array_split(x_train, num_of_batches, axis=1)
    batches_y = np.array_split(y_train, num_of_batches, axis=1)

    for num_of_iteration in range(0, num_iterations):

        batch_num = num_of_iteration % num_of_batches  # current batch number
        current_batch_x, current_batch_y = batches_x[batch_num], batches_y[batch_num]  # get current batches

        AL, caches = l_model_forward(current_batch_x, parameters, is_batchnorm)
        grads = l_model_backward(AL, current_batch_y, caches)
        parameters = update_parameters(parameters, grads, learning_rate)

        # each hundred iterations compute costs for current batch and validation set, and accuracy for validation and
        # train sets
        if num_of_iteration % 100 == 0:
            AL_val, caches_val = l_model_forward(x_val, parameters, is_batchnorm)
            costs_val.append(compute_cost(AL_val, y_val))
            val_accuracy = predict(x_val, y_val, parameters)

            AL_batch, caches_train = l_model_forward(current_batch_x, parameters, is_batchnorm)
            costs_train.append(compute_cost(AL_batch, current_batch_y))
            log.append(f"Iteration: {num_of_iteration}, Cost: {costs_train[-1]}")
            train_accuracy = predict(x_train, y_train, parameters)

            print(
                f"Epoch: {epoch}, Iteration: {num_of_iteration}, batch_loss: {costs_train[-1]}, train_accuracy: {train_accuracy}, val_loss: {costs_val[-1]}, validation_accuracy: {val_accuracy}")

            # stopping criterion - no improvement on the validation set - threshold = 0.05
            if len(costs_val) > 2 and (costs_val[-1] - costs_val[-2] >= 0.005) and (costs_val[-2] - costs_val[-3] >= 0.005):
                print("Early stopping reached.")
                break

        # count epochs
        if num_of_iteration % num_of_batches == 0 and num_of_iteration > 0:
            epoch += 1

    print(f"val_accuracy {val_accuracy}")
    plot_model_history(costs_val, "Validation")

    print(f"train_accuracy {train_accuracy}")
    plot_model_history(costs_train, "Training")

    print(*log, sep="\n")  # for report

    return parameters, costs_train


def predict(X, Y, parameters):
    """
    Description: The function receives an input data and the true labels and calculates the accuracy of the trained neural network on the data.

    :param X: the input data, a numpy array of shape (height*width, number_of_examples)
    :param Y: the “real” labels of the data, a vector of shape (num_of_classes, number of examples)
    :param parameters: a python dictionary containing the DNN architecture’s parameters
    :return: accuracy – the accuracy measure of the neural net on the provided data
    (i.e. the percentage of the samples for which the correct label receives the highest confidence score).
     Use the softmax function to normalize the output values.
    """
    is_batchnorm = True    # change to True when batchnorm is needed
    AL, caches = l_model_forward(X, parameters, is_batchnorm)
    y_predict = (AL == np.amax(AL, axis=0)).astype(int)  # the class with the maximum prob is the predicted class
    accuracy = np.sum(y_predict * Y) / AL.shape[1]  # sum number of correct predictions
    return accuracy


def plot_model_history(costs_list, type):
    """
    Description: (auxiliary function) Plot graph of cost per 100 iterations
    :param costs_list:
    :param type: str - validation or training
    """
    x_index = range(0, len(costs_list)*100, 100)
    plt.plot(x_index, costs_list)
    plt.title(f'{type} Model Costs')
    plt.ylabel('Costs')
    plt.xlabel('Iterations')
    plt.show()


"""Part 4 - MNIST classification - W/out batchnorm"""

# load MNIST data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(path='mnist.npz')

# normalize train and test sets
x_train = x_train / 255
x_test = x_test / 255

# flatten matrices
x_train = x_train.reshape(x_train.shape[0], 784).T
x_test = x_test.reshape(x_test.shape[0], 784).T

# encode y vectors to one hot vector
num_of_classes = len(np.unique(y_train))
y_train_one_hot = np.zeros((num_of_classes, y_train.shape[0]))
y_train_one_hot[y_train, np.arange(y_train.shape[0])] = 1
y_test_one_hot = np.zeros((num_of_classes, y_test.shape[0]))
y_test_one_hot[y_test, np.arange(y_test.shape[0])] = 1

# run network
layers_dims = [x_train.shape[0], 20, 7, 5, 10]
learning_rate = 0.009
num_iterations = 50000
batch_size = 32

start = time.time()

parameters, costs = l_layer_model(x_train, y_train_one_hot, layers_dims, learning_rate, num_iterations, batch_size)
accuracy = predict(x_test, y_test_one_hot, parameters)

print(f"Total time: {(time.time() - start)} seconds.")
print(f"Test Accuracy: {accuracy}")