Backends

@runtime_checkable
class ImageDataBackend(Protocol):
    """Protocol for lazy image data access.

    Implementations wrap different storage formats (in-memory tensors, NIfTI
    files via nibabel, NIfTI-Zarr via dask) behind a uniform, lazy I/O
    interface. This is an I/O adapter layer, not a lazy computation framework:
    it speeds up metadata reads and region slicing, but does not defer
    arithmetic or transforms.

    Contract:
        - `shape` is always 4D `(C, I, J, K)`, even for 3D NIfTI (channel
          dimension 1).
        - `affine` is a `float64` `torch.Tensor` of shape `(4, 4)`.
        - `dtype` reports the on-disk (or in-memory) `numpy` dtype.
        - `__getitem__` accepts the same indexing as
          [`Image.__getitem__`][torchio.Image.__getitem__], always returns a
          4D `torch.Tensor` in `(C, I, J, K)` layout, and never drops axes
          (integer indices keep a size-1 dimension).
        - `to_tensor` materializes the full volume as a `torch.Tensor`
          preserving the on-disk dtype where PyTorch supports it.
    """

    @property
    def shape(self) -> TypeTensorShape:
        """Shape as (C, I, J, K)."""
        ...

    @property
    def affine(self) -> TypeAffineMatrix:
        """$4 \\times 4$ affine matrix as a float64 tensor."""
        ...

    @property
    def dtype(self) -> np.dtype:
        """Data type of the image on disk."""
        ...

    def __getitem__(self, slices: SliceIndex) -> Tensor:
        """Slice the data, returning a 4D `(C, I, J, K)` tensor.

        Integer indices keep a size-1 dimension (axes are never dropped).
        Tensor-backed images preserve their device and dtype; lazy backends
        read only the requested region and convert it to a tensor.
        """
        ...

    def to_tensor(self) -> Tensor:
        """Materialize the full data as a tensor preserving the on-disk dtype."""
        ...

@runtime_checkable
class LazyReader(Protocol):
    """A custom reader that can build a lazy backend.

    Readers passed to `Image` are normally simple callables returning
    `(tensor, affine)`, which always load the whole volume. A reader that also
    implements `create_backend` opts in to lazy access: `Image.shape`,
    `affine`, `dtype`, and slicing then go through the returned backend without
    materializing the full tensor.
    """

    def create_backend(self, request: BackendRequest) -> ImageDataBackend:
        """Build a lazy backend for *request*."""
        ...

@dataclass(frozen=True)
class BackendRequest:
    """Description of an image source used to resolve a lazy backend.

    A request decouples backend selection from the `Image` class: resolvers
    and custom backends receive a `BackendRequest` instead of an `Image`.

    Attributes:
        path: Resolved filesystem (or fsspec) path to the image, if any.
        remote_zarr_uri: Remote NIfTI-Zarr URI, if the source is remote.
        zarr_store: An open Zarr store object, if the source is a store.
        affine: Optional $4 \\times 4$ affine override to apply to the backend.
        reader_kwargs: Extra keyword arguments forwarded to the loader.
        reader: The reader configured on the `Image`. Custom readers that
            implement [`LazyReader`][torchio.data.backends.LazyReader] can
            build a lazy backend instead of loading the whole volume.
    """

    path: Path | None = None
    remote_zarr_uri: str | None = None
    zarr_store: Any = None
    affine: TypeAffineMatrix | None = None
    reader_kwargs: Mapping[str, Any] = field(default_factory=dict)
    reader: Any = None

class NibabelBackend:
    """Backend wrapping a nibabel image for lazy NIfTI access.

    Data is accessed through nibabel's `dataobj` proxy, which supports
    memory-mapped reads. Shape and affine are read from the header
    without loading data.

    This backend also works with NIfTI-Zarr files loaded via `niizarr`,
    since `zarr2nii` returns a standard `nibabel.Nifti1Image` whose
    `dataobj` is a dask array.

    Args:
        nii: A nibabel image (typically from `nib.load()` or `zarr2nii()`).
        affine: Optional $4 \\times 4$ affine that overrides the affine stored
            in the NIfTI header. Used when the user passes an explicit affine
            to `Image`, so that `image.affine` and `image.dataobj.affine` agree.
    """

    __slots__ = ("_affine_override", "_nii", "_shape")

    def __init__(
        self,
        nii: nib.spatialimages.SpatialImage,
        affine: TypeAffineMatrix | None = None,
    ) -> None:
        self._nii = nii
        self._affine_override = (
            torch.as_tensor(affine, dtype=torch.float64) if affine is not None else None
        )
        header_shape = nii.header.get_data_shape()
        ndim = len(header_shape)
        if ndim == 3:
            si, sj, sk = header_shape
            self._shape: TypeTensorShape = (1, int(si), int(sj), int(sk))
        elif ndim == 4:
            si, sj, sk, c = header_shape
            self._shape = (int(c), int(si), int(sj), int(sk))
        elif ndim == 5 and header_shape[3] == 1:
            # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C)
            si, sj, sk, _, c = header_shape
            self._shape = (int(c), int(si), int(sj), int(sk))
        else:
            msg = f"Expected 3D or 4D NIfTI, got {ndim}D"
            raise ValueError(msg)

    @property
    def shape(self) -> TypeTensorShape:
        return self._shape

    @property
    def affine(self) -> TypeAffineMatrix:
        if self._affine_override is not None:
            return self._affine_override
        return torch.as_tensor(
            self._nii.header.get_best_affine(),
            dtype=torch.float64,
        )

    @property
    def dtype(self) -> np.dtype:
        return self._nii.header.get_data_dtype()

    def __getitem__(self, slices: SliceIndex) -> Tensor:
        """Slice in (C, I, J, K) space, returning a 4D tensor.

        The (C, I, J, K) index is translated to the on-disk layout, which is
        (I, J, K) for 3D, (I, J, K, C) for 4D, and (I, J, K, 1, C) for 5D
        vector NIfTI. Integer indices keep their axis (size 1), so the result
        is always 4D.
        """
        sc, si, sj, sk = normalize_index(slices)
        ndim_on_disk = len(self._nii.header.get_data_shape())
        if ndim_on_disk == 3:
            # On disk (I, J, K); channel axis is synthetic (size 1).
            data = np.asarray(self._nii.dataobj[si, sj, sk])
            array = rearrange(data, "i j k -> 1 i j k")[sc]
        elif ndim_on_disk == 4:
            # On disk (I, J, K, C).
            data = np.asarray(self._nii.dataobj[si, sj, sk, sc])
            array = rearrange(data, "i j k c -> c i j k")
        elif ndim_on_disk == 5:
            # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C).
            data = np.asarray(self._nii.dataobj[si, sj, sk, :, sc])
            array = rearrange(data, "i j k 1 c -> c i j k")
        else:
            msg = f"Expected 3D, 4D, or 5D NIfTI, got {ndim_on_disk}D"
            raise ValueError(msg)
        from .io import _numpy_to_tensor

        array = np.ascontiguousarray(array)
        if not array.flags.writeable:
            # Proxy reads (e.g. memory-mapped NIfTI) can be read-only, which
            # PyTorch does not support; copy so the resulting tensor is safe
            # to mutate.
            array = array.copy()
        return _numpy_to_tensor(array)

    def to_tensor(self) -> Tensor:
        """Materialize the full image preserving the on-disk dtype."""
        data = np.asarray(self._nii.dataobj)
        ndim = data.ndim
        if ndim == 3:
            data = rearrange(data, "i j k -> 1 i j k")
        elif ndim == 4:
            data = rearrange(data, "i j k c -> c i j k")
        elif ndim == 5 and data.shape[3] == 1:
            # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C)
            data = rearrange(data, "i j k 1 c -> c i j k")
        else:
            msg = f"Expected 3D or 4D data, got {ndim}D"
            raise ValueError(msg)
        from .io import _numpy_to_tensor

        return _numpy_to_tensor(np.ascontiguousarray(data))

class ZarrBackend:
    """Backend wrapping a NIfTI-Zarr file for chunked lazy access.

    NIfTI-Zarr files are loaded via `niizarr.zarr2nii()`, which returns
    a nibabel image with a dask array as its `dataobj`. This backend
    delegates to `NibabelBackend` for the actual data access.

    Requires the `nifti-zarr` package.

    Args:
        path: Path to a `.nii.zarr` directory or a remote `.nii.zarr` URI.
        affine: Optional $4 \\times 4$ affine that overrides the affine stored
            in the NIfTI-Zarr metadata.
        **kwargs: Extra keyword arguments forwarded to `niizarr.zarr2nii()`.
    """

    __slots__ = ("_nibabel_backend",)

    def __init__(
        self,
        path: str | object,
        affine: TypeAffineMatrix | None = None,
        **kwargs: Any,
    ) -> None:
        from ..external.imports import get_niizarr

        niizarr = get_niizarr()
        nii = niizarr.zarr2nii(str(path), **kwargs)
        self._nibabel_backend = NibabelBackend(nii, affine=affine)

    @property
    def shape(self) -> TypeTensorShape:
        return self._nibabel_backend.shape

    @property
    def affine(self) -> TypeAffineMatrix:
        return self._nibabel_backend.affine

    @property
    def dtype(self) -> np.dtype:
        return self._nibabel_backend.dtype

    def __getitem__(self, slices: SliceIndex) -> Tensor:
        return self._nibabel_backend[slices]

    def to_tensor(self) -> Tensor:
        return self._nibabel_backend.to_tensor()

class TensorBackend:
    """Backend wrapping an in-memory PyTorch tensor.

    Used for images created from tensors or NumPy arrays (NumPy arrays are
    converted to tensors first).

    Args:
        data: 4D tensor with shape (C, I, J, K).
        affine: $4 \\times 4$ affine tensor. Identity if not given.
    """

    __slots__ = ("_affine", "_data")

    def __init__(
        self,
        data: Tensor,
        affine: TypeAffineMatrix | None = None,
    ) -> None:
        self._data = data
        if affine is not None:
            self._affine = affine
        else:
            self._affine = torch.eye(4, dtype=torch.float64)

    @property
    def shape(self) -> TypeTensorShape:
        s = self._data.shape
        return (int(s[0]), int(s[1]), int(s[2]), int(s[3]))

    @property
    def affine(self) -> TypeAffineMatrix:
        return self._affine

    @property
    def dtype(self) -> np.dtype:
        # Map torch dtype to numpy for protocol compatibility
        return torch.empty(0, dtype=self._data.dtype).numpy().dtype

    def __getitem__(self, slices: SliceIndex) -> Tensor:
        """Slice the tensor, preserving device, dtype, and 4D layout."""
        return self._data[normalize_index(slices)]

    def to_tensor(self) -> Tensor:
        return self._data.clone()