# Caustics: a uniform field of rays is thrown through one hidden refracting surface.
# Nothing is drawn where it appears; the picture is the histogram of where rays land.
import numpy as np
from PIL import Image
import os

SEED = 7
N = 1200
RAYS = 24_000_000
HERE = os.path.dirname(os.path.abspath(__file__))

# Flexoki swatches, dark to light ramp
RAMP = ["100F0F", "1C1B1A", "282726", "343331", "163B66", "205EA6", "24837B",
        "3AA99F", "879A39", "D0A215", "DA702C", "F2F0E5", "FFFCF0"]
PAL = np.array([[int(c[i:i+2], 16) for i in (0, 2, 4)] for c in RAMP], dtype=np.uint8)

rng = np.random.default_rng(SEED)

def surface_grad(x, y):
    # hidden height field: one large off-centre lens plus a slow ripple; returns gradient
    gx = np.zeros_like(x); gy = np.zeros_like(y)
    # main lens, elongated and tilted
    cx, cy, sx, sy, a = 0.40, 0.58, 0.30, 0.20, 1.0
    c, s = np.cos(0.6), np.sin(0.6)
    u = (x - cx) * c + (y - cy) * s
    v = -(x - cx) * s + (y - cy) * c
    g = a * np.exp(-(u**2 / (2*sx**2) + v**2 / (2*sy**2)))
    dgu = -g * u / sx**2; dgv = -g * v / sy**2
    gx += dgu * c - dgv * s; gy += dgu * s + dgv * c
    # secondary lobe that folds the net
    cx, cy, r = 0.78, 0.30, 0.16
    g = 0.45 * np.exp(-((x-cx)**2 + (y-cy)**2) / (2*r**2))
    gx += -g * (x-cx) / r**2; gy += -g * (y-cy) / r**2
    # slow ripple
    k = 2*np.pi*2.3
    gx += 0.05 * k * np.cos(k*x + 1.3) * np.sin(0.7*k*y)
    gy += 0.05 * 0.7*k * np.sin(k*x + 1.3) * np.cos(0.7*k*y)
    return gx, gy

H = np.zeros((N, N), dtype=np.float64)
D = 0.62  # throw distance
chunk = 2_000_000
for _ in range(RAYS // chunk):
    x = rng.random(chunk); y = rng.random(chunk)
    gx, gy = surface_grad(x, y)
    lx = x + D * gx * 0.9
    ly = y + D * gy * 0.9
    h, _, _ = np.histogram2d(ly, lx, bins=N, range=[[0, 1], [0, 1]])
    H += h

# light blur (separable) so density reads as light, not noise
k = np.array([1, 4, 6, 4, 1], dtype=float); k /= k.sum()
for ax in (0, 1):
    H = np.apply_along_axis(lambda m: np.convolve(m, k, mode="same"), ax, H)

m = H / np.percentile(H, 99.9)
t = np.clip(m, 0, 1) ** 0.55
idx = np.clip((t * (len(RAMP) - 1) + 0.5).astype(int), 0, len(RAMP) - 1)
img = PAL[idx[::-1]]
Image.fromarray(img, "RGB").save(os.path.join(HERE, "render-1.png"))