__all__ = ['weingarten_image_curvature']

import numpy as np

import ants

from ants.decorators import image_method

from ants.internal import get_lib_fn

@image_method

def weingarten_image_curvature(image, sigma=1.0, opt='mean'):

    """

    Uses the weingarten map to estimate image mean or gaussian curvature

    ANTsR function: `weingartenImageCurvature`

    Arguments

    ---------

    image : ANTsImage

        image from which curvature is calculated

    sigma : scalar

        smoothing parameter

    opt : string

        mean by default, otherwise `gaussian` or `characterize`

    Returns

    -------

    ANTsImage

    Example

    -------

    >>> import ants

    >>> image = ants.image_read(ants.get_ants_data('mni')).resample_image((3,3,3))

    >>> imagecurv = ants.weingarten_image_curvature(image)

    """

    if image.dimension not in {2,3}:

        raise ValueError('image must be 2D or 3D')

    if image.dimension == 2:

        d = image.shape

        temp = np.zeros(list(d)+[10])

        for k in range(1,7):

            voxvals = image[:d[0],:d[1]].numpy()

            temp[:d[0],:d[1],k] = voxvals

        temp = ants.from_numpy(temp)

        myspc = image.spacing

        myspc = list(myspc) + [min(myspc)]

        temp.set_spacing(myspc)

        temp = temp.clone('float')

    else:

        temp = image.clone('float')

    optnum = 0

    if opt == 'gaussian': 

        optnum = 6

    if opt == 'characterize':

        optnum = 5

    libfn = get_lib_fn('weingartenImageCurvature')

    mykout = libfn(temp.pointer, sigma, optnum)

    mykout = ants.from_pointer(mykout)

    if image.dimension == 3:

        return mykout

    elif image.dimension == 2:

        subarr = ants.from_numpy(mykout.numpy()[:,:,4])

        return ants.copy_image_info(image, subarr)