data_split.py

"""Subject/condition-aware data splits for BrainRL.

The Le Petit Prince derivative ships with ``participant_run_info.json`` that
maps each subject's four runs to one of four conditions
(``single_m``, ``single_f``, ``mixed_m``, ``mixed_f``). For the hackathon we
typically train on a single condition (e.g. ``single_m``) and split subjects
into train/test groups so generalization claims are honest.

This module is intentionally tiny: pure stdlib, no dependency on the
environment, so it can be used both at preparation time and at episode time.
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_PARTICIPANT_INFO = PROJECT_ROOT / "configs" / "participant_run_info.json"
LEGACY_PARTICIPANT_INFO = PROJECT_ROOT.parent / "derivatives" / "participant_run_info.json"
VALID_CONDITIONS: tuple[str, ...] = ("single_m", "single_f", "mixed_m", "mixed_f")


@dataclass(frozen=True)
class SubjectRunPair:
    """One concrete (subject, run, condition) tuple usable as an episode."""

    subject_id: str
    run_id: str
    condition: str

    def episode_seed(self, base_seed: int = 0) -> int:
        h = abs(hash((self.subject_id, self.run_id, self.condition, int(base_seed))))
        return h % (2**31 - 1)

    def as_dict(self) -> dict[str, str]:
        return {
            "subject_id": self.subject_id,
            "run_id": self.run_id,
            "condition": self.condition,
        }


# ---------------------------------------------------------------------------
# Loading + filtering
# ---------------------------------------------------------------------------

def load_participant_run_info(path: str | Path | None = None) -> dict[str, dict[str, str]]:
    """Load the raw subject -> run -> condition mapping."""

    if path:
        info_path = Path(path)
    elif os.getenv("BRAINRL_CONFIG_DIR"):
        info_path = Path(os.environ["BRAINRL_CONFIG_DIR"]) / "participant_run_info.json"
    else:
        info_path = DEFAULT_PARTICIPANT_INFO
    info_path = info_path.expanduser()
    if path is None and not info_path.exists() and LEGACY_PARTICIPANT_INFO.exists():
        info_path = LEGACY_PARTICIPANT_INFO
    if not info_path.exists():
        raise FileNotFoundError(
            f"participant_run_info.json not found at {info_path}. "
            "Pass --participant-info to point at the correct file."
        )
    with info_path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    return {
        str(subject): {str(run): str(cond) for run, cond in runs.items()}
        for subject, runs in payload.items()
    }


def pairs_for_condition(
    info: dict[str, dict[str, str]],
    condition: str,
) -> list[SubjectRunPair]:
    """Return every (subject, run) tuple matching ``condition``."""

    if condition not in VALID_CONDITIONS:
        raise ValueError(
            f"Unknown condition={condition!r}. Expected one of {VALID_CONDITIONS}."
        )
    pairs: list[SubjectRunPair] = []
    for subject, runs in info.items():
        for run_id, cond in runs.items():
            if cond == condition:
                pairs.append(
                    SubjectRunPair(subject_id=subject, run_id=run_id, condition=cond)
                )
    pairs.sort(key=lambda p: (p.subject_id, p.run_id))
    return pairs


# ---------------------------------------------------------------------------
# Subject parsing
# ---------------------------------------------------------------------------

_SUB_RE = re.compile(r"^sub-(\d+)$")


def _normalize_subject(token: str) -> str:
    token = token.strip()
    if not token:
        return ""
    if token.isdigit():
        return f"sub-{int(token):02d}"
    match = _SUB_RE.match(token)
    if match:
        return f"sub-{int(match.group(1)):02d}"
    return token


def parse_subject_spec(spec: str | None, available: Iterable[str]) -> list[str]:
    """Parse a CLI subject spec into a sorted list of subject ids.

    Accepts:
      * ``"sub-01:sub-20"``  – inclusive range
      * ``"sub-01,sub-05,sub-09"`` – explicit comma list
      * ``"01:20"`` / ``"01,05,09"`` – shorthand without ``sub-`` prefix
      * ``None`` / ``""`` / ``"all"`` – every available subject
    """

    available_set = sorted({s for s in available})
    if not spec or spec.lower() == "all":
        return list(available_set)

    if ":" in spec:
        start_raw, end_raw = spec.split(":", 1)
        start = _normalize_subject(start_raw)
        end = _normalize_subject(end_raw)
        ordered = sorted(available_set)
        try:
            start_idx = ordered.index(start)
            end_idx = ordered.index(end)
        except ValueError as exc:
            raise ValueError(
                f"Subject range {spec!r} does not match available subjects: {ordered}"
            ) from exc
        if start_idx > end_idx:
            start_idx, end_idx = end_idx, start_idx
        return ordered[start_idx : end_idx + 1]

    items = [_normalize_subject(t) for t in spec.split(",") if t.strip()]
    missing = [s for s in items if s not in available_set]
    if missing:
        raise ValueError(
            f"Subjects {missing} not in available list {available_set}."
        )
    return sorted(items)


def parse_exclude_spec(spec: str | None) -> set[str]:
    """Parse an exclusion spec into a set of normalized subject ids.

    Same syntax as ``parse_subject_spec`` but does not require the subjects
    to exist in any participant list - excluding a missing id is a no-op so
    that scripts stay robust if the corrupted-subject list drifts.
    """

    if not spec:
        return set()
    if spec.lower() == "none":
        return set()

    out: set[str] = set()
    if ":" in spec:
        start_raw, end_raw = spec.split(":", 1)
        start = _normalize_subject(start_raw)
        end = _normalize_subject(end_raw)
        try:
            start_idx = int(start.split("-")[1])
            end_idx = int(end.split("-")[1])
        except (IndexError, ValueError) as exc:
            raise ValueError(f"Bad exclusion range {spec!r}.") from exc
        if start_idx > end_idx:
            start_idx, end_idx = end_idx, start_idx
        for i in range(start_idx, end_idx + 1):
            out.add(f"sub-{i:02d}")
        return out

    for token in spec.split(","):
        token = token.strip()
        if not token:
            continue
        out.add(_normalize_subject(token))
    return out


# ---------------------------------------------------------------------------
# Splits
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ConditionSplit:
    condition: str
    train_pairs: list[SubjectRunPair]
    test_pairs: list[SubjectRunPair]
    train_subjects: list[str]
    test_subjects: list[str]
    excluded_subjects: list[str]

    def pairs_for(self, split: str) -> list[SubjectRunPair]:
        split = split.lower()
        if split == "train":
            return list(self.train_pairs)
        if split == "test":
            return list(self.test_pairs)
        if split == "all":
            return list(self.train_pairs) + list(self.test_pairs)
        raise ValueError(f"Unknown split={split!r}; expected train/test/all.")

    def summary(self) -> dict[str, object]:
        return {
            "condition": self.condition,
            "n_train_subjects": len(self.train_subjects),
            "n_test_subjects": len(self.test_subjects),
            "n_train_pairs": len(self.train_pairs),
            "n_test_pairs": len(self.test_pairs),
            "train_subjects": self.train_subjects,
            "test_subjects": self.test_subjects,
            "excluded_subjects": self.excluded_subjects,
        }


def build_condition_split(
    *,
    condition: str,
    participant_info_path: str | Path | None = None,
    train_subjects_spec: str | None = None,
    test_subjects_spec: str | None = None,
    exclude_subjects: str | Iterable[str] | None = None,
    train_frac: float = 0.75,
) -> ConditionSplit:
    """High-level helper used by all CLI scripts.

    ``exclude_subjects`` can be a comma list / range / iterable of subject
    ids (e.g. ``"sub-03,sub-18"``). Excluded subjects are dropped before any
    train/test logic and never appear in either pair list, so corrupted or
    held-out subjects can be skipped consistently across the whole stack.
    """

    info = load_participant_run_info(participant_info_path)
    all_pairs = pairs_for_condition(info, condition)
    if not all_pairs:
        raise ValueError(f"No (subject, run) pairs found for condition={condition!r}.")

    if isinstance(exclude_subjects, str) or exclude_subjects is None:
        excluded = parse_exclude_spec(exclude_subjects if isinstance(exclude_subjects, str) else None)
    else:
        excluded = {_normalize_subject(s) for s in exclude_subjects if s}

    if excluded:
        all_pairs = [p for p in all_pairs if p.subject_id not in excluded]
        if not all_pairs:
            raise ValueError(
                f"All subjects for condition={condition!r} were excluded by "
                f"exclude_subjects={sorted(excluded)}."
            )

    available_subjects = sorted({p.subject_id for p in all_pairs})

    if train_subjects_spec or test_subjects_spec:
        train_subjects = [
            s for s in parse_subject_spec(train_subjects_spec, available_subjects)
            if s not in excluded
        ]
        if test_subjects_spec:
            test_subjects = [
                s for s in parse_subject_spec(test_subjects_spec, available_subjects)
                if s not in excluded
            ]
        else:
            test_subjects = [s for s in available_subjects if s not in set(train_subjects)]
        if not train_subjects:
            train_subjects = available_subjects[: max(1, int(len(available_subjects) * train_frac))]
        if not test_subjects:
            test_subjects = [s for s in available_subjects if s not in set(train_subjects)]
    else:
        cutoff = max(1, int(round(len(available_subjects) * train_frac)))
        train_subjects = available_subjects[:cutoff]
        test_subjects = available_subjects[cutoff:]

    train_set = set(train_subjects)
    test_set = set(test_subjects)
    train_pairs = [p for p in all_pairs if p.subject_id in train_set]
    test_pairs = [p for p in all_pairs if p.subject_id in test_set]

    return ConditionSplit(
        condition=condition,
        train_pairs=train_pairs,
        test_pairs=test_pairs,
        train_subjects=sorted(train_subjects),
        test_subjects=sorted(test_subjects),
        excluded_subjects=sorted(excluded),
    )