#!/usr/bin/env python3
"""
Parse ablation training logs and compute NIMA on rendered test frames.
Outputs a single comparison table.
"""
import os, re, glob, json, sys
import numpy as np

BASE = "/home/szha0669/storage/blur_slam_exp"
SCENE = "scene0005_00"
ITERS = 15000

EXPS = {
    "A_baseline":  ("bags_A_baseline",  False, False),
    "B_depth":     ("bags_B_depth",     False, True),
    "C_nima":      ("bags_C_nima",      True,  False),
    "D_nima_depth":("bags_D_nima_depth",True,  True),
}

def parse_log(log_path, iters):
    """Extract PSNR/SSIM/LPIPS from training log at final test iteration."""
    results = {}
    if not os.path.exists(log_path):
        return results
    pat = re.compile(
        r'\[ITER %d\] Evaluating test: L1 ([\d.]+) PSNR ([\d.]+) SSIM ([\d.]+) LPIPS ([\d.]+)' % iters
    )
    with open(log_path) as f:
        for line in f:
            m = pat.search(line)
            if m:
                results = {
                    'l1': float(m.group(1)),
                    'psnr': float(m.group(2)),
                    'ssim': float(m.group(3)),
                    'lpips': float(m.group(4)),
                }
    return results

def compute_nima(render_dir):
    """Run pyiqa NIMA on rendered PNG files, return mean score."""
    import torch
    import pyiqa
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    nima = pyiqa.create_metric('nima-koniq', device=device)
    files = sorted(glob.glob(os.path.join(render_dir, '*.png')))
    if not files:
        return None
    scores = []
    for f in files:
        try:
            s = float(nima(f).item())
            scores.append(s)
        except Exception:
            pass
    return float(np.mean(scores)) if scores else None

print(f"\n{'Exp':<18} {'NIMA':>7} {'PSNR':>7} {'SSIM':>7} {'LPIPS':>7}  NIMA  Depth-TV")
print("-" * 70)

for exp_key, (exp_dir, has_nima, has_depth) in EXPS.items():
    model_dir = os.path.join(BASE, "outputs/ablation", exp_dir, "scannet", SCENE)
    log_path  = os.path.join(BASE, f"outputs/logs/ablation_{exp_key}.log")
    render_dir = os.path.join(model_dir, "test", "ours_%d" % ITERS, "renders")

    metrics = parse_log(log_path, ITERS)
    nima_score = compute_nima(render_dir) if os.path.isdir(render_dir) else None

    def fmt(v, fmt_str="{:.4f}"):
        return fmt_str.format(v) if v is not None else "  N/A "

    print(f"Exp {exp_key:<14} {fmt(nima_score):>7} {fmt(metrics.get('psnr'), '{:.2f}'):>7} "
          f"{fmt(metrics.get('ssim')):>7} {fmt(metrics.get('lpips')):>7}  "
          f"{'yes' if has_nima else 'no ':3}  {'yes' if has_depth else 'no'}")

print()