///|
/// Deterministic benchmark fixtures for regression tests and demos.
pub struct BenchmarkCase {
  name : String
  episodes : Int
  steps_per_episode : Int
  reward_scale : Double
}

///|
pub fn BenchmarkCase::new(
  name : String,
  episodes : Int,
  steps_per_episode : Int,
  reward_scale : Double,
) -> BenchmarkCase {
  { name, episodes, steps_per_episode, reward_scale }
}

///|
pub fn BenchmarkCase::name(self : BenchmarkCase) -> String {
  self.name
}

///|
pub fn BenchmarkCase::episodes(self : BenchmarkCase) -> Int {
  self.episodes
}

///|
pub fn BenchmarkCase::steps_per_episode(self : BenchmarkCase) -> Int {
  self.steps_per_episode
}

///|
pub fn BenchmarkCase::reward_scale(self : BenchmarkCase) -> Double {
  self.reward_scale
}

///|
pub fn BenchmarkCase::total_steps(self : BenchmarkCase) -> Int {
  self.episodes * self.steps_per_episode
}

///|
pub fn benchmark_cases() -> Array[BenchmarkCase] {
  [
    BenchmarkCase::new("sparse-reward", 8, 32, 1.0),
    BenchmarkCase::new("dense-reward", 8, 32, 0.25),
    BenchmarkCase::new("long-horizon", 4, 256, 0.1),
    BenchmarkCase::new("short-horizon", 32, 8, 2.0),
  ]
}

///|
pub fn[S, A] generate_benchmark_dataset(
  case : BenchmarkCase,
  state : S,
  action : A,
) -> TrajectoryDataset[S, A] {
  let dataset = TrajectoryDataset::new()
  let mut episode_index = 0
  while episode_index < case.episodes() {
    let mut episode = Episode::new()
    let mut step = 0
    while step < case.steps_per_episode() {
      let terminal = step == case.steps_per_episode() - 1
      let reward = case.reward_scale() * (if terminal { 1.0 } else { 0.0 })
      let transition = Transition::new(state, action, reward, state, terminal)
      episode = episode.push(transition)
      step = step + 1
    }
    let _ = dataset.add(episode)
    episode_index = episode_index + 1
  }
  dataset
}

///|
pub fn[S, A] benchmark_summary(
  case : BenchmarkCase,
  state : S,
  action : A,
) -> EvaluationReport {
  let dataset = generate_benchmark_dataset(case, state, action)
  evaluate_episodes(dataset.episodes())
}

///|
pub fn benchmark_workload_size(cases : Array[BenchmarkCase]) -> Int {
  let mut total = 0
  for case in cases {
    total = total + case.total_steps()
  }
  total
}

///|
pub fn benchmark_case_by_name(
  cases : Array[BenchmarkCase],
  name : String,
) -> BenchmarkCase? {
  for case in cases {
    if case.name() == name {
      return Some(case)
    }
  }
  None
}

///|
pub fn benchmark_reward_scales(cases : Array[BenchmarkCase]) -> Array[Double] {
  let result : Array[Double] = []
  for case in cases {
    result.push(case.reward_scale())
  }
  result
}