#!/bin/bash
# Launch all 19 inference runs to produce descriptive-named CSVs in candidate_csvs/
#
# Usage:  bash scripts/launch_all_predicts.sh
#         (auto-locates ROOT as the parent dir of this script)
#
#         # or override the workspace root explicitly:
#         ROOT=/my/extract/path bash scripts/launch_all_predicts.sh
#
# Assumes the workspace at $ROOT holds:
#   - 19 LoRA adapters in $ROOT/checkpoints/<descriptive_name>/
#   - Qwen base models in $ROOT/hub/{Qwen3.5-27B,Qwen3.6-27B,Qwen3-32B}/
#   - vLLM env activated (`conda activate vllm`)
#   - At least 1 × 80GB GPU (H100/A100). 8 GPUs is the sweet spot (~2h); a
#     single GPU also works (~20-30h, purely sequential). Concurrency is
#     auto-capped at min(8, visible GPU count).
#   - Test JSONLs already present under $ROOT/LF/data/
#       msrh_rag_test_k3_AfriE5_TV.json                    (K=3 base — NoFewshots only)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json         (K=3 fewshot)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k4.json         (K=4 fewshot)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k5.json         (K=5 fewshot)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k5_v8.json      (K=5 v8)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k7.json         (K=7 fewshot)
#       msrh_rag_test_k3_AfriE5_TV_fewshot_k7_v8.json      (K=7 v8)
#
# Outputs:  $ROOT/candidate_csvs/<descriptive_name>.csv  (19 files, ready for ensemble)
#
# Scheduling policy:
#   - up to 8 predicts running concurrently, one per GPU
#   - staggered launches (40 s apart) to avoid concurrent vLLM engine-init races
#   - each launch is assigned to the LEAST-BUSY GPU (< 5 GB used), not a fixed rotation
#   - per-proc TORCHINDUCTOR_CACHE_DIR / TRITON_CACHE_DIR isolate compile caches
#   - failed adapters (no CSV or JSONL != 2618 rows) are retried once at the end
#
# NOTE: per-cand (max_model_len, max_num_seqs) match the original H100 chain scripts
# that produced go.csv. Changing them alters vLLM paged-attention block sizing and
# can shift greedy tokens.
set -e
ROOT="${ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
CKPT_ROOT=$ROOT/checkpoints
HUB_ROOT=$ROOT/hub
TEST_DIR=$ROOT/LF/data
OUT_DIR=$ROOT/candidate_csvs
LOG_DIR=$ROOT/training_logs/predict
mkdir -p $OUT_DIR $LOG_DIR

# 19 predict jobs — pipe-delimited: NAME|BASE|TEST_JSON|MAX_LEN|MAX_SEQS
JOBS=(
"Qwen3.5-27B-3fewshots-bs64-3eps-ckpt-1200|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json|6144|64"
"Qwen3.5-27B-3fewshots-bs64-3eps-ckpt-1100|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json|6144|64"
"Qwen3.5-27B-3fewshots-bs64-5eps-ckpt-1200|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json|6144|64"
"Qwen3.5-27B-4fewshots-bs64-3eps-ckpt-1600|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k4.json|6144|64"
"Qwen3.5-27B-5fewshots-bs64-3eps-v8prompt-ckpt-1500|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k5_v8.json|7168|64"
"Qwen3.5-27B-7fewshots-bs64-3eps-ckpt-1600|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k7.json|9216|32"
"Qwen3.5-27B-7fewshots-bs64-3eps-ckpt-1200|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k7.json|9216|32"
"Qwen3.5-27B-NoFewshots-bs64-5eps-ckpt-2800|Qwen3.5-27B|msrh_rag_test_k3_AfriE5_TV.json|6144|32"
"Qwen3.6-27B-3fewshots-bs64-3eps-ckpt-1600|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json|6144|32"
"Qwen3.6-27B-4fewshots-bs64-3eps-ckpt-1400|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k4.json|6144|64"
"Qwen3.6-27B-5fewshots-bs64-3eps-ckpt-1200|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k5.json|7168|64"
"Qwen3.6-27B-5fewshots-bs64-3eps-ckpt-1000|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k5.json|7168|64"
"Qwen3.6-27B-7fewshots-bs64-3eps-ckpt-1600|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV_fewshot_k7_v8.json|10240|32"
"Qwen3.6-27B-NoFewshots-bs64-5eps-ckpt-2600|Qwen3.6-27B|msrh_rag_test_k3_AfriE5_TV.json|6144|32"
"Qwen3-32B-3fewshots-bs64-3eps-ckpt-1400|Qwen3-32B|msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json|6144|64"
"Qwen3-32B-5fewshots-bs64-3eps-ckpt-1700|Qwen3-32B|msrh_rag_test_k3_AfriE5_TV_fewshot_k5.json|7168|64"
"Qwen3-32B-7fewshots-bs64-3eps-ckpt-1600|Qwen3-32B|msrh_rag_test_k3_AfriE5_TV_fewshot_k7.json|9216|32"
"Qwen3-32B-7fewshots-bs64-3eps-ckpt-1200|Qwen3-32B|msrh_rag_test_k3_AfriE5_TV_fewshot_k7.json|9216|32"
"Qwen3-32B-NoFewshots-bs64-4eps-ckpt-6500|Qwen3-32B|msrh_rag_test_k3_AfriE5_TV.json|6144|32"
)

# ── helpers ──────────────────────────────────────────────────────────
# Auto-detect visible GPU count (respects CUDA_VISIBLE_DEVICES if set).
if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then
  NUM_GPUS=$(echo "$CUDA_VISIBLE_DEVICES" | tr ',' '\n' | grep -c .)
  VISIBLE_GPUS=($(echo "$CUDA_VISIBLE_DEVICES" | tr ',' ' '))
else
  NUM_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null | wc -l)
  VISIBLE_GPUS=($(seq 0 $((NUM_GPUS - 1))))
fi
[ "$NUM_GPUS" -lt 1 ] && { echo "✗ no NVIDIA GPUs visible"; exit 1; }
MAX_CONCURRENT=$(( NUM_GPUS < 8 ? NUM_GPUS : 8 ))
echo "Detected $NUM_GPUS visible GPU(s) [${VISIBLE_GPUS[*]}]; up to $MAX_CONCURRENT predicts concurrently."

# Pick the first visible GPU with < 5 GB used; return -1 if none free.
free_gpu () {
  local g
  for g in "${VISIBLE_GPUS[@]}"; do
    local mem=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i $g 2>/dev/null | tr -d ' ')
    [ -z "$mem" ] && continue
    [ "$mem" -lt 5000 ] && { echo $g; return; }
  done
  echo -1
}

predict_one () {
  local NAME="$1" BASE="$2" TEST_JSON="$3" MAX_LEN="$4" MAX_SEQS="$5" GPU="$6"
  local ADAPTER=$CKPT_ROOT/$NAME
  local WORK=$ROOT/_predict_workdir/$NAME
  local IND=/tmp/tinductor_${NAME}
  local TRI=/tmp/ttriton_${NAME}
  mkdir -p $WORK $IND $TRI

  if [ ! -f "$ADAPTER/adapter_model.safetensors" ]; then
    echo "  ✗ MISSING adapter: $ADAPTER (skipping)"; return
  fi
  if [ ! -f "$TEST_DIR/$TEST_JSON" ]; then
    echo "  ✗ MISSING test JSON: $TEST_DIR/$TEST_JSON (skipping $NAME)"; return
  fi

  echo "  → [GPU $GPU] $NAME (max_len=$MAX_LEN max_seqs=$MAX_SEQS)"
  (
    TORCHINDUCTOR_CACHE_DIR=$IND TRITON_CACHE_DIR=$TRI \
      CUDA_VISIBLE_DEVICES=$GPU python $ROOT/scripts/vllm_predict_extra.py \
        --base "$HUB_ROOT/$BASE" --adapter "$ADAPTER" \
        --rag_test "$TEST_DIR/$TEST_JSON" \
        --out_jsonl "$WORK/predictions.jsonl" \
        --max_lora_rank 128 --max_new 512 \
        --max_model_len $MAX_LEN --mem_util 0.90 --max_num_seqs $MAX_SEQS \
        --temperature 0.0 --top_p 1.0 --best_of 1 --no_think \
        > "$LOG_DIR/$NAME.log" 2>&1 \
    && python $ROOT/scripts/jsonl_rowidx_to_csv.py \
        --jsonl "$WORK/predictions.jsonl" \
        --out_csv "$OUT_DIR/$NAME.csv" \
        >> "$LOG_DIR/$NAME.log" 2>&1 \
    && echo "  ✓ $NAME.csv" \
    || echo "  ✗ $NAME FAILED (see $LOG_DIR/$NAME.log)"
  ) &
}

# ── main scheduling loop: stagger + free-GPU assignment ──────────────
launch_pass () {
  local -n QUEUE=$1
  local pass_name="$2"
  echo "=== $pass_name (queue=${#QUEUE[@]}) ==="
  for JOB in "${QUEUE[@]}"; do
    while true; do
      local RUN=$(pgrep -c -f vllm_predict_extra 2>/dev/null)
      [ -z "$RUN" ] && RUN=0
      if [ "$RUN" -ge "$MAX_CONCURRENT" ]; then sleep 20; continue; fi
      local G=$(free_gpu)
      if [ "$G" = "-1" ]; then sleep 20; continue; fi
      IFS="|" read -r NAME BASE TEST MLEN MSEQS <<< "$JOB"
      predict_one "$NAME" "$BASE" "$TEST" "$MLEN" "$MSEQS" "$G"
      sleep 40   # stagger to avoid concurrent init races
      break
    done
  done
  wait
  echo "  ✓ $pass_name done"
}

# ── verify + retry ───────────────────────────────────────────────────
missing_jobs () {
  local -a MISS=()
  for JOB in "${JOBS[@]}"; do
    IFS="|" read -r NAME BASE TEST MLEN MSEQS <<< "$JOB"
    local CSV=$OUT_DIR/$NAME.csv
    local JSONL=$ROOT/_predict_workdir/$NAME/predictions.jsonl
    # OK if CSV exists AND JSONL has 2618 lines
    if [ -s "$CSV" ] && [ -s "$JSONL" ] && [ "$(wc -l < "$JSONL")" -eq 2618 ]; then
      continue
    fi
    MISS+=("$JOB")
  done
  printf '%s\n' "${MISS[@]}"
}

# Pass 1: launch all 19
PASS1=("${JOBS[@]}")
launch_pass PASS1 "PASS 1 (all 19)"

# Verify + retry up to 1 more pass on the ones that didn't produce a valid CSV
mapfile -t MISS < <(missing_jobs)
if [ "${#MISS[@]}" -gt 0 ] && [ -n "${MISS[0]:-}" ]; then
  echo
  echo "=== Retrying ${#MISS[@]} failed adapter(s) ==="
  # Give any lingering GPU state a moment to release
  sleep 30
  RETRY=("${MISS[@]}")
  launch_pass RETRY "PASS 2 (retry)"
fi

# Final report
echo
FINAL_MISS=$(missing_jobs | wc -l)
if [ "$FINAL_MISS" -eq 0 ] || [ "$FINAL_MISS" -eq 1 -a -z "$(missing_jobs)" ]; then
  echo "=== All 19 predictions complete ==="
else
  echo "=== $FINAL_MISS adapter(s) STILL FAILED — check $LOG_DIR/*.log ==="
  missing_jobs | awk -F'|' '{print "  ✗ " $1}'
fi
ls -lh $OUT_DIR/ 2>/dev/null | tail -22
echo
echo "Next step:  python $ROOT/scripts/build_ensemble.py"