#!/bin/bash

if [ $# -ne 5 ]; then
  echo "Usage: $0 <directory> <pattern> <N> <num_iters> <output_dir>"
  exit 1
fi

directory="$1"
pattern="$2"
N="$3"
num_iters="$4"
output_dir="$5"

# Check if the directory exists
if [ ! -d "$directory" ]; then
  echo "Error: Directory not found."
  exit 1
fi

# Create the output directory if it doesn't exist
mkdir -p "$output_dir"

# Get the list of files and subdirectories in the directory
contents=($(find "$directory" -type f -name "$pattern"))

# Calculate the total number of items
num_items=${#contents[@]}

# Calculate the approximate number of items in each list
items_per_list=$((num_items / N))
remainder=$((num_items % N))

# Initialize variables
start_index=0
background_pids=()

# Split the items into N lists and execute the Python command in parallel
for ((i = 1; i <= N; i++)); do
  end_index=$((start_index + items_per_list - 1))

  # If there's a remainder, distribute the remaining items
  if [ $i -le $remainder ]; then
    end_index=$((end_index + 1))
  fi

  # Create a sublist from start_index to end_index
  sublist=("${contents[@]:start_index:end_index - start_index + 1}")

  # Generate the output path
  output_path="$output_dir/$i.pt"

  # Generate the Python command using the sublist, num_iters, and output_path
  cuda_visible_devices=$((i - 1))
  cmd="CUDA_VISIBLE_DEVICES=$cuda_visible_devices python generate_zero_shot_segments.py --input_pattern ${sublist[*]} --num_iter $num_iters --output $output_path"

  # Execute the command in the background
#  if [ $i -eq 3 ]; then
#    eval "$cmd" &
#  fi
   eval "$cmd" &

  # Record the background process ID
  background_pids+=($!)

  # Update the start_index for the next sublist
  start_index=$((end_index + 1))
done

# Wait for all background processes to finish
for pid in "${background_pids[@]}"; do
  wait "$pid"
done