import numpy as np

import utils

def signal_filter_split(conf, signal_data_raw, signal_fs):
    # 滤波
    # 50Hz陷波滤波器
    # signal_data = utils.butterworth(data=signal_data, _type="bandpass", low_cut=0.5, high_cut=45, order=10, sample_rate=signal_fs)
    print("Applying 50Hz notch filter...")
    signal_data = utils.notch_filter(data=signal_data_raw, notch_freq=50.0, quality_factor=30.0, sample_rate=signal_fs)

    resp_data_0 = utils.butterworth(data=signal_data, _type="lowpass", low_cut=50, order=10, sample_rate=signal_fs)
    resp_fs = conf["resp"]["downsample_fs_1"]
    resp_data_1 = utils.downsample_signal_fast(original_signal=resp_data_0, original_fs=signal_fs, target_fs=resp_fs)
    resp_data_1 = utils.average_filter(raw_data=resp_data_1, sample_rate=resp_fs, window_size_sec=20)
    resp_data_2 = utils.butterworth(data=resp_data_1, _type=conf["resp_filter"]["filter_type"],
                                    low_cut=conf["resp_filter"]["low_cut"],
                                    high_cut=conf["resp_filter"]["high_cut"], order=conf["resp_filter"]["order"],
                                    sample_rate=resp_fs)
    print("Begin plotting signal data...")

    # fig = plt.figure(figsize=(12, 8))
    # # 绘制三个图raw_data、resp_data_1、resp_data_2
    # ax0 = fig.add_subplot(3, 1, 1)
    # ax0.plot(np.linspace(0, len(signal_data) // signal_fs, len(signal_data)), signal_data, color='blue')
    # ax0.set_title('Raw Signal Data')
    # ax1 = fig.add_subplot(3, 1, 2, sharex=ax0)
    # ax1.plot(np.linspace(0, len(resp_data_1) // resp_fs, len(resp_data_1)), resp_data_1, color='orange')
    # ax1.set_title('Resp Data after Average Filtering')
    # ax2 = fig.add_subplot(3, 1, 3, sharex=ax0)
    # ax2.plot(np.linspace(0, len(resp_data_1) // resp_fs, len(resp_data_2)), resp_data_2, color='green')
    # ax2.set_title('Resp Data after Butterworth Filtering')
    # plt.tight_layout()
    # plt.show()

    bcg_data = utils.butterworth(data=signal_data, _type=conf["bcg_filter"]["filter_type"],
                                 low_cut=conf["bcg_filter"]["low_cut"],
                                 high_cut=conf["bcg_filter"]["high_cut"], order=conf["bcg_filter"]["order"],
                                 sample_rate=signal_fs)


    return signal_data, resp_data_2, resp_fs, bcg_data, signal_fs



def rpeak2hr(rpeak_indices, signal_length):
    hr_signal = np.zeros(signal_length)
    for i in range(1, len(rpeak_indices)):
        rri = rpeak_indices[i] - rpeak_indices[i - 1]
        if rri == 0:
            continue
        hr = 60 * 1000 / rri  # 心率，单位：bpm
        if hr > 120:
            hr = 120
        elif hr < 30:
            hr = 30
        hr_signal[rpeak_indices[i - 1]:rpeak_indices[i]] = hr
    # 填充最后一个R峰之后的心率值
    if len(rpeak_indices) > 1:
        hr_signal[rpeak_indices[-1]:] = hr_signal[rpeak_indices[-2]]
    return hr_signal