import os, subprocess

SPLATATLAS = "/root/autodl-tmp/SplatAtlas"
BASE = "/root/autodl-tmp/SplatAtlas/outputs"
SRC = "/root/autodl-tmp/dataset/360/bonsai"
RES = 2
SCENE = "bonsai"

ENV_MAP = {
    "lightgaussian": ("/root/autodl-tmp/envs/lightgaussian_official", "/root/autodl-tmp/LightGaussian_official"),
    "3dgsmcmc": ("/root/autodl-tmp/envs/3dgs_mcmc", "/root/autodl-tmp/3dgs-mcmc")
}

def get_py_and_env(method):
    env_path, extra_py = ENV_MAP[method]
    python_bin = f"{env_path}/bin/python"
    env = os.environ.copy()
    env["PYTHONPATH"] = f"{SPLATATLAS}:{extra_py}:{env.get('PYTHONPATH', '')}"
    env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" # 防止算 LPIPS 爆显存
    return python_bin, env

targets = ["lightgaussian", "3dgsmcmc"]

for method in targets:
    m_dir = f"{BASE}/{method}_{SCENE}"
    metrics_f = f"{m_dir}/metrics_test_iter30000.json"
    ply_f = f"{m_dir}/point_cloud/iteration_30000/point_cloud.ply"
    
    if not os.path.exists(ply_f):
        print(f"[-] SKIPPING {method}_{SCENE}: PLY 未找到，说明训练没跑完。")
        continue
        
    print(f"\n[*] FIXING: {method}_{SCENE}")
    python_bin, custom_env = get_py_and_env(method)
    
    # 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)