# Parity test for the full FLUX transformer against a tiny random-init
# diffusers FluxTransformer2DModel (fixture generated by
# tools/gen_fixtures_flux_model.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", "flux_model.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/flux_model.safetensors"
if (!file.exists(fixture_path)) exit_file("flux model 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())
  )))
}

model <- flux_transformer(
  in_channels = 4L,
  num_layers = 1L,
  num_single_layers = 1L,
  attention_head_dim = 8L,
  num_attention_heads = 2L,
  joint_attention_dim = 16L,
  pooled_projection_dim = 12L,
  axes_dims_rope = c(2L, 2L, 4L)
)
model$eval()

# Strict both ways: every fixture key must land, every param must fill.
# This proves the R module tree matches the diffusers state dict exactly.
weights <- fx[grep("^model\\.", names(fx))]
names(weights) <- sub("^model\\.", "", names(weights))
dests <- c(model$named_parameters(), model$named_buffers())
expect_equal(character(0), setdiff(names(weights), names(dests)))
expect_equal(character(0), setdiff(names(dests), names(weights)))
torch::with_no_grad({
  for (name in names(weights)) dests[[name]]$copy_(weights[[name]])
})

ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L)
rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 4L))

out <- torch::with_no_grad(model(
  hidden_states = fx$hidden,
  encoder_hidden_states = fx$encoder,
  pooled_projections = fx$pooled,
  timestep = fx$timestep,
  image_rotary_emb = rope
))
expect_equal(as.integer(out$shape), as.integer(fx$out$shape))
expect_true(max_abs_diff(out, fx$out) < 1e-4)
