termichart_data/

indicators.rs

use termichart_core::{Candle, Indicator, Point};

// ---------------------------------------------------------------------------
// SMA
// ---------------------------------------------------------------------------

/// Simple Moving Average indicator.
pub struct SmaIndicator {
    period: usize,
    name: String,
}

impl SmaIndicator {
    pub fn new(period: usize) -> Self {
        Self {
            period,
            name: format!("SMA({})", period),
        }
    }
}

impl Indicator for SmaIndicator {
    fn name(&self) -> &str {
        &self.name
    }

    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        if candles.len() < self.period || self.period == 0 {
            return Vec::new();
        }

        let mut points = Vec::with_capacity(candles.len() - self.period + 1);

        // Compute the initial window sum.
        let mut sum: f64 = candles[..self.period].iter().map(|c| c.close).sum();
        points.push(Point {
            x: candles[self.period - 1].time,
            y: sum / self.period as f64,
        });

        // Slide the window.
        for i in self.period..candles.len() {
            sum += candles[i].close - candles[i - self.period].close;
            points.push(Point {
                x: candles[i].time,
                y: sum / self.period as f64,
            });
        }

        points
    }
}

// ---------------------------------------------------------------------------
// EMA
// ---------------------------------------------------------------------------

/// Exponential Moving Average indicator.
pub struct EmaIndicator {
    period: usize,
    name: String,
}

impl EmaIndicator {
    pub fn new(period: usize) -> Self {
        Self {
            period,
            name: format!("EMA({})", period),
        }
    }
}

impl Indicator for EmaIndicator {
    fn name(&self) -> &str {
        &self.name
    }

    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        if candles.len() < self.period || self.period == 0 {
            return Vec::new();
        }

        let multiplier = 2.0 / (self.period as f64 + 1.0);
        let mut points = Vec::with_capacity(candles.len() - self.period + 1);

        // First EMA value = SMA of the first `period` candles.
        let sma: f64 =
            candles[..self.period].iter().map(|c| c.close).sum::<f64>() / self.period as f64;
        let mut prev_ema = sma;

        points.push(Point {
            x: candles[self.period - 1].time,
            y: prev_ema,
        });

        // Subsequent values use the EMA formula.
        for candle in &candles[self.period..] {
            let ema = (candle.close - prev_ema) * multiplier + prev_ema;
            points.push(Point {
                x: candle.time,
                y: ema,
            });
            prev_ema = ema;
        }

        points
    }
}

// ---------------------------------------------------------------------------
// VWAP
// ---------------------------------------------------------------------------

/// Volume-Weighted Average Price indicator.
pub struct VwapIndicator;

impl VwapIndicator {
    pub fn new() -> Self {
        Self
    }
}

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

impl Indicator for VwapIndicator {
    fn name(&self) -> &str {
        "VWAP"
    }

    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        let mut points = Vec::with_capacity(candles.len());
        let mut cumulative_tp_vol = 0.0_f64;
        let mut cumulative_vol = 0.0_f64;

        for candle in candles {
            let typical_price = (candle.high + candle.low + candle.close) / 3.0;
            cumulative_tp_vol += typical_price * candle.volume;
            cumulative_vol += candle.volume;

            let vwap = if cumulative_vol != 0.0 {
                cumulative_tp_vol / cumulative_vol
            } else {
                0.0
            };

            points.push(Point {
                x: candle.time,
                y: vwap,
            });
        }

        points
    }
}

// ---------------------------------------------------------------------------
// Bollinger Bands
// ---------------------------------------------------------------------------

/// Bollinger Bands indicator: middle band (SMA), upper band, lower band.
pub struct BollingerBandsIndicator {
    period: usize,
    std_dev_mult: f64,
    name: String,
}

impl BollingerBandsIndicator {
    pub fn new(period: usize, std_dev_mult: f64) -> Self {
        Self {
            period,
            std_dev_mult,
            name: format!("BB({},{})", period, std_dev_mult),
        }
    }

    /// Returns (middle, upper, lower) bands as three sets of points.
    pub fn compute_bands(&self, candles: &[Candle]) -> (Vec<Point>, Vec<Point>, Vec<Point>) {
        if candles.len() < self.period || self.period == 0 {
            return (Vec::new(), Vec::new(), Vec::new());
        }

        let mut middle = Vec::with_capacity(candles.len() - self.period + 1);
        let mut upper = Vec::with_capacity(candles.len() - self.period + 1);
        let mut lower = Vec::with_capacity(candles.len() - self.period + 1);

        for i in (self.period - 1)..candles.len() {
            let start = i + 1 - self.period;
            let slice = &candles[start..=i];
            let mean: f64 = slice.iter().map(|c| c.close).sum::<f64>() / self.period as f64;
            let variance: f64 =
                slice.iter().map(|c| (c.close - mean).powi(2)).sum::<f64>() / self.period as f64;
            let std_dev = variance.sqrt();

            let time = candles[i].time;
            middle.push(Point { x: time, y: mean });
            upper.push(Point {
                x: time,
                y: mean + self.std_dev_mult * std_dev,
            });
            lower.push(Point {
                x: time,
                y: mean - self.std_dev_mult * std_dev,
            });
        }

        (middle, upper, lower)
    }
}

impl Indicator for BollingerBandsIndicator {
    fn name(&self) -> &str {
        &self.name
    }

    /// Returns the middle band points (SMA). Use compute_bands() for all three bands.
    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        let (middle, _, _) = self.compute_bands(candles);
        middle
    }
}

// ---------------------------------------------------------------------------
// RSI
// ---------------------------------------------------------------------------

/// Relative Strength Index indicator.
pub struct RsiIndicator {
    period: usize,
    name: String,
}

impl RsiIndicator {
    pub fn new(period: usize) -> Self {
        Self {
            period,
            name: format!("RSI({})", period),
        }
    }
}

impl Indicator for RsiIndicator {
    fn name(&self) -> &str {
        &self.name
    }

    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        if candles.len() < self.period + 1 || self.period == 0 {
            return Vec::new();
        }

        let mut points = Vec::with_capacity(candles.len() - self.period);

        // Calculate initial average gain and loss
        let mut avg_gain = 0.0_f64;
        let mut avg_loss = 0.0_f64;

        for i in 1..=self.period {
            let change = candles[i].close - candles[i - 1].close;
            if change > 0.0 {
                avg_gain += change;
            } else {
                avg_loss += change.abs();
            }
        }

        avg_gain /= self.period as f64;
        avg_loss /= self.period as f64;

        let rs = if avg_loss == 0.0 {
            100.0
        } else {
            avg_gain / avg_loss
        };
        let rsi = 100.0 - (100.0 / (1.0 + rs));
        points.push(Point {
            x: candles[self.period].time,
            y: rsi,
        });

        // Subsequent values use smoothed averages
        for i in (self.period + 1)..candles.len() {
            let change = candles[i].close - candles[i - 1].close;
            let (gain, loss) = if change > 0.0 {
                (change, 0.0)
            } else {
                (0.0, change.abs())
            };

            avg_gain = (avg_gain * (self.period as f64 - 1.0) + gain) / self.period as f64;
            avg_loss = (avg_loss * (self.period as f64 - 1.0) + loss) / self.period as f64;

            let rs = if avg_loss == 0.0 {
                100.0
            } else {
                avg_gain / avg_loss
            };
            let rsi = 100.0 - (100.0 / (1.0 + rs));
            points.push(Point {
                x: candles[i].time,
                y: rsi,
            });
        }

        points
    }
}

// ---------------------------------------------------------------------------
// MACD
// ---------------------------------------------------------------------------

/// MACD (Moving Average Convergence Divergence) indicator.
pub struct MacdIndicator {
    fast_period: usize,
    slow_period: usize,
    signal_period: usize,
    name: String,
}

impl MacdIndicator {
    pub fn new(fast: usize, slow: usize, signal: usize) -> Self {
        Self {
            fast_period: fast,
            slow_period: slow,
            signal_period: signal,
            name: format!("MACD({},{},{})", fast, slow, signal),
        }
    }

    /// Returns (macd_line, signal_line, histogram) as three sets of points.
    pub fn compute_full(&self, candles: &[Candle]) -> (Vec<Point>, Vec<Point>, Vec<Point>) {
        let fast_ema = EmaIndicator::new(self.fast_period);
        let slow_ema = EmaIndicator::new(self.slow_period);

        let fast_points = fast_ema.compute(candles);
        let slow_points = slow_ema.compute(candles);

        if fast_points.is_empty() || slow_points.is_empty() {
            return (Vec::new(), Vec::new(), Vec::new());
        }

        // Align by time -- find matching points
        let mut macd_line = Vec::new();
        let mut slow_idx = 0;
        for fp in &fast_points {
            while slow_idx < slow_points.len() && slow_points[slow_idx].x < fp.x {
                slow_idx += 1;
            }
            if slow_idx < slow_points.len() && (slow_points[slow_idx].x - fp.x).abs() < 0.001 {
                macd_line.push(Point {
                    x: fp.x,
                    y: fp.y - slow_points[slow_idx].y,
                });
            }
        }

        if macd_line.len() < self.signal_period {
            return (macd_line, Vec::new(), Vec::new());
        }

        // Compute signal line as EMA of MACD line
        let multiplier = 2.0 / (self.signal_period as f64 + 1.0);
        let sma: f64 = macd_line[..self.signal_period]
            .iter()
            .map(|p| p.y)
            .sum::<f64>()
            / self.signal_period as f64;
        let mut prev_signal = sma;

        let mut signal_line = Vec::with_capacity(macd_line.len() - self.signal_period + 1);
        signal_line.push(Point {
            x: macd_line[self.signal_period - 1].x,
            y: prev_signal,
        });

        for p in &macd_line[self.signal_period..] {
            let signal = (p.y - prev_signal) * multiplier + prev_signal;
            signal_line.push(Point {
                x: p.x,
                y: signal,
            });
            prev_signal = signal;
        }

        // Compute histogram
        let mut histogram = Vec::new();
        let mut sig_idx = 0;
        for mp in &macd_line {
            while sig_idx < signal_line.len() && signal_line[sig_idx].x < mp.x {
                sig_idx += 1;
            }
            if sig_idx < signal_line.len() && (signal_line[sig_idx].x - mp.x).abs() < 0.001 {
                histogram.push(Point {
                    x: mp.x,
                    y: mp.y - signal_line[sig_idx].y,
                });
            }
        }

        (macd_line, signal_line, histogram)
    }
}

impl Indicator for MacdIndicator {
    fn name(&self) -> &str {
        &self.name
    }

    /// Returns the MACD line points. Use compute_full() for all three components.
    fn compute(&self, candles: &[Candle]) -> Vec<Point> {
        let (macd_line, _, _) = self.compute_full(candles);
        macd_line
    }
}

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

    fn sample_candles() -> Vec<Candle> {
        vec![
            Candle { time: 1.0, open: 10.0, high: 12.0, low: 9.0, close: 11.0, volume: 100.0 },
            Candle { time: 2.0, open: 11.0, high: 13.0, low: 10.0, close: 12.0, volume: 150.0 },
            Candle { time: 3.0, open: 12.0, high: 14.0, low: 11.0, close: 13.0, volume: 200.0 },
            Candle { time: 4.0, open: 13.0, high: 15.0, low: 12.0, close: 14.0, volume: 120.0 },
            Candle { time: 5.0, open: 14.0, high: 16.0, low: 13.0, close: 15.0, volume: 180.0 },
        ]
    }

    #[test]
    fn sma_basic() {
        let candles = sample_candles();
        let sma = SmaIndicator::new(3);
        let points = sma.compute(&candles);
        // SMA(3) starts at index 2: (11+12+13)/3 = 12.0
        assert_eq!(points.len(), 3);
        assert!((points[0].y - 12.0).abs() < 1e-10);
        assert_eq!(points[0].x, 3.0);
        // Second: (12+13+14)/3 = 13.0
        assert!((points[1].y - 13.0).abs() < 1e-10);
        // Third: (13+14+15)/3 = 14.0
        assert!((points[2].y - 14.0).abs() < 1e-10);
    }

    #[test]
    fn sma_insufficient_data() {
        let candles = sample_candles();
        let sma = SmaIndicator::new(10);
        assert!(sma.compute(&candles).is_empty());
    }

    #[test]
    fn ema_basic() {
        let candles = sample_candles();
        let ema = EmaIndicator::new(3);
        let points = ema.compute(&candles);
        assert_eq!(points.len(), 3);
        // First EMA = SMA(3) = 12.0
        assert!((points[0].y - 12.0).abs() < 1e-10);
        // multiplier = 2/(3+1) = 0.5
        // Second: (14 - 12) * 0.5 + 12 = 13.0
        assert!((points[1].y - 13.0).abs() < 1e-10);
        // Third: (15 - 13) * 0.5 + 13 = 14.0
        assert!((points[2].y - 14.0).abs() < 1e-10);
    }

    #[test]
    fn ema_name() {
        let ema = EmaIndicator::new(20);
        assert_eq!(ema.name(), "EMA(20)");
    }

    #[test]
    fn vwap_basic() {
        let candles = sample_candles();
        let vwap = VwapIndicator::new();
        let points = vwap.compute(&candles);
        assert_eq!(points.len(), 5);

        // First candle: tp = (12+9+11)/3 = 32/3 ≈ 10.6667
        // vwap = 10.6667 * 100 / 100 = 10.6667
        let tp1 = (12.0 + 9.0 + 11.0) / 3.0;
        assert!((points[0].y - tp1).abs() < 1e-10);
    }

    #[test]
    fn vwap_name() {
        let vwap = VwapIndicator::new();
        assert_eq!(vwap.name(), "VWAP");
    }

    #[test]
    fn bb_basic() {
        let candles = sample_candles();
        let bb = BollingerBandsIndicator::new(3, 2.0);
        let (middle, upper, lower) = bb.compute_bands(&candles);
        assert_eq!(middle.len(), 3);
        assert_eq!(upper.len(), 3);
        assert_eq!(lower.len(), 3);
        // Middle should be SMA
        assert!((middle[0].y - 12.0).abs() < 1e-10);
        // Upper > middle > lower
        assert!(upper[0].y > middle[0].y);
        assert!(lower[0].y < middle[0].y);
    }

    #[test]
    fn rsi_basic() {
        let candles = sample_candles();
        let rsi = RsiIndicator::new(3);
        let points = rsi.compute(&candles);
        assert!(!points.is_empty());
        // RSI should be between 0 and 100
        for p in &points {
            assert!(p.y >= 0.0 && p.y <= 100.0);
        }
    }

    #[test]
    fn rsi_all_gains() {
        // All prices increasing -> RSI should be near 100
        let candles: Vec<Candle> = (0..10)
            .map(|i| Candle {
                time: i as f64,
                open: i as f64,
                high: i as f64 + 1.0,
                low: i as f64,
                close: i as f64 + 1.0,
                volume: 100.0,
            })
            .collect();
        let rsi = RsiIndicator::new(3);
        let points = rsi.compute(&candles);
        if let Some(last) = points.last() {
            assert!(last.y > 90.0);
        }
    }

    #[test]
    fn macd_basic() {
        // Need more data for MACD
        let candles: Vec<Candle> = (0..30)
            .map(|i| Candle {
                time: i as f64,
                open: 100.0 + (i as f64).sin() * 5.0,
                high: 105.0 + (i as f64).sin() * 5.0,
                low: 95.0 + (i as f64).sin() * 5.0,
                close: 100.0 + (i as f64).sin() * 5.0 + 0.5,
                volume: 100.0,
            })
            .collect();
        let macd = MacdIndicator::new(12, 26, 9);
        let (macd_line, _signal, _hist) = macd.compute_full(&candles);
        // With 30 candles we should get some MACD points
        assert!(!macd_line.is_empty());
    }

    #[test]
    fn macd_name() {
        let macd = MacdIndicator::new(12, 26, 9);
        assert_eq!(macd.name(), "MACD(12,26,9)");
    }

    #[test]
    fn bb_insufficient_data() {
        let candles = vec![sample_candles()[0]];
        let bb = BollingerBandsIndicator::new(5, 2.0);
        let (m, u, l) = bb.compute_bands(&candles);
        assert!(m.is_empty());
        assert!(u.is_empty());
        assert!(l.is_empty());
    }
}