updates to the indicators and fix minor bugs

1 files changed, 26 insertions, 11 deletions

diff --git a/src/main.rs b/src/main.rs
index 24f4ccc..8ca9df5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,36 +7,42 @@ use std::path::PathBuf;
 use std::process::ExitCode;
 
 const TEMPLATE: &str = include_str!("../web/template.html");
-const USAGE: &str = "usage: stock <input.csv> [-o out.html] [--rsi 14] [--kd 9] [--swing 5]";
-// --swing: minimum reversal in percent for a price swing (ZigZag) used by divergence detection
+const USAGE: &str = "usage: stock <input.csv> [-o out.html] [--rsi 14] [--kd 9] [--atr 2.5 | --swing PCT] [--smooth N]";
+// Swing detection for divergences: ZigZag with a reversal threshold of k*ATR(14) (--atr, default 2.5)
+// or a fixed percentage (--swing). --smooth N locates swings on an N-bar EMA, then snaps them to real highs/lows.
 
 struct Args {
     input: PathBuf,
     output: Option<PathBuf>,
     rsi: usize,
     kd: usize,
-    swing: f64,
+    atr_k: f64,
+    swing: Option<f64>,
+    smooth: usize,
 }
 
 fn parse_args() -> Result<Args, String> {
     let mut it = std::env::args().skip(1);
-    let (mut input, mut output, mut rsi, mut kd, mut swing) = (None, None, 14, 9, 5.0);
+    let (mut input, mut output, mut rsi, mut kd) = (None, None, 14, 9);
+    let (mut atr_k, mut swing, mut smooth) = (2.5, None, 0);
     while let Some(a) = it.next() {
         let mut val = |name: &str| it.next().ok_or(format!("{name} needs a value"));
         match a.as_str() {
             "-o" => output = Some(PathBuf::from(val("-o")?)),
             "--rsi" => rsi = val("--rsi")?.parse().map_err(|e| format!("--rsi: {e}"))?,
             "--kd" => kd = val("--kd")?.parse().map_err(|e| format!("--kd: {e}"))?,
-            "--swing" => swing = val("--swing")?.parse().map_err(|e| format!("--swing: {e}"))?,
+            "--swing" => swing = Some(val("--swing")?.parse().map_err(|e| format!("--swing: {e}"))?),
+            "--atr" => atr_k = val("--atr")?.parse().map_err(|e| format!("--atr: {e}"))?,
+            "--smooth" => smooth = val("--smooth")?.parse().map_err(|e| format!("--smooth: {e}"))?,
             "-h" | "--help" => return Err(String::new()),
             _ if a.starts_with('-') => return Err(format!("unknown option {a}")),
             _ => input = Some(PathBuf::from(a)),
         }
     }
-    if rsi == 0 || kd == 0 || !(swing > 0.0 && swing < 100.0) {
-        return Err("periods must be positive and --swing within (0, 100)".into());
+    if rsi == 0 || kd == 0 || atr_k <= 0.0 || swing.is_some_and(|p| !(p > 0.0 && p < 100.0)) {
+        return Err("periods must be positive, --atr > 0, --swing within (0, 100)".into());
     }
-    Ok(Args { input: input.ok_or("missing input.csv")?, output, rsi, kd, swing })
+    Ok(Args { input: input.ok_or("missing input.csv")?, output, rsi, kd, atr_k, swing, smooth })
 }
 
 fn run() -> Result<PathBuf, String> {
@@ -45,9 +51,18 @@ fn run() -> Result<PathBuf, String> {
     let rsi = indicators::rsi(&candles, args.rsi);
     let (k, d) = indicators::kd(&candles, args.kd);
 
-    let swings = signals::zigzag(&candles, args.swing);
-    let mut divs = signals::divergences(&swings, &rsi, "RSI");
-    divs.extend(signals::divergences(&swings, &k, "K"));
+    let thr = match args.swing {
+        Some(p) => signals::Threshold::Percent(p),
+        None => signals::Threshold::Atr { k: args.atr_k, atr: indicators::atr(&candles, 14) },
+    };
+    let swings = if args.smooth > 0 {
+        let found = signals::zigzag(&indicators::ema_candles(&candles, args.smooth), &thr);
+        signals::snap_swings(&candles, &found, args.smooth)
+    } else {
+        signals::zigzag(&candles, &thr)
+    };
+    let mut divs = signals::divergences(&candles, &swings, &thr, &rsi, "RSI", &signals::Filters::rsi());
+    divs.extend(signals::divergences(&candles, &swings, &thr, &k, "K", &signals::Filters::k()));
     divs.sort_by_key(|d| d.confirmed);
     print!("{}", signals::forward_stats(&candles, &divs, &[5, 10, 20]));
     let payload = json!({