#!/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="${EVENT_CHECKPOINT_ROOT:-checkpoints/easycom_floor_event_probe}"
RESULT_ROOT="${EVENT_RESULT_ROOT:-results/easycom_floor_event_probe}"
TRAIN_SESSIONS="${TRAIN_SESSIONS:-4 5 6 7 8 9 10 11}"
DEV_SESSIONS="${DEV_SESSIONS:-12}"
EVAL_SESSIONS="${EVAL_SESSIONS:-12}"
TEST_SESSIONS="${TEST_SESSIONS:-1 2 3}"
VARIANTS="${VARIANTS:-aggregate-visible aggregate-visible-address-control aggregate-visible-addressee}"
EPOCHS="${EPOCHS:-20}"
EXPECTED_MODE="${EXPECTED_MODE:-open-stream-replay}"
DECISION_TIMING="${DECISION_TIMING:-pre-boundary}"
PYTHON_BIN="${PYTHON_BIN:-python}"

read -r -a train_sessions <<<"${TRAIN_SESSIONS}"
read -r -a dev_sessions <<<"${DEV_SESSIONS}"
read -r -a eval_sessions <<<"${EVAL_SESSIONS}"
read -r -a test_sessions <<<"${TEST_SESSIONS}"
read -r -a variants <<<"${VARIANTS}"

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 \
      aggregate-visible \
      aggregate-visible-address-control \
      aggregate-visible-addressee \
    --history-seconds 2.0 \
    --expected-mode "${EXPECTED_MODE}" \
    --reuse-existing
}

smoke_probes() {
  "${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}" \
    --sessions Session_11 Session_12 \
    --conditions \
      aggregate-visible \
      aggregate-visible-address-control \
      aggregate-visible-addressee \
    --history-seconds 2.0 \
    --expected-mode "${EXPECTED_MODE}" \
    --reuse-existing
  for variant in "${variants[@]}"; do
    "${PYTHON_BIN}" train_easycom_floor_event_probe.py \
      --feature-root "${FEATURE_ROOT}" \
      --target-root "${TARGET_ROOT}" \
      --oracle-feature-root "${ORACLE_ROOT}/${variant}" \
      --checkpoint "${CHECKPOINT_ROOT}_smoke/${variant}_seed0.pth" \
      --result-dir "${RESULT_ROOT}_smoke/${variant}_seed0" \
      --train-sessions Session_11 \
      --dev-sessions Session_12 \
      --eval-sessions Session_12 \
      --epochs 1 \
      --batch-size 128 \
      --decision-timing "${DECISION_TIMING}" \
      --expected-mode "${EXPECTED_MODE}" \
      --seed 0 \
      --device cuda:0
  done
}

train_and_screen() {
  summary_args=()
  for seed in "${seeds[@]}"; do
    for variant in "${variants[@]}"; do
      checkpoint="${CHECKPOINT_ROOT}/${variant}_seed${seed}.pth"
      result="${RESULT_ROOT}/${variant}_seed${seed}_dev"
      echo "$(date): direct event probe ${variant}, seed=${seed}"
      "${PYTHON_BIN}" train_easycom_floor_event_probe.py \
        --feature-root "${FEATURE_ROOT}" \
        --target-root "${TARGET_ROOT}" \
        --oracle-feature-root "${ORACLE_ROOT}/${variant}" \
        --checkpoint "${checkpoint}" \
        --result-dir "${result}" \
        --train-sessions "${train_sessions[@]}" \
        --dev-sessions "${dev_sessions[@]}" \
        --eval-sessions "${eval_sessions[@]}" \
        --epochs "${EPOCHS}" \
        --decision-timing "${DECISION_TIMING}" \
        --expected-mode "${EXPECTED_MODE}" \
        --seed "${seed}" \
        --device cuda:0
      summary_args+=(--run "${variant}_seed${seed}=${result}")
    done
  done
  "${PYTHON_BIN}" summarize_easycom_floor_event_probes.py \
    "${summary_args[@]}" \
    --output "${RESULT_ROOT}/comparison_dev.csv"
}

test_frozen() {
  summary_args=()
  for seed in "${seeds[@]}"; do
    for variant in "${variants[@]}"; do
      checkpoint="${CHECKPOINT_ROOT}/${variant}_seed${seed}.pth"
      result="${RESULT_ROOT}/${variant}_seed${seed}_test"
      echo "$(date): testing frozen direct event probe ${variant}, seed=${seed}"
      "${PYTHON_BIN}" train_easycom_floor_event_probe.py \
        --feature-root "${FEATURE_ROOT}" \
        --target-root "${TARGET_ROOT}" \
        --oracle-feature-root "${ORACLE_ROOT}/${variant}" \
        --checkpoint "${checkpoint}" \
        --result-dir "${result}" \
        --eval-sessions "${test_sessions[@]}" \
        --eval-only \
        --decision-timing "${DECISION_TIMING}" \
        --expected-mode "${EXPECTED_MODE}" \
        --seed "${seed}" \
        --device cuda:0
      summary_args+=(--run "${variant}_seed${seed}=${result}")
    done
  done
  "${PYTHON_BIN}" summarize_easycom_floor_event_probes.py \
    "${summary_args[@]}" \
    --output "${RESULT_ROOT}/comparison_test.csv"
}

case "${stage}" in
  smoke)
    smoke_probes
    ;;
  prepare)
    prepare_oracles
    ;;
  screen)
    train_and_screen
    ;;
  all)
    prepare_oracles
    train_and_screen
    ;;
  test)
    test_frozen
    ;;
  *)
    echo "Usage: $0 {smoke|prepare|screen|all|test} [seed ...]" >&2
    exit 2
    ;;
esac