#!/usr/bin/env bash
#
# Assemble GIFs from the PNG frames written by NeRF.cpp during/after training.
#
# Produces in <output_dir> (and optionally copies optimized versions to <dest_dir>):
#   rgb.gif                 - final 360° RGB orbit of the trained model
#   depth.gif               - final 360° depth orbit
#   training_progress.gif   - object rotating while sharpening across training
#                             iterations, labelled with iter number
#
# If the orbit frames are missing but a checkpoint is present, set RENDER_DATA to the
# scene directory and the frames are rendered from that checkpoint first. Set
# RENDER_SIZE to the --size the checkpoint was trained at (it sets the focal length).
#
# Usage: bash scripts/make_gifs.sh <output_dir> [n_final_frames] [max_iter] [plot_freq] [n_preview_frames] [dest_dir]
set -euo pipefail

OUT="${1:?usage: scripts/make_gifs.sh <output_dir> [n_final_frames] [max_iter] [plot_freq] [n_preview_frames] [dest_dir]}"
N_FINAL="${2:-30}"
MAX_ITER="${3:-6000}"
PLOT_FREQ="${4:-100}"
N_PREVIEW="${5:-5}"
DEST="${6:-}"
RENDER_DATA="${RENDER_DATA:-}"
RENDER_SIZE="${RENDER_SIZE:-160}"
NERF_BIN="${NERF_BIN:-./build/NeRF.cpp}"

# ImageMagick v7 uses 'magick'; fall back to 'convert' on v6.
if command -v magick >/dev/null 2>&1; then IM="magick"; else IM="convert"; fi

# Re-render the orbit from a checkpoint when the frames are absent.
if [ ! -f "$OUT/frame_final_0.png" ] && [ -f "$OUT/checkpoint.pt" ]; then
  if [ -n "$RENDER_DATA" ]; then
    echo "[..] no orbit frames; rendering $N_FINAL from $OUT/checkpoint.pt"
    "$NERF_BIN" "$RENDER_DATA" "$OUT" --render "$OUT/checkpoint.pt" \
      --final-frames "$N_FINAL" --size "$RENDER_SIZE"
  else
    echo "[WARN] no frame_final_*.png, but $OUT/checkpoint.pt exists."
    echo "       Set RENDER_DATA=<scene dir> to render the orbit from it first."
  fi
fi

cd "$OUT"

# --- 1. Final RGB orbit -----------------------------------------------------
rgb_frames=()
for ((i=0; i<N_FINAL; i++)); do
  f="frame_final_${i}.png"
  [ -f "$f" ] && rgb_frames+=("$f")
done
if [ "${#rgb_frames[@]}" -gt 0 ]; then
  $IM -delay 6 -loop 0 "${rgb_frames[@]}" rgb.gif
  echo "[OK] wrote rgb.gif (${#rgb_frames[@]} frames)"
else
  echo "[WARN] no frame_final_*.png found; skipping rgb.gif"
fi

# --- 2. Final depth orbit ---------------------------------------------------
depth_frames=()
for ((i=0; i<N_FINAL; i++)); do
  f="frame_depth_final_${i}.png"
  [ -f "$f" ] && depth_frames+=("$f")
done
if [ "${#depth_frames[@]}" -gt 0 ]; then
  $IM -delay 6 -loop 0 "${depth_frames[@]}" depth.gif
  echo "[OK] wrote depth.gif (${#depth_frames[@]} frames)"
else
  echo "[WARN] no frame_depth_*final*.png found; skipping depth.gif"
fi

# --- 3. Training-progress orbit --------------------------------------------
# Step through saved iterations; rotate azimuth per stage so the object spins
# while sharpening. Annotate each frame with the iteration number.
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
idx=0
prog_frames=()
for ((it=0; it<=MAX_ITER; it+=PLOT_FREQ)); do
  az=$(( (it / PLOT_FREQ) % N_PREVIEW ))
  src="frame_${it}_${az}.png"
  [ -f "$src" ] || continue
  dst="$tmp/$(printf '%05d' "$idx").png"
  $IM "$src" -gravity NorthWest -pointsize 18 -fill white \
      -stroke black -strokewidth 2 -annotate +8+6 "iter ${it}" \
      -stroke none -fill white -annotate +8+6 "iter ${it}" "$dst"
  prog_frames+=("$dst")
  idx=$((idx+1))
done
if [ "${#prog_frames[@]}" -gt 0 ]; then
  $IM -delay 10 -loop 0 "${prog_frames[@]}" -delay 150 "${prog_frames[-1]}" \
      training_progress.gif
  echo "[OK] wrote training_progress.gif (${#prog_frames[@]} stages)"
else
  echo "[WARN] no frame_<iter>_<az>.png found; skipping training_progress.gif"
fi

optimize_copy() {
  local name="$1"
  [ -f "$name.gif" ] || return 0
  if [ -n "$DEST" ] && command -v ffmpeg >/dev/null 2>&1; then
    mkdir -p "$DEST"
    ffmpeg -y -loglevel error -i "$name.gif" \
      -vf "fps=10,scale=200:-1:flags=lanczos,split[s0][s1];[s0]palettegen=max_colors=128:stats_mode=diff[p];[s1][p]paletteuse=dither=bayer:bayer_scale=3" \
      "$DEST/$name.gif"
    echo "[OK] optimized copy -> $DEST/$name.gif"
  elif [ -n "$DEST" ]; then
    mkdir -p "$DEST"
    cp "$name.gif" "$DEST/$name.gif"
    echo "[OK] copied -> $DEST/$name.gif (install ffmpeg for palette optimization)"
  fi
}

optimize_copy rgb
optimize_copy depth
optimize_copy training_progress

echo "Done. GIFs in: $OUT${DEST:+ (optimized copies in $DEST)}"