#!/usr/bin/env bash
set -euo pipefail

# Render eval PNG frames from final phase2 checkpoints for the four Waymo scenes.
# Mapping:
#   GPU 0 -> scene 103
#   GPU 1 -> scene 016
#   GPU 2 -> scene 013
#   GPU 3 -> scene 002

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

CONDA_ENV="${CONDA_ENV:-s3gs}"
CONDA_BIN="${CONDA_BIN:-}"
if [[ -z "$CONDA_BIN" ]]; then
  if command -v conda >/dev/null 2>&1; then
    CONDA_BIN="$(command -v conda)"
  elif [[ -x /workspace/miniconda3/bin/conda ]]; then
    CONDA_BIN="/workspace/miniconda3/bin/conda"
  elif [[ -x /opt/miniforge3/bin/conda ]]; then
    CONDA_BIN="/opt/miniforge3/bin/conda"
  else
    echo "ERROR: conda executable not found. Set CONDA_BIN=/path/to/conda" >&2
    exit 1
  fi
fi

mkdir -p logs/render_waymo4_eval_images

echo "[$(date '+%F %T')] root: $ROOT_DIR"
echo "[$(date '+%F %T')] conda env: $CONDA_ENV"
echo "[$(date '+%F %T')] conda bin: $CONDA_BIN"

run_one() {
  local gpu="$1"
  local sid="$2"
  local sid3
  sid3=$(printf "%03d" "$((10#$sid))")

  local src="data/processed/waymo_eccv_v1_10/training/${sid3}"
  local model_path="work_dirs/waymo_eccv_v1_10/nvs/p2/${sid3}"
  local ckpt="${model_path}/chkpnt_fine_50000.pth"
  local log="logs/render_waymo4_eval_images/${sid3}_phase2_50000_gpu${gpu}.log"

  if [[ ! -f "${src}/frame_info.json" ]]; then
    echo "[GPU ${gpu}] ERROR: missing processed scene: ${src}" >&2
    return 1
  fi
  if [[ ! -f "${ckpt}" ]]; then
    echo "[GPU ${gpu}] ERROR: missing checkpoint: ${ckpt}" >&2
    return 1
  fi

  : > "$log"
  echo "[$(date '+%F %T')] [GPU ${gpu}] render eval images scene ${sid3} ckpt ${ckpt}" | tee -a "$log"

  "$CONDA_BIN" run --no-capture-output -n "$CONDA_ENV" env \
    PYTHONPATH="$PWD" \
    CUDA_VISIBLE_DEVICES="${gpu}" \
    python train.py \
      -s "${src}" \
      --expname waymo \
      --configs arguments/stage2_nvs.py \
      --model_path "${model_path}" \
      --start_checkpoint "${ckpt}" \
      --eval_only \
      --eval_save_images \
    >> "$log" 2>&1

  echo "[$(date '+%F %T')] [GPU ${gpu}] done scene ${sid3}" | tee -a "$log"
}

declare -a PIDS=()
declare -a NAMES=()

cleanup() {
  echo "[$(date '+%F %T')] interrupt received; stopping child jobs..." >&2
  for pid in "${PIDS[@]:-}"; do
    kill "$pid" 2>/dev/null || true
  done
}
trap cleanup INT TERM

launch_one() {
  local gpu="$1"
  local sid="$2"
  run_one "$gpu" "$sid" &
  PIDS+=("$!")
  NAMES+=("scene_${sid}_gpu_${gpu}")
}

launch_one 0 103
launch_one 1 016
launch_one 2 013
launch_one 3 002

status=0
for idx in "${!PIDS[@]}"; do
  pid="${PIDS[$idx]}"
  name="${NAMES[$idx]}"
  if wait "$pid"; then
    echo "[$(date '+%F %T')] ${name} finished successfully"
  else
    rc=$?
    echo "[$(date '+%F %T')] ${name} failed with exit code ${rc}" >&2
    status=$rc
  fi
done

if [[ "$status" -eq 0 ]]; then
  echo "[$(date '+%F %T')] all 4 eval-image render jobs finished"
else
  echo "[$(date '+%F %T')] one or more eval-image render jobs failed" >&2
fi
exit "$status"