# Brute force in space and time, independent of the formulas: dither a periodic block of frames, round,
# filter each frame by the spatial taps (direct shifts, no FFT), blend over frames in the time domain (a
# sliding box, or the decaying average run recursively until it settles), compare with the level.
import numpy as np

def spatial(img, taps):
    """img[..., y, x], periodic; out = sum_q h_q img(p + q) (the same sign convention as filters.response)."""
    out = np.zeros_like(img)
    for dx, dy, w in taps:
        out += w * np.roll(img, (-dy, -dx), axis=(-2, -1))
    return out

def temporal(seq, G):
    """seq[t, ...], periodic in t. A box: the mean of frames t, t-1, ..., t-k+1. A decay: run
    y_t = a y_{t-1} + (1 - a) x_t around the period until it has settled (to 1e-13)."""
    if G.taps is not None:
        out = np.zeros_like(seq)
        for s, w in G.taps: out += w * np.roll(seq, s, axis=0)
        return out
    a = G.a; T = seq.shape[0]; y = np.zeros_like(seq[0])
    laps = int(np.ceil(np.log(1e-13) / np.log(a) / T)) + 1
    for _ in range(laps):
        for t in range(T): y = a * y + (1 - a) * seq[t]
    out = np.empty_like(seq)
    for t in range(T): y = a * y + (1 - a) * seq[t]; out[t] = y
    return out

def lattice_block(A, B, C, N, n, T):
    t, y, x = np.mgrid[0:T, 0:n, 0:n]
    return (((A * x + B * y + C * t) % N) + .5) / N

def rms(u, taps, G, levels=1024, W=None):
    """u[t, y, x]: the threshold block (periodic). Mean over levels and the block of e^2 seen through
    the spatial taps and the temporal blend G. With W (a 0/1 function of tau), instead the power of the
    spatially filtered error at the temporal frequencies W selects (by FFT along t: the shimmer)."""
    e2 = 0
    T = u.shape[0]
    if W is not None: wsel = W(np.fft.fftfreq(T))[:, None, None]
    for f in (np.arange(levels) + .5) / levels:
        e = spatial(np.floor(f + u) - f, taps)
        if W is None:
            e2 += (temporal(e, G) ** 2).mean()
        else:
            E = np.fft.fft(e, axis=0) / T
            e2 += (np.abs(E) ** 2 * wsel).sum(0).mean()
    return e2 / levels