import os
from tempfile import mktemp

import numpy as np

import ants

__all__ = ['average_images']


def average_images( x, normalize=True, mask=None, imagetype=0, sum_image_threshold=3, return_sum_image=False, verbose=False ):
    """
    average a list of images

    images will be resampled automatically to the largest image space;
    this is not a registration so images should be in the same physical
    space to begin with.

    x : a list containing either filenames or antsImages 

    normalize : boolean

    mask : None or integer; this will perform a masked averaging which can 
        be useful when images have only partial coverage. integer greater 
        than zero will perform morphological closing.

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

    sum_image_threshold : integer
        only average regions with overlap greater than or equal to this value

    return_sum_image : boolean
        returns the average and the image that show ROI overlap; primarily for debugging

    verbose : boolean
        will print progress

    Returns
    -------
    ANTsImage

    Example
    -------
    >>> import ants
    >>> x0=[ ants.get_data('r16'), ants.get_data('r27'), ants.get_data('r62'), ants.get_data('r64') ]
    >>> x1=[]
    >>> for k in range(len(x0)):
    >>>     x1.append( ants.image_read( x0[k] ) )
    >>> avg=ants.average_images(x0)
    >>> avg1=ants.average_images(x1)
    >>> avg2=ants.average_images(x1,mask=0)
	>>> avg3=ants.average_images(x1,mask=1,normalize=True)
    """
    import numpy as np

    def gli( y, normalize=False ):
        if isinstance(y,str):
            y=ants.image_read(y)
        if normalize:
            y=y/y.mean()
        return y

    biggest=0
    biggestind=0
    for k in range( len( x ) ):
        locimg = gli( x[k], False )
        sz=np.prod( locimg.shape )
        if sz > biggest:
            biggest=sz
            biggestind=k

    avg = gli( x[biggestind], False ) * 0
    scl = float( 1.0 / len(x))
    if mask is not None:
        sumimg = gli( x[biggestind], False ) * 0

    for k in range( len( x ) ):
        if verbose and k % 20 == 0:
            print( str(k)+'...', end='',flush=True)
        locimg = gli( x[k], normalize )
        temp = ants.resample_image_to_target( locimg, avg, interp_type='linear', imagetype=imagetype )
        avg = avg + temp
        if mask is not None:
            fgmask = ants.threshold_image(temp,'Otsu',1)
            if mask > 0:
                fgmask = ants.morphology(fgmask,"close",mask)
            sumimg = sumimg + fgmask

    if return_sum_image:
        return avg * scl, sumimg
    if mask is None:
        avg = avg * scl
    else:
        nonzero = sumimg > sum_image_threshold
        tozero = sumimg <= sum_image_threshold
        avg[nonzero] = avg[nonzero] / sumimg[nonzero]
        avg[tozero] = 0
    return avg