__all__ = ['pad_image']

import math

import ants

from ants.decorators import image_method

from ants.internal import get_lib_fn

@image_method

def pad_image(image, shape=None, pad_width=None, value=0.0, return_padvals=False):

    """

    Pad an image to have the given shape or to be isotropic.

    Arguments

    ---------

    image : ANTsImage

        image to pad

    shape : tuple

        - if shape is given, the image will be padded in each dimension

          until it has this shape

        - if shape and pad_width are both None, the image will be padded along

          each dimension to match the largest existing dimension so that it has

          isotropic dimensions.

    pad_width : list of integers or list-of-list of integers

        How much to pad in each direction. If a single list is

        supplied (e.g., [4,4,4]), then the image will be padded by half

        that amount on both sides. If a list-of-list is supplied

        (e.g., [(0,4),(0,4),(0,4)]), then the image will be

        padded unevenly on the different sides

    pad_value : scalar

        value with which image will be padded

    Example

    -------

    >>> import ants

    >>> img = ants.image_read(ants.get_data('r16'))

    >>> img2 = ants.pad_image(img, shape=(300,300))

    >>> mni = ants.image_read(ants.get_data('mni'))

    >>> mni2 = ants.pad_image(mni)

    >>> mni3 = ants.pad_image(mni, pad_width=[(0,4),(0,4),(0,4)])

    >>> mni4 = ants.pad_image(mni, pad_width=(4,4,4))

    """

    inpixeltype = image.pixeltype

    ndim = image.dimension

    if image.pixeltype != 'float':

        image = image.clone('float')

    if pad_width is None:

        if shape is None:

            shape = [max(image.shape)] * image.dimension

        lower_pad_vals = [math.floor(max(ns-os,0)/2) for os,ns in zip(image.shape, shape)]

        upper_pad_vals = [math.ceil(max(ns-os,0)/2) for os,ns in zip(image.shape, shape)]

    else:

        if shape is not None:

            raise ValueError('Cannot give both `shape` and `pad_width`. Pick one!')

        if len(pad_width) != image.dimension:

            raise ValueError('Must give pad width for each image dimension')

        lower_pad_vals = []

        upper_pad_vals = []

        for p in pad_width:

            if isinstance(p, (list, tuple)):

                lower_pad_vals.append(p[0])

                upper_pad_vals.append(p[1])

            else:

                lower_pad_vals.append(math.floor(p/2))

                upper_pad_vals.append(math.ceil(p/2))

    libfn = get_lib_fn('padImage')

    itkimage = libfn(image.pointer, lower_pad_vals, upper_pad_vals, value)

    new_image = ants.from_pointer(itkimage).clone(inpixeltype)

    if return_padvals:

        return new_image, lower_pad_vals, upper_pad_vals

    else:

        return new_image