#!/bin/bash

# Multi-Language GRPO Training Script
# This script trains a model on multiple programming languages (Python, C, C++, Java, JavaScript)

# Set CUDA device (modify this to select which GPU to use)
# Examples:
#   export CUDA_VISIBLE_DEVICES=0           # Use only GPU 0
#   export CUDA_VISIBLE_DEVICES=1           # Use only GPU 1
#   export CUDA_VISIBLE_DEVICES=0,1         # Use GPU 0 and 1
#   export CUDA_VISIBLE_DEVICES=0,1,2,3     # Use all 4 GPUs
# export CUDA_VISIBLE_DEVICES=0  # Default: use GPU 0

# Set project directory
export PROJECT_DIR="/ocean/projects/cis240137p/yfei4/secureCodeGen"
export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH}"

# Set environment variables for API access
export AI_GATEWAY_API_KEY="${AI_GATEWAY_API_KEY}"

# Training configuration
MODEL_ABBREVIATION="gemma-7B"
MODEL_NAME="/ocean/projects/cis240137p/yfei4/secureCodeGen/models/merged_PSC-2_codegemma-7b-it_sft_2-epochs"
MAX_SEQ_LENGTH=1280
LOAD_IN_4BIT="true"

# Training hyperparameters
LEARNING_RATE=1e-5
NUM_EPOCHS=10
BATCH_SIZE=1
GRAD_ACCUM_STEPS=4  # Increased to maintain effective batch size
MAX_STEPS=-1  # -1 means use num_epochs
# SAVE_STRATEGY="epoch"
# SAVE_STEPS=100
SAVE_STRATEGY="steps"
SAVE_STEPS=100
SAVE_TOTAL_LIMIT=3

# GRPO configuration (Optimized for 79GB GPU to avoid OOM)
NUM_GENERATIONS=4  # Reduced from 4 to 2 - most critical for memory!
MAX_PROMPT_LENGTH=512  # Reduced from 512
MAX_COMPLETION_LENGTH=768  # Reduced from 1024
TEMPERATURE=0.9
TOP_P=1.0

# Data paths
TRAIN_DATA="${PROJECT_DIR}/data/rl_training_data_complete.json"
EVAL_DATA=""  # Optional

# Reward configuration
SECURITY_WEIGHT=0.6
CAPABILITY_WEIGHT=0.4
TIMEOUT=10
MAX_TEST_CASES=25

# Reward component weights
REWARD_WEIGHT_FORMAT=0.2
REWARD_WEIGHT_FUNCTION=0.2
REWARD_WEIGHT_RUNNABLE=0.2
REWARD_WEIGHT_CAPABILITY=0.2
REWARD_WEIGHT_SECURITY=0.2

# Format pattern for code extraction
FORMAT_PATTERN='<think>(.+?)</think>\n<code>(.+?)</code>'

# Output configuration
TIMESTAMP=$(date +%Y%m%d-%H%M%S)
EXPERIMENT_NAME="${MODEL_ABBREVIATION}-multilang-grpo-${TIMESTAMP}"
OUTPUT_DIR="${PROJECT_DIR}/outputs/${MODEL_ABBREVIATION}_multilang_grpo_${TIMESTAMP}"
LOGGING_DIR="${OUTPUT_DIR}/logs"

# Wandb configuration
USE_WANDB="--use_wandb"  # Remove this flag to disable wandb
WANDB_PROJECT="llm-multilang-code-gen-security"
WANDB_ENTITY="arihants-carnegie-mellon-university"
WANDB_RUN_NAME="${EXPERIMENT_NAME}"

# Resume from checkpoint (optional)
# RESUME_FROM_CHECKPOINT="true"  # Auto-find latest
# RESUME_FROM_CHECKPOINT="${OUTPUT_DIR}/checkpoint-100"  # Specific checkpoint
RESUME_FROM_CHECKPOINT=""  # Don't resume

# Create output directory
mkdir -p ${OUTPUT_DIR}
mkdir -p ${LOGGING_DIR}

# Save this script to output directory
cp $0 ${OUTPUT_DIR}/training_script.sh

# Run training
python ${PROJECT_DIR}/rl/train_multilang_grpo.py \
    --model_name ${MODEL_NAME} \
    --max_seq_length ${MAX_SEQ_LENGTH} \
    --load_in_4bit ${LOAD_IN_4BIT} \
    --learning_rate ${LEARNING_RATE} \
    --num_train_epochs ${NUM_EPOCHS} \
    --per_device_train_batch_size ${BATCH_SIZE} \
    --gradient_accumulation_steps ${GRAD_ACCUM_STEPS} \
    --max_steps ${MAX_STEPS} \
    --save_strategy ${SAVE_STRATEGY} \
    --save_steps ${SAVE_STEPS} \
    --save_total_limit ${SAVE_TOTAL_LIMIT} \
    --num_generations ${NUM_GENERATIONS} \
    --max_prompt_length ${MAX_PROMPT_LENGTH} \
    --max_completion_length ${MAX_COMPLETION_LENGTH} \
    --temperature ${TEMPERATURE} \
    --top_p ${TOP_P} \
    --train_dataset_path ${TRAIN_DATA} \
    --security_weight ${SECURITY_WEIGHT} \
    --capability_weight ${CAPABILITY_WEIGHT} \
    --timeout ${TIMEOUT} \
    --max_test_cases ${MAX_TEST_CASES} \
    --reward_weight_format ${REWARD_WEIGHT_FORMAT} \
    --reward_weight_function ${REWARD_WEIGHT_FUNCTION} \
    --reward_weight_runnable ${REWARD_WEIGHT_RUNNABLE} \
    --reward_weight_capability ${REWARD_WEIGHT_CAPABILITY} \
    --reward_weight_security ${REWARD_WEIGHT_SECURITY} \
    --format_pattern "${FORMAT_PATTERN}" \
    --output_dir ${OUTPUT_DIR} \
    --logging_dir ${LOGGING_DIR} \
    --experiment_name ${EXPERIMENT_NAME} \
    ${USE_WANDB} \
    --wandb_project ${WANDB_PROJECT} \
    --wandb_entity ${WANDB_ENTITY} \
    --wandb_run_name ${WANDB_RUN_NAME} \
    --wandb_bash_script $0 \
    ${RESUME_FROM_CHECKPOINT:+--resume_from_checkpoint ${RESUME_FROM_CHECKPOINT}}

echo "Training completed! Model saved to: ${OUTPUT_DIR}"