import os
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx
plt.style.use('ggplot')

# ==== DESCRIPTION OF THE NETWORK ====
"""
A simple supervised neural network composed of a single 
layer containing 10 neurons. We have at our disposal a set
of (x, y) experimental values. The underlying function is a
sine wave but the y values contain some noise. The objective 
is for the network to learn the underlying sine model and be
able to generalize to unseen data. The input is an x value and its
corresponding y value. The output is the predicted value y_pred. 
"""

# Let's create some toy data

# We are going to say that we have seen 1000 values of some underlying representation that we aim to discover
n_observations = 1000

# Instead of having an image as our input, we're going to have values from -3 to 3.  This is going to be the
# input to our network.
xs = np.linspace(-3, 3, n_observations)

# From this input, we're going to teach our network to represent a function that looks like a sine wave.
#  To make it difficult, we are going to create a noisy representation of a sine wave by adding uniform noise.
# So our true representation is a sine wave, but we are going to make it difficult by adding some noise to the function,
# and try to have our algorithm discover the underlying cause of the data, which is the sine wave without any noise.
ys = np.sin(xs) + np.random.uniform(-0.5, 0.5, n_observations)
plt.scatter(xs, ys, alpha=0.15, marker='+')


def distance(p1, p2):
    return tf.abs(p1 - p2)


X = tf.placeholder(tf.float32, name='X')
Y = tf.placeholder(tf.float32, name='Y')

# We're going to multiply our input by 10 values, creating an "inner layer"
# of n_neurons neurons.
n_neurons = 10
W = tf.Variable(tf.random_normal([1, n_neurons]), name='W')

# and allow for n_neurons additions on each of those neurons
b = tf.Variable(tf.constant(0, dtype=tf.float32, shape=[n_neurons]), name='b')

# Instead of just multiplying, we'll put our n_neuron multiplications through a non-linearity, the tanh function.
h = tf.nn.tanh(tf.matmul(tf.expand_dims(X, 1), W) + b, name='h')

Y_pred = tf.reduce_sum(h, 1)


def train(X, Y, Y_pred, n_iterations=100, batch_size=200, learning_rate=0.02):
    cost = tf.reduce_mean(distance(Y_pred, Y))
    optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
    fig, ax = plt.subplots(1, 1)
    ax.scatter(xs, ys, alpha=0.15, marker='+')
    ax.set_xlim([-4, 4])
    ax.set_ylim([-2, 2])
    with tf.Session() as sess:
        # Here we tell TensorFlow that we want to initialize all
        # the variables in the graph so we can use them
        # This will set W and b to their initial random normal value.
        sess.run(tf.global_variables_initializer())

        # We now run a loop over epochs
        prev_training_cost = 0.0
        for it_i in range(n_iterations):
            indexes = np.random.permutation(range(len(xs)))
            n_batches = len(indexes) // batch_size
            for batch_i in range(n_batches):
                # For every batch, calculate the gradient, and update the weight and bias
                indexes_i = indexes[batch_i * batch_size: (batch_i + 1) * batch_size]
                sess.run(optimizer, feed_dict={X: xs[indexes_i], Y: ys[indexes_i]})

            training_cost = sess.run(cost, feed_dict={X: xs, Y: ys})

            if it_i % 10 == 0:
                ys_pred = Y_pred.eval(feed_dict={X: xs}, session=sess)
                ax.plot(xs, ys_pred, 'k', alpha=it_i / n_iterations)
                print(training_cost)

            # Allow the training to quit if we've reached a minimum (if the new cost isn’t very different
            # from the previous, that means we’re not going to go anywhere else, we’re turning around the minimum)
            if np.abs(prev_training_cost - training_cost) < 0.000001:
                break

            # Keep track of the training cost
            prev_training_cost = training_cost

    fig.show()
    plt.draw()


# Retrain with our new Y_pred
train(X, Y, Y_pred)
