#!/usr/bin/env python3
"""Compare two DDColor ONNX files on photos: ab max/mean difference plus side-by-side colourised JPGs.
Usage: compare_ddcolor.py REF.onnx TEST.onnx OUTDIR IMG [IMG ...]"""
import sys, time
from pathlib import Path
import cv2, numpy as np, onnxruntime as ort

def preprocess(img_bgr):  # DDColor ColorizationPipeline preprocessing
    img = (img_bgr / 255.0).astype(np.float32)
    l = cv2.cvtColor(cv2.resize(img, (512, 512)), cv2.COLOR_BGR2Lab)[:, :, :1]
    rgb = cv2.cvtColor(np.concatenate((l, 0 * l, 0 * l), -1), cv2.COLOR_LAB2RGB)
    return rgb.transpose(2, 0, 1)[None].astype(np.float32)

def colourise(img_bgr, ab):
    h, w = img_bgr.shape[:2]
    l = cv2.cvtColor((img_bgr / 255.0).astype(np.float32), cv2.COLOR_BGR2Lab)[:, :, :1]
    ab = cv2.resize(ab[0].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_NEAREST)
    return (cv2.cvtColor(np.concatenate((l, ab), -1), cv2.COLOR_LAB2BGR) * 255).round().clip(0, 255).astype(np.uint8)

ref_s = ort.InferenceSession(sys.argv[1], providers=["CPUExecutionProvider"])
tst_s = ort.InferenceSession(sys.argv[2], providers=["CPUExecutionProvider"])
out = Path(sys.argv[3]); out.mkdir(parents=True, exist_ok=True)
for p in sys.argv[4:]:
    img = cv2.imread(p); x = preprocess(img)
    a = ref_s.run(["output"], {"input": x})[0]
    t = time.time(); b = tst_s.run(["output"], {"input": x})[0]; dt = time.time() - t
    d = np.abs(a - b)
    ca, cb = colourise(img, a), colourise(img, b)
    pix = np.abs(ca.astype(int) - cb.astype(int))
    print(f"{Path(p).name}: ab max {d.max():.3f} mean {d.mean():.4f} p99 {np.percentile(d, 99):.3f} | "
          f"8-bit RGB max {pix.max()} mean {pix.mean():.3f} | test run {dt:.2f}s")
    stem = Path(p).stem.split()[0]
    gap = np.full((ca.shape[0], 16, 3), 255, np.uint8)
    side = np.concatenate((ca, gap, cb), 1)
    s = 1400 / side.shape[1]
    cv2.imwrite(str(out / f"{stem}_fp32_vs_test.jpg"), cv2.resize(side, None, fx=s, fy=s, interpolation=cv2.INTER_AREA), [cv2.IMWRITE_JPEG_QUALITY, 92])