import os
from PIL import Image
import numpy as np

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import model_from_json
from keras.preprocessing.image import ImageDataGenerator

from imageio import imread, imwrite
from skimage.transform import resize

IMG_SIZE = 24

def collect():
	train_datagen = ImageDataGenerator(
			rescale=1./255,
			shear_range=0.2,
			horizontal_flip=True, 
		)

	val_datagen = ImageDataGenerator(
			rescale=1./255,
			shear_range=0.2,
			horizontal_flip=True,		)

	train_generator = train_datagen.flow_from_directory(
	    directory="dataset/train",
	    target_size=(IMG_SIZE, IMG_SIZE),
	    color_mode="grayscale",
	    batch_size=32,
	    class_mode="binary",
	    shuffle=True,
	    seed=42
	)

	val_generator = val_datagen.flow_from_directory(
	    directory="dataset/val",
	    target_size=(IMG_SIZE, IMG_SIZE),
	    color_mode="grayscale",
	    batch_size=32,
	    class_mode="binary",
	    shuffle=True,
	    seed=42
	)
	return train_generator, val_generator


def save_model(model):
	model_json = model.to_json()
	with open("model.json", "w") as json_file:
		json_file.write(model_json)
	# serialize weights to HDF5
	model.save_weights("model.h5")

def load_model():
	json_file = open('model.json', 'r')
	loaded_model_json = json_file.read()
	json_file.close()
	loaded_model = model_from_json(loaded_model_json)
	# load weights into new model
	loaded_model.load_weights("model.h5")
	loaded_model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
	return loaded_model

def train(train_generator, val_generator):
	STEP_SIZE_TRAIN=train_generator.n//train_generator.batch_size
	STEP_SIZE_VALID=val_generator.n//val_generator.batch_size

	print('[LOG] Intialize Neural Network')
	
	model = Sequential()

	model.add(Conv2D(filters=6, kernel_size=(3, 3), activation='relu', input_shape=(IMG_SIZE,IMG_SIZE,1)))
	model.add(AveragePooling2D())

	model.add(Conv2D(filters=16, kernel_size=(3, 3), activation='relu'))
	model.add(AveragePooling2D())

	model.add(Flatten())

	model.add(Dense(units=120, activation='relu'))

	model.add(Dense(units=84, activation='relu'))

	model.add(Dense(units=1, activation = 'sigmoid'))


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

	model.fit_generator(generator=train_generator,
	                    steps_per_epoch=STEP_SIZE_TRAIN,
	                    validation_data=val_generator,
	                    validation_steps=STEP_SIZE_VALID,
	                    epochs=20
	)
	save_model(model)

def predict(img, model):
	img = Image.fromarray(img, 'RGB').convert('L')
	print(img)
	img = resize(img, (IMG_SIZE,IMG_SIZE)).astype('float32')/255
	print(img)
	img = img.reshape(1,IMG_SIZE,IMG_SIZE,1)
	prediction = model.predict(img)
	if prediction < 0.1:
		prediction = 'closed'
	elif prediction > 0.9:
		prediction = 'open'
	else:
		prediction = 'idk'
	return prediction

def evaluate(X_test, y_test):
	model = load_model()
	print('Evaluate model')
	loss, acc = model.evaluate(X_test, y_test, verbose = 0)
	print(acc * 100)

if __name__ == '__main__':	
	train_generator , val_generator = collect()
	train(train_generator,val_generator)

tensorflow.keras.layers.AveragePooling2D tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.models.model_from_json tensorflow.keras.models.Sequential tensorflow.keras.layers.Flatten
[(19, 'keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'shear_range': '(0.2)', 'horizontal_flip': '(True)'}), False, 'from keras.preprocessing.image import ImageDataGenerator\n'), (25, 'keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'shear_range': '(0.2)', 'horizontal_flip': '(True)'}), False, 'from keras.preprocessing.image import ImageDataGenerator\n'), (63, 'tensorflow.keras.models.model_from_json', 'model_from_json', (['loaded_model_json'], {}), False, 'from tensorflow.keras.models import model_from_json\n'), (75, 'tensorflow.keras.models.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras.models import Sequential\n'), (77, 'tensorflow.keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(6)', 'kernel_size': '(3, 3)', 'activation': '"""relu"""', 'input_shape': '(IMG_SIZE, IMG_SIZE, 1)'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (78, 'tensorflow.keras.layers.AveragePooling2D', 'AveragePooling2D', ([], {}), False, 'from tensorflow.keras.layers import AveragePooling2D\n'), (80, 'tensorflow.keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(16)', 'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (81, 'tensorflow.keras.layers.AveragePooling2D', 'AveragePooling2D', ([], {}), False, 'from tensorflow.keras.layers import AveragePooling2D\n'), (83, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Flatten\n'), (85, 'tensorflow.keras.layers.Dense', 'Dense', ([], {'units': '(120)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense\n'), (87, 'tensorflow.keras.layers.Dense', 'Dense', ([], {'units': '(84)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense\n'), (89, 'tensorflow.keras.layers.Dense', 'Dense', ([], {'units': '(1)', 'activation': '"""sigmoid"""'}), False, 'from tensorflow.keras.layers import Dense\n'), (103, 'PIL.Image.fromarray', 'Image.fromarray', (['img', '"""RGB"""'], {}), False, 'from PIL import Image\n'), (105, 'skimage.transform.resize', 'resize', (['img', '(IMG_SIZE, IMG_SIZE)'], {}), False, 'from skimage.transform import resize\n')]