# Simple neural network example
# Taken from https://medium.com/@jaschaephraim/elementary-neural-networks-with-tensorflow-c2593ad3d60b

import tensorflow as tf

# Training data for a network to learn the AND mapping
T, F = 1., -1.
bias = 1.
train_in = [
    [T, T, bias],
    [T, F, bias],
    [F, T, bias],
    [F, F, bias],
]
train_out = [
    [T],
    [F],
    [F],
    [F],
]

# Vector of weights
w = tf.Variable(tf.random_normal([3, 1]))

# Activation function: step function
def step(x):
    is_greater = tf.greater(x, 0)
    as_float = tf.to_float(is_greater)
    doubled = tf.multiply(as_float, 2)
    return tf.subtract(doubled, 1)
    
# Output and MSE
output = step(tf.matmul(train_in, w))
error = tf.subtract(train_out, output)
mse = tf.reduce_mean(tf.square(error))

# Training evaluation
delta = tf.matmul(train_in, error, transpose_a=True)
train = tf.assign(w, tf.add(w, delta))

# Setting up the session
sess = tf.Session()
sess.run(tf.global_variables_initializer())

# Running the model
err, target = 1, 0
epoch, max_epochs = 0, 10
while err > target and epoch < max_epochs:
    epoch += 1
    err, _ = sess.run([mse, train])
    print('epoch:', epoch, 'mse:', err)
    
out = sess.run(output)
print('output after training: ', out)