# Parity tests for the LTX-2.3 RoPE port against diffusers reference
# fixtures (generated by tools/gen_fixtures_rope.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_ltx23.safetensors",
  package = "diffuseR")
if (fixture_path == "") {
  # Running from the source tree (e.g. during development)
  fixture_path <- "fixtures/rope_ltx23.safetensors"
}
if (!file.exists(fixture_path)) exit_file("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())
  )))
}

# --- apply_split_rotary_emb ---------------------------------------------------

out4 <- ltx23_apply_split_rotary_emb(fx$split4_x, list(fx$split4_cos, fx$split4_sin))
expect_equal(as.integer(out4$shape), as.integer(fx$split4_out$shape))
expect_true(max_abs_diff(out4, fx$split4_out) < 1e-5)

# 3D input with 4D freqs reshapes per-head and back
out3 <- ltx23_apply_split_rotary_emb(fx$split3_x, list(fx$split4_cos, fx$split4_sin))
expect_equal(as.integer(out3$shape), as.integer(fx$split3_out$shape))
expect_true(max_abs_diff(out3, fx$split3_out) < 1e-5)

# --- apply_interleaved_rotary_emb ----------------------------------------------

outi <- ltx23_apply_interleaved_rotary_emb(fx$inter_x, list(fx$inter_cos, fx$inter_sin))
expect_true(max_abs_diff(outi, fx$inter_out) < 1e-5)

# --- video coords + split freqs -------------------------------------------------

rope_v <- ltx23_rotary_pos_embed(
  dim = 64L, scale_factors = c(8L, 32L, 32L), modality = "video",
  rope_type = "split", num_attention_heads = 4L
)
vc <- rope_v$prepare_video_coords(
  batch_size = 2L, num_frames = 3L, height = 4L, width = 6L,
  device = "cpu", fps = 24.0
)
expect_equal(as.integer(vc$shape), as.integer(fx$video_coords$shape))
expect_true(max_abs_diff(vc, fx$video_coords) < 1e-6)

vf <- rope_v(vc)
expect_true(max_abs_diff(vf[[1]], fx$video_cos) < 1e-6)
expect_true(max_abs_diff(vf[[2]], fx$video_sin) < 1e-6)

# fp16 input stays fp16 and matches reference values
outv16 <- ltx23_apply_split_rotary_emb(fx$video_x_f16, vf)
expect_equal(outv16$dtype$.type(), "Half")
expect_true(max_abs_diff(outv16, fx$video_out_f16) < 1e-3)

# --- audio coords + split freqs -------------------------------------------------

rope_a <- ltx23_rotary_pos_embed(
  dim = 32L, scale_factors = c(8L, 32L, 32L), modality = "audio",
  rope_type = "split", num_attention_heads = 4L
)
ac <- rope_a$prepare_audio_coords(batch_size = 2L, num_frames = 5L, device = "cpu")
expect_equal(as.integer(ac$shape), as.integer(fx$audio_coords$shape))
expect_true(max_abs_diff(ac, fx$audio_coords) < 1e-6)

af <- rope_a(ac)
expect_true(max_abs_diff(af[[1]], fx$audio_cos) < 1e-6)
expect_true(max_abs_diff(af[[2]], fx$audio_sin) < 1e-6)

# --- interleaved freqs -----------------------------------------------------------

rope_i <- ltx23_rotary_pos_embed(
  dim = 64L, modality = "video", rope_type = "interleaved",
  num_attention_heads = 4L
)
ifr <- rope_i(vc)
expect_true(max_abs_diff(ifr[[1]], fx$video_inter_cos) < 1e-6)
expect_true(max_abs_diff(ifr[[2]], fx$video_inter_sin) < 1e-6)
