from __future__ import print_function
import os
import sys
import numpy as np
from skimage.transform import resize
from skimage.io import imsave, imread
import SimpleITK as sitk
import scipy.spatial
from tools import *
import warnings
warnings.filterwarnings("ignore")



def prepare_png_dataset_from_nii(src_nii_file_path, dest_png_folder, normalization = True, histeq = False, augment_rotate = 0, zoom_factor = 0, bg_removal_with_mask_path =  True):
	if not os.path.exists(src_nii_file_path):
		print("Soubor: '"+src_nii_file_path+"' neexistuje")
		return False
	nii_image = get_3D_nii_image(src_nii_file_path)
	
	if augment_rotate != 0 and augment_rotate != False:
		nii_image = image_3d_rotate(nii_image, augment_rotate)	# rotace o 10 stupnu

	if zoom_factor != 0 and zoom_factor != False:
		nii_image = image_3d_clipped_zoom(nii_image, zoom_factor)	# rotace o 10 stupnu

	# v pripade snimku IR je mozne odstranit pozadi
	if bg_removal_with_mask_path is not False  and 'reg_IR.nii.gz' in (src_nii_file_path):
		if bg_removal_with_mask_path is True:
			bg_removal_with_mask_path = os.path.abspath(os.path.join(os.path.dirname(src_nii_file_path) + '/../segm.nii.gz'))
		p("_Editace IR: " + src_nii_file_path)
		p("_Maska IR: " + bg_removal_with_mask_path)
		nii_image_mask = get_3D_nii_image(bg_removal_with_mask_path)
		if augment_rotate != 0 and augment_rotate != False:
			nii_image_mask = image_3d_rotate(nii_image_mask, augment_rotate)	# rotace o 10 stupnu
		if zoom_factor != 0 and zoom_factor != False:
			nii_image_mask = image_3d_clipped_zoom(nii_image_mask, zoom_factor)	# rotace o 10 stupnu
		nii_image = remove_background_by_mask(nii_image, nii_image_mask)


	if not os.path.exists(dest_png_folder):
	    os.makedirs(dest_png_folder)

	if not os.path.exists(dest_png_folder):
		print("Cílová složka: '"+dest_png_folder+"' neexistuje")
		sys.exit()
	if histeq:
		save_3D_png_image_histeq(nii_image, dest_png_folder)
	else:
		save_3D_png_image(nii_image, dest_png_folder, normalization)


"""
	Transformuje nii data do slozek s png obrazkama, tj. napr.: z: /dataset/1/pre/FLAIR.nii.gz ---->  do: /data/train/1_FLAIR/(000.png - 047.png)
"""


def run_dataset(drive_cwd, dataset_path):
	#dataset_path = '/dataset/'
	images_train = ['1', '4', '5', '7', '14', '070']
	images_mask = ['1', '4', '5', '7', '14', '070']
	images_test = ['148']
	histeq = True

	#image_types = ['/pre/FLAIR.nii.gz', '/pre/IR.nii.gz', '/pre/reg_T1.nii.gz']
	image_types = ['/pre/FLAIR.nii.gz', '/pre/reg_IR.nii.gz', '/pre/reg_T1.nii.gz']
	image_png_folder_suffixes = ['_FLAIR', '_IR', '_T1']
	image_mask = '/segm.nii.gz'


	dest_data_test  = '/data/test/'
	dest_data_train = '/data/train/'
	dest_data_masks = '/data/masks/'

	# train
	for  img in images_train:
		index = 0
		for img_type in image_types:
			result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
			dest_path = os.path.join(drive_cwd + dest_data_train + img + image_png_folder_suffixes[index])
			print("Zpracovávám: " +result_path)
			if img_type == '/pre/FLAIR.nii.gz' and histeq:
				prepare_png_dataset_from_nii(result_path, dest_path, True, histeq)
			else:
				prepare_png_dataset_from_nii(result_path, dest_path)
			index = index + 1

	# test
	for  img in images_test:
		index = 0
		for img_type in image_types:
			result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
			dest_path = os.path.join(drive_cwd + dest_data_test + img + image_png_folder_suffixes[index])
			print("Zpracovávám: " +result_path)
			prepare_png_dataset_from_nii(result_path, dest_path)
			index = index + 1

	# mask
	for  img in images_mask:
		index = 0
		for img_type in image_types:
			result_path = os.path.join(drive_cwd + dataset_path + img + image_mask)
			dest_path = os.path.join(drive_cwd + dest_data_masks + img + image_png_folder_suffixes[index])
			print("Zpracovávám: " +result_path)
			prepare_png_dataset_from_nii(result_path, dest_path, False)
			index = index + 1


	############################### AUGMENTACE - ROTACE ###################################X
	print("--------------- zacatek augmentace rotace --------------------------")
	#dataset_path = '/dataset/'
	

	for rotation, aug_prefix in zip( [7, -7, 12, -12, 17, -17] ,  ['AUG_ROT_', 'AUG_ROT2_', 'AUG_ROT3_', 'AUG_ROT4_', 'AUG_ROT5_', 'AUG_ROT6_']):
		print("augmentace rotace: " + str(rotation))

		images_train = ['1', '4', '5', '7', '14', '070']
		images_mask = ['1', '4', '5', '7', '14', '070']
		images_test = ['148'] # ['070'] #

		histeq = True

		#image_types = ['/pre/FLAIR.nii.gz', '/pre/IR.nii.gz', '/pre/reg_T1.nii.gz']
		image_types = ['/pre/FLAIR.nii.gz', '/pre/reg_IR.nii.gz', '/pre/reg_T1.nii.gz']
		image_types_test = ['/pre/reg_T1.nii.gz']
		image_png_folder_suffixes = ['_FLAIR', '_IR', '_T1']
		image_mask = '/segm.nii.gz'


		dest_data_test  = '/data/test/'
		dest_data_train = '/data/train/'
		dest_data_masks = '/data/masks/'

		# train
		for  img in images_train:

			index = 0
			for img_type in image_types:
				result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
				dest_path = os.path.join(drive_cwd + dest_data_train + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				if img_type == '/pre/FLAIR.nii.gz' and histeq:
					prepare_png_dataset_from_nii(result_path, dest_path, True, histeq, rotation)
				else:
					prepare_png_dataset_from_nii(result_path, dest_path, True, False, rotation)
				index = index + 1

		# test
		for  img in images_train:
			index = 0
			for img_type in image_types_test:
				result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
				dest_path = os.path.join(drive_cwd + dest_data_test + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				prepare_png_dataset_from_nii(result_path, dest_path, True, False, rotation)
				index = index + 1

		# mask
		for  img in images_mask:
			index = 0
			for img_type in image_types:
				result_path = os.path.join(drive_cwd + dataset_path + img + image_mask)
				dest_path = os.path.join(drive_cwd + dest_data_masks + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				prepare_png_dataset_from_nii(result_path, dest_path, False, False, rotation)
				index = index + 1


	############################### AUGMENTACE - zvetseni - ZOOM ###################################X
	print("--------------- zacatek augmentace zvetseni - ZOOM --------------------------")
	#dataset_path = '/dataset/'
	

	for zoom_factor, aug_prefix in zip( [1.05, 1.1, 1.17, 1.21] ,  ['AUG_ZOOM_', 'AUG_ZOOM2_', 'AUG_ZOOM3_', 'AUG_ZOOM4_']):
		print("augmentace zvetseni: " + str(zoom_factor))

		images_train = ['1', '4', '5', '7', '14', '070']
		images_mask = ['1', '4', '5', '7', '14', '070']
		images_test =  ['148']

		histeq = True

		#image_types = ['/pre/FLAIR.nii.gz', '/pre/IR.nii.gz', '/pre/reg_T1.nii.gz']
		image_types = ['/pre/FLAIR.nii.gz', '/pre/reg_IR.nii.gz', '/pre/reg_T1.nii.gz']
		image_types_test = [ '/pre/reg_T1.nii.gz']
		image_png_folder_suffixes = ['_FLAIR', '_IR', '_T1']
		image_mask = '/segm.nii.gz'


		dest_data_test  = '/data/test/'
		dest_data_train = '/data/train/'
		dest_data_masks = '/data/masks/'

		# train
		for  img in images_train:
			index = 0
			for img_type in image_types:
				result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
				dest_path = os.path.join(drive_cwd + dest_data_train + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				if img_type == '/pre/FLAIR.nii.gz' and histeq:
					prepare_png_dataset_from_nii(result_path, dest_path, True, histeq, 0, zoom_factor)
				else:
					prepare_png_dataset_from_nii(result_path, dest_path, True, False, 0, zoom_factor)
				index = index + 1

		# test
		for  img in images_test:
			index = 0
			for img_type in image_types_test:
				result_path = os.path.join(drive_cwd + dataset_path + img + img_type)
				dest_path = os.path.join(drive_cwd + dest_data_test + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				prepare_png_dataset_from_nii(result_path, dest_path, True, False, 0, zoom_factor)
				index = index + 1

		# mask
		for  img in images_mask:
			index = 0
			for img_type in image_types:
				result_path = os.path.join(drive_cwd + dataset_path + img + image_mask)
				dest_path = os.path.join(drive_cwd + dest_data_masks + aug_prefix + str(img) + image_png_folder_suffixes[index])
				print("Zpracovávám: " +result_path)
				prepare_png_dataset_from_nii(result_path, dest_path, False, False, 0, zoom_factor)
				index = index + 1


if __name__ == '__main__':
	run_dataset(os.getcwd(), '/dataset/')
