# Full-forward parity test for the Z-Image transformer against a tiny
# random-init diffusers reference model (generated by
# tools/gen_fixtures_zimage_model.py). Includes a strict bidirectional
# key census.

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", "zimage_model.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/zimage_model.safetensors"
if (!file.exists(fixture_path)) exit_file("zimage model 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())
  )))
}

sd <- fx[grep("^sd\\.", names(fx))]
names(sd) <- sub("^sd\\.", "", names(sd))

model <- zimage_transformer(
  in_channels = 4L, dim = 48L, n_layers = 2L, n_refiner_layers = 1L,
  n_heads = 2L, cap_feat_dim = 24L, axes_dims = c(8L, 8L, 8L)
)

# Strict bidirectional key census: every checkpoint key lands, every
# module parameter fills, and the context refiner carries no adaLN
dests <- c(model$named_parameters(), model$named_buffers())
expect_equal(sort(names(dests)), sort(names(sd)))
expect_false(any(grepl("context_refiner.*adaLN", names(dests))))
expect_true(any(grepl("noise_refiner.*adaLN", names(dests))))

torch::with_no_grad({
  for (name in names(sd)) dests[[name]]$copy_(sd[[name]])
})
model$eval()

out <- torch::with_no_grad(model(fx$x, fx$t, fx$cap))
expect_equal(out$shape, c(4L, 1L, 12L, 20L))
expect_true(max_abs_diff(out, fx$out) < 1e-4)

# Chunked attention path matches
out_chunked <- torch::with_no_grad(model(fx$x, fx$t, fx$cap, chunk_size = 16L))
expect_true(max_abs_diff(out_chunked, out) < 1e-5)
