from __future__ import print_function
import sys
import scipy.signal
import numpy as np
import matplotlib
matplotlib.use('WebAgg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

def plot(title, figure_number, response, p=None):
    fig = plt.figure(figure_number, dpi=78, figsize=(10,4))

    w, h, lag = response
    w /= np.pi * 2
    index_max_h = np.argmax(h)
    w_at_max = w[index_max_h]
    lag_at_max = lag[index_max_h]
    amp = 20 * np.log10(abs(h))
    if lag[0] == 0:
        lag[0] = np.nan # avoid ugly vertical line at left edge of graph
    
    if p is not None:
        critical_w = 1./p
        real_critical_w = 0
        prev_amp = amp[0]
        for i in range(len(amp)):
            if prev_amp > -3 and amp[i] < -3:
                real_critical_w = w[i]
                break
            if prev_amp < -3 and amp[i] > -3:
                real_critical_w = w[i]
                break
            prev_amp = amp[i]
        if real_critical_w != 0:
            title += ", real critical freq %.3f = %.1f bars" % (real_critical_w, 1./real_critical_w,)
    plt.suptitle(title)

    ax = plt.subplot2grid((1, 2),(0, 0))
    ax.xaxis.set_major_locator(ticker.MultipleLocator(.1))
    ax.xaxis.set_minor_locator(ticker.MultipleLocator(.01))
    ax.set_ylim(max(-40, min(amp)-3), max(3, max(amp)+3))
    if p is not None:
        plt.axvline(critical_w)
        amp_at_p = amp[min(len(amp) - 1, np.searchsorted(w, critical_w))]
        plt.title("Amplitude at period %d is %.2f" % (p, amp_at_p))
    elif index_max_h > 0 and index_max_h < len(h) - 1:
        plt.title("Max response at frequency %.3f = %.1f bars" % (w_at_max, float(1./w_at_max)))
    plt.plot(w, amp, 'b')
    plt.ylabel('Amplitude [dB]')
    plt.xlabel('Frequency')

    ax = plt.subplot2grid((1, 2),(0, 1))
    ax.set_ylim(-.9, 10.1, auto=False)
    ax.xaxis.set_major_locator(ticker.MultipleLocator(.1))
    ax.xaxis.set_minor_locator(ticker.MultipleLocator(.01))
    plt.axhline(0)
    if index_max_h > 0 and index_max_h < len(h) - 1:
        plt.axvline(w_at_max)
        plt.title("Lag at max response %.2f" % (lag_at_max,))
    elif p is not None:
        plt.axvline(critical_w)
        lag_at_p = lag[min(len(lag) - 1, np.searchsorted(w, critical_w))]
        plt.title("Lag at period %d is %.2f" % (p, lag_at_p,))
    plt.plot(w, lag, 'r')
    plt.ylabel('Lag samples')
    plt.xlabel('Frequency')

    plt.tight_layout(pad=2)

def ema(period, unused=None):
    title = "EMA(%d)" % period
    alpha = 2. / (period + 1)
    numerator = [alpha] # coeffs for data[0], data[1], data[2], ...
    denominator = [1, -(1 - alpha)] # first elem always 1, then -coeffs for filt[1], filt[2], ...
    return title, numerator, denominator, period

def sma(period, unused=None):
    title = "Simple moving average with Period = %d" % period
    numerator = [1./period for _ in range(period)]
    denominator = [1]
    return title, numerator, denominator, period

def hp2(period, unused=None):
    title = "Two-pole highpass filter with Period = %d" % period
    f = 0.707106781 * 3.141592654 / period
    alpha = 1 + (np.sin(f) - 1) / np.cos(f)
    b = 1 - alpha / 2
    c = 1 - alpha
    numerator = [b*b, -2*b*b, b*b]
    denominator = [1, -2*c, c*c]
    return title, numerator, denominator, period

def corrected_hp2(period, unused=None):
    title = "Corrected two-pole highpass filter with Period = %d" % period
    f = 2 * 0.707106781 * 3.141592654 / period
    alpha = 1 + (np.sin(f) - 1) / np.cos(f)
    b = 1 - alpha / 2
    c = 1 - alpha
    numerator = [b*b, -2*b*b, b*b]
    denominator = [1, -2*c, c*c]
    return title, numerator, denominator, period

def butterworth_hp2(period, unused=None):
    title = "Two-pole butterworth highpass filter with Period = %d" % period
    t = np.tan(np.pi / period)
    k = 1 / (1 + t*np.sqrt(2) + t*t)
    numerator = [k, -2*k, k]
    denominator = [1, k*(2*t*t - 2), k*(1 - t*np.sqrt(2) + t*t)]
    return title, numerator, denominator, period

def smooth(period, unused=None):
    title = "Super smoother with Period = %d" % period
    f = 1.414213562 * 3.141592654 / period
    a = np.exp(-f)
    c2 = 2 * a * np.cos(f)
    c3 = -(a**2)
    c1 = 1 - c2 - c3
    numerator = [c1/2, c1/2]
    denominator = [1, -c2, -c3]
    return title, numerator, denominator, period

def roof(hp, lp):
    title = "Roofing filter with HP = %d and LP = %d" % (hp, lp)
    _, numeratorHP, denominatorHP, _ = hp2(hp)
    _, numeratorSM, denominatorSM, _ = smooth(lp)
    numerator = np.convolve(numeratorHP, numeratorSM)
    denominator = np.convolve(denominatorHP, denominatorSM)
    return title, numerator, denominator, None

def corrected_roof(hp, lp):
    title = "Corrected roofing filter with HP = %d and LP = %d" % (hp, lp)
    _, numeratorHP, denominatorHP, _ = corrected_hp2(hp)
    _, numeratorSM, denominatorSM, _ = smooth(lp)
    numerator = np.convolve(numeratorHP, numeratorSM)
    denominator = np.convolve(denominatorHP, denominatorSM)
    return title, numerator, denominator, None

def butterworth_roof(hp, lp):
    title = "Butterworth Roofing filter with HP = %d and LP = %d" % (hp, lp)
    _, numeratorHP, denominatorHP, _ = butterworth_hp2(hp)
    _, numeratorSM, denominatorSM, _ = smooth(lp)
    numerator = np.convolve(numeratorHP, numeratorSM)
    denominator = np.convolve(denominatorHP, denominatorSM)
    return title, numerator, denominator, None

def calc_and_plot(filter_type, figure_number, p, q):
    if p is None:
        print("== Missing period for", filter_type)
        return
    title, numerator, denominator, p = globals()[filter_type](p, q)
    w, h = scipy.signal.freqz(b=numerator, a=denominator)
    w, lag = scipy.signal.group_delay((numerator, denominator))
    response = (w, h, lag)
    plot(title, figure_number, response, p)
    return figure_number + 1

if __name__ == "__main__":
    figure_number = 1
    filter_type, p_, q_ = [None] * 3
    for f in sys.argv[1:]:
        try:
            val = int(f)
            if p_ is None: p_ = val
            elif q_ is None: q_ = val
        except:
            if filter_type is not None:
                figure_number = calc_and_plot(filter_type, figure_number, p_, q_)
            filter_type = f
            p_, q_ = None, None
    figure_number = calc_and_plot(filter_type, figure_number, p_, q_)
    plt.show()