""" 

Copyright (c) 2016, Jose Dolz .All rights reserved.

Redistribution and use in source and binary forms, with or without modification,

are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice,

       this list of conditions and the following disclaimer.

    2. Redistributions in binary form must reproduce the above copyright notice,

       this list of conditions and the following disclaimer in the documentation

       and/or other materials provided with the distribution.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,

    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES

    OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND

    NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT

    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,

    WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING

    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR

    OTHER DEALINGS IN THE SOFTWARE.

Jose Dolz. Dec, 2016.

email: jose.dolz.upv@gmail.com

LIVIA Department, ETS, Montreal.

"""

import numpy as np

import time

import os

import pdb

from Modules.General.Evaluation import computeDice

from Modules.General.Utils import getImagesSet

from Modules.General.Utils import load_model_from_gzip_file

from Modules.IO.ImgOperations.imgOp import applyUnpadding

from Modules.IO.loadData import load_imagesSinglePatient

from Modules.IO.saveData import saveImageAsNifti

from Modules.IO.saveData import saveImageAsMatlab

from Modules.IO.sampling import *

from Modules.Parsers.parsersUtils import parserConfigIni

def segmentVolume(myNetworkModel,

                  i_d,

                  imageNames_Test,

                  names_Test,

                  groundTruthNames_Test,

                  roiNames_Test,

                  imageType,

                  padInputImagesBool,

                  receptiveField, 

                  sampleSize_Test,

                  strideVal,

                  batch_Size,

                  task # Validation (0) or testing (1)

                  ):

        # Get info from the network model        

        networkName        = myNetworkModel.networkName

        folderName         = myNetworkModel.folderName

        n_classes          = myNetworkModel.n_classes

        sampleSize_Test    = myNetworkModel.sampleSize_Test

        receptiveField     = myNetworkModel.receptiveField  

        outputShape        = myNetworkModel.lastLayer.outputShapeTest[2:] 

        batch_Size         = myNetworkModel.batch_Size

        padInputImagesBool = True

        # Get half sample size

        sampleHalf = []

        for h_i in range(3):

            sampleHalf.append((receptiveField[h_i]-1)/2)

        # Load the images to segment

        [imgSubject,  

        gtLabelsImage, 

        roi, 

        paddingValues] = load_imagesSinglePatient(i_d,

                                                  imageNames_Test,

                                                  groundTruthNames_Test,

                                                  roiNames_Test,

                                                  padInputImagesBool,

                                                  receptiveField, 

                                                  sampleSize_Test,

                                                  imageType, 

                                                  )

        # Get image dimensions                                                    

        imgDims = list(imgSubject.shape)

        [ sampleCoords ] = sampleWholeImage(imgSubject,

                                            roi,

                                            sampleSize_Test,

                                            strideVal,

                                            batch_Size

                                            )

        numberOfSamples = len(sampleCoords)

        sampleID = 0

        numberOfBatches = numberOfSamples/batch_Size

        #The probability-map that will be constructed by the predictions.

        probMaps = np.zeros([n_classes]+imgDims, dtype = "float32")

        # Run over all the batches 

        for b_i in xrange(numberOfBatches) :

            # Get samples for batch b_i

            sampleCoords_b = sampleCoords[ b_i*batch_Size : (b_i+1)*batch_Size ]

            [imgSamples] = extractSamples(imgSubject,

                                          sampleCoords_b,

                                          sampleSize_Test,

                                          receptiveField)

            # Load the data of the batch on the GPU

            myNetworkModel.testingData_x.set_value(imgSamples, borrow=True)

            # Call the testing Theano function            

            predictions = myNetworkModel.networkModel_Test(0)

            predOutput = predictions[-1]

            # --- Now we can generate the probability maps from the predictions ----

            # Run over all the regions

            for r_i in xrange(batch_Size) :

                sampleCoords_i = sampleCoords[sampleID]

                coords = [ sampleCoords_i[0][0], sampleCoords_i[1][0], sampleCoords_i[2][0] ]

                # Get the min and max coords

                xMin = coords[0] + sampleHalf[0]

                xMax = coords[0] + sampleHalf[0] + strideVal[0]

                yMin = coords[1] + sampleHalf[1]

                yMax = coords[1] + sampleHalf[1] + strideVal[1]

                zMin = coords[2] + sampleHalf[2]

                zMax = coords[2] + sampleHalf[2] + strideVal[2]

                probMaps[:,xMin:xMax, yMin:yMax, zMin:zMax] = predOutput[r_i]

                sampleID += 1

        # Release data

        myNetworkModel.testingData_x.set_value(np.zeros([1,1,1,1,1], dtype="float32"))

        # Segmentation has been done in this point.

        # Now: Save the data

        # Get the segmentation from the probability maps ---

        segmentationImage = np.argmax(probMaps, axis=0) 

        #Save Result:

        npDtypeForPredictedImage = np.dtype(np.int16)

        suffixToAdd = "_Segm"

        # Apply unpadding if specified

        if padInputImagesBool == True:

            segmentationRes = applyUnpadding(segmentationImage, paddingValues)

        else:

            segmentationRes = segmentationImage

        # Generate folders to store the model

        BASE_DIR = os.getcwd()

        path_Temp = os.path.join(BASE_DIR,'outputFiles')

        # For the predictions

        predlFolderName = os.path.join(path_Temp,myNetworkModel.folderName)

        predlFolderName = os.path.join(predlFolderName,'Pred')

        if task == 0:

            predTestFolderName = os.path.join(predlFolderName,'Validation')

        else:

            predTestFolderName = os.path.join(predlFolderName,'Testing')

        nameToSave = predTestFolderName + '/Segmentation_'+ names_Test[i_d]

        # Save Segmentation image

        print(" ... Saving segmentation result..."),

        if imageType == 0: # nifti

            imageTypeToSave = np.dtype(np.int16)

            saveImageAsNifti(segmentationRes,

                             nameToSave,

                             imageNames_Test[i_d],

                             imageTypeToSave)

        else: # Matlab

            # Cast to int8 for saving purposes

            saveImageAsMatlab(segmentationRes.astype('int8'),

                              nameToSave)

        # Save the prob maps for each class (except background)

        for c_i in xrange(1, n_classes) :

            nameToSave = predTestFolderName + '/ProbMap_class_'+ str(c_i) + '_' + names_Test[i_d] 

            probMapClass = probMaps[c_i,:,:,:]

            # Apply unpadding if specified

            if padInputImagesBool == True:

                probMapClassRes = applyUnpadding(probMapClass, paddingValues)

            else:

                probMapClassRes = probMapClass

            print(" ... Saving prob map for class {}...".format(str(c_i))),

            if imageType == 0: # nifti

                imageTypeToSave = np.dtype(np.float32)

                saveImageAsNifti(probMapClassRes,

                                 nameToSave,

                                 imageNames_Test[i_d],

                                 imageTypeToSave)

            else:

                # Cast to float32 for saving purposes

                saveImageAsMatlab(probMapClassRes.astype('float32'),

                                  nameToSave)

        # If segmentation done during evaluation, get dice

        if task == 0:

            print(" ... Computing Dice scores: ")

            DiceArray = computeDice(segmentationImage,gtLabelsImage)

            for d_i in xrange(len(DiceArray)):

                print(" -------------- DSC (Class {}) : {}".format(str(d_i+1),DiceArray[d_i]))

""" Main segmentation function """

def startTesting(networkModelName,

                 configIniName

                 ) :

    padInputImagesBool = True # from config ini

    print " ******************************************  STARTING SEGMENTATION ******************************************"

    print " **********************  Starting segmentation **********************"

    myParserConfigIni = parserConfigIni()

    myParserConfigIni.readConfigIniFile(configIniName,2)

    print " -------- Images to segment -------------"

    print " -------- Reading Images names for segmentation -------------"

    # -- Get list of images used for testing -- #

    (imageNames_Test, names_Test) = getImagesSet(myParserConfigIni.imagesFolder,myParserConfigIni.indexesToSegment)  # Images

    (groundTruthNames_Test, gt_names_Test) = getImagesSet(myParserConfigIni.GroundTruthFolder,myParserConfigIni.indexesToSegment) # Ground truth

    (roiNames_Test, roi_names_Test) = getImagesSet(myParserConfigIni.ROIFolder,myParserConfigIni.indexesToSegment) # ROI

    # --------------- Load my LiviaNet3D object  --------------- 

    print (" ... Loading model from {}".format(networkModelName))

    myLiviaNet3D = load_model_from_gzip_file(networkModelName)

    print " ... Network architecture successfully loaded...."

    # Get info from the network model        

    networkName        = myLiviaNet3D.networkName

    folderName         = myLiviaNet3D.folderName

    n_classes          = myLiviaNet3D.n_classes

    sampleSize_Test    = myLiviaNet3D.sampleSize_Test

    receptiveField     = myLiviaNet3D.receptiveField  

    outputShape        = myLiviaNet3D.lastLayer.outputShapeTest[2:] 

    batch_Size         = myLiviaNet3D.batch_Size

    padInputImagesBool = myParserConfigIni.applyPadding

    imageType          = myParserConfigIni.imageTypes

    numberImagesToSegment = len(imageNames_Test)

    strideValues = myLiviaNet3D.lastLayer.outputShapeTest[2:]

    # Run over the images to segment   

    for i_d in xrange(numberImagesToSegment) :

        print("**********************  Segmenting subject: {} ....total: {}/{}...**********************".format(names_Test[i_d],str(i_d+1),str(numberImagesToSegment)))

        segmentVolume(myLiviaNet3D,

                  i_d,

                  imageNames_Test,  # Full path

                  names_Test,       # Only image name

                  groundTruthNames_Test,

                  roiNames_Test,

                  imageType,

                  padInputImagesBool,

                  receptiveField, 

                  sampleSize_Test,

                  strideValues,

                  batch_Size,

                  1 # Validation (0) or testing (1)

                  )

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