#  Copyright 2017 Intel Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""I/O functionality."""

__all__ = [
    "load_boolean_mask",
    "load_images",
    "load_images_from_dir",
    "load_labels",
    "save_as_nifti_file",
]

from pathlib import Path
from typing import Callable, Iterable, List, Union

import logging
import nibabel as nib
import numpy as np

from nibabel.nifti1 import Nifti1Pair
from nibabel.spatialimages import SpatialImage

from .image import SingleConditionSpec

logger = logging.getLogger(__name__)


def load_images_from_dir(in_dir: Union[str, Path], suffix: str = "nii.gz",
                         ) -> Iterable[SpatialImage]:
    """Load images from directory.

    For efficiency, returns an iterator, not a sequence, so the results cannot
    be accessed by indexing.

    For every new iteration through the images, load_images_from_dir must be
    called again.

    Parameters
    ----------
    in_dir:
        Path to directory.
    suffix:
        Only load images with names that end like this.

    Yields
    ------
    SpatialImage
        Image.

    Examples
    --------
    >>> images = list(io.load_images_from_dir("face_scene", "bet.nii.gz"))
    """
    if isinstance(in_dir, str):
        in_dir = Path(in_dir)
    files = sorted(in_dir.glob("*" + suffix))
    for f in files:
        logger.debug(
            'Starting to read file %s', f
        )
        yield nib.load(str(f))


def load_images(image_paths: Iterable[Union[str, Path]]
                ) -> Iterable[SpatialImage]:
    """Load images from paths.

    For efficiency, returns an iterator, not a sequence, so the results cannot
    be accessed by indexing.

    For every new iteration through the images, load_images must be called
    again.

    Parameters
    ----------
    image_paths:
        Paths to images.

    Yields
    ------
    SpatialImage
        Image.
    """
    for image_path in image_paths:
        if isinstance(image_path, Path):
            string_path = str(image_path)
        else:
            string_path = image_path
        logger.debug(
            'Starting to read file %s', string_path
        )
        yield nib.load(string_path)


def load_boolean_mask(path: Union[str, Path],
                      predicate: Callable[[np.ndarray], np.ndarray] = None
                      ) -> np.ndarray:
    """Load boolean nibabel.SpatialImage mask.

    Parameters
    ----------
    path
        Mask path.
    predicate
        Callable used to create boolean values, e.g. a threshold function
        ``lambda x: x > 50``.

    Returns
    -------
    np.ndarray
        Boolean array corresponding to mask.
    """
    if not isinstance(path, str):
        path = str(path)
    data = nib.load(path).get_data()
    if predicate is not None:
        mask = predicate(data)
    else:
        mask = data.astype(np.bool)
    return mask


def load_labels(path: Union[str, Path]) -> List[SingleConditionSpec]:
    """Load labels files.

    Parameters
    ----------
    path
        Path of labels file.

    Returns
    -------
    List[SingleConditionSpec]
        List of SingleConditionSpec stored in labels file.
    """
    condition_specs = np.load(str(path))
    return [c.view(SingleConditionSpec) for c in condition_specs]


def save_as_nifti_file(data: np.ndarray, affine: np.ndarray,
                       path: Union[str, Path]) -> None:
    """Create a Nifti file and save it.

    Parameters
    ----------
    data
        Brain data.
    affine
        Affine of the image, usually inherited from an existing image.
    path
        Output filename.
    """
    if not isinstance(path, str):
        path = str(path)
    img = Nifti1Pair(data, affine)
    nib.nifti1.save(img, path)