# Independent re-implementation of Starlock's glow path (checker; written from the GLSL, not the run's code).
import numpy as np
F = np.float32
TWO_PI = 2 * np.pi

def frag_xy(W, H, ox=0., oy=0.):
    """gl_FragCoord at pixel centres, row 0 = bottom; plus a screen offset."""
    jj, ii = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
    return (ii.astype(F) + F(.5) + F(ox)).astype(F), (jj.astype(F) + F(.5) + F(oy)).astype(F)

def fr32(v):
    v = np.asarray(v, F)
    return (v - np.floor(v)).astype(F)

def sin_hash(x, y, model='turn'):
    x = np.asarray(x, F); y = np.asarray(y, F)
    d = (F(12.9898) * x).astype(F) + (F(78.233) * y).astype(F)
    d = d.astype(F)
    if model == 'exact':
        s = np.sin(d.astype(np.float64)).astype(F)
    elif model == 'turn':      # sin from the fp32 fraction of a turn, d * (1/2pi) rounded to fp32
        t = fr32((d * F(1 / TWO_PI)).astype(F))
        s = np.sin(TWO_PI * t.astype(np.float64)).astype(F)
    elif model == 'turn_div':  # same, but d / fp32(2pi)
        t = fr32((d / F(TWO_PI)).astype(F))
        s = np.sin(TWO_PI * t.astype(np.float64)).astype(F)
    return fr32((s * F(43758.5453)).astype(F)).astype(np.float64)

def ign(x, y):
    x = np.asarray(x, F); y = np.asarray(y, F)
    t = fr32(((x * F(.06711056)).astype(F) + (y * F(.00583715)).astype(F)).astype(F))
    return fr32((F(52.9829189) * t).astype(F)).astype(np.float64)

def sstep(e0, e1, x):
    t = np.clip((x - e0) / (e1 - e0), 0, 1); return t * t * (3 - 2 * t)

def alpha(kind, u, v, extra):
    if kind == 'glow':
        r2 = ((u - .5) ** 2 + (v - .5) ** 2) * 4.
        return np.maximum(0., (np.exp(-4 * r2) - .0183) / .9817) + extra * np.exp(-38 * r2)
    w, fill = extra
    r = 2 * np.sqrt((u - .5) ** 2 + (v - .5) ** 2)
    a = np.exp(-((r - .84) / w) ** 2) + fill * sstep(.84, 0., r) * .3
    return a * (1 - sstep(.96, 1., r))

def layers(scene, W, H):
    jj, ii = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
    px, py = ii + .5, jj + .5
    out = []
    for kind, cx, cy, w, h, col, op, extra in scene:
        u = (px - cx + w / 2) / w; v = (py - cy + h / 2) / h
        ins = (u >= 0) & (u <= 1) & (v >= 0) & (v <= 1)
        a = alpha(kind, u, v, extra)
        out.append((np.array(col, float)[None, None] * (a * op)[..., None], a * op, ins))
    return out

def q8(v):
    return np.clip(np.floor(v * 255. + .5), 0, 255)

def target(scene, W, H, origin, dither):
    """dither: function(k, x, y, aop) -> offset in [-.5,.5)/255 per pixel, or None."""
    x, y = frag_xy(W, H, *origin)
    t = np.zeros((H, W, 3))
    for k, (rgb, aop, ins) in enumerate(layers(scene, W, H)):
        src = rgb if dither is None else rgb + dither(k, x, y, aop)[..., None]
        nt = q8(t / 255. + np.clip(src, 0, 1))
        t = np.where(ins[..., None], nt, t)
    return t

def float_scene(scene, W, H):
    t = np.zeros((H, W, 3))
    for rgb, aop, ins in layers(scene, W, H):
        t += np.where(ins[..., None], rgb, 0)
    return t

def kuw(img, R, st, linear=False):
    if R == 0: return img.copy()
    H, W, _ = img.shape
    P = R * st
    pd = np.pad(img, ((P, P), (P, P), (0, 0)), mode='edge')
    rr = R * R + .5
    S1 = np.zeros((8, H, W, 3)); S2 = np.zeros((8, H, W, 3)); Wt = np.zeros(8)
    for j in range(-R, R + 1):
        for i in range(-R, R + 1):
            d = i * i + j * j
            if d > rr: continue
            c = pd[P + j * st:P + j * st + H, P + i * st:P + i * st + W]
            g = np.exp(-d / rr * 2)
            if d == 0:
                ks = range(8)
            else:
                ang = np.float32(np.arctan2(np.float32(j), np.float32(i)))
                ks = [int(np.floor((float(ang) + 3.14159265) / .78539816)) % 8]
            for k in ks:
                S1[k] += g * c; S2[k] += g * c * c; Wt[k] += g
    num = 0.; den = 0.
    for k in range(8):
        mu = S1[k] / Wt[k]
        var = np.abs(S2[k] / Wt[k] - mu * mu).sum(-1, keepdims=True)
        wk = np.ones_like(var) if linear else 1 / (1 + (1000 * var) ** 4)
        num = num + wk * mu; den = den + wk
    return num / den

A_IN = np.array([[.59719, .35458, .04823], [.07600, .90834, .01566], [.02840, .13383, .83777]])
A_OUT = np.array([[1.60475, -.53108, -.07367], [-.10208, 1.10813, -.00605], [-.00327, -.07276, 1.07602]])
def tonemap(c, exp_=1.15):
    c = np.einsum('ij,...j->...i', A_IN, c * exp_ / .6)
    c = (c * (c + .0245786) - .000090537) / (c * (.983729 * c + .4329510) + .238081)
    c = np.einsum('ij,...j->...i', A_OUT, c)
    c = np.clip(c, 0, 1)
    return np.where(c <= .0031308, 12.92 * c, 1.055 * np.power(c, .41666) - .055)