# Parity tests for the FLUX.2 transformer blocks against diffusers
# reference fixtures (generated by tools/gen_fixtures_flux2_dit.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_path <- system.file("tinytest", "fixtures", "dit_flux2.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/dit_flux2.safetensors"
if (!file.exists(fixture_path)) exit_file("flux2 dit fixtures missing")

fx <- safetensors::safe_load_file(fixture_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())
  )))
}

load_named_weights <- function(module, weights) {
  dests <- c(module$named_parameters(), module$named_buffers())
  missing_dest <- setdiff(names(weights), names(dests))
  if (length(missing_dest)) {
    stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", "))
  }
  unfilled <- setdiff(names(dests), names(weights))
  if (length(unfilled)) {
    stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", "))
  }
  torch::with_no_grad({
    for (name in names(weights)) dests[[name]]$copy_(weights[[name]])
  })
  invisible(module)
}

fixture_group <- function(prefix) {
  keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE)
  w <- fx[keys]
  names(w) <- sub(paste0("^", prefix, "\\."), "", keys)
  w
}

rope <- list(fx$rope_cos, fx$rope_sin)
DIM <- 16L
HEADS <- 2L
HEAD_DIM <- 8L

# --- flux2_modulation --------------------------------------------------------------

mod2 <- flux2_modulation(DIM, mod_param_sets = 2L)
mod2$eval()
load_named_weights(mod2, fixture_group("mod2"))
expect_true(max_abs_diff(torch::with_no_grad(mod2(fx$temb)), fx$mod2_out) < 1e-5)

mod1 <- flux2_modulation(DIM, mod_param_sets = 1L)
mod1$eval()
load_named_weights(mod1, fixture_group("mod1"))
expect_true(max_abs_diff(torch::with_no_grad(mod1(fx$temb)), fx$mod1_out) < 1e-5)

# --- flux2_feed_forward (SwiGLU) ------------------------------------------------------

ff <- flux2_feed_forward(DIM, DIM, mult = 3.0)
ff$eval()
load_named_weights(ff, fixture_group("ff"))
expect_true(max_abs_diff(torch::with_no_grad(ff(fx$joint_x)), fx$ff_out) < 1e-5)

# --- flux2_double_block ----------------------------------------------------------------

dbl <- flux2_double_block(DIM, HEADS, HEAD_DIM, mlp_ratio = 3.0)
dbl$eval()
load_named_weights(dbl, fixture_group("dbl"))
db <- torch::with_no_grad(dbl(
  hidden_states = fx$img_x,
  encoder_hidden_states = fx$txt_x,
  temb_mod_img = fx$dbl_mod_img,
  temb_mod_txt = fx$dbl_mod_txt,
  image_rotary_emb = rope
))
expect_true(max_abs_diff(db[[1]], fx$dbl_enc_out) < 1e-5)
expect_true(max_abs_diff(db[[2]], fx$dbl_hid_out) < 1e-5)

# --- flux2_single_block ------------------------------------------------------------------

sgl <- flux2_single_block(DIM, HEADS, HEAD_DIM, mlp_ratio = 3.0)
sgl$eval()
load_named_weights(sgl, fixture_group("sgl"))
sg <- torch::with_no_grad(sgl(
  hidden_states = fx$joint_x,
  temb_mod = fx$sgl_mod,
  image_rotary_emb = rope
))
expect_true(max_abs_diff(sg, fx$sgl_out) < 1e-5)
