#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")));

// FFT handler
static arm_rfft_fast_instance_f32 fft_handler;

// Публичные переменные (см. dsp_manager.h)
uint8_t fft_ready = 0;
uint16_t fft_index = 0;
volatile uint16_t notch_report_hz[3] = {0, 0, 0};

// Внутренние состояния и телеметрия
static float32_t notch_freq_smoothed[3] = {0.0f, 0.0f, 0.0f};
volatile float32_t motor_erpm_avg = 0.0f;
volatile float32_t motor_mech_rpm = 0.0f;
volatile float32_t motor_elec_hz = 0.0f;

void DSP_SetMotorErpm(const uint16_t eRPM[4], uint8_t pole_pairs) {
    // ERPM используем только как debug-метрику, не как источник notch-частот
    float base_hz = (float)eRPM[3] / 60.0f;

    motor_erpm_avg = (float)eRPM[3];
    motor_elec_hz = base_hz;
    motor_mech_rpm = (float)eRPM[3] / (float)pole_pairs;
}

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 (fft_ready) return; // Ждем, пока обработают прошлую пачку

    fft_input[fft_index++] = sample;

    if (fft_index >= FFT_SIZE) {
        fft_index = 0;
        fft_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);

    // Ищем 3 сильнейших пика в спектре и на них строим notch'и.
    // Телеметрия ERPM здесь не участвует, чтобы не уводить фильтры в неверные Гц.
    float32_t peak_freqs[3] = {0.0f, 0.0f, 0.0f};
    const uint32_t start_bin = (uint32_t)((50.0f * FFT_SIZE) / DSP_SAMPLE_RATE_HZ);

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

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

            // Не даём второму и третьему ноту садиться на уже выбранные пики.
            uint8_t too_close = 0;
            for (int j = 0; j < k; j++) {
                if (peak_freqs[j] > 1.0f && fabsf(freq_hz - peak_freqs[j]) < 40.0f) {
                    too_close = 1;
                    break;
                }
            }
            if (too_close) {
                continue;
            }

            if (magnitudes[i] > max_val) {
                max_val = magnitudes[i];
                max_idx = (int32_t)i;
            }
        }

        if (max_idx >= 0) {
            peak_freqs[k] = ((float32_t)max_idx * DSP_SAMPLE_RATE_HZ) / FFT_SIZE;
        }
    }

    // Лёгкое сглаживание только чтобы частоты не дрожали между соседними бинами
    for (int k = 0; k < 3; k++) {
        if (peak_freqs[k] > 1.0f) {
            if (notch_freq_smoothed[k] < 1.0f) {
                notch_freq_smoothed[k] = peak_freqs[k];
            } else {
                const float32_t alpha = 0.25f;
                notch_freq_smoothed[k] = (alpha * peak_freqs[k]) + ((1.0f - alpha) * notch_freq_smoothed[k]);
            }
        }
    }

    // 5. ПЕРЕНАСТРОЙКА ФИЛЬТРОВ
    float32_t Notch_Q = 2.5f;

    biquad_init_notch(&notch1, notch_freq_smoothed[0], Notch_Q, DSP_SAMPLE_RATE_HZ);
    biquad_init_notch(&notch2, notch_freq_smoothed[1], Notch_Q, DSP_SAMPLE_RATE_HZ);
    biquad_init_notch(&notch3, notch_freq_smoothed[2], Notch_Q, DSP_SAMPLE_RATE_HZ);

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

    // Готово
    fft_ready = 0;
}