import numpy as np
import matplotlib.pyplot as plt



class CICFilter:
    def __init__(self, decimation_factor, delay, order):
        self.decimation_factor = decimation_factor
        self.delay = delay
        self.order = order

    def integrate(self, input_signal):
      output = np.copy(input_signal)
      for _ in range(self.order):
          output = np.cumsum(output)
      return output

    def comb(self, decimated_signal):
        output = np.copy(decimated_signal)
        for _ in range(self.order):
            delayed = np.zeros_like(output)
            delayed[self.delay:] = output[:-self.delay]
            output = output - delayed
        return output
    
    def decimate(self, integrated_signal):
        # Децимация: выбираем каждый N-й элемент
        return integrated_signal[::self.decimation_factor]
    

    def filter(self, input_signal):
        # 1. Каскад интеграторов
        integrated_signal = self.integrate(input_signal)

        # 2. Децимация
        decimated_signal = self.decimate(integrated_signal)

        # 3. Каскад comb-фильтров
        output_signal = self.comb(decimated_signal)

        return output_signal
    
    def cascade_multiple_cic_filters(self, input_signal, num_filters):
        output_signal = input_signal
        for _ in range(num_filters):
            output_signal = self.filter(output_signal)
        return output_signal
    

    
    def get_filter_freq_resp(self, num_points=1024, is_cascade=False, num_filters=1, show_before_decimation=False):
        # Подаём дельта-функцию на вход фильтру
        impulse = np.zeros(num_points)
        impulse[0] = 1.0  # Дельта-функция

        if show_before_decimation:
            # Получаем только характеристику интеграторов без децимации и Comb-фильтров
            output_signal = np.copy(impulse)
            for _ in range(num_filters):
                output_signal = self.integrate(output_signal)
        else:
            # Стандартная реализация - полный CIC фильтр с децимацией
            if is_cascade:
                output_signal = self.cascade_multiple_cic_filters(impulse, num_filters)
            else:
                output_signal = self.filter(impulse)

        # Дополняем выходной сигнал нулями до исходной длины
        if len(output_signal) < num_points:
            padded_output = np.zeros(num_points)
            padded_output[:len(output_signal)] = output_signal
            output_signal = padded_output

        # Вычисляем частотную характеристику с помощью БПФ
        freq_response = np.fft.fft(output_signal, num_points)

        # Нормализуем частотную характеристику
        freq_response /= np.max(np.abs(freq_response))

        # Создаем массив частот с исходной частотой дискретизации fs=50e6
        fs = 50e6
        freq = np.fft.fftfreq(num_points, d=1/fs*self.decimation_factor)

        # Возвращаем частоты и амплитуду частотной характеристики
        return freq, np.abs(freq_response)
    
    def get_theoretical_freq_resp(self, num_points=1024, fs=50e6):
        freq = np.fft.fftfreq(num_points, d=1/fs)

        # Вычисляем нормализованные частоты
        f_norm = np.abs(freq) / fs

        # Расчет по формуле
        R = self.decimation_factor
        M = self.delay
        N = self.order

        # Избегаем деления на ноль при f=0
        mag = np.ones(num_points)
        nonzero = f_norm != 0

        # Угол для sin(π*f*R*M)
        angle_num = np.pi * f_norm[nonzero] * R * M
        # Угол для sin(π*f*M)
        angle_den = np.pi * f_norm[nonzero] 

        # |H(f)| = |sin(π*f*R*M)/(R*sin(π*f*M))|^N
        mag[nonzero] = np.abs(np.sin(angle_num) / ( np.sin(angle_den))) ** N

        # Нормализуем амплитудную характеристику
        mag /= np.max(mag)

        return freq, mag


def main():
    decimation_factor = 4
    delay = 2
    order = 4
    num_samples = 10000
    fs = 50e6  # Sampling frequency
    
    # Создаем CIC фильтр
    cic_filter = CICFilter(decimation_factor, delay, order)
    
    # Получаем теоретическую АЧХ (весь диапазон)
    theory_freq, theory_mag = cic_filter.get_theoretical_freq_resp(num_points=2048, fs=fs)
    
    # Получаем АЧХ после децимации
    freq_after_dec, resp_after_dec = cic_filter.get_filter_freq_resp( num_points=2048, is_cascade=False, num_filters=1)
    
    # Строим график
    plt.figure(figsize=(12, 8))
    
    # Выделяем положительные частоты для теоретической АЧХ
    mask_pos = theory_freq >= 0
    theory_freq_pos = theory_freq[mask_pos]
    theory_mag_pos = theory_mag[mask_pos]
    
    # Новая частота Найквиста после децимации
    nyquist_new = fs/(2*decimation_factor)
    
    # Выделяем частоты для АЧХ после децимации только до новой частоты Найквиста
    mask_nyquist = (freq_after_dec >= 0) & (freq_after_dec <= nyquist_new)
    freq_pos_after = freq_after_dec[mask_nyquist]
    resp_pos_after = resp_after_dec[mask_nyquist]
    
    # Преобразуем в дБ с защитой от нуля
    db_theory = 20 * np.log10(np.maximum(theory_mag_pos, 1e-10))
    db_after = 20 * np.log10(np.maximum(resp_pos_after, 1e-10))
    plt.plot(theory_freq_pos, db_theory, 'r-', label='Теоретическая АЧХ (весь диапазон)')
    plt.plot(freq_pos_after, db_after, 'b-', linewidth=2, label='АЧХ после децимации')
    
    # Отмечаем новую частоту Найквиста
    plt.axvline(x=nyquist_new, color='k', linestyle=':', 
                label=f'Новая частота Найквиста ({nyquist_new/1e6:.2f} МГц)')
    
    # Затеняем область выше новой частоты Найквиста
    plt.axvspan(nyquist_new, fs/2, alpha=0.2, color='gray', 
                label='Область наложения спектра (aliasing)')
    
    plt.title(f'Сравнение частотных характеристик CIC-фильтра R={decimation_factor}, D={delay}, N={order}')
    plt.xlabel('Частота [Гц]')
    plt.ylabel('Амплитуда [дБ]')
    plt.grid(True)
    plt.legend()
    plt.xlim(0, fs/2)  # Отображаем только до fs/2
    plt.ylim(-300, 10)
    plt.tight_layout()
    plt.show()
 
    
if __name__ == "__main__":
    main()
