# Starlock's real glow path (pipeline.py, from the 10-22 entry: glow shaders, the 8-bit linear target, the
# painterly pass, ACES, sRGB, the screen write with its own dither), run frame after frame with dithers that
# change each frame, then seen by each eye model (the eye averages light, so frames are blended in
# luminance). Measured as the 10-22 entry did (band: low-frequency error in L*; grain: visible-scale noise in
# L*), plus how visible the grain and the frame-to-frame change are in JND (vis.py's generous model, on
# luminance contrast, at 60 and 30 fps).
#   python3 pipe_t.py            -> out/pipe.txt (about an hour on 4 cores)
import sys
import numpy as np
from multiprocessing import Pool
import pipeline as P, dithers as D, vis, st

FRAME = 320                 # render pixels on a side (half the 10-22 entry's, for time: 16 frames each)
ORIGIN = (60., 1100.)
PPD = 64.                   # a phone at a pixel ratio of 2, 30 cm (as the 10-22 entry)
T, WARM = 42, 10            # frames rendered; frames before WARM only feed the eye's memory (32 measured: 1.9 Hz bins)
PHI = 0.6180339887498949
Y0 = 1e-3                   # light the screen never gets below (its black and the room reflected in it), as a share of white
VIOLET, EMBER, WHITE = (.55, .3, 1.), (.75, .45, 1.), (1., .94, 1.)

def scenes():
    c = FRAME / 2; s = FRAME / 640
    pent = [(c + 120 * s * np.cos(2 * np.pi * k / 5 + .3), c + 120 * s * np.sin(2 * np.pi * k / 5 + .3)) for k in range(5)]
    return {
        'glow': [('glow', c, c, 620 * s, 620 * s, VIOLET, .12, 0.)],
        'faint': [('glow', c, c, 620 * s, 620 * s, VIOLET, .04, 0.)],
        'stack': [('glow', x, y, (420 + 40 * k) * s, (420 + 40 * k) * s, (VIOLET, EMBER)[k % 2], .05 + .015 * k, 0.) for k, (x, y) in enumerate(pent)],
        'ring': [('ring', c, c, 620 * s, 620 * s, WHITE, .15, (.06, 0.)), ('ring', c, c, 420 * s, 420 * s, VIOLET, .3, (.1, .3))],
    }

def layer_alpha(kind, u, v, extra):
    return P.glow_alpha(u, v, extra) if kind == 'glow' else P.ring_alpha(u, v, *extra)

# The dither each layer and the screen write use at frame t: u(x, y, t) in [0, 1). The layers' noise is
# offset per layer (37, 91) as the game does (uSeed); the screen write's is IGN on its side (the game's now).
def field(kind, x, y, t):
    fr = lambda v: v - np.floor(v)
    if kind == 'still': return D.ign(x, y)
    if kind == 'half': return fr(D.ign(x, y) + .5 * (t % 2))
    if kind == 'golden': return fr(D.ign(x, y) + PHI * t)
    if kind == 'ign5588': return D.ign(x + 5.588238 * (t % 64), y + 5.588238 * (t % 64))
    if kind == 'white': return np.random.default_rng(1000 + t + int(x.flat[0])).random(x.shape)
    raise ValueError(kind)

CANDS = ['still', 'half', 'golden', 'ign5588']

def render(scene, kind, R, dpr, t):
    W = H = FRAME
    gx, gy = P.coords(W, H); gx = gx + np.float32(ORIGIN[0]); gy = gy + np.float32(ORIGIN[1])
    tgt = np.zeros((H, W, 3))
    for k, (kd, cx, cy, w, h, col, op, extra) in enumerate(scene):
        u, v, inside = P.quad_uv(W, H, cx, cy, w, h)
        rgb = np.array(col)[None, None] * (layer_alpha(kd, u, v, extra) * op)[..., None]
        d = field(kind, gx + 37 * k, gy + 91 * k, t) - .5
        rgb = rgb + (d / 255.)[..., None]
        tgt = np.where(inside[..., None], P.to_unorm8(tgt / 255. + np.clip(rgb, 0, 1)), tgt)
    disp = P.srgb(P.aces(P.kuwahara(tgt / 255., R, dpr)))
    disp = disp + ((field(kind, gy, gx, t) - .5) / 255.)[..., None]   # the screen write's dither, x and y swapped
    return P.to_unorm8(disp)

def ideal(scene, R, dpr):
    W = H = FRAME; t = np.zeros((H, W, 3))
    for kd, cx, cy, w, h, col, op, extra in scene:
        u, v, inside = P.quad_uv(W, H, cx, cy, w, h)
        t += np.where(inside[..., None], np.array(col)[None, None] * (layer_alpha(kd, u, v, extra) * op)[..., None], 0)
    return P.srgb(P.aces(P.kuwahara(t, R, dpr)))

def luminance(disp):
    c = np.clip(disp, 0, 1)
    lin = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
    return lin @ np.array([0.2126, 0.7152, 0.0722])

def lstar(Y):
    return np.where(Y > 216 / 24389, 116 * np.cbrt(Y) - 16, Y * 24389 / 27)

def blur(img, sigma):
    r = int(np.ceil(4 * sigma)); k = np.exp(-np.arange(-r, r + 1) ** 2 / (2 * sigma * sigma)); k /= k.sum()
    p = np.pad(img, r, mode='reflect')
    p = np.apply_along_axis(lambda a: np.convolve(a, k, 'valid'), 0, p)
    return np.apply_along_axis(lambda a: np.convolve(a, k, 'valid'), 1, p)

def eye_blend(Ys, G):
    """Luminance frames Ys[t] blended by an eye model; frames WARM .. T-1."""
    if G.taps is not None:
        return np.array([sum(w * Ys[t - s] for s, w in G.taps) for t in range(WARM, T)])
    y = Ys[0].copy(); out = []   # (as if the frames before had been like the first)
    for t in range(T):
        y = G.a * y + (1 - G.a) * Ys[t]
        if t >= WARM: out.append(y)
    return np.array(out)

B = 16
crop = lambda a: a[..., B:-B, B:-B]

def errors(Yseen, Yi):
    """band and grain (the 10-22 entry's measures) of each seen frame against the ideal, averaged."""
    sb, sv = .187 * PPD / 2, .187 * PPD / 8
    bi, vi = lstar(blur(Yi, sb)), lstar(blur(Yi, sv))
    band = grain = 0
    for Y in Yseen:
        eb = lstar(blur(Y, sb)) - bi; ev = lstar(blur(Y, sv)) - vi
        band += np.mean(crop(eb) ** 2); grain += np.mean(crop(ev - eb) ** 2)
    return np.sqrt(band / len(Yseen)), np.sqrt(grain / len(Yseen))

def seen(Ys, fps, model='bright'):
    """(grain, shimmer) in JND: luminance contrast against the local mean (blurred, plus the screen's floor
    Y0), split into each pixel's mean over the frames (the grain: what holds still, including any banding)
    and the frame-to-frame change (the shimmer), each through a Hann window, its spectrum weighted by vis.py's
    sensitivity at its spatial and temporal frequency. The same yardstick for a still dither and a moving one."""
    Y = crop(Ys[WARM:]); Yi = crop(YI)
    loc = blur(Yi, 2.0) + Y0
    m = Y.mean(0)
    g = (m - Yi) / loc; c = (Y - m) / loc
    n = c.shape[1]; Tn = c.shape[0]
    ws = np.hanning(n + 2)[1:-1]; w2 = ws[:, None] * ws[None, :]; wt = np.hanning(Tn + 2)[1:-1]
    fy = np.fft.fftfreq(n)[:, None]; fx = np.fft.fftfreq(n)[None, :]; ft = np.fft.fftfreq(Tn)[:, None, None]
    cpd = np.hypot(fx, fy) * PPD
    gm = (g * w2).sum() / w2.sum()   # (the mean offset, as the window sees it, counted once, apart: left in, it would leak into its neighbours)
    G = 2 * np.abs(np.fft.fft2((g - gm) * w2)) ** 2 / (w2 ** 2).sum() / (n * n); G[0, 0] = gm * gm
    C = 2 * np.abs(np.fft.fftn(c * w2 * wt[:, None, None])) ** 2 / ((w2 ** 2).sum() * (wt ** 2).sum()) / (n * n * Tn)
    S = vis.MODELS[model]
    return float(np.sqrt((G * S(cpd, 0.0) ** 2).sum())), float(np.sqrt((C * S(cpd[None], np.abs(ft) * fps) ** 2).sum()))
    # (2: a cosine's contrast is twice its two halves, as in vis.py)

EYES = [st.still(), st.box(2), st.box(4), st.decay(40), st.decay(80)]

def case(args):
    sname, kind, R = args
    sc = scenes()[sname]
    global YI
    Yi = YI = luminance(ideal(sc, R, 2))
    Ys = np.array([luminance(render(sc, kind, R, 2, t) / 255.) for t in range(T)])
    res = [errors(eye_blend(Ys, G), Yi) for G in EYES]
    return sname, kind, R, res, seen(Ys, 60), seen(Ys, 30), seen(Ys, 60, 'dark'), seen(Ys, 30, 'dark')

if __name__ == '__main__':
    jobs = [(s, k, R) for R in (5, 0) for s in scenes() for k in CANDS]
    with Pool(4) as pool:
        out = pool.map(case, jobs, chunksize=1)
    print('band / grain in L* (lower is better), each eye model; then in JND (the generous model in vis.py, and its dark variant): the grain, the shimmer at 60 fps and at 30 fps')
    for R in (5, 0):
        print(f'\npaint {"full" if R else "off"}, pixel ratio 2')
        print(f'  {"scene":6s} {"dither":8s} ' + ' '.join(f'{G.name:>13s}' for G in EYES) + '   JND: grain  shimmer 60   30 | dark: grain  60   30')
        for s, k, RR, res, s60, s30, d60, d30 in out:
            if RR != R: continue
            print(f'  {s:6s} {k:8s} ' + ' '.join(f'{b:6.3f} {g:6.3f}' for b, g in res) + f'   {s60[0]:10.2f} {s60[1]:11.2f} {s30[1]:5.2f} | {d60[0]:11.2f} {d60[1]:4.2f} {d30[1]:4.2f}')