// ddpg_update.mbt — DDPG training-time helpers: continuous replay buffer
// and per-step gradient update.

///|
/// Continuous replay buffer. Stores float-valued states/actions/rewards
/// rather than discrete indexes. Capacity-overflow is FIFO.
pub struct ContinuousReplayBuffer {
  capacity : Int
  state_dim : Int
  action_dim : Int
  states : Array[Float]
  actions : Array[Float]
  rewards : Array[Float]
  next_states : Array[Float]
  dones : Array[Float]
  mut size : Int
  mut cursor : Int
}

///|
/// Construct a new continuous replay buffer.
pub fn ContinuousReplayBuffer::new(
  capacity : Int,
  state_dim : Int,
  action_dim : Int,
) -> ContinuousReplayBuffer {
  let n = capacity
  let s_len = n * state_dim
  let a_len = n * action_dim
  {
    capacity: n,
    state_dim,
    action_dim,
    states: Array::make(s_len, 0.0F),
    actions: Array::make(a_len, 0.0F),
    rewards: Array::make(n, 0.0F),
    next_states: Array::make(s_len, 0.0F),
    dones: Array::make(n, 0.0F),
    size: 0,
    cursor: 0,
  }
}

///|
/// Number of transitions currently stored.
pub fn ContinuousReplayBuffer::len(self : ContinuousReplayBuffer) -> Int {
  self.size
}

///|
/// Append a transition. Overwrites the oldest slot when full.
pub fn ContinuousReplayBuffer::push(
  self : ContinuousReplayBuffer,
  state : Array[Float],
  action : Array[Float],
  reward : Float,
  next_state : Array[Float],
  done : Bool,
) -> Unit {
  let i = self.cursor
  let s_off = i * self.state_dim
  let a_off = i * self.action_dim
  for k in 0..= self.capacity {
    self.cursor = 0
  }
  if self.size < self.capacity {
    self.size = self.size + 1
  }
}

///|
/// Sample a mini-batch. Returns 5 parallel arrays (states, actions,
/// rewards, next_states, dones). The caller supplies its own RNG.
pub fn ContinuousReplayBuffer::sample(
  self : ContinuousReplayBuffer,
  batch_size : Int,
  rng : Xoshiro,
) -> (Array[Float], Array[Float], Array[Float], Array[Float], Array[Float]) {
  let states : Array[Float] = Array::make(batch_size * self.state_dim, 0.0F)
  let actions : Array[Float] = Array::make(batch_size * self.action_dim, 0.0F)
  let rewards : Array[Float] = Array::make(batch_size, 0.0F)
  let next_states : Array[Float] = Array::make(batch_size * self.state_dim, 0.0F)
  let dones : Array[Float] = Array::make(batch_size, 0.0F)
  for k in 0..= self.size {
      self.size - 1
    } else {
      idx
    }
    let s_off = safe * self.state_dim
    let a_off = safe * self.action_dim
    let b_s_off = k * self.state_dim
    let b_a_off = k * self.action_dim
    for j in 0.. Float {
  let (states_b, actions_b, rewards_b, next_states_b, dones_b) = batch
  let batch_size = rewards_b.length()
  let s_dim = ddpg.actor.state_dim
  let a_dim = ddpg.actor.action_dim
  let hidden_dim = ddpg.actor.hidden
  let mut total_abs_td = 0.0F
  let hidden_q1 : Array[Float] = Array::make(hidden_dim, 0.0F)
  let hidden_q2 : Array[Float] = Array::make(hidden_dim, 0.0F)
  let hidden_q1_t : Array[Float] = Array::make(hidden_dim, 0.0F)
  let hidden_q2_t : Array[Float] = Array::make(hidden_dim, 0.0F)
  let hidden_actor_t : Array[Float] = Array::make(hidden_dim, 0.0F)
  for k in 0..