%pythoncode %{

import numpy as np
class NDArrayITKBase(np.ndarray):
  """A numpy array that provides a view on the data associated with an optional itk "base" object."""

  def __new__(cls, input_array, itk_base=None):
      obj = np.asarray(input_array).view(cls)
      obj.itk_base = itk_base
      return obj

  def __array_finalize__(self, obj):
      if obj is None: return
      self.itk_base = getattr(obj, 'itk_base', None)

try:
    from distributed.protocol import dask_serialize, dask_deserialize
    from typing import Dict, List, Tuple
except (ImportError, RuntimeError):
    pass
else:
    @dask_serialize.register(NDArrayITKBase)
    def serialize(ndarray_itk_base: NDArrayITKBase) -> Tuple[Dict, List[bytes]]:
        dumps = dask_serialize.dispatch(np.ndarray)
        return dumps(ndarray_itk_base)

    @dask_deserialize.register(NDArrayITKBase)
    def deserialize(header: Dict, frames: List[bytes]) -> NDArrayITKBase:
        loads = dask_deserialize.dispatch(np.ndarray)
        return NDArrayITKBase(loads(header, frames))

def _get_numpy_pixelid(itk_Image_type):
    """Returns a ITK PixelID given a numpy array."""

    # This is a Mapping from numpy array types to itk pixel types.
    _np_itk = {"UC":np.uint8,
               "US":np.uint16,
               "UI":np.uint32,
               "UL":np.uint64,
               "ULL":np.uint64,
               "SC":np.int8,
               "SS":np.int16,
               "SI":np.int32,
               "SL":np.int64,
               "SLL":np.int64,
               "F":np.float32,
               "D":np.float64,
               "PF2":np.float32,
               "PF3":np.float32,
                }
    import os
    if os.name == 'nt':
        _np_itk['UL'] = np.uint32
        _np_itk['SL'] = np.int32
    try:
        return _np_itk[itk_Image_type]
    except KeyError as e:
        raise e
%}