# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Convolutional Neural Network Estimator for MNIST, built with tf.layers."""
from __future__ import absolute_import, division, print_function

import argparse
import json
import os

import numpy as np
import tensorflow as tf


def cnn_model_fn(features, labels, mode):
    """Model function for CNN."""
    # Input Layer
    # Reshape X to 4-D tensor: [batch_size, width, height, channels]
    # MNIST images are 28x28 pixels, and have one color channel
    input_layer = tf.reshape(features['x'], [-1, 28, 28, 1])

    # Convolutional Layer #1
    # Computes 32 features using a 5x5 filter with ReLU activation.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 28, 28, 1]
    # Output Tensor Shape: [batch_size, 28, 28, 32]
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu
    )

    # Pooling Layer #1
    # First max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 28, 28, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 32]
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)

    # Convolutional Layer #2
    # Computes 64 features using a 5x5 filter.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 14, 14, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 64]
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu
    )

    # Pooling Layer #2
    # Second max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 14, 14, 64]
    # Output Tensor Shape: [batch_size, 7, 7, 64]
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)

    # Flatten tensor into a batch of vectors
    # Input Tensor Shape: [batch_size, 7, 7, 64]
    # Output Tensor Shape: [batch_size, 7 * 7 * 64]
    pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])

    # Dense Layer
    # Densely connected layer with 1024 neurons
    # Input Tensor Shape: [batch_size, 7 * 7 * 64]
    # Output Tensor Shape: [batch_size, 1024]
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)

    # Add dropout operation; 0.6 probability that element will be kept
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)

    # Logits layer
    # Input Tensor Shape: [batch_size, 1024]
    # Output Tensor Shape: [batch_size, 10]
    logits = tf.layers.dense(inputs=dropout, units=10)

    predictions = {
        # Generate predictions (for PREDICT and EVAL mode)
        'classes': tf.argmax(input=logits, axis=1),
        # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
        # `logging_hook`.
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
    }
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    # Calculate Loss (for both TRAIN and EVAL modes)
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    # Configure the Training Op (for TRAIN mode)
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        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)

    # Add evaluation metrics (for EVAL mode)
    eval_metric_ops = {
        'accuracy': tf.metrics.accuracy(
            labels=labels, predictions=predictions['classes'])}
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)


def _load_training_data(base_dir):
    x_train = np.load(os.path.join(base_dir, 'train_data.npy'))
    y_train = np.load(os.path.join(base_dir, 'train_labels.npy'))
    return x_train, y_train


def _load_testing_data(base_dir):
    x_test = np.load(os.path.join(base_dir, 'eval_data.npy'))
    y_test = np.load(os.path.join(base_dir, 'eval_labels.npy'))
    return x_test, y_test


def _parse_args():
    parser = argparse.ArgumentParser()

    # Data, model, and output directories.
    # model_dir is always passed in from SageMaker.
    # By default this is a S3 path under the default bucket.
    parser.add_argument('--model_dir', type=str)
    parser.add_argument('--sm-model-dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAINING'))
    parser.add_argument('--hosts', type=list, default=json.loads(os.environ.get('SM_HOSTS')))
    parser.add_argument('--current-host', type=str, default=os.environ.get('SM_CURRENT_HOST'))

    return parser.parse_known_args()


def serving_input_fn():
    inputs = {'x': tf.placeholder(tf.float32, [None, 784])}
    return tf.estimator.export.ServingInputReceiver(inputs, inputs)


if __name__ == '__main__':
    args, _ = _parse_args()

    train_data, train_labels = _load_training_data(args.train)
    eval_data, eval_labels = _load_testing_data(args.train)

    # Create the Estimator
    mnist_classifier = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=args.model_dir)

    # Set up logging for predictions
    # Log the values in the 'Softmax' tensor with label 'probabilities'
    tensors_to_log = {'probabilities': 'softmax_tensor'}
    logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)

    # Train the model
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={'x': train_data},
        y=train_labels,
        batch_size=100,
        num_epochs=None,
        shuffle=True
    )

    # Evaluate the model and print results
    eval_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={'x': eval_data},
        y=eval_labels,
        num_epochs=1,
        shuffle=False
    )

    train_spec = tf.estimator.TrainSpec(train_input_fn, max_steps=20000)
    eval_spec = tf.estimator.EvalSpec(eval_input_fn)
    tf.estimator.train_and_evaluate(mnist_classifier, train_spec, eval_spec)

    if args.current_host == args.hosts[0]:
        mnist_classifier.export_savedmodel(args.sm_model_dir, serving_input_fn)

[((29, 18, 29, 60), 'tensorflow.reshape', 'tf.reshape', ({(29, 29, 29, 42): "features['x']", (29, 44, 29, 59): '[-1, 28, 28, 1]'}, {}), "(features['x'], [-1, 28, 28, 1])", True, 'import tensorflow as tf\n'), ((36, 12, 42, 5), 'tensorflow.layers.conv2d', 'tf.layers.conv2d', (), '', True, 'import tensorflow as tf\n'), ((48, 12, 48, 78), 'tensorflow.layers.max_pooling2d', 'tf.layers.max_pooling2d', (), '', True, 'import tensorflow as tf\n'), ((55, 12, 61, 5), 'tensorflow.layers.conv2d', 'tf.layers.conv2d', (), '', True, 'import tensorflow as tf\n'), ((67, 12, 67, 78), 'tensorflow.layers.max_pooling2d', 'tf.layers.max_pooling2d', (), '', True, 'import tensorflow as tf\n'), ((72, 17, 72, 52), 'tensorflow.reshape', 'tf.reshape', ({(72, 28, 72, 33): 'pool2', (72, 35, 72, 51): '[-1, 7 * 7 * 64]'}, {}), '(pool2, [-1, 7 * 7 * 64])', True, 'import tensorflow as tf\n'), ((78, 12, 78, 81), 'tensorflow.layers.dense', 'tf.layers.dense', (), '', True, 'import tensorflow as tf\n'), ((81, 14, 82, 77), 'tensorflow.layers.dropout', 'tf.layers.dropout', (), '', True, 'import tensorflow as tf\n'), ((87, 13, 87, 54), 'tensorflow.layers.dense', 'tf.layers.dense', (), '', True, 'import tensorflow as tf\n'), ((100, 11, 100, 79), 'tensorflow.losses.sparse_softmax_cross_entropy', 'tf.losses.sparse_softmax_cross_entropy', (), '', True, 'import tensorflow as tf\n'), ((114, 11, 115, 62), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', (), '', True, 'import tensorflow as tf\n'), ((131, 13, 131, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((147, 11, 147, 67), 'tensorflow.estimator.export.ServingInputReceiver', 'tf.estimator.export.ServingInputReceiver', ({(147, 52, 147, 58): 'inputs', (147, 60, 147, 66): 'inputs'}, {}), '(inputs, inputs)', True, 'import tensorflow as tf\n'), ((157, 23, 157, 94), 'tensorflow.estimator.Estimator', 'tf.estimator.Estimator', (), '', True, 'import tensorflow as tf\n'), ((162, 19, 162, 86), 'tensorflow.train.LoggingTensorHook', 'tf.train.LoggingTensorHook', (), '', True, 'import tensorflow as tf\n'), ((165, 21, 171, 5), 'tensorflow.estimator.inputs.numpy_input_fn', 'tf.estimator.inputs.numpy_input_fn', (), '', True, 'import tensorflow as tf\n'), ((174, 20, 179, 5), 'tensorflow.estimator.inputs.numpy_input_fn', 'tf.estimator.inputs.numpy_input_fn', (), '', True, 'import tensorflow as tf\n'), ((181, 17, 181, 72), 'tensorflow.estimator.TrainSpec', 'tf.estimator.TrainSpec', (), '', True, 'import tensorflow as tf\n'), ((182, 16, 182, 52), 'tensorflow.estimator.EvalSpec', 'tf.estimator.EvalSpec', ({(182, 38, 182, 51): 'eval_input_fn'}, {}), '(eval_input_fn)', True, 'import tensorflow as tf\n'), ((183, 4, 183, 76), 'tensorflow.estimator.train_and_evaluate', 'tf.estimator.train_and_evaluate', ({(183, 36, 183, 52): 'mnist_classifier', (183, 54, 183, 64): 'train_spec', (183, 66, 183, 75): 'eval_spec'}, {}), '(mnist_classifier, train_spec, eval_spec)', True, 'import tensorflow as tf\n'), ((91, 19, 91, 50), 'tensorflow.argmax', 'tf.argmax', (), '', True, 'import tensorflow as tf\n'), ((94, 25, 94, 69), 'tensorflow.nn.softmax', 'tf.nn.softmax', (), '', True, 'import tensorflow as tf\n'), ((97, 15, 97, 77), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', (), '', True, 'import tensorflow as tf\n'), ((104, 20, 104, 74), 'tensorflow.train.GradientDescentOptimizer', 'tf.train.GradientDescentOptimizer', (), '', True, 'import tensorflow as tf\n'), ((108, 15, 108, 82), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', (), '', True, 'import tensorflow as tf\n'), ((112, 20, 113, 62), 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', (), '', True, 'import tensorflow as tf\n'), ((119, 22, 119, 62), 'os.path.join', 'os.path.join', ({(119, 35, 119, 43): 'base_dir', (119, 45, 119, 61): '"""train_data.npy"""'}, {}), "(base_dir, 'train_data.npy')", False, 'import os\n'), ((120, 22, 120, 64), 'os.path.join', 'os.path.join', ({(120, 35, 120, 43): 'base_dir', (120, 45, 120, 63): '"""train_labels.npy"""'}, {}), "(base_dir, 'train_labels.npy')", False, 'import os\n'), ((125, 21, 125, 60), 'os.path.join', 'os.path.join', ({(125, 34, 125, 42): 'base_dir', (125, 44, 125, 59): '"""eval_data.npy"""'}, {}), "(base_dir, 'eval_data.npy')", False, 'import os\n'), ((126, 21, 126, 62), 'os.path.join', 'os.path.join', ({(126, 34, 126, 42): 'base_dir', (126, 44, 126, 61): '"""eval_labels.npy"""'}, {}), "(base_dir, 'eval_labels.npy')", False, 'import os\n'), ((146, 19, 146, 58), 'tensorflow.placeholder', 'tf.placeholder', ({(146, 34, 146, 44): 'tf.float32', (146, 46, 146, 57): '[None, 784]'}, {}), '(tf.float32, [None, 784])', True, 'import tensorflow as tf\n'), ((137, 60, 137, 90), 'os.environ.get', 'os.environ.get', ({(137, 75, 137, 89): '"""SM_MODEL_DIR"""'}, {}), "('SM_MODEL_DIR')", False, 'import os\n'), ((138, 53, 138, 90), 'os.environ.get', 'os.environ.get', ({(138, 68, 138, 89): '"""SM_CHANNEL_TRAINING"""'}, {}), "('SM_CHANNEL_TRAINING')", False, 'import os\n'), ((140, 60, 140, 93), 'os.environ.get', 'os.environ.get', ({(140, 75, 140, 92): '"""SM_CURRENT_HOST"""'}, {}), "('SM_CURRENT_HOST')", False, 'import os\n'), ((107, 24, 107, 50), 'tensorflow.train.get_global_step', 'tf.train.get_global_step', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((139, 65, 139, 91), 'os.environ.get', 'os.environ.get', ({(139, 80, 139, 90): '"""SM_HOSTS"""'}, {}), "('SM_HOSTS')", False, 'import os\n')]