# Path: brainiak/image.py
# def mask_image(image: SpatialImage, mask: np.ndarray, data_type: type = None
#                ) -> np.ndarray:
#     """Mask image after optionally casting its type.
# 
#     Parameters
#     ----------
#     image
#         Image to mask. Can include time as the last dimension.
#     mask
#         Mask to apply. Must have the same shape as the image data.
#     data_type
#         Type to cast image to.
# 
#     Returns
#     -------
#     np.ndarray
#         Masked image.
# 
#     Raises
#     ------
#     ValueError
#         Image data and masks have different shapes.
#     """
#     image_data = image.get_data()
#     if image_data.shape[:3] != mask.shape:
#         raise ValueError("Image data and mask have different shapes.")
#     if data_type is not None:
#         cast_data = image_data.astype(data_type)
#     else:
#         cast_data = image_data
#     return cast_data[mask]
# 
# def mask_images(images: Iterable[SpatialImage], mask: np.ndarray,
#                 image_type: type = None) -> Iterable[np.ndarray]:
#     """Mask images.
# 
#     Parameters
#     ----------
#     images:
#         Images to mask.
#     mask:
#         Mask to apply.
#     image_type:
#         Type to cast images to.
# 
#     Yields
#     ------
#     np.ndarray
#         Masked image.
#     """
#     for images in multimask_images(images, (mask,), image_type):
#         yield images[0]
# 
# class MaskedMultiSubjectData(np.ndarray):
#     """Array with shape n_TRs, n_voxels, n_subjects."""
#     @classmethod
#     def from_masked_images(cls: Type[T], masked_images: Iterable[np.ndarray],
#                            n_subjects: int) -> T:
#         """Create a new instance of MaskedMultiSubjecData from masked images.
# 
#         Parameters
#         ----------
#         masked_images : iterator
#             Images from multiple subjects to stack along 3rd dimension
#         n_subjects : int
#             Number of subjects; must match the number of images
# 
#         Returns
#         -------
#         T
#             A new instance of MaskedMultiSubjectData
# 
#         Raises
#         ------
#         ValueError
#             Images have different shapes.
# 
#             The number of images differs from n_subjects.
#         """
#         images_iterator = iter(masked_images)
#         first_image = next(images_iterator)
#         first_image_shape = first_image.T.shape
#         result = np.empty((first_image_shape[0], first_image_shape[1],
#                            n_subjects))
#         for n_images, image in enumerate(itertools.chain([first_image],
#                                                          images_iterator)):
#             image = image.T
#             if image.shape != first_image_shape:
#                 raise ValueError("Image {} has different shape from first "
#                                  "image: {} != {}".format(n_images,
#                                                           image.shape,
#                                                           first_image_shape))
#             result[:, :, n_images] = image
#         n_images += 1
#         if n_images != n_subjects:
#             raise ValueError("n_subjects != number of images: {} != {}"
#                              .format(n_subjects, n_images))
#         return result.view(cls)
# 
# def multimask_images(images: Iterable[SpatialImage],
#                      masks: Sequence[np.ndarray], image_type: type = None
#                      ) -> Iterable[Sequence[np.ndarray]]:
#     """Mask images with multiple masks.
# 
#     Parameters
#     ----------
#     images:
#         Images to mask.
#     masks:
#         Masks to apply.
#     image_type:
#         Type to cast images to.
# 
#     Yields
#     ------
#     Sequence[np.ndarray]
#         For each mask, a masked image.
#     """
#     for image in images:
#         yield [mask_image(image, mask, image_type) for mask in masks]
# 
# class SingleConditionSpec(ConditionSpec):
#     """ConditionSpec with a single condition applicable to an epoch."""
# 
#     def extract_labels(self) -> np.ndarray:
#         """Extract condition labels.
# 
#         Returns
#         -------
#         np.ndarray
#             The condition label of each epoch.
#         """
#         condition_idxs, epoch_idxs, _ = np.where(self)
#         _, unique_epoch_idxs = np.unique(epoch_idxs, return_index=True)
#         return condition_idxs[unique_epoch_idxs]


from typing import Iterable, Sequence
from nibabel.nifti1 import Nifti1Pair
from nibabel.spatialimages import SpatialImage
from brainiak.image import (mask_image, mask_images, MaskedMultiSubjectData,
                            multimask_images, SingleConditionSpec)
import numpy as np
import pytest
#  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.






@pytest.fixture
def masked_multi_subject_data(masked_images):
    return np.stack(masked_images, axis=-1)


class TestMaskedMultiSubjectData:
    def test_from_masked_images(self, masked_images,
                                masked_multi_subject_data):
        result = MaskedMultiSubjectData.from_masked_images(masked_images,
                                                           len(masked_images))
        assert np.array_equal(np.moveaxis(result, 1, 0),
                              masked_multi_subject_data)


@pytest.fixture
def condition_spec() -> SingleConditionSpec:
    return np.array([[[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
                     [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 1, 1, 1, 1, 0]]],

# Path: brainiak/image.py
# def mask_image(image: SpatialImage, mask: np.ndarray, data_type: type = None
#                ) -> np.ndarray:
#     """Mask image after optionally casting its type.
# 
#     Parameters
#     ----------
#     image
#         Image to mask. Can include time as the last dimension.
#     mask
#         Mask to apply. Must have the same shape as the image data.
#     data_type
#         Type to cast image to.
# 
#     Returns
#     -------
#     np.ndarray
#         Masked image.
# 
#     Raises
#     ------
#     ValueError
#         Image data and masks have different shapes.
#     """
#     image_data = image.get_data()
#     if image_data.shape[:3] != mask.shape:
#         raise ValueError("Image data and mask have different shapes.")
#     if data_type is not None:
#         cast_data = image_data.astype(data_type)
#     else:
#         cast_data = image_data
#     return cast_data[mask]
# 
# def mask_images(images: Iterable[SpatialImage], mask: np.ndarray,
#                 image_type: type = None) -> Iterable[np.ndarray]:
#     """Mask images.
# 
#     Parameters
#     ----------
#     images:
#         Images to mask.
#     mask:
#         Mask to apply.
#     image_type:
#         Type to cast images to.
# 
#     Yields
#     ------
#     np.ndarray
#         Masked image.
#     """
#     for images in multimask_images(images, (mask,), image_type):
#         yield images[0]
# 
# class MaskedMultiSubjectData(np.ndarray):
#     """Array with shape n_TRs, n_voxels, n_subjects."""
#     @classmethod
#     def from_masked_images(cls: Type[T], masked_images: Iterable[np.ndarray],
#                            n_subjects: int) -> T:
#         """Create a new instance of MaskedMultiSubjecData from masked images.
# 
#         Parameters
#         ----------
#         masked_images : iterator
#             Images from multiple subjects to stack along 3rd dimension
#         n_subjects : int
#             Number of subjects; must match the number of images
# 
#         Returns
#         -------
#         T
#             A new instance of MaskedMultiSubjectData
# 
#         Raises
#         ------
#         ValueError
#             Images have different shapes.
# 
#             The number of images differs from n_subjects.
#         """
#         images_iterator = iter(masked_images)
#         first_image = next(images_iterator)
#         first_image_shape = first_image.T.shape
#         result = np.empty((first_image_shape[0], first_image_shape[1],
#                            n_subjects))
#         for n_images, image in enumerate(itertools.chain([first_image],
#                                                          images_iterator)):
#             image = image.T
#             if image.shape != first_image_shape:
#                 raise ValueError("Image {} has different shape from first "
#                                  "image: {} != {}".format(n_images,
#                                                           image.shape,
#                                                           first_image_shape))
#             result[:, :, n_images] = image
#         n_images += 1
#         if n_images != n_subjects:
#             raise ValueError("n_subjects != number of images: {} != {}"
#                              .format(n_subjects, n_images))
#         return result.view(cls)
# 
# def multimask_images(images: Iterable[SpatialImage],
#                      masks: Sequence[np.ndarray], image_type: type = None
#                      ) -> Iterable[Sequence[np.ndarray]]:
#     """Mask images with multiple masks.
# 
#     Parameters
#     ----------
#     images:
#         Images to mask.
#     masks:
#         Masks to apply.
#     image_type:
#         Type to cast images to.
# 
#     Yields
#     ------
#     Sequence[np.ndarray]
#         For each mask, a masked image.
#     """
#     for image in images:
#         yield [mask_image(image, mask, image_type) for mask in masks]
# 
# class SingleConditionSpec(ConditionSpec):
#     """ConditionSpec with a single condition applicable to an epoch."""
# 
#     def extract_labels(self) -> np.ndarray:
#         """Extract condition labels.
# 
#         Returns
#         -------
#         np.ndarray
#             The condition label of each epoch.
#         """
#         condition_idxs, epoch_idxs, _ = np.where(self)
#         _, unique_epoch_idxs = np.unique(epoch_idxs, return_index=True)
#         return condition_idxs[unique_epoch_idxs]


# Path: tests/image/test_image.py
from typing import Iterable, Sequence
from nibabel.nifti1 import Nifti1Pair
from nibabel.spatialimages import SpatialImage
from brainiak.image import (mask_image, mask_images, MaskedMultiSubjectData,
                            multimask_images, SingleConditionSpec)
import numpy as np
import pytest

#  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.






@pytest.fixture
def masked_multi_subject_data(masked_images):
    return np.stack(masked_images, axis=-1)


class TestMaskedMultiSubjectData:
    def test_from_masked_images(self, masked_images,
                                masked_multi_subject_data):
        result = MaskedMultiSubjectData.from_masked_images(masked_images,
                                                           len(masked_images))
        assert np.array_equal(np.moveaxis(result, 1, 0),
                              masked_multi_subject_data)


@pytest.fixture
def condition_spec() -> SingleConditionSpec:
    return np.array([[[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
                     [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 1, 1, 1, 1, 0]]],

                    dtype=np.int8).view(SingleConditionSpec)