"""
Kelly Kapowski algorithm with computing cortical thickness
"""

__all__ = ['kelly_kapowski']

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


def kelly_kapowski(s, g, w, its=45, r=0.025, m=1.5, gm_label=2, wm_label=3, **kwargs):
    """
    Compute cortical thickness using the DiReCT algorithm.

    Diffeomorphic registration-based cortical thickness based on probabilistic
    segmentation of an image.  This is an optimization algorithm.


    Arguments
    ---------
    s : ANTsimage
        segmentation image

    g : ANTsImage
        gray matter probability image

    w : ANTsImage
        white matter probability image

    its : integer
        convergence params - controls iterations

    r : scalar
        gradient descent update parameter

    m : scalar
        gradient field smoothing parameter

    gm_label : integer
        label for gray matter in the segmentation image

    wm_label : integer
        label for white matter in the segmentation image

    kwargs : keyword arguments
        anything else, see KellyKapowski help in ANTs

    Returns
    -------
    ANTsImage

    Example
    -------
    >>> import ants
    >>> img = ants.image_read( ants.get_ants_data('r16') ,2)
    >>> img = ants.resample_image(img, (64,64),1,0)
    >>> mask = ants.get_mask( img )
    >>> segs = ants.kmeans_segmentation( img, k=3, kmask = mask)
    >>> thick = ants.kelly_kapowski(s=segs['segmentation'], g=segs['probabilityimages'][1],
                                    w=segs['probabilityimages'][2], its=45,
                                    r=0.5, m=1)
    """
    if ants.is_image(s):
        s = s.clone('unsigned int')

    d = s.dimension
    outimg = g.clone() * 0.0
    kellargs = {'d': d,
                's': "[{},{},{}]".format(get_pointer_string(s),gm_label,wm_label),
                'g': g,
                'w': w,
                'c': "[{}]".format(its),
                'r': r,
                'm': m,
                'o': outimg}
    for k, v in kwargs.items():
        kellargs[k] = v

    processed_kellargs = process_arguments(kellargs)

    libfn = get_lib_fn('KellyKapowski')
    libfn(processed_kellargs)

    # Check thickness is not still all zeros
    if outimg.sum() == 0.0:
        raise RuntimeError("KellyKapowski failed to compute thickness")

    return outimg