# Lattice dithers d(x, y) = ((A x + B y + C) mod N + 1/2) / N seen through a filter: the exact
# mean-square error over flat levels, the bounds, and an exhaustive search over (A, B).
#
# J_N(A, B) = 1 / (12 N^2) + sum_{m=1}^{N-1} |H(m (A, B) / N)|^2 / (2 N^2 sin^2(pi m / N))
# (as N grows with m fixed the weight tends to 1 / (2 pi^2 m^2): the irrational-lattice formula).
import numpy as np
import filters as F

def J(A, B, N, H2):
    """Exact mean-square error (in 8-bit levels^2) of the lattice (A, B)/N; H2 = |H|^2 on the N x N grid.
    A, B may be arrays (vectorised over candidates)."""
    A = np.atleast_1d(A); B = np.atleast_1d(B); m = np.arange(1, N)[:, None]
    w = 1 / (2 * N * N * np.sin(np.pi * m / N) ** 2)
    return 1 / (12 * N * N) + (H2[(m * B) % N, (m * A) % N] * w).sum(0)

def search(H2, N, chunk=4096):
    """J for every (A, B) in [0, N)^2, as an N x N array indexed [B, A]. Pairs with gcd(A, B, N) > 1 make a
    lattice with fewer than N levels, which the formula doesn't describe: they get inf (they are searched
    as the smaller N they really are)."""
    AB = np.arange(N * N); out = np.empty(N * N)
    for s in range(0, N * N, chunk):
        out[s:s + chunk] = J(AB[s:s + chunk] % N, AB[s:s + chunk] // N, N, H2)
    A, B = AB % N, AB // N
    out[np.gcd(np.gcd(A, B), N) > 1] = np.inf
    return out.reshape(N, N)

def best(taps, Ns):
    """The best lattice over every N in Ns: (rms, A, B, N)."""
    top = (np.inf, 0, 0, 0)
    for N in Ns:
        S = search(F.grid_response(taps, N), N)
        b, a = np.unravel_index(np.argmin(S), S.shape)
        if np.sqrt(S[b, a]) < top[0]: top = (np.sqrt(S[b, a]), int(a), int(b), N)
    return top

def bounds(taps, M=2048):
    """Three floors on the RMS error. Any threshold array: sqrt(sum h^3 / 12). Any irrational lattice:
    sqrt(sum h^3 / 6). Any dither at all: any_floor(min over nonzero frequencies of |H|), about |H|min / sqrt(6)
    (the minimum found on an M x M grid, then refined by a local search)."""
    h = np.array([w for _, _, w in taps])
    H2 = F.grid_response(taps, M); H2[0, 0] = np.inf
    j, i = np.unravel_index(np.argmin(H2), H2.shape)
    fx, fy = (i if i < M / 2 else i - M) / M, (j if j < M / 2 else j - M) / M
    step = 1 / M
    for _ in range(40):  # refine the minimum off the grid
        cand = [(fx + a * step, fy + b * step) for a in (-1, 0, 1) for b in (-1, 0, 1)]
        vals = [abs(F.response(taps, x, y)) for x, y in cand]
        fx, fy = cand[int(np.argmin(vals))]; step /= 2
    hmin = abs(F.response(taps, fx, fy))
    return dict(array=np.sqrt((h ** 3).sum() / 12), lattice=np.sqrt((h ** 3).sum() / 6), any=any_floor(hmin), hmin=hmin, at=(fx, fy))

def any_floor(m, n=200000):
    """The floor for any dither, any rounding rule, any pattern. At level f the output is an integer image
    whose mean is some f' (ideally f'= f); by Parseval its filtered mean-square error is
    (f' - f)^2 + sum over nonzero frequencies of |H|^2 times its power spectrum, and an integer image with
    mean f' has variance at least f'(1 - f'). So the error is at least min over f' of
    (f' - f)^2 + m^2 f'(1 - f'), m = min |H|; averaged over f in [0, 1)."""
    f = (np.arange(n) + .5) / n
    fp = np.clip((2 * f - m * m) / (2 - 2 * m * m), 0, 1)
    return np.sqrt((m * m * fp * (1 - fp) + (f - fp) ** 2).mean())