1 files changed, 98 insertions, 0 deletions
diff --git a/src/indicators.rs b/src/indicators.rs
new file mode 100644
index 0000000..4378757
--- /dev/null
+++ b/src/indicators.rs
@@ -0,0 +1,98 @@
+use crate::data::Candle;
+use ta::indicators::FastStochastic;
+use ta::{DataItem, Next};
+
+/// Wilder RSI from closes: first average = simple mean of the first `period`
+/// changes, then avg = (prev * (period-1) + current) / period. (`ta`'s RSI uses a
+/// standard EMA instead, which differs from charting software.) The first
+/// `period` values are `None`.
+pub fn rsi(candles: &[Candle], period: usize) -> Vec<Option<f64>> {
+    let n = period as f64;
+    let mut out = vec![None; candles.len()];
+    let (mut gain, mut loss) = (0.0, 0.0);
+    for i in 1..candles.len() {
+        let ch = candles[i].close - candles[i - 1].close;
+        let (g, l) = (ch.max(0.0), (-ch).max(0.0));
+        if i <= period {
+            gain += g / n;
+            loss += l / n;
+        } else {
+            gain = (gain * (n - 1.0) + g) / n;
+            loss = (loss * (n - 1.0) + l) / n;
+        }
+        if i >= period {
+            out[i] = Some(if loss == 0.0 { if gain == 0.0 { 50.0 } else { 100.0 } } else { 100.0 - 100.0 / (1.0 + gain / loss) });
+        }
+    }
+    out
+}
+
+/// KD: RSV from `ta`'s FastStochastic over high/low/close, then
+/// K = 2/3 K_prev + 1/3 RSV, D = 2/3 D_prev + 1/3 K, both seeded at 50.
+pub fn kd(candles: &[Candle], period: usize) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
+    let mut rsv_ind = FastStochastic::new(period).expect("period > 0");
+    let (mut k, mut d) = (50.0, 50.0);
+    let mut ks = Vec::with_capacity(candles.len());
+    let mut ds = Vec::with_capacity(candles.len());
+    for (i, c) in candles.iter().enumerate() {
+        let item = DataItem::builder()
+            .open(c.open)
+            .high(c.high)
+            .low(c.low)
+            .close(c.close)
+            .volume(c.volume)
+            .build()
+            .expect("validated candle");
+        let rsv = rsv_ind.next(&item);
+        if i + 1 < period {
+            ks.push(None);
+            ds.push(None);
+            continue;
+        }
+        k = k * 2.0 / 3.0 + rsv / 3.0;
+        d = d * 2.0 / 3.0 + k / 3.0;
+        ks.push(Some(k));
+        ds.push(Some(d));
+    }
+    (ks, ds)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn candle(h: f64, l: f64, c: f64) -> Candle {
+        Candle { date: "2024-01-01".into(), open: c, high: h, low: l, close: c, volume: 1.0 }
+    }
+
+    #[test]
+    fn kd_hand_computed() {
+        // period 2. i=1: window highs {12,14}->14, lows {8,9}->8, close 13 -> RSV = 5/6*100.
+        let cs = [candle(12.0, 8.0, 10.0), candle(14.0, 9.0, 13.0)];
+        let (k, d) = kd(&cs, 2);
+        assert_eq!(k[0], None);
+        let rsv = 5.0 / 6.0 * 100.0;
+        let k1 = 50.0 * 2.0 / 3.0 + rsv / 3.0;
+        assert!((k[1].unwrap() - k1).abs() < 1e-9);
+        assert!((d[1].unwrap() - (50.0 * 2.0 / 3.0 + k1 / 3.0)).abs() < 1e-9);
+    }
+
+    #[test]
+    fn rsi_wilder_hand_computed() {
+        // period 2, closes 10, 11, 10, 12: first avg gain=.5 loss=.5 (i=2) -> RSI 50;
+        // i=3: gain=(.5*1+2)/2=1.25, loss=(.5*1+0)/2=.25 -> RS 5 -> RSI 100-100/6
+        let cs = [10.0, 11.0, 10.0, 12.0].map(|c| candle(c, c, c));
+        let r = rsi(&cs, 2);
+        assert!((r[2].unwrap() - 50.0).abs() < 1e-9);
+        assert!((r[3].unwrap() - (100.0 - 100.0 / 6.0)).abs() < 1e-9);
+    }
+
+    #[test]
+    fn rsi_all_gains_is_100() {
+        let cs: Vec<_> = (1..=20).map(|i| candle(i as f64 + 1.0, i as f64, i as f64)).collect();
+        let r = rsi(&cs, 14);
+        assert_eq!(r[13], None);
+        assert!(r[14].is_some());
+        assert!((r[19].unwrap() - 100.0).abs() < 1e-9);
+    }
+}