#ifndef PSK_MODULATOR_H
#define PSK_MODULATOR_H


#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdint>
#include <numeric>
#include <stdexcept>
#include <vector>
#include <fftw3.h>

static constexpr double CARRIER_FREQ = 1800.0;
static constexpr size_t SYMBOL_RATE = 2400;
static constexpr double ROLLOFF_FACTOR = 0.35;
static constexpr double SCALE_FACTOR = 32767.0;

class PSKModulator {
public:
    PSKModulator(const double _sample_rate, const bool _is_frequency_hopping, const size_t _num_taps) 
        : sample_rate(validateSampleRate(_sample_rate)), gain(1.0/sqrt(2.0)), is_frequency_hopping(_is_frequency_hopping), samples_per_symbol(static_cast<size_t>(sample_rate / SYMBOL_RATE)), num_taps(_num_taps) {
            initializeSymbolMap();
    }

    std::vector<int16_t> modulate(const std::vector<uint8_t>& symbols) const {
        std::vector<double> baseband_I(symbols.size() * samples_per_symbol);
        std::vector<double> baseband_Q(symbols.size() * samples_per_symbol);
        size_t symbol_index = 0;

        for (const auto& symbol : symbols) {
            if (symbol >= symbolMap.size()) {
                throw std::out_of_range("Invalid symbol value for 8-PSK modulation. Symbol must be between 0 and 7.");
            }
            const std::complex<double> target_symbol = symbolMap[symbol];

            for (size_t i = 0; i < samples_per_symbol; ++i) {
                baseband_I[symbol_index * samples_per_symbol + i] = target_symbol.real();
                baseband_Q[symbol_index * samples_per_symbol + i] = target_symbol.imag();
            }

            symbol_index++;
        }

        // Filter the I/Q phase components
        auto filter_taps = generateSRRCTaps(num_taps, sample_rate, SYMBOL_RATE, ROLLOFF_FACTOR);
        auto filtered_I = applyFilter(baseband_I, filter_taps);
        auto filtered_Q = applyFilter(baseband_Q, filter_taps);

        std::vector<std::complex<double>> baseband_components(filtered_I.size());
        for (size_t i = 0; i < filtered_I.size(); i++) {
            std::complex<double> baseband_component = {filtered_I[i], filtered_Q[i]};
            baseband_components[i] = baseband_component;
        }

        // Combine the I and Q components
        std::vector<double> passband_signal;
        passband_signal.reserve(baseband_components.size());

        double carrier_phase = 0.0;
        double carrier_phase_increment = 2 * M_PI * CARRIER_FREQ / sample_rate;
        for (const auto& sample : baseband_components) {
            double carrier_cos = std::cos(carrier_phase);
            double carrier_sin = -std::sin(carrier_phase);
            double passband_value = sample.real() * carrier_cos + sample.imag() * carrier_sin;
            passband_signal.emplace_back(passband_value * 32767.0); // Scale to int16_t
            carrier_phase += carrier_phase_increment;
            if (carrier_phase >= 2 * M_PI)
                carrier_phase -= 2 * M_PI;
        }

        std::vector<int16_t> final_signal;
        final_signal.reserve(passband_signal.size());

        for (const auto& sample : passband_signal) {
            int16_t value = static_cast<int16_t>(sample);
            value = std::clamp(value, (int16_t)-32768, (int16_t)32767);
            final_signal.emplace_back(value);
        }

        return final_signal;
    }

private:
    const double sample_rate;        ///< The sample rate of the system.
    const double gain;               ///< The gain of the modulated signal.
    size_t samples_per_symbol; ///< Number of samples per symbol, calculated to match symbol duration with cycle.
    const size_t num_taps;
    std::vector<std::complex<double>> symbolMap;  ///< The mapping of tribit symbols to I/Q components.
    const bool is_frequency_hopping;                    ///< Whether to use frequency hopping methods. Not implemented (yet?)

    static inline double validateSampleRate(const double rate) {
        if (rate <= 2 * CARRIER_FREQ) {
            throw std::out_of_range("Sample rate must be above the Nyquist frequency (PSKModulator.h)");
        }
        return rate;
    }

    inline void initializeSymbolMap() {
        symbolMap = {
            {gain * std::cos(2.0*M_PI*(0.0/8.0)), gain * std::sin(2.0*M_PI*(0.0/8.0))}, // 0 (000) corresponds to I = 1.0, Q = 0.0
            {gain * std::cos(2.0*M_PI*(1.0/8.0)), gain * std::sin(2.0*M_PI*(1.0/8.0))}, // 1 (001) corresponds to I = cos(45), Q = sin(45)
            {gain * std::cos(2.0*M_PI*(2.0/8.0)), gain * std::sin(2.0*M_PI*(2.0/8.0))}, // 2 (010) corresponds to I = 0.0, Q = 1.0
            {gain * std::cos(2.0*M_PI*(3.0/8.0)), gain * std::sin(2.0*M_PI*(3.0/8.0))}, // 3 (011) corresponds to I = cos(135), Q = sin(135)
            {gain * std::cos(2.0*M_PI*(4.0/8.0)), gain * std::sin(2.0*M_PI*(4.0/8.0))}, // 4 (100) corresponds to I = -1.0, Q = 0.0
            {gain * std::cos(2.0*M_PI*(5.0/8.0)), gain * std::sin(2.0*M_PI*(5.0/8.0))}, // 5 (101) corresponds to I = cos(225), Q = sin(225)
            {gain * std::cos(2.0*M_PI*(6.0/8.0)), gain * std::sin(2.0*M_PI*(6.0/8.0))}, // 6 (110) corresponds to I = 0.0, Q = -1.0
            {gain * std::cos(2.0*M_PI*(7.0/8.0)), gain * std::sin(2.0*M_PI*(7.0/8.0))}  // 7 (111) corresponds to I = cos(315), Q = sin(315)
        };
    }

    std::vector<double> generateSRRCTaps(size_t num_taps,double sample_rate, double symbol_rate, double rolloff) const {
        std::vector<double> freq_response(num_taps, 0.0);
        std::vector<double> taps(num_taps);

        double fn = symbol_rate / 2.0;
        double f_step = sample_rate / num_taps;

        for (size_t i = 0; i < num_taps / 2; i++) {
            double f = i * f_step;

            if (f <= fn * (1 - rolloff)) {
                freq_response[i] = 1.0;
            } else if (f <= fn * (1 + rolloff)) {
                freq_response[i] = 0.5 * (1 - std::sin(M_PI * (f - fn * (1 - rolloff)) / (2 * rolloff * fn)));
            } else {
                freq_response[i] = 0.0;
            }
        }

        for (size_t i = num_taps / 2; i < num_taps; i++) {
            freq_response[i] = freq_response[num_taps - i - 1];
        }

        fftw_complex* freq_domain = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * num_taps);
        for (size_t i = 0; i < num_taps; i++) {
            freq_domain[i][0] = freq_response[i];
            freq_domain[i][1] = 0.0;
        }

        std::vector<double> time_domain_taps(num_taps, 0.0);

        fftw_plan plan = fftw_plan_dft_c2r_1d(num_taps, freq_domain, time_domain_taps.data(), FFTW_ESTIMATE);
        fftw_execute(plan);
        fftw_destroy_plan(plan);
        fftw_free(freq_domain);

        double norm_factor = std::sqrt(std::accumulate(time_domain_taps.begin(), time_domain_taps.end(), 0.0, [](double sum, double val) { return sum + val * val; }));
        for (auto& tap : time_domain_taps) {
            tap /= norm_factor;
        }

        return time_domain_taps;
    }

    std::vector<double> applyFilter(const std::vector<double>& signal, const std::vector<double>& filter_taps) const {
        std::vector<double> filtered_signal(signal.size(), 0.0);

        size_t filter_length = filter_taps.size();
        size_t half_filter_length = filter_length / 2;

        // Convolve the signal with the filter taps
        for (size_t i = 0; i < signal.size(); ++i) {
            double filtered_val = 0.0;

            for (size_t j = 0; j < filter_length; ++j) {
                size_t signal_index = i + j - half_filter_length;
                if (signal_index < signal.size()) {
                    filtered_val += filter_taps[j] * signal[signal_index];
                }
            }

            filtered_signal[i] = filtered_val;
        }

        return filtered_signal;
    }
};

#endif