# neural network in TensorFlow very simple example.
#

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# create random training data again
Nclass = 500
D = 2  # dimensionality of input
M = 3  # hidden layer size
K = 3  # number of classes

X1 = np.random.randn(Nclass, D) + np.array([0, -2])
X2 = np.random.randn(Nclass, D) + np.array([2, 2])
X3 = np.random.randn(Nclass, D) + np.array([-2, 2])
X = np.vstack([X1, X2, X3]).astype(np.float32)

Y = np.array([0] * Nclass + [1] * Nclass + [2] * Nclass)

# let's see what it looks like
plt.scatter(X[:, 0], X[:, 1], c=Y, s=100, alpha=0.5)
plt.show()

N = len(Y)
# turn Y into an indicator matrix for training
T = np.zeros((N, K))
for i in range(N):
    T[i, Y[i]] = 1


# tensor flow variables are not the same as regular Python variables
def init_weights(shape):  # we need a function for getting the weights
    return tf.Variable(tf.random_normal(shape, stddev=0.01))


def forward(X, W1, b1, W2, b2):
    Z = tf.nn.sigmoid(tf.matmul(X, W1) + b1)
    return tf.matmul(Z,
                     W2) + b2  # return not the softmax but just the activation --> thats one thing different about numpy and tensorflow, is that, when we do the cost function later, we want the logits which is these values returned here as the input, not the output of the softmax.


# create tf place holders, these represent ur X and Y data
tfX = tf.placeholder(tf.float32, [None, D])
tfY = tf.placeholder(tf.float32, [None, K])

# create the weights, similar to numpy
W1 = init_weights([D, M])  # create symbolic variables
b1 = init_weights([M])
W2 = init_weights([M, K])
b2 = init_weights([K])

py_x = forward(tfX, W1, b1, W2, b2)

cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=py_x, labels=tfY))  # compute costs
# WARNING: This op expects unscaled logits,
# since it performs a softmax on logits
# internally for efficiency.
# Do not call this op with the output of softmax,
# as it will produce incorrect results.

train_op = tf.train.GradientDescentOptimizer(0.03).minimize(cost)  # construct an optimizer
# input parameter is the learning rate
# 0.5 -> 0.967
# 0.4 -> 0.974

predict_op = tf.argmax(py_x, 1)
# input parameter is the axis on which to choose the max

# just stuff that has to be done
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

for i in range(100000):
    sess.run(train_op, feed_dict={tfX: X, tfY: T})
    pred = sess.run(predict_op, feed_dict={tfX: X, tfY: T})
    if i % 10 == 0:
        print(np.mean(Y == pred))
