# Draws the art for demos/mirrorball:  every picture is made from shapes and arithmetic here, so it is
# ours and can be drawn again.
#
# The club is an equirectangular picture, the whole room round the camera:  the 3D scene lights with
# it and every facet of the mirror ball reflects a piece of it, so it is mostly dark with bright
# coloured lights round the walls and along a truss overhead, which is what makes a mirror ball
# sparkle.  The rest are sprites and textures for the oldschool effects.
#
# Usage:  python3 util/mirrorballArt.py
import math
import os

import numpy as np
from PIL import Image, ImageDraw, ImageFont
from scipy import ndimage


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

SS        = 4                     # Sprites are drawn this many times larger and shrunk, for smooth edges
FONT_WIDE = "/usr/share/fonts/truetype/lato/Lato-BlackItalic.ttf"
GOLD_TOP  = np.array([255, 250, 210])
GOLD_MID  = np.array([200, 130, 30])
GOLD_LOW  = np.array([255, 210, 120])
GOLD_DARK = np.array([90, 40, 10])
DISCO     = [(255, 40, 80), (255, 160, 20), (255, 240, 60), (60, 255, 120), (40, 200, 255), (160, 60, 255)]


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


# The room:  near-black walls, a band of coloured lamps round them at eye height and a second band
# higher up, a truss of white lamps across the ceiling, and a floor that faintly carries the colours.
def club():
	w, h = 2048, 1024
	lon = (np.arange(w) + 0.5) / w * 2 * math.pi - math.pi
	lat = math.pi / 2 - (np.arange(h) + 0.5) / h * math.pi
	LON, LAT = np.meshgrid(lon, lat)
	img = np.zeros((h, w, 3))
	img[:] = np.array([8, 6, 16])
	img += (np.clip(LAT, 0, None) / (math.pi / 2))[..., None] * np.array([10, 6, 24])

	def lamps(count, latDeg, size, bright, colours):
		out = np.zeros_like(img)
		for i in range(count):
			c = np.array(colours[i % len(colours)])
			lo = -math.pi + (i + 0.5) / count * 2 * math.pi
			dl = np.angle(np.exp(1j * (LON - lo)))
			d = np.sqrt((dl * np.cos(LAT)) ** 2 + (LAT - math.radians(latDeg)) ** 2) / math.radians(size)
			out += (np.exp(-d * d * 3) * bright + np.exp(-d * 1.2) * 0.25)[..., None] * c
		return out

	img += lamps(24, 4, 2.2, 1.6, DISCO)
	img += lamps(16, 28, 3.0, 1.2, DISCO[::-1])
	img += lamps(12, 70, 2.0, 2.0, [(255, 255, 255)])
	floor = LAT < -0.05
	img[floor] = img[floor] * 0.35 + np.array([14, 8, 22])
	return Image.fromarray(np.clip(img, 0, 255).astype(np.uint8), "RGB")


# A floor tile:  white with a bevelled edge and a dark seam, as both colour and emissive map, so the
# tile's own emissive colour lights it.
def tile(size=128):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	d = np.minimum(np.minimum(x, size - x), np.minimum(y, size - y))
	v = np.clip((d - 3) / 10, 0, 1) * 0.8 + 0.2 * (d > 3)
	v *= 1 - 0.25 * ((x + y) / (2 * size))
	img = np.zeros((size, size, 4))
	img[..., :3] = (np.clip(v, 0, 1) * 255)[..., None]
	img[..., 3] = 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


def chromeText(text, size, pad=24):
	f = ImageFont.truetype(FONT_WIDE, 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]
	upper = GOLD_TOP + (GOLD_MID - GOLD_TOP) * (t / 0.5)
	lower = GOLD_DARK + (GOLD_LOW - GOLD_DARK) * ((t - 0.5) / 0.5)
	grad = np.repeat(np.where(t < 0.5, upper, lower)[:, None, :], w * SS, axis=1)
	outer = ndimage.maximum_filter(m, size=5 * SS)
	ring = ndimage.maximum_filter(m, size=2 * SS)
	halo = ndimage.gaussian_filter(outer, 7 * SS)
	img = np.zeros((h * SS, w * SS, 4))
	img[..., :3] = np.array([255, 160, 40])
	img[..., 3] = np.clip(halo * 1.4, 0, 1) * 170
	img[..., :3] = np.where(outer[..., None] > 0.5, [30, 10, 4], img[..., :3])
	img[..., 3] = np.where(outer > 0.5, 255, img[..., 3])
	img[..., :3] = np.where(ring[..., None] > 0.5, [255, 255, 230], img[..., :3])
	img[..., :3] = np.where(m[..., None] > 0.5, grad, img[..., :3])
	big = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8), "RGBA")
	return big.resize((w, h), Image.LANCZOS)


# The twister's four faces side by side, 64 pixels each:  stripes, a checker, chevrons, and dots,
# each in its own disco colour, so a turning bar shows which face is which.
def twister():
	fw, fh = 64, 64
	img = np.zeros((fh, fw * 4, 4))
	y, x = np.mgrid[0:fh, 0:fw] + 0.5
	pats = [
		((np.floor(y / 8) % 2) == 0),
		((np.floor(x / 16) + np.floor(y / 16)) % 2 == 0),
		((np.abs((x - 32)) + y) % 32 < 16),
		(((x % 16 - 8) ** 2 + (y % 16 - 8) ** 2) < 30),
	]
	cols = [DISCO[0], DISCO[2], DISCO[4], DISCO[5]]
	for i, (p, c) in enumerate(zip(pats, cols)):
		face = np.where(p[..., None], np.array(c), np.array(c) * 0.35)
		edge = (x < 2) | (x > fw - 2)
		face[edge] = [255, 255, 255]
		img[:, i * fw:(i + 1) * fw, :3] = face
	img[..., 3] = 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# One copper bar:  a vertical gradient, white in the middle to dark at its edges, tinted per bar.
def copper(h=32):
	y = (np.arange(h) + 0.5) / h
	v = np.sin(y * math.pi) ** 1.5
	img = np.zeros((h, 4, 4))
	img[..., :3] = (v * 255)[:, None, None]
	img[..., 3] = 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# The rotozoomer's tile:  a disco checker with a four-pointed star, colourful enough to read when it
# is turned and squeezed.
def rotoTile(size=128):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	c = ((np.floor(x / 64) + np.floor(y / 64)) % 2)
	img = np.zeros((size, size, 3))
	img[:] = np.where(c[..., None] > 0, np.array([40, 10, 60]), np.array([255, 120, 30]))
	cx = (x % 64) - 32
	cy = (y % 64) - 32
	star = (np.abs(cx) * np.abs(cy) < 30) & (np.abs(cx) + np.abs(cy) < 30)
	img[star] = [255, 250, 200]
	return Image.fromarray(img.astype(np.uint8), "RGB")


# The boing ball's skin:  red and white in sixteen by eight, as an equirectangular texture.
def boing(w=256, h=128):
	y, x = np.mgrid[0:h, 0:w]
	c = ((x // (w // 16)) + (y // (h // 8))) % 2
	img = np.where(c[..., None] > 0, np.array([230, 20, 30]), np.array([250, 250, 250]))
	return Image.fromarray(img.astype(np.uint8), "RGB")


def bob(size=64):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	r = size / 2 - 1
	dx = (x - size / 2) / r
	dy = (y - size / 2) / r
	d2 = dx * dx + dy * dy
	nz = np.sqrt(np.clip(1 - d2, 0, 1))
	light = np.clip(dx * -0.4 + dy * -0.5 + nz * 0.77, 0, 1)
	v = 0.18 + 0.7 * light + 0.6 * light ** 24
	a = np.clip((1 - np.sqrt(d2)) * r, 0, 1) * (d2 <= 1)
	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")


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)
	img = np.zeros((size, size, 4))
	img[..., :3] = 255
	img[..., 3] = np.clip(1 - d, 0, 1) ** power * 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


# A star-shaped sparkle, for the light the ball throws.
def sparkle(size=64):
	y, x = np.mgrid[0:size, 0:size] + 0.5
	dx = np.abs(x - size / 2) / (size / 2)
	dy = np.abs(y - size / 2) / (size / 2)
	v = np.exp(-dx * 14) * np.exp(-dy * 2.5) + np.exp(-dy * 14) * np.exp(-dx * 2.5) + np.exp(-(dx * dx + dy * dy) * 18)
	img = np.zeros((size, size, 4))
	img[..., :3] = 255
	img[..., 3] = np.clip(v, 0, 1) * 255
	return Image.fromarray(img.astype(np.uint8), "RGBA")


def main():
	os.makedirs(OUT, exist_ok=True)
	save(club(), "club.png")
	save(tile(), "tile.png")
	save(chromeText("MIRRORBALL", 96), "logo.png")
	save(chromeText("SINGE", 120), "singe.png")
	save(twister(), "twister.png")
	save(copper(), "copper.png")
	save(rotoTile(), "roto.png")
	save(boing(), "boing.png")
	save(bob(), "bob.png")
	save(glow(), "glow.png")
	save(sparkle(), "sparkle.png")


if __name__ == "__main__":
	main()