"""HCP task-fMRI data loaders (Emotion, Gambling, Language).

The three HCP tasks share nearly identical data-loading logic. They differ
in directory layout, task name, TR, and how labels are parsed from event
files. This module provides a single :class:`HCPDataLoader` that handles
all three via the ``task`` parameter, plus convenience aliases.
"""

from __future__ import annotations

import glob
import os
from typing import List, 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

_TASK_DEFAULTS = {
    "emotion": {
        "task_prefix": "emotionalfaces",
        "default_runs": [1, 2, 3, 4, 5],
        "tr": 2.0,
    },
    "gambling": {
        "task_prefix": "GAMBLING",
        "default_runs": ["LR", "RL"],
        "tr": 0.72,
    },
    "language": {
        "task_prefix": "LANGUAGE",
        "default_runs": ["LR", "RL"],
        "tr": 0.72,
    },
}


class HCPDataLoader(BaseFMRIDataLoader):
    """Unified loader for HCP task-fMRI datasets.

    Parameters
    ----------
    data_dir : str
        Root data directory.
    subjects : list of int
        Subject numbers.
    task : str
        One of ``"emotion"``, ``"gambling"``, ``"language"``.
    mask_file : str, optional
    mask_threshold : float, optional
    runs : list, optional
        Override the default run identifiers.
    """

    def __init__(
        self,
        data_dir: str,
        subjects: list,
        task: str = "emotion",
        mask_file: Optional[str] = None,
        mask_threshold: Optional[float] = 0.9,
        runs: Optional[Sequence] = None,
    ):
        task = task.lower()
        if task not in _TASK_DEFAULTS:
            raise ValueError(f"Unknown HCP task '{task}'. Use: {list(_TASK_DEFAULTS)}")

        cfg = _TASK_DEFAULTS[task]
        super().__init__(
            data_dir=data_dir,
            subjects=[int(s) for s in subjects],
            mask_file=mask_file,
            mask_threshold=mask_threshold,
            runs=runs or cfg["default_runs"],
        )
        self.task = task
        self._task_prefix = cfg["task_prefix"]
        self._tr = cfg["tr"]

        self._set_reference_img()

    def _set_reference_img(self):
        subj = self.subjects[0]
        if self.task == "emotion":
            run = self.runs[0]
            path = os.path.join(
                self.data_dir,
                f"Sub-{subj:02d}",
                f"wrsub-{subj:02d}_task-{self._task_prefix}_run-{run}_bold.nii",
            )
            self._reference_img = index_img(path, 0)
        else:
            task_upper = self._task_prefix
            run = self.runs[0]
            pattern = os.path.join(
                self.data_dir,
                f"subj{subj}/MNINonLinear/Results/tfMRI_{task_upper}_{run}",
                f"tfMRI_{task_upper}_{run}.nii.gz",
            )
            files = glob.glob(pattern)
            if files:
                self._reference_img = index_img(load_img(files[0]), 0)

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

        for subject in self.subjects:
            if self.task == "emotion":
                X_sub, y_sub = self._load_emotion(subject)
            elif self.task == "gambling":
                X_sub, y_sub = self._load_hcp_block(subject, "GAMBLING")
            else:
                X_sub, y_sub = self._load_hcp_block(subject, "LANGUAGE")

            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

    # --- Emotion (non-HCP directory layout) ---

    def _load_emotion(self, subject: int):
        labels_runs = {}
        for run in self.runs:
            labels = pd.read_csv(
                os.path.join(
                    self.data_dir,
                    "onsetime",
                    f"task-{self._task_prefix}_run-{run}_events.tsv",
                ),
                sep="\t",
            ).head(13)
            labels["onset"] = labels["onset"].astype(float) / self._tr
            labels["duration"] = labels["duration"].astype(float) / self._tr
            labels_runs[run] = [
                labels["trial_type"][j]
                for j in range(len(labels))
                for _ in range(int(labels["duration"][j]))
            ]

        data_runs = {}
        for run in self.runs:
            path = os.path.join(
                self.data_dir,
                f"Sub-{subject:02d}",
                f"wrsub-{subject:02d}_task-{self._task_prefix}_run-{run}_bold.nii",
            )
            data_runs[run] = self._load_and_preprocess(path)

        X = np.concatenate([data_runs[r] for r in self.runs], axis=0)
        y = np.concatenate([labels_runs[r] for r in self.runs], axis=0)
        return X, y

    # --- Gambling / Language (HCP directory layout with blocks.csv) ---

    def _load_hcp_block(self, subject: int, task_name: str):
        data_runs = {}
        for run in self.runs:
            path = os.path.join(
                self.data_dir,
                f"subj{subject}/MNINonLinear/Results/tfMRI_{task_name}_{run}",
                f"tfMRI_{task_name}_{run}.nii.gz",
            )
            data_runs[run] = self._load_and_preprocess(path)

        labels_runs = {}
        for run in self.runs:
            csv_path = os.path.join(
                self.data_dir,
                f"subj{subject}/MNINonLinear/Results/tfMRI_{task_name}_{run}/EVs",
                "blocks.csv",
            )
            labels = pd.read_csv(csv_path)
            labels["onset"] = labels["onset"].astype(float) / self._tr
            labels["duration"] = labels["duration"].astype(float) / self._tr

            n_vols = data_runs[run].shape[0]
            vol_labels = ["rest"] * n_vols
            for i in range(n_vols):
                for _, row in labels.iterrows():
                    if row["onset"] <= i < row["onset"] + row["duration"]:
                        vol_labels[i] = row["trial_type"]
                        break
            labels_runs[run] = vol_labels

        X = np.concatenate([data_runs[r] for r in self.runs], axis=0)
        y = np.concatenate([labels_runs[r] for r in self.runs], axis=0)
        return X, y

    # --- Shared preprocessing ---

    def _load_and_preprocess(self, path: str) -> np.ndarray:
        img = load_img(path)
        if self.mask is not None:
            masker = NiftiMasker(mask_img=self.mask, standardize=True)
            return masker.fit_transform(img)
        return clean_img(img).get_fdata().transpose(3, 0, 1, 2)


# ---- Convenience aliases ----

class EmotionDataLoader(HCPDataLoader):
    """Shortcut for ``HCPDataLoader(..., task="emotion")``."""

    def __init__(self, data_dir, subjects, **kwargs):
        super().__init__(data_dir=data_dir, subjects=subjects, task="emotion", **kwargs)


class GamblingDataLoader(HCPDataLoader):
    """Shortcut for ``HCPDataLoader(..., task="gambling")``."""

    def __init__(self, data_dir, subjects, **kwargs):
        super().__init__(data_dir=data_dir, subjects=subjects, task="gambling", **kwargs)


class LanguageDataLoader(HCPDataLoader):
    """Shortcut for ``HCPDataLoader(..., task="language")``."""

    def __init__(self, data_dir, subjects, **kwargs):
        super().__init__(data_dir=data_dir, subjects=subjects, task="language", **kwargs)
