"""Draw the computed plots: impulse, comb, allpass, density, damping.

Every curve here is calculated, not traced. The comb comes from the transfer
function of one feedback delay line. The damping curve comes from the shelf
pair in src/dsp/MultibandDamping.cpp, evaluated at the constants
src/dsp/FDN.cpp passes it. The impulse response is the one drawing that is
schematic, and its caption on the page says so.

Writes impulse.svg, comb.svg, allpass.svg, density.svg and damping.svg
beside this file.
"""
import sys

# figtheme sits beside this file; its bytecode would land in the page
# bundle and be published with it.
sys.dont_write_bytecode = True

import numpy as np
import matplotlib

matplotlib.use("Agg")
matplotlib.rcParams["svg.fonttype"] = "path"
import matplotlib.pyplot as plt

from figtheme import (TEAL, INDIGO, ROSE, AMBER, INK, MUTED, GRID, LINES,
                      save, stems)

FS = 48000.0

# src/dsp/tables/DelaySets.hpp, kLateHall48k. Samples at 48 kHz.
HALL_LATE = [961, 1259, 2893, 3103, 3367, 4560]


def frame(size, xlabel, ylabel):
    """One plot area, ruled lightly, with the top and right sides taken off."""
    fig, ax = plt.subplots(figsize=size)
    ax.grid(True, color=GRID, lw=0.6, zorder=0)
    ax.set_axisbelow(True)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(MUTED)
    ax.tick_params(which="both", colors=MUTED, labelsize=8)
    ax.set_xlabel(xlabel, color=INK, fontsize=9)
    ax.set_ylabel(ylabel, color=INK, fontsize=9)
    return fig, ax


def finish(fig, name):
    fig.tight_layout(pad=0.4)
    save(fig, name)


def impulse():
    """A room impulse response in three parts, drawn to scale in time.

    The direct sound arrives first and alone. A handful of reflections
    follow it, each one a separate arrival that the ear can still count.
    After about 60 ms the arrivals overlap and the picture becomes a
    decaying wash. The heights are drawn, not measured.
    """
    rng = np.random.default_rng(7)
    fig, ax = frame((6.4, 2.6), "time after the direct sound (ms)",
                    "amplitude")

    ax.vlines(0.0, 0.0, 1.0, color=INK, lw=1.8, zorder=3)
    ax.annotate("direct sound", xy=(0.0, 1.0), xytext=(14, 0.92),
                color=INK, fontsize=8.5,
                arrowprops=dict(arrowstyle="-", color=MUTED, lw=0.8))

    early_t = np.array([7.5, 11.2, 15.9, 19.4, 24.8, 29.1, 34.6, 41.0,
                        47.3, 53.8])
    early_a = np.array([0.52, 0.44, 0.47, 0.35, 0.38, 0.29, 0.31, 0.24,
                        0.22, 0.19])
    ax.vlines(early_t, 0.0, early_a, color=TEAL, lw=1.4, zorder=3)
    ax.annotate("early reflections", xy=(24.8, 0.38), xytext=(30, 0.66),
                color=TEAL, fontsize=8.5,
                arrowprops=dict(arrowstyle="-", color=TEAL, lw=0.8))

    # The tail: arrivals too dense to count, so the band is filled and only
    # a sample of the individual arrivals is drawn inside it.
    t = np.linspace(60.0, 400.0, 260)
    env = 0.30 * np.exp(-t / 130.0)
    ax.fill_between(t, -env, env, color=ROSE, alpha=0.16, lw=0, zorder=2)
    a = env * rng.uniform(-1.0, 1.0, size=t.size)
    stems(ax, t, a, color=ROSE, lw=0.4, alpha=0.7, zorder=3)
    ax.plot(t, env, color=ROSE, lw=1.1, zorder=4)
    ax.plot(t, -env, color=ROSE, lw=1.1, zorder=4)
    ax.annotate("late tail", xy=(190, 0.30 * np.exp(-190 / 130.0)),
                xytext=(210, 0.52), color=ROSE, fontsize=8.5,
                arrowprops=dict(arrowstyle="-", color=ROSE, lw=0.8))

    ax.axhline(0.0, color=MUTED, lw=0.8, zorder=1)
    ax.set_xlim(-12, 400)
    ax.set_ylim(-0.45, 1.12)
    finish(fig, "impulse.svg")


def comb_response(delay_samples, g, freqs):
    """Magnitude of 1 / (1 - g z^-M) at the given frequencies, in decibels."""
    z_inv = np.exp(-2j * np.pi * freqs * delay_samples / FS)
    return 20.0 * np.log10(np.abs(1.0 / (1.0 - g * z_inv)))


def comb():
    """One delay line rings at evenly spaced frequencies. Six do not.

    The upper panel is the magnitude response of a single feedback delay
    line of 961 samples. The lower panel marks where each of the six Hall
    lines has a peak, so the combined spacing can be read off directly.
    """
    freqs = np.linspace(20.0, 500.0, 2400)
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6.4, 3.8),
                                   height_ratios=[3, 1], sharex=True)

    for ax in (ax1, ax2):
        ax.grid(True, color=GRID, lw=0.6, zorder=0)
        ax.set_axisbelow(True)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color(MUTED)
        ax.tick_params(which="both", colors=MUTED, labelsize=8)

    ax1.plot(freqs, comb_response(HALL_LATE[0], 0.9, freqs),
             color=INDIGO, lw=1.2)
    ax1.set_ylabel("gain (dB)", color=INK, fontsize=9)
    ax1.set_ylim(-10, 22)
    spacing = FS / HALL_LATE[0]
    ax1.annotate(f"peaks every {spacing:.1f} Hz",
                 xy=(3 * spacing, 19.6), xytext=(185, 12.0),
                 color=INDIGO, fontsize=8.5,
                 arrowprops=dict(arrowstyle="-", color=INDIGO, lw=0.8))

    for row, (delay, colour) in enumerate(zip(HALL_LATE, LINES)):
        peaks = np.arange(1, int(500.0 * delay / FS) + 1) * FS / delay
        ax2.vlines(peaks, row, row + 0.8, color=colour, lw=0.9)
    ax2.set_yticks(np.arange(6) + 0.4)
    ax2.set_yticklabels([str(d) for d in HALL_LATE], fontsize=7.5)
    ax2.set_ylabel("line length\n(samples)", color=INK, fontsize=9)
    ax2.set_xlabel("frequency (Hz)", color=INK, fontsize=9)
    ax2.set_xlim(20, 500)

    fig.tight_layout(pad=0.4)
    save(fig, "comb.svg")


def allpass_chain(x, ms, g=0.70):
    """The eight sections of src/dsp/AllpassDiffuser.hpp, in series.

    v[n] = x[n] + g v[n-M], y[n] = v[n-M] - g v[n], which is equation (4)
    on the page. The left channel's delays, in milliseconds.
    """
    for m in ms:
        delay = max(1, int(round(m * 0.001 * FS)))
        buf = np.zeros(delay)
        w = 0
        out = np.empty_like(x)
        for n, sample in enumerate(x):
            d = buf[w]
            v = sample + g * d
            buf[w] = v
            w = (w + 1) % delay
            out[n] = d - g * v
        x = out
    return x


def late_field(x, delays, gain=0.985):
    """Six lines and the even Householder, at one gain for every line.

    The gain is uniform and close to one because this measures how fast
    arrivals accumulate, not how the tail decays. Damping would only take
    energy out of the count.
    """
    n_lines = len(delays)
    bufs = [np.zeros(d) for d in delays]
    ptr = [0] * n_lines
    out = np.zeros_like(x)
    v = np.zeros(n_lines)
    for n, sample in enumerate(x):
        for i in range(n_lines):
            v[i] = bufs[i][ptr[i]]
        out[n] = v.sum() / np.sqrt(n_lines)
        mixed = (v - (2.0 / n_lines) * v.sum()) * gain
        for i in range(n_lines):
            bufs[i][ptr[i]] = mixed[i] + sample * 0.5
            ptr[i] = (ptr[i] + 1) % delays[i]
    return out


def allpass():
    """One click in, and what comes out of one section and of eight.

    The point the prose cannot make on its own: an allpass leaves every
    level alone and still changes the sound completely, because it changes
    when each frequency leaves. A click is every frequency arriving at once.
    Afterwards they are spread, and the peak falls even though the energy
    does not.

    The counts are for the 45 ms the figure shows. Over 400 ms the chain has
    produced 10911 arrivals and its energy comes back to 1.0000 exactly,
    which is the same statement as |H| = 1 read in the time domain.
    """
    ms_left = [1.7, 2.3, 3.1, 4.1, 5.3, 6.7, 8.1, 9.3]
    n = int(0.045 * FS)
    impulse = np.zeros(n)
    impulse[0] = 1.0
    one = allpass_chain(impulse, ms_left[:1])
    eight = allpass_chain(impulse, ms_left)
    t = np.arange(n) / FS * 1000.0

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6.4, 3.6), sharex=True)
    for ax in (ax1, ax2):
        ax.grid(True, color=GRID, lw=0.6, zorder=0)
        ax.set_axisbelow(True)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color(MUTED)
        ax.tick_params(which="both", colors=MUTED, labelsize=8)
        ax.axhline(0.0, color=MUTED, lw=0.8)

    # The panels do not share a vertical scale. At the top one the lower
    # panel is a flat line, which hides the thing it is there to show.
    ax1.set_ylim(-1.05, 1.05)
    ax2.set_ylim(-0.12, 0.12)

    ax1.vlines(0, 0, 1, color=INK, lw=1.8)
    stems(ax1, t, one, color=TEAL, lw=0.9)
    ax1.text(0.6, 0.86, "the click going in", color=INK, fontsize=8.5)
    ax1.text(12, 0.86, "one section: 25 arrivals, tallest 0.70",
             color=TEAL, fontsize=8.5)
    ax1.set_ylabel("amplitude", color=INK, fontsize=9)

    stems(ax2, t, eight, color=ROSE, lw=0.7)
    ax2.text(12, 0.098, "eight in series: 1612 arrivals in this window, tallest 0.08",
             color=ROSE, fontsize=8.5)
    ax2.text(12, 0.077, "note the scale: this panel is magnified nine times",
             color=MUTED, fontsize=7.5)
    ax2.set_ylabel("amplitude", color=INK, fontsize=9)
    ax2.set_xlabel("time (ms)", color=INK, fontsize=9)
    ax2.set_xlim(-1, 45)

    fig.tight_layout(pad=0.4)
    save(fig, "allpass.svg")


# kRockstar's late field, docs/airwindows-notes.md, transcribed from his
# header in 24 kHz ticks. Six stages of six, read in order.
AIRWINDOWS_36 = [17, 204, 173, 29, 151, 372,
                 1090, 132, 71, 215, 1558, 61,
                 32, 71, 938, 1695, 129, 11,
                 131, 1682, 103, 107, 109, 113,
                 800, 60, 40, 157, 149, 1580,
                 157, 1645, 28, 6, 179, 600]
TICK = 24000.0


def cascade(x, delays=None, gain=3 ** 6 * 0.0013425):
    """Six stages of six lines, each applying the even Householder.

    Structure only. No undersampling reconstruction, no damping, no room
    tone, and none of the wandering regeneration, because this measures how
    fast arrivals accumulate and each of those only removes some.

    The default gain is his own arithmetic. His matrix is 3(I - (2/6)J), so
    six stages multiply by 3**6 = 729, and his reg6n is 0.0013425. The
    product is 0.979, just under unity, which is what says the structure
    modelled here is the one the constants were chosen for.
    """
    delays = delays or AIRWINDOWS_36
    stage = [[np.zeros(delays[k * 6 + i]) for i in range(6)] for k in range(6)]
    ptr = [[0] * 6 for _ in range(6)]
    out = np.zeros_like(x)
    for n, sample in enumerate(x):
        v = [np.array([stage[k][i][ptr[k][i]] for i in range(6)])
             for k in range(6)]
        out[n] = v[5].sum() / np.sqrt(6)
        for k in range(6):
            mixed = v[k] - (2.0 / 6.0) * v[k].sum()
            last = k == 5
            dest = (k + 1) % 6
            for i in range(6):
                stage[dest][i][ptr[dest][i]] = (
                    mixed[i] * (gain if last else 1.0)
                    + (sample * 0.5 if last else 0.0))
        for k in range(6):
            for i in range(6):
                ptr[k][i] = (ptr[k][i] + 1) % delays[k * 6 + i]
    return out


def density():
    """How fast separate arrivals accumulate, with and without diffusion.

    An arrival is counted when a sample exceeds one thousandth of the
    response's peak. The window slides in 10 ms steps, and the count is
    scaled to a second so the number is comparable across windows.
    """
    seconds = 0.26
    ms_left = [1.7, 2.3, 3.1, 4.1, 5.3, 6.7, 8.1, 9.3]

    impulse = np.zeros(int(seconds * FS))
    impulse[0] = 1.0
    bare = late_field(impulse, HALL_LATE)
    diffused = late_field(allpass_chain(impulse, ms_left), HALL_LATE)

    # His core runs at half the rate, so it is measured on its own clock and
    # the counts are put on the same axis of real time afterwards.
    tick_impulse = np.zeros(int(seconds * TICK))
    tick_impulse[0] = 1.0
    airwindows = cascade(tick_impulse)

    def per_second(y, rate):
        window = int(0.010 * rate)
        mid, count = [], []
        floor = np.abs(y).max() / 1000.0
        for start in range(0, len(y) - window, window):
            mid.append((start + window / 2) / rate * 1000.0)
            count.append((np.abs(y[start:start + window]) > floor).sum()
                         / (window / rate))
        return mid, count

    centres, bare_c = per_second(bare, FS)
    _, diff_c = per_second(diffused, FS)
    aw_centres, aw_c = per_second(airwindows, TICK)
    counts = (bare_c, diff_c)

    # A window before the first arrival holds nothing, and a logarithmic
    # axis cannot draw nothing. Masked, so each line starts where its first
    # arrival does instead of falling out of the bottom of the frame.
    def seen(values):
        a = np.array(values, dtype=float)
        a[a <= 0] = np.nan
        return a

    fig, ax = frame((6.4, 3.2), "time after the input (ms)",
                    "separate arrivals per second")
    ax.plot(centres, seen(counts[0]), color=AMBER, lw=1.5,
            label="SD-Reverb, six lines alone")
    ax.plot(aw_centres, seen(aw_c), color=INDIGO, lw=1.5,
            label="Airwindows, thirty-six lines in six stages")
    ax.plot(centres, seen(counts[1]), color=TEAL, lw=1.5,
            label="SD-Reverb, six lines behind the eight allpass sections")
    for ceiling, note in ((FS, "one per sample at 48 kHz"),
                          (TICK, "one per tick at 24 kHz")):
        ax.axhline(ceiling, color=MUTED, lw=0.9, ls=(0, (4, 3)))
        ax.text(248, ceiling * 1.1, note, color=MUTED, fontsize=7.5, ha="right")
    ax.set_yscale("log")
    ax.set_xlim(0, 250)
    ax.set_ylim(20, FS * 6)
    leg = ax.legend(frameon=False, fontsize=8.5, loc="lower right")
    for text in leg.get_texts():
        text.set_color(INK)
    finish(fig, "density.svg")


def one_pole(hz, freqs):
    """The lowpass whose state the shelf pair keeps, as a frequency response.

    MultibandDamping.cpp builds the coefficient as 1 - exp(-2 pi f / fs) and
    runs it as lp += c * (x - lp), which is the filter c / (1 - (1-c) z^-1).
    """
    c = 1.0 - np.exp(-2.0 * np.pi * hz / FS)
    z_inv = np.exp(-2j * np.pi * freqs / FS)
    return c / (1.0 - (1.0 - c) * z_inv)


def loop_gain(freqs, delay, rt60_mid, low_mult, high_mult):
    """The gain one delay line applies per round trip, at each frequency."""
    g_mid = 10.0 ** (-3.0 * delay / (rt60_mid * FS))
    g_low = 10.0 ** (-3.0 * delay / (rt60_mid * low_mult * FS))
    g_high = 10.0 ** (-3.0 * delay / (rt60_mid * high_mult * FS))
    low_rel, high_rel = g_low / g_mid, g_high / g_mid

    shelf_low = 1.0 + (low_rel - 1.0) * one_pole(250.0, freqs)
    shelf_high = 1.0 + (high_rel - 1.0) * (1.0 - one_pole(4000.0, freqs))
    return np.abs(g_mid * shelf_low * shelf_high)


def damping():
    """Decay time against frequency, for the smallest and largest room.

    The high-band multiplier falls from 0.9 at Size 0 to 0.55 at Size 1
    (src/dsp/FDN.cpp, recomputeDamping). The low-band multiplier is 1.3 at
    Tone Tilt 0. The line length is held at 2893 samples so that the curve
    shows the damping alone.
    """
    freqs = np.geomspace(20.0, 20000.0, 700)
    delay, rt60_mid = HALL_LATE[2], 2.4
    fig, ax = frame((6.4, 2.9), "frequency (Hz)", "decay time to -60 dB (s)")

    for high_mult, colour, label in ((0.9, TEAL, "Size 0, high band 0.9"),
                                     (0.55, ROSE, "Size 1, high band 0.55")):
        gain = loop_gain(freqs, delay, rt60_mid, 1.3, high_mult)
        rt60 = -3.0 * delay / (FS * np.log10(gain))
        ax.plot(freqs, rt60, color=colour, lw=1.5, label=label)

    ax.axhline(rt60_mid, color=MUTED, lw=0.9, ls=(0, (4, 3)))
    ax.annotate(f"mid-band target {rt60_mid} s", xy=(700, rt60_mid),
                xytext=(700, rt60_mid + 0.28), color=MUTED, fontsize=8.5)
    ax.set_xscale("log")
    ax.set_xlim(20, 20000)
    ax.set_ylim(0, 3.6)
    leg = ax.legend(frameon=False, fontsize=8.5, loc="lower left")
    for text in leg.get_texts():
        text.set_color(INK)
    finish(fig, "damping.svg")


if __name__ == "__main__":
    impulse()
    comb()
    allpass()
    density()
    damping()
