from numpy import abs as np_abs, ceil, floor, ndarray, argmax, dot, linalg, where, arange
from numpy.fft import fft, fftfreq
from scipy.signal import detrend, butter, filtfilt


def fft1(x, fs):
    """
    x: 输入信号
    fs: 采样频率

    返回：
    magnitude[mask]：0-1hz信号频谱
    freq[mask]:0-1hz信号对应频率数组
    """
    N = len(x)
    # 计算FFT
    fft_result = fft(x)
    # 计算幅度谱（取前N//2点，避免镜像）
    magnitude = np_abs(fft_result)[:N // 2] * 2 / N  # 归一化
    # 生成频率轴
    freqs = fftfreq(N, 1 / fs)[:N // 2]
    # 提取0-1Hz范围内的数据
    mask = freqs <= 1.0  # 筛选条件
    return magnitude[mask] / max(magnitude[mask]), freqs[mask]


def lpf(input_signal, Fs, Fc):
    """
    巴特沃斯低通滤波器 (零相位滤波)

    参数:
        input_signal : 输入信号
        Fs : float - 采样频率 (Hz)
        Fc : float - 截止频率 (Hz)

    返回:
        output : ndarray - 滤波后信号
    """
    # 滤波器参数
    order = 3  # 3阶滤波器
    nyq = 0.5 * Fs  # 奈奎斯特频率

    # 检查截止频率是否有效
    if Fc <= 0 or Fc >= nyq:
        raise ValueError("截止频率必须满足 0 < Fc < {} Hz".format(nyq))

    # 去趋势处理
    data = detrend(input_signal, type='linear')

    # 生成巴特沃斯滤波器系数
    normal_cutoff = Fc / nyq
    b, a = butter(order, normal_cutoff, btype='low', analog=False)

    # 零相位滤波
    output = filtfilt(b, a, data)

    return output


def bpf(input_signal, fs, f_low, f_high):
    """
        巴特沃斯带通滤波器 (零相位滤波)

        参数:
            signal : 输入信号
            fs : float - 采样频率 (Hz)
            f_low : float - 下截止频率 (Hz)
            f_high:float - 上截止频率（Hz)

        返回:
            output : ndarray - 滤波后信号
        """
    # 滤波器参数
    order = 2  # 3阶滤波器
    nyq = 0.5 * fs  # 奈奎斯特频率

    # 检查截止频率是否有效
    if f_low <= 0 or f_high >= nyq:
        raise ValueError("截止频率必须满足 0 < Fc < {} Hz".format(nyq))

    if f_low >= f_high:
        raise ValueError("f_low must be less than f_high")

    # 去趋势处理
    data = detrend(input_signal, type='linear')

    # 生成巴特沃斯滤波器系数
    fc1 = f_low / nyq
    fc2 = f_high / nyq
    b, a = butter(order, [fc1, fc2], btype='bandpass', analog=False)

    # 零相位滤波
    output = filtfilt(b, a, data)

    return output


def detect_peak_bottoms(frequencies, spectrum, peak_idx, left_bound, right_bound):
    """
    检测信号呼吸峰谷
    参数：
    frequencies: 频率数组
    spectrum: 频谱数组
    peak_idx: 峰值索引
    left_bound: 左边界
    right_bound: 右边界
    返回：
    f_low：峰谷索引
    f_high：峰谷索引
    """
    normalized = spectrum / max(spectrum)  # 归一化频谱
    n = len(frequencies)
    peak_idx = int(peak_idx * n)
    left_bound = int(left_bound * n)
    right_bound = int(right_bound * n)

    # 左侧峰底检测
    left_min = int(ceil(left_bound))
    for i in range(peak_idx, left_bound - 1, -1):
        if i == 0:
            break  # 避免索引超出范围
        if normalized[i] < 0.2 * normalized[peak_idx]:
            left_min = i
            break
        if normalized[i] < normalized[i - 1]:
            left_min = i

    # 右侧峰底检测
    right_min = int(floor(right_bound))
    for i in range(peak_idx, right_bound + 1):
        if i == n - 1:
            break  # 避免索引超出范围
        if normalized[i] < 0.2 * normalized[peak_idx]:
            right_min = i
            break
        if normalized[i] > normalized[i + 1] and i != right_bound:
            right_min = i + 1

    # 结果
    f_low = max(frequencies[left_min], 0.1)
    f_high = min(frequencies[right_min], 1)

    return f_low, f_high

def spectral_cosine_similarity(signal1, signal2, fs1, fs2, freq_range):
    """
    计算两个信号在指定频段内的频谱余弦相似度

    参数：
        signal1: 信号1
        signal2: 信号2
        fs1: 信号1的采样率 (Hz)
        fs2: 信号2的采样率 (Hz)
        freq_range: 目标频段 [f_low, f_high] (单位Hz)

    返回：
        similarity: 指定频段内的频谱余弦相似度 (0~1)
    """

    # 参数检查
    if len(freq_range) != 2 or freq_range[0] >= freq_range[1]:
        raise ValueError('freq_range应为升序的二元向量 [f_low, f_high]')

    f_low, f_high = freq_range

    # 计算FFT并获取频率轴
    # 信号1
    nfft1 = len(signal1)
    fft_1 = fft(signal1)
    freq1 = arange(nfft1) * (fs1 / nfft1)  # 完整频率轴
    mag1 = np_abs(fft_1)

    # 信号2
    nfft2 = len(signal2)
    fft_2 = fft(signal2)
    freq2 = arange(nfft2) * (fs2 / nfft2)
    mag2 = np_abs(fft_2)

    # 提取目标频段的幅度谱
    idx1 = where((freq1 >= f_low) & (freq1 <= f_high))[0]
    idx2 = where((freq2 >= f_low) & (freq2 <= f_high))[0]

    # 对齐长度（取最小公共频段）
    min_len = min(len(idx1), len(idx2))
    mag1_band = mag1[idx1[:min_len]]
    mag2_band = mag2[idx2[:min_len]]

    # 计算余弦相似度
    dot_product = dot(mag1_band, mag2_band)
    norm1 = linalg.norm(mag1_band)
    norm2 = linalg.norm(mag2_band)

    if norm1 == 0 or norm2 == 0:
        similarity = 0.0  # 避免除以零
    else:
        similarity = dot_product / (norm1 * norm2)

    return similarity


def get_bandpass_bcgsignal(
        bcg_signal: ndarray,
        tho_signal: ndarray,
        bcg_sample_rate: int,
        tho_sample_rate: int,
        use_custom_band: bool = False,
        low_cutoff: float = 0.1,
        high_cutoff: float = 1.0
):
    """
    绘制带通滤波后的BCG信号和原始THO信号对比图

    Args:
        bcg_signal: 原始BCG信号数组
        tho_signal: 原始THO信号数组
        bcg_sample_rate: BCG信号采样率(Hz)
        tho_sample_rate: THO信号采样率(Hz)
        use_custom_band: 是否使用自定义频带 (默认自动检测)
        low_cutoff: 自定义低截止频率(Hz)
        high_cutoff: 自定义高截止频率(Hz)

    Returns:
        filtered_bcg :带通滤波后的BCG信号（BDR信号）

    """
    # 参数校验
    if len(bcg_signal) == 0 or len(tho_signal) == 0:
        raise ValueError("输入信号不能为空数组")
    if low_cutoff >= high_cutoff:
        raise ValueError("截止频率范围不合法")

    # 频谱分析
    tho_spectrum, tho_freq = fft1(tho_signal, tho_sample_rate)
    bcg_spectrum, bcg_freq = fft1(bcg_signal, bcg_sample_rate)

    # 确定通带范围
    if use_custom_band:
        band_low = low_cutoff
        band_high = high_cutoff
    else:
        peak_freq = tho_freq[argmax(tho_spectrum)]
        band_low, band_high = detect_peak_bottoms(
            frequencies=bcg_freq,
            spectrum=bcg_spectrum,
            peak_idx=peak_freq,
            left_bound=max(peak_freq - 0.2, 0.1),
            right_bound=min(peak_freq + 0.2, 1.0)
        )

    # 带通滤波
    filtered_bcg = bpf(
        bcg_signal,
        bcg_sample_rate,
        band_low,
        band_high
    )

    return filtered_bcg, band_low, band_high, bcg_spectrum, bcg_freq, tho_spectrum, tho_freq

