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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
GPU_AHA="${GPU_AHA:-0}"
GPU_L2A="${GPU_L2A:-1}"
cd "$RECIPE"

echo "AHA = per-token, per-KV-head gate; assigned visible GPU $GPU_AHA."
echo "L2A-style = per-token, head-shared gate; assigned visible GPU $GPU_L2A."
echo "Each variant is a separate world-size-1 run with global batch size 1."

python "$SCRIPT_DIR/validate_release.py"
python -m unittest -v tests.test_router_granularity

if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA"
  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A"
else
  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA" & aha_pid=$!
  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A" & l2a_pid=$!
  status=0
  wait "$aha_pid" || status=1
  wait "$l2a_pid" || status=1
  [[ "$status" == 0 ]] || exit "$status"
fi

echo "AHA selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/aha/stage2/checkpoint-25"
echo "L2A-style selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/l2a_style/stage2/checkpoint-25"