///|
pub struct PrioritizedSample[S, A] {
  index : Int
  priority : Double
  probability : Double
  weight : Double
  transition : Transition[S, A]
}

///|
/// Basic proportional prioritized replay.
pub struct PrioritizedReplayBuffer[S, A] {
  capacity : Int
  mut len : Int
  mut head : Int
  data : Array[Transition[S, A]?]
  priorities : Array[Double]
  mut max_priority : Double
  alpha : Double
  beta : Double
  epsilon : Double
}

///|
pub fn[S, A] PrioritizedReplayBuffer::new(
  capacity : Int,
) -> PrioritizedReplayBuffer[S, A] {
  PrioritizedReplayBuffer::with_params(capacity, 0.6, 0.4, 1.0e-6)
}

///|
pub fn[S, A] PrioritizedReplayBuffer::with_params(
  capacity : Int,
  alpha : Double,
  beta : Double,
  epsilon : Double,
) -> PrioritizedReplayBuffer[S, A] {
  let data : Array[Transition[S, A]?] = []
  let priorities : Array[Double] = []
  let mut i = 0
  while i < capacity {
    data.push(None)
    priorities.push(0.0)
    i = i + 1
  }
  {
    capacity,
    len: 0,
    head: 0,
    data,
    priorities,
    max_priority: 1.0,
    alpha,
    beta,
    epsilon,
  }
}

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

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

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

///|
pub fn[S, A] PrioritizedReplayBuffer::priority_sum(
  self : PrioritizedReplayBuffer[S, A],
) -> Double {
  let mut total = 0.0
  let mut i = 0
  while i < self.len {
    let slot = (self.head + i) % self.capacity
    total = total + self.priorities[slot]
    i = i + 1
  }
  total
}

///|
pub fn[S, A] PrioritizedReplayBuffer::priority_max(
  self : PrioritizedReplayBuffer[S, A],
) -> Double {
  self.max_priority
}

///|
pub fn[S, A] PrioritizedReplayBuffer::clear(
  self : PrioritizedReplayBuffer[S, A],
) -> PrioritizedReplayBuffer[S, A] {
  let mut i = 0
  while i < self.data.length() {
    self.data[i] = None
    self.priorities[i] = 0.0
    i = i + 1
  }
  { ..self, len: 0, head: 0, max_priority: 1.0 }
}

///|
pub fn[S, A] PrioritizedReplayBuffer::push(
  self : PrioritizedReplayBuffer[S, A],
  transition : Transition[S, A],
  priority? : Double,
) -> Int {
  if self.capacity == 0 {
    return 0
  }

  let raw_priority = sanitize_priority(priority.unwrap_or(self.max_priority))
  let mut old_priority = 0.0
  if self.len == self.capacity {
    old_priority = self.priorities[self.head]
  }

  if self.len < self.capacity {
    self.data[self.len] = Some(transition)
    self.priorities[self.len] = raw_priority
    self.len = self.len + 1
  } else {
    let slot = self.head
    self.data[slot] = Some(transition)
    self.priorities[slot] = raw_priority
    self.head = (self.head + 1) % self.capacity
  }

  if priority is Some(_) {
    self.max_priority = self.recompute_priority_max()
  } else if raw_priority > self.max_priority {
    self.max_priority = raw_priority
  } else if self.len == self.capacity &&
    old_priority >= self.max_priority &&
    raw_priority < old_priority {
    self.max_priority = self.recompute_priority_max()
  }

  self.len
}

///|
pub fn[S, A] PrioritizedReplayBuffer::extend_episode(
  self : PrioritizedReplayBuffer[S, A],
  episode : Episode[S, A],
) -> Int {
  for transition in episode.transitions() {
    let _ = self.push(transition)
  }
  self.len
}

///|
pub fn[S, A] PrioritizedReplayBuffer::get(
  self : PrioritizedReplayBuffer[S, A],
  logical_index : Int,
) -> (Transition[S, A], Double)? {
  if logical_index < 0 || logical_index >= self.len || self.capacity == 0 {
    None
  } else {
    let slot = (self.head + logical_index) % self.capacity
    match self.data[slot] {
      Some(transition) => Some((transition, self.priorities[slot]))
      None => None
    }
  }
}

///|
pub fn[S, A] PrioritizedReplayBuffer::update_priority(
  self : PrioritizedReplayBuffer[S, A],
  logical_index : Int,
  priority : Double,
) -> Bool {
  if logical_index < 0 || logical_index >= self.len {
    return false
  }

  let slot = (self.head + logical_index) % self.capacity
  let old_priority = self.priorities[slot]
  let new_priority = sanitize_priority(priority)
  self.priorities[slot] = new_priority
  if new_priority > self.max_priority {
    self.max_priority = new_priority
  } else if old_priority >= self.max_priority && new_priority < old_priority {
    self.max_priority = self.recompute_priority_max()
  }
  true
}

///|
/// Priorities are non-negative masses. Clamping here keeps max-priority
/// bookkeeping and sampling probabilities consistent for invalid caller input.
fn sanitize_priority(priority : Double) -> Double {
  if priority < 0.0 {
    0.0
  } else {
    priority
  }
}

///|
fn[S, A] PrioritizedReplayBuffer::recompute_priority_max(
  self : PrioritizedReplayBuffer[S, A],
) -> Double {
  let mut maximum = 0.0
  let mut i = 0
  while i < self.len {
    let slot = (self.head + i) % self.capacity
    if self.priorities[slot] > maximum {
      maximum = self.priorities[slot]
    }
    i = i + 1
  }
  maximum
}

///|
pub fn[S, A] PrioritizedReplayBuffer::sample_batch(
  self : PrioritizedReplayBuffer[S, A],
  batch_size : Int,
  seed : Int,
) -> Array[PrioritizedSample[S, A]] {
  let target = if batch_size < self.len { batch_size } else { self.len }
  let samples : Array[PrioritizedSample[S, A]] = []
  if target == 0 {
    return samples
  }

  let effective : Array[Double] = []
  let mut total = 0.0
  let mut i = 0
  while i < self.len {
    let slot = (self.head + i) % self.capacity
    let mass = @math.pow(
      normalize_positive(self.priorities[slot]) + self.epsilon,
      self.alpha,
    )
    effective.push(mass)
    total = total + mass
    i = i + 1
  }

  if total <= 0.0 {
    total = self.len.to_double()
    let mut j = 0
    while j < effective.length() {
      effective[j] = 1.0
      j = j + 1
    }
  }

  let rng = ReplayRng::new(seed)
  let mut sample_idx = 0
  while sample_idx < target {
    let threshold = rng.next_unit() * total
    let mut cumulative = 0.0
    let mut selected = 0
    let mut k = 0
    while k < effective.length() {
      cumulative = cumulative + effective[k]
      if cumulative >= threshold {
        selected = k
        break
      }
      k = k + 1
      selected = k
    }

    match self.get(selected) {
      Some((transition, raw_priority)) => {
        let probability = effective[selected] / total
        let raw_weight = if probability <= 0.0 {
          0.0
        } else {
          @math.pow(self.len.to_double() * probability, -self.beta)
        }
        samples.push({
          index: selected,
          priority: raw_priority,
          probability,
          weight: raw_weight,
          transition,
        })
      }
      None => ()
    }

    sample_idx = sample_idx + 1
  }

  let mut max_weight = 0.0
  let mut n = 0
  while n < samples.length() {
    let weight = samples[n].weight
    if weight > max_weight {
      max_weight = weight
    }
    n = n + 1
  }

  if max_weight > 0.0 {
    let mut m = 0
    while m < samples.length() {
      samples[m] = { ..samples[m], weight: samples[m].weight / max_weight }
      m = m + 1
    }
  }

  samples
}

///|
pub fn[S, A] PrioritizedReplayBuffer::to_array(
  self : PrioritizedReplayBuffer[S, A],
) -> Array[Transition[S, A]] {
  let result : Array[Transition[S, A]] = []
  let mut i = 0
  while i < self.len {
    match self.get(i) {
      Some((transition, _)) => result.push(transition)
      None => ()
    }
    i = i + 1
  }
  result
}