from scipy import signal
from waveform_analysis import A_weighting
from math import pi
import numpy
import matplotlib.pyplot as plt


# Create our X axis values, the frequencies to test at:
gs = 2 * pi * numpy.geomspace(10, 24000, 1000)


# This is the ideal A-weighting filter:
z, p, k = A_weighting(48000, output='zpk')
_, hA = signal.freqz_zpk(z, p, k, gs, fs=48000)


# And here is our three-stage implementation:
z, p, k = signal.iirfilter(1, [640, 10000], btype='bandpass', analog=False, fs=48000, output='zpk')
_, hB = signal.freqz_zpk(z, p, k * 1.17, gs, fs=48000)

z, p, k = signal.iirfilter(1, 12000, btype='lowpass', analog=False, fs=48000, output='zpk')
_, hL = signal.freqz_zpk(z, p, k, gs, fs=48000)

z, p, k = signal.iirfilter(1, 120, btype='highpass', analog=False, fs=48000, output='zpk')
_, hH = signal.freqz_zpk(z, p, k, gs, fs=48000)


# Lastly, our SOS approximation:
sos = [[ 0.5,  0,   -0.5,  1, -1.062,  0.14],
       [ 0.5,  0.5,  0,    1,  0,      0   ],
       [ 1.0, -1.0,  0,    1, -0.985,  0   ]]
wS, hS = signal.sosfreqz(sos, fs=48000)


# Plot everything:
plt.semilogx(gs, 20 * numpy.log10(hA), label='A')
plt.semilogx(gs, 20 * numpy.log10(hB * hL * hH), label='bp+lp+hp')
plt.semilogx(wS, 20 * numpy.log10(hS), label='sos')

plt.title('Frequency response')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Amplitude [dB]')
plt.ylim(-30, 10)
plt.xlim(20, 24000)
plt.grid(True, color='0.7', linestyle='-', which='major', axis='both')
plt.grid(True, color='0.9', linestyle='-', which='minor', axis='both')
plt.legend()
plt.show()