__all__ = ['slice_image']

import math
import numpy as np
import ants
from ants.decorators import image_method
from ants.internal import get_lib_fn

@image_method
def slice_image(image, axis, idx, collapse_strategy=0):
    """
    Slice an image.

    Arguments
    ---------
    axis: integer 
        Which axis.

    idx: integer
        Which slice number.    

    collapse_strategy:  integer
        Collapse strategy for sub-matrix: 0, 1, or 2.  0: collapse to sub-matrix 
        if positive-definite.  Otherwise throw an exception. Default.  1: Collapse 
        to identity.  2:  Collapse to sub-matrix if positive definite. Otherwise
        collapse to identity.

    Example
    -------
    >>> import ants
    >>> mni = ants.image_read(ants.get_data('mni'))
    >>> mni2 = ants.slice_image(mni, axis=1, idx=100)
    """
    if image.has_components:
        ilist = ants.split_channels(image)
        if image.dimension == 2:
            return np.stack(tuple([i.slice_image(axis, idx, collapse_strategy) for i in ilist]), axis=-1)
        else:
            return ants.merge_channels([i.slice_image(axis, idx, collapse_strategy) for i in ilist])
    
    if axis == -1:
        axis = image.dimension - 1
        
    if axis > (image.dimension - 1) or axis < 0:
        raise Exception('The axis must be between 0 and image.dimension - 1')
        
    if image.dimension == 2:
        if axis == 0:
            return image[idx,:]
        elif axis == 1:
            return image[:,idx]
        raise Exception('Parameters not understood for 2D image.')
        
    if collapse_strategy != 0 and collapse_strategy != 1 and collapse_strategy != 2:
        raise ValueError('collapse_strategy must be 0, 1, or 2.') 

    inpixeltype = image.pixeltype
    ndim = image.dimension
    if image.pixeltype != 'float':
        image = image.clone('float')

    libfn = get_lib_fn('sliceImage')
    itkimage = libfn(image.pointer, axis, idx, collapse_strategy)

    return ants.from_pointer(itkimage).clone(inpixeltype)