# Parity tests for the FLUX.2 position ids, 4-axis RoPE, latent
# patchify/pack/unpack chain, and empirical mu against diffusers
# reference fixtures (generated by tools/gen_fixtures_flux2.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", "rope_flux2.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/rope_flux2.safetensors"
if (!file.exists(fixture_path)) exit_file("flux2 rope 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())
  )))
}

# --- position ids ---------------------------------------------------------------

txt_ids <- flux2_prepare_text_ids(7L)
expect_equal(as.integer(txt_ids$shape), as.integer(fx$txt_ids$shape))
expect_true(max_abs_diff(txt_ids, fx$txt_ids) == 0)

lat_ids <- flux2_prepare_latent_ids(6L, 10L)
expect_equal(as.integer(lat_ids$shape), as.integer(fx$latent_ids$shape))
expect_true(max_abs_diff(lat_ids, fx$latent_ids) == 0)

# --- 4-axis RoPE through flux_pos_embed (theta 2000) ---------------------------

pf <- flux_pos_embed(fx$ids, axes_dim = c(32L, 32L, 32L, 32L), theta = 2000)
expect_equal(as.integer(pf[[1]]$shape), as.integer(fx$pos_full_cos$shape))
expect_true(max_abs_diff(pf[[1]], fx$pos_full_cos) < 1e-6)
expect_true(max_abs_diff(pf[[2]], fx$pos_full_sin) < 1e-6)

pt <- flux_pos_embed(fx$ids, axes_dim = c(2L, 2L, 2L, 2L), theta = 2000)
expect_true(max_abs_diff(pt[[1]], fx$pos_tiny_cos) < 1e-6)
expect_true(max_abs_diff(pt[[2]], fx$pos_tiny_sin) < 1e-6)

# --- patchify / pack / unpack chain ----------------------------------------------

patched <- flux2_patchify_latents(fx$z32)
expect_equal(as.integer(patched$shape), as.integer(fx$patched$shape))
expect_true(max_abs_diff(patched, fx$patched) == 0)

expect_true(max_abs_diff(flux2_unpatchify_latents(fx$patched), fx$unpatched) == 0)
expect_true(max_abs_diff(flux2_unpatchify_latents(fx$patched), fx$z32) == 0)

packed <- flux2_pack_latents(fx$patched)
expect_equal(as.integer(packed$shape), as.integer(fx$packed$shape))
expect_true(max_abs_diff(packed, fx$packed) == 0)

unpacked <- flux2_unpack_latents_with_ids(fx$packed,
  flux2_prepare_latent_ids(6L, 10L), height = 6L, width = 10L)
expect_equal(as.integer(unpacked$shape), as.integer(fx$unpacked$shape))
expect_true(max_abs_diff(unpacked, fx$unpacked) == 0)
expect_true(max_abs_diff(unpacked, fx$patched) == 0)

# --- empirical mu ------------------------------------------------------------------

seqs <- as.numeric(fx$mu_seq)
steps <- as.numeric(fx$mu_steps)
vals <- as.numeric(fx$mu_vals)
for (i in seq_along(seqs)) {
  expect_true(abs(flux2_empirical_mu(seqs[i], steps[i]) - vals[i]) < 1e-12,
    info = sprintf("seq %d steps %d", seqs[i], steps[i]))
}
