use rust_decimal::prelude::*;
use rust_decimal::Decimal;
use std::collections::VecDeque;

#[derive(Debug, PartialEq)]
pub enum BollingerBandPosition {
    AboveUpper,
    BelowLower,
    WithinBands,
}

impl BollingerBandPosition {
    pub fn to_numeric(&self) -> Decimal {
        match self {
            BollingerBandPosition::AboveUpper => Decimal::new(1, 0),
            BollingerBandPosition::WithinBands => Decimal::ZERO,
            BollingerBandPosition::BelowLower => Decimal::new(-1, 0),
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct TrendValue {
    alpha: Decimal,
    value: VecDeque<Decimal>,
    ema: VecDeque<Decimal>,
    period: usize,
    name: String,
}

impl TrendValue {
    pub fn new(name: &str, period: usize) -> Self {
        let alpha = Decimal::from(2) / (Decimal::from(period) + Decimal::from(1));
        TrendValue {
            alpha,
            value: VecDeque::with_capacity(period),
            ema: VecDeque::with_capacity(period),
            period,
            name: name.to_owned(),
        }
    }

    pub fn update_value(&mut self, new_value: Decimal) {
        if self.value.len() > self.period {
            self.value.pop_front();
        }
        self.value.push_back(new_value);

        if self.ema.len() > self.period {
            self.ema.pop_front();
        }
        let ema = new_value * self.alpha
            + self.current_ema().unwrap_or(new_value) * (Decimal::from(1) - self.alpha);
        self.ema.push_back(ema);
    }

    pub fn current_ema(&self) -> Option<Decimal> {
        if self.value.len() < self.period {
            log::trace!(
                "{} is None: {}/{}",
                self.name,
                self.value.len(),
                self.period
            );
            None
        } else {
            self.ema.back().copied()
        }
    }

    pub fn current_value(&self) -> Option<Decimal> {
        self.value.back().copied()
    }

    fn calculate_bollinger_bands(&self, num_std_dev: f64) -> Option<(f64, f64)> {
        let len = self.value.len();
        if len < 2 {
            return None;
        }

        let values: Vec<f64> = self
            .value
            .iter()
            .map(|v| v.to_f64().unwrap_or(0.0))
            .collect();

        let mean: f64 = values.iter().sum::<f64>() / len as f64;
        let variance: f64 = values
            .iter()
            .map(|&value| {
                let diff = value - mean;
                diff * diff
            })
            .sum::<f64>()
            / len as f64;
        let std_dev = variance.sqrt();

        let upper_band = mean + num_std_dev * std_dev;
        let lower_band = mean - num_std_dev * std_dev;

        Some((upper_band, lower_band))
    }

    pub fn bollinger_band_positions(
        &self,
        num_std_devs: &[f64],
    ) -> Vec<Option<BollingerBandPosition>> {
        num_std_devs
            .iter()
            .map(|&num_std_dev| {
                if let Some(current_value) = self.current_value() {
                    if let Some((upper_band, lower_band)) =
                        self.calculate_bollinger_bands(num_std_dev)
                    {
                        let current_value_f64 = current_value.to_f64().unwrap_or(0.0);
                        return Some(if current_value_f64 > upper_band {
                            BollingerBandPosition::AboveUpper
                        } else if current_value_f64 < lower_band {
                            BollingerBandPosition::BelowLower
                        } else {
                            BollingerBandPosition::WithinBands
                        });
                    }
                }
                None
            })
            .collect()
    }
}
