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