#include "fft-beacon-finder.h"

FFTBeaconFinder::FFTBeaconFinder(int samplingrate) {
    this->samplingrate = samplingrate;

    coarse_correction = nco_crcf_create(LIQUID_NCO);
    nco_crcf_set_frequency(coarse_correction, 0.0f);

    fft = fft_create_plan(FFT_LEN, fft_in, fft_out, LIQUID_FFT_FORWARD, 0);

    pos = 0;
    next_fft_in = 0;
}

std::complex<float> FFTBeaconFinder::work(std::complex<float> in) {
    if (next_fft_in <= 0) {
        fft_in[pos] = in;
        pos += 1;
        if (pos == FFT_LEN) {
            pos = 0;
            fft_execute(fft);

            float fft_max = std::abs(fft_out[0]);
            for (int i = 0; i < FFT_LEN; i++) {
                float mag = std::abs(fft_out[i]);
                if (mag > fft_max) {
                    fft_max = mag;
                }
            }

            float max_levels = 0;
            int max_center = 0;
            for (int bin = -50; bin <= 50; bin++) {
                int center_idx = spectral_bin_to_fft_idx(bin);
                float center_val = std::abs(fft_out[center_idx]) / fft_max;
                if (center_val > 0.25) {
                    printf("Found peak candidate at %d\n", bin);
                    int left_idx = spectral_bin_to_fft_idx(bin - 127);
                    int right_idx = spectral_bin_to_fft_idx(bin + 127);
                    float left_val = std::abs(fft_out[left_idx]) / fft_max;
                    float right_val = std::abs(fft_out[right_idx]) / fft_max;

                    if (center_val + left_val + right_val > max_levels) {
                        max_levels = center_val + left_val + right_val;
                        max_center = bin;
                    }
                }
            }

            if (max_levels > 0.0) {
                float center_freq = max_center * samplingrate / FFT_LEN;
                printf("Found center at %f\n", center_freq);
                nco_crcf_set_frequency(
                    coarse_correction,
                    -(2 * M_PI * center_freq) / samplingrate);
            }

            next_fft_in = samplingrate / 4;
        }
    } else {
        next_fft_in--;
    }

    std::complex<float> y;
    // increment internal phase
    nco_crcf_step(coarse_correction);
    // compute complex exponential
    nco_crcf_cexpf(coarse_correction, &y);

    return y * in;
}

int FFTBeaconFinder::spectral_bin_to_fft_idx(int bin) {
    if (bin == 0) {
        return FFT_LEN / 2;
    } else if (bin > 0) {
        return bin - 1;
    } else {
        return bin + FFT_LEN;
    }
}