"""Abstract base class for fMRI data loaders."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import List, Optional, Sequence, Tuple, Union

import numpy as np
from nilearn.image import binarize_img, index_img, load_img, new_img_like


class BaseFMRIDataLoader(ABC):
    """Common interface for loading fMRI datasets.

    Subclasses must implement :meth:`load_volume_data` which returns
    ``(X, y)`` numpy arrays.

    Parameters
    ----------
    data_dir : str
        Root directory of the dataset.
    subjects : list of str or int
        Subject identifiers.
    mask_file : str, optional
        Path to a NIfTI mask file.
    mask_threshold : float, optional
        If provided, binarise the mask at this threshold.
    runs : list, optional
        Run identifiers to load.
    """

    def __init__(
        self,
        data_dir: str,
        subjects: Union[List, str, int],
        mask_file: Optional[str] = None,
        mask_threshold: Optional[float] = None,
        runs: Optional[Sequence] = None,
    ):
        self.data_dir = data_dir
        self.subjects = (
            [subjects] if not isinstance(subjects, (list, tuple)) else list(subjects)
        )
        self.runs = list(runs) if runs is not None else None
        self._reference_img = None

        self.mask = None
        if mask_file is not None:
            self.mask = load_img(mask_file)
            if self.mask.shape[0] > 100:
                self.mask = new_img_like(
                    self.mask, self.mask.get_fdata()[::2, ::2, ::2]
                )
            if mask_threshold is not None:
                self.mask = binarize_img(self.mask, mask_threshold)

    @property
    def reference_img(self):
        """A single 3-D reference image (Nifti1Image) for this dataset."""
        return self._reference_img

    @abstractmethod
    def load_volume_data(self) -> Tuple[np.ndarray, np.ndarray]:
        """Load fMRI data and labels.

        Returns
        -------
        X : ndarray, shape (n_samples, ...)
            If no mask: (n_samples, nx, ny, nz).
            If masked: (n_samples, n_voxels).
        y : ndarray, shape (n_samples,)
            Labels for each sample.
        """
        ...
