#!/usr/bin/env bash
# Foreground matched EasyCom pooled AV -> causal anonymous-track GRU residual.
# Usage: bash run_easycom_track_gru_residual.sh [all|pooled|track|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}
LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576}
CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_track_gru_residual}
RESULT_ROOT=${RESULT_ROOT:-results/easycom_track_gru_residual}
TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_track_gru_residual_test}
EPOCHS=${EPOCHS:-4}
if [[ -z "${WARMUP_EPOCHS:-}" ]]; then
  # Keep the four-epoch screen unchanged, but make EPOCHS=1 smoke runs valid.
  if (( EPOCHS == 1 )); then
    WARMUP_EPOCHS=0
  else
    WARMUP_EPOCHS=1.6
  fi
fi
LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256}
PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4}
PERSON_HISTORY_FRAMES=${PERSON_HISTORY_FRAMES:-16}
PERSON_TRACK_MAX_GAP_SECONDS=${PERSON_TRACK_MAX_GAP_SECONDS:-1.0}
SUMMARIZE=${SUMMARIZE:-1}
read -r -a SESSION_ARGS <<< "${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=(pooled track) ;;
  pooled|track) variants=("$variant") ;;
  test) variants=(pooled track) ;;
  *) echo "Expected all, pooled, track, or test" >&2; exit 2 ;;
esac

mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT"

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

train_and_infer_pooled() {
  local seed=$1 checkpoint result
  checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
  result="$RESULT_ROOT/pooled_seed${seed}"
  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
    echo "Refusing to overwrite $checkpoint or $result" >&2
    return 1
  fi
  python train_turn_head.py \
    --feature-root "$FEATURE_ROOT" \
    --target-root "$TARGET_ROOT" \
    --local-feature-root "$LOCAL_ROOT" \
    --local-modalities av \
    --local-fusion pooled \
    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
    --output "$checkpoint" \
    --sessions "${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
  # Screen architecture/hyperparameters only on the held-out training session.
  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
}

train_and_infer_track() {
  local seed=$1 base checkpoint result
  base="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
  checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
  result="$RESULT_ROOT/track_gru_seed${seed}"
  if [[ ! -s "$base" ]]; then
    echo "Missing matched pooled checkpoint: $base; run variant pooled or all first" >&2
    return 1
  fi
  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
    echo "Refusing to overwrite $checkpoint or $result" >&2
    return 1
  fi
  python train_turn_head.py \
    --feature-root "$FEATURE_ROOT" \
    --target-root "$TARGET_ROOT" \
    --local-feature-root "$LOCAL_ROOT" \
    --local-modalities av \
    --local-fusion track-gru-residual \
    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
    --person-attention-dim "$PERSON_ATTENTION_DIM" \
    --person-attention-heads "$PERSON_ATTENTION_HEADS" \
    --person-history-frames "$PERSON_HISTORY_FRAMES" \
    --person-track-max-gap-seconds "$PERSON_TRACK_MAX_GAP_SECONDS" \
    --initialize-pooled-checkpoint "$base" \
    --output "$checkpoint" \
    --sessions "${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
  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
}

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"
      if [[ "$selected" == track ]]; then
        checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
        result="$TEST_RESULT_ROOT/track_gru_seed${seed}"
      else
        result="$TEST_RESULT_ROOT/pooled_seed${seed}"
      fi
      if [[ ! -s "$checkpoint" ]]; then
        echo "Missing checkpoint: $checkpoint" >&2
        echo "The test mode is inference-only; run 'all' for these seeds first." >&2
        exit 1
      fi
      echo "$(date): final test ${selected}, seed=${seed}"
      infer_checkpoint "$checkpoint" "$result" 3 2 1
    done
  done
  summary_root=$TEST_RESULT_ROOT
else
  for selected in "${variants[@]}"; do
    for seed in "${seeds[@]}"; do
      echo "$(date): ${selected}, seed=${seed}"
      if [[ "$selected" == pooled ]]; then
        train_and_infer_pooled "$seed"
      else
        train_and_infer_track "$seed"
      fi
    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