#include "fft-beacon-finder.h"

#include <iostream>

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;

    last_center = 0;
    last_center_level = 0.0;

    has_sync = false;
}

std::complex<float> FFTBeaconFinder::work(std::complex<float> in) {
    if (next_fft_in <= 0 && !has_sync) {
            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;
                    }
                }

                std::cout << std::endl;

                int last_center_idx = spectral_bin_to_fft_idx(last_center);
                float last_center_val =
                    std::abs(fft_out[last_center_idx]) / fft_max;

                if (last_center_val < 0.25) {
                    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) {
                            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;

                            std::cout << bin << " " << left_val << " | "
                                      << center_val << " | " << right_val
                                      << std::endl;

                            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) {
                        last_center = max_center;

                        float center_freq =
                            float(max_center) / FFT_LEN * samplingrate;
                        nco_crcf_set_frequency(
                            coarse_correction,
                            -(2 * M_PI * center_freq) / samplingrate);

                        std::cout << "New center " << center_freq << " "
                                  << max_center << std::endl;
                        has_sync = true;
                    } else {
                        std::cout << "No center found!" << std::endl;
                    }
                } else {
                    std::cout << "Last center still valid" << std::endl;
                    has_sync = true;
                }

                next_fft_in = samplingrate / 4;
            }
    } else if(next_fft_in > 0){
        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 FFT_LEN + bin;
    }
}

void FFTBeaconFinder::reset_sync() {
    has_sync = false;
}

FFTBeaconFinder::~FFTBeaconFinder() {
    nco_crcf_destroy(coarse_correction);
    fft_destroy_plan(fft);
}