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 *
import argparse
from matplotlib import pyplot

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

# navazane soubory
from data import *
from tools import *
from models import *
from train import *


K.set_image_data_format('channels_last')  # TF dimension ordering in this code
tf.keras.backend.set_image_data_format('channels_last')
IMAGE_ORDERING = 'channels_last' # In consistency with the json file



def get_prediction_path(predict_file = 'prediction.npy'):
	prediction_dest = os.path.join(os.getcwd() + '/weights/')
	# Save prediction and weights
	if not os.path.isdir(prediction_dest):
		os.makedirs(prediction_dest)
	prediction_path = os.path.join(prediction_dest, predict_file)
	return prediction_path



# non patch prediction
def eval_fit_prediction(X_data, predicted_data, dest_folder = '/mask_patches/model_fit/', dest_format = 'nii'):
	print("predicted shape: " + str(predicted_data.shape))
	train = X_data
	number_of_images = train.shape[0]
	patch_shape = (train.shape[1],train.shape[2],train.shape[3])

	predicted_data = predicted_data.reshape((number_of_images, *patch_shape, 10))
	for image in range(0, number_of_images):
		mask_path_3d = binarylabelToBack3D(
			predicted_data[image].shape[0], 
			predicted_data[image].shape[1], 
			predicted_data[image].shape[2], 
			predicted_data[image]
		)
		print(mask_path_3d.shape)
		mask = os.path.join(os.getcwd() + dest_folder + 'img_'+str(image)+ ('.nii.gz' if dest_format == 'nii' else ''))
		if dest_format == 'nii':
			save_3D_nii_image(mask_path_3d , mask, True)
		else:
			save_3D_png_image(mask_path_3d , mask, True)



# non patch prediction
def eval_patch_prediction(TEST_data, predicted_data, patch_shape, dest_folder = '/mask_patches/model_fitgen/', dest_format = 'nii'):
	print("predicted shape: " + str(predicted_data.shape))
	orig_shape = (48,240,240)
	res_shape = extract_sliding_window_patches(np.zeros(orig_shape), patch_shape, patch_shape)
	number_of_patches_real = res_shape.shape[0]

	number_of_images = TEST_data.shape[0]
	number_of_patches = int(predicted_data.shape[0] / number_of_images)

	if number_of_patches_real == number_of_patches:
		predicted_data = np.reshape(predicted_data, (number_of_images, number_of_patches, *patch_shape, 10))
	else:
		print("je potreba zarovnani")
		# provadi se korekce pokud sedi patche
		correct_num_patches = int(number_of_patches_real * number_of_images) 
		correction_gap = abs(predicted_data.shape[0] - correct_num_patches)
		new = np.zeros((correct_num_patches, tuple_multiply(patch_shape), 10))

		for c_gap in range(0,predicted_data.shape[0]):
			new[c_gap] = predicted_data[c_gap]
		predicted_data = new
		number_of_patches = number_of_patches_real
		predicted_data = np.reshape(predicted_data, (number_of_images, number_of_patches, *patch_shape, 10))	# final reshape

	for image in range(0, number_of_images):
		tmp_image_patches = np.empty((number_of_patches, *patch_shape), dtype=predicted_data.dtype)
		for img_patch in range(0, number_of_patches):
			mask_path_3d = binarylabelToBack3D(
				predicted_data[image][img_patch].shape[0], 
				predicted_data[image][img_patch].shape[1], 
				predicted_data[image][img_patch].shape[2], 
				predicted_data[image][img_patch]
			)
			tmp_image_patches[img_patch] = mask_path_3d

		builded_mask_image = build_image_from_patches(tmp_image_patches, orig_shape)
		if not isinstance(builded_mask_image, (np.ndarray, np.generic) ):
			builded_mask_image = build_image_from_unpadded_patches(tmp_image_patches, orig_shape)
		
		mask = os.path.join(os.getcwd() + dest_folder + 'img_'+str(image)+ ('.nii.gz' if dest_format == 'nii' else ''))
		print("3D image: "+str(mask)+" done, shape: " + str(builded_mask_image.shape))

		if dest_format == 'nii':
			save_3D_nii_image(builded_mask_image , mask, False)
			#save_3D_nii_image(builded_mask_image/builded_mask_image.max() * 255, dest_mask_path, False)
		else:
			save_3D_png_image(builded_mask_image , mask, False)




def predict():
	MS = model_settings()
	###########################################
	################# predict #################
	#

	patch_shape = MS['patch_shape'] #  (8,24,24)

	stride_shape = patch_shape 		#  [int(x/2) for x in patch_shape]
	batch =  MS['batch_predict']
	#model = create_unet_model3D_FULL_IMAGE((*patch_shape, 1), MS['num_classes'], MS['unet_layers'], MS['min_unet_filters'])
	#model = get_unet2_FULL_IMAGE(*patch_shape)
	#model = get_unet_edited(patch_shape)
	model = dense_unet(patch_shape)

	X = prepare_and_load_test_data()

	try:
		if os.path.exists(get_weights_path(MS['patch_weights'])):
			model.load_weights(get_weights_path(MS['patch_weights']))
			print('vahy nacteny OK')
	except:
		print("vahy nejdou nacist")

	testing_generator = MY_Generator(batch, patch_shape, patch_shape , X, X, 'predict', 0.0, False)	# 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 eval_predict():
	MS = model_settings()
	patch_shape = MS['patch_shape']
	X = prepare_and_load_test_data()
	print("Eval data shape: ", X.shape)
	ret=np.load(get_prediction_path())
	print("predicted data shape: ", ret.shape)
	print("patch shape: ", patch_shape)
	eval_patch_prediction(X, ret, patch_shape)




def predict_fit():
	MS = model_settings()
	batch_size = 1
	# nacteni testovacich dat
	X = prepare_and_load_test_data()
	# nema definovanej kanal
	if len(X.shape) == 4:
		X = X[..., np.newaxis]
		patch_shape = (X.shape[1],X.shape[2],X.shape[3], 1)
	else:
		patch_shape = (X.shape[1],X.shape[2],X.shape[3], X.shape[4])


	model = create_unet_model3D_FULL_IMAGE2(patch_shape, 10, 5, 32)

	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")

	imgs_mask_test = model.predict(X, batch_size=batch_size, verbose=1)

	# saving prediuction data
	np.save(get_prediction_path(), imgs_mask_test)


	eval_fit_prediction(X, imgs_mask_test)




def predict_fit_mc():
	batch_size = 1
	patch_shape = (48,240,240)
	model = model_3_encoders_1_decoder((*patch_shape,1))
	
	try:
		if os.path.exists(get_weights_path()):
			model.load_weights(get_weights_path())
			print('vahy nacteny OK')
	except:
		print("vahy nejdou nacist")

	# 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)






if __name__ == '__main__':
	# zpracovani parametru
	arguments = argparse.ArgumentParser(description="UNET")
	arguments = argparse.ArgumentParser(add_help=False)
	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("--fitgeneval",  "-fge", action="store_true", dest="fitgeneval")
	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.fitgeneval:
		p("EVAL PREDICT (for patches)")
		eval_predict()
	elif args.fit:
		predict_fit()
	elif args.fitgen:
		predict()
	elif args.fit_multi:
		predict_fit_mc()
	else:
		print("lze spustit s argumentem: --fit nebo --fitgen (-f|-fg)")

