import numpy as np
import tensorflow as tf

from tensorflow.examples.tutorials.mnist import input_data

# Note that because the MNIST data matrices are transposed compared to examples 1-4,
# we now calculate x*W instead of W*x in the predictor. This means that W and b also
# is transposed compared to examples 1-4
class LogisticRegressionModel:
    def __init__(self):
        # model input
        self.x = tf.placeholder(tf.float32)
        self.y = tf.placeholder(tf.float32)
        
        # model variables
        W = tf.Variable(tf.zeros([784, 10]))
        b = tf.Variable(tf.zeros([10]))
    
        # predictor. The softmax function is a multiclass extension of the sigmoid function
        f = tf.nn.softmax(tf.matmul(self.x, W) + b)
    
        # uses multiclass Cross Entropy. The error function is also sometimes called cost or loss function
        self.error = tf.reduce_mean(-tf.reduce_sum(self.y * tf.log(f)))
        
        # the following function is similar to the function above, but is more numerically stable:
        # self.error = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=self.y, logits=tf.matmul(self.x, W) + b))
        
        # accuracy
        self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(f, 1), tf.argmax(self.y, 1)), tf.float32))
model = LogisticRegressionModel()

mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

minimize_operation = tf.train.GradientDescentOptimizer(0.01).minimize(model.error)

# create session object for running TensorFlow operations
session = tf.Session()

# initialize tf.Variable objects
session.run(tf.global_variables_initializer())

batch_size=100
number_of_batches = mnist.train.num_examples/batch_size
for epoch in range(20):
    error_total=0
    for i in range(int(number_of_batches)):
        x_train, y_train = mnist.train.next_batch(batch_size)
        _, error = session.run([minimize_operation, model.error], {model.x: x_train, model.y: y_train})
        error_total+=error
    
    print("epoch", "%s" % epoch, "error", "%s" % (error_total/number_of_batches))

    print("accuracy", session.run(model.accuracy, {model.x: mnist.test.images, model.y: mnist.test.labels}))

session.close()
