from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

tf.logging.set_verbosity(tf.logging.INFO)


def model_fn(features, labels, mode):
  input_data = tf.reshape(features['x'], (-1, 28, 28, 1))

  conv1 = tf.layers.conv2d(
    inputs=input_data,
    filters=32,
    kernel_size=(3, 3),
    padding='same',
    activation=tf.nn.relu
  )
  conv2 = tf.layers.conv2d(
    inputs=conv1,
    filters=32,
    kernel_size=(3, 3),
    padding='same',
    activation=tf.nn.relu
  )
  pool1 = tf.layers.max_pooling2d(inputs=conv2, pool_size=(2, 2), strides=2)

  conv3 = tf.layers.conv2d(
    inputs=pool1,
    filters=64,
    kernel_size=(3, 3),
    padding="same",
    activation=tf.nn.relu)
  conv4 = tf.layers.conv2d(
    inputs=conv3,
    filters=64,
    kernel_size=(3, 3),
    padding='same',
    activation=tf.nn.relu
  )
  pool2 = tf.layers.max_pooling2d(inputs=conv4, pool_size=(2, 2), strides=2)

  pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
  dropout = tf.layers.dropout(
    inputs=pool2_flat, rate=0.5, training=mode == tf.estimator.ModeKeys.TRAIN)
  dense = tf.layers.dense(inputs=dropout, units=1024, activation=tf.nn.relu)

  decode = tf.layers.dense(inputs=dense, units=pool2_flat.get_shape()[1])
  decode_reshape = tf.reshape(decode, [-1, 7, 7, 64])

  deconv1 = tf.layers.conv2d_transpose(
    inputs=decode_reshape,
    filters=64,
    kernel_size=(3, 3),
    strides=1,
    padding='same',
    activation=tf.nn.relu)
  deconv2 = tf.layers.conv2d_transpose(
    inputs=deconv1,
    filters=32,
    kernel_size=(3, 3),
    strides=2,
    padding='same',
    activation=tf.nn.relu)

  deconv3 = tf.layers.conv2d_transpose(
    inputs=deconv2,
    filters=32,
    kernel_size=(3, 3),
    strides=1,
    padding='same',
    activation=tf.nn.relu)
  deconv4 = tf.layers.conv2d_transpose(
    inputs=deconv3,
    filters=1,
    kernel_size=(3, 3),
    strides=2,
    padding='same',
    activation=tf.nn.relu)

  reconstruct = deconv4

  loss = tf.nn.l2_loss(input_data - reconstruct)

  if mode == tf.estimator.ModeKeys.TRAIN:
    optimizer = tf.train.AdamOptimizer()
    train_op = optimizer.minimize(
      loss=loss,
      global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

  return tf.estimator.EstimatorSpec(
    mode=mode, loss=loss)


def main(argv):
  mnist = tf.contrib.learn.datasets.load_dataset("mnist")
  train_data = mnist.train.images
  train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
  eval_data = mnist.test.images
  eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)

  mnist_classifier = tf.estimator.Estimator(
    model_fn=model_fn, model_dir="./log")

  train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": train_data},
    y=train_labels,
    batch_size=128,
    num_epochs=None,
    shuffle=True)

  mnist_classifier.train(
    input_fn=train_input_fn,
    steps=20000)

  eval_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": eval_data},
    y=eval_labels,
    num_epochs=1,
    shuffle=False)
  eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
  print(eval_results)


if __name__ == '__main__':
  tf.app.run()