import os, subprocess

SPLATATLAS = "/root/autodl-tmp/SplatAtlas"
BASE = "/root/autodl-tmp/SplatAtlas/outputs"

# 映射表：(环境路径, 额外 PYTHONPATH)
ENV_MAP = {
    "vanilla_3dgs": ("base", "/root/autodl-tmp/3dgs_official"),
    "analyticsplatting": ("/root/autodl-tmp/envs/ana_splatting", "/root/autodl-tmp/analyticsplatting_official"),
    "erankgs": ("/root/autodl-tmp/envs/erank_gs", "/root/autodl-tmp/erank_gs"),
    "pgsr": ("/root/autodl-tmp/envs/pgsr_v2", "/root/autodl-tmp/pgsr_official")
}

def get_py_and_env(method):
    if method not in ENV_MAP:
        return "python", os.environ.copy()
        
    env_path, extra_py = ENV_MAP[method]
    python_bin = "python" if env_path == "base" else f"{env_path}/bin/python"
    
    env = os.environ.copy()
    env["PYTHONPATH"] = f"{SPLATATLAS}:{extra_py}:{env.get('PYTHONPATH', '')}"
    # 缓解 LPIPS 评测时的 PyTorch 显存碎片化 OOM
    env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
    return python_bin, env

targets = [
    ("vanilla_3dgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2),
    ("analyticsplatting", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2),
    ("analyticsplatting", "Lego", "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/Lego", 1),
    ("erankgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2),
    ("pgsr", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2)
]
seeds = ["", "_seed1", "_seed2"]

for method, scene, src, res in targets:
    python_bin, custom_env = get_py_and_env(method)
    
    for suffix in seeds:
        m_dir = f"{BASE}/{method}{suffix}_{scene}"
        metrics_f = f"{m_dir}/metrics_test_iter30000.json"
        ply_f = f"{m_dir}/point_cloud/iteration_30000/point_cloud.ply"
        
        if os.path.exists(metrics_f):
            continue
            
        if not os.path.exists(ply_f):
            print(f"[-] SKIPPING {method}{suffix}_{scene}: PLY not found")
            continue

        print(f"\n[*] FIXING: {method}{suffix}_{scene}")
        print(f"    -> Using Python: {python_bin}")
        
        # 1. 渲染 (强制覆盖)
        render_cmd = f"{python_bin} {SPLATATLAS}/scripts/main_render.py --method {method} --source_path {src} --model_path {m_dir} --iteration 30000 --resolution {res}"
        subprocess.run(render_cmd, shell=True, env=custom_env)
        
        # 2. 评估
        renders = f"{m_dir}/renders_test_30000"
        gt = f"{m_dir}/gt_test_30000"
        depths = f"{m_dir}/depths_test_30000"
        
        eval_cmd = (
            f"{python_bin} {SPLATATLAS}/ufd_evalkit/run_eval.py "
            f"--method {method} --scene {scene} --render_dir {renders} "
            f"--gt_dir {gt} --ply_path {ply_f} --output_json {metrics_f} "
            f"--colmap_dir {src} --depth_dir {depths}"
        )
        subprocess.run(eval_cmd, shell=True, env=custom_env)