# C7 audit of pipe_t.seen()'s normalisation, on synthetic inputs whose JND is known in closed form
# (convention: a cosine of contrast amplitude c at (cpd, Hz) is c * S(cpd, Hz) JND; a constant offset c0 is c0 * S(0,0)).
# run: python3 c7_window.py
import os, sys, numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import vis
PPD = 64.
def seen_like(g, c, fps, model):   # pipe_t.seen's arithmetic, verbatim, on given grain g[y,x] and change c[t,y,x]
    n = c.shape[1]; Tn = c.shape[0]
    ws = np.hanning(n + 2)[1:-1]; w2 = ws[:, None]*ws[None, :]; wt = np.hanning(Tn + 2)[1:-1]
    fy = np.fft.fftfreq(n)[:, None]; fx = np.fft.fftfreq(n)[None, :]; ft = np.fft.fftfreq(Tn)[:, None, None]
    cpd = np.hypot(fx, fy)*PPD
    G = 2*np.abs(np.fft.fft2(g*w2))**2/(w2**2).sum()/(n*n); G[0, 0] /= 2
    C = 2*np.abs(np.fft.fftn(c*w2*wt[:, None, None]))**2/((w2**2).sum()*(wt**2).sum())/(n*n*Tn)
    S = vis.MODELS[model]
    return float(np.sqrt((G*S(cpd, 0.0)**2).sum())), float(np.sqrt((C*S(cpd[None], np.abs(ft)*fps)**2).sum()))
n = 288; Tn = 6
y, x = np.mgrid[0:n, 0:n]
S = vis.MODELS
# 1) a constant offset c0 in the grain
c0 = 0.01
gr, _ = seen_like(np.full((n, n), c0), np.zeros((Tn, n, n)), 60, 'bright')
print(f'constant offset {c0}: seen() grain {gr:.4f} JND, should be {c0*S["bright"](0., 0.):.4f}  (ratio {gr/(c0*S["bright"](0.,0.)):.3f})')
# 2) a spatial cosine in the grain
for cyc in (3, 20.5, 60):
    g = c0*np.cos(2*np.pi*cyc*x/n)
    gr, _ = seen_like(g, np.zeros((Tn, n, n)), 60, 'bright')
    print(f'grain cosine {cyc} cycles ({cyc/n*PPD:.2f} cpd): {gr:.4f}, should be {c0*S["bright"](cyc/n*PPD, 0.):.4f}')
# 3) a period-2 flip c(x) (-1)^t, as the half-step dither makes: exact JND at fps/2 vs seen() with 6 frames
for model in ('bright', 'dark'):
    for fps in (60, 30):
        for cyc in (0, 3, 40.5, 100):
            sp = np.cos(2*np.pi*cyc*x/n) if cyc else np.ones((n, n))
            amp = c0
            c = amp*sp[None]*((-1.0)**np.arange(Tn))[:, None, None]
            _, sh = seen_like(np.zeros((n, n)), c, fps, model)
            exact = amp*S[model](cyc/n*PPD, fps/2)   # amplitude * S at fps/2
            print(f'{model:6s} {fps} fps flip, spatial {cyc:5} cyc ({cyc/n*PPD:5.2f} cpd): seen() shimmer {sh:.4f}, exact {exact:.4f} (ratio {sh/exact:.3f})')
wt = np.hanning(8)[1:-1]; X = np.fft.fft(wt*(-1.0)**np.arange(6)); P = np.abs(X)**2; print('power of a windowed 6-frame flip by temporal bin (0,10,20,30,-20,-10 Hz at 60 fps):', np.round(P/P.sum(), 4))
print('S bright at 0/10/20/30 Hz:', [round(float(S['bright'](1.0, h)), 1) for h in (0, 10, 20, 30)], ' dark:', [round(float(S['dark'](1.0, h)), 1) for h in (0, 10, 20, 30)])