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

ckpt_dir <- system.file("tinytest", "fixtures", "t5_tiny_ckpt",
  package = "diffuseR")
if (ckpt_dir == "") ckpt_dir <- "fixtures/t5_tiny_ckpt"
if (!dir.exists(ckpt_dir)) exit_file("t5 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())
  )))
}

# --- relative position bucketing: exact integer parity ---------------------------

buckets_r <- diffuseR:::.t5_relative_position_bucket(
  fx$rel_positions$to(dtype = torch::torch_long())
)
expect_equal(
  as.integer(buckets_r$flatten()),
  as.integer(fx$rel_buckets$flatten())
)

# --- tiny encoder parity (padded batch, no mask) -----------------------------------

model <- load_t5_text_encoder(ckpt_dir, device = "cpu", dtype = "float32",
  verbose = FALSE)
out <- torch::with_no_grad(
  model(fx$input_ids$to(dtype = torch::torch_long()) + 1L)
)
expect_equal(as.integer(out$shape), as.integer(fx$out$shape))
expect_true(max_abs_diff(out, fx$out) < 1e-4)

# (encode_with_t5 is exercised end-to-end in the pipeline tests; real
# token ids exceed this fixture's 100-entry vocab.)
