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

MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}"
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
PREP_ROOT="${PREP_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_prepared}"
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_5fps}"
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/causal_prefix_5fps}"

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

run_worker() {
  local gpu="$1"
  shift

  local sid session metadata target output log
  for sid in "$@"; do
    session="Session_${sid}"
    metadata="${PREP_ROOT}/${session}/${session}.json"
    target="${TARGET_ROOT}/${session}.npy"
    output="${FEATURE_ROOT}/${session}.pt"
    log="${LOG_ROOT}/${session}.log"

    if [[ ! -s "${metadata}" ]]; then
      echo "Missing metadata: ${metadata}" >&2
      return 1
    fi
    if [[ ! -s "${target}" ]]; then
      echo "Missing target: ${target}" >&2
      return 1
    fi
    if [[ -s "${output}" ]]; then
      echo "Skipping existing output: ${output}"
      continue
    fi

    {
      echo "$(date): GPU ${gpu} starting ${session}"

      CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
        python streaming_turn_model.py \
          --model-id "${MODEL}" \
          --session-metadata "${metadata}" \
          --target "${target}" \
          --use-audio-in-video \
          --mode causal-prefix \
          --memory-seconds 20 \
          --video-fps 5 \
          --strict-fixed-video-grid \
          --flash-attention-2 \
          --output "${output}"

      echo "$(date): GPU ${gpu} completed ${session}"
    } >"${log}" 2>&1
  done
}

pids=()

# Two independent extraction workers per GPU. Each GPU receives approximately
# the same total number of EasyCom timesteps.
run_worker 0 12 11 </dev/null &
pids+=("$!")
run_worker 0 7 </dev/null &
pids+=("$!")

run_worker 1 6 10 </dev/null &
pids+=("$!")
run_worker 1 8 </dev/null &
pids+=("$!")

run_worker 2 4 1 </dev/null &
pids+=("$!")
run_worker 2 9 </dev/null &
pids+=("$!")

run_worker 3 5 2 </dev/null &
pids+=("$!")
run_worker 3 3 </dev/null &
pids+=("$!")

status=0
for pid in "${pids[@]}"; do
  wait "${pid}" || status=1
done

if [[ "${status}" -ne 0 ]]; then
  echo "At least one extraction worker failed. Check ${LOG_ROOT}." >&2
  exit 1
fi

echo "ALL 12 CAUSAL 5 FPS EXTRACTIONS COMPLETED"