// iwae.mbt — Importance Weighted Autoencoder (v0.108.0).
//
// The IWAE bound (Burda et al. 2016) is a tighter variational bound
// than the ELBO:
//
//   log p(x) ≥ L_K(x) = E_{z_1, ..., z_K ~ q(z|x)}
//                       [ log ( (1/K) Σ_{k=1..K} w_k ) ]
//   where w_k = p(x, z_k) / q(z_k | x)
//         = p(x | z_k) · p(z_k) / q(z_k | x)
//
// As K → ∞, L_K → log p(x). For K=1, L_K = ELBO.
//
// For Gaussian p(z) = N(0, I), Gaussian q(z|x) = N(μ, σ²), and a
// Gaussian decoder p(x|z) = N(x_recon, 1):
//
//   log w_k = log p(x|z_k) + log p(z_k) - log q(z_k|x)
//
// The IWAE bound is generally tighter than the ELBO and gives better
// posterior approximations at test time, but the gradient w.r.t. θ
// has higher variance (a known IWAE training issue).
//
// Reference: Burda et al. 2016 "Importance Weighted Autoencoders".

///|
/// Compute the IWAE bound for one sample x. `x_recon_k`, `log_var_k`
/// are flat arrays of length K × (their respective dim).
/// Returns L_K(x) (higher = better bound on log p(x)).
pub fn iwae_per_sample_bound(
  x : Array[Float],
  x_recon_k : Array[Float],
  z_k : Array[Float],
  mu : Array[Float],
  log_var_k : Array[Float],
  k_samples : Int,
) -> Float {
  if k_samples <= 0 {
    return 0.0F
  }
  let x_dim = x.length()
  let latent_dim = mu.length()
  let mut max_log_w = -1.0e30F
  // First pass: compute log w_k for each sample, track max for log-sum-exp.
  let log_ws : Array[Float] = Array::make(k_samples, 0.0F)
  for k in 0.. max_log_w {
      max_log_w = log_w_k
    }
  }
  // Second pass: log-sum-exp.
  let mut sum_exp = 0.0F
  for k in 0.. (VAEModel, Float) {
  let (mu, log_var) = vae_encode(model.encoder, x)
  // Sample K latents via reparameterization.
  let z_k : Array[Float] = Array::make(
    k_samples * model.latent_dim, 0.0F,
  )
  let log_var_k : Array[Float] = Array::make(
    k_samples * model.latent_dim, 0.0F,
  )
  let x_recon_k : Array[Float] = Array::make(
    k_samples * model.x_dim, 0.0F,
  )
  for k in 0..