#!/bin/bash

cd ./VLMEvalKit

# Function to print usage
print_usage() {
    echo "Usage: $0 [INFO] [STEP_PATTERN]"
    echo "Examples:"
    echo "  $0                          # Run default Qwen models"
    echo "  $0 QWEN25-3B-1             # Run models from PROJECT_DIR/EasyR1/checkpoints/QWEN25-3B-1"
    echo "  $0 QWEN25-3B-1 '5000|10000' # Run only steps 5000 and 10000 from the specified directory"
    echo "  $0 QWEN25-3B-1 '5000:'     # Run step 5000 and all steps after it from the specified directory"
}

# Show usage if help requested
if [[ "$1" == "-h" || "$1" == "--help" ]]; then
    print_usage
    exit 0
fi

# Always use Qwen model
MODEL_TYPE="Custom-Qwen2-VL-7B-Instruct"

# Default models
default_models=(
    "Qwen/Qwen2.5-VL-7B-Instruct"
)

# Base checkpoints directory
BASE_CHECKPOINTS_DIR="PROJECT_DIR/VerifiableSynthesis/EasyR1/checkpoints"

# Set up models based on arguments
if [ $# -eq 0 ]; then
    # No arguments - use default models
    selected_models=("${default_models[@]}")
    LOGS_DIR="./logs_vlm/BASELINE"
    using_defaults=true
else
    INFO="$1"
    STEP_PATTERN="$2"
    
    # Set checkpoints directory with INFO appended
    CHECKPOINTS_DIR="${BASE_CHECKPOINTS_DIR}/${INFO}"
    
    # Create logs directory using the INFO for organization
    LOGS_DIR="./logs_vlm/${INFO}"
    
    # Check if the checkpoints directory exists
    if [ ! -d "$CHECKPOINTS_DIR" ]; then
        echo "Error: Checkpoints directory '$CHECKPOINTS_DIR' not found!"
        exit 1
    fi
    
    # Find all step directories first and sort them
    all_steps=($(find "$CHECKPOINTS_DIR" -maxdepth 1 -type d -path "*global_step_*" | sort -V))
    
    # Check if any steps were found
    if [ ${#all_steps[@]} -eq 0 ]; then
        echo "No checkpoint steps found in '$CHECKPOINTS_DIR'!"
        exit 1
    fi
    
    # Process the step pattern if provided
    if [ -n "$STEP_PATTERN" ]; then
        # Check if the pattern contains a colon (indicating range)
        if [[ "$STEP_PATTERN" == *":"* ]]; then
            # Extract the starting step number
            start_step="${STEP_PATTERN%:}"
            
            # Filter steps that are >= start_step
            selected_steps=()
            for step in "${all_steps[@]}"; do
                # Extract the step number from the path
                step_num=$(basename "$step" | sed -E 's/.*global_step_([0-9]+).*/\1/')
                
                # Add to selected steps if >= start_step
                if [ "$step_num" -ge "$start_step" ]; then
                    selected_steps+=("$step")
                fi
            done
        else
            # Use the original pattern-based filtering for patterns like '5000|10000'
            selected_steps=($(find "$CHECKPOINTS_DIR" -maxdepth 1 -type d -path "*global_step_*" | grep -E "global_step_(${STEP_PATTERN})" | sort -V))
        fi
    else
        # No pattern provided, use all steps
        selected_steps=("${all_steps[@]}")
    fi
    
    # Check if any steps were selected after filtering
    if [ ${#selected_steps[@]} -eq 0 ]; then
        echo "No matching checkpoint steps found for pattern '$STEP_PATTERN'!"
        exit 1
    fi
    
    using_defaults=false
fi

# Define tasks/datasets
datasets=(
    "MathVista_MINI"
    "WeMath"
    "MathVerse_MINI"
    "MathVision"
    # Add more datasets as needed
)

# Print selected configuration
if [ "$using_defaults" = true ]; then
    echo "Using default models:"
    for model in "${default_models[@]}"; do
        echo "  - $model"
    done
else
    echo "Selected checkpoint steps:"
    for step in "${selected_steps[@]}"; do
        echo "  - $(basename "$step")"
    done
fi

echo "Using model type: $MODEL_TYPE"
echo "Logs directory: $LOGS_DIR"
echo "----------------------------------------"

# Iterate over models/checkpoints and datasets
if [ "$using_defaults" = true ]; then
    # Use default models
    for model in "${default_models[@]}"; do
        for dataset in "${datasets[@]}"; do
            # Define the work directory
            work_dir="$LOGS_DIR/$(basename "$model")/$dataset"
            
            echo "Evaluating Model on Dataset: $dataset"
            echo "Model Path: $model"
            echo "Work Directory: $work_dir"
            
            # Create the work directory if it doesn't exist
            mkdir -p "$work_dir"
            
            # Run the evaluation
            torchrun --nproc_per_node=8 \
                    --master_port=12345 \
                    run.py \
                    -- \
                    --data "$dataset" \
                    --model "$MODEL_TYPE" \
                    --model-path "$model" \
                    --mode infer \
                    --prompt_type reasoning \
                    --work-dir "$work_dir" \
                    --verbose
            
            echo "Evaluation completed for Model on Dataset: $dataset"
            echo "-----------------------------------------"
        done
    done
else
    # Use custom checkpoints
    for step_dir in "${selected_steps[@]}"; do
        # The actual model path includes /actor/huggingface
        model_path="${step_dir}/actor/huggingface"
        
        # Check if the huggingface path exists
        if [ ! -d "$model_path" ]; then
            echo "Warning: Model path '$model_path' not found, skipping..."
            continue
        fi
        
        step_name=$(basename "$step_dir")
        
        for dataset in "${datasets[@]}"; do
            # Define the work directory
            work_dir="$LOGS_DIR/${step_name}/$dataset"
            
            echo "Evaluating Checkpoint on Dataset: $dataset"
            echo "Model Path: $model_path"
            echo "Work Directory: $work_dir"
            
            # Create the work directory if it doesn't exist
            mkdir -p "$work_dir"
            
            # Run the evaluation
            torchrun --nproc_per_node=8 \
                    --master_port=12345 \
                    run.py \
                    --data "$dataset" \
                    --model "$MODEL_TYPE" \
                    --model-path "$model_path" \
                    --mode infer \
                    --prompt_type reasoning \
                    --work-dir "$work_dir" \
                    --verbose
            
            echo "Evaluation completed for Checkpoint on Dataset: $dataset"
            echo "-----------------------------------------"
        done
    done
fi

cd ..

echo "Done"