updates to the indicators and fix minor bugs

1 files changed, 156 insertions, 18 deletions

diff --git a/src/signals.rs b/src/signals.rs
index 4b91d8e..3019c5c 100644
--- a/src/signals.rs
+++ b/src/signals.rs
@@ -50,14 +50,31 @@ pub struct Swing {
     pub index: usize,
     pub is_high: bool,
     pub price: f64,
-    /// bar where price had reversed by `pct` from the extreme
+    /// bar where price had reversed by the threshold from the extreme
     pub confirmed_at: usize,
 }
 
+/// Minimum reversal that confirms a swing.
+pub enum Threshold {
+    /// fixed percentage of the extreme price
+    Percent(f64),
+    /// `k` times the ATR at the confirming bar (adapts to each stock's volatility)
+    Atr { k: f64, atr: Vec<f64> },
+}
+
+impl Threshold {
+    /// Reversal size in price units at bar `i` for an extreme at `price`.
+    pub fn at(&self, i: usize, price: f64) -> f64 {
+        match self {
+            Threshold::Percent(p) => price * p / 100.0,
+            Threshold::Atr { k, atr } => k * atr[i],
+        }
+    }
+}
+
 /// ZigZag swings: a swing high (low) is confirmed once price falls (rises) by at
-/// least `pct` percent from the running extreme. The last, unconfirmed extreme is dropped.
-pub fn zigzag(candles: &[Candle], pct: f64) -> Vec<Swing> {
-    let f = pct / 100.0;
+/// least the threshold from the running extreme. The last, unconfirmed extreme is dropped.
+pub fn zigzag(candles: &[Candle], thr: &Threshold) -> Vec<Swing> {
     let mut out = Vec::new();
     if candles.is_empty() {
         return out;
@@ -71,11 +88,11 @@ pub fn zigzag(candles: &[Candle], pct: f64) -> Vec<Swing> {
             None => {
                 if c.high > hi.1 { hi = (i, c.high); }
                 if c.low < lo.1 { lo = (i, c.low); }
-                if c.low <= hi.1 * (1.0 - f) && hi.0 < i {
+                if c.low <= hi.1 - thr.at(i, hi.1) && hi.0 < i {
                     out.push(Swing { index: hi.0, is_high: true, price: hi.1, confirmed_at: i });
                     dir = Some(false);
                     lo = (i, c.low);
-                } else if c.high >= lo.1 * (1.0 + f) && lo.0 < i {
+                } else if c.high >= lo.1 + thr.at(i, lo.1) && lo.0 < i {
                     out.push(Swing { index: lo.0, is_high: false, price: lo.1, confirmed_at: i });
                     dir = Some(true);
                     hi = (i, c.high);
@@ -84,7 +101,7 @@ pub fn zigzag(candles: &[Candle], pct: f64) -> Vec<Swing> {
             Some(true) => {
                 if c.high > hi.1 {
                     hi = (i, c.high);
-                } else if c.low <= hi.1 * (1.0 - f) {
+                } else if c.low <= hi.1 - thr.at(i, hi.1) {
                     out.push(Swing { index: hi.0, is_high: true, price: hi.1, confirmed_at: i });
                     dir = Some(false);
                     lo = (i, c.low);
@@ -93,7 +110,7 @@ pub fn zigzag(candles: &[Candle], pct: f64) -> Vec<Swing> {
             Some(false) => {
                 if c.low < lo.1 {
                     lo = (i, c.low);
-                } else if c.high >= lo.1 * (1.0 + f) {
+                } else if c.high >= lo.1 + thr.at(i, lo.1) {
                     out.push(Swing { index: lo.0, is_high: false, price: lo.1, confirmed_at: i });
                     dir = Some(true);
                     hi = (i, c.high);
@@ -104,6 +121,53 @@ pub fn zigzag(candles: &[Candle], pct: f64) -> Vec<Swing> {
     out
 }
 
+/// Move each swing to the real extreme (highest high / lowest low) within +-`n` bars,
+/// never past its confirmation bar (no look-ahead). Used after zigzag on smoothed prices.
+pub fn snap_swings(candles: &[Candle], swings: &[Swing], n: usize) -> Vec<Swing> {
+    swings
+        .iter()
+        .map(|s| {
+            let lo = s.index.saturating_sub(n);
+            let hi = (s.index + n).min(s.confirmed_at);
+            let pick = (lo..=hi).reduce(|a, b| {
+                let better = if s.is_high { candles[b].high > candles[a].high } else { candles[b].low < candles[a].low };
+                if better { b } else { a }
+            });
+            let index = pick.unwrap_or(s.index);
+            let price = if s.is_high { candles[index].high } else { candles[index].low };
+            Swing { index, price, ..*s }
+        })
+        .collect()
+}
+
+/// Quality filters that reject noisy divergences.
+pub struct Filters {
+    pub min_bars: usize,
+    pub max_bars: usize,
+    /// minimum indicator difference between the two swings (indicator points)
+    pub min_ind_diff: f64,
+    /// minimum price difference, as a multiple of the swing threshold
+    pub min_price_k: f64,
+    /// reject if price or indicator crosses the line joining the two swings in between
+    pub check_line: bool,
+    /// regular divergences only: (bullish needs first-swing indicator below .0,
+    /// bearish needs it above .1)
+    pub zone: Option<(f64, f64)>,
+}
+
+impl Filters {
+    pub fn rsi() -> Self {
+        Filters { min_bars: 5, max_bars: 60, min_ind_diff: 3.0, min_price_k: 0.3, check_line: true, zone: Some((45.0, 55.0)) }
+    }
+    pub fn k() -> Self {
+        Filters { min_bars: 5, max_bars: 60, min_ind_diff: 5.0, min_price_k: 0.3, check_line: true, zone: Some((30.0, 70.0)) }
+    }
+    #[cfg(test)]
+    pub fn none() -> Self {
+        Filters { min_bars: 0, max_bars: usize::MAX, min_ind_diff: 0.0, min_price_k: 0.0, check_line: false, zone: None }
+    }
+}
+
 /// Indicator extreme (min for lows, max for highs) within +-`slack` bars of `idx`.
 fn indicator_extreme(ind: &[Option<f64>], idx: usize, is_high: bool, slack: usize) -> Option<(usize, f64)> {
     let lo = idx.saturating_sub(slack);
@@ -113,6 +177,15 @@ fn indicator_extreme(ind: &[Option<f64>], idx: usize, is_high: bool, slack: usiz
     })
 }
 
+/// True if no bar strictly between the endpoints lies beyond the straight line
+/// joining them (below it for swing lows, above it for swing highs).
+fn line_clear(a: (usize, f64), b: (usize, f64), is_high: bool, value: impl Fn(usize) -> Option<f64>) -> bool {
+    (a.0 + 1..b.0).all(|j| {
+        let line = a.1 + (b.1 - a.1) * (j - a.0) as f64 / (b.0 - a.0) as f64;
+        value(j).is_none_or(|v| if is_high { v <= line } else { v >= line })
+    })
+}
+
 /// Divergences between price swings and an indicator, using consecutive swings of the
 /// same kind (low->low, high->high):
 ///   regular bullish: price lower low,  indicator higher low   (reversal up)
@@ -120,15 +193,22 @@ fn indicator_extreme(ind: &[Option<f64>], idx: usize, is_high: bool, slack: usiz
 ///   hidden  bullish: price higher low, indicator lower low    (uptrend continues)
 ///   hidden  bearish: price lower high, indicator higher high  (downtrend continues)
 pub fn divergences(
+    candles: &[Candle],
     swings: &[Swing],
+    thr: &Threshold,
     ind: &[Option<f64>],
     name: &'static str,
+    f: &Filters,
 ) -> Vec<Divergence> {
     let mut out = Vec::new();
     for is_high in [false, true] {
         let same: Vec<&Swing> = swings.iter().filter(|s| s.is_high == is_high).collect();
         for w in same.windows(2) {
             let (a, b) = (w[0], w[1]);
+            let gap = b.index - a.index;
+            if gap < f.min_bars || gap > f.max_bars {
+                continue;
+            }
             let (Some((ja, ia)), Some((jb, ib))) = (
                 indicator_extreme(ind, a.index, is_high, 2),
                 indicator_extreme(ind, b.index, is_high, 2),
@@ -137,8 +217,27 @@ pub fn divergences(
             if price_up == ind_up || b.price == a.price || ib == ia {
                 continue;
             }
+            if (ib - ia).abs() < f.min_ind_diff
+                || (b.price - a.price).abs() < f.min_price_k * thr.at(b.confirmed_at, b.price)
+            {
+                continue;
+            }
             // lows: price down & ind up is regular; highs: price up & ind down is regular.
             let regular = if is_high { price_up } else { !price_up };
+            if regular {
+                if let Some((lo_max, hi_min)) = f.zone {
+                    if (!is_high && ia >= lo_max) || (is_high && ia <= hi_min) {
+                        continue;
+                    }
+                }
+            }
+            if f.check_line
+                && !(line_clear((a.index, a.price), (b.index, b.price), is_high, |j| {
+                    Some(if is_high { candles[j].high } else { candles[j].low })
+                }) && line_clear((ja, ia), (jb, ib), is_high, |j| ind[j]))
+            {
+                continue;
+            }
             out.push(Divergence {
                 kind: if is_high { "bearish" } else { "bullish" },
                 class: if regular { "regular" } else { "hidden" },
@@ -198,10 +297,17 @@ mod tests {
     }
 
     const PRICES: [f64; 9] = [100.0, 95.0, 90.0, 95.0, 100.0, 97.0, 95.0, 100.0, 105.0];
+    const PCT: Threshold = Threshold::Percent(5.0);
+
+    fn run(p: &[f64], ind: &[Option<f64>], f: &Filters) -> Vec<Divergence> {
+        let c = flat(p);
+        let sw = zigzag(&c, &PCT);
+        divergences(&c, &sw, &PCT, ind, "RSI", f)
+    }
 
     #[test]
     fn zigzag_finds_confirmed_swings() {
-        let sw = zigzag(&flat(&PRICES), 5.0);
+        let sw = zigzag(&flat(&PRICES), &PCT);
         let lows: Vec<_> = sw.iter().filter(|s| !s.is_high).map(|s| (s.index, s.confirmed_at)).collect();
         assert_eq!(lows, vec![(2, 3), (6, 7)]);
         assert!(sw.iter().any(|s| s.is_high && s.index == 4));
@@ -210,39 +316,71 @@ mod tests {
     #[test]
     fn hidden_bullish_not_regular() {
         // price: lows 90 then 95 (higher low); indicator lows 30 then 20 (lower low)
-        let sw = zigzag(&flat(&PRICES), 5.0);
         let mut ind = vec![Some(50.0); PRICES.len()];
         ind[2] = Some(30.0);
         ind[6] = Some(20.0);
-        let d = divergences(&sw, &ind, "RSI");
+        let d = run(&PRICES, &ind, &Filters::none());
         assert_eq!(d.len(), 1);
         assert_eq!((d[0].class, d[0].kind, d[0].i1, d[0].i2), ("hidden", "bullish", 2, 6));
     }
 
     #[test]
     fn regular_bullish() {
-        // price: lower low needs 90 then 85
         let p = [100.0, 95.0, 90.0, 95.0, 100.0, 95.0, 85.0, 90.0, 95.0];
-        let sw = zigzag(&flat(&p), 5.0);
         let mut ind = vec![Some(50.0); p.len()];
         ind[2] = Some(20.0);
         ind[6] = Some(30.0);
-        let d = divergences(&sw, &ind, "K");
+        let d = run(&p, &ind, &Filters::none());
         assert_eq!(d.len(), 1);
         assert_eq!((d[0].class, d[0].kind), ("regular", "bullish"));
     }
 
     #[test]
     fn hidden_bearish() {
-        // price highs 100 then 95 (lower high); indicator highs 60 then 80 (higher high)
         let p = [90.0, 95.0, 100.0, 95.0, 90.0, 92.0, 95.0, 90.0, 85.0];
-        let sw = zigzag(&flat(&p), 5.0);
         let mut ind = vec![Some(50.0); p.len()];
         ind[2] = Some(60.0);
         ind[6] = Some(80.0);
-        let d = divergences(&sw, &ind, "RSI");
+        let d = run(&p, &ind, &Filters::none());
         assert_eq!(d.len(), 1);
         assert_eq!((d[0].class, d[0].kind), ("hidden", "bearish"));
     }
-}
 
+    #[test]
+    fn filters_reject_close_swings_but_keep_clean_ones() {
+        // swings only 4 bars apart: unfiltered fires, filtered does not
+        let mut ind = vec![Some(50.0); PRICES.len()];
+        ind[2] = Some(30.0);
+        ind[6] = Some(20.0);
+        assert_eq!(run(&PRICES, &ind, &Filters::none()).len(), 1);
+        assert!(run(&PRICES, &ind, &Filters::rsi()).is_empty());
+
+        // swing lows 9 bars apart (idx 2: 90, idx 11: 95), nothing between dips below the line
+        let p = [100.0, 95.0, 90.0, 95.0, 100.0, 105.0, 110.0, 105.0, 100.0, 98.0, 96.0, 95.0, 100.0, 105.0];
+        let mut ind = vec![Some(50.0); p.len()];
+        ind[2] = Some(30.0);
+        ind[11] = Some(20.0);
+        let d = run(&p, &ind, &Filters::rsi());
+        assert_eq!(d.len(), 1);
+        assert_eq!((d[0].class, d[0].i1, d[0].i2), ("hidden", 2, 11));
+
+        // the indicator dips below the joining line in between -> rejected
+        ind[6] = Some(5.0);
+        assert!(run(&p, &ind, &Filters::rsi()).is_empty());
+    }
+
+    #[test]
+    fn atr_threshold_scales_with_volatility() {
+        let t = Threshold::Atr { k: 2.0, atr: vec![1.0, 3.0] };
+        assert_eq!(t.at(1, 100.0), 6.0);
+    }
+
+    #[test]
+    fn snap_moves_to_real_extreme_without_lookahead() {
+        let mut c = flat(&[10.0, 9.0, 8.0, 9.0, 10.0]);
+        c[1].low = 7.0; // real low one bar before the swing
+        let sw = [Swing { index: 2, is_high: false, price: 8.0, confirmed_at: 3 }];
+        let s = snap_swings(&c, &sw, 2);
+        assert_eq!((s[0].index, s[0].price), (1, 7.0));
+    }
+}