#!/usr/bin/env python3
"""Preview the exact Ch438 source-space 3x3 binomial scanout."""
import math
import sys

from PIL import Image


def binom3(a, b, c):
    return tuple((x + 2 * y + z + 2) // 4 for x, y, z in zip(a, b, c))


def score(a, b):
    ae = 0
    se = 0
    n = a.width * a.height * 3
    for pa, pb in zip(a.getdata(), b.getdata()):
        for ca, cb in zip(pa, pb):
            d = int(ca) - int(cb)
            ae += abs(d)
            se += d * d
    return ae / n, math.sqrt(se / n)


def main(argv):
    if len(argv) != 4:
        raise SystemExit(
            "usage: preview_scanout_ch438.py <fb.png> <reference.png> <out.png>"
        )

    src = Image.open(argv[1]).convert("RGB")
    if src.size != (640, 480):
        raise SystemExit(f"expected 640x480 framebuffer, got {src.size}")
    ref = Image.open(argv[2]).convert("RGB").resize(
        (640, 480), Image.Resampling.BILINEAR
    )

    # RTL rounds after each separable pass, so mirror that order exactly.
    horizontal = Image.new("RGB", src.size)
    hp = horizontal.load()
    for y in range(480):
        for x in range(512):
            xm1 = max(0, x - 1)
            xp1 = min(511, x + 1)
            hp[x, y] = binom3(src.getpixel((xm1, y)),
                              src.getpixel((x, y)),
                              src.getpixel((xp1, y)))

    filtered = Image.new("RGB", src.size)
    fp = filtered.load()
    for y in range(32, 480):
        ym1 = max(32, y - 1)
        yp1 = min(479, y + 1)
        for x in range(512):
            fp[x, y] = binom3(horizontal.getpixel((x, ym1)),
                              horizontal.getpixel((x, y)),
                              horizontal.getpixel((x, yp1)))

    out = Image.new("RGB", (640, 480))
    px = out.load()
    for y in range(480):
        sy = 32 + (y * 14) // 15
        for x in range(640):
            sx = (x * 4) // 5
            px[x, y] = filtered.getpixel((sx, sy))

    out.save(argv[3])
    mae, rmse = score(out, ref)
    print(f"binomial3x3 MAE={mae:.4f} RMSE={rmse:.4f}")
    print(f"wrote {argv[3]}")


if __name__ == "__main__":
    main(sys.argv)