#!/bin/bash
# 一键复现 exp_q 提交：从模型权重出发，跑推理 + 后处理 → 最终 submission JSON
#
# 用法:
#   bash run_inference.sh \
#       /path/to/track_1_test.json \
#       /path/to/PortraitCraft   # 图片根目录 (images_00/ images_01/ ...)

set -e

TEST_JSON="${1:?usage: bash run_inference.sh <test_json> <images_base>}"
IMAGES_BASE="${2:?usage: bash run_inference.sh <test_json> <images_base>}"

ROOT="$(cd "$(dirname "$0")" && pwd)"
MODEL="$ROOT/model"
WORK="$ROOT/work_dir"
mkdir -p "$WORK"

echo "=================================================="
echo "exp_q reproducibility inference"
echo "model:       $MODEL"
echo "test json:   $TEST_JSON"
echo "images:      $IMAGES_BASE"
echo "=================================================="

# 检查 GPU 数量
N_GPU=$(nvidia-smi --list-gpus 2>/dev/null | wc -l)
echo "Detected $N_GPU GPUs"

if [ "$N_GPU" -ge 2 ]; then
    echo "Running std and hr in parallel on GPU 0 and 1..."
    CUDA_VISIBLE_DEVICES=0 python "$ROOT/inference/infer_score.py" \
        --model "$MODEL" \
        --test-json "$TEST_JSON" \
        --images-base "$IMAGES_BASE" \
        --output "$WORK/exp_q_std_tta.json" \
        --tta flip --max-pixels 1003520 --tp 1 \
        2>&1 | tee "$WORK/std.log" &
    PID_STD=$!

    CUDA_VISIBLE_DEVICES=1 python "$ROOT/inference/infer_score.py" \
        --model "$MODEL" \
        --test-json "$TEST_JSON" \
        --images-base "$IMAGES_BASE" \
        --output "$WORK/exp_q_hr_tta.json" \
        --tta flip --max-pixels 2007040 --tp 1 \
        2>&1 | tee "$WORK/hr.log" &
    PID_HR=$!

    wait $PID_STD
    wait $PID_HR
else
    echo "Running sequentially on GPU 0..."
    CUDA_VISIBLE_DEVICES=0 python "$ROOT/inference/infer_score.py" \
        --model "$MODEL" \
        --test-json "$TEST_JSON" \
        --images-base "$IMAGES_BASE" \
        --output "$WORK/exp_q_std_tta.json" \
        --tta flip --max-pixels 1003520 --tp 1

    CUDA_VISIBLE_DEVICES=0 python "$ROOT/inference/infer_score.py" \
        --model "$MODEL" \
        --test-json "$TEST_JSON" \
        --images-base "$IMAGES_BASE" \
        --output "$WORK/exp_q_hr_tta.json" \
        --tta flip --max-pixels 2007040 --tp 1
fi

echo "=================================================="
echo "Building submission JSON ..."
python "$ROOT/inference/make_submission.py" \
    --std "$WORK/exp_q_std_tta.json" \
    --hr  "$WORK/exp_q_hr_tta.json" \
    --test-json "$TEST_JSON" \
    --output "$ROOT/submission.json"

# 打包提交 zip
cd "$ROOT"
cp submission.json track_1_test.json
zip -j submission.zip track_1_test.json
rm track_1_test.json
echo "Created submission.zip at $ROOT/submission.zip"
echo "=================================================="