#!/bin/bash
set -e
source /root/miniconda3/etc/profile.d/conda.sh

ROOT=/root/autodl-tmp/SplatAtlas
ITER=30000

# 方法 -> 环境 + 源码根目录
declare -A ENV_MAP
declare -A CODE_MAP

ENV_MAP[lod_gs]="/root/autodl-tmp/envs/lod_gs"
CODE_MAP[lod_gs]="/root/autodl-tmp/LOD_GS"

ENV_MAP[coadaptgs]="/root/autodl-tmp/envs/co_adaptation_3dgs"
CODE_MAP[coadaptgs]="/root/autodl-tmp/CoAdaptGS"

get_source_res() {
    case "$1" in
        auditorium|ballroom|barn|caterpillar|courtroom|lighthouse|museum|palace|playground|temple|train|truck)
            echo "/root/autodl-tmp/dataset/tnt/$1 2" ;;
        bicycle|flowers|garden|stump|treehill)
            echo "/root/autodl-tmp/dataset/360/$1 4" ;;
        bonsai|counter|kitchen|room)
            echo "/root/autodl-tmp/dataset/360/$1 2" ;;
        Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship)
            echo "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/$1 1" ;;
        DrJohnson|Playroom)
            echo "/root/autodl-tmp/dataset/deepblending_clean/$1 1" ;;
        *)
            echo ""; return 1 ;;
    esac
}

MISSING=(
    "coadaptgs auditorium"
    "coadaptgs ballroom"
    "coadaptgs barn"
    "coadaptgs caterpillar"
    "coadaptgs courtroom"
    "coadaptgs lighthouse"
    "coadaptgs museum"
    "coadaptgs palace"
    "coadaptgs playground"
    "coadaptgs temple"
    "coadaptgs train"
    "coadaptgs truck"
    "coadaptgs bicycle"
    "coadaptgs bonsai"
    "coadaptgs counter"
    "coadaptgs flowers"
    "coadaptgs garden"
    "coadaptgs kitchen"
    "coadaptgs room"
    "coadaptgs stump"
    "coadaptgs treehill"
    "coadaptgs DrJohnson"
    "coadaptgs Playroom"
    "lod_gs Chair"
    "lod_gs Drums"
    "lod_gs Ficus"
    "lod_gs Hotdog"
    "lod_gs Lego"
    "lod_gs Materials"
    "lod_gs Mic"
    "lod_gs Ship"
)

run_step() {
    local name="$1"
    local cmd="$2"
    echo "  [Step] $name ..."
    if eval "$cmd"; then
        echo "  [OK]"
        return 0
    else
        echo "  [FAIL]"
        return 1
    fi
}

for entry in "${MISSING[@]}"; do
    read -r method scene <<< "$entry"
    env_path="${ENV_MAP[$method]}"
    code_root="${CODE_MAP[$method]}"
    if [ -z "$env_path" ] || [ -z "$code_root" ]; then
        echo "Unknown method or code root for $method"; continue
    fi

    read -r source_path resolution < <(get_source_res "$scene")
    if [ -z "$source_path" ]; then
        echo "Skipping $method/$scene: unknown dataset"; continue
    fi

    model_path="$ROOT/outputs/${method}_${scene}"
    ply_path="$model_path/point_cloud/iteration_${ITER}/point_cloud.ply"
    if [ ! -f "$ply_path" ]; then
        echo "Skipping $method/$scene: PLY not found"; continue
    fi

    render_dir="$model_path/renders_test_${ITER}"
    gt_dir="$model_path/gt_test_${ITER}"
    render_flag="$model_path/render_complete_${ITER}.flag"

    png_count=$(find "$render_dir" -maxdepth 1 -name "*.png" 2>/dev/null | wc -l)
    if [ "$png_count" -eq 0 ]; then
        echo "=== $method / $scene (res=$resolution) - RENDER ==="
        rm -rf "$render_dir" "$gt_dir"
    else
        echo "=== $method / $scene - renders exist ($png_count pngs) ==="
    fi

    conda activate "$env_path"

    # 关键修复：把方法自己的源码根目录放在最前面
    export PYTHONPATH="${code_root}:${ROOT}:${PYTHONPATH:+:$PYTHONPATH}"

    if [ "$png_count" -eq 0 ]; then
        run_step "Render" "
            python $ROOT/scripts/main_render.py \
                --method $method \
                --source_path $source_path \
                --model_path $model_path \
                --iteration $ITER \
                --resolution $resolution
        " || continue
    fi

    metrics_json="$model_path/metrics_test_iter${ITER}.json"
    if [ ! -f "$metrics_json" ]; then
        run_step "Eval" "
            python $ROOT/ufd_evalkit/run_eval.py \
                --method $method \
                --scene $scene \
                --render_dir $render_dir \
                --gt_dir $gt_dir \
                --ply_path $ply_path \
                --output_json $metrics_json \
                --colmap_dir $source_path
        " || continue
    else
        echo "  [Metrics json already exists]"
    fi

    conda deactivate
done

echo "All done. Now re-run fix_phase5a_table.py to update the CSV."