#!/usr/bin/env python3
"""Benchmark the OpenVINO Real-ESRGAN variants (quality + speed), no PyTorch needed.

Every model is compared against the FP16 model, which reproduces the original
PyTorch `RealESRGAN_x4plus` result to ~48 dB.

    python benchmark_models.py                       # default set, synthetic input
    python benchmark_models.py --images a.png b.png  # real images
    python benchmark_models.py --models int4 --models int8
"""

import argparse
import os
import pathlib
import time

import cv2
import numpy as np

from realesrgan_ov import MODELS, RealESRGAN, _download_model

HERE = pathlib.Path(__file__).resolve().parent


def psnr(a, b, peak=255.0):
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return 99.0 if mse == 0 else 10 * np.log10(peak ** 2 / mse)


def ssim(a, b):
    c1, c2 = 0.01 ** 2, 0.03 ** 2
    a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255.
    b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255.
    mu_a, mu_b = cv2.GaussianBlur(a, (11, 11), 1.5), cv2.GaussianBlur(b, (11, 11), 1.5)
    sa = cv2.GaussianBlur(a * a, (11, 11), 1.5) - mu_a ** 2
    sb = cv2.GaussianBlur(b * b, (11, 11), 1.5) - mu_b ** 2
    sab = cv2.GaussianBlur(a * b, (11, 11), 1.5) - mu_a * mu_b
    return (((2 * mu_a * mu_b + c1) * (2 * sab + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (sa + sb + c2))).mean()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--models', nargs='+', default=['fp16', 'int8', 'int4', 'int4-mixed', 'int4-w4a8'])
    parser.add_argument('--reference', default='fp16')
    parser.add_argument('--images', nargs='*', default=[])
    parser.add_argument('--tile', type=int, default=128)
    parser.add_argument('--tile-pad', type=int, default=16)
    parser.add_argument('--runs', type=int, default=2)
    parser.add_argument('--device', default='CPU')
    parser.add_argument('--size', type=int, nargs=2, default=[512, 256], help='synthetic input WxH')
    parser.add_argument('--token', default=os.environ.get('HF_TOKEN'), help='Hugging Face token')
    args = parser.parse_args()

    if args.images:
        inputs = [cv2.imread(p) for p in args.images]
        inputs = [im for im in inputs if im is not None]
    else:
        rng = np.random.default_rng(0)
        yy, xx = np.mgrid[0:args.size[1], 0:args.size[0]]
        base = (127 + 100 * np.sin(xx / 7.0) * np.cos(yy / 11.0))[..., None]
        inputs = [np.clip(base + rng.normal(0, 25, base.shape), 0, 255).astype(np.uint8).repeat(3, axis=2)]

    def load(key):
        path = HERE / 'models' / f'{MODELS[key][0]}.xml'
        if not path.exists():  # not downloaded yet -> fetch it from the Hub
            path = pathlib.Path(_download_model(key, args.token))
        return RealESRGAN(model=str(path), device=args.device)

    ref_up = load(args.reference) if args.reference in args.models else load(args.reference)
    refs = [ref_up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)[0] for im in inputs]

    size_mb = {}
    for key in args.models:
        up = load(key)
        xml = pathlib.Path(up.model_path)
        size_mb[key] = (xml.stat().st_size + xml.with_suffix('.bin').stat().st_size) / 1e6
        ps, ss, ts, mp = [], [], [], []
        for im, ref in zip(inputs, refs):
            up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)
            best = min([_time(up, im, args) for _ in range(args.runs)])
            out = up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)[0]
            ts.append(best)
            mp.append(im.shape[0] * im.shape[1] / best / 1e6)
            ps.append(psnr(ref, out))
            ss.append(ssim(ref, out))
        print(f'{key:12s} {size_mb[key]:6.1f} MB  PSNR {np.mean(ps):6.2f} dB  SSIM {np.mean(ss):.4f}  '
              f'{np.mean(ts):6.3f} s/img  {np.mean(mp):5.2f} MPix/s')


def _time(up, im, args):
    t0 = time.time()
    up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)
    return time.time() - t0


if __name__ == '__main__':
    main()