Source code for thesis.workflows.mrtrix3.nodes.preprocessing

"""DWI import and brain-mask preparation nodes."""

from pathlib import Path

from nipype import Node
from nipype.interfaces.mrtrix3 import MRConvert
from nipype.interfaces.utility import Function

from ..operations import generate_dwi2mask_task, prepare_brain_mask_task




def prepare_dwi_import(
    dwi_image: Path,
    bvec: Path,
    bval: Path,
    out_dir: Path,
    name: str = "dwi_import",
) -> Node:
    """Convert NIfTI DWI + FSL gradient table to ``dwi.mif``.

    The resulting ``.mif`` bundles the gradient table so downstream
    interfaces (``ResponseSD``, ``EstimateFOD``) do not need to be
    repeatedly told about ``bvecs``/``bvals``.

    Args:
        dwi_image: Path to the DWI NIfTI (typically ``data.nii.gz``).
        bvec: Path to the FSL bvecs file.
        bval: Path to the FSL bvals file.
        out_dir: Output directory; ``dwi.mif`` is written here.
        name: Nipype node name.

    Returns:
        Configured Nipype Node.
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(MRConvert(), name=name)
    # ``MRConvert.in_file`` carries an ``exists=True`` trait, which would fail
    # build-time validation if the DWI hasn't been generated yet (e.g. when
    # this sub-workflow is embedded in ``full_pipeline`` (mrtrix3 backend) and consumes
    # preprocess output that doesn't exist at build time). When the static
    # paths are absent, leave the inputs Undefined; the parent meta-workflow
    # is expected to wire ``in_file`` and ``grad_fsl`` at runtime.
    if Path(dwi_image).exists():
        node.inputs.in_file = str(dwi_image)
        node.inputs.grad_fsl = (str(bvec), str(bval))
    node.inputs.out_file = str(out_dir / "dwi.mif")
    # Outputs land in the patient's shared dir, not the Nipype work dir, so
    # restarts must overwrite stale artefacts from a previous failed run.
    node.inputs.args = "-force"
    return node





def prepare_mask_node(
    mask_path: Path,
    out_dir: Path,
    dilate_voxels: int = 1,
    name: str = "mask_prep",
) -> Node:
    """Import the FSL brain mask and dilate it for MRtrix consumption.

    Wrapped as a Function node (rather than chaining ``MRConvert`` and
    ``MaskFilter``) so the dilation pass count is preserved as a single
    visible parameter and the output filename is deterministic.

    Args:
        mask_path: Source NIfTI mask (typically ``nodif_brain_mask.nii.gz``).
        out_dir: Output directory where ``mask.mif`` is written.
        dilate_voxels: Number of dilation passes (0 disables dilation).
        name: Nipype node name.

    Returns:
        Configured Nipype Node with ``out_file`` output.
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(
        Function(
            input_names=["mask_path", "output_dir", "dilate_voxels"],
            output_names=["out_file"],
            function=prepare_brain_mask_task,
        ),
        name=name,
    )
    node.inputs.mask_path = str(mask_path)
    node.inputs.output_dir = str(out_dir)
    node.inputs.dilate_voxels = int(dilate_voxels)
    return node





def prepare_dwi2mask_node(
    out_dir: Path,
    dilate_voxels: int = 1,
    algorithm: str = "legacy",
    name: str = "mask_prep",
) -> Node:
    """Generate the brain mask from the DWI via ``dwi2mask`` (mask_source='dwi2mask').

    Drop-in alternative to :func:`prepare_mask_node`: produces the same
    ``mask.mif`` output field, but derives the mask from the imported DWI
    ``.mif`` instead of importing the FSL ``nodif_brain_mask``. The workflow
    wires ``dwi_file`` from the ``dwi_import`` node's ``out_file`` at build time.

    Args:
        out_dir: Output directory where ``mask.mif`` is written.
        dilate_voxels: Number of dilation passes (0 disables dilation).
        algorithm: ``dwi2mask`` algorithm selector (default ``"legacy"``).
        name: Nipype node name (default ``"mask_prep"`` so it is a true
            drop-in for :func:`prepare_mask_node`).

    Returns:
        Configured Nipype Node with ``out_file`` output and an unconnected
        ``dwi_file`` input (wired by the workflow).
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(
        Function(
            input_names=["dwi_file", "output_dir", "dilate_voxels", "algorithm"],
            output_names=["out_file"],
            function=generate_dwi2mask_task,
        ),
        name=name,
    )
    node.inputs.output_dir = str(out_dir)
    node.inputs.dilate_voxels = int(dilate_voxels)
    node.inputs.algorithm = str(algorithm)
    return node