#!/usr/bin/env python3
"""
Reproduce Table 2: Blur | EVSSM | A(noNIMA) | D(global)
- Blur    : original ScanNet frame NIMA (no-ref input quality)
- EVSSM   : bags_main renders NIMA (BAGS trained on EVSSM-deblurred data, best proxy)
- A(noNIMA): BAGS on bags_orig, no NIMA, no depth (bags_A_4scenes experiment)
- D(global): BAGS on bags_orig, NIMA + depth alpha=0.1 (bags_D_4scenes experiment)

Run first: scripts/run_bags_AD_4scenes.sh
"""
import os, glob, json, numpy as np, torch
import pyiqa

BASE   = "/home/szha0669/storage/blur_slam_exp"
ITERS  = 15000
SCENES = ["scene0000_00", "scene0000_01", "scene0005_00", "scene0006_00"]

device  = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
nima_fn = pyiqa.create_metric("nima-koniq", device=device)

def nima_dir(d):
    imgs = sorted(glob.glob(f"{d}/*.png") + glob.glob(f"{d}/*.jpg"))
    if not imgs:
        return None
    return float(np.mean([float(nima_fn(p).item()) for p in imgs]))

# Load pre-computed Blur + EVSSM from nima_main_results.json
main_results = {}
main_path = f"{BASE}/outputs/eval/nima_main_results.json"
if os.path.exists(main_path):
    with open(main_path) as f:
        main_results = json.load(f)

print(f"\n{'Scene':<16} {'Blur':>7} {'EVSSM':>7} {'A(noNIMA)':>11} {'D(global)':>11}")
print("-" * 58)

results = {}
for scene in SCENES:
    blur_nima  = main_results.get(scene, {}).get("orig")   # ScanNet original frames
    evssm_nima = main_results.get(scene, {}).get("bags")   # bags_main renders (proxy for EVSSM)

    # A(noNIMA): bags_A_4scenes (BAGS on bags_orig, no NIMA, no depth)
    a_dir = f"{BASE}/outputs/bags_A_4scenes/scannet/{scene}/test/ours_{ITERS}/test_preds_1"
    a_nima = nima_dir(a_dir) if os.path.isdir(a_dir) else None

    # D(global): bags_D_4scenes (BAGS on bags_orig, NIMA + depth alpha=0.1)
    d_dir = f"{BASE}/outputs/bags_D_4scenes/scannet/{scene}/test/ours_{ITERS}/test_preds_1"
    d_nima = nima_dir(d_dir) if os.path.isdir(d_dir) else None

    results[scene] = dict(blur=blur_nima, evssm=evssm_nima, A=a_nima, D=d_nima)

    def fmt(v):
        return f"{v:.3f}" if v is not None else "  N/A"

    print(f"{scene:<16} {fmt(blur_nima):>7} {fmt(evssm_nima):>7} {fmt(a_nima):>11} {fmt(d_nima):>11}")

# MEAN
print("-" * 58)
for key, label in [("blur","Blur"),("evssm","EVSSM"),("A","A(noNIMA)"),("D","D(global)")]:
    vals = [results[s][key] for s in SCENES if results[s].get(key) is not None]
    if vals:
        print(f"{'MEAN':<16} {np.mean(vals):>7.3f}  ({label})")

out = f"{BASE}/outputs/eval/table2_comparison.json"
with open(out, "w") as f:
    json.dump(results, f, indent=2)
print(f"\nSaved: {out}")