code/a1_pipeline/alignment.py

"""TR-level HRF regressor and z-score alignment utilities."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import nibabel as nib
import numpy as np
import pandas as pd


def _zscore_columns(values: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    mean = np.mean(values, axis=0, keepdims=True)
    std = np.std(values, axis=0, keepdims=True)
    std[std < eps] = 1.0
    return ((values - mean) / std).astype(np.float32)


def _trim_time_axis(values: np.ndarray, trim_start_tr: int, trim_end_tr: int) -> np.ndarray:
    if trim_start_tr < 0 or trim_end_tr < 0:
        raise ValueError("trim_start_tr and trim_end_tr must be non-negative")

    start = int(trim_start_tr)
    stop = values.shape[0] - int(trim_end_tr)

    if start >= stop:
        raise ValueError(
            "Trimming removed all time points: "
            f"n_tr={values.shape[0]}, trim_start={trim_start_tr}, trim_end={trim_end_tr}"
        )

    return values[start:stop]


def _format_tr_for_filename(tr_seconds: float) -> str:
    return f"{tr_seconds:.4f}".replace(".", "p")


def _assert_same_grid(img: nib.Nifti1Image, reference_img: nib.Nifti1Image, atol: float = 1e-5) -> None:
    if img.shape[:3] != reference_img.shape:
        raise ValueError(
            "Image shape mismatch against analysis mask grid: "
            f"image={img.shape[:3]}, reference={reference_img.shape}"
        )
    if not np.allclose(img.affine, reference_img.affine, atol=atol):
        raise ValueError("Image affine mismatch against analysis mask grid")


def _extract_masked_bold_2d(
    bold_path: Path,
    analysis_mask_bool: np.ndarray,
    analysis_mask_img: nib.Nifti1Image,
) -> np.ndarray:
    img = nib.load(str(bold_path))
    if len(img.shape) != 4:
        raise ValueError(f"Expected 4D BOLD image, got shape={img.shape} for {bold_path}")

    _assert_same_grid(img=img, reference_img=analysis_mask_img)

    data = np.asanyarray(img.dataobj)
    # Convert from (x,y,z,t) to (t,voxels)
    bold_2d = data[analysis_mask_bool, :].T.astype(np.float32)
    return bold_2d


def compute_hrf_regressor_matrix(
    word_features: np.ndarray,
    onsets_s: np.ndarray,
    offsets_s: np.ndarray,
    tr_seconds: float,
    n_volumes: int,
    hrf_model: str,
) -> np.ndarray:
    """Build TR-level HRF-convolved regressors from word-level features."""
    from nilearn.glm.first_level import compute_regressor

    if word_features.ndim != 2:
        raise ValueError("word_features must be 2D: (n_words, hidden_dim)")

    if onsets_s.shape[0] != word_features.shape[0] or offsets_s.shape[0] != word_features.shape[0]:
        raise ValueError("onset/offset lengths must match n_words in word_features")

    durations = offsets_s - onsets_s
    frame_times = np.arange(n_volumes, dtype=np.float64) * float(tr_seconds) + 0.5 * float(tr_seconds)

    regressors = np.zeros((int(n_volumes), int(word_features.shape[1])), dtype=np.float32)

    for feature_index in range(word_features.shape[1]):
        amplitudes = word_features[:, feature_index].astype(np.float64)
        exp_condition = np.array([onsets_s, durations, amplitudes], dtype=np.float64)
        signal, _ = compute_regressor(exp_condition, hrf_model, frame_times)
        regressors[:, feature_index] = signal[:, 0].astype(np.float32)

    return regressors


@dataclass(frozen=True)
class BoldAlignmentRecord:
    """Summary for one subject-run BOLD alignment artifact."""

    subject: str
    run: int
    tr_seconds: float
    n_volumes_raw: int
    n_tr_aligned: int
    n_voxels: int
    trim_start_tr: int
    trim_end_tr: int
    bold_z_path: str


@dataclass(frozen=True)
class RegressorAlignmentRecord:
    """Summary for one model-run-layer regressor artifact."""

    model_id: str
    model_slug: str
    run: int
    tr_seconds: float
    n_volumes_raw: int
    n_tr_aligned: int
    layer_idx: int
    n_features: int
    trim_start_tr: int
    trim_end_tr: int
    hrf_model: str
    regressor_z_path: str


def build_and_cache_alignment_inputs(
    manifest_df: pd.DataFrame,
    run_events_df: pd.DataFrame,
    feature_summary_df: pd.DataFrame,
    analysis_mask_path: Path,
    output_dir: Path,
    trim_start_tr: int,
    trim_end_tr: int,
    hrf_model: str,
    overwrite: bool,
    alignment_subjects: list[str] | None = None,
) -> tuple[pd.DataFrame, pd.DataFrame, dict[str, Any]]:
    """Cache z-scored BOLD targets and HRF-convolved z-scored regressors."""
    if manifest_df.empty:
        raise ValueError("Manifest is empty; cannot build alignment inputs")
    if run_events_df.empty:
        raise ValueError("run_events_df is empty; cannot build regressors")
    if feature_summary_df.empty:
        raise ValueError("feature_summary_df is empty; feature extraction must run first")

    required_manifest_columns = {"subject", "run", "tr_seconds", "n_volumes", "derivatives_bold_path"}
    missing_manifest = required_manifest_columns.difference(manifest_df.columns)
    if missing_manifest:
        raise ValueError(f"manifest_df missing required columns: {sorted(missing_manifest)}")

    required_feature_columns = {"model_id", "model_slug", "run", "features_npz_path"}
    missing_feature = required_feature_columns.difference(feature_summary_df.columns)
    if missing_feature:
        raise ValueError(f"feature_summary_df missing required columns: {sorted(missing_feature)}")

    analysis_mask_img = nib.load(str(analysis_mask_path))
    analysis_mask_bool = analysis_mask_img.get_fdata() > 0.5

    manifest_alignment_df = manifest_df
    requested_subjects = [value for value in (alignment_subjects or []) if str(value).strip()]
    if requested_subjects:
        subject_set = {str(value) for value in requested_subjects}
        manifest_alignment_df = manifest_df[manifest_df["subject"].isin(subject_set)].copy()
        if manifest_alignment_df.empty:
            raise ValueError(
                "Alignment subject filter produced no rows. "
                f"Requested subjects={sorted(subject_set)}"
            )

    output_dir = output_dir.resolve()
    bold_output_dir = output_dir / "bold_z"
    regressor_output_dir = output_dir / "regressors"
    bold_output_dir.mkdir(parents=True, exist_ok=True)
    regressor_output_dir.mkdir(parents=True, exist_ok=True)

    bold_rows: list[BoldAlignmentRecord] = []

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

        bold_z_path = bold_output_dir / f"{subject}_run-{run:02d}_bold_z.npy"
        if bold_z_path.exists() and not overwrite:
            aligned_bold = np.load(bold_z_path)
        else:
            bold_2d = _extract_masked_bold_2d(
                bold_path=bold_path,
                analysis_mask_bool=analysis_mask_bool,
                analysis_mask_img=analysis_mask_img,
            )

            if bold_2d.shape[0] != n_volumes:
                raise ValueError(
                    f"BOLD volume mismatch for {subject} run {run}: "
                    f"manifest={n_volumes}, loaded={bold_2d.shape[0]}"
                )

            aligned_bold = _trim_time_axis(
                values=bold_2d,
                trim_start_tr=trim_start_tr,
                trim_end_tr=trim_end_tr,
            )
            aligned_bold = _zscore_columns(aligned_bold)
            np.save(bold_z_path, aligned_bold)

        bold_rows.append(
            BoldAlignmentRecord(
                subject=subject,
                run=run,
                tr_seconds=tr_seconds,
                n_volumes_raw=n_volumes,
                n_tr_aligned=int(aligned_bold.shape[0]),
                n_voxels=int(aligned_bold.shape[1]),
                trim_start_tr=int(trim_start_tr),
                trim_end_tr=int(trim_end_tr),
                bold_z_path=str(bold_z_path),
            )
        )

    regressor_rows: list[RegressorAlignmentRecord] = []

    # Compute run configuration keys from manifest; regressors depend on run + TR + n_volumes.
    run_config_map: dict[int, list[tuple[float, int]]] = {}
    for run, group_df in manifest_alignment_df.groupby("run"):
        unique_pairs = sorted(
            {(float(tr), int(n_vol)) for tr, n_vol in zip(group_df["tr_seconds"], group_df["n_volumes"])},
            key=lambda pair: (pair[0], pair[1]),
        )
        run_config_map[int(run)] = unique_pairs

    for feature_row in feature_summary_df.itertuples(index=False):
        model_id = str(getattr(feature_row, "model_id"))
        model_slug = str(getattr(feature_row, "model_slug"))
        run = int(getattr(feature_row, "run"))
        features_npz_path = Path(str(getattr(feature_row, "features_npz_path")))

        if run not in run_config_map:
            continue

        run_events = run_events_df[run_events_df["run"] == run].sort_values("word_index")
        onsets_s = run_events["onset_s"].to_numpy(dtype=np.float64)
        offsets_s = run_events["offset_s"].to_numpy(dtype=np.float64)

        feature_bundle = np.load(features_npz_path)
        layer_keys = sorted(
            [key for key in feature_bundle.files if key.startswith("layer_")],
            key=lambda value: int(value.split("_")[1]),
        )

        for tr_seconds, n_volumes in run_config_map[run]:
            tr_tag = _format_tr_for_filename(tr_seconds)

            model_regressor_dir = regressor_output_dir / model_slug
            model_regressor_dir.mkdir(parents=True, exist_ok=True)

            for layer_key in layer_keys:
                layer_idx = int(layer_key.split("_")[1])
                layer_features = np.asarray(feature_bundle[layer_key], dtype=np.float32)

                if layer_features.shape[0] != onsets_s.shape[0]:
                    raise ValueError(
                        f"Word count mismatch for model={model_id}, run={run}, layer={layer_idx}: "
                        f"features={layer_features.shape[0]}, events={onsets_s.shape[0]}"
                    )

                regressor_z_path = (
                    model_regressor_dir
                    / f"run-{run:02d}_tr-{tr_tag}_nvol-{n_volumes}_layer-{layer_idx:03d}_regressor_z.npy"
                )

                if regressor_z_path.exists() and not overwrite:
                    regressors_aligned = np.load(regressor_z_path)
                else:
                    regressors = compute_hrf_regressor_matrix(
                        word_features=layer_features,
                        onsets_s=onsets_s,
                        offsets_s=offsets_s,
                        tr_seconds=tr_seconds,
                        n_volumes=n_volumes,
                        hrf_model=hrf_model,
                    )
                    regressors_aligned = _trim_time_axis(
                        values=regressors,
                        trim_start_tr=trim_start_tr,
                        trim_end_tr=trim_end_tr,
                    )
                    regressors_aligned = _zscore_columns(regressors_aligned)
                    np.save(regressor_z_path, regressors_aligned)

                regressor_rows.append(
                    RegressorAlignmentRecord(
                        model_id=model_id,
                        model_slug=model_slug,
                        run=run,
                        tr_seconds=tr_seconds,
                        n_volumes_raw=n_volumes,
                        n_tr_aligned=int(regressors_aligned.shape[0]),
                        layer_idx=layer_idx,
                        n_features=int(regressors_aligned.shape[1]),
                        trim_start_tr=int(trim_start_tr),
                        trim_end_tr=int(trim_end_tr),
                        hrf_model=hrf_model,
                        regressor_z_path=str(regressor_z_path),
                    )
                )

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

    regressor_summary_df = pd.DataFrame([asdict(row) for row in regressor_rows])
    if not regressor_summary_df.empty:
        regressor_summary_df = regressor_summary_df.sort_values(
            ["model_slug", "run", "layer_idx", "tr_seconds", "n_volumes_raw"]
        ).reset_index(drop=True)

    alignment_qc: dict[str, Any] = {
        "analysis_mask_path": str(analysis_mask_path),
        "hrf_model": hrf_model,
        "trim_start_tr": int(trim_start_tr),
        "trim_end_tr": int(trim_end_tr),
        "alignment_subjects": sorted({str(value) for value in manifest_alignment_df["subject"].tolist()}),
        "n_bold_rows": int(len(bold_summary_df)),
        "n_regressor_rows": int(len(regressor_summary_df)),
    }

    if not bold_summary_df.empty:
        alignment_qc["bold_n_tr_min"] = int(bold_summary_df["n_tr_aligned"].min())
        alignment_qc["bold_n_tr_max"] = int(bold_summary_df["n_tr_aligned"].max())
        alignment_qc["bold_n_voxels_min"] = int(bold_summary_df["n_voxels"].min())
        alignment_qc["bold_n_voxels_max"] = int(bold_summary_df["n_voxels"].max())

    if not regressor_summary_df.empty:
        alignment_qc["regressor_n_tr_min"] = int(regressor_summary_df["n_tr_aligned"].min())
        alignment_qc["regressor_n_tr_max"] = int(regressor_summary_df["n_tr_aligned"].max())
        alignment_qc["regressor_n_features_min"] = int(regressor_summary_df["n_features"].min())
        alignment_qc["regressor_n_features_max"] = int(regressor_summary_df["n_features"].max())

    return bold_summary_df, regressor_summary_df, alignment_qc