#!/usr/bin/env bash
# Box B (audio) — dataset preprocess + flatten, so the trainer can render.
#
# WHY THIS IS NEEDED FOR PURE INFERENCE: ltx_a2a_generate*.py drives
# scripts/train.py with optimization.steps=1, because validation_runner IS the
# a2a inference path. train.py builds a train dataloader before validation
# fires, so preprocessed latents must exist even when we only want to render.
#
# SUBSET is the fast path. Preprocessing all 266 pairs runs the audio VAE and
# Gemma text encoder over every row and sits on the critical path behind the
# 66 GB weight pull. The dataloader only needs to be non-empty, so a handful of
# pairs unblocks rendering in under a minute; run the full pass afterwards, in
# the background, only if we are retraining.
#
# Usage:
#   box_b_trainer_setup.sh 8      # subset of 8 train pairs — fast, render-only
#   box_b_trainer_setup.sh full   # all 266 pairs — required before any retrain
set -uo pipefail
export PATH="$HOME/.local/bin:$PATH" TMPDIR=/workspace/.tmp UV_CACHE_DIR=/workspace/.uv-cache

MODE="${1:-8}"
ASC=/workspace/Demos/LTX/acoustic-space-control
M=/workspace/models/ltx-2.5
DATA=/workspace/Demos/data/acoustic-space-v5
TOOLS=/workspace/akuspace-tools

for f in "$M/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors" \
         "$M/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" \
         "$M/vae/ltx-2.5-video-vae-bf16.safetensors" \
         "$M/vae/ltx-2.5-audio-vae-bf16.safetensors"; do
  [ -s "$f" ] || { echo "WEIGHTS NOT READY: $f"; exit 2; }
done
[ -d "$ASC/training" ] || { echo "MAC PUSH NOT LANDED: $ASC/training missing"; exit 2; }

cd "$ASC"
SRC=training/ableton-assets-grid-v5.csv
if [ "$MODE" = "full" ]; then
  MAN="$SRC"; echo "=== FULL preprocess (266 pairs) ==="
else
  MAN=training/_subset_${MODE}.csv
  head -1 "$SRC" > "$MAN"
  awk -F, 'NR>1 && $2=="train"' "$SRC" | head -"$MODE" >> "$MAN"
  echo "=== SUBSET preprocess ($(($(wc -l < "$MAN")-1)) train pairs) — render-only fast path ==="
fi

[ -d "$TOOLS" ] || uv venv "$TOOLS" --python 3.12 >/dev/null 2>&1
uv pip install --python "$TOOLS/bin/python" -q -r training/requirements.txt 2>&1 | tail -2

"$TOOLS/bin/python" training/prepare_dataset.py \
  --manifest "$MAN" --output "$DATA" --mode copy 2>&1 | tail -3
echo "PREPARE_EXIT=$?"

cd /workspace/LTX-2.5-repo/packages/ltx-trainer
/workspace/LTX-2.5-repo/.venv/bin/python scripts/process_dataset.py \
  "$DATA/dataset_train.json" \
  --audio-durations 6.0 \
  --model-path "$M/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors" \
  --text-encoder-path "$M/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" \
  --video-vae-path "$M/vae/ltx-2.5-video-vae-bf16.safetensors" \
  --audio-vae-path "$M/vae/ltx-2.5-audio-vae-bf16.safetensors" \
  --output-dir "$DATA/.precomputed" \
  --lora-trigger AKUSPACE 2>&1 | tail -5
echo "PROCESS_EXIT=$?"

# MANDATORY FLATTEN — process_dataset mirrors the dataset-JSON relative paths,
# but datasets.py::_discover_samples requires one identical rel_path across all
# three sources, so nested targets/references never match. Silent: the dirs are
# populated and the run still reports "No valid samples found".
PRE="$DATA/.precomputed"
[ -d "$PRE/audio_latents/audio/targets" ] && { mv "$PRE/audio_latents/audio/targets/"* "$PRE/audio_latents/"; rm -rf "$PRE/audio_latents/audio"; }
[ -d "$PRE/reference_audio_latents/audio/references" ] && { mv "$PRE/reference_audio_latents/audio/references/"* "$PRE/reference_audio_latents/"; rm -rf "$PRE/reference_audio_latents/audio"; }
[ -d "$PRE/conditions/audio/targets" ] && { mv "$PRE/conditions/audio/targets/"* "$PRE/conditions/"; rm -rf "$PRE/conditions/audio"; }

A=$(find "$PRE/audio_latents" -maxdepth 1 -name '*.pt' | wc -l)
R=$(find "$PRE/reference_audio_latents" -maxdepth 1 -name '*.pt' | wc -l)
C=$(find "$PRE/conditions" -maxdepth 1 -name '*.pt' | wc -l)
echo "flatten: audio=$A reference=$R conditions=$C (must be equal and non-zero)"
[ "$A" -gt 0 ] && [ "$A" -eq "$R" ] && echo "TRAINER_SETUP_DONE" || echo "TRAINER_SETUP_MISMATCH"