import pickle

import samseg as gems
from samseg.SamsegUtility import *
from samseg.utilities import requireNumpyArray
from samseg.figures import initVisualizer
from samseg.Affine import Affine
from samseg.ProbabilisticAtlas import ProbabilisticAtlas
from samseg.Samseg import Samseg

"""
Longitudinal version of samsegment
The idea is based on the generative model in the paper

  Iglesias, Juan Eugenio, et al.
  Bayesian longitudinal segmentation of hippocampal substructures in brain MRI using subject-specific atlases.
  Neuroimage 141 (2016): 542-555,
in which a subject-specific atlas is obtained by generating a random warp from the usual population atlas, and
subsequently each time point is again randomly warped from this subject-specific atlas. The intermediate
subject-specific atlas is effectively a latent variable in the model, and it's function is to encourage the
various time points to have atlas warps that are similar between themselves, without having to define a priori
what these warps should be similar to. In the implementation provided here, the stiffness of the first warp
(denoted by K0, as in the paper) is taken as the stiffness used in ordinary samseg, and the stiffness of the
second warp (denoted by K1) is controlled by the setting

  strengthOfLatentDeformationHyperprior

so that K1 = strengthOfLatentDeformationHyperprior * K0. In the Iglesias paper, the setting

  strengthOfLatentDeformationHyperprior = 1.0

was used.

The overall idea is extended here by adding also latent variables encouraging corresponding Gaussian mixture
models across time points to be similar across time -- again without having to define a priori what exactly
they should look like. For given values of these latent variables, they effectively act has hyperparameters
on the mixture model parameters, the strength of which is controlled through the setting

  strengthOfLatentGMMHyperprior

which weighs the relative strength of this hyperprior relative to the relevant (expected) data term in the
estimation procedures of the mixture model parameters at each time point. This aspect can be switched off by
setting

  strengthOfLatentGMMHyperprior = 0.0

NOTE: The general longitudinal pipeline in FreeSurfer 6.0 is described in the paper

  Reuter, Martin, et al.
  Within-subject template estimation for unbiased longitudinal image analysis.
  Neuroimage 61.4 (2012): 1402-1418,

which is based on the idea of retaining some temporal consistency by simply initializing the model fitting
procedure across all time points in exactly the same way. This is achieved by first creating a subject-specific
template that is subsequently analyzed and the result of which is then used as a (very good) initialization
in each time point. This behavior can be mimicked by setting

  initializeLatentDeformationToZero = True
  numberOfIterations = 1
  strengthOfLatentGMMHyperprior = 0.0
  strengthOfLatentDeformationHyperprior = 1.0

in the implementation provided here.
"""

eps = np.finfo( float ).eps


class SamsegLongitudinal:
    def __init__(self,
        imageFileNamesList,
        atlasDir,
        savePath,
        userModelSpecifications={},
        userOptimizationOptions={},
        visualizer=None, 
        saveHistory=False,
        saveMesh=None,
        targetIntensity=None,
        targetSearchStrings=None,
        numberOfIterations=5,
        strengthOfLatentGMMHyperprior=0.5,
        strengthOfLatentDeformationHyperprior=20.0,
        saveSSTResults=True,
        updateLatentMeans=True,
        updateLatentVariances=True,
        updateLatentMixtureWeights=True,
        updateLatentDeformation=True,
        initializeLatentDeformationToZero=False,
        imageToImageTransformMatrix=None,
        threshold=None,
        thresholdSearchString=None,
        modeNames=None,
        pallidumAsWM=True,
        saveModelProbabilities=False,
        savePosteriors=False,
        tpToBaseTransforms=None,
        ):

        # Store input parameters as class variables
        self.imageFileNamesList = imageFileNamesList
        self.numberOfTimepoints = len(self.imageFileNamesList)
        self.savePath = savePath
        self.atlasDir = atlasDir
        self.threshold = threshold
        self.thresholdSearchString = thresholdSearchString
        self.targetIntensity = targetIntensity
        self.targetSearchStrings = targetSearchStrings
        self.numberOfIterations = numberOfIterations
        self.strengthOfLatentGMMHyperprior = strengthOfLatentGMMHyperprior
        self.strengthOfLatentDeformationHyperprior = strengthOfLatentDeformationHyperprior
        self.modeNames = modeNames
        self.pallidumAsWM = pallidumAsWM
        self.savePosteriors = savePosteriors
        self.tpToBaseTransforms = tpToBaseTransforms
        self.saveModelProbabilities = saveModelProbabilities

        # Check if all time point to base transforms are identity matrices.
        # If so, we can derive a combined 4D mask during preprocessing
        self.allIdentityTransforms = True
        if tpToBaseTransforms is not None:
            for tp, transform in enumerate(self.tpToBaseTransforms):
                if not np.allclose(transform.matrix, np.eye(4)):
                    self.allIdentityTransforms = False
        else:
            self.tpToBaseTransforms = []
            # If the trasnform are None, set them to identities
            for tp in range(self.numberOfTimepoints):
                 self.tpToBaseTransforms.append(np.eye(4))

        # Set image-to-image matrix if provided
        self.imageToImageTransformMatrix = imageToImageTransformMatrix

        # Initialize some objects
        self.probabilisticAtlas = ProbabilisticAtlas()

        # Get full model specifications and optimization options (using default unless overridden by user)
        self.userModelSpecifications = userModelSpecifications
        self.userOptimizationOptions = userOptimizationOptions

        # Setup a null visualizer if necessary
        if visualizer is None:
            self.visualizer = initVisualizer(False, False)
        else:
            self.visualizer = visualizer

        self.saveModelProbabilities = saveModelProbabilities
        self.saveHistory = saveHistory
        self.saveMesh = saveMesh
        self.saveSSTResults = saveSSTResults
        self.updateLatentMeans = updateLatentMeans
        self.updateLatentVariances = updateLatentVariances
        self.updateLatentMixtureWeights = updateLatentMixtureWeights
        self.updateLatentDeformation = updateLatentDeformation
        self.initializeLatentDeformationToZero = initializeLatentDeformationToZero

        # Make sure we can write in the target/results directory
        os.makedirs(savePath, exist_ok=True)

        # Here some class variables that will be defined later
        self.sstModel = None
        self.timepointModels = None
        self.imageBuffersList = None
        self.sstFileNames = None
        self.combinedImageBuffers = None
        self.latentDeformation = None
        self.latentDeformationAtlasFileName = None
        self.latentMeans = None
        self.latentVariances = None
        self.latentMixtureWeights = None
        self.latentMeansNumberOfMeasurements = None
        self.latentVariancesNumberOfMeasurements = None
        self.latentMixtureWeightsNumberOfMeasurements = None
        self.timepointVolumesInCubicMm = []
        self.optimizationSummary = None
        self.history = None
        self.historyOfTotalCost = None
        self.historyOfTotalTimepointCost = None
        self.historyOfLatentAtlasCost = None

    def segment(self, saveWarp=False, initTransformFile=None):
        # =======================================================================================
        #
        # Main function that runs the whole longitudinal segmentation pipeline
        #
        # =======================================================================================
        self.constructAndRegisterSubjectSpecificTemplate(initTransformFile)
        self.preProcess()
        self.fitModel()
        return self.postProcess(saveWarp=saveWarp)

    def constructAndRegisterSubjectSpecificTemplate(self, initTransformFile=None):
        # =======================================================================================
        #
        # Construction and affine registration of subject-specific template (sst)
        #
        # =======================================================================================

        # Initialization transform for registration
        initTransform = None
        if initTransformFile:
            trg = self.validateTransform(sf.load_affine(initTransformFile))
            initTransform = convertRASTransformToLPS(trg.convert(space='world').matrix)

        # Generate the subject specific template (sst)
        self.sstFileNames = self.generateSubjectSpecificTemplate()
        sstDir, _ = os.path.split(self.sstFileNames[0])

        if self.imageToImageTransformMatrix is None:

            # Affine atlas registration to sst
            templateFileName = os.path.join(self.atlasDir, 'template.nii.gz')
            affineRegistrationMeshCollectionFileName = os.path.join(self.atlasDir, 'atlasForAffineRegistration.txt.gz')

            affine = Affine(imageFileName=self.sstFileNames[0],
                             meshCollectionFileName=affineRegistrationMeshCollectionFileName,
                             templateFileName=templateFileName)
            self.imageToImageTransformMatrix, _ = affine.registerAtlas(savePath=sstDir, visualizer=self.visualizer, initTransform=initTransform)


    def preProcess(self):

        # construct sstModel
        self.constructSstModel()

        # =======================================================================================
        #
        # Preprocessing (reading and masking of data)
        #
        # =======================================================================================

        templateFileName = os.path.join(self.atlasDir, 'template.nii.gz')
        self.sstModel.imageBuffers, self.sstModel.transform, self.sstModel.voxelSpacing, self.sstModel.cropping = readCroppedImages(self.sstFileNames, templateFileName, self.imageToImageTransformMatrix)

        self.imageBuffersList = []
        self.voxelSpacings = []
        self.transforms = []
        self.masks = []
        self.croppings = []
        if self.allIdentityTransforms:

            self.imageBuffersList = []
            for imageFileNames in self.imageFileNamesList:
                imageBuffers, _, _, _ = readCroppedImages(imageFileNames, templateFileName,
                                                          self.imageToImageTransformMatrix)
                self.imageBuffersList.append(imageBuffers)

            # Put everything in a big 4-D matrix to derive one consistent mask across all time points
            imageSize = self.sstModel.imageBuffers.shape[:3]
            numberOfContrasts = self.sstModel.imageBuffers.shape[-1]
            self.combinedImageBuffers = np.zeros(imageSize + (numberOfContrasts * (1 + self.numberOfTimepoints),))
            self.combinedImageBuffers[..., 0:numberOfContrasts] = self.sstModel.imageBuffers
            for timepointNumber in range(self.numberOfTimepoints):
                self.combinedImageBuffers[..., (timepointNumber + 1) * numberOfContrasts:
                                               (timepointNumber + 2) * numberOfContrasts] = self.imageBuffersList[
                    timepointNumber]

            self.combinedImageBuffers, self.sstModel.mask = maskOutBackground(self.combinedImageBuffers,
                                                                              self.sstModel.modelSpecifications.atlasFileName,
                                                                              self.sstModel.transform,
                                                                              self.sstModel.modelSpecifications.maskingProbabilityThreshold,
                                                                              self.sstModel.modelSpecifications.maskingDistance,
                                                                              self.probabilisticAtlas,
                                                                              self.sstModel.voxelSpacing)
            combinedImageBuffers = logTransform(self.combinedImageBuffers, self.sstModel.mask)

            # Retrieve the masked sst and time points
            self.sstModel.imageBuffers = combinedImageBuffers[..., 0:numberOfContrasts]
            for timepointNumber in range(self.numberOfTimepoints):
                self.imageBuffersList[timepointNumber] = combinedImageBuffers[...,
                                                         (timepointNumber + 1) * numberOfContrasts:
                                                         (timepointNumber + 2) * numberOfContrasts]
                self.masks.append(self.sstModel.mask)
                self.croppings.append(self.sstModel.cropping)
                self.voxelSpacings.append(self.sstModel.voxelSpacing)
                self.transforms.append(self.sstModel.transform)

        else:

            for timepointNumber, imageFileNames in enumerate(self.imageFileNamesList):

                # Compute transformation from population atlas (p) to time point (tp), passing through the template space (s)
                # The transformation needs to be in vox to vox space as self.imageToImageTransformMatrix
                # We need to concatenate the following transformations
                # self.imageToImageTransformMatrix -> population to template space - vox to vox transform
                # tmp_s.geom.vox2world -> template space - vox to world transform
                # self.tpToBaseTransforms[timepointNumber].inv() -> template to time point space - world to world transform
                # tmp_tp.geom.world2vox -> time point space - world to vox transform
                tmpTp = sf.load_volume(imageFileNames[0])
                tmpS = sf.load_volume(os.path.join(self.savePath, "base", "template_coregistered.mgz"))
                pToTpTransform = tmpTp.geom.world2vox @ self.tpToBaseTransforms[timepointNumber].inv() @ tmpS.geom.vox2world @ self.imageToImageTransformMatrix

                imageBuffers, transform, voxelSpacing, cropping = readCroppedImages(imageFileNames, templateFileName, pToTpTransform.matrix)

                #
                self.imageBuffersList.append(imageBuffers)
                self.voxelSpacings.append(voxelSpacing)
                self.transforms.append(transform)
                self.croppings.append(cropping)

            # Derive mask for sst model
            imageBuffer, self.sstModel.mask = maskOutBackground(self.sstModel.imageBuffers, self.sstModel.modelSpecifications.atlasFileName,
                                                                self.sstModel.transform,
                                                                self.sstModel.modelSpecifications.maskingProbabilityThreshold,
                                                                self.sstModel.modelSpecifications.maskingDistance,
                                                                self.probabilisticAtlas,
                                                                self.sstModel.voxelSpacing)
            self.sstModel.imageBuffers = logTransform(imageBuffer, self.sstModel.mask)

            # Derive one mask for each time point model
            for timepointNumber in range(self.numberOfTimepoints):
                imageBuffer, timepointMask = maskOutBackground(self.imageBuffersList[timepointNumber],
                                                               self.sstModel.modelSpecifications.atlasFileName,
                                                               self.transforms[timepointNumber],
                                                               self.sstModel.modelSpecifications.maskingProbabilityThreshold,
                                                               self.sstModel.modelSpecifications.maskingDistance,
                                                               self.probabilisticAtlas,
                                                               self.voxelSpacings[timepointNumber])
                imageBuffer = logTransform(imageBuffer, timepointMask)
                self.imageBuffersList[timepointNumber] = imageBuffer
                self.masks.append(timepointMask)

        # construct timepoint models
        self.constructTimepointModels()

        self.visualizer.show(images=self.sstModel.imageBuffers, title='sst')
        for timepointNumber in range(self.numberOfTimepoints):
            self.visualizer.show(images=self.imageBuffersList[timepointNumber], title='time point ' + str(timepointNumber))

    def fitModel(self):

        # =======================================================================================
        #
        # Parameter estimation for SST
        #
        # =======================================================================================

        self.sstModel.fitModel()

        if hasattr(self.visualizer, 'show_flag'):
            import matplotlib.pyplot as plt  # avoid importing matplotlib by default
            plt.ion()
            self.sstModel.biasField.downSampleBasisFunctions([1, 1, 1])
            sstBiasFields = self.sstModel.biasField.getBiasFields( self.sstModel.mask)
            sstData = self.sstModel.imageBuffers[self.sstModel.mask, :] - sstBiasFields[self.sstModel.mask, :]
            axsList = []
            for contrastNumber in range(self.sstModel.gmm.numberOfContrasts):
                f = plt.figure()
                numberOfAxes = 2 + self.numberOfTimepoints
                numberOfRows = np.int(np.ceil(np.sqrt(numberOfAxes)))
                numberOfColumns = np.int(np.ceil(numberOfAxes / numberOfRows))
                axs = f.subplots(numberOfRows, numberOfColumns, sharex=True)
                ax = axs.ravel()[0]
                _, bins, _ = ax.hist(self.sstModel.imageBuffers[self.sstModel.mask, contrastNumber], 100)
                ax.grid()
                ax.set_title('sst before bias field correction')
                ax = axs.ravel()[1]
                ax.hist(sstData[:, contrastNumber], bins)
                ax.grid()
                ax.set_title('sst after bias field correction')
                for timepointNumber in range(self.numberOfTimepoints):
                    ax = axs.ravel()[2 + timepointNumber]
                    ax.hist(self.imageBuffersList[timepointNumber][self.masks[timepointNumber], contrastNumber], bins)
                    ax.grid()
                    ax.set_title('time point ' + str(timepointNumber))
                axsList.append(axs)
            plt.draw()

        if self.saveHistory:
            self.history = {
                "sstMeans": self.sstModel.gmm.means,
                "sstVariances": self.sstModel.gmm.variances,
                "sstMixtureWeights": self.sstModel.gmm.mixtureWeights,
                "sstBiasFieldCoefficients": self.sstModel.biasField.coefficients,
                "sstDeformation": self.sstModel.deformation,
                "sstDeformationAtlasFileName": self.sstModel.deformationAtlasFileName,
                "sstOptimizationSummary": self.sstModel.optimizationSummary,
                "sstOptimizationHistory": self.sstModel.optimizationHistory
            }

        if self.saveSSTResults:
            sstlabels, sstnames, sstVolumesInCubicMm, sstoptimizationSummary = self.sstModel.postProcess()

            #
            if self.saveHistory:
                self.history["sstVolumesInCubicMm"] = sstVolumesInCubicMm

        # =======================================================================================
        #
        # Iterative parameter vs. latent variables estimation, using SST result for initialization
        # and/or anchoring of hyperprior strength
        #
        # =======================================================================================

        # Initialization of the time-specific model parameters
        for timepointNumber in range(self.numberOfTimepoints):
            self.timepointModels[timepointNumber].initializeGMM()
            self.timepointModels[timepointNumber].gmm.means = self.sstModel.gmm.means.copy()
            self.timepointModels[timepointNumber].gmm.variances = self.sstModel.gmm.variances.copy()
            self.timepointModels[timepointNumber].gmm.mixtureWeights = self.sstModel.gmm.mixtureWeights.copy()
            self.timepointModels[timepointNumber].initializeBiasField()

        # Initialization of the latent variables, acting as hyperparameters when viewed from the model parameters' perspective
        self.latentDeformation = self.sstModel.deformation.copy()
        self.latentDeformationAtlasFileName = self.sstModel.deformationAtlasFileName
        self.latentMeans = self.sstModel.gmm.means.copy()
        self.latentVariances = self.sstModel.gmm.variances.copy()
        self.latentMixtureWeights = self.sstModel.gmm.mixtureWeights.copy()

        if self.initializeLatentDeformationToZero:
            for timepointNumber in range(self.numberOfTimepoints):
                self.timepointModels[timepointNumber].deformation = self.latentDeformation.copy()
                self.timepointModels[timepointNumber].latentDeformationAtlasFileName = self.latentDeformationAtlasFileName
                self.latentDeformation[:] = 0

        # Strength of the hyperprior (i.e., how much the latent variables control the conditional posterior of the parameters)
        # is user-controlled.
        #
        # For the GMM part, I'm using the *average* number of voxels assigned to the components in each mixture (class) of the
        # SST segmentation, so that all the components in each mixture are well-regularized (and tiny components don't get to do
        # whatever they want)
        #
        # Note that we need to take into account the possible resolution difference between SST and each time point.
        # Here we assume that these time points have similar resolution, otherwise the mean might not be the best choice
        # Scale latent number of measurements by the voxel spacing ratio between the subject-specific template and the time point mean resolution
        meanTimePointResolution = 0
        for t in range(self.numberOfTimepoints):
            meanTimePointResolution += np.prod(self.timepointModels[t].voxelSpacing)
        meanTimePointResolution /= self.numberOfTimepoints
        voxelSpacingRatio = np.prod(self.sstModel.voxelSpacing) / meanTimePointResolution
        print("Voxel spacing ratio: " + str(voxelSpacingRatio))

        K0 = self.sstModel.modelSpecifications.K  # Stiffness population -> latent position
        K1 = self.strengthOfLatentDeformationHyperprior * K0  # Stiffness latent position -> each time point
        sstEstimatedNumberOfVoxelsPerGaussian = np.sum(self.sstModel.optimizationHistory[-1]['posteriorsAtEnd'], axis=0) * \
                                                np.prod(self.sstModel.optimizationHistory[-1]['downSamplingFactors'])
        numberOfClasses = len(self.sstModel.gmm.numberOfGaussiansPerClass)
        numberOfGaussians = sum(self.sstModel.gmm.numberOfGaussiansPerClass)
        self.latentMeansNumberOfMeasurements = np.zeros(numberOfGaussians)
        self.latentVariancesNumberOfMeasurements = np.zeros(numberOfGaussians)
        self.latentMixtureWeightsNumberOfMeasurements = np.zeros(numberOfClasses)
        for classNumber in range(numberOfClasses):
            #
            numberOfComponents = self.sstModel.gmm.numberOfGaussiansPerClass[classNumber]
            gaussianNumbers = np.array(np.sum(self.sstModel.gmm.numberOfGaussiansPerClass[:classNumber]) +
                                       np.array(range(numberOfComponents)), dtype=np.uint32)
            sstEstimatedNumberOfVoxelsInClass = np.sum(sstEstimatedNumberOfVoxelsPerGaussian[gaussianNumbers])

            self.latentMixtureWeightsNumberOfMeasurements[
                classNumber] = self.strengthOfLatentGMMHyperprior * sstEstimatedNumberOfVoxelsInClass * voxelSpacingRatio

            averageSizeOfComponents = sstEstimatedNumberOfVoxelsInClass / numberOfComponents
            self.latentMeansNumberOfMeasurements[gaussianNumbers] = self.strengthOfLatentGMMHyperprior * averageSizeOfComponents * voxelSpacingRatio
            self.latentVariancesNumberOfMeasurements[gaussianNumbers] = self.strengthOfLatentGMMHyperprior * averageSizeOfComponents * voxelSpacingRatio

        # Estimating the mode of the latentVariance posterior distribution (which is Wishart) requires a stringent condition
        # on latentVariancesNumberOfMeasurements so that the mode is actually defined
        threshold = (self.sstModel.gmm.numberOfContrasts + 2) + 1e-6
        self.latentVariancesNumberOfMeasurements[self.latentVariancesNumberOfMeasurements < threshold] = threshold

        # No point in updating latent GMM parameters if the GMM hyperprior has zero weight. The latent variances are also
        # a bit tricky, as they're technically driven to zero in that scenario -- let's try not to go there...
        if self.strengthOfLatentGMMHyperprior == 0:
            self.updateLatentMeans, self.updateLatentVariances, self.updateLatentMixtureWeights = False, False, False

        # Loop over all iterations
        self.historyOfTotalCost, self.historyOfTotalTimepointCost, self.historyOfLatentAtlasCost = [], [], []
        progressPlot = None
        iterationNumber = 0
        if self.saveHistory:
            self.history = {**self.history,
                       **{
                           "timepointMeansEvolution": [],
                           "timepointVariancesEvolution": [],
                           "timepointMixtureWeightsEvolution": [],
                           "timepointBiasFieldCoefficientsEvolution": [],
                           "timepointDeformationsEvolution": [],
                           "timepointDeformationAtlasFileNamesEvolution": [],
                           "latentMeansEvolution": [],
                           "latentVariancesEvolution": [],
                           "latentMixtureWeightsEvolution": [],
                           "latentDeformationEvolution": [],
                           "latentDeformationAtlasFileNameEvolution": []
                       }
                       }

        # Make latent atlas directory
        latentAtlasDirectory = os.path.join(self.savePath, 'latentAtlases')
        os.makedirs(latentAtlasDirectory, exist_ok=True)

        while True:

            # =======================================================================================
            #
            # Update parameters for each time point using the current latent variable estimates
            #
            # =======================================================================================

            # Create a new atlas that will be the basis to deform the individual time points from
            latentAtlasFileName = os.path.join(latentAtlasDirectory, 'latentAtlas_iteration_%02d.mgz' % (iterationNumber + 1))
            self.probabilisticAtlas.saveDeformedAtlas(self.latentDeformationAtlasFileName, latentAtlasFileName, self.latentDeformation, True)

            # Only use the last resolution level, and with the newly created atlas as atlas
            for timepointNumber in range(self.numberOfTimepoints):
                self.timepointModels[timepointNumber].optimizationOptions = self.sstModel.optimizationOptions
                self.timepointModels[timepointNumber].optimizationOptions['multiResolutionSpecification'] = [
                    self.timepointModels[timepointNumber].optimizationOptions['multiResolutionSpecification'][-1]]
                self.timepointModels[timepointNumber].optimizationOptions['multiResolutionSpecification'][0]['atlasFileName'] = latentAtlasFileName
                print(self.timepointModels[timepointNumber].optimizationOptions)

            # Loop over all time points
            totalTimepointCost = 0
            for timepointNumber in range(self.numberOfTimepoints):
                #
                self.timepointModels[timepointNumber].modelSpecifications.K = K1
                self.timepointModels[timepointNumber].gmm.hyperMeans = self.latentMeans
                self.timepointModels[timepointNumber].gmm.hyperVariances = self.latentVariances
                self.timepointModels[timepointNumber].gmm.hyperMixtureWeights = self.latentMixtureWeights
                self.timepointModels[timepointNumber].gmm.fullHyperMeansNumberOfMeasurements = self.latentMeansNumberOfMeasurements.copy()
                self.timepointModels[timepointNumber].gmm.fullHyperVariancesNumberOfMeasurements = self.latentVariancesNumberOfMeasurements.copy()
                self.timepointModels[timepointNumber].gmm.fullHyperMixtureWeightsNumberOfMeasurements = self.latentMixtureWeightsNumberOfMeasurements.copy()

                self.timepointModels[timepointNumber].estimateModelParameters(
                                            initialBiasFieldCoefficients=self.timepointModels[timepointNumber].biasField.coefficients,
                                            initialDeformation=self.timepointModels[timepointNumber].deformation,
                                            initialDeformationAtlasFileName=self.timepointModels[timepointNumber].deformationAtlasFileName,
                                            skipBiasFieldParameterEstimationInFirstIteration=False,
                                            skipGMMParameterEstimationInFirstIteration=(iterationNumber == 0)
                                            )

                totalTimepointCost += self.timepointModels[timepointNumber].optimizationHistory[-1]['historyOfCost'][-1]

                print('=================================')
                print('\n')
                print('timepointNumber: ', timepointNumber)
                print('perVoxelCost: ', self.timepointModels[timepointNumber].optimizationSummary[-1]['perVoxelCost'])
                print('\n')
                print('=================================')
                if hasattr(self.visualizer, 'show_flag'):
                    import matplotlib.pyplot as plt  # avoid importing matplotlib by default
                    plt.ion()
                    self.timepointModels[timepointNumber].biasField.downSampleBasisFunctions([1, 1, 1])
                    timepointBiasFields = self.timepointModels[timepointNumber].biasField.getBiasFields(self.masks[timepointNumber])
                    timepointData = self.imageBuffersList[timepointNumber][self.masks[timepointNumber], :] - timepointBiasFields[self.masks[timepointNumber], :]
                    for contrastNumber in range(self.sstModel.gmm.numberOfContrasts):
                        axs = axsList[contrastNumber]
                        ax = axs.ravel()[2 + timepointNumber]
                        ax.clear()
                        ax.hist(timepointData[:, contrastNumber], bins)
                        ax.grid()
                        ax.set_title('time point ' + str(timepointNumber))
                    plt.draw()

                # End loop over time points

            # =======================================================================================
            #
            # Check for convergence.
            # =======================================================================================

            # In order to also measure the deformation from the population atlas -> latent position,
            # create:
            #   (1) a mesh collection with as reference position the population reference position, and as positions
            #       the currently estimated time point positions.
            #   (2) a mesh with the current latent position
            # Note that in (1) we don't need those time positions now, but these will come in handy very soon to
            # optimize the latent position
            #
            # The parameter estimation happens in a (potentially) downsampled image grid, so it's import to work in the same space
            # when measuring and updating the latentDeformation

            transformUsedForEstimation = gems.KvlTransform(
                requireNumpyArray(self.sstModel.optimizationHistory[-1]['downSampledTransformMatrix']))
            mesh_collection = gems.KvlMeshCollection()
            mesh_collection.read(self.latentDeformationAtlasFileName)
            mesh_collection.transform(transformUsedForEstimation)
            referencePosition = mesh_collection.reference_position
            timepointPositions = []
            for timepointNumber in range(self.numberOfTimepoints):
                positionInTemplateSpace = self.probabilisticAtlas.mapPositionsFromSubjectToTemplateSpace(referencePosition,
                                                                                 transformUsedForEstimation) + \
                                          self.latentDeformation + self.timepointModels[timepointNumber].deformation
                timepointPositions.append(
                    self.probabilisticAtlas.mapPositionsFromTemplateToSubjectSpace(positionInTemplateSpace, transformUsedForEstimation))
            mesh_collection.set_positions(referencePosition, timepointPositions)


            # Read mesh in sst warp
            mesh = self.probabilisticAtlas.getMesh(latentAtlasFileName, transformUsedForEstimation)

            #
            calculator = gems.KvlCostAndGradientCalculator(mesh_collection, K0, 0.0, transformUsedForEstimation)
            latentAtlasCost, _ = calculator.evaluate_mesh_position(mesh)

            #
            totalCost = totalTimepointCost + latentAtlasCost
            print('*' * 100 + '\n')
            print('iterationNumber: ', iterationNumber)
            print('totalCost: ', totalCost)
            print('   latentAtlasCost: ', latentAtlasCost)
            print('   totalTimepointCost: ', totalTimepointCost)
            print('*' * 100 + '\n')
            self.historyOfTotalCost.append(totalCost)
            self.historyOfTotalTimepointCost.append(totalTimepointCost)
            self.historyOfLatentAtlasCost.append(latentAtlasCost)

            if hasattr(self.visualizer, 'show_flag'):
                import matplotlib.pyplot as plt  # avoid importing matplotlib by default
                plt.ion()
                if progressPlot is None:
                    plt.figure()
                    progressPlot = plt.subplot()
                progressPlot.clear()
                progressPlot.plot(self.historyOfTotalCost, color='k')
                progressPlot.plot(self.historyOfTotalTimepointCost, linestyle='-.', color='b')
                progressPlot.plot(self.historyOfLatentAtlasCost, linestyle='-.', color='r')
                progressPlot.grid()
                progressPlot.legend(['total', 'timepoints', 'latent atlas deformation'])
                plt.draw()

            if self.saveHistory:
                self.history["timepointMeansEvolution"].append(self.timepointModels[timepointNumber].gmm.means.copy())
                self.history["timepointVariancesEvolution"].append(self.timepointModels[timepointNumber].gmm.variances.copy())
                self.history["timepointMixtureWeightsEvolution"].append(self.timepointModels[timepointNumber].gmm.mixtureWeights.copy())
                self.history["timepointBiasFieldCoefficientsEvolution"].append(self.timepointModels[timepointNumber].biasField.coefficients.copy())
                self.history["timepointDeformationsEvolution"].append(self.timepointModels[timepointNumber].deformation)
                self.history["timepointDeformationAtlasFileNamesEvolution"].append(self.timepointModels[timepointNumber].deformationAtlasFileName)
                self.history["latentMeansEvolution"].append(self.latentMeans.copy())
                self.history["latentVariancesEvolution"].append(self.latentVariances.copy())
                self.history["latentMixtureWeightsEvolution"].append(self.latentMixtureWeights.copy())
                self.history["latentDeformationEvolution"].append(self.latentDeformation.copy())
                self.history["latentDeformationAtlasFileNameEvolution"].append(self.latentDeformationAtlasFileName)

            if iterationNumber >= (self.numberOfIterations - 1):
                print('Stopping')
                break

            # =======================================================================================
            #
            # Update the latent variables based on the current parameter estimates
            #
            # =======================================================================================
            self.updateLatentDeformationAtlas(mesh_collection, mesh, K0, K1, transformUsedForEstimation)
            self.updateGMMLatentVariables()

            iterationNumber += 1
            # End loop over parameter and latent variable estimation iterations

    def updateLatentDeformationAtlas(self, mesh_collection, mesh, K0, K1, transformUsedForEstimation):
        # Update the latentDeformation
        if self.updateLatentDeformation:
            # Set up calculator
            calculator = gems.KvlCostAndGradientCalculator(mesh_collection, K0, K1, transformUsedForEstimation)

            # Get optimizer and plug calculator in it
            optimizerType = 'L-BFGS'
            optimizationParameters = {
                'Verbose': self.sstModel.optimizationOptions['verbose'],
                'MaximalDeformationStopCriterion': self.sstModel.optimizationOptions['maximalDeformationStopCriterion'],
                'LineSearchMaximalDeformationIntervalStopCriterion': self.sstModel.optimizationOptions[
                    'lineSearchMaximalDeformationIntervalStopCriterion'],
                'MaximumNumberOfIterations': self.sstModel.optimizationOptions['maximumNumberOfDeformationIterations'],
                'BFGS-MaximumMemoryLength': self.sstModel.optimizationOptions['BFGSMaximumMemoryLength']
            }
            optimizer = gems.KvlOptimizer(optimizerType, mesh, calculator, optimizationParameters)

            # Run deformation optimization
            historyOfDeformationCost = []
            historyOfMaximalDeformation = []
            nodePositionsBeforeDeformation = mesh.points
            while True:
                minLogLikelihoodTimesDeformationPrior, maximalDeformation = optimizer.step_optimizer_samseg()
                print("maximalDeformation=%.4f minLogLikelihood=%.4f" % (
                    maximalDeformation, minLogLikelihoodTimesDeformationPrior))
                historyOfDeformationCost.append(minLogLikelihoodTimesDeformationPrior)
                historyOfMaximalDeformation.append(maximalDeformation)
                if maximalDeformation == 0:
                    break

            nodePositionsAfterDeformation = mesh.points
            maximalDeformationApplied = np.sqrt(
                np.max(np.sum((nodePositionsAfterDeformation - nodePositionsBeforeDeformation) ** 2, 1)))

            #
            nodePositionsBeforeDeformation = self.probabilisticAtlas.mapPositionsFromSubjectToTemplateSpace(
                nodePositionsBeforeDeformation,
                transformUsedForEstimation)
            nodePositionsAfterDeformation = self.probabilisticAtlas.mapPositionsFromSubjectToTemplateSpace(
                nodePositionsAfterDeformation,
                transformUsedForEstimation)
            estimatedUpdate = nodePositionsAfterDeformation - nodePositionsBeforeDeformation
            self.latentDeformation += estimatedUpdate

    def updateGMMLatentVariables(self):

        # Update latentMeans
        if self.updateLatentMeans:
            numberOfGaussians = np.sum(self.sstModel.gmm.numberOfGaussiansPerClass)
            numberOfContrasts = self.latentMeans.shape[-1]
            for gaussianNumber in range(numberOfGaussians):
                # Set up linear system
                lhs = np.zeros((numberOfContrasts, numberOfContrasts))
                rhs = np.zeros((numberOfContrasts, 1))
                for timepointNumber in range(self.numberOfTimepoints):
                    mean = np.expand_dims(self.timepointModels[timepointNumber].gmm.means[gaussianNumber], 1)
                    variance = self.timepointModels[timepointNumber].gmm.variances[gaussianNumber]

                    lhs += np.linalg.inv(variance)
                    rhs += np.linalg.solve(variance, mean)

                # Solve linear system
                latentMean = np.linalg.solve(lhs, rhs)
                self.latentMeans[gaussianNumber, :] = latentMean.T

        # Update latentVariances
        if self.updateLatentVariances:
            numberOfGaussians = np.sum(self.sstModel.gmm.numberOfGaussiansPerClass)
            numberOfContrasts = self.latentMeans.shape[-1]
            for gaussianNumber in range(numberOfGaussians):
                # Precision is essentially averaged
                averagePrecision = np.zeros((numberOfContrasts, numberOfContrasts))
                for timepointNumber in range(self.numberOfTimepoints):
                    variance = self.timepointModels[timepointNumber].gmm.variances[gaussianNumber]

                    averagePrecision += np.linalg.inv(variance)
                averagePrecision /= self.numberOfTimepoints

                latentVarianceNumberOfMeasurements = self.latentVariancesNumberOfMeasurements[gaussianNumber]
                latentVariance = np.linalg.inv(averagePrecision) * \
                                 (latentVarianceNumberOfMeasurements - numberOfContrasts - 2) / latentVarianceNumberOfMeasurements
                self.latentVariances[gaussianNumber] = latentVariance

        # Update latentMixtureWeights
        if self.updateLatentMixtureWeights:
            numberOfClasses = len(self.sstModel.gmm.numberOfGaussiansPerClass)
            for classNumber in range(numberOfClasses):
                numberOfComponents = self.sstModel.gmm.numberOfGaussiansPerClass[classNumber]
                averageInLogDomain = np.zeros(numberOfComponents)
                for componentNumber in range(numberOfComponents):
                    gaussianNumber = sum(self.sstModel.gmm.numberOfGaussiansPerClass[:classNumber]) + componentNumber
                    for timepointNumber in range(self.numberOfTimepoints):
                        mixtureWeight = self.timepointModels[timepointNumber].gmm.mixtureWeights[gaussianNumber]
                        averageInLogDomain[componentNumber] += np.log(mixtureWeight + eps)
                    averageInLogDomain[componentNumber] /= self.numberOfTimepoints

                # Solution is normalized version
                solution = np.exp(averageInLogDomain)
                solution /= np.sum(solution + eps)

                #
                for componentNumber in range(numberOfComponents):
                    gaussianNumber = sum(self.sstModel.gmm.numberOfGaussiansPerClass[:classNumber]) + componentNumber
                    self.latentMixtureWeights[gaussianNumber] = solution[componentNumber]

    def postProcess(self, saveWarp=False):

        # =======================================================================================
        #
        # Using estimated parameters, segment and write out results for each time point
        #
        # =======================================================================================
        #

        sstDir = os.path.join(self.savePath, 'base')
        os.makedirs(sstDir, exist_ok=True)
        baseModel = self.sstModel;
        # Save the final mesh collection
        if self.saveModelProbabilities:
            print('Saving base model probs')
            baseModel.saveGaussianProbabilities(os.path.join(sstDir, 'probabilities') )
        if saveWarp:
            baseModel.saveWarpField(os.path.join(sstDir, 'template.m3z'))

        self.timepointVolumesInCubicMm = []
        for timepointNumber in range(self.numberOfTimepoints):
            #
            timepointModel = self.timepointModels[timepointNumber]

            #
            timepointModel.deformation = self.latentDeformation + timepointModel.deformation
            timepointModel.deformationAtlasFileName = self.latentDeformationAtlasFileName
            posteriors, biasFields, nodePositions, _, _ = timepointModel.computeFinalSegmentation()

            #
            timepointDir = os.path.join(self.savePath, 'tp%03d' % (timepointNumber + 1))
            os.makedirs(timepointDir, exist_ok=True)

            #
            timepointModel.savePath = timepointDir
            volumesInCubicMm = timepointModel.writeResults(biasFields, posteriors)

            # Save the timepoint->template warp
            if saveWarp:
                timepointModel.saveWarpField(os.path.join(timepointDir, 'template.m3z'))

            # Save the Gaussian priors, (normalized) likelihoods and posteriors
            if self.saveModelProbabilities:
                timepointModel.saveGaussianProbabilities( os.path.join(timepointModel.savePath, 'probabilities') )

            # Save the final mesh collection
            if self.saveMesh:
                print('Saving the final mesh in template space')
                deformedAtlasFileName = os.path.join(timepointModel.savePath, 'mesh.txt')
                timepointModel.probabilisticAtlas.saveDeformedAtlas(timepointModel.modelSpecifications.atlasFileName,
                                                                    deformedAtlasFileName, nodePositions)
            if self.saveModelProbabilities:
                print('Saving model probs')
                timepointModel.saveGaussianProbabilities( os.path.join(timepointModel.savePath, 'probabilities') )
                                                                                                          
            # Save the history of the parameter estimation process
            if self.saveHistory:
                history = {'input': {
                    'imageFileNames': timepointModel.imageFileNames,
                    'imageToImageTransformMatrix': timepointModel.imageToImageTransformMatrix,
                    'modelSpecifications': timepointModel.modelSpecifications,
                    'optimizationOptions': timepointModel.optimizationOptions,
                    'savePath': timepointModel.savePath
                }, 'imageBuffers': timepointModel.imageBuffers, 'mask': timepointModel.mask,
                    'cropping': timepointModel.cropping,
                    'transform': timepointModel.transform.as_numpy_array,
                    'historyWithinEachMultiResolutionLevel': timepointModel.optimizationHistory,
                    "labels": timepointModel.modelSpecifications.FreeSurferLabels, "names": timepointModel.modelSpecifications.names,
                    "volumesInCubicMm": volumesInCubicMm, "optimizationSummary": timepointModel.optimizationSummary}
                with open(os.path.join(timepointModel.savePath, 'history.p'), 'wb') as file:
                    pickle.dump(history, file, protocol=pickle.HIGHEST_PROTOCOL)
            

            self.timepointVolumesInCubicMm.append(volumesInCubicMm)

        #
        self.optimizationSummary = {
            "historyOfTotalCost": self.historyOfTotalCost,
            "historyOfTotalTimepointCost": self.historyOfTotalTimepointCost,
            "historyOfLatentAtlasCost": self.historyOfLatentAtlasCost
        }

        if self.saveHistory:
            self.history["labels"] = self.sstModel.modelSpecifications.FreeSurferLabels
            self.history["names"] = self.sstModel.modelSpecifications.names
            self.history["timepointVolumesInCubicMm"] = self.timepointVolumesInCubicMm
            self.history["optimizationSummary"] = self.optimizationSummary
            with open(os.path.join(self.savePath, 'history.p'), 'wb') as file:
                pickle.dump(self.history, file, protocol=pickle.HIGHEST_PROTOCOL)

    def generateSubjectSpecificTemplate(self, saveWarp=False):

        sstDir = os.path.join(self.savePath, 'base')
        os.makedirs(sstDir, exist_ok=True)

        sstFileNames = []
        for contrastNumber, contrastImageFileNames in enumerate(zip(*self.imageFileNamesList)):

            # Read in the various time point images, and compute the average
            numberOfTimepoints = len(contrastImageFileNames)
            image0 = sf.load_volume(contrastImageFileNames[0])
            imageBuffer = image0.transform(sf.Affine(self.tpToBaseTransforms[0]))
            # Make sure that we are averaging only non zero voxels
            count = np.zeros(imageBuffer.shape)
            count[imageBuffer > 0] += 1
            for timepointNumber in range(1, numberOfTimepoints):
                tmp = sf.load_volume(contrastImageFileNames[timepointNumber]).transform(sf.Affine(self.tpToBaseTransforms[timepointNumber])).data
                imageBuffer += tmp
                count[tmp > 0] += 1
            # Make sure that we are not dividing by zero for, e.g., background voxels
            imageBuffer[count > 0] /= count[count > 0]

            # Write image to disk
            sstFilename = os.path.join(sstDir, 'mode%02d_average.mgz' % (contrastNumber + 1))
            imageBuffer.save(sstFilename)

            #
            sstFileNames.append(sstFilename)

        #
        return sstFileNames

    def constructSstModel(self):

        sstDir, _ = os.path.split(self.sstFileNames[0])

        self.sstModel = Samseg(
            imageFileNames=self.sstFileNames,
            atlasDir=self.atlasDir,
            savePath=sstDir,
            imageToImageTransformMatrix=self.imageToImageTransformMatrix,
            userModelSpecifications=self.userModelSpecifications,
            userOptimizationOptions=self.userOptimizationOptions,
            visualizer=self.visualizer,
            saveHistory=True,
            savePosteriors=self.savePosteriors,
            targetIntensity=self.targetIntensity,
            targetSearchStrings=self.targetSearchStrings,
            modeNames=self.modeNames,
            pallidumAsWM=self.pallidumAsWM
        )

    def constructTimepointModels(self):

        self.timepointModels = []

        # Construction of the cross sectional model for each time point
        for timepointNumber in range(self.numberOfTimepoints):
            self.timepointModels.append(Samseg(
                imageFileNames=self.imageFileNamesList[timepointNumber],
                atlasDir=self.atlasDir,
                savePath=self.savePath,
                imageToImageTransformMatrix=self.imageToImageTransformMatrix,
                userModelSpecifications=self.userModelSpecifications,
                userOptimizationOptions=self.userOptimizationOptions,
                visualizer=self.visualizer,
                saveHistory=True,
                targetIntensity=self.targetIntensity,
                targetSearchStrings=self.targetSearchStrings,
                modeNames=self.modeNames,
                pallidumAsWM=self.pallidumAsWM,
                saveModelProbabilities=self.saveModelProbabilities,
                savePosteriors=self.savePosteriors
            ))

            self.timepointModels[timepointNumber].mask = self.masks[timepointNumber]
            self.timepointModels[timepointNumber].imageBuffers = self.imageBuffersList[timepointNumber]
            self.timepointModels[timepointNumber].voxelSpacing = self.voxelSpacings[timepointNumber]
            self.timepointModels[timepointNumber].transform = self.transforms[timepointNumber]
            self.timepointModels[timepointNumber].cropping = self.croppings[timepointNumber]
