#!/usr/bin/env bash
# Ordered, resumable run of the whole model-3 pipeline.
#
#   bash scripts/queue.sh            # full pipeline
#   bash scripts/queue.sh embed      # embeddings only
#   STAGE_GPU_TRAIN=0 STAGE_GPU_TEST=1 bash scripts/queue.sh
#
# Every stage is resumable: re-running skips work that is already complete, so a
# kill (shared GPU, disk, anything) costs at most one 250k-row shard.
#
# Check the GPUs are actually free before launching:
#   nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv
set -euo pipefail

cd "$(dirname "$0")/.."
ROOT=$(pwd)
PY=${PY:-python3}
MODEL=${M3_MODEL:-e5-small}
GPU_TRAIN=${STAGE_GPU_TRAIN:-0}
GPU_TEST=${STAGE_GPU_TEST:-1}
LOGS=${M3_LOGS:-$ROOT/logs}
mkdir -p "$LOGS"
STAGE=${1:-all}

export M3_MODEL="$MODEL"
export PYTHONPATH="$ROOT/src:${PYTHONPATH:-}"
# Never let a stray HF cache land on /home (it hit 96% full and killed a run).
export HF_HOME=${HF_HOME:-/scratch/user4/mtp/cache/huggingface}
export TOKENIZERS_PARALLELISM=false   # fork deadlock with DataLoader workers

ts() { date +%H:%M:%S; }
run() { echo "== [$(ts)] $*"; "$@"; }

echo "== model-3 queue | model=$MODEL | train-gpu=$GPU_TRAIN test-gpu=$GPU_TEST"
echo "   work dir: ${M3_WORK:-$ROOT/work}"

if [[ "$STAGE" == "all" || "$STAGE" == "embed" ]]; then
  # All four parts are persisted (train pool, train S1, test pool, test S1).
  # At e5-small speed train S1 costs ~6 min, and caching it means the training
  # sample can be scaled past 200k later without re-encoding anything.
  # Order follows the pipeline dependency: train parts first so run_train_eval.py
  # can start while the test corpora are still encoding.
  if [[ "${SINGLE_GPU:-0}" == "1" ]]; then
    for spec in "train pool" "train s1" "test pool" "test s1"; do
      set -- $spec
      run $PY src/embed.py --split "$1" --part "$2" --gpu "$GPU_TRAIN" \
          2>&1 | tee "$LOGS/embed_$1_$2.log"
    done
  else
    run $PY src/embed.py --split train --part all --gpu "$GPU_TRAIN" \
        2>&1 | tee "$LOGS/embed_train.log" &
    P1=$!
    run $PY src/embed.py --split test --part all --gpu "$GPU_TEST" \
        2>&1 | tee "$LOGS/embed_test.log" &
    P2=$!
    wait $P1 $P2
  fi
fi

if [[ "$STAGE" == "all" || "$STAGE" == "train" ]]; then
  # CUDA_VISIBLE_DEVICES, not just --gpu: cross_encoder.py falls back to a bare
  # "cuda" (= cuda:0), so --gpu 1 alone still put the 568M reranker on GPU 0 --
  # the shared card. Masking the device makes every default land correctly.
  run env CUDA_VISIBLE_DEVICES="${TRAIN_GPUS:-$GPU_TRAIN}" $PY src/run_train_eval.py --gpu 0 2>&1 | tee "$LOGS/train_eval.log"
fi

if [[ "$STAGE" == "all" || "$STAGE" == "test" ]]; then
  ART="${M3_WORK:-$ROOT/work}/results_A"
  run env CUDA_VISIBLE_DEVICES="$GPU_TEST" $PY src/run_test.py --gpu 0 --artifacts "$ART" 2>&1 | tee "$LOGS/run_test.log"
  VAL=$(cd "$ROOT/../dataset" && pwd)/utils/validate_submission.py
  run $PY "$VAL" \
      --matching "$ROOT/output/matching_results.tsv" \
      --candidate "$ROOT/output/candidate_pairs.tsv" \
      --test-dir "$(cd "$ROOT/../dataset/dataset" && pwd)/test" \
      2>&1 | tee "$LOGS/validate.log"
fi

echo "== [$(ts)] queue finished"