code/a1_pipeline/spatial.py

"""Spatial harmonization utilities for A1 baseline bootstrap."""

from __future__ import annotations

from dataclasses import asdict, dataclass
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


@dataclass(frozen=True)
class RunMaskQC:
    """Per-run QC record used when building the analysis mask."""

    subject: str
    run: int
    derivatives_bold_path: str
    run_mask_method: str
    needs_resample_to_reference: bool
    run_mask_voxels: int


def _make_reference_3d_from_manifest(manifest_df: pd.DataFrame) -> nib.Nifti1Image:
    if manifest_df.empty:
        raise ValueError("Manifest is empty; cannot create spatial reference")

    first_path = Path(str(manifest_df.iloc[0]["derivatives_bold_path"]))
    first_img = nib.load(str(first_path))
    if len(first_img.shape) != 4:
        raise ValueError(f"Expected 4D BOLD image but got shape={first_img.shape} for {first_path}")

    shape_3d = first_img.shape[:3]
    return nib.Nifti1Image(np.zeros(shape_3d, dtype=np.uint8), first_img.affine)


def _same_grid(img_a: nib.Nifti1Image, img_b: nib.Nifti1Image, atol: float = 1e-5) -> bool:
    return img_a.shape == img_b.shape and np.allclose(img_a.affine, img_b.affine, atol=atol)


def _extract_run_mask_from_bold(
    bold_img: nib.Nifti1Image,
    method: str,
    epsilon: float,
) -> nib.Nifti1Image:
    shape = bold_img.shape
    if len(shape) != 4:
        raise ValueError(f"Expected 4D BOLD image, got shape={shape}")

    if method == "first_volume_nonzero":
        volume = np.asanyarray(bold_img.dataobj[..., 0])
        run_mask = np.abs(volume) > epsilon
    elif method == "temporal_any_nonzero":
        run_mask = np.zeros(shape[:3], dtype=bool)
        for time_index in range(shape[3]):
            volume = np.asanyarray(bold_img.dataobj[..., time_index])
            run_mask |= np.abs(volume) > epsilon
    else:
        raise ValueError(f"Unsupported run mask method: {method}")

    return nib.Nifti1Image(run_mask.astype(np.uint8), bold_img.affine)


def _resample_binary_mask_to_reference(
    binary_mask_img: nib.Nifti1Image,
    reference_3d_img: nib.Nifti1Image,
) -> nib.Nifti1Image:
    resampled = resample_from_to(
        binary_mask_img,
        (reference_3d_img.shape, reference_3d_img.affine),
        order=0,
    )
    binary_data = (resampled.get_fdata() > 0.5).astype(np.uint8)
    return nib.Nifti1Image(binary_data, reference_3d_img.affine)


def build_symmetric_analysis_mask(
    manifest_df: pd.DataFrame,
    run_mask_method: str = "first_volume_nonzero",
    epsilon: float = 1e-6,
) -> tuple[nib.Nifti1Image, dict[tuple[str, int], np.ndarray], pd.DataFrame, dict[str, Any]]:
    """Build analysis mask in a canonical grid with left-right symmetry enforced.

    Returns:
        analysis_mask_img: symmetric boolean mask image in reference grid
        run_mask_map: (subject, run) -> boolean 3D run mask in reference grid
        run_mask_qc_df: per-run mask stats including whether resampling was needed
        mask_qc: summary dictionary
    """
    if manifest_df.empty:
        raise ValueError("Manifest is empty; cannot build analysis mask")

    required_columns = {"subject", "run", "derivatives_bold_path"}
    missing = required_columns.difference(manifest_df.columns)
    if missing:
        raise ValueError(f"Manifest missing required columns: {sorted(missing)}")

    reference_3d_img = _make_reference_3d_from_manifest(manifest_df)
    reference_affine = reference_3d_img.affine

    run_mask_map: dict[tuple[str, int], np.ndarray] = {}
    run_qc_rows: list[RunMaskQC] = []

    intersection_mask: np.ndarray | None = None
    n_resampled = 0

    for row in manifest_df.itertuples(index=False):
        subject = str(getattr(row, "subject"))
        run = int(getattr(row, "run"))
        bold_path = Path(str(getattr(row, "derivatives_bold_path")))

        bold_img = nib.load(str(bold_path))
        run_mask_img = _extract_run_mask_from_bold(
            bold_img=bold_img,
            method=run_mask_method,
            epsilon=epsilon,
        )

        needs_resample = not _same_grid(run_mask_img, reference_3d_img)
        if needs_resample:
            n_resampled += 1
            run_mask_img = _resample_binary_mask_to_reference(run_mask_img, reference_3d_img)

        run_mask_bool = run_mask_img.get_fdata() > 0.5
        run_key = (subject, run)
        run_mask_map[run_key] = run_mask_bool

        if intersection_mask is None:
            intersection_mask = run_mask_bool.copy()
        else:
            intersection_mask &= run_mask_bool

        run_qc_rows.append(
            RunMaskQC(
                subject=subject,
                run=run,
                derivatives_bold_path=str(bold_path),
                run_mask_method=run_mask_method,
                needs_resample_to_reference=needs_resample,
                run_mask_voxels=int(run_mask_bool.sum()),
            )
        )

    if intersection_mask is None:
        raise RuntimeError("Could not compute intersection mask from manifest entries")

    swapped_mask = np.flip(intersection_mask, axis=0)
    symmetric_mask = intersection_mask & swapped_mask

    analysis_mask_img = nib.Nifti1Image(symmetric_mask.astype(np.uint8), reference_affine)

    shape_x = symmetric_mask.shape[0]
    mid_x = shape_x // 2
    left_voxels = int(symmetric_mask[:mid_x, :, :].sum())
    right_voxels = int(symmetric_mask[-mid_x:, :, :].sum()) if mid_x > 0 else 0

    run_mask_qc_df = pd.DataFrame([asdict(row) for row in run_qc_rows])
    if not run_mask_qc_df.empty:
        run_mask_qc_df = run_mask_qc_df.sort_values(["subject", "run"]).reset_index(drop=True)

    orientation_codes = "".join(nib.aff2axcodes(reference_affine))

    mask_qc: dict[str, Any] = {
        "run_mask_method": run_mask_method,
        "n_manifest_rows": int(len(manifest_df)),
        "n_run_masks": int(len(run_mask_map)),
        "n_resampled_run_masks": int(n_resampled),
        "reference_shape": [int(dim) for dim in analysis_mask_img.shape],
        "reference_orientation": orientation_codes,
        "intersection_voxels": int(intersection_mask.sum()),
        "symmetric_mask_voxels": int(symmetric_mask.sum()),
        "left_hemisphere_voxels": left_voxels,
        "right_hemisphere_voxels": right_voxels,
    }

    if not run_mask_qc_df.empty:
        mask_qc["run_mask_voxels_min"] = int(run_mask_qc_df["run_mask_voxels"].min())
        mask_qc["run_mask_voxels_max"] = int(run_mask_qc_df["run_mask_voxels"].max())

    return analysis_mask_img, run_mask_map, run_mask_qc_df, mask_qc