// normalizing_flow_model.mbt — NormalizingFlowModel (v0.104.0):
// full flow model with prior + sample generation + density
// estimation.
//
// A normalizing flow (Rezende & Mohamed 2015, Dinh et al. 2016/2018)
// is an invertible transformation f: R^d → R^d such that:
//   z = f(x)    (encoder / inverse flow; maps data to latent)
//   x = f^{-1}(z)  (decoder / forward flow; samples from prior p(z))
//
// Given a prior p(z) (typically a standard Normal), the density of x
// under the model is:
//   log p(x) = log p(z) + log |det ∂f^{-1}/∂x|
//            = log p(z) + log |det J_f|
//
// Reference: Rezende & Mohamed 2015 "Variational Inference with
// Normalizing Flows"; Dinh et al. 2016/2018 (Real NVP / Glow).

///|
/// NormalizingFlowModel: a Glow flow + a standard Normal prior.
/// log p(z) = -0.5 · Σ z_i² - 0.5 · d · log(2π).
pub struct NormalizingFlowModel {
  dim : Int
  flow : Glow
}

///|
/// Build a fresh NormalizingFlowModel with `n_steps` Glow steps.
pub fn NormalizingFlowModel::new(
  dim : Int,
  n_steps : Int,
  seed : UInt64,
) -> NormalizingFlowModel {
  let flow = Glow::new(dim, n_steps, seed)
  { dim, flow }
}

///|
/// Initialize the flow's ActNorms from data (must be called once
/// before the first forward pass for proper Glow initialization).
pub fn normalizing_flow_init_from_data(
  model : NormalizingFlowModel,
  data : Array[Float],
  n : Int,
) -> NormalizingFlowModel {
  let new_flow = glow_init_from_data(model.flow, data, n)
  { ..model, flow: new_flow }
}

///|
/// Forward (decoder): z → x via Glow forward. Samples x by drawing
/// z ~ N(0, I) and mapping through the flow. Returns (x, log_det)
/// where log_det = log |det ∂f/∂z|.
pub fn normalizing_flow_forward(
  model : NormalizingFlowModel,
  z : Array[Float],
) -> (Array[Float], Float) {
  glow_forward(model.flow, z)
}

///|
/// Sample x from the model: draw z ~ N(0, I) and apply Glow forward.
pub fn normalizing_flow_sample(
  model : NormalizingFlowModel,
  rng : Xoshiro,
) -> Array[Float] {
  let z : Array[Float] = Array::make(model.dim, 0.0F)
  for i in 0.. (Array[Float], Float) {
  let z = glow_inverse(model.flow, x)
  let log_det = glow_total_log_det(model.flow, x)
  (z, log_det)
}

///|
/// Total log |det ∂f/∂x| for the current model + x (used by
/// normalizing_flow_inverse / log_likelihood).
pub fn glow_total_log_det(
  flow : Glow,
  x : Array[Float],
) -> Float {
  // The inverse (x → z) change of variables has log_det = total
  // forward log_det (with sign). For the inverse, log|det| is the
  // same magnitude. We compute it as the sum of per-step log dets
  // for the inverse pass: x → z. Each GlowStep invert has the same
  // log|det| as the forward (since affine couplings and ActNorm are
  // self-inverse under |·|, and 1×1 conv has log|det W| = log|det W|).
  let mut cur = x.copy()
  let mut total_log_det = 0.0F
  // Forward direction: same log det as inverse since these are all
  // sign-symmetric. We use the forward pass for simplicity.
  for i in 0.. Float {
  let (z, log_det) = normalizing_flow_inverse(model, x)
  // log p(z) = -0.5 · Σ z_i² - 0.5 · d · log(2π)
  let mut lp_z = -0.5F * Float::from_int(model.dim) * log_2pi_v2()
  for i in 0.. Float {
  logf(2.0F * 3.14159265F)
}

///|
/// Compute the mean log-likelihood over a batch. Useful as the
/// loss to minimize during training.
pub fn normalizing_flow_mean_log_likelihood(
  model : NormalizingFlowModel,
  data : Array[Float],
  n : Int,
) -> Float {
  if n <= 0 {
    return 0.0F
  }
  let mut sum = 0.0F
  for i in 0..