# Path: brainiak/io.py
# def load_images_from_dir(in_dir: Union[str, Path], suffix: str = "nii.gz",
#                          ) -> Iterable[SpatialImage]:
# def load_images(image_paths: Iterable[Union[str, Path]]
#                 ) -> Iterable[SpatialImage]:
# def load_boolean_mask(path: Union[str, Path],
#                       predicate: Callable[[np.ndarray], np.ndarray] = None
#                       ) -> np.ndarray:
# def load_labels(path: Union[str, Path]) -> List[SingleConditionSpec]:
# def save_as_nifti_file(data: np.ndarray, affine: np.ndarray,
#                        path: Union[str, Path]) -> None:


from pathlib import Path
from typing import Iterable, Sequence
from brainiak import io
import nibabel as nib
import numpy as np
import pytest
def image_paths(in_dir: Path) -> Iterable[Path]:
    return (in_dir / "subject1_bet.nii.gz", in_dir / "subject2_bet.nii.gz")


def test_load_images_from_dir_data_shape(
        in_dir: Path,
        expected_image_data_shape: Sequence[int],
        expected_n_subjects: int
        ) -> None:
    for i, image in enumerate(io.load_images_from_dir(in_dir, "bet.nii.gz")):
        assert image.get_data().shape == (64, 64, 26, 10)
    assert i + 1 == expected_n_subjects


def test_load_images_data_shape(
        image_paths: Iterable[Path],
        expected_image_data_shape: Sequence[int],
        expected_n_subjects: int
        ) -> None:
    for i, image in enumerate(io.load_images(image_paths)):
        assert image.get_data().shape == (64, 64, 26, 10)
    assert i + 1 == expected_n_subjects


def test_load_boolean_mask(mask_path: Path) -> None:
    mask = io.load_boolean_mask(mask_path)
    assert mask.dtype == np.bool


def test_load_boolean_mask_predicate(mask_path: Path) -> None:

# Path: brainiak/io.py
# def load_images_from_dir(in_dir: Union[str, Path], suffix: str = "nii.gz",
#                          ) -> Iterable[SpatialImage]:
# def load_images(image_paths: Iterable[Union[str, Path]]
#                 ) -> Iterable[SpatialImage]:
# def load_boolean_mask(path: Union[str, Path],
#                       predicate: Callable[[np.ndarray], np.ndarray] = None
#                       ) -> np.ndarray:
# def load_labels(path: Union[str, Path]) -> List[SingleConditionSpec]:
# def save_as_nifti_file(data: np.ndarray, affine: np.ndarray,
#                        path: Union[str, Path]) -> None:


# Path: tests/io/test_io.py
from pathlib import Path
from typing import Iterable, Sequence
from brainiak import io
import nibabel as nib
import numpy as np
import pytest

def image_paths(in_dir: Path) -> Iterable[Path]:
    return (in_dir / "subject1_bet.nii.gz", in_dir / "subject2_bet.nii.gz")


def test_load_images_from_dir_data_shape(
        in_dir: Path,
        expected_image_data_shape: Sequence[int],
        expected_n_subjects: int
        ) -> None:
    for i, image in enumerate(io.load_images_from_dir(in_dir, "bet.nii.gz")):
        assert image.get_data().shape == (64, 64, 26, 10)
    assert i + 1 == expected_n_subjects


def test_load_images_data_shape(
        image_paths: Iterable[Path],
        expected_image_data_shape: Sequence[int],
        expected_n_subjects: int
        ) -> None:
    for i, image in enumerate(io.load_images(image_paths)):
        assert image.get_data().shape == (64, 64, 26, 10)
    assert i + 1 == expected_n_subjects


def test_load_boolean_mask(mask_path: Path) -> None:
    mask = io.load_boolean_mask(mask_path)
    assert mask.dtype == np.bool


def test_load_boolean_mask_predicate(mask_path: Path) -> None:

    mask = io.load_boolean_mask(mask_path, lambda x: np.logical_not(x))