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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
REPO="$(cd "$RECIPE/.." && pwd)"
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
EVAL_ROOT="${EVAL_ROOT:-$OUTPUT_ROOT/eval}"
VANILLA_MODEL="${VANILLA_DIR:-$REPO}"
AHA_MODEL="${AHA_MODEL:-$OUTPUT_ROOT/aha/stage2/checkpoint-25}"
L2A_MODEL="${L2A_MODEL:-$OUTPUT_ROOT/l2a_style/stage2/checkpoint-25}"
GPU_LIST="${GPU_LIST:-0}"
THRESHOLDS=(0.45 0.525 0.575 0.65)
BENCHMARKS=(ruler_a ruler_b babilong helmet mrcr)
IFS=',' read -r -a GPUS <<< "$GPU_LIST"
declare -A SLOT_PIDS=()
mkdir -p "$EVAL_ROOT"

for model in "$VANILLA_MODEL" "$L2A_MODEL"; do
  [[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
done

# Upload outputs to HF Hub after each config group to survive timeouts
upload_to_hub() {
  echo "=== Incremental upload to HF Hub ==="
  python -c "
from huggingface_hub import HfApi
import os
api = HfApi()
if os.path.exists('/workspace/outputs'):
    api.upload_folder(
        folder_path='/workspace/outputs',
        repo_id='keepsloading/icml_repro_scratch',
        repo_type='model',
        path_in_repo='outputs'
    )
    print('Incremental upload complete.')
else:
    print('No outputs dir yet.')
" || echo "Upload failed (non-fatal), continuing..."
}

launch() {
  local gpu="$1" method="$2" model="$3" threshold="$4" benchmark="$5"
  EVAL_ROOT="$EVAL_ROOT" RULER_LIMIT=20 BABILONG_LIMIT=50 \
    AHA_GATE_HARD_THRESHOLD="$threshold" \
    bash -x "$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
  LAST_PID="$!"
}

# Run vanilla baseline
for benchmark in "${BENCHMARKS[@]}"; do
  slot=0
  if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
    wait "${SLOT_PIDS[$slot]}"
  fi
  launch "${GPUS[$slot]}" "vanilla" "$VANILLA_MODEL" "0.5" "$benchmark"
  SLOT_PIDS[$slot]="$LAST_PID"
done
for pid in "${SLOT_PIDS[@]:-}"; do wait "$pid" || true; done
SLOT_PIDS=()
upload_to_hub

# Run each L2A threshold group
for threshold in "${THRESHOLDS[@]}"; do
  slug="${threshold/./}"
  method="token_t${slug}"
  for benchmark in "${BENCHMARKS[@]}"; do
    slot=0
    if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
      wait "${SLOT_PIDS[$slot]}"
    fi
    launch "${GPUS[$slot]}" "$method" "$L2A_MODEL" "$threshold" "$benchmark"
    SLOT_PIDS[$slot]="$LAST_PID"
  done
  for pid in "${SLOT_PIDS[@]:-}"; do wait "$pid" || true; done
  SLOT_PIDS=()
  upload_to_hub
done

python "$SCRIPT_DIR/summarize_qwen1p7b_router_granularity_20260714.py" \
  --repo "$REPO" \
  --eval-root "$EVAL_ROOT" \
  --baseline-root "$EVAL_ROOT/vanilla" \
  --input-dir "$RECIPE/data/eval_inputs" \
  --output-dir "$OUTPUT_ROOT/summary"