import numpy as np
import random
import ConfigParser

from help_functions import load_hdf5
from help_functions import visualize
from help_functions import group_images

from pre_processing import my_PreProc


#Load the original data and return the extracted patches for training/testing
def get_data_training(train_imgs_original, train_groundTruth,
					  patch_height, patch_width,
					  N_subimgs, inside_FOV, patches):

	# Load train images from hdf5 files from pre-processing
	train_imgs_original = load_hdf5(train_imgs_original)
	train_groundTruth = load_hdf5(train_groundTruth) #masks always the same


	# Normalize images
	train_imgs = my_PreProc(train_imgs_original)
	train_groundTruth = train_groundTruth/255.
	visualize(group_images(train_imgs[100:120,:,:,:],5),'imgs_train')#.show()  #check original imgs train

	# shuffle indices to shuffle data
	idx = np.random.permutation(train_imgs.shape[0])
	train_imgs = train_imgs[idx]
	train_groundTruth = train_groundTruth[idx]

	visualize(group_images(train_imgs[100:120,:,:,:],5),'imgs_train_random')
	visualize(group_images(train_groundTruth[100:120,:,:,:],5),'gTruths_train_random')

	#train_imgs = train_imgs[:,:,9:574,:]  #cut bottom and top so now it is 565*565
	#train_groundTruth = train_groundTruth[:,:,9:574,:]  #cut bottom and top so now it is 565*565
	# data_consistency_check(train_imgs,train_groundTruth)

	# Check masks are within 0-1
	assert(np.min(train_groundTruth) == 0 and np.max(train_groundTruth) == 1)

	print "train images shape:" + str(train_imgs.shape)
	print "train images range (min-max): " + str(np.min(train_imgs)) + ' - ' + str(np.max(train_imgs))
	print "train ground truths shape:" + str(train_groundTruth.shape)
	print "train ground truths range (min-max): " + str(np.min(train_groundTruth)) + ' - ' + str(np.max(train_groundTruth))

	if patches == True:
		# Extract the TRAINING patches from the full images
		patches_imgs_train, patches_groundTruths_train = extract_random(train_imgs,train_groundTruth,patch_height,patch_width,N_subimgs,inside_FOV)
		data_consistency_check(patches_imgs_train, patches_groundTruths_train)

		print "train PATCHES images shape: " + str(patches_imgs_train.shape)
		print "train PATCHES images range (min-max): " + str(np.min(patches_imgs_train)) + ' - ' + str(np.max(patches_imgs_train))
		print "train PATCHES ground truths shape: " + str(patches_groundTruths_train.shape)
		print "train PATCHES ground truths range (min-max): " + str(np.min(patches_groundTruths_train)) + ' - ' + str(np.max(patches_groundTruths_train))

		# visualize(group_images(patches_imgs_train[100:120,:,:,:],5),'imgs_train_patches')
		# visualize(group_images(patches_groundTruths_train[100:120,:,:,:],5),'gTruth_train_patches')

		return patches_imgs_train, patches_groundTruths_train
	else:
		return train_imgs, train_groundTruth

#Load the original data and return the extracted patches for training/testing
def get_data_testing(test_imgs_original, test_groudTruth, Imgs_to_test, patch_height, patch_width):

	# Load test images from hdf5 files from pre-processing
	test_imgs_original = load_hdf5(test_imgs_original)
	test_groundTruths = load_hdf5(test_groudTruth)

	# Normalize test ground truths
	test_imgs = my_PreProc(test_imgs_original)
	test_groundTruths = test_groundTruths/255.

	#extend both images and masks so they can be divided exactly by the patches dimensions
	test_imgs = test_imgs[0:Imgs_to_test,:,:,:]
	test_groundTruths = test_groundTruths[0:Imgs_to_test,:,:,:]
	test_imgs = paint_border(test_imgs,patch_height,patch_width)
	test_groundTruths = paint_border(test_groundTruths,patch_height,patch_width)

	data_consistency_check(test_imgs, test_groundTruths)

	# Check masks are within 0-1
	assert(np.max(test_groundTruths)==1  and np.min(test_groundTruths)==0)

	print "train images shape:" + str(test_imgs.shape)
	print "train images range (min-max): " + str(np.min(test_imgs)) + ' - ' + str(np.max(test_imgs))
	print "train ground truths shape:" + str(test_groundTruths.shape)
	print "train ground truths range (min-max): " + str(np.min(test_groundTruths)) + ' - ' + str(np.max(test_groundTruths))

	#extract the TEST patches from the full images
	patches_imgs_test = extract_ordered(test_imgs,patch_height,patch_width)
	patches_groundTruths_test = extract_ordered(test_groundTruths,patch_height,patch_width)
	data_consistency_check(patches_imgs_test, patches_groundTruths_test)

	print "test PATCHES images shape: " + patches_imgs_test.shape
	print "test PATCHES images range (min-max): " + str(np.min(patches_imgs_test)) + ' - ' + str(np.max(patches_imgs_test))
	print "test PATCHES ground truths shape: " + patches_groundTruths_test.shape
	print "test PATCHES ground truths range (min-max): " + str(np.min(patches_groundTruths_test)) + ' - ' + str(np.max(patches_groundTruths_test))

	return patches_imgs_test, patches_groundTruths_test




# Load the original data and return the extracted patches for testing
# return the ground truth in its original shape
def get_data_testing_overlap(test_imgs_original,
							test_groudTruth,
							Imgs_to_test,
							patch_height, patch_width,
							stride_height, stride_width, patches):
	### test
	test_imgs_original = load_hdf5(test_imgs_original)
	test_imgs_groundTruth = load_hdf5(test_groudTruth)

	test_imgs = my_PreProc(test_imgs_original)
	test_imgs_groundTruth = test_imgs_groundTruth/255.

	#extend both images and masks so they can be divided exactly by the patches dimensions
	# test_imgs = test_imgs[0:Imgs_to_test,:,:,:]
	# test_imgs_groundTruth = test_imgs_groundTruth[0:Imgs_to_test,:,:,:]
	# test_imgs = paint_border_overlap(test_imgs, patch_height, patch_width, stride_height, stride_width)

	#check masks are within 0-1
	assert(np.max(test_imgs_groundTruth)==1  and np.min(test_imgs_groundTruth)==0)

	print "\ntest images shape:"
	print test_imgs.shape
	print "\ntest mask shape:"
	print test_imgs_groundTruth.shape
	print "test images range (min-max): " +str(np.min(test_imgs)) +' - '+str(np.max(test_imgs))
	print "test masks are within 0-1\n"

	if patches == True:
		#extract the TEST patches from the full images
		patches_imgs_test = extract_ordered_overlap(test_imgs,patch_height,patch_width,stride_height,stride_width)

		print "\ntest PATCHES images shape:"
		print patches_imgs_test.shape
		print "test PATCHES images range (min-max): " +str(np.min(patches_imgs_test)) +' - '+str(np.max(patches_imgs_test))

		return patches_imgs_test, test_imgs.shape[2], test_imgs.shape[3], test_imgs_groundTruth
		#return patches_imgs_test, test_imgs.shape[2], test_imgs.shape[3]
	else:
		return test_imgs, test_imgs.shape[2], test_imgs.shape[3], test_imgs_groundTruth


# data consinstency check - checks if imgs shape == masks shape
def data_consistency_check(imgs,masks):
	assert(len(imgs.shape)==len(masks.shape))
	assert(imgs.shape[0]==masks.shape[0])
	assert(imgs.shape[2]==masks.shape[2])
	assert(imgs.shape[3]==masks.shape[3])
	assert(masks.shape[1]==1)
	assert(imgs.shape[1]==1 or imgs.shape[1]==3)


#extract patches randomly in the full training images
#  -- Inside OR in full image
def extract_random(full_imgs,full_masks, patch_h,patch_w, N_patches, inside=True):
	print "FULL IMG SHAPE", full_imgs.shape[0]
	if (N_patches%full_imgs.shape[0] != 0):
		print "N_patches: plase enter a multiple of" + full_imgs.shape[0]
		exit()
	assert (len(full_imgs.shape)==4 and len(full_masks.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	print full_masks.shape
	# THEO assert (full_masks.shape[1]==1)   #masks only black and white
	assert (full_imgs.shape[2] == full_masks.shape[2] and full_imgs.shape[3] == full_masks.shape[3])
	patches = np.empty((N_patches,full_imgs.shape[1],patch_h,patch_w))
	patches_masks = np.empty((N_patches,full_masks.shape[1],patch_h,patch_w))
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	# (0,0) in the center of the image
	patch_per_img = int(N_patches/full_imgs.shape[0])  #N_patches equally divided in the full images
	print "patches per full image: " +str(patch_per_img)
	iter_tot = 0   #iter over the total numbe rof patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		k=0
		while k <patch_per_img:
			x_center = random.randint(0+int(patch_w/2),img_w-int(patch_w/2))
			# print "x_center " +str(x_center)
			y_center = random.randint(0+int(patch_h/2),img_h-int(patch_h/2))
			# print "y_center " +str(y_center)
			#check whether the patch is fully contained in the FOV
			if inside==True:
				if is_patch_inside_FOV(x_center,y_center,img_w,img_h,patch_h)==False:
					continue
			patch = full_imgs[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
			patch_mask = full_masks[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
			patches[iter_tot]=patch
			patches_masks[iter_tot]=patch_mask
			iter_tot +=1   #total
			k+=1  #per full_img
	return patches, patches_masks


#check if the patch is fully contained in the FOV
def is_patch_inside_FOV(x,y,img_w,img_h,patch_h):
	x_ = x - int(img_w/2) # origin (0,0) shifted to image center
	y_ = y - int(img_h/2)  # origin (0,0) shifted to image center
	R_inside = 270 - int(patch_h * np.sqrt(2.0) / 2.0) #radius is 270 (from DRIVE db docs), minus the patch diagonal (assumed it is a square #this is the limit to contain the full patch in the FOV
	radius = np.sqrt((x_*x_)+(y_*y_))
	if radius < R_inside:
		return True
	else:
		return False


#Divide all the full_imgs in pacthes
def extract_ordered(full_imgs, patch_h, patch_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	N_patches_h = int(img_h/patch_h) #round to lowest int
	if (img_h%patch_h != 0):
		print "warning: " +str(N_patches_h) +" patches in height, with about " +str(img_h%patch_h) +" pixels left over"
	N_patches_w = int(img_w/patch_w) #round to lowest int
	if (img_h%patch_h != 0):
		print "warning: " +str(N_patches_w) +" patches in width, with about " +str(img_w%patch_w) +" pixels left over"
	print "number of patches per image: " +str(N_patches_h*N_patches_w)
	N_patches_tot = (N_patches_h*N_patches_w)*full_imgs.shape[0]
	patches = np.empty((N_patches_tot,full_imgs.shape[1],patch_h,patch_w))

	iter_tot = 0   #iter over the total number of patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		for h in range(N_patches_h):
			for w in range(N_patches_w):
				patch = full_imgs[i,:,h*patch_h:(h*patch_h)+patch_h,w*patch_w:(w*patch_w)+patch_w]
				patches[iter_tot]=patch
				iter_tot +=1   #total
	assert (iter_tot==N_patches_tot)
	return patches  #array with all the full_imgs divided in patches


def paint_border_overlap(full_imgs, patch_h, patch_w, stride_h, stride_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	leftover_h = (img_h-patch_h)%stride_h  #leftover on the h dim
	leftover_w = (img_w-patch_w)%stride_w  #leftover on the w dim
	if (leftover_h != 0):  #change dimension of img_h
		print "\nthe side H is not compatible with the selected stride of " +str(stride_h)
		print "img_h " +str(img_h) + ", patch_h " +str(patch_h) + ", stride_h " +str(stride_h)
		print "(img_h - patch_h) MOD stride_h: " +str(leftover_h)
		print "So the H dim will be padded with additional " +str(stride_h - leftover_h) + " pixels"
		tmp_full_imgs = np.zeros((full_imgs.shape[0],full_imgs.shape[1],img_h+(stride_h-leftover_h),img_w))
		tmp_full_imgs[0:full_imgs.shape[0],0:full_imgs.shape[1],0:img_h,0:img_w] = full_imgs
		full_imgs = tmp_full_imgs
	if (leftover_w != 0):   #change dimension of img_w
		print "the side W is not compatible with the selected stride of " +str(stride_w)
		print "img_w " +str(img_w) + ", patch_w " +str(patch_w) + ", stride_w " +str(stride_w)
		print "(img_w - patch_w) MOD stride_w: " +str(leftover_w)
		print "So the W dim will be padded with additional " +str(stride_w - leftover_w) + " pixels"
		tmp_full_imgs = np.zeros((full_imgs.shape[0],full_imgs.shape[1],full_imgs.shape[2],img_w+(stride_w - leftover_w)))
		tmp_full_imgs[0:full_imgs.shape[0],0:full_imgs.shape[1],0:full_imgs.shape[2],0:img_w] = full_imgs
		full_imgs = tmp_full_imgs
	print "new full images shape: \n" +str(full_imgs.shape)
	return full_imgs

#Divide all the full_imgs in pacthes
def extract_ordered_overlap(full_imgs, patch_h, patch_w,stride_h,stride_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	assert ((img_h-patch_h)%stride_h==0 and (img_w-patch_w)%stride_w==0)
	N_patches_img = ((img_h-patch_h)//stride_h+1)*((img_w-patch_w)//stride_w+1)  #// --> division between integers
	N_patches_tot = N_patches_img*full_imgs.shape[0]
	print "Number of patches on h : " +str(((img_h-patch_h)//stride_h+1))
	print "Number of patches on w : " +str(((img_w-patch_w)//stride_w+1))
	print "number of patches per image: " +str(N_patches_img) +", totally for this dataset: " +str(N_patches_tot)
	patches = np.empty((N_patches_tot,full_imgs.shape[1],patch_h,patch_w))
	iter_tot = 0   #iter over the total number of patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		for h in range((img_h-patch_h)//stride_h+1):
			for w in range((img_w-patch_w)//stride_w+1):
				patch = full_imgs[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]
				patches[iter_tot]=patch
				iter_tot +=1   #total
	assert (iter_tot==N_patches_tot)
	return patches  #array with all the full_imgs divided in patches


def recompone_overlap(preds, img_h, img_w, stride_h, stride_w):
	assert (len(preds.shape)==4)  #4D arrays
	assert (preds.shape[1]==1 or preds.shape[1]==3)  #check the channel is 1 or 3
	patch_h = preds.shape[2]
	patch_w = preds.shape[3]
	N_patches_h = (img_h-patch_h)//stride_h+1
	N_patches_w = (img_w-patch_w)//stride_w+1
	N_patches_img = N_patches_h * N_patches_w
	print "N_patches_h: " +str(N_patches_h)
	print "N_patches_w: " +str(N_patches_w)
	print "N_patches_img: " +str(N_patches_img)
	assert (preds.shape[0]%N_patches_img==0)
	N_full_imgs = preds.shape[0]//N_patches_img
	print "According to the dimension inserted, there are " +str(N_full_imgs) +" full images (of " +str(img_h)+"x" +str(img_w) +" each)"
	full_prob = np.zeros((N_full_imgs,preds.shape[1],img_h,img_w))  #itialize to zero mega array with sum of Probabilities
	full_sum = np.zeros((N_full_imgs,preds.shape[1],img_h,img_w))

	k = 0 #iterator over all the patches
	for i in range(N_full_imgs):
		for h in range((img_h-patch_h)//stride_h+1):
			for w in range((img_w-patch_w)//stride_w+1):
				full_prob[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]+=preds[k]
				full_sum[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]+=1
				k+=1
	assert(k==preds.shape[0])
	assert(np.min(full_sum)>=1.0)  #at least one
	final_avg = full_prob/full_sum
	print final_avg.shape
	assert(np.max(final_avg)<=1.0) #max value for a pixel is 1.0
	assert(np.min(final_avg)>=0.0) #min value for a pixel is 0.0
	return final_avg


#Recompone the full images with the patches
def recompone(data,N_h,N_w):
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	assert(len(data.shape)==4)
	N_pacth_per_img = N_w*N_h
	assert(data.shape[0]%N_pacth_per_img == 0)
	N_full_imgs = data.shape[0]/N_pacth_per_img
	patch_h = data.shape[2]
	patch_w = data.shape[3]
	N_pacth_per_img = N_w*N_h
	#define and start full recompone
	full_recomp = np.empty((N_full_imgs,data.shape[1],N_h*patch_h,N_w*patch_w))
	k = 0  #iter full img
	s = 0  #iter single patch
	while (s<data.shape[0]):
		#recompone one:
		single_recon = np.empty((data.shape[1],N_h*patch_h,N_w*patch_w))
		for h in range(N_h):
			for w in range(N_w):
				single_recon[:,h*patch_h:(h*patch_h)+patch_h,w*patch_w:(w*patch_w)+patch_w]=data[s]
				s+=1
		full_recomp[k]=single_recon
		k+=1
	assert (k==N_full_imgs)
	return full_recomp


#Extend the full images becasue patch divison is not exact
def paint_border(data,patch_h,patch_w):
	assert (len(data.shape)==4)  #4D arrays
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	img_h=data.shape[2]
	img_w=data.shape[3]
	new_img_h = 0
	new_img_w = 0
	if (img_h%patch_h)==0:
		new_img_h = img_h
	else:
		new_img_h = ((int(img_h)/int(patch_h))+1)*patch_h
	if (img_w%patch_w)==0:
		new_img_w = img_w
	else:
		new_img_w = ((int(img_w)/int(patch_w))+1)*patch_w
	new_data = np.zeros((data.shape[0],data.shape[1],new_img_h,new_img_w))
	new_data[:,:,0:img_h,0:img_w] = data[:,:,:,:]
	return new_data


#return only the pixels contained in the FOV, for both images and masks
def pred_only_FOV(data_imgs,data_masks):
	print data_masks.shape
	print data_imgs.shape
	# assert (len(data_imgs.shape)==4 and len(data_masks.shape)==4)  #4D arrays
	assert (data_imgs.shape[0]==data_masks.shape[0])
	assert (data_imgs.shape[2]==data_masks.shape[2])
	assert (data_imgs.shape[3]==data_masks.shape[3])
	# THEO assert (data_imgs.shape[1]==1 and data_masks.shape[1]==1)  #check the channel is 1
	height = data_imgs.shape[2]
	width = data_imgs.shape[3]
	new_pred_imgs = []
	new_pred_masks = []
	for i in range(data_imgs.shape[0]):  #loop over the full images
		for x in range(width):
			for y in range(height):
				# if inside_FOV_DRIVE(i,x,y,original_imgs_border_masks)==True:
				new_pred_imgs.append(data_imgs[i,:,y,x])
				new_pred_masks.append(data_masks[i,:,y,x])
	new_pred_imgs = np.asarray(new_pred_imgs)
	new_pred_masks = np.asarray(new_pred_masks)
	return new_pred_imgs, new_pred_masks

#function to set to black everything outside the FOV, in a full image
def kill_border(data, original_imgs_border_masks):
	assert (len(data.shape)==4)  #4D arrays
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	height = data.shape[2]
	width = data.shape[3]
	for i in range(data.shape[0]):  #loop over the full images
		for x in range(width):
			for y in range(height):
				if inside_FOV_DRIVE(i,x,y,original_imgs_border_masks)==False:
					data[i,:,y,x]=0.0


# def inside_FOV_DRIVE(i, x, y, DRIVE_masks):

# 	return False # victor

# 	assert (len(DRIVE_masks.shape)==4)  #4D arrays
# 	# THEO assert (DRIVE_masks.shape[1]==1)  #DRIVE masks is black and white
# 	# DRIVE_masks = DRIVE_masks/255.  #NOOO!! otherwise with float numbers takes forever!!

# 	if (x >= DRIVE_masks.shape[3] or y >= DRIVE_masks.shape[2]): #my image bigger than the original
# 		return False

# 	if (DRIVE_masks[i,0,y,x]>0):  #0==black pixels
# 		# print DRIVE_masks[i,0,y,x]  #verify it is working right
# 		return True
# 	else:
# 		return False

69635042b5bb88070028c0873c503ba1b7da6686
import numpy as np
import random
import ConfigParser

from help_functions import load_hdf5
from help_functions import visualize
from help_functions import group_images

from pre_processing import my_PreProc


#Load the original data and return the extracted patches for training/testing
def get_data_training(train_imgs_original, train_groundTruth,
					  patch_height, patch_width,
					  N_subimgs, inside_FOV, patches):

	# Load train images from hdf5 files from pre-processing
	train_imgs_original = load_hdf5(train_imgs_original)
	train_groundTruth = load_hdf5(train_groundTruth) #masks always the same


	# Normalize images
	train_imgs = my_PreProc(train_imgs_original)
	train_groundTruth = train_groundTruth/255.
	visualize(group_images(train_imgs[100:120,:,:,:],5),'imgs_train')#.show()  #check original imgs train

	# shuffle indices to shuffle data
	idx = np.random.permutation(train_imgs.shape[0])
	train_imgs = train_imgs[idx]
	train_groundTruth = train_groundTruth[idx]

	visualize(group_images(train_imgs[100:120,:,:,:],5),'imgs_train_random')
	visualize(group_images(train_groundTruth[100:120,:,:,:],5),'gTruths_train_random')

	#train_imgs = train_imgs[:,:,9:574,:]  #cut bottom and top so now it is 565*565
	#train_groundTruth = train_groundTruth[:,:,9:574,:]  #cut bottom and top so now it is 565*565
	# data_consistency_check(train_imgs,train_groundTruth)

	# Check masks are within 0-1
	assert(np.min(train_groundTruth) == 0 and np.max(train_groundTruth) == 1)

	print "train images shape:" + str(train_imgs.shape)
	print "train images range (min-max): " + str(np.min(train_imgs)) + ' - ' + str(np.max(train_imgs))
	print "train ground truths shape:" + str(train_groundTruth.shape)
	print "train ground truths range (min-max): " + str(np.min(train_groundTruth)) + ' - ' + str(np.max(train_groundTruth))

	if patches == True:
		# Extract the TRAINING patches from the full images
		patches_imgs_train, patches_groundTruths_train = extract_random(train_imgs,train_groundTruth,patch_height,patch_width,N_subimgs,inside_FOV)
		data_consistency_check(patches_imgs_train, patches_groundTruths_train)

		print "train PATCHES images shape: " + str(patches_imgs_train.shape)
		print "train PATCHES images range (min-max): " + str(np.min(patches_imgs_train)) + ' - ' + str(np.max(patches_imgs_train))
		print "train PATCHES ground truths shape: " + str(patches_groundTruths_train.shape)
		print "train PATCHES ground truths range (min-max): " + str(np.min(patches_groundTruths_train)) + ' - ' + str(np.max(patches_groundTruths_train))

		# visualize(group_images(patches_imgs_train[100:120,:,:,:],5),'imgs_train_patches')
		# visualize(group_images(patches_groundTruths_train[100:120,:,:,:],5),'gTruth_train_patches')

		return patches_imgs_train, patches_groundTruths_train
	else:
		return train_imgs, train_groundTruth

#Load the original data and return the extracted patches for training/testing
def get_data_testing(test_imgs_original, test_groudTruth, Imgs_to_test, patch_height, patch_width):

	# Load test images from hdf5 files from pre-processing
	test_imgs_original = load_hdf5(test_imgs_original)
	test_groundTruths = load_hdf5(test_groudTruth)

	# Normalize test ground truths
	test_imgs = my_PreProc(test_imgs_original)
	test_groundTruths = test_groundTruths/255.

	#extend both images and masks so they can be divided exactly by the patches dimensions
	test_imgs = test_imgs[0:Imgs_to_test,:,:,:]
	test_groundTruths = test_groundTruths[0:Imgs_to_test,:,:,:]
	test_imgs = paint_border(test_imgs,patch_height,patch_width)
	test_groundTruths = paint_border(test_groundTruths,patch_height,patch_width)

	data_consistency_check(test_imgs, test_groundTruths)

	# Check masks are within 0-1
	assert(np.max(test_groundTruths)==1  and np.min(test_groundTruths)==0)

	print "train images shape:" + str(test_imgs.shape)
	print "train images range (min-max): " + str(np.min(test_imgs)) + ' - ' + str(np.max(test_imgs))
	print "train ground truths shape:" + str(test_groundTruths.shape)
	print "train ground truths range (min-max): " + str(np.min(test_groundTruths)) + ' - ' + str(np.max(test_groundTruths))

	#extract the TEST patches from the full images
	patches_imgs_test = extract_ordered(test_imgs,patch_height,patch_width)
	patches_groundTruths_test = extract_ordered(test_groundTruths,patch_height,patch_width)
	data_consistency_check(patches_imgs_test, patches_groundTruths_test)

	print "test PATCHES images shape: " + patches_imgs_test.shape
	print "test PATCHES images range (min-max): " + str(np.min(patches_imgs_test)) + ' - ' + str(np.max(patches_imgs_test))
	print "test PATCHES ground truths shape: " + patches_groundTruths_test.shape
	print "test PATCHES ground truths range (min-max): " + str(np.min(patches_groundTruths_test)) + ' - ' + str(np.max(patches_groundTruths_test))

	return patches_imgs_test, patches_groundTruths_test




# Load the original data and return the extracted patches for testing
# return the ground truth in its original shape
def get_data_testing_overlap(test_imgs_original,
							test_groudTruth,
							Imgs_to_test,
							patch_height, patch_width,
							stride_height, stride_width, patches):
	### test
	test_imgs_original = load_hdf5(test_imgs_original)
	test_imgs_groundTruth = load_hdf5(test_groudTruth)

	test_imgs = my_PreProc(test_imgs_original)
	test_imgs_groundTruth = test_imgs_groundTruth/255.

	#extend both images and masks so they can be divided exactly by the patches dimensions
	# test_imgs = test_imgs[0:Imgs_to_test,:,:,:]
	# test_imgs_groundTruth = test_imgs_groundTruth[0:Imgs_to_test,:,:,:]
	# test_imgs = paint_border_overlap(test_imgs, patch_height, patch_width, stride_height, stride_width)

	#check masks are within 0-1
	assert(np.max(test_imgs_groundTruth)==1  and np.min(test_imgs_groundTruth)==0)

	print "\ntest images shape:"
	print test_imgs.shape
	print "\ntest mask shape:"
	print test_imgs_groundTruth.shape
	print "test images range (min-max): " +str(np.min(test_imgs)) +' - '+str(np.max(test_imgs))
	print "test masks are within 0-1\n"

	if patches == True:
		#extract the TEST patches from the full images
		patches_imgs_test = extract_ordered_overlap(test_imgs,patch_height,patch_width,stride_height,stride_width)

		print "\ntest PATCHES images shape:"
		print patches_imgs_test.shape
		print "test PATCHES images range (min-max): " +str(np.min(patches_imgs_test)) +' - '+str(np.max(patches_imgs_test))

		return patches_imgs_test, test_imgs.shape[2], test_imgs.shape[3], test_imgs_groundTruth
		#return patches_imgs_test, test_imgs.shape[2], test_imgs.shape[3]
	else:
		return test_imgs, test_imgs.shape[2], test_imgs.shape[3], test_imgs_groundTruth


# data consinstency check - checks if imgs shape == masks shape
def data_consistency_check(imgs,masks):
	assert(len(imgs.shape)==len(masks.shape))
	assert(imgs.shape[0]==masks.shape[0])
	assert(imgs.shape[2]==masks.shape[2])
	assert(imgs.shape[3]==masks.shape[3])
	assert(masks.shape[1]==1)
	assert(imgs.shape[1]==1 or imgs.shape[1]==3)


#extract patches randomly in the full training images
#  -- Inside OR in full image
def extract_random(full_imgs,full_masks, patch_h,patch_w, N_patches, inside=True):
	print "FULL IMG SHAPE", full_imgs.shape[0]
	if (N_patches%full_imgs.shape[0] != 0):
		print "N_patches: plase enter a multiple of" + full_imgs.shape[0]
		exit()
	assert (len(full_imgs.shape)==4 and len(full_masks.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	print full_masks.shape
	# THEO assert (full_masks.shape[1]==1)   #masks only black and white
	assert (full_imgs.shape[2] == full_masks.shape[2] and full_imgs.shape[3] == full_masks.shape[3])
	patches = np.empty((N_patches,full_imgs.shape[1],patch_h,patch_w))
	patches_masks = np.empty((N_patches,full_masks.shape[1],patch_h,patch_w))
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	# (0,0) in the center of the image
	patch_per_img = int(N_patches/full_imgs.shape[0])  #N_patches equally divided in the full images
	print "patches per full image: " +str(patch_per_img)
	iter_tot = 0   #iter over the total numbe rof patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		k=0
		while k <patch_per_img:
			x_center = random.randint(0+int(patch_w/2),img_w-int(patch_w/2))
			# print "x_center " +str(x_center)
			y_center = random.randint(0+int(patch_h/2),img_h-int(patch_h/2))
			# print "y_center " +str(y_center)
			#check whether the patch is fully contained in the FOV
			if inside==True:
				if is_patch_inside_FOV(x_center,y_center,img_w,img_h,patch_h)==False:
					continue
			patch = full_imgs[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
			patch_mask = full_masks[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
			patches[iter_tot]=patch
			patches_masks[iter_tot]=patch_mask
			iter_tot +=1   #total
			k+=1  #per full_img
	return patches, patches_masks


#check if the patch is fully contained in the FOV
def is_patch_inside_FOV(x,y,img_w,img_h,patch_h):
	x_ = x - int(img_w/2) # origin (0,0) shifted to image center
	y_ = y - int(img_h/2)  # origin (0,0) shifted to image center
	R_inside = 270 - int(patch_h * np.sqrt(2.0) / 2.0) #radius is 270 (from DRIVE db docs), minus the patch diagonal (assumed it is a square #this is the limit to contain the full patch in the FOV
	radius = np.sqrt((x_*x_)+(y_*y_))
	if radius < R_inside:
		return True
	else:
		return False


#Divide all the full_imgs in pacthes
def extract_ordered(full_imgs, patch_h, patch_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	N_patches_h = int(img_h/patch_h) #round to lowest int
	if (img_h%patch_h != 0):
		print "warning: " +str(N_patches_h) +" patches in height, with about " +str(img_h%patch_h) +" pixels left over"
	N_patches_w = int(img_w/patch_w) #round to lowest int
	if (img_h%patch_h != 0):
		print "warning: " +str(N_patches_w) +" patches in width, with about " +str(img_w%patch_w) +" pixels left over"
	print "number of patches per image: " +str(N_patches_h*N_patches_w)
	N_patches_tot = (N_patches_h*N_patches_w)*full_imgs.shape[0]
	patches = np.empty((N_patches_tot,full_imgs.shape[1],patch_h,patch_w))

	iter_tot = 0   #iter over the total number of patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		for h in range(N_patches_h):
			for w in range(N_patches_w):
				patch = full_imgs[i,:,h*patch_h:(h*patch_h)+patch_h,w*patch_w:(w*patch_w)+patch_w]
				patches[iter_tot]=patch
				iter_tot +=1   #total
	assert (iter_tot==N_patches_tot)
	return patches  #array with all the full_imgs divided in patches


def paint_border_overlap(full_imgs, patch_h, patch_w, stride_h, stride_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	leftover_h = (img_h-patch_h)%stride_h  #leftover on the h dim
	leftover_w = (img_w-patch_w)%stride_w  #leftover on the w dim
	if (leftover_h != 0):  #change dimension of img_h
		print "\nthe side H is not compatible with the selected stride of " +str(stride_h)
		print "img_h " +str(img_h) + ", patch_h " +str(patch_h) + ", stride_h " +str(stride_h)
		print "(img_h - patch_h) MOD stride_h: " +str(leftover_h)
		print "So the H dim will be padded with additional " +str(stride_h - leftover_h) + " pixels"
		tmp_full_imgs = np.zeros((full_imgs.shape[0],full_imgs.shape[1],img_h+(stride_h-leftover_h),img_w))
		tmp_full_imgs[0:full_imgs.shape[0],0:full_imgs.shape[1],0:img_h,0:img_w] = full_imgs
		full_imgs = tmp_full_imgs
	if (leftover_w != 0):   #change dimension of img_w
		print "the side W is not compatible with the selected stride of " +str(stride_w)
		print "img_w " +str(img_w) + ", patch_w " +str(patch_w) + ", stride_w " +str(stride_w)
		print "(img_w - patch_w) MOD stride_w: " +str(leftover_w)
		print "So the W dim will be padded with additional " +str(stride_w - leftover_w) + " pixels"
		tmp_full_imgs = np.zeros((full_imgs.shape[0],full_imgs.shape[1],full_imgs.shape[2],img_w+(stride_w - leftover_w)))
		tmp_full_imgs[0:full_imgs.shape[0],0:full_imgs.shape[1],0:full_imgs.shape[2],0:img_w] = full_imgs
		full_imgs = tmp_full_imgs
	print "new full images shape: \n" +str(full_imgs.shape)
	return full_imgs

#Divide all the full_imgs in pacthes
def extract_ordered_overlap(full_imgs, patch_h, patch_w,stride_h,stride_w):
	assert (len(full_imgs.shape)==4)  #4D arrays
	assert (full_imgs.shape[1]==1 or full_imgs.shape[1]==3)  #check the channel is 1 or 3
	img_h = full_imgs.shape[2]  #height of the full image
	img_w = full_imgs.shape[3] #width of the full image
	assert ((img_h-patch_h)%stride_h==0 and (img_w-patch_w)%stride_w==0)
	N_patches_img = ((img_h-patch_h)//stride_h+1)*((img_w-patch_w)//stride_w+1)  #// --> division between integers
	N_patches_tot = N_patches_img*full_imgs.shape[0]
	print "Number of patches on h : " +str(((img_h-patch_h)//stride_h+1))
	print "Number of patches on w : " +str(((img_w-patch_w)//stride_w+1))
	print "number of patches per image: " +str(N_patches_img) +", totally for this dataset: " +str(N_patches_tot)
	patches = np.empty((N_patches_tot,full_imgs.shape[1],patch_h,patch_w))
	iter_tot = 0   #iter over the total number of patches (N_patches)
	for i in range(full_imgs.shape[0]):  #loop over the full images
		for h in range((img_h-patch_h)//stride_h+1):
			for w in range((img_w-patch_w)//stride_w+1):
				patch = full_imgs[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]
				patches[iter_tot]=patch
				iter_tot +=1   #total
	assert (iter_tot==N_patches_tot)
	return patches  #array with all the full_imgs divided in patches


def recompone_overlap(preds, img_h, img_w, stride_h, stride_w):
	assert (len(preds.shape)==4)  #4D arrays
	assert (preds.shape[1]==1 or preds.shape[1]==3)  #check the channel is 1 or 3
	patch_h = preds.shape[2]
	patch_w = preds.shape[3]
	N_patches_h = (img_h-patch_h)//stride_h+1
	N_patches_w = (img_w-patch_w)//stride_w+1
	N_patches_img = N_patches_h * N_patches_w
	print "N_patches_h: " +str(N_patches_h)
	print "N_patches_w: " +str(N_patches_w)
	print "N_patches_img: " +str(N_patches_img)
	assert (preds.shape[0]%N_patches_img==0)
	N_full_imgs = preds.shape[0]//N_patches_img
	print "According to the dimension inserted, there are " +str(N_full_imgs) +" full images (of " +str(img_h)+"x" +str(img_w) +" each)"
	full_prob = np.zeros((N_full_imgs,preds.shape[1],img_h,img_w))  #itialize to zero mega array with sum of Probabilities
	full_sum = np.zeros((N_full_imgs,preds.shape[1],img_h,img_w))

	k = 0 #iterator over all the patches
	for i in range(N_full_imgs):
		for h in range((img_h-patch_h)//stride_h+1):
			for w in range((img_w-patch_w)//stride_w+1):
				full_prob[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]+=preds[k]
				full_sum[i,:,h*stride_h:(h*stride_h)+patch_h,w*stride_w:(w*stride_w)+patch_w]+=1
				k+=1
	assert(k==preds.shape[0])
	assert(np.min(full_sum)>=1.0)  #at least one
	final_avg = full_prob/full_sum
	print final_avg.shape
	assert(np.max(final_avg)<=1.0) #max value for a pixel is 1.0
	assert(np.min(final_avg)>=0.0) #min value for a pixel is 0.0
	return final_avg


#Recompone the full images with the patches
def recompone(data,N_h,N_w):
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	assert(len(data.shape)==4)
	N_pacth_per_img = N_w*N_h
	assert(data.shape[0]%N_pacth_per_img == 0)
	N_full_imgs = data.shape[0]/N_pacth_per_img
	patch_h = data.shape[2]
	patch_w = data.shape[3]
	N_pacth_per_img = N_w*N_h
	#define and start full recompone
	full_recomp = np.empty((N_full_imgs,data.shape[1],N_h*patch_h,N_w*patch_w))
	k = 0  #iter full img
	s = 0  #iter single patch
	while (s<data.shape[0]):
		#recompone one:
		single_recon = np.empty((data.shape[1],N_h*patch_h,N_w*patch_w))
		for h in range(N_h):
			for w in range(N_w):
				single_recon[:,h*patch_h:(h*patch_h)+patch_h,w*patch_w:(w*patch_w)+patch_w]=data[s]
				s+=1
		full_recomp[k]=single_recon
		k+=1
	assert (k==N_full_imgs)
	return full_recomp


#Extend the full images becasue patch divison is not exact
def paint_border(data,patch_h,patch_w):
	assert (len(data.shape)==4)  #4D arrays
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	img_h=data.shape[2]
	img_w=data.shape[3]
	new_img_h = 0
	new_img_w = 0
	if (img_h%patch_h)==0:
		new_img_h = img_h
	else:
		new_img_h = ((int(img_h)/int(patch_h))+1)*patch_h
	if (img_w%patch_w)==0:
		new_img_w = img_w
	else:
		new_img_w = ((int(img_w)/int(patch_w))+1)*patch_w
	new_data = np.zeros((data.shape[0],data.shape[1],new_img_h,new_img_w))
	new_data[:,:,0:img_h,0:img_w] = data[:,:,:,:]
	return new_data


#return only the pixels contained in the FOV, for both images and masks
def pred_only_FOV(data_imgs,data_masks):
	print data_masks.shape
	print data_imgs.shape
	# assert (len(data_imgs.shape)==4 and len(data_masks.shape)==4)  #4D arrays
	assert (data_imgs.shape[0]==data_masks.shape[0])
	assert (data_imgs.shape[2]==data_masks.shape[2])
	assert (data_imgs.shape[3]==data_masks.shape[3])
	# THEO assert (data_imgs.shape[1]==1 and data_masks.shape[1]==1)  #check the channel is 1
	height = data_imgs.shape[2]
	width = data_imgs.shape[3]
	new_pred_imgs = []
	new_pred_masks = []
	for i in range(data_imgs.shape[0]):  #loop over the full images
		for x in range(width):
			for y in range(height):
				# if inside_FOV_DRIVE(i,x,y,original_imgs_border_masks)==True:
				new_pred_imgs.append(data_imgs[i,:,y,x])
				new_pred_masks.append(data_masks[i,:,y,x])
	new_pred_imgs = np.asarray(new_pred_imgs)
	new_pred_masks = np.asarray(new_pred_masks)
	return new_pred_imgs, new_pred_masks

#function to set to black everything outside the FOV, in a full image
def kill_border(data, original_imgs_border_masks):
	assert (len(data.shape)==4)  #4D arrays
	assert (data.shape[1]==1 or data.shape[1]==3)  #check the channel is 1 or 3
	height = data.shape[2]
	width = data.shape[3]
	for i in range(data.shape[0]):  #loop over the full images
		for x in range(width):
			for y in range(height):
				if inside_FOV_DRIVE(i,x,y,original_imgs_border_masks)==False:
					data[i,:,y,x]=0.0


# def inside_FOV_DRIVE(i, x, y, DRIVE_masks):

# 	return False # victor

# 	assert (len(DRIVE_masks.shape)==4)  #4D arrays
# 	# THEO assert (DRIVE_masks.shape[1]==1)  #DRIVE masks is black and white
# 	# DRIVE_masks = DRIVE_masks/255.  #NOOO!! otherwise with float numbers takes forever!!

# 	if (x >= DRIVE_masks.shape[3] or y >= DRIVE_masks.shape[2]): #my image bigger than the original
# 		return False

# 	if (DRIVE_masks[i,0,y,x]>0):  #0==black pixels
# 		# print DRIVE_masks[i,0,y,x]  #verify it is working right
# 		return True
# 	else:
# 		return False
