# Parity tests for the Z-Image transformer blocks against diffusers
# reference fixtures (generated by tools/gen_fixtures_zimage_dit.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", "dit_zimage.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/dit_zimage.safetensors"
if (!file.exists(fixture_path)) exit_file("zimage dit 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())
  )))
}

load_named_weights <- function(module, weights) {
  dests <- c(module$named_parameters(), module$named_buffers())
  missing_dest <- setdiff(names(weights), names(dests))
  if (length(missing_dest)) {
    stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", "))
  }
  unfilled <- setdiff(names(dests), names(weights))
  if (length(unfilled)) {
    stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", "))
  }
  torch::with_no_grad({
    for (name in names(weights)) dests[[name]]$copy_(weights[[name]])
  })
  invisible(module)
}

subset_prefix <- function(fx, prefix) {
  keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE)
  out <- fx[keys]
  names(out) <- sub(paste0("^", prefix, "\\."), "", keys)
  out
}

freqs <- list(fx$freqs_cos, fx$freqs_sin)

# --- modulated block (noise refiner / main trunk) --------------------------------

blk <- zimage_block(dim = 32L, n_heads = 2L, modulation = TRUE)
load_named_weights(blk, subset_prefix(fx, "mod"))
blk$eval()
out <- torch::with_no_grad(blk(fx$x, freqs, adaln_input = fx$adaln))
expect_equal(out$shape, c(1L, 24L, 32L))
expect_true(max_abs_diff(out, fx$mod_out) < 1e-5)

# --- unmodulated block (context refiner) ------------------------------------------

ublk <- zimage_block(dim = 32L, n_heads = 2L, modulation = FALSE)
load_named_weights(ublk, subset_prefix(fx, "unmod"))
ublk$eval()
uout <- torch::with_no_grad(ublk(fx$x, freqs))
expect_true(max_abs_diff(uout, fx$unmod_out) < 1e-5)

# The unmodulated block must carry no adaLN weights
expect_false(any(grepl("adaLN", names(ublk$named_parameters()))))

# --- chunked attention matches unchunked -------------------------------------------

out_chunked <- torch::with_no_grad(
  blk(fx$x, freqs, adaln_input = fx$adaln, chunk_size = 7L))
expect_true(max_abs_diff(out_chunked, out) < 1e-6)

# --- final layer -------------------------------------------------------------------

fin <- zimage_final_layer(hidden_size = 32L, out_channels = 64L)
load_named_weights(fin, subset_prefix(fx, "final"))
fin$eval()
fout <- torch::with_no_grad(fin(fx$x, fx$adaln))
expect_equal(fout$shape, c(1L, 24L, 64L))
expect_true(max_abs_diff(fout, fx$final_out) < 1e-5)

# --- timestep embedder --------------------------------------------------------------

temb <- zimage_t_embedder(out_size = 32L, mid_size = 48L)
load_named_weights(temb, subset_prefix(fx, "temb"))
temb$eval()
tout <- torch::with_no_grad(temb(fx$temb_in))
expect_true(max_abs_diff(tout, fx$temb_out) < 1e-5)

# --- cap embedder (Sequential RMSNorm + Linear, keys 0/1) ----------------------------

cap <- torch::nn_sequential(
  diffuseR:::ltx23_rms_norm(24L, eps = 1e-5),
  torch::nn_linear(24L, 32L, bias = TRUE)
)
load_named_weights(cap, subset_prefix(fx, "cap"))
cap$eval()
cout <- torch::with_no_grad(cap(fx$cap_in))
expect_true(max_abs_diff(cout, fx$cap_out) < 1e-6)
