text: import logging
  from typing import Union, Tuple
  
  import matplotlib as mpl
  import matplotlib.pyplot as plt
  import numpy as np
  import numpy.ma as ma
  import pandas as pd
  import SimpleITK as sITK
  import xarray as xr
  from numpy.polynomial.polynomial import Polynomial
  from scipy import ndimage as ndi
  from scipy.stats import norm, zscore
  from skimage import io, measure, transform
  from skimage.measure import label, regionprops
  from skimage.transform import AffineTransform, warp
  
  
  def measure_under_labels(
      imgs: xr.DataArray,
      masks: xr.DataArray,
      ref_wvl: str = "410",
      ratio_numerator="410",
      ratio_denominator="470",
  ):
      """Measure the intensities of each channel under the label image"""
      df = []
      imgs = imgs.where(imgs.wavelength != "TL", drop=True)
  
      for a in imgs.animal.values:
          for tp in imgs.timepoint.values:
              for p in imgs.pair.values:
                  for wvl in imgs.wavelength.values:
                      img_selector = dict(animal=a, timepoint=tp, pair=p, wavelength=wvl)
  
                      if "wavelength" in masks.dims:
                          seg_frame = masks.sel(
                              animal=a, timepoint=tp, pair=p, wavelength=ref_wvl
                          )
                      else:
                          # single wavelength was passed
                          seg_frame = masks.sel(animal=a, timepoint=tp, pair=p)
  
                      labels = measure.label(seg_frame)
  
                      sub_df = pd.DataFrame(
                          measure.regionprops_table(
                              labels,
                              intensity_image=imgs.sel(**img_selector).values,
                              properties=["label", "mean_intensity", "area"],
                          )
                      )
  
                      sub_df["animal"] = a
                      sub_df["timepoint"] = tp
                      sub_df["pair"] = p
                      sub_df["wavelength"] = wvl
                      sub_df["strain"] = imgs.sel(**img_selector).strain.values
  
                      df.append(sub_df)
  
      df = pd.concat(df)
      df = df.set_index(["animal", "timepoint", "pair", "wavelength", "label"]).unstack(
          "wavelength"
      )
      df[("mean_intensity", "r")] = (
          df["mean_intensity"][ratio_numerator] / df["mean_intensity"][ratio_denominator]
      )
      df[("area", "r")] = df[("area", ratio_numerator)]
      df[("strain", "r")] = df[("strain", ratio_numerator)]
  
      df = df.stack("wavelength")
  
      return df
  
  
  def subtract_medians(
      data: xr.DataArray, image_data: xr.DataArray = None
  ) -> xr.DataArray:
      """
      Subtract medians from data, optionally calculating them from a separate piece of
      data.
  
      Parameters
      ----------
      data
          the data to subtract the median from.
      image_data
          the data to calculate the median with. Must include the dimensions `x` and `y`.
          If specified, all other dimensions must be identical to those in `data`. If
          not specified, the medians will be calculated with `data`.
      """
  
      if image_data is None:
          image_data = data
  
      submed = data.copy()
      submed.values = np.maximum(data - image_data.median(dim=["x", "y"]), 0)
      return submed
  
  
  def get_lr_bounds(
      rot_seg_stack: xr.DataArray, pad: int = 0, ref_wvl: str = "410", ref_pair: int = 0
  ) -> np.ndarray:
      """
      Get the left and right boundaries of the rotated pharynxes
  
      Parameters
      ----------
      rot_seg_stack
          the rotated segmented pharynxes
      pad
          the amount of padding on the left/right of the  bounds
      ref_wvl
          the wavelength to use for calculating bounds
      ref_pair
          the pair to use for calculating bounds
  
      Returns
      -------
      bounds
          An (m, 2) array where m = number of animals, the first column is the left bound
          and the second column is the right bound
      """
      imgs = rot_seg_stack.sel(wavelength=ref_wvl, pair=ref_pair)
      bounds = np.zeros((imgs.animal.size, 2))  # (animal, (l, r))
      for i, img in enumerate(imgs):
          _, l, _, r = measure.regionprops(measure.label(img))[0].bbox
          bounds[i, :] = [l - pad, r + pad - 1]
      return bounds.astype(np.int)
  
  
  def center_and_rotate_pharynxes(
      fl_images: xr.DataArray, seg_images: xr.DataArray
  ) -> Tuple[xr.DataArray, xr.DataArray]:
      """
      Given a fluorescence stack and a pharyngeal mask stack, center and rotate each frame
      of both the FL and mask such that the pharynx is in the center of the image, with
      its anterior on the left.
  
      Parameters
      ----------
      fl_images
          The fluorescence images to rotate and align
      seg_images
          The segmented images to rotate and align
  
      Returns
      -------
      (rotated_fl_stack, rotated_seg_stack)
          A 2-tuple where the first item is the rotated fluorescence stack and the second
          is the rotated mask stack
      """
      img_center_y, img_center_x = (
          fl_images.y.size // 2,
          fl_images.x.size // 2,
      )
  
      fl_rotated_stack = fl_images.copy()
      seg_rotated_stack = seg_images.copy()
  
      # STACK_ITERATION
      for img_idx in range(fl_images.animal.size):
          for wvl in fl_images.wavelength.data:
              for pair in fl_images.pair.data:
                  for tp in fl_images.timepoint.values:
                      # Optimization potential here...
                      # this recalculates all region properties for the reference each time
                      img = fl_images.isel(animal=img_idx).sel(
                          wavelength=wvl, pair=pair, timepoint=tp
                      )
  
                      try:
                          # Old data, had wavelength attached
                          ref_seg = seg_images.isel(animal=img_idx, wavelength=0).sel(
                              pair=pair, timepoint=tp
                          )
                      except ValueError:
                          ref_seg = seg_images.isel(animal=img_idx).sel(
                              pair=pair, timepoint=tp
                          )
  
                      try:
                          props = measure.regionprops(measure.label(ref_seg))[0]
                      except IndexError:
                          raise ValueError(
                              f"No binary objects found in image @ [idx={img_idx} ; wvl={wvl} ; pair={pair}]"
                          )
  
                      # pharynx_center_y, pharynx_center_x = props.centroid
                      pharynx_center_y, pharynx_center_x = np.mean(
                          np.nonzero(ref_seg), axis=1
                      )
                      pharynx_orientation = props.orientation
  
                      translation_matrix = transform.EuclideanTransform(
                          translation=(
                              -(img_center_x - pharynx_center_x),
                              -(img_center_y - pharynx_center_y),
                          )
                      )
  
                      rotated_img = rotate(
                          img.data, translation_matrix, pharynx_orientation
                      )
                      rotated_seg = rotate(
                          ref_seg.data,
                          translation_matrix,
                          pharynx_orientation,
                          order=0,
                          preserve_range=True,
                      )
  
                      fl_rotated_stack.loc[dict(wavelength=wvl, pair=pair, timepoint=tp)][
                          img_idx
                      ] = rotated_img
  
                      seg_rotated_stack.loc[dict(pair=pair, timepoint=tp)][
                          img_idx
                      ] = rotated_seg
  
      fl_rotated_stack.values = fl_rotated_stack.values.astype(fl_images.dtype)
      return fl_rotated_stack, seg_rotated_stack
  
  
  def extract_largest_binary_object(
      bin_img: Union[xr.DataArray, np.ndarray]
  ) -> Union[xr.DataArray, np.ndarray]:
      """
      Extracts the largest binary object from the given binary image
  
      Parameters
      ----------
      bin_img
          The binary image to process
  
      Returns
      -------
      bin_img
          The binary image containing only the largest binary object from the input
  
      """
      labels = measure.label(bin_img)
      if labels.max() == 0:
          # If there are no objects in the image... simply return the image
          return bin_img
  
      return labels == np.argmax(np.bincount(labels.flat)[1:]) + 1
  
  
  def get_area_of_largest_object(mask: np.ndarray) -> int:
      """Returns the area (px) of the largest object in a binary image
  
      Parameters
      ----------
      mask : np.ndarray
          the binary image
  
      Returns
      -------
      int
          the area of the largest object
      """
      try:
          return measure.regionprops(measure.label(mask))[0].area
      except IndexError:
          return 0
  
  
  def segment_pharynx(
      fl_img: xr.DataArray, target_area: int = 450, area_range: int = 100
  ) -> xr.DataArray:
      """Generate a mask for the given image containing a pharynx.
  
      Parameters
      ----------
      fl_img : xr.DataArray
          a fluorescent image containing a single pharynx
      target_area : int, optional
          the presumptive area (in px) of a pharynx, by default 450
      area_range : int, optional
          the acceptable range (in px) above/below the target_area, by default 100
  
      Returns
      -------
      xr.DataArray
          an image containing the segmented pharynx (dtype: np.uint8). Pixels of value=1
          indicate the pharynx, pixels of value=0 indicate the background.
      """
      # target_area = 450  # experimentally derived
      # area_range = 100
      min_area = target_area - area_range
      max_area = target_area + area_range
  
      max_iter = 300
  
      p = 0.15
      t = fl_img.max() * p
      mask = fl_img > t
  
      area = get_area_of_largest_object(mask)
  
      i = 0
      while (min_area > area) or (area > max_area):
          if i >= max_iter:
              return mask
          area = get_area_of_largest_object(mask)
  
          logging.debug(f"Setting p={p}")
          if area > max_area:
              p = p + 0.01
          if area < min_area:
              p = p - 0.01
          i = i + 1
  
          t = fl_img.max() * p
          mask = fl_img > t
  
          if p < 0:
              # break out if loop gets stuck w/ sensible default
              logging.warning("Caught infinite loop")
              return fl_img > (fl_img.max() * 0.15)
          if p > 0.9:
              logging.warning("Caught infinite loop")
              return fl_img > (fl_img.max() * 0.15)
  
      mask = extract_largest_binary_object(mask).astype(np.uint8)
  
      return mask
  
  
  def segment_pharynxes(
      fl_stack: xr.DataArray,
      wvl: str = "410",
      target_area: int = 450,
      area_range: int = 100,
  ) -> xr.DataArray:
      """Segment a hyperstack of pharynxes
  
      Parameters
      ----------
      fl_stack : xr.DataArray
          the fluorescent images to segment
      wvl : str, optional
          the wavelength to segment, by default "410"
      target_area : int, optional
          the presumptive area of a pharynx, in pixels, by default 450
      area_range : int, optional
          the acceptable range of pharyngeal areas, by default 100
  
      Returns
      -------
      xr.DataArray
          the masks for the specified wavelength
      """
  
      to_segment = fl_stack.sel(wavelength=wvl)
      seg = xr.apply_ufunc(
          segment_pharynx,
          to_segment,
          input_core_dims=[["y", "x"]],
          output_core_dims=[["y", "x"]],
          vectorize=True,
          kwargs={"target_area": target_area, "area_range": area_range},
      )
      return seg
  
  
  def rotate(
      img: Union[np.ndarray, xr.DataArray],
      tform,
      orientation,
      order=1,
      preserve_range=True,
  ):
      """
      Rotate the given image with the given translation matrix and orientation angle
  
      Parameters
      ----------
      img
          the image to rotate
      tform
          the translation matrix to apply
      orientation
          the angle of orientation (radians)
      order
          the order of the interpolation
      preserve_range
          preserve the input data range
  
      Returns
      -------
      rotated
          the translated and rotated image
  
      """
  
      return transform.rotate(
          transform.warp(
              img, tform, preserve_range=preserve_range, mode="wrap", order=order
          ),
          np.degrees(np.pi / 2 - orientation),
          mode="edge",
          order=order,
      )
  
  
  def calculate_midlines(rot_seg_stack: xr.DataArray, degree: int = 4) -> xr.DataArray:
      """
      Calculate a midline for each animal in the given stack
  
      Parameters
      ----------
      rot_seg_stack
          The rotated mask with which midlines should be calculated.
  
      degree
          The degree of the polynomial fit
  
      Returns
      -------
      midlines
          a DataArray containing the midline objects
  
      See Also
      --------
      calculate_midline
      """
      return xr.apply_ufunc(
          calculate_midline,
          rot_seg_stack,
          input_core_dims=[["y", "x"]],
          vectorize=True,
          keep_attrs=True,
          kwargs={"degree": degree},
      )
  
  
  def calculate_midline(
      rot_seg_img: Union[np.ndarray, xr.DataArray], degree: int = 4, pad: int = 10
  ) -> Polynomial:
      """
      Calculate a the midline for a single image by fitting a polynomial to the segmented
      pharynx
  
      Parameters
      ----------
      rot_seg_img: Union[np.ndarray, xr.DataArray]
          The rotated masked pharynx image
      degree
          the degree of the polynomial
      pad
          the number of pixels to "pad" the domain of the midline with respect to the
          boundaries of the segmentation mask
  
      Returns
      -------
      Polynomial
          the estimated midline
  
      Notes
      -----
      Right now this only works for images that this have been centered and aligned with their
      anterior-posterior along the horizontal.
      """
      import warnings
  
      with warnings.catch_warnings():
          warnings.simplefilter("ignore")
          try:
              rp = measure.regionprops(measure.label(rot_seg_img))[0]
              xs, ys = rp.coords[:, 1], rp.coords[:, 0]
              left_bound, _, _, right_bound = rp.bbox
  
              return Polynomial.fit(
                  xs, ys, degree, domain=[left_bound - pad, right_bound + pad]
              )
          except IndexError:
              # Indicates trying to measure on TL for example
              return None
  
  
  def measure_under_midline(
      fl: xr.DataArray,
      mid: Polynomial,
      n_points: int = 100,
      thickness: float = 0.0,
      order=1,
      norm_scale=1,
      flatten=True,
  ) -> np.ndarray:
      """
      Measure the intensity profile of the given image under the given midline at the given x-coordinates.
  
      Parameters
      ----------
      flatten
      norm_scale
      order
          the interpolation order
      fl
          The fluorescence image to measure
      mid
          The midline under which to measure
      n_points
          The number of points to measure under
      thickness
          The thickness of the line to measure under.
  
      Notes
      -----
      Using thickness is slower, depending on the amount of thickness
  
      On my machine (2GHz Intel Core i5), as of 12/4/19:
          0-thickness:
              492 µs ± 16.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
          2-thickness:
              1.99 ms ± 65.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
          10-thickness:
              3.89 ms ± 92.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  
      Returns
      -------
      zs: np.ndarray
          The intensity profile of the image measured under the midline at the given
          x-coordinates.
  
      """
      # Make sure the image orientation matches with the expected order of map_coordinates
      try:
          if thickness == 0:
              xs, ys = mid.linspace(n=n_points)
              fl = np.asarray(fl)
              return ndi.map_coordinates(fl, np.stack([xs, ys]), order=order)
          else:
              # Gets a bit wonky, but makes sense
  
              # We need to get the normal lines from each point in the midline
              # then measure under those lines.
  
              # First, get the coordinates of the midline
              xs, ys = mid.linspace(n=n_points)
  
              # Now, we get the angles of each normal vector
              der = mid.deriv()
              normal_slopes = -1 / der(xs)
              normal_thetas = np.arctan(normal_slopes)
  
              # We get the x and y components of the start/end of the normal vectors
              mag = thickness / 2
              x0 = np.cos(normal_thetas) * mag
              y0 = np.sin(normal_thetas) * mag
  
              x1 = np.cos(normal_thetas) * -mag
              y1 = np.sin(normal_thetas) * -mag
  
              # These are the actual coordinates of the starts/ends of the normal vectors as they move
              # from (x,y) coordinates in the midline
              xs0 = xs + x0
              xs1 = xs + x1
              ys0 = ys + y0
              ys1 = ys + y1
  
              # We need to measure in a consistent direction along the normal line
              # if y0 < y1, we're going to be measuring in an opposite direction along the line... so we need flip the coordinates
              for y0, y1, x0, x1, i in zip(ys0, ys1, xs0, xs1, range(len(xs0))):
                  if y0 < y1:
                      tx = xs0[i]
                      xs0[i] = xs1[i]
                      xs1[i] = tx
  
                      ty = ys0[i]
                      ys0[i] = ys1[i]
                      ys1[i] = ty
  
              n_line_pts = thickness
  
              all_xs = np.linspace(xs0, xs1, n_line_pts)
              all_ys = np.linspace(ys0, ys1, n_line_pts)
  
              straightened = ndi.map_coordinates(fl, [all_xs, all_ys], order=order)
  
              if flatten:
                  # Create a normal distribution centered around 0 with the given scale (see scipy.norm.pdf)
                  # the distribution is then tiled to be the same shape as the straightened pharynx
                  # then, this resultant matrix is the weights for averaging
                  w = np.tile(
                      norm.pdf(np.linspace(-1, 1, n_line_pts), scale=norm_scale),
                      (n_points, 1),
                  ).T
                  profile = np.average(straightened, axis=0, weights=w)
  
                  return profile
              else:
                  return straightened
      except AttributeError:
          # This happens if the image is TL. Then it will have `None` instead of
          # a midline object
          pass
      except Exception as e:
          # Here, something actually went wrong
          logging.warning(f"measuring under midline failed with error {e}")
  
      return np.zeros((1, n_points))
  
  
  def measure_under_midlines(
      fl_stack: xr.DataArray,
      midlines: xr.DataArray,
      n_points: int = 300,
      order=1,
      thickness=0,
  ) -> xr.DataArray:
      """
      Measure under all midlines in stack
  
      Parameters
      ----------
      order
      fl_stack
          The fluorescence stack under which to measure
      midlines: dict
          A DataArray containing the midlines
      n_points: int
          the number of points to sample under the midline
      thickness: float
          the thickness of the midline to measure under
  
      Returns
      -------
      profile_data: xr.DataArray
          the intensity profiles for each image in the stack
      """
      measurements = xr.apply_ufunc(
          measure_under_midline,
          fl_stack,
          midlines,
          input_core_dims=[["x", "y"], []],
          output_core_dims=[["position"]],
          vectorize=True,
          keep_attrs=True,
          kwargs={
              "n_points": n_points,
              "thickness": thickness,
              "order": order,
              "flatten": True,
          },
      )
  
      measurements = measurements.assign_coords(
          {"position": np.linspace(0, 1, measurements.position.size)},
      )
      try:
          measurements = measurements.assign_coords(time=fl_stack.time)
      except AttributeError:
          pass
  
      return measurements
  
  
  def shift(image: np.ndarray, vector: np.ndarray) -> np.ndarray:
      """
      Translate the image according to the given movement vector
  
      Parameters
      ----------
      image
          the image to translate
      vector :
          translation parameters ``(dx, dy)``
  
      Returns
      -------
      img: np.ndarray
          the translated image
  
      """
      tform = AffineTransform(translation=vector)
      shifted = warp(image, tform, mode="wrap", preserve_range=True)
      shifted = shifted.astype(image.dtype)
      return shifted
  
  
  def normalize_images_by_wvl_pair(
      fl_imgs: xr.DataArray, profiles: xr.DataArray, percent_to_clip: float = 2.0
  ):
      """
      Normalize images by subtracting mean profile then min-max rescaling to [0, 1]
      Parameters
      ----------
      fl_imgs
          the images to normalize
      profiles
          the intensity profiles corresponding to the images
      percent_to_clip
          how much to clip the profile when calculating mean/min/max, expressed as a percentage of the length of the profile
  
      Returns
      -------
      xr.DataArray
          the normalized images
      """
      idx_to_clip = int(profiles.shape[-1] * percent_to_clip / 100)
      profiles = profiles[:, idx_to_clip:-idx_to_clip]
  
      norm_fl = fl_imgs.copy().astype(np.float)
      for pair in fl_imgs.pair:
          for wvl in fl_imgs.wavelength.values:
              if wvl not in profiles.wavelength.values:
                  continue
              for animal in range(fl_imgs.animal.size):
                  prof = profiles.sel(wavelength=wvl, pair=pair).isel(animal=animal)
                  img = fl_imgs.sel(wavelength=wvl, pair=pair)[animal].astype(np.float)
  
                  # First, center according to mean
                  img = img - prof.mean()
  
                  # Then rescale to [0, 1]
                  img = (img - prof.min()) / (prof.max() - prof.min())
  
                  norm_fl.loc[dict(wavelength=wvl, pair=pair)][animal] = img
  
      return norm_fl
  
  
  def normalize_images_single_wvl(
      fl_imgs: Union[np.ndarray, xr.DataArray],
      profiles: Union[np.ndarray, xr.DataArray],
      percent_to_clip: float = 2.0,
  ) -> Union[np.ndarray, xr.DataArray]:
      """
      Normalize single wavelength image stack by subtracting the mean of the corresponding
      intensity profile, then min-max rescaling to [0, 1]
  
      Parameters
      ----------
      fl_imgs
          an array-like structure of shape (frame, row, col)
      profiles
          an array-like structure of shape (frame, position_along_midline)
      percent_to_clip
          how much to clip the profile when calculating mean/min/max, expressed as a percentage of the length of the profile
  
      Returns
      -------
      Union[np.ndarray, xr.DataArray]
          normalized images
      """
  
      if fl_imgs.ndim != 3:
          raise ValueError("images must have shape (frame, row, col)")
      if profiles.ndim != 2:
          raise ValueError("profiles must have shape (frame, position_along_midline)")
  
      normed_imgs = fl_imgs.copy().astype(np.float32)
  
      idx_to_clip = int(profiles.shape[-1] * percent_to_clip / 100)
      profiles = profiles[:, idx_to_clip:-idx_to_clip]
  
      prof_means = np.mean(profiles, axis=1)
  
      profiles = profiles - prof_means
      normed_imgs = normed_imgs - prof_means
  
      prof_mins = np.min(profiles, axis=1)
      prof_maxs = np.max(profiles, axis=1)
  
      normed_imgs = (normed_imgs - prof_mins) / (prof_maxs - prof_mins)
  
      return normed_imgs
  
  
  def z_normalize_with_masks(imgs, masks):
      """
      Perform z-normalization [0] on the entire image (relative to the content within the
      masks).
  
      That is to say, we center the pixels (within the mask) such that their mean is 0,
      and ensure their standard deviation is ~1.
  
      This allows us to see spatial patterns within the masked region (even if pixels
      outside of the masked region fall very far above or below those inside) by setting
      the colormap center around 0.
  
      [0] - https://jmotif.github.io/sax-vsm_site/morea/algorithm/znorm.html
      """
      masked = ma.masked_array(imgs, np.logical_not(masks))
      mu = np.mean(masked, axis=(-2, -1), keepdims=True)
      sigma = np.std(masked, axis=(-2, -1), keepdims=True)
  
      return (imgs - mu) / sigma
  
  
  def create_normed_rgb_ratio_stack(
      r_imgs, seg_imgs, vmin=-7, vmax=7, cmap="coolwarm", output_filename=None
  ):
      """
      Z-normalize the images (relative to the masks), then transform them into RGB with
      the given colormap
      """
      r_znormed = z_normalize_with_masks(r_imgs, seg_imgs)
      # noinspection PyUnresolvedReferences
      normalizer = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
      if isinstance(cmap, str):
          cmap = plt.get_cmap(cmap)
      # TODO generalize dtype? for now, 32-bit only
      rgb_img = cmap(normalizer(r_znormed))[:, :, :, :3].astype(np.float16)
  
      if output_filename is not None:
          io.imsave(output_filename, rgb_img)
  
      return rgb_img
  
  
  def get_bbox(m, pad=5):
      try:
          y_min, x_min, y_max, x_max = np.array(regionprops(label(m))[0].bbox)
  
          y_min = max(int(y_min - (pad / 2)), 0)
          x_min = max(int(x_min - (pad / 2)), 0)
          y_max = min(int(y_max + (pad / 2)), m.shape[0])
          x_max = min(int(x_max + (pad / 2)), m.shape[1])
  
          return np.array([y_min, x_min, y_max, x_max]).astype(np.float)
      except IndexError:
          return [np.nan, np.nan, np.nan, np.nan]
  
  
  def bspline_intra_modal_registration(
      fixed_image,
      moving_image,
      fixed_image_mask=None,
      point_width=5.0,
  ):
  
      registration_method = sITK.ImageRegistrationMethod()
  
      # Determine the number of BSpline control points using the physical spacing we want
      # for the control grid.
      grid_physical_spacing = [
          point_width,
          point_width,
          point_width,
      ]
      image_physical_size = [
          size * spacing
          for size, spacing in zip(fixed_image.GetSize(), fixed_image.GetSpacing())
      ]
      mesh_size = [
          int(image_size / grid_spacing + 0.5)
          for image_size, grid_spacing in zip(image_physical_size, grid_physical_spacing)
      ]
  
      initial_transform = sITK.BSplineTransformInitializer(
          image1=fixed_image, transformDomainMeshSize=mesh_size, order=2
      )
      registration_method.SetInitialTransform(initial_transform)
  
      registration_method.SetMetricAsMeanSquares()
  
      if fixed_image_mask:
          registration_method.SetMetricFixedMask(fixed_image_mask)
  
      registration_method.SetInterpolator(sITK.sitkLinear)
      registration_method.SetOptimizerAsLBFGSB(
          gradientConvergenceTolerance=1e-5, numberOfIterations=10
      )
  
      return registration_method.Execute(fixed_image, moving_image)
  
  
  def register_image(fixed, moving, mask=None, point_width=5.0):
      z_fixed = zscore(fixed.values)
      z_moving = zscore(moving.values)
  
      if mask is not None:
          mask = sITK.GetImageFromArray(mask * 255)
  
      tx = bspline_intra_modal_registration(
          sITK.GetImageFromArray(z_fixed),
          sITK.GetImageFromArray(z_moving),
          fixed_image_mask=mask,
          point_width=point_width,
      )
  
      reg_moving = sITK.GetArrayFromImage(
          sITK.Resample(
              sITK.GetImageFromArray(moving),
              sITK.GetImageFromArray(fixed),
              tx,
              sITK.sitkLinear,
          )
      )
  
      return reg_moving
  
  
  def crop(img, bbox):
      y_min, x_min, y_max, x_max = bbox.values.astype(np.int)
  
      return img[y_min:y_max, x_min:x_max]
  
  
  def register_all_images(
      imgs,
      masks,
      bbox_pad=10,
      point_width=6.0,
      fixed_wvl="410",
      moving_wvl="470",
      mask_wvl="410",
  ):
      bboxes = xr.apply_ufunc(
          get_bbox,
          masks,
          input_core_dims=[["y", "x"]],
          output_core_dims=[["pos"]],
          vectorize=True,
          kwargs={"pad": bbox_pad},
      ).assign_coords({"pos": ["min_row", "max_row", "min_col", "max_col"]})
  
      reg_imgs = imgs.copy()
  
      for animal in imgs.animal:
          for pair in imgs.pair:
              for timepoint in imgs.timepoint:
                  fixed = imgs.sel(
                      animal=animal, pair=pair, timepoint=timepoint, wavelength=fixed_wvl
                  )
                  moving = imgs.sel(
                      animal=animal, pair=pair, timepoint=timepoint, wavelength=moving_wvl
                  )
                  mask = masks.sel(animal=animal, pair=pair, timepoint=timepoint)
                  bbox = bboxes.sel(animal=animal, pair=pair, timepoint=timepoint)
  
                  # crop image
                  crop_fixed = crop(fixed, bbox)
                  crop_moving = crop(moving, bbox)
                  crop_mask = crop(mask, bbox)
  
                  # register image
                  reg_moving = register_image(
                      crop_fixed, crop_moving, mask=crop_mask, point_width=point_width
                  )
  
                  # paste cropped images back into correct location (from bbox)
                  y_min, x_min, y_max, x_max = bbox.values.astype(np.int)
                  reg_imgs.loc[
                      dict(
                          animal=animal,
                          pair=pair,
                          timepoint=timepoint,
                          wavelength=moving_wvl,
                      )
                  ][y_min:y_max, x_min:x_max] = reg_moving
  
      return reg_imgs
  