//! N-Step Experience Buffer for Multi-Step Temporal Difference Learning
//!
//! This module implements n-step bootstrapping for DQN, which improves credit
//! assignment by using n-step returns instead of 1-step TD targets.
//!
//! ## Algorithm
//!
//! N-step return:
//! ```text
//! G_t^(n) = R_t + γR_{t+1} + γ²R_{t+2} + ... + γ^{n-1}R_{t+n-1} + γ^n V(S_{t+n})
//! ```
//!
//! ## Benefits
//!
//! - **Faster credit assignment**: Rewards propagate faster (n steps vs 1)
//! - **Better sample efficiency**: +15-25% improvement typical
//! - **Reduced bias-variance tradeoff**: n=3-5 optimal for most tasks
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml::dqn::nstep_buffer::NStepBuffer;
//! use ml::dqn::Experience;
//!
//! let mut buffer = NStepBuffer::new(3, 0.99);
//! 
//! // Add experiences - returns n-step experience when buffer is full
//! for exp in experiences {
//!     if let Some(nstep_exp) = buffer.add(exp) {
//!         // Use n-step experience for training
//!         replay_buffer.push(nstep_exp);
//!     }
//! }
//! 
//! // Flush remaining experiences at episode end
//! for exp in buffer.flush() {
//!     replay_buffer.push(exp);
//! }
//! ```

use std::collections::VecDeque;
use crate::Experience;

/// N-step experience buffer for multi-step temporal difference learning
///
/// Accumulates experiences over n steps and computes n-step returns using
/// gamma discounting. Returns transformed experiences with bootstrapped
/// rewards for improved credit assignment.
#[derive(Debug, Clone)]
pub struct NStepBuffer {
    /// Number of steps to accumulate (1-10, typical: 3)
    n: usize,
    /// Discount factor (gamma) for future rewards
    gamma: f64,
    /// Buffer of recent experiences (max size: n)
    buffer: VecDeque<Experience>,
}

impl NStepBuffer {
    /// Create new n-step buffer
    ///
    /// # Arguments
    ///
    /// * `n` - Number of steps (1-10, recommended: 3-5)
    /// * `gamma` - Discount factor (0.95-0.99 typical)
    ///
    /// # Panics
    ///
    /// Panics if n < 1 or n > 10 (sanity check)
    pub fn new(n: usize, gamma: f64) -> Self {
        assert!((1..=10).contains(&n), "n must be in range 1-10 (got {})", n);
        assert!(gamma > 0.0 && gamma <= 1.0, "gamma must be in (0, 1] (got {})", gamma);
        
        Self {
            n,
            gamma,
            buffer: VecDeque::with_capacity(n),
        }
    }

    /// Add transition and return n-step experience if ready
    ///
    /// # Algorithm
    ///
    /// 1. Add new experience to buffer
    /// 2. If buffer has n experiences:
    ///    - Pop oldest experience
    ///    - Compute n-step return: `R_t` + γR_{t+1} + ... + γ^{n-1}R_{t+n-1}
    ///    - Create n-step experience with accumulated reward and `next_state` from newest experience
    /// 3. Return `Some(n_step_experience)` if ready, None if still accumulating
    ///
    /// # Returns
    ///
    /// * `Some(Experience)` - N-step experience ready for training
    /// * `None` - Still accumulating (buffer size < n)
    pub fn add(&mut self, exp: Experience) -> Option<Experience> {
        self.buffer.push_back(exp);

        if self.buffer.len() < self.n {
            return None;  // Not enough steps yet
        }

        // We now have exactly n experiences - compute n-step return
        let first = self.buffer.pop_front()?;

        // Special case: n=1 (no multi-step accumulation)
        // After pop_front, buffer is empty, so use first experience's next_state
        if self.n == 1 {
            return Some(Experience {
                state: first.state,
                action: first.action,
                reward: first.reward,
                next_state: first.next_state,
                done: first.done,
                timestamp: first.timestamp,
            });
        }

        // Compute n-step return: R_t + γR_{t+1} + ... + γ^{n-1}R_{t+n-1}
        // Note: rewards are i32 (fixed-point), so we use f64 for discounting
        let mut n_step_reward_f64 = first.reward as f64;
        let mut discount = self.gamma;

        for exp in self.buffer.iter() {
            n_step_reward_f64 += discount * exp.reward as f64;
            discount *= self.gamma;
        }

        let n_step_reward = n_step_reward_f64.round() as i32;

        // Get next state and done flag from newest experience (n steps ahead)
        let newest = self.buffer.back()?;

        // Create n-step experience
        Some(Experience {
            state: first.state,
            action: first.action,
            reward: n_step_reward,
            next_state: newest.next_state.clone(),
            done: newest.done,
            timestamp: first.timestamp,
        })
    }

    /// Flush remaining experiences at episode end
    ///
    /// When an episode terminates, we still have 1 to n-1 experiences in the buffer
    /// that haven't been processed yet. This method computes truncated n-step returns
    /// for these experiences (using whatever steps are available).
    ///
    /// # Returns
    ///
    /// Vector of remaining experiences with truncated n-step returns
    pub fn flush(&mut self) -> Vec<Experience> {
        let mut result = Vec::new();
        
        while let Some(first) = self.buffer.pop_front() {

            // Compute truncated n-step return (fewer than n steps)
            let mut n_step_reward_f64 = first.reward as f64;
            let mut discount = self.gamma;

            for exp in self.buffer.iter() {
                n_step_reward_f64 += discount * exp.reward as f64;
                discount *= self.gamma;
            }

            let n_step_reward = n_step_reward_f64.round() as i32;
            
            // Get next state and done flag from newest experience (or current if buffer empty)
            let (next_state, done) = if let Some(newest) = self.buffer.back() {
                (newest.next_state.clone(), newest.done)
            } else {
                (first.next_state.clone(), first.done)
            };
            
            result.push(Experience {
                state: first.state,
                action: first.action,
                reward: n_step_reward,
                next_state,
                done,
                timestamp: first.timestamp,
            });
        }
        
        result
    }

    /// Clear the buffer (called at start of new episode)
    pub fn clear(&mut self) {
        self.buffer.clear();
    }

    /// Get current buffer size
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

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

    /// Get n-step parameter
    pub const fn n_steps(&self) -> usize {
        self.n
    }

    /// Get gamma parameter
    pub const fn gamma(&self) -> f64 {
        self.gamma
    }
}

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

    fn create_test_experience(state_val: f32, action: u8, reward: f32, done: bool) -> Experience {
        Experience {
            state: vec![state_val; 32],
            action,
            reward: (reward * 100.0).round() as i32,  // Convert to fixed-point (100x scale)
            next_state: vec![state_val + 1.0; 32],
            done,
            timestamp: chrono::Utc::now().timestamp_nanos_opt().unwrap() as u64,
        }
    }

    #[test]
    fn test_nstep_buffer_creation() {
        let buffer = NStepBuffer::new(3, 0.99);
        assert_eq!(buffer.n_steps(), 3);
        assert_eq!(buffer.gamma(), 0.99);
        assert!(buffer.is_empty());
    }

    #[test]
    #[should_panic(expected = "n must be in range 1-10")]
    fn test_invalid_n_too_small() {
        NStepBuffer::new(0, 0.99);
    }

    #[test]
    #[should_panic(expected = "n must be in range 1-10")]
    fn test_invalid_n_too_large() {
        NStepBuffer::new(11, 0.99);
    }

    #[test]
    #[should_panic(expected = "gamma must be in (0, 1]")]
    fn test_invalid_gamma() {
        NStepBuffer::new(3, 1.5);
    }

    #[test]
    fn test_n1_returns_immediately() {
        // n=1 should return experience immediately (no accumulation)
        let mut buffer = NStepBuffer::new(1, 0.99);
        
        let exp1 = create_test_experience(1.0, 0, 1.0, false);
        let result = buffer.add(exp1.clone());
        
        assert!(result.is_some());
        let nstep = result.unwrap();
        assert_eq!(nstep.reward, exp1.reward);
        assert_eq!(nstep.state, exp1.state);
    }

    #[test]
    fn test_n3_accumulation() {
        // n=3 should accumulate 3 experiences before returning
        let mut buffer = NStepBuffer::new(3, 0.99);
        
        let exp1 = create_test_experience(1.0, 0, 1.0, false);
        let exp2 = create_test_experience(2.0, 1, 2.0, false);
        let exp3 = create_test_experience(3.0, 0, 3.0, false);
        
        // First two add() calls should return None
        assert!(buffer.add(exp1.clone()).is_none());
        assert!(buffer.add(exp2.clone()).is_none());
        
        // Third add() should return n-step experience
        let result = buffer.add(exp3.clone());
        assert!(result.is_some());
        
        let nstep = result.unwrap();
        
        // Verify n-step return: R_1 + γR_2 + γ²R_3 = 1.0 + 0.99*2.0 + 0.99²*3.0
        // In fixed-point (×100): 100 + 0.99*200 + 0.99²*300 = 100 + 198 + 293.91 ≈ 592
        let expected_reward_f64: f64 = 1.0 + 0.99*2.0 + 0.99*0.99*3.0;
        let expected_reward = (expected_reward_f64 * 100.0).round() as i32;
        assert_eq!(nstep.reward, expected_reward);
        
        // Verify state is from first experience
        assert_eq!(nstep.state, exp1.state);
        
        // Verify next_state is from last experience
        assert_eq!(nstep.next_state, exp3.next_state);
    }

    #[test]
    fn test_gamma_discounting() {
        let mut buffer = NStepBuffer::new(3, 0.5);  // γ=0.5 for easy verification
        
        let exp1 = create_test_experience(1.0, 0, 4.0, false);
        let exp2 = create_test_experience(2.0, 1, 2.0, false);
        let exp3 = create_test_experience(3.0, 0, 1.0, false);
        
        buffer.add(exp1);
        buffer.add(exp2);
        let result = buffer.add(exp3);
        
        assert!(result.is_some());
        let nstep = result.unwrap();
        
        // Verify: 4.0 + 0.5*2.0 + 0.25*1.0 = 4.0 + 1.0 + 0.25 = 5.25
        // In fixed-point (×100): 400 + 0.5*200 + 0.25*100 = 400 + 100 + 25 = 525
        let expected_reward_f64: f64 = 4.0 + 0.5*2.0 + 0.25*1.0;
        let expected_reward = (expected_reward_f64 * 100.0).round() as i32;
        assert_eq!(nstep.reward, expected_reward);
    }

    #[test]
    fn test_flush_with_remaining_experiences() {
        let mut buffer = NStepBuffer::new(3, 0.99);
        
        let exp1 = create_test_experience(1.0, 0, 1.0, false);
        let exp2 = create_test_experience(2.0, 1, 2.0, false);
        
        buffer.add(exp1.clone());
        buffer.add(exp2.clone());
        
        // Flush should return truncated n-step experiences
        let flushed = buffer.flush();
        
        assert_eq!(flushed.len(), 2);
        
        // First flushed experience: R_1 + γR_2 = 1.0 + 0.99*2.0
        // In fixed-point (×100): 100 + 0.99*200 = 100 + 198 = 298
        let expected_reward1_f64: f64 = 1.0 + 0.99*2.0;
        let expected_reward1 = (expected_reward1_f64 * 100.0).round() as i32;
        assert_eq!(flushed[0].reward, expected_reward1);
        assert_eq!(flushed[0].state, exp1.state);
        assert_eq!(flushed[0].next_state, exp2.next_state);

        // Second flushed experience: just R_2 = 2.0 × 100 = 200
        assert_eq!(flushed[1].reward, exp2.reward);
        assert_eq!(flushed[1].state, exp2.state);
        
        // Buffer should be empty after flush
        assert!(buffer.is_empty());
    }

    #[test]
    fn test_clear() {
        let mut buffer = NStepBuffer::new(3, 0.99);
        
        buffer.add(create_test_experience(1.0, 0, 1.0, false));
        buffer.add(create_test_experience(2.0, 1, 2.0, false));
        
        assert_eq!(buffer.len(), 2);
        
        buffer.clear();
        assert!(buffer.is_empty());
    }

    #[test]
    fn test_done_flag_propagation() {
        let mut buffer = NStepBuffer::new(3, 0.99);
        
        let exp1 = create_test_experience(1.0, 0, 1.0, false);
        let exp2 = create_test_experience(2.0, 1, 2.0, false);
        let exp3 = create_test_experience(3.0, 0, 3.0, true);  // Terminal state
        
        buffer.add(exp1);
        buffer.add(exp2);
        let result = buffer.add(exp3);
        
        assert!(result.is_some());
        let nstep = result.unwrap();
        
        // Done flag should be from last experience
        assert!(nstep.done);
    }

    #[test]
    fn test_continuous_streaming() {
        // Test that buffer continues to work after first n-step return
        let mut buffer = NStepBuffer::new(3, 0.99);
        
        // Add 5 experiences, should get 3 n-step returns
        let experiences = vec![
            create_test_experience(1.0, 0, 1.0, false),
            create_test_experience(2.0, 1, 2.0, false),
            create_test_experience(3.0, 0, 3.0, false),
            create_test_experience(4.0, 1, 4.0, false),
            create_test_experience(5.0, 0, 5.0, false),
        ];
        
        let mut nstep_count = 0;
        for exp in experiences {
            if buffer.add(exp).is_some() {
                nstep_count += 1;
            }
        }
        
        assert_eq!(nstep_count, 3);
        assert_eq!(buffer.len(), 2);  // 2 experiences remaining in buffer
    }
}
