#!/usr/bin/env python3.5

import cred
import boto3
import model
import h5py
import numpy as np 
import tensorflow as tf
import os

def readImages_hdf5(filename):
		'''Reads hdf5 file.
			 Parameter
			 ---------
			 filename : the name of the hdf5 file
		'''
		file = h5py.File(filename + '.h5', "r+") #open the hdf5 file.
		
		hdf5_images = np.array(file["/images"]).astype("uint8") #read the images as np array
		hdf5_labels = np.array(file["/meta"]).astype("uint8") #read the images as np array
		
		return hdf5_images, hdf5_labels



s3 = boto3.resource('s3', region_name='ap-southeast-1', aws_access_key_id=cred.ACCESS_KEY, aws_secret_access_key=cred.SECRET_KEY)
bucket = s3.Bucket('sg-batch-test')

model_save_path = '/opt/ml/model/'

mnist_file = open('mnist.h5', 'wb')

object_mnist = bucket.Object('mnist.h5')

object_mnist.download_fileobj(mnist_file)

mnist_file.close()

images,labels = readImages_hdf5('mnist')
images = images.astype('float32')
labels = labels.astype('int32')


testmodel = model.model()

sess = tf.InteractiveSession()

sess.run(tf.global_variables_initializer())

saver = tf.train.Saver(tf.global_variables())

for idx in range(100):

	label = np.zeros((1,10))
	label[0][labels[idx]] = 1.0

	image = np.reshape(images[idx], (1,28,28,1))

	_, _loss,accuracy = sess.run([testmodel.optimizer, testmodel.loss, testmodel.accuracy], feed_dict={
			testmodel.X : image,
			testmodel.Y : label,
			testmodel.dropout : 1.0
		})

# if not os.path.exists(model_save_path):
#   os.makedirs(model_save_path)

for key in bucket.objects.filter():
	key.delete()

saver.save(sess, model_save_path+'model.ckpt')

sess.close()

print("Training Success!")
listing = os.listdir(model_save_path)
print(listing)