// ml/src/features/barrier_optimization.rs
//
// Barrier Optimization Engine
// Optimizes triple barrier parameters via grid search + Sharpe maximization

use std::fmt;
use std::time::Instant;

/// Parameters for triple barrier labeling
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BarrierParams {
    pub profit_factor: f64,
    pub stop_factor: f64,
    pub time_horizon: usize,
}

impl BarrierParams {
    /// Create new barrier parameters with validation
    pub fn new(profit_factor: f64, stop_factor: f64, time_horizon: usize) -> Self {
        assert!(
            profit_factor > 0.0,
            "profit_factor must be positive, got {}",
            profit_factor
        );
        assert!(
            stop_factor > 0.0,
            "stop_factor must be positive, got {}",
            stop_factor
        );
        assert!(
            time_horizon >= 1,
            "time_horizon must be at least 1, got {}",
            time_horizon
        );

        Self {
            profit_factor,
            stop_factor,
            time_horizon,
        }
    }
}

impl Default for BarrierParams {
    fn default() -> Self {
        Self {
            profit_factor: 2.0,
            stop_factor: 1.0,
            time_horizon: 10,
        }
    }
}

/// Result of barrier optimization
#[derive(Debug, Clone)]
pub struct OptimizationResult {
    pub best_params: BarrierParams,
    pub best_sharpe: f64,
    pub evaluations: usize,
    pub duration_ms: u128,
}

impl fmt::Display for OptimizationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "OptimizationResult {{ \
             profit_factor: {:.2}, \
             stop_factor: {:.2}, \
             time_horizon: {}, \
             sharpe: {:.4}, \
             evaluations: {}, \
             duration_ms: {} }}",
            self.best_params.profit_factor,
            self.best_params.stop_factor,
            self.best_params.time_horizon,
            self.best_sharpe,
            self.evaluations,
            self.duration_ms
        )
    }
}

/// Barrier parameter optimizer using grid search
#[derive(Debug)]
pub struct BarrierOptimizer {
    profit_range: Vec<f64>,
    stop_range: Vec<f64>,
    horizon_range: Vec<usize>,
}

impl BarrierOptimizer {
    /// Create optimizer with default search ranges
    pub fn new() -> Self {
        Self {
            profit_range: vec![1.0, 1.5, 2.0, 2.5, 3.0],
            stop_range: vec![0.5, 1.0, 1.5, 2.0],
            horizon_range: vec![5, 10, 20, 30],
        }
    }

    /// Create optimizer with custom search ranges
    pub const fn with_ranges(
        profit_range: Vec<f64>,
        stop_range: Vec<f64>,
        horizon_range: Vec<usize>,
    ) -> Self {
        Self {
            profit_range,
            stop_range,
            horizon_range,
        }
    }

    /// Get profit factor search range
    pub fn profit_range(&self) -> &[f64] {
        &self.profit_range
    }

    /// Get stop factor search range
    pub fn stop_range(&self) -> &[f64] {
        &self.stop_range
    }

    /// Get time horizon search range
    pub fn horizon_range(&self) -> &[usize] {
        &self.horizon_range
    }

    /// Get total number of parameter combinations
    pub fn total_combinations(&self) -> usize {
        self.profit_range.len() * self.stop_range.len() * self.horizon_range.len()
    }

    /// Optimize barrier parameters using grid search
    ///
    /// # Arguments
    /// * `prices` - Historical price data for backtesting
    ///
    /// # Returns
    /// Optimal parameters and Sharpe ratio
    pub fn optimize(&self, prices: &[f64]) -> OptimizationResult {
        let start = Instant::now();

        let mut best_sharpe = f64::NEG_INFINITY;
        let mut best_params = BarrierParams::default();
        let mut evaluations = 0;

        // Grid search over all parameter combinations
        for &profit in &self.profit_range {
            for &stop in &self.stop_range {
                for &horizon in &self.horizon_range {
                    let params = BarrierParams::new(profit, stop, horizon);
                    let sharpe = self.backtest_params(&params, prices);

                    evaluations += 1;

                    if sharpe > best_sharpe && sharpe.is_finite() {
                        best_sharpe = sharpe;
                        best_params = params;
                    }
                }
            }
        }

        // If no valid Sharpe found, use 0.0
        if !best_sharpe.is_finite() {
            best_sharpe = 0.0;
        }

        let duration = start.elapsed();

        OptimizationResult {
            best_params,
            best_sharpe,
            evaluations,
            duration_ms: duration.as_millis(),
        }
    }

    /// Backtest specific barrier parameters and return Sharpe ratio
    ///
    /// # Arguments
    /// * `params` - Barrier parameters to test
    /// * `prices` - Historical price data
    ///
    /// # Returns
    /// Sharpe ratio for the given parameters
    pub fn backtest_params(&self, params: &BarrierParams, prices: &[f64]) -> f64 {
        // Handle edge cases
        if prices.is_empty() || prices.len() < 2 {
            return 0.0;
        }

        // Filter out NaN and infinite values
        // Preallocate: worst case is all prices are clean
        let mut clean_prices = Vec::with_capacity(prices.len());
        clean_prices.extend(prices.iter().copied().filter(|&p| p.is_finite()));

        if clean_prices.len() < 2 {
            return 0.0;
        }

        // Generate trading signals using triple barrier logic
        let returns = self.simulate_triple_barrier_trading(params, &clean_prices);

        // Calculate Sharpe ratio
        self.calculate_sharpe(&returns)
    }

    /// Simulate triple barrier trading strategy
    ///
    /// For each entry point, we:
    /// 1. Calculate profit target: entry * (1 + `profit_factor` * volatility)
    /// 2. Calculate stop loss: entry * (1 - `stop_factor` * volatility)
    /// 3. Hold for up to `time_horizon` periods
    /// 4. Exit when price hits barrier or horizon reached
    fn simulate_triple_barrier_trading(&self, params: &BarrierParams, prices: &[f64]) -> Vec<f64> {
        // Preallocate capacity: estimate max trades as prices.len() / time_horizon
        // This prevents unbounded growth and reduces allocations
        let estimated_trades = prices.len().saturating_div(params.time_horizon.max(1));
        let mut returns = Vec::with_capacity(estimated_trades);

        // Need at least time_horizon + 1 prices for meaningful backtest
        if prices.len() < params.time_horizon + 1 {
            return returns;
        }

        // Calculate rolling volatility (using simple std dev over 20 periods)
        let vol_window = 20.min(prices.len() / 2);

        // Iterate through potential entry points
        let mut i = vol_window;
        while i < prices.len() - params.time_horizon {
            let entry_price = prices[i];

            // Calculate volatility from recent price changes
            let volatility = self.calculate_volatility(&prices[i.saturating_sub(vol_window)..=i]);

            if volatility <= 0.0 || !volatility.is_finite() {
                i += 1;
                continue;
            }

            // Set barriers
            let profit_target = entry_price * (1.0 + params.profit_factor * volatility);
            let stop_loss = entry_price * (1.0 - params.stop_factor * volatility);

            // Simulate holding period
            let mut exit_price = entry_price;
            let max_horizon = (i + params.time_horizon).min(prices.len() - 1);

            for &current_price in &prices[(i + 1)..=max_horizon] {

                // Check if profit target hit
                if current_price >= profit_target {
                    exit_price = profit_target;
                    break;
                }

                // Check if stop loss hit
                if current_price <= stop_loss {
                    exit_price = stop_loss;
                    break;
                }

                // Update exit price (will use this if horizon reached)
                exit_price = current_price;
            }

            // Calculate return
            let trade_return = (exit_price - entry_price) / entry_price;
            returns.push(trade_return);

            // Move to next entry point (skip ahead to avoid overlapping trades)
            i += params.time_horizon;
        }

        returns
    }

    /// Calculate volatility as standard deviation of returns
    fn calculate_volatility(&self, prices: &[f64]) -> f64 {
        if prices.len() < 2 {
            return 0.0;
        }

        // Calculate returns
        // Preallocate: max returns is prices.len() - 1 (one per window)
        let mut returns = Vec::with_capacity(prices.len().saturating_sub(1));
        returns.extend(
            prices
                .windows(2)
                .map(|w| (w[1] - w[0]) / w[0])
                .filter(|r| r.is_finite()),
        );

        if returns.is_empty() {
            return 0.0;
        }

        // Calculate standard deviation
        let mean = returns.iter().sum::<f64>() / returns.len() as f64;
        let variance =
            returns.iter().map(|r| (r - mean).powi(2)).sum::<f64>() / returns.len() as f64;
        variance.sqrt()
    }

    /// Calculate Sharpe ratio from returns
    ///
    /// Sharpe = (`mean_return` - `risk_free_rate`) / `std_dev_return`
    /// Assuming `risk_free_rate` = 0 for simplicity
    pub fn calculate_sharpe(&self, returns: &[f64]) -> f64 {
        if returns.is_empty() {
            return 0.0;
        }

        // Filter out non-finite values
        // Preallocate: worst case is all returns are clean
        let mut clean_returns = Vec::with_capacity(returns.len());
        clean_returns.extend(returns.iter().copied().filter(|r| r.is_finite()));

        if clean_returns.is_empty() {
            return 0.0;
        }

        // Calculate mean return
        let mean_return = clean_returns.iter().sum::<f64>() / clean_returns.len() as f64;

        // Calculate standard deviation of returns
        let variance = clean_returns
            .iter()
            .map(|r| (r - mean_return).powi(2))
            .sum::<f64>()
            / clean_returns.len() as f64;

        let std_dev = variance.sqrt();

        // Handle zero volatility case
        if std_dev < 1e-10 {
            // If std_dev is near zero and mean is positive, return large Sharpe
            // If std_dev is near zero and mean is zero/negative, return 0
            if mean_return > 1e-10 {
                return 100.0; // Cap at reasonable value
            } else {
                return 0.0;
            }
        }

        // Calculate Sharpe ratio (annualized: multiply by sqrt(252) for daily data)
        // For now, return non-annualized Sharpe
        mean_return / std_dev
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_barrier_params_creation() {
        let params = BarrierParams::new(2.0, 1.0, 10);
        assert_eq!(params.profit_factor, 2.0);
        assert_eq!(params.stop_factor, 1.0);
        assert_eq!(params.time_horizon, 10);
    }

    #[test]
    #[should_panic]
    fn test_barrier_params_invalid_profit() {
        BarrierParams::new(-1.0, 1.0, 10);
    }

    #[test]
    fn test_optimizer_creation() {
        let optimizer = BarrierOptimizer::new();
        assert_eq!(optimizer.total_combinations(), 80); // 5 * 4 * 4
    }

    #[test]
    fn test_calculate_sharpe_basic() {
        let optimizer = BarrierOptimizer::new();
        let returns = vec![0.01, 0.02, -0.005, 0.015, 0.008];
        let sharpe = optimizer.calculate_sharpe(&returns);
        assert!(sharpe.is_finite());
    }

    #[test]
    fn test_calculate_volatility() {
        let optimizer = BarrierOptimizer::new();
        let prices = vec![100.0, 101.0, 99.0, 102.0, 98.0];
        let vol = optimizer.calculate_volatility(&prices);
        assert!(vol > 0.0);
        assert!(vol.is_finite());
    }

    #[test]
    fn test_optimize_simple() {
        let optimizer = BarrierOptimizer::new();
        let prices: Vec<f64> = (0..50).map(|i| 100.0 + i as f64).collect();
        let result = optimizer.optimize(&prices);
        assert!(result.best_sharpe.is_finite());
        assert_eq!(result.evaluations, 80);
    }
}
