1use std::collections::VecDeque;
4
5use crate::error::{Error, Result};
6use crate::ohlcv::Candle;
7use crate::traits::Indicator;
8
9#[derive(Debug, Clone, Copy, PartialEq)]
12pub struct AndrewsPitchforkOutput {
13 pub median: f64,
15 pub upper: f64,
17 pub lower: f64,
19}
20
21#[derive(Debug, Clone, Copy)]
23struct Pivot {
24 index: f64,
25 price: f64,
26 is_high: bool,
27}
28
29#[derive(Debug, Clone)]
69pub struct AndrewsPitchfork {
70 strength: usize,
71 window: VecDeque<Candle>,
72 pivots: Vec<Pivot>,
73 count: usize,
74 last: Option<AndrewsPitchforkOutput>,
75}
76
77impl AndrewsPitchfork {
78 pub fn new(strength: usize) -> Result<Self> {
85 if strength == 0 {
86 return Err(Error::PeriodZero);
87 }
88 Ok(Self {
89 strength,
90 window: VecDeque::with_capacity(2 * strength + 1),
91 pivots: Vec::new(),
92 count: 0,
93 last: None,
94 })
95 }
96
97 pub const fn strength(&self) -> usize {
99 self.strength
100 }
101
102 pub const fn value(&self) -> Option<AndrewsPitchforkOutput> {
104 self.last
105 }
106
107 fn record_pivot(&mut self, pivot: Pivot) {
109 if let Some(last) = self.pivots.last_mut() {
110 if last.is_high == pivot.is_high {
111 let more_extreme = if pivot.is_high {
113 pivot.price > last.price
114 } else {
115 pivot.price < last.price
116 };
117 if more_extreme {
118 *last = pivot;
119 }
120 return;
121 }
122 }
123 self.pivots.push(pivot);
124 if self.pivots.len() > 3 {
125 self.pivots.remove(0);
126 }
127 }
128
129 fn project(&self, tc: f64) -> Option<AndrewsPitchforkOutput> {
130 let [p0, p1, p2] = self.pivots.as_slice() else {
131 return None;
132 };
133 let mid_t = f64::midpoint(p1.index, p2.index);
134 let mid_p = f64::midpoint(p1.price, p2.price);
135 let slope = (mid_p - p0.price) / (mid_t - p0.index);
136 let median = p0.price + slope * (tc - p0.index);
137 let off1 = p1.price - (p0.price + slope * (p1.index - p0.index));
138 let off2 = p2.price - (p0.price + slope * (p2.index - p0.index));
139 Some(AndrewsPitchforkOutput {
140 median,
141 upper: median + off1.max(off2),
142 lower: median + off1.min(off2),
143 })
144 }
145}
146
147impl Indicator for AndrewsPitchfork {
148 type Input = Candle;
149 type Output = AndrewsPitchforkOutput;
150
151 fn update(&mut self, candle: Candle) -> Option<AndrewsPitchforkOutput> {
152 self.count += 1;
153 let span = 2 * self.strength + 1;
154 if self.window.len() == span {
155 self.window.pop_front();
156 }
157 self.window.push_back(candle);
158 if self.window.len() == span {
159 let center = self.window[self.strength];
160 let is_high = self
161 .window
162 .iter()
163 .enumerate()
164 .all(|(i, c)| i == self.strength || c.high < center.high);
165 let is_low = self
166 .window
167 .iter()
168 .enumerate()
169 .all(|(i, c)| i == self.strength || c.low > center.low);
170 let center_index = (self.count - 1 - self.strength) as f64;
172 if is_high && !is_low {
173 self.record_pivot(Pivot {
174 index: center_index,
175 price: center.high,
176 is_high: true,
177 });
178 } else if is_low && !is_high {
179 self.record_pivot(Pivot {
180 index: center_index,
181 price: center.low,
182 is_high: false,
183 });
184 }
185 }
186 let tc = (self.count - 1) as f64;
187 if let Some(out) = self.project(tc) {
188 self.last = Some(out);
189 return Some(out);
190 }
191 None
192 }
193
194 fn reset(&mut self) {
195 self.window.clear();
196 self.pivots.clear();
197 self.count = 0;
198 self.last = None;
199 }
200
201 fn warmup_period(&self) -> usize {
202 2 * self.strength + 1
203 }
204
205 fn is_ready(&self) -> bool {
206 self.last.is_some()
207 }
208
209 fn name(&self) -> &'static str {
210 "AndrewsPitchfork"
211 }
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 use crate::traits::BatchExt;
218
219 fn c(high: f64, low: f64) -> Candle {
220 Candle::new_unchecked(
221 f64::midpoint(high, low),
222 high,
223 low,
224 f64::midpoint(high, low),
225 1_000.0,
226 0,
227 )
228 }
229
230 fn zigzag() -> Vec<Candle> {
232 let mut out = Vec::new();
233 for i in 0..120 {
234 let base = 100.0 + (f64::from(i) * 0.5).sin() * 10.0;
235 out.push(c(base + 1.0, base - 1.0));
236 }
237 out
238 }
239
240 #[test]
241 fn rejects_zero_strength() {
242 assert!(matches!(AndrewsPitchfork::new(0), Err(Error::PeriodZero)));
243 }
244
245 #[test]
246 fn accessors_and_metadata() {
247 let p = AndrewsPitchfork::new(2).unwrap();
248 assert_eq!(p.strength(), 2);
249 assert_eq!(p.warmup_period(), 5);
250 assert_eq!(p.name(), "AndrewsPitchfork");
251 assert!(!p.is_ready());
252 assert_eq!(p.value(), None);
253 }
254
255 #[test]
256 fn none_before_three_pivots() {
257 let mut p = AndrewsPitchfork::new(2).unwrap();
258 let out = p.batch(&[c(101.0, 99.0), c(102.0, 100.0), c(101.0, 99.0)]);
260 assert!(out.iter().all(Option::is_none));
261 }
262
263 #[test]
264 fn eventually_emits_on_swings() {
265 let mut p = AndrewsPitchfork::new(2).unwrap();
266 let out = p.batch(&zigzag());
267 assert!(
268 out.iter().any(Option::is_some),
269 "a swinging series should form a pitchfork"
270 );
271 assert!(p.is_ready());
272 }
273
274 #[test]
275 fn upper_at_or_above_lower() {
276 let mut p = AndrewsPitchfork::new(2).unwrap();
277 for o in p.batch(&zigzag()).into_iter().flatten() {
278 assert!(
279 o.upper >= o.lower,
280 "upper {} below lower {}",
281 o.upper,
282 o.lower
283 );
284 }
285 }
286
287 #[test]
288 fn reset_clears_state() {
289 let mut p = AndrewsPitchfork::new(2).unwrap();
290 p.batch(&zigzag());
291 assert!(p.is_ready());
292 p.reset();
293 assert!(!p.is_ready());
294 assert_eq!(p.value(), None);
295 assert_eq!(p.strength(), 2);
296 }
297
298 #[test]
299 fn record_pivot_keeps_more_extreme_same_kind() {
300 let mut p = AndrewsPitchfork::new(2).unwrap();
301 p.record_pivot(Pivot {
302 index: 0.0,
303 price: 100.0,
304 is_high: true,
305 });
306 p.record_pivot(Pivot {
308 index: 1.0,
309 price: 105.0,
310 is_high: true,
311 });
312 assert_eq!(p.pivots.len(), 1);
313 assert_eq!(p.pivots[0].price, 105.0);
314 p.record_pivot(Pivot {
316 index: 2.0,
317 price: 102.0,
318 is_high: true,
319 });
320 assert_eq!(p.pivots.len(), 1);
321 assert_eq!(p.pivots[0].price, 105.0);
322 p.record_pivot(Pivot {
324 index: 3.0,
325 price: 90.0,
326 is_high: false,
327 });
328 assert_eq!(p.pivots.len(), 2);
329 p.record_pivot(Pivot {
331 index: 4.0,
332 price: 85.0,
333 is_high: false,
334 });
335 assert_eq!(p.pivots[1].price, 85.0);
336 p.record_pivot(Pivot {
338 index: 5.0,
339 price: 88.0,
340 is_high: false,
341 });
342 assert_eq!(p.pivots[1].price, 85.0);
343 }
344
345 #[test]
346 fn batch_equals_streaming() {
347 let candles = zigzag();
348 let batch = AndrewsPitchfork::new(2).unwrap().batch(&candles);
349 let mut b = AndrewsPitchfork::new(2).unwrap();
350 let streamed: Vec<_> = candles.iter().map(|x| b.update(*x)).collect();
351 assert_eq!(batch, streamed);
352 }
353}