///|
/// A compact summary of rewards in an episode or a collection of transitions.
/// The summary is deliberately independent of the state and action types, so it
/// can be used with strings, records, tensors, or application-specific values.
pub struct RewardStats {
  count : Int
  sum : Double
  mean : Double
  minimum : Double
  maximum : Double
  positive_count : Int
  negative_count : Int
  zero_count : Int
}

///|
pub fn RewardStats::empty() -> RewardStats {
  {
    count: 0,
    sum: 0.0,
    mean: 0.0,
    minimum: 0.0,
    maximum: 0.0,
    positive_count: 0,
    negative_count: 0,
    zero_count: 0,
  }
}

///|
pub fn RewardStats::count(self : RewardStats) -> Int {
  self.count
}

///|
pub fn RewardStats::sum(self : RewardStats) -> Double {
  self.sum
}

///|
pub fn RewardStats::mean(self : RewardStats) -> Double {
  self.mean
}

///|
pub fn RewardStats::minimum(self : RewardStats) -> Double {
  self.minimum
}

///|
pub fn RewardStats::maximum(self : RewardStats) -> Double {
  self.maximum
}

///|
pub fn RewardStats::positive_count(self : RewardStats) -> Int {
  self.positive_count
}

///|
pub fn RewardStats::negative_count(self : RewardStats) -> Int {
  self.negative_count
}

///|
pub fn RewardStats::zero_count(self : RewardStats) -> Int {
  self.zero_count
}

///|
fn[S, A] reward_stats_from_transitions(
  transitions : Array[Transition[S, A]],
) -> RewardStats {
  if transitions.length() == 0 {
    return RewardStats::empty()
  }
  let mut sum = 0.0
  let mut minimum = transitions[0].reward
  let mut maximum = transitions[0].reward
  let mut positive_count = 0
  let mut negative_count = 0
  let mut zero_count = 0
  for transition in transitions {
    let reward = transition.reward
    sum = sum + reward
    if reward < minimum {
      minimum = reward
    }
    if reward > maximum {
      maximum = reward
    }
    if reward > 0.0 {
      positive_count = positive_count + 1
    } else if reward < 0.0 {
      negative_count = negative_count + 1
    } else {
      zero_count = zero_count + 1
    }
  }
  {
    count: transitions.length(),
    sum,
    mean: sum / transitions.length().to_double(),
    minimum,
    maximum,
    positive_count,
    negative_count,
    zero_count,
  }
}

///|
pub fn[S, A] Episode::reward_stats(self : Episode[S, A]) -> RewardStats {
  reward_stats_from_transitions(self.transitions)
}

///|
/// Validate common data invariants before an episode enters a training buffer.
/// The returned array is empty when the episode is safe to consume.
pub fn[S, A] Episode::validate(self : Episode[S, A]) -> Array[String] {
  let issues : Array[String] = []
  if self.transitions.length() == 0 {
    issues.push("episode is empty")
  }
  if self.terminated && self.truncated {
    issues.push("episode cannot be both terminated and truncated")
  }
  let mut terminal_seen = false
  let mut i = 0
  while i < self.transitions.length() {
    let transition = self.transitions[i]
    if terminal_seen {
      issues.push(
        "transition appears after a terminal transition at index \\{i}",
      )
    }
    if transition.done {
      terminal_seen = true
      if i != self.transitions.length() - 1 {
        issues.push("terminal transition must be the last transition")
      }
    }
    i = i + 1
  }
  if self.terminated && !terminal_seen {
    issues.push("episode marked terminated without a terminal transition")
  }
  if !self.terminated && terminal_seen {
    issues.push(
      "terminal transition exists but episode is not marked terminated",
    )
  }
  issues
}

///|
/// Compute returns while treating a value as the continuation value after the
/// last transition. This is useful for truncated rollouts in bootstrapped RL.
pub fn[S, A] Episode::discounted_returns_with_bootstrap(
  self : Episode[S, A],
  gamma : Double,
  bootstrap : Double,
) -> Array[Double] {
  let len = self.transitions.length()
  let returns : Array[Double] = []
  let mut i = 0
  while i < len {
    returns.push(0.0)
    i = i + 1
  }
  let mut running = bootstrap
  let mut idx = len
  while idx > 0 {
    idx = idx - 1
    let transition = self.transitions[idx]
    running = transition.reward + gamma * running
    returns[idx] = running
    if transition.done {
      running = 0.0
    }
  }
  returns
}

///|
pub fn[S, A] Episode::terminal_index(self : Episode[S, A]) -> Int? {
  let mut i = 0
  while i < self.transitions.length() {
    if self.transitions[i].done {
      return Some(i)
    }
    i = i + 1
  }
  None
}

///|
pub fn[S, A] Episode::count_terminals(self : Episode[S, A]) -> Int {
  let mut count = 0
  for transition in self.transitions {
    if transition.done {
      count = count + 1
    }
  }
  count
}

///|
pub fn[S, A] Episode::rewards(self : Episode[S, A]) -> Array[Double] {
  let result : Array[Double] = []
  for transition in self.transitions {
    result.push(transition.reward)
  }
  result
}

///|
pub fn[S, A] Episode::states(self : Episode[S, A]) -> Array[S] {
  let result : Array[S] = []
  for transition in self.transitions {
    result.push(transition.state)
  }
  result
}

///|
pub fn[S, A] Episode::actions(self : Episode[S, A]) -> Array[A] {
  let result : Array[A] = []
  for transition in self.transitions {
    result.push(transition.action)
  }
  result
}