# Analyses a demo's music and writes what the demo synchronises to as a Lua data file:  the beats,
# where the bars start, envelopes of the drums, the bass, the loudness and the brightness sixty times
# a second, the times of the drum hits, and the pitch class of the strongest note.
#
# The beats come from a fixed grid when the track is on one (a produced track usually is), found by
# folding the whole track's onsets at every candidate period and keeping the sharpest;  a track that
# drifts gets the tracker's beats, each moved to the strongest fine onset near it and smoothed through
# its neighbours.  Beat one of the bar is where the loudness and the bass jump most often, since a
# song changes section on a downbeat.
#
# Each envelope is a string of base64-alphabet characters, one a frame, 0 to 63, so it sits in a Lua
# long string with no escapes and the demo decodes it once at startup.
#
# Needs librosa, which Singe does not:  python3 -m venv env && env/bin/pip install librosa
#
# Usage:  env/bin/python util/demoSync.py music.mp3 out.lua
import sys

import librosa
import numpy as np

RATE     = 60                       # Envelope samples a second:  one a frame
SR       = 22050
HOP      = 256
LEVELS   = 64
ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
DECAY    = 0.86                     # Per sample:  a hit's envelope halves in about 75 ms
NOTES    = "C C# D D# E F F# G G# A A# B".split()


def beatGrid(y, yp, dur):
	onset = librosa.onset.onset_strength(y=yp, sr=SR, hop_length=HOP)
	ot    = librosa.frames_to_time(np.arange(len(onset)), sr=SR, hop_length=HOP)
	_, b  = librosa.beat.beat_track(onset_envelope=onset, sr=SR, hop_length=HOP, tightness=400)
	bt    = librosa.frames_to_time(b, sr=SR, hop_length=HOP)
	w     = np.maximum(0, onset - np.median(onset))
	guess = np.median(np.diff(bt))
	best  = None
	for period in np.arange(guess * 0.97, guess * 1.03, 0.00001):
		hist, _ = np.histogram(np.mod(ot, period) / period, bins=100, range=(0, 1), weights=w)
		hist = np.convolve(np.r_[hist[-2:], hist, hist[:2]], np.ones(5) / 5, "valid")
		if best is None or hist.max() > best[0]:
			best = (hist.max(), period, (np.argmax(hist) + 0.5) / 100 * period)
	_, period, offset = best
	resid = bt - (offset + period * np.round((bt - offset) / period))
	if np.median(np.abs(resid)) <= 0.015:
		print("  a fixed grid:  %.3f bpm, beat 0 at %.4f s;  tracked beats off it by %.1f ms (median)" % (60 / period, offset, np.median(np.abs(resid)) * 1000))
		return offset + period * np.arange(int((dur - offset) / period) + 1)
	fine = librosa.onset.onset_strength(y=yp, sr=SR, hop_length=64)
	ft   = librosa.frames_to_time(np.arange(len(fine)), sr=SR, hop_length=64)
	_, b = librosa.beat.beat_track(onset_envelope=onset, sr=SR, hop_length=HOP, tightness=100, start_bpm=60 / period)
	ref  = []
	for t in librosa.frames_to_time(b, sr=SR, hop_length=HOP):
		m = np.where(np.abs(ft - t) < 0.05)[0]
		ref.append(ft[m[np.argmax(fine[m])]])
	ref    = np.array(ref)
	n      = np.arange(len(ref))
	smooth = np.empty(len(ref))
	for k in n:
		lo = max(0, k - 8)
		hi = min(len(ref), k + 9)
		smooth[k] = np.polyval(np.polyfit(n[lo:hi], ref[lo:hi], 1), k)
	c     = np.polyfit(n[-24:], smooth[-24:], 1)
	extra = []
	k     = len(smooth)
	while np.polyval(c, k) < dur:
		extra.append(np.polyval(c, k))
		k += 1
	print("  no fixed grid:  %d beats tracked and smoothed (%.1f ms median off), %d carried to the end;  %.2f to %.2f bpm" % (
		len(smooth), np.median(np.abs(ref - smooth)) * 1000, len(extra), 60 / np.median(np.diff(smooth[:32])), 60 / np.median(np.diff(smooth[-32:]))))
	return np.r_[smooth, extra]


def flux(M, freqs, lo, hi):
	b = np.log1p(M[(freqs >= lo) & (freqs < hi)] * 10)
	return np.maximum(0, np.diff(b, axis=1, prepend=b[:, :1])).sum(axis=0)


def band(M, freqs, lo, hi):
	return np.sqrt((M[(freqs >= lo) & (freqs < hi)] ** 2).sum(axis=0))


def normal(curve):
	return np.clip(curve / (np.percentile(curve, 99.5) + 1e-9), 0, 1)


# Onto the RATE grid, keeping each sample's peak so a short hit is not averaged away.
def resample(curve, ft, frames):
	out = np.zeros(frames)
	idx = np.clip((ft * RATE).astype(int), 0, frames - 1)
	np.maximum.at(out, idx, curve)
	return out


def decayed(curve):
	out  = np.empty_like(curve)
	held = 0.0
	for i, v in enumerate(curve):
		held   = max(v, held * DECAY)
		out[i] = held
	return out


# Hits:  local peaks over a moving median, at least gap seconds apart.
def hits(curve, ft, gap, k):
	win  = int(1.0 * SR / HOP)
	pad  = np.pad(curve, win // 2, mode="edge")
	med  = np.array([np.median(pad[i:i + win]) for i in range(len(curve))])
	std  = np.array([np.std(pad[i:i + win]) for i in range(len(curve))])
	pick = []
	last = -1.0
	for i in range(1, len(curve) - 1):
		if curve[i] >= curve[i - 1] and curve[i] > curve[i + 1] and curve[i] > med[i] + k * std[i] and ft[i] - last >= gap:
			pick.append(ft[i])
			last = ft[i]
	return np.array(pick)


def packed(curve):
	q = np.clip(np.round(curve * (LEVELS - 1)), 0, LEVELS - 1).astype(int)
	return "".join(ALPHABET[v] for v in q)


def numbers(values, fmt):
	text  = []
	line  = []
	for v in values:
		line.append(fmt % v)
		if len(line) == 12:
			text.append("\t" + ", ".join(line) + ",")
			line = []
	if line:
		text.append("\t" + ", ".join(line) + ",")
	return "\n".join(text)


def main():
	source, target = sys.argv[1], sys.argv[2]
	print(source)
	y, _   = librosa.load(source, sr=SR, mono=True)
	dur    = len(y) / SR
	yh, yp = librosa.effects.hpss(y)
	grid   = beatGrid(y, yp, dur)

	S     = np.abs(librosa.stft(y, n_fft=2048, hop_length=HOP))
	Sp    = np.abs(librosa.stft(yp, n_fft=2048, hop_length=HOP))
	Sh    = np.abs(librosa.stft(yh, n_fft=2048, hop_length=HOP))
	freqs = librosa.fft_frequencies(sr=SR, n_fft=2048)
	ft    = librosa.frames_to_time(np.arange(S.shape[1]), sr=SR, hop_length=HOP)

	kick  = normal(flux(Sp, freqs, 30, 150))
	snare = normal(flux(Sp, freqs, 1500, 5000))
	hat   = normal(flux(Sp, freqs, 7000, 11000))
	bass  = normal(band(Sh, freqs, 40, 250))
	rms   = normal(librosa.feature.rms(y=y, hop_length=HOP)[0])
	cent  = librosa.feature.spectral_centroid(S=S, sr=SR)[0]
	high  = np.clip((cent - 1000) / 3500, 0, 1)

	# Beat one:  where loudness and bass jump most often from the beat before.
	def over(curve, a, b):
		m = (ft >= a) & (ft < b)
		return curve[m].mean() if m.any() else 0

	level = np.array([over(rms, grid[i], grid[i + 1]) + over(bass, grid[i], grid[i + 1]) for i in range(len(grid) - 1)])
	jump  = np.r_[0, np.abs(np.diff(level))]
	big   = jump > np.percentile(jump, 95)
	votes = [int(big[p::4].sum()) for p in range(4)]
	phase = int(np.argmax(votes))
	print("  beat one is beat %d of the grid (votes %s)" % (phase, votes))

	# The strongest pitch class in the upper harmonic part, and how clear it is.
	chroma = librosa.feature.chroma_stft(S=Sh * (freqs[:, None] > 300), sr=SR)
	note   = np.argmax(chroma, axis=0)
	clear  = chroma.max(axis=0) / (chroma.sum(axis=0) + 1e-9)

	frames = int(dur * RATE) + 1
	env    = {}
	for name, curve, fall in (("kick", kick, True), ("snare", snare, True), ("hat", hat, True), ("bass", bass, False), ("level", rms, False), ("bright", high, False)):
		c = resample(curve, ft, frames)
		if fall:
			c = decayed(c)
		else:
			c = np.convolve(c, np.ones(5) / 5, "same")
		env[name] = c
	idx       = np.clip(np.round(np.arange(frames) / RATE * SR / HOP).astype(int), 0, len(note) - 1)
	noteFrame = note[idx]
	clearness = np.clip((clear[idx] - 0.12) / 0.25, 0, 1)

	kickHits  = hits(kick, ft, 0.12, 1.5)
	snareHits = hits(snare, ft, 0.12, 1.5)
	hatHits   = hits(hat, ft, 0.07, 1.2)
	print("  %d frames;  %d kick hits, %d snare hits, %d hat hits;  %d beats, %.1f s" % (frames, len(kickHits), len(snareHits), len(hatHits), len(grid), dur))

	with open(target, "w") as out:
		out.write("-- Written by util/demoSync.py from %s:  what the demo synchronises to.  Times are seconds of\n" % source.split("/")[-1])
		out.write("-- music;  envelopes are %d samples a second, one character each, 0 to 63 in the base64 alphabet.\n" % RATE)
		out.write("return {\n")
		out.write("\tduration = %.3f,\n" % dur)
		out.write("\trate     = %d,\n" % RATE)
		out.write("\tbarStart = %d,  -- The beat (counting from 1) that begins the first whole bar\n" % (phase + 1))
		out.write("\tbeats = {\n%s\n\t},\n" % numbers(grid, "%.4f"))
		for name, values in (("kicks", kickHits), ("snares", snareHits), ("hats", hatHits)):
			out.write("\t%s = {\n%s\n\t},\n" % (name, numbers(values, "%.3f")))
		for name in ("kick", "snare", "hat", "bass", "level", "bright"):
			out.write("\t%s = [[%s]],\n" % (name, packed(env[name])))
		out.write("\tnote = [[%s]],\n" % "".join(ALPHABET[v] for v in noteFrame))
		out.write("\tnoteClear = [[%s]],\n" % packed(clearness))
		out.write("}\n")


if __name__ == "__main__":
	main()