# Parity test for the full FLUX.2 transformer against a tiny random-init
# diffusers Flux2Transformer2DModel (fixture generated by
# tools/gen_fixtures_flux2_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", "flux2_model.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/flux2_model.safetensors"
if (!file.exists(fixture_path)) exit_file("flux2 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 <- flux2_transformer(
  in_channels = 8L,
  num_layers = 1L,
  num_single_layers = 1L,
  attention_head_dim = 8L,
  num_attention_heads = 2L,
  joint_attention_dim = 24L,
  mlp_ratio = 3.0,
  axes_dims_rope = c(2L, 2L, 2L, 2L),
  rope_theta = 2000
)
model$eval()

# Strict both ways: 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, 2L, 2L), theta = 2000)

out <- torch::with_no_grad(model(
  hidden_states = fx$hidden,
  encoder_hidden_states = fx$encoder,
  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)

# Position-id builders match the fixture ids exactly
expect_true(max_abs_diff(flux2_prepare_text_ids(7L), fx$txt_ids) == 0)
expect_true(max_abs_diff(flux2_prepare_latent_ids(6L, 10L), fx$img_ids) == 0)
