# Parity tests for the Z-Image position ids, 3-axis RoPE, patchify chain,
# timestep sinusoid, and static-shift scheduler against diffusers
# reference fixtures (generated by tools/gen_fixtures_zimage.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_zimage.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/rope_zimage.safetensors"
if (!file.exists(fixture_path)) exit_file("zimage 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())
  )))
}

# --- padding lengths ------------------------------------------------------------

pad_len <- diffuseR:::zimage_pad_len
expect_equal(pad_len(37L), 27L)
expect_equal(pad_len(60L), 4L)
expect_equal(pad_len(64L), 0L)
expect_equal(pad_len(4096L), 0L)

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

# Caption: 37 tokens pad to 64; ramp 1..64 on axis 1 built over the
# padded length
cap_ids <- zimage_cap_pos_ids(37L + pad_len(37L))
expect_equal(cap_ids$shape, c(64L, 3L))
expect_true(max_abs_diff(cap_ids, fx$cap_pos_ids) == 0)

# Image: 6x10 grid = 60 tokens pad to 64; axis-1 start just past the
# padded caption, pads at the origin
img_ids <- zimage_img_pos_ids(6L, 10L, start0 = 65L)
expect_equal(img_ids$shape, c(64L, 3L))
expect_true(max_abs_diff(img_ids, fx$img_pos_ids) == 0)

# --- RoPE frequencies (theta 256, f32-angle cast) --------------------------------

freqs <- zimage_pos_embed(fx$rope_ids)
expect_equal(freqs[[1]]$shape, c(9L, 128L))
expect_true(max_abs_diff(freqs[[1]], fx$rope_cos) < 1e-6)
expect_true(max_abs_diff(freqs[[2]], fx$rope_sin) < 1e-6)

# Position (0,0,0) is the identity rotation
expect_true(max_abs_diff(freqs[[1]][9, ],
  torch::torch_ones(128L)) == 0)
expect_true(max_abs_diff(freqs[[2]][9, ],
  torch::torch_zeros(128L)) == 0)

# --- RoPE application (reuses the FLUX interleaved apply) -------------------------

# Fixture x is [B, S, H, D] (processor layout); flux_apply_rotary_emb
# takes [B, H, S, D]
x_bhsd <- fx$rope_x$permute(c(1L, 3L, 2L, 4L))
out <- flux_apply_rotary_emb(x_bhsd, freqs)
expect_true(max_abs_diff(out, fx$rope_out$permute(c(1L, 3L, 2L, 4L))) < 1e-5)

# --- patchify / unpatchify --------------------------------------------------------

tokens <- zimage_patchify(fx$img)
expect_equal(tokens$shape, c(60L, 64L))
expect_true(max_abs_diff(tokens, fx$img_tokens_padded[1:60, ]) == 0)

unpat <- zimage_unpatchify(fx$unpat_tokens, c(1L, 12L, 20L))
expect_equal(unpat$shape, c(16L, 1L, 12L, 20L))
expect_true(max_abs_diff(unpat, fx$unpat_out) == 0)

# Roundtrip identity
rt <- zimage_unpatchify(zimage_patchify(fx$img), c(1L, 12L, 20L))
expect_true(max_abs_diff(rt, fx$img) == 0)

# --- timestep sinusoid -------------------------------------------------------------

t_emb <- ltx23_get_timestep_embedding(fx$t_emb_in, 256L,
  flip_sin_to_cos = TRUE, downscale_freq_shift = 0)
# 1e-4, not 1e-5: sin(t * exp(...)) at t ~ 1000 amplifies a 1-ULP
# libtorch difference in exp() to ~6e-5 (CRAN torch vs the fixture's
# libtorch); flag/math bugs are O(1) so the assertion still catches them
expect_true(max_abs_diff(t_emb, fx$t_emb_out) < 1e-4)

# --- scheduler: static shift 3.0 on linspace(1, 1/N, N) ---------------------------

for (n in c(4L, 8L)) {
  sched <- flowmatch_scheduler_create(shift = 3.0, use_dynamic_shifting = FALSE)
  sched <- flowmatch_set_timesteps(sched, num_inference_steps = n,
    sigmas = seq(1, 1 / n, length.out = n))
  expect_true(max_abs_diff(sched$sigmas,
    fx[[paste0("sched_sigmas_", n)]]) < 1e-6)
  expect_true(max_abs_diff(sched$timesteps,
    fx[[paste0("sched_timesteps_", n)]]) < 1e-3)
}
