#include "dsp_manager.h"
#include "imu.h"
#include <math.h>

// Буферы для расчета
static float32_t fft_input[FFT_SIZE] __attribute__((section(".sram2")));
static float32_t fft_output[FFT_SIZE] __attribute__((section(".sram2")));
static float32_t magnitudes[FFT_SIZE / 2] __attribute__((section(".sram2")));
static float32_t hann_window[FFT_SIZE] __attribute__((section(".sram2")));

// Коэффициенты биквадратного фильтра
float32_t b[3] = {1.0f, -2.0f, 1.0f}; // Примерные значения
float32_t a[3] = {1.0f, -1.8f, 0.81f};

// Буфер состояния фильтра
float32_t x[3] = {0};
float32_t y[3] = {0};

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

// Определение переменных
uint8_t dsp_buffer_ready = 0;
uint16_t sample_count = 0;
volatile uint16_t dsp_notch_freqs[3] = {0, 0, 0};

// Переменные для сглаживания частоты (чтобы не прыгала)
static float32_t filtered_freqs[3] = {0.0f, 0.0f, 0.0f};
const float32_t SMOOTH_ALPHA = 0.2f; // Коэффициент плавности (0.1 - очень медленно, 0.9 - мгновенно)
const float32_t FAST_TRACK_ALPHA = 0.55f;
const float32_t PEAK_THRESHOLD = 2200.0f;
const float32_t FAST_TRACK_DELTA_HZ = 25.0f;
const float32_t TRACK_WINDOW_HZ = 35.0f;
static uint8_t freq_seen_frames[3] = {0, 0, 0};
static uint8_t freq_missing_frames[3] = {0, 0, 0};
const uint8_t REQUIRED_STABLE_FRAMES = 2;
const uint8_t DISABLE_AFTER_MISSING_FRAMES = 5;

void DSP_Init(void) {
    // Инициализируем структуру БПФ
    arm_rfft_fast_init_f32(&fft_handler, FFT_SIZE);
    
    // Генерируем окно Ханна (делается один раз)
    const float32_t sin_delta = 0.0246374509f;
    const float32_t cos_delta = 0.9996964600f;
    float32_t sin_angle = 0.0f;
    float32_t cos_angle = 1.0f;

    for (int i = 0; i < FFT_SIZE; i++) {
        hann_window[i] = 0.5f * (1.0f - cos_angle);

        float32_t next_cos = (cos_angle * cos_delta) - (sin_angle * sin_delta);
        float32_t next_sin = (sin_angle * cos_delta) + (cos_angle * sin_delta);
        cos_angle = next_cos;
        sin_angle = next_sin;
    }
}

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 current_iteration_freqs[3] = {0.0f, 0.0f, 0.0f};
    const uint32_t start_bin = (uint32_t)((50.0f * FFT_SIZE) / DSP_SAMPLE_RATE_HZ);
    const float32_t freq_per_bin = DSP_SAMPLE_RATE_HZ / FFT_SIZE;
    const uint32_t track_window_bins = (uint32_t)(TRACK_WINDOW_HZ / freq_per_bin);

    for (int k = 0; k < 3; k++) {
        float32_t max_val = 0.0f;
        int32_t max_idx = -1;

        uint32_t search_start_bin = start_bin;
        uint32_t search_end_bin = (FFT_SIZE / 2);

        if (filtered_freqs[k] > 1.0f) {
            int32_t center_bin = (int32_t)(filtered_freqs[k] / freq_per_bin);

            if (center_bin < (int32_t)track_window_bins) {
                search_start_bin = start_bin;
            } else {
                search_start_bin = (uint32_t)(center_bin - (int32_t)track_window_bins);
            }

            search_end_bin = (uint32_t)(center_bin + (int32_t)track_window_bins);
            if (search_end_bin > (FFT_SIZE / 2)) {
                search_end_bin = (FFT_SIZE / 2);
            }
        }

        for (uint32_t i = search_start_bin; i < search_end_bin; i++) {
            float32_t freq_hz = ((float32_t)i * DSP_SAMPLE_RATE_HZ) / FFT_SIZE;

            if (filtered_freqs[k] <= 1.0f) {
                uint8_t too_close_to_locked = 0;
                for (int locked = 0; locked < k; locked++) {
                    if (filtered_freqs[locked] > 1.0f && fabsf(freq_hz - filtered_freqs[locked]) < 40.0f) {
                        too_close_to_locked = 1;
                        break;
                    }
                }

                if (too_close_to_locked) {
                    continue;
                }
            }

            uint8_t too_close = 0;
            for (int j = 0; j < k; j++) {
                if (current_iteration_freqs[j] > 1.0f && fabsf(freq_hz - current_iteration_freqs[j]) < 40.0f) { // Разнос 40Гц
                    too_close = 1;
                    break;
                }
            }

            if (!too_close && magnitudes[i] > max_val) {
                max_val = magnitudes[i];
                max_idx = i;
            }
        }

        // Если нашли пик выше порога
        if (max_idx != -1 && max_val > PEAK_THRESHOLD) {
            current_iteration_freqs[k] = ((float32_t)max_idx * DSP_SAMPLE_RATE_HZ) / FFT_SIZE;
        } else if (filtered_freqs[k] > 1.0f) {
            // Если активный notch не нашел точку в окне, удерживаем текущую частоту,
            // а не переводим его в гонку за соседним пиком.
            current_iteration_freqs[k] = filtered_freqs[k];
        }
    }

    // Применяем задержку подтверждения и сглаживание к частотам
    for (int k = 0; k < 3; k++) {
        if (current_iteration_freqs[k] > 1.0f) {
            freq_missing_frames[k] = 0;

            if (freq_seen_frames[k] < 255) {
                freq_seen_frames[k]++;
            }

            if (filtered_freqs[k] < 1.0f) {
                if (freq_seen_frames[k] >= REQUIRED_STABLE_FRAMES) {
                    filtered_freqs[k] = current_iteration_freqs[k];
                }
            } else {
                float32_t alpha = SMOOTH_ALPHA;

                if (fabsf(current_iteration_freqs[k] - filtered_freqs[k]) > FAST_TRACK_DELTA_HZ) {
                    alpha = FAST_TRACK_ALPHA;
                }

                filtered_freqs[k] = (alpha * current_iteration_freqs[k]) +
                                    ((1.0f - alpha) * filtered_freqs[k]);
            }
        } else {
            freq_seen_frames[k] = 0;

            if (freq_missing_frames[k] < 255) {
                freq_missing_frames[k]++;
            }

            if (freq_missing_frames[k] >= DISABLE_AFTER_MISSING_FRAMES) {
                filtered_freqs[k] = 0.0f;
            }
        }
    }

    // 5. ПЕРЕНАСТРОЙКА ФИЛЬТРОВ
    float32_t Notch_Q = 2.5f; // Чуть шире яма, чтобы лучше удерживать пик и ловить дрейф

    // Вызываем инициализацию только если частота > 0
    biquad_init_notch(&notch1, filtered_freqs[0], Notch_Q, 1000.0f);
    biquad_init_notch(&notch2, filtered_freqs[1], Notch_Q, 1000.0f);
    biquad_init_notch(&notch3, filtered_freqs[2], Notch_Q, 1000.0f);

    // В телеметрию
    dsp_notch_freqs[0] = (uint16_t)filtered_freqs[0];
    dsp_notch_freqs[1] = (uint16_t)filtered_freqs[1];
    dsp_notch_freqs[2] = (uint16_t)filtered_freqs[2];

    dsp_buffer_ready = 0; 
}

// Реализация функции Biquad_Filter
float32_t Biquad_Filter(float32_t input) {
    float32_t output = b[0] * input + b[1] * x[1] + b[2] * x[2] - a[1] * y[1] - a[2] * y[2];

    x[2] = x[1];
    x[1] = input;
    y[2] = y[1];
    y[1] = output;

    return output;
}