# End-to-end plumbing test for the LTX-2.3 pipeline with tiny
# random-weight components (no model downloads; CPU-safe).

if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
  exit_file("torch not fully installed")
}

library(diffuseR)

torch::torch_manual_seed(99)

# Tiny components wired to the pipeline's structural constants:
# 128 video latent channels, 8 audio channels x 16 latent mel bins.
transformer <- ltx23_transformer(
  in_channels = 128L, out_channels = 128L,
  num_attention_heads = 2L, attention_head_dim = 8L,
  cross_attention_dim = 16L,
  audio_in_channels = 128L, audio_out_channels = 128L,
  audio_num_attention_heads = 2L, audio_attention_head_dim = 4L,
  audio_cross_attention_dim = 8L,
  num_layers = 1L
)
transformer$eval()

connectors <- ltx23_text_connectors(
  caption_channels = 8L, text_proj_in_factor = 3L,
  video_connector_num_attention_heads = 2L,
  video_connector_attention_head_dim = 8L,
  video_connector_num_layers = 1L,
  video_connector_num_learnable_registers = 4L,
  audio_connector_num_attention_heads = 2L,
  audio_connector_attention_head_dim = 4L,
  audio_connector_num_layers = 1L,
  audio_connector_num_learnable_registers = 4L,
  video_hidden_dim = 16L, audio_hidden_dim = 8L
)
connectors$eval()

vae <- ltx23_video_vae(
  latent_channels = 128L,
  block_out_channels = c(8L, 8L, 8L, 8L),
  decoder_block_out_channels = c(4L, 8L, 8L, 16L),
  layers_per_block = c(1L, 1L, 1L, 1L, 1L),
  decoder_layers_per_block = c(1L, 1L, 1L, 1L, 1L)
)
vae$eval()

audio_vae <- ltx23_audio_vae(
  base_channels = 128L, ch_mult = c(1L, 1L, 1L), num_res_blocks = 1L,
  latent_channels = 8L, mel_bins = 64L
)
audio_vae$eval()

vocoder <- ltx23_vocoder_with_bwe(
  in_channels = 128L, hidden_channels = 16L,
  upsample_kernel_sizes = c(4L, 4L), upsample_factors = c(2L, 2L),
  resnet_kernel_sizes = c(3L), resnet_dilations = list(c(1L, 3L)),
  bwe_in_channels = 128L, bwe_hidden_channels = 8L,
  bwe_upsample_kernel_sizes = c(8L, 4L), bwe_upsample_factors = c(4L, 2L),
  bwe_resnet_kernel_sizes = c(3L), bwe_resnet_dilations = list(c(1L, 3L)),
  filter_length = 8L, hop_length = 2L, window_length = 8L,
  num_mel_channels = 64L,
  input_sampling_rate = 100L, output_sampling_rate = 400L
)
vocoder$eval()

pipe <- structure(
  list(
    transformer = transformer,
    connectors = connectors,
    vae = vae,
    audio_vae = audio_vae,
    vocoder = vocoder
  ),
  class = "ltx23_pipeline"
)

stub_embeds <- list(
  prompt_embeds = torch::torch_randn(1L, 8L, 8L, 3L),
  prompt_attention_mask = torch::torch_ones(1L, 8L)
)

res <- txt2vid_ltx2(
  prompt = "a tiny test",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L,
  device = "cpu",
  dtype = "float32",
  verbose = FALSE
)

# Video: [frames, height, width, 3] in [0, 1]
expect_equal(dim(res$video), c(9L, 64L, 64L, 3L))
expect_true(all(res$video >= 0 & res$video <= 1))
expect_true(all(is.finite(res$video)))

# Audio: [2, samples] in [-1, 1]; latent L = round(9/24*25) = 9 ->
# mel frames 4*9-3 = 33 (padded to hop multiple 34) -> vocoder x4 -> BWE x4
expect_equal(nrow(res$audio), 2L)
expect_true(ncol(res$audio) > 100L)
expect_true(all(abs(res$audio) <= 1))
expect_true(all(is.finite(res$audio)))
expect_equal(res$sample_rate, 48000L)

# Guardrails
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, guidance_scale = 3,
    device = "cpu", dtype = "float32"),
  pattern = "guidance_scale"
)
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, width = 50L,
    device = "cpu", dtype = "float32"),
  pattern = "multiples of 32"
)
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, num_frames = 10L,
    device = "cpu", dtype = "float32"),
  pattern = "8k"
)

# WAV writer round trip (header + size)
wav_path <- tempfile(fileext = ".wav")
write_wav(res$audio, wav_path, sample_rate = 48000L)
expect_true(file.exists(wav_path))
expect_equal(readBin(wav_path, "raw", 4), charToRaw("RIFF"))
expect_equal(file.size(wav_path), 44 + 2 * 2 * ncol(res$audio))
unlink(wav_path)

# MP4 mux when av is available
if (requireNamespace("av", quietly = TRUE)) {
  mp4 <- tempfile(fileext = ".mp4")
  save_video_ltx23(res$video, mp4, fps = 24, audio = res$audio,
    sample_rate = 48000L, verbose = FALSE)
  expect_true(file.exists(mp4) && file.size(mp4) > 0)
  unlink(mp4)
}

# --- Prefix conditioning smoke tests -------------------------------------------------

# i2v: start image conditions frame 0; pipeline runs end to end
start_img <- array(runif(64 * 64 * 3), dim = c(64L, 64L, 3L))
res_i2v <- txt2vid_ltx2(
  prompt = "tiny i2v",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  image = start_img,
  decode_audio = FALSE,
  verbose = FALSE
)
expect_equal(dim(res_i2v$video), c(9L, 64L, 64L, 3L))
expect_true(all(is.finite(res_i2v$video)))

# Same seed without conditioning gives a different video (mask engaged)
res_t2v <- txt2vid_ltx2(
  prompt = "tiny i2v",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  decode_audio = FALSE,
  verbose = FALSE
)
expect_true(max(abs(res_i2v$video - res_t2v$video)) > 1e-4)

# Continuation: 9-frame tail array as the frozen prefix
tail_arr <- array(runif(9 * 64 * 64 * 3), dim = c(9L, 64L, 64L, 3L))
res_cont <- txt2vid_ltx2(
  prompt = "tiny continuation",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  condition_video = tail_arr, conditioning_frames = 9L,
  decode_audio = FALSE,
  verbose = FALSE
)
expect_equal(dim(res_cont$video), c(17L, 64L, 64L, 3L))
expect_true(all(is.finite(res_cont$video)))

# Guardrails
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
    image = start_img, condition_video = tail_arr,
    device = "cpu", dtype = "float32"),
  pattern = "only one"
)

# --- Latent-space chaining seams -----------------------------------------------------

# latent_shape rides in the result: 17 frames @ 64x64 -> [3, 2, 2]
expect_equal(res_cont$latent_shape, c(3L, 2L, 2L))

# condition_latents reproduces the condition_video path exactly when
# fed the same encoded tail (the pixel path is just one producer)
enc <- ltx23_encode_video_frames(
  vae, ltx23_preprocess_frames(tail_arr, 64L, 64L))
res_lat <- txt2vid_ltx2(
  prompt = "tiny continuation",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  condition_latents = enc,
  decode_audio = FALSE,
  verbose = FALSE
)
expect_true(max(abs(res_lat$video - res_cont$video)) < 1e-6)

# Tail slicing: last 2 latent frames, matching a manual unpack + narrow
tail_lat <- ltx23_tail_latents(res_cont, k = 2L)
expect_equal(as.integer(tail_lat$shape), c(1L, 128L, 2L, 2L, 2L))
manual <- diffuseR:::ltx23_unpack_video_latents(res_cont$latents, 3L, 2L, 2L)$
  narrow(3L, 2L, 2L)
expect_equal(as.numeric((tail_lat - manual)$abs()$max()), 0)

# Feeding the sliced tail back runs end to end (chunk N -> chunk N+1)
res_next <- txt2vid_ltx2(
  prompt = "tiny chunk 3",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
  seed = 8L, device = "cpu", dtype = "float32",
  condition_latents = tail_lat,
  decode_audio = FALSE,
  verbose = FALSE
)
expect_equal(dim(res_next$video), c(17L, 64L, 64L, 3L))
expect_true(all(is.finite(res_next$video)))

# trim_frames delivers head-free pixels; latents keep the full sequence
res_trim <- txt2vid_ltx2(
  prompt = "tiny continuation",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  condition_video = tail_arr, trim_frames = 9L,
  decode_audio = FALSE,
  verbose = FALSE
)
expect_equal(dim(res_trim$video), c(8L, 64L, 64L, 3L))
expect_equal(max(abs(res_trim$video -
                     res_cont$video[-(1:9), , , , drop = FALSE])), 0)
expect_equal(res_trim$latent_shape, c(3L, 2L, 2L))
expect_equal(as.integer(ltx23_tail_latents(res_trim)$shape[3]), 2L)

# resident= is accepted (inert on CPU, where phase_offload is off)
res_res <- txt2vid_ltx2(
  prompt = "tiny resident",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  resident = "transformer",
  decode_audio = FALSE,
  verbose = FALSE
)
expect_true(all(is.finite(res_res$video)))

# Guardrails
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
    condition_latents = enc, image = start_img,
    device = "cpu", dtype = "float32"),
  pattern = "only one"
)
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
    condition_latents = torch::torch_randn(1L, 128L, 2L, 3L, 3L),
    device = "cpu", dtype = "float32"),
  pattern = "condition_latents"
)
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, num_frames = 9L,
    trim_frames = 9L, device = "cpu", dtype = "float32"),
  pattern = "trim_frames"
)
expect_error(
  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
    resident = "flux_capacitor", device = "cpu", dtype = "float32"),
  pattern = "resident"
)
expect_error(ltx23_tail_latents(res_cont$latents), pattern = "latent_shape")
expect_error(ltx23_tail_latents(res_cont, k = 5L), pattern = "k must be")

# --- connector_embeds bypass ---------------------------------------------------------

# Precomputed connector outputs reproduce the internal connectors path
# exactly, and satisfy the text requirement without prompt_embeds
conn_out <- torch::with_no_grad(connectors(
  stub_embeds$prompt_embeds, stub_embeds$prompt_attention_mask))
cemb <- list(video_text_embedding = conn_out$video_text_embedding,
             audio_text_embedding = conn_out$audio_text_embedding,
             attention_mask = conn_out$attention_mask)
res_ce <- txt2vid_ltx2(
  prompt = "tiny bypass",
  pipeline = pipe,
  connector_embeds = cemb,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  decode_audio = FALSE,
  verbose = FALSE
)
res_full <- txt2vid_ltx2(
  prompt = "tiny bypass",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  decode_audio = FALSE,
  verbose = FALSE
)
expect_equal(dim(res_ce$video), c(9L, 64L, 64L, 3L))
expect_true(max(abs(res_ce$video - res_full$video)) < 1e-6)


# --- Audio-conditioned generation (lip-sync plumbing) --------------------------------

wav_in <- matrix(runif(2L * 6000L, -0.5, 0.5), nrow = 2L)
res_audio <- txt2vid_ltx2(
  prompt = "tiny audio-driven",
  pipeline = pipe,
  prompt_embeds = stub_embeds,
  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
  seed = 7L, device = "cpu", dtype = "float32",
  image = start_img,
  audio = wav_in,
  verbose = FALSE
)
expect_equal(dim(res_audio$video), c(9L, 64L, 64L, 3L))
expect_true(all(is.finite(res_audio$video)))
# The original audio rides through untouched
expect_equal(res_audio$sample_rate, 16000L)
expect_equal(res_audio$audio, wav_in)
