#!/usr/bin/env python3
import json
from pathlib import Path
from collections import defaultdict

ROOT = Path("/root/autodl-tmp/SplatAtlas")
OUTPUTS = ROOT / "outputs"

# 31 methods (完整列表)
METHODS = [
    "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm", "lapisgs", "reactgs", "vanilla_3dgs",
    "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs",
    "coadaptgs", "erankgs", "gaussianpro", "minisplatting", "opti3dgs", "steepgs",
    "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting", "pixelgs",
    "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs"
]

# 31 canonical scenes
SCENES = [
    "Auditorium", "Ballroom", "Barn", "Caterpillar", "Courtroom", "Lighthouse", "Museum", "Palace", "Playground", "Temple", "Train", "Truck",
    "Bicycle", "Bonsai", "Counter", "Flowers", "Garden", "Kitchen", "Room", "Stump", "Treehill",
    "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship",
    "DrJohnson", "Playroom"
]

def check_cell(method, scene):
    run_dir = OUTPUTS / f"{method}_{scene}"
    if not run_dir.exists():
        return {"status": "NO_RUN_DIR"}

    renders_dir = run_dir / "renders_test_30000"
    gt_dir = run_dir / "gt_test_30000"
    metrics_file = run_dir / "metrics_test_iter30000.json"

    n_renders = len(list(renders_dir.glob("*.png"))) if renders_dir.exists() else 0
    n_gt = len(list(gt_dir.glob("*.png"))) if gt_dir.exists() else 0

    metrics_status = "NO_METRICS"
    psnr = ssim = lpips = None
    if metrics_file.exists():
        try:
            obj = json.load(open(metrics_file))
            photo = obj.get("photometric", {})
            psnr = photo.get("PSNR")
            ssim = photo.get("SSIM")
            lpips = photo.get("LPIPS")

            if psnr is not None and not (isinstance(psnr, float) and str(psnr) == "nan"):
                metrics_status = "VALID"
            else:
                metrics_status = "NAN"
        except:
            metrics_status = "JSON_ERROR"

    return {
        "status": "OK",
        "n_renders": n_renders,
        "n_gt": n_gt,
        "metrics_status": metrics_status,
        "psnr": psnr,
        "ssim": ssim,
        "lpips": lpips
    }

print("=" * 120)
print(f"{'method':14s} {'scene':12s} {'run?':5s} {'renders':8s} {'gt':6s} {'metrics':10s} {'PSNR':>8s} {'SSIM':>8s} {'LPIPS':>8s}")
print("=" * 120)

summary = defaultdict(int)
for method in METHODS:
    for scene in SCENES:
        res = check_cell(method, scene)
        summary[res["status"]] += 1
        if res["status"] == "OK":
            summary[res["metrics_status"]] += 1

        psnr_str = f"{res.get('psnr', ''):.2f}" if res.get("psnr") not in (None, "nan") else ""
        ssim_str = f"{res.get('ssim', ''):.4f}" if res.get("ssim") not in (None, "nan") else ""
        lpips_str = f"{res.get('lpips', ''):.4f}" if res.get("lpips") not in (None, "nan") else ""

        print(f"{method:14s} {scene:12s} {res['status']:5s} {res.get('n_renders', 0):>8d} {res.get('n_gt', 0):>6d} {res['metrics_status']:10s} {psnr_str:>8s} {ssim_str:>8s} {lpips_str:>8s}")

print("\n" + "=" * 120)
print("SUMMARY")
print("=" * 120)
for k, v in sorted(summary.items()):
    print(f"{k:20s}: {v}")
print(f"Total cells checked: {len(METHODS) * len(SCENES)}")