foxhunt/crates/ml-dqn/src/experience.rs

//! 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,
    /// Market regime at transition time: 0=Trending, 1=Ranging, 2=Volatile
    pub regime: u8,
}

/// Classify market regime from a state feature vector.
///
/// Uses ADX (index 40) and CUSUM (index 41) to determine the market regime:
/// - 0 = Trending (ADX > 0.25)
/// - 1 = Ranging (default / low ADX, low CUSUM)
/// - 2 = Volatile (|CUSUM| > 0.7)
///
/// Returns 1 (Ranging) as safe default when the state vector is too short.
pub fn classify_regime_from_state(state: &[f32]) -> u8 {
    let adx = if state.len() > 40 { state[40] } else { 0.0 };
    let cusum = if state.len() > 41 { state[41] } else { 0.0 };
    if adx > 0.25 {
        0 // Trending
    } else if cusum.abs() > 0.7 {
        2 // Volatile
    } else {
        1 // Ranging
    }
}

/// Count regime distribution over a batch of state feature vectors.
/// Returns `(trending_pct, ranging_pct, volatile_pct)` as percentages `[0.0, 100.0]`.
pub fn regime_distribution(states: &[Vec<f32>]) -> (f64, f64, f64) {
    if states.is_empty() {
        return (0.0, 100.0, 0.0); // default: all ranging
    }
    let mut counts = [0usize; 3]; // trending, ranging, volatile
    for state in states {
        let regime = classify_regime_from_state(state) as usize;
        if regime < 3 {
            counts[regime] += 1;
        }
    }
    let total = states.len() as f64;
    (
        counts[0] as f64 / total * 100.0,
        counts[1] as f64 / total * 100.0,
        counts[2] as f64 / total * 100.0,
    )
}

/// Count regime distribution over flat state buffer (f32).
/// `states_flat`: `[num_samples * state_dim]`, row-major.
pub fn regime_distribution_flat(states_flat: &[f32], state_dim: usize) -> (f64, f64, f64) {
    if states_flat.is_empty() || state_dim == 0 {
        return (0.0, 100.0, 0.0);
    }
    let num_samples = states_flat.len() / state_dim;
    let mut counts = [0usize; 3];
    for i in 0..num_samples {
        let start = i * state_dim;
        let end = (start + state_dim).min(states_flat.len());
        let state = &states_flat[start..end];
        let regime = classify_regime_from_state(state) as usize;
        if regime < 3 {
            counts[regime] += 1;
        }
    }
    let total = num_samples as f64;
    (
        counts[0] as f64 / total * 100.0,
        counts[1] as f64 / total * 100.0,
        counts[2] as f64 / total * 100.0,
    )
}

/// Count regime distribution over flat state buffer (f64).
/// `states_flat`: `[num_samples * state_dim]`, row-major.
/// Converts each row to f32 on the fly for `classify_regime_from_state`.
pub fn regime_distribution_flat_f64(states_flat: &[f64], state_dim: usize) -> (f64, f64, f64) {
    if states_flat.is_empty() || state_dim == 0 {
        return (0.0, 100.0, 0.0);
    }
    let num_samples = states_flat.len() / state_dim;
    let mut counts = [0usize; 3];
    // Only ADX (idx 40) and CUSUM (idx 41) are read, so convert on the fly
    for i in 0..num_samples {
        let start = i * state_dim;
        let adx = if state_dim > 40 {
            states_flat.get(start + 40).copied().unwrap_or(0.0)
        } else {
            0.0
        };
        let cusum = if state_dim > 41 {
            states_flat.get(start + 41).copied().unwrap_or(0.0)
        } else {
            0.0
        };
        let regime = if adx > 0.25 {
            0 // Trending
        } else if cusum.abs() > 0.7 {
            2 // Volatile
        } else {
            1 // Ranging
        };
        counts[regime] += 1;
    }
    let total = num_samples as f64;
    (
        counts[0] as f64 / total * 100.0,
        counts[1] as f64 / total * 100.0,
        counts[2] as f64 / total * 100.0,
    )
}

impl Experience {
    /// Create a new experience
    pub fn new(state: Vec<f32>, action: u8, reward: f32, next_state: Vec<f32>, done: bool) -> Self {
        let regime = classify_regime_from_state(&state);
        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,
            regime,
        }
    }

    /// 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]);
    }

    #[test]
    fn test_regime_distribution_empty() {
        let (t, r, v) = regime_distribution(&[]);
        assert!((t - 0.0).abs() < f64::EPSILON);
        assert!((r - 100.0).abs() < f64::EPSILON);
        assert!((v - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_regime_distribution_all_trending() {
        // ADX at index 40 > 0.25 => Trending
        let state = {
            let mut s = vec![0.0_f32; 42];
            s[40] = 0.5; // high ADX
            s
        };
        let states = vec![state.clone(), state.clone(), state];
        let (t, r, v) = regime_distribution(&states);
        assert!((t - 100.0).abs() < f64::EPSILON);
        assert!((r - 0.0).abs() < f64::EPSILON);
        assert!((v - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_regime_distribution_mixed() {
        let trending = {
            let mut s = vec![0.0_f32; 42];
            s[40] = 0.5; // ADX > 0.25
            s
        };
        let ranging = vec![0.0_f32; 42]; // low ADX, low CUSUM
        let volatile = {
            let mut s = vec![0.0_f32; 42];
            s[41] = 0.9; // |CUSUM| > 0.7
            s
        };
        let states = vec![trending, ranging.clone(), ranging, volatile];
        let (t, r, v) = regime_distribution(&states);
        assert!((t - 25.0).abs() < f64::EPSILON);
        assert!((r - 50.0).abs() < f64::EPSILON);
        assert!((v - 25.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_regime_distribution_flat_basic() {
        // 2 samples, state_dim=42
        let mut flat = vec![0.0_f32; 84];
        flat[40] = 0.5; // sample 0: Trending (ADX > 0.25)
        // sample 1: all zeros => Ranging
        let (t, r, v) = regime_distribution_flat(&flat, 42);
        assert!((t - 50.0).abs() < f64::EPSILON);
        assert!((r - 50.0).abs() < f64::EPSILON);
        assert!((v - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_regime_distribution_flat_f64_basic() {
        // 3 samples, state_dim=42
        let mut flat = vec![0.0_f64; 126];
        flat[40] = 0.5;          // sample 0: Trending
        // sample 1: all zeros => Ranging
        flat[2 * 42 + 41] = 0.9; // sample 2: Volatile
        let (t, r, v) = regime_distribution_flat_f64(&flat, 42);
        assert!((t - 100.0 / 3.0).abs() < 0.1);
        assert!((r - 100.0 / 3.0).abs() < 0.1);
        assert!((v - 100.0 / 3.0).abs() < 0.1);
    }

    #[test]
    fn test_regime_distribution_flat_empty() {
        let (t, r, v) = regime_distribution_flat(&[], 42);
        assert!((t - 0.0).abs() < f64::EPSILON);
        assert!((r - 100.0).abs() < f64::EPSILON);
        assert!((v - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_regime_distribution_flat_f64_empty() {
        let (t, r, v) = regime_distribution_flat_f64(&[], 42);
        assert!((t - 0.0).abs() < f64::EPSILON);
        assert!((r - 100.0).abs() < f64::EPSILON);
        assert!((v - 0.0).abs() < f64::EPSILON);
    }
}