# Builds the first library demo of the game's sound effects into build/sfx-sets/demo1 (for the sound review), next to
# the synthesized set in build/sfx-sets/synth. Each game sound is cut from one library track (or a few, layered): the track is split into its separate
# sounds by loudness, one is picked, trimmed and faded, optionally pitched, matched to the loudness of the game's current
# sound and bitcrushed. SOURCES.md beside the demo lists where each one came from, and compare.html plays each beside
# the game's current sound.
#
# python tools/sfx_library_demo.py [library folder]

import pathlib
import subprocess
import sys
import tempfile
import wave

import numpy as np

ROOT = pathlib.Path(__file__).resolve().parent.parent
SFX = ROOT / "build" / "sfx-sets" / "synth"
DEMO = ROOT / "build" / "sfx-sets" / "demo1"
LIBRARY = pathlib.Path.home() / "Trash/universal-studios-sound-effects-library/Universal Studios Sound Effects Library (OGG)"
RATE = 44100
OUT_RATE = 22050
CRUSH = "acrusher=bits=8:mode=log:samples=4:aa=0:mix=1"
FRAME = 0.002
GAP = 0.05
THRESHOLD = 0.12
SIGNIFICANT = 0.3
TRUNCATED_FADE = 0.35
TAIL = 0.02
FADE = 0.012
LOUDNESS_WINDOW = 0.05
PEAK = 0.95

# Each game sound: the library track's number, which of its sounds to take (by order among those at least a share as
# loud as its loudest, -1 for the last, or "loud" for the loudest) or a start in seconds, the longest it may run, a speed (which pitches it too) and a gain. A list layers several, each with an "at"
# offset in seconds.
PICKS = {
    # Movement and feedback
    "jump": {"f": "US02-79-1", "n": 0.22, "sp": 1.3},
    "land": {"f": "US02-52-1", "n": 0.15, "sp": 1.3},
    "step": {"f": "US02-52-2", "n": 0.06, "sp": 1.8, "g": 0.6},
    "crawl": {"f": "US02-80-1", "n": 0.1, "sp": 1.4, "g": 0.7},
    "duck": {"f": "US03-21-1", "n": 0.15},
    "unduck": {"f": "US03-21-1", "n": 0.12, "sp": 1.5},
    "roll": {"f": "US02-51-4", "e": "loud", "n": 0.12, "sp": 1.3},
    "dash": {"f": "US02-79-2", "n": 0.2, "sp": 1.2},
    "climb": {"f": "US02-62-1", "n": 0.14},
    "fall": {"f": "US02-87-2", "n": 0.7, "sp": 1.2},
    "bounce": {"f": "US02-68-1", "n": 0.3, "sp": 1.2},
    "pound": {"f": "US03-47-1", "n": 0.55},
    "slam": {"f": "US02-50-3", "n": 0.7},
    "death": {"f": "US02-73-1", "n": 0.9},
    "finish": {"f": "US02-22-3", "e": 1, "n": 1.0},
    "break": {"f": "US02-88-1", "n": 0.45},
    "ready": {"f": "US02-77-1", "n": 0.5, "sp": 1.3},
    "charge": {"f": "US03-69-2", "s": 3.0, "n": 0.8, "sp": 2.0},
    "select": {"f": "US02-64-1", "n": 0.06, "sp": 1.2, "g": 0.7},
    "swap": {"f": "US02-61-2", "n": 0.1},
    "equip": {"f": "US04-50-3", "n": 0.3},
    "unequip": {"f": "US04-41-1", "n": 0.26},
    "level_up": {"f": "US02-76-1", "n": 0.4},
    "level_down": {"f": "US02-86-3", "n": 0.4},
    "echo": {"f": "US02-77-2", "n": 0.3, "sp": 1.5, "g": 0.7},
    "ricochet": {"f": "US04-62-1", "n": 0.3},
    "zap": {"f": "US03-12-1", "n": 0.15},
    # Weapons firing
    "slash": {"f": "US04-30-1", "n": 0.2, "sp": 1.3},
    "slash_heavy": {"f": "US03-74-1", "n": 0.45},
    "swing": {"f": "US02-60-1", "n": 0.3},
    "swing_heavy": {"f": "US02-60-2", "n": 0.6},
    "poke": {"f": "US03-79-1", "n": 0.15, "sp": 1.2},
    "poke_heavy": {"f": "US03-21-3", "n": 0.35},
    "whip": {"f": "US03-89-1", "n": 0.25},
    "whip_heavy": {"f": "US03-89-3", "n": 0.35},
    "whip_blast": {"f": "US04-11-1", "n": 0.5},
    "grapple": {"f": "US04-42-1", "n": 0.2},
    "grapple_heavy": {"f": "US04-18-1", "n": 0.3},
    "latch": {"f": "US03-35-1", "n": 0.2},
    "boomerang": {"f": "US02-79-3", "n": 0.3},
    "boomerang_heavy": {"f": "US02-79-4", "n": 0.5, "sp": 0.85},
    "bow": {"f": "US02-44-1", "n": 0.3},
    "bow_heavy": {"f": "US02-44-2", "n": 0.5, "sp": 0.9},
    "blast": {"f": "US03-59-4", "n": 0.15},
    "beam": {"f": "US03-59-2", "n": 0.45},
    "uzi": {"f": "US04-72-1", "n": 0.08, "g": 0.7},
    "uzi_heavy": {"f": "US04-81-1", "n": 0.3},
    "shotgun": [{"f": "US04-59-1", "n": 0.45}, {"f": "US04-51-1", "n": 0.35, "at": 0.35, "g": 0.6}],
    "shotgun_heavy": [{"f": "US04-59-3", "n": 0.5}, {"f": "US04-51-2", "n": 0.35, "at": 0.45, "g": 0.6}],
    "flame": {"f": "US03-17-1", "n": 0.15, "sp": 1.4, "g": 0.6},
    "fireball": {"f": "US03-17-2", "n": 0.5},
    "rocket": {"f": "US04-02-1", "n": 0.45},
    "rocket_heavy": {"f": "US04-07-2", "n": 0.8},
    "explosion": {"f": "US04-12-1", "n": 0.9},
    "mine": {"f": "US02-56-2", "e": "loud", "n": 0.25, "sp": 1.3},
    "detonate": {"f": "US03-82-1", "n": 0.12},
    "detonate_empty": {"f": "US04-51-3", "n": 0.08},
    "flail": {"f": "US04-50-4", "n": 0.3},
    "flail_heavy": {"f": "US03-39-1", "n": 0.45},
    "flail_land": {"f": "US03-39-2", "n": 0.25},
    "punch": {"f": "US02-55-2", "n": 0.1},
    "uppercut": {"f": "US02-87-1", "n": 0.45, "sp": 1.8},
    # Shots landing on walls
    "impact_arrow": {"f": "US02-46-1", "n": 0.2},
    "impact_blaster": {"f": "US03-08-1", "n": 0.12},
    "impact_bullet": {"f": "US04-71-3", "n": 0.08},
    "impact_boomerang": {"f": "US02-70-2", "n": 0.18},
    "impact_mine": {"f": "US03-36-1", "n": 0.15},
    # Hits, by type, light then heavy
    "hit_blunt": {"f": "US03-40-1", "n": 0.15},
    "hit_blunt_heavy": {"f": "US03-20-1", "n": 0.3},
    "hit_slash": {"f": "US03-78-1", "n": 0.18},
    "hit_slash_heavy": {"f": "US03-76-1", "n": 0.3},
    "hit_pierce": {"f": "US03-71-1", "n": 0.12},
    "hit_pierce_heavy": {"f": "US03-79-2", "n": 0.25},
    "hit_lash": {"f": "US03-91-2", "n": 0.1},
    "hit_lash_heavy": {"f": "US03-91-1", "e": "loud", "n": 0.2},
    "hit_bullet": {"f": "US04-71-2", "n": 0.06},
    "hit_bullet_heavy": {"f": "US03-20-3", "n": 0.35},
    "hit_energy": {"f": "US03-59-1", "e": -1, "n": 0.14},
    "hit_energy_heavy": {"f": "US03-08-2", "n": 0.3},
    "hit_fire": {"f": "US03-45-3", "n": 0.25},
    "hit_fire_heavy": {"f": "US03-11-1", "n": 0.47},
    "hit_explosive": {"f": "US04-11-2", "n": 0.3},
    "hit_explosive_heavy": {"f": "US04-01-1", "n": 0.5},
    "hit_magic": {"f": "US02-66-2", "n": 0.2},
    "hit_magic_heavy": {"f": "US02-22-4", "e": 1, "n": 0.35},
}


def decode(path):
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", str(path), "-ac", "1", "-ar", str(RATE), "-f", "f32le", "-"],
                         capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype=np.float32).astype(np.float64)


def track_find(number):
    matches = [path for path in LIBRARY.glob(f"{number} *.ogg")]
    assert len(matches) == 1, f"{number}: {len(matches)} tracks"
    return matches[0]


# The separate sounds in a track: runs of frames louder than a share of the track's peak, merged across short gaps.
# Each is (start, end, peak) in samples.
def events_get(signal):
    size = int(RATE * FRAME)
    frames = np.abs(signal[: len(signal) // size * size]).reshape(-1, size).max(axis=1)
    floor = np.percentile(frames, 10)
    above = frames > max(frames.max() * THRESHOLD, floor * 5)
    events = []
    for i, isAbove in enumerate(above):
        if not isAbove:
            continue
        if events and i - events[-1][1] <= GAP / FRAME:
            events[-1][1] = i
        else:
            events.append([i, i])
    return [(start * size, (end + 1) * size, frames[start:end + 1].max()) for start, end in events]


# One sound from a track: from just before it starts until it has died away (or the next sound starts, or the length
# runs out), with its end faded (over its last third when the length cut it short).
def cut(pick):
    signal = decode(track_find(pick["f"]))
    if "s" in pick:
        start = int(pick["s"] * RATE)
        end = len(signal)
    else:
        events = events_get(signal)
        loudest = max(event[2] for event in events)
        significant = [event for event in events if event[2] >= loudest * SIGNIFICANT]
        choice = pick.get("e", 0)
        start, eventEnd, peak = max(significant, key=lambda event: event[2]) if choice == "loud" else significant[choice]
        start = max(0, start - int(0.004 * RATE))
        following = [event[0] for event in events if event[0] > eventEnd]
        end = following[0] if following else len(signal)
        quiet = np.nonzero(np.abs(signal[eventEnd:end]) > peak * TAIL)[0]
        end = min(end, eventEnd + (quiet[-1] if len(quiet) else 0) + int(FADE * RATE))
    speed = pick.get("sp", 1)
    limit = start + int(pick["n"] * speed * RATE)
    isTruncated = limit < end
    piece = signal[start:min(end, limit)]
    if speed != 1:
        piece = np.interp(np.arange(0, len(piece) - 1, speed), np.arange(len(piece)), piece)
    fade = min(len(piece), max(int(FADE * RATE), int(len(piece) * TRUNCATED_FADE * isTruncated)))
    piece[len(piece) - fade:] *= np.linspace(1, 0, fade)
    return piece * pick.get("g", 1), start / RATE


def loudness(signal):
    size = int(LOUDNESS_WINDOW * RATE)
    if len(signal) <= size:
        return np.sqrt(np.mean(signal ** 2))
    power = np.convolve(signal ** 2, np.ones(size) / size, mode="valid")
    return np.sqrt(power.max())


def sound_build(name, picks):
    picks = picks if isinstance(picks, list) else [picks]
    pieces = [(cut(pick), pick.get("at", 0)) for pick in picks]
    length = max(int(at * RATE) + len(piece) for (piece, _), at in pieces)
    mixed = np.zeros(length)
    for (piece, _), at in pieces:
        mixed[int(at * RATE):int(at * RATE) + len(piece)] += piece
    target = loudness(decode(SFX / f"{name}.ogg"))
    mixed *= target / max(loudness(mixed), 1e-9)
    mixed *= min(1, PEAK / max(np.abs(mixed).max(), 1e-9))
    return mixed, [(pick, start) for pick, ((_, start), _) in zip(picks, pieces)]


def save(name, signal):
    with tempfile.TemporaryDirectory() as directory:
        wavPath = pathlib.Path(directory) / "sound.wav"
        with wave.open(str(wavPath), "wb") as file:
            file.setnchannels(1)
            file.setsampwidth(2)
            file.setframerate(RATE)
            file.writeframes((np.clip(signal, -1, 1) * 32767).astype(np.int16).tobytes())
        subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", str(wavPath), "-af", CRUSH, "-ar", str(OUT_RATE), "-ac", "1",
                        "-c:a", "libvorbis", "-q:a", "4", str(DEMO / f"{name}.ogg")], check=True)


def main():
    global LIBRARY
    if len(sys.argv) > 1:
        LIBRARY = pathlib.Path(sys.argv[1])
    DEMO.mkdir(exist_ok=True)
    names = sorted(path.stem for path in SFX.glob("*.ogg"))
    missing = [name for name in names if name not in PICKS]
    if missing:
        print("no pick, left out:", " ".join(missing))
    names = [name for name in names if name in PICKS]

    lines = ["# Demo sound sources", "", "Cut from the Universal Studios Sound Effects Library, bitcrushed "
             f"(`{CRUSH}`, {OUT_RATE} Hz mono).", "", "| Sound | Track | From | Speed |", "|---|---|---|---|"]
    for name in names:
        signal, sources = sound_build(name, PICKS[name])
        save(name, signal)
        for pick, start in sources:
            lines.append(f"| {name} | {track_find(pick['f']).stem} | {start:.2f}s | {pick.get('sp', 1)}x |")
        print(name, f"{len(signal) / RATE:.2f}s")
    (DEMO / "SOURCES.md").write_text("\n".join(lines) + "\n")
    compare_save(names)


# A page to audition the demo against the game's current sounds, side by side.
def compare_save(names):
    rows = []
    for name in names:
        picks = PICKS[name] if isinstance(PICKS[name], list) else [PICKS[name]]
        sources = "<br>".join(track_find(pick["f"]).stem for pick in picks)
        rows.append(f"<tr><td>{name}</td><td><audio controls preload=none src='../synth/{name}.ogg'></audio></td>"
                    f"<td><audio controls preload=none src='{name}.ogg'></audio></td><td>{sources}</td></tr>")
    page = ("<!DOCTYPE html><meta charset=utf-8><title>Sound demo</title><style>body{font:14px sans-serif;background:#111;"
            "color:#ddd}td{padding:4px 8px;border-bottom:1px solid #333}audio{height:28px}</style><table><tr><th>Sound</th>"
            "<th>Synth</th><th>Demo</th><th>Source</th></tr>" + "".join(rows) + "</table>")
    (DEMO / "compare.html").write_text(page)


if __name__ == "__main__":
    main()