"""
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,
    }
