# Parity tests for the FLUX.2 VAE decoder (post_quant_conv + decoder
# body + BatchNorm latent statistics) against diffusers reference
# fixtures (generated by tools/gen_fixtures_flux2_vae.py).

if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
  exit_file("torch not fully installed")
}
if (!requireNamespace("safetensors", quietly = TRUE)) {
  exit_file("safetensors not installed")
}

library(diffuseR)

fixture <- function(name) {
  p <- system.file("tinytest", "fixtures", name, package = "diffuseR")
  if (p == "") p <- file.path("fixtures", name)
  p
}
vae_path <- fixture("vae_flux2_tiny.safetensors")
io_path <- fixture("vae_flux2_io.safetensors")
if (!file.exists(io_path)) exit_file("flux2 vae fixtures missing")

fx <- safetensors::safe_load_file(io_path, framework = "torch")

max_abs_diff <- function(a, b) {
  as.numeric(torch::torch_max(torch::torch_abs(
    a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
  )))
}

dec <- load_flux2_vae_decoder(vae_path, latent_channels = 32L,
  block_channels = c(32L, 32L, 16L, 8L), norm_groups = 8L,
  verbose = FALSE)

# BN statistics loaded from the checkpoint
expect_true(max_abs_diff(dec$bn$running_mean, fx$bn_mean) == 0)
expect_true(max_abs_diff(dec$bn$running_var, fx$bn_var) == 0)

# Decode parity (post_quant_conv + decoder body)
img <- torch::with_no_grad(dec(fx$latent))
expect_equal(as.integer(img$shape), as.integer(fx$image$shape))
expect_true(max_abs_diff(img, fx$image) < 1e-5)

# BN normalize / denormalize parity + roundtrip
norm <- flux2_bn_normalize(fx$patched, dec$bn$running_mean,
  dec$bn$running_var)
expect_true(max_abs_diff(norm, fx$normalized) < 1e-5)
denorm <- flux2_bn_normalize(fx$normalized, dec$bn$running_mean,
  dec$bn$running_var, inverse = TRUE)
expect_true(max_abs_diff(denorm, fx$denormalized) < 1e-5)
expect_true(max_abs_diff(denorm, fx$patched) < 1e-5)
