"""GTZAN music-genre fMRI dataset loader."""

from __future__ import annotations

import glob
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 GTZANDataLoader(BaseFMRIDataLoader):
    """Load the GTZAN music-genre fMRI dataset (fMRIPrep preprocessed).

    Parameters
    ----------
    data_dir : str
        Derivative directory (e.g. ``"data/ds003720-derivative"``).
    subjects : list of int
        Subject numbers.
    mask_file : str, optional
    mask_threshold : float, optional
    runs : list of int, optional
        Run numbers to include (default: 1-12).
    """

    def __init__(
        self,
        data_dir: str = "data/ds003720-derivative",
        subjects: list = None,
        mask_file: Optional[str] = None,
        mask_threshold: Optional[float] = 0.9,
        runs: Optional[Sequence[int]] = None,
    ):
        super().__init__(
            data_dir=data_dir,
            subjects=[int(s) for s in (subjects or [1])],
            mask_file=mask_file,
            mask_threshold=mask_threshold,
            runs=runs or list(range(1, 13)),
        )

        nii_files = glob.glob(
            os.path.join(
                self.data_dir,
                f"sub-{self.subjects[0]:03d}/func",
                "*MNI152NLin2009cAsym_desc-preproc_bold.nii",
            )
        )
        if nii_files:
            self._reference_img = index_img(load_img(nii_files[0]), 0)

    def load_volume_data(self):
        X_all, y_all = [], []

        for subject in self.subjects:
            func_dir = os.path.join(self.data_dir, f"sub-{subject:03d}", "func")
            nii_files = glob.glob(
                os.path.join(func_dir, "*MNI152NLin2009cAsym_desc-preproc_bold.nii")
            )
            runs = [
                f.split("task-", 1)[1].split("_space", 1)[0] for f in nii_files
            ]
            runs = [r for r in runs if int(r.split("_run-", 1)[1]) in self.runs]

            labels_runs = {}
            for run in runs:
                events = pd.read_csv(
                    os.path.join(func_dir, f"sub-{subject:03d}_task-{run}_events.tsv"),
                    sep="\t",
                )
                events["onset"] = events["onset"].astype(float) / 1.5
                events["duration"] = events["duration"].astype(float) / 1.5
                labels_runs[run] = np.repeat(
                    events["genre"].values, events["duration"].values.astype(int)
                )

            data_runs = {}
            for run in runs:
                path = os.path.join(
                    func_dir,
                    f"sub-{subject:03d}_task-{run}_space-MNI152NLin2009cAsym_desc-preproc_bold.nii",
                )
                img = load_img(path)
                if self.mask is not None:
                    masker = NiftiMasker(mask_img=self.mask, standardize=True)
                    data_runs[run] = masker.fit_transform(img)
                else:
                    data_runs[run] = clean_img(img).get_fdata().transpose(3, 0, 1, 2)

            X_sub = np.concatenate([data_runs[r] for r in runs], axis=0)
            y_sub = np.concatenate([labels_runs[r] for r in runs], axis=0)
            X_all.append(X_sub)
            y_all.append(y_sub)

        X = np.concatenate(X_all, axis=0)
        y = np.concatenate(y_all, axis=0)
        assert X.shape[0] == y.shape[0]
        return X, y
