//! PPO strategy adapters for the [`ValidatableStrategy`] trait.
//!
//! Wraps a [`PPO`] agent (MLP or LSTM variant) so it can be trained and evaluated
//! through the walk-forward validation harness, deflated Sharpe ratio, and PBO
//! pipelines.
//!
//! Two adapters are provided:
//! - [`PpoStrategy`]: wraps a standard MLP-based PPO.
//! - [`PpoLstmStrategy`]: wraps an LSTM-based PPO with hidden state management.

use std::cell::RefCell;

use rand::Rng;

use crate::common::action::FactoredAction;
use crate::ppo::gae::compute_gae;
#[cfg(test)]
use crate::ppo::gae::GAEConfig;
use crate::ppo::hidden_state_manager::HiddenStateManager;
use crate::ppo::ppo::{ActorNetwork, CriticNetwork, PPOConfig, PPO};
use crate::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
use ml_validation::{TimeSeriesData, ValidatableStrategy};
use crate::MLError;

// ---------------------------------------------------------------------------
// PpoStrategy (MLP only)
// ---------------------------------------------------------------------------

/// Adapter that wraps an MLP-based [`PPO`] agent behind the [`ValidatableStrategy`] trait.
///
/// Uses `RefCell` for interior mutability because several PPO operations require
/// `&mut self` while [`ValidatableStrategy::evaluate`] takes `&self`.
pub struct PpoStrategy {
    config: PPOConfig,
    /// Interior mutability: PPO methods require `&mut self`
    /// but `ValidatableStrategy::evaluate` only receives `&self`.
    ppo: RefCell<PPO>,
}

impl std::fmt::Debug for PpoStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PpoStrategy")
            .field("config", &self.config)
            .field("ppo", &"<PPO>")
            .finish()
    }
}

impl PpoStrategy {
    /// Create a new `PpoStrategy` from the given configuration.
    ///
    /// # Errors
    ///
    /// Returns [`MLError`] if the underlying [`PPO`] cannot be constructed.
    pub fn new(config: PPOConfig) -> Result<Self, MLError> {
        let ppo = PPO::new(config.clone())?;
        Ok(Self {
            config,
            ppo: RefCell::new(ppo),
        })
    }
}

impl ValidatableStrategy for PpoStrategy {
    /// Train the MLP PPO on the provided time series data.
    ///
    /// Iterates over consecutive bars, collecting trajectory steps via the
    /// actor's `sample_action` and critic's `forward`, then computes GAE
    /// advantages and runs a single PPO update.
    fn train(&mut self, data: &TimeSeriesData) -> Result<(), MLError> {
        let n = data.len().saturating_sub(1);
        if n == 0 {
            return Ok(());
        }

        let mut trajectory = Trajectory::new();
        let state_dim = self.config.state_dim;

        for i in 0..n {
            let features = data.features.get(i).cloned().unwrap_or_default();
            let state = pad_or_truncate(&features, state_dim);

            let ppo = self.ppo.borrow();
            let (action, log_prob) = ppo.actor.sample_action(&state)?;

            // Critic forward on host via forward_host
            let value_vec = ppo.critic.forward_host(&state, 1)?;
            let value = value_vec.first().copied().unwrap_or(0.0);
            drop(ppo);

            let reward = data.returns.get(i).copied().unwrap_or(0.0) as f32;
            let done = i == n - 1;

            trajectory.add_step(TrajectoryStep::new(state, action, log_prob, value, reward, done));
        }

        // Compute GAE advantages
        let trajectories = vec![trajectory];
        let (advantages, returns) = compute_gae(&trajectories, &self.config.gae_config)?;

        // Build batch and run one PPO update
        let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
        self.ppo.borrow_mut().update(&mut batch)?;

        Ok(())
    }

    /// Evaluate the MLP PPO on out-of-sample data.
    ///
    /// For each bar, computes action probabilities via the actor, takes the
    /// argmax action, and maps it to PnL using the factored action's target
    /// exposure (e.g., Long100 -> +1.0, Flat -> 0.0, Short100 -> -1.0).
    fn evaluate(&self, data: &TimeSeriesData) -> Result<Vec<f64>, MLError> {
        let num_returns = data.returns.len();
        let mut pnl = Vec::with_capacity(num_returns);
        let state_dim = self.config.state_dim;

        for i in 0..num_returns {
            let features = data.features.get(i).cloned().unwrap_or_default();
            let state = pad_or_truncate(&features, state_dim);

            let ppo = self.ppo.borrow();
            // Use actor's forward_host to get logits, then argmax for greedy action
            let logits = ppo.actor.forward_host(&state, 1)?;
            drop(ppo);

            let action_idx = argmax(&logits);
            let bar_return = data.returns.get(i).copied().unwrap_or(0.0);

            let bar_pnl = if let Ok(action) = FactoredAction::from_index(action_idx) {
                action.target_exposure() * bar_return
            } else {
                0.0
            };

            pnl.push(bar_pnl);
        }

        Ok(pnl)
    }

    fn name(&self) -> &str {
        "PPO"
    }

    /// Reset the PPO agent to a fresh state by reconstructing it from config.
    fn reset(&mut self) -> Result<(), MLError> {
        let new_ppo = PPO::new(self.config.clone())?;
        self.ppo = RefCell::new(new_ppo);
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// PpoLstmStrategy (LSTM)
// ---------------------------------------------------------------------------

/// Adapter that wraps an LSTM-based [`PPO`] agent behind the [`ValidatableStrategy`] trait.
///
/// Maintains hidden state across bars during both training and evaluation,
/// resetting at episode boundaries.
pub struct PpoLstmStrategy {
    config: PPOConfig,
    ppo: RefCell<PPO>,
}

impl std::fmt::Debug for PpoLstmStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PpoLstmStrategy")
            .field("config", &self.config)
            .field("ppo", &"<PPO-LSTM>")
            .finish()
    }
}

impl PpoLstmStrategy {
    /// Create a new `PpoLstmStrategy` from the given configuration.
    ///
    /// The config must have `use_lstm = true`.
    ///
    /// # Errors
    ///
    /// Returns [`MLError`] if the underlying LSTM PPO cannot be constructed.
    pub fn new(config: PPOConfig) -> Result<Self, MLError> {
        if !config.use_lstm {
            return Err(MLError::ConfigError(
                "PpoLstmStrategy requires config.use_lstm = true".to_owned(),
            ));
        }
        let ppo = PPO::new(config.clone())?;
        Ok(Self {
            config,
            ppo: RefCell::new(ppo),
        })
    }
}

impl ValidatableStrategy for PpoLstmStrategy {
    /// Train the LSTM PPO on the provided time series data.
    ///
    /// Maintains hidden states across bars via a [`HiddenStateManager`],
    /// using the LSTM forward paths for both actor and critic networks.
    fn train(&mut self, data: &TimeSeriesData) -> Result<(), MLError> {
        let n = data.len().saturating_sub(1);
        if n == 0 {
            return Ok(());
        }

        let mut trajectory = Trajectory::new();
        let state_dim = self.config.state_dim;
        let hidden_dim = self.config.lstm_hidden_dim;
        let num_layers = self.config.lstm_num_layers;

        // Create a local hidden state manager for trajectory collection
        let mut hsm = HiddenStateManager::with_defaults(num_layers, 1, hidden_dim)?;

        for i in 0..n {
            let features = data.features.get(i).cloned().unwrap_or_default();
            let state = pad_or_truncate(&features, state_dim);

            // Get hidden states (host-side f32 slices)
            let (policy_h, policy_c) = hsm.get_policy_state();
            let (value_h, value_c) = hsm.get_value_state();

            let ppo = self.ppo.borrow();

            // Actor forward (LSTM path) — returns (logits, new_h, new_c) as host Vec<f32>
            let (logits, new_policy_h, new_policy_c) = match &ppo.actor {
                ActorNetwork::LSTM(ref lstm_net) => {
                    lstm_net.forward(&state, policy_h, policy_c, 1)?
                }
                _ => {
                    return Err(MLError::ModelError(
                        "PpoLstmStrategy requires LSTM actor network".to_owned(),
                    ));
                }
            };

            // Critic forward (LSTM path) — returns (value, new_h, new_c) as host Vec<f32>
            let (value_out, new_value_h, new_value_c) = match &ppo.critic {
                CriticNetwork::LSTM(ref lstm_net) => {
                    lstm_net.forward(&state, value_h, value_c, 1)?
                }
                _ => {
                    return Err(MLError::ModelError(
                        "PpoLstmStrategy requires LSTM critic network".to_owned(),
                    ));
                }
            };
            drop(ppo);

            // Update hidden states (already host-side Vec<f32>, no detach needed)
            hsm.update_policy_state(new_policy_h, new_policy_c)?;
            hsm.update_value_state(new_value_h, new_value_c)?;

            // Sample action from logits via softmax + categorical
            let (action, log_prob) = sample_action_from_logits_host(&logits)?;

            let value = value_out.first().copied().unwrap_or(0.0);

            let reward = data.returns.get(i).copied().unwrap_or(0.0) as f32;
            let done = i == n - 1;

            trajectory.add_step(TrajectoryStep::new(state, action, log_prob, value, reward, done));
        }

        // Compute GAE advantages
        let trajectories = vec![trajectory];
        let (advantages, returns) = compute_gae(&trajectories, &self.config.gae_config)?;

        // Build batch and run one PPO update
        let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
        self.ppo.borrow_mut().update(&mut batch)?;

        Ok(())
    }

    /// Evaluate the LSTM PPO on out-of-sample data.
    ///
    /// Maintains hidden states across bars, using LSTM forward paths to
    /// compute action logits, then takes the argmax action.
    fn evaluate(&self, data: &TimeSeriesData) -> Result<Vec<f64>, MLError> {
        let num_returns = data.returns.len();
        let mut pnl = Vec::with_capacity(num_returns);
        let state_dim = self.config.state_dim;
        let hidden_dim = self.config.lstm_hidden_dim;
        let num_layers = self.config.lstm_num_layers;

        // Create a local hidden state manager for evaluation
        let mut hsm = HiddenStateManager::with_defaults(num_layers, 1, hidden_dim)?;

        for i in 0..num_returns {
            let features = data.features.get(i).cloned().unwrap_or_default();
            let state = pad_or_truncate(&features, state_dim);

            let (policy_h, policy_c) = hsm.get_policy_state();

            let ppo = self.ppo.borrow();

            // Actor forward (LSTM path) for greedy evaluation
            let (logits, new_policy_h, new_policy_c) = match &ppo.actor {
                ActorNetwork::LSTM(ref lstm_net) => {
                    lstm_net.forward(&state, policy_h, policy_c, 1)?
                }
                _ => {
                    return Err(MLError::ModelError(
                        "PpoLstmStrategy requires LSTM actor network".to_owned(),
                    ));
                }
            };
            drop(ppo);

            // Update hidden state for next step (host-side, no detach needed)
            hsm.update_policy_state(new_policy_h, new_policy_c)?;

            // Greedy action: softmax then argmax (host-side)
            let probs = softmax_host(&logits);
            let action_idx = argmax(&probs);
            let bar_return = data.returns.get(i).copied().unwrap_or(0.0);

            let bar_pnl = if let Ok(action) = FactoredAction::from_index(action_idx) {
                action.target_exposure() * bar_return
            } else {
                0.0
            };

            pnl.push(bar_pnl);
        }

        Ok(pnl)
    }

    fn name(&self) -> &str {
        "PPO-LSTM"
    }

    /// Reset the LSTM PPO agent to a fresh state by reconstructing it from config.
    fn reset(&mut self) -> Result<(), MLError> {
        let new_ppo = PPO::new(self.config.clone())?;
        self.ppo = RefCell::new(new_ppo);
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Pad or truncate a feature vector to exactly `target_dim` elements.
fn pad_or_truncate(features: &[f32], target_dim: usize) -> Vec<f32> {
    let mut result = vec![0.0_f32; target_dim];
    let copy_len = features.len().min(target_dim);
    if let Some(dst) = result.get_mut(..copy_len) {
        dst.copy_from_slice(features.get(..copy_len).unwrap_or_default());
    }
    result
}

/// Return the index of the maximum element in a slice.
fn argmax(values: &[f32]) -> usize {
    values
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(idx, _)| idx)
        .unwrap_or(0)
}

/// Compute softmax on a host-side f32 slice.
fn softmax_host(logits: &[f32]) -> Vec<f32> {
    if logits.is_empty() {
        return Vec::new();
    }
    let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits.iter().map(|&x| (x - max_val).exp()).collect();
    let sum: f32 = exps.iter().sum();
    if sum.abs() < f32::EPSILON {
        // Uniform fallback to avoid division by zero
        vec![1.0 / logits.len() as f32; logits.len()]
    } else {
        exps.iter().map(|&e| e / sum).collect()
    }
}

/// Sample an action from host-side logits via softmax + categorical sampling.
///
/// Returns `(FactoredAction, log_prob)`.
fn sample_action_from_logits_host(logits: &[f32]) -> Result<(FactoredAction, f32), MLError> {
    let probs = softmax_host(logits);

    // Categorical sampling
    let mut rng = rand::thread_rng();
    let sample: f32 = rng.gen();
    let mut cumulative = 0.0;

    for (i, &prob) in probs.iter().enumerate() {
        cumulative += prob;
        if sample <= cumulative {
            let action = FactoredAction::from_index(i)?;
            const EPSILON: f32 = 1e-8;
            let log_prob = (prob + EPSILON).ln();
            return Ok((action, log_prob));
        }
    }

    // Fallback to last valid action
    let last_idx = probs.len().saturating_sub(1);
    let action = FactoredAction::from_index(last_idx)?;
    const EPSILON: f32 = 1e-8;
    let log_prob = (probs.get(last_idx).copied().unwrap_or(EPSILON) + EPSILON).ln();
    Ok((action, log_prob))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{TimeZone, Utc};

    /// Helper: create N UTC timestamps starting from 2024-01-01.
    fn make_timestamps(n: usize) -> Vec<chrono::DateTime<Utc>> {
        (0..n)
            .map(|i| {
                Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0)
                    .single()
                    .unwrap_or_else(Utc::now)
                    + chrono::Duration::days(i as i64)
            })
            .collect()
    }

    /// Helper: create N dummy feature rows with small deterministic values.
    fn make_features(n: usize, dim: usize) -> Vec<Vec<f32>> {
        (0..n)
            .map(|i| {
                (0..dim)
                    .map(|j| ((i * 7 + j * 3) % 100) as f32 * 0.01)
                    .collect()
            })
            .collect()
    }

    /// Helper: create monotonically increasing prices.
    fn make_prices(n: usize) -> Vec<f64> {
        (0..n).map(|i| 100.0 + i as f64 * 0.5).collect()
    }

    fn make_test_data(n: usize, feature_dim: usize) -> TimeSeriesData {
        TimeSeriesData::new(make_timestamps(n), make_features(n, feature_dim), make_prices(n))
            .unwrap_or_else(|e| panic!("Failed to create test data: {}", e))
    }

    fn make_mlp_config() -> PPOConfig {
        PPOConfig {
            state_dim: 10,
            num_actions: 45,
            policy_hidden_dims: vec![16, 8],
            value_hidden_dims: vec![16, 8],
            batch_size: 4,
            mini_batch_size: 4,
            num_epochs: 2,
            policy_learning_rate: 1e-3,
            value_learning_rate: 1e-3,
            clip_epsilon: 0.2,
            value_loss_coeff: 1.0,
            entropy_coeff: 0.01,
            max_grad_norm: 0.5,
            gae_config: GAEConfig::default(),
            use_lstm: false,
            ..PPOConfig::default()
        }
    }

    fn make_lstm_config() -> PPOConfig {
        let mut config = make_mlp_config();
        config.use_lstm = true;
        config.lstm_hidden_dim = 16;
        config.lstm_num_layers = 1;
        config.lstm_sequence_length = 8;
        config
    }

    // --- MLP tests ---

    #[test]
    fn test_ppo_strategy_creation() {
        let config = make_mlp_config();
        let result = PpoStrategy::new(config);
        assert!(result.is_ok(), "PpoStrategy::new failed: {:?}", result.err());
    }

    #[test]
    fn test_ppo_strategy_train_and_evaluate() {
        let config = make_mlp_config();
        let mut strategy =
            PpoStrategy::new(config).unwrap_or_else(|e| panic!("Failed to create strategy: {}", e));

        let data = make_test_data(30, 10);

        // Train should succeed
        let train_result = strategy.train(&data);
        assert!(train_result.is_ok(), "train failed: {:?}", train_result.err());

        // Evaluate should return one PnL per return period (data.len() - 1)
        let eval_result = strategy.evaluate(&data);
        assert!(eval_result.is_ok(), "evaluate failed: {:?}", eval_result.err());

        let pnl = eval_result.unwrap_or_else(|e| panic!("evaluate error: {}", e));
        assert_eq!(
            pnl.len(),
            29,
            "Expected 29 PnL values (one per return), got {}",
            pnl.len()
        );
    }

    #[test]
    fn test_ppo_strategy_reset() {
        let config = make_mlp_config();
        let mut strategy =
            PpoStrategy::new(config).unwrap_or_else(|e| panic!("Failed to create strategy: {}", e));

        let data = make_test_data(30, 10);

        // Train
        let train_result = strategy.train(&data);
        assert!(train_result.is_ok(), "first train failed: {:?}", train_result.err());

        // Reset
        let reset_result = strategy.reset();
        assert!(reset_result.is_ok(), "reset failed: {:?}", reset_result.err());

        // Retrain after reset
        let retrain_result = strategy.train(&data);
        assert!(retrain_result.is_ok(), "retrain failed: {:?}", retrain_result.err());
    }

    #[test]
    fn test_ppo_strategy_name() {
        let config = make_mlp_config();
        let strategy =
            PpoStrategy::new(config).unwrap_or_else(|e| panic!("Failed to create strategy: {}", e));
        assert_eq!(strategy.name(), "PPO");
    }

    // --- LSTM tests ---

    #[test]
    fn test_ppo_lstm_strategy_creation() {
        let config = make_lstm_config();
        let result = PpoLstmStrategy::new(config);
        assert!(
            result.is_ok(),
            "PpoLstmStrategy::new failed: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_ppo_lstm_strategy_train_and_evaluate() {
        let config = make_lstm_config();
        let mut strategy = PpoLstmStrategy::new(config)
            .unwrap_or_else(|e| panic!("Failed to create LSTM strategy: {}", e));

        let data = make_test_data(30, 10);

        // Train should succeed
        let train_result = strategy.train(&data);
        assert!(
            train_result.is_ok(),
            "LSTM train failed: {:?}",
            train_result.err()
        );

        // Evaluate should return one PnL per return period
        let eval_result = strategy.evaluate(&data);
        assert!(
            eval_result.is_ok(),
            "LSTM evaluate failed: {:?}",
            eval_result.err()
        );

        let pnl = eval_result.unwrap_or_else(|e| panic!("LSTM evaluate error: {}", e));
        assert_eq!(
            pnl.len(),
            29,
            "Expected 29 PnL values (one per return), got {}",
            pnl.len()
        );
    }

    #[test]
    fn test_ppo_lstm_strategy_reset() {
        let config = make_lstm_config();
        let mut strategy = PpoLstmStrategy::new(config)
            .unwrap_or_else(|e| panic!("Failed to create LSTM strategy: {}", e));

        let data = make_test_data(30, 10);

        // Train
        let train_result = strategy.train(&data);
        assert!(train_result.is_ok(), "LSTM first train failed: {:?}", train_result.err());

        // Reset
        let reset_result = strategy.reset();
        assert!(reset_result.is_ok(), "LSTM reset failed: {:?}", reset_result.err());

        // Retrain after reset
        let retrain_result = strategy.train(&data);
        assert!(retrain_result.is_ok(), "LSTM retrain failed: {:?}", retrain_result.err());
    }

    #[test]
    fn test_ppo_lstm_strategy_name() {
        let config = make_lstm_config();
        let strategy = PpoLstmStrategy::new(config)
            .unwrap_or_else(|e| panic!("Failed to create LSTM strategy: {}", e));
        assert_eq!(strategy.name(), "PPO-LSTM");
    }

    #[test]
    fn test_ppo_lstm_requires_use_lstm_flag() {
        let mut config = make_mlp_config();
        config.use_lstm = false;
        let result = PpoLstmStrategy::new(config);
        assert!(
            result.is_err(),
            "PpoLstmStrategy should reject config with use_lstm=false"
        );
    }
}
