# Every space-time lattice (A, B, C)/N up to N levels, through Starlock's pass (paint on, a tap per CSS pixel)
# or with paint off (a device pixel at a pixel ratio of 2; the identity, and the identity plus a slight eye
# blur), under each eye model: the best overall; the best whose shimmer is no easier to see than the grain of
# the best still lattice (vis.py's generous model, at 60 fps, and at 30 fps too); and the one the eye sees
# least of in all (grain and shimmer pooled, by the same model at 60 fps).
# Prints out/search.txt. Usage: python3 search.py [Nmax] (default 64; a few minutes).
import sys
import numpy as np
import filters as F, st, vis

NMAX = int(sys.argv[1]) if len(sys.argv) > 1 else 64
EYES = [st.still(), st.box(2), st.box(3), st.box(4), st.box(6), st.decay(20), st.decay(40), st.decay(80)]
CHECKS = [('bright', 60), ('bright', 30), ('dark', 60), ('dark', 30)]

def visibility(N, taps, deg, Cs):
    """(grain, shimmer) L*[B, A, C] for every lattice with N levels, for each check."""
    m = np.arange(1, N)[:, None, None]; w = 1 / (2 * N * N * np.sin(np.pi * m / N) ** 2)
    a = np.arange(N)[None, None, :]; b = np.arange(N)[None, :, None]
    fx, fy = vis.wrapf(m * a / N), vis.wrapf(m * b / N)
    P = w * np.abs(F.response(taps, fx, fy)) ** 2            # [m, B, A]
    cpd = np.hypot(fx, fy) / deg
    hz = np.abs(vis.wrapf(np.arange(1, N)[:, None] * Cs[None, :] / N))   # [m, C] in cycles per frame
    out = []
    for model, fps in CHECKS:
        S = vis.MODELS[model]
        g = np.zeros((N, N, len(Cs))); sh = np.zeros((N, N, len(Cs)))
        for j, c in enumerate(Cs):
            V = 2 * P * S(cpd, hz[:, j][:, None, None] * fps) ** 2
            still = (hz[:, j] == 0)
            g[:, :, j] = V[still].sum(0) + S(0.0, 0.0) ** 2 / (12 * N * N); sh[:, :, j] = V[~still].sum(0)
        out.append((np.sqrt(g), np.sqrt(sh)))
    return out

def run(name, taps, deg, Ns, vtaps=None):
    print(f'== {name}')
    data = []
    for N in Ns:
        Cs = np.arange(N)
        Js, _ = st.search(N, F.grid_response(taps, N), EYES, [])
        data.append((N, Js, visibility(N, vtaps or taps, deg, Cs)))
    # the grain of the best still lattice (C = 0) under each check: the bar a shimmer must stay under
    bar = []
    for c in range(len(CHECKS)):
        best = (np.inf, 0)
        for N, Js, Ls in data:
            v = np.where(np.isfinite(Js[0][:, :, 0]), Js[0][:, :, 0], np.inf); k = np.unravel_index(np.argmin(v), v.shape)
            if v[k] < best[0]: best = (v[k], Ls[c][0][k[0], k[1], 0])
        bar.append(best[1])
    print('  the best still lattice\'s grain, L* by check (' + ', '.join(f'{m} {f} fps' for m, f in CHECKS) + '): ' + ' '.join(f'{b:.2f}' for b in bar))
    rows = [[(np.inf,)] * 3 for _ in EYES]
    for N, Js, Ls in data:
        calm60 = Ls[0][1] <= bar[0]; calm30 = calm60 & (Ls[1][1] <= bar[1])
        for i, Jg in enumerate(Js):
            for j, mask in enumerate((np.ones_like(calm60), calm60, calm30)):
                v = np.where(mask & np.isfinite(Jg), Jg, np.inf); k = np.unravel_index(np.argmin(v), v.shape)
                if np.sqrt(v[k]) < rows[i][j][0] - 1e-7:
                    b, a, c = k; rows[i][j] = (np.sqrt(v[k]), a, b, c, N, Ls[0][0][k], Ls[0][1][k], Ls[1][1][k])
    f = lambda r: f'({r[1]},{r[2]},{r[3]})/{r[4]} {r[0]:.4f} [{r[5]:.1f} {r[6]:.1f} {r[7]:.1f}]'
    print('  eye          best overall                          calm at 60 fps                        calm at 30 fps too')
    for G, r in zip(EYES, rows):
        print(f'  {G.name:12s} {f(r[0]):37s} {f(r[1]):37s} {f(r[2])}')
    # least seen in all: grain and shimmer pooled, at 60 fps (bright), and what it gives at 30
    for c, (m, fps) in enumerate(CHECKS):
        best = (np.inf,)
        for N, Js, Ls in data:
            tot = np.hypot(Ls[c][0], Ls[c][1]); tot = np.where(np.isfinite(Js[0]), tot, np.inf)
            k = np.unravel_index(np.argmin(tot), tot.shape)
            if tot[k] < best[0] - 1e-9: b, a, cc = k; best = (tot[k], a, b, cc, N, Ls[c][0][k], Ls[c][1][k], np.sqrt(Js[1][k]))
        print(f'  least seen in all, {m} at {fps} fps: ({best[1]},{best[2]},{best[3]})/{best[4]}  L* {best[0]:.2f} (grain {best[5]:.2f}, shimmer {best[6]:.2f}); 2-frame box error {best[7]:.4f}')

print('[a b c] = L* (the background, in levels, below which it passes 1 JND), generous model: grain at 60 fps, shimmer at 60 fps, shimmer at 30 fps')
ID = [(0, 0, 1.0)]
run('Starlock full pass (per tap, a CSS pixel)', F.kuwahara(5), vis.CSS_DEG, range(2, NMAX + 1))
run('paint off, per device pixel, seen as is', ID, vis.DEV_DEG, range(2, NMAX + 1))
run('paint off, per device pixel, through a slight eye blur (Gaussian, 1 device pixel)', F.gaussian(1.0), vis.DEV_DEG, range(2, NMAX + 1), vtaps=ID)