# The page's before-and-after footage: two folders of lossless frames (film.mjs.txt, one per checkout of the
# game) put side by side, at 390 x 844 a side (frames filmed at a pixel ratio of 2 are scaled down by averaging),
# encoded once (VP9 in WebM for the page), and a poster frame. Then the encoding is checked against the lossless frames, because
# video compression is exactly what adds or hides banding in dark gradients:
#   kept    the rings are the difference between before and after; after a blur wider than the grain (4 pixels),
#           that difference in the decoded video should match the lossless one (correlation, and RMS ratio near 1)
#   added   the encoding's own banding: contour pixels (a step between two flat runs of 8, the measure poolsteps.py
#           uses) in each half, lossless against decoded; and how far the decoded "after" half strays from the
#           lossless one after the same blur, in L*
#   python3 footage.py <before-dir> <after-dir> <out.webm> <poster.jpg> [crf] [top bottom]
# (VP9 in WebM, which every current browser plays, Playwright's Chromium included, so the clip can be checked as a
# browser decodes it; .mp4 gives H.264. top/bottom crop each side to rows of the 390 x 844 page.) The check above
# decodes with ffmpeg, whose colour conversion differs a little from a browser's; for the clip as a browser shows it:
#   node --input-type=module - footage-rift.webm <dir> < footage-browser.mjs.txt
#   python3 footage.py browser <before-dir> <after-dir> <dir> [top bottom]       (out/footage-browser.txt)
import sys, os, glob, subprocess, tempfile
import numpy as np
from PIL import Image

def load(d):
    return sorted(glob.glob(os.path.join(d, 'f*.png')))

TOP, BOTTOM = 0, 844
def side(fa, fb):
    a = Image.open(fa).convert('RGB'); b = Image.open(fb).convert('RGB')
    if a.size[0] != 390:
        a = a.resize((390, 844), Image.BOX); b = b.resize((390, 844), Image.BOX)
    a = a.crop((0, TOP, 390, BOTTOM)); b = b.crop((0, TOP, 390, BOTTOM))
    im = Image.new('RGB', (782, BOTTOM - TOP)); im.paste(a, (0, 0)); im.paste(b, (392, 0))
    return im

def lstar(rgb):
    c = rgb / 255.
    lin = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
    Y = lin @ np.array([0.2126, 0.7152, 0.0722])
    return np.where(Y > 216 / 24389, 116 * np.cbrt(Y) - 16, Y * 24389 / 27)

def blur(img, s=4.0):
    r = int(np.ceil(3 * s)); k = np.exp(-np.arange(-r, r + 1) ** 2 / (2 * s * s)); k /= k.sum()
    p = np.pad(img, r, mode='edge')
    p = np.apply_along_axis(lambda v: np.convolve(v, k, 'valid'), 0, p)
    return np.apply_along_axis(lambda v: np.convolve(v, k, 'valid'), 1, p)

def steps(img):
    """Contour pixels: a step between two neighbours with a flat run of 8 on each side (poolsteps.py's measure)."""
    code = (np.asarray(img, np.int64) @ np.array([1, 256, 65536])); n = 0
    for c in (code, code.T):
        d = np.diff(c, axis=1); st = d != 0; fl = np.ones_like(st)
        for k in range(1, 8): fl[:, k:] &= d[:, :-k] == 0; fl[:, :-k] &= d[:, k:] == 0
        n += int((st & fl).sum())
    return n

def main():
    global TOP, BOTTOM
    bdir, adir, out, poster = sys.argv[1:5]; crf = sys.argv[5] if len(sys.argv) > 5 else '16'
    if len(sys.argv) > 7: TOP, BOTTOM = int(sys.argv[6]), int(sys.argv[7])
    fb, fa = load(bdir), load(adir); n = min(len(fb), len(fa))
    tmp = tempfile.mkdtemp()
    for i in range(n): side(fb[i], fa[i]).save(os.path.join(tmp, f's{i:04d}.png'))
    side(fb[n // 2], fa[n // 2]).save(poster, quality=90)
    # (sRGB frames to BT.709 limited-range YUV, tagged so browsers convert back the same way)
    codec = (['-c:v', 'libvpx-vp9', '-crf', crf, '-b:v', '0', '-row-mt', '1', '-deadline', 'good', '-cpu-used', '0', '-g', '30']
             if out.endswith('.webm') else ['-c:v', 'libx264', '-preset', 'veryslow', '-tune', 'grain', '-crf', crf, '-profile:v', 'high', '-g', '30', '-movflags', '+faststart'])
    subprocess.run(['ffmpeg', '-v', 'error', '-y', '-framerate', '30', '-i', os.path.join(tmp, 's%04d.png'),
                    '-vf', 'scale=out_color_matrix=bt709:out_range=tv,format=yuv420p', *codec,
                    '-colorspace', 'bt709', '-color_primaries', 'bt709', '-color_trc', 'iec61966-2-1', '-color_range', 'tv',
                    '-an', out], check=True)
    dec = os.path.join(tmp, 'dec'); os.makedirs(dec)
    subprocess.run(['ffmpeg', '-v', 'error', '-i', out, '-vf', 'scale=in_color_matrix=bt709:in_range=tv', os.path.join(dec, 'd%04d.png')], check=True)
    kept = []; added = []; rings = []; cont = []
    for i in range(0, n, 6):
        src = np.asarray(Image.open(os.path.join(tmp, f's{i:04d}.png')), float)
        d = np.asarray(Image.open(os.path.join(dec, f'd{i + 1:04d}.png')).convert('RGB'), float)
        Ls, Ld = blur(lstar(src)), blur(lstar(d))
        diff_s = Ls[:, :390] - Ls[:, 392:]; diff_d = Ld[:, :390] - Ld[:, 392:]
        rings.append(np.sqrt((diff_s ** 2).mean()))
        kept.append((np.corrcoef(diff_s.ravel(), diff_d.ravel())[0, 1], np.sqrt((diff_d ** 2).mean()) / np.sqrt((diff_s ** 2).mean())))
        added.append(np.sqrt(((Ld[:, 392:] - Ls[:, 392:]) ** 2).mean()))
        cont.append([steps(src[:, :390]), steps(d[:, :390]), steps(src[:, 392:]), steps(d[:, 392:])])
    k = np.array(kept)
    print(f'{os.path.basename(out)}: {n} frames, rows {TOP}-{BOTTOM}, {os.path.getsize(out) / 1e6:.2f} MB, {"VP9" if out.endswith(".webm") else "H.264"} CRF {crf}')
    print(f'  rings (before minus after, blurred 4 px): RMS {np.mean(rings):.3f} L*')
    print(f'  kept: correlation {k[:, 0].mean():.3f}, RMS ratio {k[:, 1].mean():.3f}')
    c = np.array(cont).mean(0)
    print(f'  contour pixels a frame, lossless -> decoded: before half {c[0]:.0f} -> {c[1]:.0f}, after half {c[2]:.0f} -> {c[3]:.0f}')
    print(f'  the after half after the blur, decoded against lossless: RMS {np.mean(added):.3f} L*')

def browser():
    """Screenshots of Chromium playing the clip (footage-browser.mjs.txt) against the lossless frames: each matched to
    the source frame it is closest to, then the same measures as the encode check."""
    global TOP, BOTTOM
    bdir, adir, ddir = sys.argv[2:5]
    if len(sys.argv) > 6: TOP, BOTTOM = int(sys.argv[5]), int(sys.argv[6])
    fb, fa = load(bdir), load(adir)
    print(f'the footage as Chromium shows it (screenshots while it plays), against the lossless frames (rows {TOP}-{BOTTOM}, blurred 4 px)')
    rows = []
    small = [np.asarray(side(fb[j], fa[j]).resize((196, (BOTTOM - TOP) // 4)), float) for j in range(len(fb))]
    for f in sorted(glob.glob(os.path.join(ddir, 'b*.png'))):
        d = np.asarray(Image.open(f).convert('RGB'), float); Ld = blur(lstar(d))
        ds_ = np.asarray(Image.open(f).convert('RGB').resize((196, (BOTTOM - TOP) // 4)), float)
        near = sorted(range(len(fb)), key=lambda j: ((small[j] - ds_) ** 2).mean())[:3]
        best = min((np.sqrt(((blur(lstar(np.asarray(side(fb[j], fa[j]), float))) - Ld) ** 2).mean()), j) for j in near)
        j = best[1]; src = np.asarray(side(fb[j], fa[j]), float); Ls = blur(lstar(src))
        ds, dd = Ls[:, :390] - Ls[:, 392:], Ld[:, :390] - Ld[:, 392:]
        r = (j, np.sqrt(((Ld - Ls)[:, 392:] ** 2).mean()), (Ld - Ls)[:, 392:].mean(), np.corrcoef(ds.ravel(), dd.ravel())[0, 1],
             steps(src[:, :390]), steps(d[:, :390]), steps(src[:, 392:]), steps(d[:, 392:]))
        rows.append(r)
        print(f'  frame {j:3d}: after half RMS {r[1]:.3f} L* (mean {r[2]:+.3f}); rings kept, correlation {r[3]:.3f}; contour pixels, lossless -> browser: before {r[4]} -> {r[5]}, after {r[6]} -> {r[7]}')
    m = np.array(rows).mean(0)
    print(f'  mean: after half RMS {m[1]:.3f} L*; correlation {m[3]:.3f}; contour pixels before {m[4]:.0f} -> {m[5]:.0f}, after {m[6]:.0f} -> {m[7]:.0f}')

if __name__ == '__main__':
    browser() if sys.argv[1] == 'browser' else main()