#!/usr/bin/env bash
# Global Omni with Qwen audio, traditional audio, and high-rate person evidence.
# Usage: bash run_easycom_highrate_person_av.sh \
#   [all|global|qwen_audio|traditional_audio|highrate|test] [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}
QWEN_LOCAL_ROOT=${QWEN_LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576}
HIGHRATE_LOCAL_ROOT=${HIGHRATE_LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_muvap_asd_25hz_v2}
CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_highrate_person_av}
RESULT_ROOT=${RESULT_ROOT:-results/easycom_highrate_person_av_screen}
TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_highrate_person_av_test}
EPOCHS=${EPOCHS:-4}
WARMUP_EPOCHS=${WARMUP_EPOCHS:-1.6}
LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256}
PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4}
SUMMARIZE=${SUMMARIZE:-1}
read -r -a TRAIN_SESSION_ARGS <<< "${TRAIN_SESSIONS:-4 5 6 7 8 9 10 11 12}"
read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}"
read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}"
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=(global qwen_audio traditional_audio highrate) ;;
  global|qwen_audio|traditional_audio|highrate) variants=("$variant") ;;
  test) variants=(global qwen_audio traditional_audio highrate) ;;
  *) echo "Expected all, global, qwen_audio, traditional_audio, highrate, or test" >&2; exit 2 ;;
esac

mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT"

run_inference() {
  local checkpoint=$1 local_root=$2 output=$3
  shift 3
  if [[ -e "$output/metrics.json" ]]; then
    echo "Refusing to overwrite $output/metrics.json" >&2
    return 1
  fi
  local local_args=()
  if [[ -n "$local_root" ]]; then
    local_args=(--local-feature-root "$local_root")
  fi
  python infer_turn_head.py \
    --feature-root "$FEATURE_ROOT" \
    --target-root "$TARGET_ROOT" \
    "${local_args[@]}" \
    --checkpoint "$checkpoint" \
    --output-dir "$output" \
    --sessions "$@" \
    --expected-mode open-stream-replay \
    --device cuda:0
}

if [[ "$variant" == test ]]; then
  mkdir -p "$TEST_RESULT_ROOT"
  for selected in "${variants[@]}"; do
    for seed in "${seeds[@]}"; do
      checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth"
      [[ -s "$checkpoint" ]] || { echo "Missing checkpoint: $checkpoint" >&2; exit 1; }
      if [[ "$selected" == global ]]; then
        local_root=
      elif [[ "$selected" == qwen_audio ]]; then
        local_root=$QWEN_LOCAL_ROOT
      else
        local_root=$HIGHRATE_LOCAL_ROOT
      fi
      run_inference "$checkpoint" "$local_root" \
        "$TEST_RESULT_ROOT/${selected}_seed${seed}" 3 2 1
    done
  done
  summary_root=$TEST_RESULT_ROOT
else
  for selected in "${variants[@]}"; do
    for seed in "${seeds[@]}"; do
      checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth"
      result="$RESULT_ROOT/${selected}_seed${seed}"
      initialization_args=()
      local_args=()
      if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
        echo "Refusing to overwrite $checkpoint or $result" >&2
        exit 1
      fi
      if [[ "$selected" == global ]]; then
        local_root=
        modalities=av
        fusion=pooled
      elif [[ "$selected" == qwen_audio ]]; then
        local_root=$QWEN_LOCAL_ROOT
        modalities=audio
        fusion=pooled
      elif [[ "$selected" == traditional_audio ]]; then
        local_root=$HIGHRATE_LOCAL_ROOT
        modalities=audio
        fusion=pooled
      else
        local_root=$HIGHRATE_LOCAL_ROOT
        modalities=av
        fusion=highrate-person-av-attention
        audio_checkpoint="$CHECKPOINT_ROOT/traditional_audio_seed${seed}.pth"
        [[ -s "$audio_checkpoint" ]] || {
          echo "Missing matched audio checkpoint: $audio_checkpoint" >&2
          echo "Run: bash run_easycom_highrate_person_av.sh traditional_audio $seed" >&2
          exit 1
        }
        initialization_args=(--initialize-pooled-checkpoint "$audio_checkpoint")
      fi
      if [[ -n "$local_root" ]]; then
        local_args=(
          --local-feature-root "$local_root"
          --local-modalities "$modalities"
        )
      fi
      echo "$(date): ${selected}, seed=${seed}"
      python train_turn_head.py \
        --feature-root "$FEATURE_ROOT" \
        --target-root "$TARGET_ROOT" \
        "${local_args[@]}" \
        --local-fusion "$fusion" \
        --local-projection-dim "$LOCAL_PROJECTION_DIM" \
        --person-attention-dim "$PERSON_ATTENTION_DIM" \
        --person-attention-heads "$PERSON_ATTENTION_HEADS" \
        "${initialization_args[@]}" \
        --output "$checkpoint" \
        --sessions "${TRAIN_SESSION_ARGS[@]}" \
        --dev-sessions "${DEV_SESSION_ARGS[@]}" \
        --epochs "$EPOCHS" \
        --batch-size 512 \
        --learning-rate 7e-5 \
        --weight-decay 5e-5 \
        --warmup-epochs "$WARMUP_EPOCHS" \
        --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
      run_inference "$checkpoint" "$local_root" "$result" "${SCREEN_SESSION_ARGS[@]}"
    done
  done
  summary_root=$RESULT_ROOT
fi

if [[ "$SUMMARIZE" == 1 ]]; then
  python summarize_local_baselines.py \
    --root "$summary_root" \
    --dataset easycom \
    --protocol causal_valid_only \
    --paired-deltas
fi