data_loader.py

"""Candidate loading for the BrainRL OpenEnv environment.

Two candidate modes are supported:

* ``atlas_parcels`` (default) reads the frozen manifest produced by
  ``prepare_parcels.py`` at ``configs/parcel_candidates.json``. This is the
  hackathon-target setup with ~200 Schaefer-style parcels.
* ``roi_priors`` reads the small ``configs/region_priors.json`` fixture and is
  kept as the easy-curriculum smoke-test path.

The loader never touches NIfTI files; pruning happens once in
``prepare_parcels.py`` and the environment treats the resulting manifest as
read-only at episode time.
"""

from __future__ import annotations

import json
import os
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_CONFIG_DIR = PROJECT_ROOT / "configs"
DEFAULT_CONFIG_PATH = DEFAULT_CONFIG_DIR / "subset_config.yaml"
DEFAULT_PARCEL_MANIFEST = DEFAULT_CONFIG_DIR / "parcel_candidates.json"
DEFAULT_PRIOR_PATH = DEFAULT_CONFIG_DIR / "region_priors.json"


@dataclass(frozen=True)
class RegionCandidate:
    """Compact metadata + scoring parameters for one candidate region/parcel."""

    region_id: str
    label: str
    hemisphere: str
    semantic_prior: float
    base_r2: float
    cost: float
    redundancy_group: str
    notes: str = ""
    atlas: str = ""
    network: str = ""
    sub_region: str = ""
    n_voxels: int = 0
    prune_score: float = 0.0

    def as_dict(self) -> dict[str, Any]:
        return {
            "region_id": self.region_id,
            "label": self.label,
            "hemisphere": self.hemisphere,
            "semantic_prior": float(self.semantic_prior),
            "base_r2": float(self.base_r2),
            "cost": float(self.cost),
            "redundancy_group": self.redundancy_group,
            "notes": self.notes,
            "atlas": self.atlas,
            "network": self.network,
            "sub_region": self.sub_region,
            "n_voxels": int(self.n_voxels),
            "prune_score": float(self.prune_score),
        }

    def prompt_dict(self) -> dict[str, Any]:
        """Compact view used by the prompt builder (no internal hints)."""

        return {
            "region_id": self.region_id,
            "label": self.label,
            "hemisphere": self.hemisphere,
            "network": self.network,
            "semantic_prior": float(round(self.semantic_prior, 4)),
            "base_r2_hint": float(round(self.base_r2, 4)),
            "cost": float(round(self.cost, 4)),
            "n_voxels": int(self.n_voxels),
        }


@dataclass(frozen=True)
class BrainSubset:
    """Loaded candidate set + task parameters."""

    dataset_name: str
    candidates: list[RegionCandidate]
    selection_budget: int
    cost_penalty: float
    source: str
    candidate_mode: str
    atlas: str
    prompt_top_k: int

    @property
    def n_regions(self) -> int:
        return len(self.candidates)


# ---------------------------------------------------------------------------
# Config / I/O helpers
# ---------------------------------------------------------------------------

def _read_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def config_dir() -> Path:
    """Resolve config directory, allowing HF Dataset bundles to override it."""

    override = os.getenv("BRAINRL_CONFIG_DIR")
    return Path(override).expanduser() if override else DEFAULT_CONFIG_DIR


def default_config_path() -> Path:
    return config_dir() / "subset_config.yaml"


def default_parcel_manifest() -> Path:
    return config_dir() / "parcel_candidates.json"


def default_prior_path() -> Path:
    return config_dir() / "region_priors.json"


def _resolve_config_path(raw_path: str | Path, *, base_dir: Path) -> Path:
    path = Path(raw_path).expanduser()
    if path.is_absolute():
        return path
    candidates = (
        base_dir / path,
        PROJECT_ROOT / path,
        base_dir.parent / path,
    )
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return base_dir / path


def _parse_simple_yaml(path: Path) -> dict[str, Any]:
    """Tiny YAML reader so we don't drag PyYAML into the runtime deps."""

    if not path.exists():
        return {}

    config: dict[str, Any] = {}
    current_list_key: str | None = None
    with path.open("r", encoding="utf-8") as handle:
        for raw_line in handle:
            line = raw_line.split("#", 1)[0].rstrip()
            if not line.strip():
                continue
            if line.startswith("  - ") and current_list_key:
                config.setdefault(current_list_key, []).append(line[4:].strip())
                continue
            if ":" not in line:
                continue
            key, raw_value = line.split(":", 1)
            key = key.strip()
            value = raw_value.strip()
            if value == "":
                config[key] = []
                current_list_key = key
                continue
            current_list_key = None
            if value.lower() in {"true", "false"}:
                config[key] = value.lower() == "true"
            else:
                try:
                    config[key] = int(value)
                except ValueError:
                    try:
                        config[key] = float(value)
                    except ValueError:
                        config[key] = value.strip('"').strip("'")
    return config


# ---------------------------------------------------------------------------
# Atlas-parcel mode
# ---------------------------------------------------------------------------

def _candidates_from_parcel_manifest(payload: dict[str, Any]) -> list[RegionCandidate]:
    parcels = payload.get("candidates", [])
    candidates: list[RegionCandidate] = []
    for entry in parcels:
        candidates.append(
            RegionCandidate(
                region_id=str(entry["region_id"]),
                label=str(entry.get("label", entry["region_id"])),
                hemisphere=str(entry.get("hemisphere", "unknown")),
                semantic_prior=float(entry.get("semantic_prior", 0.5)),
                base_r2=float(entry.get("base_r2", 0.04)),
                cost=float(entry.get("cost", 1.0)),
                redundancy_group=str(entry.get("redundancy_group", "association")),
                notes=str(entry.get("notes", entry.get("network", ""))),
                atlas=str(entry.get("atlas", payload.get("atlas", ""))),
                network=str(entry.get("network", "")),
                sub_region=str(entry.get("sub_region", "")),
                n_voxels=int(entry.get("n_voxels", 0)),
                prune_score=float(entry.get("prune_score", 0.0)),
            )
        )
    return candidates


def _load_atlas_parcels(config: dict[str, Any]) -> tuple[BrainSubset, str]:
    base_dir = Path(str(config.get("_config_dir", config_dir()))).expanduser()
    manifest_path = _resolve_config_path(
        config.get("parcel_manifest_path", str(default_parcel_manifest())),
        base_dir=base_dir,
    )
    if not manifest_path.exists():
        raise FileNotFoundError(
            f"Parcel manifest not found at {manifest_path}. "
            "Run `python prepare_parcels.py` (see Makefile target `prepare`)."
        )
    payload = _read_json(manifest_path)
    candidates = _candidates_from_parcel_manifest(payload)

    max_candidates = int(config.get("max_candidates", payload.get("max_candidates", len(candidates))))
    candidates = candidates[: max(1, max_candidates)]

    selection_budget = int(
        config.get("selection_budget", payload.get("selection_budget", min(20, len(candidates))))
    )
    selection_budget = max(1, min(selection_budget, len(candidates)))

    prompt_top_k = int(config.get("prompt_top_k", payload.get("prompt_top_k", 30)))
    prompt_top_k = max(1, min(prompt_top_k, len(candidates)))

    cost_penalty = float(config.get("cost_penalty", payload.get("cost_penalty", 0.002)))
    atlas = str(payload.get("atlas", config.get("atlas_name", "schaefer200")))

    subset = BrainSubset(
        dataset_name=str(config.get("dataset_name", "le_petit_prince_atlas_parcels")),
        candidates=candidates,
        selection_budget=selection_budget,
        cost_penalty=cost_penalty,
        source=f"parcel_manifest:{manifest_path.name}",
        candidate_mode="atlas_parcels",
        atlas=atlas,
        prompt_top_k=prompt_top_k,
    )
    return subset, str(manifest_path)


# ---------------------------------------------------------------------------
# ROI-priors (legacy easy curriculum)
# ---------------------------------------------------------------------------

def _redundancy_group_for_roi(region_id: str) -> str:
    if "STS" in region_id or region_id == "TP":
        return "auditory_temporal"
    if region_id.startswith("BA"):
        return "inferior_frontal"
    return "association"


def _load_roi_priors(config: dict[str, Any]) -> tuple[BrainSubset, str]:
    base_dir = Path(str(config.get("_config_dir", config_dir()))).expanduser()
    prior_path = _resolve_config_path(
        config.get("prior_path", str(default_prior_path())),
        base_dir=base_dir,
    )
    if not prior_path.exists():
        raise FileNotFoundError(f"ROI priors file not found at {prior_path}.")
    payload = _read_json(prior_path)
    priors = payload.get("priors", [])

    rng = random.Random(int(config.get("seed", 42)))
    candidates: list[RegionCandidate] = []
    for prior in priors:
        region_id = str(prior["region_id"])
        semantic_prior = float(prior.get("semantic_prior", 0.5))
        jitter = rng.uniform(-0.01, 0.015)
        base_r2 = max(0.005, min(0.2, 0.025 + 0.09 * semantic_prior + jitter))
        candidates.append(
            RegionCandidate(
                region_id=region_id,
                label=str(prior.get("label", region_id)),
                hemisphere=str(prior.get("hemisphere", "left")),
                semantic_prior=semantic_prior,
                base_r2=base_r2,
                cost=float(prior.get("cost", 1.0)),
                redundancy_group=_redundancy_group_for_roi(region_id),
                notes=str(prior.get("notes", "")),
                atlas="language_rois",
                network="language",
                n_voxels=0,
                prune_score=semantic_prior,
            )
        )

    candidates.sort(key=lambda item: item.semantic_prior, reverse=True)
    max_regions = int(config.get("max_regions", len(candidates)))
    candidates = candidates[: max(1, max_regions)]

    requested_budget = int(config.get("selection_budget", min(5, len(candidates))))
    selection_budget = max(1, min(requested_budget, len(candidates)))
    subset = BrainSubset(
        dataset_name=str(config.get("dataset_name", "le_petit_prince_small_roi")),
        candidates=candidates,
        selection_budget=selection_budget,
        cost_penalty=float(config.get("cost_penalty", 0.002)),
        source=f"roi_priors:{prior_path.name}",
        candidate_mode="roi_priors",
        atlas="language_rois",
        prompt_top_k=int(config.get("prompt_top_k", len(candidates))),
    )
    return subset, str(prior_path)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def load_brain_subset(config_path: Path | str | None = None) -> BrainSubset:
    """Load a candidate set based on subset_config.yaml.

    Falls back to roi_priors mode if the parcel manifest is missing so the
    repository remains runnable without first running ``prepare_parcels.py``.
    """

    resolved_config_path = Path(config_path).expanduser() if config_path else default_config_path()
    config = _parse_simple_yaml(resolved_config_path)
    config["_config_dir"] = str(resolved_config_path.parent)
    mode = str(config.get("candidate_mode", "atlas_parcels")).strip().lower()

    if mode == "atlas_parcels":
        try:
            subset, _ = _load_atlas_parcels(config)
            return subset
        except FileNotFoundError:
            # Graceful fallback so smoke tests work without a prepared manifest.
            subset, _ = _load_roi_priors(config)
            return subset

    if mode == "roi_priors":
        subset, _ = _load_roi_priors(config)
        return subset

    raise ValueError(
        f"Unknown candidate_mode={mode!r}. Expected 'atlas_parcels' or 'roi_priors'."
    )


def candidate_table(subset: BrainSubset) -> list[dict[str, Any]]:
    """Return JSON-serializable candidate metadata."""

    return [candidate.as_dict() for candidate in subset.candidates]