import numpy as np
import sys
import os
import keras
import SimpleITK as sitk
import tensorflow as tf
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 keras.models import *
from keras.layers import *
from tools import tuple_multiply

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"



def get_unet(img_depth, img_rows, img_cols, num_of_classes = 10):
	inputs = Input((img_depth, img_rows, img_cols, 1))
	x = 0.5
	conv1 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(inputs)
	conv1 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(conv1)
	pool1 = MaxPooling3D(pool_size=(2, 2, 2))(conv1)

	conv2 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(pool1)
	conv2 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(conv2)
	pool2 = MaxPooling3D(pool_size=(2, 2, 2))(conv2)

	conv3 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(pool2)
	conv3 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(conv3)
	pool3 = MaxPooling3D(pool_size=(2, 2, 2))(conv3)

	conv4 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(pool3)
	conv4 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(conv4)
	pool4 = MaxPooling3D(pool_size=(2, 2, 2))(conv4)

	conv5 = Conv3D(int(x*512 ), (2,2,2), activation='relu', padding='same')(pool4)
	conv5 = Conv3D(int(x*512 ), (2,2,2), activation='relu', padding='same')(conv5)

	print(conv4.shape)
	print(conv5.shape)

	up6 = concatenate([Conv3DTranspose( int(x*256  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv5), conv4], axis=4)
	conv6 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(up6)
	conv6 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(conv6)

	up7 = concatenate([Conv3DTranspose( int(x*128  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv6), conv3], axis=4)
	conv7 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(up7)
	conv7 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(conv7)

	up8 = concatenate([Conv3DTranspose( int(x*64  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv7), conv2], axis=4)
	conv8 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(up8)
	conv8 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(conv8)

	up9 = concatenate([Conv3DTranspose( int(x*32  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv8), conv1], axis=4)
	conv9 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(up9)
	conv9 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(conv9)
	conv9 = Conv3D(num_of_classes,  (2,2,2), activation='relu', padding='same')(conv9)									

	conv10 = Reshape((num_of_classes, img_depth * img_rows * img_cols), input_shape = (num_of_classes, img_depth, img_rows, img_cols))(conv9)
	permute = Permute((2,1))(conv10)
	outputs = Activation('softmax')(permute)


	model = Model(inputs=[inputs], outputs=[outputs])
	model.summary()
	model.compile(optimizer=Adam(lr=1e-5), loss='categorical_crossentropy', metrics=['accuracy'])
	return model


def create_unet_model3D(input_image_size,
						n_labels=10,
						layers=3,
						lowest_resolution=16,
						convolution_kernel_size=(3,3,3),
						deconvolution_kernel_size=(2,2,2),
						pool_size=(2,2,2),
						strides=(2,2,2),
						init_lr=0.0001,
						IMAGE_ORDERING =  "channels_last"
						):
	"""
	Create a 3D Unet model
	Example
	-------
	unet_model = create_unet_model3D( (128,128,128,1), 1, 4)
	"""
	layers = np.arange(layers)
	number_of_classification_labels = n_labels
	
	inputs = Input(shape=input_image_size)

	## ENCODING PATH ##

	encoding_convolution_layers = []
	pool = None
	for i in range(len(layers)):
		number_of_filters = lowest_resolution * 2**(layers[i])

		if i == 0:
			conv = Conv3D(filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(inputs)
		else:
			conv = Conv3D(filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(pool)

		encoding_convolution_layers.append(Conv3D(filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(conv))

		if i < len(layers)-1:
			pool = MaxPooling3D(pool_size=pool_size, data_format=IMAGE_ORDERING)(encoding_convolution_layers[i])

	## DECODING PATH ##
	outputs = encoding_convolution_layers[len(layers)-1]
	for i in range(1,len(layers)):
		number_of_filters = lowest_resolution * 2**(len(layers)-layers[i]-1)
		tmp_deconv = Conv3DTranspose(filters=number_of_filters, kernel_size=deconvolution_kernel_size, padding='same', data_format=IMAGE_ORDERING)(outputs)
		tmp_deconv = UpSampling3D(size=pool_size, data_format=IMAGE_ORDERING)(tmp_deconv)
		outputs = Concatenate(axis=4)([tmp_deconv, encoding_convolution_layers[len(layers)-i-1]])
		outputs = Conv3D(filters=number_of_filters, kernel_size=convolution_kernel_size, activation='relu', padding='same', data_format=IMAGE_ORDERING)(outputs)
		outputs = Conv3D(filters=number_of_filters, kernel_size=convolution_kernel_size, activation='relu', padding='same', data_format=IMAGE_ORDERING)(outputs)

	outputs = Conv3D(filters=number_of_classification_labels, kernel_size=(1,1,1), activation='softmax', data_format=IMAGE_ORDERING)(outputs)
	outputs = Reshape((number_of_classification_labels, input_image_size[0]*input_image_size[1]*input_image_size[2]), input_shape = input_image_size)(outputs)
	outputs = Permute((2,1))(outputs)

	unet_model = Model(inputs=inputs, outputs=outputs)
	unet_model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=init_lr), metrics=['accuracy'])

	return unet_model



def create_unet_model3D_FULL_IMAGE(input_image_size,
						n_labels=10,
						layers=4,
						lowest_resolution=32,
						convolution_kernel_size=(4,4,4),
						deconvolution_kernel_size=(4,4,4),
						dilate_rate=(1,1,1),
						pool_size=(2,2,2),
						strides=(1,1,1),
						init_lr=0.00001,
						IMAGE_ORDERING =  "channels_last"
	):

	layers = np.arange(layers)
	number_of_classification_labels = n_labels
	
	inputs = Input(shape=input_image_size)
	encoding_convolution_layers = []
	pool = None
	for i in range(len(layers)):
		number_of_filters = lowest_resolution * 2**(layers[i])

		if i == 0:
			conv = Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(inputs)
		else:
			conv = Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(pool)

		#conv = Dropout(0.2)(conv) # drop 0.5
		conv = BatchNormalization()(conv)
		encoding_convolution_layers.append(Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=convolution_kernel_size,activation='relu',padding='same', data_format=IMAGE_ORDERING)(conv))

		if i < len(layers)-1:
			pool = MaxPooling3D(pool_size=pool_size, data_format=IMAGE_ORDERING)(encoding_convolution_layers[i])

	outputs = encoding_convolution_layers[len(layers)-1]
	for i in range(1,len(layers)):
		number_of_filters = lowest_resolution * 2**(len(layers)-layers[i]-1)
		tmp_deconv = Conv3DTranspose(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=deconvolution_kernel_size, padding='same', data_format=IMAGE_ORDERING)(outputs)
		tmp_deconv = UpSampling3D(size=pool_size, data_format=IMAGE_ORDERING)(tmp_deconv)
		outputs = Concatenate(axis=4)([tmp_deconv, encoding_convolution_layers[len(layers)-i-1]])
		outputs = BatchNormalization()(outputs)
		outputs = Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=convolution_kernel_size, activation='relu', padding='same', data_format=IMAGE_ORDERING)(outputs)
		#outputs = Dropout(0.2)(outputs) # drop 0.5
		outputs = Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=convolution_kernel_size, activation='relu', padding='same', data_format=IMAGE_ORDERING)(outputs)

	#outputs = BatchNormalization()(outputs)
	#outputs = Dropout(0.2)(outputs) # drop 0.5
	outputs = Conv3D(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_classification_labels, kernel_size=(1,1,1), activation='softmax', data_format=IMAGE_ORDERING)(outputs)
	outputs = Reshape((number_of_classification_labels, input_image_size[0] * input_image_size[1] * input_image_size[2] ), input_shape = input_image_size)(outputs)
	outputs = Permute((2,1))(outputs)

	unet_model = Model(inputs=inputs, outputs=outputs)
	unet_model.summary()
	unet_model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=init_lr), metrics=['accuracy', 'categorical_crossentropy'])
	return unet_model




def create_unet_model3D2(input_image_size,
						n_labels=10,
						layers=4,
						lowest_resolution=32,
						convolution_kernel_size=(4,4,4),
						deconvolution_kernel_size=(4,4,4),
						dilate_rate=(1,1,1),
						pool_size=(2,2,2),
						strides=(1,1,1),
						init_lr=0.00001,
						IMAGE_ORDERING =  "channels_last"
	):

	layers = np.arange(layers)
	number_of_classification_labels = n_labels
	
	inputs = Input(shape=input_image_size)

	## ENCODING PATH ##
	encoding_convolution_layers = []
	pool = None
	for i in range(len(layers)):
		number_of_filters = lowest_resolution * 2**(layers[i])
		if i == 0:
			conv = Conv3D_block(inputs, number_of_filters, kernel=convolution_kernel_size, activation='relu')
		else:
			conv = Conv3D_block(pool, number_of_filters, kernel=convolution_kernel_size, activation='relu')
		conv = Conv3D_block(conv, number_of_filters, kernel=convolution_kernel_size, activation='relu')
		encoding_convolution_layers.append(conv)
		if i < len(layers)-1:
			pool = MaxPooling3D(pool_size=pool_size, data_format=IMAGE_ORDERING)(encoding_convolution_layers[i])


	## DECODING PATH ##
	outputs = encoding_convolution_layers[len(layers)-1]
	for i in range(1,len(layers)):
		number_of_filters = lowest_resolution * 2**(len(layers)-layers[i]-1)
		tmp_deconv = Conv3DTranspose(kernel_initializer = 'he_normal', dilation_rate=dilate_rate, filters=number_of_filters, kernel_size=deconvolution_kernel_size, padding='same', data_format=IMAGE_ORDERING)(outputs)
		tmp_deconv = UpSampling3D(size=pool_size, data_format=IMAGE_ORDERING)(tmp_deconv)
		outputs = Concatenate(axis=4)([tmp_deconv, encoding_convolution_layers[len(layers)-i-1]])
		outputs = Conv3D_block(outputs, number_of_filters, kernel=convolution_kernel_size, activation='relu')
		outputs = Conv3D_block(outputs, number_of_filters, kernel=convolution_kernel_size, activation='relu')

	outputs = Conv3D_block(outputs, number_of_classification_labels, kernel=(1,1,1), activation='softmax')
	outputs = Reshape((number_of_classification_labels, input_image_size[0] * input_image_size[1] * input_image_size[2] ), input_shape = input_image_size)(outputs)
	outputs = Permute((2,1))(outputs)

	unet_model = Model(inputs=inputs, outputs=outputs)
	unet_model.summary()
	unet_model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=init_lr), metrics=['accuracy', 'categorical_crossentropy'])

	return unet_model



def Conv3D_block(input_layer, n_filters, kernel=(3, 3, 3), strides=(1, 1, 1), padding='same', activation=None,  batch_normalization=True, kernel_initializer = 'he_normal', data_format="channels_last"):
    layer = Conv3D(n_filters, kernel, strides=strides, padding=padding, kernel_initializer=kernel_initializer, data_format=data_format)(input_layer)
    if batch_normalization:
        layer = BatchNormalization()(layer)
    
    if activation is None:
        return Activation('relu')(layer)
    else:
        return Activation(activation)(layer)





def get_unet(img_depth, img_rows, img_cols, num_of_classes = 10):
	inputs = Input((img_depth, img_rows, img_cols, 1))
	x = 0.5
	conv1 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(inputs)
	conv1 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(conv1)
	pool1 = MaxPooling3D(pool_size=(2, 2, 2))(conv1)

	conv2 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(pool1)
	conv2 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(conv2)
	pool2 = MaxPooling3D(pool_size=(2, 2, 2))(conv2)

	conv3 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(pool2)
	conv3 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(conv3)
	pool3 = MaxPooling3D(pool_size=(2, 2, 2))(conv3)

	conv4 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(pool3)
	conv4 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(conv4)
	pool4 = MaxPooling3D(pool_size=(2, 2, 2))(conv4)

	conv5 = Conv3D(int(x*512 ), (2,2,2), activation='relu', padding='same')(pool4)
	conv5 = Conv3D(int(x*512 ), (2,2,2), activation='relu', padding='same')(conv5)

	print(conv4.shape)
	print(conv5.shape)
	up6 = concatenate([Conv3DTranspose( int(x*256  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv5), conv4], axis=4)
	conv6 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(up6)
	conv6 = Conv3D(int(x*256 ), (2,2,2), activation='relu', padding='same')(conv6)

	up7 = concatenate([Conv3DTranspose( int(x*128  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv6), conv3], axis=4)
	conv7 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(up7)
	conv7 = Conv3D(int(x*128 ), (2,2,2), activation='relu', padding='same')(conv7)

	up8 = concatenate([Conv3DTranspose( int(x*64  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv7), conv2], axis=4)
	conv8 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(up8)
	conv8 = Conv3D(int(x*64 ), (2,2,2), activation='relu', padding='same')(conv8)

	up9 = concatenate([Conv3DTranspose( int(x*32  ), (2, 2, 2), strides=(2, 2, 2), padding='same')(conv8), conv1], axis=4)
	conv9 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(up9)
	conv9 = Conv3D(int(x*32 ), (2,2,2), activation='relu', padding='same')(conv9)
	conv9 = Conv3D(num_of_classes,  (2,2,2), activation='relu', padding='same')(conv9)									

	conv10 = Reshape((num_of_classes, img_depth * img_rows * img_cols), input_shape = (num_of_classes, img_depth, img_rows, img_cols))(conv9)
	
	print(inputs.shape)
	print(conv10.shape)
	permute = Permute((2,1))(conv10)
	outputs = Activation('softmax')(permute)


	model = Model(inputs=[inputs], outputs=[outputs])
	model.summary()
	model.compile(optimizer=Adam(lr=1e-5), loss='categorical_crossentropy', metrics=['accuracy'])

	return model


def get_unet_edited(input_shape, num_of_classes =10, kernel_size = (3, 3, 3), pool_stride = (2, 2, 2)):
	x = 1
	def f(filters):
		return int(filters * x)

	inputs = Input((*input_shape, 1))
	#conv1 = Conv3D( f(32), (3, 3, 3), activation='relu', padding='same')(inputs)
	#conv1 = Conv3D( f(32), (3, 3, 3), activation='relu', padding='same')(conv1)
	#pool1 = MaxPooling3D(pool_size=(2, 2, 2))(conv1)
	#print ("pool1 shape:",pool1.shape)

	conv2 = Conv3D( f(64), (3, 3, 3), activation='relu', padding='same')(inputs)
	conv2 = Conv3D( f(64), (3, 3, 3), activation='relu', padding='same')(conv2)
	pool2 = MaxPooling3D(pool_size=pool_stride)(conv2)

	conv3 = Conv3D( f(128), (3, 3, 3), activation='relu', padding='same')(pool2)
	conv3 = Conv3D( f(128), (3, 3, 3), activation='relu', padding='same')(conv3)
	pool3 = MaxPooling3D(pool_size=pool_stride)(conv3)

	conv4 = Conv3D( f(256), (3, 3, 3), activation='relu', padding='same')(pool3)
	conv4 = Conv3D( f(256), (3, 3, 3), activation='relu', padding='same')(conv4)
	pool4 = MaxPooling3D(pool_size=pool_stride)(conv4)


	conv5 = Conv3D( f(512), (3, 3, 3), activation='relu', padding='same')(pool4)
	conv5 = Conv3D( f(512), (3, 3, 3), activation='relu', padding='same')(conv5)
	up6 = concatenate([Conv3DTranspose( f(256), (2, 2, 2), strides=pool_stride, padding='same')(conv5), conv4], axis=4)
	conv6 = Conv3D( f(256), (3, 3, 3), activation='relu', padding='same')(up6)
	conv6 = Conv3D( f(256), (3, 3, 3), activation='relu', padding='same')(conv6)


	up7 = concatenate([Conv3DTranspose( f(128), (2, 2, 2), strides=pool_stride, padding='same')(conv6), conv3], axis=4)
	conv7 = Conv3D( f(128), (3, 3, 3), activation='relu', padding='same')(up7)
	conv7 = Conv3D( f(128), (3, 3, 3), activation='relu', padding='same')(conv7)

	up8 = concatenate([Conv3DTranspose( f(64), (2, 2, 2), strides=pool_stride, padding='same')(conv7), conv2], axis=4)
	conv8 = Conv3D( f(64), (3, 3, 3), activation='relu', padding='same')(up8)
	conv8 = Conv3D( f(64), (3, 3, 3), activation='relu', padding='same')(conv8)

	#up9 = concatenate([Conv3DTranspose( f(32), (2, 2, 2), strides=(1, 1, 1), padding='same')(conv8), conv1], axis=4)
	#conv9 = Conv3D( f(32), (3, 3, 3), activation='relu', padding='same')(up9)
	#conv9 = Conv3D( f(32), (3, 3, 3), activation='relu', padding='same')(conv9)

	outputs = Conv3D(num_of_classes, (1, 1, 1), activation='softmax')(conv8)

	outputs = Reshape((num_of_classes, input_shape[0]*input_shape[1]*input_shape[2]), input_shape = input_shape)(outputs)
	outputs = Permute((2,1))(outputs)
	activation = Activation('softmax')(outputs)

	model = Model(inputs=[inputs], outputs=[outputs])
	#model.summary()
	model.compile(optimizer=Adam(lr=1e-5), loss=categorical_focal_loss_fixed, metrics=['accuracy' ]) #loss="mean_absolute_error", metrics=["accuracy"])  #loss=dice_coef_loss, metrics=[dice_coef])  #loss='categorical_crossentropy', metrics=['accuracy']) # 

	return model


def categorical_focal_loss_fixed(y_true, y_pred):
    """
    :param y_true: A tensor of the same shape as `y_pred`
    :param y_pred: A tensor resulting from a softmax
    :return: Output tensor.
    """
    gamma=2.
    alpha=.25
    # Scale predictions so that the class probas of each sample sum to 1
    y_pred /= K.sum(y_pred, axis=-1, keepdims=True)

    # Clip the prediction value to prevent NaN's and Inf's
    epsilon = K.epsilon()
    y_pred = K.clip(y_pred, epsilon, 1. - epsilon)

    # Calculate Cross Entropy
    cross_entropy = -y_true * K.log(y_pred)

    # Calculate Focal Loss
    loss = alpha * K.pow(1 - y_pred, gamma) * cross_entropy

    # Sum the losses in mini_batch
    return K.sum(loss, axis=1)


def get_unet2_FULL_IMAGE( img_rows, img_cols, img_depth, num_of_classes = 10, batch_normalization = True, IMAGE_ORDERING = "channels_last", pool = (1, 1, 1)):
	x = 1/2
	input_shape = (img_rows, img_cols, img_depth,  1)
	inputs = Input(input_shape)
	print( "inputs shape:",inputs.shape)

	conv1 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(inputs)
	conv1 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv1)
	pool1 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv1)
	print ("pool1 shape:",pool1.shape)

	conv2 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool1)
	conv2 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv2)
	pool2 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv2)
	print ("pool2 shape:",pool2.shape)

	conv3 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool2)
	conv3 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv3)
	pool3 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv3)
	print ("pool3 shape:",pool3.shape)
	
	conv4 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool3)
	conv4 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv4)
	pool4 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv4)
	print( "pool4 shape:",pool4.shape)

	conv5 = Conv3D(int(x*512), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool4)
	conv5 = Conv3D(int(x*512), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv5)
	conv5 = Dropout(0.2)(conv5) # drop 0.5

	################### 
	
	up6 = Conv3DTranspose(int(x*256), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv5)
	up6 = concatenate([up6, conv4], axis=4)
	conv6 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up6)
	conv6 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv6)
	print( "up6 shape:",up6.shape)

	up7 = concatenate([Conv3DTranspose(int(x*128), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv6), conv3], axis=4)
	conv7 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up7)
	conv7 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv7)
	print ("conv7 shape:",conv7.shape)

	up8 = concatenate([Conv3DTranspose(int(x*64), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv7), conv2], axis=4)
	conv8 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up8)
	conv8 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv8)
	print ("conv8 shape:",conv8.shape)

	up9 = concatenate([Conv3DTranspose(int(x*32), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv8), conv1], axis=4)
	conv9 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up9)
	conv9 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv9)
	conv9 = Conv3D(num_of_classes,  kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv9)									
	print ("conv9 shape: ",conv9.shape)
	conv9 = BatchNormalization()(conv9) if batch_normalization else conv9 ## NEW

	# output
	conv10 = Reshape((num_of_classes, img_depth * img_rows * img_cols), input_shape=input_shape)(conv9) # mozna je problem tady, zkusit mergnout pouze radke a sloupce a hloubku necat samostatne
	print ("conv10 shape:",conv10.shape)

	permute = Permute((2,1))(conv10)
	print ("permute shape:",permute.shape)

	activation = Activation('softmax')(permute)
	print ("inputs shape: ",inputs.shape)
	print ("outputs shape: ",activation.shape)

	model = Model(inputs=[inputs], outputs=[activation])
	model.summary()
	model.compile(optimizer=Adam(lr=1e-5), loss='categorical_crossentropy', metrics=['accuracy' ]) #loss="mean_absolute_error", metrics=["accuracy"])  #loss=dice_coef_loss, metrics=[dice_coef])  #loss='categorical_crossentropy', metrics=['accuracy']) # 
	return model;



# muj upravnej z ultra nerve segmentation
def get_unet2(img_depth, img_rows, img_cols, num_of_classes = 10, batch_normalization = True, IMAGE_ORDERING = "channels_last"):
	x = 1/4
	inputs = Input((img_depth, img_rows, img_cols, 1))
	print( "inputs shape:",inputs.shape)

	conv1 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(inputs)
	conv1 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv1)
	pool1 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv1)
	print ("pool1 shape:",pool1.shape)

	conv2 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool1)
	conv2 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv2)
	pool2 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv2)
	print ("pool2 shape:",pool2.shape)

	conv3 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool2)
	conv3 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv3)
	pool3 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv3)
	print ("pool3 shape:",pool3.shape)
	
	conv4 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool3)
	conv4 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv4)
	pool4 = MaxPooling3D(pool_size=(1, 1, 1), data_format=IMAGE_ORDERING )(conv4)
	print( "pool4 shape:",pool4.shape)

	conv5 = Conv3D(int(x*512), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(pool4)
	conv5 = Conv3D(int(x*512), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv5)
	conv5 = Dropout(0.2)(conv5) # drop 0.5

	################### 
	
	up6 = Conv3DTranspose(int(x*256), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv5)
	up6 = concatenate([up6, conv4], axis=4)
	conv6 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up6)
	conv6 = Conv3D(int(x*256), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv6)
	print( "up6 shape:",up6.shape)

	up7 = concatenate([Conv3DTranspose(int(x*128), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv6), conv3], axis=4)
	conv7 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up7)
	conv7 = Conv3D(int(x*128), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv7)
	print ("conv7 shape:",conv7.shape)

	up8 = concatenate([Conv3DTranspose(int(x*64), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv7), conv2], axis=4)
	conv8 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up8)
	conv8 = Conv3D(int(x*64), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv8)
	print ("conv8 shape:",conv8.shape)

	up9 = concatenate([Conv3DTranspose(int(x*32), (1, 1, 1), strides=(1, 1, 1), padding='same', data_format=IMAGE_ORDERING )(conv8), conv1], axis=4)
	conv9 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(up9)
	conv9 = Conv3D(int(x*32), kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv9)
	conv9 = Conv3D(num_of_classes,  kernel_size=(2,2,2), activation='relu', padding='same', data_format=IMAGE_ORDERING )(conv9)									
	print ("conv9 shape: ",conv9.shape)
	conv9 = BatchNormalization()(conv9) if batch_normalization else conv9 ## NEW

	# output
	conv10 = Reshape((num_of_classes, img_depth * img_rows * img_cols), input_shape = (num_of_classes, img_depth, img_rows, img_cols))(conv9) # mozna je problem tady, zkusit mergnout pouze radke a sloupce a hloubku necat samostatne
	print ("conv10 shape:",conv10.shape)

	permute = Permute((2,1))(conv10)
	print ("permute shape:",permute.shape)

	activation = Activation('softmax')(permute)
	print ("inputs shape: ",inputs.shape)
	print ("outputs shape: ",activation.shape)

	model = Model(inputs=[inputs], outputs=[activation])
	model.summary()
	model.compile(optimizer=Adam(lr=1e-5), loss='categorical_crossentropy', metrics=['accuracy']) #loss="mean_absolute_error", metrics=["accuracy"])  #loss=dice_coef_loss, metrics=[dice_coef])  #loss='categorical_crossentropy', metrics=['accuracy']) # 
	return model;



def model_3_encoders_1_decoder(input_image_shape = (48, 240, 240, 1), number_of_classes = 10, do=0.2):
	def f(filters):
		return int(filters * 5.0)
	# dopsat
	return True


########################## DENSE UNET #####################


def dense_unet(input_image_size, classes = 10, ki="he_normal",  pool_deconv_size = (2, 2, 2)):
	inputs = Input((*input_image_size, 1))
	conv11 = Conv3D(32, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(inputs)
	conc11 = concatenate([inputs, conv11], axis=4)
	conv12 = Conv3D(32, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc11)
	conc12 = concatenate([inputs, conv12], axis=4)
	pool1 = MaxPooling3D(pool_size=pool_deconv_size)(conc12)

	conv21 = Conv3D(64, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(pool1)
	conc21 = concatenate([pool1, conv21], axis=4)
	conv22 = Conv3D(64, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc21)
	conc22 = concatenate([pool1, conv22], axis=4)
	pool2 = MaxPooling3D(pool_size=pool_deconv_size)(conc22)

	conv31 = Conv3D(128, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(pool2)
	conc31 = concatenate([pool2, conv31], axis=4)
	conv32 = Conv3D(128, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc31)
	conc32 = concatenate([pool2, conv32], axis=4)
	pool3 = MaxPooling3D(pool_size=pool_deconv_size)(conc32)

	conv41 = Conv3D(256, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(pool3)
	conc41 = concatenate([pool3, conv41], axis=4)
	conv42 = Conv3D(256, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc41)
	conc42 = concatenate([pool3, conv42], axis=4)
	pool4 = MaxPooling3D(pool_size=pool_deconv_size)(conc42)

	conv51 = Conv3D(512, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(pool4)
	conc51 = concatenate([pool4, conv51], axis=4)
	conv52 = Conv3D(512, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc51)
	conc52 = concatenate([pool4, conv52], axis=4)

	up6 = concatenate([Conv3DTranspose(256, (2, 2, 2), strides=pool_deconv_size, padding='same', kernel_initializer=ki)(conc52), conc42], axis=4)
	conv61 = Conv3D(256, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(up6)
	conc61 = concatenate([up6, conv61], axis=4)
	conv62 = Conv3D(256, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc61)
	conc62 = concatenate([up6, conv62], axis=4)

	up7 = concatenate([Conv3DTranspose(128, (2, 2, 2), strides=pool_deconv_size, padding='same', kernel_initializer=ki)(conc62), conv32], axis=4)
	conv71 = Conv3D(128, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(up7)
	conc71 = concatenate([up7, conv71], axis=4)
	conv72 = Conv3D(128, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc71)
	conc72 = concatenate([up7, conv72], axis=4)

	up8 = concatenate([Conv3DTranspose(64, (2, 2, 2), strides=pool_deconv_size, padding='same', kernel_initializer=ki)(conc72), conv22], axis=4)
	conv81 = Conv3D(64, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(up8)
	conc81 = concatenate([up8, conv81], axis=4)
	conv82 = Conv3D(64, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc81)
	conc82 = concatenate([up8, conv82], axis=4)

	up9 = concatenate([Conv3DTranspose(32, (2, 2, 2), strides=pool_deconv_size, padding='same', kernel_initializer=ki)(conc82), conv12], axis=4)
	conv91 = Conv3D(32, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(up9)
	conc91 = concatenate([up9, conv91], axis=4)
	conv92 = Conv3D(32, (3, 3, 3), activation='relu', padding='same', kernel_initializer=ki)(conc91)
	conc92 = concatenate([up9, conv92], axis=4)

	conv10 = Conv3D(classes, (1, 1, 1), activation='softmax')(conc92)
	conv10 = Reshape((classes, input_image_size[0]*input_image_size[1]*input_image_size[2]), input_shape = input_image_size)(conv10)
	conv10 = Permute((2,1))(conv10)

	model = Model(inputs=[inputs], outputs=[conv10])
	model.summary()
	model.compile(optimizer=Adam(lr=1e-5, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.000000199), loss='categorical_crossentropy', metrics=['accuracy'])
	return model

###############################################
