#!/usr/bin/env python3
"""Compute simple RGB PSNR/SSIM between two image folders."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import cv2
import numpy as np


def images(path: Path) -> list[Path]:
    return sorted([*path.glob("*.png"), *path.glob("*.jpg"), *path.glob("*.jpeg")])


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2)
    if mse == 0:
        return float("inf")
    return 20 * math.log10(255.0 / math.sqrt(mse))


def ssim_gray(a: np.ndarray, b: np.ndarray) -> float:
    # Lightweight global SSIM for fast prototype reporting.
    a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY).astype(np.float64)
    b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY).astype(np.float64)
    c1 = (0.01 * 255) ** 2
    c2 = (0.03 * 255) ** 2
    mu_a, mu_b = a.mean(), b.mean()
    var_a, var_b = a.var(), b.var()
    cov = ((a - mu_a) * (b - mu_b)).mean()
    return ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (var_a + var_b + c2))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ref", type=Path, required=True)
    parser.add_argument("--pred", type=Path, required=True)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()

    ref_paths = images(args.ref)
    pred_paths = images(args.pred)
    if len(ref_paths) != len(pred_paths):
        raise ValueError(f"image count mismatch: {len(ref_paths)} vs {len(pred_paths)}")

    rows = []
    for ref, pred in zip(ref_paths, pred_paths):
        a = cv2.imread(str(ref), cv2.IMREAD_COLOR)
        b = cv2.imread(str(pred), cv2.IMREAD_COLOR)
        if a is None or b is None:
            raise ValueError(f"failed to read pair {ref} {pred}")
        if a.shape != b.shape:
            b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_LINEAR)
        rows.append({"name": ref.name, "psnr": psnr(a, b), "ssim": ssim_gray(a, b)})

    summary = {
        "num_images": len(rows),
        "mean_psnr": float(np.mean([r["psnr"] for r in rows])),
        "mean_ssim": float(np.mean([r["ssim"] for r in rows])),
        "ref": str(args.ref),
        "pred": str(args.pred),
    }
    print(json.dumps(summary, indent=2))
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps({"summary": summary, "per_image": rows}, indent=2) + "\n")


if __name__ == "__main__":
    main()