#!/usr/bin/env nix-script
#!>python
#! python3 | matplotlib numpy scipy pyaudio

##  Real-time sound spectrum analysis.

import pyaudio
import scipy.fftpack
import scipy.interpolate
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# pyAudio settings
settings = {
    "format": pyaudio.paInt16,
    "channels": 2,
    "rate": 48000,
    "input": True,
    "frames_per_buffer": 1024
}
audio = pyaudio.PyAudio()
stream = audio.open(**settings)

# plt configurations
plt.style.use('seaborn-v0_8-dark')
fig, (spectrum_log, spectrum, wave) = plt.subplots(nrows=3, ncols=1)
fig.dpi = 200


def init():
    plt.subplots_adjust(hspace=0.5, left=0.1)
    global curve1, curve2, curve3

    # Plot settings
    spectrum_log.set_ylabel("a / dB")
    spectrum_log.set_xlabel("ƒ / Hz")
    spectrum_log.ticklabel_format(style="sci", scilimits=(0, 0),
                                  axis="x", useMathText=True)
    spectrum_log.grid(True)
    curve1, = spectrum_log.plot([0, 8 * 10 ** 3], [0, 10 ** 2], "b-")

    spectrum.yaxis.set_visible(False)
    spectrum_log.set_ylabel("a")
    spectrum.set_xlabel("ƒ / Hz")
    spectrum.ticklabel_format(style="sci", scilimits=(0, 0),
                              axis="x", useMathText=True)
    spectrum.grid(True)
    curve2, = spectrum.plot([0, 8 * 10 ** 3], [0, 10 ** 5])

    wave.set_xlim(0, 8 * 10 ** 3)
    wave.xaxis.set_visible(False)
    wave.yaxis.set_visible(False)
    curve3, = wave.plot([0, 8 * 10 ** 3], [-2 * 10 ** 3, 2 * 10 ** 3], "g-")

    return curve1, curve2, curve3


def update(frame):
    # Acquire data from microphone
    block = np.array(np.frombuffer(
            stream.read(settings["frames_per_buffer"],
                        exception_on_overflow=False),
            dtype="int16"))

    # FFT of registred block
    fftx = scipy.fftpack.rfftfreq(
        settings["frames_per_buffer"],
        1 / settings["rate"])

    # Calculate log version
    ffty = abs(scipy.fftpack.fft(block)[0:len(fftx)])
    ffty_log = 10 * np.log10(ffty)

    # Data interpolation
    wave = scipy.interpolate.interp1d(fftx, block[0:len(fftx)])(fftx)

    # Update plot data
    curve1.set_xdata(fftx)
    curve2.set_xdata(fftx)
    curve3.set_xdata(fftx)

    curve1.set_ydata(ffty_log)
    curve2.set_ydata(ffty)
    curve3.set_ydata(wave)

    return curve1, curve2, curve3

ani = FuncAnimation(fig, func=update, init_func=init,
                    interval=10,
                    blit=True, cache_frame_data=False)
plt.show()