"""VKMC (musical pitch) fMRI dataset loader."""

from __future__ import annotations

import os
from typing import Optional, Sequence

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

from autofmri.data.base import BaseFMRIDataLoader


class VKMCDataLoader(BaseFMRIDataLoader):
    """Load the VKMC music-cognition fMRI dataset.

    Parameters
    ----------
    data_dir : str
        Root data directory (e.g. ``"data/2021_Reconstruction_dataset"``).
    subjects : list of str
        Subject IDs like ``["sub-001", "sub-003"]``.
    mask_file : str, optional
    mask_threshold : float, optional
    label_type : str
        Label column: ``"midi"`` or ``"timbre"``.
    augment : int
        Temporal interpolation factor (1 = no augmentation).
    """

    def __init__(
        self,
        data_dir: str,
        subjects: list,
        mask_file: Optional[str] = None,
        mask_threshold: Optional[float] = None,
        label_type: str = "midi",
        augment: int = 0,
    ):
        super().__init__(
            data_dir=data_dir,
            subjects=[s if isinstance(s, str) else str(s) for s in subjects],
            mask_file=mask_file,
            mask_threshold=mask_threshold,
        )
        self.label_type = label_type
        self.augment = augment

    def load_volume_data(
        self,
        runs: Sequence[int] = (1, 2, 3, 4, 5, 6, 7, 8),
        sessions: Sequence[int] = (1, 2),
    ):
        X_parts, y_parts = [], []

        for subject in self.subjects:
            onset_dfs = {
                f"session-{s}": pd.read_csv(
                    os.path.join(self.data_dir, "logfiles-AB", f"{subject}_ses-{s}.csv")
                )
                for s in sessions
            }

            func_dir = os.path.join(self.data_dir, subject, "func")
            for niifile in sorted(os.listdir(func_dir)):
                if not niifile.endswith(".nii"):
                    continue

                run = niifile.split("run-")[1].split(".")[0]
                session = niifile.split("ses-")[1].split("_")[0]

                if session not in map(str, sessions) or run not in map(str, runs):
                    continue

                img = load_img(os.path.join(func_dir, niifile))
                if self._reference_img is None:
                    self._reference_img = index_img(img, 0)

                if self.mask is not None:
                    masker = NiftiMasker(
                        mask_img=self.mask, standardize="zscore_sample"
                    )
                    data = masker.fit_transform(img)
                else:
                    data = clean_img(img).get_fdata().transpose(3, 0, 1, 2)

                volume_times = np.arange(0, data.shape[0] * 2, 2)
                labels = self._assign_labels(
                    onset_dfs[f"session-{session}"], int(run), volume_times
                )

                X_parts.append(data)
                y_parts.append(labels)

        X = np.concatenate(X_parts, axis=0)
        y = np.concatenate(y_parts, axis=0)

        if self.augment > 1:
            zoom = (self.augment, 1) if X.ndim == 2 else (self.augment, 1, 1, 1)
            X = scipy.ndimage.zoom(X, zoom, order=1)
            y = np.repeat(y, self.augment)

        X = np.nan_to_num(X)
        assert X.shape[0] == y.shape[0]
        return X, y

    def load_data(self):
        """Load beta-map (LSA) data."""
        beta_dir = os.path.join(self.data_dir, "LSA_all") if "LSA_all" not in self.data_dir else self.data_dir

        if self.mask is not None:
            masker = NiftiMasker(mask_img=self.mask, standardize="zscore_sample")
        else:
            masker = NiftiMasker(standardize="zscore_sample")

        X, y = [], []
        for subject in self.subjects:
            subj_dir = os.path.join(beta_dir, subject)
            if not os.path.isdir(subj_dir):
                continue
            for f in sorted(os.listdir(subj_dir)):
                if not f.endswith(".nii"):
                    continue
                parts = f.split("_")
                label_map = {
                    "midi": parts[2].split("-")[1] if len(parts) > 2 else "",
                    "timbre": parts[1] if len(parts) > 1 else "",
                    "instrument": parts[0],
                }
                img = load_img(os.path.join(subj_dir, f))
                if self._reference_img is None:
                    self._reference_img = img
                X.append(img.get_fdata())
                y.append(label_map.get(self.label_type, ""))

        X = np.nan_to_num(np.array(X, dtype=np.float32).squeeze())
        y = np.array(y)
        assert X.shape[0] == y.shape[0]
        return X, y

    def _assign_labels(self, onset_df, run: int, volume_times):
        df = onset_df[onset_df["run"] == run]
        labels = []
        for t in volume_times:
            match = df.loc[
                (df["sound.started"] <= t) & (df["sound.started"] + 6 >= t)
            ]
            if not match.empty:
                labels.append(match.iloc[-1][f"sound.{self.label_type}"])
            else:
                labels.append("rest")
        return labels
