# The dithers games use now, scored by simulation (they are not lattices, so there's no exact formula): a
# block of frames of each, rounded at a sweep of flat levels, filtered in space (on a canvas larger than the
# window kept, so nothing wraps), blended over frames by each eye model, and the shimmer measured by a
# windowed FFT along time with each pixel's mean taken out.
import numpy as np
import st, vis

PHI = 0.6180339887498949

def fr(v):
    return v - np.floor(v)

def ign_xy(x, y):
    """Jimenez's interleaved gradient noise, in float32 as a GPU computes it."""
    x = x.astype(np.float32); y = y.astype(np.float32)
    f = lambda v: (v - np.floor(v)).astype(np.float32)
    return f(np.float32(52.9829189) * f(x * np.float32(0.06711056) + y * np.float32(0.00583715))).astype(float)

def make(name, n, T, seed=1, tile=None):
    """Frames [T, n, n] of a named dither, pixel centres at (x + 1/2, y + 1/2) like gl_FragCoord."""
    y, x = np.mgrid[0:n, 0:n] + .5
    y = y + 1100; x = x + 60          # (somewhere on a phone screen, like the entries before)
    out = np.empty((T, n, n))
    rng = np.random.default_rng(seed)
    for t in range(T):
        if name == 'ign still': out[t] = ign_xy(x, y)
        elif name == 'ign + 5.588238 t': out[t] = ign_xy(x + 5.588238 * (t % 64), y + 5.588238 * (t % 64))
        elif name == 'ign + golden t': out[t] = fr(ign_xy(x, y) + PHI * t)
        elif name == 'white': out[t] = rng.random((n, n))
        elif name == 'blue still': out[t] = tile[0][(y.astype(int)) % tile.shape[1], (x.astype(int)) % tile.shape[2]]
        elif name == 'blue + golden t': out[t] = fr(tile[0][(y.astype(int)) % tile.shape[1], (x.astype(int)) % tile.shape[2]] + PHI * t)
        elif name == 'stbn': out[t] = tile[t % tile.shape[0]][(y.astype(int)) % tile.shape[1], (x.astype(int)) % tile.shape[2]]
        elif name == 'ign + half every other frame': out[t] = fr(ign_xy(x, y) + .5 * (t % 2))
        elif name == 'lattice + golden t':
            A, B, C, N = tile; out[t] = fr((((A * x.astype(int) + B * y.astype(int)) % N) + .5) / N + PHI * t)
        elif name.startswith('lattice'):
            A, B, C, N = tile
            xi, yi = x.astype(int), y.astype(int)
            out[t] = (((A * xi + B * yi + C * t) % N) + .5) / N
        else: raise ValueError(name)
    return out

def spatial_fft(frames, taps, keep):
    """Filter each frame by the taps (out(p) = sum h_q in(p + q)); return the central keep x keep window."""
    T, n, _ = frames.shape
    K = np.zeros((n, n))
    for dx, dy, w in taps: K[(-dy) % n, (-dx) % n] += w
    FK = np.fft.rfft2(K)
    out = np.fft.irfft2(np.fft.rfft2(frames) * FK, s=(n, n))
    o = (n - keep) // 2
    return out[:, o:o + keep, o:o + keep]

def eye(seq, G, warm):
    """seq[t]: blend over frames causally (no wrap): a box over the last k, or the decay, from frame 0;
    return frames warm .. T-1."""
    T = seq.shape[0]
    if G.taps is not None:
        out = np.zeros_like(seq[warm:])
        for s, w in G.taps: out += w * seq[warm - s:T - s]
        return out
    y = np.zeros_like(seq[0]); out = []
    for t in range(T):
        y = G.a * y + (1 - G.a) * seq[t]
        if t >= warm: out.append(y.copy())
    return np.array(out)

def shimmer(seq, W):
    """Power of seq[t] (per pixel, mean removed) at the temporal frequencies W selects, by a Hann-windowed FFT."""
    T = seq.shape[0]
    x = seq - seq.mean(0)
    w = np.hanning(T)[:, None, None]
    X = np.fft.fft(x * w, axis=0)
    sel = W(np.fft.fftfreq(T))[:, None, None]
    return float((np.abs(X) ** 2 * sel).sum(0).mean() / (T * (w ** 2).sum() / T) / T)

def seen(e, deg, checks):
    """(grain, shimmer) squared JND per check, from frames e[t, y, x] that need not repeat: the grain from
    each pixel's mean over the frames, the shimmer from what's left, both through a Hann window (in space for
    the grain; in space and time for the shimmer), normalised by the window's energy."""
    T, n, _ = e.shape
    still = e.mean(0); rest = e - still
    ws = np.hanning(n + 2)[1:-1]; w2 = ws[:, None] * ws[None, :]; wt = np.hanning(T + 2)[1:-1]
    fy = np.fft.fftfreq(n)[:, None]; fx = np.fft.fftfreq(n)[None, :]; ft = np.fft.fftfreq(T)[:, None, None]
    cpd = np.hypot(fx, fy) / deg
    G = np.abs(np.fft.fft2(still * w2)) ** 2 / (w2 ** 2).sum() / (n * n)
    R = np.abs(np.fft.fftn(rest * w2 * wt[:, None, None])) ** 2 / ((w2 ** 2).sum() * (wt ** 2).sum()) / (n * n * T)
    G *= 2; G[0, 0] /= 2; R *= 2   # (a cosine's contrast is twice its two halves: vis.py's 2 P)
    out = []
    for model, fps in checks:
        S = vis.MODELS[model]
        out.append(((G * S(cpd, 0.0) ** 2).sum(), (R * S(cpd[None], np.abs(ft) * fps) ** 2).sum()))
    return out

def score(frames, taps, eyes, bands, levels=48, keep=96, warm=48, deg=None, checks=()):
    """Mean over flat levels of: the error power of one frame; kept by each eye model; the shimmer in each
    band; and the (grain, shimmer) JND of each check. Square roots of the means."""
    kept = np.zeros(len(eyes)); sh = np.zeros(len(bands)); still = 0; js = np.zeros((len(checks), 2))
    for f in (np.arange(levels) + .5) / levels:
        e = spatial_fft(np.floor(f + frames) - f, taps, keep)
        still += (e[warm:] ** 2).mean()
        for i, G in enumerate(eyes): kept[i] += (eye(e, G, warm) ** 2).mean()
        for i, W in enumerate(bands): sh[i] += shimmer(e[warm:], W)
        if checks: js += np.array(seen(e[warm:], deg, checks))
    return np.sqrt(still / levels), np.sqrt(kept / levels), np.sqrt(sh / levels), np.sqrt(js / levels)