import os
import subprocess
import yaml
import glob

def build_scene_map():
    """反向解析 configs 目录下的所有 yaml，提取真实的源路径映射"""
    scene_map = {}
    config_files = glob.glob("/root/autodl-tmp/SplatAtlas/configs/*.yaml")
    for cfg_file in config_files:
        try:
            with open(cfg_file, 'r') as f:
                cfg = yaml.safe_load(f)
            if "datasets" in cfg:
                for ds in cfg["datasets"]:
                    base = ds.get("base_path", "")
                    for sc in ds.get("scenes", []):
                        scene_map[sc] = os.path.join(base, sc)
        except Exception:
            pass
    return scene_map

def main():
    base_dir = "/root/autodl-tmp/SplatAtlas/outputs"
    iteration = 30000
    
    # 建立 Scene -> Source Path 的绝对真理映射
    scene_map = build_scene_map()

    for folder in os.listdir(base_dir):
        if not folder.startswith("2dgs_"): continue
        
        m_path = os.path.join(base_dir, folder)
        method, scene = folder.split('_', 1)
        
        # 精准获取源路径，增加后备启发式搜索以防万一
        src = scene_map.get(scene, "")
        if not src or not os.path.exists(src):
            roots = [
                "/root/autodl-tmp/dataset/synthetic_nerf/Synthetic_NeRF",
                "/root/autodl-tmp/dataset/mipnerf360",
                "/root/autodl-tmp/dataset/mipnerf360/360_v2", # Mip-NeRF 360 的常见嵌套
                "/root/autodl-tmp/dataset/deepblending_clean"
            ]
            for r in roots:
                p = os.path.join(r, scene)
                if os.path.exists(p):
                    src = p
                    break
        
        if not src or not os.path.exists(src):
            print(f"❌ [Fatal] 彻底找不到场景 {scene} 的数据集目录，跳过。")
            continue

        # 补测基础评估
        metrics_json = os.path.join(m_path, f"metrics_test_iter{iteration}.json")
        if not os.path.exists(metrics_json):
            print(f"🛠️ [Eval] 正在抢救基础评测: {folder}")
            cmd = (f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py "
                   f"--method {method} --scene {scene} "
                   f"--render_dir {m_path}/renders_test_{iteration} "
                   f"--gt_dir {m_path}/gt_test_{iteration} "
                   f"--ply_path {m_path}/point_cloud/iteration_{iteration}/point_cloud.ply "
                   f"--output_json {metrics_json} --colmap_dir {src} --skip_quarantine")
            subprocess.run(cmd, shell=True)

        # 补测物理诊断
        physics_json = os.path.join(m_path, f"offline_physics_{iteration}.json")
        if not os.path.exists(physics_json):
            print(f"🛠️ [Physics] 正在抢救物理诊断: {folder}")
            cmd = (f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py "
                   f"--method {method} --source_path {src} --model_path {m_path} --iteration {iteration}")
            subprocess.run(cmd, shell=True)

    # 聚合大榜单
    print("📈 [GrandTable] 正在重构汇总榜单...")
    subprocess.run("python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", shell=True)

if __name__ == "__main__":
    main()