#!/bin/bash

output_folder="my_games" # your folder containing the output games and animations
avr_folder="AVR_Eval_Agent" # your folder containing the github code
api_key="YOUR-OpenRouter-API" # your openrouter api key
chromium_path=${HOME}/chromium # your path to chromium

# Assuming 4 GPUs
CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen3-32B \
  --dtype bfloat16 \
  --api-key token-1 \
  --max-model-len 32768 \
  --gpu_memory_utilization 0.9 \
  --tensor_parallel_size 2 \
  --port 8001 &
CUDA_VISIBLE_DEVICES=2,3 vllm serve Qwen/Qwen2.5-Omni-7B \
  --dtype bfloat16 \
  --api-key token-2 \
  --max-model-len 32768 \
  --gpu_memory_utilization 0.9 \
  --tensor_parallel_size 2 \
  --trust-remote-code \
  --port 8002 &


# Wait for the server to start
echo "Waiting for VLLM server to start..."
sleep 8m # it takes about 3.5min-ish
echo "Starting experiments"

# Start local server
cd ${output_folder}
python3 -m http.server &

# Array to store PIDs of background processes
display_num=1
pids=()

# With feedback and assets
# Training Qwen3-32B on the 5 different games of video_games_short.csv

for row_index in $(seq 1 5); do # for the 5 games
  # Audio-pulse setup
  OUTPUT=$(${avr_folder}/setup_pulseaudio_sink.sh ${display_num})
  MODULE_ID=$(echo "$OUTPUT" | grep MODULE_ID | cut -d'=' -f2)
  MONITOR_SOURCE=$(echo "$OUTPUT" | grep MONITOR_SOURCE | cut -d'=' -f2)
  current_dir=${output_folder}/games_${row_index}_qwen3_Omnifeedback_10to12iter_seed1
  mkdir $current_dir
  cd $current_dir
  nohup python ${avr_folder}/video_game_builder.py \
    --use_vllm_server \
    --model_path Qwen/Qwen3-32B \
    --vllm_server_url "http://localhost:8001" \
    --api_key token-1 \
    --assets_agent_model_path Qwen/Qwen3-32B \
    --assets_agent_vllm_server_url "http://localhost:8001" \
    --assets_agent_api_key token-1 \
    --evaluator_model_path Qwen/Qwen2.5-Omni-7B \
    --evaluator_vllm_server_url "http://localhost:8002" \
    --evaluator_api_key token-2 \
    --server_timeout 1800 \
    --feedback_max_tokens 3200 \
    --evaluation_max_tokens 8000 \
    --content_type video-game \
    --dataset ${avr_folder}/data/video_games_short.csv \
    --row_index ${row_index} \
    --chromium_path ${chromium_path} \
    --select_assets \
    --max_sample_packs 5 \
    --assets_selection individual \
    --max_assets 50 \
    --min_iterations 10 \
    --max_iterations 12 \
    --video_duration 20 \
    --video_fps 4 \
    --output_dir . \
    --asset_dir ../assets \
    --use_separate_evaluator \
    --enable_audio \
    --seed 1 \
    --display_num ${display_num} \
    --monitor_source ${MONITOR_SOURCE} \
    --top_p 0.95 \
    --repetition_penalty 1.1 \
    --temp_coding 0.6 \
    --temp_feedback 0.3 \
    --use_memory \
    --memory_len 5 \
    --multiround \
    --description_feedback \
    --coding_evaluation \
    --cdn_allowed \
    --auto_resume \
    --search_replace \
    --server_root ${output_folder} \
    2>&1 | tee output.log & pids+=($!)
  display_num=$((display_num + 1))
done

echo "Waiting for specific background processes to complete: ${pids[@]}..."
for pid in "${pids[@]}"; do
    wait "$pid" || echo "Warning: Process $pid did not exit cleanly."
done
echo "All waited-for processes have completed."
