baselines.py

"""Baseline policies for BrainRL evaluation.

These are intentionally non-RL policies that share the same OpenEnv interface
as the LLM/GRPO policy. They give the hackathon demo something concrete to
beat without dragging in any Gym/PPO machinery.
"""

from __future__ import annotations

import random
from dataclasses import dataclass
from statistics import mean
from typing import Protocol

from metrics import EpisodeMetrics, score_selection_order
from prompts import llm_prompt_policy_action, static_prompt_policy_action
from server.brain_environment import BrainRegionSelectionEnvironment


class RegionPolicy(Protocol):
    name: str

    def reset(self, env: BrainRegionSelectionEnvironment) -> None: ...

    def select_region(self, env: BrainRegionSelectionEnvironment) -> str: ...


@dataclass
class EpisodeResult:
    policy: str
    rewards: list[float]
    r2_curve: list[float]
    selected_regions: list[str]
    errors: list[str]
    subject_id: str | None = None
    run_id: str | None = None
    condition: str | None = None
    stimulus_window: int | None = None
    stimulus_summary: dict | None = None
    metrics: EpisodeMetrics | None = None

    @property
    def final_r2(self) -> float:
        return self.r2_curve[-1] if self.r2_curve else 0.0

    @property
    def total_reward(self) -> float:
        return float(sum(self.rewards))

    @property
    def priority_correlation(self) -> float:
        return float(self.metrics.priority_correlation) if self.metrics else 0.0

    @property
    def two_v_two_accuracy(self) -> float:
        return float(self.metrics.two_v_two_accuracy) if self.metrics else 0.0


class RandomPolicy:
    name = "random"

    def __init__(self, seed: int = 42):
        self._rng = random.Random(seed)

    def reset(self, env: BrainRegionSelectionEnvironment) -> None:
        return None

    def select_region(self, env: BrainRegionSelectionEnvironment) -> str:
        selected = set(env._selected_region_ids)
        remaining = [
            candidate.region_id
            for candidate in env._subset.candidates
            if candidate.region_id not in selected
        ]
        return self._rng.choice(remaining)


class SemanticPriorPolicy:
    name = "semantic_prior"

    def reset(self, env: BrainRegionSelectionEnvironment) -> None:
        self._ordered = [
            candidate.region_id
            for candidate in sorted(
                env._subset.candidates,
                key=lambda item: (item.semantic_prior, item.prune_score),
                reverse=True,
            )
        ]

    def select_region(self, env: BrainRegionSelectionEnvironment) -> str:
        selected = set(env._selected_region_ids)
        for region_id in self._ordered:
            if region_id not in selected:
                return region_id
        return self._ordered[0]


class GreedyImprovementPolicy:
    name = "greedy"

    def reset(self, env: BrainRegionSelectionEnvironment) -> None:
        return None

    def select_region(self, env: BrainRegionSelectionEnvironment) -> str:
        selected = list(env._selected_region_ids)
        selected_set = set(selected)
        best_region: str | None = None
        best_gain = float("-inf")

        for candidate in env._subset.candidates:
            if candidate.region_id in selected_set:
                continue
            next_regions = selected + [candidate.region_id]
            gain = env._score_regions(next_regions) - env._current_r2
            gain -= env._subset.cost_penalty * candidate.cost
            if gain > best_gain:
                best_gain = gain
                best_region = candidate.region_id

        if best_region is None:
            return env._subset.candidates[0].region_id
        return best_region


class PromptPolicy:
    """OpenEnv prompt-action policy.

    In static mode it uses the exact prompt interface but selects deterministically
    from prune_score/semantic_prior. In LLM mode it asks a chat model for one
    JSON action.
    """

    def __init__(self, use_llm: bool = False):
        self.use_llm = use_llm
        self.name = "llm_prompt" if use_llm else "prompt_static"

    def reset(self, env: BrainRegionSelectionEnvironment) -> None:
        return None

    def select_region(self, env: BrainRegionSelectionEnvironment) -> str:
        state = env._build_selection_state()
        if self.use_llm:
            region_id, _ = llm_prompt_policy_action(state)
            return region_id
        return static_prompt_policy_action(state)


def run_policy_episode(
    policy: RegionPolicy,
    seed: int = 42,
    *,
    subject_id: str | None = None,
    run_id: str | None = None,
    condition: str | None = None,
    stimulus_window: int | None = None,
    env: BrainRegionSelectionEnvironment | None = None,
) -> EpisodeResult:
    env = env or BrainRegionSelectionEnvironment()
    env.reset(
        seed=seed,
        subject_id=subject_id,
        run_id=run_id,
        condition=condition,
        stimulus_window=stimulus_window,
    )
    policy.reset(env)

    rewards: list[float] = []
    r2_curve: list[float] = []
    errors: list[str] = []

    done = False
    while not done:
        region_id = policy.select_region(env)
        result = env._process_action(region_id)
        rewards.append(float(result["reward"]))
        r2_curve.append(float(result["current_r2"]))
        if result.get("error"):
            errors.append(str(result["error"]))
        done = bool(result["done"])

    stim = env._stimulus_features
    candidate_ids = [candidate.region_id for candidate in env._subset.candidates]
    metrics = score_selection_order(
        candidate_ids=candidate_ids,
        selected_region_ids=env._selected_region_ids,
        target_scores=env._effective_base_r2,
    )
    return EpisodeResult(
        policy=policy.name,
        rewards=rewards,
        r2_curve=r2_curve,
        selected_regions=list(env._selected_region_ids),
        errors=errors,
        subject_id=subject_id,
        run_id=run_id,
        condition=condition,
        stimulus_window=int(stim["window_index"]) if stim else None,
        stimulus_summary={
            "window_index": int(stim["window_index"]),
            "n_windows": int(stim["n_windows"]),
            "dominant_pos": stim.get("dominant_pos"),
            "n_words": int(stim.get("n_words", 0)),
            "speech_density": float(stim.get("speech_density", 0.0)),
        }
        if stim
        else None,
        metrics=metrics,
    )


def summarize_results(
    results: list[EpisodeResult],
    example_limit: int = 5,
    split_label: str | None = None,
) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    for policy_name in sorted({result.policy for result in results}):
        policy_results = [result for result in results if result.policy == policy_name]
        example_order = " -> ".join(policy_results[0].selected_regions[:example_limit])
        if len(policy_results[0].selected_regions) > example_limit:
            example_order += f" -> (+{len(policy_results[0].selected_regions) - example_limit} more)"
        row: dict[str, object] = {
            "policy": policy_name,
            "episodes": len(policy_results),
            "mean_final_r2": mean(result.final_r2 for result in policy_results),
            "mean_priority_correlation": mean(
                result.priority_correlation for result in policy_results
            ),
            "mean_2v2_accuracy": mean(
                result.two_v_two_accuracy for result in policy_results
            ),
            "mean_total_reward": mean(result.total_reward for result in policy_results),
            "mean_regions_selected": mean(
                len(result.selected_regions) for result in policy_results
            ),
            "example_order": example_order,
        }
        if split_label:
            row["split"] = split_label
        rows.append(row)
    return rows


def r2_curves_by_policy(
    results: list[EpisodeResult],
) -> dict[str, list[list[float]]]:
    """Group r2_curve lists by policy name for plotting."""

    grouped: dict[str, list[list[float]]] = {}
    for result in results:
        grouped.setdefault(result.policy, []).append(list(result.r2_curve))
    return grouped


def default_baselines(seed: int = 42) -> list[RegionPolicy]:
    return [
        RandomPolicy(seed=seed),
        SemanticPriorPolicy(),
        PromptPolicy(use_llm=False),
        GreedyImprovementPolicy(),
    ]