//! Trajectory collection and management for PPO
//!
//! This module handles collecting trajectories from environment interactions
//! and preparing them for PPO training with proper batching and preprocessing.

use serde::{Deserialize, Serialize};

use ml_core::action_space::FactoredAction;
use ml_core::MLError;

// Re-export the GPU-native trajectory tensors for callers that want zero-Candle batches.
pub use crate::cuda_nn::CudaTrajectoryTensors;

/// Single step trajectory data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrajectoryStep {
    /// State at this step
    pub state: Vec<f32>,
    /// Action taken
    pub action: FactoredAction,
    /// Action probability (log probability)
    pub log_prob: f32,
    /// Value estimate at this state
    pub value: f32,
    /// Reward received
    pub reward: f32,
    /// Whether episode terminated
    pub done: bool,
}

impl TrajectoryStep {
    /// Create new trajectory step
    pub const fn new(
        state: Vec<f32>,
        action: FactoredAction,
        log_prob: f32,
        value: f32,
        reward: f32,
        done: bool,
    ) -> Self {
        Self {
            state,
            action,
            log_prob,
            value,
            reward,
            done,
        }
    }
}

/// Complete trajectory (episode or segment)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Trajectory {
    /// Steps in this trajectory
    pub steps: Vec<TrajectoryStep>,
    /// Total return of trajectory
    pub total_return: f32,
    /// Length of trajectory
    pub length: usize,
}

impl Trajectory {
    /// Create new empty trajectory
    pub const fn new() -> Self {
        Self {
            steps: Vec::new(),
            total_return: 0.0,
            length: 0,
        }
    }

    /// Add step to trajectory
    pub fn add_step(&mut self, step: TrajectoryStep) {
        self.total_return += step.reward;
        self.steps.push(step);
        self.length += 1;
    }

    /// Get trajectory states as flat vector
    pub fn get_states(&self) -> Vec<Vec<f32>> {
        self.steps.iter().map(|step| step.state.clone()).collect()
    }

    /// Get trajectory actions
    pub fn get_actions(&self) -> Vec<FactoredAction> {
        self.steps.iter().map(|step| step.action).collect()
    }

    /// Get trajectory log probabilities
    pub fn get_log_probs(&self) -> Vec<f32> {
        self.steps.iter().map(|step| step.log_prob).collect()
    }

    /// Get trajectory values
    pub fn get_values(&self) -> Vec<f32> {
        self.steps.iter().map(|step| step.value).collect()
    }

    /// Get trajectory rewards
    pub fn get_rewards(&self) -> Vec<f32> {
        self.steps.iter().map(|step| step.reward).collect()
    }

    /// Get trajectory done flags
    pub fn get_dones(&self) -> Vec<bool> {
        self.steps.iter().map(|step| step.done).collect()
    }

    /// Extend a flat `f32` buffer with all state data
    pub fn extend_flat_states(&self, buf: &mut Vec<f32>) {
        for step in &self.steps {
            buf.extend_from_slice(&step.state);
        }
    }

    /// Extend a buffer with all actions
    pub fn extend_actions(&self, buf: &mut Vec<FactoredAction>) {
        buf.extend(self.steps.iter().map(|s| s.action));
    }

    /// Extend a buffer with all log probabilities
    pub fn extend_log_probs(&self, buf: &mut Vec<f32>) {
        buf.extend(self.steps.iter().map(|s| s.log_prob));
    }

    /// Extend a buffer with all value estimates
    pub fn extend_values(&self, buf: &mut Vec<f32>) {
        buf.extend(self.steps.iter().map(|s| s.value));
    }

    /// Extend a buffer with all rewards
    pub fn extend_rewards(&self, buf: &mut Vec<f32>) {
        buf.extend(self.steps.iter().map(|s| s.reward));
    }

    /// Extend a buffer with all done flags
    pub fn extend_dones(&self, buf: &mut Vec<bool>) {
        buf.extend(self.steps.iter().map(|s| s.done));
    }

    /// Check if trajectory is complete (ends with done=true)
    pub fn is_complete(&self) -> bool {
        self.steps.last().map(|step| step.done).unwrap_or(false)
    }

    /// Compute discounted returns for this trajectory
    pub fn compute_returns(&self, gamma: f32) -> Vec<f32> {
        let mut returns = vec![0.0; self.length];
        let mut running_return = 0.0;

        for i in (0..self.length).rev() {
            let Some(step) = self.steps.get(i) else {
                continue;
            };

            if step.done {
                running_return = 0.0;
            }
            running_return = step.reward + gamma * running_return;

            if let Some(ret) = returns.get_mut(i) {
                *ret = running_return;
            }
        }

        returns
    }
}

impl Default for Trajectory {
    fn default() -> Self {
        Self::new()
    }
}

/// Batch of trajectories for training
#[derive(Debug, Clone)]
pub struct TrajectoryBatch {
    /// All trajectories in batch
    pub trajectories: Vec<Trajectory>,
    /// Flattened states from all trajectories
    pub states: Vec<Vec<f32>>,
    /// Pre-flattened state data for zero-copy tensor creation.
    pub states_flat: Vec<f32>,
    /// Flattened actions from all trajectories
    pub actions: Vec<FactoredAction>,
    /// Flattened log probabilities
    pub log_probs: Vec<f32>,
    /// Flattened values
    pub values: Vec<f32>,
    /// Flattened rewards
    pub rewards: Vec<f32>,
    /// Flattened done flags
    pub dones: Vec<bool>,
    /// Computed advantages
    pub advantages: Vec<f32>,
    /// Computed returns
    pub returns: Vec<f32>,
}

impl TrajectoryBatch {
    /// Create batch from trajectories.
    pub fn from_trajectories(
        trajectories: Vec<Trajectory>,
        advantages: Vec<f32>,
        returns: Vec<f32>,
    ) -> Self {
        let total_steps: usize = trajectories.iter().map(|t| t.length).sum();

        let mut states = Vec::with_capacity(total_steps);
        let mut actions = Vec::with_capacity(total_steps);
        let mut log_probs = Vec::with_capacity(total_steps);
        let mut values = Vec::with_capacity(total_steps);
        let mut rewards = Vec::with_capacity(total_steps);
        let mut dones = Vec::with_capacity(total_steps);

        let state_dim = trajectories
            .first()
            .and_then(|t| t.steps.first())
            .map(|s| s.state.len())
            .unwrap_or(0);
        let mut states_flat = Vec::with_capacity(total_steps * state_dim);

        for trajectory in &trajectories {
            states.extend(trajectory.steps.iter().map(|step| step.state.clone()));
            trajectory.extend_flat_states(&mut states_flat);
            trajectory.extend_actions(&mut actions);
            trajectory.extend_log_probs(&mut log_probs);
            trajectory.extend_values(&mut values);
            trajectory.extend_rewards(&mut rewards);
            trajectory.extend_dones(&mut dones);
        }

        Self {
            trajectories,
            states,
            states_flat,
            actions,
            log_probs,
            values,
            rewards,
            dones,
            advantages,
            returns,
        }
    }

    /// Convert to GPU-native `CudaTrajectoryTensors`.
    pub fn to_cuda_tensors(
        &self,
        ctx: &crate::cuda_nn::GpuContext,
        state_dim: usize,
    ) -> Result<CudaTrajectoryTensors, MLError> {
        if self.states_flat.len() != self.states.len() * state_dim {
            return Err(MLError::TensorOperationError(
                "states_flat not populated; call from_trajectories first".to_owned(),
            ));
        }

        CudaTrajectoryTensors::from_batch(
            ctx,
            &self.states_flat,
            &self.actions,
            &self.log_probs,
            &self.values,
            &self.advantages,
            &self.returns,
            state_dim,
        )
    }

    /// Get total number of steps in batch
    pub fn total_steps(&self) -> usize {
        self.states.len()
    }

    /// Get number of trajectories
    pub fn num_trajectories(&self) -> usize {
        self.trajectories.len()
    }

    /// Normalize advantages (zero mean, unit variance)
    pub fn normalize_advantages(&mut self) -> Result<(), MLError> {
        if self.advantages.is_empty() {
            return Ok(());
        }

        let mean = self.advantages.iter().sum::<f32>() / self.advantages.len() as f32;
        let var = self
            .advantages
            .iter()
            .map(|a| (a - mean).powi(2))
            .sum::<f32>()
            / self.advantages.len() as f32;
        let std = (var + 1e-8).sqrt();

        for advantage in &mut self.advantages {
            *advantage = (*advantage - mean) / std;
        }

        Ok(())
    }

    /// Create mini-batch index ranges for training.
    ///
    /// Returns `(start, end)` pairs. Callers should use
    /// `CudaTrajectoryTensors::sub_batch()` for GPU-side `DtoD` slicing.
    pub fn create_mini_batch_ranges(&self, mini_batch_size: usize) -> Vec<MiniBatchRange> {
        let total_steps = self.total_steps();
        let mut ranges = Vec::new();

        for start in (0..total_steps).step_by(mini_batch_size) {
            let end = (start + mini_batch_size).min(total_steps);
            ranges.push(MiniBatchRange { start, end });
        }

        ranges
    }

    /// Convert trajectories into sequence ranges for BPTT.
    ///
    /// Returns `(start, end, actual_length)` tuples. Callers should use
    /// `CudaTrajectoryTensors::sub_batch()` for GPU-side `DtoD` slicing.
    pub fn to_sequence_ranges(&self, sequence_length: usize) -> Vec<SequenceRange> {
        if sequence_length == 0 {
            return Vec::new();
        }

        let mut sequences = Vec::new();
        let mut global_idx = 0;

        for trajectory in &self.trajectories {
            let episode_length = trajectory.length;
            let mut local_idx = 0;

            while local_idx < episode_length {
                let remaining = episode_length - local_idx;
                let seq_len = remaining.min(sequence_length);
                let start = global_idx + local_idx;
                let end = start + seq_len;

                sequences.push(SequenceRange {
                    start,
                    end,
                    actual_length: seq_len,
                });
                local_idx += seq_len;
            }

            global_idx += episode_length;
        }

        sequences
    }
}

/// Tensors for trajectory batch (GPU-native via `CudaTrajectoryTensors`)
///
/// This type alias exists for backward compatibility. Callers should migrate
/// to `CudaTrajectoryTensors` directly.
pub type TrajectoryTensors = CudaTrajectoryTensors;

/// GPU-friendly mini-batch range (start..end indices into parent batch).
///
/// Use with `CudaTrajectoryTensors::sub_batch(range.start, range.end, stream)`
/// for zero-CPU mini-batch extraction.
#[derive(Debug, Clone, Copy)]
pub struct MiniBatchRange {
    /// Start index (inclusive) in the parent batch
    pub start: usize,
    /// End index (exclusive) in the parent batch
    pub end: usize,
}

impl MiniBatchRange {
    /// Number of samples in this mini-batch
    pub const fn len(&self) -> usize {
        self.end - self.start
    }

    /// Whether the range is empty
    pub const fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

/// GPU-friendly sequence range (start..end indices into parent batch) for BPTT.
///
/// Use with `CudaTrajectoryTensors::sub_batch(range.start, range.end, stream)`
/// for zero-CPU sequence extraction.
#[derive(Debug, Clone, Copy)]
pub struct SequenceRange {
    /// Start index (inclusive) in the parent batch
    pub start: usize,
    /// End index (exclusive) in the parent batch
    pub end: usize,
    /// Actual number of valid steps (may equal `end - start`)
    pub actual_length: usize,
}

impl SequenceRange {
    /// Get the actual length of this sequence
    pub const fn length(&self) -> usize {
        self.actual_length
    }
}

/// Collect trajectories from environment
pub fn collect_trajectories<F>(
    mut collect_fn: F,
    num_trajectories: usize,
    max_steps_per_trajectory: usize,
) -> Result<Vec<Trajectory>, MLError>
where
    F: FnMut() -> Result<Trajectory, MLError>,
{
    let mut trajectories = Vec::with_capacity(num_trajectories);

    for _ in 0..num_trajectories {
        let trajectory = collect_fn()?;

        if trajectory.length > max_steps_per_trajectory {
            return Err(MLError::ValidationError {
                message: format!(
                    "Trajectory too long: {} > {}",
                    trajectory.length, max_steps_per_trajectory
                ),
            });
        }

        trajectories.push(trajectory);
    }

    Ok(trajectories)
}

#[cfg(test)]
#[allow(clippy::doc_markdown)]
mod tests {
    use super::*;
    use ml_core::action_space::{ExposureLevel, OrderType, Urgency};
    use ml_core::trading_action::TradingAction;

    fn fa(ta: TradingAction) -> FactoredAction {
        FactoredAction::from_trading_action(ta)
    }

    #[test]
    fn test_trajectory_creation() {
        let mut trajectory = Trajectory::new();
        assert_eq!(trajectory.length, 0);
        assert_eq!(trajectory.total_return, 0.0);

        let step = TrajectoryStep::new(
            vec![1.0, 2.0, 3.0],
            fa(TradingAction::Buy),
            -0.5,
            10.0,
            1.0,
            false,
        );

        trajectory.add_step(step);
        assert_eq!(trajectory.length, 1);
        assert_eq!(trajectory.total_return, 1.0);
    }

    #[test]
    fn test_trajectory_returns_computation() {
        let mut trajectory = Trajectory::new();

        trajectory.add_step(TrajectoryStep::new(
            vec![0.0], fa(TradingAction::Buy), 0.0, 0.0, 1.0, false,
        ));
        trajectory.add_step(TrajectoryStep::new(
            vec![1.0], fa(TradingAction::Sell), 0.0, 0.0, 2.0, false,
        ));
        trajectory.add_step(TrajectoryStep::new(
            vec![2.0], fa(TradingAction::Hold), 0.0, 0.0, 3.0, true,
        ));

        let returns = trajectory.compute_returns(0.9);
        assert_eq!(returns.len(), 3);
        assert!((returns[2] - 3.0).abs() < 1e-6);
        assert!((returns[1] - 4.7).abs() < 1e-6);
        assert!((returns[0] - 5.23).abs() < 1e-6);
    }

    #[test]
    fn test_trajectory_batch_creation() {
        let mut traj1 = Trajectory::new();
        traj1.add_step(TrajectoryStep::new(
            vec![1.0], fa(TradingAction::Buy), 0.0, 0.0, 1.0, false,
        ));
        traj1.add_step(TrajectoryStep::new(
            vec![2.0], fa(TradingAction::Sell), 0.0, 0.0, 2.0, true,
        ));

        let mut traj2 = Trajectory::new();
        traj2.add_step(TrajectoryStep::new(
            vec![3.0], fa(TradingAction::Hold), 0.0, 0.0, 3.0, true,
        ));

        let trajectories = vec![traj1, traj2];
        let advantages = vec![0.1, 0.2, 0.3];
        let returns = vec![1.0, 2.0, 3.0];

        let batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);

        assert_eq!(batch.total_steps(), 3);
        assert_eq!(batch.num_trajectories(), 2);
        assert_eq!(batch.states.len(), 3);
        assert_eq!(batch.actions.len(), 3);
        assert_eq!(batch.states_flat, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_advantage_normalization() -> Result<(), Box<dyn std::error::Error>> {
        let trajectories = vec![Trajectory::new()];
        let advantages = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let returns = vec![0.0; 5];

        let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
        batch.normalize_advantages()?;

        let mean = batch.advantages.iter().sum::<f32>() / batch.advantages.len() as f32;
        assert!(mean.abs() < 1e-6);

        let var =
            batch.advantages.iter().map(|a| a.powi(2)).sum::<f32>() / batch.advantages.len() as f32;
        assert!((var - 1.0).abs() < 1e-5);
        Ok(())
    }

    #[test]
    fn test_mini_batch_ranges() {
        let trajectories = vec![Trajectory::new()];
        let advantages = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let returns = vec![0.0; 5];
        let states = vec![vec![1.0]; 5];
        let buy_action = FactoredAction::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal);
        let actions = vec![buy_action; 5];
        let log_probs = vec![0.0; 5];
        let values = vec![0.0; 5];
        let dones = vec![false; 5];

        let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
        batch.states_flat = states.iter().flatten().copied().collect();
        batch.states = states;
        batch.actions = actions;
        batch.log_probs = log_probs;
        batch.values = values;
        batch.dones = dones;

        let ranges = batch.create_mini_batch_ranges(2);
        assert_eq!(ranges.len(), 3);
        assert_eq!(ranges[0].len(), 2);
        assert_eq!(ranges[1].len(), 2);
        assert_eq!(ranges[2].len(), 1);
    }

    #[test]
    fn test_extend_flat_states() {
        let mut trajectory = Trajectory::new();
        trajectory.add_step(TrajectoryStep::new(
            vec![1.0, 2.0], fa(TradingAction::Buy), 0.0, 0.0, 0.0, false,
        ));
        trajectory.add_step(TrajectoryStep::new(
            vec![3.0, 4.0], fa(TradingAction::Sell), 0.0, 0.0, 0.0, true,
        ));

        let mut buf = Vec::new();
        trajectory.extend_flat_states(&mut buf);
        assert_eq!(buf, vec![1.0, 2.0, 3.0, 4.0]);

        let mut actions = Vec::new();
        trajectory.extend_actions(&mut actions);
        assert_eq!(actions, trajectory.get_actions());
    }

    #[test]
    fn test_flatten_states_consistency() {
        let mut traj = Trajectory::new();
        traj.add_step(TrajectoryStep::new(
            vec![1.0, 2.0, 3.0], fa(TradingAction::Buy), 0.0, 0.0, 1.0, false,
        ));
        traj.add_step(TrajectoryStep::new(
            vec![4.0, 5.0, 6.0], fa(TradingAction::Sell), 0.0, 0.0, 2.0, true,
        ));

        let batch = TrajectoryBatch::from_trajectories(
            vec![traj], vec![0.1, 0.2], vec![1.0, 2.0],
        );

        let flat_old: Vec<f32> = batch.states.iter().flatten().copied().collect();
        assert_eq!(batch.states_flat, flat_old);
        assert_eq!(batch.states_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    }
}
