sanos 0.2.1

// src/density/marginal.rs
use crate::error::{SanosError, SanosResult};
use super::DensityTolerances;

/// Discrete marginal distribution at a given maturity:
///     mu = sum_i q_i * delta_{k_i}
///
/// Invariants (validated):
/// - maturity > 0
/// - k_i > 0, strictly increasing
/// - q_i >= 0
/// - sum_i q_i == 1 (within mass tol)
/// - sum_i q_i * k_i == 1 (within mean tol)
#[derive(Debug, Clone)]
pub struct MarginalDensity {
    maturity: f64,
    atoms: Vec<(f64, f64)>, // (k, q)
}

impl MarginalDensity {
    pub fn new(maturity: f64, mut atoms: Vec<(f64, f64)>, tol: DensityTolerances) -> SanosResult<Self> {
        if !maturity.is_finite() {
            return Err(SanosError::NonFinite { field: "maturity", value: maturity });
        }
        if maturity <= 0.0 {
            return Err(SanosError::InvalidBound {
                field: "maturity",
                value: maturity,
                min: f64::MIN_POSITIVE,
                max: f64::INFINITY,
            });
        }
        if atoms.is_empty() {
            return Err(SanosError::EmptyCollection { what: "MarginalDensity.atoms" });
        }

        // Sort by strike
        atoms.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());

        let md = Self { maturity, atoms };
        md.validate(tol)?;
        Ok(md)
    }

    #[inline]
    pub fn maturity(&self) -> f64 {
        self.maturity
    }

    #[inline]
    pub fn atoms(&self) -> &[(f64, f64)] {
        &self.atoms
    }

    #[inline]
    pub fn strikes(&self) -> Vec<f64> {
        self.atoms.iter().map(|&(k, _)| k).collect()
    }

    #[inline]
    pub fn probabilities(&self) -> Vec<f64> {
        self.atoms.iter().map(|&(_, q)| q).collect()
    }

    /// Call-transform:
    ///     C(kappa) = E[(K - kappa)^+] = sum_i q_i * (k_i - kappa)^+
    pub fn call(&self, kappa: f64) -> SanosResult<f64> {
        if !kappa.is_finite() {
            return Err(SanosError::NonFinite { field: "kappa", value: kappa });
        }
        let mut acc = 0.0_f64;
        for &(k, q) in &self.atoms {
            let payoff = (k - kappa).max(0.0);
            acc += q * payoff;
        }
        Ok(acc)
    }

    pub fn validate(&self, tol: DensityTolerances) -> SanosResult<()> {
        // Validate atoms are finite and strictly increasing in k
        for (i, &(k, q)) in self.atoms.iter().enumerate() {
            if !k.is_finite() {
                return Err(SanosError::NonFinite { field: "k", value: k });
            }
            if !q.is_finite() {
                return Err(SanosError::NonFinite { field: "q", value: q });
            }
            if k <= 0.0 {
                return Err(SanosError::InvalidBound {
                    field: "k",
                    value: k,
                    min: f64::MIN_POSITIVE,
                    max: f64::INFINITY,
                });
            }
            if q < -tol.mass {
                return Err(SanosError::InvalidBound {
                    field: "q",
                    value: q,
                    min: 0.0,
                    max: f64::INFINITY,
                });
            }
            if i > 0 {
                let k_prev = self.atoms[i - 1].0;
                if k <= k_prev {
                    if (k - k_prev).abs() == 0.0 {
                        return Err(SanosError::DuplicateKey { what: "atom strike", value: k });
                    }
                    return Err(SanosError::InvalidOrdering { msg: "atom strikes must be strictly increasing" });
                }
            }
        }

        // Validate mass and mean
        let mut mass = 0.0_f64;
        let mut mean = 0.0_f64;
        for &(k, q) in &self.atoms {
            mass += q;
            mean += q * k;
        }

        if (mass - 1.0).abs() > tol.mass {
            return Err(SanosError::InvalidOrdering {
                msg: "marginal mass constraint violated (sum q != 1 within tolerance)",
            });
        }
        if (mean - 1.0).abs() > tol.mean {
            return Err(SanosError::InvalidOrdering {
                msg: "marginal mean constraint violated (sum q*k != 1 within tolerance)",
            });
        }

        Ok(())
    }
}

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

    #[test]
    fn marginal_density_validates_mass_and_mean() {
        let tol = DensityTolerances::from_tol(1e-12).unwrap();
        let m = MarginalDensity::new(1.0, vec![(0.8, 0.5), (1.2, 0.5)], tol).unwrap();
        let c_at_1 = m.call(1.0).unwrap();
        assert!(c_at_1 >= 0.0);
    }

    #[test]
    fn marginal_density_rejects_bad_mass() {
        let tol = DensityTolerances::from_tol(1e-12).unwrap();
        let err = MarginalDensity::new(1.0, vec![(1.0, 0.9)], tol).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("marginal mass constraint violated"));
    }

    #[test]
    fn marginal_density_rejects_bad_mean() {
        let tol = DensityTolerances::from_tol(1e-12).unwrap();
        let err = MarginalDensity::new(1.0, vec![(1.1, 1.0)], tol).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("marginal mean constraint violated"));
    }

    #[test]
    fn marginal_density_allows_tiny_negative_with_tolerance() {
        let tol = DensityTolerances::from_tol(1e-8).unwrap();
        let e = 1e-12;
        let m = MarginalDensity::new(1.0, vec![(0.5, -e), (1.0, 1.0 + 2.0 * e), (1.5, -e)], tol)
            .unwrap();
        assert_eq!(m.atoms().len(), 3);
    }
}