#!/usr/bin/env bash
# Extract all 12 EasyCom sessions on host GPUs 3 and 4.
# Exactly one Qwen/Decord worker runs on each GPU; sessions are processed
# sequentially within each worker. The ordering balances the official 5-Hz
# target lengths across the two GPUs.
set -euo pipefail

PROJECT_ROOT="${PROJECT_ROOT:-/data/wangx/Code/icassp2027-experiment}"
DATASET_ROOT="${DATASET_ROOT:-/data/wangx/Datasets/easycom}"
MODEL="${MODEL:-/data/wangx/Models/Qwen2.5-Omni-3B}"

VIDEO_ROOT="${VIDEO_ROOT:-${DATASET_ROOT}/extracted_easycom/Main/Video_Compressed}"
TARGET_ROOT="${TARGET_ROOT:-${DATASET_ROOT}/easycom_target_frame}"
FEATURE_ROOT="${FEATURE_ROOT:-${DATASET_ROOT}/omni_features/open_stream_2fps_20s_decord}"
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_2x4090_decord}"

mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"

for gpu in 3 4; do
  nvidia-smi -i "${gpu}" --query-gpu=index,name,memory.total \
    --format=csv,noheader
done

for sid in {1..12}; do
  test -d "${VIDEO_ROOT}/Session_${sid}" || {
    echo "Missing video directory: ${VIDEO_ROOT}/Session_${sid}" >&2
    exit 1
  }
  test -s "${TARGET_ROOT}/Session_${sid}.npy" || {
    echo "Missing target: ${TARGET_ROOT}/Session_${sid}.npy" >&2
    exit 1
  }
done

echo "Features: ${FEATURE_ROOT}"
echo "Logs:     ${LOG_ROOT}"
echo "GPU 3: Session_12 Session_5 Session_9 Session_8 Session_10 Session_1"
echo "GPU 4: Session_6 Session_4 Session_3 Session_2 Session_7 Session_11"

exec nice -n 5 env \
  OMP_NUM_THREADS=4 \
  MKL_NUM_THREADS=4 \
  OPENBLAS_NUM_THREADS=4 \
  PROJECT_ROOT="${PROJECT_ROOT}" \
  MODEL="${MODEL}" \
  VIDEO_ROOT="${VIDEO_ROOT}" \
  TARGET_ROOT="${TARGET_ROOT}" \
  FEATURE_ROOT="${FEATURE_ROOT}" \
  LOG_ROOT="${LOG_ROOT}" \
  GPUS="3 4" \
  SESSIONS="12 6 5 4 9 3 8 2 10 7 1 11" \
  MEMORY_SECONDS=20 \
  NATIVE_BLOCK_SECONDS=2 \
  VIDEO_FPS=2 \
  FRAME_CACHE_SIZE=64 \
  MODE=open-stream-replay \
  DTYPE=bfloat16 \
  DECODER=decord \
  DECODER_GPU_ID=0 \
  bash "${PROJECT_ROOT}/run_extract_easycom_direct.sh"