import numpy as np

import random

from glob import glob

import os

import SimpleITK as sitk

from evaluation_metrics import *

from model import Unet_model

class Prediction(object):

    def __init__(self, batch_size_test,load_model_path):

        self.batch_size_test=batch_size_test

        unet=Unet_model(img_shape=(240,240,4),load_model_weights=load_model_path)

        self.model=unet.model

        print ('U-net CNN compiled!\n')

    def predict_volume(self, filepath_image,show):

        '''

        segment the input volume

        INPUT   (1) str 'filepath_image': filepath of the volume to predict 

                (2) bool 'show': True to ,

        OUTPUt  (1) np array of the predicted volume

                (2) np array of the corresping ground truth

        '''

        #read the volume

        flair = glob( filepath_image + '/*_flair.nii.gz')

        t2 = glob( filepath_image + '/*_t2.nii.gz')

        gt = glob( filepath_image + '/*_seg.nii.gz')

        t1s = glob( filepath_image + '/*_t1.nii.gz')

        t1c = glob( filepath_image + '/*_t1ce.nii.gz')

        t1=[scan for scan in t1s if scan not in t1c]

        if (len(flair)+len(t2)+len(gt)+len(t1)+len(t1c))<5:

            print("there is a problem here!!! the problem lies in this patient :")

        scans_test = [flair[0], t1[0], t1c[0], t2[0], gt[0]]

        test_im = [sitk.GetArrayFromImage(sitk.ReadImage(scans_test[i])) for i in range(len(scans_test))]

        test_im=np.array(test_im).astype(np.float32)

        test_image = test_im[0:4]

        gt=test_im[-1]

        gt[gt==4]=3

        #normalize each slice following the same scheme used for training

        test_image=self.norm_slices(test_image)

        #transform teh data to channels_last keras format

        test_image = test_image.swapaxes(0,1)

        test_image=np.transpose(test_image,(0,2,3,1))

        if show:

            verbose=1

        else:

            verbose=0

        # predict classes of each pixel based on the model

        prediction = self.model.predict(test_image,batch_size=self.batch_size_test,verbose=verbose)   

        prediction = np.argmax(prediction, axis=-1)

        prediction=prediction.astype(np.uint8)

        #reconstruct the initial target values .i.e. 0,1,2,4 for prediction and ground truth

        prediction[prediction==3]=4

        gt[gt==3]=4

        return np.array(prediction),np.array(gt)

    def evaluate_segmented_volume(self, filepath_image,save,show,save_path):

        '''

        computes the evaluation metrics on the segmented volume

        INPUT   (1) str 'filepath_image': filepath to test image for segmentation, including file extension

                (2) bool 'save': whether to save to disk or not

                (3) bool 'show': If true, prints the evaluation metrics

        OUTPUT np array of all evaluation metrics

        '''

        predicted_images,gt= self.predict_volume(filepath_image,show)

        if save:

            tmp=sitk.GetImageFromArray(predicted_images)

            sitk.WriteImage ( tmp,'predictions/{}.nii.gz'.format(save_path) )

        #compute the evaluation metrics 

        Dice_complete=DSC_whole(predicted_images,gt)

        Dice_enhancing=DSC_en(predicted_images,gt)

        Dice_core=DSC_core(predicted_images,gt)

        Sensitivity_whole=sensitivity_whole(predicted_images,gt)

        Sensitivity_en=sensitivity_en(predicted_images,gt)

        Sensitivity_core=sensitivity_core(predicted_images,gt)

        Specificity_whole=specificity_whole(predicted_images,gt)

        Specificity_en=specificity_en(predicted_images,gt)

        Specificity_core=specificity_core(predicted_images,gt)

        Hausdorff_whole=hausdorff_whole(predicted_images,gt)

        Hausdorff_en=hausdorff_en(predicted_images,gt)

        Hausdorff_core=hausdorff_core(predicted_images,gt)

        if show:

            print("************************************************************")

            print("Dice complete tumor score : {:0.4f}".format(Dice_complete))

            print("Dice core tumor score (tt sauf vert): {:0.4f}".format(Dice_core))

            print("Dice enhancing tumor score (jaune):{:0.4f} ".format(Dice_enhancing))

            print("**********************************************")

            print("Sensitivity complete tumor score : {:0.4f}".format(Sensitivity_whole))

            print("Sensitivity core tumor score (tt sauf vert): {:0.4f}".format(Sensitivity_core))

            print("Sensitivity enhancing tumor score (jaune):{:0.4f} ".format(Sensitivity_en))

            print("***********************************************")

            print("Specificity complete tumor score : {:0.4f}".format(Specificity_whole))

            print("Specificity core tumor score (tt sauf vert): {:0.4f}".format(Specificity_core))

            print("Specificity enhancing tumor score (jaune):{:0.4f} ".format(Specificity_en))

            print("***********************************************")

            print("Hausdorff complete tumor score : {:0.4f}".format(Hausdorff_whole))

            print("Hausdorff core tumor score (tt sauf vert): {:0.4f}".format(Hausdorff_core))

            print("Hausdorff enhancing tumor score (jaune):{:0.4f} ".format(Hausdorff_en))

            print("***************************************************************\n\n")

        return np.array((Dice_complete,Dice_core,Dice_enhancing,Sensitivity_whole,Sensitivity_core,Sensitivity_en,Specificity_whole,Specificity_core,Specificity_en,Hausdorff_whole,Hausdorff_core,Hausdorff_en))#))

    def predict_multiple_volumes (self, filepath_volumes,save,show):

        results,Ids=[],[]

        for patient in filepath_volumes:

            tmp1=patient.split('/')

            print("Volume ID: " ,tmp1[-2]+'/'+tmp1[-1])

            tmp=self.evaluate_segmented_volume(patient,save=save,show=show,save_path=os.path.basename(patient))

            #save the results of each volume

            results.append(tmp)

            #save each ID for later use

            Ids.append(str(tmp1[-2]+'/'+tmp1[-1]))

        res=np.array(results)     

        print("mean : ",np.mean(res,axis=0))

        print("std : ",np.std(res,axis=0))

        print("median : ",np.median(res,axis=0))

        print("25 quantile : ",np.percentile(res,25,axis=0))

        print("75 quantile : ",np.percentile(res,75,axis=0))

        print("max : ",np.max(res,axis=0))

        print("min : ",np.min(res,axis=0))

        np.savetxt('Results.out', res)

        np.savetxt('Volumes_ID.out', Ids,fmt='%s')

    def norm_slices(self,slice_not):

        '''

            normalizes each slice, excluding gt

            subtracts mean and div by std dev for each slice

            clips top and bottom one percent of pixel intensities

        '''

        normed_slices = np.zeros(( 4,155, 240, 240))

        for slice_ix in range(4):

            normed_slices[slice_ix] = slice_not[slice_ix]

            for mode_ix in range(155):

                normed_slices[slice_ix][mode_ix] = self._normalize(slice_not[slice_ix][mode_ix])

        return normed_slices    

    def _normalize(self,slice):

        b = np.percentile(slice, 99)

        t = np.percentile(slice, 1)

        slice = np.clip(slice, t, b)

        image_nonzero = slice[np.nonzero(slice)]

        if np.std(slice)==0 or np.std(image_nonzero) == 0:

            return slice

        else:

            tmp= (slice - np.mean(image_nonzero)) / np.std(image_nonzero)

            tmp[tmp==tmp.min()]=-9

            return tmp

if __name__ == "__main__":

    #set arguments

    model_to_load="models_saved/ResUnet.04_0.646.hdf5" 

    #paths for the testing data

    path_HGG = glob('Brats2017/Brats17TrainingData/HGG/**')

    path_LGG = glob('Brats2017/Brats17TrainingData/LGG/**')

    test_path=path_HGG+path_LGG

    np.random.seed(2022)

    np.random.shuffle(test_path)

    #compile the model

    brain_seg_pred = Prediction(batch_size_test=2 ,load_model_path=model_to_load)

    #predicts each volume and save the results in np array

    brain_seg_pred.predict_multiple_volumes(test_path[200:290],save=False,show=True)