#!/usr/bin/env python3
import os, json, cv2, numpy as np, torch, pyiqa

BASE    = "/home/szha0669/storage/blur_slam_exp"
TUM_RGB = f"{BASE}/data/TUM_RGBD/rgbd_dataset_freiburg1_desk/rgb"
OUT     = f"{BASE}/outputs/logs/nima_tum_fr1desk_scores.json"
rgb_files = sorted(os.listdir(TUM_RGB))

device = torch.device("cuda")
nima   = pyiqa.create_metric("nima-koniq", device=device)

scores = []
for i, f in enumerate(rgb_files):
    s = float(nima(f"{TUM_RGB}/{f}"))
    scores.append({"fi": i, "nima": round(s, 5), "file": f})
    if (i+1) % 100 == 0:
        print(f"{i+1}/{len(rgb_files)}", flush=True)

json.dump(scores, open(OUT, "w"), indent=2)
print(f"Saved {len(scores)} scores → {OUT}")

scores_s = sorted(scores, key=lambda x: -x["nima"])

def lv(fi):
    img = cv2.imread(f"{TUM_RGB}/{rgb_files[fi]}")
    return cv2.Laplacian(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), cv2.CV_64F).var()

# Top-14 with min spacing=10
selected = []
for x in scores_s:
    fi = x["fi"]
    if all(abs(fi - s) >= 10 for s in selected):
        selected.append(fi)
        if len(selected) == 14: break

selected.sort()
print(f"\nTop-14 GT frames (min spacing=10): {selected}")
for fi in selected:
    s = next(x["nima"] for x in scores if x["fi"] == fi)
    print(f"  fi={fi:3d}  NIMA={s:.4f}  LapVar={lv(fi):.1f}")