"""Draw the spectrum of white noise against the spectrum of pink noise.

One figure. Both signals are generated by the page's own equations rather
than by a library: equation (1) is the xorshift step, equation (2) turns the
state into a sample, and equations (3) and (4) are the three-section pink
filter. The spectrum is an average of periodograms over many windows,
because one window of a random signal is itself random and reads as grass
rather than as a slope.

The script fits a straight line to the pink spectrum between 30 and 12000
hertz and prints the slope, which is the number the caption quotes. Nothing
here asserts the slope. It is measured from the signal the equations
produce, so a wrong coefficient shows up as a wrong slope.

The colours become custom properties on the way out, so the drawing follows
the page between light and dark. See the note on `save`.

Writes spectrum.svg beside this file.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
matplotlib.rcParams["svg.fonttype"] = "path"
import matplotlib.pyplot as plt

TEAL, INDIGO, ROSE = "#00706B", "#3A49A6", "#A93356"
INK, MUTED, GRID = "#0D1620", "#5A6874", "#D3DBE3"

PROPERTY = {TEAL: "--teal", INDIGO: "--indigo", ROSE: "--rose",
            INK: "--ink", MUTED: "--muted", GRID: "--line"}

FS = 48000.0
MASK = 0xFFFFFFFF


def white(n, seed=2463534242):
    """Equations (1) and (2), one sample at a time.

    The mask after every shift is what a 32-bit register does for free and
    Python does not, since its integers have no width. Without it the state
    grows without bound and the period is lost.
    """
    s = seed
    out = np.empty(n)
    for i in range(n):
        s ^= (s << 13) & MASK
        s ^= s >> 17
        s ^= (s << 5) & MASK
        out[i] = (s >> 8) / 2.0 ** 23 - 1.0
    return out


def pink(x):
    """Equations (3) and (4): three one-pole sections and their sum."""
    b0 = b1 = b2 = 0.0
    out = np.empty_like(x)
    for i, xi in enumerate(x):
        b0 = 0.99765 * b0 + xi * 0.0990460
        b1 = 0.96300 * b1 + xi * 0.2965164
        b2 = 0.57000 * b2 + xi * 1.0526913
        out[i] = 0.12 * (b0 + b1 + b2 + xi * 0.1848)
    return out


def spectrum(x, size=4096):
    """Average the periodograms of consecutive windows.

    One window of a random signal is itself random, so a single transform
    draws a band of grass about ten decibels wide. Averaging many windows
    narrows that band without moving the line it scatters around.
    """
    win = np.hanning(size)
    count = len(x) // size
    acc = np.zeros(size // 2 + 1)
    for k in range(count):
        seg = x[k * size:(k + 1) * size] * win
        acc += np.abs(np.fft.rfft(seg)) ** 2
    acc /= count
    freq = np.fft.rfftfreq(size, 1.0 / FS)
    return freq[1:], 10.0 * np.log10(acc[1:])


def save(fig, name):
    """Write the figure, then make it the page's own drawing.

    Matplotlib writes every colour into the SVG as a literal, so a figure
    saved the ordinary way carries fixed values and cannot follow the page
    between light and dark. Each literal becomes var(--prop, literal) here,
    which keeps the file readable on its own and lets the page recolour it.
    """
    import re

    fig.tight_layout()
    # transparent, or Matplotlib paints an opaque white rectangle behind
    # everything and the figure reads as a light card on a dark page.
    fig.savefig(name, transparent=True, bbox_inches="tight", pad_inches=0.04)
    text = open(name, encoding="utf-8").read()

    stem = re.sub(r"[^A-Za-z0-9]+", "-", name.rsplit(".", 1)[0]).strip("-")
    text = re.sub(r'(\sid=")([^"]+)"', lambda m: f'{m.group(1)}{stem}-{m.group(2)}"', text)
    text = re.sub(r'((?:xlink:)?href="#)([^"]+)"',
                  lambda m: f'{m.group(1)}{stem}-{m.group(2)}"', text)
    text = re.sub(r'(url\(#)([^)]+)\)', lambda m: f'{m.group(1)}{stem}-{m.group(2)})', text)

    for literal, prop in PROPERTY.items():
        text = re.sub(re.escape(literal), f"var({prop}, {literal.lower()})",
                      text, flags=re.IGNORECASE)

    with open(name, "w", encoding="utf-8") as f:
        f.write(text)
    print(f"{name} written")


# Equation (5), for both sample rates the page quotes.
for rate in (48000.0, 44100.0):
    corners = [-rate * np.log(a) / (2 * np.pi) for a in (0.99765, 0.96300, 0.57000)]
    print(f"corners at {rate:.0f}: " + ", ".join(f"{c:.1f}" for c in corners))

N = 4096 * 240
w = white(N)
p = pink(w)

print(f"white: mean {w.mean():+.4f}, rms {np.sqrt((w ** 2).mean()):.4f}, "
      f"predicted rms {1 / np.sqrt(3):.4f}")

fw, sw = spectrum(w)
fp, sp = spectrum(p)

# Line up the two curves at 1 kHz so the figure compares slopes, not levels.
ref = np.argmin(np.abs(fw - 1000.0))
sw = sw - sw[ref]
sp = sp - sp[ref]

# The slope, fitted on the decades where the three sections are doing their
# work. Below 30 Hz the lowest section has flattened out, and above 12 kHz
# the highest one has.
band = (fp >= 30.0) & (fp <= 12000.0)
slope, _ = np.polyfit(np.log2(fp[band]), sp[band], 1)
print(f"pink slope between 30 Hz and 12 kHz: {slope:.2f} dB per octave")

fig, ax = plt.subplots(figsize=(6.4, 3.5))

ax.plot(fw, sw, lw=0.8, color=MUTED, alpha=0.85, zorder=2)
ax.plot(fp, sp, lw=0.9, color=ROSE, alpha=0.9, zorder=3)

# The ideal three-decibel slope, for the eye to judge the fit against.
guide = np.array([20.0, 20000.0])
ax.plot(guide, -3.0 * np.log2(guide / 1000.0), lw=1.4, color=TEAL,
        linestyle=(0, (5, 3)), zorder=4)

ax.annotate("white", (11000.0, 2.4), ha="center", va="bottom",
            fontsize=10, color=MUTED)
ax.annotate("pink", (70.0, 14.5), ha="left", va="bottom",
            fontsize=10, color=ROSE)
ax.annotate("3 dB per octave", (330.0, 3.6), ha="left", va="top",
            fontsize=9.5, color=TEAL)

ax.set_xscale("log")
ax.set_xlim(20.0, 20000.0)
ax.set_ylim(-20.0, 24.0)
ax.set_xlabel("frequency (Hz)", fontsize=10, color=INK)
ax.set_ylabel("power (dB, 0 at 1 kHz)", fontsize=10, color=INK)
ax.set_xticks([20, 100, 1000, 10000])
ax.set_xticklabels(["20", "100", "1k", "10k"])
ax.set_yticks([-18, -12, -6, 0, 6, 12, 18])
# which="both", or the minor ticks a logarithmic axis adds keep
# Matplotlib's default black and stay black on a dark page.
ax.tick_params(which="both", colors=MUTED, labelsize=9)
ax.grid(True, which="major", lw=0.6, color=GRID, alpha=0.7)
for side in ("top", "right"):
    ax.spines[side].set_visible(False)
for side in ("left", "bottom"):
    ax.spines[side].set_color(GRID)

save(fig, "spectrum.svg")