# Precalculates the demos' per-pixel effects as videos:  the oldest trick in the demoscene, done with
# numpy instead of lookup tables in assembler.  Singe then plays them, scrubs them, turns them, and
# puts them on 3D surfaces.
#
#     python3 util/demoVideo.py tunnel   out.mp4 [palette] [frames]
#     python3 util/demoVideo.py plasma   out.mp4 [palette] [frames] [fps]
#     python3 util/demoVideo.py noise    out.mp4 [frames]
#
# tunnel draws exactly one period of travel down a texture-mapped tunnel in frames frames, every one a
# keyframe, so a demo can seek to any frame at once and move through the tunnel at whatever speed the
# music says;  the last frame leads back into the first.  plasma is an ordinary looping video whose
# every term repeats a whole number of times in its length, so it loops without a seam.  noise is a
# tape's snow and tracking bands, also looping.
import math
import subprocess
import sys

import numpy as np

W = 640
H = 480

PALETTES = {
	# Neon:  black, violet, magenta, orange, and a cyan accent.
	"neon": [(0.00, (4, 0, 16)), (0.25, (60, 10, 120)), (0.50, (255, 30, 170)), (0.70, (255, 140, 60)), (0.85, (255, 230, 150)), (1.00, (40, 230, 255))],
	# Disco:  deep red, gold, and white.
	"disco": [(0.00, (20, 0, 10)), (0.30, (150, 20, 40)), (0.55, (255, 150, 30)), (0.80, (255, 230, 140)), (1.00, (255, 255, 255))],
	# Arcade:  blue, green, and yellow.
	"arcade": [(0.00, (0, 0, 20)), (0.30, (10, 40, 160)), (0.55, (20, 200, 120)), (0.80, (240, 240, 60)), (1.00, (255, 255, 255))],
}


def ramp(name, t):
	stops = PALETTES[name]
	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


def encoder(out, fps, intra):
	args = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", "%dx%d" % (W, H), "-r", str(fps), "-i", "-",
		"-c:v", "libx264", "-preset", "slow", "-crf", "18", "-pix_fmt", "yuv420p", "-an"]
	if intra:
		args += ["-g", "1", "-keyint_min", "1"]
	else:
		args += ["-g", str(int(float(fps)))]
	return subprocess.Popen(args + [out], stdin=subprocess.PIPE)


# Down a tunnel:  each pixel's angle round the axis and its depth (inverse distance from the centre)
# index a texture of neon rings and spokes.  Depth runs one texture period over the frames.
def tunnel(out, palette="neon", frames=120):
	y, x = np.mgrid[0:H, 0:W].astype(np.float64)
	x = (x - W / 2 + 0.5) / (H / 2)
	y = (y - H / 2 + 0.5) / (H / 2)
	r = np.sqrt(x * x + y * y) + 1e-6
	a = np.arctan2(y, x) / (2 * math.pi)
	depth = 0.35 / r
	shade = np.clip(r * 1.6, 0, 1) ** 0.8
	pipe = encoder(out, 50, True)
	RINGS = 4
	SPOKES = 16
	for f in range(frames):
		p = f / frames
		u = depth + p / RINGS
		v = a + depth * 0.15
		ring = 0.5 + 0.5 * np.cos(u * RINGS * 2 * math.pi * RINGS)
		edge = np.abs(np.mod(u * RINGS * RINGS, 1.0) - 0.5) * 2
		spoke = np.abs(np.mod(v * SPOKES, 1.0) - 0.5) * 2
		check = ((np.floor(u * RINGS * RINGS) + np.floor(v * SPOKES)) % 2)
		glowRing = np.exp(-(1 - edge) * 9)
		glowSpoke = np.exp(-(1 - spoke) * 12)
		value = 0.18 + 0.22 * check + 0.75 * glowRing + 0.45 * glowSpoke * ring
		col = ramp(palette, np.clip(value * shade, 0, 1))
		# Far away the tunnel fades into a hot white core.
		core = np.exp(-r * 9)[..., None]
		col = col * (1 - core) + core * np.array([255, 235, 255])
		pipe.stdin.write(np.clip(col, 0, 255).astype(np.uint8).tobytes())
	pipe.stdin.close()
	pipe.wait()


# Plasma:  four interfering waves, each a whole number of turns over the loop, run through a palette
# that itself cycles a whole number of times.
def plasma(out, palette="neon", frames=400, fps=50):
	y, x = np.mgrid[0:H, 0:W].astype(np.float64)
	x /= W
	y /= H
	pipe = encoder(out, fps, False)
	for f in range(frames):
		p = f / frames * 2 * math.pi
		v = np.sin(x * 9 + p)
		v += np.sin((y * 7 + p * 2) * 1.0)
		v += np.sin((x * 6 + y * 5) + p * 3)
		cx = x - 0.5 + 0.3 * math.sin(p)
		cy = y - 0.5 + 0.3 * math.cos(p * 2)
		v += np.sin(np.sqrt(cx * cx + cy * cy) * 18 - p * 4)
		t = (np.sin(v * math.pi * 0.5 + p) + 1) / 2
		col = ramp(palette, t)
		pipe.stdin.write(np.clip(col, 0, 255).astype(np.uint8).tobytes())
	pipe.stdin.close()
	pipe.wait()


# A tape's snow, with a couple of tracking bands rolling up the picture, at a quarter of the size
# and scaled up by the player.
def noise(out, frames=100):
	global W, H
	W, H = 320, 240
	rng = np.random.default_rng(11)
	pipe = encoder(out, 25, False)
	rows = np.arange(H)[:, None]
	for f in range(frames):
		snow = rng.random((H, W)) ** 1.6 * 200
		# Grain in streaks:  each row's noise smeared sideways a little.
		snow = (snow + np.roll(snow, 1, axis=1) + np.roll(snow, 2, axis=1)) / 3
		p = f / frames
		band = np.exp(-((rows - (H * (1 - p) * 2) % (H + 40) + 20) ** 2) / 60.0)
		band2 = np.exp(-((rows - (H * (1 - p) * 2 + H / 2) % (H + 40) + 20) ** 2) / 20.0)
		v = snow * (0.55 + 0.45 * band + 0.6 * band2) + band * 60
		col = np.stack([v * 0.95, v * 0.9, v * 1.05], -1)
		pipe.stdin.write(np.clip(col, 0, 255).astype(np.uint8).tobytes())
	pipe.stdin.close()
	pipe.wait()


def main():
	kind = sys.argv[1]
	out = sys.argv[2]
	if kind == "tunnel":
		tunnel(out, sys.argv[3] if len(sys.argv) > 3 else "neon", int(sys.argv[4]) if len(sys.argv) > 4 else 120)
	elif kind == "plasma":
		plasma(out, sys.argv[3] if len(sys.argv) > 3 else "neon", int(sys.argv[4]) if len(sys.argv) > 4 else 400, sys.argv[5] if len(sys.argv) > 5 else 50)
	elif kind == "noise":
		noise(out, int(sys.argv[3]) if len(sys.argv) > 3 else 100)
	print("wrote", out)


if __name__ == "__main__":
	main()