#!/usr/bin/env python3
"""
Evaluate deblurring quality: PSNR/SSIM per scene, save visual comparison grids.

Compares: blur | turtle | evssm | vdiff  vs  sharp GT

Outputs:
  outputs/eval/deblur/per_scene/<scene>_metrics.json   — per-frame + avg metrics
  outputs/eval/deblur/per_scene/<scene>_grid.jpg       — visual grid (4 cols x N rows)
  outputs/eval/deblur/summary.json                      — all scenes aggregated
  outputs/eval/deblur/summary.md                        — markdown table

Usage:
  python eval_deblur_quality.py
  python eval_deblur_quality.py --scenes scene0000_00 scene0001_00
  python eval_deblur_quality.py --n-vis 6
"""
import argparse, json, os, sys
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2

BASE = Path(__file__).parent.parent

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--blur-root",   default=str(BASE / "data/scannet_blur_proto/vddiff/test/blur"))
    p.add_argument("--sharp-root",  default=str(BASE / "data/scannet_blur_proto/vddiff/test/sharp"))
    p.add_argument("--turtle-root", default=str(BASE / "data/turtle_deblurred"))
    p.add_argument("--evssm-root",  default=str(BASE / "data/evssm_deblurred"))
    p.add_argument("--vdiff-root",  default=str(BASE / "data/vdiff_deblurred"))
    p.add_argument("--out-dir",     default=str(BASE / "outputs/eval/deblur"))
    p.add_argument("--scenes",      nargs="*", default=None)
    p.add_argument("--n-vis",       type=int, default=4, help="frames per scene in visual grid")
    return p.parse_args()


def psnr(img1, img2):
    mse = np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2)
    if mse < 1e-10:
        return 100.0
    return 20.0 * np.log10(255.0 / np.sqrt(mse))


def ssim(img1, img2):
    from skimage.metrics import structural_similarity
    return float(structural_similarity(img1, img2, channel_axis=2, data_range=255))


def load_frames(directory):
    """Return sorted list of (filename, np.array HxWx3 uint8)."""
    directory = Path(directory)
    files = sorted(directory.glob("*.png")) + sorted(directory.glob("*.jpg"))
    frames = []
    for f in files:
        arr = np.array(Image.open(f).convert("RGB"))
        frames.append((f.name, arr))
    return frames


def make_label(text, w, h=28, bg=(40,40,40), fg=(255,255,255)):
    img = np.full((h, w, 3), bg, dtype=np.uint8)
    pil = Image.fromarray(img)
    draw = ImageDraw.Draw(pil)
    draw.text((6, 5), text, fill=fg)
    return np.array(pil)


def save_visual_grid(frames_by_model, scene, out_path, n_vis=4):
    """
    frames_by_model: dict of label -> list of np.array HxWx3
    Saves a grid with columns = models, rows = selected frames.
    """
    models = list(frames_by_model.keys())
    n_frames = min(len(v) for v in frames_by_model.values())
    idxs = np.linspace(0, n_frames - 1, min(n_vis, n_frames), dtype=int)

    rows = []
    for idx in idxs:
        cols = []
        for m in models:
            img = frames_by_model[m][idx]
            label = make_label(f"{m} | frame {idx}", img.shape[1])
            cols.append(np.vstack([label, img]))
        rows.append(np.hstack(cols))
    grid = np.vstack(rows)
    Image.fromarray(grid).save(out_path, quality=92)


def evaluate_scene(scene, blur_root, sharp_root, turtle_root, evssm_root, vdiff_root, out_dir, n_vis):
    roots = {
        "blur":   Path(blur_root)   / scene,
        "turtle": Path(turtle_root) / scene,
        "evssm":  Path(evssm_root)  / scene,
        "vdiff":  Path(vdiff_root)  / scene,
        "sharp":  Path(sharp_root)  / scene,
    }

    available = {k: v for k, v in roots.items() if v.exists()}
    if "sharp" not in available:
        print(f"  [skip] {scene}: no sharp GT")
        return None
    if "blur" not in available:
        print(f"  [skip] {scene}: no blur input")
        return None

    sharp_frames = load_frames(available["sharp"])
    if not sharp_frames:
        print(f"  [skip] {scene}: empty sharp dir")
        return None

    sharp_dict = {name: arr for name, arr in sharp_frames}
    metrics = {}
    vis_frames = {"sharp": [arr for _, arr in sharp_frames]}

    for model in ["blur", "turtle", "evssm", "vdiff"]:
        if model not in available:
            print(f"  [warn] {scene}: {model} not available, skipping")
            continue
        model_frames = load_frames(available[model])
        if not model_frames:
            print(f"  [warn] {scene}: {model} dir empty, skipping")
            continue

        vis_frames[model] = [arr for _, arr in model_frames]

        per_frame = []
        for fname, gt in sharp_frames:
            # match by index if names differ
            idx = list(sharp_dict.keys()).index(fname) if fname in sharp_dict else None
            if idx is not None and idx < len(model_frames):
                pred = model_frames[idx][1]
                pred_resized = cv2.resize(pred, (gt.shape[1], gt.shape[0]))
                p = psnr(pred_resized, gt)
                s = ssim(pred_resized, gt)
                per_frame.append({"frame": fname, "psnr": round(p, 4), "ssim": round(s, 4)})

        if per_frame:
            avg_psnr = float(np.mean([x["psnr"] for x in per_frame]))
            avg_ssim = float(np.mean([x["ssim"] for x in per_frame]))
            metrics[model] = {
                "psnr_mean": round(avg_psnr, 4),
                "ssim_mean": round(avg_ssim, 4),
                "n_frames": len(per_frame),
                "per_frame": per_frame,
            }
            print(f"  {model:8s}  PSNR={avg_psnr:.2f}  SSIM={avg_ssim:.4f}  ({len(per_frame)} frames)")

    out_dir = Path(out_dir) / "per_scene"
    out_dir.mkdir(parents=True, exist_ok=True)

    with open(out_dir / f"{scene}_metrics.json", "w") as f:
        json.dump({"scene": scene, "metrics": metrics}, f, indent=2)

    # visual grid — reorder columns: blur | turtle | evssm | vdiff | sharp
    ordered = {k: vis_frames[k] for k in ["blur", "turtle", "evssm", "vdiff", "sharp"] if k in vis_frames}
    if ordered:
        save_visual_grid(ordered, scene, out_dir / f"{scene}_grid.jpg", n_vis=n_vis)
        print(f"  grid -> {out_dir}/{scene}_grid.jpg")

    return metrics


def main():
    args = parse_args()

    blur_root = Path(args.blur_root)
    if args.scenes:
        scenes = args.scenes
    else:
        scenes = sorted(p.name for p in blur_root.iterdir() if p.is_dir() and p.name.startswith("scene"))

    print(f"Evaluating {len(scenes)} scenes...")
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    all_metrics = {}
    for scene in scenes:
        print(f"\n[{scene}]")
        m = evaluate_scene(scene, args.blur_root, args.sharp_root,
                           args.turtle_root, args.evssm_root, args.vdiff_root, str(out_dir), args.n_vis)
        if m:
            all_metrics[scene] = m

    # summary JSON
    summary = {}
    for model in ["blur", "turtle", "evssm", "vdiff"]:
        psnrs = [all_metrics[s][model]["psnr_mean"] for s in all_metrics if model in all_metrics[s]]
        ssims = [all_metrics[s][model]["ssim_mean"] for s in all_metrics if model in all_metrics[s]]
        if psnrs:
            summary[model] = {
                "psnr_mean": round(float(np.mean(psnrs)), 4),
                "ssim_mean": round(float(np.mean(ssims)), 4),
                "n_scenes": len(psnrs),
            }

    with open(out_dir / "summary.json", "w") as f:
        json.dump({"per_scene": all_metrics, "overall": summary}, f, indent=2)

    # markdown table
    lines = ["# Deblurring Quality (vs Sharp GT)\n",
             "| Scene | Blur PSNR | EVSSM PSNR | VD-Diff PSNR | Turtle PSNR | Blur SSIM | EVSSM SSIM | VD-Diff SSIM | Turtle SSIM |",
             "|-------|-----------|------------|--------------|-------------|-----------|------------|--------------|-------------|"]
    for scene in sorted(all_metrics):
        m = all_metrics[scene]
        def g(model, key): return f"{m[model][key]:.4f}" if model in m else "—"
        lines.append(f"| {scene} | {g('blur','psnr_mean')} | {g('evssm','psnr_mean')} | {g('vdiff','psnr_mean')} | {g('turtle','psnr_mean')} | {g('blur','ssim_mean')} | {g('evssm','ssim_mean')} | {g('vdiff','ssim_mean')} | {g('turtle','ssim_mean')} |")

    lines.append(f"\n**Overall (mean across scenes):**")
    for model in ["blur", "evssm", "vdiff", "turtle"]:
        if model in summary:
            s = summary[model]
            lines.append(f"- {model}: PSNR={s['psnr_mean']:.4f}  SSIM={s['ssim_mean']:.4f}  (n={s['n_scenes']})")

    md = "\n".join(lines)
    with open(out_dir / "summary.md", "w") as f:
        f.write(md)

    print(f"\n{'='*60}")
    print("Summary:")
    for model in ["blur", "evssm", "vdiff", "turtle"]:
        s = summary.get(model)
        if not s:
            continue
        print(f"  {model:8s}  PSNR={s['psnr_mean']:.4f}  SSIM={s['ssim_mean']:.4f}  (n={s['n_scenes']})")
    print(f"\nOutputs -> {out_dir}/")



if __name__ == "__main__":
    main()