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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing/bin/python}"
METRICS_PYTHON_BIN="${METRICS_PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/vbench_eval/bin/python}"
VBENCH_ROOT="${VBENCH_ROOT:-/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
EXPERIMENT_NAME="${EXPERIMENT_NAME:?EXPERIMENT_NAME is required}"
FULL_CHECKPOINT="${FULL_CHECKPOINT:?FULL_CHECKPOINT is required}"
PREDICTOR_CHECKPOINT="${PREDICTOR_CHECKPOINT:-}"
INCLUDE_FFFF="${INCLUDE_FFFF:-0}"
PROMPT_FILE="${PROMPT_FILE:-$ROOT_DIR/prompts/MovieGenVideoBench_extended_101_200.txt}"
NVME_ROOT="${NVME_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_dmd_validation}"
PROJECT_VIDEO_ROOT="${PROJECT_VIDEO_ROOT:-$ROOT_DIR/videos}"
RESULT_ROOT="${RESULT_ROOT:-$ROOT_DIR/evaluation_results/predictor_v4_dmd}"
REFERENCE_FFFF="${REFERENCE_FFFF:-}"

cd "$ROOT_DIR"
IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
NPROC_PER_NODE="${#GPU_ARRAY[@]}"
mkdir -p "$NVME_ROOT" "$PROJECT_VIDEO_ROOT" "$RESULT_ROOT"
for required in "$PYTHON_BIN" "$METRICS_PYTHON_BIN" "$FULL_CHECKPOINT" "$PROMPT_FILE"; do
  if [[ ! -e "$required" ]]; then
    echo "错误：找不到 $required" >&2
    exit 1
  fi
done
if [[ -n "$PREDICTOR_CHECKPOINT" && ! -f "$PREDICTOR_CHECKPOINT" ]]; then
  echo "错误：找不到 $PREDICTOR_CHECKPOINT" >&2
  exit 1
fi

variants=(fppf fppp)
if [[ "$INCLUDE_FFFF" == "1" ]]; then
  variants=(ffff fppf fppp)
fi

video_dir() {
  echo "$NVME_ROOT/${EXPERIMENT_NAME}_$1_seed0"
}

run_inference() {
  local variant="$1"
  local output
  output="$(video_dir "$variant")"
  if "$PYTHON_BIN" scripts/validate_generated_videos.py "$output" \
      >/dev/null 2>&1; then
    echo "[$(date -u +%FT%TZ)] 跳过已完成的 $variant 推理。"
    return
  fi
  echo "[$(date -u +%FT%TZ)] 启动 $EXPERIMENT_NAME $variant 验证集推理。"
  if [[ "$variant" == "ffff" ]]; then
    GPU_IDS="$GPU_IDS" \
    CHECKPOINT_PATH="$FULL_CHECKPOINT" \
    DATA_PATH="$PROMPT_FILE" \
    OUTPUT_FOLDER="$output" \
    bash scripts/run_moviegen_first100_dmd.sh
  else
    GPU_IDS="$GPU_IDS" \
    CHECKPOINT_PATH="$FULL_CHECKPOINT" \
    PREDICTOR_CHECKPOINT="$PREDICTOR_CHECKPOINT" \
    PREDICTOR_SCHEDULE="$variant" \
    DATA_PATH="$PROMPT_FILE" \
    OUTPUT_FOLDER="$output" \
    bash scripts/run_moviegen_first100_dmd.sh
  fi
  "$PYTHON_BIN" scripts/validate_generated_videos.py "$output"
  local project_output="$PROJECT_VIDEO_ROOT/${EXPERIMENT_NAME}_${variant}_seed0"
  mkdir -p "$project_output"
  rsync -a "$output/" "$project_output/"
}

run_vbench() {
  local variant="$1"
  local output_dir="$VBENCH_ROOT/evaluation_results/${EXPERIMENT_NAME}_${variant}_moviegen_101_200_seed0"
  if find "$output_dir" -maxdepth 1 -type f -name '*_eval_results.json' \
      -print -quit 2>/dev/null | grep -q .; then
    echo "[$(date -u +%FT%TZ)] 跳过已有的 $variant VBench。"
    return
  fi
  VIDEO_DIR="$(video_dir "$variant")" \
  OUTPUT_DIR="$output_dir" \
  GPU_IDS="$GPU_IDS" \
  bash scripts/run_vbench_custom_6_metrics.sh
}

for variant in "${variants[@]}"; do
  run_inference "$variant"
  run_vbench "$variant"
done

if [[ "$INCLUDE_FFFF" == "1" ]]; then
  pair_reference="$(video_dir ffff)"
else
  pair_reference="$REFERENCE_FFFF"
fi
if [[ -n "$pair_reference" ]]; then
  for variant in fppf fppp; do
    CUDA_VISIBLE_DEVICES="$GPU_IDS" \
    "$METRICS_PYTHON_BIN" -m torch.distributed.run \
      --standalone \
      --nproc_per_node="$NPROC_PER_NODE" \
      scripts/compare_video_metrics.py \
      --reference_dir "$pair_reference" \
      --comparison_dir "$(video_dir "$variant")" \
      --output_json "$RESULT_ROOT/${EXPERIMENT_NAME}_ffff_vs_${variant}.json"
  done
fi

touch "$RESULT_ROOT/${EXPERIMENT_NAME}.complete"
echo "[$(date -u +%FT%TZ)] $EXPERIMENT_NAME 验证集评估完成。"