// realnvp.mbt — Real NVP normalizing flow (v0.102.0): stack of
// AffineCouplingLayers with fixed permutation between layers.
//
// Real NVP (Dinh et al. 2016) chains multiple affine coupling layers,
// with a fixed permutation of the input dimensions between layers so
// that different "halves" of the input get transformed at each
// layer. The composition of bijections is a bijection, and the total
// log |det J| is the sum of the per-layer log |det| terms.
//
// Reference: Dinh et al. 2016 "Density Estimation with Real NVP".

///|
/// Real NVP flow: N stacked affine coupling layers with a fixed
/// permutation between layers.
pub struct RealNVP {
  n_layers : Int
  layers : Array[AffineCouplingLayer]
  // Permutation applied between layers. Default: reverse the input
  // dimensions so that x₁ becomes x₂ in the next layer.
  perm : Array[Int]
  // Input dim (must equal 2 × half_dim for all layers).
  dim : Int
}

///|
/// Build a fresh Real NVP flow with `n_layers` coupling layers + a
/// default reverse permutation between layers. Each layer has its
/// own seed (offset).
pub fn RealNVP::new(
  dim : Int,
  n_layers : Int,
  seed : UInt64,
) -> RealNVP {
  let half = dim / 2
  let layers : Array[AffineCouplingLayer] = Array::make(
    n_layers, AffineCouplingLayer::new(half, seed),
  )
  for i in 0.. Array[Float] {
  let out : Array[Float] = Array::make(x.length(), 0.0F)
  for i in 0.. (Array[Float], Float) {
  let mut cur = x.copy()
  let mut total_log_det = 0.0F
  for i in 0.. Array[Float] {
  let mut cur = y.copy()
  // Walk layers in reverse.
  let mut i = flow.n_layers - 1
  while i >= 0 {
    // Reverse-permute first (since coupling_inverse expects the
    // permuted layout).
    if i >= 0 {
      // Apply the inverse permutation: if perm maps i → perm[i], then
      // the inverse permutation maps perm[i] → i. Build it once.
      let inv_perm : Array[Int] = Array::make(flow.dim, 0)
      for j in 0.. Float {
  let mut cur = x.copy()
  let mut total_log_det = 0.0F
  for i in 0..