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