# Checks the floor's certificates with numpy alone (no LP solver): python3 check_floor.py certs.json
#
# The claim: through Starlock's painterly pass (filters.kuwahara(5), per tap), any dither whose output at
# each pixel is one of the two levels around the input (any q = floor(v + d), any d in [0, 1), patterns
# free to change with the level) leaves an RMS error, over flat levels, of at least the number printed.
#
# Why a certificate proves it. At level f the rounded image is a 0/1 field b with some mean f'. Its error
# is (f' - f)^2 + sum_xi |H(xi)|^2 P(xi), P the power spectrum of b - f' (Parseval). A certificate is a set
# of windows (point sets) with symmetric weights W and vectors rho, giving lam_r = sum over pairs (p, q)
# of a window with q - p = r of W_pq. If |H|^2 >= lam_hat - eps everywhere, then
#   sum |H|^2 P >= sum_r lam_r C(r) - eps f'(1-f')       (C the autocovariance)
#              = E[sum_W x'Wx] - f'^2 sum lam - eps f'(1-f')
#              >= sum_W min_x (x'Wx - rho.x) + f' sum rho - f'^2 sum lam - eps f'(1-f'),
# the last step because x'Wx - rho.x is at least its minimum over all 0/1 configurations of the window.
# This script recomputes those minima exactly, finds eps on a fine grid with a proved margin (the
# minimum of a trigonometric polynomial g is at least its minimum over a grid of spacing s minus
# (1/2) max|g''| (s/sqrt 2)^2, since the gradient vanishes at the true minimum), and integrates over f,
# taking for each f the best certificate and the worst mean f' exactly (floor_rms; a field's complement has
# the same covariance, so the bound at f' also holds at 1 - f').
import sys, os, json, itertools
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import filters as F

def taps_K(taps):
    K = {}
    for dx, dy, w in taps:
        for ex, ey, v in taps:
            r = (ex - dx, ey - dy); K[r] = K.get(r, 0) + w * v
    return K

def certificate_bound(c, K, G):
    lam = {}; const = 0.0; lin = 0.0
    for win in c['windows']:
        P = [tuple(p) for p in win['points']]; W = np.array(win['W']); rho = np.array(win['rho']); n = len(P)
        assert np.allclose(W, W.T)
        for a in range(n):
            for b in range(n):
                r = (P[b][0] - P[a][0], P[b][1] - P[a][1]); lam[r] = lam.get(r, 0) + W[a, b]
        X = np.array(list(itertools.product((0, 1), repeat=n)), float)
        const += np.min(np.einsum('ij,jk,ik->i', X, W, X) - X @ rho)
        lin += rho.sum()
    quad = sum(lam.values())
    # eps: how far lam_hat exceeds |H|^2 anywhere
    coef = dict(K)
    for r, v in lam.items(): coef[r] = coef.get(r, 0) - v
    A = np.zeros((G, G))
    for (dx, dy), v in coef.items(): A[dy % G, dx % G] += v
    g = np.real(np.fft.fft2(A))  # g = |H|^2 - lam_hat on the grid (coefficients are symmetric)
    second = 4 * np.pi ** 2 * sum(abs(v) * (r[0] ** 2 + r[1] ** 2) for r, v in coef.items())
    margin = 0.5 * second * 2 * (0.5 / G) ** 2
    eps = max(0.0, margin - g.min())
    return const, lin, quad, eps

def floor_rms(rows, nf=1_000_000):
    """sqrt of the integral over f of LB(f) = min over f' in [0, 1] of (f' - f)^2 + Phi(f'), with
    Phi(f') = max(0, max_i piece_i(h)), h = min(f', 1 - f'), piece_i(h) = c_i + b_i h - q_i h^2 - e_i h (1 - h).
    Exact minimisation: on each half [0, 1/2], [1/2, 1] every piece is a quadratic in f' and (f' - f)^2 is
    added to all of them alike, so the minimum is at an end, at a crossing of two pieces (these don't depend
    on f), or at one piece's own stationary point. Then a lower Riemann sum: on each of nf cells LB is at
    least its value at the midpoint minus (cell / 2) times |dLB/df| = 2 |f - f'*| <= 2."""
    P = [(0.0, 0.0, 0.0)] + [(c, b - e, q - e) for c, b, q, e in rows]   # piece(h) = c + B h - Q h^2; the 0 branch too
    def coeffs(half):  # piece as a polynomial in f': h = f' (half 0) or 1 - f' (half 1) -> (k0, k1, k2): k0 + k1 x + k2 x^2
        out = []
        for c, B, Q in P:
            if half == 0: out.append((c, B, -Q))
            else: out.append((c + B - Q, -B + 2 * Q, -Q))
        return out
    cand = [0.0, 0.5, 1.0]
    for half, (lo, hi) in ((0, (0, .5)), (1, (.5, 1))):
        cs = coeffs(half)
        for i in range(len(cs)):
            for j in range(i + 1, len(cs)):
                d0, d1, d2 = cs[i][0] - cs[j][0], cs[i][1] - cs[j][1], cs[i][2] - cs[j][2]
                roots = np.roots([d2, d1, d0]) if abs(d2) > 1e-15 else ([-d0 / d1] if abs(d1) > 1e-15 else [])
                cand += [float(r.real) for r in np.atleast_1d(roots) if abs(np.imag(r)) < 1e-12 and lo <= r.real <= hi]
    cand = np.array(sorted(set(cand)))
    def Phi(x):
        h = np.minimum(x, 1 - x)
        return np.maximum(0, np.max([c + B * h - Q * h * h for c, B, Q in P], axis=0))
    phic = Phi(cand)
    total = 0.0; chunk = 20000
    for s0 in range(0, nf, chunk):
        f = (np.arange(s0, min(nf, s0 + chunk)) + .5) / nf
        best = np.min((cand[None, :] - f[:, None]) ** 2 + phic[None, :], axis=1)
        for half in (0, 1):  # each piece's own stationary point, where it is the max
            lo, hi = (0, .5) if half == 0 else (.5, 1)
            for k0, k1, k2 in coeffs(half):
                # minimise (x - f)^2 + k0 + k1 x + k2 x^2: x = (2 f - k1) / (2 + 2 k2)
                x = np.clip((2 * f - k1) / (2 + 2 * k2), lo, hi)
                best = np.minimum(best, (x - f) ** 2 + Phi(x))
        total += np.sum(best - 1.0 / nf)   # minus (cell/2) * 2: the slope correction
    return np.sqrt(total / nf)

if __name__ == '__main__':
    certs = json.load(open(sys.argv[1])); G = int(sys.argv[2]) if len(sys.argv) > 2 else 4096
    K = taps_K(F.kuwahara(5))
    rows = []
    for c in certs:
        const, lin, quad, eps = certificate_bound(c, K, G)
        rows.append((const, lin, quad, eps))
        f = c['f']; print(f"certificate for f' = {f:.3f}: bound {const + lin * f - quad * f * f - eps * f * (1 - f):.6f} (eps {eps:.1e})")
    print(f'floor: RMS >= {floor_rms(rows):.5f}')