#include <opencv2/opencv.hpp>

void dump(std::string const& name, std::vector<double>& v)
{
    std::cout << name << " = [";
    for (auto val : v) {
        std::cout << val << ", ";
    }
    std::cout << "]\n";
}

cv::Mat1d gen_time_steps(double sample_freq, int32_t sample_count)
{
    cv::Mat1d t;
    for (int32_t i(0); i < sample_count; ++i) {
        t.push_back(i / sample_freq);
    }
    return t.reshape(1, 1); // Make it single row
}


// Generate complex sine wave
cv::Mat2d gen_sine_wave(double freq, cv::Mat1d const& t)
{
    cv::Mat2d y;
    for (auto ts : t) {
        y.push_back(cv::Vec2d(sin(ts * (2 * CV_PI) * freq), 0));
    }
    return y.reshape(2, 1); // Make it single row
}

// Get magnitude from a complex spectrum
cv::Mat1d get_magnitudes(cv::Mat2d fd)
{
    std::vector<cv::Mat1d> ri(2);
    cv::split(fd, ri);
    cv::Mat1d mag;
    cv::sqrt(ri[0].mul(ri[0]) + ri[1].mul(ri[1]), mag);
    return mag;
}

double get_peak_freq(cv::Mat1d const& spectrum, double sampling_rate, int32_t sample_count)
{
    cv::Point2i maxloc;
    cv::minMaxLoc(spectrum, nullptr, nullptr, nullptr, &maxloc);

    return maxloc.x * (sampling_rate / sample_count);
}

double test(double freq, double sampling_rate, int32_t sample_count)
{
    // Create time vector
    cv::Mat1d t(gen_time_steps(sampling_rate, sample_count));

    // Creating sin signal with 1Hz period
    cv::Mat2d y(gen_sine_wave(freq, t));

    // Compute DFT
    cv::Mat2d fd;
    cv::dft(y, fd);
    fd = fd.colRange(0, fd.cols / 2); // Other half is mirrored

    cv::Mat1d spectrum(get_magnitudes(fd));

    return get_peak_freq(spectrum, sampling_rate, sample_count);
}

int main()
{
    std::vector<int32_t> window_sizes{ 256, 512, 1024 }; // samples
    std::vector<double> sampling_rates{ 16, 32, 64 }; // Hz
    std::vector<double> freqs{ 1./16, 1./8, 1./4, 1./2, 1.0
        , 1.5, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 7.5, 8.0
        , 12.0, 16 - (1./4), 16 - (1./8), 16 - (1./16), 16.0 }; // Hz

    for (int32_t window_size : window_sizes) {
        std::cout << "Window size = " << window_size << "\n";
        for (double sampling_rate : sampling_rates) {
            std::cout << "* Sampling rate = " << sampling_rate << " Hz\n";
            for (double freq : freqs) {
                std::cout << "** f = " << freq << " : detected f = "
                    << test(freq, sampling_rate, window_size) << " Hz\n";
            }
        }
    }

    return 0;
}