# FLUX checkpoint opening + quantization round trip on the tiny sharded
# diffusers checkpoint (fixture generated by tools/gen_fixtures_flux_model.py,
# checked in). Everything runs on CPU.

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)

ckpt_dir <- system.file("tinytest", "fixtures", "flux_tiny_ckpt",
  package = "diffuseR")
if (ckpt_dir == "") ckpt_dir <- "fixtures/flux_tiny_ckpt"
if (!dir.exists(ckpt_dir)) exit_file("flux tiny checkpoint missing")

fixture_path <- system.file("tinytest", "fixtures", "flux_model.safetensors",
  package = "diffuseR")
if (fixture_path == "") fixture_path <- "fixtures/flux_model.safetensors"
if (!file.exists(fixture_path)) exit_file("flux 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())
  )))
}
cosine_sim <- function(a, b) {
  a <- a$to(dtype = torch::torch_float32())$flatten()
  b <- b$to(dtype = torch::torch_float32())$flatten()
  as.numeric(torch::torch_dot(a, b) / (a$norm() * b$norm()))
}

# --- cast-set census ------------------------------------------------------------
# Tiny config: 1 double block x 14 + 1 single block x 6 = 20 cast weights.
# (Full schnell: 19 x 14 + 38 x 6 = 494.)

ckpt <- flux_open_checkpoint(ckpt_dir)
expect_true(inherits(ckpt, "ltx23_checkpoint"))
expect_equal(sum(flux_is_quant_key(ckpt$keys)), 20L)
expect_equal(ckpt$config$num_layers, 1L)

# The sharded index resolves every key to a readable tensor
t0 <- ckpt$handle$get_tensor("x_embedder.weight")
expect_equal(as.integer(t0$shape), c(16L, 4L))

# --- full-precision load through the sharded index --------------------------------

model <- flux_load_transformer(ckpt, device = "cpu", dtype = "float32",
  verbose = FALSE)
ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L)
rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 4L))
out_full <- torch::with_no_grad(model(
  hidden_states = fx$hidden,
  encoder_hidden_states = fx$encoder,
  pooled_projections = fx$pooled,
  timestep = fx$timestep,
  image_rotary_emb = rope
))
expect_true(max_abs_diff(out_full, fx$out) < 1e-4)

# --- NF4 round trip -----------------------------------------------------------------

nf4_dir <- file.path(tempdir(), "flux-tiny-nf4")
unlink(nf4_dir, recursive = TRUE)
manifest <- flux_quantize(ckpt_dir, output_dir = nf4_dir, format = "nf4",
  verbose = FALSE)
expect_equal(manifest$cast, 20L)
expect_true(file.exists(file.path(nf4_dir, "manifest.json")))

nf4_ckpt <- flux_open_quantized(nf4_dir)
expect_equal(nf4_ckpt$format, "nf4")
model_nf4 <- flux_load_transformer(nf4_ckpt, device = "cpu", verbose = FALSE)
out_nf4 <- torch::with_no_grad(model_nf4(
  hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()),
  encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()),
  pooled_projections = fx$pooled$to(dtype = torch::torch_bfloat16()),
  timestep = fx$timestep,
  image_rotary_emb = rope
))
expect_true(all(is.finite(as.numeric(out_nf4$to(dtype = torch::torch_float32())))))
expect_true(cosine_sim(out_nf4, out_full) > 0.98)
ltx23_release_dequant_buffers()

# --- fp8 round trip (needs an F8-capable safetensors build) ---------------------------

# Ask the same gate flux_quantize() asks, rather than re-probing here: a
# private probe can disagree with .st_can_write() whenever the
# diffuseR.st_caps override is set, and then this block runs only to hit
# the hard error the override was meant to simulate.
f8_ok <- diffuseR:::.st_can_write("float8_e4m3fn")

if (f8_ok) {
  fp8_dir <- file.path(tempdir(), "flux-tiny-fp8")
  unlink(fp8_dir, recursive = TRUE)
  manifest8 <- flux_quantize(ckpt_dir, output_dir = fp8_dir, format = "fp8",
    verbose = FALSE)
  expect_equal(manifest8$cast, 20L)

  fp8_ckpt <- flux_open_quantized(fp8_dir)
  model_fp8 <- flux_load_transformer(fp8_ckpt, device = "cpu", pin = FALSE,
    verbose = FALSE)
  out_fp8 <- torch::with_no_grad(model_fp8(
    hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()),
    encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()),
    pooled_projections = fx$pooled$to(dtype = torch::torch_bfloat16()),
    timestep = fx$timestep,
    image_rotary_emb = rope
  ))
  expect_true(cosine_sim(out_fp8, out_full) > 0.99)
}

options(diffuseR.block_gc = NULL)
