#!/bin/bash
# 특정 문제에 대해 TTRLVR + AZR 통합 학습 실행

# 기본 설정
GPU_ID=${1:-5}
PROBLEM_ID=${2:-"Mbpp/7"}
ROUNDS=${3:-3}

# 환경 변수 설정
export CUDA_VISIBLE_DEVICES=$GPU_ID
export VLLM_ATTENTION_BACKEND=FLASH_ATTN
export RAY_memory_monitor_refresh_ms=0
export RAY_LOGGING_LEVEL=DEBUG
export HYDRA_FULL_ERROR=1
export PYTHONPATH="${PYTHONPATH}:/home/ubuntu/RLVR/verl:/home/ubuntu/RLVR/TestTime-RLVR-v2"

# 로그 디렉토리
LOG_DIR="/home/ubuntu/RLVR/TestTime-RLVR-v2/logs"
mkdir -p $LOG_DIR

# 타임스탬프
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
SAFE_PROBLEM_ID=$(echo $PROBLEM_ID | tr '/' '_')
LOG_FILE="$LOG_DIR/single_problem_${SAFE_PROBLEM_ID}_${TIMESTAMP}.log"

echo "🎯 특정 문제 학습 시작"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "📌 문제 ID: $PROBLEM_ID"
echo "🔄 라운드 수: $ROUNDS"
echo "🖥️  GPU: $GPU_ID"
echo "📝 로그 파일: $LOG_FILE"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""

# train_ttrlvr_azr.py 실행
python /home/ubuntu/RLVR/TestTime-RLVR-v2/test/train_ttrlvr_azr.py \
    --benchmark mbpp \
    --problem-id "$PROBLEM_ID" \
    --rounds $ROUNDS \
    --gpu $GPU_ID \
    --model "Qwen/Qwen2.5-7B" \
    2>&1 | tee $LOG_FILE

echo ""
echo "✅ 학습 완료. 로그 확인: $LOG_FILE"