# Draws the art for demos/neonCoast:  every picture is made from shapes and arithmetic here, so it is
# ours and can be drawn again.
#
# The sky is the important one.  It is an equirectangular picture, the whole sphere around the
# camera, and the 3D scene both shows it and lights with it:  a chrome object reflects whatever is in
# it.  So besides the sunset and the stars it carries the striped sun straight ahead (the sky's centre
# faces the camera's default view) and, in the lower half, the neon grid in perspective, as if a
# floor ran to the horizon in every direction.  Chrome then reflects the sun and the grid, which is
# the whole look.
#
# The car is rendered in 3D by util/neonCoastCar.sh.
#
# Usage:  python3 util/neonCoastArt.py
import math
import os

import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageFont


REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT  = os.path.join(REPO, "demos", "neonCoast", "art")

SKY_W      = 2048
SKY_H      = 1024
SUN_LAT    = 9.0                  # Degrees above the horizon of the sun's centre
SUN_RADIUS = 13.0                 # Degrees
GRID_CELL  = 4.0                  # The painted floor grid's cell, in camera heights
DENSITY    = 2.25                 # The overlay sprites' pixels to a design unit (demos/kit.lua, K.sprite):  1920 by 1080
FONT_BOLD  = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
FONT_WIDE  = "/usr/share/fonts/truetype/lato/Lato-BlackItalic.ttf"

# The palette:  night at the top, through violet and magenta, to orange at the horizon.
SKY_STOPS = [
	(0.00, (6, 3, 22)),
	(0.45, (26, 8, 60)),
	(0.72, (88, 18, 110)),
	(0.88, (200, 40, 120)),
	(1.00, (255, 120, 70)),
]
MAGENTA = np.array([255, 40, 200])
CYAN    = np.array([40, 230, 255])


def save(image, name):
	image.save(os.path.join(OUT, name))
	print("%-16s %d x %d" % (name, image.width, image.height))


def ramp(stops, t):
	t = np.clip(t, 0, 1)
	out = np.zeros(t.shape + (3,))
	for (t0, c0), (t1, c1) in zip(stops[:-1], stops[1:]):
		m = (t >= t0) & (t <= t1)
		k = ((t - t0) / (t1 - t0))[m][:, None]
		out[m] = np.array(c0) * (1 - k) + np.array(c1) * k
	return out


# The sun's own look, wherever it is drawn:  yellow at the top to magenta at the bottom, with
# horizontal bands cut out of its lower half that thicken toward the bottom.  u, v run -1 to 1.
def sunMask(u, v):
	inside = (u * u + v * v) <= 1.0
	band   = np.ones_like(u, dtype=bool)
	for k in range(7):
		centre = 0.12 + k * 0.14
		half   = 0.012 + k * 0.009
		band  &= ~((v > centre - half) & (v < centre + half))
	return inside & band


def sunColour(v):
	return ramp([(0.0, (255, 240, 120)), (0.5, (255, 140, 60)), (1.0, (255, 40, 150))], (v + 1) / 2)


def sky():
	lon = (np.arange(SKY_W) + 0.5) / SKY_W * 2 * math.pi - math.pi
	lat = math.pi / 2 - (np.arange(SKY_H) + 0.5) / SKY_H * math.pi
	LON, LAT = np.meshgrid(lon, lat)
	img = np.zeros((SKY_H, SKY_W, 3))

	# Above the horizon:  the gradient by height, and stars thinning toward the horizon.
	up = LAT >= 0
	img[up] = ramp(SKY_STOPS, 1 - LAT[up] / (math.pi / 2))
	rng = np.random.default_rng(7)
	for _ in range(2600):
		x = rng.integers(0, SKY_W)
		y = int(rng.random() ** 1.8 * SKY_H * 0.42)
		b = rng.random() ** 2 * 255
		img[y, x] = np.maximum(img[y, x], [b, b, min(255, b * 1.1)])

	# Behind the camera, a city's cyan glow along the horizon:  the half of the sky a chrome face
	# toward the camera reflects, so the chrome shades cyan above and magenta below.
	city = np.exp(-np.maximum(LAT, 0) / math.radians(14)) * np.clip(-np.cos(LON), 0, 1) ** 0.7 * up
	img = img * (1 - 0.85 * city[..., None]) + city[..., None] * np.array([60, 210, 255]) * 0.95

	# The sun straight ahead, its glow first.
	du = (LON / math.radians(SUN_RADIUS))
	dv = -(LAT - math.radians(SUN_LAT)) / math.radians(SUN_RADIUS)
	dist = np.sqrt(du * du + dv * dv)
	glow = np.exp(-np.maximum(dist - 1, 0) * 2.2) * (LAT > -0.01)
	img += glow[..., None] * np.array([255, 70, 120]) * 0.55
	mask = sunMask(du, dv)
	img[mask] = sunColour(dv[mask])

	# Below the horizon:  a dark floor and the grid on it in perspective, fading into haze.
	down = LAT < 0
	dirX = np.cos(LAT) * np.sin(LON)
	dirZ = -np.cos(LAT) * np.cos(LON)
	t    = 1.0 / np.maximum(-np.sin(LAT), 1e-4)
	# The camera stands in the middle of a cell, so no line runs straight under it.
	gx   = t * dirX / GRID_CELL + 0.5
	gz   = t * dirZ / GRID_CELL + 0.5
	fx   = np.abs(gx - np.round(gx))
	fz   = np.abs(gz - np.round(gz))
	# Lines a constant width on the floor get thinner with distance;  their brightness is kept by
	# widening the soft edge as they shrink, then faded into the horizon's haze.
	width = 0.03 + 0.002 * t
	line  = np.clip(1 - np.minimum(fx, fz) / width, 0, 1) ** 1.5
	haze  = np.clip((t - 6) / 60, 0, 1)
	floor = np.array([14, 4, 30]) * (1 - haze[..., None]) + np.array([120, 30, 110]) * haze[..., None]
	lineColour = MAGENTA * (1 - haze[..., None] * 0.7)
	below = floor + line[..., None] * lineColour
	img[down] = below[down]
	img = np.clip(img, 0, 255).astype(np.uint8)
	return Image.fromarray(img, "RGB")


# The floor tile:  a cell of the neon grid, glowing magenta lines on black, with alpha, used as the
# floor's colour and its emissive map at once.  Tiled, the line at each edge meets its neighbour's.
def gridTile(colour, size=256, width=5.0, glow=18.0):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	d = np.minimum(np.minimum(x, size - x), np.minimum(y, size - y))
	core = np.clip(1 - (d - width / 2), 0, 1)
	halo = np.exp(-d / glow) * 0.5
	v = np.clip(core + halo, 0, 1)
	img = np.zeros((size, size, 4))
	img[..., :3] = np.array(colour) * v[..., None]
	img[..., 3] = 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# A heightmap for the mountains:  a ridge line with a valley in the middle where the sun sits,
# rising toward both ends, with ridged noise for the peaks.  Red carries the height.
def mountains(w=256, h=64):
	rng = np.random.default_rng(3)
	x = np.linspace(-1, 1, w)
	z = np.linspace(0, 1, h)
	X, Z = np.meshgrid(x, z)
	base = np.clip(np.abs(X) * 1.4 - 0.15, 0, 1)
	noise = np.zeros_like(X)
	amp = 1.0
	freq = 3.0
	for _ in range(5):
		phase = rng.random(4) * 10
		n = np.sin(X * freq * 3.1 + phase[0]) * np.sin(Z * freq * 2.3 + phase[1]) + 0.5 * np.sin((X + Z) * freq * 4.7 + phase[2])
		noise += (1 - np.abs(n)) * amp
		amp *= 0.5
		freq *= 2.0
	ridge = np.exp(-((Z - 0.5) ** 2) / 0.05)
	hgt = np.clip(base * (0.4 + 0.6 * noise / 1.9) * ridge, 0, 1)
	img = (hgt * 255).astype(np.uint8)
	return Image.fromarray(np.stack([img, img, img], -1), "RGB")


SS = 4                            # Sprites are drawn this many times larger and shrunk, for smooth edges


def big(w, h):
	return Image.new("RGBA", (w * SS, h * SS), (0, 0, 0, 0))


def shrink(image, w, h):
	return image.resize((w, h), Image.LANCZOS)


def pts(points):
	return [(x * SS, y * SS) for x, y in points]


def palm():
	w, h = 200, 230
	im = big(w, h)
	d = ImageDraw.Draw(im)
	trunk = []
	for i in range(41):
		t = i / 40
		x = 100 + 18 * math.sin(t * 2.2) - 10 * t
		trunk.append((x, 40 + t * 180))
	for i, ((x0, y0), (x1, y1)) in enumerate(zip(trunk[:-1], trunk[1:])):
		wdt = 7 + (y0 - 40) / 180 * 5
		# A brown trunk, ringed darker every few steps the way a palm's bark is.
		shade = 0.75 if (i % 4 == 0) else 1.0
		d.line(pts([(x0, y0), (x1, y1)]), fill=(int(120 * shade), int(78 * shade), int(44 * shade), 255), width=int(wdt * SS))
	top = trunk[0]
	green = (34, 120, 52, 255)         # The fronds:  green, lit from behind by the sunset
	for k in range(9):
		ang = -math.pi * 0.98 + k * math.pi * 0.98 / 8 + (0.12 if k % 2 else -0.08)
		length = 78 + 12 * math.sin(k * 1.7)
		spine = []
		for n in range(17):
			t = n / 16
			spine.append((top[0] + math.cos(ang) * length * t, top[1] + math.sin(ang) * length * t * 0.55 + t * t * 52))
		# The frond as one feathered leaf:  its outline runs out along one side and back along the
		# other, the edge in and out at every step so the leaflets show, widest a third of the way out.
		left = []
		right = []
		for n in range(len(spine) - 1):
			(x0, y0), (x1, y1) = spine[n], spine[n + 1]
			dx = x1 - x0
			dy = y1 - y0
			m = math.hypot(dx, dy) + 1e-9
			px, py = -dy / m, dx / m
			t = n / (len(spine) - 1)
			width = 15 * math.sin(math.pi * min(1, t * 1.25)) * (1.0 if n % 2 == 0 else 0.45) + 1
			droop = width * 0.5
			left.append((x0 + px * width, y0 + py * width + droop))
			right.append((x0 - px * width, y0 - py * width + droop))
		d.polygon(pts(left + [spine[-1]] + right[::-1]), fill=green)
		d.line(pts(spine), fill=(24, 90, 38, 255), width=int(1.5 * SS))
	d.ellipse(pts([(top[0] - 7, top[1] - 5), (top[0] + 7, top[1] + 7)]), fill=(90, 60, 34, 255))
	# A magenta rim on the side away from the viewer's left:  the sunset behind it.
	a = np.array(im)
	alpha = a[..., 3] > 0
	edge = alpha & ~np.roll(alpha, -SS, axis=1)
	a[edge] = [255, 120, 170, 255]
	return shrink(Image.fromarray(a), round(w * DENSITY), round(h * DENSITY))


def lamp():
	w, h = 64, 200
	im = big(w, h)
	d = ImageDraw.Draw(im)
	d.rectangle(pts([(30, 20), (35, 199)]), fill=(20, 14, 30, 255))
	d.line(pts([(33, 22), (8, 14)]), fill=(20, 14, 30, 255), width=4 * SS)
	d.ellipse(pts([(0, 8), (18, 20)]), fill=(255, 245, 210, 255))
	return shrink(im, round(w * DENSITY), round(h * DENSITY))


# The night road a row at a time, as in the arcade racers:  four kinds of row (light and dark road,
# light and dark ground) at HAZE levels each, fading into the violet haze at the horizon.
STRIP_W = 256
HAZE = 8
HAZE_RGB = np.array([110, 28, 100])


def roadStrip():
	img = np.zeros((4 * HAZE, STRIP_W, 4), dtype=np.float64)
	x = np.arange(STRIP_W)
	for kind in range(4):
		for level in range(HAZE):
			if kind < 2:
				row = np.tile(np.array([16, 10, 26], dtype=np.float64), (STRIP_W, 1))
				edge = (x < 10) | (x >= STRIP_W - 10)
				row[edge] = [255, 40, 200]
				inner = ((x >= 10) & (x < 14)) | ((x >= STRIP_W - 14) & (x < STRIP_W - 10))
				row[inner] = [120, 20, 100]
				if kind == 0:
					for lane in (1, 2):
						c = 10 + (STRIP_W - 20) * lane // 3
						row[(x >= c - 2) & (x < c + 2)] = [40, 230, 255]
			else:
				row = np.tile(np.array([10, 4, 22] if kind == 2 else [20, 6, 36], dtype=np.float64), (STRIP_W, 1))
			k = level / (HAZE - 1) * 0.8
			row = row * (1 - k) + HAZE_RGB * k
			img[kind * HAZE + level, :, :3] = row
			img[kind * HAZE + level, :, 3] = 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# A shaded ball for vector bobs:  white and grey, so spriteSetTint colours it.
def bob(size=round(64 * DENSITY)):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	cx = cy = size / 2
	r = size / 2 - 1
	dx = (x - cx) / r
	dy = (y - cy) / r
	d2 = dx * dx + dy * dy
	inside = d2 <= 1
	nz = np.sqrt(np.clip(1 - d2, 0, 1))
	light = np.clip(dx * -0.4 + dy * -0.5 + nz * 0.77, 0, 1)
	spec = np.clip(light, 0, 1) ** 24
	v = 0.18 + 0.7 * light + 0.6 * spec
	a = np.clip((1 - np.sqrt(d2)) * r, 0, 1) * inside
	img = np.zeros((size, size, 4))
	img[..., :3] = np.clip(v, 0, 1)[..., None] * 255
	img[..., 3] = a * 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# A soft round glow, white, for additive light:  bloom on the overlay.
def glow(size=128, power=2.2):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	d = np.sqrt((x - size / 2) ** 2 + (y - size / 2) ** 2) / (size / 2)
	v = np.clip(1 - d, 0, 1) ** power
	img = np.zeros((size, size, 4))
	img[..., :3] = 255
	img[..., 3] = v * 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# Lettering in eighties chrome:  a sky-to-sunset gradient with a hard horizon line through the middle,
# a white outline, a dark outline outside that, and a soft magenta glow behind.
def chromeText(text, size, font=FONT_WIDE, pad=24):
	size = round(size * DENSITY)
	pad  = round(pad * DENSITY)
	f = ImageFont.truetype(font, size * SS)
	box = f.getbbox(text)
	w = (box[2] - box[0]) // SS + pad * 2
	h = (box[3] - box[1]) // SS + pad * 2
	mask = Image.new("L", (w * SS, h * SS), 0)
	ImageDraw.Draw(mask).text((pad * SS - box[0], pad * SS - box[1]), text, font=f, fill=255)
	m = np.array(mask).astype(np.float64) / 255
	rows = np.linspace(0, 1, h * SS)
	top = (pad * SS) / (h * SS)
	t = np.clip((rows - top) / (1 - 2 * top), 0, 1)[:, None]
	sky = np.array([20, 40, 120]) + (np.array([200, 245, 255]) - np.array([20, 40, 120])) * (t / 0.5)
	sunset = np.array([90, 20, 60]) + (np.array([255, 120, 200]) - np.array([90, 20, 60])) * ((t - 0.5) / 0.5)
	rowColour = np.where(t < 0.5, sky, sunset)
	grad = np.repeat(rowColour[:, None, :], w * SS, axis=1)
	from scipy import ndimage
	outer = ndimage.maximum_filter(m, size=round(5 * SS * DENSITY))
	ring = ndimage.maximum_filter(m, size=round(2 * SS * DENSITY))
	halo = ndimage.gaussian_filter(outer, 7 * SS * DENSITY)
	img = np.zeros((h * SS, w * SS, 4))
	img[..., :3] = np.array([255, 40, 200]) * 1.0
	img[..., 3] = np.clip(halo * 1.4, 0, 1) * 180
	img[..., :3] = np.where(outer[..., None] > 0.5, [20, 4, 30], img[..., :3])
	img[..., 3] = np.where(outer > 0.5, 255, img[..., 3])
	img[..., :3] = np.where(ring[..., None] > 0.5, [255, 255, 255], img[..., :3])
	img[..., :3] = np.where(m[..., None] > 0.5, grad, img[..., :3])
	return shrink(Image.fromarray(np.clip(img, 0, 255).astype(np.uint8), "RGBA"), w, h)


def flag():
	w, h = 512, 320
	im = Image.new("RGBA", (w, h), (20, 6, 40, 255))
	d = ImageDraw.Draw(im)
	for i in range(8):
		y = int(h * 0.62 + i * i * 2.2)
		d.rectangle([0, y, w, y + 2], fill=(255, 40, 200, 255))
	d.ellipse([w // 2 - 90, 40, w // 2 + 90, 220], fill=(255, 160, 70, 255))
	logo = chromeText("NEON COAST", 64)
	logo = logo.resize((w - 40, int(logo.height * (w - 40) / logo.width)), Image.LANCZOS)
	im.alpha_composite(logo, (20, h // 2 - logo.height // 2 + 20))
	return im


# A building's face:  dark glass in a grid of windows, most lit warm, some cyan or magenta, some dark,
# with a floor line every row.  Used as the colour and the emissive map, so the windows glow.
def facade(w=128, h=256):
	rng = np.random.default_rng(12)
	img = np.zeros((h, w, 4))
	img[..., :3] = [10, 6, 18]
	img[..., 3] = 255
	for y in range(4, h - 4, 12):
		img[y - 2, :, :3] = [24, 14, 36]
		for x in range(4, w - 4, 10):
			r = rng.random()
			if r < 0.55:
				c = [255, 200, 120]
			elif r < 0.68:
				c = [80, 230, 255]
			elif r < 0.75:
				c = [255, 60, 200]
			else:
				continue
			k = 0.55 + 0.45 * rng.random()
			img[y:y + 6, x:x + 6, :3] = np.array(c) * k
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# Terrain for the run over the mountains:  ridges either side of a valley down the middle, repeating
# exactly along its length so the tiles can leapfrog for ever.  Red carries the height.
def terrain(width=640, depth=256, span=3.5):
	# span is the tile's width in tile depths:  the ridges keep the size they have along the run however
	# wide the tile is, so a wider tile is more mountains rather than the same ones stretched.
	z = np.arange(depth) / depth * 2 * math.pi
	x = np.linspace(-span, span, width)
	X, Z = np.meshgrid(x, z)
	h = np.zeros_like(X)
	for k, (fx, fz, amp) in enumerate(((3, 2, 1.0), (7, 5, 0.5), (13, 9, 0.25), (23, 14, 0.12))):
		h += amp * (1 - np.abs(np.sin(X * fx * 1.7 + k) * np.sin(Z * fz + k * 1.3)))
	valley = np.clip((np.abs(X) - 0.12) / 0.5, 0, 1) ** 1.4
	h = h / h.max() * valley
	img = (np.clip(h, 0, 1) * 255).astype(np.uint8)
	return Image.fromarray(np.stack([img, img, img], -1), "RGB")


def main():
	os.makedirs(OUT, exist_ok=True)
	save(sky(), "sky.png")
	save(gridTile(MAGENTA), "grid.png")
	save(gridTile(CYAN, width=3.0, glow=10.0), "gridCyan.png")
	save(mountains(), "mountains.png")
	save(palm(), "palm.png")
	save(lamp(), "lamp.png")
	save(roadStrip(), "road.png")
	save(bob(), "bob.png")
	save(glow(), "glow.png")
	save(chromeText("NEON COAST", 96), "logo.png")
	save(chromeText("SINGE", 120), "singe.png")
	save(flag(), "flag.png")
	save(facade(), "facade.png")
	save(terrain(), "terrain.png")


if __name__ == "__main__":
	main()