"""
Atropos segmentation
"""



__all__ = ['atropos']

import os
import glob
import warnings
from tempfile import mktemp


import ants
from ants.internal import get_lib_fn, get_pointer_string, process_arguments


def atropos(a, x, i='Kmeans[3]', m='[0.2,1x1]', c='[5,0]',
            priorweight=0.25, **kwargs):
    """
    A finite mixture modeling (FMM) segmentation approach with possibilities
    for specifying prior constraints. These prior constraints include the
    specification of a prior label image, prior probability images (one for
    each class), and/or an MRF prior to enforce spatial smoothing of the
    labels. Similar algorithms include FAST and SPM. atropos can also perform
    multivariate segmentation if you pass a list of images in: e.g. a=(img1,img2).

    ANTsR function: `atropos`

    Arguments
    ---------
    a : ANTsImage or list/tuple of ANTsImage types
        One or more scalar images to segment. If priors are not used,
        the intensities of the first image are used to order the classes
        in the segmentation output, from lowest to highest intensity. Otherwise
        the order of the classes is dictated by the order of the prior images.

    x : ANTsImage
        mask image.

    i : string
        initialization usually KMeans[N] for N classes or a list of N prior
        probability images. See Atropos in ANTs for full set of options.

    m : string
        mrf parameters as a string, usually "[smoothingFactor,radius]" where
        smoothingFactor determines the amount of smoothing and radius determines
        the MRF neighborhood, as an ANTs style neighborhood vector eg "1x1x1"
        for a 3D image. The radius must match the dimensionality of the image,
        eg 1x1 for 2D and The default in ANTs is smoothingFactor=0.3 and
        radius=1. See Atropos for more options.

    c : string
        convergence parameters, "[numberOfIterations,convergenceThreshold]".
        A threshold of 0 runs the full numberOfIterations, otherwise Atropos
        tests convergence by comparing the mean maximum posterior probability
        over the whole region of interest defined by the mask x.

    priorweight : scalar
        usually 0 (priors used for initialization only), 0.25 or 0.5.

    kwargs : keyword arguments
        more parameters, see Atropos help in ANTs

    Returns
    -------
    dictionary with the following key/value pairs:
        `segmentation`: ANTsImage
            actually segmented image
        `probabilityimages` : list of ANTsImage types
            one image for each segmentation class

    Example
    -------
    >>> import ants
    >>> img = ants.image_read(ants.get_ants_data('r16'))
    >>> img = ants.resample_image(img, (64,64), 1, 0)
    >>> mask = ants.get_mask(img)
    >>> ants.atropos( a = img, m = '[0.2,1x1]', c = '[2,0]',  i = 'kmeans[3]', x = mask )
    >>> seg2 = ants.atropos( a = img, m = '[0.2,1x1]', c = '[2,0]', i = seg['probabilityimages'], x = mask, priorweight=0.25 )
    """
    probs = mktemp(prefix='antsr', suffix='prob%02d.nii.gz')
    tdir = probs.replace(os.path.basename(probs),'')
    probsbase = os.path.basename(probs)
    searchpattern = probsbase.replace('%02d', '*')

    ct = 0
    if isinstance(i, (list,tuple)) and (len(i) > 1):
        while ct < len(i):
            probchar = str(ct+1)
            if ct < 9:
                probchar = '0%s' % probchar
            tempfn = probs.replace('%02d', probchar)
            ants.image_write(i[ct], tempfn)
            ct += 1
        i = 'PriorProbabilityImages[%s,%s,%s]' % (str(len(i)), probs, str(priorweight))

    if isinstance(a, (list,tuple)):
        outimg = a[0].clone('unsigned int')
    else:
        outimg = a.clone('unsigned int')

    mydim = outimg.dimension
    outs = '[%s,%s]' % (get_pointer_string(outimg), probs)
    mymask = x.clone('unsigned int')

    if (not isinstance(a, (list,tuple))) or (len(a) == 1):
        myargs = {
            'd': mydim,
            'a': a,
            'm-MULTINAME-0': m,
            'o': outs,
            'c': c,
            'm-MULTINAME-1': m,
            'i': i,
            'x': mymask
        }
        for k, v in kwargs.items():
            myargs[k] = v

    elif isinstance(a, (list, tuple)):
        if len(a) > 6:
            print('more than 6 input images not really supported, using first 6')
            a = a[:6]
        myargs = {
            'd': mydim,
            'm-MULTINAME-0': m,
            'o': outs,
            'c': c,
            'm-MULTINAME-1': m,
            'i': i,
            'x': mymask
        }
        for aa_idx, aa in enumerate(a):
            myargs['a-MULTINAME-%i'%aa_idx] = aa

    processed_args = process_arguments(myargs)
    libfn = get_lib_fn('Atropos')
    retval = libfn(processed_args)

    if retval != 0:
        raise Exception(f"Atropos exited with non-zero status {retval}. Run with verbose=1 to see error messages")

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

    if probsout is None or len(probsout) == 0:
        raise Exception('No Atropos output probability images found. Run with verbose=1 to see error messages')

    probsout.sort()
    probimgs = [ants.image_read(probsout[0])]
    for idx in range(1, len(probsout)):
        probimgs.append(ants.image_read(probsout[idx]))

    outimg = outimg.clone('float')
    return {'segmentation': outimg,
            'probabilityimages': probimgs}