use super::{Indicator, IncrementalIndicator, IndicatorResult, IndicatorError, PriceData};
use super::common::RollingWindow;

/// Average True Range (ATR) Indicator
/// 
/// Measures volatility by calculating the average of true ranges over a period
/// True Range = max(High - Low, |High - Previous Close|, |Low - Previous Close|)
pub struct ATR {
    period: usize,
    atr_value: Option<f64>,
    prev_close: Option<f64>,
    true_ranges: RollingWindow<f64>,
    sum: f64,
    initialized: bool,
}

impl ATR {
    pub fn new(period: usize) -> Result<Self, IndicatorError> {
        if period == 0 {
            return Err(IndicatorError::InvalidParameter(
                "Period must be greater than 0".to_string()
            ));
        }
        
        Ok(Self {
            period,
            atr_value: None,
            prev_close: None,
            true_ranges: RollingWindow::new(period),
            sum: 0.0,
            initialized: false,
        })
    }
    
    /// Calculate True Range
    fn calculate_true_range(high: f64, low: f64, prev_close: Option<f64>) -> f64 {
        let high_low = high - low;
        
        match prev_close {
            Some(prev) => {
                let high_close = (high - prev).abs();
                let low_close = (low - prev).abs();
                high_low.max(high_close).max(low_close)
            }
            None => high_low,
        }
    }
    
    /// Calculate ATR for a series of price data
    pub fn calculate_series(
        high: &[f64],
        low: &[f64],
        close: &[f64],
        period: usize
    ) -> Result<Vec<f64>, IndicatorError> {
        if high.len() != low.len() || high.len() != close.len() {
            return Err(IndicatorError::InvalidParameter(
                "High, low, and close arrays must have the same length".to_string()
            ));
        }
        
        if high.len() < period + 1 {
            return Err(IndicatorError::InsufficientData {
                required: period + 1,
                actual: high.len(),
            });
        }
        
        let mut true_ranges = Vec::with_capacity(high.len() - 1);
        
        // Calculate true ranges
        for i in 1..high.len() {
            let tr = Self::calculate_true_range(high[i], low[i], Some(close[i - 1]));
            true_ranges.push(tr);
        }
        
        let mut atr_values = Vec::with_capacity(true_ranges.len() - period + 1);
        
        // Calculate initial ATR as SMA of first period true ranges
        let initial_atr: f64 = true_ranges[0..period].iter().sum::<f64>() / period as f64;
        atr_values.push(initial_atr);
        
        // Calculate subsequent ATRs using Wilder's smoothing
        let mut atr = initial_atr;
        for i in period..true_ranges.len() {
            // Wilder's smoothing: ATR = ((n-1) * ATR + TR) / n
            atr = ((period - 1) as f64 * atr + true_ranges[i]) / period as f64;
            atr_values.push(atr);
        }
        
        Ok(atr_values)
    }
}

impl Indicator for ATR {
    fn calculate(&mut self, data: &PriceData) -> Result<IndicatorResult, IndicatorError> {
        if data.high.len() != data.low.len() || data.high.len() != data.close.len() {
            return Err(IndicatorError::InvalidParameter(
                "Price data arrays must have the same length".to_string()
            ));
        }
        
        let atr_values = Self::calculate_series(
            &data.high,
            &data.low,
            &data.close,
            self.period
        )?;
        
        Ok(IndicatorResult::Series(atr_values))
    }
    
    fn reset(&mut self) {
        self.atr_value = None;
        self.prev_close = None;
        self.true_ranges.clear();
        self.sum = 0.0;
        self.initialized = false;
    }
    
    fn is_ready(&self) -> bool {
        self.initialized
    }
}

impl IncrementalIndicator for ATR {
    fn update(&mut self, value: f64) -> Result<Option<f64>, IndicatorError> {
        // For single value update, we assume it's the close price
        // In real usage, you'd need to provide high/low/close separately
        self.update_hlc(value, value, value)
    }
    
    fn current(&self) -> Option<f64> {
        self.atr_value
    }
}

impl ATR {
    /// Update with high, low, close values
    pub fn update_hlc(&mut self, high: f64, low: f64, close: f64) -> Result<Option<f64>, IndicatorError> {
        if let Some(prev_close) = self.prev_close {
            let tr = Self::calculate_true_range(high, low, Some(prev_close));
            
            if !self.initialized {
                // Still building initial window
                self.true_ranges.push(tr);
                self.sum += tr;
                
                if self.true_ranges.is_full() {
                    // Calculate initial ATR
                    let initial_atr = self.sum / self.period as f64;
                    self.atr_value = Some(initial_atr);
                    self.initialized = true;
                }
            } else {
                // Update ATR using Wilder's smoothing
                if let Some(current_atr) = self.atr_value {
                    let new_atr = ((self.period - 1) as f64 * current_atr + tr) / self.period as f64;
                    self.atr_value = Some(new_atr);
                }
            }
        }
        
        self.prev_close = Some(close);
        Ok(self.atr_value)
    }
    
    /// Get ATR as a percentage of price
    pub fn atr_percent(&self, price: f64) -> Option<f64> {
        self.atr_value.map(|atr| (atr / price) * 100.0)
    }
    
    /// Calculate stop loss based on ATR multiple
    pub fn calculate_stop_loss(&self, entry_price: f64, is_long: bool, atr_multiple: f64) -> Option<f64> {
        self.atr_value.map(|atr| {
            if is_long {
                entry_price - (atr * atr_multiple)
            } else {
                entry_price + (atr * atr_multiple)
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_atr_calculation() {
        let high = vec![
            48.70, 48.72, 48.90, 48.87, 48.82,
            49.05, 49.20, 49.35, 49.92, 50.19,
            50.12, 49.66, 49.88, 50.19, 50.36
        ];
        let low = vec![
            47.79, 48.14, 48.39, 48.37, 48.24,
            48.64, 48.94, 48.86, 49.50, 49.87,
            49.20, 48.90, 49.43, 49.73, 49.26
        ];
        let close = vec![
            48.16, 48.61, 48.75, 48.63, 48.74,
            49.03, 49.07, 49.32, 49.91, 50.13,
            49.53, 49.50, 49.75, 50.03, 50.29
        ];
        
        let atr_values = ATR::calculate_series(&high, &low, &close, 14).unwrap();
        
        assert!(!atr_values.is_empty());
        
        // ATR should always be positive
        for atr in &atr_values {
            assert!(*atr > 0.0);
        }
    }
    
    #[test]
    fn test_true_range_calculation() {
        // Test case where high-low is largest
        let tr = ATR::calculate_true_range(50.0, 45.0, Some(47.0));
        assert_eq!(tr, 5.0);
        
        // Test case where high-prev_close is largest
        let tr = ATR::calculate_true_range(50.0, 48.0, Some(45.0));
        assert_eq!(tr, 5.0);
        
        // Test case where prev_close-low is largest
        let tr = ATR::calculate_true_range(48.0, 45.0, Some(50.0));
        assert_eq!(tr, 5.0);
    }
    
    #[test]
    fn test_incremental_atr() {
        let mut atr = ATR::new(5).unwrap();
        
        // Need at least one previous close
        assert_eq!(atr.update_hlc(10.0, 9.0, 9.5).unwrap(), None);
        
        // Build up window
        assert_eq!(atr.update_hlc(10.2, 9.1, 9.8).unwrap(), None);
        assert_eq!(atr.update_hlc(10.5, 9.3, 10.0).unwrap(), None);
        assert_eq!(atr.update_hlc(10.3, 9.5, 9.7).unwrap(), None);
        assert_eq!(atr.update_hlc(10.1, 9.2, 9.6).unwrap(), None);
        
        // Should have ATR value now
        let atr_value = atr.update_hlc(10.0, 9.0, 9.5).unwrap();
        assert!(atr_value.is_some());
        assert!(atr_value.unwrap() > 0.0);
    }
}