import numpy as np
from scipy.signal import find_peaks
# ==================== 干扰雷达类 ====================
class JammerRadar:
    def __init__(self, usrp, rx, tx):
        self.usrp = usrp
        self.rx = rx
        self.tx = tx

    def generate_jamming_signal(self, algorithm: callable, bandwidth: float, sample_rate: float, duration: float, gain: float) -> np.ndarray:
        """
        自主生成干扰信号
        :param algorithm: 干扰算法
        :param bandwidth: 信号带宽 (Hz)
        :param sample_rate: 采样率 (Hz)
        :param duration: 信号持续时间 (秒)
        :param gain: 增益 (dB)
        """
        # 生成基带噪声信号
        num_samples = int(duration * sample_rate)
        base_signal = np.random.randn(num_samples) * 0.1 * gain
        
        # 应用干扰算法
        jamming_signal = algorithm(
            base_signal=base_signal,
            bandwidth=bandwidth,
            sample_rate=sample_rate
        )
        return jamming_signal.astype(np.complex64)

    def execute_jamming(self, algorithm: callable, bandwidth: float, sample_rate: float, duration: float, gain: float) -> np.ndarray:
        """
        执行自主干扰生成
        :param algorithm: 干扰算法
        :param bandwidth: 信号带宽 (Hz)
        :param sample_rate: 采样率 (Hz)
        :param duration: 信号持续时间 (秒)
        :param gain: 增益 (dB)
        """
        return self.generate_jamming_signal(
            algorithm=algorithm,
            bandwidth=bandwidth,
            sample_rate=sample_rate,
            duration=duration,
            gain=gain
        )

    #封装一个发送信号的函数
    def send_signal(self, tx_signal, duration, center_freq, sample_rate, gain):
        # 发送信号
        self.usrp.send_waveform(tx_signal, duration, center_freq, sample_rate, [self.tx], gain)
        print('干扰雷达已发送信号')

    # 封装一个接收信号的函数
    def recv_signal(self, num_samples, sample_rate, center_freq):
        # 获取设备支持的最接近采样率
        valid_rate = self._get_valid_sample_rate(sample_rate)
        
        if valid_rate != sample_rate:
            print(f'警告：调整采样率 {sample_rate/1e6}MHz -> {valid_rate/1e6}MHz')
            
        rx_signal = self.usrp.recv_num_samps(num_samples, center_freq, valid_rate, [self.rx])
        print('干扰雷达已接收信号')
        return rx_signal

    def _get_valid_sample_rate(self, requested_rate):
        """获取设备支持的有效采样率"""
        # USRP N210支持的有效采样率范围（单位：Hz）
        MIN_RATE = 1e6   # 1 MHz
        MAX_RATE = 25e6  # 25 MHz
        
        # 调整到最近的整数MHz
        adjusted_rate = round(requested_rate / 1e6) * 1e6
        
        # 限制在设备支持的范围内
        return max(MIN_RATE, min(adjusted_rate, MAX_RATE))