# encoding:utf-8

"""
@ date:             2020-09-16
@ author:           jingxian
@ illustration:     Pre-processing
"""

from scipy.signal import butter, lfilter, filtfilt

def Butterworth(data, sample_rate, type, low_cut=0.0, high_cut=0.0, order=10):
    """
    :param type:      Type of Butter. filter, lowpass, bandpass, ...
    :param lowcut:    Low cutoff frequency
    :param highcut:   High cutoff frequency
    :param order:     Order of filter
    :return:          Signal after filtering
    """
    if type == "lowpass":  # 低通滤波处理
        b, a = butter(order, low_cut / (sample_rate * 0.5), btype='lowpass', output='ba')
        return filtfilt(b, a, data)
    elif type == "bandpass":  # 带通滤波处理
        low = low_cut / (sample_rate * 0.5)
        high = high_cut / (sample_rate * 0.5)
        b, a = butter(order, [low, high], btype='bandpass', output='ba')
        return filtfilt(b, a, data)
    elif type == "highpass":  # 高通滤波处理
        b, a = butter(order, high_cut / (sample_rate * 0.5), btype='highpass', output='ba')
        return filtfilt(b, a, data)
    else:  # 警告,滤波器类型必须有
        raise ValueError("Please choose a type of fliter")