//! Alpha Block Z — PIN (Probability of Informed Trading) via EM-MLE.
//!
//! ## Model (Easley, Kiefer & O'Hara 1996, *J. Finance*)
//!
//! Daily / per-window observations of (B_t, S_t) buy/sell trade counts are
//! modeled as a mixture:
//!
//! With probability `α`: information event occurs.
//!   With probability `δ`: bad news (informed sellers active).
//!   With probability `1−δ`: good news (informed buyers active).
//! With probability `1−α`: no information event.
//!
//! Uninformed traders arrive Poisson at rate `ε` per side regardless.
//! Informed traders arrive Poisson at rate `μ` (one side only, news-dependent).
//!
//! The PIN — the unconditional fraction of orders coming from informed flow:
//!
//! ```text
//!   PIN = (α · μ) / (α · μ + 2ε)
//! ```
//!
//! ## Likelihood
//!
//! Per-window log-likelihoods (constants dropped):
//!
//! ```text
//!   ln L(B, S | no event)        = -2ε    + B·ln(ε)        + S·ln(ε)
//!   ln L(B, S | good news event) = -(2ε+μ) + B·ln(ε+μ)      + S·ln(ε)
//!   ln L(B, S | bad news event)  = -(2ε+μ) + B·ln(ε)        + S·ln(ε+μ)
//! ```
//!
//! Posteriors for the latent state z ∈ {no, good, bad}: standard EM.
//!
//! ## Output features (3 dims)
//!
//! 1. `pin` — PIN value in [0, 1]
//! 2. `mu` — informed arrival rate
//! 3. `epsilon` — uninformed arrival rate (per side)
//!
//! ## Refit cadence
//!
//! Maintained as a rolling window of (B, S) counts per bar. EM re-fitted every
//! `refit_interval` bars. EM uses up to `max_iters` iterations or until
//! parameter change < `tol`.

use std::collections::VecDeque;

/// Alpha Block Z — Streaming PIN estimator with periodic EM refits.
#[derive(Debug, Clone)]
pub struct PinEstimator {
    window_size: usize,
    refit_interval: usize,
    max_iters: usize,
    tol: f64,

    counts: VecDeque<(u64, u64)>,
    cached_alpha: f64,
    cached_delta: f64,
    cached_mu: f64,
    cached_epsilon: f64,

    obs_since_refit: usize,
}

impl PinEstimator {
    pub fn new(window_size: usize, refit_interval: usize) -> Self {
        Self {
            window_size,
            refit_interval: refit_interval.max(1),
            max_iters: 100,
            tol: 1e-4,
            counts: VecDeque::with_capacity(window_size),
            cached_alpha: 0.3,
            cached_delta: 0.5,
            cached_mu: 1.0,
            cached_epsilon: 1.0,
            obs_since_refit: 0,
        }
    }

    /// Alpha defaults: 200-bar window, refit every 50 bars.
    pub fn with_defaults() -> Self {
        Self::new(200, 50)
    }

    /// Update with one bar's (buy_count, sell_count) and return
    /// `[PIN, μ, ε]` from the most recent fit.
    pub fn update_bar(&mut self, buy_count: u64, sell_count: u64) -> [f64; 3] {
        self.counts.push_back((buy_count, sell_count));
        if self.counts.len() > self.window_size {
            self.counts.pop_front();
        }
        self.obs_since_refit += 1;
        if self.obs_since_refit >= self.refit_interval && self.counts.len() >= 30 {
            self.refit_em();
            self.obs_since_refit = 0;
        }
        self.features()
    }

    /// Force a refit now (e.g. before producing alpha fxcache).
    pub fn force_refit(&mut self) {
        if self.counts.len() >= 30 {
            self.refit_em();
            self.obs_since_refit = 0;
        }
    }

    pub fn features(&self) -> [f64; 3] {
        let pin = self.cached_alpha * self.cached_mu
            / (self.cached_alpha * self.cached_mu + 2.0 * self.cached_epsilon).max(1e-12);
        [pin.clamp(0.0, 1.0), self.cached_mu, self.cached_epsilon]
    }

    /// EM iterations on the cached (B, S) window. Updates cached params.
    fn refit_em(&mut self) {
        if self.counts.is_empty() {
            return;
        }
        let mut alpha = self.cached_alpha;
        let mut delta = self.cached_delta;
        let mut mu = self.cached_mu;
        let mut epsilon = self.cached_epsilon;

        for _iter in 0..self.max_iters {
            let (new_alpha, new_delta, new_mu, new_epsilon) =
                em_step(&self.counts, alpha, delta, mu, epsilon);

            let max_diff = [
                (new_alpha - alpha).abs(),
                (new_delta - delta).abs(),
                (new_mu - mu).abs() / mu.max(1e-9),
                (new_epsilon - epsilon).abs() / epsilon.max(1e-9),
            ]
            .iter()
            .copied()
            .fold(0.0_f64, f64::max);

            alpha = new_alpha;
            delta = new_delta;
            mu = new_mu;
            epsilon = new_epsilon;

            if max_diff < self.tol {
                break;
            }
        }

        self.cached_alpha = alpha;
        self.cached_delta = delta;
        self.cached_mu = mu;
        self.cached_epsilon = epsilon;
    }

    pub fn reset(&mut self) {
        self.counts.clear();
        self.cached_alpha = 0.3;
        self.cached_delta = 0.5;
        self.cached_mu = 1.0;
        self.cached_epsilon = 1.0;
        self.obs_since_refit = 0;
    }
}

impl Default for PinEstimator {
    fn default() -> Self {
        Self::with_defaults()
    }
}

/// Per-window log-likelihoods `[ln L(no), ln L(good), ln L(bad)]`.
/// Constants (factorials) dropped — only relative magnitudes matter for posteriors.
fn log_likelihoods(b: u64, s: u64, mu: f64, epsilon: f64) -> [f64; 3] {
    let b_f = b as f64;
    let s_f = s as f64;
    let ln_eps = (epsilon + 1e-12).ln();
    let ln_eps_mu = (epsilon + mu + 1e-12).ln();
    [
        -2.0 * epsilon + (b_f + s_f) * ln_eps,
        -(2.0 * epsilon + mu) + b_f * ln_eps_mu + s_f * ln_eps,
        -(2.0 * epsilon + mu) + b_f * ln_eps + s_f * ln_eps_mu,
    ]
}

/// One EM step: E-step computes per-window posteriors `(γ_no, γ_good, γ_bad)`;
/// M-step computes new parameter estimates.
fn em_step(
    counts: &VecDeque<(u64, u64)>,
    alpha: f64,
    delta: f64,
    mu: f64,
    epsilon: f64,
) -> (f64, f64, f64, f64) {
    let t = counts.len() as f64;

    let mut sum_gamma_event = 0.0_f64;
    let mut sum_gamma_bad = 0.0_f64;
    let mut sum_b = 0.0_f64;
    let mut sum_s = 0.0_f64;
    let mut sum_gamma_good_b = 0.0_f64;
    let mut sum_gamma_bad_s = 0.0_f64;

    let ln_pi_no = (1.0 - alpha + 1e-12).ln();
    let ln_pi_good = (alpha * (1.0 - delta) + 1e-12).ln();
    let ln_pi_bad = (alpha * delta + 1e-12).ln();

    for &(b, s) in counts.iter() {
        let log_l = log_likelihoods(b, s, mu, epsilon);
        let log_post = [
            ln_pi_no + log_l[0],
            ln_pi_good + log_l[1],
            ln_pi_bad + log_l[2],
        ];
        // Log-sum-exp normalization
        let m = log_post.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        let exp_diff: [f64; 3] = [
            (log_post[0] - m).exp(),
            (log_post[1] - m).exp(),
            (log_post[2] - m).exp(),
        ];
        let z = exp_diff[0] + exp_diff[1] + exp_diff[2];
        if !z.is_finite() || z < 1e-12 {
            // numerical issue — skip this observation
            continue;
        }
        let gamma_no = exp_diff[0] / z;
        let gamma_good = exp_diff[1] / z;
        let gamma_bad = exp_diff[2] / z;

        sum_gamma_event += gamma_good + gamma_bad;
        sum_gamma_bad += gamma_bad;
        sum_b += b as f64;
        sum_s += s as f64;
        sum_gamma_good_b += gamma_good * (b as f64);
        sum_gamma_bad_s += gamma_bad * (s as f64);
        let _ = gamma_no; // not used in M-step (encoded via complement)
    }

    let new_alpha = (sum_gamma_event / t).clamp(1e-4, 1.0 - 1e-4);
    let new_delta = if sum_gamma_event > 1e-6 {
        (sum_gamma_bad / sum_gamma_event).clamp(1e-4, 1.0 - 1e-4)
    } else {
        delta
    };

    // M-step for μ: μ̂ = [Σ γ_good·B + γ_bad·S] / [Σ (γ_good + γ_bad)]
    let new_mu = if sum_gamma_event > 1e-6 {
        ((sum_gamma_good_b + sum_gamma_bad_s) / sum_gamma_event).max(1e-4)
    } else {
        mu
    };

    // M-step for ε: total uninformed flow = total flow − informed flow
    // expected informed contribution = sum_gamma_event × μ
    let total_flow = sum_b + sum_s;
    let informed_flow = sum_gamma_event * new_mu;
    let uninformed_flow = (total_flow - informed_flow).max(0.0);
    let new_epsilon = (uninformed_flow / (2.0 * t)).max(1e-4);

    (new_alpha, new_delta, new_mu, new_epsilon)
}

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

    #[test]
    fn test_cold_start_features_pin_near_default() {
        let pe = PinEstimator::with_defaults();
        let [pin, mu, eps] = pe.features();
        assert!(pin >= 0.0 && pin <= 1.0, "PIN must be in [0,1], got {pin}");
        assert!(mu > 0.0);
        assert!(eps > 0.0);
    }

    #[test]
    fn test_pin_bounded_under_balanced_input() {
        // Honest test: PIN is **weakly identified** on near-symmetric data
        // (the EKO mixture has multiple local optima that fit equally well
        // when no clear-news days are present). Our load-bearing check is
        // just that PIN stays in [0, 1] and doesn't explode — the asymmetric
        // case (`test_strongly_asymmetric_flow_higher_pin`) is what proves
        // PIN actually responds to informed-flow signal.
        let mut pe = PinEstimator::new(200, 1);
        pe.max_iters = 200;
        for i in 0..100 {
            let b = 45 + ((i * 7) % 11) as u64;
            let s = 45 + ((i * 13) % 11) as u64;
            pe.update_bar(b, s);
        }
        let [pin, mu, eps] = pe.features();
        assert!((0.0..=1.0).contains(&pin), "PIN must be in [0,1], got {pin}");
        assert!(mu > 0.0 && mu.is_finite(), "μ must be positive finite, got {mu}");
        assert!(eps > 0.0 && eps.is_finite(), "ε must be positive finite, got {eps}");
    }

    #[test]
    fn test_strongly_asymmetric_flow_higher_pin() {
        // Alternate: half the days are (200, 10) and half are (10, 200) —
        // strong news days. PIN should be substantially > 0.
        let mut pe = PinEstimator::new(200, 1);
        pe.max_iters = 300;
        for i in 0..100 {
            if i % 2 == 0 {
                pe.update_bar(200, 10); // good news days
            } else {
                pe.update_bar(10, 200); // bad news days
            }
        }
        let [pin, _, _] = pe.features();
        assert!(pin > 0.4, "strong news-driven flow → high PIN, got {pin}");
    }

    #[test]
    fn test_log_likelihood_neg_for_zero_eps_or_mu() {
        let ll = log_likelihoods(10, 10, 0.0, 0.0);
        // log(eps + 1e-12) ≈ -27 — all elements very negative
        assert!(ll.iter().all(|&x| x < 0.0), "tiny rates → very negative log-L");
    }

    #[test]
    fn test_force_refit_does_nothing_with_under_30_obs() {
        let mut pe = PinEstimator::new(200, 50);
        for _ in 0..20 {
            pe.update_bar(10, 10);
        }
        let before = pe.features();
        pe.force_refit();
        let after = pe.features();
        assert_eq!(before, after, "<30 obs → force_refit no-op");
    }

    #[test]
    fn test_em_converges_within_max_iters() {
        // Stationary data: EM should converge.
        let mut pe = PinEstimator::new(200, 1);
        pe.max_iters = 100;
        for _ in 0..50 {
            pe.update_bar(20, 30);
        }
        // After several refits, params should stabilize
        let snap1 = pe.features();
        for _ in 0..10 {
            pe.update_bar(20, 30); // more of same
        }
        let snap2 = pe.features();
        // Should not change drastically
        for i in 0..3 {
            assert!(
                (snap1[i] - snap2[i]).abs() < 0.1,
                "stationary data → params should be stable; diff[{i}] = {}",
                (snap1[i] - snap2[i]).abs()
            );
        }
    }

    #[test]
    fn test_reset_clears_state() {
        let mut pe = PinEstimator::new(200, 50);
        for _ in 0..100 {
            pe.update_bar(50, 30);
        }
        pe.reset();
        let [_, _, eps] = pe.features();
        assert_eq!(eps, 1.0, "reset should restore default ε=1.0");
    }
}
