Source code for thesis.workflows.mrtrix3.operations.mask_prep

"""Brain-mask preparation for the MRtrix3 workflow.

Imports the FSL ``nodif_brain_mask.nii.gz`` to MRtrix's native ``.mif``
format and optionally dilates it. The dilation step prevents the white-
matter gaps that ``dwi2mask legacy`` produces in the cerebral peduncle and
other deep WM (a verified failure mode that killed tractography during
the manual MRtrix3 validation run).

Used as the body of a Nipype Function node so the workflow graph captures
the dilation parameter and output path explicitly.
"""

# NOTE: Nipype Function nodes execute in subprocesses and cannot pickle the
# project's loguru logger. Use ``print`` for diagnostics inside this body.




def prepare_brain_mask_task(
    mask_path: str,
    output_dir: str,
    dilate_voxels: int = 1,
) -> str:
    """Import the FSL brain mask into ``mask.mif`` with optional dilation.

    Args:
        mask_path: Absolute path to the source NIfTI brain mask.
        output_dir: Directory where ``mask.mif`` is written.
        dilate_voxels: Number of dilation passes (0 disables dilation).

    Returns:
        Absolute path to the produced ``.mif`` mask.
    """
    import shutil
    import subprocess
    from pathlib import Path

    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    src = Path(mask_path)
    if not src.exists():
        raise FileNotFoundError(f"Brain mask not found: {src}")

    converted = out_dir / "mask_raw.mif"
    final = out_dir / "mask.mif"

    if shutil.which("mrconvert") is None:
        raise RuntimeError("mrconvert not found on PATH; ensure MRtrix3 is installed.")

    convert_cmd = ["mrconvert", str(src), str(converted), "-force"]
    print(f"[mrtrix3] {' '.join(convert_cmd)}")
    subprocess.run(convert_cmd, check=True)

    if dilate_voxels and dilate_voxels > 0:
        if shutil.which("maskfilter") is None:
            raise RuntimeError("maskfilter not found on PATH; ensure MRtrix3 is installed.")
        dilate_cmd = [
            "maskfilter",
            str(converted),
            "dilate",
            str(final),
            "-npass",
            str(int(dilate_voxels)),
            "-force",
        ]
        print(f"[mrtrix3] {' '.join(dilate_cmd)}")
        subprocess.run(dilate_cmd, check=True)
        try:
            converted.unlink()
        except OSError:
            pass
    else:
        converted.replace(final)

    return str(final)





def generate_dwi2mask_task(
    dwi_file: str,
    output_dir: str,
    dilate_voxels: int = 1,
    algorithm: str = "legacy",
) -> str:
    """Generate a brain mask from the DWI via MRtrix3 ``dwi2mask``.

    Alternative to :func:`prepare_brain_mask_task`: instead of importing and
    dilating the FSL ``nodif_brain_mask``, this derives the mask directly from
    the (already ``.mif``) DWI series with ``dwi2mask``. An optional dilation
    pass widens the mask the same way :func:`prepare_brain_mask_task` does, so
    downstream ACT does not lose deep white matter.

    Args:
        dwi_file: Absolute path to the DWI ``.mif`` (carries its gradient table).
        output_dir: Directory where ``mask.mif`` is written.
        dilate_voxels: Number of dilation passes (0 disables dilation).
        algorithm: ``dwi2mask`` algorithm selector (default ``"legacy"``).

    Returns:
        Absolute path to the produced ``.mif`` mask.
    """
    import shutil
    import subprocess
    from pathlib import Path

    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    src = Path(dwi_file)
    if not src.exists():
        raise FileNotFoundError(f"DWI not found: {src}")

    raw = out_dir / "mask_raw.mif"
    final = out_dir / "mask.mif"

    if shutil.which("dwi2mask") is None:
        raise RuntimeError("dwi2mask not found on PATH; ensure MRtrix3 is installed.")

    mask_cmd = ["dwi2mask", str(algorithm), str(src), str(raw), "-force"]
    print(f"[mrtrix3] {' '.join(mask_cmd)}")
    subprocess.run(mask_cmd, check=True)

    if dilate_voxels and dilate_voxels > 0:
        if shutil.which("maskfilter") is None:
            raise RuntimeError("maskfilter not found on PATH; ensure MRtrix3 is installed.")
        dilate_cmd = [
            "maskfilter",
            str(raw),
            "dilate",
            str(final),
            "-npass",
            str(int(dilate_voxels)),
            "-force",
        ]
        print(f"[mrtrix3] {' '.join(dilate_cmd)}")
        subprocess.run(dilate_cmd, check=True)
        try:
            raw.unlink()
        except OSError:
            pass
    else:
        raw.replace(final)

    return str(final)