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

stage="${1:-screen}"
shift || true
if (( $# )); then
  seeds=("$@")
else
  seeds=(0)
fi

: "${DATASET_ROOT:?Set DATASET_ROOT to the extracted EasyCom root}"
: "${FEATURE_ROOT:?Set FEATURE_ROOT to the verified global Omni features}"
: "${TARGET_ROOT:?Set TARGET_ROOT to EasyCom target_perframe files}"
: "${LOCAL_ROOT:?Set LOCAL_ROOT to the complete Qwen local AV sidecars}"

ORACLE_ROOT="${ORACLE_ROOT:-features/easycom_oracle_information_2s}"
CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-checkpoints/easycom_oracle_information_4ep}"
RESULT_ROOT="${RESULT_ROOT:-results/easycom_oracle_information_4ep}"
TRAIN_SESSIONS="${TRAIN_SESSIONS:-4 5 6 7 8 9 10 11 12}"
DEV_SESSIONS="${DEV_SESSIONS:-12}"
SCREEN_SESSIONS="${SCREEN_SESSIONS:-12}"
TEST_SESSIONS="${TEST_SESSIONS:-1 2 3}"
EPOCHS="${EPOCHS:-4}"
WARMUP_EPOCHS="${WARMUP_EPOCHS:-1}"
LAYER_FUSION="${LAYER_FUSION:-concat-mlp}"
EXPECTED_MODE="${EXPECTED_MODE:-open-stream-replay}"
VARIANTS="${VARIANTS:-audio visibility-only aggregate-visible per-person shuffled-id delayed}"
ORACLE_CONDITIONS="${ORACLE_CONDITIONS:-visibility-only aggregate-visible per-person shuffled-id delayed}"
PYTHON_BIN="${PYTHON_BIN:-python}"

read -r -a train_sessions <<<"${TRAIN_SESSIONS}"
read -r -a dev_sessions <<<"${DEV_SESSIONS}"
read -r -a screen_sessions <<<"${SCREEN_SESSIONS}"
read -r -a test_sessions <<<"${TEST_SESSIONS}"
read -r -a variants <<<"${VARIANTS}"
read -r -a oracle_conditions <<<"${ORACLE_CONDITIONS}"

layer_args=()
if [[ "${LAYER_FUSION}" != "single" ]]; then
  layer_args=(--layer-fusion "${LAYER_FUSION}")
fi

prepare_oracles() {
  "${PYTHON_BIN}" prepare_easycom_oracle_features.py \
    --dataset-root "${DATASET_ROOT}" \
    --feature-root "${FEATURE_ROOT}" \
    --target-root "${TARGET_ROOT}" \
    --local-feature-root "${LOCAL_ROOT}" \
    --output-root "${ORACLE_ROOT}" \
    --conditions "${oracle_conditions[@]}" \
    --history-seconds 2.0 \
    --delay-seconds 1.0 \
    --expected-mode "${EXPECTED_MODE}" \
    --reuse-existing
}

variant_local_args() {
  local variant="$1"
  if [[ "${variant}" == "audio" ]]; then
    printf '%s\0' \
      --local-feature-root "${LOCAL_ROOT}" \
      --local-modalities audio \
      --local-fusion pooled
  else
    printf '%s\0' \
      --local-feature-root "${ORACLE_ROOT}/${variant}" \
      --local-modalities av \
      --local-fusion oracle-frame-set-attention \
      --person-attention-dim 64 \
      --person-attention-heads 4
  fi
}

train_and_screen() {
  mkdir -p "${CHECKPOINT_ROOT}" "${RESULT_ROOT}"
  for seed in "${seeds[@]}"; do
    for variant in "${variants[@]}"; do
      checkpoint="${CHECKPOINT_ROOT}/${variant}_seed${seed}.pth"
      result="${RESULT_ROOT}/${variant}_seed${seed}"
      local_feature_root="${ORACLE_ROOT}/${variant}"
      if [[ "${variant}" == "audio" ]]; then
        local_feature_root="${LOCAL_ROOT}"
      fi
      if [[ -e "${checkpoint}" || -e "${result}" ]]; then
        echo "Refusing to overwrite ${checkpoint} or ${result}" >&2
        exit 1
      fi
      local_args=()
      while IFS= read -r -d '' item; do
        local_args+=("${item}")
      done < <(variant_local_args "${variant}")
      echo "$(date): training ${variant}, seed=${seed}"
      "${PYTHON_BIN}" train_turn_head.py \
        --feature-root "${FEATURE_ROOT}" \
        --target-root "${TARGET_ROOT}" \
        --output "${checkpoint}" \
        --sessions "${train_sessions[@]}" \
        --dev-sessions "${dev_sessions[@]}" \
        --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 \
        "${layer_args[@]}" \
        "${local_args[@]}" \
        --local-projection-dim 256 \
        --expected-mode "${EXPECTED_MODE}" \
        --seed "${seed}" \
        --device cuda:0

      echo "$(date): screening ${variant}, seed=${seed} on ${screen_sessions[*]}"
      "${PYTHON_BIN}" infer_turn_head.py \
        --feature-root "${FEATURE_ROOT}" \
        --target-root "${TARGET_ROOT}" \
        --checkpoint "${checkpoint}" \
        --output-dir "${result}" \
        --sessions "${screen_sessions[@]}" \
        --local-feature-root "${local_feature_root}" \
        --expected-mode "${EXPECTED_MODE}" \
        --device cuda:0
    done
  done
  "${PYTHON_BIN}" summarize_local_baselines.py \
    --root "${RESULT_ROOT}" \
    --dataset easycom \
    --protocol causal_valid_only
}

test_selected() {
  mkdir -p "${RESULT_ROOT}_test"
  for seed in "${seeds[@]}"; do
    for variant in "${variants[@]}"; do
      checkpoint="${CHECKPOINT_ROOT}/${variant}_seed${seed}.pth"
      result="${RESULT_ROOT}_test/${variant}_seed${seed}"
      local_feature_root="${ORACLE_ROOT}/${variant}"
      if [[ "${variant}" == "audio" ]]; then
        local_feature_root="${LOCAL_ROOT}"
      fi
      if [[ ! -f "${checkpoint}" ]]; then
        echo "Missing checkpoint: ${checkpoint}" >&2
        exit 1
      fi
      if [[ -e "${result}" ]]; then
        echo "Refusing to overwrite ${result}" >&2
        exit 1
      fi
      "${PYTHON_BIN}" infer_turn_head.py \
        --feature-root "${FEATURE_ROOT}" \
        --target-root "${TARGET_ROOT}" \
        --checkpoint "${checkpoint}" \
        --output-dir "${result}" \
        --sessions "${test_sessions[@]}" \
        --local-feature-root "${local_feature_root}" \
        --expected-mode "${EXPECTED_MODE}" \
        --device cuda:0
    done
  done
  "${PYTHON_BIN}" summarize_local_baselines.py \
    --root "${RESULT_ROOT}_test" \
    --dataset easycom \
    --protocol causal_valid_only
}

case "${stage}" in
  prepare)
    prepare_oracles
    ;;
  screen)
    train_and_screen
    ;;
  all)
    prepare_oracles
    train_and_screen
    ;;
  test)
    test_selected
    ;;
  summarize)
    "${PYTHON_BIN}" summarize_local_baselines.py \
      --root "${RESULT_ROOT}" \
      --dataset easycom \
      --protocol causal_valid_only
    ;;
  *)
    echo "Usage: $0 {prepare|screen|all|test|summarize} [seed ...]" >&2
    exit 2
    ;;
esac