#!/bin/bash
#SBATCH --job-name=semantic_encode
#SBATCH --output=/app/suno/slurm/logs/semantic_encode_%j.txt
#SBATCH --error=/app/suno/slurm/logs/semantic_encode_%j.err
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8

# Create logs directory if it doesn't exist
mkdir -p /app/suno/slurm/logs

# Set environment variables
export CUDA_LAUNCH_BLOCKING=0
export NCCL_DEBUG=WARN
export TORCH_DISTRIBUTED_DEBUG=OFF
export TORCH_CPP_LOG_LEVEL=WARNING
export OMP_NUM_THREADS=1

# Get master node information
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12835
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER
export SLURM_NTASKS_PER_NODE=8

# Set working directory
WORK_PATH=/home/tony/Work/tony/Preference
echo "Working from $WORK_PATH"
cd $WORK_PATH

# Kill any existing processes to avoid conflicts
pkill -f 'semantic_encode.py'

# Set paths - can be overridden by command line arguments
JSONL_PATH="${1:-/app2/suno/data/dpo/sft/sft_metas_tr_v11.jsonl}"
OUTPUT_DIR="${2:-/app2/suno/data/dpo/sft/semantic_codes_v11}"

# Create output directory if it doesn't exist
mkdir -p $OUTPUT_DIR

echo "========================================"
echo "Semantic Encoding Job Information:"
echo "========================================"
echo "JSONL Path: $JSONL_PATH"
echo "Output Directory: $OUTPUT_DIR"
echo "Node: $MASTER_ADDR"
echo "GPUs: 8"
echo "========================================"

# Count total entries first
echo "Counting total entries in JSONL..."
TOTAL_ENTRIES=$(python -c "
count = 0
with open('$JSONL_PATH', 'r') as f:
    for line in f:
        if line.strip():
            count += 1
print(count)
")
echo "Total entries to process: $TOTAL_ENTRIES"
ENTRIES_PER_GPU=$((TOTAL_ENTRIES / 8))
echo "Entries per GPU: ~$ENTRIES_PER_GPU"
echo "========================================"

# Launch 8 processes, one for each GPU
echo "Launching 8 parallel encoding processes..."
for GPU_ID in {0..7}; do
    echo "Starting GPU $GPU_ID (chunk $GPU_ID)"
    
    # Launch the process in the background
    CUDA_VISIBLE_DEVICES=$GPU_ID python semantic_encode.py \
        --jsonl_path "$JSONL_PATH" \
        --output_dir "$OUTPUT_DIR" \
        --chunk_id $GPU_ID \
        --total_chunks 8 \
        2>&1 | tee /app/suno/slurm/logs/semantic_encode_gpu${GPU_ID}_${SLURM_JOB_ID}.log &
    
    # Small delay to avoid simultaneous model loading issues
    sleep 2
done

# Monitor progress
echo "========================================"
echo "Monitoring progress..."
echo "========================================"

# Wait and show progress every 30 seconds
while true; do
    # Check if any python processes are still running
    if ! pgrep -f "semantic_encode.py" > /dev/null; then
        break
    fi
    
    # Count processed files
    NPZ_COUNT=$(find $OUTPUT_DIR -name "*.npz" 2>/dev/null | wc -l)
    PROGRESS=$((NPZ_COUNT * 100 / TOTAL_ENTRIES))
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Progress: $NPZ_COUNT/$TOTAL_ENTRIES files ($PROGRESS%)"
    
    sleep 30
done

# Wait for all background processes to complete
wait

# Final statistics
echo "========================================"
echo "Encoding Complete! Final Statistics:"
echo "========================================"

python -c "
import json
import os
from pathlib import Path
import numpy as np

jsonl_path = '$JSONL_PATH'
output_dir = '$OUTPUT_DIR'

# Count total entries in JSONL
total_entries = 0
entry_ids = set()
with open(jsonl_path, 'r') as f:
    for line in f:
        if line.strip():
            try:
                entry = json.loads(line)
                entry_ids.add(entry.get('id'))
                total_entries += 1
            except:
                pass

# Count NPZ files
npz_files = list(Path(output_dir).glob('*.npz'))
npz_count = len(npz_files)
npz_ids = {f.stem for f in npz_files}

# Find missing entries
missing_ids = entry_ids - npz_ids
success_rate = npz_count/total_entries*100 if total_entries > 0 else 0

print(f'Total entries in JSONL: {total_entries}')
print(f'Total NPZ files created: {npz_count}')
print(f'Success rate: {success_rate:.2f}%')
print(f'Missing entries: {len(missing_ids)}')

# Check file sizes and shapes
if npz_files:
    sizes = []
    shapes = []
    for f in npz_files[:100]:  # Sample first 100
        sizes.append(f.stat().st_size)
        try:
            data = np.load(f)
            shapes.append(data['codes'].shape)
        except:
            pass
    
    avg_size = sum(sizes) / len(sizes) / 1024  # KB
    print(f'\\nFile Statistics (sample of {len(sizes)} files):')
    print(f'  Average file size: {avg_size:.2f} KB')
    print(f'  Min size: {min(sizes)/1024:.2f} KB')
    print(f'  Max size: {max(sizes)/1024:.2f} KB')
    
    if shapes:
        print(f'\\nShape Statistics:')
        print(f'  Sample shapes: {shapes[:5]}')
        avg_frames = sum(s[0] for s in shapes) / len(shapes)
        print(f'  Average frames: {avg_frames:.1f}')

# Save missing IDs for debugging
if missing_ids:
    missing_file = Path(output_dir) / 'missing_ids.txt'
    with open(missing_file, 'w') as f:
        for mid in sorted(missing_ids):
            f.write(f'{mid}\\n')
    print(f'\\nMissing IDs saved to: {missing_file}')
"

echo "========================================"
echo "Job completed at: $(date)"
echo "Log files available at: /app/suno/slurm/logs/"
echo "========================================"