#!/bin/bash
#SBATCH --job-name=encode_crow_missing
#SBATCH --nodes=4                   # 4 nodes for faster completion
#SBATCH --ntasks-per-node=4         # 4 GPUs per node
#SBATCH --gres=gpu:4                # Request 4 GPUs per node
#SBATCH --cpus-per-task=8           # CPUs per GPU task
#SBATCH --mem=256G                  # Memory per node
#SBATCH --time=24:00:00             # Max 24 hours
#SBATCH --output=logs/encode_crow_missing_%j.out
#SBATCH --error=logs/encode_crow_missing_%j.err
# #SBATCH --partition=gpu           # Uncomment if your cluster requires a partition

# Exit on error
set -e

echo "=========================================="
echo "SLURM Job: Encode Missing Crow T1 Files"
echo "=========================================="
echo "Job ID: $SLURM_JOB_ID"
echo "Nodes: $SLURM_JOB_NUM_NODES"
echo "GPUs per node: 4"
echo "Total GPUs: $(($SLURM_JOB_NUM_NODES * 4))"
echo "Encoding missing stem/infill conditions"
echo "=========================================="
echo ""

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

# Environment setup
echo "Setting up environment..."
source ~/.bashrc

# Activate conda environment (adjust environment name as needed)
# Uncomment and modify the line below based on your setup
# conda activate your_env_name

# Configuration
AUDIO_LIST="/home/tony/Data/Preference/crow_t1/crow_t1_missing_encodings.jsonl"
OUTPUT_DIR="/app2/suno/data/semantic_code/crow"
BATCH_SIZE=64

echo "Configuration:"
echo "  Audio List: $AUDIO_LIST"
echo "  Output:     $OUTPUT_DIR"
echo "  Batch Size: $BATCH_SIZE"
echo ""

# Check if audio list exists
if [ ! -f "$AUDIO_LIST" ]; then
    echo "ERROR: Missing encodings list not found: $AUDIO_LIST"
    echo "Please run find_missing_encodings.py first:"
    echo "  python3 /home/tony/Work/tony/RealGen/find_missing_encodings.py"
    exit 1
fi

# Count files to encode
NUM_FILES=$(wc -l < $AUDIO_LIST)
echo "Files to encode: $NUM_FILES"
echo ""

if [ $NUM_FILES -eq 0 ]; then
    echo "No files to encode! All files are already encoded."
    exit 0
fi

# Create output directory
mkdir -p $OUTPUT_DIR

# Launch encoding across all GPUs
echo "Launching multi-GPU encoding..."
echo "Start time: $(date)"
echo ""

srun python3 /home/tony/Work/tony/RealGen/encode_semantic_codes.py \
    --audio_list $AUDIO_LIST \
    --output_dir $OUTPUT_DIR \
    --batch_size $BATCH_SIZE

echo ""
echo "=========================================="
echo "Encoding completed!"
echo "End time: $(date)"
echo "Encoded files saved to: $OUTPUT_DIR"
echo "=========================================="
echo ""
echo "To verify all files are encoded, run:"
echo "  python3 /home/tony/Work/tony/RealGen/find_missing_encodings.py"