#!/usr/bin/env bash
# Extract open-stream Thinker states for all 12 EasyCom sessions via
# streaming_turn_model_directly.py (original clips, no prepared 2 FPS MP4).
#
# Recommended for full sessions: --memory-seconds 20 (bounded KV).
# Do not default to --memory-seconds 0 on multi-minute sessions.
set -euo pipefail

PROJECT_ROOT="${PROJECT_ROOT:-/mnt/58462F4C462F29E8/Code/icassp2027-experiment}"
MODEL="${MODEL:-/mnt/58462F4C462F29E8/Models/Qwen2.5-Omni-3B}"
VIDEO_ROOT="${VIDEO_ROOT:-/mnt/58462F4C462F29E8/Datasets/easycom/Main/Video_Compressed}"
TARGET_ROOT="${TARGET_ROOT:-/mnt/58462F4C462F29E8/Code/EgoSpeak/dataset/EasyCom/target_perframe}"
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"

MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}"
VIDEO_FPS="${VIDEO_FPS:-2}"
READOUT_LAST_N_LAYERS="${READOUT_LAST_N_LAYERS:-1}"
MODE="${MODE:-open-stream-replay}"
DTYPE="${DTYPE:-bfloat16}"
DECODER="${DECODER:-decord}"
DECODER_GPU_ID="${DECODER_GPU_ID:-0}"
FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}"
# Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time.
GPUS="${GPUS:-0}"
# Optional: only run selected session ids, e.g. SESSIONS="1 2 3"
SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}"

# StreamConfig requires memory / native_block to be an integer (e.g. 20/2=10).
# Values like 15 yield 7.5 and crash before loading the model.
python - "${MEMORY_SECONDS}" "${NATIVE_BLOCK_SECONDS}" <<'PY'
import sys
memory, block = float(sys.argv[1]), float(sys.argv[2])
ratio = memory / block
if abs(ratio - round(ratio)) > 1e-6:
    raise SystemExit(
        f"MEMORY_SECONDS={memory} is invalid: memory/native_block={ratio} "
        f"(native_block={block}). Use a multiple of {block:g}, e.g. 16, 18, 20."
    )
print(f"memory_seconds={memory:g} native_block={block:g} blocks={int(round(ratio))}")
PY

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

run_one() {
  local gpu="$1"
  local sid="$2"
  local session="Session_${sid}"
  local session_dir="${VIDEO_ROOT}/${session}"
  local target="${TARGET_ROOT}/${session}.npy"
  local output="${FEATURE_ROOT}/${session}.pt"
  local log="${LOG_ROOT}/${session}.log"

  if [[ ! -d "${session_dir}" ]]; then
    echo "Missing session dir: ${session_dir}" >&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}"
    return 0
  fi

  {
    echo "$(date): GPU ${gpu} starting ${session}"
    echo "model=${MODEL}"
    echo "session_dir=${session_dir}"
    echo "target=${target}"
    echo "output=${output}"
    echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
    echo "readout_last_n_layers=${READOUT_LAST_N_LAYERS}"
    echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}"

    if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
      python streaming_turn_model_directly.py \
        --model-id "${MODEL}" \
        --easycom-session-dir "${session_dir}" \
        --target "${target}" \
        --use-audio-in-video \
        --mode "${MODE}" \
        --memory-seconds "${MEMORY_SECONDS}" \
        --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
        --video-fps "${VIDEO_FPS}" \
        --readout-last-n-layers "${READOUT_LAST_N_LAYERS}" \
        --frame-cache-size "${FRAME_CACHE_SIZE}" \
        --decoder "${DECODER}" \
        --decoder-gpu-id "${DECODER_GPU_ID}" \
        --dtype "${DTYPE}" \
        --flash-attention-2 \
        --device cuda:0 \
        --output "${output}"
    then
      echo "$(date): GPU ${gpu} FAILED ${session}"
      exit 1
    fi

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

# Round-robin sessions across GPUs.
mapfile -t GPU_ARR < <(echo "${GPUS}" | tr ' ' '\n' | sed '/^$/d')
mapfile -t SID_ARR < <(echo "${SESSIONS}" | tr ' ' '\n' | sed '/^$/d')

declare -A WORKERS=()
for gpu in "${GPU_ARR[@]}"; do
  WORKERS["${gpu}"]=""
done

i=0
for sid in "${SID_ARR[@]}"; do
  gpu="${GPU_ARR[$((i % ${#GPU_ARR[@]}))]}"
  WORKERS["${gpu}"]+="${sid} "
  i=$((i + 1))
done

pids=()
status=0
for gpu in "${GPU_ARR[@]}"; do
  sids="${WORKERS[${gpu}]}"
  [[ -z "${sids// }" ]] && continue
  (
    for sid in ${sids}; do
      run_one "${gpu}" "${sid}" || exit 1
    done
  ) &
  pids+=("$!")
done

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

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

echo "ALL EasyCom direct extractions completed -> ${FEATURE_ROOT}"