# A port of Starlock's cutscene glow path to numpy, step for step:
#   glow shader (+ dither) -> 8-bit linear target (additive) -> painterly pass (generalised
#   Kuwahara, radius 0/2/5) -> ACES filmic (exposure 1.15) -> sRGB -> 8-bit canvas.
# The shader and pass are copied from public/ideas/starlock/cutscene3d.js and reveal3d.js;
# three.js r186's tone-mapping and colour-space chunks are copied from its source.
import numpy as np

f32 = np.float32

def coords(W, H):
    """gl_FragCoord.xy at each pixel centre (row 0 is the bottom row, as GL reads)."""
    y, x = np.mgrid[0:H, 0:W].astype(f32) + f32(0.5)
    return x, y

# --- the glows (cutscene3d.js) -------------------------------------------------------
def glow_alpha(u, v, core=0.0):
    r2 = ((u - .5) ** 2 + (v - .5) ** 2) * 4
    return np.maximum(0., (np.exp(-r2 * 4.) - .0183) / .9817) + core * np.exp(-r2 * 38.)

def ring_alpha(u, v, width=.08, fill=0.0):
    r = np.hypot(u - .5, v - .5) * 2
    s = np.clip((r - .84) / (0. - .84), 0, 1); s = s * s * (3 - 2 * s)  # smoothstep(.84, 0., r)
    e = np.clip((r - .96) / (1. - .96), 0, 1); e = e * e * (3 - 2 * e)
    return (np.exp(-((r - .84) / width) ** 2) + fill * s * .3) * (1 - e)

def quad_uv(W, H, x0, y0, w, h):
    """vUv for a quad centred at (x0, y0) of size w x h pixels; None outside it."""
    x, y = coords(W, H)
    u = (x - (x0 - w / 2)) / w
    v = (y - (y0 - h / 2)) / h
    inside = (u >= 0) & (u <= 1) & (v >= 0) & (v <= 1)
    return u.astype(np.float64), v.astype(np.float64), inside

# --- the 8-bit target ------------------------------------------------------------------
def to_unorm8(x):
    return np.clip(np.floor(np.asarray(x, np.float64) * 255. + .5), 0, 255)

def additive(layers, W, H, base=None):
    """layers: list of (rgb float (H,W,3) shader output incl. dither, inside mask). Each draw adds
    to the stored 8-bit value and is stored again (GL blends at the target's precision; the
    source is clamped to [0, 1] before the add, as for any fixed-point target)."""
    t = np.zeros((H, W, 3)) if base is None else base.copy()
    for rgb, inside in layers:
        s = to_unorm8(t / 255. + np.clip(rgb, 0, 1))
        t = np.where(inside[..., None], s, t)
    return t  # codes 0..255

# --- the painterly pass (reveal3d.js kuwaharaShader) --------------------------------------
def kuwahara(img, R, stride=1):
    """img: float (H,W,3) in 0..1 (the target read back). Clamp-to-edge sampling. The pass steps
    one CSS pixel per tap (uTexel = pixelRatio / width), so on a 2x screen it reads every 2nd texel."""
    if R == 0:
        return img.copy()
    H, W, _ = img.shape
    P = R * stride
    pad = np.pad(img, ((P, P), (P, P), (0, 0)), mode='edge')
    r2 = float(R * R) + .5
    m = np.zeros((8, H, W, 3)); s = np.zeros((8, H, W, 3)); w = np.zeros(8)
    for j in range(-R, R + 1):
        for i in range(-R, R + 1):
            d = float(i * i + j * j)
            if d > r2: continue
            c = pad[P + j * stride:P + j * stride + H, P + i * stride:P + i * stride + W]
            g = np.exp(-d / r2 * 2.)
            if d < .5:
                m += c * g; s += c * c * g; w += g; continue
            # (a tap on a sector's edge, along an axis or a diagonal, lands where the GPU's float rounding
            # puts it: in the higher sector on SwiftShader. In float64 it would fall just short, hence the
            # nudge. Another GPU may round the other way; this is a property of the game's filter.)
            k = int(np.floor((np.arctan2(j, i) + 3.14159265) / .78539816 + 1e-5) % 8)
            m[k] += c * g; s[k] += c * c * g; w[k] += g
    num = np.zeros((H, W, 3)); den = np.zeros((H, W, 1))
    for k in range(8):
        mk = m[k] / max(w[k], 1e-4)
        vk = np.abs(s[k] / max(w[k], 1e-4) - mk * mk)
        a = 1. / (1. + (1000. * vk.sum(-1, keepdims=True)) ** 4)
        num += mk * a; den += a
    return num / den

# --- three.js r186: ACESFilmicToneMapping, sRGBTransferOETF ---------------------------------
ACES_IN = np.array([[0.59719, 0.35458, 0.04823], [0.07600, 0.90834, 0.01566], [0.02840, 0.13383, 0.83777]])
ACES_OUT = np.array([[1.60475, -0.53108, -0.07367], [-0.10208, 1.10813, -0.00605], [-0.00327, -0.07276, 1.07602]])

def aces(c, exposure=1.15):
    c = c * (exposure / 0.6)
    c = c @ ACES_IN.T
    a = c * (c + 0.0245786) - 0.000090537
    b = c * (0.983729 * c + 0.4329510) + 0.238081
    c = (a / b) @ ACES_OUT.T
    return np.clip(c, 0, 1)

def srgb(c):
    return np.where(c <= 0.0031308, c * 12.92, np.power(np.maximum(c, 0), 0.41666) * 1.055 - 0.055)

def post(target_codes, R, exposure=1.15, stride=1):
    """The painterly pass from the 8-bit target codes: float display values in 0..1 (before the
    canvas's own 8-bit write)."""
    return srgb(aces(kuwahara(target_codes / 255., R, stride), exposure))

def ideal(rgb_float, R, exposure=1.15, stride=1):
    """The same path with no quantisation anywhere and no dither: what the glow should look like."""
    return srgb(aces(kuwahara(rgb_float, R, stride), exposure))