import numpy as np
from scipy import signal, fft, linalg
from typing import Tuple, List


class AntiJammingSignalAlgo:

    @staticmethod
    def polarization_filter(rx_signal, sample_rate, polarization_params):
        """极化滤波算法"""
        # 应用极化参数
        angle_rad = np.deg2rad(polarization_params['angle'])
        ellipticity = polarization_params['ellipticity']
        
        # 构建极化滤波矩阵
        filter_matrix = np.array([
            [np.cos(angle_rad), -ellipticity * np.sin(angle_rad)],
            [np.sin(angle_rad), ellipticity * np.cos(angle_rad)]
        ])
        
        # 应用极化滤波
        polarized_signal = filter_matrix @ rx_signal.reshape(2, -1)
        return polarized_signal[0].real

    SAMPLING_RATE = 10e6

    @staticmethod
    def frequency_agility(tx_signal, hop_sequence, sample_rate=SAMPLING_RATE):
        """频率捷变算法
        Args:
            tx_signal: 原始发射信号
            sample_rate: 采样率
            hop_sequence: 跳频序列 (包含频率点和驻留时间)
        """
        processed_signal = np.zeros_like(tx_signal)
        t = 0
        
        for freq, dwell in hop_sequence:
            n_samples = 100
            freq_shift = np.exp(1j*2*np.pi*freq/sample_rate*np.arange(n_samples))
            segment = tx_signal[t:t+n_samples] * freq_shift
            processed_signal[t:t+n_samples] = segment
            t += n_samples
            if t >= len(tx_signal):
                break
        return processed_signal

    @staticmethod
    def waveform_agility(tx_signal, waveform_params, sample_rate=SAMPLING_RATE):
        """波形捷变算法
        Args:
            tx_signal: 原始发射信号
            sample_rate: 采样率
            waveform_params: 波形参数 {
                'type': 'LFM'/'NLFM'/'PhaseCode',
                'bandwidth': 带宽(Hz),
                'duration': 持续时间(s)
            }
        """
        t = np.linspace(0, waveform_params['duration'], int(sample_rate * waveform_params['duration']))
        
        if waveform_params['type'] == 'LFM':
            waveform = signal.chirp(t, f0=0, f1=waveform_params['bandwidth'], t1=waveform_params['duration'])
        elif waveform_params['type'] == 'NLFM':
            # 非线性调频波形
            beta = 2.5
            waveform = np.cos(2*np.pi* waveform_params['bandwidth'] * (t + beta*t**3))
        elif waveform_params['type'] == 'PhaseCode':
            # 相位编码波形
            code_length = 13
            phase_code = np.exp(1j*np.pi*np.random.randint(0,2,code_length))
            waveform = np.kron(phase_code, np.ones(int(sample_rate*waveform_params['duration']/code_length)))
        else:
            raise ValueError("不支持的波形类型")
        
        return tx_signal * waveform[:len(tx_signal)]