import numpy as np
import sys
import os
import keras
import SimpleITK as sitk

from keras.models import Sequential
import tensorflow as tf
# tf.device('/gpu:2')
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from keras.models import Model
from keras.layers import Input, concatenate, Conv3D, MaxPooling3D, Conv3DTranspose, AveragePooling3D, ZeroPadding3D, BatchNormalization, Activation, Dropout, Reshape, Permute
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint,  CSVLogger, EarlyStopping
from keras.utils import plot_model
from keras import backend as K
from sklearn.feature_extraction.image import extract_patches as sk_extract_patches
from skimage.util.shape import view_as_windows
from sklearn.feature_extraction.image import _make_edges_3d
from keras.models import *
from keras.layers import *
from keras.models import Sequential
from keras.callbacks import *
from matplotlib import pyplot
import argparse
import time

# potlaceni chyb
import warnings
warnings.filterwarnings("ignore")

# navazane soubory
from data import *
from tools import *
from models import *
from predict import *
from generator import *



K.set_image_data_format('channels_last')  # TF dimension ordering in this code
tf.keras.backend.set_image_data_format('channels_last')

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="3" # The GPU id to use, usually either "0" or "1"



def get_weights_path(weights_name = 'keras_script.h5'):
	model_dest = os.path.join(os.getcwd() + '/weights/')
	# Save model and weights
	if not os.path.isdir(model_dest):
		os.makedirs(model_dest)
	model_path = os.path.join(model_dest, weights_name)
	return model_path


def model_settings():
	return {
		'orig_image_shape': (48,240,240),
		'patch_shape': (48,48,48),
		'stride_shape': (24,24,24),	
		'batch': 10,
		'batch_predict': 4,
		'patch_weights' : 'patch_weights.h5',	# MS['patch_weights']
		'patch_weights_mc' : 'patch_weights_mc.h5',	# MS['patch_weights']
		'fit_weights' : 'fit_weights.h5',	# MS['patch_weights']
		'fit_weights_mc' : 'fit_weight_mcs.h5',	# MS['patch_weights']
		# unet settings
		'unet_layers' : 4,
		'min_unet_filters' : 16,
		'num_classes': 10
	}




def train():
	MS = model_settings()
	# input params
	patch_shape = MS['patch_shape'] 	#  (8,24,24)
	stride_shape =  MS['stride_shape'] 	#[int(x/2) for x in patch_shape]
	batch = MS['batch']

	# Design model
	#model = create_unet_model3D((*patch_shape,1), 10, 3, 64 )
	model = dense_unet(patch_shape)
	#model = create_unet_model3D_FULL_IMAGE((*patch_shape, 1), MS['num_classes'], MS['unet_layers'], MS['min_unet_filters'])
	#model = get_unet_edited(patch_shape)
	#model = get_unet2_FULL_IMAGE(*patch_shape)
	#model = get_unet2(*patch_shape, 10)

	# load data
	X,Y= prepare_and_load_train_data()
	training_generator = MY_Generator(batch, patch_shape, stride_shape, X, Y, 'fit', 0.2, False)
	validation_generator = MY_Generator(batch, patch_shape, stride_shape, X, Y, 'val', 0.2, False)

	# load weights
	try:
		if os.path.exists(get_weights_path(MS['patch_weights'])):
			model.load_weights(get_weights_path(MS['patch_weights']))
			print('TRAIN BY PATCHES: vahy nacteny OK')
	except:
		print("vahy nejdou nacist")
	#weights = get_weights_path(MS['patch_weights'])
	p("NUM of patches per image: "+ str(training_generator.get_number_of_patches()))

	# Train model on dataset
	history = model.fit_generator(generator=training_generator,
						validation_data=validation_generator,
						use_multiprocessing=False,
						steps_per_epoch = training_generator.get_steps_per_epoch(),
						epochs = 300,
						verbose=1,
						workers=1,
						callbacks=[
								ModelCheckpoint(get_weights_path(MS['patch_weights']), monitor='val_loss', verbose=1, save_best_only=True),
								CSVLogger(os.path.join(os.getcwd() + '/log.txt'), append=True, separator=';'), 
								CheckPointCallback(5),
								EarlyStopping(patience=10, verbose=1)#,
								#TensorBoard(log_dir=os.path.join(os.getcwd() + '/graph'), histogram_freq=0, batch_size=batch, write_graph=True, write_images=False, update_freq='epoch')
							]
						)
	try:
		# ploting
		pyplot.plot(history.history['loss'])
		pyplot.plot(history.history['val_loss'])
		pyplot.title('Model train vs validation loss')
		pyplot.ylabel('loss')
		pyplot.xlabel('epoch')
		pyplot.legend(['train', 'validation'], loc='upper right')
		pyplot.show()
		pyplot.savefig('train_fit_gen_fig.png')

		model.save(get_weights_path(MS['patch_weights']))
		print('Saved trained model at %s ' % get_weights_path(MS['patch_weights']))
	except:
		p('Nepovedlo se vypsat historii trénování')

	###########################################
	################# predict #################
	X = prepare_and_load_test_data()
	model.load_weights(get_weights_path(MS['patch_weights']))

	testing_generator = MY_Generator(batch, patch_shape, patch_shape , X, X, 'predict')	# testovaci data budou mit stejnej patch shape i stride, aby tam nebyl prekryv

	ret = model.predict_generator(generator=testing_generator,
						use_multiprocessing=False,
						steps = testing_generator.get_steps_per_epoch(),
						verbose=1,
						workers=1)

	print('predict_generator: ', ret.shape)
	np.save(get_prediction_path(), ret)

	###########################################
	################# eval #################
	eval_patch_prediction(X, ret, patch_shape)





def train_mc():
	MS = model_settings()
	# input params
	patch_shape = MS['patch_shape'] 	#  (8,24,24)
	stride_shape =  MS['stride_shape'] 	#[int(x/2) for x in patch_shape]
	batch = MS['batch']

	# Design model
	model = create_unet_model3D2((*patch_shape, 3), MS['num_classes'], MS['unet_layers']+1, MS['min_unet_filters']*2)

	# load data
	X,Y= prepare_and_load_train_data_multichannel()
	training_generator = MY_Generator(batch, (*patch_shape, 3), stride_shape, X, Y, 'fit', 0.2, False)
	validation_generator = MY_Generator(batch, (*patch_shape, 3), stride_shape, X, Y, 'val', 0.2, False)

	# load weights
	try:
		if os.path.exists(get_weights_path(MS['patch_weights_mc'])):
			model.load_weights(get_weights_path(MS['patch_weights_mc']))
			print('TRAIN BY PATCHES: vahy nacteny OK')
	except:
		print("vahy nejdou nacist")
	#weights = get_weights_path(MS['patch_weights_mc'])
	p("NUM of patches per image: "+ str(training_generator.get_number_of_patches()))

	# Train model on dataset
	history = model.fit_generator(generator=training_generator,
						validation_data=validation_generator,
						use_multiprocessing=False,
						steps_per_epoch = training_generator.get_steps_per_epoch(),
						epochs = 300,
						verbose=1,
						workers=1,
						callbacks=[
								ModelCheckpoint(get_weights_path(MS['patch_weights_mc']), monitor='val_loss', verbose=1, save_best_only=True),
								CSVLogger(os.path.join(os.getcwd() + '/log.txt'), append=True, separator=';'), 
								CheckPointCallback(5),
								EarlyStopping(patience=10, verbose=1)#,
								#TensorBoard(log_dir=os.path.join(os.getcwd() + '/graph'), histogram_freq=0, batch_size=batch, write_graph=True, write_images=False, update_freq='epoch')
							]
						)
	try:
		# ploting
		pyplot.plot(history.history['loss'])
		pyplot.plot(history.history['val_loss'])
		pyplot.title('Model train vs validation loss')
		pyplot.ylabel('loss')
		pyplot.xlabel('epoch')
		pyplot.legend(['train', 'validation'], loc='upper right')
		pyplot.show()
		pyplot.savefig('train_fit_gen_mc_fig.png')

		model.save(get_weights_path(MS['patch_weights_mc']))
		print('Saved trained model at %s ' % get_weights_path(MS['patch_weights_mc']))
	except:
		p('Nepovedlo se vypsat historii trénování')

	###########################################
	################# predict #################
	X = prepare_and_load_test_data_multichannel()
	X = np.moveaxis(X,  4, 0)
	X = X[0]
	model.load_weights(get_weights_path(MS['patch_weights_mc']))

	testing_generator = MY_Generator(batch, patch_shape, patch_shape , X, X, 'predict')	# testovaci data budou mit stejnej patch shape i stride, aby tam nebyl prekryv

	ret = model.predict_generator(generator=testing_generator,
						use_multiprocessing=False,
						steps = testing_generator.get_steps_per_epoch(),
						verbose=1,
						workers=1)

	print('predict_generator: ', ret.shape)
	np.save(get_prediction_path(), ret)

	###########################################
	################# eval #################
	eval_patch_prediction(X, ret, patch_shape)







def train_fit():
	MS = model_settings()
	X,Y= prepare_and_load_train_data()
	batch_size = 1
	patch_shape = (48,240,240)
	patch_shape_reshape =  patch_shape # ( 240,240,48 )

	flatten_patch_shape = patch_shape_reshape[0] * patch_shape_reshape[1] * patch_shape_reshape[2]	
	imgs_mask = np.reshape(Y, (X.shape[0], flatten_patch_shape, 1))
	Y = keras.utils.to_categorical(imgs_mask, num_classes=10)

	# pokud vyjde tedka stejnej tak to shape
	reshape_format = (X.shape[0] , *patch_shape_reshape)

	X = X.reshape(reshape_format)
	##Y = Y.reshape((*reshape_format, 10))

	X = X[..., np.newaxis]


	#model = create_unet_model3D_FULL_IMAGE((*patch_shape_reshape,1), 10, 4, 32)
	#model = get_unet2(*patch_shape)
	model = get_unet_edited((48,240,240))
	#model = get_unet2_FULL_IMAGE(*patch_shape, 10)

	model_path = get_weights_path(MS['fit_weights'])

	try:
		if os.path.exists(model_path):
			model.load_weights(model_path)
			print('vahy nacteny OK')
	except:
		print("vahy nejdou nacist")

	print(X.shape)
	print(Y.shape)

	history = model.fit(
		X, 
		Y, 
		batch_size=batch_size, 
		nb_epoch=300, 
		verbose=1, 
		shuffle=True, 
		validation_split=0.1, 
		callbacks=[
				ModelCheckpoint(model_path, monitor='val_loss',verbose=1, save_best_only=True), 
				CSVLogger(os.path.join(os.getcwd() + '/train_fit.txt'), append=True, separator=';'), 
				EarlyStopping(patience=7, verbose=1)
			]
		)

	# ploting
	pyplot.plot(history.history['loss'])
	pyplot.plot(history.history['val_loss'])
	pyplot.title('model train vs validation loss')
	pyplot.ylabel('loss')
	pyplot.xlabel('epoch')
	pyplot.legend(['train', 'validation'], loc='upper right')
	pyplot.show()
	pyplot.savefig('train_fit_fig.png')

	# model save weights
	model.save(model_path)
	print('Saved trained model at %s ' % model_path)

	# load test data
	X = prepare_and_load_test_data()
	X = X[..., np.newaxis]
	imgs_mask_test = model.predict(X, batch_size=batch_size, verbose=1)
	
	# eval
	eval_fit_prediction(X, imgs_mask_test)




def train_fit_mc():
	MS = model_settings()
	X, Y= prepare_and_load_train_data_multichannel()
	batch_size = 1
	patch_shape = (48,240,240)

	imgs_mask = np.reshape(Y, (X.shape[0], tuple_multiply(patch_shape), 1))
	Y = keras.utils.to_categorical(imgs_mask, num_classes=10)

	X1, X2, X3 = split_3channels_3d_images_into_arrays(X)

	X1 = X1[..., np.newaxis]
	X2 = X2[..., np.newaxis]
	X3 = X3[..., np.newaxis]

	model_path = get_weights_path(MS['fit_weights_mc'])
	d("end")
	model = model_3_encoders_1_decoder((*patch_shape,1))

	try:
		if os.path.exists(model_path):
			model.load_weights(model_path)
			print('vahy nacteny OK')
	except:
		print("vahy nejdou nacist")

	
	history = model.fit(
		[X1, X2, X3], 
		Y, 
		batch_size=batch_size, 
		nb_epoch=300, 
		verbose=1, 
		shuffle=True, 
		validation_split=0.1, 
		callbacks=[
				ModelCheckpoint(model_path, monitor='val_loss',verbose=1, save_best_only=True), 
				CSVLogger(os.path.join(os.getcwd() + '/log.txt'), append=True, separator=';'), 
				EarlyStopping(patience=7, verbose=1),
				CheckPointCallback_mc(7)
			]
		)

	# ploting
	pyplot.plot(history.history['loss'])
	pyplot.plot(history.history['val_loss'])
	pyplot.title('Train loss vs validation loss')
	pyplot.ylabel('loss')
	pyplot.xlabel('epoch')
	pyplot.legend(['train', 'validation'], loc='upper right')
	pyplot.show()
	pyplot.savefig('train_fit_mc_fig.png')
	
	# model save weights
	model.save(model_path)
	print('Saved trained model at %s ' % model_path)
	

	# load test data
	X = prepare_and_load_test_data_multichannel()
	X1, X2, X3 = split_3channels_3d_images_into_arrays(X)

	X1 = X1[..., np.newaxis]
	X2 = X2[..., np.newaxis]
	X3 = X3[..., np.newaxis]

	imgs_mask_test = model.predict([X1, X2, X3], batch_size=batch_size, verbose=1)
	
	# eval
	eval_fit_prediction(X, imgs_mask_test)


class CheckPointCallback(keras.callbacks.Callback):
	def __init__(self, N):
		self.N = N
		self.nb_ep = 1
		self.train_timer = time.time()

	def on_epoch_end(self, epoch, logs={}):
		if  self.nb_ep % self.N == 0:	# kazdou n-tou epou budem ukladat vystup
			try:
				MS = model_settings()
				X = prepare_and_load_test_data()
				testing_generator = MY_Generator(MS['batch'], MS['patch_shape'], MS['patch_shape'], X, X, 'predict', 0.2, False)	
				ret = self.model.predict_generator(generator=testing_generator,
									use_multiprocessing=False,
									steps = testing_generator.get_steps_per_epoch(),
									verbose=1,
									workers=1)
				print('on_epoch_end: predict_generator output shape: ', ret.shape)
				np.save(get_prediction_path(), ret)
			except:
				print("NEPOVEDLO SE PROVEST predikci")
			try:
				eval_patch_prediction(X, ret, MS['patch_shape'])
			except:
				print("NEPOVEDLO SE PROVEST EVAL")

		print('TIME: aktuální délká výpočtu: ',   (time.time() - self.train_timer)/(60*60), ' hod.' )
		self.nb_ep = self.nb_ep + 1
		return

class CheckPointCallback_mc(keras.callbacks.Callback):
	def __init__(self, N):
		self.N = N
		self.nb_ep = 1
		self.train_timer = time.time()

	def on_epoch_end(self, epoch, logs={}):
		if self.nb_ep % self.N == 0:	# kazdou n-tou epou budem ukladat vystup

			# load test data
			X = prepare_and_load_test_data_multichannel()
			X1, X2, X3 = split_3channels_3d_images_into_arrays(X)
			batch_size = 1
			X1 = X1[..., np.newaxis]
			X2 = X2[..., np.newaxis]
			X3 = X3[..., np.newaxis]

			imgs_mask_test = self.model.predict([X1, X2, X3], batch_size=batch_size, verbose=1)
			print('on_epoch_end: predict_generator output shape: ', imgs_mask_test.shape)
			try:
				# eval
				eval_fit_prediction(X, imgs_mask_test)
			except:
				print("NEPOVEDLO SE PROVEST EVAL")

		self.nb_ep = self.nb_ep + 1
		return



if __name__ == '__main__':


	# zpracovani parametru
	arguments = argparse.ArgumentParser(description="UNET")
	arguments = argparse.ArgumentParser(add_help=False)
	arguments.add_argument("--mc", "-mc", action="store_true", dest="multich")
	arguments.add_argument("--fitmc", "-fm", action="store_true", dest="fit_multi")
	arguments.add_argument("--fit", "-f", action="store_true", dest="fit")
	arguments.add_argument("--fitgen", "-fg", action="store_true", dest="fitgen")
	arguments.add_argument('--gpu', action="store", dest="gpu")
	args = arguments.parse_args()## naparsovani argumentu

	if args.gpu:
		if str(args.gpu) in ["0", "1", "2", "3" ]:
			os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
			os.environ["CUDA_VISIBLE_DEVICES"]= str(args.gpu) 	# The GPU id to use, usually either "0" or "1"
	else:
		os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
		os.environ["CUDA_VISIBLE_DEVICES"]= "2"	# The GPU id to use, usually either "0" or "1"

	print("CHOOSED GPU: " + str(os.environ["CUDA_VISIBLE_DEVICES"]))

	if args.fit:
		train_fit()
	elif args.fitgen:
		if args.multich:
			train_mc()
		else:
			train()
	elif args.fit_multi:
		train_fit_mc()
	else:
		print("lze spustit s argumentem: --fit | --fitgen | --fitmc tj. zkratky: (-f|-fg|-fm) a dale volba GPU: --gpu <cislo:0-3>")
