///|
/// A contiguous view of an episode used by recurrent policies and sequence models.
pub struct EpisodeWindow[S, A] {
  offset : Int
  items : Array[Transition[S, A]]
  terminal : Bool
}

///|
pub fn[S, A] EpisodeWindow::offset(self : EpisodeWindow[S, A]) -> Int {
  self.offset
}

///|
pub fn[S, A] EpisodeWindow::length(self : EpisodeWindow[S, A]) -> Int {
  self.items.length()
}

///|
pub fn[S, A] EpisodeWindow::items(
  self : EpisodeWindow[S, A],
) -> Array[Transition[S, A]] {
  self.items
}

///|
pub fn[S, A] EpisodeWindow::is_terminal(self : EpisodeWindow[S, A]) -> Bool {
  self.terminal
}

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

///|
pub fn[S, A] EpisodeWindow::total_reward(self : EpisodeWindow[S, A]) -> Double {
  let mut result = 0.0
  for item in self.items {
    result = result + item.reward
  }
  result
}

///|
pub fn[S, A] EpisodeWindow::first(
  self : EpisodeWindow[S, A],
) -> Transition[S, A]? {
  if self.items.length() == 0 {
    None
  } else {
    Some(self.items[0])
  }
}

///|
pub fn[S, A] EpisodeWindow::last(
  self : EpisodeWindow[S, A],
) -> Transition[S, A]? {
  if self.items.length() == 0 {
    None
  } else {
    Some(self.items[self.items.length() - 1])
  }
}

///|
pub fn[S, A] Episode::window(
  self : Episode[S, A],
  offset : Int,
  length : Int,
) -> EpisodeWindow[S, A] {
  let items : Array[Transition[S, A]] = []
  if length <= 0 {
    return { offset, items, terminal: false }
  }
  let start = if offset < 0 { 0 } else { offset }
  let mut i = start
  let end = if start + length > self.transitions.length() {
    self.transitions.length()
  } else {
    start + length
  }
  let mut terminal = false
  while i < end {
    let item = self.transitions[i]
    items.push(item)
    if item.done {
      terminal = true
    }
    i = i + 1
  }
  { offset: start, items, terminal }
}

///|
pub fn[S, A] Episode::windows(
  self : Episode[S, A],
  length : Int,
  stride : Int,
) -> Array[EpisodeWindow[S, A]] {
  let result : Array[EpisodeWindow[S, A]] = []
  if length <= 0 || stride <= 0 {
    return result
  }
  let mut offset = 0
  while offset < self.len() {
    let window = self.window(offset, length)
    if window.length() == length || offset == 0 {
      result.push(window)
    }
    offset = offset + stride
  }
  result
}

///|
pub fn[S, A] EpisodeWindow::discounted_reward(
  self : EpisodeWindow[S, A],
  gamma : Double,
) -> Double {
  discounted_sum(self.rewards(), gamma)
}

///|
pub fn[S, A] EpisodeWindow::contains_terminal(
  self : EpisodeWindow[S, A],
) -> Bool {
  for item in self.items {
    if item.done {
      return true
    }
  }
  false
}

///|
pub fn[S, A] EpisodeWindow::all_non_terminal(
  self : EpisodeWindow[S, A],
) -> Bool {
  !self.contains_terminal()
}

///|
pub fn[S, A] EpisodeWindow::state_count(self : EpisodeWindow[S, A]) -> Int {
  self.items.length() + (if self.items.length() == 0 { 0 } else { 1 })
}

///|
pub fn[S, A] EpisodeWindow::is_empty(self : EpisodeWindow[S, A]) -> Bool {
  self.items.length() == 0
}

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