#ifndef FSK_DEMODULATOR_H
#define FSK_DEMODULATOR_H

#include <cmath>
#include <cstdint>
#include <functional>
#include <memory>
#include <vector>

#include "BitStreamWriter.h"

class FSKDemodulatorConfig {
    public:
    int freq_lo;
    int freq_hi;
    int sample_rate;
    int baud_rate;
    std::shared_ptr<BitStreamWriter> bitstreamwriter;
};

namespace milstd {
    class FSKDemodulator {
        public:
        FSKDemodulator(const FSKDemodulatorConfig& s) : freq_lo(s.freq_lo), freq_hi(s.freq_hi), sample_rate(s.sample_rate), baud_rate(s.baud_rate), bit_writer(s.bitstreamwriter) {
            initialize();
        }

        void demodulate(const std::vector<int16_t>& samples) {
            size_t nb = samples.size();

            for (size_t i = 0; i < nb; i++) {
                filter_buf[buf_ptr++] = samples[i];

                if (buf_ptr == filter_buf.size()) {
                    std::copy(filter_buf.begin() + filter_buf.size() - filter_size, filter_buf.end(), filter_buf.begin());
                    buf_ptr = filter_size;
                }

                int corr;
                int sum = 0;

                corr = dotProduct(&filter_buf[buf_ptr - filter_size], filter_hi_i.data(), filter_size);
                sum += corr * corr;

                corr = dotProduct(&filter_buf[buf_ptr - filter_size], filter_hi_q.data(), filter_size);
                sum += corr * corr;

                corr = dotProduct(&filter_buf[buf_ptr - filter_size], filter_lo_i.data(), filter_size);
                sum -= corr * corr;

                corr = dotProduct(&filter_buf[buf_ptr - filter_size], filter_lo_q.data(), filter_size);
                sum -= corr * corr;

                int new_sample = (sum > 0) ? 1 : 0;

                if (last_sample != new_sample) {
                    last_sample = new_sample;
                    if (baud_pll < 0.5)
                        baud_pll += baud_pll_adj;
                    else
                        baud_pll -= baud_pll_adj;
                }

                baud_pll += baud_incr;

                if (baud_pll >= 1.0) {
                    baud_pll -= 1.0;
                    bit_writer->putBit(last_sample);
                }
            }
        }

        private:
        int freq_lo;
        int freq_hi;
        int sample_rate;
        int baud_rate;
        std::shared_ptr<BitStreamWriter> bit_writer;

        int filter_size;
        std::vector<double> filter_lo_i;
        std::vector<double> filter_lo_q;
        std::vector<double> filter_hi_i;
        std::vector<double> filter_hi_q;
        std::vector<double> filter_buf;
        size_t buf_ptr;

        double baud_incr;
        double baud_pll;
        double baud_pll_adj;
        int last_sample;

        void initialize() {
            baud_incr = static_cast<double>(baud_rate) / sample_rate;
            baud_pll = 0.0;
            baud_pll_adj = baud_incr / 4;

            filter_size = sample_rate / baud_rate;

            filter_buf.resize(filter_size * 2, 0.0);
            buf_ptr = filter_size;

            last_sample = 0;

            filter_lo_i.resize(filter_size);
            filter_lo_q.resize(filter_size);
            filter_hi_i.resize(filter_size);
            filter_hi_q.resize(filter_size);
            
            for (int i = 0; i < filter_size; i++) {
                double phase_lo = 2.0 * M_PI * freq_lo * i / sample_rate;
                filter_lo_i[i] = std::cos(phase_lo);
                filter_lo_q[i] = std::sin(phase_lo);
                
                double phase_hi = 2.0 * M_PI * freq_hi * i / sample_rate;
                filter_hi_i[i] = std::cos(phase_hi);
                filter_hi_q[i] = std::sin(phase_hi);
            }
        }

        double dotProduct(const double* x, const double* y, size_t size) {
            double sum = 0.0;
            for (size_t i = 0; i < size; i++) {
                sum += x[i] * y[i];
            }
            return sum;
        }
    };
} // namespace milstd

#endif /* FSK_DEMODULATOR_H */