"""Haxby 2001 fMRI dataset loader."""

from __future__ import annotations

import os
from typing import Optional, Sequence

import numpy as np
import pandas as pd
from nilearn.image import clean_img, index_img, load_img
from nilearn.maskers import NiftiMasker

from autofmri.data.base import BaseFMRIDataLoader


class HaxbyDataLoader(BaseFMRIDataLoader):
    """Load the Haxby 2001 visual-object recognition fMRI dataset.

    Supports two loading modes:

    * **Volume mode** (``load_volume_data``): loads a single 4-D NIfTI
      (``preprocessed/swubold.nii``) and reads labels from ``labels.txt``.
    * **First-level mode** (``load_first_level_data``): loads individual
      beta-map NIfTI files where the filename encodes the label.

    Parameters
    ----------
    data_dir : str
        Path to the subject directory, e.g. ``"data/haxby2001/subj4"``.
    subjects : str or int or list
        Subject identifier(s).
    mask_file : str, optional
        Path to a mask NIfTI file.
    mask_threshold : float, optional
        Threshold for binarising the mask.
    """

    def __init__(
        self,
        data_dir: str,
        subjects=None,
        mask_file: Optional[str] = None,
        mask_threshold: Optional[float] = None,
    ):
        super().__init__(
            data_dir=data_dir,
            subjects=subjects or ["subj"],
            mask_file=mask_file,
            mask_threshold=mask_threshold,
        )

    def load_volume_data(self):
        path_4d = os.path.join(self.data_dir, "preprocessed", "swubold.nii")
        img_4d = load_img(path_4d)

        if self.mask is not None:
            masker = NiftiMasker(mask_img=self.mask, standardize=True)
            X = masker.fit_transform(img_4d)
        else:
            img_4d = clean_img(img_4d)
            X = np.array(
                [index_img(img_4d, i).get_fdata() for i in range(img_4d.shape[-1])]
            )

        y = pd.read_csv(
            os.path.join(self.data_dir, "labels.txt"), sep=" ", header=None
        ).values[1:, 0]

        assert X.shape[0] == y.shape[0], (
            f"Volume count ({X.shape[0]}) != label count ({y.shape[0]})"
        )

        self._reference_img = index_img(path_4d, 0)
        return X, y

    def load_first_level_data(self):
        """Load first-level output where each file is ``<label>_<run>.nii.gz``."""
        X, y = [], []
        for f in sorted(os.listdir(self.data_dir)):
            if not f.endswith(".nii.gz"):
                continue
            img = load_img(os.path.join(self.data_dir, f))
            if self.mask is not None:
                masker = NiftiMasker(mask_img=self.mask, standardize=True)
                data = masker.fit_transform(img)
            else:
                data = img.get_fdata()
            X.append(data if isinstance(data, np.ndarray) else data)
            y.append(f.split("_")[0])

        X = np.array(X, dtype=np.float32).squeeze()
        y = np.array(y)
        self._reference_img = load_img(
            os.path.join(self.data_dir, sorted(os.listdir(self.data_dir))[0])
        )
        return X, y
