///|
/// Dataset-level filtering and deterministic ordering helpers.
pub fn[S, A] TrajectoryDataset::filter_by_length(
  self : TrajectoryDataset[S, A],
  minimum : Int,
  maximum : Int,
) -> TrajectoryDataset[S, A] {
  let result = TrajectoryDataset::new()
  for episode in self.episodes() {
    if episode.len() >= minimum && episode.len() <= maximum {
      result.add_unchecked(episode)
    }
  }
  result
}

///|
pub fn[S, A] TrajectoryDataset::filter_by_reward(
  self : TrajectoryDataset[S, A],
  minimum : Double,
  maximum : Double,
) -> TrajectoryDataset[S, A] {
  let result = TrajectoryDataset::new()
  for episode in self.episodes() {
    let reward = episode.total_reward()
    if reward >= minimum && reward <= maximum {
      result.add_unchecked(episode)
    }
  }
  result
}

///|
pub fn[S, A] TrajectoryDataset::longest_episode(
  self : TrajectoryDataset[S, A],
) -> Episode[S, A]? {
  if self.episodes().length() == 0 {
    return None
  }
  let mut best = self.episodes()[0]
  for episode in self.episodes() {
    if episode.len() > best.len() {
      best = episode
    }
  }
  Some(best)
}

///|
pub fn[S, A] TrajectoryDataset::shortest_episode(
  self : TrajectoryDataset[S, A],
) -> Episode[S, A]? {
  if self.episodes().length() == 0 {
    return None
  }
  let mut best = self.episodes()[0]
  for episode in self.episodes() {
    if episode.len() < best.len() {
      best = episode
    }
  }
  Some(best)
}

///|
pub fn[S, A] TrajectoryDataset::terminal_rate(
  self : TrajectoryDataset[S, A],
) -> Double {
  if self.episode_count() == 0 {
    0.0
  } else {
    self.terminal_episode_count().to_double() / self.episode_count().to_double()
  }
}

///|
pub fn[S, A] TrajectoryDataset::truncation_rate(
  self : TrajectoryDataset[S, A],
) -> Double {
  if self.episode_count() == 0 {
    0.0
  } else {
    self.truncated_episode_count().to_double() /
    self.episode_count().to_double()
  }
}

///|
pub fn[S, A] TrajectoryDataset::mean_episode_length(
  self : TrajectoryDataset[S, A],
) -> Double {
  if self.episode_count() == 0 {
    0.0
  } else {
    self.transition_count().to_double() / self.episode_count().to_double()
  }
}

///|
pub fn[S, A] TrajectoryDataset::episode_reward_range(
  self : TrajectoryDataset[S, A],
) -> (Double, Double)? {
  if self.episode_count() == 0 {
    return None
  }
  let mut minimum = self.episodes()[0].total_reward()
  let mut maximum = minimum
  for episode in self.episodes() {
    let reward = episode.total_reward()
    if reward < minimum {
      minimum = reward
    }
    if reward > maximum {
      maximum = reward
    }
  }
  Some((minimum, maximum))
}

///|
pub fn[S, A] TrajectoryDataset::count_with_terminal(
  self : TrajectoryDataset[S, A],
) -> Int {
  let mut result = 0
  for episode in self.episodes() {
    if episode.terminal_index() is Some(_) {
      result = result + 1
    }
  }
  result
}

///|
pub fn[S, A] TrajectoryDataset::count_with_reward_at_least(
  self : TrajectoryDataset[S, A],
  threshold : Double,
) -> Int {
  let mut result = 0
  for episode in self.episodes() {
    if episode.total_reward() >= threshold {
      result = result + 1
    }
  }
  result
}

///|
pub fn[S, A] TrajectoryDataset::reward_per_step(
  self : TrajectoryDataset[S, A],
) -> Double {
  if self.transition_count() == 0 {
    0.0
  } else {
    self.reward_stats().sum() / self.transition_count().to_double()
  }
}

///|
pub fn[S, A] TrajectoryDataset::append(
  self : TrajectoryDataset[S, A],
  other : TrajectoryDataset[S, A],
) -> Unit {
  for episode in other.episodes() {
    self.add_unchecked(episode)
  }
}

///|
pub fn[S, A] TrajectoryDataset::take_episodes(
  self : TrajectoryDataset[S, A],
  count : Int,
) -> Array[Episode[S, A]] {
  let result : Array[Episode[S, A]] = []
  let limit = if count < 0 {
    0
  } else if count > self.episode_count() {
    self.episode_count()
  } else {
    count
  }
  let mut i = 0
  while i < limit {
    result.push(self.episodes()[i])
    i = i + 1
  }
  result
}

///|
pub fn[S, A] TrajectoryDataset::drop_episodes(
  self : TrajectoryDataset[S, A],
  count : Int,
) -> TrajectoryDataset[S, A] {
  let result = TrajectoryDataset::new()
  let start = if count < 0 {
    0
  } else if count > self.episode_count() {
    self.episode_count()
  } else {
    count
  }
  let mut i = start
  while i < self.episode_count() {
    result.add_unchecked(self.episodes()[i])
    i = i + 1
  }
  result
}