# A simple neural network using TensorFlow
# Created for (self-) education purposes
# Configured to be in Jupyter
# Author: Andrew M. Chap
# Last edited August 2018

import tensorflow as tf
import numpy as np
import random
import time

# -----------------------------
# Define training function
# -----------------------------
def train_function(x):
    N = x.shape[0]
    y = np.zeros((N,1))
    for i in range(0,N):
        if x[i,0] + x[i,1] < 1:
            if x[i,0] > x[i,1]:
                y[i,0] = x[i,1]/x[i,0]
            else:
                y[i,0] = x[i,0] + x[i,1]
        else:
            if x[i,0] > x[i,1]:
                y[i,0] = x[i,0] - x[i,1]
            else:
                y[i,0] = x[i,0]*x[i,1]
    return y

# -----------------------------
# Set network parameters
# -----------------------------
# Number of nodes in each neural layer, including 
# input layer (first) and output layer (last)
layers = [2,40,40,40,1]
# More suboptimal heuristics follow:
epsilon = 0.02   # learning rate for gradient descent
numPasses = 50000
batchSize = 10000
regLambda = 1/(200*batchSize)

# Number of frames in our output video
numPlots = 201
# At which iterations we plot
plotTimes = np.linspace(0,1,numPlots)
plotIterations = np.round(plotTimes*numPasses)
# Create grid on which to plot model output
nx = 100 # number of grid points in x0 and x1
x0_grid,x1_grid = np.meshgrid(np.linspace(0,1,nx),
                              np.linspace(0,1,nx))
X_mesh = np.column_stack((x0_grid.ravel(),
                          x1_grid.ravel()))
# empty array for storing our 
# network's output during training
Y_hypothesis = np.empty([nx,nx,numPlots]) 
# true data for plotting error
Y_train = train_function(X_mesh)
Y_train = Y_train.reshape([nx,nx])
# Set up error tracking
Y_error = np.empty(numPlots)
# track computation time
compTime = np.zeros(numPlots)

# -----------------------------
# Create neural network
# -----------------------------
np.random.seed(0) # for repeatability
class Network:
    def __init__(self,
                 batchSize,
                 train_function,
                 layers,
                 epsilon):
        self.layers = layers
        self.N = len(layers)
        self.W = []
        self.B = []
        self.A = []
        self.epsilon = epsilon
        self.eDumps = 0
        self.error = 0
        self.batchSize = batchSize
        self.X_place = tf.placeholder(
            dtype = tf.float32,
            shape = [batchSize,self.layers[0]])        
        self.Y_place = tf.placeholder(
            dtype = tf.float32,
            shape=[batchSize,self.layers[-1]])        
        
        ii = 0
        for layer0, layer1 in zip(self.layers[:-1],
                                  self.layers[1:]):
            self.W.append(
                tf.Variable(
                    tf.random_normal(
                        shape=[layer0,layer1],
                        stddev=1/np.sqrt(layer0)
                        )))
            self.B.append(
                tf.Variable(
                    tf.zeros([layer1])
                    ))
            if ii == 0:
                A_prev = self.X_place
            else:
                A_prev = self.A[-1]
            
            if ii == self.N-2:
                self.Hypothesis = tf.tanh(
                        tf.matmul(A_prev, self.W[-1])
                        + self.B[-1]
                    )
            else:
                self.A.append(
                    tf.tanh(
                        tf.matmul(A_prev,self.W[-1])
                        + self.B[-1]
                        )
                    )
            ii += 1
        

network = Network(
            batchSize=batchSize,
            train_function=train_function,
            layers=layers,
            epsilon=epsilon)


cost = tf.reduce_mean(
        (
            abs(network.Y_place - network.Hypothesis)
        )
    )

#output = 

train_step = tf.train. \
             GradientDescentOptimizer(network.epsilon). \
             minimize(cost)


# -----------------------------
# Train neural network
# -----------------------------
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
# For when we want to plot the output of our network
feed_dict_test = {network.X_place: X_mesh} 
plotNumber = 0 # plotting iterator
tStart = time.time()
for i in range(numPasses):
    
    #  Use random training data input
    X = np.random.rand(network.batchSize,network.layers[0])
    
    # Generate training data output
    Y = train_function(X)
    
    feed_dict = {network.X_place: X, network.Y_place: Y}
    sess.run(train_step, feed_dict=feed_dict)
    
    # Output progress
    if i in plotIterations:
        compTime[plotNumber] = time.time() - tStart
        Y_error[plotNumber] = sess.run(cost,
                                   feed_dict=feed_dict)
        Y_hypothesis[:,:,plotNumber] = \
            sess.run(network.Hypothesis,
                     feed_dict=feed_dict_test) \
            .reshape([nx,nx])
        print("i = {} of {}, error = {:0.3f}" \
            .format(i,numPasses,Y_error[plotNumber]))
        plotNumber += 1
        
        
# -----------------------------
# Set up plot
# -----------------------------
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation
# Allows \mathrm in plots
matplotlib.rcParams['text.usetex'] = True 
# Display the animation in Jupyter
matplotlib.rc('animation', html='html5')  


plt.close()
labelsize = 14
fig,axes = plt.subplots(1,3)
fig.dpi = 120
c_space = np.linspace(0,1,11) # for colorbar
fig.set_size_inches(5,2.7)
contourplots = []
for ax in axes:

    # Fill all contours with training data grid
    contourplots.append(ax.contourf(x0_grid,x1_grid,
                                    Y_train,c_space))

    ax.set_xlim([0,1])
    ax.set_ylim([0,1])
    ax.set_xticks([0,1])

    # Share axes between subplots to save real estate
    if ax == axes[0]:
        ax.set_yticks([0,1])
        ax.set_ylabel(r'$X_1$',rotation=0,
                      fontsize=labelsize,
                      verticalalignment='center')
    else:
        ax.set_yticks([])
        ax.tick_params(labelleft=False)
    if ax == axes[1]:
        ax.set_xlabel(r'$X_0$',
                      fontsize = labelsize,
                      labelpad=-8)

axes[0].set_title(r'$\textrm{Training data}$')
axes[1].set_title(r'$\textrm{Hypothesis}$')
axes[2].set_title(r'$\left|\textrm{Error}\right|$')
fig.tight_layout()

# make space for colorbar, error plot, and annotation
subplots_adjust_right = 0.85
subplots_adjust_bottom = 0.45
fig.subplots_adjust(bottom=subplots_adjust_bottom,
                    right=subplots_adjust_right)

# Add in a colorbar with same top and bottom as subplots
boxbottom = axes[0]._position._points[0,1]
boxtop =    axes[0]._position._points[1,1]
leftmostside = axes[0]._position._points[0,0]
boxheight = boxtop - boxbottom
boxgap = axes[1]._position._points[0,0] - \
         axes[0]._position._points[1,0]
colorbaroffset = subplots_adjust_right + boxgap
colorbarwidth = 0.03
cbar_ax = fig.add_axes([colorbaroffset, boxbottom,
                        colorbarwidth, boxheight])
cbar = fig.colorbar(contourplots[0], cax=cbar_ax,
                    boundaries = [0.0, 0.5, 1.0])
cbar.ax.text(subplots_adjust_right,1.05,r'$Y$',
             rotation=0,fontsize=labelsize,
             horizontalalignment='right')

# lower and upper bounds for error plot
emin = 0.01
emax = 1
# Add in error plot
eplot = fig.add_axes([0.18, 0.17, .45, .15])
eplot.set_xlabel(r'$\textrm{Epoch}$',labelpad=-4)
eplot.set_ylabel(r'$\langle \textrm{Error} \rangle$',
                 rotation=0,labelpad=14,
                 verticalalignment='center')
eplot.semilogy(plotIterations,Y_error)
eplot.set_xlim([0,plotIterations[-1]])
eplot.set_ylim([emin,emax])
eplot.set_yticklabels(['','0.01','0.1','1'])
xticks = np.round(np.linspace(0,numPasses,5)).astype(int)
eplot.set_xticks(xticks)
xticklabels = xticks.astype(str)
xticklabels[2] = ''
eplot.set_xticklabels(xticklabels)
eplot.grid(True)

# Add in computation time annotation
ctx = 0.97
cty0 = 0.32
cty1 = cty0 - 0.08
eplot.annotate(r'$\underline{\textrm{Computation time:}}$',
                xy=(ctx, cty0), xycoords='figure fraction',
                horizontalalignment='right',
                verticalalignment='top',
                fontsize=13,color=[0,0,0])
comptimestring = r'$\textrm{{{} \small{{(HH:MM:SS)}}}}$'
eplot.annotate(comptimestring.format(
                time.strftime('%H:%M:%S',time.gmtime(0))),
                xy=(ctx, cty1), xycoords='figure fraction',
                horizontalalignment='right', 
                verticalalignment='top',
                fontsize=13,color=[0,0,0])

plt.show()

# -----------------------------
# Animate progress and results
# -----------------------------
def update(i):
    print('animating {} of {}'.format(i,numPlots))
    # replace hypothesis and error contour plots 
    # with new data
    contourplots[1] = axes[1].contourf(x0_grid,x1_grid,
                                       Y_hypothesis[:,:,i],
                                       c_space)
    contourplots[2] = axes[2].contourf(
                          x0_grid,x1_grid,
                          np.abs(Y_hypothesis[:,:,i]
                                 -Y_train),
                          c_space)
    # Update y-data for error plot
    Y_error_plot = Y_error.copy()
    # nan out the values we "haven't gotten to yet"
    Y_error_plot[i+1:] = np.nan 
    eplot.lines[0].set_ydata(Y_error_plot)
    # Update computation time string
    eplot.texts[1].set_text(comptimestring.format(
                            time.strftime('%H:%M:%S',
                            time.gmtime(compTime[i]))))

    return contourplots[1].collections + \
           contourplots[2].collections


anim = matplotlib.animation.FuncAnimation(
           fig, update, frames=numPlots,
           interval=60, blit=True, repeat=True)
anim # Animate into video