# Code reproduced from
# http://www.scipy.org/Cookbook/ButterworthBandpass

try:
    from scipy.signal import butter, lfilter

    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False


def butter_bandpass(lowcut, highcut, fs, order=5):
    if not SCIPY_AVAILABLE:
        raise ImportError("scipy is required for butter_bandpass. Install with: pip install scipy")
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype="band")
    return b, a


def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    if not SCIPY_AVAILABLE:
        raise ImportError(
            "scipy is required for butter_bandpass_filter. Install with: pip install scipy"
        )
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.signal import freqz

    # Sample rate and desired cutoff frequencies (in Hz).
    fs = 5000.0
    lowcut = 500.0
    highcut = 1250.0

    # Plot the frequency response for a few different orders.
    plt.figure(1)
    plt.clf()
    for order in [3, 6, 9]:
        b, a = butter_bandpass(lowcut, highcut, fs, order=order)
        w, h = freqz(b, a, worN=2000)
        plt.plot((fs * 0.5 / np.pi) * w, abs(h), label="order = %d" % order)

    plt.plot([0, 0.5 * fs], [np.sqrt(0.5), np.sqrt(0.5)], "--", label="sqrt(0.5)")
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Gain")
    plt.grid(True)
    plt.legend(loc="best")

    # Filter a noisy signal.
    T = 0.05
    nsamples = T * fs
    t = np.linspace(0, T, nsamples, endpoint=False)
    a = 0.02
    f0 = 600.0
    x = 0.1 * np.sin(2 * np.pi * 1.2 * np.sqrt(t))
    x += 0.01 * np.cos(2 * np.pi * 312 * t + 0.1)
    x += a * np.cos(2 * np.pi * f0 * t + 0.11)
    x += 0.03 * np.cos(2 * np.pi * 2000 * t)
    plt.figure(2)
    plt.clf()
    plt.plot(t, x, label="Noisy signal")

    y = butter_bandpass_filter(x, lowcut, highcut, fs, order=6)
    plt.plot(t, y, label="Filtered signal (%g Hz)" % f0)
    plt.xlabel("time (seconds)")
    plt.hlines([-a, a], 0, T, linestyles="--")
    plt.grid(True)
    plt.axis("tight")
    plt.legend(loc="upper left")

    plt.show()
