Source code for thesis.workflows.transforms.operations

"""Common, reusable ANTs transform operations.

These are pure functions that can be called from any Nipype Function node or
used directly.  Every function that calls ANTs must be self-contained (no thesis
logger) because Nipype Function nodes run in isolated subprocesses.
"""

from __future__ import annotations

from pathlib import Path
from typing import List, Optional

import nibabel as nib
from nibabel.nifti1 import Nifti1Header

__all__ = [
    "apply_transform_ants",
    "get_transform_type_from_paths",
    "build_output_filename",
    "validate_transform_inputs",
]

_WARP_KEYWORDS = ("Warp", "warp", "field", "Field")
_AFFINE_KEYWORDS = ("Affine", "affine", "GenericAffine")
_RIGID_KEYWORDS = ("Rigid", "rigid")




def apply_transform_ants(
    input_image: Path,
    output_image: Path,
    transforms: List[Path],
    reference_image: Path,
    interpolation: str = "Linear",
    invert_transform_flags: Optional[List[bool]] = None,
) -> Path:
    """Apply a sequence of ANTs transforms to a single image.

    Wraps ANTs ApplyTransforms and additionally copies qform → sform so that
    downstream FSL tools (which read sform) work correctly.

    Args:
        input_image: Source image path.
        output_image: Destination path for the transformed image.
        transforms: Ordered list of transform paths (warp fields, affine mats, …).
        reference_image: Reference image that defines the target voxel grid.
        interpolation: ANTs interpolation mode (e.g. ``"Linear"``,
            ``"NearestNeighbor"``, ``"BSpline"``).
        invert_transform_flags: Optional per-transform inversion flags.  Defaults
            to all ``False`` when not provided.

    Returns:
        Path to the written output image.

    Raises:
        FileNotFoundError: If ``input_image``, any entry in ``transforms``, or
            ``reference_image`` does not exist.
        RuntimeError: If the underlying ANTs command fails.
    """
    from nipype.interfaces import ants  # noqa: PLC0415

    if not input_image.exists():
        raise FileNotFoundError(
            f"Input image not found: '{input_image}'. "
            "Check that the file exists before running the transform node."
        )
    for t in transforms:
        if not Path(t).exists():
            raise FileNotFoundError(f"Transform file not found: '{t}'.")
    if not reference_image.exists():
        raise FileNotFoundError(f"Reference image not found: '{reference_image}'.")

    output_image.parent.mkdir(parents=True, exist_ok=True)

    flags = invert_transform_flags if invert_transform_flags else [False] * len(transforms)

    applier = ants.ApplyTransforms()
    applier.inputs.input_image = str(input_image.resolve())
    applier.inputs.reference_image = str(reference_image.resolve())
    applier.inputs.output_image = str(output_image)
    applier.inputs.transforms = [str(Path(p).resolve()) for p in transforms]
    applier.inputs.interpolation = interpolation
    applier.inputs.invert_transform_flags = list(flags)
    applier.run()

    # ANTs ApplyTransforms sets sform_code=0 in its output leaving only the
    # qform populated.  FSL tools read the sform matrix directly; a zero/unset
    # sform causes an "inv(): matrix is singular" crash.  Copy qform → sform.
    loaded = nib.load(str(output_image))
    if isinstance(loaded.header, Nifti1Header) and loaded.header.get_sform(coded=True)[1] == 0:
        qform, qcode = loaded.header.get_qform(coded=True)
        loaded.header.set_sform(qform, code=max(qcode, 1))
        nib.save(loaded, str(output_image))

    return output_image





def get_transform_type_from_paths(transform_paths: List[Path]) -> str:
    """Infer a transform type label from the names of transform files.

    The label is used in output filenames to document which transforms were
    applied.  Heuristics are based on common ANTs naming conventions.

    Args:
        transform_paths: List of transform file paths to inspect.

    Returns:
        One of ``"SyN"``, ``"Affine"``, ``"Rigid"``, or ``"transformed"``
        (fallback when no recognisable pattern is found).

    Example:
        >>> paths = [Path("sub_1Warp.nii.gz"), Path("sub_0GenericAffine.mat")]
        >>> get_transform_type_from_paths(paths)
        'SyN'
    """
    names = [p.name for p in transform_paths]
    has_warp = any(any(kw in n for kw in _WARP_KEYWORDS) for n in names)
    has_rigid = any(any(kw in n for kw in _RIGID_KEYWORDS) for n in names)
    has_affine = any(any(kw in n for kw in _AFFINE_KEYWORDS) for n in names) or (
        any(n.endswith(".mat") for n in names) and not has_rigid
    )

    if has_warp:
        return "SyN"
    if has_rigid:
        return "Rigid"
    if has_affine:
        return "Affine"
    return "transformed"





def build_output_filename(
    input_path: Path,
    transform_type: str,
    direction: str,
    target_space: str,
) -> str:
    """Build a descriptive output filename that encodes transform provenance.

    The result follows the pattern
    ``{stem}_{transform_type}_{direction}_in_{target_space}_space.nii.gz``.

    Args:
        input_path: Original source image path; the stem is reused.
        transform_type: Short transform label, e.g. ``"SyN"`` or ``"Affine"``.
        direction: Transform direction, e.g. ``"template_to_patient"``.
        target_space: Short name of the target space, e.g. ``"patient"`` or
            ``"template"``.

    Returns:
        Filename string (no directory component) such as
        ``"mean_left_SyN_template_to_patient_in_patient_space.nii.gz"``.

    Example:
        >>> build_output_filename(
        ...     Path("cohort/atlas/mean_left.nii.gz"),
        ...     "SyN",
        ...     "template_to_patient",
        ...     "patient",
        ... )
        'mean_left_SyN_template_to_patient_in_patient_space.nii.gz'
    """
    stem = input_path.stem
    if stem.endswith(".nii"):
        stem = stem[:-4]
    return f"{stem}_{transform_type}_{direction}_in_{target_space}_space.nii.gz"





def validate_transform_inputs(
    input_files: List[Path],
    transform_paths: List[Path],
    reference_image: Path,
) -> List[str]:
    """Return a list of error messages for missing or invalid transform inputs.

    An empty list means all files are present and the workflow can proceed.

    Args:
        input_files: Source images that will be transformed.
        transform_paths: Warp fields / affine mats that will be applied.
        reference_image: Reference image for resampling.

    Returns:
        List of human-readable error strings.  Empty if everything is valid.

    Example:
        >>> errors = validate_transform_inputs(
        ...     [Path("mean.nii.gz")],
        ...     [Path("warp.nii.gz"), Path("affine.mat")],
        ...     Path("ref.nii.gz"),
        ... )
    """
    errors: List[str] = []
    for f in input_files:
        if not f.exists():
            errors.append(f"Input file not found: '{f}'")
    for t in transform_paths:
        if not t.exists():
            errors.append(f"Transform file not found: '{t}'")
    if not reference_image.exists():
        errors.append(f"Reference image not found: '{reference_image}'")
    return errors