#!/bin/bash

# Bash loop version of SLURM script for PerVA dataset
# Usage: ./run_PerVA_loop.sh 23

# Check if seed argument is provided
if [ $# -ne 1 ]; then
    echo "Usage: $0 <seed>"
    echo "Example: $0 23"
    exit 1
fi

# Set parameters
seed=$1
data='PerVA'
concept_list_file="data/perva-data/${data}_concept_list.txt"

# Check if concept list file exists
if [ ! -f "$concept_list_file" ]; then
    echo "Error: Concept list file not found: $concept_list_file"
    exit 1
fi

# Create logs directory if it doesn't exist
mkdir -p logs

# Activate conda environment (adjust path as needed)
# source /gpfs/projects/ehpc171/ddas/envs/mini_cpm/bin/activate
# Uncomment and adjust the above line for your environment

# Set environment variable (if needed)
export HF_HUB_OFFLINE=1

# Get total number of lines in concept list
total_concepts=$(wc -l < "$concept_list_file")
echo "Found $total_concepts concepts to process"

# Process each concept (equivalent to SLURM array 0-328, so 329 concepts)
for i in $(seq 0 328); do
    # Get the concept entry (line number = i + 1)
    line_number=$((i + 1))
    
    # Check if line exists
    if [ $line_number -gt $total_concepts ]; then
        echo "Warning: Line $line_number exceeds total concepts ($total_concepts). Skipping."
        continue
    fi
    
    # Extract concept information
    sks_entry=$(sed -n "${line_number}p" "$concept_list_file")
    
    if [ -z "$sks_entry" ]; then
        echo "Warning: Empty line at position $line_number. Skipping."
        continue
    fi
    
    category_name=$(echo "$sks_entry" | cut -d, -f1)
    sks_name=$(echo "$sks_entry" | cut -d, -f2)
    
    echo "[$i/$((total_concepts-1))] Processing category: $category_name, concept: $sks_name"
    
    # Create log file name similar to SLURM format
    log_file="logs/perva_loop_${seed}_${i}.out"
    
    # Run the main script and redirect output to log file
    python src/main.py --concept_name "$sks_name" --category_name "$category_name" \
        --seed "$seed" --data_name "$data" --rerank_early --refined_k 3 \
        --attribute_based_step_by_step --two_step --input_type text --only_recall \
        --model_name mini_cpm > "$log_file" 2>&1
    
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo "[$i/$((total_concepts-1))] Successfully completed: $sks_name"
    else
        echo "[$i/$((total_concepts-1))] Error processing: $sks_name (check $log_file)"
    fi
    
    # Optional: Add a small delay to prevent overwhelming the system
    # sleep 1
done

echo "All concepts processed! Check logs/ directory for individual outputs."