#!/usr/bin/env python3
"""Optional sample-free stereo WAV renderer for the generated score.

Requires NumPy; the MIDI generator itself has no third-party dependencies.
Usage: python3 render_preview.py --out ./midnight-circuit
This is a transparent synthesis demonstration, not a mastered production.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import wave

import numpy as np

SR = 32000
TAU = 2 * np.pi
TRACK_GAIN = {"drums": .37, "bass": .29, "chords": .11, "lead": .12, "pad": .045}
TRACK_PAN = {"bass": 0, "chords": -.27, "lead": .18, "pad": 0}


def envelope(t, gate, attack, release):
    return np.minimum(t/max(attack, 1/SR), 1) * np.clip((gate+release-t)/release, 0, 1)


def synth(n, spb, rng):
    track, pitch = n["track"], n["pitch"]
    gate = n["duration"] * spb
    if track == "drums":
        lengths = {36: .36, 38: .23, 39: .18, 37: .11, 42: .062, 46: .21}
        t = np.arange(round(lengths[pitch]*SR), dtype=np.float64)/SR
        noise = rng.standard_normal(len(t))
        # First differences give the hats/snare a reproducible bright noise bed.
        bright = np.concatenate(([0.0], np.diff(noise))) * .5
        if pitch == 36:
            phase = TAU*(49*t + 93*.017*(1-np.exp(-t/.017)))
            wave_data = np.tanh(1.25*np.sin(phase))*np.exp(-t/0.085)
            wave_data += .12*bright*np.exp(-t/.003)
            wave_data *= np.minimum(t/.0015, 1)
            pan = 0
        elif pitch == 38:
            body = .48*np.sin(TAU*185*t)*np.exp(-t/.041)
            body += .17*np.sin(TAU*330*t)*np.exp(-t/.025)
            wave_data = body + .43*bright*np.exp(-t/.049)
            pan = -.025
        elif pitch == 39:
            bursts = sum(np.exp(-np.maximum(t-o, 0)/.016)*(t>=o) for o in (0, .012, .026))
            wave_data = .30*bright*bursts
            pan = .17
        elif pitch == 37:
            wave_data = (.45*np.sin(TAU*820*t)+.20*np.sin(TAU*1780*t))*np.exp(-t/.014)
            pan = -.30
        else:
            decay = .014 if pitch == 42 else .052
            metal = sum(np.sin(TAU*f*t) for f in (3870, 5210, 7180))/3
            wave_data = (.32*bright + .08*metal)*np.exp(-t/decay)
            pan = .24 if pitch == 42 else -.18
        wave_data *= np.minimum((len(t)-np.arange(len(t)))/(SR*.012), 1)
    else:
        release = {"bass": .05, "chords": .16, "lead": .14, "pad": .22}[track]
        t = np.arange(round((gate+release)*SR), dtype=np.float64)/SR
        freq = 440 * 2 ** ((pitch-69)/12)
        phase = TAU*freq*t
        if track == "bass":
            wave_data = .88*np.sin(phase)
            wave_data += .19*np.sin(phase*2)*np.exp(-t*8)
            wave_data += .10*np.sin(phase*3)*np.exp(-t*12)
            wave_data *= envelope(t, gate, .006, release) * (.78+.22*np.exp(-t/.045))
        elif track == "chords":
            # A light bell tine at onset, then a mellow harmonic sustain.
            wave_data = np.sin(phase + .42*np.sin(phase*2)*np.exp(-t*8))
            wave_data += .13*np.sin(phase*3)*np.exp(-t*6)
            wave_data *= envelope(t, gate, .007, release)*np.exp(-t*1.7)
        elif track == "lead":
            # FM pluck, with a softly detuned fundamental for width through pan.
            wave_data = .76*np.sin(phase + .70*np.sin(phase*2)*np.exp(-t*10))
            wave_data += .15*np.sin(phase*1.002)
            wave_data *= envelope(t, gate, .008, release)*np.exp(-t*2.0)
        else:
            wave_data = .60*np.sin(phase*.999)+.40*np.sin(phase*1.001)
            wave_data += .09*np.sin(phase*2)
            wave_data *= envelope(t, gate, .23, release)
            wave_data *= .92+.08*np.sin(TAU*.27*t+pitch)
        pan = TRACK_PAN[track]
        if track in ("chords", "pad"):
            pan = np.clip(pan + ((pitch % 5)-2)*.12, -.7, .7)
    velocity = (n["velocity"]/127)**1.35
    return (wave_data * velocity * TRACK_GAIN[track]).astype(np.float32), float(pan)


def place(bus, signal, start, pan, gain=1):
    stop = min(len(bus), start+len(signal))
    if stop <= start:
        return
    length = stop-start
    # Constant-power stereo panning, centered sub bass.
    angle = (pan+1)*np.pi/4
    bus[start:stop, 0] += signal[:length]*(gain*np.cos(angle))
    bus[start:stop, 1] += signal[:length]*(gain*np.sin(angle))


def render(notes, bars, spb, path):
    rng = np.random.default_rng(260907)
    frames = round(bars*4*spb*SR)
    mix = np.zeros((frames, 2), dtype=np.float32)
    bass = np.zeros_like(mix)
    pad = np.zeros_like(mix)
    kick_starts = []
    for n in notes:
        signal, pan_value = synth(n, spb, rng)
        start = round(n["beat"]*spb*SR)
        dest = bass if n["track"] == "bass" else pad if n["track"] == "pad" else mix
        place(dest, signal, start, pan_value)
        if n["track"] == "drums" and n["pitch"] == 36:
            kick_starts.append(start)
        if n["track"] in ("chords", "lead"):
            # Two audible, tempo-related repeats with decreasing gains.
            place(mix, signal, start+round(spb*.75*SR), -pan_value, .18)
            place(mix, signal, start+round(spb*1.5*SR), pan_value, .075)
    duck = np.ones(frames, dtype=np.float32)
    recovery_t = np.arange(round(.21*SR))/SR
    recovery = 1-.33*np.exp(-recovery_t/.062)
    for start in kick_starts:
        end = min(frames, start+len(recovery))
        duck[start:end] = np.minimum(duck[start:end], recovery[:end-start])
    mix += bass*duck[:, None]
    mix += pad*(.65+.35*duck[:, None])
    # Remove tiny DC offsets and fade the score's exact 64/16-bar boundaries.
    mix -= mix.mean(axis=0)
    fade_in = round(.03*SR)
    fade_out = round(.42*SR)
    mix[:fade_in] *= np.linspace(0, 1, fade_in)[:, None]
    mix[-fade_out:] *= np.linspace(1, 0, fade_out)[:, None]
    peak_before = float(np.max(np.abs(mix)))
    rms_before = float(np.sqrt(np.mean(mix.astype(np.float64)**2)))
    # Match rendered examples to an RMS target, preserving at least 1 dB peak
    # headroom. This is static gain, with no limiter or hidden clipping.
    gain = min(.1/max(rms_before, 1e-12), 10**(-1/20)/max(peak_before, 1e-12))
    mix *= gain
    pcm = np.rint(mix*32767).astype("<i2")
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(2)
        wav.setsampwidth(2)
        wav.setframerate(SR)
        wav.writeframes(pcm.tobytes())
    return {"file": path.name, "frames": frames, "sample_rate": SR,
            "seconds": frames/SR, "peak_dbfs": 20*np.log10(float(np.max(np.abs(mix)))),
            "rms_dbfs": 20*np.log10(float(np.sqrt(np.mean(mix.astype(np.float64)**2)))),
            "static_gain": gain}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, default=Path("midnight-circuit"))
    args = parser.parse_args()
    score = json.loads((args.out/"score.json").read_text())
    spb = round(60_000_000/score["bpm"])/1_000_000
    results = []
    jobs = (("notes", 64, "midnight-circuit-preview.wav"),
            ("basic_notes", 16, "comparison-basic-16-bars.wav"),
            ("revised_notes", 16, "comparison-revised-16-bars.wav"))
    for key, bars, filename in jobs:
        result = render(score[key], bars, spb, args.out/filename)
        results.append(result)
        print(json.dumps(result), flush=True)
    (args.out/"render-metrics.json").write_text(json.dumps(results, indent=2)+"\n")


if __name__ == "__main__":
    main()
