# /// script
# requires-python = ">=3.11"
# dependencies = ["coreai-core==1.0.0b3", "coreai-torch==0.4.3", "spandrel==0.4.1", "numpy", "pillow"]
# [tool.uv]
# index-url = "https://pypi.org/simple"
# prerelease = "allow"
# index-strategy = "unsafe-best-match"
# ///
"""Core AI export vs the original PyTorch weights (float32), tile by tile, on stock photos."""
import asyncio, sys, time
import numpy as np, torch
from PIL import Image
from spandrel import ModelLoader
import coreai.runtime as rt

ckpt, path, tile = sys.argv[1], sys.argv[2], int(sys.argv[3])
images = sys.argv[4:]
desc = ModelLoader().load_from_file(ckpt)
net = desc.model.eval().float()
luma = torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1)
def reference(x):
    if desc.input_channels == 1: x = (x * luma).sum(dim=1, keepdim=True)
    y = net(x)
    if desc.output_channels == 1: y = y.repeat(1, 3, 1, 1)
    return y.clamp(0, 1)
def psnr(a, b):
    mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
    return 99.0 if mse == 0 else 10 * np.log10(1.0 / mse)
async def main():
    rng = np.random.default_rng(0); tiles = []
    for p in images:
        im = np.asarray(Image.open(p).convert("RGB")).astype(np.float32) / 255
        h, w, _ = im.shape
        for _ in range(3):
            y = int(rng.integers(0, max(1, h - tile))); x = int(rng.integers(0, max(1, w - tile)))
            c = im[y:y+tile, x:x+tile]
            if c.shape[:2] == (tile, tile): tiles.append(c)
    m = await rt.AIModel.load(path)
    fn = m.load_function("main")
    scores, times = [], []
    for c in tiles:
        x = torch.from_numpy(c).permute(2, 0, 1)[None].contiguous()
        with torch.no_grad(): ref = reference(x)[0].permute(1, 2, 0).numpy()
        t = time.time()
        out = await fn({"x": rt.NDArray(data=x.half().contiguous())})
        times.append(time.time() - t)
        got = out["output"].numpy().astype(np.float32)[0].transpose(1, 2, 0)
        scores.append(psnr(ref, got))
    print(f"{path.split('/')[-1]:34} scale x{desc.scale}  PSNR vs PyTorch fp32 over {len(scores)} tiles: "
          f"min {min(scores):.1f} dB, mean {np.mean(scores):.1f} dB; {np.median(times)*1000:.0f} ms/tile")
asyncio.run(main())