# Parity tests for the LTX-2.3 DiT port against diffusers reference
# fixtures (generated by tools/gen_fixtures_dit.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", "dit_ltx23.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/dit_ltx23.safetensors"
if (!file.exists(fixture_path)) exit_file("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())
  )))
}

# Copy fixture weights (diffusers state-dict names) into an R module.
# Errors if any name has no destination or any destination stays unfilled.
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)
}

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

# --- Scratch buffers keep score and output roles distinct -------------------------
release_attn_buffers <- diffuseR:::.ltx23_release_attn_buffers
get_attn_buffer <- diffuseR:::.ltx23_get_attn_buffer
release_attn_buffers()
collision_shape <- c(1L, 2L, 8L, 8L)
score_buf <- get_attn_buffer(
  "scores", collision_shape, torch::torch_float32(), torch::torch_device("cpu")
)
output_buf <- get_attn_buffer(
  "output", collision_shape, torch::torch_float32(), torch::torch_device("cpu")
)
score_buf$fill_(1)
output_buf$fill_(2)
score_buf$add_(1)
expect_equal(as.numeric(output_buf$mean()), 2)
release_attn_buffers()
# --- Fused SDPA matches the chunked manual fallback ---------------------------
sdpa <- diffuseR:::.ltx23_sdpa
sdpa_manual <- diffuseR:::.ltx23_sdpa_manual
fused_sdpa <- diffuseR:::.ltx23_fused_sdpa
fused_sdpa_fn <- diffuseR:::.ltx23_fused_sdpa_fn
attn_scratch <- get(".ltx23_attn_scratch", envir = asNamespace("diffuseR"))

torch::torch_manual_seed(23)
q <- torch::torch_randn(c(1L, 2L, 5L, 8L))
k <- torch::torch_randn(c(1L, 2L, 7L, 8L))
v <- torch::torch_randn(c(1L, 2L, 7L, 6L))
mask <- torch::torch_zeros(c(1L, 1L, 5L, 7L))
mask[, , , 7] <- -10000

manual_mask <- sdpa_manual(q, k, v, mask, chunk_size = 2L)$clone()
manual_nomask <- sdpa_manual(q, k, v, chunk_size = 2L)$clone()
if (!is.null(fused_sdpa_fn())) {
  release_attn_buffers()
  fused_mask <- fused_sdpa(q, k, v, mask)
  fused_nomask <- fused_sdpa(q, k, v)
  expect_true(max_abs_diff(fused_mask, manual_mask) < 1e-5)
  expect_true(max_abs_diff(fused_nomask, manual_nomask) < 1e-5)

  old <- options(diffuseR.ltx23_fused_sdpa = TRUE)
  auto_mask <- sdpa(q, k, v, mask)
  options(old)
  expect_true(max_abs_diff(auto_mask, fused_mask) < 1e-6)
  expect_equal(length(ls(attn_scratch)), 0L)
}

release_attn_buffers()
old <- options(diffuseR.ltx23_fused_sdpa = FALSE)
fallback_mask <- sdpa(q, k, v, mask)$clone()
options(old)
expect_true(max_abs_diff(fallback_mask, manual_mask) < 1e-5)
expect_true(length(ls(attn_scratch)) > 0L)
release_attn_buffers()

# --- Gated attention with split RoPE ------------------------------------------

attn <- ltx23_attention(
  query_dim = 16L, heads = 2L, dim_head = 8L,
  rope_type = "split", apply_gated_attention = TRUE
)
attn$eval()
load_named_weights(attn, fixture_group("attn"))

rope <- list(fx$attn_cos, fx$attn_sin)
torch::with_no_grad({
  out <- attn(fx$attn_x, attention_mask = fx$attn_mask, query_rotary_emb = rope)
  out_nomask <- attn(fx$attn_x, query_rotary_emb = rope)
})
expect_true(max_abs_diff(out, fx$attn_out) < 1e-4)
expect_true(max_abs_diff(out_nomask, fx$attn_out_nomask) < 1e-4)

# Query chunking is numerically identical
attn$attn_chunk <- 2L
torch::with_no_grad({
  out_chunked <- attn(fx$attn_x, attention_mask = fx$attn_mask, query_rotary_emb = rope)
})
expect_true(max_abs_diff(out_chunked, fx$attn_out) < 1e-5)
attn$attn_chunk <- NULL

# --- AdaLN single (9 mod params) -----------------------------------------------

ada <- ltx23_ada_layer_norm_single(16L, num_mod_params = 9L)
ada$eval()
load_named_weights(ada, fixture_group("ada"))
torch::with_no_grad({
  ada_res <- ada(fx$ada_t, hidden_dtype = torch::torch_float32())
})
expect_true(max_abs_diff(ada_res[[1]], fx$ada_out) < 1e-4)
expect_true(max_abs_diff(ada_res[[2]], fx$ada_emb) < 1e-4)

# --- Tiny full model -------------------------------------------------------------

model <- ltx23_transformer(
  in_channels = 4L,
  out_channels = 4L,
  num_attention_heads = 2L,
  attention_head_dim = 8L,
  cross_attention_dim = 16L,
  audio_in_channels = 4L,
  audio_out_channels = 4L,
  audio_num_attention_heads = 2L,
  audio_attention_head_dim = 4L,
  audio_cross_attention_dim = 8L,
  num_layers = 2L,
  rope_type = "split"
)
model$eval()
load_named_weights(model, fixture_group("model"))

run_model <- function(...) {
  torch::with_no_grad({
    model(
      hidden_states = fx$m_hidden,
      audio_hidden_states = fx$m_audio_hidden,
      encoder_hidden_states = fx$m_enc,
      audio_encoder_hidden_states = fx$m_audio_enc,
      timestep = fx$m_timestep,
      sigma = fx$m_sigma,
      num_frames = 2L, height = 3L, width = 4L, fps = 24.0,
      audio_num_frames = 5L,
      use_cross_timestep = TRUE,
      ...
    )
  })
}

res <- run_model(
  encoder_attention_mask = fx$m_enc_mask,
  audio_encoder_attention_mask = fx$m_enc_mask
)
expect_equal(as.integer(res$sample$shape), as.integer(fx$m_out_video$shape))
expect_true(max_abs_diff(res$sample, fx$m_out_video) < 1e-4)
expect_true(max_abs_diff(res$audio_sample, fx$m_out_audio) < 1e-4)

# Isolated modalities (a2v/v2a off)
res_iso <- run_model(isolate_modalities = TRUE)
expect_true(max_abs_diff(res_iso$sample, fx$m_out_video_iso) < 1e-4)
expect_true(max_abs_diff(res_iso$audio_sample, fx$m_out_audio_iso) < 1e-4)

# STG perturbation on block index 1 (0-based, i.e. the second block)
res_stg <- run_model(spatio_temporal_guidance_blocks = c(1L))
expect_true(max_abs_diff(res_stg$sample, fx$m_out_video_stg) < 1e-4)
expect_true(max_abs_diff(res_stg$audio_sample, fx$m_out_audio_stg) < 1e-4)

# --- Key mapper ------------------------------------------------------------------

expect_equal(
  ltx23_map_dit_key("model.diffusion_model.patchify_proj.weight"),
  "proj_in.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.audio_patchify_proj.bias"),
  "audio_proj_in.bias"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.weight"),
  "time_embed.emb.timestep_embedder.linear_1.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.audio_adaln_single.linear.bias"),
  "audio_time_embed.linear.bias"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.prompt_adaln_single.linear.weight"),
  "prompt_adaln.linear.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.audio_prompt_adaln_single.linear.weight"),
  "audio_prompt_adaln.linear.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.av_ca_a2v_gate_adaln_single.linear.weight"),
  "av_cross_attn_video_a2v_gate.linear.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.transformer_blocks.0.attn1.q_norm.weight"),
  "transformer_blocks.0.attn1.norm_q.weight"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.transformer_blocks.5.scale_shift_table_a2v_ca_video"),
  "transformer_blocks.5.video_a2v_cross_attn_scale_shift_table"
)
expect_equal(
  ltx23_map_dit_key("model.diffusion_model.transformer_blocks.0.ff.net.0.proj.weight"),
  "transformer_blocks.0.ff.net.0.proj.weight"
)
expect_equal(ltx23_map_dit_key("model.diffusion_model.scale_shift_table"), "scale_shift_table")

# --- In-place tanh GELU (large-activation feed-forward path) ------------------------

ff <- ltx23_feed_forward(8L)
x_ff <- torch::torch_randn(1L, 37L, 8L)
torch::with_no_grad({
  ref_ff <- ff(x_ff)
  # Force the in-place chunked path (threshold below this tensor's numel,
  # chunk smaller than the intermediate so several chunks run)
  op <- options(diffuseR.gelu_inplace_min = 1)
  inp_ff <- ff(x_ff)
  options(op)
})
expect_true(as.numeric((ref_ff - inp_ff)$abs()$max()) < 1e-6)

h_g <- torch::torch_randn(3L, 41L)
g_ref <- torch::nnf_gelu(h_g, approximate = "tanh")
torch::with_no_grad(
  g_inp <- diffuseR:::.ltx23_gelu_tanh_inplace(h_g$clone(), chunk_elements = 16)
)
expect_true(as.numeric((g_ref - g_inp)$abs()$max()) < 1e-6)
