from matplotlib.mlab import detrend
from numpy import argmax

from func.utils.resp_quality_label_filter import fft1, bpf, spectral_cosine_similarity


def get_slice(biosignal, seg_idx, duration, fs):
    """数据切片"""
    start = (seg_idx - 1) * duration * fs
    end = min(seg_idx * duration * fs, len(biosignal))
    return biosignal[start:end]

def pre_process(biosignal,fs,fc):
    """数据预处理"""
    biosignal = detrend(biosignal)
    biosignal = bpf(biosignal, fs, 0.1,fc)
    return biosignal

def calculate_spectral_cohere(BCG, THO, fs_bcg,fs_tho):
    """计算频谱相似度"""
    fft_THO, freq_tho = fft1(THO, fs_tho)
    peak_freq = freq_tho[argmax(fft_THO)]
    f_low = min(peak_freq-0.2,0.1)
    f_high = max(peak_freq+0.2,0.1)
    # 计算频谱相似度
    Similarity = spectral_cosine_similarity(BCG, THO, fs_bcg, fs_tho,[f_low,f_high])
    return Similarity

def evaluate_quality(BCG,THO,fs_bcg,fs_tho,artifact_flag,thresholds):
    """评估信号质量"""
    # 计算信号余弦相似度
    cohere_value = calculate_spectral_cohere(BCG, THO, fs_bcg, fs_tho)
    # 质量判断逻辑
    if cohere_value > thresholds[1]:
        return 1 if artifact_flag == 0 else -1
    elif cohere_value < thresholds[0]:
        BCG = pre_process(BCG, fs_bcg, 0.7)
        THO = pre_process(THO, fs_tho, 0.7)
        cohere_value2 = calculate_spectral_cohere(BCG, THO, fs_bcg, fs_tho)
        return 0 if cohere_value2 < thresholds[0] else -1
    return -1