# Copyright 2024 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""Track data container analogous to AnnData."""

from collections.abc import Sequence
import copy
import dataclasses
import enum
from typing import Any, Union

from alphagenome import typing
from alphagenome.data import genome
from alphagenome.data import ontology
from jaxtyping import Bool, Float32, Int32  # pylint: disable=g-multiple-import, g-importing-member
import numpy as np
import pandas as pd

# Required columns: name, strand.
# Optional standardized columns: cell_type, assay, padding.
TrackMetadata = pd.DataFrame
PositionalIndex = slice | genome.Interval | int
TrackIndex = np.ndarray | Sequence[int] | Sequence[str] | slice | int | str
Index = PositionalIndex | tuple[PositionalIndex, TrackIndex] | TrackIndex


@enum.unique
class AggregationType(enum.Enum):
  """Aggregation types for downsampling/upsampling track resolutions.

  SUM: Sum pooling, where values within a bin are summed. This is recommended
    for continuous tracks where the total value within a bin is meaningful (e.g.
    read counts, coverage).
  MAX: Max pooling, where the maximum value within a bin is selected. This is
    recommended for binary tracks (e.g., gene masks, regions of interest).
  """

  SUM = 'sum'
  MAX = 'max'


@typing.jaxtyped
@dataclasses.dataclass(frozen=True)
class TrackData:
  """Container for storing track values and metadata.

  `TrackData` stores multiple genomic tracks at the same resolution, stacked
  into an ND matrix of shape (positional_bins, num_tracks). It also contains
  metadata information as a pandas DataFrame with `num_tracks` rows.

  Metadata DataFrame has two main required columns:

    * name: The name of the track.
    * strand: The strand of the track ('+', '-', or '.').

  Other columns are optional.

  Valid shapes of `TrackData.values` are:

    * [positional_bins]
    * [positional_bins, num_tracks]
    * [positional_bins, positional_bins, num_tracks]
    * ...

  `TrackData` can store both model predictions and raw data. It can
  optionally hold information about the `genome.Interval` from which the data
  were derived and `.uns` for storing additional unstructured data.

  In addition to being a container, `TrackData` provides functionality for
  common aggregation and slicing operations.

  Attributes:
    values: A numpy array of floats or integers representing the track values.
      Positional axes have the same length. Example valid shapes are:
      [num_tracks], [positional_bins, num_tracks], and [positional_bins,
      positional_bins, num_tracks].
    metadata: A pandas DataFrame containing metadata for each track. The
      DataFrame must have at least two columns: 'name' and 'strand'.
    resolution: The resolution of the track data in base pairs.
    interval: An optional `Interval` object representing the genomic region.
    uns: An optional dictionary to store additional unstructured data.

  Raises:
    ValueError: If the number of tracks in `values` does not match the  number
      of rows in `metadata`, or if `metadata` contains duplicate (name, strand)
      pairs, or if the positional axes have different lengths, or if the
      interval width does not match the expected width.
  """

  # Use Union due to https://github.com/patrick-kidger/jaxtyping/issues/73.
  values: Union[
      Float32[np.ndarray, '*positional_bins num_tracks'],
      Int32[np.ndarray, '*positional_bins num_tracks'],
      Bool[np.ndarray, '*positional_bins num_tracks'],
  ]
  metadata: TrackMetadata
  resolution: int = 1
  interval: genome.Interval | None = None
  uns: dict[str, Any] | None = (
      # Unstructured data dict, analagous to anndata.AnnData.uns.
      None
  )

  def __post_init__(self):
    """Validates the consistency of the data."""
    if self.values.shape[-1] != len(self.metadata):
      raise ValueError(
          f'value number of tracks {self.values.shape[-1]} and '
          f'metadata {len(self.metadata)} do not match.'
      )

    if self.positional_axes:
      if len(set(np.array(self.values.shape)[self.positional_axes])) != 1:
        raise ValueError('All positional axes must have the same length.')

      if self.interval and self.interval.width != self.width:
        raise ValueError(
            f'Interval width must match expected width. {self.interval.width=},'
            f' {self.width=}'
        )

    if not {'name', 'strand'}.issubset(self.metadata.columns):
      raise ValueError('Metadata must contain columns "name" and "strand".')

    if self.metadata[['name', 'strand']].duplicated().any():
      raise ValueError(
          'Metadata contain duplicated values for (name, strand) tuples.'
      )

  @property
  def positional_axes(self) -> list[int]:
    """Returns a list of the positional axes."""
    return list(range(self.values.ndim - 1))

  @property
  def num_tracks(self) -> int:
    """Returns the number of tracks."""
    return self.values.shape[-1]

  @property
  def width(self) -> int:
    """Returns the interval width covered by the tracks."""
    if self.positional_axes:
      return self.values.shape[0] * self.resolution
    else:
      return 0

  @property
  def names(self) -> np.ndarray:
    """Returns an array of track names (not necessarily unique)."""
    return self.metadata['name'].values

  @property
  def strands(self) -> np.ndarray:
    """Returns an array of track strands."""
    return self.metadata['strand'].values

  @property
  def ontology_terms(self) -> Sequence[ontology.OntologyTerm | None] | None:
    """Returns a list of ontology terms (if available)."""
    if 'ontology_curie' in self.metadata.columns:
      return [
          ontology.from_curie(curie) if curie is not None else None
          for curie in self.metadata['ontology_curie'].values
      ]
    else:
      return None

  def copy(self) -> 'TrackData':
    """Returns a deep copy of the `TrackData` object."""
    if self.interval:
      interval = self.interval.copy()
    else:
      interval = None
    return TrackData(
        self.values.copy(),
        resolution=self.resolution,
        metadata=self.metadata.copy(),
        interval=interval,
        uns=copy.deepcopy(self.uns),
    )

  def bin_index(self, relative_position: int) -> int:
    """Returns the bin index for a relative position.

    Args:
      relative_position: The relative position within the interval.

    Returns:
      The corresponding bin index.
    """
    return relative_position // self.resolution

  def slice_by_positions(self, start: int, end: int) -> 'TrackData':
    """Slices the track data along the positional axes.

    The slicing follows Python slicing conventions (0 indexed, and includes
    elements up to end-1).

    Args:
      start: The 1-bp resolution start position for slicing.
      end: The 1-bp resolution end position for slicing.

    Returns:
      A new `TrackData` object with the sliced values.

    Raises:
      ValueError: If (end - start) is greater than the width, or if (end -
      start) is not divisible by the resolution.
    """
    if (end - start) > self.width:
      raise ValueError(
          'When slicing track data, (end - start) must be less than or '
          'equal to width.'
      )

    if (end - start) % self.resolution != 0:
      raise ValueError(
          f'end - start needs to be to be divisible by {self.resolution=}'
      )

    sl = slice(self.bin_index(start), self.bin_index(end))
    slice_list = [slice(None)] * self.values.ndim
    for i in self.positional_axes:
      slice_list[i] = sl

    interval = self.interval
    if interval:
      interval = genome.Interval(
          interval.chromosome,
          interval.start + start,
          interval.start + end,
          strand=interval.strand,
          name=interval.name,
          info=interval.info,
      )

    return TrackData(
        self.values[tuple(slice_list)],
        resolution=self.resolution,
        metadata=self.metadata,
        interval=interval,
        uns=self.uns,
    )

  def slice_by_interval(
      self, interval: genome.Interval, match_resolution: bool = False
  ) -> 'TrackData':
    """Slices the track data using a `genome.Interval`.

    Args:
      interval: The interval to slice to.
      match_resolution: If True, the interval will first be extended to make
        sure the width is divisible by resolution.

    Returns:
      A new `TrackData` object sliced to the interval.

    Raises:
      ValueError: If `.interval` is not specified or if the specified interval
        is not fully contained within the current interval.
    """
    if self.interval is None:
      raise ValueError(
          '.interval is needs to be specified for slice_by_interval.'
      )
    if not self.interval.contains(interval):
      raise ValueError(
          f'Interval {self.interval=} does not fully contain {interval=}.'
      )
    start = interval.start - self.interval.start
    end = interval.end - self.interval.start

    if match_resolution and self.resolution != 1:
      start = int(np.floor(start / self.resolution) * self.resolution)
      end = int(np.ceil(end / self.resolution) * self.resolution)
    return self.slice_by_positions(start, end)

  def pad(self, start_pad: int, end_pad: int) -> 'TrackData':
    """Pads the track data along positional axes.

    Args:
      start_pad: The amount of padding to add at the beginning.
      end_pad: The amount of padding to add at the end.

    Returns:
      A new `TrackData` object with padded values.

    Raises:
      ValueError: If `start_pad` or `end_pad` is not divisible by the
      resolution.
    """
    if start_pad == 0 and end_pad == 0:
      return self
    if start_pad % self.resolution != 0:
      raise ValueError(f'start_pad needs to be divisible by {self.resolution}')
    if end_pad % self.resolution != 0:
      raise ValueError(f'end_pad needs to be divisible by {self.resolution}')

    pad = [(0, 0)] * self.values.ndim
    for axis in self.positional_axes:
      pad[axis] = (start_pad // self.resolution, end_pad // self.resolution)

    return TrackData(
        np.pad(self.values, tuple(pad)),
        resolution=self.resolution,
        metadata=self.metadata,
        interval=None,  # Padding invalidates the interval.
        uns=self.uns,
    )

  def resize(self, width: int) -> 'TrackData':
    """Resizes the track data by cropping or padding with a fixed center.

    Args:
      width: The desired width in base pairs.

    Returns:
      A new `TrackData` object with resized values.

    Raises:
      ValueError: If `width` is not divisible by the resolution.
    """
    if width == self.width:
      return self
    elif width > self.width:
      if width % self.resolution != 0:
        raise ValueError(f'width needs to be divisible by {self.resolution}')
      pad_amount = (width - self.width) // self.resolution
      pad_start = (pad_amount // 2 + pad_amount % 2) * self.resolution
      pad_end = (pad_amount // 2) * self.resolution
      return self.pad(pad_start, pad_end)
    else:
      crop_amount = (self.width - width) // self.resolution
      start = (crop_amount // 2 + crop_amount % 2) * self.resolution
      return self.slice_by_positions(start, start + width)

  def upsample(
      self,
      resolution: int,
      aggregation_type: AggregationType = AggregationType.SUM,
  ) -> 'TrackData':
    """Upsamples the track data to a higher resolution by repeating existing values.

    Args:
      resolution: The desired resolution in base pairs.
      aggregation_type: The aggregation method to use for pooling the values.

    Returns:
      A new `TrackData` object with upsampled values.

    Raises:
      ValueError: If `resolution` is not lower than the current resolution
        or not divisible by the current resolution.
    """
    if resolution == self.resolution:
      return self
    if resolution > self.resolution:
      raise ValueError(f'Resolution must be lower than {self.resolution}')
    repeat = self.resolution // resolution
    if self.resolution % resolution != 0:
      raise ValueError(f'Resolution not divisible by {resolution}')

    values = self.values
    for axis in self.positional_axes:
      values = np.repeat(values, repeat, axis=axis)
      match aggregation_type:
        case AggregationType.SUM:
          values = values / repeat
        case AggregationType.MAX:
          pass
    return TrackData(
        values,
        resolution=resolution,
        metadata=self.metadata,
        interval=self.interval,
        uns=self.uns,
    )

  def downsample(
      self,
      resolution: int,
      aggregation_type: AggregationType = AggregationType.SUM,
  ) -> 'TrackData':
    """Downsamples the track data to a lower resolution using sum pooling.

    Args:
      resolution: The desired resolution in base pairs.
      aggregation_type: The aggregation method to use for pooling the values.

    Returns:
      A new `TrackData` object with downsampled values.

    Raises:
      ValueError: If `resolution` is not greater than the current resolution
        or not divisible by the current resolution.
    """
    if resolution == self.resolution:
      return self
    if resolution < self.resolution:
      raise ValueError(f'Resolution must be greater than {self.resolution}')
    if resolution % self.resolution != 0:
      raise ValueError(f'Resolution not divisible by {resolution}')
    pool_width = resolution // self.resolution

    values = self.values
    for axis in self.positional_axes:
      # Bring axis of interest to the front, reshape and aggregate, and reswap
      values = np.swapaxes(values, 0, axis)
      shape = list(values.shape)
      reshaped_values = values.reshape(
          [shape[0] // pool_width, pool_width] + shape[1:]
      )
      match aggregation_type:
        case AggregationType.SUM:
          values = reshaped_values.sum(axis=1)
        case AggregationType.MAX:
          values = reshaped_values.max(axis=1)
      values = np.swapaxes(values, 0, axis)

    return TrackData(
        values,
        resolution=resolution,
        metadata=self.metadata,
        interval=self.interval,
        uns=self.uns,
    )

  def change_resolution(
      self,
      resolution: int,
      aggregation_type: AggregationType = AggregationType.SUM,
  ) -> 'TrackData':
    """Changes the resolution of the track data.

    Args:
      resolution: The desired resolution in base pairs.
      aggregation_type: The aggregation method to use for pooling the values.

    Returns:
      A new `TrackData` object with the new resolution.
    """
    if resolution >= self.resolution:
      return self.downsample(resolution, aggregation_type)
    else:
      return self.upsample(resolution, aggregation_type)

  def filter_tracks(self, mask: np.ndarray | list[bool]) -> 'TrackData':
    """Filters tracks by a boolean mask.

    Args:
      mask: A boolean mask to select tracks.

    Returns:
      A new `TrackData` object with the filtered tracks.
    """
    return TrackData(
        self.values[..., mask],
        resolution=self.resolution,
        metadata=self.metadata.iloc[mask],
        interval=self.interval,
        uns=self.uns,
    )

  def filter_to_positive_strand(self) -> 'TrackData':
    """Filters tracks to the positive DNA strand."""
    return self.filter_tracks(self.strands == genome.STRAND_POSITIVE)

  def filter_to_negative_strand(self) -> 'TrackData':
    """Filters tracks to the negative DNA strand."""
    return self.filter_tracks(self.strands == genome.STRAND_NEGATIVE)

  def filter_to_nonnegative_strand(self) -> 'TrackData':
    """Filters tracks to the non-negative DNA strands (positive and unstranded)."""
    return self.filter_tracks(self.strands != genome.STRAND_NEGATIVE)

  def filter_to_nonpositive_strand(self) -> 'TrackData':
    """Filters tracks to the non-positive DNA strands (negative and unstranded)."""
    return self.filter_tracks(self.strands != genome.STRAND_POSITIVE)

  def filter_to_stranded(self) -> 'TrackData':
    """Filters tracks to stranded tracks (excluding unstranded)."""
    return self.filter_tracks(self.strands != genome.STRAND_UNSTRANDED)

  def filter_to_unstranded(self) -> 'TrackData':
    """Filters tracks to unstranded tracks."""
    return self.filter_tracks(self.strands == genome.STRAND_UNSTRANDED)

  def select_tracks_by_index(
      self, idx: np.ndarray | Sequence[int]
  ) -> 'TrackData':
    """Selects tracks by numerical index.

    Args:
      idx: A list or array of numerical indices to select tracks.

    Returns:
      A new `TrackData` object with the selected tracks.
    """
    return TrackData(
        self.values[..., idx],
        resolution=self.resolution,
        metadata=self.metadata.iloc[idx],
        interval=self.interval,
        uns=self.uns,
    )

  def select_tracks_by_name(
      self, names: np.ndarray | Sequence[str]
  ) -> 'TrackData':
    """Selects tracks by name.

    Args:
      names: A list or array of track names to select.

    Returns:
      A new `TrackData` object with the selected tracks.
    """
    track_idx = pd.Series(np.arange(self.num_tracks), index=self.names)
    return self.select_tracks_by_index(track_idx.loc[names].values)

  def __getitem__(self, index: Index) -> 'TrackData':
    """Retrieves a subset of TrackData using positional and/or track indices.

    This method allows slicing `TrackData` similar to numpy arrays or pandas
    DataFrames. The index can be a single value or a tuple.

    Args:
      index: A single index or a tuple of indices. If a single index, it's
        treated as a positional index if the `TrackData` has positional axes,
        otherwise as a track index. If a tuple, the first element specifies the
        positional slice, and the second element specifies the track index.
        Positional indices can be `int`, `slice`, or `genome.Interval`, and this
        slice is applied to all positional axes of the `values` array. Track
        indices can be `int`, `str` (track name), `slice`, `Sequence[int]`, or
        `Sequence[str]` (track names).

    Returns:
      A new `TrackData` object containing the selected subset.

    Raises:
      IndexError: If a slice step is not 1 for positional indexing.
      IndexError: If an unsupported index type is provided.
    """
    if isinstance(index, tuple):
      position_index, track_index = index
    elif self.positional_axes:
      position_index, track_index = index, None
    else:
      position_index, track_index = None, index
      if isinstance(track_index, genome.Interval):
        raise IndexError(
            'Track indexing by interval is supported only when there are'
            ' positional axes.'
        )

    tdata = self
    match position_index:
      case None:
        pass
      case int():
        tdata = tdata.slice_by_positions(position_index, position_index + 1)
      case slice():
        if position_index.step is not None and position_index.step != 1:
          raise IndexError('Slice step must be 1 for positional indexing.')
        if position_index != slice(None):
          tdata = tdata.slice_by_positions(
              position_index.start, position_index.stop
          )
      case genome.Interval():
        tdata = tdata.slice_by_interval(position_index)
      case _:
        raise IndexError(
            f'Unsupported positional index type: {type(position_index)}'
        )

    match track_index:
      case None:
        pass
      case str():
        tdata = tdata.select_tracks_by_name([track_index])
      case int():
        tdata = tdata.select_tracks_by_index([track_index])
      case slice():
        if track_index != slice(None):
          indices = np.arange(tdata.num_tracks)[track_index]
          tdata = tdata.select_tracks_by_index(indices)
      case np.ndarray() if np.issubdtype(track_index.dtype, np.character):
        tdata = tdata.select_tracks_by_name(track_index)
      case np.ndarray():
        tdata = tdata.select_tracks_by_index(track_index)
      case Sequence():
        track_index_arr = np.asarray(track_index)
        if np.issubdtype(track_index_arr.dtype, np.character):
          tdata = tdata.select_tracks_by_name(track_index_arr)
        else:
          tdata = tdata.select_tracks_by_index(track_index_arr)
      case _:
        raise IndexError(f'Unsupported track index type: {type(track_index)}')
    return tdata

  def groupby(self, column: str) -> dict[str, 'TrackData']:
    """Splits tracks into groups based on a metadata column.

    This method splits the tracks in the `TrackData` object into separate
    `TrackData` objects based on the unique values in the specified metadata
    column. It returns a dictionary where the keys are the  unique values in
    the column, and the values are new `TrackData` objects containing the
    tracks corresponding to each key.

    Args:
      column: The name of the metadata column to split by.

    Returns:
      A dictionary mapping unique values in the column to `TrackData` objects
      containing the corresponding tracks.
    """
    output = {}
    for key in self.metadata[column].unique():
      mask = (self.metadata[column] == key).values
      output[key] = self.filter_tracks(mask)
    return output

  def _reverse_complement_idx(self) -> np.ndarray:
    """Gets indices for reverse complementing the tracks.

    Returns:
      An array of indices that reorders the tracks to achieve reverse
      complementation.

    Raises:
      ValueError: If not all stranded tracks have both '+' and '-' strands,
         or if the number of '+' and '-' stranded tracks is not equal.
    """
    df_strands = pd.DataFrame({
        'name': self.names,
        'strand': self.strands,
        'old_idx': np.arange(self.num_tracks),
    })
    df_strands = df_strands[df_strands.strand != genome.STRAND_UNSTRANDED]
    df_strands.sort_values(['strand', 'name'], inplace=True)
    if np.all(df_strands.groupby('name').size() != 2):
      raise ValueError('Not all stranded tracks have both + and - strand.')
    if (df_strands.strand == genome.STRAND_POSITIVE).sum() != (
        df_strands.strand == genome.STRAND_NEGATIVE
    ).sum():
      raise ValueError(
          'We need to have the exact same number of + and - stranded tracks'
      )
    new_idx = df_strands.old_idx.values.reshape((2, -1))[::-1].ravel()
    # Swap strands by idx.
    idx = np.arange(self.num_tracks)
    idx[df_strands.old_idx.values] = new_idx
    return idx

  def reverse_complement(self) -> 'TrackData':
    """Reverse complements the track data and interval if present.

    Returns:
      A new `TrackData` object with reverse complemented tracks.
    """
    if self.interval:
      # Note that the interval needs to be stranded in order to perform
      # this operation.
      interval = self.interval.swap_strand()
    else:
      interval = None

    idx = self._reverse_complement_idx()
    slices = [slice(None)] * self.values.ndim
    slices[-1] = idx
    for axis in self.positional_axes:
      slices[axis] = slice(None, None, -1)

    return TrackData(
        self.values[tuple(slices)],
        resolution=self.resolution,
        metadata=self.metadata.iloc[idx],
        interval=interval,
        uns=self.uns,
    )

  def _check_track_data_compatibility(self, other: 'TrackData') -> None:
    """Checks if two `TrackData` objects are compatible for sum/diff.

    Args:
      other: The other `TrackData` object to compare.

    Raises:
      TypeError: If `other` is not a `TrackData` object.
      ValueError: If the intervals, resolutions, shapes, or metadata
        shapes don't match between the two objects.
    """
    if not isinstance(other, TrackData):
      raise TypeError(
          f'Unsupported type "{type(other)}". Must be a TrackData object'
      )
    if self.interval != other.interval:
      raise ValueError('Intervals must match for the two TrackData objects.')
    if self.resolution != other.resolution:
      raise ValueError('Resolutions must match for the two TrackData objects.')
    if self.values.shape != other.values.shape:
      raise ValueError('Shapes must match for the two TrackData objects.')
    if self.metadata.shape != other.metadata.shape:
      raise ValueError(
          'Metadata shapes must match for the two TrackData objects.'
      )

  def __add__(self, other: 'TrackData') -> 'TrackData':
    """Adds the values of two `TrackData` objects.

    Args:
      other: The `TrackData` object to add.

    Returns:
      A new `TrackData` object with the summed values.

    Raises:
      ValueError: If the objects are not compatible (see
        `_check_track_data_compatibility`).
      TypeError: If `other` is not a `TrackData` object.
    """
    self._check_track_data_compatibility(other)
    new_values = self.values + other.values
    return TrackData(
        values=new_values,
        metadata=self.metadata,
        resolution=self.resolution,
        interval=self.interval,
        uns=self.uns,
    )

  def __sub__(self, other: 'TrackData') -> 'TrackData':
    """Subtracts the values of two `TrackData` objects.

    Args:
      other: The `TrackData` object to subtract.

    Returns:
      A new `TrackData` object with the difference of the values.

    Raises:
      ValueError: If the objects are not compatible (see
        `_check_track_data_compatibility`).
      TypeError: If `other` is not a `TrackData` object.
    """
    self._check_track_data_compatibility(other)
    new_values = self.values - other.values
    return TrackData(
        values=new_values,
        metadata=self.metadata,
        resolution=self.resolution,
        interval=self.interval,
        uns=self.uns,
    )


def concat(
    track_datas: Sequence[TrackData],
    extra_metadata_name_and_keys: (
        tuple[str, Sequence[str | int | float]] | None
    ) = None,
) -> TrackData:
  """Concatenates multiple `TrackData` objects along the track dimension.

  This function combines multiple `TrackData` objects into a single object
  by concatenating their values and metadata. The resulting `TrackData`
  object will have the same resolution and interval as the input objects.

  Args:
    track_datas: A sequence of `TrackData` objects to concatenate. All objects
      must have the same resolution, interval, and width.
    extra_metadata_name_and_keys: An optional tuple specifying a new metadata
      column to add. The first element is the column name, and the second is a
      sequence of values to populate the column.

  Returns:
    A new `TrackData` object containing the concatenated data.

  Raises:
    ValueError: If the input `TrackData` objects have different resolutions,
      intervals, or widths, or if the length of
      `extra_metadata_name_and_keys[1]` does not match the length of
      `track_datas`.
  """
  if len(set(x.resolution for x in track_datas)) != 1:
    raise ValueError('Track data contain multiple resolutions')
  if len(set(str(x.interval) for x in track_datas)) != 1:
    raise ValueError('Track data contain multiple intervals')
  if len(set(x.width for x in track_datas)) != 1:
    raise ValueError('Track data are of different width')
  if extra_metadata_name_and_keys:
    if len(extra_metadata_name_and_keys[1]) != len(track_datas):
      raise ValueError(
          'Second element of new_metadata_name_and_keys must be the same'
          ' length as track_datas'
      )

  concatenated_metadata = (
      pd.concat(
          [x.metadata for x in track_datas],
          keys=extra_metadata_name_and_keys[1]
          if extra_metadata_name_and_keys
          else None,
          names=[extra_metadata_name_and_keys[0]]
          if extra_metadata_name_and_keys
          else None,
      )
      .reset_index(level=0, drop=extra_metadata_name_and_keys is None)
      .reset_index(drop=True)
  )
  return TrackData(
      np.concatenate(
          [x.values for x in track_datas], axis=track_datas[0].values.ndim - 1
      ),
      resolution=track_datas[0].resolution,
      metadata=concatenated_metadata,
      interval=track_datas[0].interval,
      uns=None,
  )


def interleave(
    track_datas: Sequence[TrackData], name_prefixes: Sequence[str]
) -> TrackData:
  """Interleaves multiple `TrackData` objects by alternating rows.

  This function combines multiple `TrackData` objects into a single object
  by interleaving their rows and metadata. This interleaves operation alternates
  between the trackdatas, like shuffling cards, i.e., 'abcd' interleaved with
  'efgh' would be "aebfcgdh". The resulting `TrackData` object will have the
  same resolution and interval as the input objects, but the number of tracks
  will be the sum of the tracks in the input objects.

  Args:
    track_datas: A sequence of `TrackData` objects to interleave. All objects
      must have the same shape, resolution, and interval. The order in this list
      will determine the interleaving order.
    name_prefixes: A sequence of name prefixes to add to the track names in the
      metadata to ensure uniqueness of (name, strand) pairs.

  Returns:
    A new `TrackData` object containing the interleaved data.

  Raises:
    ValueError: If the input `TrackData` objects have different shapes,
      resolutions, or intervals.
  """
  # Checks on the track data.
  shapes = set(data.values.shape for data in track_datas)
  if len(shapes) != 1:
    raise ValueError(
        'Cannot interleave track data which have different shapes. '
        f'Detected shapes: {shapes}'
    )

  if any(data.resolution != track_datas[0].resolution for data in track_datas):
    raise ValueError(
        'Cannot interleave track data which have different resolutions. '
        f'Detected shapes: {shapes}'
    )

  if any(data.interval != track_datas[0].interval for data in track_datas):
    raise ValueError(
        'Cannot interleave track data which have different intervals. '
    )

  # Interleave arrays.
  shape = list(track_datas[0].values.shape)
  shape[-1] = shape[-1] * len(track_datas)
  interleaved_data = np.empty(tuple(shape), dtype=track_datas[0].values.dtype)

  for i, data in enumerate(track_datas):
    interleaved_data[..., i :: len(track_datas)] = data.values

  # Interleave metadata.
  metadatas = []
  for prefix, data in zip(name_prefixes, track_datas):
    metadata = data.metadata.copy()
    metadata['name'] = prefix + metadata['name']
    metadatas.append(metadata)

  # Assign a new index idx that, when sorted, will produce an interleave.
  interleaved_metadata = (
      pd.concat([
          metadata.assign(
              idx=np.arange(
                  stop=(len(metadata) * len(track_datas)),
                  step=len(track_datas),
              )
              + i
          )
          for i, metadata in enumerate(metadatas)
      ])
      .sort_values('idx')
      .drop('idx', axis=1)
      .reset_index(drop=True)
  )

  return TrackData(
      values=interleaved_data,
      metadata=interleaved_metadata,
      resolution=track_datas[0].resolution,
      interval=track_datas[0].interval,
      uns={'num_interleaved_trackdatas': len(track_datas)},
  )
