#!/usr/bin/env python3
"""Evaluate the 4-scene GT-pose retraining campaign (BAGS + TriSplat,
GT pose + GT depth, depth_loss_alpha=0.1, 15K iters) on the NIMA-koniq
'sharp frame' subset selected by select_sharp_i2slam.py (top-15% NIMA
scoring of raw pre-EVSSM frames).

Unlike the COLMAP-pose 50K runs (metrics_i2slam_sharp_4scenes_50k.py),
these GT-pose scenes register 100% of frames in sequential evssm_idx
order, so the train/test (--eval --llffhold 8) split and each frame's
render_idx within its split are exactly recoverable from evssm_idx:
  - evssm_idx % 8 == 0  -> test split,  render_idx = evssm_idx // 8
  - else                -> train split, render_idx = evssm_idx - evssm_idx//8 - 1
(verified against actual train/test counts for all 4 scenes).

Computes PSNR / SSIM / LPIPS / NIMA, comparing renders against the
EVSSM-deblurred frames (the same 'sharp-frame' ground truth definition
used by the COLMAP-pose evaluation, for direct comparability)."""
import os, json, numpy as np
from PIL import Image
import torch, pyiqa, lpips as lpips_lib
from skimage.metrics import structural_similarity as ssim_fn
from skimage.metrics import peak_signal_noise_ratio as psnr_fn

BASE = "/home/szha0669/storage/blur_slam_exp"
ITER = 15000
HOLD = 8
SELECTION = json.load(open(f"{BASE}/outputs/logs/i2slam_sharp_frame_selection.json"))
SCENES = ["scene0024_01", "scene0031_00", "scene0736_00", "tum_fr2_xyz"]

device   = torch.device("cuda")
nima_fn  = pyiqa.create_metric("nima-koniq", device=device)
lpips_fn = lpips_lib.LPIPS(net='vgg').to(device)

def t(img):
    return (torch.from_numpy(np.array(img).astype(np.float32) / 255.)
            .permute(2, 0, 1).unsqueeze(0).to(device) * 2 - 1)

def split_and_render_idx(evssm_idx):
    if evssm_idx % HOLD == 0:
        return "test", evssm_idx // HOLD
    return "train", evssm_idx - evssm_idx // HOLD - 1

def eval_scene(out_root, scene, frames):
    rows = []
    gt_dir = f"{BASE}/data/evssm_deblurred_i2slam/{scene}"
    for fr in frames:
        ei = fr["evssm_idx"]
        split, ri = split_and_render_idx(ei)
        pred_p = os.path.join(out_root, scene, split, f"ours_{ITER}", "test_preds_2", f"{ri:05d}.png")
        gt_p   = os.path.join(gt_dir, f"{ei:06d}.png")
        if not (os.path.exists(pred_p) and os.path.exists(gt_p)):
            continue
        pred = Image.open(pred_p).convert('RGB')
        gt   = Image.open(gt_p).convert('RGB').resize(pred.size, Image.LANCZOS)
        pa, ga = np.array(pred), np.array(gt)
        with torch.no_grad():
            lp = float(lpips_fn(t(pred), t(gt)).item())
        rows.append(dict(
            psnr=psnr_fn(ga, pa, data_range=255),
            ssim=ssim_fn(ga, pa, channel_axis=2, data_range=255),
            lpips=lp,
            nima=float(nima_fn(pred_p)),
        ))
    if not rows:
        return None, 0
    return {k: float(np.mean([x[k] for x in rows])) for k in rows[0]}, len(rows)

def collect(out_root):
    return {scene: eval_scene(out_root, scene, SELECTION[scene]) for scene in SCENES}

def print_table(label, per_scene):
    print(f"\n=== {label} (iteration {ITER}, GT pose+depth, alpha=0.1, sharp-frame test set) ===")
    print(f"{'scene':>14}  {'n':>4}  {'PSNR':>7} {'SSIM':>7} {'LPIPS':>7} {'NIMA':>7}")
    print("-" * 56)
    valid = []
    for scene, (r, n) in per_scene.items():
        if r is None:
            print(f"{scene:>14}  {'N/A':>4}")
            continue
        valid.append(r)
        print(f"{scene:>14}  {n:>4}  {r['psnr']:7.3f} {r['ssim']:7.4f} {r['lpips']:7.4f} {r['nima']:7.4f}")
    if valid:
        avg = {k: float(np.mean([x[k] for x in valid])) for k in valid[0]}
        print("-" * 56)
        print(f"{'AVERAGE':>14}  {'':>4}  {avg['psnr']:7.3f} {avg['ssim']:7.4f} {avg['lpips']:7.4f} {avg['nima']:7.4f}")
    return per_scene

bags_results = collect(f"{BASE}/outputs/bags_i2slam_gtall")
tri_results  = collect(f"{BASE}/outputs/trigsplat_i2slam_gtall")

print_table("BAGS / MipSplat  (GT pose+depth)", bags_results)
print_table("TriSplat + BPN   (GT pose+depth)", tri_results)

out_json = {
    "bags_gtpose":  {s: r for s, (r, n) in bags_results.items()},
    "trisplat_gtpose": {s: r for s, (r, n) in tri_results.items()},
}
out_path = f"{BASE}/outputs/logs/metrics_4scenes_gtpose_sharp.json"
json.dump(out_json, open(out_path, "w"), indent=2)
print(f"\nSaved -> {out_path}")