"""
Conversion utilities.
"""
# Author: Gael Varoquaux, Alexandre Abraham, Philippe Gervais
# License: simplified BSD
import warnings
import os.path
import glob

import nilearn as ni
import numpy as np
import itertools

from joblib import Memory

from .cache_mixin import cache
from .niimg import _safe_get_data, load_niimg
from .path_finding import _resolve_globbing

from .exceptions import DimensionError
from .niimg import _get_data


def _check_fov(img, affine, shape):
    """ Return True if img's field of view correspond to given
        shape and affine, False elsewhere.
    """
    img = check_niimg(img)
    return (img.shape[:3] == shape and
            np.allclose(img.affine, affine))


def _check_same_fov(*args, **kwargs):
    """Returns True if provided images has the same field of view (shape and
    affine). Returns False or raise an error elsewhere, depending on the
    `raise_error` argument. This function can take an unlimited number of
    images as arguments or keyword arguments and raise a user-friendly
    ValueError if asked.

    Parameters
    ----------
    args : images
        Images to be checked. Images passed without keywords will be labelled
        as img_#1 in the error message (replace 1 with the appropriate index).

    kwargs : images
        Images to be checked. In case of error, images will be reference by
        their keyword name in the error message.

    raise_error : boolean, optional
        If True, an error will be raised in case of error.

    """
    raise_error = kwargs.pop('raise_error', False)
    for i, arg in enumerate(args):
        kwargs['img_#%i' % i] = arg
    errors = []
    for (a_name, a_img), (b_name, b_img) in itertools.combinations(
            kwargs.items(), 2):
        if not a_img.shape[:3] == b_img.shape[:3]:
            errors.append((a_name, b_name, 'shape'))
        if not np.allclose(a_img.affine, b_img.affine):
            errors.append((a_name, b_name, 'affine'))
    if len(errors) > 0 and raise_error:
        raise ValueError('Following field of view errors were detected:\n' +
                         '\n'.join(['- %s and %s do not have the same %s' % e
                                    for e in errors]))
    return (len(errors) == 0)


def _index_img(img, index):
    from ..image import new_img_like  # avoid circular imports

    """Helper function for check_niimg_4d."""
    return new_img_like(
        img, _get_data(img)[:, :, :, index], img.affine,
        copy_header=True)


def _iter_check_niimg(niimgs, ensure_ndim=None, atleast_4d=False,
                      target_fov=None, dtype=None,
                      memory=Memory(location=None),
                      memory_level=0, verbose=0):
    """Iterate over a list of niimgs and do sanity checks and resampling

    Parameters
    ----------
    niimgs : list of niimg or glob pattern
        Image to iterate over.

    ensure_ndim : integer, optional
        If specified, an error is raised if the data does not have the
        required dimension.

    atleast_4d : boolean, optional
        If True, any 3D image is converted to a 4D single scan.
        Default=False.

    target_fov : tuple of affine and shape, optional
       If specified, images are resampled to this field of view.

    dtype : {dtype, "auto"}, optional
        Data type toward which the data should be converted. If "auto", the
        data will be converted to int32 if dtype is discrete and float32 if it
        is continuous.

    memory : instance of joblib.Memory or string, optional
        Used to cache the masking process.
        By default, no caching is done. If a string is given, it is the
        path to the caching directory.
        Default=Memory(location=None).

    memory_level : integer, optional
        Rough estimator of the amount of memory used by caching. Higher value
        means more memory for caching. Default=0.

    verbose : integer, optional
        Indicate the level of verbosity. By default, nothing is printed.
        Default=0.

    See also
    --------
        check_niimg, check_niimg_3d, check_niimg_4d

    """
    # If niimgs is a string, use glob to expand it to the matching filenames.
    niimgs = _resolve_globbing(niimgs)

    ref_fov = None
    resample_to_first_img = False
    ndim_minus_one = ensure_ndim - 1 if ensure_ndim is not None else None
    if target_fov is not None and target_fov != "first":
        ref_fov = target_fov
    i = -1
    for i, niimg in enumerate(niimgs):
        try:
            niimg = check_niimg(
                niimg, ensure_ndim=ndim_minus_one, atleast_4d=atleast_4d,
                dtype=dtype)
            if i == 0:
                ndim_minus_one = len(niimg.shape)
                if ref_fov is None:
                    ref_fov = (niimg.affine, niimg.shape[:3])
                    resample_to_first_img = True

            if not _check_fov(niimg, ref_fov[0], ref_fov[1]):
                if target_fov is not None:
                    from nilearn import image  # we avoid a circular import
                    if resample_to_first_img:
                        warnings.warn('Affine is different across subjects.'
                                      ' Realignement on first subject '
                                      'affine forced')
                    niimg = cache(image.resample_img, memory,
                                  func_memory_level=2,
                                  memory_level=memory_level)(
                            niimg, target_affine=ref_fov[0],
                            target_shape=ref_fov[1])
                else:
                    raise ValueError(
                        "Field of view of image #%d is different from "
                        "reference FOV.\n"
                        "Reference affine:\n%r\nImage affine:\n%r\n"
                        "Reference shape:\n%r\nImage shape:\n%r\n"
                        % (i, ref_fov[0], niimg.affine, ref_fov[1],
                           niimg.shape))
            yield niimg
        except DimensionError as exc:
            # Keep track of the additional dimension in the error
            exc.increment_stack_counter()
            raise
        except TypeError as exc:
            img_name = ''
            if isinstance(niimg, str):
                img_name = " (%s) " % niimg

            exc.args = (('Error encountered while loading image #%d%s'
                         % (i, img_name),) + exc.args)
            raise

    # Raising an error if input generator is empty.
    if i == -1:
        raise ValueError("Input niimgs list is empty.")


def check_niimg(niimg, ensure_ndim=None, atleast_4d=False, dtype=None,
                return_iterator=False,
                wildcards=True):
    """Check that niimg is a proper 3D/4D niimg. Turn filenames into objects.

    Parameters
    ----------
    niimg : Niimg-like object
        See http://nilearn.github.io/manipulating_images/input_output.html
        If niimg is a string, consider it as a path to Nifti image and
        call nibabel.load on it. The '~' symbol is expanded to the user home
        folder.
        If it is an object, check if the affine attribute present and that
        nilearn.image.get_data returns a result, raise TypeError otherwise.

    ensure_ndim : integer {3, 4}, optional
        Indicate the dimensionality of the expected niimg. An
        error is raised if the niimg is of another dimensionality.

    atleast_4d : boolean, optional
        Indicates if a 3d image should be turned into a single-scan 4d niimg.
        Default=False.

    dtype : {None, dtype, "auto"}, optional
        Data type toward which the data should be converted. If "auto", the
        data will be converted to int32 if dtype is discrete and float32 if it
        is continuous. If None, data will not be converted to a new data type.
        Default=None.

    return_iterator : boolean, optional
        Returns an iterator on the content of the niimg file input.
        Default=False.

    wildcards : boolean, optional
        Use niimg as a regular expression to get a list of matching input
        filenames.
        If multiple files match, the returned list is sorted using an ascending
        order.
        If no file matches the regular expression, a ValueError exception is
        raised.
        Default=True.

    Returns
    -------
    result : 3D/4D Niimg-like object
        Result can be nibabel.Nifti1Image or the input, as-is. It is guaranteed
        that the returned object has an affine attribute and that its data can
        be retrieved with nilearn.image.get_data.

    Notes
    -----
    In nilearn, special care has been taken to make image manipulation easy.
    This method is a kind of pre-requisite for any data processing method in
    nilearn because it checks if data have a correct format and loads them if
    necessary.

    Its application is idempotent.

    See also
    --------
        _iter_check_niimg, check_niimg_3d, check_niimg_4d

    """
    from ..image import new_img_like  # avoid circular imports

    if isinstance(niimg, str):
        if wildcards and ni.EXPAND_PATH_WILDCARDS:
            # Ascending sorting + expand user path
            filenames = sorted(glob.glob(os.path.expanduser(niimg)))

            # processing filenames matching globbing expression
            if len(filenames) >= 1 and glob.has_magic(niimg):
                niimg = filenames  # iterable case
            # niimg is an existing filename
            elif [niimg] == filenames:
                niimg = filenames[0]
            # No files found by glob
            elif glob.has_magic(niimg):
                # No files matching the glob expression, warn the user
                message = ("No files matching the entered niimg expression: "
                            "'%s'.\n You may have left wildcards usage "
                           "activated: please set the global constant "
                           "'nilearn.EXPAND_PATH_WILDCARDS' to False to "
                           "deactivate this behavior.") % niimg
                raise ValueError(message)
            else:
                raise ValueError("File not found: '%s'" % niimg)
        elif not os.path.exists(niimg):
            raise ValueError("File not found: '%s'" % niimg)

    # in case of an iterable
    if hasattr(niimg, "__iter__") and not isinstance(niimg, str):
        if return_iterator:
            return _iter_check_niimg(niimg, ensure_ndim=ensure_ndim,
                                     dtype=dtype)
        return concat_niimgs(niimg, ensure_ndim=ensure_ndim, dtype=dtype)

    # Otherwise, it should be a filename or a SpatialImage, we load it
    niimg = load_niimg(niimg, dtype=dtype)

    if ensure_ndim == 3 and len(niimg.shape) == 4 and niimg.shape[3] == 1:
        # "squeeze" the image.
        data = _safe_get_data(niimg)
        affine = niimg.affine
        niimg = new_img_like(niimg, data[:, :, :, 0], affine)
    if atleast_4d and len(niimg.shape) == 3:
        data = _get_data(niimg).view()
        data.shape = data.shape + (1, )
        niimg = new_img_like(niimg, data, niimg.affine)

    if ensure_ndim is not None and len(niimg.shape) != ensure_ndim:
        raise DimensionError(len(niimg.shape), ensure_ndim)

    if return_iterator:
        return (_index_img(niimg, i) for i in range(niimg.shape[3]))

    return niimg


def check_niimg_3d(niimg, dtype=None):
    """Check that niimg is a proper 3D niimg-like object and load it.

    Parameters
    ----------
    niimg : Niimg-like object
        See http://nilearn.github.io/manipulating_images/input_output.html
        If niimg is a string, consider it as a path to Nifti image and
        call nibabel.load on it.
        If it is an object, check if the affine attribute present and that
        nilearn.image.get_data returns a result, raise TypeError otherwise.

    dtype : {dtype, "auto"}, optional
        Data type toward which the data should be converted. If "auto", the
        data will be converted to int32 if dtype is discrete and float32 if it
        is continuous.

    Returns
    -------
    result : 3D Niimg-like object
        Result can be nibabel.Nifti1Image or the input, as-is. It is guaranteed
        that the returned object has an affine attribute and that its data can
        be retrieved with nilearn.image.get_data.

    Notes
    -----
    In nilearn, special care has been taken to make image manipulation easy.
    This method is a kind of pre-requisite for any data processing method in
    nilearn because it checks if data have a correct format and loads them if
    necessary.

    Its application is idempotent.

    """
    return check_niimg(niimg, ensure_ndim=3, dtype=dtype)


def check_niimg_4d(niimg, return_iterator=False, dtype=None):
    """Check that niimg is a proper 4D niimg-like object and load it.

    Parameters
    ----------
    niimg : 4D Niimg-like object
        See http://nilearn.github.io/manipulating_images/input_output.html
        If niimgs is an iterable, checks if data is really 4D. Then,
        considering that it is a list of niimg and load them one by one.
        If niimg is a string, consider it as a path to Nifti image and
        call nibabel.load on it.
        If it is an object, check if the affine attribute present and that
        nilearn.image.get_data returns a result, raise TypeError otherwise.

    dtype : {dtype, "auto"}, optional
        Data type toward which the data should be converted. If "auto", the
        data will be converted to int32 if dtype is discrete and float32 if it
        is continuous.

    return_iterator : boolean, optional
        If True, an iterator of 3D images is returned. This reduces the memory
        usage when `niimgs` contains 3D images.
        If False, a single 4D image is returned. When `niimgs` contains 3D
        images they are concatenated together.
        Default=False.

    Returns
    -------
    niimg: 4D nibabel.Nifti1Image or iterator of 3D nibabel.Nifti1Image

    Notes
    -----
    This function is the equivalent to check_niimg_3d() for Niimg-like objects
    with a session level.

    Its application is idempotent.

    """
    return check_niimg(niimg, ensure_ndim=4, return_iterator=return_iterator,
                       dtype=dtype)


def concat_niimgs(niimgs, dtype=np.float32, ensure_ndim=None,
                  memory=Memory(location=None), memory_level=0,
                  auto_resample=False, verbose=0):
    """Concatenate a list of 3D/4D niimgs of varying lengths.

    The niimgs list can contain niftis/paths to images of varying dimensions
    (i.e., 3D or 4D) as well as different 3D shapes and affines, as they
    will be matched to the first image in the list if auto_resample=True.

    Parameters
    ----------
    niimgs : iterable of Niimg-like objects or glob pattern
        See http://nilearn.github.io/manipulating_images/input_output.html
        Niimgs to concatenate.

    dtype : numpy dtype, optional
        The dtype of the returned image. Default=np.float32.

    ensure_ndim : integer, optional
        Indicate the dimensionality of the expected niimg. An
        error is raised if the niimg is of another dimensionality.

    auto_resample : boolean, optional
        Converts all images to the space of the first one.
        Default=False.

    verbose : int, optional
        Controls the amount of verbosity (0 means no messages).
        Default=0.

    memory : instance of joblib.Memory or string, optional
        Used to cache the resampling process.
        By default, no caching is done. If a string is given, it is the
        path to the caching directory.
        Default=Memory(location=None).

    memory_level : integer, optional
        Rough estimator of the amount of memory used by caching. Higher value
        means more memory for caching. Default=0.

    Returns
    -------
    concatenated : nibabel.Nifti1Image
        A single image.

    See Also
    --------
    nilearn.image.index_img

    """
    from ..image import new_img_like  # avoid circular imports

    target_fov = 'first' if auto_resample else None

    # We remove one to the dimensionality because of the list is one dimension.
    ndim = None
    if ensure_ndim is not None:
        ndim = ensure_ndim - 1

    # If niimgs is a string, use glob to expand it to the matching filenames.
    niimgs = _resolve_globbing(niimgs)

    # First niimg is extracted to get information and for new_img_like
    first_niimg = None

    iterator, literator = itertools.tee(iter(niimgs))
    try:
        first_niimg = check_niimg(next(literator), ensure_ndim=ndim)
    except StopIteration:
        raise TypeError('Cannot concatenate empty objects')
    except DimensionError as exc:
        # Keep track of the additional dimension in the error
        exc.increment_stack_counter()
        raise

    # If no particular dimensionality is asked, we force consistency wrt the
    # first image
    if ndim is None:
        ndim = len(first_niimg.shape)

    if ndim not in [3, 4]:
        raise TypeError('Concatenated images must be 3D or 4D. You gave a '
                        'list of %dD images' % ndim)

    lengths = [first_niimg.shape[-1] if ndim == 4 else 1]
    for niimg in literator:
        # We check the dimensionality of the niimg
        try:
            niimg = check_niimg(niimg, ensure_ndim=ndim)
        except DimensionError as exc:
            # Keep track of the additional dimension in the error
            exc.increment_stack_counter()
            raise
        lengths.append(niimg.shape[-1] if ndim == 4 else 1)

    target_shape = first_niimg.shape[:3]
    if dtype == None:
        dtype = _get_data(first_niimg).dtype
    data = np.ndarray(target_shape + (sum(lengths), ),
                      order="F", dtype=dtype)
    cur_4d_index = 0
    for index, (size, niimg) in enumerate(zip(lengths, _iter_check_niimg(
            iterator, atleast_4d=True, target_fov=target_fov,
            memory=memory, memory_level=memory_level))):

        if verbose > 0:
            if isinstance(niimg, str):
                nii_str = "image " + niimg
            else:
                nii_str = "image #" + str(index)
            print("Concatenating {0}: {1}".format(index + 1, nii_str))

        data[..., cur_4d_index:cur_4d_index + size] = _get_data(niimg)
        cur_4d_index += size

    return new_img_like(first_niimg, data, first_niimg.affine, copy_header=True)
