# Parity tests for the Qwen3 encoder port against HF transformers
# reference fixtures (generated by tools/gen_fixtures_qwen3.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", "qwen3_flux2.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/qwen3_flux2.safetensors"
if (!file.exists(fixture_path)) exit_file("qwen3 fixtures missing")

ckpt_dir <- system.file("tinytest", "fixtures", "qwen3_tiny_ckpt",
  package = "diffuseR")
if (ckpt_dir == "") ckpt_dir <- "fixtures/qwen3_tiny_ckpt"
if (!dir.exists(ckpt_dir)) exit_file("qwen3 tiny checkpoint 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 <- load_qwen3_text_encoder(ckpt_dir, device = "cpu",
  dtype = "float32", verbose = FALSE)

ids <- fx$input_ids$to(dtype = torch::torch_long()) + 1L
mask <- fx$attention_mask$to(dtype = torch::torch_long())

states <- torch::with_no_grad(model(ids, attention_mask = mask,
  out_layers = c(1L, 2L, 3L)))
expect_equal(as.integer(states[[1]]$shape), as.integer(fx$h1$shape))
expect_true(max_abs_diff(states[[1]], fx$h1) < 1e-4)
expect_true(max_abs_diff(states[[2]], fx$h2) < 1e-4)
expect_true(max_abs_diff(states[[3]], fx$h3) < 1e-4)

# Per-token concatenation matches the pipeline's stack/permute/reshape
stacked <- torch::torch_cat(states, dim = -1L)
expect_equal(as.integer(stacked$shape), as.integer(fx$stacked$shape))
expect_true(max_abs_diff(stacked, fx$stacked) < 1e-4)

# The mask matters: dropping it must change the padded batch's states
states_nomask <- torch::with_no_grad(model(ids, out_layers = c(3L)))
expect_true(max_abs_diff(states_nomask[[1]], fx$h3) > 1e-4)
