//! Hindsight Experience Replay (HER) for enhanced data efficiency
//!
//! HER relabels failed experiences with achieved goals to improve learning efficiency
//! by 5-10x. Originally proposed in "Hindsight Experience Replay" (Andrychowicz et al., 2017).
//!
//! Key benefits:
//! - 5-10x data efficiency improvement
//! - Learn from failure experiences
//! - Better generalization across goals
//! - Faster convergence in sparse reward settings

use parking_lot::RwLock;
use rand::prelude::*;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};

use super::experience::{Experience, ExperienceBatch};
use super::prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig};
use ml_core::MLError;

/// Configuration for Hindsight Experience Replay
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HindsightReplayConfig {
    /// Base prioritized replay buffer configuration
    pub base_config: PrioritizedReplayConfig,

    /// Dimension of the goal portion in the state vector
    /// For trading: goal = target return/profit
    pub goal_dim: usize,

    /// Fraction of batch samples from HER (rest from normal replay)
    /// Range: [0.0, 1.0], recommended: 0.5-0.8
    pub her_ratio: f64,

    /// HER strategy: "final" (use final achieved goal) or "future" (use future achieved goal)
    pub her_strategy: HindsightStrategy,

    /// Number of future goals to sample per experience (for "future" strategy)
    pub k_future: usize,

    /// Batch size for sampling
    pub batch_size: usize,
}

impl Default for HindsightReplayConfig {
    fn default() -> Self {
        Self {
            base_config: PrioritizedReplayConfig::default(),
            goal_dim: 1, // Single goal dimension (target return)
            her_ratio: 0.5, // 50% HER samples
            her_strategy: HindsightStrategy::Final,
            k_future: 4, // Sample 4 future goals
            batch_size: 32,
        }
    }
}

/// HER relabeling strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum HindsightStrategy {
    /// Use final achieved state as goal
    Final,
    /// Sample k future achieved states as goals
    Future,
    /// Sample from episode trajectory
    Episode,
    /// Random goal from buffer
    Random,
}

/// Statistics for HER buffer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HindsightReplayStats {
    /// Total normal samples taken
    pub normal_samples: u64,
    /// Total HER samples taken
    pub her_samples: u64,
    /// Total experiences relabeled
    pub relabeled_count: u64,
    /// Average reward improvement from relabeling
    pub avg_reward_improvement: f32,
}

/// Hindsight Experience Replay buffer
///
/// Wraps a prioritized replay buffer and adds HER functionality.
/// Automatically relabels experiences with achieved goals during sampling.
#[derive(Debug)]
pub struct HindsightReplayBuffer {
    /// Configuration
    config: HindsightReplayConfig,

    /// Base prioritized replay buffer
    base_buffer: PrioritizedReplayBuffer,

    /// Episode boundaries (start indices of each episode)
    episode_boundaries: RwLock<Vec<usize>>,

    /// Statistics
    normal_samples: AtomicU64,
    her_samples: AtomicU64,
    relabeled_count: AtomicU64,
    total_reward_improvement: RwLock<f32>,
}

impl HindsightReplayBuffer {
    /// Create a new HER buffer
    pub fn new(config: HindsightReplayConfig) -> Result<Self, MLError> {
        let base_buffer = PrioritizedReplayBuffer::new(config.base_config.clone())?;

        Ok(Self {
            config,
            base_buffer,
            episode_boundaries: RwLock::new(vec![0]),
            normal_samples: AtomicU64::new(0),
            her_samples: AtomicU64::new(0),
            relabeled_count: AtomicU64::new(0),
            total_reward_improvement: RwLock::new(0.0),
        })
    }

    /// Add experience to buffer
    pub fn push(&self, experience: Experience) -> Result<(), MLError> {
        let is_done = experience.done;

        // Push experience first
        self.base_buffer.push(experience)?;

        // Mark episode boundary on terminal states (after push, so boundary is at correct index)
        if is_done {
            let mut boundaries = self.episode_boundaries.write();
            boundaries.push(self.base_buffer.len());
        }

        Ok(())
    }

    /// Sample batch with HER relabeling
    pub fn sample_with_hindsight(&self, batch_size: usize) -> Result<(ExperienceBatch, Vec<f32>, Vec<usize>), MLError> {
        if !self.base_buffer.can_sample(batch_size) {
            return Err(MLError::InvalidInput(
                "Not enough experiences in buffer".to_owned(),
            ));
        }

        // Calculate split between normal and HER samples
        let her_size = (batch_size as f64 * self.config.her_ratio) as usize;
        let normal_size = batch_size.saturating_sub(her_size);

        // Sample normal experiences
        let mut all_experiences = Vec::with_capacity(batch_size);
        let mut all_weights = Vec::with_capacity(batch_size);
        let mut all_indices = Vec::with_capacity(batch_size);

        if normal_size > 0 {
            let (normal_exp, normal_weights, normal_indices) = self.base_buffer.sample(normal_size)?;
            all_experiences.extend(normal_exp);
            all_weights.extend(normal_weights);
            all_indices.extend(normal_indices);
            self.normal_samples.fetch_add(normal_size as u64, Ordering::Relaxed);
        }

        // Sample and relabel HER experiences
        if her_size > 0 {
            let (her_exp, her_weights, her_indices) = self.sample_her_experiences(her_size)?;
            all_experiences.extend(her_exp);
            all_weights.extend(her_weights);
            all_indices.extend(her_indices);
            self.her_samples.fetch_add(her_size as u64, Ordering::Relaxed);
        }

        // Shuffle while preserving correspondence
        let mut combined: Vec<_> = all_experiences.into_iter()
            .zip(all_weights)
            .zip(all_indices)
            .collect();
        combined.shuffle(&mut thread_rng());

        let (experiences_with_weights, indices): (Vec<_>, Vec<_>) = combined.into_iter().unzip();
        let (experiences, weights): (Vec<_>, Vec<_>) = experiences_with_weights.into_iter().unzip();

        Ok((ExperienceBatch::new(experiences), weights, indices))
    }

    /// Sample experiences and apply HER relabeling
    fn sample_her_experiences(&self, count: usize) -> Result<(Vec<Experience>, Vec<f32>, Vec<usize>), MLError> {
        // Sample base experiences
        let (base_experiences, weights, indices) = self.base_buffer.sample(count)?;
        let mut her_experiences = Vec::with_capacity(count);

        for experience in base_experiences {
            // Relabel based on strategy
            let relabeled = match self.config.her_strategy {
                HindsightStrategy::Final => self.relabel_with_final_goal(&experience)?,
                HindsightStrategy::Future => self.relabel_with_future_goal(&experience)?,
                HindsightStrategy::Episode => self.relabel_with_episode_goal(&experience)?,
                HindsightStrategy::Random => self.relabel_with_random_goal(&experience)?,
            };

            her_experiences.push(relabeled);
            self.relabeled_count.fetch_add(1, Ordering::Relaxed);
        }

        Ok((her_experiences, weights, indices))
    }

    /// Relabel experience with final achieved goal
    fn relabel_with_final_goal(&self, exp: &Experience) -> Result<Experience, MLError> {
        // Extract achieved goal from next_state (final state of trajectory)
        let achieved_goal = self.extract_achieved_goal(&exp.next_state)?;
        self.relabel_goal(exp, &achieved_goal)
    }

    /// Relabel experience with future achieved goal
    fn relabel_with_future_goal(&self, exp: &Experience) -> Result<Experience, MLError> {
        // Sample a random future goal from the buffer
        let buffer_size = self.base_buffer.len();
        if buffer_size == 0 {
            return Ok(exp.clone());
        }

        // Sample from buffer to get future experience
        let (future_exps, _, _) = self.base_buffer.sample(1)?;
        if let Some(future_exp) = future_exps.first() {
            let achieved_goal = self.extract_achieved_goal(&future_exp.next_state)?;
            self.relabel_goal(exp, &achieved_goal)
        } else {
            Ok(exp.clone())
        }
    }

    /// Relabel experience with goal from same episode
    fn relabel_with_episode_goal(&self, exp: &Experience) -> Result<Experience, MLError> {
        // Sample from buffer to get episode experience
        let buffer_size = self.base_buffer.len();
        if buffer_size == 0 {
            return Ok(exp.clone());
        }

        let (random_exps, _, _) = self.base_buffer.sample(1)?;
        if let Some(random_exp) = random_exps.first() {
            let achieved_goal = self.extract_achieved_goal(&random_exp.next_state)?;
            self.relabel_goal(exp, &achieved_goal)
        } else {
            Ok(exp.clone())
        }
    }

    /// Relabel experience with random goal from buffer
    fn relabel_with_random_goal(&self, exp: &Experience) -> Result<Experience, MLError> {
        let buffer_size = self.base_buffer.len();
        if buffer_size == 0 {
            return Ok(exp.clone());
        }

        let (random_exps, _, _) = self.base_buffer.sample(1)?;
        if let Some(random_exp) = random_exps.first() {
            let achieved_goal = self.extract_achieved_goal(&random_exp.next_state)?;
            self.relabel_goal(exp, &achieved_goal)
        } else {
            Ok(exp.clone())
        }
    }

    /// Extract achieved goal from state
    fn extract_achieved_goal(&self, state: &[f32]) -> Result<Vec<f32>, MLError> {
        if state.len() < self.config.goal_dim {
            return Err(MLError::InvalidInput(
                format!("State dimension {} < goal dimension {}", state.len(), self.config.goal_dim)
            ));
        }

        // Goal is first goal_dim elements of state
        Ok(state[..self.config.goal_dim].to_vec())
    }

    /// Relabel experience with new goal
    fn relabel_goal(&self, exp: &Experience, achieved_goal: &[f32]) -> Result<Experience, MLError> {
        if achieved_goal.len() != self.config.goal_dim {
            return Err(MLError::InvalidInput(
                format!("Achieved goal dimension {} != goal dimension {}", achieved_goal.len(), self.config.goal_dim)
            ));
        }

        // Create new state with relabeled goal
        let mut new_state = exp.state.clone();
        let mut new_next_state = exp.next_state.clone();

        // Replace goal portion with achieved goal
        new_state[..self.config.goal_dim].copy_from_slice(achieved_goal);
        new_next_state[..self.config.goal_dim].copy_from_slice(achieved_goal);

        // Compute new reward based on goal achievement
        let achieved_this_step = self.extract_achieved_goal(&exp.next_state)?;
        let goal_achieved = self.compute_goal_distance(&achieved_this_step, achieved_goal) < 0.01;
        let new_reward = if goal_achieved { 1.0 } else { -0.01 }; // Sparse reward

        // Track reward improvement
        let reward_improvement = new_reward - exp.reward_f32();
        let mut total_improvement = self.total_reward_improvement.write();
        *total_improvement += reward_improvement;

        Ok(Experience::new(
            new_state,
            exp.action,
            new_reward,
            new_next_state,
            exp.done,
        ))
    }

    /// Compute distance between achieved and desired goal
    fn compute_goal_distance(&self, achieved: &[f32], desired: &[f32]) -> f32 {
        achieved.iter()
            .zip(desired.iter())
            .map(|(a, d)| (a - d).powi(2))
            .sum::<f32>()
            .sqrt()
    }

    /// Get buffer statistics
    pub fn stats(&self) -> HindsightReplayStats {
        let normal = self.normal_samples.load(Ordering::Relaxed);
        let her = self.her_samples.load(Ordering::Relaxed);
        let relabeled = self.relabeled_count.load(Ordering::Relaxed);
        let total_improvement = *self.total_reward_improvement.read();

        HindsightReplayStats {
            normal_samples: normal,
            her_samples: her,
            relabeled_count: relabeled,
            avg_reward_improvement: if relabeled > 0 {
                total_improvement / relabeled as f32
            } else {
                0.0
            },
        }
    }

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

    /// Check if buffer can be sampled
    pub fn can_sample(&self) -> bool {
        self.base_buffer.can_sample(self.config.batch_size)
    }

    /// Update priorities for experiences
    pub fn update_priorities(&self, indices: &[usize], priorities: &[f32]) -> Result<(), MLError> {
        self.base_buffer.update_priorities(indices, priorities)
    }

    /// Mark episode boundary (call when episode ends)
    pub fn mark_episode_boundary(&self) {
        let mut boundaries = self.episode_boundaries.write();
        boundaries.push(self.base_buffer.len());
    }

    /// Get minimum number of experiences needed before sampling
    pub const fn min_experiences(&self) -> usize {
        self.config.batch_size
    }
}

#[cfg(test)]
#[allow(
    clippy::modulo_arithmetic,
    clippy::manual_range_contains,
    clippy::assertions_on_result_states
)]
mod tests {
    use super::*;

    fn create_test_experience(state_val: f32, reward: f32, done: bool) -> Experience {
        let state = vec![state_val, 0.5, 0.5]; // [goal, feature1, feature2]
        let next_state = vec![state_val + 0.1, 0.6, 0.6];
        Experience::new(state, 0, reward, next_state, done)
    }

    #[test]
    fn test_hindsight_replay_creation() {
        let config = HindsightReplayConfig::default();
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        assert_eq!(buffer.size(), 0);
        assert!(!buffer.can_sample());
    }

    #[test]
    fn test_push_experience() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            batch_size: 16, // Set smaller batch size so 20 experiences can be sampled
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add experiences
        for i in 0..20 {
            let exp = create_test_experience(i as f32 * 0.1, -1.0, i == 19);
            buffer.push(exp).unwrap();
        }

        assert_eq!(buffer.size(), 20);
        assert!(buffer.can_sample());
    }

    #[test]
    fn test_sample_with_hindsight() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            her_ratio: 0.5,
            goal_dim: 1,
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add episode of experiences
        for i in 0..50 {
            let exp = create_test_experience(i as f32 * 0.1, -1.0, i == 49);
            buffer.push(exp).unwrap();
        }

        // Sample batch
        let (batch, weights, indices) = buffer.sample_with_hindsight(32).unwrap();
        assert_eq!(batch.batch_size, 32);
        assert!(batch.is_valid());
        assert_eq!(weights.len(), 32);
        assert_eq!(indices.len(), 32);

        // Verify stats
        let stats = buffer.stats();
        assert!(stats.normal_samples > 0);
        assert!(stats.her_samples > 0);
        assert!(stats.relabeled_count > 0);
    }

    #[test]
    fn test_her_ratio_distribution() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            her_ratio: 0.8, // 80% HER samples
            goal_dim: 1,
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add experiences
        for i in 0..100 {
            let exp = create_test_experience(i as f32 * 0.01, -1.0, i % 20 == 19);
            buffer.push(exp).unwrap();
        }

        // Sample multiple batches
        for _ in 0..10 {
            let _ = buffer.sample_with_hindsight(20).unwrap();
        }

        let stats = buffer.stats();
        let total_samples = stats.normal_samples + stats.her_samples;
        let her_fraction = stats.her_samples as f64 / total_samples as f64;

        // Should be close to 0.8
        assert!(her_fraction >= 0.75 && her_fraction <= 0.85);
    }

    #[test]
    fn test_goal_relabeling() {
        let config = HindsightReplayConfig {
            goal_dim: 1,
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        let exp = create_test_experience(0.5, -1.0, false);
        let achieved_goal = vec![0.8];

        let relabeled = buffer.relabel_goal(&exp, &achieved_goal).unwrap();

        // Goal portion should be updated
        assert_eq!(relabeled.state[0], 0.8);
        assert_eq!(relabeled.next_state[0], 0.8);

        // Non-goal features should be preserved
        assert_eq!(relabeled.state[1], exp.state[1]);
        assert_eq!(relabeled.state[2], exp.state[2]);
    }

    #[test]
    fn test_extract_achieved_goal() {
        let config = HindsightReplayConfig {
            goal_dim: 2,
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        let state = vec![0.5, 0.7, 0.3, 0.9]; // First 2 are goal
        let goal = buffer.extract_achieved_goal(&state).unwrap();

        assert_eq!(goal.len(), 2);
        assert_eq!(goal[0], 0.5);
        assert_eq!(goal[1], 0.7);
    }

    #[test]
    fn test_goal_distance_computation() {
        let config = HindsightReplayConfig::default();
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        let achieved = vec![0.5];
        let desired = vec![0.5];
        let dist = buffer.compute_goal_distance(&achieved, &desired);
        assert!(dist < 0.001); // Should be ~0

        let achieved2 = vec![0.5];
        let desired2 = vec![0.8];
        let dist2 = buffer.compute_goal_distance(&achieved2, &desired2);
        assert!((dist2 - 0.3).abs() < 0.001); // Should be 0.3
    }

    #[test]
    fn test_hindsight_strategies() {
        let strategies = vec![
            HindsightStrategy::Final,
            HindsightStrategy::Future,
            HindsightStrategy::Episode,
            HindsightStrategy::Random,
        ];

        for strategy in strategies {
            let config = HindsightReplayConfig {
                base_config: PrioritizedReplayConfig {
                    capacity: 1000,
                    ..Default::default()
                },
                her_strategy: strategy,
                goal_dim: 1,
                her_ratio: 1.0, // 100% HER for testing
                ..Default::default()
            };
            let buffer = HindsightReplayBuffer::new(config).unwrap();

            // Add experiences
            for i in 0..50 {
                let exp = create_test_experience(i as f32 * 0.1, -1.0, i % 10 == 9);
                buffer.push(exp).unwrap();
            }

            // Should be able to sample with any strategy
            let (batch, _, _) = buffer.sample_with_hindsight(20).unwrap();
            assert_eq!(batch.batch_size, 20);
        }
    }

    #[test]
    fn test_episode_boundaries() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add 3 episodes of 10 experiences each
        for episode in 0..3 {
            for i in 0..10 {
                let done = i == 9;
                let exp = create_test_experience(
                    (episode * 10 + i) as f32 * 0.01,
                    -1.0,
                    done
                );
                buffer.push(exp).unwrap();
            }
        }

        let boundaries = buffer.episode_boundaries.read();
        assert_eq!(boundaries.len(), 4); // Start + 3 episode ends
        assert_eq!(boundaries[0], 0);
        assert_eq!(boundaries[1], 10);
        assert_eq!(boundaries[2], 20);
        assert_eq!(boundaries[3], 30);
    }

    #[test]
    fn test_reward_improvement_tracking() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            her_ratio: 1.0, // 100% HER
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add failing experiences (negative reward)
        for i in 0..50 {
            let exp = create_test_experience(i as f32 * 0.01, -1.0, i % 10 == 9);
            buffer.push(exp).unwrap();
        }

        // Sample with HER (should improve rewards)
        for _ in 0..5 {
            let _ = buffer.sample_with_hindsight(10).unwrap();
        }

        let stats = buffer.stats();
        // Some experiences should have improved rewards
        assert!(stats.relabeled_count > 0);
    }

    #[test]
    fn test_invalid_goal_dimension() {
        let config = HindsightReplayConfig {
            goal_dim: 5, // Larger than state dimension
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        let state = vec![0.1, 0.2]; // Only 2 dimensions
        let result = buffer.extract_achieved_goal(&state);

        assert!(result.is_err());
    }

    #[test]
    fn test_empty_buffer_sampling() {
        let config = HindsightReplayConfig::default();
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        let result = buffer.sample_with_hindsight(10);
        assert!(result.is_err());
    }

    #[test]
    fn test_stats_accumulation() {
        let config = HindsightReplayConfig {
            base_config: PrioritizedReplayConfig {
                capacity: 1000,
                ..Default::default()
            },
            her_ratio: 0.5,
            ..Default::default()
        };
        let buffer = HindsightReplayBuffer::new(config).unwrap();

        // Add experiences
        for i in 0..100 {
            let exp = create_test_experience(i as f32 * 0.01, -1.0, i % 10 == 9);
            buffer.push(exp).unwrap();
        }

        // Sample multiple times
        let num_samples = 10;
        for _ in 0..num_samples {
            let _ = buffer.sample_with_hindsight(20).unwrap();
        }

        let stats = buffer.stats();
        assert_eq!(stats.normal_samples + stats.her_samples, num_samples * 20);
    }
}
