# C4: exhaustive search over every (A,B,C)/N, N <= 64, gcd(A,B,C,N)=1, with the (verified) closed form, my own
# implementation: |H|^2 on the N x N grid from an FFT of the collapsed taps. Also the best still lattice with
# at most M levels for each M, to test "best box-k = best still / k".  run: python3 c4_search.py [Nmax]
import sys, math, numpy as np, vtaps
NMAX = int(sys.argv[1]) if len(sys.argv) > 1 else 64
taps = vtaps.kuw(5)
def G2box(k):
    def f(t):
        s = np.sin(np.pi*t); out = np.ones_like(t); nz = np.abs(s) > 1e-12
        out[nz] = (np.sin(np.pi*k*t[nz])/(k*s[nz]))**2; return out
    return f
def G2dec(ms):
    a = math.exp(-1000/60/ms); return lambda t: np.abs((1 - a)/(1 - a*np.exp(-2j*np.pi*t)))**2
EYES = {'one frame': lambda t: np.ones_like(t), 'box 2': G2box(2), 'box 3': G2box(3), 'box 4': G2box(4), 'box 6': G2box(6),
        'decay 20': G2dec(20), 'decay 40': G2dec(40), 'decay 80': G2dec(80)}
best = {k: (np.inf,) for k in EYES}; still_by_N = {}
for N in range(2, NMAX + 1):
    K = np.zeros((N, N))
    for (dx, dy), w in taps.items(): K[dy % N, dx % N] += w
    H2 = np.abs(np.fft.fft2(K))**2         # H2[j, i] = |H(i/N, j/N)|^2 up to sign convention (|.| is symmetric)
    m = np.arange(1, N); w = 1/(2*N*N*np.sin(np.pi*m/N)**2)
    a = np.arange(N)
    P = w[:, None, None]*H2[(m[:, None, None]*a[None, :, None]) % N, (m[:, None, None]*a[None, None, :]) % N]  # [m, B, A]
    tau = (m[:, None]*a[None, :] % N)/N    # [m, C]
    g = np.gcd(np.gcd(a[:, None, None], a[None, :, None]), np.gcd(a[None, None, :], N))  # [B, A, C]
    for name, G2 in EYES.items():
        J = 1/(12*N*N) + np.einsum('mba,mc->bac', P, G2(tau)); J[g > 1] = np.inf
        if name == 'one frame': still_by_N[N] = math.sqrt(J[:, :, 0].min())
        k = np.unravel_index(np.argmin(J), J.shape)
        if math.sqrt(J[k]) < best[name][0] - 1e-12: best[name] = (math.sqrt(J[k]), k[1], k[0], k[2], N)
for name, b in best.items(): print(f'{name:10s} {b[0]:.5f} ({b[1]},{b[2]},{b[3]})/{b[4]}')
bs = best['one frame'][0]
for k in (2, 3, 4, 6):
    sm = min(v for M, v in still_by_N.items() if M <= NMAX//k)
    print(f'box {k}: best {best[f"box {k}"][0]:.5f}; best still/k {bs/k:.5f}; best still with <= {NMAX//k} levels, /k: {sm/k:.5f}')