#!/bin/bash
#SBATCH --job-name=dq_filter
#SBATCH --output=/home/tony/slurm/logs/run_dq_filter_%j.txt
#SBATCH --error=/home/tony/slurm/logs/run_dq_filter_%j.err
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8

# Create logs directory if it doesn't exist
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
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


TRAIN_PATH=/home/tony/Work/tony/Preference
echo working from $TRAIN_PATH
cd $TRAIN_PATH
pkill -f 'spawn_main'
# Calculate chunk size based on total dataset size
# We'll pass this as an argument to our script
START_IDX="${1:-0}"    # Default to 0 if not specified
INPUT_FILE_NAME="${2:-interesting_clips_dorado_t1_20250929.pkl}"    # Default if not specified
echo "START_IDX: $START_IDX"
echo "INPUT_FILE_NAME: $INPUT_FILE_NAME"

# Extract the dataset identifier from the filename (e.g., carp_t1 from interesting_clips_carp_t1_20250828.pkl)
# This regex extracts the part between "interesting_clips_" and the date/version suffix
DATASET_ID=$(echo "$INPUT_FILE_NAME" | sed -E 's/interesting_clips_([^_]+_[^_]+)_.*/\1/')
echo "DATASET_ID: $DATASET_ID"

# Launch 8 processes, one for each GPU
for GPU_ID in {0..7}; do
# for GPU_ID in {8..15}; do
# for GPU_ID in {16..23}; do
# for GPU_ID in {24..31}; do
    # Calculate the start index for this GPU
    # Each GPU will handle its own portion of the data
    PROCESS_START_IDX=$((START_IDX + GPU_ID))
    echo "Processing GPU $GPU_ID with start index $PROCESS_START_IDX"
    # Launch the process in the background
    # --input_file_name "interesting_clips_upv2_u2_20250425_full.pkl" \
    CUDA_VISIBLE_DEVICES=$GPU_ID /home/tony/anaconda3/envs/suno_env_dev/bin/python s_all_filter_full.py \
        --input_file_name "$INPUT_FILE_NAME" \
        --job_idx $PROCESS_START_IDX &
done

# Wait for all background processes to complete
wait

/home/tony/anaconda3/envs/suno_env_dev/bin/python -c "
import json
# Use the extracted dataset ID for the output path
dataset_id = '$DATASET_ID'
base_path = f'/home/tony/Data/Preference/{dataset_id}'
with open(f'{base_path}/full_pair_quality.json', 'r') as f:
   result = json.load(f)
print('original', len(result))
for job_idx in range(8):
    with open(f'{base_path}/full_pair_quality_{job_idx}.json', 'r') as fp:
        current_result = json.load(fp)
        result.update(current_result)
print('updated', len(result))

with open(f'{base_path}/full_pair_quality.json', 'w') as f:
    json.dump(result, f, indent=4)
"
echo "DONE!"