# -*- coding: utf-8 -*-
# VYHODNOCENI DLE MRBRAINS18
# AUTHOR: MRBRAINS18
# EDIT: XSLUS01, VUT FEKT
import difflib
import numpy as np
import os
import SimpleITK as sitk
import scipy.spatial
from tools import *
import sys
import pandas as pd


	
labels = {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',
		  }
labels_short = {
		  1: 'CGM',
		  2: 'BG',
		  3: 'WM',
		  4: 'WML',
		  5: 'CFS',
		  6: 'V',
		  7: 'C',
		  8: 'BS',
		  # The two labels below are ignored:
		  #9: 'Infarction',
		  #10: 'Other',
		  }



def dictToLabel(index):
	return labels_short[index]



	
def getResultFilename(participantDir):
	"""Find the filename of the result image.
	
	This should be result.nii.gz or result.nii. If these files are not present,
	it tries to find the closest filename."""
	files = os.listdir(participantDir)
	
	if not files:
		raise Exception("No results in "+ participantDir)
	
	resultFilename = None
	if 'result.nii.gz' in files:
		resultFilename = os.path.join(participantDir, 'result.nii.gz')
	elif 'result.nii' in files:
		resultFilename = os.path.join(participantDir, 'result.nii')
	else:
		# Find the filename that is closest to 'result.nii.gz'
		maxRatio = -1
		for f in files:
			currentRatio = difflib.SequenceMatcher(a = f, b = 'result.nii.gz').ratio()
			
			if currentRatio > maxRatio:
				resultFilename = os.path.join(participantDir, f)
				maxRatio = currentRatio
				
	return resultFilename
	

def getImages(testFilename, resultFilename):
	"""Return the test and result images, thresholded and pathology masked."""
	testImage   = sitk.ReadImage(testFilename)
	resultImage = sitk.ReadImage(resultFilename)
	
	# Check for equality
	assert testImage.GetSize() == resultImage.GetSize()
	
	# Get meta data from the test-image, needed for some sitk methods that check this
	resultImage.CopyInformation(testImage)    
	
	# Remove pathology from the test and result images, since we don't evaluate on that
	pathologyImage = sitk.BinaryThreshold(testImage, 9, 11, 0, 1)  # pathology == 9 or 10
	
	maskedTestImage   = sitk.Mask(testImage,   pathologyImage)     # tissue    == 1 --  8
	maskedResultImage = sitk.Mask(resultImage, pathologyImage)
	
	# Force integer
	if not 'integer' in maskedResultImage.GetPixelIDTypeAsString():
		maskedResultImage = sitk.Cast(maskedResultImage, sitk.sitkUInt8)
			
	return maskedTestImage, maskedResultImage
	
	
def getDSC(testImage, resultImage):    
	"""Compute the Dice Similarity Coefficient."""        
	dsc = dict()
	for k in labels.keys():
		testArray   = sitk.GetArrayFromImage(sitk.BinaryThreshold(  testImage, k, k, 1, 0)).flatten()
		resultArray = sitk.GetArrayFromImage(sitk.BinaryThreshold(resultImage, k, k, 1, 0)).flatten()
		
		# similarity = 1.0 - dissimilarity
		# scipy.spatial.distance.dice raises a ZeroDivisionError if both arrays contain only zeros.
		try:
			dsc[k] = 1.0 - scipy.spatial.distance.dice(testArray, resultArray)
		except ZeroDivisionError:
			dsc[k] = None
	
	return dsc

		
def getHausdorff(testImage, resultImage):
	"""Compute the 95% Hausdorff distance."""    
	hd = dict()
	for k in labels.keys():
		lTestImage   = sitk.BinaryThreshold(  testImage, k, k, 1, 0)
		lResultImage = sitk.BinaryThreshold(resultImage, k, k, 1, 0)
		
		# Hausdorff distance is only defined when something is detected
		statistics = sitk.StatisticsImageFilter()
		statistics.Execute(lTestImage)
		lTestSum = statistics.GetSum()
		statistics.Execute(lResultImage)
		lResultSum = statistics.GetSum()
		if lTestSum == 0 or lResultSum == 0:
			hd[k] = None
			continue
								
		# Edge detection is done by ORIGINAL - ERODED, keeping the outer boundaries of lesions. Erosion is performed in 2D
		eTestImage   = sitk.BinaryErode(lTestImage, (1,1,0))
		eResultImage = sitk.BinaryErode(lResultImage, (1,1,0))
		
		hTestImage   = sitk.Subtract(lTestImage, eTestImage)
		hResultImage = sitk.Subtract(lResultImage, eResultImage)    
		
		hTestArray   = sitk.GetArrayFromImage(hTestImage)
		hResultArray = sitk.GetArrayFromImage(hResultImage)   
			
		# Convert voxel location to world coordinates. Use the coordinate system of the test image
		# np.nonzero   = elements of the boundary in numpy order (zyx)
		# np.flipud    = elements in xyz order
		# np.transpose = create tuples (x,y,z)
		# testImage.TransformIndexToPhysicalPoint converts (xyz) to world coordinates (in mm)
		# (Simple)ITK does not accept all Numpy arrays; therefore we need to convert the coordinate tuples into a Python list before passing them to TransformIndexToPhysicalPoint().
		testCoordinates   = [testImage.TransformIndexToPhysicalPoint(x.tolist()) for x in np.transpose( np.flipud( np.nonzero(hTestArray) ))]
		resultCoordinates = [testImage.TransformIndexToPhysicalPoint(x.tolist()) for x in np.transpose( np.flipud( np.nonzero(hResultArray) ))]
				
		# Use a kd-tree for fast spatial search
		def getDistancesFromAtoB(a, b):    
			kdTree = scipy.spatial.KDTree(a, leafsize=100)
			return kdTree.query(b, k=1, eps=0, p=2)[0]
		
		# Compute distances from test to result and vice versa. 
		dTestToResult = getDistancesFromAtoB(testCoordinates, resultCoordinates)
		dResultToTest = getDistancesFromAtoB(resultCoordinates, testCoordinates)
		hd[k] = max(np.percentile(dTestToResult, 95), np.percentile(dResultToTest, 95))
		
	return hd


def getVS(testImage, resultImage):   
	"""Volume similarity.
	
	VS = 1 - abs(A - B) / (A + B)
	
	A = ground truth in ML
	B = participant segmentation in ML
	"""    
	# Compute statistics of both images
	testStatistics   = sitk.StatisticsImageFilter()
	resultStatistics = sitk.StatisticsImageFilter()
	
	vs = dict()
	for k in labels.keys():
		testStatistics.Execute(sitk.BinaryThreshold(testImage, k, k, 1, 0))
		resultStatistics.Execute(sitk.BinaryThreshold(resultImage, k, k, 1, 0))
		
		numerator = abs(testStatistics.GetSum() - resultStatistics.GetSum())
		denominator = testStatistics.GetSum() + resultStatistics.GetSum()               
		
		if denominator > 0:        
			vs[k] = 1 - float(numerator) / denominator
		else:
			vs[k] = None
		
	return vs
	

def do(mask_src_file, predicted_image):
	"""Main function"""    
	#resultFilename = getResultFilename(participantDir)  
		
	testImage, resultImage = getImages(mask_src_file, predicted_image)
	
	dsc = getDSC(testImage, resultImage)
	h95 = getHausdorff(testImage, resultImage)
	vs  = getVS(testImage, resultImage)
	
	print('Dice',                dsc,       '(higher is better, max=1)')
	p('-----------')
	print('HD',                  h95, 'mm',  '(lower is better, min=0)')
	p('-----------')
	print('VS',                   vs,       '(higher is better, max=1)')

	return dsc, h95, vs





def eval_fragments_result():
	# src predicts
	src_img_3_AUG_ROT2_070  	= os.path.join(os.getcwd() + '/evals/img_3_AUG_ROT2_070.nii.gz')
	src_img_5_AUG_ROT2_1  		= os.path.join(os.getcwd() + '/evals/img_5_AUG_ROT2_1.nii.gz')
	src_img_21_AUG_ROT5_070  	= os.path.join(os.getcwd() + '/evals/img_21_AUG_ROT5_070.nii.gz')
	src_img_24_AUG_ROT5_4  		= os.path.join(os.getcwd() + '/evals/img_24_AUG_ROT5_4.nii.gz')
	src_img_40_AUG_ZOOM3_148  	= os.path.join(os.getcwd() + '/evals/img_40_AUG_ZOOM3_148.nii.gz')

	tmp_dest_prediction_norm = os.path.join(os.getcwd() + '/builded_image/prediction_normalized.nii.gz')
	tmp_dest_prediction = os.path.join(os.getcwd() + '/builded_image/prediction.nii.gz')
	tmp_dest_mask = os.path.join(os.getcwd() + '/builded_image/mask.nii.gz')

	# src masks
	src_mask_1_segm  	= os.path.join(os.getcwd() + '/evals/mask_1_segm.nii.gz')
	src_mask_4_segm  	= os.path.join(os.getcwd() + '/evals/mask_4_segm.nii.gz')
	src_mask_070_segm  	= os.path.join(os.getcwd() + '/evals/mask_070_segm.nii.gz')
	src_mask_148_segm  	= os.path.join(os.getcwd() + '/evals/mask_148_segm.nii.gz')


	img_3_AUG_ROT2_070 = get_3D_nii_image(src_img_3_AUG_ROT2_070,None, sitk.sitkUInt8)
	img_3_AUG_ROT2_070_mask = image_3d_rotate(get_3D_nii_image(src_mask_070_segm,None, sitk.sitkUInt8), -7)
	img_3_AUG_ROT2_070 = remove_background_by_mask(img_3_AUG_ROT2_070, img_3_AUG_ROT2_070_mask)
	
	# SNAHA a zvirazneni
	tmp = img_3_AUG_ROT2_070
	n=40
	#↓tmp[tmp == 9] = 6*n 
	tmp[tmp == 8] = 6*n  
	tmp[tmp == 7] = 4*n  
	tmp[tmp == 6] = 4*n  
	tmp[tmp == 5] = 3*n  
	tmp[tmp == 4] = 2*n  
	tmp[tmp == 3] = 1*n  
	tmp[tmp == 2] = 0.5*n  
	tmp[tmp == 1] =  10

	Image, Origin, Spacing = get_3D_nii_image_full(src_mask_070_segm)
	save_3D_nii_image_with_spacing(tmp, Origin, Spacing, tmp_dest_prediction_norm)	# ukladam predikci
	save_3D_nii_image(img_3_AUG_ROT2_070, tmp_dest_prediction, False)	# ukladam predikci
	save_3D_nii_image(img_3_AUG_ROT2_070_mask, tmp_dest_mask, False)	# ukladam masku (otocenou protoze predikce je taky)

	mask = tmp_dest_mask
	prediction = tmp_dest_prediction

	if not os.path.exists(os.path.join(os.getcwd() + '/builded_image/dsc.npy')):
		dsc, h95, vs = do(mask, prediction)
		np.save(os.path.join(os.getcwd() + '/builded_image/dsc.npy'), dsc )
		np.save(os.path.join(os.getcwd() + '/builded_image/h95.npy'), h95)
		np.save(os.path.join(os.getcwd() + '/builded_image/vs.npy'), vs) 
	else:
		dsc=np.load(os.path.join(os.getcwd() + '/builded_image/dsc.npy'))
		h95=np.load(os.path.join(os.getcwd() + '/builded_image/h95.npy'))
		vs=np.load(os.path.join(os.getcwd() + '/builded_image/vs.npy'))
		
		dsc= dsc.item()
		h95= h95.item()
		vs= vs.item()

	CGM = []
	BG = []
	WM = []
	WML = []
	CFS = []
	V = []
	C = []
	BS = []

	CGM.append(round(dsc[1], 2))
	CGM.append(round(h95[1], 2))
	CGM.append(round(vs[1], 2))

	BG.append(round(dsc[2], 2))
	BG.append(round(h95[2], 2))
	BG.append(round(vs[2], 2))

	WM.append(round(dsc[3], 2))
	WM.append(round(h95[3], 2))
	WM.append(round(vs[3], 2))

	WML.append(round(dsc[4], 2))
	WML.append(round(h95[4], 2))
	WML.append(round(vs[4], 2))

	CFS.append(round(dsc[5], 2))
	CFS.append(round(h95[5], 2))
	CFS.append(round(vs[5], 2))

	V.append(round(dsc[6], 2))
	V.append(round(h95[6], 2))
	V.append(round(vs[6], 2))
	
	C.append(round(dsc[7], 2))
	C.append(round(h95[7], 2))
	C.append(round(vs[7], 2))

	BS.append(round(dsc[8], 2))
	BS.append(round(h95[8], 2))
	BS.append(round(vs[8], 2))

	p(CGM)
	p(BG)
	p(WM)
	p(WML)
	p(CFS)
	p(V)
	p(C)
	p(BS)
	df = pd.DataFrame({
						'#': ['DSC', 'H', 'VS'],
						'CGM': CGM,
						'BG': BG,
						'WM': WM,
						'WML': WML,
						'CFS': CFS,
						'V': V,
						'C': C,
						'BS': BS,
					})

	with open('doc/mrbrains_eval.tex','w') as tf:
		tf.write(df.to_latex(bold_rows = False, index=False, column_format='|c|c|c|c|c|c|c|c|c|'))

	return True


def eval_fit_result():
	# src predicts
	patch_65_4  	= os.path.join(os.getcwd() + '/evals/fit/patch_65_4.nii.gz')

	tmp_dest_prediction_norm = os.path.join(os.getcwd() + '/builded_image/fit/prediction_normalized.nii.gz')
	tmp_dest_prediction = os.path.join(os.getcwd() + '/builded_image/fit/prediction.nii.gz')
	tmp_dest_mask = os.path.join(os.getcwd() + '/builded_image/fit/mask.nii.gz')

	# src masks
	src_mask_1_segm  	= os.path.join(os.getcwd() + '/evals/fit/segm.nii.gz')

	patch_65_4 = get_3D_nii_image(patch_65_4,None, sitk.sitkUInt8)
	#patch_65_4_mask = get_3D_nii_image(src_mask_1_segm,None, sitk.sitkUInt8)
	patch_65_4_mask, Origin, Spacing = get_3D_nii_image_full(src_mask_1_segm)
	patch_65_4 = remove_background_by_mask(patch_65_4, patch_65_4_mask)
	
	tttttmp = patch_65_4
	patch_65_4 = (patch_65_4/255) * 9
	save_3D_nii_image(patch_65_4, tmp_dest_prediction, False)
	patch_65_4 = get_3D_nii_image(tmp_dest_prediction,None, sitk.sitkUInt8)
	
	save_3D_nii_image_with_spacing(tttttmp, Origin, Spacing, tmp_dest_prediction_norm)	# ukladam predikci
	save_3D_nii_image_with_spacing(patch_65_4, Origin, Spacing, tmp_dest_prediction)	# ukladam predikci

	mask =src_mask_1_segm
	prediction = tmp_dest_prediction

	if not os.path.exists(os.path.join(os.getcwd() + '/builded_image/fit/dsc.npy')):
		dsc, h95, vs = do(mask, prediction)
		np.save(os.path.join(os.getcwd() + '/builded_image/fit/dsc.npy'), dsc )
		np.save(os.path.join(os.getcwd() + '/builded_image/fit/h95.npy'), h95)
		np.save(os.path.join(os.getcwd() + '/builded_image/fit/vs.npy'), vs) 
	else:
		dsc=np.load(os.path.join(os.getcwd() + '/builded_image/fit/dsc.npy'))
		h95=np.load(os.path.join(os.getcwd() + '/builded_image/fit/h95.npy'))
		vs=np.load(os.path.join(os.getcwd() + '/builded_image/fit/vs.npy'))
		
		dsc= dsc.item()
		h95= h95.item()
		vs= vs.item()

	CGM = []
	BG = []
	WM = []
	WML = []
	CFS = []
	V = []
	C = []
	BS = []

	CGM.append(round(dsc[1], 2))
	CGM.append(round(h95[1], 2))
	CGM.append(round(vs[1], 2))

	BG.append(round(dsc[2], 2))
	BG.append(round(h95[2], 2))
	BG.append(round(vs[2], 2))

	WM.append(round(dsc[3], 2))
	WM.append(round(h95[3], 2))
	WM.append(round(vs[3], 2))

	WML.append(round(dsc[4], 2))
	WML.append(round(h95[4], 2))
	WML.append(round(vs[4], 2))

	CFS.append(round(dsc[5], 2))
	CFS.append(round(h95[5], 2))
	CFS.append(round(vs[5], 2))

	V.append(round(dsc[6], 2))
	V.append(round(h95[6], 2))
	V.append(round(vs[6], 2))
	
	C.append(round(dsc[7], 2))
	C.append(round(h95[7], 2))
	C.append(round(vs[7], 2))

	BS.append(round(dsc[8], 2))
	BS.append((round(h95[8], 2) if h95[8] is not None else '--'))
	BS.append(round(vs[8], 2))

	p(CGM)
	p(BG)
	p(WM)
	p(WML)
	p(CFS)
	p(V)
	p(C)
	p(BS)

	df = pd.DataFrame({
						'#': ['DSC', 'H', 'VS'],
						'CGM': CGM,
						'BG': BG,
						'WM': WM,
						'WML': WML,
						'CFS': CFS,
						'V': V,
						'C': C,
						'BS': BS,
					})

	with open('doc/mrbrains_eval_fit.tex','w') as tf:
		tf.write(df.to_latex(bold_rows = False, index=False, column_format='|c|c|c|c|c|c|c|c|c|'))

	return True








if __name__ == "__main__":
	try:
		eval_fit_result()
		eval_fragments_result()
	except:
		print("Doslo k chybe")

