# Default GPU devices
DEFAULT_GPUS="0 1 2 3"

# Get GPU devices from the first argument or use default
GPU_DEVICES=${1:-$DEFAULT_GPUS}

# Check if sweep configuration file is provided
if [ -z "$2" ]; then
  echo "Usage: $0 [GPU_DEVICES] <SWEEP_CONFIG>"
  exit 1
fi

SWEEP_CONFIG=$2

# Extract WANDB_ENTITY from the default config file if not already set
if [ -z "$WANDB_ENTITY" ]; then
  WANDB_ENTITY_FROM_CONFIG=$(grep 'wandb_entity:' conf/run/defaults.yaml | awk '{print $2}')
  if [ -n "$WANDB_ENTITY_FROM_CONFIG" ]; then
    export WANDB_ENTITY=$WANDB_ENTITY_FROM_CONFIG
  else
    echo "Error: wandb_entity not found in conf/run/defaults.yaml and WANDB_ENTITY environment variable is not set."
    exit 1
  fi
fi

# Run wandb sweep and capture all output
WANDB_OUTPUT=$(wandb sweep $SWEEP_CONFIG 2>&1 --project emergence_sparse_attention)

echo "$WANDB_OUTPUT"

# Extract the sweep ID
SWEEP_ID=$(echo "$WANDB_OUTPUT" | awk '/Creating sweep with ID:/ {print $NF}')

echo "Sweep ID: $SWEEP_ID"

if [ -z "$SWEEP_ID" ]; then
  echo "Failed to start sweep. Please check your configuration."
  exit 1
fi

echo "WandB entity is: $WANDB_ENTITY"

# Extract the full agent command
AGENT_CMD="wandb agent ${WANDB_ENTITY}/emergence_sparse_attention/$SWEEP_ID"

echo "Starting agents with command: $AGENT_CMD"

for DEVICE_NUM in $GPU_DEVICES
do
  CUDA_VISIBLE_DEVICES=$DEVICE_NUM $AGENT_CMD &
  sleep 5
done