#!/usr/bin/env bash
# Foreground only. Usage: bash run_easycom_local_baselines.sh [all|baseline|audio|visual|av] [seed ...]
set -euo pipefail
cd "$(dirname "$0")"

FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack}
TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}
LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s}
CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_local_b1}
RESULT_ROOT=${RESULT_ROOT:-results/easycom_local_b1}
LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4}
export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4}

variant=${1:-all}
if (( $# )); then shift; fi
if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi
case "$variant" in
  all) variants=(baseline audio visual av) ;;
  baseline|audio|visual|av) variants=("$variant") ;;
  *) echo "Expected all, baseline, audio, visual, or av" >&2; exit 2 ;;
esac

for variant in "${variants[@]}"; do
  train_local=()
  infer_local=()
  if [[ "$variant" != baseline ]]; then
    train_local=(--local-feature-root "$LOCAL_ROOT" --local-modalities "$variant"
                 --local-projection-dim "$LOCAL_PROJECTION_DIM")
    infer_local=(--local-feature-root "$LOCAL_ROOT")
  fi
  for seed in "${seeds[@]}"; do
    checkpoint="$CHECKPOINT_ROOT/${variant}_seed${seed}.pth"
    result="$RESULT_ROOT/${variant}_seed${seed}"
    if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
      echo "Refusing to overwrite $checkpoint or $result; select new output roots." >&2
      exit 1
    fi
    python train_turn_head.py \
      --feature-root "$FEATURE_ROOT" --target-root "$TARGET_ROOT" \
      --output "$checkpoint" --sessions 4 5 6 7 8 9 10 11 12 \
      --epochs 4 --batch-size 512 --learning-rate 7e-5 --weight-decay 5e-5 \
      --warmup-epochs 1.6 --warmup-factor 0.3 --dropout 0.2 \
      --loss lstr-mce --transition-objective none --layer-fusion concat-mlp \
      --expected-mode open-stream-replay --seed "$seed" --device cuda:0 \
      "${train_local[@]}"
    python infer_turn_head.py \
      --feature-root "$FEATURE_ROOT" --target-root "$TARGET_ROOT" \
      --checkpoint "$checkpoint" --output-dir "$result" --sessions 3 2 1 \
      --expected-mode open-stream-replay --device cuda:0 "${infer_local[@]}"
  done
done