# Setup library
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import numpy as np
import PIL.Image as Image
from PIL import ImageFile
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras import layers
import matplotlib.pylab as plt
import efficientnet.tfkeras as efn

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=""

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
	try:
		for gpu in gpus:
			tf.config.experimental.set_memory_growth(gpu, True)
		logical_gpus = tf.config.experimental.list_logical_devices('GPU')
		print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
	except RuntimeError as e:
		print(e)


class transfer_learning_fit(object):
	def __init__(self, config, weights):
		self.weights = weights
		self.image_shape = (config['image_shape'], config['image_shape'])
		self.batch_size = config['batch_size']
		self.learning_rate = config['learning_rate']
		self.epochs = config['epochs']
		self.optimizer = config['optimizer']
		self.model_link = config['model']
		self.class_names = np.array(['book', 'laptop', 'phone', 'wash', 'water'])

		tf.random.set_seed(2020)

	def image_generator(self):
		image_gen_train = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255,
																	rotation_range=15,
																	horizontal_flip=True,
																	brightness_range=[0.7,1.0])
		image_gen_val = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255)
		return image_gen_train, image_gen_val

	def gen_train_val_data(self):
		gen_train, gen_val = self.image_generator()

		train_data_dir = os.path.abspath('INPUT YOUR TRANING DATA SET PATH')
		train_data_gen = gen_train.flow_from_directory(directory=str(train_data_dir),
											batch_size=self.batch_size,
											color_mode='rgb',
											shuffle=True,
											target_size=self.image_shape,
											classes=list(self.class_names))

		return train_data_gen

	def select_optimizer(self, opti, lr):
		if opti == 'adam':
			return tf.keras.optimizers.Adam(learning_rate=lr)

	def set_model(self, vector_layer):
		#efficient_net = efn.EfficientNetB0(
		#	weights=None,
		#	input_shape=self.image_shape+(3,),
		#	include_top=False,
		#	pooling='max'
		#)

		#model = tf.keras.Sequential([
		#	efficient_net,
		#	layers.Dense(5, activation='softmax')
		#])

		mobilenet_v2 = tf.keras.applications.MobileNetV2(
			weights=None,
			input_shape=self.image_shape+(3,),
			include_top=False,
			pooling='max'
		)

		model = tf.keras.Sequential([
			mobilenet_v2,
			layers.Dense(5, activation='softmax')
		])

		return model

	def build_model(self):
		feature_vector_url = self.model_link
		feature_vector_layer = hub.KerasLayer(feature_vector_url,
										input_shape=self.image_shape+(3,))
		
		feature_vector_layer.trainable = True

		made_model = self.set_model(feature_vector_layer)

		print(made_model.summary())

		made_model.compile(
			optimizer=self.select_optimizer(self.optimizer, self.learning_rate),
			loss='categorical_crossentropy',
			metrics=['acc'])

		return made_model, feature_vector_layer

	def train_model_tosave(self, weight):
		callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3)

		if weight == list():
			local_model, feature_layer = self.build_model()
			gen_train_data = self.gen_train_val_data()
			local_model.fit_generator(gen_train_data, epochs=self.epochs, callbacks=[callback])
		else:
			local_model, feature_layer = self.build_model()
			gen_train_data = self.gen_train_val_data()
			local_model.set_weights(weight)
			local_model.fit_generator(gen_train_data, epochs=self.epochs, callbacks=[callback])
			
		return local_model.get_weights()

	def get_weight_finetune_model(self, expath, feature_layer, gtrain_data):
		reloaded_model = tf.keras.models.load_model(expath)
		
		feature_layer.trainable = True

		callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3)

		reloaded_model.compile(
			optimizer=self.select_optimizer(self.optimizer, self.learning_rate*0.1),
			loss='categorical_crossentropy',
			metrics=['accuracy'])
		reloaded_model.fit_generator(gtrain_data, epochs=self.epochs+(self.epochs*2),
						initial_epoch=self.epochs, callbacks=[callback])

		return reloaded_model.get_weights() # Dense layer weight는 제외하고 반환

	def manage_train(self):
		get_weights = list()
		training_weight = self.train_model_tosave(self.weights)
			
		return training_weight

tensorflow.keras.models.load_model tensorflow.config.experimental.list_logical_devices tensorflow.keras.preprocessing.image.ImageDataGenerator tensorflow.config.experimental.set_memory_growth tensorflow.keras.layers.Dense tensorflow.config.experimental.list_physical_devices tensorflow.keras.callbacks.EarlyStopping tensorflow.keras.optimizers.Adam numpy.array tensorflow.random.set_seed tensorflow.keras.applications.MobileNetV2
[(17, 'tensorflow.config.experimental.list_physical_devices', 'tf.config.experimental.list_physical_devices', (['"""GPU"""'], {}), True, 'import tensorflow as tf\n'), (22, 'tensorflow.config.experimental.list_logical_devices', 'tf.config.experimental.list_logical_devices', (['"""GPU"""'], {}), True, 'import tensorflow as tf\n'), (37, 'numpy.array', 'np.array', (["['book', 'laptop', 'phone', 'wash', 'water']"], {}), True, 'import numpy as np\n'), (39, 'tensorflow.random.set_seed', 'tf.random.set_seed', (['(2020)'], {}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'tf.keras.preprocessing.image.ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'rotation_range': '(15)', 'horizontal_flip': '(True)', 'brightness_range': '[0.7, 1.0]'}), True, 'import tensorflow as tf\n'), (46, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'tf.keras.preprocessing.image.ImageDataGenerator', ([], {'rescale': '(1.0 / 255)'}), True, 'import tensorflow as tf\n'), (52, 'os.path.abspath', 'os.path.abspath', (['"""INPUT YOUR TRANING DATA SET PATH"""'], {}), False, 'import os\n'), (79, 'tensorflow.keras.applications.MobileNetV2', 'tf.keras.applications.MobileNetV2', ([], {'weights': 'None', 'input_shape': '(self.image_shape + (3,))', 'include_top': '(False)', 'pooling': '"""max"""'}), True, 'import tensorflow as tf\n'), (95, 'tensorflow_hub.KerasLayer', 'hub.KerasLayer', (['feature_vector_url'], {'input_shape': '(self.image_shape + (3,))'}), True, 'import tensorflow_hub as hub\n'), (112, 'tensorflow.keras.callbacks.EarlyStopping', 'tf.keras.callbacks.EarlyStopping', ([], {'monitor': '"""loss"""', 'patience': '(3)'}), True, 'import tensorflow as tf\n'), (127, 'tensorflow.keras.models.load_model', 'tf.keras.models.load_model', (['expath'], {}), True, 'import tensorflow as tf\n'), (131, 'tensorflow.keras.callbacks.EarlyStopping', 'tf.keras.callbacks.EarlyStopping', ([], {'monitor': '"""loss"""', 'patience': '(3)'}), True, 'import tensorflow as tf\n'), (21, 'tensorflow.config.experimental.set_memory_growth', 'tf.config.experimental.set_memory_growth', (['gpu', '(True)'], {}), True, 'import tensorflow as tf\n'), (64, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'lr'}), True, 'import tensorflow as tf\n'), (88, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(5)'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras import layers\n')]