#!/bin/bash
#SBATCH --job-name=encode_crow_t1
#SBATCH --nodes=8                   # Number of nodes (adjust as needed)
#SBATCH --ntasks-per-node=8         # 8 GPUs per node
#SBATCH --gres=gpu:8                # Request 8 GPUs per node
#SBATCH --cpus-per-task=8           # CPUs per GPU task
#SBATCH --mem=256G                  # Memory per node
#SBATCH --time=48:00:00             # Max 48 hours (large dataset ~2.7M clips)
#SBATCH --output=logs/encode_crow_t1_%j.out
#SBATCH --error=logs/encode_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"
echo "=========================================="
echo "Job ID: $SLURM_JOB_ID"
echo "Nodes: $SLURM_JOB_NUM_NODES"
echo "Tasks per node: $SLURM_NTASKS_PER_NODE"
echo "Total tasks: $SLURM_NTASKS"
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

# Set environment variables for distributed training
export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n 1)
export MASTER_PORT=29500
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=ib0

echo "Master address: $MASTER_ADDR"
echo "Master port: $MASTER_PORT"
echo ""

# 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
N_CODEBOOKS=""  # Empty = use all available (50 codebooks from centroids)

echo "Configuration:"
echo "  Audio List: $AUDIO_LIST"
echo "  Output:     $OUTPUT_DIR"
echo "  Batch Size: $BATCH_SIZE"
echo "  Codebooks:  All available (50)"
echo "  Centroids:  /home/minz/temp/mert_768d_centroids_4000_50.npy (hardcoded in script)"
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 output directory
mkdir -p $OUTPUT_DIR

# Count total files to encode
TOTAL_FILES=$(wc -l < $AUDIO_LIST)
echo "Total audio files to encode: $TOTAL_FILES"
echo ""

# Launch distributed encoding job
echo "Launching distributed encoding..."
echo ""

if [ -z "$N_CODEBOOKS" ]; then
    # Use all available codebooks (no --n_codebooks argument)
    srun python3 /home/tony/Work/tony/RealGen/encode_semantic_codes.py \
        --audio_list $AUDIO_LIST \
        --output_dir $OUTPUT_DIR \
        --batch_size $BATCH_SIZE
else
    # Use specific number of codebooks
    srun python3 /home/tony/Work/tony/RealGen/encode_semantic_codes.py \
        --audio_list $AUDIO_LIST \
        --output_dir $OUTPUT_DIR \
        --batch_size $BATCH_SIZE \
        --n_codebooks $N_CODEBOOKS
fi

echo ""
echo "=========================================="
echo "Job completed!"
echo "=========================================="