use crate::{
    N_INF, approx,
    factor::{VarId, VarStore},
    gaussian::Gaussian,
    ln_interval, ln_sf,
};

/// EP truncation factor on a diff variable.
///
/// Implements the rectified-Gaussian approximation that turns a diff
/// distribution into a "this team rank-beats that team" or "tied" likelihood.
/// Stores its outgoing message to the diff variable so the cavity computation
/// produces the correct EP message on each propagation.
#[derive(Debug)]
pub struct TruncFactor {
    pub diff: VarId,
    pub margin: f64,
    pub tie: bool,
    /// Outgoing message to the diff variable (initial: `N_INF`, the EP identity).
    pub(crate) msg: Gaussian,
    /// Cached evidence (linear, not log) computed from the cavity on first propagation.
    pub(crate) log_evidence_cached: Option<f64>,
}

impl TruncFactor {
    #[must_use]
    pub fn new(diff: VarId, margin: f64, tie: bool) -> Self {
        Self {
            diff,
            margin,
            tie,
            msg: N_INF,
            log_evidence_cached: None,
        }
    }
}

impl TruncFactor {
    /// Propagate this factor's message, optionally damping the update in
    /// natural-parameter space. `alpha = 1.0` matches `Factor::propagate`
    /// exactly; `alpha < 1.0` writes `α·new_msg + (1−α)·old_msg`.
    pub(crate) fn propagate_with_alpha(&mut self, vars: &mut VarStore, alpha: f64) -> (f64, f64) {
        let marginal = vars.get(self.diff);
        let cavity = marginal.cavity(self.msg);

        if self.log_evidence_cached.is_none() {
            self.log_evidence_cached = Some(cavity_log_evidence(cavity, self.margin, self.tie));
        }

        let trunc = approx(cavity, self.margin, self.tie);
        let new_msg = trunc.cavity(cavity);

        let damped = self.msg.damp_natural(new_msg, alpha);
        let old_msg = self.msg;
        self.msg = damped;

        // marginal_new = cavity * stored_msg. With alpha = 1.0 this equals
        // `trunc` (since cavity * new_msg = trunc by construction); with
        // alpha < 1.0 it reflects the partially-applied update.
        vars.set(self.diff, cavity.ep_product(damped));

        old_msg.delta(damped)
    }
}

/// Undamped wrappers, used by this module's tests. Inference drives these
/// factors through `propagate_with_alpha` and reads the cached log evidence
/// directly, so these are not on any production path.
#[cfg(test)]
impl TruncFactor {
    pub(crate) fn propagate(&mut self, vars: &mut VarStore) -> (f64, f64) {
        self.propagate_with_alpha(vars, 1.0)
    }
}

/// `ln P(diff > margin)` for a win, `ln P(|diff| < margin)` for a tie.
///
/// Computed in log space throughout. Two earlier shapes both lost the tail:
/// `1 - cdf(..)` cancelled away every digit of an unlikely outcome, and even
/// once that was fixed the linear probability underflows to zero past about 38
/// sigma, where clamping reported -708 nats regardless of the truth. An upset
/// is the observation a log-evidence figure exists to notice, so it has to stay
/// exact precisely where it is smallest.
fn cavity_log_evidence(diff: Gaussian, margin: f64, tie: bool) -> f64 {
    let (mu, sigma) = (diff.mu(), diff.sigma());

    let value = if tie {
        ln_interval(-margin, margin, mu, sigma)
    } else {
        ln_sf(margin, mu, sigma)
    };

    // A degenerate cavity is the only route to a non-finite result; keep the
    // old floor for it rather than letting -inf poison the whole history's sum.
    if value.is_finite() {
        value
    } else {
        libm::log(f64::MIN_POSITIVE)
    }
}

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

    #[test]
    fn idempotent_after_convergence() {
        // After enough iterations, propagate should return ~0 delta.
        let mut vars = VarStore::new();
        let diff = vars.alloc(Gaussian::from_ms(2.0, 3.0));

        let mut f = TruncFactor::new(diff, 0.0, false);

        // Propagate many times; delta should drop toward 0.
        let mut last = (f64::INFINITY, f64::INFINITY);
        for _ in 0..20 {
            last = f.propagate(&mut vars);
        }
        assert!(last.0 < 1e-10, "expected converged delta, got {}", last.0);
        assert!(last.1 < 1e-10);
    }

    #[test]
    fn evidence_cached_on_first_propagate() {
        let mut vars = VarStore::new();
        let diff = vars.alloc(Gaussian::from_ms(2.0, 3.0));

        let mut f = TruncFactor::new(diff, 0.0, false);
        assert!(f.log_evidence_cached.is_none());

        f.propagate(&mut vars);
        assert!(f.log_evidence_cached.is_some());
        let first = f.log_evidence_cached.unwrap();

        // Evidence should be P(diff > 0) for diff ~ N(2, 9) ≈ 0.748
        assert!(first.exp() > 0.7);
        assert!(first.exp() < 0.8);

        // Subsequent propagations don't change it.
        f.propagate(&mut vars);
        assert_eq!(f.log_evidence_cached.unwrap(), first);
    }

    /// The defect this guards: `1 - cdf` collapsed to zero for a surprising
    /// result, the clamp turned that into `f64::MIN_POSITIVE`, and
    /// `log_evidence` reported ln of *that* — about -708 whatever the truth
    /// was. An upset is the observation a model-comparison score exists to
    /// notice, so it was wrong exactly where it mattered.
    #[test]
    fn evidence_of_an_upset_is_not_flattened_to_the_clamp_floor() {
        // diff ~ N(-9, 1) with margin 0: the favoured side lost by nine sigma.
        let evidence = cavity_log_evidence(Gaussian::from_ms(-9.0, 1.0), 0.0, false).exp();

        assert!(
            evidence > f64::MIN_POSITIVE,
            "evidence collapsed onto the clamp floor: {evidence}"
        );
        // P(X > 0) for X ~ N(-9, 1) is the standard normal tail at 9 sigma.
        assert!(
            (evidence - 1.128_588e-19).abs() / 1.128_588e-19 < 1e-6,
            "expected ~1.13e-19, got {evidence}"
        );
        assert!(
            (evidence.ln() + 43.628).abs() < 1e-2,
            "log evidence {} should be about -43.6, not -708",
            evidence.ln()
        );
    }

    /// Evidence must stay finite and positive however extreme the mismatch,
    /// since `log_evidence` sums across the whole history and one `-inf` or
    /// `NaN` poisons all of it.
    ///
    /// Finiteness alone is too weak a bar — the clamped version was finite too,
    /// and wrong by hundreds of nats. `log_evidence_tracks_the_analytic_tail`
    /// below is the assertion that actually holds this up.
    #[test]
    fn evidence_stays_positive_and_finite_at_any_separation() {
        for mu in [-300.0f64, -50.0, -9.0, 0.0, 9.0, 50.0, 300.0] {
            for tie in [false, true] {
                let ln_e = cavity_log_evidence(Gaussian::from_ms(mu, 1.0), 1.0, tie);
                assert!(
                    ln_e.is_finite() && ln_e <= 0.0,
                    "mu={mu} tie={tie}: log evidence {ln_e} is not a log-probability"
                );
            }
        }
    }

    /// The clamp used to floor everything past ~38 sigma at `ln(MIN_POSITIVE)`
    /// = -708, however far out the real observation was. In log space the
    /// answer is a polynomial and stays exact: at 1000 sigma the truth is about
    /// -500_000 nats, and -708 is not a rounding error.
    #[test]
    fn log_evidence_tracks_the_analytic_tail() {
        for mu in [-40.0f64, -60.0, -100.0, -1000.0] {
            // P(diff > 0) for diff ~ N(mu, 1), mu far below zero.
            let got = cavity_log_evidence(Gaussian::from_ms(mu, 1.0), 0.0, false);

            // ln Phi(mu) ~ -mu^2/2 - ln(-mu) - ln(sqrt(2 pi)) for mu << 0.
            let z = -mu;
            let approx = -0.5 * z * z - z.ln() - (2.0 * std::f64::consts::PI).sqrt().ln();

            assert!(
                got < libm::log(f64::MIN_POSITIVE),
                "mu={mu}: {got} is still stuck on the old clamp floor"
            );
            assert!(
                (got - approx).abs() / approx.abs() < 1e-3,
                "mu={mu}: got {got}, asymptotic expectation {approx}"
            );
        }
    }

    #[test]
    fn tie_evidence_uses_two_sided() {
        let mut vars = VarStore::new();
        let diff = vars.alloc(Gaussian::from_ms(0.0, 2.0));

        let mut f = TruncFactor::new(diff, 1.0, true);
        f.propagate(&mut vars);

        // For diff ~ N(0, 4), tie=true with margin=1: P(-1 < diff < 1) ≈ 0.383
        let ev = f.log_evidence_cached.unwrap().exp();
        assert!(ev > 0.35 && ev < 0.42);
    }

    #[test]
    fn propagate_with_alpha_one_matches_undamped_propagate() {
        let mut vars_a = VarStore::new();
        let diff_a = vars_a.alloc(Gaussian::from_ms(2.0, 3.0));
        let mut f_a = TruncFactor::new(diff_a, 0.0, false);
        let delta_a = f_a.propagate(&mut vars_a);
        let result_a = vars_a.get(diff_a);

        let mut vars_b = VarStore::new();
        let diff_b = vars_b.alloc(Gaussian::from_ms(2.0, 3.0));
        let mut f_b = TruncFactor::new(diff_b, 0.0, false);
        let delta_b = f_b.propagate_with_alpha(&mut vars_b, 1.0);
        let result_b = vars_b.get(diff_b);

        assert_eq!(result_a.pi(), result_b.pi());
        assert_eq!(result_a.tau(), result_b.tau());
        assert_eq!(delta_a, delta_b);
        assert_eq!(f_a.msg.pi(), f_b.msg.pi());
        assert_eq!(f_a.msg.tau(), f_b.msg.tau());
    }

    #[test]
    fn propagate_with_alpha_half_blends_msg_in_natural_params() {
        // Run undamped to capture (initial_msg, undamped_new_msg).
        let mut vars_full = VarStore::new();
        let diff_full = vars_full.alloc(Gaussian::from_ms(2.0, 3.0));
        let mut f_full = TruncFactor::new(diff_full, 0.0, false);
        let initial_msg_pi = f_full.msg.pi();
        let initial_msg_tau = f_full.msg.tau();
        f_full.propagate(&mut vars_full);
        let undamped_msg_pi = f_full.msg.pi();
        let undamped_msg_tau = f_full.msg.tau();

        // Run damped at α = 0.5 from the same initial state.
        let mut vars_half = VarStore::new();
        let diff_half = vars_half.alloc(Gaussian::from_ms(2.0, 3.0));
        let mut f_half = TruncFactor::new(diff_half, 0.0, false);
        f_half.propagate_with_alpha(&mut vars_half, 0.5);

        let expected_pi = 0.5 * undamped_msg_pi + 0.5 * initial_msg_pi;
        let expected_tau = 0.5 * undamped_msg_tau + 0.5 * initial_msg_tau;
        assert!((f_half.msg.pi() - expected_pi).abs() < 1e-12);
        assert!((f_half.msg.tau() - expected_tau).abs() < 1e-12);
    }
}