legume-numeric 0.8.11

use candle_core::{DType, Result, Tensor};

use super::regression_linear::RegressionSGVB;
use super::traits::{
    AnalyticalKL, BlackBoxLikelihood, LocalReparamModel, Prior, VariationalDistribution,
};

/// Configuration for SGVB estimator.
#[derive(Debug, Clone)]
pub struct SGVBConfig {
    /// Number of Monte Carlo samples S for gradient estimation
    pub num_samples: usize,
    /// KL divergence weight β for annealing (0 = no KL, 1 = full ELBO)
    pub kl_weight: f64,
}

impl Default for SGVBConfig {
    fn default() -> Self {
        Self {
            num_samples: 10,
            kl_weight: 1.0,
        }
    }
}

impl SGVBConfig {
    /// Create a new SGVB configuration.
    pub fn new(num_samples: usize) -> Self {
        Self {
            num_samples,
            kl_weight: 1.0,
        }
    }
}

/// Antithetic sampling for variance reduction in the local reparameterization trick.
///
/// Draws S/2 noise vectors epsilon, pairs each with its mirror -epsilon so the empirical
/// mean of noise is exactly zero. This strictly reduces variance of the MC
/// log likelihood estimate with no hyperparameters to tune.
///
/// # References
///
/// - Hammersley, J. M. & Morton, K. W. (1956). A new Monte Carlo technique:
///   antithetic variates. *Math. Proc. Cambridge Phil. Soc.*, 52(3), 449-475.
/// - Kucukelbir, A., Tran, D., Ranganath, R., Gelman, A. & Blei, D. M. (2017).
///   Automatic differentiation variational inference. *JMLR*, 18(14), 1-45.
/// - Ruiz, F. R., Titsias, M. K. & Blei, D. M. (2016). The generalized
///   reparameterization gradient. *NeurIPS*.
/// - Roeder, G., Wu, Y. & Duvenaud, D. (2017). Sticking the landing: simple,
///   lower-variance gradient estimators for variational inference. *NeurIPS*.
///
/// Returns epsilon of shape (num_samples, n, k).
pub fn antithetic_epsilon(
    num_samples: usize,
    n: usize,
    k: usize,
    device: &candle_core::Device,
    dtype: DType,
) -> Result<Tensor> {
    let half_s = num_samples / 2;
    if half_s > 0 {
        let eps_half = Tensor::randn(0f32, 1f32, (half_s, n, k), device)?.to_dtype(dtype)?;
        let eps_neg = eps_half.neg()?;
        if num_samples % 2 == 1 {
            let eps_extra = Tensor::randn(0f32, 1f32, (1, n, k), device)?.to_dtype(dtype)?;
            Tensor::cat(&[eps_half, eps_neg, eps_extra], 0)
        } else {
            Tensor::cat(&[eps_half, eps_neg], 0)
        }
    } else {
        Tensor::randn(0f32, 1f32, (1, n, k), device)?.to_dtype(dtype)
    }
}

/// Compute ELBO loss for any model implementing `LocalReparamModel`.
///
/// loss = -E[log p(y|eta)] + beta * KL
pub fn generic_local_reparam_loss<M: LocalReparamModel, L: BlackBoxLikelihood>(
    model: &M,
    likelihood: &L,
    num_samples: usize,
    kl_weight: f64,
) -> Result<Tensor> {
    let sample = model.forward(num_samples)?;
    let llik = likelihood.log_likelihood(&[&sample.eta])?;
    let llik = if llik.rank() > 1 { llik.sum(1)? } else { llik };

    let elbo = (llik.mean(0)? - (sample.kl * kl_weight)?)?;
    elbo.neg()
}

/// Compute loss using the local reparameterization trick.
///
/// Delegates to `generic_local_reparam_loss` via the `LocalReparamModel` trait.
pub fn local_reparam_loss<V, P, L>(
    model: &RegressionSGVB<V, P>,
    likelihood: &L,
    num_samples: usize,
    kl_weight: f64,
) -> Result<Tensor>
where
    V: VariationalDistribution,
    P: Prior + AnalyticalKL,
    L: BlackBoxLikelihood,
{
    generic_local_reparam_loss(model, likelihood, num_samples, kl_weight)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::candle::sgvb::{GaussianPrior, GaussianRegressionSGVB};
    use candle_core::{DType, Device, Tensor};
    use candle_nn::{Optimizer, VarBuilder, VarMap};

    /// Simple Gaussian likelihood for testing
    struct GaussianLikelihood {
        y: Tensor,     // observations (n, k)
        ln_sigma: f64, // log noise std
    }

    impl GaussianLikelihood {
        fn new(y: Tensor, sigma: f64) -> Self {
            Self {
                y,
                ln_sigma: sigma.ln(),
            }
        }
    }

    impl BlackBoxLikelihood for GaussianLikelihood {
        fn log_likelihood(&self, etas: &[&Tensor]) -> Result<Tensor> {
            // etas[0]: (S, n, k), y: (n, k)
            // log N(y; eta, sigma^2) = -0.5 * [(y - eta)^2 / sigma^2 + 2*ln(sigma) + ln(2pi)]
            let eta = etas[0];
            let sigma_sq = (2.0 * self.ln_sigma).exp();
            let ln_2pi = (2.0 * std::f64::consts::PI).ln();
            let const_term = 2.0 * self.ln_sigma + ln_2pi;

            let diff_sq = eta.broadcast_sub(&self.y)?.sqr()?;
            let log_prob = (((diff_sq / sigma_sq)? + const_term)? * (-0.5))?;

            // Sum over (n, k) dimensions, return (S,)
            log_prob.sum(2)?.sum(1)
        }
    }

    #[test]
    fn test_local_reparam_loss_recovery() -> Result<()> {
        let device = Device::Cpu;
        let dtype = DType::F32;

        let n = 150;
        let p = 50;
        let k = 1;

        let x = Tensor::randn(0f32, 1f32, (n, p), &device)?;

        // y = X[:,0] * 2.0 + noise
        let true_coef = 2.0f64;
        let x_first = x.narrow(1, 0, 1)?;
        let noise = Tensor::randn(0f32, 0.5f32, (n, k), &device)?;
        let y = (x_first * true_coef)?.add(&noise)?;

        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, dtype, &device);

        let likelihood = GaussianLikelihood::new(y, 0.5);
        let prior = GaussianPrior::new(vb.pp("prior"), 1.0)?;
        let config = SGVBConfig::new(50);
        let model = GaussianRegressionSGVB::new(vb.pp("model"), x, k, prior, config)?;

        let mut optimizer = candle_nn::AdamW::new_lr(varmap.all_vars(), 0.01)?;

        for i in 0..200 {
            let loss = local_reparam_loss(&model, &likelihood, 50, 1.0)?;
            optimizer.backward_step(&loss)?;

            if i % 50 == 0 {
                let loss_val: f32 = loss.to_scalar()?;
                println!("local_reparam iter {}: loss = {:.4}", i, loss_val);
            }
        }

        let coef_mean = model.coef_mean()?;
        let coef_first: f32 = coef_mean.get(0)?.get(0)?.to_scalar()?;

        let mut other_sum = 0.0f32;
        for i in 1..p {
            let val: f32 = coef_mean.get(i)?.get(0)?.to_scalar()?;
            other_sum += val.abs();
        }
        let other_mean = other_sum / (p - 1) as f32;

        println!("\nLocal reparam coefficients:");
        println!("  First (true={:.1}): {:.4}", true_coef, coef_first);
        println!("  Others mean abs: {:.4}", other_mean);

        assert!(
            coef_first > 1.0,
            "First coef should be > 1.0, got {}",
            coef_first
        );
        assert!(
            coef_first.abs() > other_mean * 2.0,
            "First coef should dominate others"
        );

        Ok(())
    }
}