import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

# optionally, use an alternate style for plotting.
# this reconfigures colors, fonts, padding, etc.
# i personally prefer the look of ggplot, but the contrast is generally worse.
plt.style.use('ggplot')

# create log-spaced points from 20 to 20480 Hz.
# we'll compute our y values using these later.
# increasing precision will result in a smoother plot.
precision = 4096
xs = np.arange(0, precision) / precision
xs = 20 * 1024**xs

# decide on a sample-rate and a center-frequency (the -3dB point).
fs = 44100
fc = 1000


def response_setup(ax, ymin=-24, ymax=24, major_ticks_every=6, minor_ticks_split=2):
    """configure axes for magnitude plotting within the range of human hearing."""
    ax.set_xlim(20, 20000)
    ax.set_ylim(ymin, ymax)
    ax.grid(True, 'both') # enable major and minor gridlines on both axes
    ax.set_xlabel('frequency (Hz)')
    ax.set_ylabel('magnitude (dB)')
    # manually set y tick positions.
    # this is usually better than relying on matplotlib to choose good values.
    from matplotlib import ticker
    ax.set_yticks(tuple(range(ymin, ymax + 1, major_ticks_every)))
    minor_ticker = ticker.AutoMinorLocator(minor_ticks_split)
    ax.yaxis.set_minor_locator(minor_ticker)


def modified_bilinear(b, a, w0):
    """
    the modified bilinear transform defers specifying the center frequency
    of an analog filter to the conversion to a digital filter:
    this s-plane to z-plane transformation.
    it also compensates for the frequency-warping effect.

    b and a are the normalized numerator and denominator coefficients (respectively)
    of the polynomial in the s-plane, and w0 is the angular center frequency.
    effectively, w0 = 2 * pi * fc / fs.

    the modified bilinear transform is specified by RBJ [1] as:
        s := (1 - z^-1) / (1 + z^-1) / tan(w0 / 2)
    as opposed to the traditional bilinear transform:
        s := (1 - z^-1) / (1 + z^-1)

    [1]: http://musicdsp.org/files/Audio-EQ-Cookbook.txt
    """

    # in my opinion, this is a more convenient way of creating digital filters.
    # scipy doesn't implement the modified bilinear transform,
    # so i'll just write the expanded form for second-order filters here.
    # TODO: include the equation that expanded/simplified to this

    # use padding to allow for first-order filters.
    # untested
    if len(b) == 2:
        b = (b[0], b[1], 0)
    if len(a) == 2:
        a = (a[0], a[1], 0)

    assert len(b) == 3 and len(a) == 3, "unsupported order of filter"

    cw = np.cos(w0)
    sw = np.sin(w0)
    zb = np.array((
           b[0]*(1 - cw) + b[2]*(1 + cw) - b[1]*sw,
        2*(b[0]*(1 - cw) - b[2]*(1 + cw)),
           b[0]*(1 - cw) + b[2]*(1 + cw) + b[1]*sw,
    ))
    za = np.array((
           a[0]*(1 - cw) + a[2]*(1 + cw) - a[1]*sw,
        2*(a[0]*(1 - cw) - a[2]*(1 + cw)),
           a[0]*(1 - cw) + a[2]*(1 + cw) + a[1]*sw,
    ))
    return zb, za


def digital_magnitude(xs, cascade, fc, fs, decibels=True):
    """
    compute the digital magnitude response
    of an analog SOS filter (cascade)
    at the points given at xs
    and the center frequency fc
    for the sampling rate of fs.
    """

    to_magnitude_decibels = lambda x: np.real(np.log10(np.abs(x)**2) * 10)

    # compute angular frequencies
    ws = 2 * np.pi * xs / fs
    w0 = 2 * np.pi * fc / fs

    digital_cascade = [modified_bilinear(*ba, w0=w0) for ba in cascade]

    # we don't need the first result freqz returns
    # because it'll just be the worN we're passing.
    responses = [sig.freqz(*ba, worN=ws)[1] for ba in digital_cascade]

    if decibels:
        mags = [to_magnitude_decibels(response) for response in responses]
        mags = np.array(mags)
        composite = np.sum(mags, axis=0)
    else:
        # untested
        mags = [np.abs(response) for response in responses]
        mags = np.array(mags)
        composite = np.prod(mags, axis=0)

    # we could also compute phase using (untested):
    # -np.arctan2(np.imag(response), np.real(response))

    return composite, mags


# compute the butterworth coefficients for later.
# we'll request them in SOS format: second-order sections.
# that means we'll get an array of second-order filters
# that combine to produce one higher-order filter.
# note that we specify 1 as the frequency;  we'll give the actual frequency
# later, when we transform it to a digital filter.
butterworth_2 = sig.butter(N=2, Wn=1, analog=True, output='sos')
butterworth_6 = sig.butter(N=6, Wn=1, analog=True, output='sos')
# scipy returns the b and a coefficients in a flat array,
# but we need them separate.
butterworth_2 = butterworth_2.reshape(-1, 2, 3)
butterworth_6 = butterworth_6.reshape(-1, 2, 3)
butterworth_2_times3 = butterworth_2.repeat(3, axis=0)


# set up our first plot.
fig, ax = plt.subplots()
response_setup(ax, -30, 3)
ax.set_title("comparison of digital butterworth filters (fc=1kHz, fs=44.1kHz)")

# compute the magnitude of the filters at our x positions.
# we don't want the individual magnitudes here,
# so assign them to the dummy variable _.
ys1, _ = digital_magnitude(xs, butterworth_2, fc, fs)
ys2, _ = digital_magnitude(xs, butterworth_2_times3, fc, fs)
ys3, _ = digital_magnitude(xs, butterworth_6, fc, fs)

# our x axis is frequency, so it should be spaced logarithmically.
# our y axis is decibels, which is already a logarithmic unit.
# thus, semilogx is the plotting function to use here.
ax.semilogx(xs, ys1, label="second order")
ax.semilogx(xs, ys2, label="second order, three times")
ax.semilogx(xs, ys3, label="sixth order")

ax.legend() # display the legend, given by labels when plotting.
fig.show()
fig.savefig('butterworth_comparison.png')


# set up our second plot.
fig, ax = plt.subplots()
response_setup(ax, -12, 12)
ax.set_title("digital butterworth decomposition (fc=1kHz, fs=44.1kHz)")

cascade = butterworth_6
composite, mags = digital_magnitude(xs, cascade, fc, fs)
for i, mag in enumerate(mags):
    # Q has an inverse relation to the 1st-degree coefficient
    # of the analog filter's denominator, so we can compute it that way.
    Q = 1 / cascade[i][1][1]
    ax.semilogx(xs, mag, label="second order (Q = {:.6f})".format(Q))
ax.semilogx(xs, composite, label="composite")

ax.legend()
fig.show()
fig.savefig('butterworth_composite.png')


# as a bonus, here's how the modified bilinear transform is especially useful:
# we can specify all the basic second-order analog filter types
# with minimal code!

lowpass2    = lambda A, Q: ((1,     0, 0),
                            (1,   1/Q, 1))
highpass2   = lambda A, Q: ((0,     0, 1),
                            (1,   1/Q, 1))

allpass2    = lambda A, Q: ((1,   1/Q, 1),
                            (1,   1/Q, 1))

# peaking filters are also (perhaps for the better) known as bell filters.
peaking2    = lambda A, Q: ((1,   A/Q, 1),
                            (1, 1/A/Q, 1))
# TODO: add classical "analog" peaking filter, where bandwidth and amplitude are interlinked

# there are two common bandpass variants:
# one where the bandwidth and amplitude are interlinked,
bandpass2a  = lambda A, Q: ((0,  -A/Q, 0),
                            (1, 1/A/Q, 1))
# and one where they aren't.
bandpass2b  = lambda A, Q: ((0,-A*A/Q, 0),
                            (1,   1/Q, 1))

notch2      = lambda A, Q: ((1,     0, 1),
                            (1,   1/Q, 1))

lowshelf2   = lambda A, Q: ((  A,   np.sqrt(A)/Q, 1),
                            (1/A, 1/np.sqrt(A)/Q, 1))

highshelf2  = lambda A, Q: ((1,   np.sqrt(A)/Q,   A),
                            (1, 1/np.sqrt(A)/Q, 1/A))


# here's an example of how to plot these.
# we specify the gain in decibels and bandwidth in octaves,
# then convert these to A and Q.
fig, ax = plt.subplots()
response_setup(ax, -30, 30)

invsqrt2 = 1 / np.sqrt(2) # for bandwidth <-> Q calculation
# note that A is squared, so the decibel math here is a little different.
designer = lambda f, decibels, bandwidth: f(10**(decibels / 40.), invsqrt2 / bandwidth)

ax.set_title("example of interlinked bandpass amplitude/bandwidth")
cascade = [
    designer(bandpass2a, 0, 2),
    designer(bandpass2a, -18, 2),
    designer(bandpass2a, 18, 2),
]

composite, mags = digital_magnitude(xs, cascade, fc, fs)
for mag in mags:
    ax.semilogx(xs, mag)
#ax.semilogx(xs, composite)

#ax.legend()
fig.show()
fig.savefig('filter_example.png')