#!/usr/bin/env bash
# Re-runnable pipeline: COLMAP output → Nerfstudio → splatfacto → .ply
# Skips COLMAP (assumes my_scene/sparse/0/ already exists).
set -euo pipefail
cd "$(dirname "$0")"

source venv/bin/activate
# Use the full Homebrew ffmpeg (nerfstudio's bundled one lacks split/fps filters)
export PATH="/opt/homebrew/opt/ffmpeg/bin:$PATH"

SCENE=my_scene
NS_DATA=$SCENE/ns_data
EXPORT_DIR=$SCENE/exports
PLY=$EXPORT_DIR/splat.ply

# ── 1. Verify COLMAP output ────────────────────────────────────────────────
echo ""
echo "=== Step 1: Verifying COLMAP output ==="

if [ ! -d "$SCENE/sparse" ] || [ -z "$(ls -A $SCENE/sparse 2>/dev/null)" ]; then
  echo "ERROR: $SCENE/sparse/ not found or empty. Run main.py + match_crossvideo.py first."
  exit 1
fi

# Pick the model with the most registered images
BEST_MODEL=$(python3 -c "
import struct, os, sys
best_dir, best_imgs = '', 0
for m in sorted(os.listdir('$SCENE/sparse')):
    d = '$SCENE/sparse/' + m
    f = d + '/images.bin'
    if not os.path.isfile(f): continue
    with open(f,'rb') as fh: n = struct.unpack('<Q', fh.read(8))[0]
    if n > best_imgs: best_imgs, best_dir = n, d
print(best_dir)
")

if [ -z "$BEST_MODEL" ]; then
  echo "ERROR: no valid COLMAP model found in $SCENE/sparse/"
  exit 1
fi

for f in cameras.bin images.bin points3D.bin; do
  if [ ! -f "$BEST_MODEL/$f" ]; then
    echo "ERROR: missing $BEST_MODEL/$f"
    exit 1
  fi
done

NUM_IMGS=$(python3 -c "import struct; f=open('$BEST_MODEL/images.bin','rb'); print(struct.unpack('<Q',f.read(8))[0])")
NUM_PTS=$(python3 -c "import struct; f=open('$BEST_MODEL/points3D.bin','rb'); print(struct.unpack('<Q',f.read(8))[0])")
echo "  Best model: $BEST_MODEL (images=$NUM_IMGS  points3D=$NUM_PTS)"

num_models=$(find "$SCENE/sparse" -mindepth 1 -maxdepth 1 -type d | wc -l | tr -d ' ')
if [ "$num_models" -gt 1 ]; then
  echo "  WARNING: $num_models disconnected models — using largest ($BEST_MODEL)."
  echo "           Run match_crossvideo.py and re-map to attempt a full stitch."
fi

# ── 2. Convert COLMAP → Nerfstudio format ─────────────────────────────────
echo ""
echo "=== Step 2: ns-process-data (COLMAP → Nerfstudio) ==="
ns-process-data images \
  --data "$(pwd)/$SCENE/images" \
  --output-dir "$(pwd)/$NS_DATA" \
  --skip-colmap \
  --colmap-model-path "$(pwd)/$BEST_MODEL"

# ── 3. Train splatfacto with browser viewer ────────────────────────────────
echo ""
echo "=== Step 3: Training splatfacto ==="
echo ""
echo "  ┌──────────────────────────────────────────────────────┐"
echo "  │  Live viewer (fly around during training):           │"
echo "  │  http://localhost:7007                               │"
echo "  └──────────────────────────────────────────────────────┘"
echo ""
ns-train splatfacto \
  --data "$NS_DATA" \
  --vis viewer \
  --viewer.quit-on-train-completion True

# ── 4. Find the latest training output ────────────────────────────────────
TRAIN_OUT=$(ls -td outputs/*/splatfacto/*/  2>/dev/null | head -1)
if [ -z "$TRAIN_OUT" ]; then
  echo "ERROR: could not find training output folder under outputs/"
  exit 1
fi
CONFIG_PATH="$TRAIN_OUT/config.yml"
echo "  Training output: $TRAIN_OUT"

# ── 5. Export .ply ────────────────────────────────────────────────────────
echo ""
echo "=== Step 4: Exporting Gaussian splat to .ply ==="
mkdir -p "$EXPORT_DIR"
ns-export gaussian-splat \
  --load-config "$CONFIG_PATH" \
  --output-dir "$EXPORT_DIR"

# ── 6. Final summary ──────────────────────────────────────────────────────
echo ""
echo "======================================================================"
if [ -f "$PLY" ]; then
  echo "  .ply exported: $(pwd)/$PLY"
else
  echo "  WARNING: splat.ply not found at $PLY — check $EXPORT_DIR/"
fi
echo ""
echo "  View during training : http://localhost:7007"
echo ""
echo "  View final .ply (Option A — recommended):"
echo "    Drag $(pwd)/$PLY into:"
echo "    https://playcanvas.com/supersplat/editor"
echo "    Runs 100% in-browser; the file stays on your machine."
echo ""
echo "  View final .ply (Option B — fully offline):"
echo "    python3 -m http.server 8080 --directory \$(dirname $PLY)"
echo "    Then open http://localhost:8080/splat.ply in gsplat viewer"
echo "    (requires a separate gsplat.js page — Option A is simpler)"
echo "======================================================================"