# Checker for claims 2 and 3: own pipeline (vpipe.py), own banding measures.
#  A: angle-averaged radial luminance profile -> L* error (rings), for radially symmetric scenes
#  B: Mannos-Sakrison CSF-weighted luminance error, limited to < 4 cpd, linearised to L* (all scenes)
import sys, json, itertools
import numpy as np
from multiprocessing import Pool
import os; sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import experiment as E          # only for the scene list, frame size and origin (inputs)
import vpipe as V

FR = E.FRAME; ORG = E.ORIGIN
def D_old(model='turn', gate=True, seeded=False):
    def f(k, x, y, aop):
        if seeded: x, y = x + 37 * k, y + 91 * k
        d = (V.sin_hash(x, y, model) - .5) / 255.
        return d * (aop >= .002) if gate else d
    return f
def D_ign(gate=False, seeded=True):
    def f(k, x, y, aop):
        if seeded: x, y = (x + np.float32(37 * k)).astype(np.float32), (y + np.float32(91 * k)).astype(np.float32)
        d = (V.ign(x, y) - .5) / 255.
        return d * (aop >= .002) if gate else d
    return f
CANDS = {
    'none': (None, False),
    'old (M2 sin, gate, shared)': (D_old(), False),
    'old, exact sin': (D_old('exact'), False),
    'sin, no gate': (D_old(gate=False), False),
    'sin, gate, seeded': (D_old(seeded=True), False),
    'sin, no gate, seeded': (D_old(gate=False, seeded=True), False),
    'ign, gate, seeded': (D_ign(gate=True), False),
    'ign, no gate, shared': (D_ign(seeded=False), False),
    'ign, no gate, seeded': (D_ign(), False),
    'SHIPPED (ign seeded + out)': (D_ign(), True),
    'old + out': (D_old(), True),
}

def run(sname, cname, R, dpr):
    sc = E.scenes()[sname]
    dith, out = CANDS[cname]
    t = V.target(sc, FR, FR, ORG, dith)
    disp = V.tonemap(V.kuw(t / 255., R, dpr))
    if out:
        x, y = V.frag_xy(FR, FR, *ORG)
        disp = disp + ((V.ign(y, x) - .5) / 255.)[..., None]
    return V.q8(disp)

def ideals(sname, R, dpr):
    f = V.float_scene(E.scenes()[sname], FR, FR)
    return {'I1 same Kuwahara': V.tonemap(V.kuw(f, R, dpr)),
            'I2 no filter': V.tonemap(f),
            'I3 linear disc blur': V.tonemap(V.kuw(f, R, dpr, linear=True))}

def lin(c):
    c = np.clip(c, 0, 1); return np.where(c <= .04045, c / 12.92, ((c + .055) / 1.055) ** 2.4)
def Y(disp): return lin(disp) @ np.array([.2126, .7152, .0722])
def Ls(y): return np.where(y > 216 / 24389, 116 * np.cbrt(y) - 16, y * 24389 / 27)
def dLs(y):  # derivative of L* wrt Y
    return np.where(y > 216 / 24389, 116 / 3 * np.cbrt(np.maximum(y, 1e-9)) ** -2, 24389 / 27)

def radial(Yo, Yi, rmax=300):
    jj, ii = np.mgrid[0:FR, 0:FR]
    r = np.hypot(ii + .5 - FR / 2, jj + .5 - FR / 2)
    b = np.floor(r).astype(int); m = b < rmax
    n = np.bincount(b[m], minlength=rmax)
    mo = np.bincount(b[m], Yo[m], rmax) / np.maximum(n, 1); mi = np.bincount(b[m], Yi[m], rmax) / np.maximum(n, 1)
    e = Ls(mo) - Ls(mi)
    return float(np.sqrt((n * e * e).sum() / n.sum())), float(np.abs(e[5:]).max())

def csf_err(Yo, Yi, ppd, cut=4.):
    fy = np.fft.fftfreq(FR)[:, None]; fx = np.fft.fftfreq(FR)[None, :]
    f = np.hypot(fx, fy) * ppd   # cycles per degree
    A = 2.6 * (.0192 + .114 * f) * np.exp(-(.114 * f) ** 1.1); A /= A.max()
    A = A * (f < cut)
    dY = np.real(np.fft.ifft2(np.fft.fft2(Yo - Yi) * A))
    Yb = np.real(np.fft.ifft2(np.fft.fft2(Yi) * (np.hypot(fx, fy) * ppd < cut)))
    e = dLs(np.maximum(Yb, 0)) * dY
    B = 32
    return float(np.sqrt(np.mean(e[B:-B, B:-B] ** 2)))

def job(a):
    sname, R, dpr = a
    I = ideals(sname, R, dpr)
    res = {}
    for c in CANDS:
        o = run(sname, c, R, dpr); Yo = Y(o / 255.)
        for iname, idisp in I.items():
            Yi = Y(idisp)
            res[(c, iname)] = dict(csf=csf_err(Yo, Yi, E.PPD[dpr]), rad=radial(Yo, Yi)[0] if sname != 'stack' else None,
                                   black=float(o[(idisp.max(-1) == 0)].max()) if (idisp.max(-1) == 0).any() else None)
    return sname, R, dpr, res

if __name__ == '__main__':
    jobs = [(s, 5, 2) for s in E.scenes()] + [(s, R, 2) for s in ('stack', 'glow') for R in (0, 2)]
    with Pool(4) as p:
        out = p.map(job, jobs)
    dump = []
    for sname, R, dpr, res in out:
        print(f'\n== {sname}  R={R} dpr={dpr}')
        for iname in ('I1 same Kuwahara', 'I2 no filter', 'I3 linear disc blur'):
            print(f'  ideal {iname}')
            for c in CANDS:
                m = res[(c, iname)]
                print(f"    {c:30s} csf<4cpd {m['csf']:.4f}" + (f"  radial {m['rad']:.4f}" if m['rad'] is not None else ''))
                dump.append(dict(scene=sname, R=R, dpr=dpr, cand=c, ideal=iname, **m))
            o, n = res[('old (M2 sin, gate, shared)', iname)], res[('SHIPPED (ign seeded + out)', iname)]
            print(f"    RATIO old/shipped: csf {o['csf'] / n['csf']:.2f}" + (f"  radial {o['rad'] / n['rad']:.2f}" if o['rad'] else ''))
    json.dump(dump, open('banding.json', 'w'), indent=0)