path: root/src/signals.rs

use crate::data::Candle;
use serde::Serialize;

#[derive(Serialize)]
pub struct Cross {
    pub index: usize,
    /// "golden" (K crosses above D) or "death" (K crosses below D)
    pub kind: &'static str,
    /// true when the cross happens in the extreme zone (<20 golden, >80 death)
    pub strong: bool,
}

#[derive(Serialize)]
pub struct Divergence {
    /// "bullish" or "bearish"
    pub kind: &'static str,
    /// "regular" (reversal signal) or "hidden" (trend-continuation signal)
    pub class: &'static str,
    /// "RSI" or "K"
    pub indicator: &'static str,
    /// bars of the two price swings
    pub i1: usize,
    pub i2: usize,
    /// price at the swings (swing low's Low / swing high's High)
    pub p1: f64,
    pub p2: f64,
    /// bars where the indicator has its extreme near each swing
    pub j1: usize,
    pub j2: usize,
    /// bar at which the second swing was confirmed (earliest the signal was knowable)
    pub confirmed: usize,
}

pub fn crosses(k: &[Option<f64>], d: &[Option<f64>]) -> Vec<Cross> {
    let mut out = Vec::new();
    for i in 1..k.len() {
        if let (Some(k0), Some(d0), Some(k1), Some(d1)) = (k[i - 1], d[i - 1], k[i], d[i]) {
            if k0 <= d0 && k1 > d1 {
                out.push(Cross { index: i, kind: "golden", strong: k1 < 20.0 });
            } else if k0 >= d0 && k1 < d1 {
                out.push(Cross { index: i, kind: "death", strong: k1 > 80.0 });
            }
        }
    }
    out
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Swing {
    pub index: usize,
    pub is_high: bool,
    pub price: f64,
    /// 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 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;
    }
    // Some(true): tracking a high (rising); Some(false): tracking a low; None: undecided.
    let mut dir: Option<bool> = None;
    let mut hi = (0, candles[0].high);
    let mut lo = (0, candles[0].low);
    for (i, c) in candles.iter().enumerate().skip(1) {
        match dir {
            None => {
                if c.high > hi.1 { hi = (i, c.high); }
                if c.low < lo.1 { lo = (i, c.low); }
                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 + 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);
                }
            }
            Some(true) => {
                if c.high > hi.1 {
                    hi = (i, c.high);
                } 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);
                }
            }
            Some(false) => {
                if c.low < lo.1 {
                    lo = (i, c.low);
                } 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);
                }
            }
        }
    }
    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);
    let hi = (idx + slack).min(ind.len() - 1);
    (lo..=hi).filter_map(|j| ind[j].map(|v| (j, v))).reduce(|a, b| {
        if (b.1 > a.1) == is_high && b.1 != a.1 { b } else { a }
    })
}

/// 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)
///   regular bearish: price higher high, indicator lower high  (reversal down)
///   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),
            ) else { continue };
            let (price_up, ind_up) = (b.price > a.price, ib > ia);
            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" },
                indicator: name,
                i1: a.index, i2: b.index, p1: a.price, p2: b.price,
                j1: ja, j2: jb, confirmed: b.confirmed_at,
            });
        }
    }
    out.sort_by_key(|d| d.confirmed);
    out
}

/// Rough forward-return summary per (class, kind, indicator): count, then for each
/// horizon the mean close-to-close return after the confirmation bar and the share of
/// signals that moved in the signal's direction.
pub fn forward_stats(candles: &[Candle], divs: &[Divergence], horizons: &[usize]) -> String {
    use std::collections::BTreeMap;
    let mut groups: BTreeMap<(&str, &str, &str), Vec<&Divergence>> = BTreeMap::new();
    for d in divs {
        groups.entry((d.class, d.kind, d.indicator)).or_default().push(d);
    }
    let mut out = String::new();
    for ((class, kind, ind), ds) in groups {
        out += &format!("{class:8} {kind:8} {ind:4} n={:<3}", ds.len());
        for &h in horizons {
            let rets: Vec<f64> = ds
                .iter()
                .filter(|d| d.confirmed + h < candles.len())
                .map(|d| {
                    let r = candles[d.confirmed + h].close / candles[d.confirmed].close - 1.0;
                    if kind == "bullish" { r } else { -r } // positive = moved the signalled way
                })
                .collect();
            if rets.is_empty() {
                out += &format!("  +{h}b: n/a");
            } else {
                let mean = rets.iter().sum::<f64>() / rets.len() as f64 * 100.0;
                let win = rets.iter().filter(|&&r| r > 0.0).count() as f64 / rets.len() as f64 * 100.0;
                out += &format!("  +{h}b: {mean:+.2}% win {win:.0}% (n={})", rets.len());
            }
        }
        out.push('\n');
    }
    out
}

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

    fn flat(closes: &[f64]) -> Vec<Candle> {
        closes
            .iter()
            .map(|&c| Candle { date: "2024-01-01".into(), open: c, high: c, low: c, close: c, volume: 1.0 })
            .collect()
    }

    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), &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));
    }

    #[test]
    fn hidden_bullish_not_regular() {
        // price: lows 90 then 95 (higher low); indicator lows 30 then 20 (lower low)
        let mut ind = vec![Some(50.0); PRICES.len()];
        ind[2] = Some(30.0);
        ind[6] = Some(20.0);
        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() {
        let p = [100.0, 95.0, 90.0, 95.0, 100.0, 95.0, 85.0, 90.0, 95.0];
        let mut ind = vec![Some(50.0); p.len()];
        ind[2] = Some(20.0);
        ind[6] = Some(30.0);
        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() {
        let p = [90.0, 95.0, 100.0, 95.0, 90.0, 92.0, 95.0, 90.0, 85.0];
        let mut ind = vec![Some(50.0); p.len()];
        ind[2] = Some(60.0);
        ind[6] = Some(80.0);
        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));
    }
}