//! Experience replay data structures

use std::time::{SystemTime, UNIX_EPOCH};

// CANONICAL TYPE IMPORTS - Use common::Decimal
use serde::{Deserialize, Serialize};

/// Experience tuple for `DQN` replay buffer
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Experience {
    /// Current state representation
    pub state: Vec<f32>,
    /// Action taken (as integer index)
    pub action: u8,
    /// Reward received (scaled to fixed-point)
    pub reward: i32,
    /// Next state representation
    pub next_state: Vec<f32>,
    /// Whether this was a terminal state
    pub done: bool,
    /// Experience timestamp
    pub timestamp: u64,
}

impl Experience {
    /// Create a new experience
    pub fn new(state: Vec<f32>, action: u8, reward: f32, next_state: Vec<f32>, done: bool) -> Self {
        Self {
            state,
            action,
            reward: (reward * 1_000_000.0) as i32, // Scale to fixed-point (100x more precision for DSR values)
            next_state,
            done,
            timestamp: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos() as u64,
        }
    }

    /// Get reward as f32
    pub fn reward_f32(&self) -> f32 {
        self.reward as f32 / 1_000_000.0
    }

    /// Check if experience is valid
    pub fn is_valid(&self) -> bool {
        !self.state.is_empty()
            && !self.next_state.is_empty()
            && self.state.len() == self.next_state.len()
    }
}

/// Batch of experiences for training
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExperienceBatch {
    /// Batch of experiences
    pub experiences: Vec<Experience>,
    /// Number of experiences in batch
    pub batch_size: usize,
}

impl ExperienceBatch {
    /// Create a new batch from experiences
    pub fn new(experiences: Vec<Experience>) -> Self {
        let batch_size = experiences.len();
        Self {
            experiences,
            batch_size,
        }
    }

    /// Create empty batch
    pub const fn empty() -> Self {
        Self {
            experiences: Vec::new(),
            batch_size: 0,
        }
    }

    /// Check if batch is valid
    pub fn is_valid(&self) -> bool {
        self.batch_size == self.experiences.len() && self.experiences.iter().all(|e| e.is_valid())
    }

    /// Convert batch to tensor format for training
    pub fn to_tensors(&self) -> (Vec<Vec<f32>>, Vec<u8>, Vec<f32>, Vec<Vec<f32>>, Vec<bool>) {
        let states = self.experiences.iter().map(|e| e.state.clone()).collect();
        let actions = self.experiences.iter().map(|e| e.action).collect();
        let rewards = self.experiences.iter().map(|e| e.reward_f32()).collect();
        let next_states = self
            .experiences
            .iter()
            .map(|e| e.next_state.clone())
            .collect();
        let dones = self.experiences.iter().map(|e| e.done).collect();

        (states, actions, rewards, next_states, dones)
    }

    /// Add experience to batch
    pub fn add(&mut self, experience: Experience) {
        self.experiences.push(experience);
        self.batch_size = self.experiences.len();
    }

    /// Get batch size
    pub const fn len(&self) -> usize {
        self.batch_size
    }

    /// Check if batch is empty
    pub const fn is_empty(&self) -> bool {
        self.batch_size == 0
    }
}

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

    #[test]
    fn test_experience_creation() {
        let state = vec![1.0, 2.0, 3.0];
        let next_state = vec![1.1, 2.1, 3.1];
        let experience = Experience::new(state.clone(), 2, 0.5, next_state.clone(), false);

        assert_eq!(experience.state, state);
        assert_eq!(experience.action, 2);
        assert_eq!(experience.reward, 500000); // 0.5 * 1_000_000
        assert_eq!(experience.next_state, next_state);
        assert!(!experience.done);
        assert!(experience.is_valid());
    }

    #[test]
    fn test_experience_batch() {
        let experiences = vec![
            Experience::new(vec![1.0, 2.0], 0, 0.1, vec![1.1, 2.1], false),
            Experience::new(vec![2.0, 3.0], 1, 0.2, vec![2.1, 3.1], false),
        ];

        let batch = ExperienceBatch::new(experiences);
        assert_eq!(batch.batch_size, 2);
        assert!(batch.is_valid());

        let (states, actions, rewards, _next_states, _dones) = batch.to_tensors();
        assert_eq!(states.len(), 2);
        assert_eq!(actions, vec![0, 1]);
        assert_eq!(rewards, vec![0.1, 0.2]);
    }
}
