slicer_bridge.py

"""Slicer integration for abcdmicro resources.

This module provides resources that wrap 3D Slicer MRML scene objects,
allowing seamless integration between abcdmicro's Resource abstraction
and Slicer's visualization and analysis capabilities.

SlicerVolumeResource is considered "loaded" (is_loaded=True) because the
data lives in memory (Slicer's VTK structures). This is analogous to
InMemoryVolumeResource, just with VTK as the backing storage instead of
numpy arrays.

Usage:
    This module should only be imported within 3D Slicer's Python environment.
    It requires the `slicer` module to be available.

Example (in Slicer Python console):
    >>> from abcdmicro.io import NiftiVolumeResource
    >>> from abcdmicro.slicer_bridge import SlicerVolumeResource
    >>>
    >>> # Load a NIfTI from disk and display in Slicer
    >>> nifti = NiftiVolumeResource("/path/to/volume.nii.gz")
    >>> slicer_vol = SlicerVolumeResource.from_resource(nifti, name="MyVolume")
    >>>
    >>> # SlicerVolumeResource is "loaded" - data is in Slicer's memory
    >>> slicer_vol.is_loaded  # True
    >>> slicer_vol.load() is slicer_vol  # True (no-op, already loaded)
    >>>
    >>> # Access data as numpy (converts from VTK on each call)
    >>> array = slicer_vol.get_array()
    >>> affine = slicer_vol.get_affine()
    >>>
    >>> # For repeated numpy access, convert once to InMemoryVolumeResource
    >>> in_mem = slicer_vol.to_numpy_resource()
    >>> in_mem.get_array()  # No conversion, data already in numpy
    >>>
    >>> # Wrap an existing volume from Slicer's scene
    >>> existing = SlicerVolumeResource.from_scene_by_name("MyVolume")
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar

import numpy as np
from numpy.typing import NDArray

from abcdmicro.resource import (
    InMemoryVolumeResource,
    VolumeResource,
)

if TYPE_CHECKING:
    import slicer
    from vtkmodules.vtkCommonDataModel import vtkImageData


def _check_slicer_available() -> None:
    """Raise ImportError if not running inside Slicer."""
    try:
        import slicer  # noqa: F401
    except ImportError as e:
        msg = (
            "The slicer module is not available. "
            "This module must be used within 3D Slicer's Python environment."
        )
        raise ImportError(msg) from e


def _numpy_to_vtk_image(array: NDArray[np.number]) -> vtkImageData:
    """Convert a numpy array to a vtkImageData object.

    In Slicer's convention, vtkImageData should have origin (0,0,0),
    spacing (1,1,1), and identity direction. The actual spatial transform
    is stored in the volume node's IJKToRAS matrix.

    Args:
        array: The volume data array (3D or 4D)

    Returns:
        vtkImageData with the volume data (identity transform)
    """
    import vtk
    from vtk.util import numpy_support

    # Ensure array is in Fortran order for VTK
    array = np.asfortranarray(array)

    # Create vtkImageData
    image_data = vtk.vtkImageData()

    # Set dimensions (VTK uses x, y, z order)
    if array.ndim == 3:
        image_data.SetDimensions(array.shape[0], array.shape[1], array.shape[2])
    elif array.ndim == 4:
        # 4D array: treat as 3D with multiple components
        image_data.SetDimensions(array.shape[0], array.shape[1], array.shape[2])
    else:
        msg = f"Array must be 3D or 4D, got {array.ndim}D"
        raise ValueError(msg)

    # Slicer convention: vtkImageData has identity transform
    # Actual spatial info goes in the volume node's IJKToRAS matrix
    image_data.SetSpacing(1.0, 1.0, 1.0)
    image_data.SetOrigin(0.0, 0.0, 0.0)

    # Convert numpy array to VTK array
    if array.ndim == 4:
        # Flatten the 4th dimension into components
        flat_array = array.reshape(-1, order="F")
        n_components = array.shape[3]
    else:
        flat_array = array.flatten(order="F")
        n_components = 1

    vtk_array = numpy_support.numpy_to_vtk(
        flat_array, deep=True, array_type=vtk.VTK_FLOAT
    )
    vtk_array.SetNumberOfComponents(n_components)

    image_data.GetPointData().SetScalars(vtk_array)

    return image_data


def _vtk_image_to_numpy(image_data: vtkImageData) -> NDArray[np.number]:
    """Convert vtkImageData to a numpy array.

    Args:
        image_data: The VTK image data

    Returns:
        Numpy array with the volume data
    """
    from vtk.util import numpy_support

    dims = image_data.GetDimensions()
    scalars = image_data.GetPointData().GetScalars()
    n_components = scalars.GetNumberOfComponents()

    # Convert to numpy
    array = numpy_support.vtk_to_numpy(scalars)

    if n_components > 1:
        # Reshape to 4D
        array = array.reshape(dims[0], dims[1], dims[2], n_components, order="F")
    else:
        # Reshape to 3D
        array = array.reshape(dims[0], dims[1], dims[2], order="F")

    return array


def _affine_to_ijk_to_ras_matrix(affine: NDArray[np.floating]) -> Any:
    """Convert a numpy affine to a Slicer IJKToRAS vtkMatrix4x4.

    Args:
        affine: 4x4 numpy affine matrix

    Returns:
        vtkMatrix4x4 for use with Slicer volume nodes
    """
    import vtk

    matrix = vtk.vtkMatrix4x4()
    for i in range(4):
        for j in range(4):
            matrix.SetElement(i, j, affine[i, j])
    return matrix


def _ijk_to_ras_matrix_to_affine(matrix: Any) -> NDArray[np.floating]:
    """Convert a Slicer IJKToRAS vtkMatrix4x4 to a numpy affine.

    Args:
        matrix: vtkMatrix4x4 from Slicer volume node

    Returns:
        4x4 numpy affine matrix
    """
    affine = np.eye(4)
    for i in range(4):
        for j in range(4):
            affine[i, j] = matrix.GetElement(i, j)
    return affine


@dataclass
class SlicerVolumeResource(VolumeResource):
    """A volume resource that lives in the 3D Slicer MRML scene.

    This resource wraps a vtkMRMLScalarVolumeNode (or vtkMRMLVectorVolumeNode
    for 4D data) in the Slicer scene. The volume is immediately visible in
    Slicer's viewers.

    SlicerVolumeResource is considered "loaded" (is_loaded = True) because
    the data is already in memory (VTK structures), just in a different format
    than numpy arrays. No disk I/O is needed to access the data.

    Attributes:
        node_id: The MRML node ID in the Slicer scene
        _node: Cached reference to the MRML node (not persisted)
    """

    is_loaded: ClassVar[bool] = True  # Data is in memory (VTK structures)

    node_id: str
    """The unique MRML node ID in the Slicer scene"""

    _node: Any = field(default=None, repr=False, compare=False)
    """Cached reference to the MRML node"""

    def __post_init__(self) -> None:
        _check_slicer_available()

    def _get_node(self) -> Any:
        """Get the MRML node, fetching from scene if needed."""
        if self._node is None:
            import slicer

            self._node = slicer.mrmlScene.GetNodeByID(self.node_id)
            if self._node is None:
                msg = f"MRML node with ID '{self.node_id}' not found in scene"
                raise ValueError(msg)
        return self._node

    def load(self) -> SlicerVolumeResource:
        """Return self since SlicerVolumeResource is already loaded.

        The data lives in Slicer's VTK structures which are in memory.
        Use to_numpy_resource() if you need an InMemoryVolumeResource.

        Returns:
            Self (no-op for already-loaded resources)
        """
        return self

    def to_numpy_resource(self) -> InMemoryVolumeResource:
        """Convert to an InMemoryVolumeResource.

        This copies the data from VTK structures into numpy arrays.
        Use this when you need to pass data to code that expects
        InMemoryVolumeResource specifically.

        Returns:
            InMemoryVolumeResource with copied volume data
        """
        import vtk

        node = self._get_node()

        # Get the image data
        image_data = node.GetImageData()
        array = _vtk_image_to_numpy(image_data)

        # Get the IJKToRAS matrix (affine)
        ijk_to_ras = vtk.vtkMatrix4x4()
        node.GetIJKToRASMatrix(ijk_to_ras)
        affine = _ijk_to_ras_matrix_to_affine(ijk_to_ras)

        # Gather metadata
        metadata: dict[str, Any] = {
            "slicer_node_name": node.GetName(),
            "slicer_node_id": self.node_id,
        }

        return InMemoryVolumeResource(array=array, affine=affine, metadata=metadata)

    def get_array(self) -> NDArray[np.number]:
        """Get the volume data array.

        This converts from VTK to numpy on each call. For repeated access,
        consider using to_numpy_resource() once and reusing the result.
        """
        node = self._get_node()
        image_data = node.GetImageData()
        return _vtk_image_to_numpy(image_data)

    def get_affine(self) -> NDArray[np.floating]:
        """Get the 4x4 affine matrix."""
        import vtk

        node = self._get_node()
        ijk_to_ras = vtk.vtkMatrix4x4()
        node.GetIJKToRASMatrix(ijk_to_ras)
        return _ijk_to_ras_matrix_to_affine(ijk_to_ras)

    def get_metadata(self) -> dict[str, Any]:
        """Get volume metadata."""
        node = self._get_node()
        return {
            "slicer_node_name": node.GetName(),
            "slicer_node_id": self.node_id,
        }

    def get_node(self) -> Any:
        """Get the underlying MRML node.

        Returns:
            The vtkMRMLScalarVolumeNode or vtkMRMLVectorVolumeNode
        """
        return self._get_node()

    @staticmethod
    def from_resource(
        vol: VolumeResource,
        name: str | None = None,
        show: bool = True,
    ) -> SlicerVolumeResource:
        """Create a Slicer volume node from any VolumeResource.

        This adds the volume to the Slicer scene and optionally displays it
        in the viewers.

        Args:
            vol: The source VolumeResource to convert
            name: Name for the Slicer node (default: "abcdmicro_volume")
            show: If True, display the volume in Slicer viewers

        Returns:
            SlicerVolumeResource wrapping the new MRML node
        """
        import slicer

        _check_slicer_available()

        if name is None:
            name = "abcdmicro_volume"

        # Load the resource if needed
        if not vol.is_loaded:
            vol = vol.load()

        array = vol.get_array()
        affine = vol.get_affine()

        # Determine node type based on array dimensions
        if array.ndim == 4:
            # Vector volume for 4D data
            node = slicer.mrmlScene.AddNewNodeByClass(
                "vtkMRMLVectorVolumeNode", name
            )
        else:
            # Scalar volume for 3D data
            node = slicer.mrmlScene.AddNewNodeByClass(
                "vtkMRMLScalarVolumeNode", name
            )

        # Set the IJKToRAS matrix
        ijk_to_ras = _affine_to_ijk_to_ras_matrix(affine)
        node.SetIJKToRASMatrix(ijk_to_ras)

        # Set the image data (affine is already in IJKToRAS matrix)
        image_data = _numpy_to_vtk_image(array)
        node.SetAndObserveImageData(image_data)

        # Optionally show in viewers
        if show:
            slicer.util.setSliceViewerLayers(background=node)

        return SlicerVolumeResource(node_id=node.GetID(), _node=node)

    @staticmethod
    def from_node(node: Any) -> SlicerVolumeResource:
        """Wrap an existing MRML volume node.

        Args:
            node: A vtkMRMLScalarVolumeNode or vtkMRMLVectorVolumeNode

        Returns:
            SlicerVolumeResource wrapping the node
        """
        _check_slicer_available()

        node_id = node.GetID()
        if node_id is None:
            msg = "Node must be added to the scene before wrapping"
            raise ValueError(msg)

        return SlicerVolumeResource(node_id=node_id, _node=node)

    @staticmethod
    def from_scene_by_name(name: str) -> SlicerVolumeResource:
        """Get a volume from the Slicer scene by name.

        Args:
            name: The name of the volume node in the scene

        Returns:
            SlicerVolumeResource wrapping the found node

        Raises:
            ValueError: If no node with the given name is found
        """
        import slicer

        _check_slicer_available()

        node = slicer.util.getNode(name)
        if node is None:
            msg = f"No node named '{name}' found in scene"
            raise ValueError(msg)

        return SlicerVolumeResource.from_node(node)