#!/bin/bash
#SBATCH --job-name=semantic_encode_16n
#SBATCH --output=/app/suno/slurm/logs/semantic_encode_16n_%A_%a.txt
#SBATCH --error=/app/suno/slurm/logs/semantic_encode_16n_%A_%a.err
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --array=0-15%16  # 16 array jobs, all can run simultaneously

# 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 node information
export NODE_ID=$SLURM_ARRAY_TASK_ID
export TOTAL_NODES=16
export GPUS_PER_NODE=8
export TOTAL_CHUNKS=$((TOTAL_NODES * GPUS_PER_NODE))  # 128 total chunks

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

# 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 Multi-Node Job Information:"
echo "========================================"
echo "Array Task ID (Node): $NODE_ID / $TOTAL_NODES"
echo "JSONL Path: $JSONL_PATH"
echo "Output Directory: $OUTPUT_DIR"
echo "Total chunks across all nodes: $TOTAL_CHUNKS"
echo "Chunks for this node: $((NODE_ID * GPUS_PER_NODE)) - $(((NODE_ID + 1) * GPUS_PER_NODE - 1))"
echo "========================================"

# Only count entries on the first node to avoid redundant work
if [ $NODE_ID -eq 0 ]; then
    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_CHUNK=$((TOTAL_ENTRIES / TOTAL_CHUNKS))
    echo "Entries per chunk: ~$ENTRIES_PER_CHUNK"
    
    # Save this info for other nodes
    echo "$TOTAL_ENTRIES" > "${OUTPUT_DIR}/.total_entries"
fi
echo "========================================"

# Launch 8 processes on this node, one for each GPU
echo "Node $NODE_ID: Launching 8 parallel encoding processes..."
for LOCAL_GPU_ID in {0..7}; do
    # Calculate global chunk ID
    GLOBAL_CHUNK_ID=$((NODE_ID * GPUS_PER_NODE + LOCAL_GPU_ID))
    
    echo "Node $NODE_ID, GPU $LOCAL_GPU_ID: Starting chunk $GLOBAL_CHUNK_ID"
    
    # Launch the process in the background
    CUDA_VISIBLE_DEVICES=$LOCAL_GPU_ID python semantic_encode.py \
        --jsonl_path "$JSONL_PATH" \
        --output_dir "$OUTPUT_DIR" \
        --chunk_id $GLOBAL_CHUNK_ID \
        --total_chunks $TOTAL_CHUNKS \
        --local_gpu_id $LOCAL_GPU_ID \
        2>&1 | tee /app/suno/slurm/logs/semantic_encode_node${NODE_ID}_gpu${LOCAL_GPU_ID}_${SLURM_ARRAY_JOB_ID}.log &
    
    # Small delay to avoid simultaneous model loading issues
    sleep 1
done

# Monitor progress for this node
echo "========================================"
echo "Node $NODE_ID: Monitoring progress..."
echo "========================================"

# Wait for all processes on this node to complete
wait

# Report completion for this node
echo "========================================"
echo "Node $NODE_ID: Completed!"
echo "========================================"

# If this is the last node, generate final statistics
if [ $NODE_ID -eq $((TOTAL_NODES - 1)) ]; then
    # Wait a bit for other nodes to finish writing
    sleep 30
    
    echo "========================================"
    echo "Generating Final Statistics (from last node)..."
    echo "========================================"
    
    python -c "
import json
import os
from pathlib import Path
import numpy as np

jsonl_path = '$JSONL_PATH'
output_dir = '$OUTPUT_DIR'
total_nodes = $TOTAL_NODES
total_chunks = $TOTAL_CHUNKS

# 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'Final Statistics for {total_nodes} nodes, {total_chunks} chunks:')
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
    
    if sizes:
        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(len(s) if len(s.shape) == 1 else 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}')

# Save summary
summary = {
    'total_nodes': total_nodes,
    'total_chunks': total_chunks,
    'total_entries': total_entries,
    'npz_count': npz_count,
    'success_rate': success_rate,
    'missing_count': len(missing_ids)
}
summary_file = Path(output_dir) / 'encoding_summary.json'
with open(summary_file, 'w') as f:
    json.dump(summary, f, indent=2)
print(f'\\nSummary saved to: {summary_file}')
"
fi

echo "========================================"
echo "Node $NODE_ID completed at: $(date)"
echo "========================================"