# Parity tests for the FLUX RoPE port and latent pack/unpack against
# diffusers reference fixtures (generated by tools/gen_fixtures_flux.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", "rope_flux.safetensors",
  package = "diffuseR")
if (fixture_path == "") {
  # Running from the source tree (e.g. during development)
  fixture_path <- "fixtures/rope_flux.safetensors"
}
if (!file.exists(fixture_path)) exit_file("flux 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())
  )))
}

# --- latent image ids ---------------------------------------------------------

ids_r <- flux_prepare_latent_image_ids(8L, 12L)
expect_equal(as.integer(ids_r$shape), as.integer(fx$img_ids$shape))
expect_true(max_abs_diff(ids_r, fx$img_ids) == 0)

# --- flux_pos_embed: full-size axes (16, 56, 56) -------------------------------

pf <- flux_pos_embed(fx$ids, axes_dim = c(16L, 56L, 56L))
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)

# Text ids (all zero) get the identity rotation: cos 1, sin 0
txt_cos <- pf[[1]][1:7, ]
txt_sin <- pf[[2]][1:7, ]
expect_true(max_abs_diff(txt_cos, torch::torch_ones_like(txt_cos)) == 0)
expect_true(max_abs_diff(txt_sin, torch::torch_zeros_like(txt_sin)) == 0)

# --- flux_pos_embed: tiny axes (2, 2, 4) ---------------------------------------

pt <- flux_pos_embed(fx$ids, axes_dim = c(2L, 2L, 4L))
expect_equal(as.integer(pt[[1]]$shape), as.integer(fx$pos_tiny_cos$shape))
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)

# --- flux_apply_rotary_emb ------------------------------------------------------

rot <- flux_apply_rotary_emb(fx$rot_x, list(fx$pos_tiny_cos, fx$pos_tiny_sin))
expect_equal(as.integer(rot$shape), as.integer(fx$rot_out$shape))
expect_true(max_abs_diff(rot, fx$rot_out) < 1e-5)

# fp16 input stays fp16 and matches reference values
rot16 <- flux_apply_rotary_emb(fx$rot_x_f16, list(fx$pos_tiny_cos, fx$pos_tiny_sin))
expect_equal(rot16$dtype$.type(), "Half")
expect_true(max_abs_diff(rot16, fx$rot_out_f16) < 1e-3)

# --- latent pack / unpack -------------------------------------------------------

packed <- flux_pack_latents(fx$pack_in)
expect_equal(as.integer(packed$shape), as.integer(fx$pack_out$shape))
expect_true(max_abs_diff(packed, fx$pack_out) == 0)

unpacked <- flux_unpack_latents(fx$pack_out, height = 64L, width = 96L,
  vae_scale_factor = 8L)
expect_equal(as.integer(unpacked$shape), as.integer(fx$unpack_out$shape))
expect_true(max_abs_diff(unpacked, fx$unpack_out) == 0)

# R-side roundtrip identity: unpack(pack(x)) == x
expect_true(max_abs_diff(unpacked, fx$pack_in) == 0)
