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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
GPU_COUNT="$(nvidia-smi -L | wc -l)"
if [[ "$GPU_COUNT" -lt 1 ]]; then
  echo "At least one NVIDIA GPU is required." >&2
  exit 2
fi
GPU_AHA="${GPU_AHA:-0}"
if [[ -z "${GPU_L2A+x}" ]]; then
  GPU_L2A=0
  if [[ "$GPU_COUNT" -ge 2 ]]; then
    GPU_L2A=1
  fi
fi

for gpu in "$GPU_AHA" "$GPU_L2A"; do
  if [[ ! "$gpu" =~ ^[0-9]+$ ]] || (( gpu >= GPU_COUNT )); then
    echo "Invalid GPU index $gpu; nvidia-smi reports $GPU_COUNT visible GPU(s)." >&2
    exit 2
  fi
done

echo "Training mode: one GPU per experimental variant (no DDP/FSDP/DeepSpeed)."
echo "AHA variant GPU: $GPU_AHA"
echo "L2A-style variant GPU: $GPU_L2A"
if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
  echo "The two variants will run sequentially on the same GPU."
else
  echo "The two variants will run concurrently on two independent GPUs."
fi
if [[ "$GPU_COUNT" -gt 2 ]]; then
  echo "Visible GPUs: $GPU_COUNT; this recipe intentionally does not use all GPUs."
fi

docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
docker run --rm --gpus all --ipc=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -e GPU_AHA="$GPU_AHA" -e GPU_L2A="$GPU_L2A" \
  -v "$REPO:/workspace" \
  -w /workspace/recipe \
  "$IMAGE" bash scripts/train_both.sh