import tensorflow as tf
import numpy

def one_hot(y, num_classes=10):
  """Given a numpy matrix of shape [N, 1], ret one-hot matrix [N, num_classes]"""
  arr = numpy.zeros(shape=(y.shape[0], num_classes), dtype='float32')
  arr[range(y.shape[0]), y[:]] = 1.0
  return arr

mnist = tf.keras.datasets.mnist
(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

x_train = x_train[:10000]
y_train = one_hot(y_train[:10000])
x_test = x_test[:1000]
y_test = one_hot(y_test[:1000])

x_train = numpy.expand_dims(x_train, 3)
x_test = numpy.expand_dims(x_test, 3)


x = tf.placeholder(tf.float32, [None, 28, 28, 1])
y = tf.placeholder(tf.float32, [None, 10])

is_training = tf.placeholder(tf.bool, [])

## Convolutional Regime
l2_reg = tf.contrib.layers.l2_regularizer(1e-5)
net = x
for i in range(2):
  net = tf.contrib.layers.conv2d(
      net, 32, 5, padding='VALID', activation_fn=None,
      weights_regularizer=l2_reg)
  net = tf.contrib.layers.batch_norm(net, is_training=is_training, decay=0.9)
  net = tf.nn.relu(net)
  net = tf.contrib.layers.max_pool2d(net, kernel_size=(2,2), stride=2)

## Fully-connected Regime
net = tf.reshape(net, (-1, numpy.prod(net.shape[1:])))

net = tf.contrib.layers.fully_connected(
    net, 32, activation_fn=None, weights_regularizer=l2_reg)
net = tf.contrib.layers.batch_norm(net, is_training=is_training, decay=0.9)
net = tf.nn.relu(net)

net = tf.contrib.layers.fully_connected(
    net, 10, activation_fn=None, weights_regularizer=l2_reg)

## Add Loss
tf.losses.softmax_cross_entropy(onehot_labels=y, logits=net)
# Above is equiv. to: tf.losses.add_loss(tf.nn.softmax_cross_entropy_with_logits(...))


## Training
opt = tf.train.AdamOptimizer(learning_rate=0.005)
train_op = tf.contrib.training.create_train_op(tf.losses.get_total_loss(), opt)

import IPython; IPython.embed()

sess = tf.Session()
sess.run(tf.global_variables_initializer())
for j in range(10):  # Train for 5 epochs.
  batch_size = 200
  indices = numpy.random.permutation(x_train.shape[0])
  for si in range(0, x_train.shape[0], batch_size):
    se = min(si + batch_size, x_train.shape[0])
    sess.run(train_op, {
        is_training: True,
        x: x_train[indices[si:se]],
        y: y_train[indices[si:se]],
    })
  
  probs = sess.run(net, {is_training: False, x: x_test})
  accuracy = numpy.mean(probs.argmax(axis=1) == y_test.argmax(axis=1))
  print('After epoch %i: test accuracy=%f' % (j, accuracy))