portfolio_analyzer/

sd.rs

/// SahastaDirectionality FSM — Rust implementation of the paper/ThinkScript model.
use crate::indicators::*;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum GroupA {
    Plus,
    Minus,
    Tilde,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum GroupB {
    SemiPlus,
    SemiMinus,
    SemiTilde,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SDState {
    pub group_a: GroupA,
    pub group_b: GroupB, // only meaningful when group_a == Tilde
    pub group_c: u8,     // 0=none, 1=***, 2=**., 3=*.., 4=..., 5=**, 6=*., 7=.., 8=*, 9=.
}

impl SDState {
    pub fn score(&self) -> f64 {
        // Half-continuous encoding from the paper hypothesis. Group A is
        // definitive; Group B is consulted only for an uncertain Group A.
        let base = match self.group_a {
            GroupA::Plus => 1.0,
            GroupA::Minus => 0.0,
            GroupA::Tilde => match self.group_b {
                GroupB::SemiPlus => 0.75,
                GroupB::SemiMinus => 0.25,
                GroupB::SemiTilde => 0.50,
            },
        };
        let (stars, dots) = match self.group_c {
            1 => (3, 0), // ***
            2 => (2, 1), // **.
            3 => (1, 2), // *..
            4 => (0, 3), // ...
            5 => (2, 0), // **
            6 => (1, 1), // *.
            7 => (0, 2), // ..
            8 => (1, 0), // *
            9 => (0, 1), // .
            _ => (0, 0),
        };
        let star_sign = if base > 0.5 {
            1.0
        } else if base < 0.5 {
            -1.0
        } else {
            0.0
        };
        let dot_sign = match self.group_a {
            GroupA::Plus => 1.0,
            GroupA::Minus => -1.0,
            GroupA::Tilde => 0.0,
        };
        (base + stars as f64 * 0.05 * star_sign + dots as f64 * 0.025 * dot_sign).clamp(0.0, 1.0)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SDSqueezeDiagnostics {
    pub state: SDState,
    pub score: f64,
    pub compression: u8,
    pub highest_compression: u8,
    pub is_fired: bool,
    pub fired_start: bool,
    pub squeeze_start: bool,
    pub is_active_squeeze: bool,
    pub momentum_changed: bool,
    pub time_expired: bool,
    pub bars_since_squeeze_start: usize,
    pub bars_since_fired: usize,
    pub average_total_squeeze_ticks: f64,
    pub average_mid_squeeze_ticks: f64,
    pub average_low_squeeze_ticks: f64,
    pub expected_squeeze_ticks: f64,
    pub elapsed_squeeze_ticks: f64,
    pub remaining_squeeze_ticks: f64,
    /// Continuous, scale-free inputs that produced the finite SD state.  They
    /// are exposed for causal diagnostic backtests; `score` remains the
    /// separately documented state-machine collapse.
    pub ao_over_close: f64,
    pub ao_delta_over_close: f64,
    pub deriv_price_over_close: f64,
    pub deriv_ema8_over_close: f64,
    pub deriv_ema21_over_close: f64,
    pub ema8_minus_ema21_over_close: f64,
    pub close_minus_ema21_over_close: f64,
}

pub struct SahastaDirectionality {
    bars: Vec<OHLCV>,
    length: usize,
}

fn running_fired_age(is_fired: &[bool]) -> Vec<usize> {
    let mut ages = Vec::with_capacity(is_fired.len());
    let mut age = 0usize;
    for (i, fired) in is_fired.iter().copied().enumerate() {
        if !fired {
            age = 0;
        } else if i == 0 || !is_fired[i - 1] {
            age = 1;
        } else {
            age += 1;
        }
        ages.push(age);
    }
    ages
}

fn classify_group_b(group_a: &GroupA, semi_plus: bool, semi_minus: bool) -> GroupB {
    if group_a != &GroupA::Tilde {
        GroupB::SemiTilde
    } else if semi_plus {
        // The chart's lexical expression checks semiPlus first.
        GroupB::SemiPlus
    } else if semi_minus {
        GroupB::SemiMinus
    } else {
        GroupB::SemiTilde
    }
}

impl SahastaDirectionality {
    pub fn new(bars: Vec<OHLCV>) -> Self {
        Self { bars, length: 21 }
    }

    pub fn compute_state(&self) -> Option<SDState> {
        self.compute_diagnostics()
            .map(|diagnostics| diagnostics.state)
    }

    pub fn compute_diagnostics(&self) -> Option<SDSqueezeDiagnostics> {
        let bars = &self.bars;
        let n = bars.len();
        if n < self.length * 2 + 10 {
            return None;
        }

        let close: Vec<f64> = bars.iter().map(|b| b.close).collect();

        // Bollinger Bands (length=21, 2 std)
        let bb21 = bollinger_bands(&close, 21, 2.0);
        // Keltner Channels at factors 1.0, 1.5, 2.0. The chart and scanner
        // ThinkScripts pass their common length (21) to all core channels.
        let kc_hi = keltner_channels(bars, self.length, 1.0);
        let kc_mid = keltner_channels(bars, self.length, 1.5);
        let kc_lo = keltner_channels(bars, self.length, 2.0);
        // Compression
        let comp = compute_compression(&bb21, &kc_hi, &kc_mid, &kc_lo);

        // EMAs
        let ema8 = ema(&close, 8);
        let ema21 = ema(&close, 21);

        // Derivatives
        let d_close: Vec<f64> = std::iter::once(0.0)
            .chain(close.windows(2).map(|w| w[1] - w[0]))
            .collect();
        let d_ema8: Vec<f64> = std::iter::once(0.0)
            .chain(ema8.windows(2).map(|w| w[1] - w[0]))
            .collect();
        let d_ema21: Vec<f64> = std::iter::once(0.0)
            .chain(ema21.windows(2).map(|w| w[1] - w[0]))
            .collect();

        let deriv_price = smf_derivative(&d_close, 21);
        let deriv_ema8 = smf_derivative(&d_ema8, 21);
        let deriv_ema21 = smf_derivative(&d_ema21, 21);

        // Awesome Oscillator
        let ao = awesome_oscillator(bars);

        // AO direction at current bar
        let ao_up = ao[n - 1] > ao[n - 2];
        let ao_down = ao[n - 1] < ao[n - 2];

        // ── Squeeze state machine ──────────────────────────────────────────────
        let mut highest_comp = vec![0u8; n];
        let mut is_fired = vec![false; n];
        let mut sq_start = vec![false; n];

        for i in 0..n {
            let c = comp[i];
            let c_prev = if i > 0 { comp[i - 1] } else { 0 };
            let hc_prev = if i > 0 { highest_comp[i - 1] } else { 0 };

            let sq_start_base = c > 0 && c_prev == 0;
            let unfired_sq = c > c_prev && c_prev < hc_prev && hc_prev > 0;
            sq_start[i] = sq_start_base || unfired_sq;

            highest_comp[i] = if sq_start[i] || c > hc_prev {
                c
            } else {
                hc_prev
            };
            is_fired[i] = c < highest_comp[i] && highest_comp[i] > 0;
        }

        // firedStart: transition False→True
        let mut fired_start = vec![false; n];
        for i in 1..n {
            fired_start[i] = is_fired[i] && !is_fired[i - 1];
        }

        // barsSinceFired
        let bars_since_fired = running_fired_age(&is_fired);

        // barsSinceSqueezeStart tracks the current compression lifecycle even
        // before it has fired. This is the execution clock used by SD Timer.
        let mut bars_since_squeeze_start = vec![0usize; n];
        let mut sq_count = 0usize;
        for i in 0..n {
            if sq_start[i] {
                sq_count = 1;
            } else if comp[i] > 0 || is_fired[i] {
                sq_count = sq_count.saturating_add(1);
            } else {
                sq_count = 0;
            }
            bars_since_squeeze_start[i] = sq_count;
        }

        // momChanged
        let mut mom_changed = vec![false; n];
        for i in 1..n {
            if fired_start[i] {
                mom_changed[i] = false;
            } else if is_fired[i] && {
                let a_up = ao[i] > ao[i - 1];
                let a_down = ao[i] < ao[i - 1];
                let p_up = i >= 2 && ao[i - 1] > ao[i - 2];
                let p_down = i >= 2 && ao[i - 1] < ao[i - 2];
                (a_up && !p_up) || (a_down && !p_down)
            } {
                mom_changed[i] = true;
            } else {
                mom_changed[i] = mom_changed[i - 1];
            }
        }

        // Dynamic squeeze duration estimates over the same 125-bar window as
        // the ThinkScript/Python implementation. These are per-timeframe values:
        // the caller converts ticks to minutes using that timeframe's bar size.
        let mut n_total = vec![0.0f64; n];
        let mut n_mid = vec![0.0f64; n];
        let mut n_low = vec![0.0f64; n];
        let theta: Vec<f64> = (0..n)
            .map(|i| {
                let start = i.saturating_sub(124);
                let w = &comp[start..=i];
                let bars_in_compression = w.iter().filter(|&&c| c > 0).count() as f64;
                let squeeze_starts = sq_start[start..=i].iter().filter(|&&s| s).count() as f64;
                let total = bars_in_compression / squeeze_starts.max(1.0);

                let bars_at_mid = w.iter().filter(|&&c| c == 2).count() as f64;
                let mut entering_mid = 0.0f64;
                for j in start..=i {
                    if comp[j] == 2 && (j == 0 || comp[j - 1] != 2) {
                        entering_mid += 1.0;
                    }
                }
                let mid = bars_at_mid / entering_mid.max(1.0);

                let bars_at_low = w.iter().filter(|&&c| c == 1).count() as f64;
                let mut entering_low = 0.0f64;
                for j in start..=i {
                    if comp[j] == 1 && (j == 0 || comp[j - 1] != 1) {
                        entering_low += 1.0;
                    }
                }
                let low = bars_at_low / entering_low.max(1.0);

                n_total[i] = total;
                n_mid[i] = mid;
                n_low[i] = low;

                match comp[i] {
                    2 => {
                        if mid >= total / 3.0 {
                            mid
                        } else {
                            low
                        }
                    }
                    1 => {
                        if low >= total / 3.0 {
                            low
                        } else {
                            total
                        }
                    }
                    _ => total,
                }
            })
            .collect();

        // isActiveSqueeze
        let is_active: Vec<bool> = (0..n)
            .map(|i| {
                is_fired[i] && bars_since_fired[i] <= theta[i].ceil() as usize && !mom_changed[i]
            })
            .collect();

        // ── Group A ────────────────────────────────────────────────────────────
        let i = n - 1;
        let dir_plus = is_active[i] && ao_up && deriv_ema8[i] >= 0.0;
        let dir_minus = is_active[i] && ao_down && deriv_ema8[i] <= 0.0;

        let group_a = if dir_plus {
            GroupA::Plus
        } else if dir_minus {
            GroupA::Minus
        } else {
            GroupA::Tilde
        };

        // ── Group B (semi-directionality, only meaningful when GroupA==Tilde) ─
        let bb20 = bollinger_bands(&close, 20, 2.0);
        let kc20 = keltner_channels(bars, 20, 2.0);

        let diff_ema21 = if ema21[i] != 0.0 {
            (close[i] - ema21[i]).abs() / ema21[i]
        } else {
            0.0
        };
        let cond1 = diff_ema21 >= 0.005 && deriv_ema21[i] <= deriv_price[i];
        let cond1_plus = cond1 && close[i] < ema21[i];
        let cond1_minus = cond1 && close[i] > ema21[i];

        let cond2_plus = !kc_lo.upper[i].is_nan()
            && !bb21.upper[i].is_nan()
            && (close[i] > kc_lo.upper[i] || close[i] > bb21.upper[i]);
        let cond2_minus = !kc20.lower[i].is_nan()
            && !bb20.lower[i].is_nan()
            && (close[i] < kc20.lower[i] || close[i] < bb20.lower[i]);

        let group_b = classify_group_b(
            &group_a,
            cond1_plus || cond2_plus,
            cond1_minus || cond2_minus,
        );

        // ── Group C (squeeze lexicography) ────────────────────────────────────
        let eff_highest = if is_active[i] {
            highest_comp[i]
        } else {
            comp[i]
        };
        let c = comp[i];
        let group_c: u8 = match (eff_highest, c) {
            (3, 3) => 1,
            (3, 2) => 2,
            (3, 1) => 3,
            (3, 0) => 4,
            (2, 2) => 5,
            (2, 1) => 6,
            (2, 0) => 7,
            (1, 1) => 8,
            (1, 0) => 9,
            _ => 0,
        };

        Some(SDState {
            group_a,
            group_b,
            group_c,
        })
        .map(|state| {
            let expected = theta[i].max(0.0);
            let elapsed = if is_fired[i] {
                bars_since_fired[i] as f64
            } else {
                bars_since_squeeze_start[i] as f64
            };
            let remaining = if comp[i] > 0 || is_fired[i] {
                (expected - elapsed).max(0.0)
            } else {
                0.0
            };
            let score = state.score();
            let scale = close[i].abs().max(f64::EPSILON);
            SDSqueezeDiagnostics {
                state,
                score,
                compression: comp[i],
                highest_compression: highest_comp[i],
                is_fired: is_fired[i],
                fired_start: fired_start[i],
                squeeze_start: sq_start[i],
                is_active_squeeze: is_active[i],
                momentum_changed: mom_changed[i],
                time_expired: is_fired[i] && bars_since_fired[i] > theta[i].ceil() as usize,
                bars_since_squeeze_start: bars_since_squeeze_start[i],
                bars_since_fired: bars_since_fired[i],
                average_total_squeeze_ticks: n_total[i],
                average_mid_squeeze_ticks: n_mid[i],
                average_low_squeeze_ticks: n_low[i],
                expected_squeeze_ticks: expected,
                elapsed_squeeze_ticks: elapsed,
                remaining_squeeze_ticks: remaining,
                ao_over_close: ao[i] / scale,
                ao_delta_over_close: (ao[i] - ao[i - 1]) / scale,
                deriv_price_over_close: deriv_price[i] / scale,
                deriv_ema8_over_close: deriv_ema8[i] / scale,
                deriv_ema21_over_close: deriv_ema21[i] / scale,
                ema8_minus_ema21_over_close: (ema8[i] - ema21[i]) / scale,
                close_minus_ema21_over_close: (close[i] - ema21[i]) / scale,
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn state(group_a: GroupA, group_b: GroupB, group_c: u8) -> SDState {
        SDState {
            group_a,
            group_b,
            group_c,
        }
    }

    #[test]
    fn paper_encoding_keeps_definitive_states_at_the_bounds() {
        for group_c in 0..=9 {
            assert_eq!(state(GroupA::Plus, GroupB::SemiTilde, group_c).score(), 1.0);
            assert_eq!(
                state(GroupA::Minus, GroupB::SemiTilde, group_c).score(),
                0.0
            );
        }
    }

    #[test]
    fn paper_encoding_amplifies_only_stars_for_semi_directionality() {
        let stars = [0, 3, 2, 1, 0, 2, 1, 0, 1, 0];
        for (group_c, stars) in stars.into_iter().enumerate() {
            let push = stars as f64 * 0.05;
            assert_eq!(
                state(GroupA::Tilde, GroupB::SemiPlus, group_c as u8).score(),
                0.75 + push
            );
            assert_eq!(
                state(GroupA::Tilde, GroupB::SemiMinus, group_c as u8).score(),
                0.25 - push
            );
            assert_eq!(
                state(GroupA::Tilde, GroupB::SemiTilde, group_c as u8).score(),
                0.5
            );
        }
    }

    #[test]
    fn fired_clock_resets_exactly_when_thinkscript_is_fired_is_false() {
        assert_eq!(
            running_fired_age(&[false, true, true, false, false, true, true]),
            vec![0, 1, 2, 0, 0, 1, 2]
        );
    }

    #[test]
    fn group_b_is_gated_by_group_a_and_uses_thinkscript_precedence() {
        assert_eq!(
            classify_group_b(&GroupA::Plus, true, true),
            GroupB::SemiTilde
        );
        assert_eq!(
            classify_group_b(&GroupA::Tilde, true, true),
            GroupB::SemiPlus
        );
        assert_eq!(
            classify_group_b(&GroupA::Tilde, false, true),
            GroupB::SemiMinus
        );
    }
}