__all__ = ['functional_lung_segmentation']

import ants

def functional_lung_segmentation(image,

                                 mask=None,

                                 number_of_iterations=2,

                                 number_of_atropos_iterations=5,

                                 mrf_parameters="[0.7,2x2x2]",

                                 number_of_clusters = 6,

                                 cluster_centers = None,

                                 bias_correction = "n4",

                                 verbose=True):

    """

    Ventilation-based segmentation of hyperpolarized gas lung MRI.

    Lung segmentation into classes based on ventilation as described in

    this paper:

         https://pubmed.ncbi.nlm.nih.gov/21837781/

    Arguments

    ---------

    image : ANTs image

        Input proton-weighted MRI.

    mask : ANTs image

        Mask image designating the region to segment.  0/1 = background/foreground.

    number_of_iterations : integer

        Number of Atropos <--> bias correction iterations (outer loop).

    number_of_atropos_iterations : integer

        Number of Atropos iterations (inner loop).  If number_of_atropos_iterations = 0,

        this is equivalent to K-means with no MRF priors.

    mrf_parameters : string

        Parameters for MRF in Atropos.

    number_of_clusters : integer

        Number of tissue classes.

    cluster_centers: array or tuple

        Initialization centers for k-means.

    bias_correction: string

        Apply "n3", "n4", or no bias correction (default = "n4").

    verbose : boolean

        Print progress to the screen.

    Returns

    -------

    Dictionary with segmentation image, probability images, and

    processed image.

    Example

    -------

    >>> import ants

    >>> image = ants.image_read(ants.get_data("mni")).resample_image((4,4,4))

    >>> mask = image.get_mask()

    >>> seg = ants.functional_lung_segmentation(image, mask, verbose=True,

                                                number_of_iterations=1,

                                                number_of_clusters=2,

                                                number_of_atropos_iterations=1)

    """

    if image.dimension != 3:

        raise ValueError("Function only works for 3-D images.")

    if mask is None:

        raise ValueError("Mask is missing.")

    if number_of_iterations < 1:

        raise ValueError("number_of_iterations must be >= 1.")

    def generate_pure_tissue_n4_weight_mask(probability_images):

        number_of_probability_images = len(probability_images)

        pure_tissue_mask = probability_images[0] * 0

        for i in range(number_of_probability_images):

            negation_image = probability_images[0] * 0 + 1

            for j in range(number_of_probability_images):

                if i != j:

                    negation_image = negation_image * (probability_images[j] * -1.0 + 1.0)

                pure_tissue_mask = pure_tissue_mask + negation_image * probability_images[i]

        return(pure_tissue_mask)

    weight_mask = None

    # This is a multiplicative factor for both the image

    # and the cluster centers which shouldn't effect the

    # user.  Otherwise, we'll get a singular covariance

    # complaint from Atropos.

    image_scale_factor = 1000.0

    number_of_atropos_n4_iterations = number_of_iterations

    for i in range(number_of_atropos_n4_iterations):

        if verbose == True:

            print("Outer iteration: ", i, " out of ", number_of_atropos_n4_iterations)

        preprocessed_image = ants.image_clone(image)

        quantiles = (preprocessed_image.quantile(0.0), preprocessed_image.quantile(0.995))

        preprocessed_image[preprocessed_image < quantiles[0]] = quantiles[0]

        preprocessed_image[preprocessed_image > quantiles[1]] = quantiles[1]

        if verbose == True:

            print("Outer: bias correction.")

        if bias_correction.lower() == "n4":

            preprocessed_image = ants.n4_bias_field_correction(preprocessed_image,

                mask=mask, shrink_factor=2, convergence={'iters': [50, 50, 50, 50], 'tol': 0.0000000001},

                spline_param=200, return_bias_field=False, weight_mask=weight_mask, verbose=verbose)

        elif bias_correction.lower == "n3":

            preprocessed_image = ants.n3_bias_field_correction(preprocessed_image, downsample_factor=2)

        preprocessed_image = ((preprocessed_image - preprocessed_image.min())

                              /(preprocessed_image.max() - preprocessed_image.min()))

        preprocessed_image *= image_scale_factor

        if verbose == True:

            print("Outer: Atropos segmentation.")

        atropos_initialization = "Kmeans[" + str(number_of_clusters) + "]"

        if cluster_centers is not None:

            if len(cluster_centers) != number_of_clusters:

                raise ValueError("number_of_clusters should match the vector size of the cluster_centers.")

            else:

                scaled_cluster_centers = image_scale_factor * cluster_centers

                cluster_centers_string = 'x'.join(str(s) for s in set(scaled_cluster_centers))

                atropos_initialization = "Kmeans[" + str(number_of_clusters) + "," + cluster_centers_string + "]"

        posterior_formulation = "Socrates[0]"

        if i > 0:

            atropos_initialization = atropos_output['probabilityimages']

            posterior_formulation = "Socrates[1]"

        iterations = "[" + str(number_of_atropos_iterations) + ",0]"

        atropos_verbose = 0

        if verbose == True:

            atropos_verbose = 1

        atropos_output = ants.atropos(preprocessed_image, x=mask, i=atropos_initialization,

            m=mrf_parameters, c=iterations, priorweight=0.0, v=atropos_verbose, p=posterior_formulation)

        weight_mask = generate_pure_tissue_n4_weight_mask(atropos_output['probabilityimages'][1:number_of_clusters])

    masked_segmentation_image = atropos_output['segmentation'] * mask

    masked_probability_images = list()

    for i in range(len(atropos_output['probabilityimages'])):

        masked_probability_images.append(atropos_output['probabilityimages'][i] * mask)

    return_dict = {'segmentation_image' : masked_segmentation_image,

                   'probability_images' : masked_probability_images,

                   'processed_image' : (preprocessed_image / image_scale_factor)}

    return(return_dict)