#!/bin/bash
#
# Batch SV4D generation across 8 GPUs.
# Each GPU gets a chunk of videos and processes them sequentially.
#
# Usage (run from project root):
#   bash scripts/sampling/batch_sv4d.sh ./cogvideo_batch ./outputs_4d
#   bash scripts/sampling/batch_sv4d.sh ./cogvideo_batch ./outputs_4d 8 50   # num_gpus, num_steps

INPUT_DIR="${1:?Usage: bash scripts/sampling/batch_sv4d.sh <input_dir> <output_dir> [num_gpus] [num_steps]}"
OUTPUT_DIR="${2:?Usage: bash scripts/sampling/batch_sv4d.sh <input_dir> <output_dir> [num_gpus] [num_steps]}"
NUM_GPUS="${3:-8}"
NUM_STEPS="${4:-50}"

MODEL_PATH="checkpoints/sv4d2.safetensors"
SCRIPT="scripts/sampling/simple_video_sample_4d2.py"

mapfile -t ALL_VIDEOS < <(find "$INPUT_DIR" -maxdepth 1 -name "*.mp4" | sort)
TOTAL=${#ALL_VIDEOS[@]}

if [ "$TOTAL" -eq 0 ]; then
    echo "No .mp4 files found in $INPUT_DIR"
    exit 1
fi

echo "Found $TOTAL videos, distributing across $NUM_GPUS GPUs"
mkdir -p "$OUTPUT_DIR"

process_gpu() {
    local gpu_id=$1
    shift
    local videos=("$@")
    local count=${#videos[@]}

    for i in "${!videos[@]}"; do
        local vpath="${videos[$i]}"
        local vname=$(basename "${vpath%.*}")
        local out_dir="${OUTPUT_DIR}/${vname}"

        if [ -f "${out_dir}/_done" ]; then
            echo "[GPU $gpu_id] ($((i+1))/$count) Skip done: $vname"
            continue
        fi

        echo "[GPU $gpu_id] ($((i+1))/$count) Processing: $vname"

        CUDA_VISIBLE_DEVICES=$gpu_id python $SCRIPT \
            "--input_path=$vpath" \
            "--model_path=$MODEL_PATH" \
            "--output_folder=$out_dir" \
            "--num_steps=$NUM_STEPS"

        if [ $? -eq 0 ]; then
            touch "${out_dir}/_done"
            echo "[GPU $gpu_id] ($((i+1))/$count) Done: $vname"
        else
            echo "[GPU $gpu_id] ($((i+1))/$count) FAILED: $vname"
        fi
    done
}

# Distribute videos round-robin across GPUs
for ((g=0; g<NUM_GPUS; g++)); do
    GPU_VIDEOS=()
    for ((j=g; j<TOTAL; j+=NUM_GPUS)); do
        GPU_VIDEOS+=("${ALL_VIDEOS[$j]}")
    done
    if [ ${#GPU_VIDEOS[@]} -gt 0 ]; then
        process_gpu $g "${GPU_VIDEOS[@]}" &
        echo "Launched GPU $g with ${#GPU_VIDEOS[@]} videos"
    fi
done

echo "All $NUM_GPUS GPU workers launched. Waiting..."
wait
echo "All done! Results in $OUTPUT_DIR/<video_name>/sv4d2/"