#include "dsp_manager.h"
#include "imu.h"

// Буферы для расчета
static float32_t fft_input[FFT_SIZE];
static float32_t fft_output[FFT_SIZE];
static float32_t magnitudes[FFT_SIZE / 2];

// Буфер для окна Ханна (чтобы убрать шумы по краям выборки)
static float32_t hann_window[FFT_SIZE];

static uint16_t sample_count = 0;
uint8_t dsp_buffer_ready = 0;

// Структура БПФ из библиотеки
static arm_rfft_fast_instance_f32 fft_handler;

// Частоты текущих 3 подавляющих фильтров
float active_notch_freqs[3] = {0.0f, 0.0f, 0.0f};

//альфа
static float32_t smoothed_freqs[3] = {0.0f, 0.0f, 0.0f};
#define FREQ_ALPHA 0.05f  // Чем МЕНЬШЕ значение, тем дольше фильтр "едет" к новой частоте

void DSP_Init(void) {
    // Инициализируем структуру БПФ
    arm_rfft_fast_init_f32(&fft_handler, FFT_SIZE);
    
    // Генерируем окно Ханна (делается один раз)
    for (int i = 0; i < FFT_SIZE; i++) {
        hann_window[i] = 0.5f * (1.0f - arm_cos_f32(2.0f * 3.14159f * i / (1023.0f)));
    }
}

void DSP_AddSample(float32_t sample) {
    if (dsp_buffer_ready) return; // Ждем, пока обработают прошлую пачку

    fft_input[sample_count++] = sample;

    if (sample_count >= FFT_SIZE) {
        sample_count = 0;
        dsp_buffer_ready = 1; // Сигнализируем в main
    }
}

void DSP_Process(void) {
    // 1. Применяем окно Ханна
    arm_mult_f32(fft_input, hann_window, fft_input, FFT_SIZE);

    // 2. САМО БПФ
    arm_rfft_fast_f32(&fft_handler, fft_input, fft_output, 0);

    // 3. Считаем амплитуды
    arm_cmplx_mag_f32(fft_output, magnitudes, FFT_SIZE / 2);

    // 4. Поиск 3-х независимых самых мощных пиков
    float32_t top_freq_indices[3] = {0};
    float32_t top_mags[3] = {0};

    // Индексы для поиска от 50 Гц до 450 Гц
    // index = freq * FFT_SIZE / fs = freq * 512 / 1000
    uint32_t start_idx = 25; // ~50 Гц (50 * 512 / 1000 = 25)
    uint32_t end_idx = 230;  // ~450 Гц (450 * 512 / 1000 = 230)

    for (int k = 0; k < 3; k++) {
        float32_t max_m = 0;
        uint32_t max_i = 0;

        // Ищем глобальный максимум
        for (uint32_t i = start_idx; i < end_idx; i++) {
            if (magnitudes[i] > max_m) {
                max_m = magnitudes[i];
                max_i = i;
            }
        }

        top_mags[k] = max_m;
        top_freq_indices[k] = (float32_t)max_i;

        // "Зануляем" гору вокруг найденного пика (±10 бинов, ~±20 Гц)
        // Чтобы следующий фильтр не прицепился к «склону» того же самого пика
        if (max_i > 0) {
            uint32_t clear_start = (max_i > 10) ? (max_i - 10) : 0;
            uint32_t clear_end   = (max_i + 10 < (FFT_SIZE / 2)) ? (max_i + 10) : ((FFT_SIZE / 2) - 1);
            for (uint32_t j = clear_start; j <= clear_end; j++) {
                magnitudes[j] = 0.0f; 
            }
        }
    }

    // --- 5. ПЕРЕНАСТРОЙКА ТРЕХ КАСКАДОВ FMAC ---
    const float fs = 1000.0f;     // Частота дискретизации
    const float Q = 3.5f;        // Добротность
    const float bin_to_hz = fs / (float)FFT_SIZE; 

    for (int i = 0; i < 3; i++) {
        float mag = top_mags[i];
        float new_freq = top_freq_indices[i] * bin_to_hz;
        
        // Hysteresis (ГИСТЕРЕЗИС):
        // Если фильтр сейчас ВЫКЛЮЧЕН (active_notch_freqs == 0)
        if (active_notch_freqs[i] == 0) {
            if (mag > 4000.0f) {
                // Включаем фильтр! Чтобы не полз с нуля, присваиваем частоту сразу:
                smoothed_freqs[i] = new_freq;
                active_notch_freqs[i] = new_freq;
            }
        } 
        // Если фильтр сейчас ВКЛЮЧЕН
        else {
            if (mag < 2000.0f) {
                // Выключаем фильтр, так как амплитуда сильно упала
                active_notch_freqs[i] = 0;
            } else {
                // Продолжаем отслеживать с Альфой (EMA)
                smoothed_freqs[i] = (smoothed_freqs[i] * (1.0f - FREQ_ALPHA)) + (new_freq * FREQ_ALPHA);
                active_notch_freqs[i] = smoothed_freqs[i];
            }
        }
        
        // Применяем настройки
        if (active_notch_freqs[i] > 0) {
            float real_freq = active_notch_freqs[i];
            
            // Математика Notch-фильтра
            float w0 = 2.0f * 3.14159265f * real_freq / fs;
            float alpha = arm_sin_f32(w0) / (2.0f * Q);
            float cosw0 = arm_cos_f32(w0);
            float a0 = 1.0f + alpha;

            float b0 = 1.0f / a0;
            float b1 = -2.0f * cosw0 / a0;
            float b2 = 1.0f / a0;
            float a1 = -2.0f * cosw0 / a0;
            float a2 = (1.0f - alpha) / a0;

            Update_FMAC_Coeffs(i, b0, b1, b2, a1, a2);
        } else {
            // Bypass
            Update_FMAC_Coeffs(i, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f);
        }
    }

    dsp_buffer_ready = 0; // Разрешаем новый сбор данных
}