# C1, independently: the exact mean-square error of a space-time lattice by the residue method (no FFT, no
# level sampling): every tap (spatial and temporal) only shifts the residue r = (Ax+By+Ct) mod N, so the eye's
# output at a point is a circular correlation of an N-vector of collapsed weights with the 0/1 output q(r),
# which is constant on N+1 level intervals; (y - v)^2 is integrated exactly over each. Compared with the
# claimed closed form (implemented here from scratch) on random lattices, filters and eye models.
# run: python3 c1_exact.py
import numpy as np, math, vtaps
rng = np.random.default_rng(7)

def spatial_vec(A, B, N, taps):
    u = np.zeros(N)
    for (dx, dy), w in taps.items(): u[(A*dx + B*dy) % N] += w
    return u
def temporal_vec(C, N, eye):
    u = np.zeros(N)
    if eye[0] == 'box':
        for s in range(eye[1]): u[(-C*s) % N] += 1/eye[1]
    elif eye[0] == 'decay':
        a = eye[1]; P = N // math.gcd(C, N)
        for s in range(P): u[(-C*s) % N] += (1 - a)*a**s/(1 - a**P)
    else: u[0] = 1
    return u
def cconv(a, b):
    N = len(a); out = np.zeros(N)
    for i in np.nonzero(a)[0]: out += a[i]*np.roll(b, i)
    return out
def exact_J(A, B, C, N, taps, eye):
    u = cconv(spatial_vec(A, B, N, taps), temporal_vec(C, N, eye))   # u[delta]
    tot = 0.0; r = np.arange(N)
    for c in range(N + 1):
        lo, hi = max(0, (c - .5)/N), min(1, (c + .5)/N)
        if hi <= lo: continue
        q = (r >= N - c).astype(float)
        y = np.array([np.dot(u, q[(r0 + r) % N]) for r0 in range(N)])
        tot += np.mean(((hi - y)**3 - (lo - y)**3)/3)
    return tot
def exact_shimmer(A, B, C, N, taps, W):
    us = spatial_vec(A, B, N, taps); r = np.arange(N); g = math.gcd(C, N); P = N//g
    if P == 1: return 0.0
    tot = 0.0; k = np.arange(P); sel = W(k/P)
    for c in range(N + 1):
        lo, hi = max(0, (c - .5)/N), min(1, (c + .5)/N)
        if hi <= lo: continue
        q = (r >= N - c).astype(float)
        z = np.array([np.dot(us, q[(r0 + r) % N]) for r0 in range(N)])
        acc = 0.0
        for rs in range(g):   # each coset is visited by the pixels equally often
            ts = z[(rs + C*np.arange(P)) % N]
            Z = np.fft.fft(ts)/P; acc += (np.abs(Z)**2*sel).sum()
        tot += (hi - lo)*acc/g
    return tot
def formula(A, B, C, N, taps, G2, extra=True):
    m = np.arange(1, N)
    H2 = np.abs(vtaps.H(taps, m*A/N, m*B/N))**2
    s = (H2*G2((m*C % N)/N)/(2*N*N*np.sin(np.pi*m/N)**2)).sum()
    return s + (1/(12*N*N) if extra else 0)
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(a): return lambda t: np.abs((1 - a)/(1 - a*np.exp(-2j*np.pi*t)))**2
def Wband(fc, fps=60):
    return lambda t: ((np.abs((t + .5) % 1 - .5) > 1e-12) & (np.abs((t + .5) % 1 - .5)*fps <= fc + 1e-9)).astype(float)

if __name__ == '__main__':
    kuw = vtaps.kuw(5); ident = {(0, 0): 1.0}
    rnd = {(int(rng.integers(-3, 4)), int(rng.integers(-3, 4))): float(w) for w in rng.random(9)}
    s = sum(rnd.values()); rnd = {k: v/s for k, v in rnd.items()}
    filt = {'kuwahara': kuw, 'identity': ident, 'random asym': rnd}
    worst = 0; worst_s = 0; n = 0
    for trial in range(400):
        N = int(rng.integers(2, 65))
        while True:
            A, B, C = (int(x) for x in rng.integers(0, N, 3))
            if math.gcd(math.gcd(A, B), math.gcd(C, N)) == 1: break
        fname = list(filt)[trial % 3]; taps = filt[fname]
        kind = trial % 4
        if kind == 0: eye = ('box', int(rng.integers(1, 7))); G2 = G2box(eye[1])
        elif kind == 1:
            ms = float(rng.choice([20, 40, 80, 10])); a = math.exp(-1000/60/ms); eye = ('decay', a); G2 = G2dec(a)
        elif kind == 2: eye = ('still',); G2 = lambda t: np.ones_like(t)
        else: eye = ('box', N); G2 = G2box(N)
        e = exact_J(A, B, C, N, taps, eye); f = formula(A, B, C, N, taps, G2)
        rel = abs(e - f)/e; worst = max(worst, rel); n += 1
        if rel > 1e-9: print('MISMATCH', A, B, C, N, fname, eye, e, f)
        if trial % 5 == 0:
            W = Wband(float(rng.choice([10, 20, 30])))
            es = exact_shimmer(A, B, C, N, taps, W); fs = formula(A, B, C, N, taps, W, extra=False)
            d = abs(es - fs)/max(es, 1e-30) if es > 1e-15 else abs(fs); worst_s = max(worst_s, d)
            if d > 1e-9: print('SHIMMER MISMATCH', A, B, C, N, fname, es, fs)
    print(f'{n} lattices: worst relative difference J {worst:.2e}; shimmer {worst_s:.2e}')
    # the named ones
    for (A, B, C, N), eye, G2, nm in [((54, 11, 0, 56), ('still',), lambda t: np.ones_like(t), 'one frame'),
                                      ((25, 5, 26, 52), ('box', 2), G2box(2), 'box 2')]:
        print(nm, (A, B, C, N), f'{math.sqrt(exact_J(A, B, C, N, kuw, eye)):.5f}')