__all__ = ['mask_image']

import numpy as np

import ants
from ants.decorators import image_method

@image_method
def mask_image(image, mask, level=1, binarize=False):
    """
    Mask an input image by a mask image.  If the mask image has multiple labels,
    it is possible to specify which label(s) to mask at.

    ANTsR function: `maskImage`

    Arguments
    ---------
    image : ANTsImage
        Input image.

    mask : ANTsImage
        Mask or label image.

    level : scalar or tuple of scalars
        Level(s) at which to mask image. If vector or list of values, output image is non-zero at all locations where label image matches any of the levels specified.

    binarize : boolean
        whether binarize the output image

    Returns
    -------
    ANTsImage

    Example
    -------
    >>> import ants
    >>> myimage = ants.image_read(ants.get_ants_data('r16'))
    >>> mask = ants.get_mask(myimage)
    >>> myimage_mask = ants.mask_image(myimage, mask, 3)
    >>> seg = ants.kmeans_segmentation(myimage, 3)
    >>> myimage_mask = ants.mask_image(myimage, seg['segmentation'], (1,3))
    """
    leveluse = level
    if type(leveluse) is np.ndarray:
        leveluse = level.tolist()
    if type(leveluse) is int or type(leveluse) is float:
        leveluse = [level]
    image_out = image.clone() * 0
    for mylevel in leveluse:
        temp = ants.threshold_image(mask, mylevel, mylevel)
        if binarize:
            image_out = image_out + temp
        else:
            image_out = image_out + temp * image
    return image_out