//! DQN strategy adapter for the [`ValidatableStrategy`] trait.
//!
//! Wraps a [`DQN`] agent so it can be trained and evaluated through the
//! walk-forward validation harness, deflated Sharpe ratio, and PBO pipelines.

use std::cell::RefCell;
use std::cmp::min;

use crate::dqn::{DQNConfig, Experience, DQN};
use ml_validation::{TimeSeriesData, ValidatableStrategy};
use crate::MLError;

/// Adapter that wraps a [`DQN`] agent behind the [`ValidatableStrategy`] trait.
///
/// Uses `RefCell` for interior mutability because [`DQN::select_action`] requires
/// `&mut self` (it increments step counters and tracks actions), while
/// [`ValidatableStrategy::evaluate`] takes `&self`.
pub struct DqnStrategy {
    config: DQNConfig,
    /// Interior mutability: `DQN::select_action` requires `&mut self`
    /// but `ValidatableStrategy::evaluate` only receives `&self`.
    dqn: RefCell<DQN>,
}

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

impl DqnStrategy {
    /// Create a new `DqnStrategy` from the given configuration.
    ///
    /// # Errors
    ///
    /// Returns [`MLError`] if the underlying [`DQN`] cannot be constructed
    /// (e.g. invalid dimensions, GPU initialisation failure).
    pub fn new(config: DQNConfig) -> Result<Self, MLError> {
        let dqn = DQN::new(config.clone())?;
        Ok(Self {
            config,
            dqn: RefCell::new(dqn),
        })
    }

    /// Create a new `DqnStrategy` on a specific device.
    ///
    /// Shares the caller's device instead of creating a new cuBLAS handle.
    #[cfg(test)]
    pub(crate) fn new_on_device(config: DQNConfig, device: ml_core::device::MlDevice) -> Result<Self, MLError> {
        let dqn = DQN::new_on_device(config.clone(), device)?;
        Ok(Self {
            config,
            dqn: RefCell::new(dqn),
        })
    }
}

impl ValidatableStrategy for DqnStrategy {
    /// Train the DQN on the provided time series data.
    ///
    /// Iterates over consecutive bar pairs, collecting experiences into the
    /// replay buffer, then runs a bounded number of gradient-descent steps.
    fn train(&mut self, data: &TimeSeriesData) -> Result<(), MLError> {
        let n = data.len().saturating_sub(1);

        // Collect experiences
        for i in 0..n {
            let state = data.features.get(i).cloned().unwrap_or_default();
            let next_state = data.features.get(i + 1).cloned().unwrap_or_default();

            let action = self.dqn.borrow_mut().select_action(&state)?;
            let reward = data.returns.get(i).copied().unwrap_or(0.0) as f32;
            let done = i == data.len().saturating_sub(2);

            let exp = Experience::new(
                state,
                action.exposure as u8, // DQN exposure index (0-4), not factored index (0-44)
                reward,
                next_state,
                done,
            );
            self.dqn.borrow().store_experience(exp)?;
        }

        // Run bounded training steps
        let batch_size = if self.config.batch_size == 0 {
            1
        } else {
            self.config.batch_size
        };
        let max_steps = min(data.len() / batch_size, 100);

        for _ in 0..max_steps {
            match self.dqn.borrow_mut().train_step(None) {
                Ok(_) => {}
                Err(e) => {
                    let msg = format!("{}", e);
                    if msg.contains("Not enough") {
                        break;
                    }
                    return Err(e);
                }
            }
        }

        Ok(())
    }

    /// Evaluate the DQN on out-of-sample data.
    ///
    /// For each return period, selects an action and maps it to a PnL value.
    /// Returns one PnL value per return in `data.returns` (length = data.len() - 1).
    ///
    /// Action-to-exposure mapping depends on `num_actions`:
    /// - **num_actions <= 5** (simple): direct mapping (0=short, mid=flat, last=long)
    /// - **num_actions = 45** (factored): uses FactoredAction exposure decomposition
    fn evaluate(&self, data: &TimeSeriesData) -> Result<Vec<f64>, MLError> {
        let num_returns = data.returns.len();
        let mut pnl = Vec::with_capacity(num_returns);

        let num_actions = self.config.num_actions;

        for i in 0..num_returns {
            let state = data.features.get(i).cloned().unwrap_or_default();
            let action = self.dqn.borrow_mut().select_action(&state)?;
            let bar_return = data.returns.get(i).copied().unwrap_or(0.0);

            let bar_pnl = if num_actions <= 5 {
                // Simple action space: map action index directly to direction.
                // For num_actions=3: 0=short, 1=flat, 2=long
                // For num_actions=5: 0=short100, 1=short50, 2=flat, 3=long50, 4=long100
                let action_idx = action.exposure as usize;
                let mid = num_actions / 2;
                if action_idx < mid {
                    -bar_return // Short
                } else if action_idx > mid {
                    bar_return // Long
                } else {
                    0.0 // Flat
                }
            } else {
                // Factored action space (45 actions): use exposure level
                let exposure_idx = action.exposure as usize;
                match exposure_idx {
                    0 | 1 => -bar_return, // Short exposure
                    3 | 4 => bar_return,  // Long exposure
                    _ => 0.0,             // Flat
                }
            };

            pnl.push(bar_pnl);
        }

        Ok(pnl)
    }

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

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

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

    static SHARED_CUDA: OnceLock<ml_core::device::MlDevice> = OnceLock::new();

    fn shared_device() -> ml_core::device::MlDevice {
        SHARED_CUDA
            .get_or_init(|| {
                ml_core::device::MlDevice::cuda(0).expect("CUDA required for tests")
            })
            .clone()
    }

    /// 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 random-ish 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 (ensures positive log returns).
    fn make_prices(n: usize) -> Vec<f64> {
        (0..n).map(|i| 100.0 + i as f64 * 0.5).collect()
    }

    fn make_test_config() -> DQNConfig {
        let mut config = DQNConfig::default();
        config.state_dim = 10;
        config.num_actions = 3;
        config.hidden_dims = vec![16, 8];
        config.batch_size = 4;
        config.min_replay_size = 4;
        config.warmup_steps = 0;
        config.use_iqn = false;
        config.epsilon_start = 0.5;
        config
    }

    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))
    }

    #[test]
    fn test_dqn_strategy_creation() {
        let config = make_test_config();
        let result = DqnStrategy::new_on_device(config, shared_device());
        assert!(result.is_ok(), "DqnStrategy::new failed: {:?}", result.err());
    }

    #[test]
    fn test_dqn_strategy_train_and_evaluate() {
        let config = make_test_config();
        let mut strategy = DqnStrategy::new_on_device(config, shared_device())
            .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_dqn_strategy_reset() {
        let config = make_test_config();
        let mut strategy = DqnStrategy::new_on_device(config, shared_device())
            .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_dqn_strategy_name() {
        let config = make_test_config();
        let strategy = DqnStrategy::new_on_device(config, shared_device())
            .unwrap_or_else(|e| panic!("Failed to create strategy: {}", e));
        assert_eq!(strategy.name(), "DQN");
    }
}
