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

# Frozen-Qwen head equivalent of EgoSpeak's run_LSTR_{train,inference}.sh.
# Override any path or device by setting the corresponding environment value.
ACTION="${1:-all}"
if [[ "${ACTION}" != "train" && "${ACTION}" != "infer" && "${ACTION}" != "all" ]]; then
  echo "Usage: $0 [train|infer|all]" >&2
  exit 2
fi
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/test_qwen_omni}"
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_2fps}"
TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
CHECKPOINT="${CHECKPOINT:-${PROJECT_ROOT}/checkpoints/causal_prefix_2fps_head.pth}"
RESULT_ROOT="${RESULT_ROOT:-${PROJECT_ROOT}/results/causal_prefix_2fps_head}"
DEVICE="${DEVICE:-cuda:0}"

cd "${PROJECT_ROOT}"

if [[ "${ACTION}" == "train" || "${ACTION}" == "all" ]]; then
  python train_turn_head.py \
    --feature-root "${FEATURE_ROOT}" \
    --target-root "${TARGET_ROOT}" \
    --output "${CHECKPOINT}" \
    --device "${DEVICE}"
fi

if [[ "${ACTION}" == "infer" || "${ACTION}" == "all" ]]; then
  python infer_turn_head.py \
    --feature-root "${FEATURE_ROOT}" \
    --target-root "${TARGET_ROOT}" \
    --checkpoint "${CHECKPOINT}" \
    --output-dir "${RESULT_ROOT}" \
    --device "${DEVICE}"
fi