#!/usr/bin/env python3
"""compare_examples.py - every 12-bit example scored on one footing.

    ~/.venvs/fgc-plots/bin/python compare_examples.py <examples-dir> [<out.json>]

Decodes each 12-bit example .jxl from `raw12_bench examples` (needs imagecodecs
and OpenCV), turns the mosaic kinds into colour with the same edge-aware
de-Bayer, and compares every result with the lossless 12-bit colour image of
the same frame (c12_ea_lossless.jxl, exact by construction):
  PSNR  on the 12-bit scale, all three channels
  SSIM  on luma, 12-bit, 11x11 Gaussian window (standard constants)
So compression is judged on what the downstream algorithms would receive, not
against each kind's own reference.
"""
import glob
import json
import os
import sys

import cv2
import imagecodecs
import numpy as np

d = sys.argv[1]
ref_path = os.path.join(d, "c12_ea_lossless.jxl")


def decode(path):
    a = imagecodecs.jpegxl_decode(open(path, "rb").read())
    return a


def to_rgb12(name, a):
    """Decoded array -> HxWx3 float RGB on the 12-bit scale."""
    if a.ndim == 3:
        return a.astype(np.float64)
    m = a.astype(np.uint16)
    h, w = m.shape
    if "planes" in name:  # undo the 2x2 tiling of R, G1 / G2, B
        hh, hw = h // 2, w // 2
        mo = np.empty_like(m)
        mo[0::2, 0::2] = m[:hh, :hw]
        mo[0::2, 1::2] = m[:hh, hw:]
        mo[1::2, 0::2] = m[hh:, :hw]
        mo[1::2, 1::2] = m[hh:, hw:]
        m = mo
    return cv2.cvtColor(m, cv2.COLOR_BayerBG2RGB_EA).astype(np.float64)


def luma(rgb):
    return 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]


def ssim(a, b, peak=4095.0):
    c1, c2 = (0.01 * peak) ** 2, (0.03 * peak) ** 2
    g = lambda x: cv2.GaussianBlur(x, (11, 11), 1.5)
    mu_a, mu_b = g(a), g(b)
    va, vb, cov = g(a * a) - mu_a ** 2, g(b * b) - mu_b ** 2, g(a * b) - mu_a * mu_b
    s = ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2))
    return float(s.mean())


ref = decode(ref_path)
if ref.max() > 4095:  # decoded to 16-bit range: bring back to 12-bit
    ref = ref / 16.0
ref = ref.astype(np.float64)
ref_l = luma(ref)
res = []
for path in sorted(glob.glob(os.path.join(d, "*.jxl"))):
    name = os.path.splitext(os.path.basename(path))[0]
    if not any(k in name for k in ("c12", "raw12")):
        continue  # 8-bit examples are other captures, not comparable pixel by pixel
    a = decode(path)
    if a.max() > 4095:
        a = a / 16.0
    rgb = to_rgb12(name, a)
    mse = float(np.mean((rgb - ref) ** 2))
    psnr = 99.0 if mse == 0 else 10 * np.log10(4095.0 ** 2 / mse)
    r = dict(name=name, bytes=os.path.getsize(path), psnr_db=round(psnr, 2),
             ssim=round(ssim(luma(rgb), ref_l), 5))
    res.append(r)
    print(json.dumps(r))
if len(sys.argv) > 2:
    json.dump(res, open(sys.argv[2], "w"), indent=1)