#!/bin/bash
#SBATCH --job-name=test_crow_t1
#SBATCH --nodes=1                   # Single node for testing
#SBATCH --ntasks-per-node=1         # Single GPU for testing
#SBATCH --gres=gpu:1                # Request 1 GPU
#SBATCH --cpus-per-task=8           # CPUs per GPU task
#SBATCH --mem=64G                   # Memory
#SBATCH --time=02:00:00             # Max 2 hours for testing
#SBATCH --output=logs/test_crow_t1_%j.out
#SBATCH --error=logs/test_crow_t1_%j.err
# #SBATCH --partition=gpu           # Uncomment if your cluster requires a partition

# Exit on error
set -e

echo "=========================================="
echo "SLURM Job: Crow T1 Semantic Encoding Test"
echo "=========================================="
echo "Job ID: $SLURM_JOB_ID"
echo "Testing on small sample (~100 files)"
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_audio_list.jsonl"
OUTPUT_DIR="/app2/suno/data/semantic_code/crow"
BATCH_SIZE=16
SAMPLE_SIZE=100

echo "Configuration:"
echo "  Audio List: $AUDIO_LIST"
echo "  Output:     $OUTPUT_DIR"
echo "  Batch Size: $BATCH_SIZE"
echo "  Sample:     First $SAMPLE_SIZE files"
echo ""

# Check if audio list exists
if [ ! -f "$AUDIO_LIST" ]; then
    echo "ERROR: Audio list not found: $AUDIO_LIST"
    echo "Please run create_crow_t1_audio_list.py first to generate the audio list."
    exit 1
fi

# Create temporary sample file
TEMP_SAMPLE="/tmp/crow_t1_test_sample_${SLURM_JOB_ID}.jsonl"
echo "Creating sample file with first $SAMPLE_SIZE entries..."
head -n $SAMPLE_SIZE $AUDIO_LIST > $TEMP_SAMPLE

# Create output directory
mkdir -p $OUTPUT_DIR

# Launch encoding on single GPU
echo "Launching single-GPU encoding on sample..."
echo ""

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

# Clean up temporary file
rm -f $TEMP_SAMPLE

echo ""
echo "=========================================="
echo "Test completed!"
echo "Encoded files saved to: $OUTPUT_DIR"
echo "=========================================="
echo ""
echo "If test successful, submit full job with:"
echo "  sbatch submit_encode_crow_t1.sh"