antecedent-prob 0.2.0

//! Prior specifications.
//!
//! Priors are recorded as assumptions; they do not create nonparametric
//! identification.
//!
//! SPDX-License-Identifier: MIT OR Apache-2.0

use std::sync::Arc;

use antecedent_core::{PriorAssumption, VariableId};

use crate::error::ProbError;

/// Floor applied when converting effect-draw SD into a prior scale.
const EFFECT_PRIOR_SD_FLOOR: f64 = 1e-12;

/// Contrast coding for categorical predictors (required for Bayesian GLMs).
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum ContrastCoding {
    /// Treatment (dummy) coding with a designated reference level.
    Treatment,
    /// Sum (deviation) coding.
    Sum,
}

/// Gaussian prior on a scalar effect functional (e.g. ATE).
///
/// Used for cross-design transfer: moments come from source effect draws or
/// stored summaries, then map onto a target coefficient (identity-link bridge).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct EffectPrior {
    /// Prior mean of the effect functional.
    pub mean: f64,
    /// Prior SD of the effect functional (must be finite and > 0).
    pub sd: f64,
}

impl EffectPrior {
    /// Construct from mean / SD with validation.
    ///
    /// # Errors
    ///
    /// Non-finite mean or non-positive / non-finite SD.
    pub fn new(mean: f64, sd: f64) -> Result<Self, ProbError> {
        let p = Self { mean, sd };
        p.validate()?;
        Ok(p)
    }

    /// Sample mean / SD from effect draws (population SD with Bessel correction when `n > 1`).
    ///
    /// SD is floored at a tiny positive value so conjugate scale stays valid.
    ///
    /// # Errors
    ///
    /// Empty draws or non-finite moments.
    pub fn from_effect_draws(draws: &[f64]) -> Result<Self, ProbError> {
        if draws.is_empty() {
            return Err(ProbError::InvalidPrior { message: "from_effect_draws: empty draws" });
        }
        let n = draws.len() as f64;
        let mean = draws.iter().sum::<f64>() / n;
        if !mean.is_finite() {
            return Err(ProbError::InvalidPrior { message: "from_effect_draws: non-finite mean" });
        }
        let sd = if draws.len() == 1 {
            EFFECT_PRIOR_SD_FLOOR
        } else {
            let var = draws
                .iter()
                .map(|&x| {
                    let d = x - mean;
                    d * d
                })
                .sum::<f64>()
                / (n - 1.0);
            var.sqrt().max(EFFECT_PRIOR_SD_FLOOR)
        };
        if !sd.is_finite() {
            return Err(ProbError::InvalidPrior { message: "from_effect_draws: non-finite sd" });
        }
        Self::new(mean, sd)
    }

    /// Validate finite mean and positive finite SD.
    ///
    /// # Errors
    ///
    /// Invalid parameters.
    pub fn validate(self) -> Result<(), ProbError> {
        if !self.mean.is_finite() {
            return Err(ProbError::InvalidPrior { message: "effect prior mean must be finite" });
        }
        if !(self.sd > 0.0) || !self.sd.is_finite() {
            return Err(ProbError::InvalidPrior {
                message: "effect prior sd must be finite and > 0",
            });
        }
        Ok(())
    }
}

/// Gaussian coefficient prior for conjugate / NIG linear models.
///
/// Under the conjugate Normal–Inv-Gamma (and known-σ² Normal) backends,
/// `variance[i]` is the diagonal entry of the *scale* matrix `V0` in
/// `β | σ² ~ N(mean, σ² · diag(V0))` — not an absolute prior variance of `β`.
/// Absolute prior variance of coefficient `i` is therefore `σ² · variance[i]`.
#[derive(Clone, Debug, PartialEq)]
pub struct GaussianCoefficientPrior {
    /// Prior mean per coefficient (length = p), or a single shared mean.
    pub mean: Arc<[f64]>,
    /// Diagonal of conjugate scale `V0` (length = p); see struct docs.
    pub variance: Arc<[f64]>,
}

impl GaussianCoefficientPrior {
    /// Isotropic weakly informative prior: mean 0, V0 diagonal `scale²`
    /// (absolute prior variance of β is `σ² · scale²` under conjugate models).
    #[must_use]
    pub fn isotropic(n_coef: usize, scale: f64) -> Self {
        let var = scale * scale;
        Self { mean: Arc::from(vec![0.0; n_coef]), variance: Arc::from(vec![var; n_coef]) }
    }

    /// Shared mean / V0-diagonal broadcast to `n_coef` coefficients.
    ///
    /// # Errors
    ///
    /// Non-positive variance or zero coefficients.
    pub fn shared(n_coef: usize, mean: f64, variance: f64) -> Result<Self, ProbError> {
        if n_coef == 0 {
            return Err(ProbError::InvalidPrior { message: "n_coef must be > 0" });
        }
        if !(variance > 0.0) {
            return Err(ProbError::InvalidPrior { message: "variance must be > 0" });
        }
        Ok(Self {
            mean: Arc::from(vec![mean; n_coef]),
            variance: Arc::from(vec![variance; n_coef]),
        })
    }

    /// Number of coefficients.
    #[must_use]
    pub fn len(&self) -> usize {
        self.mean.len()
    }

    /// Whether empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.mean.is_empty()
    }

    /// Precision (1/variance) vector.
    #[must_use]
    pub fn precision(&self) -> Vec<f64> {
        self.variance.iter().map(|&v| 1.0 / v).collect()
    }

    /// Validate lengths match.
    ///
    /// # Errors
    ///
    /// Length mismatch or non-positive variance.
    pub fn validate(&self) -> Result<(), ProbError> {
        if self.mean.len() != self.variance.len() {
            return Err(ProbError::InvalidPrior { message: "mean and variance length mismatch" });
        }
        if self.mean.is_empty() {
            return Err(ProbError::InvalidPrior { message: "empty coefficient prior" });
        }
        for &v in self.variance.iter() {
            if !(v > 0.0) && v.is_finite() {
                return Err(ProbError::InvalidPrior { message: "variance must be > 0" });
            }
            if !v.is_finite() {
                return Err(ProbError::InvalidPrior { message: "variance must be finite" });
            }
        }
        Ok(())
    }
}

/// Inv-Gamma prior on residual variance (conjugate Gaussian linear).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct InvGammaPrior {
    /// Shape α > 0.
    pub shape: f64,
    /// Scale β > 0 (mean = β/(α−1) for α > 1).
    pub scale: f64,
}

impl InvGammaPrior {
    /// Weakly informative default.
    #[must_use]
    pub const fn weakly_informative() -> Self {
        Self { shape: 1e-3, scale: 1e-3 }
    }

    /// Validate.
    ///
    /// # Errors
    ///
    /// Non-positive shape or scale.
    pub fn validate(self) -> Result<(), ProbError> {
        if !(self.shape > 0.0) || !(self.scale > 0.0) {
            return Err(ProbError::InvalidPrior {
                message: "InvGamma shape and scale must be > 0",
            });
        }
        Ok(())
    }
}

/// A named prior specification entry.
#[derive(Clone, Debug, PartialEq)]
pub enum PriorSpec {
    /// Gaussian coefficient prior for a linear / GLM mechanism.
    GaussianCoefficients(GaussianCoefficientPrior),
    /// Residual variance prior (conjugate Gaussian).
    ResidualInvGamma(InvGammaPrior),
    /// Fixed residual variance (known σ²).
    KnownResidualVariance(f64),
}

impl PriorSpec {
    /// Convert to a [`PriorAssumption`] for the assumption record.
    #[must_use]
    pub fn as_assumption(&self) -> PriorAssumption {
        match self {
            Self::GaussianCoefficients(_) => PriorAssumption {
                id: Arc::from("gaussian_coefficients"),
                description: Arc::from("Gaussian prior on regression coefficients"),
            },
            Self::ResidualInvGamma(_) => PriorAssumption {
                id: Arc::from("residual_inv_gamma"),
                description: Arc::from("Inverse-Gamma prior on residual variance"),
            },
            Self::KnownResidualVariance(_) => PriorAssumption {
                id: Arc::from("known_residual_variance"),
                description: Arc::from("Known residual variance (no prior uncertainty)"),
            },
        }
    }

    /// Validate this prior.
    ///
    /// # Errors
    ///
    /// Invalid parameters.
    pub fn validate(&self) -> Result<(), ProbError> {
        match self {
            Self::GaussianCoefficients(p) => p.validate(),
            Self::ResidualInvGamma(p) => p.validate(),
            Self::KnownResidualVariance(v) => {
                if !(*v > 0.0) || !v.is_finite() {
                    return Err(ProbError::InvalidPrior {
                        message: "known residual variance must be finite and > 0",
                    });
                }
                Ok(())
            }
        }
    }
}

/// Residual-variance model for Gaussian linear targets (HMC / Laplace).
///
/// Resolved once from [`PriorSet`] before sampling or optimization. At most one
/// residual specification may appear in the prior set.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum GaussianVarianceModel {
    /// Fixed known residual variance for the entire run.
    Known {
        /// Residual variance σ² > 0.
        sigma2: f64,
    },
    /// Inverse-gamma prior on σ²; HMC state includes `λ = log(σ²)`.
    InvGamma {
        /// Shape α₀ > 0.
        shape: f64,
        /// Scale β₀ > 0.
        scale: f64,
    },
}

impl GaussianVarianceModel {
    /// Resolve the residual model from a validated prior set.
    ///
    /// `KnownResidualVariance` → [`Self::Known`]; `ResidualInvGamma` →
    /// [`Self::InvGamma`]; an omitted residual specification defaults to
    /// [`InvGammaPrior::weakly_informative`].
    ///
    /// # Errors
    ///
    /// More than one residual specification, or invalid known / InvGamma params.
    pub fn from_prior_set(prior: &PriorSet) -> Result<Self, ProbError> {
        let mut known: Option<f64> = None;
        let mut inv_gamma: Option<InvGammaPrior> = None;
        for spec in &prior.specs {
            match spec {
                PriorSpec::KnownResidualVariance(v) => {
                    if known.is_some() || inv_gamma.is_some() {
                        return Err(ProbError::InvalidPrior {
                            message: "PriorSet must contain at most one residual variance specification",
                        });
                    }
                    known = Some(*v);
                }
                PriorSpec::ResidualInvGamma(p) => {
                    if known.is_some() || inv_gamma.is_some() {
                        return Err(ProbError::InvalidPrior {
                            message: "PriorSet must contain at most one residual variance specification",
                        });
                    }
                    inv_gamma = Some(*p);
                }
                PriorSpec::GaussianCoefficients(_) => {}
            }
        }
        if let Some(sigma2) = known {
            if !(sigma2 > 0.0) || !sigma2.is_finite() {
                return Err(ProbError::InvalidPrior {
                    message: "known residual variance must be finite and > 0",
                });
            }
            return Ok(Self::Known { sigma2 });
        }
        let ig = inv_gamma.unwrap_or_else(InvGammaPrior::weakly_informative);
        ig.validate()?;
        Ok(Self::InvGamma { shape: ig.shape, scale: ig.scale })
    }

    /// Unconstrained state dimension for coefficients of length `ncols`.
    #[must_use]
    pub const fn state_dim(self, ncols: usize) -> usize {
        match self {
            Self::Known { .. } => ncols,
            Self::InvGamma { .. } => ncols.saturating_add(1),
        }
    }

    /// Whether draws include a residual-variance column.
    #[must_use]
    pub const fn include_sigma2(self) -> bool {
        matches!(self, Self::InvGamma { .. })
    }
}

/// Collection of priors for an inference run.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct PriorSet {
    /// Ordered prior entries.
    pub specs: Vec<PriorSpec>,
    /// Explicit contrast coding when categorical predictors are present.
    pub contrast: Option<ContrastCoding>,
    /// Variables that are categorical and require the declared contrast.
    pub categorical: Vec<VariableId>,
    /// Extra prior-restriction assumptions (e.g. external bank mapping ids).
    pub restrictions: Vec<PriorAssumption>,
}

impl PriorSet {
    /// Empty set.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Weakly informative Gaussian coefficient prior of width `scale` for `n_coef`.
    #[must_use]
    pub fn weakly_informative(n_coef: usize) -> Self {
        Self {
            specs: vec![
                PriorSpec::GaussianCoefficients(GaussianCoefficientPrior::isotropic(n_coef, 10.0)),
                PriorSpec::ResidualInvGamma(InvGammaPrior::weakly_informative()),
            ],
            contrast: None,
            categorical: Vec::new(),
            restrictions: Vec::new(),
        }
    }

    /// Push a prior spec.
    pub fn push(&mut self, spec: PriorSpec) {
        self.specs.push(spec);
    }

    /// Require an explicit contrast when categoricals are present.
    ///
    /// # Errors
    ///
    /// Categoricals listed without a contrast coding.
    pub fn validate_contrasts(&self) -> Result<(), ProbError> {
        if !self.categorical.is_empty() && self.contrast.is_none() {
            return Err(ProbError::InvalidPrior {
                message: "categorical predictors require explicit contrast coding",
            });
        }
        Ok(())
    }

    /// Validate all specs.
    ///
    /// # Errors
    ///
    /// Invalid specs, missing contrast, or more than one residual variance spec.
    pub fn validate(&self) -> Result<(), ProbError> {
        for s in &self.specs {
            s.validate()?;
        }
        self.validate_contrasts()?;
        let mut n_residual = 0usize;
        for s in &self.specs {
            match s {
                PriorSpec::ResidualInvGamma(_) | PriorSpec::KnownResidualVariance(_) => {
                    n_residual = n_residual.saturating_add(1);
                }
                PriorSpec::GaussianCoefficients(_) => {}
            }
        }
        if n_residual > 1 {
            return Err(ProbError::InvalidPrior {
                message: "PriorSet must contain at most one residual variance specification",
            });
        }
        Ok(())
    }

    /// First Gaussian coefficient prior, if any.
    #[must_use]
    pub fn gaussian_coefficients(&self) -> Option<&GaussianCoefficientPrior> {
        self.specs.iter().find_map(|s| match s {
            PriorSpec::GaussianCoefficients(p) => Some(p),
            _ => None,
        })
    }

    /// Residual Inv-Gamma prior, if any.
    #[must_use]
    pub fn residual_inv_gamma(&self) -> Option<InvGammaPrior> {
        self.specs.iter().find_map(|s| match s {
            PriorSpec::ResidualInvGamma(p) => Some(*p),
            _ => None,
        })
    }

    /// Known residual variance, if any.
    #[must_use]
    pub fn known_residual_variance(&self) -> Option<f64> {
        self.specs.iter().find_map(|s| match s {
            PriorSpec::KnownResidualVariance(v) => Some(*v),
            _ => None,
        })
    }
}

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

    #[test]
    fn weakly_informative_validates() {
        let p = PriorSet::weakly_informative(3);
        p.validate().unwrap();
        assert_eq!(p.gaussian_coefficients().unwrap().len(), 3);
    }

    #[test]
    fn categorical_requires_contrast() {
        let mut p = PriorSet::weakly_informative(2);
        p.categorical.push(VariableId::from_raw(0));
        assert!(p.validate().is_err());
        p.contrast = Some(ContrastCoding::Treatment);
        p.validate().unwrap();
    }

    #[test]
    fn effect_prior_from_draws_moments() {
        let draws = [1.0, 3.0, 5.0];
        let p = EffectPrior::from_effect_draws(&draws).unwrap();
        assert!((p.mean - 3.0).abs() < 1e-12);
        // sample sd of [1,3,5] = 2
        assert!((p.sd - 2.0).abs() < 1e-12);
    }

    #[test]
    fn effect_prior_rejects_empty_and_nonfinite() {
        assert!(EffectPrior::from_effect_draws(&[]).is_err());
        assert!(EffectPrior::new(f64::NAN, 1.0).is_err());
        assert!(EffectPrior::new(0.0, 0.0).is_err());
        assert!(EffectPrior::new(0.0, -1.0).is_err());
    }

    #[test]
    fn effect_prior_single_draw_floors_sd() {
        let p = EffectPrior::from_effect_draws(&[2.5]).unwrap();
        assert!((p.mean - 2.5).abs() < 1e-12);
        assert!(p.sd > 0.0);
    }

    #[test]
    fn residual_specs_must_be_unique() {
        let mut p = PriorSet::new();
        p.push(PriorSpec::GaussianCoefficients(GaussianCoefficientPrior::isotropic(1, 1.0)));
        p.push(PriorSpec::KnownResidualVariance(1.0));
        p.push(PriorSpec::ResidualInvGamma(InvGammaPrior::weakly_informative()));
        assert!(p.validate().is_err());
        assert!(GaussianVarianceModel::from_prior_set(&p).is_err());
    }

    #[test]
    fn variance_model_defaults_to_weak_inv_gamma() {
        let mut p = PriorSet::new();
        p.push(PriorSpec::GaussianCoefficients(GaussianCoefficientPrior::isotropic(2, 1.0)));
        p.validate().unwrap();
        let model = GaussianVarianceModel::from_prior_set(&p).unwrap();
        let weak = InvGammaPrior::weakly_informative();
        assert_eq!(model, GaussianVarianceModel::InvGamma { shape: weak.shape, scale: weak.scale });
        assert_eq!(model.state_dim(2), 3);
        assert!(model.include_sigma2());
    }

    #[test]
    fn variance_model_known() {
        let mut p = PriorSet::new();
        p.push(PriorSpec::KnownResidualVariance(2.5));
        let model = GaussianVarianceModel::from_prior_set(&p).unwrap();
        assert_eq!(model, GaussianVarianceModel::Known { sigma2: 2.5 });
        assert_eq!(model.state_dim(4), 4);
        assert!(!model.include_sigma2());
    }
}