from __future__ import print_function

import os
import sys
import numpy as np
from skimage.transform import resize
from skimage.io import imsave, imread
from skimage import data
from tools import *
import argparse

data_path = 'data/'
train_folder = 'train'
test_folder = 'test'
mask_folder = 'masks'

# velikosti obrazku
image_rows = 240
image_cols = 240
img_depth = 48
img_channels = 1

# velikosti obrazku(rezu) po zmenseni
IMG_HEIGHT = 240
IMG_WIDTH = 240

IMG_COUNT  = 1
# allowe image format
IMG_FORMAT = ".png"


labels = {
		  0: 'Background',
		  1: 'Cortical gray matter',
          2: 'Basal ganglia',
          3: 'White matter',
          4: 'White matter lesions',
          5: 'Cerebrospinal fluid in the extracerebral space',
          6: 'Ventricles',
          7: 'Cerebellum',
          8: 'Brain stem',
          # The two labels below are ignored:
          9: 'Infarction',
          #10: 'Other',
          }


def get_number_of_classes():
	return int(len(labels.keys()))


def create_train_data():

	train_data_path = os.path.join(data_path, train_folder)
	mask_data_path = os.path.join(data_path, mask_folder)
	train_data_dirs = os.listdir(train_data_path)

	#zde mam pocet polozek v prvni slozce, zbyle slozkymusi mit stejny pocet
	train_data_first_folder_path = os.path.join(train_data_path, train_data_dirs[0])

	# pocet slozek s trenovacimi daty
	dirs = int(len(train_data_dirs))

	#celkovy pocet obrazku
	total = dirs *  int(len(os.listdir(train_data_first_folder_path)))

	# hloubka = pocet snimku 3D obrazu, rows & cols = rastrova mrizka
	imgs = np.ndarray((dirs, img_depth, image_rows, image_cols), dtype=np.uint8)
	imgs_mask = np.ndarray((dirs, img_depth, image_rows, image_cols), dtype=np.uint8)

	print('-'*30)
	print('Creating training images...')
	print('-'*30)

	# osetreni
	if os.listdir(train_data_path)!=os.listdir(mask_data_path):
		p(os.listdir(mask_data_path))
		p(os.listdir(train_data_path))
		sys.exit("Slozky ve slozce train a masks se musi shodovat!")


	# klice trid
	classes = list()
	for i in labels.keys():
	    classes.append(i)


	# nacitam podadresars a soubory v rootu
	jt=0
	kt=0
	for root, subdirs, files in os.walk(train_data_path):
		# prochazim podadresare
		j = 0
		k = 0
		subdirs.sort()
		for subdir in subdirs:
			# nastaveni cest na slozku train a masks
			# je zde predpoklad, ze slozky masek a trenovacich dat jsou shodne
			tmp_train_data_path = os.path.join(train_data_path, subdir)
			tmp_mask_data_path = os.path.join(mask_data_path, subdir)
			
			# jelikoz obrazky trenovacich dat a masek musi mit stejny nazev, tak je nactem jen z trenovaci slozky
			images_train = sorted(os.listdir(tmp_train_data_path))
			images_mask = sorted(os.listdir(tmp_mask_data_path))

			if images_train!=images_mask:
				p(images_train)
				p(images_mask)
				sys.exit("obrazky spatne serazeny")

			# prochazim obrazky a budeme je zmensovat a ukladat do pozadoveneho souboru
			img_index_j=0 #   index
			for img_name_train in images_train:
				#musim overit, zdali ma slubor spravny format tj. png
				if IMG_FORMAT not in img_name_train.lower():
					print("Obrázek "+img_name_train+" neni podporován")
				# cesta k trenovacim datum
				tmp_train_data_path_image = os.path.join(tmp_train_data_path, img_name_train)
				# trenovaci data
				img =  imread(tmp_train_data_path_image, as_gray=True ) # img.shape
				#img = img.astype(np.uint8)
				
				# zmenseni na pozadovanou velikost
				#img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
				# nacteni do pole
				img = np.array([img])
				# ulozeni do vystupniho pole
				imgs[j][img_index_j] = img;

				if img_index_j % 10 == 0:
					print('Done: {0}/{1} images_train'.format(jt, total))
				img_index_j += 1
				jt+= 1

			
			# zde se zpracovavaji masky, muzou mit ruzna jmena
			img_index_k=0 #   index
			for img_name_mask in images_mask:
				if IMG_FORMAT not in img_name_mask.lower():
					print("Obrázek "+img_name_mask+" neni podporován")
				tmp_mask_data_path_image = os.path.join(tmp_mask_data_path, img_name_mask)
				# masky
				img_mask = imread(tmp_mask_data_path_image, as_gray=True)
				#img_mask = resize(img_mask, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
				# nacteni do pole
				
				# spojeni masky a obrazku
				#img_mask = binarylabel(image_rows, image_cols, img_mask, get_number_of_classes())

				#mask = np.zeros((IMG_HEIGHT, IMG_WIDTH, 5), dtype=np.uint8)
				#img_mask = np.expand_dims(img_mask, axis=3)
				#img_mask[:,:,0] = classes
				img_mask = np.array([img_mask]) # img_mask.shape
				imgs_mask[k][img_index_k] = img_mask;

				if img_index_k % 10 == 0:
					print('Done: {0}/{1} images_mask'.format(kt, total))
				img_index_k += 1
				kt+= 1

			j+=1
			k+=1
	

	#p(imgs.shape)
	#p(imgs_mask.shape)
	#p(len(images_mask))

	#imgs_mask = np.reshape(imgs_mask, (dirs, len(images_mask) * image_rows * image_cols, get_number_of_classes()))

	print ('------------------')
	print('Done: {0}/{1} train images! '.format(total, total))
	
	# shape: (6, 48, 240, 240, 10)
	print('Done: {0}/{1} mask images! '.format(total, total))
	print('Loading done.')

	print("train images shape: " + str(imgs.shape))
	print("mask images shape: " + str(imgs_mask.shape))

	np.save('imgs_train.npy', imgs)
	np.save('imgs_mask_train.npy', imgs_mask)
	print('Saving to .npy files done.')
	return True

def create_test_data():
	
	test_data_path = os.path.join(data_path, test_folder)
	test_data_dirs = os.listdir(test_data_path)

	number_of_dirs = int(len(test_data_dirs))
	imgs = np.ndarray((number_of_dirs, img_depth, image_rows, image_cols), dtype=np.uint8)

	print('-'*30)
	print('Creating testing images...')
	print('-'*30)

	for root, subdirs, files in os.walk(test_data_path):
		# prochazim podadresare
		j = 0
		subdirs.sort()
		for subdir in subdirs:
			print("Zpracovavam slozku: " + str(subdir))
			image_path = os.path.join(test_data_path , subdir)
			imgs[j] = get_3D_png_image(image_path)
			j+=1

	np.save('imgs_test.npy', imgs)
	print('Saving to .npy files done.')
	return True


def load_train_data():
	imgs_train = np.load('imgs_train.npy')
	imgs_mask_train = np.load('imgs_mask_train.npy')
	return imgs_train, imgs_mask_train


def load_test_data():
	imgs_test = np.load('imgs_test.npy')
	return imgs_test


def prepare_and_load_train_data():
	imgs_train, imgs_mask_train = load_train_data()
	imgs_train = imgs_train.astype('float32')
	imgs_train /= 255

	mean = np.mean(imgs_train)  # mean for data centering
	std = np.std(imgs_train)  # std for data normalization
	#imgs_train -= mean
	#imgs_train /= std

	print("max train val: "  + str(imgs_train.max()))
	print("min train val: "  + str(imgs_train.min()))
	return imgs_train, imgs_mask_train


def prepare_and_load_test_data():
	imgs_test = load_test_data()
	imgs_test = imgs_test.astype('float32')
	imgs_test /= 255
	return imgs_test


## MULTICHANNEL ##

def create_train_data_multichannel(masks_one_channel_only = False):
	train_data_path = os.path.join(data_path, train_folder)
	#channels_prefixes = ['1_', '4_',  '5_', '7_', '14_', '070_', 'AUG_ROT_1_', 'AUG_ROT_4_', 'AUG_ROT_5_', 'AUG_ROT_7_', 'AUG_ROT_14_', 'AUG_ROT_070_']
	channels_prefixes = [
		'1_', '4_',  '5_', '7_', '14_', '070_', 
		'AUG_ROT_1', 'AUG_ROT_4', 'AUG_ROT_5', 'AUG_ROT_7', 'AUG_ROT_14', 'AUG_ROT_070', 
		'AUG_ROT2_1', 'AUG_ROT2_4', 'AUG_ROT2_5', 'AUG_ROT2_7', 'AUG_ROT2_14', 'AUG_ROT2_070', 
		'AUG_ROT3_1', 'AUG_ROT3_4', 'AUG_ROT3_5', 'AUG_ROT3_7', 'AUG_ROT3_14', 'AUG_ROT3_070', 
		'AUG_ROT4_1', 'AUG_ROT4_4', 'AUG_ROT4_5', 'AUG_ROT4_7', 'AUG_ROT4_14', 'AUG_ROT4_070', 
		'AUG_ROT5_1', 'AUG_ROT5_4', 'AUG_ROT5_5', 'AUG_ROT5_7', 'AUG_ROT5_14', 'AUG_ROT5_070', 
		'AUG_ROT6_1', 'AUG_ROT6_4', 'AUG_ROT6_5', 'AUG_ROT6_7', 'AUG_ROT6_14', 'AUG_ROT6_070',

		'AUG_ZOOM_1', 'AUG_ZOOM_4', 'AUG_ZOOM_5', 'AUG_ZOOM_7', 'AUG_ZOOM_14', 'AUG_ZOOM_070',
		'AUG_ZOOM2_1', 'AUG_ZOOM2_4', 'AUG_ZOOM2_5', 'AUG_ZOOM2_7', 'AUG_ZOOM2_14', 'AUG_ZOOM2_070',
		'AUG_ZOOM3_1', 'AUG_ZOOM3_4', 'AUG_ZOOM3_5', 'AUG_ZOOM3_7', 'AUG_ZOOM3_14', 'AUG_ZOOM3_070',
		'AUG_ZOOM4_1', 'AUG_ZOOM4_4', 'AUG_ZOOM4_5', 'AUG_ZOOM4_7', 'AUG_ZOOM4_14', 'AUG_ZOOM4_070',
	]
	images = process_all_3d_png_in_dir_with_channel(train_data_path, channels_prefixes)
	images = np.moveaxis(images, 1, 4)		# channels last
	np.save('imgs_train_mc.npy', images)
	p(images.shape)
	p("DONE: Trénovací sada vytvořena.")
	p("")

	mask_data_path = os.path.join(data_path, mask_folder)
	#channels_prefixes = ['1_', '4_',  '5_', '7_', '14_', '070_', 'AUG_ROT_1_', 'AUG_ROT_4_', 'AUG_ROT_5_', 'AUG_ROT_7_', 'AUG_ROT_14_', 'AUG_ROT_070_']
	channels_prefixes = [
		'1_', '4_',  '5_', '7_', '14_', '070_', 
		'AUG_ROT_1', 'AUG_ROT_4', 'AUG_ROT_5', 'AUG_ROT_7', 'AUG_ROT_14', 'AUG_ROT_070', 
		'AUG_ROT2_1', 'AUG_ROT2_4', 'AUG_ROT2_5', 'AUG_ROT2_7', 'AUG_ROT2_14', 'AUG_ROT2_070', 
		'AUG_ROT3_1', 'AUG_ROT3_4', 'AUG_ROT3_5', 'AUG_ROT3_7', 'AUG_ROT3_14', 'AUG_ROT3_070', 
		'AUG_ROT4_1', 'AUG_ROT4_4', 'AUG_ROT4_5', 'AUG_ROT4_7', 'AUG_ROT4_14', 'AUG_ROT4_070', 
		'AUG_ROT5_1', 'AUG_ROT5_4', 'AUG_ROT5_5', 'AUG_ROT5_7', 'AUG_ROT5_14', 'AUG_ROT5_070', 
		'AUG_ROT6_1', 'AUG_ROT6_4', 'AUG_ROT6_5', 'AUG_ROT6_7', 'AUG_ROT6_14', 'AUG_ROT6_070',

		'AUG_ZOOM_1', 'AUG_ZOOM_4', 'AUG_ZOOM_5', 'AUG_ZOOM_7', 'AUG_ZOOM_14', 'AUG_ZOOM_070',
		'AUG_ZOOM2_1', 'AUG_ZOOM2_4', 'AUG_ZOOM2_5', 'AUG_ZOOM2_7', 'AUG_ZOOM2_14', 'AUG_ZOOM2_070',
		'AUG_ZOOM3_1', 'AUG_ZOOM3_4', 'AUG_ZOOM3_5', 'AUG_ZOOM3_7', 'AUG_ZOOM3_14', 'AUG_ZOOM3_070',
		'AUG_ZOOM4_1', 'AUG_ZOOM4_4', 'AUG_ZOOM4_5', 'AUG_ZOOM4_7', 'AUG_ZOOM4_14', 'AUG_ZOOM4_070',
	]



	masks = process_all_3d_png_in_dir_with_channel(mask_data_path, channels_prefixes, [], 3, (1 if masks_one_channel_only else 3))
	masks = np.moveaxis(masks, 1, 4)		# channels last
	# pokud bude channel jen jeden, tak ho oddelam
	if masks_one_channel_only:
		if masks.shape[4] == 1:
			masks = masks.squeeze(axis=4)
		np.save('imgs_mask_train_mc.npy', masks)
	else:
		np.save('imgs_mask_train_mc.npy', masks)
	p(masks.shape)
	p("DONE: Sada masek vytvořena.")
	p("")
	return images


def prepare_and_load_train_data_multichannel():
	if not os.path.exists('imgs_train_mc.npy') or not os.path.exists('imgs_mask_train_mc.npy') :
		d("ERRO: data nebyla ještě vytvořena")

	imgs_train_mc = np.load('imgs_train_mc.npy')
	imgs_mask_train_mc = np.load('imgs_mask_train_mc.npy')

	imgs_train_mc = imgs_train_mc.astype('float32')
	imgs_train_mc /= 255

	mean = np.mean(imgs_train_mc)  # mean for data centering
	std = np.std(imgs_train_mc)  # std for data normalization
	#imgs_train_mc -= mean
	#imgs_train_mc /= std

	print("max train val: "  + str(imgs_train_mc.max()))
	print("min train val: "  + str(imgs_train_mc.min()))
	return imgs_train_mc, imgs_mask_train_mc


def create_test_data_multichannel():
	test_data_path = os.path.join(data_path, test_folder)
	channels_prefixes = ['148_', 'AUG_ROT_148_']
	images = process_all_3d_png_in_dir_with_channel(test_data_path, channels_prefixes)
	images = np.moveaxis(images, 1, 4)		# channels last
	np.save('imgs_test_mc.npy', images)
	p(images.shape)
	p("DONE: Testovací sada vytvořena.")
	p("")

	return images


def prepare_and_load_test_data_multichannel():
	if not os.path.exists('imgs_test_mc.npy'):
		d("ERRO: data nebyla ještě vytvořena")

	imgs_test_mc = np.load('imgs_test_mc.npy')
	imgs_test_mc = imgs_test_mc.astype('float32')
	imgs_test_mc /= 255
	mean = np.mean(imgs_test_mc)  # mean for data centering
	std = np.std(imgs_test_mc)  # std for data normalization
	#imgs_test_mc -= mean
	#imgs_test_mc /= std
	return imgs_test_mc




"""
	Transformuje data ze slozek: /data/train/1_FLAIR/(000.png - 047.png) do numpy matic .npy
"""
if __name__ == '__main__':


	# zpracovani parametru
	arguments = argparse.ArgumentParser(description="UNET")
	arguments = argparse.ArgumentParser(add_help=False)
	arguments.add_argument("--create", "-c", action="store_true", dest="create_data")
	arguments.add_argument("--load", "-l", action="store_true", dest="load_data")
	arguments.add_argument("--multi", "-m", action="store_true", dest="multichannel")
	args = arguments.parse_args()## naparsovani argumentu



	if args.create_data:
		if args.multichannel:
			create_train_data_multichannel()
			create_test_data_multichannel()
		else:
			create_train_data()
			create_test_data()

	elif args.load_data:
		if args.multichannel:
			X,Y =prepare_and_load_train_data_multichannel()
			XT = prepare_and_load_test_data_multichannel()
			p('Trénovací data: ' + str(X.shape) + "| Masky: "  + str(Y.shape))
			p('Testovací data: ' + str(XT.shape))
		else:
			X,Y = prepare_and_load_train_data()
			XT = prepare_and_load_test_data()
			p('Trénovací data: ' + str(X.shape) + "| Masky: "  + str(Y.shape))
			p('Testovací data: ' + str(XT.shape))
	else:
		print("Lze pouzit parametry: --create | --load | --multi")

