from tempfile import mktemp

import os

import ants

from ants.internal import get_lib_fn, process_arguments

__all__ = ['average_affine_transform', 

           'average_affine_transform_no_rigid']

def _average_affine_transform_driver(transformlist, referencetransform=None, funcname="AverageAffineTransform"):

    """

    takes a list of transforms (files at the moment)

    and returns the average

    """

    # AverageAffineTransform deals with transform files,

    # so this function will need to deal with already

    # loaded files. Doesn't look like the magic

    # available for images has been added for transforms.

    res_temp_file = mktemp(suffix='.mat')

    # could do some stuff here to cope with transform lists that

    # aren't files

    # load one of the transforms to figure out the dimension

    tf = ants.read_transform(transformlist[0])

    if referencetransform is None:

        args = [tf.dimension, res_temp_file] + transformlist

    else:

        args = [tf.dimension, res_temp_file] + ['-R',  referencetransform] + transformlist

    pargs = process_arguments(args)

    print(pargs)

    libfun = get_lib_fn(funcname)

    status = libfun(pargs)

    res = ants.read_transform(res_temp_file)

    os.remove(res_temp_file)

    return res

def average_affine_transform(transformlist, referencetransform=None):

    return _average_affine_transform_driver(transformlist, referencetransform, "AverageAffineTransform")

def average_affine_transform_no_rigid(transformlist, referencetransform=None):

    return _average_affine_transform_driver(transformlist, referencetransform, "AverageAffineTransformNoRigid")