import numpy as np, mytaps as T
from scipy.optimize import minimize
from scipy.integrate import quad
def minH(t, M=4096):
    K = np.zeros((M, M))
    for (dx, dy), w in t.items(): K[dy % M, dx % M] += w
    A = np.abs(np.fft.fft2(K)); A[0, 0] = np.inf
    idx = np.argsort(A, axis=None)[:40]; best = (np.inf, None)
    for id_ in idx:
        j, i = np.unravel_index(id_, A.shape)
        x0 = [(i if i < M/2 else i-M)/M, (j if j < M/2 else j-M)/M]
        r = minimize(lambda p: abs(T.H(t, p[0], p[1])), x0, method='Nelder-Mead', options=dict(xatol=1e-10, fatol=1e-12))
        if r.fun < best[0]: best = (r.fun, r.x)
    return best, A.min()
def anyfloor(m):
    def g(f):
        fp = min(max((2*f - m*m)/(2 - 2*m*m), 0), 1)
        # brute check of the minimiser over f' in R (including outside [0,1] with frac variance)
        return (fp-f)**2 + m*m*fp*(1-fp)
    return np.sqrt(quad(g, 0, 1, limit=200)[0])
def anyfloor_brute(m, nf=2001):
    fps = np.linspace(-2, 3, 50001); fr = fps - np.floor(fps)
    tot = 0
    for f in (np.arange(nf)+.5)/nf: tot += np.min((fps-f)**2 + m*m*fr*(1-fr))
    return np.sqrt(tot/nf)
if __name__ == '__main__':
    t = T.kuw(5); (m, at), gm = minH(t); print('Starlock min|H|', m, at, 'grid', gm)
    print('any floor', anyfloor(m), 'brute over f in R', anyfloor_brute(m), 'm/sqrt6', m/np.sqrt(6))
    for c in (0.05, .1, .2, .4):
        for s in (1., 2.):
            (mm, at), _ = minH(T.gauss_centre(c, s), 1024); print(c, s, mm, anyfloor(mm))
