stimulus_loader.py

"""Stimulus context loader for BrainRL prompts and rewards.

The Le Petit Prince dataset ships per-condition word annotation CSVs at
``<data_root>/data/annotation/<condition>_word_information.csv`` with one row
per spoken word and columns ``word, onset, offset, duration, logfreq, pos,
td, bu, lc``. This module slices those rows into compact "stimulus windows"
that the OpenEnv environment threads into each episode so:

* the LLM can see what the subject is hearing (top words + POS mix +
  speech density) and start to learn a mapping from stimulus features to
  brain parcels, and
* the reward function can apply a small, stimulus-conditioned bias to
  parcel groups (auditory_temporal / inferior_frontal / association),
  giving GRPO a learnable "space ↔ word" signal across episodes.

Design goals:
    * Pure stdlib (csv + hashlib) so we don't pull in pandas at runtime.
    * Lazy + cached: each CSV is parsed at most once per process.
    * Robust: returns ``None`` instead of raising if the CSV is missing
      so existing smoke tests still pass without the dataset attached.
"""

from __future__ import annotations

import csv
import os
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

PROJECT_ROOT = Path(__file__).resolve().parent
# Default annotation root: ds005345/data/annotation, sitting next to the repo.
_DEFAULT_ANNOTATION_DIR = PROJECT_ROOT.parent / "data" / "annotation"


# ``time`` columns in the word CSV are integer counts. Spot-checking the
# dataset (single_male: 1515 words spread over ~10 minutes of audio with
# offsets up to ~60000) confirms the unit is centiseconds (10 ms).
_TIME_UNIT_SECONDS = 0.01

# Universal POS tags that count as content words for stimulus features.
_CONTENT_POS = {"NOUN", "VERB", "ADJ", "ADV", "PROPN", "NUM"}

# Mapping from BrainRL condition codes to annotation file basenames.
# ``mixed_*`` sessions play single-male and single-female audio simultaneously
# and have only an acoustic CSV (no word-level annotation), so we fall back to
# the male annotation as the dominant track for those conditions.
_CONDITION_TO_FILE: dict[str, str] = {
    "single_m": "single_male_word_information.csv",
    "single_f": "single_female_word_information.csv",
    "mixed_m": "single_male_word_information.csv",
    "mixed_f": "single_female_word_information.csv",
}


@dataclass(frozen=True)
class WordRow:
    """Single annotated word from the stimulus CSV."""

    word: str
    onset: float
    offset: float
    duration: float
    logfreq: float
    pos: str

    @property
    def is_content(self) -> bool:
        return self.pos.upper() in _CONTENT_POS


# ---------------------------------------------------------------------------
# Loading + caching
# ---------------------------------------------------------------------------

_CACHE: dict[str, list[WordRow]] = {}


def annotation_dir() -> Path:
    """Resolve the annotation directory from env var or default location."""

    override = os.getenv("BRAINRL_STIMULUS_DIR")
    if override:
        return Path(override).expanduser()
    data_dir = os.getenv("BRAINRL_DATA_DIR")
    if data_dir:
        candidate = Path(data_dir).expanduser() / "annotation"
        if candidate.exists():
            return candidate
    config_dir = os.getenv("BRAINRL_CONFIG_DIR")
    if config_dir:
        candidate = Path(config_dir).expanduser().parent / "annotation"
        if candidate.exists():
            return candidate
    return _DEFAULT_ANNOTATION_DIR


def _resolve_csv_path(condition: str, *, root: Path | None = None) -> Path | None:
    file_name = _CONDITION_TO_FILE.get(condition.lower())
    if not file_name:
        return None
    base = (root or annotation_dir()).expanduser()
    candidate = base / file_name
    return candidate if candidate.exists() else None


def load_word_rows(condition: str, *, root: Path | None = None) -> list[WordRow]:
    """Return cached word rows for ``condition`` (empty list if unavailable)."""

    cache_key = f"{(root or annotation_dir()).resolve()}::{condition}"
    if cache_key in _CACHE:
        return _CACHE[cache_key]

    csv_path = _resolve_csv_path(condition, root=root)
    if csv_path is None:
        _CACHE[cache_key] = []
        return _CACHE[cache_key]

    rows: list[WordRow] = []
    with csv_path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for raw in reader:
            try:
                onset = float(raw.get("onset", 0))
                offset = float(raw.get("offset", 0))
                duration = float(raw.get("duration", offset - onset))
                logfreq = float(raw.get("logfreq", 0) or 0.0)
            except (TypeError, ValueError):
                continue
            word = (raw.get("word") or "").strip()
            pos = (raw.get("pos") or "").strip().upper()
            if not word:
                continue
            rows.append(
                WordRow(
                    word=word,
                    onset=onset,
                    offset=offset,
                    duration=duration,
                    logfreq=logfreq,
                    pos=pos,
                )
            )

    _CACHE[cache_key] = rows
    return rows


# ---------------------------------------------------------------------------
# Window summarization
# ---------------------------------------------------------------------------


def n_windows(condition: str, *, window_size: int = 30, root: Path | None = None) -> int:
    """Number of stimulus windows available for ``condition``."""

    rows = load_word_rows(condition, root=root)
    if not rows:
        return 0
    window_size = max(1, int(window_size))
    return max(1, (len(rows) + window_size - 1) // window_size)


def _stable_window_index(*, condition: str, key: Iterable[object], n: int) -> int:
    """Deterministic mapping from any context tuple to a window index in [0, n).

    Uses Python's ``hash`` with the same trick we use for episode_seed:
    works well across runs of a single process. We do not rely on
    cross-process determinism here because the loader is also called from
    fresh ``BrainRegionSelectionEnvironment`` instances at episode time.
    """

    if n <= 0:
        return 0
    seed_bits = abs(hash((condition, *tuple(key))))
    return seed_bits % int(n)


def summarize_window(
    condition: str | None,
    *,
    window_index: int | None = None,
    window_size: int = 30,
    top_words: int = 10,
    deterministic_key: Iterable[object] | None = None,
    root: Path | None = None,
) -> dict | None:
    """Build a compact stimulus-window dict for prompts and reward shaping.

    Returns ``None`` when the condition has no annotation CSV (e.g. when
    the dataset is not mounted). Callers should treat ``None`` as
    "no stimulus context available" and fall back to the prior behaviour.
    """

    if not condition:
        return None
    rows = load_word_rows(condition, root=root)
    if not rows:
        return None

    window_size = max(1, int(window_size))
    total_windows = max(1, (len(rows) + window_size - 1) // window_size)

    if window_index is None:
        if deterministic_key is None:
            deterministic_key = (condition,)
        window_index = _stable_window_index(
            condition=condition, key=deterministic_key, n=total_windows
        )
    window_index = max(0, min(int(window_index), total_windows - 1))

    start = window_index * window_size
    end = min(start + window_size, len(rows))
    window_rows = rows[start:end]
    if not window_rows:
        return None

    pos_counter: Counter[str] = Counter()
    content_count = 0
    logfreqs: list[float] = []
    durations: list[float] = []
    for row in window_rows:
        pos_counter[row.pos or "X"] += 1
        if row.is_content:
            content_count += 1
        logfreqs.append(row.logfreq)
        durations.append(row.duration)

    start_time_s = window_rows[0].onset * _TIME_UNIT_SECONDS
    end_time_s = window_rows[-1].offset * _TIME_UNIT_SECONDS
    span_seconds = max(end_time_s - start_time_s, 1e-6)
    speech_seconds = sum(durations) * _TIME_UNIT_SECONDS
    density = float(min(1.0, max(0.0, speech_seconds / span_seconds)))

    dominant_pos = pos_counter.most_common(1)[0][0] if pos_counter else "NA"
    mean_logfreq = float(sum(logfreqs) / len(logfreqs)) if logfreqs else 0.0
    mean_duration_s = float(
        (sum(durations) / len(durations)) * _TIME_UNIT_SECONDS
    ) if durations else 0.0

    # ``top_words`` is intentionally compact: just the first ``top_words``
    # words in playback order with their POS so the prompt stays readable
    # while still letting the LLM see actual stimulus content.
    sample = [
        {"w": row.word, "pos": row.pos or "X"}
        for row in window_rows[: max(1, int(top_words))]
    ]

    return {
        "condition": condition,
        "window_index": int(window_index),
        "n_windows": int(total_windows),
        "n_words": int(len(window_rows)),
        "n_content_words": int(content_count),
        "pos_counts": dict(pos_counter),
        "dominant_pos": dominant_pos,
        "mean_logfreq": float(round(mean_logfreq, 3)),
        "mean_duration_s": float(round(mean_duration_s, 3)),
        "start_time_s": float(round(start_time_s, 3)),
        "end_time_s": float(round(end_time_s, 3)),
        "speech_density": float(round(density, 3)),
        "top_words": sample,
    }


# ---------------------------------------------------------------------------
# Reward shaping helpers
# ---------------------------------------------------------------------------


def stimulus_bias_for_group(
    redundancy_group: str,
    features: dict | None,
) -> float:
    """Map a stimulus window to a per-parcel-group reward multiplier.

    The intent is to give the RL policy a learnable signal: when the
    current window is dominated by content nouns and high speech density
    (lots of acoustic information), auditory_temporal parcels should pay
    out a bit more; when it is heavy on grammatical/function words, the
    inferior_frontal group is preferred; otherwise the broader
    association group benefits. The multiplier is bounded to a small
    range so it shapes choices without overwhelming the base reward.
    """

    if not features:
        return 1.0

    n_words = max(1, int(features.get("n_words", 1)))
    pos = features.get("pos_counts", {}) or {}
    nouns = float(pos.get("NOUN", 0)) / n_words
    verbs = float(pos.get("VERB", 0)) / n_words
    function_share = (
        float(pos.get("PART", 0))
        + float(pos.get("ADP", 0))
        + float(pos.get("DET", 0))
        + float(pos.get("PRON", 0))
        + float(pos.get("CCONJ", 0))
        + float(pos.get("SCONJ", 0))
    ) / n_words
    density = float(features.get("speech_density", 0.5))

    if redundancy_group == "auditory_temporal":
        bias = 1.0 + 0.25 * density + 0.10 * nouns
    elif redundancy_group == "inferior_frontal":
        bias = 1.0 + 0.20 * function_share + 0.10 * verbs
    elif redundancy_group == "association":
        bias = 1.0 + 0.10 * (nouns + verbs)
    else:
        bias = 1.0

    return float(max(0.7, min(1.4, bias)))


def compact_for_prompt(features: dict | None) -> dict | None:
    """Trim a window summary to the fields safe to paste into the prompt."""

    if not features:
        return None
    return {
        "condition": features.get("condition"),
        "window": (
            f"{int(features.get('window_index', 0)) + 1}/"
            f"{int(features.get('n_windows', 1))}"
        ),
        "n_words": int(features.get("n_words", 0)),
        "dominant_pos": features.get("dominant_pos"),
        "pos_mix": features.get("pos_counts", {}),
        "mean_logfreq": features.get("mean_logfreq"),
        "speech_density": features.get("speech_density"),
        "duration_s": float(round(
            float(features.get("end_time_s", 0.0))
            - float(features.get("start_time_s", 0.0)),
            3,
        )),
        "top_words": features.get("top_words", []),
    }