#!/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.50 0.525 0.55 0.575 0.60 0.625 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" "$AHA_MODEL" "$L2A_MODEL"; do
  [[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
done

jobs=()
for benchmark in "${BENCHMARKS[@]}"; do
  jobs+=("vanilla|$VANILLA_MODEL|0.5|$benchmark")
done
for threshold in "${THRESHOLDS[@]}"; do
  slug="${threshold/./}"
  for benchmark in "${BENCHMARKS[@]}"; do
    jobs+=("token_kv_head_t${slug}|$AHA_MODEL|$threshold|$benchmark")
    jobs+=("token_t${slug}|$L2A_MODEL|$threshold|$benchmark")
  done
done

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" \
    "$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
  LAST_PID="$!"
}

for i in "${!jobs[@]}"; do
  slot=$((i % ${#GPUS[@]}))
  if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
    wait "${SLOT_PIDS[$slot]}"
  fi
  IFS='|' read -r method model threshold benchmark <<< "${jobs[$i]}"
  launch "${GPUS[$slot]}" "$method" "$model" "$threshold" "$benchmark"
  SLOT_PIDS[$slot]="$LAST_PID"
done

status=0
for pid in "${SLOT_PIDS[@]}"; do
  wait "$pid" || status=1
done
[[ "$status" == 0 ]] || exit "$status"

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"