

__all__ = ['apply_transforms',
           'apply_transforms_to_points']

import os

import ants
from ants.internal import get_lib_fn, process_arguments


def apply_transforms(fixed, moving, transformlist,
                     interpolator='linear', imagetype=0,
                     whichtoinvert=None, compose=None,
                     defaultvalue=0, singleprecision=False, verbose=False, **kwargs):
    """
    Apply a transform list to map an image from one domain to another.
    In image registration, one computes mappings between (usually) pairs
    of images. These transforms are often a sequence of increasingly
    complex maps, e.g. from translation, to rigid, to affine to deformation.
    The list of such transforms is passed to this function to interpolate one
    image domain into the next image domain, as below. The order matters
    strongly and the user is advised to familiarize with the standards
    established in examples.

    ANTsR function: `antsApplyTransforms`

    Arguments
    ---------
    fixed : ANTsImage
        fixed image defining domain into which the moving image is transformed. The output will
        have the same pixel type as this image.

    moving : AntsImage
        moving image to be mapped to fixed space.

    transformlist : list of strings
        list of transforms generated by ants.registration where each transform is a filename.

    interpolator : string
        Choice of interpolator. Supports partial matching.
            linear
            nearestNeighbor
            multiLabel for label images (deprecated, prefer genericLabel)
            gaussian
            bSpline
            cosineWindowedSinc
            welchWindowedSinc
            hammingWindowedSinc
            lanczosWindowedSinc
            genericLabel use this for label images

    imagetype : integer
        choose 0/1/2/3 mapping to scalar/vector/tensor/time-series

    whichtoinvert : list of booleans (optional)
        Must be same length as transformlist.
        whichtoinvert[i] is True if transformlist[i] is a matrix,
        and the matrix should be inverted. If transformlist[i] is a
        warp field, whichtoinvert[i] must be False.
        If the transform list is a matrix followed by a warp field,
        whichtoinvert defaults to (True,False). Otherwise it defaults
        to [False]*len(transformlist)).

    compose : string (optional)
        if it is a string pointing to a valid file location,
        this will force the function to return a composite transformation filename.

    defaultvalue : scalar
        Default voxel value for mappings outside the image domain.

    singleprecision : boolean
        if True, use float32 for computations. This is useful for reducing memory
        usage for large datasets, at the cost of precision.

    verbose : boolean
        print command and run verbose application of transform.

    kwargs : keyword arguments
        extra parameters

    Returns
    -------
    ANTsImage or string (transformation filename)

    Example
    -------
    >>> import ants
    >>> fixed = ants.image_read( ants.get_ants_data('r16') )
    >>> moving = ants.image_read( ants.get_ants_data('r64') )
    >>> fixed = ants.resample_image(fixed, (64,64), 1, 0)
    >>> moving = ants.resample_image(moving, (64,64), 1, 0)
    >>> mytx = ants.registration(fixed=fixed , moving=moving ,
                                 type_of_transform = 'SyN' )
    >>> mywarpedimage = ants.apply_transforms( fixed=fixed, moving=moving,
                                               transformlist=mytx['fwdtransforms'] )
    """

    if not isinstance(transformlist, (tuple, list)) and (transformlist is not None):
        transformlist = [transformlist]

    accepted_interpolators = {"linear", "nearestNeighbor", "multiLabel", "gaussian",
                        "bSpline", "cosineWindowedSinc", "welchWindowedSinc",
                        "hammingWindowedSinc", "lanczosWindowedSinc", "genericLabel"}

    if interpolator not in accepted_interpolators:
        raise ValueError('interpolator not supported - see %s' % accepted_interpolators)

    args = [fixed, moving, transformlist, interpolator]

    output_pixel_type = 'float' if singleprecision else 'double'

    if not isinstance(fixed, str):
        if ants.is_image(fixed) and ants.is_image(moving):
            for tl_path in transformlist:
                if not os.path.exists(tl_path):
                    raise Exception('Transform %s does not exist' % tl_path)

            inpixeltype = fixed.pixeltype
            fixed = fixed.clone(output_pixel_type)
            moving = moving.clone(output_pixel_type)
            warpedmovout = moving.clone(output_pixel_type)
            f = fixed
            m = moving
            if (moving.dimension == 4) and (fixed.dimension == 3) and (imagetype == 0):
                raise Exception('Set imagetype 3 to transform time series images.')

            wmo = warpedmovout
            mytx = []
            if whichtoinvert is None or (isinstance(whichtoinvert, (tuple,list)) and (sum([w is not None for w in whichtoinvert])==0)):
                if (len(transformlist) == 2) and ('.mat' in transformlist[0]) and ('.mat' not in transformlist[1]):
                    whichtoinvert = (True, False)
                else:
                    whichtoinvert = tuple([False]*len(transformlist))

            if len(whichtoinvert) != len(transformlist):
                raise ValueError('Transform list and inversion list must be the same length')

            for i in range(len(transformlist)):
                ismat = False
                if '.mat' in transformlist[i]:
                    ismat = True
                if whichtoinvert[i] and (not ismat):
                    raise ValueError('Cannot invert transform %i (%s) because it is not a matrix' % (i, transformlist[i]))
                if whichtoinvert[i]:
                    mytx = mytx + ['-t', '[%s,1]' % (transformlist[i])]
                else:
                    mytx = mytx + ['-t', transformlist[i]]

            if compose is None:
                args = ['-d', fixed.dimension,
                        '-i', m,
                        '-o', wmo,
                        '-r', f,
                        '-n', interpolator]
                args = args + mytx
            if compose:
                tfn = '%scomptx.nii.gz' % compose if not compose.endswith('.h5') else compose
            else:
                tfn = 'NA'
            if compose is not None:
                mycompo = '[%s,1]' % tfn
                args = ['-d', fixed.dimension,
                        '-i', m,
                        '-o', mycompo,
                        '-r', f,
                        '-n', interpolator]
                args = args + mytx

            myargs = process_arguments(args)

            myverb = int(verbose)
            if verbose:
                print(myargs)

            processed_args = myargs + ['-z', str(1), '-v', str(myverb), '--float', str(int(singleprecision)), '-e', str(imagetype), '-f', str(defaultvalue)]
            libfn = get_lib_fn('antsApplyTransforms')
            libfn(processed_args)

            if compose is None:
                return warpedmovout.clone(inpixeltype)
            else:
                if os.path.exists(tfn):
                    return tfn
                else:
                    return None

        else:
            return 1
    else:
        args = args + ['-z', str(1), '--float', str(int(singleprecision)), '-e', imagetype, '-f', defaultvalue]
        processed_args = process_arguments(args)
        libfn = get_lib_fn('antsApplyTransforms')
        libfn(processed_args)






def apply_transforms_to_points( dim, points, transformlist,
                     whichtoinvert=None, verbose=False ):
    """
     Apply a transform list to map a pointset from one domain to
     another. In registration, one computes mappings between pairs of
     domains. These transforms are often a sequence of increasingly
     complex maps, e.g. from translation, to rigid, to affine to
     deformation.  The list of such transforms is passed to this
     function to interpolate one image domain into the next image
     domain, as below.  The order matters strongly and the user is
     advised to familiarize with the standards established in examples.
     Importantly, point mapping goes the opposite direction of image
     mapping, for both reasons of convention and engineering.

    ANTsR function: `antsApplyTransformsToPoints`

    Arguments
    ---------
    dim: integer
         dimensionality of the transformation.

    points: data frame
          moving point set with n-points in rows of at least dim
          columns - we maintain extra information in additional
          columns. this should be a data frame with columns names x, y, z, t.

    transformlist : list of strings
        list of transforms generated by ants.registration where each transform is a filename.

    whichtoinvert : list of booleans (optional)
        Must be same length as transformlist.
        whichtoinvert[i] is True if transformlist[i] is a matrix,
        and the matrix should be inverted. If transformlist[i] is a
        warp field, whichtoinvert[i] must be False.
        If the transform list is a matrix followed by a warp field,
        whichtoinvert defaults to (True,False). Otherwise it defaults
        to [False]*len(transformlist)).

    verbose : boolean

    Returns
    -------
    data frame of transformed points

    Example
    -------
    >>> import ants
    >>> fixed = ants.image_read( ants.get_ants_data('r16') )
    >>> moving = ants.image_read( ants.get_ants_data('r27') )
    >>> reg = ants.registration( fixed, moving, 'Affine' )
    >>> d = {'x': [128, 127], 'y': [101, 111]}
    >>> pts = pd.DataFrame(data=d)
    >>> ptsw = ants.apply_transforms_to_points( 2, pts, reg['fwdtransforms'])
    """

    if not isinstance(transformlist, (tuple, list)) and (transformlist is not None):
        transformlist = [transformlist]

    args = [dim, points, transformlist, whichtoinvert]

    for tl_path in transformlist:
        if not os.path.exists(tl_path):
            raise Exception('Transform %s does not exist' % tl_path)

    mytx = []

    if whichtoinvert is None or (isinstance(whichtoinvert, (tuple,list)) and (sum([w is not None for w in whichtoinvert])==0)):
        if (len(transformlist) == 2) and ('.mat' in transformlist[0]) and ('.mat' not in transformlist[1]):
            whichtoinvert = (True, False)
        else:
            whichtoinvert = tuple([False]*len(transformlist))

    if len(whichtoinvert) != len(transformlist):
        raise ValueError('Transform list and inversion list must be the same length')

    for i in range(len(transformlist)):
        ismat = False
        if '.mat' in transformlist[i]:
            ismat = True
        if whichtoinvert[i] and (not ismat):
            raise ValueError('Cannot invert transform %i (%s) because it is not a matrix' % (i, transformlist[i]))
        if whichtoinvert[i]:
            mytx = mytx + ['-t', '[%s,1]' % (transformlist[i])]
        else:
            mytx = mytx + ['-t', transformlist[i]]
    if dim == 2:
        pointsSub = points[['x','y']]
    if dim == 3:
        pointsSub = points[['x','y','z']]
    if dim == 4:
        pointsSub = points[['x','y','z','t']]
    pointImage = ants.make_image( pointsSub.shape, pointsSub.values.flatten())
    pointsOut = pointImage.clone()
    args = ['-d', dim,
            '-i', pointImage,
            '-o', pointsOut ]
    args = args + mytx
    myargs = process_arguments(args)

    myverb = int(verbose)
    if verbose:
        print(myargs)

    processed_args = myargs + [ '-f', str(1), '--precision', str(0)]
    libfn = get_lib_fn('antsApplyTransformsToPoints')
    libfn(processed_args)
    mynp = pointsOut.numpy()
    pointsOutDF = points.copy()
    pointsOutDF['x'] = mynp[:,0]
    if dim >= 2:
        pointsOutDF['y'] = mynp[:,1]
    if dim >= 3:
        pointsOutDF['z'] = mynp[:,2]
    if dim >= 4:
        pointsOutDF['t'] = mynp[:,3]
    return pointsOutDF
