# How visible a dither's shimmer is: each piece of the error that changes over time, weighted by the eye's
# sensitivity at its spatial frequency (cycles per degree) and its temporal frequency (Hz), pooled by
# root-sum-of-squares into just-noticeable differences (JND) on a background of L levels.
#
# The sensitivity model is a deliberately generous envelope ('bright'): a peak sensitivity of 250 (contrast
# threshold 0.4%) at every spatial frequency up to 8 cpd, Mannos and Sakrison's fall-off above it, times
# Watson's 1986 temporal shape (a fit to de Lange's bright-field flicker data) held at its peak below 10 Hz
# and for still patterns. A product like this overstates fine, fast flicker, which the eye sees worse than
# either factor alone says, and bright-field numbers overstate what the eye sees in Starlock's dark scenes.
# 'dark' is the same with the temporal axis stretched by 1 / 0.638: Krajancich, Kellnhofer and Wetzstein's
# (2021) luminance scaling of flicker fusion, (0.204 log10(l / 1488 Td) + 1), at about 25 trolands (a dark
# glow on a phone). Kelly's 1979 space-time sensitivity is here too, for comparison only: as a formula for
# drifting gratings it goes to zero at the lowest spatial frequencies, where real flicker is easiest to see.
#
# Geometry: a phone at 30 cm, 390 CSS pixels across about 71 mm, so a CSS pixel is 0.035 degrees; the game
# caps the pixel ratio at 2, so a device pixel is 0.0175 degrees. The painterly pass reads one tap per CSS
# pixel; with paint off, the pattern is seen per device pixel.
import numpy as np

CSS_DEG = 0.035
DEV_DEG = CSS_DEG / 2

def kelly(cpd, hz):
    cpd = np.maximum(np.asarray(cpd, float), 1e-6); hz = np.asarray(hz, float)
    v = np.maximum(hz / cpd, 1e-9)
    k = 6.1 + 7.3 * np.abs(np.log10(1.7 * v / 3)) ** 3
    rmax = 45.9 / (1.7 * v + 2)
    s = k * 1.14 * 1.7 * v * (2 * np.pi * cpd * .67) ** 2 * np.exp(-4 * np.pi * .67 * cpd / rmax)
    return s

def _watson(hz):
    t = 4.3e-3
    return np.abs((1 + 2j * np.pi * hz * t) ** -9 - .9 * (1 + 2j * np.pi * hz * 1.33 * t) ** -10)
_WPEAK = _watson(np.linspace(9, 11, 2001)).max()

def safe(cpd, hz):
    cpd = np.asarray(cpd, float); hz = np.asarray(hz, float)
    ms = lambda r: 2.6 * (.0192 + .114 * r) * np.exp(-(.114 * r) ** 1.1)
    peak = ms(8.0)
    sp = np.where(cpd <= 8, 1.0, ms(np.maximum(cpd, 8)) / peak)
    tm = np.where(hz <= 10, 1.0, _watson(np.maximum(hz, 10)) / _WPEAK)
    return 250 * sp * tm

def dark(cpd, hz):
    return safe(cpd, np.asarray(hz, float) / .638)

MODELS = {'bright': safe, 'dark': dark, 'kelly': kelly}

def wrapf(f):
    return (np.asarray(f, float) + .5) % 1 - .5

def lattice_L(A, B, C, N, taps, deg, fps, model):
    """(grain, shimmer): the background level (in the dither's levels) at which the part of the error that
    holds still, and the part that changes, of lattice (A, B, C)/N seen through the taps (one per `deg`
    degrees) at `fps`, reaches 1 JND: sqrt(sum 2 P_m S_m^2), P_m the power of harmonic m after the filter.
    Any background brighter than this hides it. The grain includes the level error itself (1 / (12 N^2), at
    zero frequency), which across a gradient is the banding."""
    import filters as F
    m = np.arange(1, N)
    w = 1 / (2 * N * N * np.sin(np.pi * m / N) ** 2)
    fx, fy = wrapf(m * A / N), wrapf(m * B / N)
    P = w * np.abs(F.response(taps, fx, fy)) ** 2
    cpd = np.hypot(fx, fy) / deg
    hz = np.abs(wrapf(m * C / N)) * fps
    V = 2 * P * MODELS[model](cpd, hz) ** 2
    still = hz == 0
    g = V[still].sum() + MODELS[model](0.0, 0.0) ** 2 / (12 * N * N)
    return float(np.sqrt(g)), float(np.sqrt(V[~still].sum()))

def frames_L(e, deg, fps, model):
    """(grain, shimmer) from a block of error frames e[t, y, x] (already filtered in space, periodic in all
    three axes): the 3-D FFT's power at each (fx, fy, ft), weighted by S^2, with ft = 0 and ft != 0 apart.
    Average the squares over levels before taking roots."""
    T, n, _ = e.shape
    E = np.fft.fftn(e) / e.size
    ft = np.fft.fftfreq(T)[:, None, None]; fy = np.fft.fftfreq(n)[None, :, None]; fx = np.fft.fftfreq(n)[None, None, :]
    S = MODELS[model](np.hypot(fx, fy) / deg, np.abs(ft) * fps)
    V = 2 * np.abs(E) ** 2 * S ** 2
    V[0, 0, 0] /= 2   # (the zero frequency is its own conjugate)
    return float(np.sqrt(V[0].sum())), float(np.sqrt(V[1:].sum()))