Skip to main content

ipfrs_tensorlogic/
bayesian_updater.rs

1//! Bayesian belief updating with conjugate priors, likelihood functions,
2//! and posterior inference.
3//!
4//! This module provides a production-grade [`BayesianUpdateEngine`] that
5//! supports the four canonical conjugate-prior / likelihood pairs:
6//!
7//! | Prior      | Likelihood  |
8//! |------------|-------------|
9//! | Beta       | Bernoulli   |
10//! | Gaussian   | Gaussian    |
11//! | Dirichlet  | Categorical |
12//! | Gamma      | Poisson     |
13//!
14//! # Example
15//!
16//! ```
17//! use ipfrs_tensorlogic::bayesian_updater::{
18//!     BayesianUpdateEngine, Prior, Observation,
19//! };
20//!
21//! let mut engine = BayesianUpdateEngine::new(64);
22//!
23//! // Start with a uniform Beta(1,1) prior and observe 7 successes in 10 trials.
24//! let prior = Prior::Beta { alpha: 1.0, beta: 1.0 };
25//! let obs   = Observation::Bernoulli { successes: 7, trials: 10 };
26//!
27//! let posterior = engine.update(prior, &obs).expect("example: should succeed in docs");
28//! // Posterior should be Beta(8, 4)
29//! if let Prior::Beta { alpha, beta } = &posterior.updated {
30//!     assert!((alpha - 8.0).abs() < 1e-10);
31//!     assert!((beta  - 4.0).abs() < 1e-10);
32//! }
33//! ```
34
35use std::collections::VecDeque;
36use thiserror::Error;
37
38// ──────────────────────────────────────────────────────────────────────────────
39// Error type
40// ──────────────────────────────────────────────────────────────────────────────
41
42/// Errors that can arise during Bayesian updating.
43#[derive(Debug, Error, Clone, PartialEq)]
44pub enum BayesError {
45    /// The prior distribution family does not match the observation type.
46    #[error(
47        "prior/observation mismatch: cannot update {prior_type} prior with {obs_type} observation"
48    )]
49    PriorObservationMismatch {
50        /// Name of the prior family (e.g. "Beta").
51        prior_type: String,
52        /// Name of the observation type (e.g. "Gaussian").
53        obs_type: String,
54    },
55
56    /// One or more parameter values are invalid (e.g. non-positive concentration).
57    #[error("invalid parameters: {0}")]
58    InvalidParameters(String),
59
60    /// The requested operation is not supported for this combination of types.
61    #[error("unsupported operation: {0}")]
62    UnsupportedOperation(String),
63
64    /// A numerical error occurred (e.g. NaN or infinity).
65    #[error("numerical error: {0}")]
66    NumericalError(String),
67}
68
69// ──────────────────────────────────────────────────────────────────────────────
70// Core enums and structs
71// ──────────────────────────────────────────────────────────────────────────────
72
73/// A conjugate prior distribution.
74#[derive(Debug, Clone, PartialEq)]
75pub enum Prior {
76    /// Beta(α, β) — conjugate prior for Bernoulli/Binomial likelihoods.
77    Beta {
78        /// Pseudo-successes (must be > 0).
79        alpha: f64,
80        /// Pseudo-failures (must be > 0).
81        beta: f64,
82    },
83    /// Normal(μ, σ²) — conjugate prior for Gaussian likelihoods with known
84    /// observation variance.
85    Gaussian {
86        /// Prior mean.
87        mean: f64,
88        /// Prior variance (must be > 0).
89        variance: f64,
90    },
91    /// Dirichlet(α₁, …, αₖ) — conjugate prior for Categorical likelihoods.
92    Dirichlet {
93        /// Concentration parameters (all must be > 0).
94        alphas: Vec<f64>,
95    },
96    /// Gamma(shape, rate) — conjugate prior for Poisson likelihoods.
97    Gamma {
98        /// Shape parameter (must be > 0).
99        shape: f64,
100        /// Rate parameter (must be > 0).
101        rate: f64,
102    },
103}
104
105impl Prior {
106    /// Human-readable name for error messages.
107    fn type_name(&self) -> &'static str {
108        match self {
109            Prior::Beta { .. } => "Beta",
110            Prior::Gaussian { .. } => "Gaussian",
111            Prior::Dirichlet { .. } => "Dirichlet",
112            Prior::Gamma { .. } => "Gamma",
113        }
114    }
115
116    /// Validate that all parameters are in their legal ranges.
117    fn validate(&self) -> Result<(), BayesError> {
118        match self {
119            Prior::Beta { alpha, beta } => {
120                if *alpha <= 0.0 || alpha.is_nan() {
121                    return Err(BayesError::InvalidParameters(format!(
122                        "Beta alpha must be > 0, got {alpha}"
123                    )));
124                }
125                if *beta <= 0.0 || beta.is_nan() {
126                    return Err(BayesError::InvalidParameters(format!(
127                        "Beta beta must be > 0, got {beta}"
128                    )));
129                }
130            }
131            Prior::Gaussian { variance, .. } => {
132                if *variance <= 0.0 || variance.is_nan() {
133                    return Err(BayesError::InvalidParameters(format!(
134                        "Gaussian variance must be > 0, got {variance}"
135                    )));
136                }
137            }
138            Prior::Dirichlet { alphas } => {
139                if alphas.is_empty() {
140                    return Err(BayesError::InvalidParameters(
141                        "Dirichlet alphas must be non-empty".to_string(),
142                    ));
143                }
144                for (i, &a) in alphas.iter().enumerate() {
145                    if a <= 0.0 || a.is_nan() {
146                        return Err(BayesError::InvalidParameters(format!(
147                            "Dirichlet alpha[{i}] must be > 0, got {a}"
148                        )));
149                    }
150                }
151            }
152            Prior::Gamma { shape, rate } => {
153                if *shape <= 0.0 || shape.is_nan() {
154                    return Err(BayesError::InvalidParameters(format!(
155                        "Gamma shape must be > 0, got {shape}"
156                    )));
157                }
158                if *rate <= 0.0 || rate.is_nan() {
159                    return Err(BayesError::InvalidParameters(format!(
160                        "Gamma rate must be > 0, got {rate}"
161                    )));
162                }
163            }
164        }
165        Ok(())
166    }
167}
168
169/// An observed datum used to update a prior.
170#[derive(Debug, Clone, PartialEq)]
171pub enum Observation {
172    /// Outcome of a series of Bernoulli trials.
173    Bernoulli {
174        /// Number of successes observed.
175        successes: u64,
176        /// Total number of trials.
177        trials: u64,
178    },
179    /// Sufficient statistics from a Gaussian-distributed sample.
180    Gaussian {
181        /// Sample mean.
182        sample_mean: f64,
183        /// Sample variance (must be > 0).
184        sample_variance: f64,
185        /// Sample size.
186        n: u64,
187    },
188    /// Category counts from a Categorical/Multinomial sample.
189    Categorical {
190        /// Observed count for each category (length must equal Dirichlet dim).
191        counts: Vec<u64>,
192    },
193    /// Sufficient statistics from a Poisson process.
194    Poisson {
195        /// Total events observed.
196        total_events: u64,
197        /// Total elapsed time (must be > 0).
198        total_time: f64,
199    },
200}
201
202impl Observation {
203    /// Human-readable name for error messages.
204    fn type_name(&self) -> &'static str {
205        match self {
206            Observation::Bernoulli { .. } => "Bernoulli",
207            Observation::Gaussian { .. } => "Gaussian",
208            Observation::Categorical { .. } => "Categorical",
209            Observation::Poisson { .. } => "Poisson",
210        }
211    }
212}
213
214/// The result of a single Bayesian update step.
215#[derive(Debug, Clone)]
216pub struct Posterior {
217    /// The prior used in this update.
218    pub prior: Prior,
219    /// Human-readable description of the likelihood model.
220    pub likelihood_type: String,
221    /// The updated (posterior) distribution.
222    pub updated: Prior,
223    /// Log marginal likelihood ln p(observation | model).
224    pub log_marginal: f64,
225}
226
227/// A credible interval `[lower, upper]` at the stated probability level.
228#[derive(Debug, Clone, PartialEq)]
229pub struct CredibleInterval {
230    /// Lower bound of the interval.
231    pub lower: f64,
232    /// Upper bound of the interval.
233    pub upper: f64,
234    /// Nominal probability mass contained (e.g. `0.95` for a 95% CI).
235    pub probability: f64,
236}
237
238// ──────────────────────────────────────────────────────────────────────────────
239// Pure-Rust math helpers
240// ──────────────────────────────────────────────────────────────────────────────
241
242/// Stirling-series approximation for ln Γ(x), accurate to ~1e-13 for x > 0.5.
243/// Uses the recurrence Γ(x) = (x-1)·Γ(x-1) to shift small arguments up.
244fn ln_gamma(x: f64) -> f64 {
245    // Reflect small values upward via the recurrence ln Γ(x) = ln Γ(x+1) – ln x
246    if x < 0.5 {
247        // Use reflection: ln Γ(x) = ln π - ln sin(πx) - ln Γ(1-x)
248        // But for our use-cases x is always positive, so just recurse.
249        return ln_gamma(x + 1.0) - x.ln();
250    }
251    if x < 7.0 {
252        // Shift x into the Stirling-series regime
253        return ln_gamma(x + 1.0) - x.ln();
254    }
255    // Stirling's series: ln Γ(x) ≈ (x-0.5)·ln x - x + 0.5·ln(2π) + 1/(12x) - 1/(360x³) + …
256    let half_ln_two_pi = 0.918_938_533_204_672_8_f64; // 0.5*ln(2π)
257    let inv_x = 1.0 / x;
258    let inv_x2 = inv_x * inv_x;
259    (x - 0.5) * x.ln() - x
260        + half_ln_two_pi
261        + inv_x * (1.0 / 12.0 - inv_x2 * (1.0 / 360.0 - inv_x2 / 1260.0))
262}
263
264/// ln B(a, b) = ln Γ(a) + ln Γ(b) – ln Γ(a+b).
265fn log_beta(a: f64, b: f64) -> f64 {
266    ln_gamma(a) + ln_gamma(b) - ln_gamma(a + b)
267}
268
269/// ln normalisation constant of the Dirichlet distribution:
270/// Σᵢ ln Γ(αᵢ) – ln Γ(Σᵢ αᵢ).
271fn log_dirichlet_norm(alphas: &[f64]) -> f64 {
272    let sum: f64 = alphas.iter().sum();
273    let sum_lg: f64 = alphas.iter().map(|&a| ln_gamma(a)).sum();
274    sum_lg - ln_gamma(sum)
275}
276
277/// Digamma function ψ(x) = d/dx ln Γ(x).
278///
279/// Uses the asymptotic series for x > 6 and the recurrence ψ(x) = ψ(x+1) – 1/x
280/// for smaller arguments.
281fn digamma(x: f64) -> f64 {
282    if x < 6.0 {
283        // Recurrence: ψ(x) = ψ(x+1) - 1/x
284        return digamma(x + 1.0) - 1.0 / x;
285    }
286    // Asymptotic expansion: ψ(x) ≈ ln x - 1/(2x) - 1/(12x²) + 1/(120x⁴) - 1/(252x⁶)
287    let inv_x = 1.0 / x;
288    let inv_x2 = inv_x * inv_x;
289    x.ln() - 0.5 * inv_x - inv_x2 * (1.0 / 12.0 - inv_x2 * (1.0 / 120.0 - inv_x2 / 252.0))
290}
291
292/// Z-score for a given two-tailed credible probability using a rational
293/// approximation to the probit function (accurate to ~1e-4).
294///
295/// Reference: Abramowitz & Stegun 26.2.17.
296fn z_score(probability: f64) -> f64 {
297    // We need the upper tail quantile for (1 + p) / 2
298    let p = (1.0 + probability) / 2.0;
299    // Rational approximation to Φ⁻¹(p) for 0.5 < p < 1
300    if (p - 0.5).abs() < 1e-10 {
301        return 0.0;
302    }
303    let t = (-2.0 * (1.0 - p).ln()).sqrt();
304    let c0 = 2.515_517;
305    let c1 = 0.802_853;
306    let c2 = 0.010_328;
307    let d1 = 1.432_788;
308    let d2 = 0.189_269;
309    let d3 = 0.001_308;
310    let numer = c0 + c1 * t + c2 * t * t;
311    let denom = 1.0 + d1 * t + d2 * t * t + d3 * t * t * t;
312    t - numer / denom
313}
314
315/// Guard against NaN/Inf in a computed f64.
316fn check_finite(val: f64, context: &str) -> Result<f64, BayesError> {
317    if val.is_finite() {
318        Ok(val)
319    } else {
320        Err(BayesError::NumericalError(format!(
321            "{context}: computed non-finite value {val}"
322        )))
323    }
324}
325
326// ──────────────────────────────────────────────────────────────────────────────
327// Engine
328// ──────────────────────────────────────────────────────────────────────────────
329
330/// Bayesian belief updating engine.
331///
332/// Maintains a bounded history of [`Posterior`] results and supports both
333/// single-step updates and sequential folding of multiple observations.
334#[derive(Debug)]
335pub struct BayesianUpdateEngine {
336    /// Ordered history of posteriors, newest at the back.
337    history: VecDeque<Posterior>,
338    /// Maximum number of posteriors to retain.
339    max_history: usize,
340}
341
342impl BayesianUpdateEngine {
343    /// Create a new engine with the given history capacity.
344    pub fn new(max_history: usize) -> Self {
345        Self {
346            history: VecDeque::with_capacity(max_history.min(1024)),
347            max_history,
348        }
349    }
350
351    // ── Core update ──────────────────────────────────────────────────────────
352
353    /// Perform a single Bayesian update, returning the posterior.
354    ///
355    /// The method validates parameters and dispatches to the appropriate
356    /// conjugate-update formula.
357    pub fn update(
358        &mut self,
359        prior: Prior,
360        observation: &Observation,
361    ) -> Result<Posterior, BayesError> {
362        prior.validate()?;
363
364        let posterior = match (&prior, observation) {
365            // ── Beta–Bernoulli ────────────────────────────────────────────────
366            (Prior::Beta { alpha, beta }, Observation::Bernoulli { successes, trials }) => {
367                if successes > trials {
368                    return Err(BayesError::InvalidParameters(format!(
369                        "successes ({successes}) cannot exceed trials ({trials})"
370                    )));
371                }
372                let s = *successes as f64;
373                let f = (*trials - *successes) as f64;
374                let alpha_post = alpha + s;
375                let beta_post = beta + f;
376                let log_marginal = check_finite(
377                    log_beta(alpha_post, beta_post) - log_beta(*alpha, *beta),
378                    "Beta-Bernoulli log_marginal",
379                )?;
380                Posterior {
381                    prior: prior.clone(),
382                    likelihood_type: "Bernoulli".to_string(),
383                    updated: Prior::Beta {
384                        alpha: alpha_post,
385                        beta: beta_post,
386                    },
387                    log_marginal,
388                }
389            }
390
391            // ── Gaussian–Gaussian (normal-normal conjugate) ───────────────────
392            (
393                Prior::Gaussian {
394                    mean: prior_mean,
395                    variance: prior_var,
396                },
397                Observation::Gaussian {
398                    sample_mean,
399                    sample_variance,
400                    n,
401                },
402            ) => {
403                if *sample_variance <= 0.0 || sample_variance.is_nan() {
404                    return Err(BayesError::InvalidParameters(format!(
405                        "sample_variance must be > 0, got {sample_variance}"
406                    )));
407                }
408                if *n == 0 {
409                    return Err(BayesError::InvalidParameters(
410                        "n must be > 0 for Gaussian observation".to_string(),
411                    ));
412                }
413                let n_f = *n as f64;
414                // Posterior precision = 1/σ²_0 + n/σ²_obs
415                let post_prec = 1.0 / prior_var + n_f / sample_variance;
416                let post_var = 1.0 / post_prec;
417                let post_mean =
418                    post_var * (prior_mean / prior_var + n_f * sample_mean / sample_variance);
419
420                // ln p(x̄ | model) ≈ -0.5 * ln(2π * (σ²_0 + σ²_obs / n))
421                let effective_var = prior_var + sample_variance / n_f;
422                let log_marginal = check_finite(
423                    -0.5 * (std::f64::consts::TAU * effective_var).ln(),
424                    "Gaussian-Gaussian log_marginal",
425                )?;
426
427                Posterior {
428                    prior: prior.clone(),
429                    likelihood_type: "Gaussian".to_string(),
430                    updated: Prior::Gaussian {
431                        mean: post_mean,
432                        variance: post_var,
433                    },
434                    log_marginal,
435                }
436            }
437
438            // ── Dirichlet–Categorical ─────────────────────────────────────────
439            (Prior::Dirichlet { alphas }, Observation::Categorical { counts }) => {
440                if alphas.len() != counts.len() {
441                    return Err(BayesError::InvalidParameters(format!(
442                        "Dirichlet dim {} != Categorical counts dim {}",
443                        alphas.len(),
444                        counts.len()
445                    )));
446                }
447                let alphas_post: Vec<f64> = alphas
448                    .iter()
449                    .zip(counts.iter())
450                    .map(|(&a, &c)| a + c as f64)
451                    .collect();
452
453                let log_marginal = check_finite(
454                    log_dirichlet_norm(&alphas_post) - log_dirichlet_norm(alphas),
455                    "Dirichlet-Categorical log_marginal",
456                )?;
457
458                Posterior {
459                    prior: prior.clone(),
460                    likelihood_type: "Categorical".to_string(),
461                    updated: Prior::Dirichlet {
462                        alphas: alphas_post,
463                    },
464                    log_marginal,
465                }
466            }
467
468            // ── Gamma–Poisson ─────────────────────────────────────────────────
469            (
470                Prior::Gamma { shape, rate },
471                Observation::Poisson {
472                    total_events,
473                    total_time,
474                },
475            ) => {
476                if *total_time <= 0.0 || total_time.is_nan() {
477                    return Err(BayesError::InvalidParameters(format!(
478                        "total_time must be > 0, got {total_time}"
479                    )));
480                }
481                let k = *total_events as f64;
482                let shape_post = shape + k;
483                let rate_post = rate + total_time;
484
485                // ln p(data | model) = lgamma(shape') – lgamma(shape)
486                //                    + shape * ln(rate) – shape' * ln(rate')
487                let log_marginal = check_finite(
488                    ln_gamma(shape_post) - ln_gamma(*shape) + shape * rate.ln()
489                        - shape_post * rate_post.ln(),
490                    "Gamma-Poisson log_marginal",
491                )?;
492
493                Posterior {
494                    prior: prior.clone(),
495                    likelihood_type: "Poisson".to_string(),
496                    updated: Prior::Gamma {
497                        shape: shape_post,
498                        rate: rate_post,
499                    },
500                    log_marginal,
501                }
502            }
503
504            // ── Mismatch ──────────────────────────────────────────────────────
505            _ => {
506                return Err(BayesError::PriorObservationMismatch {
507                    prior_type: prior.type_name().to_string(),
508                    obs_type: observation.type_name().to_string(),
509                });
510            }
511        };
512
513        // Store in history
514        if self.history.len() >= self.max_history && self.max_history > 0 {
515            self.history.pop_front();
516        }
517        if self.max_history > 0 {
518            self.history.push_back(posterior.clone());
519        }
520
521        Ok(posterior)
522    }
523
524    // ── Sequential update ────────────────────────────────────────────────────
525
526    /// Apply a sequence of observations left-to-right, using each posterior
527    /// as the prior for the next update.
528    ///
529    /// Returns the final posterior, or an error at the first failing update.
530    pub fn sequential_update(
531        &mut self,
532        prior: Prior,
533        observations: &[Observation],
534    ) -> Result<Posterior, BayesError> {
535        if observations.is_empty() {
536            return Err(BayesError::InvalidParameters(
537                "observations slice must not be empty".to_string(),
538            ));
539        }
540
541        let mut current_prior = prior;
542        let mut last_posterior: Option<Posterior> = None;
543
544        for obs in observations {
545            let posterior = self.update(current_prior, obs)?;
546            current_prior = posterior.updated.clone();
547            last_posterior = Some(posterior);
548        }
549
550        // SAFETY: observations is non-empty, so last_posterior is Some.
551        last_posterior.ok_or_else(|| {
552            BayesError::NumericalError("unexpected empty observation sequence".to_string())
553        })
554    }
555
556    // ── Credible interval ────────────────────────────────────────────────────
557
558    /// Compute a symmetric credible interval for a posterior distribution.
559    ///
560    /// Uses normal / Wilson approximations — suitable for moderate-to-large
561    /// concentration parameters.
562    ///
563    /// # Arguments
564    /// * `posterior` – the posterior distribution.
565    /// * `probability` – the desired probability mass (e.g. `0.95`).
566    pub fn credible_interval(
567        posterior: &Prior,
568        probability: f64,
569    ) -> Result<CredibleInterval, BayesError> {
570        if !(0.0 < probability && probability < 1.0) {
571            return Err(BayesError::InvalidParameters(format!(
572                "probability must be in (0, 1), got {probability}"
573            )));
574        }
575        posterior.validate()?;
576
577        let z = z_score(probability);
578
579        match posterior {
580            Prior::Beta { alpha, beta } => {
581                let n = alpha + beta;
582                let center = alpha / n;
583                let half_width = z * (center * (1.0 - center) / n).sqrt();
584                let lower = (center - half_width).max(0.0);
585                let upper = (center + half_width).min(1.0);
586                Ok(CredibleInterval {
587                    lower,
588                    upper,
589                    probability,
590                })
591            }
592
593            Prior::Gaussian { mean, variance } => {
594                let half_width = z * variance.sqrt();
595                Ok(CredibleInterval {
596                    lower: mean - half_width,
597                    upper: mean + half_width,
598                    probability,
599                })
600            }
601
602            Prior::Gamma { shape, rate } => {
603                // Normal approximation: mean = shape/rate, var = shape/rate²
604                let mean = shape / rate;
605                let std_dev = (shape / (rate * rate)).sqrt();
606                let half_width = z * std_dev;
607                let lower = (mean - half_width).max(0.0);
608                let upper = mean + half_width;
609                Ok(CredibleInterval {
610                    lower,
611                    upper,
612                    probability,
613                })
614            }
615
616            Prior::Dirichlet { alphas } => {
617                // Return interval for the category with the highest concentration
618                let sum: f64 = alphas.iter().sum();
619                let max_alpha = alphas.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
620                let center = max_alpha / sum;
621                let half_width = z * (center * (1.0 - center) / (sum + 1.0)).sqrt();
622                let lower = (center - half_width).max(0.0);
623                let upper = (center + half_width).min(1.0);
624                Ok(CredibleInterval {
625                    lower,
626                    upper,
627                    probability,
628                })
629            }
630        }
631    }
632
633    // ── MAP estimate ─────────────────────────────────────────────────────────
634
635    /// Return the maximum a posteriori (MAP) estimate for a distribution.
636    ///
637    /// | Distribution | MAP |
638    /// |---|---|
639    /// | Beta(α,β) | (α-1)/(α+β-2) if α>1 && β>1, else α/(α+β) |
640    /// | Gaussian(μ,σ²) | μ |
641    /// | Gamma(k,r) | (k-1)/r if k>1, else 0 |
642    /// | Dirichlet(α) | argmax(αᵢ) / Σαᵢ |
643    pub fn map_estimate(posterior: &Prior) -> f64 {
644        match posterior {
645            Prior::Beta { alpha, beta } => {
646                if *alpha > 1.0 && *beta > 1.0 {
647                    (alpha - 1.0) / (alpha + beta - 2.0)
648                } else {
649                    alpha / (alpha + beta)
650                }
651            }
652            Prior::Gaussian { mean, .. } => *mean,
653            Prior::Gamma { shape, rate } => {
654                if *shape > 1.0 {
655                    (shape - 1.0) / rate
656                } else {
657                    0.0
658                }
659            }
660            Prior::Dirichlet { alphas } => {
661                let sum: f64 = alphas.iter().sum();
662                let max_alpha = alphas.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
663                max_alpha / sum
664            }
665        }
666    }
667
668    // ── KL divergence ────────────────────────────────────────────────────────
669
670    /// Compute the KL divergence KL(p ‖ q) between two distributions of the
671    /// same family.
672    ///
673    /// Supported pairs: Beta vs Beta, Gaussian vs Gaussian.
674    /// All other combinations return [`BayesError::UnsupportedOperation`].
675    pub fn kl_divergence(p: &Prior, q: &Prior) -> Result<f64, BayesError> {
676        match (p, q) {
677            (
678                Prior::Beta {
679                    alpha: ap,
680                    beta: bp,
681                },
682                Prior::Beta {
683                    alpha: aq,
684                    beta: bq,
685                },
686            ) => {
687                // KL(Beta(αp,βp) ‖ Beta(αq,βq))
688                // = log B(αq,βq) - log B(αp,βp)
689                //   + (αp - αq) ψ(αp) + (βp - βq) ψ(βp)
690                //   - (αp + βp - αq - βq) ψ(αp + βp)
691                let psi_ap = digamma(*ap);
692                let psi_bp = digamma(*bp);
693                let psi_ap_bp = digamma(ap + bp);
694                let kl = log_beta(*aq, *bq) - log_beta(*ap, *bp)
695                    + (ap - aq) * psi_ap
696                    + (bp - bq) * psi_bp
697                    - ((ap + bp) - (aq + bq)) * psi_ap_bp;
698                check_finite(kl, "KL(Beta‖Beta)")
699            }
700
701            (
702                Prior::Gaussian {
703                    mean: mp,
704                    variance: vp,
705                },
706                Prior::Gaussian {
707                    mean: mq,
708                    variance: vq,
709                },
710            ) => {
711                // KL(N(μp,σp²) ‖ N(μq,σq²))
712                // = 0.5 * (ln(σq²/σp²) + σp²/σq² + (μp-μq)²/σq² - 1)
713                let kl = 0.5 * ((vq / vp).ln() + vp / vq + (mp - mq) * (mp - mq) / vq - 1.0);
714                check_finite(kl, "KL(Gaussian‖Gaussian)")
715            }
716
717            _ => Err(BayesError::UnsupportedOperation(format!(
718                "KL divergence not implemented for {} vs {}",
719                p.type_name(),
720                q.type_name()
721            ))),
722        }
723    }
724
725    // ── History accessors ────────────────────────────────────────────────────
726
727    /// Immutable reference to the update history (oldest first).
728    pub fn history(&self) -> &VecDeque<Posterior> {
729        &self.history
730    }
731
732    /// Clear the update history.
733    pub fn clear_history(&mut self) {
734        self.history.clear();
735    }
736}
737
738// ──────────────────────────────────────────────────────────────────────────────
739// Tests
740// ──────────────────────────────────────────────────────────────────────────────
741
742#[cfg(test)]
743mod tests {
744    use super::{
745        digamma, ln_gamma, log_beta, log_dirichlet_norm, z_score, BayesError, BayesianUpdateEngine,
746        Observation, Prior,
747    };
748
749    // ── Math helper tests ────────────────────────────────────────────────────
750
751    #[test]
752    fn ln_gamma_integer_values() {
753        // Γ(n) = (n-1)! for positive integers
754        // ln Γ(1) = 0
755        assert!((ln_gamma(1.0)).abs() < 1e-8);
756        // ln Γ(2) = 0
757        assert!((ln_gamma(2.0)).abs() < 1e-8);
758        // ln Γ(3) = ln 2 ≈ 0.6931
759        assert!((ln_gamma(3.0) - 2.0_f64.ln()).abs() < 1e-8);
760        // ln Γ(5) = ln 24
761        assert!((ln_gamma(5.0) - 24.0_f64.ln()).abs() < 1e-7);
762    }
763
764    #[test]
765    fn ln_gamma_half() {
766        // Γ(1/2) = √π, so ln Γ(0.5) = 0.5 ln π
767        let expected = 0.5 * std::f64::consts::PI.ln();
768        assert!((ln_gamma(0.5) - expected).abs() < 1e-6);
769    }
770
771    #[test]
772    fn log_beta_symmetry() {
773        // B(a,b) == B(b,a)
774        let diff = (log_beta(2.0, 5.0) - log_beta(5.0, 2.0)).abs();
775        assert!(diff < 1e-12);
776    }
777
778    #[test]
779    fn log_beta_known_value() {
780        // B(1,1) = 1  →  ln B(1,1) = 0  (Stirling approximation, tolerance 1e-6)
781        assert!(
782            log_beta(1.0, 1.0).abs() < 1e-6,
783            "log_beta(1,1) = {}",
784            log_beta(1.0, 1.0)
785        );
786    }
787
788    #[test]
789    fn log_dirichlet_norm_two_dim_equals_log_beta() {
790        // For k=2, Dirichlet norm = ln B(a1, a2)
791        let a = 3.0_f64;
792        let b = 7.0_f64;
793        let dir = log_dirichlet_norm(&[a, b]);
794        let lb = log_beta(a, b);
795        assert!((dir - lb).abs() < 1e-10);
796    }
797
798    #[test]
799    fn digamma_known_value() {
800        // ψ(1) = -γ ≈ -0.5772156649
801        let expected = -0.577_215_664_9_f64;
802        assert!((digamma(1.0) - expected).abs() < 1e-4);
803    }
804
805    #[test]
806    fn digamma_recurrence_property() {
807        // ψ(x+1) - ψ(x) = 1/x
808        let x = 4.5_f64;
809        let diff = digamma(x + 1.0) - digamma(x);
810        assert!((diff - 1.0 / x).abs() < 1e-8);
811    }
812
813    #[test]
814    fn z_score_95_percent() {
815        // 95% CI should give z ≈ 1.96
816        let z = z_score(0.95);
817        assert!((z - 1.96).abs() < 0.01);
818    }
819
820    #[test]
821    fn z_score_99_percent() {
822        // 99% CI should give z ≈ 2.576
823        let z = z_score(0.99);
824        assert!((z - 2.576).abs() < 0.01);
825    }
826
827    // ── Beta–Bernoulli update ────────────────────────────────────────────────
828
829    #[test]
830    fn beta_bernoulli_uniform_prior() {
831        let mut engine = BayesianUpdateEngine::new(10);
832        let prior = Prior::Beta {
833            alpha: 1.0,
834            beta: 1.0,
835        };
836        let obs = Observation::Bernoulli {
837            successes: 7,
838            trials: 10,
839        };
840        let post = engine
841            .update(prior, &obs)
842            .expect("test: TD update should succeed");
843        match &post.updated {
844            Prior::Beta { alpha, beta } => {
845                assert!((alpha - 8.0).abs() < 1e-10);
846                assert!((beta - 4.0).abs() < 1e-10);
847            }
848            _ => panic!("wrong variant"),
849        }
850    }
851
852    #[test]
853    fn beta_bernoulli_all_successes() {
854        let mut engine = BayesianUpdateEngine::new(10);
855        let prior = Prior::Beta {
856            alpha: 2.0,
857            beta: 3.0,
858        };
859        let obs = Observation::Bernoulli {
860            successes: 5,
861            trials: 5,
862        };
863        let post = engine
864            .update(prior, &obs)
865            .expect("test: TD update should succeed");
866        match post.updated {
867            Prior::Beta { alpha, beta } => {
868                assert!((alpha - 7.0).abs() < 1e-10);
869                assert!((beta - 3.0).abs() < 1e-10);
870            }
871            _ => panic!("wrong variant"),
872        }
873    }
874
875    #[test]
876    fn beta_bernoulli_zero_successes() {
877        let mut engine = BayesianUpdateEngine::new(10);
878        let prior = Prior::Beta {
879            alpha: 1.0,
880            beta: 1.0,
881        };
882        let obs = Observation::Bernoulli {
883            successes: 0,
884            trials: 5,
885        };
886        let post = engine
887            .update(prior, &obs)
888            .expect("test: TD update should succeed");
889        match post.updated {
890            Prior::Beta { alpha, beta } => {
891                assert!((alpha - 1.0).abs() < 1e-10);
892                assert!((beta - 6.0).abs() < 1e-10);
893            }
894            _ => panic!("wrong variant"),
895        }
896    }
897
898    #[test]
899    fn beta_bernoulli_log_marginal_finite() {
900        let mut engine = BayesianUpdateEngine::new(10);
901        let prior = Prior::Beta {
902            alpha: 2.0,
903            beta: 2.0,
904        };
905        let obs = Observation::Bernoulli {
906            successes: 3,
907            trials: 6,
908        };
909        let post = engine
910            .update(prior, &obs)
911            .expect("test: TD update should succeed");
912        assert!(post.log_marginal.is_finite());
913    }
914
915    #[test]
916    fn beta_bernoulli_successes_exceed_trials_error() {
917        let mut engine = BayesianUpdateEngine::new(10);
918        let prior = Prior::Beta {
919            alpha: 1.0,
920            beta: 1.0,
921        };
922        let obs = Observation::Bernoulli {
923            successes: 11,
924            trials: 10,
925        };
926        let result = engine.update(prior, &obs);
927        assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
928    }
929
930    // ── Gaussian–Gaussian update ─────────────────────────────────────────────
931
932    #[test]
933    fn gaussian_gaussian_update_basic() {
934        let mut engine = BayesianUpdateEngine::new(10);
935        let prior = Prior::Gaussian {
936            mean: 0.0,
937            variance: 1.0,
938        };
939        let obs = Observation::Gaussian {
940            sample_mean: 2.0,
941            sample_variance: 1.0,
942            n: 1,
943        };
944        let post = engine
945            .update(prior, &obs)
946            .expect("test: TD update should succeed");
947        match post.updated {
948            Prior::Gaussian { mean, variance } => {
949                // post_var = 1/(1+1) = 0.5, post_mean = 0.5*(0+2) = 1.0
950                assert!((variance - 0.5).abs() < 1e-10);
951                assert!((mean - 1.0).abs() < 1e-10);
952            }
953            _ => panic!("wrong variant"),
954        }
955    }
956
957    #[test]
958    fn gaussian_gaussian_large_n_pulls_to_sample() {
959        let mut engine = BayesianUpdateEngine::new(10);
960        let prior = Prior::Gaussian {
961            mean: 0.0,
962            variance: 100.0,
963        };
964        let obs = Observation::Gaussian {
965            sample_mean: 5.0,
966            sample_variance: 1.0,
967            n: 1000,
968        };
969        let post = engine
970            .update(prior, &obs)
971            .expect("test: TD update should succeed");
972        match post.updated {
973            Prior::Gaussian { mean, .. } => {
974                // With large n, posterior mean ≈ sample mean
975                assert!((mean - 5.0).abs() < 0.1);
976            }
977            _ => panic!("wrong variant"),
978        }
979    }
980
981    #[test]
982    fn gaussian_gaussian_log_marginal_finite() {
983        let mut engine = BayesianUpdateEngine::new(10);
984        let prior = Prior::Gaussian {
985            mean: 1.0,
986            variance: 2.0,
987        };
988        let obs = Observation::Gaussian {
989            sample_mean: 1.5,
990            sample_variance: 0.5,
991            n: 10,
992        };
993        let post = engine
994            .update(prior, &obs)
995            .expect("test: TD update should succeed");
996        assert!(post.log_marginal.is_finite());
997    }
998
999    #[test]
1000    fn gaussian_gaussian_zero_n_error() {
1001        let mut engine = BayesianUpdateEngine::new(10);
1002        let prior = Prior::Gaussian {
1003            mean: 0.0,
1004            variance: 1.0,
1005        };
1006        let obs = Observation::Gaussian {
1007            sample_mean: 1.0,
1008            sample_variance: 1.0,
1009            n: 0,
1010        };
1011        let result = engine.update(prior, &obs);
1012        assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1013    }
1014
1015    // ── Dirichlet–Categorical update ─────────────────────────────────────────
1016
1017    #[test]
1018    fn dirichlet_categorical_update_basic() {
1019        let mut engine = BayesianUpdateEngine::new(10);
1020        let prior = Prior::Dirichlet {
1021            alphas: vec![1.0, 1.0, 1.0],
1022        };
1023        let obs = Observation::Categorical {
1024            counts: vec![3, 2, 5],
1025        };
1026        let post = engine
1027            .update(prior, &obs)
1028            .expect("test: TD update should succeed");
1029        match post.updated {
1030            Prior::Dirichlet { alphas } => {
1031                assert!((alphas[0] - 4.0).abs() < 1e-10);
1032                assert!((alphas[1] - 3.0).abs() < 1e-10);
1033                assert!((alphas[2] - 6.0).abs() < 1e-10);
1034            }
1035            _ => panic!("wrong variant"),
1036        }
1037    }
1038
1039    #[test]
1040    fn dirichlet_categorical_dim_mismatch_error() {
1041        let mut engine = BayesianUpdateEngine::new(10);
1042        let prior = Prior::Dirichlet {
1043            alphas: vec![1.0, 1.0],
1044        };
1045        let obs = Observation::Categorical {
1046            counts: vec![1, 2, 3],
1047        };
1048        let result = engine.update(prior, &obs);
1049        assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1050    }
1051
1052    #[test]
1053    fn dirichlet_categorical_log_marginal_finite() {
1054        let mut engine = BayesianUpdateEngine::new(10);
1055        let prior = Prior::Dirichlet {
1056            alphas: vec![2.0, 3.0, 5.0],
1057        };
1058        let obs = Observation::Categorical {
1059            counts: vec![10, 15, 25],
1060        };
1061        let post = engine
1062            .update(prior, &obs)
1063            .expect("test: TD update should succeed");
1064        assert!(post.log_marginal.is_finite());
1065    }
1066
1067    #[test]
1068    fn dirichlet_categorical_zero_counts_no_change() {
1069        let mut engine = BayesianUpdateEngine::new(10);
1070        let prior = Prior::Dirichlet {
1071            alphas: vec![2.0, 3.0],
1072        };
1073        let obs = Observation::Categorical { counts: vec![0, 0] };
1074        let post = engine
1075            .update(prior, &obs)
1076            .expect("test: TD update should succeed");
1077        match post.updated {
1078            Prior::Dirichlet { alphas } => {
1079                assert!((alphas[0] - 2.0).abs() < 1e-10);
1080                assert!((alphas[1] - 3.0).abs() < 1e-10);
1081            }
1082            _ => panic!("wrong variant"),
1083        }
1084    }
1085
1086    // ── Gamma–Poisson update ─────────────────────────────────────────────────
1087
1088    #[test]
1089    fn gamma_poisson_update_basic() {
1090        let mut engine = BayesianUpdateEngine::new(10);
1091        let prior = Prior::Gamma {
1092            shape: 1.0,
1093            rate: 1.0,
1094        };
1095        let obs = Observation::Poisson {
1096            total_events: 5,
1097            total_time: 2.0,
1098        };
1099        let post = engine
1100            .update(prior, &obs)
1101            .expect("test: TD update should succeed");
1102        match post.updated {
1103            Prior::Gamma { shape, rate } => {
1104                assert!((shape - 6.0).abs() < 1e-10);
1105                assert!((rate - 3.0).abs() < 1e-10);
1106            }
1107            _ => panic!("wrong variant"),
1108        }
1109    }
1110
1111    #[test]
1112    fn gamma_poisson_log_marginal_finite() {
1113        let mut engine = BayesianUpdateEngine::new(10);
1114        let prior = Prior::Gamma {
1115            shape: 2.0,
1116            rate: 0.5,
1117        };
1118        let obs = Observation::Poisson {
1119            total_events: 10,
1120            total_time: 5.0,
1121        };
1122        let post = engine
1123            .update(prior, &obs)
1124            .expect("test: TD update should succeed");
1125        assert!(post.log_marginal.is_finite());
1126    }
1127
1128    #[test]
1129    fn gamma_poisson_zero_time_error() {
1130        let mut engine = BayesianUpdateEngine::new(10);
1131        let prior = Prior::Gamma {
1132            shape: 1.0,
1133            rate: 1.0,
1134        };
1135        let obs = Observation::Poisson {
1136            total_events: 5,
1137            total_time: 0.0,
1138        };
1139        let result = engine.update(prior, &obs);
1140        assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1141    }
1142
1143    // ── Mismatch errors ──────────────────────────────────────────────────────
1144
1145    #[test]
1146    fn mismatch_beta_gaussian_obs() {
1147        let mut engine = BayesianUpdateEngine::new(10);
1148        let prior = Prior::Beta {
1149            alpha: 1.0,
1150            beta: 1.0,
1151        };
1152        let obs = Observation::Gaussian {
1153            sample_mean: 0.5,
1154            sample_variance: 1.0,
1155            n: 10,
1156        };
1157        let result = engine.update(prior, &obs);
1158        assert!(matches!(
1159            result,
1160            Err(BayesError::PriorObservationMismatch { .. })
1161        ));
1162    }
1163
1164    #[test]
1165    fn mismatch_gaussian_bernoulli_obs() {
1166        let mut engine = BayesianUpdateEngine::new(10);
1167        let prior = Prior::Gaussian {
1168            mean: 0.0,
1169            variance: 1.0,
1170        };
1171        let obs = Observation::Bernoulli {
1172            successes: 3,
1173            trials: 5,
1174        };
1175        let result = engine.update(prior, &obs);
1176        assert!(matches!(
1177            result,
1178            Err(BayesError::PriorObservationMismatch { .. })
1179        ));
1180    }
1181
1182    // ── Sequential update ────────────────────────────────────────────────────
1183
1184    #[test]
1185    fn sequential_update_equivalent_to_batch() {
1186        // For Beta-Bernoulli, two sequential updates should equal one combined update
1187        let mut engine = BayesianUpdateEngine::new(64);
1188        let prior = Prior::Beta {
1189            alpha: 1.0,
1190            beta: 1.0,
1191        };
1192        let obs = vec![
1193            Observation::Bernoulli {
1194                successes: 3,
1195                trials: 5,
1196            },
1197            Observation::Bernoulli {
1198                successes: 2,
1199                trials: 4,
1200            },
1201        ];
1202        let seq_post = engine
1203            .sequential_update(prior.clone(), &obs)
1204            .expect("test: should succeed");
1205
1206        // Batch equivalent: alpha' = 1 + 5, beta' = 1 + 4
1207        let mut engine2 = BayesianUpdateEngine::new(64);
1208        let batch_obs = Observation::Bernoulli {
1209            successes: 5,
1210            trials: 9,
1211        };
1212        let batch_post = engine2
1213            .update(prior, &batch_obs)
1214            .expect("test: TD update should succeed");
1215
1216        match (&seq_post.updated, &batch_post.updated) {
1217            (
1218                Prior::Beta {
1219                    alpha: a1,
1220                    beta: b1,
1221                },
1222                Prior::Beta {
1223                    alpha: a2,
1224                    beta: b2,
1225                },
1226            ) => {
1227                assert!((a1 - a2).abs() < 1e-10);
1228                assert!((b1 - b2).abs() < 1e-10);
1229            }
1230            _ => panic!("wrong variant"),
1231        }
1232    }
1233
1234    #[test]
1235    fn sequential_update_empty_error() {
1236        let mut engine = BayesianUpdateEngine::new(10);
1237        let prior = Prior::Beta {
1238            alpha: 1.0,
1239            beta: 1.0,
1240        };
1241        let result = engine.sequential_update(prior, &[]);
1242        assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1243    }
1244
1245    // ── Credible interval ────────────────────────────────────────────────────
1246
1247    #[test]
1248    fn credible_interval_beta_bounds() {
1249        let post = Prior::Beta {
1250            alpha: 8.0,
1251            beta: 4.0,
1252        };
1253        let ci =
1254            BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1255        assert!(ci.lower >= 0.0);
1256        assert!(ci.upper <= 1.0);
1257        assert!(ci.lower < ci.upper);
1258        assert!((ci.probability - 0.95).abs() < 1e-10);
1259    }
1260
1261    #[test]
1262    fn credible_interval_gaussian_symmetric() {
1263        let post = Prior::Gaussian {
1264            mean: 5.0,
1265            variance: 1.0,
1266        };
1267        let ci =
1268            BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1269        let center = (ci.lower + ci.upper) / 2.0;
1270        assert!((center - 5.0).abs() < 1e-8);
1271        // Half-width ≈ 1.96
1272        let hw = (ci.upper - ci.lower) / 2.0;
1273        assert!((hw - 1.96).abs() < 0.01);
1274    }
1275
1276    #[test]
1277    fn credible_interval_invalid_probability() {
1278        let post = Prior::Gaussian {
1279            mean: 0.0,
1280            variance: 1.0,
1281        };
1282        assert!(BayesianUpdateEngine::credible_interval(&post, 0.0).is_err());
1283        assert!(BayesianUpdateEngine::credible_interval(&post, 1.0).is_err());
1284        assert!(BayesianUpdateEngine::credible_interval(&post, -0.1).is_err());
1285    }
1286
1287    #[test]
1288    fn credible_interval_gamma() {
1289        let post = Prior::Gamma {
1290            shape: 9.0,
1291            rate: 3.0,
1292        };
1293        let ci =
1294            BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1295        // mean = 3.0; lower should be positive
1296        assert!(ci.lower >= 0.0);
1297        assert!(ci.upper > ci.lower);
1298    }
1299
1300    #[test]
1301    fn credible_interval_dirichlet() {
1302        let post = Prior::Dirichlet {
1303            alphas: vec![10.0, 2.0, 3.0],
1304        };
1305        let ci =
1306            BayesianUpdateEngine::credible_interval(&post, 0.90).expect("test: should succeed");
1307        assert!(ci.lower >= 0.0);
1308        assert!(ci.upper <= 1.0);
1309    }
1310
1311    // ── MAP estimate ─────────────────────────────────────────────────────────
1312
1313    #[test]
1314    fn map_beta_mode() {
1315        // Beta(3,3): mode = (3-1)/(3+3-2) = 2/4 = 0.5
1316        let p = Prior::Beta {
1317            alpha: 3.0,
1318            beta: 3.0,
1319        };
1320        let map = BayesianUpdateEngine::map_estimate(&p);
1321        assert!((map - 0.5).abs() < 1e-10);
1322    }
1323
1324    #[test]
1325    fn map_beta_uniform_fallback() {
1326        // Beta(1,1) → alpha <= 1, use alpha/(alpha+beta) = 0.5
1327        let p = Prior::Beta {
1328            alpha: 1.0,
1329            beta: 1.0,
1330        };
1331        let map = BayesianUpdateEngine::map_estimate(&p);
1332        assert!((map - 0.5).abs() < 1e-10);
1333    }
1334
1335    #[test]
1336    fn map_gaussian_is_mean() {
1337        let p = Prior::Gaussian {
1338            mean: 3.7,
1339            variance: 2.0,
1340        };
1341        assert!((BayesianUpdateEngine::map_estimate(&p) - 3.7).abs() < 1e-10);
1342    }
1343
1344    #[test]
1345    fn map_gamma_mode() {
1346        // Gamma(5, 2): mode = (5-1)/2 = 2.0
1347        let p = Prior::Gamma {
1348            shape: 5.0,
1349            rate: 2.0,
1350        };
1351        assert!((BayesianUpdateEngine::map_estimate(&p) - 2.0).abs() < 1e-10);
1352    }
1353
1354    #[test]
1355    fn map_gamma_shape_one_gives_zero() {
1356        let p = Prior::Gamma {
1357            shape: 1.0,
1358            rate: 2.0,
1359        };
1360        assert!((BayesianUpdateEngine::map_estimate(&p) - 0.0).abs() < 1e-10);
1361    }
1362
1363    #[test]
1364    fn map_dirichlet_argmax_proportion() {
1365        let p = Prior::Dirichlet {
1366            alphas: vec![1.0, 5.0, 2.0],
1367        };
1368        // max alpha = 5.0, sum = 8.0 → map = 5/8
1369        let expected = 5.0 / 8.0;
1370        assert!((BayesianUpdateEngine::map_estimate(&p) - expected).abs() < 1e-10);
1371    }
1372
1373    // ── KL divergence ────────────────────────────────────────────────────────
1374
1375    #[test]
1376    fn kl_beta_self_is_zero() {
1377        let p = Prior::Beta {
1378            alpha: 3.0,
1379            beta: 5.0,
1380        };
1381        let kl = BayesianUpdateEngine::kl_divergence(&p, &p).expect("test: should succeed");
1382        assert!(kl.abs() < 1e-8);
1383    }
1384
1385    #[test]
1386    fn kl_gaussian_self_is_zero() {
1387        let p = Prior::Gaussian {
1388            mean: 2.0,
1389            variance: 3.0,
1390        };
1391        let kl = BayesianUpdateEngine::kl_divergence(&p, &p).expect("test: should succeed");
1392        assert!(kl.abs() < 1e-10);
1393    }
1394
1395    #[test]
1396    fn kl_gaussian_asymmetry() {
1397        let p = Prior::Gaussian {
1398            mean: 0.0,
1399            variance: 1.0,
1400        };
1401        let q = Prior::Gaussian {
1402            mean: 1.0,
1403            variance: 2.0,
1404        };
1405        let kl_pq = BayesianUpdateEngine::kl_divergence(&p, &q).expect("test: should succeed");
1406        let kl_qp = BayesianUpdateEngine::kl_divergence(&q, &p).expect("test: should succeed");
1407        // KL is asymmetric in general
1408        assert!((kl_pq - kl_qp).abs() > 1e-6);
1409        // Both should be non-negative
1410        assert!(kl_pq >= 0.0);
1411        assert!(kl_qp >= 0.0);
1412    }
1413
1414    #[test]
1415    fn kl_known_gaussian_value() {
1416        // KL(N(0,1) ‖ N(1,1)) = 0.5 * (0 + 1 + 1 - 1) = 0.5
1417        let p = Prior::Gaussian {
1418            mean: 0.0,
1419            variance: 1.0,
1420        };
1421        let q = Prior::Gaussian {
1422            mean: 1.0,
1423            variance: 1.0,
1424        };
1425        let kl = BayesianUpdateEngine::kl_divergence(&p, &q).expect("test: should succeed");
1426        assert!((kl - 0.5).abs() < 1e-10);
1427    }
1428
1429    #[test]
1430    fn kl_unsupported_pair_error() {
1431        let p = Prior::Beta {
1432            alpha: 1.0,
1433            beta: 1.0,
1434        };
1435        let q = Prior::Gamma {
1436            shape: 1.0,
1437            rate: 1.0,
1438        };
1439        let result = BayesianUpdateEngine::kl_divergence(&p, &q);
1440        assert!(matches!(result, Err(BayesError::UnsupportedOperation(_))));
1441    }
1442
1443    // ── History management ───────────────────────────────────────────────────
1444
1445    #[test]
1446    fn history_bounded_by_max() {
1447        let mut engine = BayesianUpdateEngine::new(3);
1448        for i in 0..5_u64 {
1449            let prior = Prior::Beta {
1450                alpha: 1.0,
1451                beta: 1.0,
1452            };
1453            let obs = Observation::Bernoulli {
1454                successes: i % 3,
1455                trials: 5,
1456            };
1457            engine
1458                .update(prior, &obs)
1459                .expect("test: TD update should succeed");
1460        }
1461        assert_eq!(engine.history().len(), 3);
1462    }
1463
1464    #[test]
1465    fn history_clear() {
1466        let mut engine = BayesianUpdateEngine::new(10);
1467        let prior = Prior::Beta {
1468            alpha: 1.0,
1469            beta: 1.0,
1470        };
1471        let obs = Observation::Bernoulli {
1472            successes: 3,
1473            trials: 5,
1474        };
1475        engine
1476            .update(prior, &obs)
1477            .expect("test: TD update should succeed");
1478        assert!(!engine.history().is_empty());
1479        engine.clear_history();
1480        assert!(engine.history().is_empty());
1481    }
1482
1483    #[test]
1484    fn history_zero_capacity_no_store() {
1485        let mut engine = BayesianUpdateEngine::new(0);
1486        let prior = Prior::Beta {
1487            alpha: 1.0,
1488            beta: 1.0,
1489        };
1490        let obs = Observation::Bernoulli {
1491            successes: 3,
1492            trials: 5,
1493        };
1494        engine
1495            .update(prior, &obs)
1496            .expect("test: TD update should succeed");
1497        assert!(engine.history().is_empty());
1498    }
1499
1500    // ── Prior validation ─────────────────────────────────────────────────────
1501
1502    #[test]
1503    fn invalid_beta_alpha_zero() {
1504        let mut engine = BayesianUpdateEngine::new(10);
1505        let prior = Prior::Beta {
1506            alpha: 0.0,
1507            beta: 1.0,
1508        };
1509        let obs = Observation::Bernoulli {
1510            successes: 1,
1511            trials: 2,
1512        };
1513        assert!(matches!(
1514            engine.update(prior, &obs),
1515            Err(BayesError::InvalidParameters(_))
1516        ));
1517    }
1518
1519    #[test]
1520    fn invalid_gamma_rate_negative() {
1521        let mut engine = BayesianUpdateEngine::new(10);
1522        let prior = Prior::Gamma {
1523            shape: 1.0,
1524            rate: -1.0,
1525        };
1526        let obs = Observation::Poisson {
1527            total_events: 5,
1528            total_time: 1.0,
1529        };
1530        assert!(matches!(
1531            engine.update(prior, &obs),
1532            Err(BayesError::InvalidParameters(_))
1533        ));
1534    }
1535
1536    // ── likelihood_type field ────────────────────────────────────────────────
1537
1538    #[test]
1539    fn likelihood_type_labels() {
1540        let mut engine = BayesianUpdateEngine::new(10);
1541
1542        let p1 = engine
1543            .update(
1544                Prior::Beta {
1545                    alpha: 1.0,
1546                    beta: 1.0,
1547                },
1548                &Observation::Bernoulli {
1549                    successes: 1,
1550                    trials: 2,
1551                },
1552            )
1553            .expect("test: should succeed");
1554        assert_eq!(p1.likelihood_type, "Bernoulli");
1555
1556        let p2 = engine
1557            .update(
1558                Prior::Gaussian {
1559                    mean: 0.0,
1560                    variance: 1.0,
1561                },
1562                &Observation::Gaussian {
1563                    sample_mean: 1.0,
1564                    sample_variance: 1.0,
1565                    n: 5,
1566                },
1567            )
1568            .expect("test: should succeed");
1569        assert_eq!(p2.likelihood_type, "Gaussian");
1570
1571        let p3 = engine
1572            .update(
1573                Prior::Dirichlet {
1574                    alphas: vec![1.0, 1.0],
1575                },
1576                &Observation::Categorical { counts: vec![3, 2] },
1577            )
1578            .expect("test: should succeed");
1579        assert_eq!(p3.likelihood_type, "Categorical");
1580
1581        let p4 = engine
1582            .update(
1583                Prior::Gamma {
1584                    shape: 1.0,
1585                    rate: 1.0,
1586                },
1587                &Observation::Poisson {
1588                    total_events: 3,
1589                    total_time: 1.0,
1590                },
1591            )
1592            .expect("test: should succeed");
1593        assert_eq!(p4.likelihood_type, "Poisson");
1594    }
1595}