code/a1_pipeline/constants.py

"""Constants and shared defaults for A1 SI baseline bootstrap."""

from __future__ import annotations

from dataclasses import dataclass

DEFAULT_ALLOWED_RUNS: tuple[int, ...] = (1, 2, 3, 4)
DEFAULT_TIME_SCALE_SECONDS: float = 0.01
DEFAULT_HRF_LAG_SECONDS: float = 5.0

# Locked baseline models from project scope.
DEFAULT_MODEL_IDS: tuple[str, ...] = (
    "Qwen/Qwen3-0.6B",
    "meta-llama/Llama-3.2-1B",
    "GSAI-ML/LLaDA-8B-Instruct",
)

# Run-level word timing tables available in ds005345 annotations.
DEFAULT_RUN_WORD_TABLE_MAP: dict[int, str | None] = {
    1: "single_female_word_information.csv",
    2: "single_male_word_information.csv",
    3: None,
    4: None,
}

# Fallback sources for mixed-condition runs when dedicated mixed word timings are absent.
DEFAULT_MIXED_RUN_WORD_TABLE_FALLBACK_MAP: dict[int, str] = {
    3: "single_female_word_information.csv",
    4: "single_male_word_information.csv",
}


@dataclass(frozen=True)
class RunCondition:
    """Condition labels attached to each run in the baseline mapping."""

    condition_fixed: str
    speaker_stream: str
    condition_fallback: str | None = None


DEFAULT_RUN_CONDITION_MAP: dict[int, RunCondition] = {
    1: RunCondition(condition_fixed="single_female", speaker_stream="female"),
    2: RunCondition(condition_fixed="single_male", speaker_stream="male"),
    3: RunCondition(
        condition_fixed="mixed_female",
        speaker_stream="mixed_female",
        condition_fallback="mixed",
    ),
    4: RunCondition(
        condition_fixed="mixed_male",
        speaker_stream="mixed_male",
        condition_fallback="mixed",
    ),
}


def resolve_run_condition(run: int, enable_mixed_fallback: bool) -> tuple[str, str, str, bool]:
    """Resolve fixed/effective condition labels for one run.

    Returns:
        (condition_fixed, condition_effective, speaker_stream, used_mixed_fallback)
    """
    if run not in DEFAULT_RUN_CONDITION_MAP:
        raise KeyError(f"Run {run} is not present in condition map")

    info = DEFAULT_RUN_CONDITION_MAP[run]
    use_fallback = bool(enable_mixed_fallback and info.condition_fallback is not None)
    condition_effective = info.condition_fallback if use_fallback else info.condition_fixed

    return info.condition_fixed, condition_effective, info.speaker_stream, use_fallback