code/run_a1_fit.py

#!/usr/bin/env python
"""Fit and evaluate A1 alignment caches for target masks.

This script consumes Phase 6 artifacts from `run_a1_bootstrap.py` and computes
ridge-regression performance for:
- Protocol C: cross-subject subject-holdout on a shared canonical stimulus

It supports three target-mask modes:
- `run_top10`: per-run-condition top-10% ISC masks from Swati outputs
- `run_top25`: per-run-condition top-25% ISC masks from Swati outputs
- `core_roi`: legacy evaluation with the 7 core language ROIs
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any

import nibabel as nib
from nibabel.processing import resample_from_to
import numpy as np
import pandas as pd

from a1_pipeline.io_utils import ensure_directory, write_json
from a1_pipeline.participant_runs import (
    load_participant_run_map,
    resolve_subject_actual_run,
    resolve_subject_canonical_run,
)
from a1_pipeline.targets import CORE_ROI_NAMES, load_core_roi_masks


TARGET_MASK_MODE_CHOICES: tuple[str, ...] = ("run_top10", "run_top25", "core_roi")
RUN_MASK_PERCENT_BY_MODE: dict[str, int] = {
    "run_top10": 10,
    "run_top25": 25,
}
SWATI_CONDITION_BY_CANONICAL_RUN: dict[int, str] = {
    1: "single_f",
    2: "single_m",
    3: "mixed_f",
    4: "mixed_m",
}


def _resolve_num_fit_workers(requested: str, n_layers: int) -> tuple[int, int]:
    """Resolve --num-fit-workers to (n_workers, blas_threads_per_worker).

    'auto' picks min(n_layers, max(1, cpu_count // 2)) so each worker still
    has multiple BLAS threads available for matmul.
    """
    cpu_count = os.cpu_count() or 4

    token = str(requested).strip().lower()
    if token in {"", "auto"}:
        n_workers = max(1, min(int(n_layers), max(1, cpu_count // 2)))
    else:
        try:
            value = int(token)
        except ValueError as exc:
            raise ValueError(f"Invalid --num-fit-workers={requested!r}") from exc
        n_workers = max(1, min(value, int(n_layers)))

    blas_threads = max(1, cpu_count // n_workers)
    return n_workers, blas_threads


def _parse_csv_int_list(raw: str) -> list[int]:
    tokens = [token.strip() for token in raw.split(",") if token.strip()]
    return [int(token) for token in tokens]


def _parse_subjects(raw: str) -> list[str] | None:
    value = raw.strip().lower()
    if value == "all":
        return None
    tokens = [token.strip() for token in raw.split(",") if token.strip()]
    return sorted(set(tokens))


def _parse_layers(raw: str) -> list[int] | None:
    value = raw.strip().lower()
    if value == "all":
        return None
    layers = sorted(set(_parse_csv_int_list(raw)))
    if any(layer < 0 for layer in layers):
        raise ValueError("Layer indices must be non-negative")
    return layers


def _parse_protocols(raw: str) -> set[str]:
    values = {token.strip().upper() for token in raw.split(",") if token.strip()}
    if not values:
        raise ValueError("At least one protocol is required")

    allowed = {"C"}
    unknown = values.difference(allowed)
    if unknown:
        raise ValueError(f"Unsupported protocol(s): {sorted(unknown)}")

    return values


def _parse_subject_field(raw: Any) -> list[str]:
    text = str(raw).strip()
    if text == "" or text.lower() == "nan":
        return []
    return [token.strip() for token in text.split(",") if token.strip()]


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Fit and evaluate A1 cached alignment inputs")
    parser.add_argument(
        "--bootstrap-output-dir",
        type=str,
        default=None,
        help="Output directory from run_a1_bootstrap.py (required)",
    )
    parser.add_argument(
        "--model-slug",
        type=str,
        default="Qwen_Qwen3-0.6B",
        help="Model slug in alignment_regressor_summary.csv",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default=None,
        help="Optional human-readable model ID for reporting",
    )
    parser.add_argument(
        "--alpha",
        type=float,
        default=1e3,
        help="Ridge regularization coefficient",
    )
    parser.add_argument(
        "--subjects",
        type=str,
        default="all",
        help="Comma-separated subject IDs or 'all'",
    )
    parser.add_argument(
        "--layer-indices",
        type=str,
        default="all",
        help="Comma-separated layer indices or 'all'",
    )
    parser.add_argument(
        "--protocols",
        type=str,
        default="C",
        help="Protocols to run: C (cross-subject shared-space).",
    )
    parser.add_argument(
        "--target-mask-mode",
        type=str,
        default="run_top10",
        choices=sorted(TARGET_MASK_MODE_CHOICES),
        help=(
            "Target mask family to evaluate. 'run_top10' selects the ISC top-10%% mask for the "
            "canonical stimulus behind each actual run, 'run_top25' keeps the 25%% mask option, "
            "and 'core_roi' keeps the legacy 7-ROI evaluation."
        ),
    )
    parser.add_argument(
        "--run-mask-dir",
        "--run-top10-mask-dir",
        "--run-top25-mask-dir",
        dest="run_mask_dir",
        type=str,
        default=None,
        help=(
            "Directory containing Swati ISC run-conditioned masks. Can point either to the Swati "
            "output root or directly to its isc_group folder."
        ),
    )
    parser.add_argument(
        "--roi-mask-dir",
        type=str,
        default=None,
        help="Legacy core ROI mask directory, used only when --target-mask-mode core_roi.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Fit result output directory (default: <bootstrap-output-dir>/fit_results/<model-slug>)",
    )
    parser.add_argument(
        "--participant-run-info",
        type=str,
        default=str((Path(__file__).resolve().parent / "assets" / "participant_run_info.json")),
        help=(
            "JSON file mapping each subject run to canonical stimulus condition "
            "(default: code/assets/participant_run_info.json)"
        ),
    )
    parser.add_argument(
        "--num-fit-workers",
        type=str,
        default="auto",
        help=(
            "Number of CPU workers to evaluate layers in parallel for protocol C. "
            "'auto' uses min(n_layers, cpu_count // 2). Use 1 to force serial."
        ),
    )
    return parser


def _resolve_core_paths(bootstrap_output_dir: Path) -> dict[str, Path]:
    csv_dir = bootstrap_output_dir / "csv"

    def _pick_csv(name: str) -> Path:
        csv_candidate = csv_dir / name
        if csv_candidate.exists():
            return csv_candidate
        return bootstrap_output_dir / name

    return {
        "bootstrap_summary": bootstrap_output_dir / "bootstrap_summary.json",
        "analysis_mask": bootstrap_output_dir / "analysis_mask.nii.gz",
        "bold_summary": _pick_csv("alignment_bold_summary.csv"),
        "regressor_summary": _pick_csv("alignment_regressor_summary.csv"),
        "protocol_c": _pick_csv("protocol_c_cross_subject_folds.csv"),
    }


def _check_required_files(path_map: dict[str, Path], protocols: set[str]) -> None:
    required_names = {
        "bootstrap_summary",
        "analysis_mask",
        "bold_summary",
        "regressor_summary",
    }
    if "C" in protocols:
        required_names.add("protocol_c")

    missing = [name for name in sorted(required_names) if not path_map[name].exists()]
    if missing:
        lines = [f"{name}: {path_map[name]}" for name in missing]
        raise FileNotFoundError("Missing required files:\n" + "\n".join(lines))


def _resolve_core_roi_mask_dir(args: argparse.Namespace, bootstrap_summary: dict[str, Any]) -> Path:
    if args.roi_mask_dir:
        return Path(args.roi_mask_dir).resolve()

    if "roi_mask_dir" in bootstrap_summary:
        return Path(str(bootstrap_summary["roi_mask_dir"])).resolve()

    return (Path(__file__).resolve().parent / "assets" / "roi_masks")


def _canonical_run_target_name(canonical_run: int, top_percent: int) -> str:
    if canonical_run not in SWATI_CONDITION_BY_CANONICAL_RUN:
        raise KeyError(
            f"Unsupported canonical run for target mask selection: {canonical_run}. "
            f"Allowed runs: {sorted(SWATI_CONDITION_BY_CANONICAL_RUN)}"
        )
    return f"{SWATI_CONDITION_BY_CANONICAL_RUN[canonical_run]}_top{int(top_percent)}"


def _analysis_mask_components(
    analysis_mask_path: Path,
) -> tuple[nib.Nifti1Image, np.ndarray, np.ndarray]:
    analysis_img = nib.load(str(analysis_mask_path))
    analysis_mask_bool = analysis_img.get_fdata() > 0.5
    analysis_flat = analysis_mask_bool.ravel(order="C")
    analysis_flat_indices = np.flatnonzero(analysis_flat)
    return analysis_img, analysis_mask_bool, analysis_flat_indices


def _mask_to_vector_indices(
    mask_bool: np.ndarray,
    analysis_mask_bool: np.ndarray,
    analysis_flat_indices: np.ndarray,
    mask_name: str,
) -> np.ndarray:
    mask_in_analysis = analysis_mask_bool & mask_bool
    mask_flat_indices = np.flatnonzero(mask_in_analysis.ravel(order="C"))

    if mask_flat_indices.size == 0:
        raise ValueError(f"Target mask {mask_name} has zero voxels inside analysis mask")

    return np.searchsorted(analysis_flat_indices, mask_flat_indices).astype(np.int64)


def _build_core_roi_index_maps(
    analysis_mask_path: Path,
    roi_mask_dir: Path,
) -> dict[int, dict[str, np.ndarray]]:
    analysis_img, analysis_mask_bool, analysis_flat_indices = _analysis_mask_components(
        analysis_mask_path=analysis_mask_path,
    )

    core_roi_masks, _ = load_core_roi_masks(roi_mask_dir=roi_mask_dir, reference_img=analysis_img)

    roi_index_map: dict[str, np.ndarray] = {}

    for roi_name in CORE_ROI_NAMES:
        roi_bool = core_roi_masks[roi_name]
        roi_index_map[roi_name] = _mask_to_vector_indices(
            mask_bool=roi_bool,
            analysis_mask_bool=analysis_mask_bool,
            analysis_flat_indices=analysis_flat_indices,
            mask_name=roi_name,
        )

    return {
        canonical_run: dict(roi_index_map)
        for canonical_run in sorted(SWATI_CONDITION_BY_CANONICAL_RUN)
    }


def _normalize_run_mask_dir(path: Path, top_percent: int) -> Path | None:
    candidates = [path, path / "isc_group"]
    required_files = [
        f"{condition}_isc_top{int(top_percent)}.nii.gz"
        for condition in SWATI_CONDITION_BY_CANONICAL_RUN.values()
    ]

    for candidate in candidates:
        if not candidate.exists() or not candidate.is_dir():
            continue
        if all((candidate / filename).exists() for filename in required_files):
            return candidate.resolve()

    return None


def _resolve_run_mask_dir(args: argparse.Namespace, top_percent: int) -> Path:
    script_path = Path(__file__).resolve()
    workspace_root = script_path.parents[3]

    candidate_roots: list[Path] = []
    if args.run_mask_dir:
        candidate_roots.append(Path(args.run_mask_dir).expanduser().resolve())

    candidate_roots.extend(
        [
            workspace_root / "data" / "isc_group",
            workspace_root / "data",
            workspace_root / "swati" / "TEAM-9" / "output",
            workspace_root / "swati" / "TEAM-9" / "output" / "isc_group",
        ]
    )

    for candidate in candidate_roots:
        normalized = _normalize_run_mask_dir(candidate, top_percent=top_percent)
        if normalized is not None:
            return normalized

    searched = "\n".join(str(path) for path in candidate_roots)
    raise FileNotFoundError(
        f"Could not locate Swati run-top{int(top_percent)} masks. Checked:\n" + searched
    )


def _build_run_mask_index_maps(
    analysis_mask_path: Path,
    run_mask_dir: Path,
    top_percent: int,
) -> dict[int, dict[str, np.ndarray]]:
    analysis_img, analysis_mask_bool, analysis_flat_indices = _analysis_mask_components(
        analysis_mask_path=analysis_mask_path,
    )

    index_maps: dict[int, dict[str, np.ndarray]] = {}
    for canonical_run, condition in sorted(SWATI_CONDITION_BY_CANONICAL_RUN.items()):
        mask_name = _canonical_run_target_name(canonical_run, top_percent=top_percent)
        mask_path = run_mask_dir / f"{condition}_isc_top{int(top_percent)}.nii.gz"
        if not mask_path.exists():
            raise FileNotFoundError(
                f"Missing run-top{int(top_percent)} mask for {condition}: {mask_path}"
            )

        mask_img = nib.load(str(mask_path))
        resampled = resample_from_to(
            mask_img,
            (analysis_img.shape, analysis_img.affine),
            order=0,
        )
        mask_bool = resampled.get_fdata() > 0.0

        index_maps[canonical_run] = {
            mask_name: _mask_to_vector_indices(
                mask_bool=mask_bool,
                analysis_mask_bool=analysis_mask_bool,
                analysis_flat_indices=analysis_flat_indices,
                mask_name=mask_name,
            )
        }

    return index_maps


def _build_bold_path_map(bold_summary_df: pd.DataFrame) -> dict[tuple[str, int], Path]:
    required_columns = {"subject", "run", "bold_z_path"}
    missing = required_columns.difference(bold_summary_df.columns)
    if missing:
        raise ValueError(f"alignment_bold_summary.csv missing columns: {sorted(missing)}")

    duplicate_check = bold_summary_df.groupby(["subject", "run"]).size()
    duplicates = duplicate_check[duplicate_check > 1]
    if not duplicates.empty:
        raise ValueError(
            "Expected unique bold_z_path per (subject,run). Found duplicates for: "
            + ", ".join([f"({subject},{run})" for subject, run in duplicates.index.tolist()])
        )

    mapping: dict[tuple[str, int], Path] = {}
    for row in bold_summary_df.itertuples(index=False):
        subject = str(getattr(row, "subject"))
        run = int(getattr(row, "run"))
        path = Path(str(getattr(row, "bold_z_path")))
        mapping[(subject, run)] = path

    return mapping


def _build_regressor_path_map(
    regressor_summary_df: pd.DataFrame,
    model_slug: str,
) -> tuple[dict[tuple[int, int], Path], str]:
    required_columns = {"model_id", "model_slug", "run", "layer_idx", "regressor_z_path"}
    missing = required_columns.difference(regressor_summary_df.columns)
    if missing:
        raise ValueError(f"alignment_regressor_summary.csv missing columns: {sorted(missing)}")

    model_df = regressor_summary_df[regressor_summary_df["model_slug"] == model_slug].copy()
    if model_df.empty:
        available = sorted(set(str(value) for value in regressor_summary_df["model_slug"].tolist()))
        raise ValueError(
            f"No regressors found for model_slug={model_slug}. Available slugs: {available}"
        )

    selected_model_id = str(model_df.iloc[0]["model_id"])

    duplicate_check = model_df.groupby(["run", "layer_idx"]).size()
    duplicates = duplicate_check[duplicate_check > 1]
    if not duplicates.empty:
        details = ", ".join([f"(run={run},layer={layer})" for run, layer in duplicates.index.tolist()])
        raise ValueError(
            "Multiple regressors found for run/layer pairs. "
            "This fitter expects one regressor per (run,layer) for the chosen model. "
            f"Conflicts: {details}"
        )

    mapping: dict[tuple[int, int], Path] = {}
    for row in model_df.itertuples(index=False):
        run = int(getattr(row, "run"))
        layer = int(getattr(row, "layer_idx"))
        path = Path(str(getattr(row, "regressor_z_path")))
        mapping[(run, layer)] = path

    return mapping, selected_model_id


def _load_subject_roi_runs(
    subject: str,
    runs: list[int],
    bold_path_map: dict[tuple[str, int], Path],
    participant_run_map: dict[str, dict[int, int]],
    target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]],
) -> dict[int, dict[str, np.ndarray]]:
    out: dict[int, dict[str, np.ndarray]] = {}

    for run in runs:
        key = (subject, run)
        if key not in bold_path_map:
            raise KeyError(f"Missing BOLD cache for subject={subject}, run={run}")

        bold_path = bold_path_map[key]
        bold_mmap = np.load(bold_path, mmap_mode="r")
        canonical_run = resolve_subject_canonical_run(
            participant_run_map=participant_run_map,
            subject=subject,
            run=run,
        )
        if canonical_run not in target_index_map_by_canonical_run:
            raise KeyError(
                f"Missing target mask indices for canonical_run={canonical_run}, subject={subject}, run={run}"
            )

        roi_run: dict[str, np.ndarray] = {}
        for roi_name, roi_indices in target_index_map_by_canonical_run[canonical_run].items():
            roi_run[roi_name] = np.asarray(bold_mmap[:, roi_indices], dtype=np.float32)

        out[run] = roi_run

    return out


def _ridge_projection_matrix(x_train: np.ndarray, alpha: float) -> np.ndarray:
    if alpha <= 0:
        raise ValueError("alpha must be positive")

    x_train64 = np.asarray(x_train, dtype=np.float64)

    xtx = x_train64.T @ x_train64
    reg = np.eye(xtx.shape[0], dtype=np.float64) * float(alpha)

    return np.linalg.solve(xtx + reg, x_train64.T)


def _safe_corr_per_voxel(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> np.ndarray:
    yt = np.asarray(y_true, dtype=np.float64)
    yp = np.asarray(y_pred, dtype=np.float64)

    yt_centered = yt - np.mean(yt, axis=0, keepdims=True)
    yp_centered = yp - np.mean(yp, axis=0, keepdims=True)

    numerator = np.sum(yt_centered * yp_centered, axis=0)
    denom = np.sqrt(np.sum(yt_centered**2, axis=0) * np.sum(yp_centered**2, axis=0))

    corr = np.full(denom.shape, np.nan, dtype=np.float64)
    valid = denom > eps
    corr[valid] = numerator[valid] / denom[valid]

    return corr


def _safe_r2_per_voxel(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> np.ndarray:
    yt = np.asarray(y_true, dtype=np.float64)
    yp = np.asarray(y_pred, dtype=np.float64)

    sse = np.sum((yt - yp) ** 2, axis=0)
    yt_mean = np.mean(yt, axis=0, keepdims=True)
    sst = np.sum((yt - yt_mean) ** 2, axis=0)

    r2 = np.full(sst.shape, np.nan, dtype=np.float64)
    valid = sst > eps
    r2[valid] = 1.0 - (sse[valid] / sst[valid])

    return r2


def _two_v_two_accuracy(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> tuple[float, int]:
    yt = np.asarray(y_true, dtype=np.float64)
    yp = np.asarray(y_pred, dtype=np.float64)

    if yt.shape != yp.shape:
        raise ValueError(
            "2v2 accuracy expects y_true and y_pred to share shape, got "
            f"{yt.shape} and {yp.shape}"
        )
    if yt.ndim != 2:
        raise ValueError(f"2v2 accuracy expects 2D matrices, got ndim={yt.ndim}")
    if yt.shape[0] < 2 or yt.shape[1] == 0:
        return float("nan"), 0

    yt_centered = yt - np.mean(yt, axis=1, keepdims=True)
    yp_centered = yp - np.mean(yp, axis=1, keepdims=True)

    yt_norm = np.linalg.norm(yt_centered, axis=1)
    yp_norm = np.linalg.norm(yp_centered, axis=1)

    valid_rows = (yt_norm > eps) & (yp_norm > eps)
    if int(np.sum(valid_rows)) < 2:
        return float("nan"), 0

    yt_unit = yt_centered[valid_rows] / yt_norm[valid_rows, None]
    yp_unit = yp_centered[valid_rows] / yp_norm[valid_rows, None]

    similarity = yt_unit @ yp_unit.T
    diagonal = np.diag(similarity)

    # Margin > 0 means matched assignment beats swapped assignment for a pair.
    pair_margin = diagonal[:, None] + diagonal[None, :] - similarity - similarity.T
    pair_idx = np.triu_indices(pair_margin.shape[0], k=1)
    margins = pair_margin[pair_idx]

    n_pairs = int(margins.size)
    if n_pairs == 0:
        return float("nan"), 0

    wins = float(np.sum(margins > eps))
    ties = float(np.sum(np.abs(margins) <= eps))
    accuracy = (wins + 0.5 * ties) / float(n_pairs)

    return float(accuracy), n_pairs


def _score_matrix(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
    corr = _safe_corr_per_voxel(y_true=y_true, y_pred=y_pred)
    r2 = _safe_r2_per_voxel(y_true=y_true, y_pred=y_pred)
    two_v_two_accuracy, n_2v2_pairs = _two_v_two_accuracy(y_true=y_true, y_pred=y_pred)

    finite_corr = corr[np.isfinite(corr)]
    finite_r2 = r2[np.isfinite(r2)]

    return {
        "mean_corr": float(np.mean(finite_corr)) if finite_corr.size > 0 else float("nan"),
        "median_corr": float(np.median(finite_corr)) if finite_corr.size > 0 else float("nan"),
        "mean_r2": float(np.mean(finite_r2)) if finite_r2.size > 0 else float("nan"),
        "two_v_two_accuracy": float(two_v_two_accuracy),
        "n_2v2_pairs": int(n_2v2_pairs),
        "n_voxels_scored": int(finite_corr.size),
    }


def _collect_protocol_c_subjects(cross_subject_df: pd.DataFrame) -> list[str]:
    if cross_subject_df.empty:
        return []

    subjects: set[str] = set()
    for row in cross_subject_df.itertuples(index=False):
        subjects.update(_parse_subject_field(getattr(row, "train_subjects")))
        subjects.update(_parse_subject_field(getattr(row, "test_subjects")))

    return sorted(subjects)


def _evaluate_protocol_c(
    layers: list[int],
    alpha: float,
    cross_subject_df: pd.DataFrame,
    bold_path_map: dict[tuple[str, int], Path],
    regressor_path_map: dict[tuple[int, int], Path],
    target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]],
    participant_run_map: dict[str, dict[int, int]],
    model_slug: str,
    model_id: str,
    n_workers: int = 1,
    blas_threads_per_worker: int = 0,
) -> pd.DataFrame:
    if cross_subject_df.empty:
        return pd.DataFrame()

    all_subjects = _collect_protocol_c_subjects(cross_subject_df)
    canonical_runs = sorted({int(value) for value in cross_subject_df["canonical_run"].tolist()})

    required_actual_runs: dict[str, list[int]] = {}
    for subject in all_subjects:
        actual_runs = sorted(
            {
                resolve_subject_actual_run(
                    participant_run_map=participant_run_map,
                    subject=subject,
                    canonical_run=canonical_run,
                )
                for canonical_run in canonical_runs
            }
        )
        required_actual_runs[subject] = actual_runs

    subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]] = {}
    for subject in all_subjects:
        subject_roi_cache[subject] = _load_subject_roi_runs(
            subject=subject,
            runs=required_actual_runs[subject],
            bold_path_map=bold_path_map,
            participant_run_map=participant_run_map,
            target_index_map_by_canonical_run=target_index_map_by_canonical_run,
        )

    if blas_threads_per_worker <= 0:
        blas_threads_per_worker = max(1, (os.cpu_count() or 4) // max(1, n_workers))

    def _layer_task(layer_idx: int) -> list[dict[str, Any]]:
        return _evaluate_protocol_c_layer(
            layer_idx=int(layer_idx),
            alpha=alpha,
            cross_subject_df=cross_subject_df,
            regressor_path_map=regressor_path_map,
            canonical_runs=canonical_runs,
            subject_roi_cache=subject_roi_cache,
            target_index_map_by_canonical_run=target_index_map_by_canonical_run,
            participant_run_map=participant_run_map,
            model_slug=model_slug,
            model_id=model_id,
            blas_threads=blas_threads_per_worker,
        )

    if n_workers > 1 and len(layers) > 1:
        from joblib import Parallel, delayed

        print(
            f"[fit] Protocol C: parallel layers across {n_workers} threads "
            f"(layers={len(layers)}, blas_threads/worker={blas_threads_per_worker})",
            flush=True,
        )
        results = Parallel(n_jobs=int(n_workers), prefer="threads")(
            delayed(_layer_task)(layer_idx) for layer_idx in layers
        )
    else:
        results = [_layer_task(layer_idx) for layer_idx in layers]

    rows: list[dict[str, Any]] = []
    for layer_rows in results:
        rows.extend(layer_rows)

    df = pd.DataFrame(rows)
    if not df.empty:
        df = df.sort_values(["layer_idx", "fold_id", "subject", "roi_name"]).reset_index(drop=True)
    return df


def _evaluate_protocol_c_noise_ceiling(
    cross_subject_df: pd.DataFrame,
    bold_path_map: dict[tuple[str, int], Path],
    target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]],
    participant_run_map: dict[str, dict[int, int]],
) -> pd.DataFrame:
    if cross_subject_df.empty:
        return pd.DataFrame()

    all_subjects = _collect_protocol_c_subjects(cross_subject_df)
    canonical_runs = sorted({int(value) for value in cross_subject_df["canonical_run"].tolist()})

    required_actual_runs: dict[str, list[int]] = {}
    for subject in all_subjects:
        actual_runs = sorted(
            {
                resolve_subject_actual_run(
                    participant_run_map=participant_run_map,
                    subject=subject,
                    canonical_run=canonical_run,
                )
                for canonical_run in canonical_runs
            }
        )
        required_actual_runs[subject] = actual_runs

    subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]] = {}
    for subject in all_subjects:
        subject_roi_cache[subject] = _load_subject_roi_runs(
            subject=subject,
            runs=required_actual_runs[subject],
            bold_path_map=bold_path_map,
            participant_run_map=participant_run_map,
            target_index_map_by_canonical_run=target_index_map_by_canonical_run,
        )

    rows: list[dict[str, Any]] = []
    for fold_row in cross_subject_df.itertuples(index=False):
        fold_id = str(getattr(fold_row, "fold_id"))
        canonical_run = int(getattr(fold_row, "canonical_run"))
        condition_label = str(getattr(fold_row, "condition_label"))
        train_subjects = _parse_subject_field(getattr(fold_row, "train_subjects"))
        test_subjects = _parse_subject_field(getattr(fold_row, "test_subjects"))

        if not train_subjects or not test_subjects:
            raise ValueError(f"Protocol C fold {fold_id} must have non-empty train and test subjects")

        for roi_name in target_index_map_by_canonical_run[canonical_run].keys():
            y_train_blocks: list[np.ndarray] = []
            for subject in train_subjects:
                actual_run = resolve_subject_actual_run(
                    participant_run_map=participant_run_map,
                    subject=subject,
                    canonical_run=canonical_run,
                )
                y_train_subject = subject_roi_cache[subject][actual_run][roi_name]
                y_train_blocks.append(y_train_subject)

            y_train_stack = np.stack(y_train_blocks, axis=0)
            y_pred = np.mean(y_train_stack, axis=0, dtype=np.float64)

            for subject in test_subjects:
                actual_run = resolve_subject_actual_run(
                    participant_run_map=participant_run_map,
                    subject=subject,
                    canonical_run=canonical_run,
                )
                y_test = subject_roi_cache[subject][actual_run][roi_name]
                if y_test.shape != y_pred.shape:
                    raise ValueError(
                        "Noise ceiling test TR/voxel mismatch for "
                        f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, "
                        f"roi={roi_name}. y_pred={y_pred.shape}, y_test={y_test.shape}"
                    )

                scores = _score_matrix(y_true=y_test, y_pred=y_pred)

                rows.append(
                    {
                        "protocol": "C_cross_subject_noise_ceiling",
                        "subject": subject,
                        "run": int(actual_run),
                        "canonical_run": int(canonical_run),
                        "condition_label": condition_label,
                        "fold_id": fold_id,
                        "train_subjects": ",".join(train_subjects),
                        "test_subjects": ",".join(test_subjects),
                        "n_train_subjects": int(len(train_subjects)),
                        "n_test_subjects": int(len(test_subjects)),
                        "roi_name": roi_name,
                        "n_train_tr": int(y_pred.shape[0] * len(train_subjects)),
                        "n_test_tr": int(y_test.shape[0]),
                        "n_voxels_roi": int(y_test.shape[1]),
                        **scores,
                    }
                )

    df = pd.DataFrame(rows)
    if not df.empty:
        df = df.sort_values(["fold_id", "subject", "roi_name"]).reset_index(drop=True)
    return df


def _evaluate_protocol_c_layer(
    layer_idx: int,
    alpha: float,
    cross_subject_df: pd.DataFrame,
    regressor_path_map: dict[tuple[int, int], Path],
    canonical_runs: list[int],
    subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]],
    target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]],
    participant_run_map: dict[str, dict[int, int]],
    model_slug: str,
    model_id: str,
    blas_threads: int,
) -> list[dict[str, Any]]:
    """Evaluate Protocol C for a single layer. Returns list of result rows."""
    from threadpoolctl import threadpool_limits

    rows: list[dict[str, Any]] = []

    with threadpool_limits(limits=int(blas_threads)):
        x_by_canonical_run: dict[int, np.ndarray] = {}
        for canonical_run in canonical_runs:
            key = (canonical_run, layer_idx)
            if key not in regressor_path_map:
                raise KeyError(
                    "Missing regressor cache for "
                    f"canonical_run={canonical_run}, layer={layer_idx}, model={model_slug}"
                )
            x_by_canonical_run[canonical_run] = np.asarray(np.load(regressor_path_map[key]), dtype=np.float32)

        projector_cache: dict[tuple[int, int], tuple[np.ndarray, int]] = {}

        for fold_row in cross_subject_df.itertuples(index=False):
            fold_id = str(getattr(fold_row, "fold_id"))
            canonical_run = int(getattr(fold_row, "canonical_run"))
            condition_label = str(getattr(fold_row, "condition_label"))
            train_subjects = _parse_subject_field(getattr(fold_row, "train_subjects"))
            test_subjects = _parse_subject_field(getattr(fold_row, "test_subjects"))

            if not train_subjects or not test_subjects:
                raise ValueError(f"Protocol C fold {fold_id} must have non-empty train and test subjects")

            x_test = x_by_canonical_run[canonical_run]
            projector_key = (canonical_run, len(train_subjects))
            if projector_key not in projector_cache:
                x_train = np.vstack([x_test for _ in train_subjects])
                projector_cache[projector_key] = (
                    _ridge_projection_matrix(x_train=x_train, alpha=alpha),
                    int(x_train.shape[0]),
                )

            projector, n_train_tr = projector_cache[projector_key]

            for roi_name in target_index_map_by_canonical_run[canonical_run].keys():
                y_train_blocks: list[np.ndarray] = []
                for subject in train_subjects:
                    actual_run = resolve_subject_actual_run(
                        participant_run_map=participant_run_map,
                        subject=subject,
                        canonical_run=canonical_run,
                    )
                    y_train_subject = subject_roi_cache[subject][actual_run][roi_name]
                    if y_train_subject.shape[0] != x_test.shape[0]:
                        raise ValueError(
                            "Protocol C train TR mismatch between regressors and BOLD for "
                            f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, "
                            f"layer={layer_idx}, roi={roi_name}. "
                            f"x_train_tr_per_subject={x_test.shape[0]}, y_train_tr={y_train_subject.shape[0]}"
                        )
                    y_train_blocks.append(y_train_subject)

                y_train = np.vstack(y_train_blocks)
                if y_train.shape[0] != n_train_tr:
                    raise ValueError(
                        "Protocol C pooled train TR mismatch between regressors and BOLD for "
                        f"fold={fold_id}, layer={layer_idx}, roi={roi_name}. "
                        f"x_train_tr={n_train_tr}, y_train_tr={y_train.shape[0]}"
                    )

                weights = projector @ np.asarray(y_train, dtype=np.float64)
                y_pred = np.asarray(x_test, dtype=np.float64) @ weights

                for subject in test_subjects:
                    actual_run = resolve_subject_actual_run(
                        participant_run_map=participant_run_map,
                        subject=subject,
                        canonical_run=canonical_run,
                    )
                    y_test = subject_roi_cache[subject][actual_run][roi_name]
                    if y_test.shape[0] != x_test.shape[0]:
                        raise ValueError(
                            "Protocol C test TR mismatch between regressors and BOLD for "
                            f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, "
                            f"layer={layer_idx}, roi={roi_name}. "
                            f"x_test_tr={x_test.shape[0]}, y_test_tr={y_test.shape[0]}"
                        )

                    scores = _score_matrix(y_true=y_test, y_pred=y_pred)

                    rows.append(
                        {
                            "protocol": "C_cross_subject_shared_space",
                            "model_id": model_id,
                            "model_slug": model_slug,
                            "subject": subject,
                            "run": int(actual_run),
                            "canonical_run": int(canonical_run),
                            "condition_label": condition_label,
                            "fold_id": fold_id,
                            "train_subjects": ",".join(train_subjects),
                            "test_subjects": ",".join(test_subjects),
                            "n_train_subjects": int(len(train_subjects)),
                            "n_test_subjects": int(len(test_subjects)),
                            "layer_idx": int(layer_idx),
                            "roi_name": roi_name,
                            "alpha": float(alpha),
                            "n_train_tr": int(n_train_tr),
                            "n_test_tr": int(x_test.shape[0]),
                            "n_voxels_roi": int(y_test.shape[1]),
                            **scores,
                        }
                    )

    return rows


def _summarize_layers(scores_df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    if scores_df.empty:
        return pd.DataFrame(), pd.DataFrame()

    layer_summary = (
        scores_df.groupby(["protocol", "model_id", "model_slug", "layer_idx", "roi_name"], as_index=False)
        .agg(
            mean_corr=("mean_corr", "mean"),
            std_corr=("mean_corr", "std"),
            mean_r2=("mean_r2", "mean"),
            mean_2v2_accuracy=("two_v_two_accuracy", "mean"),
            std_2v2_accuracy=("two_v_two_accuracy", "std"),
            n_records=("mean_corr", "size"),
        )
        .sort_values(["protocol", "roi_name", "layer_idx"])
        .reset_index(drop=True)
    )

    best_rows: list[pd.Series] = []
    for (_, roi_name), group_df in layer_summary.groupby(["protocol", "roi_name"]):
        idx = int(group_df["mean_corr"].idxmax())
        best_rows.append(layer_summary.loc[idx])

    best_layer_summary = pd.DataFrame(best_rows).reset_index(drop=True)
    best_layer_summary = best_layer_summary.sort_values(["protocol", "roi_name"]).reset_index(drop=True)

    return layer_summary, best_layer_summary


def _summarize_noise_ceiling(scores_df: pd.DataFrame) -> pd.DataFrame:
    if scores_df.empty:
        return pd.DataFrame()

    return (
        scores_df.groupby(["protocol", "roi_name"], as_index=False)
        .agg(
            mean_corr=("mean_corr", "mean"),
            std_corr=("mean_corr", "std"),
            mean_r2=("mean_r2", "mean"),
            mean_2v2_accuracy=("two_v_two_accuracy", "mean"),
            std_2v2_accuracy=("two_v_two_accuracy", "std"),
            n_records=("mean_corr", "size"),
        )
        .sort_values(["protocol", "roi_name"])
        .reset_index(drop=True)
    )


def _compare_best_layer_to_noise_ceiling(
    best_layer_df: pd.DataFrame,
    noise_ceiling_summary_df: pd.DataFrame,
) -> pd.DataFrame:
    if best_layer_df.empty or noise_ceiling_summary_df.empty:
        return pd.DataFrame()

    ceiling_df = noise_ceiling_summary_df.rename(
        columns={
            "mean_corr": "noise_ceiling_mean_corr",
            "std_corr": "noise_ceiling_std_corr",
            "mean_r2": "noise_ceiling_mean_r2",
            "mean_2v2_accuracy": "noise_ceiling_mean_2v2_accuracy",
            "std_2v2_accuracy": "noise_ceiling_std_2v2_accuracy",
            "n_records": "noise_ceiling_n_records",
        }
    )

    merged = best_layer_df.merge(ceiling_df, on="roi_name", how="left")
    if merged.empty:
        return merged

    for metric in ["mean_corr", "mean_r2", "mean_2v2_accuracy"]:
        denom = np.asarray(merged[f"noise_ceiling_{metric}"], dtype=np.float64)
        numer = np.asarray(merged[metric], dtype=np.float64)
        ratio = np.full(denom.shape, np.nan, dtype=np.float64)
        valid = np.isfinite(denom) & (np.abs(denom) > 1e-12)
        ratio[valid] = numer[valid] / denom[valid]
        merged[f"fraction_of_noise_ceiling_{metric}"] = ratio

    return merged


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    bootstrap_output_dir = Path(args.bootstrap_output_dir).resolve()
    protocols = _parse_protocols(args.protocols)
    path_map = _resolve_core_paths(bootstrap_output_dir=bootstrap_output_dir)
    _check_required_files(path_map=path_map, protocols=protocols)

    with path_map["bootstrap_summary"].open("r", encoding="utf-8") as handle:
        bootstrap_summary = json.load(handle)

    target_mask_mode = str(args.target_mask_mode).strip().lower()
    if target_mask_mode not in TARGET_MASK_MODE_CHOICES:
        raise ValueError(
            f"Unsupported --target-mask-mode={target_mask_mode!r}. "
            f"Allowed: {sorted(TARGET_MASK_MODE_CHOICES)}"
        )

    roi_mask_dir: Path | None = None
    run_mask_dir: Path | None = None
    target_top_percent: int | None = None
    if target_mask_mode == "core_roi":
        roi_mask_dir = _resolve_core_roi_mask_dir(args=args, bootstrap_summary=bootstrap_summary)
        if not roi_mask_dir.exists():
            raise FileNotFoundError(f"ROI mask dir not found: {roi_mask_dir}")
    else:
        target_top_percent = RUN_MASK_PERCENT_BY_MODE[target_mask_mode]
        run_mask_dir = _resolve_run_mask_dir(args=args, top_percent=target_top_percent)

    output_dir = (
        Path(args.output_dir).resolve()
        if args.output_dir
        else bootstrap_output_dir / "fit_results" / str(args.model_slug)
    )
    ensure_directory(output_dir)

    bold_summary_df = pd.read_csv(path_map["bold_summary"])
    regressor_summary_df = pd.read_csv(path_map["regressor_summary"])
    cross_subject_df = (
        pd.read_csv(path_map["protocol_c"]) if path_map["protocol_c"].exists() else pd.DataFrame()
    )

    requested_subjects = _parse_subjects(args.subjects)

    if requested_subjects is not None and "C" in protocols:
        raise ValueError("Protocol C currently requires --subjects all so fold membership stays valid")

    if requested_subjects is not None:
        subject_set = set(requested_subjects)
        bold_summary_df = bold_summary_df[bold_summary_df["subject"].isin(subject_set)].copy()

    if cross_subject_df.empty and "C" in protocols:
        raise ValueError("Protocol C was requested but no cross-subject rows were available")

    bold_path_map = _build_bold_path_map(bold_summary_df=bold_summary_df)
    regressor_path_map, detected_model_id = _build_regressor_path_map(
        regressor_summary_df=regressor_summary_df,
        model_slug=str(args.model_slug),
    )

    model_id = str(args.model_id) if args.model_id else detected_model_id

    available_layers = sorted({layer for (_, layer) in regressor_path_map.keys()})
    requested_layers = _parse_layers(args.layer_indices)

    if requested_layers is None:
        layers = available_layers
    else:
        layer_set = set(available_layers)
        missing_layers = [layer for layer in requested_layers if layer not in layer_set]
        if missing_layers:
            raise ValueError(
                f"Requested layers are unavailable for model {args.model_slug}: {missing_layers}. "
                f"Available: {available_layers}"
            )
        layers = requested_layers

    if not layers:
        raise ValueError("No layers selected for fitting")

    protocol_c_subjects = _collect_protocol_c_subjects(cross_subject_df) if "C" in protocols else []
    subjects = sorted(set(protocol_c_subjects))
    if not subjects:
        raise ValueError("No subjects available after filtering")

    participant_run_info_path = Path(args.participant_run_info).resolve()
    required_runs = sorted({int(value) for value in bold_summary_df["run"].tolist()})
    if not required_runs:
        raise ValueError("No required runs detected for participant run mapping")

    participant_run_map = load_participant_run_map(
        participant_run_info_path=participant_run_info_path,
        subjects=subjects,
        runs=required_runs,
    )

    if target_mask_mode == "core_roi":
        target_index_map_by_canonical_run = _build_core_roi_index_maps(
            analysis_mask_path=path_map["analysis_mask"],
            roi_mask_dir=roi_mask_dir,
        )
        target_mask_source_dir = roi_mask_dir
    else:
        target_index_map_by_canonical_run = _build_run_mask_index_maps(
            analysis_mask_path=path_map["analysis_mask"],
            run_mask_dir=run_mask_dir,
            top_percent=target_top_percent,
        )
        target_mask_source_dir = run_mask_dir

    canonical_runs_used = sorted(
        {
            canonical_run
            for subject in subjects
            for canonical_run in participant_run_map[subject].values()
        }
    )

    protocol_c_df = pd.DataFrame()
    noise_ceiling_df = pd.DataFrame()

    n_fit_workers, blas_threads_per_worker = _resolve_num_fit_workers(
        requested=str(args.num_fit_workers),
        n_layers=int(len(layers)),
    )
    if n_fit_workers > 1:
        print(
            f"[fit] Layer-parallel fitting enabled: workers={n_fit_workers}, "
            f"blas_threads/worker={blas_threads_per_worker}, layers={len(layers)}, "
            f"cpu_count={os.cpu_count()}",
            flush=True,
        )

    if "C" in protocols:
        protocol_c_df = _evaluate_protocol_c(
            layers=layers,
            alpha=float(args.alpha),
            cross_subject_df=cross_subject_df,
            bold_path_map=bold_path_map,
            regressor_path_map=regressor_path_map,
            target_index_map_by_canonical_run=target_index_map_by_canonical_run,
            participant_run_map=participant_run_map,
            model_slug=str(args.model_slug),
            model_id=model_id,
            n_workers=n_fit_workers,
            blas_threads_per_worker=blas_threads_per_worker,
        )
        protocol_c_df.to_csv(output_dir / "protocol_c_core_roi_scores.csv", index=False)
        noise_ceiling_df = _evaluate_protocol_c_noise_ceiling(
            cross_subject_df=cross_subject_df,
            bold_path_map=bold_path_map,
            target_index_map_by_canonical_run=target_index_map_by_canonical_run,
            participant_run_map=participant_run_map,
        )
        noise_ceiling_df.to_csv(output_dir / "noise_ceiling_protocol_c_core_roi_scores.csv", index=False)

    combined_df = protocol_c_df.copy()
    combined_df.to_csv(output_dir / "core_roi_scores_all.csv", index=False)

    layer_summary_df, best_layer_df = _summarize_layers(scores_df=combined_df)
    layer_summary_df.to_csv(output_dir / "core_roi_layer_summary.csv", index=False)
    best_layer_df.to_csv(output_dir / "core_roi_best_layer_summary.csv", index=False)
    noise_ceiling_summary_df = _summarize_noise_ceiling(scores_df=noise_ceiling_df)
    noise_ceiling_summary_df.to_csv(output_dir / "noise_ceiling_core_roi_summary.csv", index=False)
    best_vs_noise_ceiling_df = _compare_best_layer_to_noise_ceiling(
        best_layer_df=best_layer_df,
        noise_ceiling_summary_df=noise_ceiling_summary_df,
    )
    best_vs_noise_ceiling_df.to_csv(output_dir / "best_layer_vs_noise_ceiling.csv", index=False)

    run_summary = {
        "bootstrap_output_dir": str(bootstrap_output_dir),
        "output_dir": str(output_dir),
        "model_slug": str(args.model_slug),
        "model_id": model_id,
        "alpha": float(args.alpha),
        "subjects": subjects,
        "n_subjects": int(len(subjects)),
        "layers": [int(value) for value in layers],
        "n_layers": int(len(layers)),
        "protocols": sorted(protocols),
        "target_mask_mode": target_mask_mode,
        "target_top_percent": int(target_top_percent) if target_top_percent is not None else None,
        "target_mask_source_dir": str(target_mask_source_dir),
        "roi_mask_dir": str(roi_mask_dir) if roi_mask_dir is not None else None,
        "run_mask_dir": str(run_mask_dir) if run_mask_dir is not None else None,
        "run_top10_mask_dir": str(run_mask_dir) if target_top_percent == 10 and run_mask_dir is not None else None,
        "run_top25_mask_dir": str(run_mask_dir) if target_top_percent == 25 and run_mask_dir is not None else None,
        "participant_run_info_path": str(participant_run_info_path),
        "required_runs": [int(value) for value in required_runs],
        "canonical_runs_used": [int(value) for value in canonical_runs_used],
        "target_names_by_canonical_run": {
            str(canonical_run): list(target_index_map_by_canonical_run[canonical_run].keys())
            for canonical_run in sorted(target_index_map_by_canonical_run)
        },
        "n_protocol_c_rows": int(len(protocol_c_df)),
        "n_noise_ceiling_rows": int(len(noise_ceiling_df)),
        "n_total_rows": int(len(combined_df)),
    }
    write_json(output_dir / "fit_run_summary.json", run_summary)

    print("=" * 72)
    print("A1 fit complete")
    print(f"Model slug: {args.model_slug}")
    print(f"Model id: {model_id}")
    print(f"Subjects: {len(subjects)}")
    print(f"Layers: {len(layers)}")
    print(f"Alpha: {float(args.alpha)}")
    print(f"Target mask mode: {target_mask_mode}")
    print(f"Target mask source: {target_mask_source_dir}")
    print(f"Protocol C rows: {len(protocol_c_df)}")
    print(f"Noise ceiling rows: {len(noise_ceiling_df)}")
    print(f"Output directory: {output_dir}")
    print("=" * 72)


if __name__ == "__main__":
    main()