"""

Joint Label Fusion algorithm

"""

__all__ = ["joint_label_fusion", "local_joint_label_fusion"]

import os

import numpy as np

import warnings

from pathlib import Path

from tempfile import NamedTemporaryFile

from tempfile import mktemp

import glob

import re

import math

import ants

from ants.internal import get_lib_fn, get_pointer_string, process_arguments

def joint_label_fusion(

    target_image,

    target_image_mask,

    atlas_list,

    beta=4,

    rad=2,

    label_list=None,

    rho=0.01,

    usecor=False,

    r_search=3,

    nonnegative=False,

    no_zeroes=False,

    max_lab_plus_one=False,

    output_prefix=None,

    verbose=False,

):

    """

    A multiple atlas voting scheme to customize labels for a new subject.

    This function will also perform intensity fusion. It almost directly

    calls the C++ in the ANTs executable so is much faster than other

    variants in ANTsR.

    One may want to normalize image intensities for each input image before

    passing to this function. If no labels are passed, we do intensity fusion.

    Note on computation time: the underlying C++ is multithreaded.

    You can control the number of threads by setting the environment

    variable ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS e.g. to use all or some

    of your CPUs. This will improve performance substantially.

    For instance, on a macbook pro from 2015, 8 cores improves speed by about 4x.

    ANTsR function: `jointLabelFusion`

    Arguments

    ---------

    target_image : ANTsImage

        image to be approximated

    target_image_mask : ANTsImage

        mask with value 1

    atlas_list : list of ANTsImage types

        list containing intensity images

    beta : scalar

        weight sharpness, default to 2

    rad : scalar

        neighborhood radius, default to 2

    label_list : list of ANTsImage types (optional)

        list containing images with segmentation labels

    rho : scalar

        ridge penalty increases robustness to outliers but also makes image converge to average

    usecor : boolean

        employ correlation as local similarity

    r_search : scalar

        radius of search, default is 3

    nonnegative : boolean

        constrain weights to be non-negative

    no_zeroes : boolean

        this will constrain the solution only to voxels that are always non-zero in the label list

    max_lab_plus_one : boolean

        this will add max label plus one to the non-zero parts of each label where the target mask

        is greater than one.  NOTE: this will have a side effect of adding to the original label

        images that are passed to the program.  It also guarantees that every position in the

        labels have some label, rather than none.  Ie it guarantees to explicitly parcellate the

        input data.

    output_prefix: string

        file prefix for storing output probabilityimages to disk

    verbose : boolean

        whether to show status updates

    Returns

    -------

    dictionary w/ following key/value pairs:

        `segmentation` : ANTsImage

            segmentation image

        `intensity` : ANTsImage

            intensity image

        `probabilityimages` : list of ANTsImage types

            probability map image for each label

        `segmentation_numbers` : list of numbers

            segmentation label (number, int) for each probability map

    Example

    -------

    >>> import ants

    >>> ref = ants.image_read( ants.get_ants_data('r16'))

    >>> ref = ants.resample_image(ref, (50,50),1,0)

    >>> ref = ants.iMath(ref,'Normalize')

    >>> mi = ants.image_read( ants.get_ants_data('r27'))

    >>> mi2 = ants.image_read( ants.get_ants_data('r30'))

    >>> mi3 = ants.image_read( ants.get_ants_data('r62'))

    >>> mi4 = ants.image_read( ants.get_ants_data('r64'))

    >>> mi5 = ants.image_read( ants.get_ants_data('r85'))

    >>> refmask = ants.get_mask(ref)

    >>> refmask = ants.iMath(refmask,'ME',2) # just to speed things up

    >>> ilist = [mi,mi2,mi3,mi4,mi5]

    >>> seglist = [None]*len(ilist)

    >>> for i in range(len(ilist)):

    >>>     ilist[i] = ants.iMath(ilist[i],'Normalize')

    >>>     mytx = ants.registration(fixed=ref , moving=ilist[i] ,

    >>>         type_of_transform = ('Affine') )

    >>>     mywarpedimage = ants.apply_transforms(fixed=ref,moving=ilist[i],

    >>>             transformlist=mytx['fwdtransforms'])

    >>>     ilist[i] = mywarpedimage

    >>>     seg = ants.threshold_image(ilist[i],'Otsu', 3)

    >>>     seglist[i] = ( seg ) + ants.threshold_image( seg, 1, 3 ).morphology( operation='dilate', radius=3 )

    >>> r = 2

    >>> pp = ants.joint_label_fusion(ref, refmask, ilist, r_search=2,

    >>>                     label_list=seglist, rad=[r]*ref.dimension )

    >>> pp = ants.joint_label_fusion(ref,refmask,ilist, r_search=2, rad=[r]*ref.dimension)

    """

    segpixtype = "unsigned int"

    if (label_list is None) or (np.any([l is None for l in label_list])):

        doJif = True

    else:

        doJif = False

    if not doJif:

        if len(label_list) != len(atlas_list):

            raise ValueError("len(label_list) != len(atlas_list)")

        if no_zeroes:

            for label in label_list:

                target_image_mask[label == 0] = 0

        inlabs = set()

        for label in label_list:

            values = np.unique(label[target_image_mask != 0 and label != 0])

            inlabs = inlabs.union(values)

        inlabs = sorted(inlabs)

        maxLab = max(inlabs)

        if max_lab_plus_one:

            for label in label_list:

                label[label == 0] = maxLab + 1

        mymask = target_image_mask.clone()

    else:

        mymask = target_image_mask

###### security issues with mktemp but could not figure out the right solution

###### NamedTemporaryFile creates a file with permissions:

###### -rw-------  1 stnava  staff

###### whereas mktemp gives

###### -rw-r--r--  1 stnava  staff

###### the latter is what we want - one solution is to use chmod via os but

###### am currently too lazy to change one line of code to two or more everywhere

#    osegfn = NamedTemporaryFile(prefix="antsr", suffix="myseg.nii.gz",delete=False).name

    osegfn = mktemp(prefix="antsr", suffix="myseg.nii.gz")

    # segdir = osegfn.replace(os.path.basename(osegfn),'')

    if os.path.exists(osegfn):

        os.remove(osegfn)

    if output_prefix is None:

#        probs = NamedTemporaryFile(prefix="antsr", suffix="prob%02d.nii.gz",delete=False).name

        probs = mktemp(prefix="antsr", suffix="prob%02d.nii.gz")

        probsbase = os.path.basename(probs)

        tdir = probs.replace(probsbase, "")

        searchpattern = probsbase.replace("%02d", "*")

    if output_prefix is not None:

        probs = output_prefix + "prob%02d.nii.gz"

        probpath = Path(probs).parent

        Path(probpath).mkdir(parents=True, exist_ok=True)

        probsbase = os.path.basename(probs)

        tdir = probs.replace(probsbase, "")

        searchpattern = probsbase.replace("%02d", "*")

    mydim = target_image_mask.dimension

    if not doJif:

        # not sure if these should be allocated or what their size should be

        outimg = label_list[1].clone(segpixtype)

        outimgi = target_image * 0

        outimg_ptr = get_pointer_string(outimg)

        outimgi_ptr = get_pointer_string(outimgi)

        outs = "[%s,%s,%s]" % (outimg_ptr, outimgi_ptr, probs)

    else:

        outimgi = target_image * 0

        outs = get_pointer_string(outimgi)

    mymask = mymask.clone(segpixtype)

    if (not isinstance(rad, (tuple, list))) or (len(rad) == 1):

        myrad = [rad] * mydim

    else:

        myrad = rad

    if len(myrad) != mydim:

        raise ValueError("path radius dimensionality must equal image dimensionality")

    myrad = "x".join([str(mr) for mr in myrad])

    vnum = 1 if verbose else 0

    nnum = 1 if nonnegative else 0

    mypc = "MSQ"

    if usecor:

        mypc = "PC"

    myargs = {

        "d": mydim,

        "t": target_image,

        "a": rho,

        "b": beta,

        "c": nnum,

        "p": myrad,

        "m": mypc,

        "s": r_search,

        "x": mymask,

        "o": outs,

        "v": vnum,

    }

    kct = len(myargs.keys())

    for k in range(len(atlas_list)):

        kct += 1

        myargs["g-MULTINAME-%i" % kct] = atlas_list[k]

        if not doJif:

            kct += 1

            castseg = label_list[k].clone(segpixtype)

            myargs["l-MULTINAME-%i" % kct] = castseg

    myprocessedargs = process_arguments(myargs)

    libfn = get_lib_fn("antsJointFusion")

    rval = libfn(myprocessedargs)

    if rval != 0:

        print("Warning: Non-zero return from antsJointFusion")

    if doJif:

        return outimgi

    probsout = glob.glob(os.path.join(tdir, "*" + searchpattern))

    probsout.sort()

    probimgs = []

#    print( os.system("ls -l "+probsout[0]) )

    for idx in range(len(probsout)):

        probimgs.append(ants.image_read(probsout[idx]))

    #    if len(probsout) != (len(inlabs)) and max_lab_plus_one == False:

    #        warnings.warn("Length of output probabilities != length of unique input labels")

    segmentation_numbers = [0] * len(probsout)

    for i in range(len(probsout)):

        temp = str.split(probsout[i], "prob")

        segnum = temp[len(temp) - 1].split(".nii.gz")[0]

        segmentation_numbers[i] = int(segnum)

    if max_lab_plus_one == False:

        segmat = ants.images_to_matrix(probimgs, target_image_mask)

        finalsegvec = segmat.argmax(axis=0)

        finalsegvec2 = finalsegvec.copy()

        # mapfinalsegvec to original labels

        for i in range(len(probsout)):

            temp = str.split(probsout[i], "prob")

            segnum = temp[len(temp) - 1].split(".nii.gz")[0]

            finalsegvec2[finalsegvec == i] = segnum

        outimg = ants.make_image(target_image_mask, finalsegvec2)

        return {

            "segmentation": outimg,

            "intensity": outimgi,

            "probabilityimages": probimgs,

            "segmentation_numbers": segmentation_numbers,

        }

    if max_lab_plus_one == True:

        mymaxlab = max(segmentation_numbers)

        matchings_indices = [

            i

            for i, segmentation_numbers in enumerate(segmentation_numbers)

            if segmentation_numbers == mymaxlab

        ]

        background_prob = probimgs[matchings_indices[0]]

        background_probfn = probsout[matchings_indices[0]]

        del probimgs[matchings_indices[0]]

        del probsout[matchings_indices[0]]

        del segmentation_numbers[matchings_indices[0]]

        segmat = ants.images_to_matrix(probimgs, target_image_mask)

        finalsegvec = segmat.argmax(axis=0)

        finalsegvec2 = finalsegvec.copy()

        # mapfinalsegvec to original labels

        for i in range(len(probsout)):

            temp = str.split(probsout[i], "prob")

            segnum = temp[len(temp) - 1].split(".nii.gz")[0]

            finalsegvec2[finalsegvec == i] = segnum

        outimg = ants.make_image(target_image_mask, finalsegvec2)

        # next decide what is "background" based on the sum of the first k labels vs the prob of the last one

        firstK = probimgs[0] * 0

        for i in range(len(probsout)):

            firstK = firstK + probimgs[i]

        segmat = ants.images_to_matrix([background_prob, firstK], target_image_mask)

        bkgsegvec = segmat.argmax(axis=0)

        outimg = outimg * ants.make_image(target_image_mask, bkgsegvec)

        return {

            "segmentation": outimg * ants.make_image(target_image_mask, bkgsegvec),

            "segmentation_raw": outimg,

            "intensity": outimgi,

            "probabilityimages": probimgs,

            "segmentation_numbers": segmentation_numbers,

            "background_prob": background_prob,

        }

def local_joint_label_fusion(

    target_image,

    which_labels,

    target_mask,

    initial_label,

    atlas_list,

    label_list,

    submask_dilation=10,

    type_of_transform="SyN",

    aff_metric="meansquares",

    syn_metric="mattes",

    syn_sampling=32,

    reg_iterations=(40, 20, 0),

    aff_iterations=(500, 50, 0),

    grad_step=0.2,

    flow_sigma=3,

    total_sigma=0,

    beta=4,

    rad=2,

    rho=0.1,

    usecor=False,

    r_search=3,

    nonnegative=False,

    no_zeroes=False,

    max_lab_plus_one=False,

    local_mask_transform="Similarity",

    output_prefix=None,

    verbose=False,

):

    """

    A local version of joint label fusion that focuses on a subset of labels.

    This is primarily different from standard JLF because it performs

    registration on the label subset and focuses JLF on those labels alone.

    ANTsR function: `localJointLabelFusion`

    Arguments

    ---------

    target_image : ANTsImage

        image to be labeled

    which_labels : numeric vector

        label number(s) that exist(s) in both the template and library

    target_image_mask : ANTsImage

        a mask for the target image (optional), passed to joint fusion

    initial_label : ANTsImage

        initial label set, may be same labels as library or binary.

        typically labels would be produced by a single deformable registration

        or by manual labeling.

    atlas_list : list of ANTsImage types

        list containing intensity images

    label_list : list of ANTsImage types (optional)

        list containing images with segmentation labels

    submask_dilation : integer

        amount to dilate initial mask to define region on which

        we perform focused registration

    type_of_transform : string

        A linear or non-linear registration type. Mutual information metric by default.

        See Notes below for more.

    aff_metric : string

        the metric for the affine part (GC, mattes, meansquares)

    syn_metric : string

        the metric for the syn part (CC, mattes, meansquares, demons)

    syn_sampling : scalar

        the nbins or radius parameter for the syn metric

    reg_iterations : list/tuple of integers

        vector of iterations for syn. we will set the smoothing and multi-resolution parameters based on the length of this vector.

    aff_iterations : list/tuple of integers

        vector of iterations for low-dimensional registration.

    grad_step : scalar

        gradient step size (not for all tx)

    flow_sigma : scalar

        smoothing for update field

    total_sigma : scalar

        smoothing for total field

    beta : scalar

        weight sharpness, default to 2

    rad : scalar

        neighborhood radius, default to 2

    rho : scalar

        ridge penalty increases robustness to outliers but also makes image converge to average

    usecor : boolean

        employ correlation as local similarity

    r_search : scalar

        radius of search, default is 3

    nonnegative : boolean

        constrain weights to be non-negative

    no_zeroes : boolean

        this will constrain the solution only to voxels that are always non-zero in the label list

    max_lab_plus_one : boolean

        this will add max label plus one to the non-zero parts of each label where the target mask

        is greater than one.  NOTE: this will have a side effect of adding to the original label

        images that are passed to the program.  It also guarantees that every position in the

        labels have some label, rather than none.  Ie it guarantees to explicitly parcellate the

        input data.

    local_mask_transform: string

        the type of transform for the local mask alignment - usually translation,

        rigid, similarity or affine.

    output_prefix: string

        file prefix for storing output probabilityimages to disk

    verbose : boolean

        whether to show status updates

    Returns

    -------

    dictionary w/ following key/value pairs:

        `segmentation` : ANTsImage

            segmentation image

        `intensity` : ANTsImage

            intensity image

        `probabilityimages` : list of ANTsImage types

            probability map image for each label

    """

    myregion = ants.mask_image(initial_label, initial_label, which_labels)

    if myregion.max() == 0:

        myregion = ants.threshold_image(initial_label, 1, math.inf)

    myregionb = ants.threshold_image(myregion, 1, math.inf)

    myregionAroundRegion = ants.iMath(myregionb, "MD", submask_dilation)

    if target_mask is not None:

        myregionAroundRegion = myregionAroundRegion * target_mask

    croppedImage = ants.crop_image(target_image, myregionAroundRegion)

    croppedMask = ants.crop_image(myregionAroundRegion, myregionAroundRegion)

    mycroppedregion = ants.crop_image(myregion, myregionAroundRegion)

    croppedmappedImages = []

    croppedmappedSegs = []

    if verbose is True:

        print("Begin registrations:")

    for k in range(len(atlas_list)):

        if verbose is True:

            print(str(k) + "...")

        if verbose is True:

            print( "local-seg-tx: " + local_mask_transform )

        libregion = ants.mask_image(label_list[k], label_list[k], which_labels)

        initMap = ants.registration(

            mycroppedregion, libregion, type_of_transform=local_mask_transform, aff_metric=aff_metric, aff_iterations=aff_iterations, verbose=False

        )["fwdtransforms"]

        if verbose is True:

            print( "local-img-tx: " + type_of_transform )

        localReg = ants.registration(

            croppedImage,

            atlas_list[k],

            reg_iterations=reg_iterations,

            flow_sigma=flow_sigma,

            total_sigma=total_sigma,

            grad_step=grad_step,

            type_of_transform=type_of_transform,

            syn_metric=syn_metric,

            syn_sampling=syn_sampling,

            initial_transform=initMap[0],

            verbose=False,

        )

        transformedImage = ants.apply_transforms(

            croppedImage, atlas_list[k], localReg["fwdtransforms"]

        )

        transformedLabels = ants.apply_transforms(

            croppedImage,

            label_list[k],

            localReg["fwdtransforms"],

            interpolator="nearestNeighbor",

        )

        croppedmappedImages.append(transformedImage)

        croppedmappedSegs.append(transformedLabels)

    ljlf = joint_label_fusion(

        croppedImage,

        croppedMask,

        atlas_list=croppedmappedImages,

        label_list=croppedmappedSegs,

        beta=beta,

        rad=rad,

        rho=rho,

        usecor=usecor,

        r_search=r_search,

        nonnegative=nonnegative,

        no_zeroes=no_zeroes,

        max_lab_plus_one=max_lab_plus_one,

        output_prefix=output_prefix,

        verbose=verbose,

    )

    return {

        "ljlf": ljlf,

        "croppedImage": croppedImage,

        "croppedmappedImages": croppedmappedImages,

        "croppedmappedSegs": croppedmappedSegs,

    }