import numpy as np, mytaps as T
def Jform(t, A, B, N):
    m = np.arange(1, N); s = 1/(12*N*N)
    Hm = T.H(t, m*A/N, m*B/N)
    return s + np.sum(np.abs(Hm)**2/(2*N*N*np.sin(np.pi*m/N)**2))
def conv(img, t):
    out = np.zeros_like(img)
    for (dx, dy), w in t.items(): out += w*np.roll(np.roll(img, -dy, 0), -dx, 1)  # out[p] = sum h_q img[p+q]
    return out
def brute_exact(t, d, N):
    """d: periodic threshold field, values (k+.5)/N. Exact integral over f in [0,1) of mean_p e^2."""
    br = np.concatenate([[0], (np.arange(N)+.5)/N, [1]]); tot = 0
    for a, b in zip(br[:-1], br[1:]):
        S = conv(np.floor((a+b)/2 + d), t)
        tot += np.mean(((S-a)**3 - (S-b)**3)/3)
    return tot
def lat(A, B, N, n, C=0):
    y, x = np.mgrid[0:n, 0:n]; return (((A*x + B*y + C) % N) + .5)/N
if __name__ == '__main__':
    rng = np.random.default_rng(5)
    filters = {'kuw5': T.kuw(5), 'kuw2': T.kuw(2), 'g1.5c.2': T.gauss_centre(.2, 1.5), 'box3': T.box(3), 'kuw5s2': T.kuw(5, 2)}
    # random asymmetric filter too
    rt = {(int(a), int(b)): w for a, b, w in zip(rng.integers(-3, 5, 12), rng.integers(-2, 4, 12), rng.random(12))}
    S = sum(rt.values()); filters['random'] = {k: v/S for k, v in rt.items()}
    worst = 0
    for name, t in filters.items():
        for A, B, N in [(54, 11, 56), (5, 1, 24), (3, 1, 16), (2, 1, 5), (7, 3, 12), (13, 5, 32), (1, 0, 7), (6, 4, 9)]:
            if np.gcd(np.gcd(A, B), N) > 1: continue
            n = N*int(np.ceil(40/N)); C = int(rng.integers(N))
            jb = brute_exact(t, lat(A, B, N, n, C), N); jf = Jform(t, A, B, N)
            worst = max(worst, abs(jb-jf)/jf)
            print(f'{name:8s} ({A},{B})/{N}: formula rms {np.sqrt(jf):.6f} brute {np.sqrt(jb):.6f} rel {abs(jb-jf)/jf:.1e}')
    print('worst rel diff', worst)
    # limit N -> infinity for alpha=(0.7548776662,0.5698402910)
    t = T.kuw(5); al = np.array([0.7548776662, 0.5698402910]); k = np.arange(1, 200001)
    lim = np.sum(np.abs(T.H(t, k*al[0], k*al[1]))**2/(np.pi**2*k**2))
    for N in (1024, 4096, 16384, 65536):
        print('N', N, np.sqrt(Jform(t, round(al[0]*N), round(al[1]*N), N)), 'limit', np.sqrt(lim))