"""
    coding: utf-8

    Vanishing Gradients

    We will demonstrate the difference between using sigmoid and ReLU nonlinearities in a simple
    neural network with two hidden layers. This notebook is built off of a minimal net demo done
    by Andrej Karpathy for CS 231n, which you can check out here:
    http://cs231n.github.io/neural-networks-case-study/
"""

# Setup
import os
import numpy as np
import matplotlib.pyplot as plt


plt.rcParams['figure.figsize'] = (10.0, 8.0)  # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'


plot_dir = 'plots'
plot_n = 0
try:
    os.makedirs(plot_dir)
except FileExistsError:
    pass


def save_plot(name):
    global plot_n
    path = os.path.join(plot_dir, '%03d-%s-%s.png' % (plot_n, rubric, name))
    print('Saving "%s"' % path)
    plt.savefig(path)
    plot_n += 1


# generate random data -- not linearly separable
np.random.seed(1)
N = 80  # number of points per class
D = 2  # dimensionality
K = 4  # number of classes
N_EPOCHS = 50000
d_theta = 2 * np.pi / K
S = 1.5
R = 0.6
delta = R / K
print('N=%d D=%d K=%d N_EPOCHS=%d' % (N, D, K, N_EPOCHS))
rubric = 'N=%d-K=%d-S=%.1f-R=%.1f' % (N, K, S, R)


X = np.zeros((N * K, D))
num_train_examples = X.shape[0]
y = np.zeros(N * K, dtype='uint8')
for j in range(K):
    ix = range(N * j, N * (j + 1))
    r = np.sqrt(np.linspace(0.0, 1, N))  # radius
    t = np.linspace(j * d_theta, j * d_theta + S * np.pi, N) + np.random.randn(N) * delta  # theta
    X[ix] = np.c_[r * np.sin(t), r * np.cos(t)]
    y[ix] = j
fig = plt.figure()
plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.Spectral)
plt.xlim([-1, 1])
plt.ylim([-1, 1])
plt.legend()
save_plot('input')


# The sigmoid function "squashes" inputs to lie between 0 and 1. Unfortunately, this means that for
# inputs with sigmoid output close to 0 or 1, the gradient with respect to those inputs are close to
# zero. This leads to the phenomenon of vanishing gradients, where gradients drop close to zero, and
# the net does not learn well.
#
# On the other hand, the relu function (max(0, x)) does not saturate with input size. Plot these
# functions to gain intuition.

def sigmoid(x):
    x = 1 / (1 + np.exp(-x))
    return x


def sigmoid_grad(x):
    return (x) * (1 - x)


def relu(x):
    return np.maximum(0, x)


# Let's try and see now how the two kinds of nonlinearities change deep neural net training in
# practice. Below, we build a very simple neural net with three layers (two hidden layers), for
# which you can swap out ReLU/ sigmoid nonlinearities.


def three_layer_net(NONLINEARITY, X, y, model, step_size, reg):
    """function to train a three layer neural net with either RELU or sigmoid nonlinearity via
        vanilla grad descent
    """
    # parameter initialization

    h = model['h']
    h2 = model['h2']
    W1 = model['W1']
    W2 = model['W2']
    W3 = model['W3']
    b1 = model['b1']
    b2 = model['b2']
    b3 = model['b3']

    # some hyper-parameters

    # gradient descent loop
    num_examples = X.shape[0]
    plot_array_1 = []
    plot_array_2 = []
    for i in range(N_EPOCHS):

        # FORWARD PROP

        if NONLINEARITY == 'RELU':
            hidden_layer = relu(np.dot(X, W1) + b1)
            hidden_layer2 = relu(np.dot(hidden_layer, W2) + b2)
            scores = np.dot(hidden_layer2, W3) + b3

        elif NONLINEARITY == 'SIGM':
            hidden_layer = sigmoid(np.dot(X, W1) + b1)
            hidden_layer2 = sigmoid(np.dot(hidden_layer, W2) + b2)
            scores = np.dot(hidden_layer2, W3) + b3

        exp_scores = np.exp(scores)
        probs = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)  # [N x K]

        #print(X.shape)
        #print(scores.shape)
        #print(np.sum(exp_scores, axis=1, keepdims=True).shape)
        #print(probs.shape)
        #assert False

        # compute the loss: average cross-entropy loss and regularization
        # v = probs[range(num_examples), y] -> 1d vector v[i] = probs[i, y[i]]]
        corect_logprobs = -np.log(probs[range(num_examples), y])
        data_loss = np.sum(corect_logprobs) / num_examples
        reg_loss = 0.5*reg*np.sum(W1*W1) + 0.5*reg*np.sum(W2*W2) + 0.5*reg*np.sum(W3*W3)
        loss = data_loss + reg_loss
        if i % 1000 == 0:
            print("iteration %d: loss %f" % (i, loss))

        # compute the gradient on scores
        dscores = probs
        dscores[range(num_examples), y] -= 1
        dscores /= num_examples

        # BACKPROP HERE
        dW3 = (hidden_layer2.T).dot(dscores)
        db3 = np.sum(dscores, axis=0, keepdims=True)

        if NONLINEARITY == 'RELU':

            # backprop ReLU nonlinearity here
            dhidden2 = np.dot(dscores, W3.T)
            dhidden2[hidden_layer2 <= 0] = 0
            dW2 = np.dot( hidden_layer.T, dhidden2)
            plot_array_2.append(np.sum(np.abs(dW2)) / np.sum(np.abs(dW2.shape)))
            db2 = np.sum(dhidden2, axis=0)
            dhidden = np.dot(dhidden2, W2.T)
            dhidden[hidden_layer <= 0] = 0

        elif NONLINEARITY == 'SIGM':

            # backprop sigmoid nonlinearity here
            dhidden2 = dscores.dot(W3.T)*sigmoid_grad(hidden_layer2)
            dW2 = (hidden_layer.T).dot(dhidden2)
            plot_array_2.append(np.sum(np.abs(dW2))/np.sum(np.abs(dW2.shape)))
            db2 = np.sum(dhidden2, axis=0)
            dhidden = dhidden2.dot(W2.T)*sigmoid_grad(hidden_layer)

        dW1 =  np.dot(X.T, dhidden)
        plot_array_1.append(np.sum(np.abs(dW1))/np.sum(np.abs(dW1.shape)))
        db1 = np.sum(dhidden, axis=0)

        # add regularization
        dW3 += reg * W3
        dW2 += reg * W2
        dW1 += reg * W1

        #option to return loss, grads -- uncomment next comment
        grads={}
        grads['W1']=dW1
        grads['W2']=dW2
        grads['W3']=dW3
        grads['b1']=db1
        grads['b2']=db2
        grads['b3']=db3
        #return loss, grads

        # update
        W1 += -step_size * dW1
        b1 += -step_size * db1
        W2 += -step_size * dW2
        b2 += -step_size * db2
        W3 += -step_size * dW3
        b3 += -step_size * db3
    # evaluate training set accuracy
    if NONLINEARITY == 'RELU':
        hidden_layer = relu(np.dot(X, W1) + b1)
        hidden_layer2 = relu(np.dot(hidden_layer, W2) + b2)
    elif NONLINEARITY == 'SIGM':
        hidden_layer = sigmoid(np.dot(X, W1) + b1)
        hidden_layer2 = sigmoid(np.dot(hidden_layer, W2) + b2)
    scores = np.dot(hidden_layer2, W3) + b3