Source code for thesis.workflows.tract_synthseg.validation

"""Validation sub-workflow for tract_synthseg.

Wraps ROI mask validation as a two-node Nipype workflow:
  roi_collector → validate_rois

Only used when ``config.validation.check_rois`` is True. Backend-agnostic:
both the ProbTrackX2 and MRtrix3 tractography sub-workflows re-publish their
final ROI sources under the same outputnode contract, so this validation
workflow is shared.
"""

from nipype import Node, Workflow
from nipype.interfaces.utility import Function

from thesis.core.config import PipelineConfig
from thesis.core.context import ProcessingContext
from thesis.core.logging import get_logger
from thesis.workflows.hcp.operations.validation import validate_warped_rois_task

logger = get_logger(__name__)


def _collect_roi_paths(seed: str, stop_mask: str, avoid_mask: str, target_mask: str) -> list[str]:
    """Collect individual ROI path strings into a flat list, skipping empty entries.

    Args:
        seed: Path to seed mask NIfTI, or empty string if absent.
        stop_mask: Path to stop mask NIfTI, or empty string if absent.
        avoid_mask: Path to avoid mask NIfTI, or empty string if absent.
        target_mask: Path to target mask NIfTI, or empty string if absent.

    Returns:
        List of non-empty path strings.
    """
    return [p for p in [seed, stop_mask, avoid_mask, target_mask] if p]




def build_validation_workflow(
    config: PipelineConfig,
    context: ProcessingContext,
    brain_mask_path: str,
) -> Workflow:
    """Build the ROI validation sub-workflow.

    Contains two nodes:

    * ``roi_collector`` — gathers individual ROI path fields into a ``List[str]``.
    * ``validate_rois`` — wraps :func:`validate_warped_rois_task`.

    External inputs (connected from the parent meta-workflow):

    * ``roi_collector.seed``
    * ``roi_collector.stop_mask``
    * ``roi_collector.avoid_mask``
    * ``roi_collector.target_mask``

    Args:
        config: PipelineConfig.  Must have a ``validation`` section with
            ``min_voxels``.
        context: ProcessingContext with ``patient_id``.
        brain_mask_path: Absolute path to the subject-space brain mask used
            for centroid plausibility check inside ``validate_warped_rois_task``.

    Returns:
        Nipype Workflow named ``validation_{patient_id}``.
    """
    pid = context.patient_id
    wf = Workflow(name=f"validation_{pid}")

    # --- roi_collector: gather individual paths into a list ---
    roi_collector = Node(
        Function(
            input_names=["seed", "stop_mask", "avoid_mask", "target_mask"],
            output_names=["roi_paths"],
            function=_collect_roi_paths,
        ),
        name="roi_collector",
    )
    # Default all inputs to empty string so absent masks are silently skipped.
    roi_collector.inputs.seed = ""
    roi_collector.inputs.stop_mask = ""
    roi_collector.inputs.avoid_mask = ""
    roi_collector.inputs.target_mask = ""

    # --- validate_rois: check voxel count, centroid, volume ---
    validate_rois = Node(
        Function(
            input_names=[
                "roi_paths",
                "reference_image",
                "min_voxels",
                "singularity_threshold",
                "volume_ratio_min",
                "volume_ratio_max",
            ],
            output_names=["roi_paths"],
            function=validate_warped_rois_task,
        ),
        name="validate_rois",
    )
    validate_rois.inputs.reference_image = brain_mask_path
    validate_rois.inputs.min_voxels = config.validation.min_voxels
    validate_rois.inputs.singularity_threshold = config.validation.singularity_threshold
    validate_rois.inputs.volume_ratio_min = config.validation.volume_ratio_min
    validate_rois.inputs.volume_ratio_max = config.validation.volume_ratio_max

    wf.connect(roi_collector, "roi_paths", validate_rois, "roi_paths")

    logger.info(
        "Built validation workflow | patient={} | brain_mask={} | min_voxels={}",
        pid,
        brain_mask_path,
        config.validation.min_voxels,
    )
    return wf