1use std::collections::VecDeque;
7
8use crate::model::Bar;
9use crate::timeframe::Timeframe;
10
11use super::smoothing::Rma;
12use super::{Indicator, IndicatorAlert, IndicatorOutput};
13
14#[derive(Debug, Clone, Copy, PartialEq)]
16pub enum ZigZagDeviationMode {
17 Percent(f64),
19 AtrMultiple(f64),
21}
22
23#[derive(Debug, Clone, Copy, PartialEq)]
25pub struct ZigZagNode {
26 pub timestamp: i64,
27 pub price: f64,
28 pub is_high: bool,
29 pub confirmed: bool,
33}
34
35pub struct AdvancedZigZagEngine {
56 depth: usize,
57 backstep: usize,
58 deviation: ZigZagDeviationMode,
59 atr: Rma,
60 prev_close: Option<f64>,
61 bars: VecDeque<Bar>,
62 bar_index: usize,
63 nodes: Vec<ZigZagNode>,
64 current_direction: i8,
65 last_confirmed_bar_index: Option<usize>,
66 alerts: Vec<IndicatorAlert>,
67}
68
69impl AdvancedZigZagEngine {
70 pub fn new(
71 depth: usize,
72 backstep: usize,
73 deviation: ZigZagDeviationMode,
74 atr_len: usize,
75 ) -> Self {
76 let depth = depth.max(1);
77 Self {
78 depth,
79 backstep,
80 deviation,
81 atr: Rma::new(atr_len.max(1)),
82 prev_close: None,
83 bars: VecDeque::with_capacity(depth * 2 + 1),
84 bar_index: 0,
85 nodes: Vec::new(),
86 current_direction: 0,
87 last_confirmed_bar_index: None,
88 alerts: Vec::new(),
89 }
90 }
91
92 pub fn with_defaults() -> Self {
93 Self::new(3, 2, ZigZagDeviationMode::Percent(1.0), 14)
94 }
95
96 pub fn nodes(&self) -> &[ZigZagNode] {
98 &self.nodes
99 }
100
101 pub fn current_leg(&self) -> Option<&ZigZagNode> {
104 self.nodes.last().filter(|n| !n.confirmed)
105 }
106
107 fn deviation_threshold(&self, atr: Option<f64>) -> f64 {
108 match self.deviation {
109 ZigZagDeviationMode::Percent(pct) => pct / 100.0,
110 ZigZagDeviationMode::AtrMultiple(mult) => {
111 match atr {
113 Some(a) if a > 0.0 => mult * a,
114 _ => f64::INFINITY, }
116 }
117 }
118 }
119
120 pub fn reduce(nodes: &[ZigZagNode], deviation_pct: f64) -> Vec<ZigZagNode> {
124 if nodes.is_empty() {
125 return Vec::new();
126 }
127 let threshold = deviation_pct / 100.0;
128 let mut reduced: Vec<ZigZagNode> = vec![nodes[0]];
129
130 for &node in &nodes[1..] {
131 let last = *reduced.last().expect("seeded with nodes[0]");
132 if node.is_high == last.is_high {
133 let replace = (node.is_high && node.price > last.price)
135 || (!node.is_high && node.price < last.price);
136 if replace {
137 *reduced.last_mut().unwrap() = node;
138 }
139 continue;
140 }
141
142 let change = if last.price != 0.0 {
143 (node.price - last.price).abs() / last.price.abs()
144 } else {
145 f64::INFINITY
146 };
147 if change >= threshold {
148 reduced.push(node);
149 }
150 }
151
152 reduced
153 }
154
155 pub fn project_to_timeframe(
159 nodes: &[ZigZagNode],
160 target_tf: Timeframe,
161 utc_offset_seconds: i32,
162 ) -> Vec<ZigZagNode> {
163 use std::collections::BTreeMap;
164
165 let mut buckets: BTreeMap<(i64, bool), ZigZagNode> = BTreeMap::new();
166 for &node in nodes {
167 let bucket = target_tf.bucket_start(node.timestamp, utc_offset_seconds);
168 let key = (bucket, node.is_high);
169 buckets
170 .entry(key)
171 .and_modify(|existing| {
172 let more_extreme = (node.is_high && node.price > existing.price)
173 || (!node.is_high && node.price < existing.price);
174 if more_extreme {
175 *existing = node;
176 }
177 })
178 .or_insert(node);
179 }
180
181 let mut projected: Vec<ZigZagNode> = buckets.into_values().collect();
182 projected.sort_by_key(|n| n.timestamp);
183 projected
184 }
185}
186
187impl Indicator for AdvancedZigZagEngine {
188 fn name(&self) -> &str {
189 "zigzag_advanced"
190 }
191
192 fn warmup_period(&self) -> usize {
193 self.depth * 2 + 1
194 }
195
196 fn reset(&mut self) {
197 self.atr.reset();
198 self.prev_close = None;
199 self.bars.clear();
200 self.bar_index = 0;
201 self.nodes.clear();
202 self.current_direction = 0;
203 self.last_confirmed_bar_index = None;
204 self.alerts.clear();
205 }
206
207 fn on_bar(&mut self, bar: &Bar) -> Option<IndicatorOutput> {
208 self.alerts.clear();
209
210 let tr = match self.prev_close {
211 Some(pc) => (bar.high - bar.low)
212 .max((bar.high - pc).abs())
213 .max((bar.low - pc).abs()),
214 None => bar.high - bar.low,
215 };
216 self.prev_close = Some(bar.close);
217 let atr = self.atr.update(tr);
218
219 self.bars.push_back(bar.clone());
220 if self.bars.len() > self.depth * 2 + 1 {
221 self.bars.pop_front();
222 }
223 let current_bar_index = self.bar_index;
224 self.bar_index += 1;
225
226 if self.bars.len() < self.depth * 2 + 1 {
227 return None;
228 }
229
230 let mid_idx = self.depth;
231 let mid_bar = self.bars[mid_idx].clone();
232 let mid_bar_index = current_bar_index - self.depth;
233
234 let is_pivot_high = self
235 .bars
236 .iter()
237 .enumerate()
238 .all(|(i, b)| i == mid_idx || b.high <= mid_bar.high);
239 let is_pivot_low = self
240 .bars
241 .iter()
242 .enumerate()
243 .all(|(i, b)| i == mid_idx || b.low >= mid_bar.low);
244
245 let threshold = self.deviation_threshold(atr);
246
247 let backstep_ok = self
248 .last_confirmed_bar_index
249 .map(|last| mid_bar_index >= last + self.backstep)
250 .unwrap_or(true);
251
252 if is_pivot_high {
253 self.try_extend(
254 true,
255 mid_bar.high,
256 mid_bar.timestamp,
257 mid_bar_index,
258 threshold,
259 backstep_ok,
260 );
261 }
262 if is_pivot_low {
263 self.try_extend(
264 false,
265 mid_bar.low,
266 mid_bar.timestamp,
267 mid_bar_index,
268 threshold,
269 backstep_ok,
270 );
271 }
272
273 let leg_price = self.nodes.last().map(|n| n.price).unwrap_or(mid_bar.close);
274 Some(
275 IndicatorOutput::new(leg_price).with_state(if self.current_leg().is_some() {
276 "running"
277 } else {
278 "confirmed"
279 }),
280 )
281 }
282
283 fn alerts(&self) -> Vec<IndicatorAlert> {
284 self.alerts.clone()
285 }
286}
287
288impl AdvancedZigZagEngine {
289 #[allow(clippy::too_many_arguments)]
290 fn try_extend(
291 &mut self,
292 is_high: bool,
293 price: f64,
294 timestamp: i64,
295 bar_index: usize,
296 threshold: f64,
297 backstep_ok: bool,
298 ) {
299 let opposite_direction = if is_high { 1 } else { -1 };
300
301 if self.current_direction == opposite_direction || self.current_direction == 0 {
302 let should_replace = match self.nodes.last() {
304 Some(last) if !last.confirmed && last.is_high == is_high => {
305 (is_high && price > last.price) || (!is_high && price < last.price)
306 }
307 _ => true,
308 };
309 if should_replace {
310 if let Some(last) = self
311 .nodes
312 .last_mut()
313 .filter(|n| !n.confirmed && n.is_high == is_high)
314 {
315 *last = ZigZagNode {
316 timestamp,
317 price,
318 is_high,
319 confirmed: false,
320 };
321 } else {
322 self.nodes.push(ZigZagNode {
323 timestamp,
324 price,
325 is_high,
326 confirmed: false,
327 });
328 }
329 self.current_direction = opposite_direction;
330 }
331 return;
332 }
333
334 let last_price = self.nodes.last().map(|n| n.price);
339 let change = match (self.deviation, last_price) {
340 (ZigZagDeviationMode::AtrMultiple(_), Some(lp)) => (price - lp).abs(),
341 (ZigZagDeviationMode::Percent(_), Some(lp)) if lp != 0.0 => {
342 (price - lp).abs() / lp.abs()
343 }
344 _ => f64::INFINITY,
345 };
346
347 if change >= threshold && backstep_ok {
348 if let Some(last) = self.nodes.last_mut() {
349 last.confirmed = true;
350 }
351 self.nodes.push(ZigZagNode {
352 timestamp,
353 price,
354 is_high,
355 confirmed: false,
356 });
357 self.current_direction = opposite_direction;
358 self.last_confirmed_bar_index = Some(bar_index);
359 self.alerts.push(IndicatorAlert::new(
360 "zigzag_pivot_confirmed",
361 if is_high {
362 "ZigZag confirmed a swing low"
363 } else {
364 "ZigZag confirmed a swing high"
365 },
366 0.6,
367 ));
368 }
369 }
370}
371
372#[cfg(test)]
373mod tests {
374 use super::*;
375
376 fn sine_bars(n: usize) -> Vec<Bar> {
377 (0..n)
378 .map(|i| {
379 let price = if (i / 5) % 2 == 0 {
380 100.0 + (i % 5) as f64 * 4.0
381 } else {
382 120.0 - (i % 5) as f64 * 4.0
383 };
384 Bar::new(i as i64 * 60, price, price + 1.0, price - 1.0, price, 100.0)
385 })
386 .collect()
387 }
388
389 #[test]
390 fn test_produces_confirmed_and_running_nodes() {
391 let mut engine = AdvancedZigZagEngine::new(2, 1, ZigZagDeviationMode::Percent(1.0), 5);
392 for bar in sine_bars(40) {
393 engine.on_bar(&bar);
394 }
395 assert!(!engine.nodes().is_empty());
396 assert!(engine.nodes().iter().any(|n| n.confirmed));
397 }
398
399 #[test]
400 fn test_backstep_suppresses_pivots_too_close_together() {
401 let lenient = {
402 let mut e = AdvancedZigZagEngine::new(2, 0, ZigZagDeviationMode::Percent(0.01), 5);
403 for bar in sine_bars(40) {
404 e.on_bar(&bar);
405 }
406 e.nodes().iter().filter(|n| n.confirmed).count()
407 };
408 let strict = {
409 let mut e = AdvancedZigZagEngine::new(2, 20, ZigZagDeviationMode::Percent(0.01), 5);
410 for bar in sine_bars(40) {
411 e.on_bar(&bar);
412 }
413 e.nodes().iter().filter(|n| n.confirmed).count()
414 };
415 assert!(
416 strict <= lenient,
417 "a large backstep must never confirm more pivots than a near-zero one"
418 );
419 }
420
421 #[test]
422 fn test_atr_mode_requires_warm_atr_before_confirming() {
423 let mut engine =
424 AdvancedZigZagEngine::new(2, 0, ZigZagDeviationMode::AtrMultiple(0.5), 100);
425 for bar in sine_bars(20) {
426 engine.on_bar(&bar);
427 }
428 assert!(engine.nodes().iter().all(|n| !n.confirmed));
431 }
432
433 #[test]
434 fn test_reduce_produces_a_coarser_recursive_level() {
435 let base = vec![
436 ZigZagNode {
437 timestamp: 0,
438 price: 100.0,
439 is_high: false,
440 confirmed: true,
441 },
442 ZigZagNode {
443 timestamp: 1,
444 price: 102.0,
445 is_high: true,
446 confirmed: true,
447 },
448 ZigZagNode {
449 timestamp: 2,
450 price: 101.0,
451 is_high: false,
452 confirmed: true,
453 },
454 ZigZagNode {
455 timestamp: 3,
456 price: 110.0,
457 is_high: true,
458 confirmed: true,
459 },
460 ZigZagNode {
461 timestamp: 4,
462 price: 95.0,
463 is_high: false,
464 confirmed: true,
465 },
466 ];
467 let coarse = AdvancedZigZagEngine::reduce(&base, 5.0);
470 assert!(coarse.len() < base.len());
471 assert_eq!(coarse.first().unwrap().price, 100.0);
472 assert_eq!(coarse.last().unwrap().price, 95.0);
473 }
474
475 #[test]
476 fn test_project_to_timeframe_keeps_most_extreme_per_bucket() {
477 let nodes = vec![
478 ZigZagNode {
479 timestamp: 0,
480 price: 100.0,
481 is_high: true,
482 confirmed: true,
483 },
484 ZigZagNode {
485 timestamp: 60,
486 price: 105.0,
487 is_high: true,
488 confirmed: true,
489 },
490 ZigZagNode {
491 timestamp: 120,
492 price: 102.0,
493 is_high: true,
494 confirmed: true,
495 },
496 ];
497 let projected = AdvancedZigZagEngine::project_to_timeframe(&nodes, Timeframe::Minute(5), 0);
499 assert_eq!(projected.len(), 1);
500 assert_eq!(
501 projected[0].price, 105.0,
502 "must keep the highest high within the bucket"
503 );
504 }
505}