# Parity tests for the FLUX transformer blocks against diffusers
# reference fixtures (generated by tools/gen_fixtures_flux_dit.py,
# checked in).

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_flux.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/dit_flux.safetensors"
if (!file.exists(fixture_path)) exit_file("flux 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())
  )))
}

# Copy fixture weights (diffusers state-dict names) into an R module.
# Errors if any name has no destination or any destination stays unfilled.
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
S_TXT <- 7L

# --- flux_ada_layer_norm_zero ---------------------------------------------------

adazero <- flux_ada_layer_norm_zero(DIM)
adazero$eval()
load_named_weights(adazero, fixture_group("adazero"))
az <- torch::with_no_grad(adazero(fx$joint_x, emb = fx$temb))
expect_true(max_abs_diff(az[[1]], fx$adazero_x_norm) < 1e-5)
expect_true(max_abs_diff(az[[2]], fx$adazero_gate_msa) < 1e-5)
expect_true(max_abs_diff(az[[3]], fx$adazero_shift_mlp) < 1e-5)
expect_true(max_abs_diff(az[[4]], fx$adazero_scale_mlp) < 1e-5)
expect_true(max_abs_diff(az[[5]], fx$adazero_gate_mlp) < 1e-5)

# --- flux_ada_layer_norm_zero_single ----------------------------------------------

adasingle <- flux_ada_layer_norm_zero_single(DIM)
adasingle$eval()
load_named_weights(adasingle, fixture_group("adasingle"))
as_ <- torch::with_no_grad(adasingle(fx$joint_x, emb = fx$temb))
expect_true(max_abs_diff(as_[[1]], fx$adasingle_x_norm) < 1e-5)
expect_true(max_abs_diff(as_[[2]], fx$adasingle_gate) < 1e-5)

# --- flux_ada_layer_norm_continuous (scale-first chunk order) ----------------------

adacont <- flux_ada_layer_norm_continuous(DIM, DIM)
adacont$eval()
load_named_weights(adacont, fixture_group("adacont"))
ac <- torch::with_no_grad(adacont(fx$joint_x, fx$temb))
expect_true(max_abs_diff(ac, fx$adacont_out) < 1e-5)

# --- flux_attention, double-stream variant ------------------------------------------

attn_d <- flux_attention(DIM, HEADS, HEAD_DIM, added_kv = TRUE)
attn_d$eval()
load_named_weights(attn_d, fixture_group("attnd"))
ad <- torch::with_no_grad(attn_d(
  hidden_states = fx$img_x,
  encoder_hidden_states = fx$txt_x,
  image_rotary_emb = rope
))
expect_equal(as.integer(ad[[1]]$shape), as.integer(fx$attnd_out$shape))
expect_true(max_abs_diff(ad[[1]], fx$attnd_out) < 1e-5)
expect_true(max_abs_diff(ad[[2]], fx$attnd_ctx_out) < 1e-5)

# --- flux_attention, pre_only variant -------------------------------------------------

attn_s <- flux_attention(DIM, HEADS, HEAD_DIM, pre_only = TRUE)
attn_s$eval()
load_named_weights(attn_s, fixture_group("attns"))
asn <- torch::with_no_grad(attn_s(
  hidden_states = fx$joint_x,
  image_rotary_emb = rope
))
expect_true(max_abs_diff(asn, fx$attns_out) < 1e-5)

# --- flux_double_block -----------------------------------------------------------------

dbl <- flux_double_block(DIM, HEADS, HEAD_DIM)
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 = fx$temb,
  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)

# --- flux_single_block -------------------------------------------------------------------
# The R block takes the pre-concatenated [txt; img] sequence and returns it
# joint; the reference returns the split streams.

sgl <- flux_single_block(DIM, HEADS, HEAD_DIM)
sgl$eval()
load_named_weights(sgl, fixture_group("sgl"))
joint_in <- torch::torch_cat(list(fx$txt_x, fx$img_x), dim = 2L)
sg <- torch::with_no_grad(sgl(
  hidden_states = joint_in,
  temb = fx$temb,
  image_rotary_emb = rope
))
expect_true(max_abs_diff(sg$narrow(2L, 1L, S_TXT), fx$sgl_enc_out) < 1e-5)
expect_true(max_abs_diff(
  sg$narrow(2L, S_TXT + 1L, sg$shape[2] - S_TXT),
  fx$sgl_hid_out
) < 1e-5)
