//! Expert Demonstration Generation for DQN Imitation Learning
//!
//! Generates expert action labels from a simple MA crossover strategy,
//! filtered by ADX trend strength. These demonstrations can seed the
//! replay buffer during early training epochs to accelerate convergence.
//!
//! ## Strategy
//!
//! - **Fast EMA** (default 20-period) crosses above **Slow EMA** (default 50-period)
//!   AND ADX > threshold => Long100 (exposure index 4 in 5-action space, mapped to 8
//!   in 45-action factored space)
//! - Fast EMA crosses below Slow EMA AND ADX > threshold => Short100 (exposure index 0)
//! - Otherwise => Flat (exposure index 2, mapped to 4 in factored space)
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml::trainers::dqn::expert_demos::ExpertDemoGenerator;
//! use ml_core::types::OHLCVBar;
//!
//! let generator = ExpertDemoGenerator::default();
//! let bars: Vec<OHLCVBar> = vec![]; // load bars
//! let actions = generator.generate_actions(&bars);
//! // actions: Vec<(usize, i32)> = [(bar_index, exposure_action_index), ...]
//! ```

use crate::types::OHLCVBar;

/// Expert demonstration generator using MA crossover filtered by ADX.
///
/// Produces (bar_index, exposure_action_index) pairs that can be inserted
/// into the replay buffer as expert demonstrations for imitation learning.
#[derive(Debug, Clone)]
pub struct ExpertDemoGenerator {
    /// Fast EMA period (default: 20).
    pub fast_period: usize,
    /// Slow EMA period (default: 50).
    pub slow_period: usize,
    /// Minimum ADX value to trigger a directional signal (default: 25.0).
    pub adx_threshold: f64,
}

impl Default for ExpertDemoGenerator {
    fn default() -> Self {
        Self {
            fast_period: 20,
            slow_period: 50,
            adx_threshold: 25.0,
        }
    }
}

/// Exposure action indices in the 5-action exposure space.
/// These map to the branching DQN exposure dimension:
///   0 = Short100, 1 = Short50, 2 = Flat, 3 = Long50, 4 = Long100
const ACTION_SHORT100: i32 = 0;
const ACTION_FLAT: i32 = 2;
const ACTION_LONG100: i32 = 4;

impl ExpertDemoGenerator {
    /// Create a new generator with custom parameters.
    pub fn new(fast_period: usize, slow_period: usize, adx_threshold: f64) -> Self {
        Self {
            fast_period,
            slow_period,
            adx_threshold,
        }
    }

    /// Generate expert action labels for each bar based on MA crossover + ADX filter.
    ///
    /// Returns `(bar_index, exposure_action_index)` pairs for every bar that has
    /// sufficient history for both the slow EMA and ADX(14) computation.
    ///
    /// # Arguments
    ///
    /// * `bars` - Chronologically sorted OHLCV bars.
    ///
    /// # Returns
    ///
    /// A vector of `(bar_index, exposure_action_index)` where:
    /// - `bar_index` is the 0-based index into `bars`
    /// - `exposure_action_index` is one of: 0 (Short100), 2 (Flat), 4 (Long100)
    pub fn generate_actions(&self, bars: &[OHLCVBar]) -> Vec<(usize, i32)> {
        if bars.is_empty() {
            return Vec::new();
        }

        // Pre-compute close prices for EMA and ADX
        let closes: Vec<f64> = bars.iter().map(|b| b.close).collect();

        // Compute fast and slow EMAs
        let fast_ema = Self::compute_ema(&closes, self.fast_period);
        let slow_ema = Self::compute_ema(&closes, self.slow_period);

        // Compute ADX(14)
        let adx_values = self.compute_adx_series(bars);

        // Minimum index where all indicators are valid
        let adx_period: usize = 14;
        // ADX needs 2*period bars of history, EMA needs slow_period bars
        let min_valid = self
            .slow_period
            .max(adx_period.saturating_mul(2))
            .saturating_add(1);

        let mut actions = Vec::with_capacity(bars.len().saturating_sub(min_valid));

        for i in min_valid..bars.len() {
            let fast = match fast_ema.get(i) {
                Some(&v) => v,
                None => continue,
            };
            let slow = match slow_ema.get(i) {
                Some(&v) => v,
                None => continue,
            };
            let adx = match adx_values.get(i) {
                Some(&v) => v,
                None => continue,
            };

            let action = if fast > slow && adx > self.adx_threshold {
                ACTION_LONG100
            } else if fast < slow && adx > self.adx_threshold {
                ACTION_SHORT100
            } else {
                ACTION_FLAT
            };

            actions.push((i, action));
        }

        actions
    }

    /// Compute the effective expert demo ratio at a given epoch, applying linear decay.
    ///
    /// The ratio decays linearly from `initial_ratio` to 0 over `decay_epochs`.
    /// After `decay_epochs`, returns 0.0 (pure RL).
    ///
    /// # Arguments
    ///
    /// * `initial_ratio` - Starting expert demo ratio (from DQNHyperparameters).
    /// * `decay_epochs` - Number of epochs for full decay.
    /// * `current_epoch` - Current training epoch (0-indexed).
    pub fn effective_ratio(initial_ratio: f64, decay_epochs: usize, current_epoch: usize) -> f64 {
        if decay_epochs == 0 || current_epoch >= decay_epochs {
            return 0.0;
        }
        initial_ratio * (1.0 - (current_epoch as f64 / decay_epochs as f64))
    }

    /// Compute Exponential Moving Average for a price series.
    ///
    /// Returns a Vec of the same length as `prices`. Values before `period` are
    /// set to simple average bootstrapping.
    fn compute_ema(prices: &[f64], period: usize) -> Vec<f64> {
        let mut ema = vec![0.0_f64; prices.len()];
        if prices.is_empty() || period == 0 {
            return ema;
        }

        let k = 2.0 / (period as f64 + 1.0);

        // Bootstrap: SMA of first `period` values
        if prices.len() < period {
            // Not enough data for even one EMA value
            return ema;
        }

        let initial_sma: f64 = prices.iter().take(period).sum::<f64>() / period as f64;
        if let Some(slot) = ema.get_mut(period.saturating_sub(1)) {
            *slot = initial_sma;
        }

        // Fill EMA forward
        for i in period..prices.len() {
            let price = match prices.get(i) {
                Some(&p) => p,
                None => continue,
            };
            let prev = match ema.get(i.wrapping_sub(1)) {
                Some(&p) => p,
                None => continue,
            };
            if let Some(slot) = ema.get_mut(i) {
                *slot = price * k + prev * (1.0 - k);
            }
        }

        ema
    }

    /// Compute ADX(14) series for all bars.
    ///
    /// Returns a Vec of the same length as `bars`. Values before sufficient
    /// history are 0.0.
    fn compute_adx_series(&self, bars: &[OHLCVBar]) -> Vec<f64> {
        let period: usize = 14;
        let mut adx_series = vec![0.0_f64; bars.len()];

        if bars.len() < period.saturating_add(1) {
            return adx_series;
        }

        // Compute TR, +DM, -DM
        let mut tr_values = Vec::with_capacity(bars.len().saturating_sub(1));
        let mut plus_dm_values = Vec::with_capacity(bars.len().saturating_sub(1));
        let mut minus_dm_values = Vec::with_capacity(bars.len().saturating_sub(1));

        for i in 1..bars.len() {
            let prev = match bars.get(i.wrapping_sub(1)) {
                Some(b) => b,
                None => continue,
            };
            let curr = match bars.get(i) {
                Some(b) => b,
                None => continue,
            };

            let hl = curr.high - curr.low;
            let hpc = (curr.high - prev.close).abs();
            let lpc = (curr.low - prev.close).abs();
            let tr = hl.max(hpc).max(lpc);
            tr_values.push(tr);

            let up_move = curr.high - prev.high;
            let down_move = prev.low - curr.low;

            let plus_dm = if up_move > down_move && up_move > 0.0 {
                up_move
            } else {
                0.0
            };
            let minus_dm = if down_move > up_move && down_move > 0.0 {
                down_move
            } else {
                0.0
            };

            plus_dm_values.push(plus_dm);
            minus_dm_values.push(minus_dm);
        }

        if tr_values.len() < period {
            return adx_series;
        }

        // Wilder smoothing
        let mut smoothed_tr: f64 = tr_values.iter().take(period).sum();
        let mut smoothed_plus_dm: f64 = plus_dm_values.iter().take(period).sum();
        let mut smoothed_minus_dm: f64 = minus_dm_values.iter().take(period).sum();
        let period_f64 = period as f64;

        let mut dx_values: Vec<(usize, f64)> = Vec::new(); // (bar_index, dx)

        // First DX at bar index = period
        if smoothed_tr.abs() > 1e-12 {
            let plus_di = 100.0 * smoothed_plus_dm / smoothed_tr;
            let minus_di = 100.0 * smoothed_minus_dm / smoothed_tr;
            let di_sum = plus_di + minus_di;
            if di_sum.abs() > 1e-12 {
                let dx = 100.0 * (plus_di - minus_di).abs() / di_sum;
                dx_values.push((period, dx));
            }
        }

        for i in period..tr_values.len() {
            let tr_val = match tr_values.get(i) {
                Some(&v) => v,
                None => continue,
            };
            let pdm_val = match plus_dm_values.get(i) {
                Some(&v) => v,
                None => continue,
            };
            let mdm_val = match minus_dm_values.get(i) {
                Some(&v) => v,
                None => continue,
            };

            smoothed_tr = smoothed_tr - (smoothed_tr / period_f64) + tr_val;
            smoothed_plus_dm = smoothed_plus_dm - (smoothed_plus_dm / period_f64) + pdm_val;
            smoothed_minus_dm = smoothed_minus_dm - (smoothed_minus_dm / period_f64) + mdm_val;

            if smoothed_tr.abs() > 1e-12 {
                let plus_di = 100.0 * smoothed_plus_dm / smoothed_tr;
                let minus_di = 100.0 * smoothed_minus_dm / smoothed_tr;
                let di_sum = plus_di + minus_di;
                if di_sum.abs() > 1e-12 {
                    let dx = 100.0 * (plus_di - minus_di).abs() / di_sum;
                    // bar index = i + 1 (because tr_values[0] corresponds to bars[1])
                    dx_values.push((i.saturating_add(1), dx));
                }
            }
        }

        if dx_values.len() < period {
            return adx_series;
        }

        // Smooth DX to get ADX
        let initial_adx: f64 =
            dx_values.iter().take(period).map(|(_, dx)| dx).sum::<f64>() / period_f64;
        let first_adx_bar = match dx_values.get(period.saturating_sub(1)) {
            Some(&(idx, _)) => idx,
            None => return adx_series,
        };
        if let Some(slot) = adx_series.get_mut(first_adx_bar) {
            *slot = initial_adx;
        }

        let mut adx = initial_adx;
        for i in period..dx_values.len() {
            let (bar_idx, dx) = match dx_values.get(i) {
                Some(&v) => v,
                None => continue,
            };
            adx = (adx * (period_f64 - 1.0) + dx) / period_f64;
            if let Some(slot) = adx_series.get_mut(bar_idx) {
                *slot = adx;
            }
        }

        // Forward-fill ADX values for bars between computed points
        let mut last_adx = 0.0_f64;
        for val in &mut adx_series {
            if *val > 0.0 {
                last_adx = *val;
            } else if last_adx > 0.0 {
                *val = last_adx;
            }
        }

        adx_series
    }
}

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

    /// Create a bar with specified OHLCV values at the given day offset.
    fn make_bar_with_prices(
        day_offset: u32,
        open: f64,
        high: f64,
        low: f64,
        close: f64,
    ) -> OHLCVBar {
        let date =
            NaiveDate::from_ymd_opt(2024, 1, 1).unwrap_or_default() + chrono::Duration::days(i64::from(day_offset));
        let time = NaiveTime::from_hms_opt(10, 0, 0).unwrap_or_default();
        let dt = date.and_time(time);
        let timestamp = Utc.from_utc_datetime(&dt);
        OHLCVBar {
            timestamp,
            open,
            high,
            low,
            close,
            volume: 1000.0,
        }
    }

    /// Generate a trending market (prices steadily increasing).
    fn make_trending_bars(count: usize) -> Vec<OHLCVBar> {
        (0..count)
            .map(|i| {
                let base = 100.0 + i as f64 * 0.5;
                make_bar_with_prices(i as u32, base, base + 1.0, base - 0.3, base + 0.4)
            })
            .collect()
    }

    /// Generate a ranging market (prices oscillating around a mean).
    fn make_ranging_bars(count: usize) -> Vec<OHLCVBar> {
        // Tight range-bound oscillation: close ≈ open, tiny HL range.
        // ADX should be very low (< 25) for this flat-line data.
        (0..count)
            .map(|i| {
                let base = 100.0 + (i as f64 * 0.3).sin() * 0.5; // tiny oscillation
                let open = base;
                let close = base + (i as f64 * 0.7).cos() * 0.1; // close ≈ open
                let high = base + 0.3;
                let low = base - 0.3;
                make_bar_with_prices(i as u32, open, high, low, close)
            })
            .collect()
    }

    #[test]
    fn test_default_generator() {
        let gen = ExpertDemoGenerator::default();
        assert_eq!(gen.fast_period, 20);
        assert_eq!(gen.slow_period, 50);
        assert!((gen.adx_threshold - 25.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_generate_actions_empty_bars() {
        let gen = ExpertDemoGenerator::default();
        let actions = gen.generate_actions(&[]);
        assert!(actions.is_empty());
    }

    #[test]
    fn test_generate_actions_insufficient_bars() {
        let gen = ExpertDemoGenerator::default();
        // Only 10 bars -- not enough for slow EMA (50) or ADX (28)
        let bars = make_trending_bars(10);
        let actions = gen.generate_actions(&bars);
        assert!(actions.is_empty(), "Expected no actions with insufficient bars");
    }

    #[test]
    fn test_generate_actions_trending_market() {
        let gen = ExpertDemoGenerator::default();
        // 200 bars of trending data -- should produce some Long100 signals
        let bars = make_trending_bars(200);
        let actions = gen.generate_actions(&bars);

        // Should have at least some actions after warmup
        assert!(
            !actions.is_empty(),
            "Expected actions for 200-bar trending market"
        );

        // All bar indices should be within bounds
        for &(idx, action) in &actions {
            assert!(idx < bars.len(), "Bar index out of bounds");
            assert!(
                action == ACTION_SHORT100 || action == ACTION_FLAT || action == ACTION_LONG100,
                "Invalid action: {action}"
            );
        }
    }

    #[test]
    fn test_generate_actions_ranging_market_mostly_flat() {
        let gen = ExpertDemoGenerator::default();
        let bars = make_ranging_bars(200);
        let actions = gen.generate_actions(&bars);

        // In a ranging market, most actions should be Flat (ADX likely < 25)
        let flat_count = actions.iter().filter(|&&(_, a)| a == ACTION_FLAT).count();
        let total = actions.len();
        if total > 0 {
            let flat_ratio = flat_count as f64 / total as f64;
            assert!(
                flat_ratio > 0.5,
                "Expected majority flat in ranging market, got {flat_ratio:.2}"
            );
        }
    }

    #[test]
    fn test_effective_ratio_decay() {
        // At epoch 0, should be full ratio
        let r0 = ExpertDemoGenerator::effective_ratio(0.5, 20, 0);
        assert!((r0 - 0.5).abs() < 1e-10);

        // At epoch 10 (halfway), should be half
        let r10 = ExpertDemoGenerator::effective_ratio(0.5, 20, 10);
        assert!((r10 - 0.25).abs() < 1e-10);

        // At epoch 20 (end), should be 0
        let r20 = ExpertDemoGenerator::effective_ratio(0.5, 20, 20);
        assert!(r20.abs() < 1e-10);

        // After decay_epochs, should be 0
        let r30 = ExpertDemoGenerator::effective_ratio(0.5, 20, 30);
        assert!(r30.abs() < 1e-10);
    }

    #[test]
    fn test_effective_ratio_zero_decay_epochs() {
        // Zero decay epochs should return 0 immediately
        let r = ExpertDemoGenerator::effective_ratio(0.5, 0, 0);
        assert!(r.abs() < 1e-10);
    }

    #[test]
    fn test_ema_computation() {
        // Simple ascending series -- EMA should follow
        let prices: Vec<f64> = (1..=100).map(|x| x as f64).collect();
        let ema = ExpertDemoGenerator::compute_ema(&prices, 10);

        // EMA at end should be close to the price but lag slightly
        let last_ema = match ema.last() {
            Some(&v) => v,
            None => panic!("Empty EMA"),
        };
        assert!(
            last_ema > 90.0 && last_ema < 100.0,
            "Expected EMA near end of ascending series, got {last_ema}"
        );
    }

    #[test]
    fn test_custom_generator() {
        let gen = ExpertDemoGenerator::new(10, 30, 20.0);
        assert_eq!(gen.fast_period, 10);
        assert_eq!(gen.slow_period, 30);
        assert!((gen.adx_threshold - 20.0).abs() < f64::EPSILON);
    }
}
