# Space-time lattice dithers d(x, y, t) = ((A x + B y + C t) mod N + 1/2) / N seen through a spatial filter h
# (Starlock's pass, or the identity with the pass off) and the eye's blending over frames g: the exact
# mean-square error, the shimmer, and searches over (A, B, C).
#
# The error of harmonic m (m = 1 .. N-1) sits at spatial frequency m (A, B) / N and temporal frequency m C / N
# (cycles per frame), with power 1 / (2 N^2 sin^2(pi m / N)) before any filter. So
#   J  = 1 / (12 N^2) + sum_m |H(m (A, B) / N)|^2 |G(m C / N)|^2 / (2 N^2 sin^2(pi m / N))   (what the eye keeps)
#   F  = sum_m |H(m (A, B) / N)|^2 W(m C / N) / (2 N^2 sin^2(pi m / N))                     (the shimmer)
# where W is 1 on the temporal frequencies a viewer sees as shimmer (0 < |tau| fps <= fc) and 0 elsewhere.
import numpy as np
import filters as F

# --- the eye's blending over frames, as a temporal response G(tau), tau in cycles per frame
def box(k):
    """A sliding box over the last k frames."""
    def G(tau):
        tau = np.asarray(tau, float); s = np.sin(np.pi * tau)
        with np.errstate(invalid='ignore', divide='ignore'):
            r = np.sin(np.pi * k * tau) / (k * s)
        return np.where(np.abs(s) < 1e-12, 1.0, np.abs(r))
    G.taps = [(t, 1 / k) for t in range(k)]; G.name = f'box {k}'
    return G

def decay(ms, fps=60):
    """A decaying average with time constant `ms`: each frame keeps a = exp(-frame / ms) of what came before."""
    a = np.exp(-1000 / fps / ms)
    def G(tau):
        return np.abs((1 - a) / (1 - a * np.exp(-2j * np.pi * np.asarray(tau, float))))
    G.a = a; G.taps = None; G.name = f'decay {ms:g} ms'
    return G

def still():
    """No blending over frames: the eye sees each frame alone (a still image, a screenshot)."""
    G = lambda tau: np.ones_like(np.asarray(tau, float)); G.taps = [(0, 1.0)]; G.name = 'one frame'
    return G

def wrap(tau):
    return np.abs((np.asarray(tau, float) + .5) % 1 - .5)

def shimmer_band(fc, fps=60):
    """W(tau): 1 where the frame-to-frame change has a temporal frequency between 0 (excluded: a still
    pattern is grain, not shimmer) and fc Hz on a display running at fps frames a second."""
    return lambda tau: ((wrap(tau) > 1e-12) & (wrap(tau) * fps <= fc + 1e-9)).astype(float)

# --- the exact formulas
def harmonics(A, B, C, N, taps):
    """For m = 1 .. N-1: the weight 1 / (2 N^2 sin^2(pi m / N)), |H|^2 at m (A, B) / N and tau = m C / N."""
    m = np.arange(1, N)
    w = 1 / (2 * N * N * np.sin(np.pi * m / N) ** 2)
    H2 = np.abs(F.response(taps, (m * A % N) / N, (m * B % N) / N)) ** 2
    return w, H2, (m * C % N) / N

def J(A, B, C, N, taps, G):
    w, H2, tau = harmonics(A, B, C, N, taps)
    return 1 / (12 * N * N) + (w * H2 * G(tau) ** 2).sum()

def Fl(A, B, C, N, taps, W):
    w, H2, tau = harmonics(A, B, C, N, taps)
    return (w * H2 * W(tau)).sum()

def search(N, H2grid, Gs, Ws, Cs=None):
    """Every (A, B, C) with N levels (gcd(A, B, C, N) = 1). H2grid = |H|^2 on the N x N grid [B-freq, A-freq]
    (filters.grid_response). Returns arrays J[g][B, A, C] for each eye model and Fl[w][B, A, C] for each band."""
    Cs = np.arange(N) if Cs is None else np.asarray(Cs)
    m = np.arange(1, N); wt = 1 / (2 * N * N * np.sin(np.pi * m / N) ** 2)
    A = np.arange(N)[None, :, None]; B = np.arange(N)[:, None, None]
    # P[m, B, A] = w_m |H(m (A, B) / N)|^2
    P = wt[:, None, None] * H2grid[(m[:, None, None] * B[None, :, :, 0]) % N, (m[:, None, None] * A[None, :, :, 0]) % N]
    tau = (m[:, None] * Cs[None, :] % N) / N          # [m, C]
    out_J = [1 / (12 * N * N) + np.einsum('mba,mc->bac', P, G(tau) ** 2) for G in Gs]
    out_F = [np.einsum('mba,mc->bac', P, W(tau)) for W in Ws]
    g = np.gcd(np.gcd(np.arange(N)[:, None, None], np.arange(N)[None, :, None]), np.gcd(Cs[None, None, :], N))
    bad = g > 1
    for x in out_J + out_F: x[bad] = np.inf
    return out_J, out_F