# Segment ventricles and return volume in cubic millimetres (= millilitres)
# ADD IMPLEMENTATION of segmentation and volume calculation below

from nipy import load_image, save_image
from nipy.core.image.image_spaces import (make_xyz_image,xyz_affine)
from nipy.algorithms.segmentation.brain_segmentation import BrainT1Segmentation

def my_segmentation_func (fullTrainingImageName,maskname,segSaveDir='./data/seg/'):

    # Segmentation name
    fullSegmentationName  = segSaveDir + fullTrainingImageName.split('/')[-1]

    # Load the mask
    # See https://sites.google.com/site/mrilateralventricle/ for more details
    maskimage = sitk.ReadImage(maskname)
    maskarray = sitk.GetArrayFromImage(maskimage)

    # Check if segmentation is already saved to disk
    if not os.path.isfile(fullSegmentationName):

        # Read the image
        img_sitk = sitk.ReadImage(fullTrainingImageName)
        img      = load_image(fullTrainingImageName)
        mask     = img.get_data() > 0 # ingoring background pixels

        # Perform tissue classification
        # ADD IMPLEMENTATION - store segmentation in labelimage
        S = BrainT1Segmentation(img.get_data(),
                        mask=mask,       # the algorithm ignores the background pixels for faster comp.
                        model='5k',
                        niters=int(50),  # number of iterations
                        beta=float(0.4), # Markov regularization parameter
                        ngb_size=int(6)) # neighborhood size of MRF
        xyz_image_array = make_xyz_image(S.label, xyz_affine(img), 'scanner').get_data()
        masked_image_array = xyz_image_array * maskarray
        labelimage = sitk.GetImageFromArray(masked_image_array)

        # Save the segmentation to disk to avoid recomputation
        sitk.WriteImage(labelimage,fullSegmentationName)

    # Load the segmentation
    labelimage = sitk.ReadImage(fullSegmentationName)

    # Compute the volume
    # ADD IMPLEMENTATION - store volume in vol
    labelarray = sitk.GetArrayFromImage(labelimage)
    vol = sum(labelarray == ????)
    return vol