#!/bin/bash
#
# Full pipeline: CogVideoX video generation → SV4D 4D generation → Gallery collection
#
# Usage: bash scripts/sampling/run_all.sh
#
# To skip a completed stage, comment out the corresponding section below.

set -e

# ============== Config (edit here) ==============
FLUX_IMAGE_DIR="/home/wanghongbo06/baipurui/generative-models/prepare/flux_outputs_batch/select_outputs_batch"
PROMPT_FILE="/home/wanghongbo06/baipurui/generative-models/prepare/image_prompts/select_prompts_3.txt"
VIDEO_DIR="/home/wanghongbo06/baipurui/generative-models/prepare/cogvideo_batch/select_cogvideo_batch_3"
SV4D_OUTPUT_DIR="/home/wanghongbo06/baipurui/generative-models/outputs/outputs_4d/outputs_bad_7"
GALLERY_DIR="/home/wanghongbo06/baipurui/generative-models/outputs/gallery/gallery_bad_7"

NUM_GPUS=8
SV4D_NUM_STEPS=4
SV4D_MODEL_PATH="checkpoints/sv4d.safetensors"
# ================================================

# echo "============================================"
# echo " Stage 1/3: CogVideoX Image-to-Video"
# echo "============================================"
# python prepare/video_generate.py \
#     --input_dir "$FLUX_IMAGE_DIR" \
#     --output_dir "$VIDEO_DIR" \
#     --prompt_file "$PROMPT_FILE" \
#     --num_gpus "$NUM_GPUS"

echo ""
echo "============================================"
echo " Stage 1.5: Sanitize filenames (remove commas)"
echo "============================================"
for f in "$VIDEO_DIR"/*.mp4; do
    [ -f "$f" ] || continue
    safe_f="${f//,/_}"
    if [ "$f" != "$safe_f" ]; then
        mv "$f" "$safe_f"
        echo "Renamed: $(basename "$f") -> $(basename "$safe_f")"
    fi
done

echo ""
echo "============================================"
echo " Stage 2/3: SV4D 4D Generation"
echo "============================================"

SV4D_SCRIPT="scripts/sampling/simple_video_sample_4d.py"

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

if [ "$TOTAL" -eq 0 ]; then
    echo "No .mp4 files found in $VIDEO_DIR, skipping SV4D."
else
    echo "Found $TOTAL videos, distributing across $NUM_GPUS GPUs"
    mkdir -p "$SV4D_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="${SV4D_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 $SV4D_SCRIPT \
                "--input_path=$vpath" \
                "--output_folder=$out_dir" \
                "--num_steps=$SV4D_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
    }

    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 "Waiting for all $NUM_GPUS GPU workers..."
    wait
    echo "SV4D generation complete."
fi

echo ""
echo "============================================"
echo " Stage 3/3: Collect Gallery"
echo "============================================"
python scripts/sampling/collect_results.py \
    --input_dir "$SV4D_OUTPUT_DIR" \
    --output_dir "$GALLERY_DIR"

echo ""
echo "============================================"
echo " All done!"
echo " Gallery: $GALLERY_DIR"
echo "============================================"