#!/usr/bin/env python3
"""Bake a soft volumetric flashlight cone as RGBA (bottom = lens, top = far end)."""
import numpy as np
from PIL import Image, ImageFilter

W, H = 1100, 1700
yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)
t = 1.0 - yy / (H - 1)                 # 0 at the lens (bottom), 1 at the far end (top)
x = xx - (W - 1) / 2.0

half = 26.0 + (250.0 - 26.0) * np.power(t, 0.92)
u = np.abs(x) / half

core = np.exp(-np.square(u / 0.34) * 2.4)
glow = np.exp(-np.square(u / 1.05) * 1.55) * 0.55
spill = np.exp(-np.square(u / 1.7) * 1.2) * 0.12

# brightest just past the lens, gentle decay with distance, tiny hot spot at the far end
length = 0.62 + 0.38 * np.exp(-t * 1.35)
length *= 1.0 + 0.10 * np.exp(-np.square((t - 0.02) / 0.05))
i = (core * 0.95 + glow + spill) * length
i = np.clip(i, 0, 1.35)

core_col = np.array([255, 238, 206], np.float32)
glow_col = np.array([255, 166, 74], np.float32)
w_core = np.clip(core * 0.95 * length, 0, 1)[..., None]
w_glow = np.clip((glow + spill) * length, 0, 1)[..., None]
rgb = core_col * w_core + glow_col * w_glow
den = np.maximum(w_core + w_glow, 1e-5)
rgb = np.clip(rgb / den, 0, 255)

a = np.clip(i, 0, 1) * 255
img = Image.fromarray(np.dstack([rgb, a]).astype(np.uint8), "RGBA")
img = img.filter(ImageFilter.GaussianBlur(5))
img.save("/root/notchlight/assets/beam.png")
print("beam", img.size)