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 *

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

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




class MY_Generator(keras.utils.Sequence):

	def __init__(self, batch_size = 120, patch_shape=(8,24,24), patch_stride=(8,24,24), data_X = '', data_Y = '', mode = 'fit', validation_data_split = 0.2, strict_batch_division = True):
		# batch_size = 120;  120 patches 
		self.mode = mode
		self.validation_data_split = validation_data_split
		self.patch_shape = (patch_shape[0], patch_shape[1], patch_shape[2])
		self.patch_stride = patch_stride
		self.batch_size = batch_size
		self.images_number = 4
		
		# pokud by mel patch shape jest 4 dimenzi, jedna se o kanal
		self.number_of_channels = patch_shape[3] if len(patch_shape) == 4 else 1

		self.shuffle = True if mode != 'predict' else False
		if self.number_of_channels > 1:
			self.X, self.Y  = self.prepare_data_mc(data_X, data_Y)
		else:
			self.X, self.Y  = self.prepare_data(data_X, data_Y)
		self.index = 0
		
		# check
		if strict_batch_division:
			if (self.images_number * self.number_of_patches) % self.batch_size != 0:
				sys.exit("celkový počet patchu neni deitelny davkou!")
		
		# generovani indexu
		self.indexes = np.arange(self.get_steps_per_epoch()).astype('uint16') 
		if self.shuffle == True:
			self.on_epoch_end()


	def __len__(self):
		return int(np.floor((self.images_number * self.number_of_patches) / self.batch_size))  # patches = 600


	def __getitem__(self, index):
		# random indexes
		index = self.indexes[index]
		################[num of batches as interval								 , ... other dims..]
		batch_x = self.X[index * self.batch_size : (index + 1) * self.batch_size]
		batch_y = self.Y[index * self.batch_size : (index + 1) * self.batch_size]
		if self.mode == 'predict':
			return batch_x
		else:
			return batch_x, batch_y


	def prepare_data(self, data_X, data_Y = []):
		# Initialization
		self.images_number = data_X.shape[0]

		"""
			validation split:
				rozdeli se pocet obrazku na pocet trenovaichc a na pocet validacnich a tomu se upravi i pocet obrazku: images_number, data_X a data_Y 
		"""
		############## VALIATION SPLIT BEGIN ##################################################
		validation_images_number = int(self.validation_data_split * self.images_number)
		train_images_number = self.images_number - validation_images_number

		if self.mode == 'fit' :
			self.images_number = train_images_number
			data_X = data_X[0:self.images_number]
			data_Y = data_Y[0:self.images_number]

		elif self.mode == 'val':
			self.images_number = validation_images_number
			data_X = data_X[train_images_number : train_images_number + self.images_number]
			data_Y = data_Y[train_images_number : train_images_number + self.images_number]
		############## VALIATION SPLIT ENDS ##################################################

		self.number_of_patches = self.compute_number_of_patches(data_X[0])

		# alokace numpy pole
		X = np.empty((self.images_number , self.number_of_patches, *self.patch_shape, self.number_of_channels), dtype = data_X.dtype)
		Y = np.empty((self.images_number , self.number_of_patches, tuple_multiply(self.patch_shape), 10), dtype = data_Y.dtype)

		for image_number in range(0, self.images_number):

			train_patches = self.extract_sliding_window_patches(data_X[image_number], self.patch_shape, self.patch_stride)
			mask_patches =  self.extract_sliding_window_patches(data_Y[image_number], self.patch_shape, self.patch_stride)
			train_patches = train_patches[..., np.newaxis]

			# Store sample
			X[image_number] = train_patches

			# mask
			mask_patches = mask_patches.reshape((self.number_of_patches, tuple_multiply(self.patch_shape)))
			Y[image_number] =  keras.utils.to_categorical(mask_patches, num_classes=10)

		XX = np.reshape(X, (self.images_number * self.number_of_patches, *self.patch_shape, self.number_of_channels))
		YY = np.reshape(Y, (self.images_number * self.number_of_patches,  tuple_multiply(self.patch_shape), 10))

		p('---- prepare data:')
		p(XX.shape)
		p(YY.shape)

		return  XX, YY



	def prepare_data_mc(self, data_X, data_Y = [], input_data_channel_position = 'last'):
		# Initialization
		self.images_number = data_X.shape[0]

		"""
			validation split:
				rozdeli se pocet obrazku na pocet trenovaichc a na pocet validacnich a tomu se upravi i pocet obrazku: images_number, data_X a data_Y 
		"""
		############## VALIATION SPLIT BEGIN ##################################################
		validation_images_number = int(self.validation_data_split * self.images_number)
		train_images_number = self.images_number - validation_images_number

		if self.mode == 'fit' :
			self.images_number = train_images_number
			data_X = data_X[0:self.images_number]
			data_Y = data_Y[0:self.images_number]

		elif self.mode == 'val':
			self.images_number = validation_images_number
			data_X = data_X[train_images_number : train_images_number + self.images_number]
			data_Y = data_Y[train_images_number : train_images_number + self.images_number]
		############## VALIATION SPLIT ENDS ##################################################

		if input_data_channel_position == 'last':
			data_X = np.moveaxis(data_X, 4, 1)		# channels first
			data_Y = np.moveaxis(data_Y, 4, 1)		# channels first

		self.number_of_patches = self.compute_number_of_patches(data_X[0][0])

		# alokace numpy pole
		X = np.empty((self.images_number , self.number_of_channels, self.number_of_patches, *self.patch_shape), dtype = data_X.dtype)
		Y = np.empty((self.images_number , self.number_of_channels, self.number_of_patches, *self.patch_shape), dtype = data_Y.dtype)

		for image_number in range(0, self.images_number):
			for channel in range(self.number_of_channels):
				print(data_X[image_number][channel].shape)
				train_patches = self.extract_sliding_window_patches(data_X[image_number][channel], self.patch_shape, self.patch_stride)
				mask_patches =  self.extract_sliding_window_patches(data_Y[image_number][channel], self.patch_shape, self.patch_stride)

				# Store sample
				X[image_number][channel] = train_patches

				# mask
				Y[image_number][channel] = mask_patches

		X = np.moveaxis(X, 1, 5)		# channels last
		Y = np.moveaxis(Y, 1, 5)		# channels last
		
		XX = np.reshape(X, (self.images_number * self.number_of_patches, *self.patch_shape, self.number_of_channels))
		YY = np.reshape(Y, (self.images_number * self.number_of_patches, *self.patch_shape, self.number_of_channels))

		flatten_patch_shape = self.patch_shape[0] * self.patch_shape[1] * self.patch_shape[2]	
		imgs_mask = np.reshape(YY, (self.images_number * self.number_of_patches, flatten_patch_shape, self.number_of_channels))
		imgs_mask = np.moveaxis(imgs_mask, 2, 0)
		imgs_mask = imgs_mask[0]
		YY = keras.utils.to_categorical(imgs_mask, num_classes=10)

		p('---- prepare data:')
		p(XX.shape)
		p(YY.shape)
		p(YY.max())
		p(YY.min())

		return  XX, YY


	def nextBatch(self):
		X, Y = self.__getitem__(self.index)
		self.index = self.index + 1

		if self.mode == 'predict':
			return X
		else:
			return X, Y


	def getitem(self, index):
		return self.__getitem__(index)


	def on_epoch_end(self):
		if self.shuffle == True:
			np.random.shuffle(self.indexes)	# Updates indexes after each epoch 


	def get_steps_per_epoch(self):
		return int(self.get_number_of_patches() * self.images_number / self.batch_size)


	def get_images_number(self):
		return self.images_number


	def get_number_of_patches(self):
		return self.number_of_patches


	def compute_number_of_patches(self, image_3d):
		train_patches = self.extract_sliding_window_patches(image_3d, self.patch_shape, self.patch_stride)
		self.number_of_patches = train_patches.shape[0]
		return self.number_of_patches


	def extract_sliding_window_patches(self, image, window_shape, stride_shape):
		sliding_window = view_as_windows(image , window_shape, stride_shape)
		self.sliding_window_shape = sliding_window.shape
		sliding_window = sliding_window.reshape(sliding_window.shape[0] * sliding_window.shape[1] * sliding_window.shape[2], *window_shape)
		return sliding_window

