# encoding:utf-8

"""
@ date:             2020-09-16
@ author:           jingxian
@ illustration:     Pre-processing
"""
from matplotlib import pyplot as plt
from numpy import array
from pandas import read_csv
from scipy.signal import butter, filtfilt, sosfiltfilt


def Butterworth_for_ECG_PreProcess(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")

def Butterworth_for_BCG_PreProcess(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":  # 低通滤波处理
        sos = butter(order, low_cut / (sample_rate * 0.5), btype='lowpass', output='sos')
        return sosfiltfilt(sos, array(data))
    elif type == "bandpass":  # 带通滤波处理
        low = low_cut / (sample_rate * 0.5)
        high = high_cut / (sample_rate * 0.5)
        sos = butter(order, [low, high], btype='bandpass', output='sos')
        return sosfiltfilt(sos, array(data))
    elif type == "highpass":  # 高通滤波处理
        sos = butter(order, high_cut / (sample_rate * 0.5), btype='highpass', output='sos')
        return sosfiltfilt(sos, array(data))
    else:  # 警告,滤波器类型必须有
        raise ValueError("Please choose a type of fliter")

def data_preprocess_for_label_check(data, n, f1, f2, fs):

    f1 = f1 / (fs / 2.0)
    f2 = f2 / (fs / 2.0)
    b, a = butter(n, [f1, f2], btype="bandpass")
    data = array(filtfilt(b, a, data))
    return data