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

SELF_FORCING_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VBENCH_ROOT="${VBENCH_ROOT:-/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench}"
PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/vbench_eval/bin/python}"
VIDEO_DIR="${VIDEO_DIR:-$SELF_FORCING_ROOT/videos/moviegenvideobench_extended_first100_dmd_seed0}"
OUTPUT_DIR="${OUTPUT_DIR:-$VBENCH_ROOT/evaluation_results/self_forcing_dmd_moviegen_first100_seed0}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
DRY_RUN="${DRY_RUN:-0}"

IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
NPROC_PER_NODE="${NPROC_PER_NODE:-${#GPU_ARRAY[@]}}"

if [[ ! -d "$VBENCH_ROOT" || ! -f "$VBENCH_ROOT/evaluate.py" ]]; then
  echo "错误：VBench 仓库无效：$VBENCH_ROOT" >&2
  exit 1
fi
if [[ ! -x "$PYTHON_BIN" ]]; then
  echo "错误：VBench Python 不可执行：$PYTHON_BIN" >&2
  exit 1
fi
if [[ ! -d "$VIDEO_DIR" && "$DRY_RUN" != "1" ]]; then
  echo "错误：视频目录不存在：$VIDEO_DIR" >&2
  exit 1
fi
if (( NPROC_PER_NODE != ${#GPU_ARRAY[@]} )); then
  echo "错误：NPROC_PER_NODE 与 GPU_IDS 的 GPU 数量不一致。" >&2
  exit 1
fi

if [[ "$DRY_RUN" != "1" ]]; then
  VIDEO_COUNT="$(find "$VIDEO_DIR" -maxdepth 1 -type f -name '*.mp4' | wc -l)"
  if (( VIDEO_COUNT != 100 )); then
    echo "错误：VBench 评测需要 100 个视频，当前找到 $VIDEO_COUNT 个。" >&2
    exit 1
  fi
fi

ARGS=(
  evaluate.py
  --videos_path "$VIDEO_DIR"
  --output_path "$OUTPUT_DIR"
  --mode custom_input
  --load_ckpt_from_local True
  --dimension
  subject_consistency
  background_consistency
  motion_smoothness
  dynamic_degree
  aesthetic_quality
  imaging_quality
)

if (( NPROC_PER_NODE == 1 )); then
  COMMAND=("$PYTHON_BIN" "${ARGS[@]}")
else
  COMMAND=(
    "$PYTHON_BIN" -m torch.distributed.run
    --standalone
    --nproc_per_node="$NPROC_PER_NODE"
    "${ARGS[@]}"
  )
fi

echo "VBench custom_input 六指标评测："
echo "  GPUs:   $GPU_IDS"
echo "  Videos: $VIDEO_DIR"
echo "  Output: $OUTPUT_DIR"
printf '  Command: '
printf '%q ' "${COMMAND[@]}"
printf '\n'

if [[ "$DRY_RUN" == "1" ]]; then
  exit 0
fi

mkdir -p "$OUTPUT_DIR"
cd "$VBENCH_ROOT"
CUDA_VISIBLE_DEVICES="$GPU_IDS" "${COMMAND[@]}" 2>&1 | tee "$OUTPUT_DIR/evaluation.log"