//! ExO-PPO (Extended Off-Policy PPO) trajectory replay
//!
//! Stores past M rollouts in a FIFO buffer for 4x sample efficiency.
//! Uses importance-weighted clipping with exponential attenuation
//! outside clip bounds (smoother than hard clipping).
//!
//! Reference: ExO-PPO (2026) -- <https://arxiv.org/abs/2602.09726>

use serde::{Deserialize, Serialize};
use std::collections::VecDeque;

/// A stored rollout with pre-computed advantages and old log-probs.
///
/// Advantages are computed ONCE per rollout (via GAE) and stored here,
/// avoiding expensive recomputation when replaying off-policy data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredRollout {
    /// State features for each timestep, flattened `[T * state_dim]`
    pub states: Vec<f32>,
    /// Actions taken at each timestep
    pub actions: Vec<u32>,
    /// Log-probabilities under the OLD policy that generated this rollout
    pub old_log_probs: Vec<f32>,
    /// Pre-computed GAE advantages (computed once, reused across replay iterations)
    pub advantages: Vec<f32>,
    /// Pre-computed discounted returns (targets for value function)
    pub returns: Vec<f32>,
    /// State dimension (number of features per timestep)
    pub state_dim: usize,
    /// Number of timesteps in this rollout
    pub num_steps: usize,
    /// Policy generation (monotonically increasing counter)
    pub generation: u64,
}

impl StoredRollout {
    /// Get state slice at timestep `t`.
    ///
    /// Returns `None` if `t` is out of bounds or arithmetic overflows.
    pub fn get_state(&self, t: usize) -> Option<&[f32]> {
        let start = t.checked_mul(self.state_dim)?;
        let end = start.checked_add(self.state_dim)?;
        self.states.get(start..end)
    }
}

/// FIFO buffer of past rollouts for ExO-PPO.
///
/// On each training iteration the agent samples from both the fresh on-policy
/// rollout AND the M stored off-policy rollouts. Importance sampling ratios
/// correct for distributional shift, with exponential attenuation outside clip
/// bounds (more stable than hard clipping).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrajectoryReplayBuffer {
    /// Buffer of past rollouts (FIFO, max size M)
    buffer: VecDeque<StoredRollout>,
    /// Maximum number of stored rollouts (M=4 by default)
    max_rollouts: usize,
    /// Current policy generation counter
    generation: u64,
    /// Clip epsilon for importance weighting
    clip_epsilon: f32,
    /// Exponential attenuation rate (alpha=5 from paper)
    attenuation_alpha: f32,
}

impl TrajectoryReplayBuffer {
    /// Create a new replay buffer with capacity `max_rollouts`.
    ///
    /// # Arguments
    /// * `max_rollouts` - Maximum stored rollouts (M=4 in paper)
    /// * `clip_epsilon` - PPO clip range epsilon (typically 0.2)
    pub fn new(max_rollouts: usize, clip_epsilon: f32) -> Self {
        Self {
            buffer: VecDeque::with_capacity(max_rollouts),
            max_rollouts,
            generation: 0,
            clip_epsilon,
            attenuation_alpha: 5.0, // Paper default
        }
    }

    /// Store a new rollout, evicting the oldest if the buffer is full (FIFO).
    pub fn store_rollout(&mut self, rollout: StoredRollout) {
        if self.buffer.len() >= self.max_rollouts {
            self.buffer.pop_front(); // Evict oldest
        }
        self.buffer.push_back(rollout);
        self.generation += 1;
    }

    /// Get current generation number.
    pub const fn generation(&self) -> u64 {
        self.generation
    }

    /// Number of stored rollouts.
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// Whether buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Get all stored rollouts for training.
    pub fn rollouts(&self) -> impl Iterator<Item = &StoredRollout> {
        self.buffer.iter()
    }

    /// Compute importance-weighted surrogate loss for a single off-policy sample.
    ///
    /// Uses exponential attenuation outside clip bounds:
    /// - Inside `[1-eps, 1+eps]`: linear passthrough (same as standard PPO)
    /// - Outside bounds: exponentially attenuated (decay rate alpha)
    ///
    /// This is smoother than hard clipping and prevents the gradient from
    /// being completely zeroed for important but off-policy samples.
    pub fn compute_is_weight(&self, new_log_prob: f32, old_log_prob: f32, advantage: f32) -> f32 {
        let ratio = (new_log_prob - old_log_prob).exp();
        let clip_low = 1.0 - self.clip_epsilon;
        let clip_high = 1.0 + self.clip_epsilon;

        if ratio >= clip_low && ratio <= clip_high {
            // Inside clip range: standard surrogate
            ratio * advantage
        } else if ratio < clip_low {
            // Below clip: exponential attenuation
            let distance = clip_low - ratio;
            let attenuated = clip_low * (-self.attenuation_alpha * distance).exp();
            attenuated * advantage
        } else {
            // Above clip: exponential attenuation
            let distance = ratio - clip_high;
            let attenuated = clip_high * (-self.attenuation_alpha * distance).exp();
            attenuated * advantage
        }
    }

    /// Batch compute IS-weighted surrogate for a full stored rollout.
    ///
    /// # Arguments
    /// * `rollout` - The stored off-policy rollout
    /// * `new_log_probs` - Log-probs under the current policy for all timesteps
    ///
    /// # Returns
    /// Mean IS-weighted surrogate loss for this rollout (0.0 if empty).
    pub fn compute_rollout_loss(&self, rollout: &StoredRollout, new_log_probs: &[f32]) -> f32 {
        if rollout.num_steps == 0 {
            return 0.0;
        }

        let mut total_loss = 0.0;
        let count = rollout
            .num_steps
            .min(new_log_probs.len())
            .min(rollout.old_log_probs.len());

        if count == 0 {
            return 0.0;
        }

        for i in 0..count {
            let new_lp = new_log_probs.get(i).copied().unwrap_or(0.0);
            let old_lp = rollout.old_log_probs.get(i).copied().unwrap_or(0.0);
            let adv = rollout.advantages.get(i).copied().unwrap_or(0.0);
            total_loss += self.compute_is_weight(new_lp, old_lp, adv);
        }

        total_loss / count as f32
    }

    /// Create a [`StoredRollout`] from training data.
    ///
    /// # Arguments
    /// * `states` - Flattened state features `[T * state_dim]`
    /// * `actions` - Actions per timestep
    /// * `old_log_probs` - Log-probs under the collecting policy
    /// * `advantages` - Pre-computed GAE advantages
    /// * `returns` - Pre-computed discounted returns
    /// * `state_dim` - Number of features per state
    /// * `generation` - Policy generation that produced this rollout
    pub fn create_rollout(
        states: Vec<f32>,
        actions: Vec<u32>,
        old_log_probs: Vec<f32>,
        advantages: Vec<f32>,
        returns: Vec<f32>,
        state_dim: usize,
        generation: u64,
    ) -> StoredRollout {
        let num_steps = actions.len();
        StoredRollout {
            states,
            actions,
            old_log_probs,
            advantages,
            returns,
            state_dim,
            num_steps,
            generation,
        }
    }
}

impl Default for TrajectoryReplayBuffer {
    fn default() -> Self {
        Self::new(4, 0.2) // M=4, clip=0.2
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: create a simple rollout with `n` steps of dimension `dim`.
    fn make_rollout(n: usize, dim: usize, generation: u64) -> StoredRollout {
        TrajectoryReplayBuffer::create_rollout(
            vec![1.0_f32; n * dim],
            vec![0_u32; n],
            vec![-0.5_f32; n], // old log-probs
            vec![1.0_f32; n],  // advantages
            vec![10.0_f32; n], // returns
            dim,
            generation,
        )
    }

    #[test]
    fn test_replay_buffer_new() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        assert_eq!(buf.len(), 0);
        assert!(buf.is_empty());
        assert_eq!(buf.generation(), 0);
    }

    #[test]
    fn test_replay_buffer_store_and_evict() {
        let mut buf = TrajectoryReplayBuffer::new(4, 0.2);

        // Store 5 rollouts into a buffer of size 4
        for gen in 0..5 {
            buf.store_rollout(make_rollout(10, 4, gen));
        }

        // Buffer should contain exactly 4 (the max)
        assert_eq!(buf.len(), 4);
        // Generation counter is 5 (incremented once per store)
        assert_eq!(buf.generation(), 5);

        // Oldest (gen=0) was evicted; first remaining should be gen=1
        let generations: Vec<u64> = buf.rollouts().map(|r| r.generation).collect();
        assert_eq!(generations, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_replay_buffer_is_weight_inside_clip() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        let advantage = 2.0_f32;

        // ratio = exp(new - old) = exp(0) = 1.0 (inside [0.8, 1.2])
        let weight = buf.compute_is_weight(-0.5, -0.5, advantage);
        // Should be ratio * advantage = 1.0 * 2.0 = 2.0
        assert!((weight - 2.0).abs() < 1e-5, "weight = {weight}");

        // ratio = exp(-0.3 - (-0.5)) = exp(0.2) ~= 1.2214
        // This is barely above 1.2, so it goes through the attenuation branch
        // Let's pick a ratio that stays inside: exp(0.1) ~= 1.1052
        let weight2 = buf.compute_is_weight(-0.4, -0.5, advantage);
        let expected_ratio = (0.1_f32).exp(); // ~1.1052
        let expected = expected_ratio * advantage;
        assert!(
            (weight2 - expected).abs() < 1e-4,
            "weight2 = {weight2}, expected = {expected}"
        );
    }

    #[test]
    fn test_replay_buffer_is_weight_below_clip() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        let advantage = 1.0_f32;

        // ratio = exp(-2.0 - 0.0) = exp(-2) ~= 0.1353 (well below 0.8)
        let weight = buf.compute_is_weight(-2.0, 0.0, advantage);
        let ratio = (-2.0_f32).exp();
        let clip_low = 0.8_f32;
        let distance = clip_low - ratio;
        let expected = clip_low * (-5.0 * distance).exp() * advantage;
        assert!(
            (weight - expected).abs() < 1e-5,
            "weight = {weight}, expected = {expected}"
        );

        // Attenuated weight should be less than clip_low * advantage
        assert!(weight < clip_low * advantage);
    }

    #[test]
    fn test_replay_buffer_is_weight_above_clip() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        let advantage = 1.0_f32;

        // ratio = exp(2.0 - 0.0) = exp(2) ~= 7.389 (well above 1.2)
        let weight = buf.compute_is_weight(2.0, 0.0, advantage);
        let ratio = (2.0_f32).exp();
        let clip_high = 1.2_f32;
        let distance = ratio - clip_high;
        let expected = clip_high * (-5.0 * distance).exp() * advantage;
        assert!(
            (weight - expected).abs() < 1e-5,
            "weight = {weight}, expected = {expected}"
        );

        // Attenuated weight should be less than clip_high * advantage
        assert!(weight < clip_high * advantage);
    }

    #[test]
    fn test_replay_buffer_is_weight_attenuation_decays() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        let advantage = 1.0_f32;

        // Two ratios below clip: one closer, one further from 0.8
        // ratio1 = exp(-0.5) ~= 0.6065 (distance from 0.8 = 0.1935)
        let w1 = buf.compute_is_weight(-0.5, 0.0, advantage);
        // ratio2 = exp(-2.0) ~= 0.1353 (distance from 0.8 = 0.6647)
        let w2 = buf.compute_is_weight(-2.0, 0.0, advantage);

        // Further from clip boundary = more attenuation = smaller weight
        assert!(
            w1.abs() > w2.abs(),
            "closer sample ({w1}) should have larger |weight| than further ({w2})"
        );
    }

    #[test]
    fn test_replay_buffer_rollout_loss() {
        let buf = TrajectoryReplayBuffer::new(4, 0.2);
        let rollout = make_rollout(3, 2, 0);

        // new_log_probs same as old => ratio = 1.0 => weight = 1.0 * adv = 1.0
        let new_log_probs = vec![-0.5_f32; 3];
        let loss = buf.compute_rollout_loss(&rollout, &new_log_probs);
        // Each sample: ratio=1.0, advantage=1.0, weight=1.0
        // Mean = 1.0
        assert!((loss - 1.0).abs() < 1e-5, "loss = {loss}, expected ~1.0");
    }

    #[test]
    fn test_replay_buffer_generation_counter() {
        let mut buf = TrajectoryReplayBuffer::new(4, 0.2);
        assert_eq!(buf.generation(), 0);

        buf.store_rollout(make_rollout(5, 2, 0));
        assert_eq!(buf.generation(), 1);

        buf.store_rollout(make_rollout(5, 2, 1));
        assert_eq!(buf.generation(), 2);

        buf.store_rollout(make_rollout(5, 2, 2));
        assert_eq!(buf.generation(), 3);

        // Generation should be monotonically increasing
        assert!(buf.generation() > 0);
    }

    #[test]
    fn test_stored_rollout_get_state() {
        let rollout = make_rollout(3, 4, 0);

        // Get state at timestep 0
        let s0 = rollout.get_state(0);
        assert!(s0.is_some());
        assert_eq!(s0.map(|s| s.len()), Some(4));

        // Get state at timestep 2 (last)
        let s2 = rollout.get_state(2);
        assert!(s2.is_some());
        assert_eq!(s2.map(|s| s.len()), Some(4));

        // Out of bounds
        let s3 = rollout.get_state(3);
        assert!(s3.is_none());

        // Very large index (overflow protection)
        let s_huge = rollout.get_state(usize::MAX);
        assert!(s_huge.is_none());
    }

    #[test]
    fn test_replay_buffer_empty() {
        let buf = TrajectoryReplayBuffer::default();
        assert!(buf.is_empty());
        assert_eq!(buf.len(), 0);
        assert_eq!(buf.generation(), 0);
        assert_eq!(buf.rollouts().count(), 0);

        // Compute loss on an empty rollout
        let empty_rollout =
            TrajectoryReplayBuffer::create_rollout(vec![], vec![], vec![], vec![], vec![], 4, 0);
        assert_eq!(empty_rollout.num_steps, 0);
        let loss = buf.compute_rollout_loss(&empty_rollout, &[]);
        assert!((loss - 0.0).abs() < 1e-10);
    }
}
