#!/usr/bin/env python3
import sys
import time
import tensorflow as tf
from tensorflow import keras
import numpy as np
from datetime import datetime
import json
import os
import boto3
import io

print('Arguments received:', str(sys.argv))
print('Env variables', os.environ)

s3 = boto3.client('s3')
cloudwatch_client = boto3.client('cloudwatch', region_name='us-east-1')
s3_bucket = sys.argv[1]
str_keras_model_path = '/app/mnist_model.h5'
str_tflite_model_path = '/app/mnist_model.tflite'

## Importing hyperparameters
result = s3.get_object(Bucket=s3_bucket, Key=sys.argv[2])
hyperparameters = json.loads(result["Body"].read().decode())
print(hyperparameters)

parameters = {}
parameters['input'] = hyperparameters

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
train_images = train_images / 255.0
test_images = test_images / 255.0

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation=tf.nn.relu),
    keras.layers.Dense(10, activation=tf.nn.softmax)
])

model.compile(optimizer='adam', 
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=int(hyperparameters['num_of_epochs']))

test_loss, test_acc = model.evaluate(test_images, test_labels)

print('Test accuracy:', test_acc)

predictions = model.predict(test_images)
print(np.argmax(predictions[0]))
print(test_labels[0])
print("Accuracy = {}".format(test_acc))
cloudwatch_client.put_metric_data(
	Namespace='production',
	MetricData=[
		{
			'MetricName' : 'Accuracy',
			'Dimensions' : [
				{
					'Name' : 'project',
					'Value' : 'training'
				}
			],
			'Timestamp': datetime.now(),
			'Value': test_acc
		}
	]
)


# Save the Keras model
model.save(str_keras_model_path)

# Convert the Keras model to TFLite model
converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file(
	str_keras_model_path, input_shapes={'input_1' : [1,299,299,3]}
)
tflite_model = converter.convert()
open(str_tflite_model_path, "wb").write(tflite_model)

# Upload the TFLite model to S3
with open(str_tflite_model_path, 'rb') as data:
	s3.upload_fileobj(data, s3_bucket, sys.argv[4])

parameters['output'] = {
	'accuracy' : test_acc,
	'model_path' : sys.argv[4]
}
print(parameters)
s3.upload_fileobj(io.BytesIO(json.dumps(parameters).encode('utf-8')), s3_bucket, sys.argv[3])

sys.exit(0)
