#define BOOST_TEST_MODULE fft_test

#include <boost/test/included/unit_test.hpp>
#include <boost/test/data/dataset.hpp>
#include <boost/test/data/monomorphic.hpp>
#include <boost/test/data/test_case.hpp>

#include "autocorrelation.h"
#include "fft.h"
#include "tuner.h"
#include "util.h"

#include <complex>
#include <chrono>
#include <cmath>
#include <ctime>
#include <exception>
#include <functional>
#include <iostream>
#include <memory>
#include <numeric>
#include <string>
#include <vector>

const std::complex<double> i(0, 1);

const double tolerance = 0.000001;

using namespace std::complex_literals;

/// CPU time clock
class Timer {
public:
	Timer(): mStart(std::clock()){}
	void start(){mStart = std::clock();}
	double elapsed() { return (double(std::clock()) - mStart) / CLOCKS_PER_SEC; }
	static double precision() {
		static double result{0};
		
		if (result != 0)
			return result;
		
		std::clock_t start = std::clock();
		while (start == std::clock()) {}
		start = std::clock();

		std::clock_t end = start;
		while (end == std::clock()) {}
		end = std::clock();

		result = (double(end) - start) / CLOCKS_PER_SEC;
		return result;
	}
private:
	std::clock_t mStart;
};

// (Slow) DFT
std::vector<std::complex<double>> dft(const std::vector<std::complex<double>> &v) {
	std::vector<std::complex<double>> result(v.size(), 0);

	const double N = v.size();

	for (int k = 0; k < v.size(); k++) {
		for (int j = 0; j < result.size(); j++) {
			result[k] += v[j] * std::exp(-i * 2. * M_PI * double(k * j) / N);
		}
	}

	return result;
}

std::vector<double> freqs{ 2, 5, 11, 17, 29 }; // known freqs for testing

void generate(std::vector<std::complex<double>>& v) {
    for (int i = 0; i < v.size(); i++) {
        v[i] = 0.;
        // sum several known sinusoids into v[]
        for(int j = 0; j < freqs.size(); j++)
            v[i] += sin(2 * M_PI * freqs[j] * i / v.size() );
    }
}

namespace CooleyTukey {

// separate even/odd elements to lower/upper halves of array respectively.
// Due to Butterfly combinations, this turns out to be the simplest way 
// to get the job done without clobbering the wrong elements.
void separate(std::vector<std::complex<double>>::iterator a, int n) {
	std::vector<std::complex<double>> b(n/2);
	for(int i=0; i<n/2; i++)    // copy all odd elements to heap storage
		b[i] = a[i*2+1];
	for(int i=0; i<n/2; i++)    // copy all even elements to lower-half of a[]
		a[i] = a[i*2];
	for(int i=0; i<n/2; i++)    // copy all odd (from heap) to upper-half of a[]
		a[i+n/2] = b[i];
}

// N must be a power-of-2, or bad things will happen.
// Currently no check for this condition.
//
// N input samples in X[] are FFT'd and results left in X[].
// Because of Nyquist theorem, N samples means 
// only first N/2 FFT results in X[] are the answer.
// (upper half of X[] is a reflection with no new information).
void fft_recursive(std::vector<std::complex<double>>::iterator X, int N) {
    if(N < 2) {
        // bottom of recursion.
        // Do nothing here, because already X[0] = x[0]
    } else {
        separate(X,N);      // all evens to lower half, all odds to upper half
        fft_recursive(X,     N/2);   // recurse even items
        fft_recursive(X+N/2, N/2);   // recurse odd  items
        // combine results of two half recursions
        for(int k=0; k<N/2; k++) {
            std::complex<double> e = X[k    ];   // even
            std::complex<double> o = X[k+N/2];   // odd
                         // w is the "twiddle-factor"
            std::complex<double> w = exp( std::complex<double>(0,-2.*M_PI*k/N) );
            X[k    ] = e + w * o;
            X[k+N/2] = e - w * o;
        }
    }
}

// Cooley-Tukey
std::vector<std::complex<double>> fft(const std::vector<std::complex<double>> &v) {
	std::vector<std::complex<double>> result(v);

	fft_recursive(std::begin(result), result.size());
	
	return result;
}
}; // namespace CooleyTukey

class Measure {
	Timer mTimer;

public:
	using Data = std::vector<std::complex<double>>;
	
protected:
	Measure* mReference; // Reference measurement: we should calculate similar to this one and be faster than this one
	
	const Data& mIn;
	Data mResult;

	std::string mName;
	double mElapsed;

public:
	Measure(const Data& in): mReference(nullptr), mIn(in), mResult(in.size()) {}

	virtual void run_impl() = 0; // implemented by subclasses

	void run() {
		mTimer.start();
		run_impl();
		mElapsed = mTimer.elapsed();
		std::cout << mName << ": " << mElapsed << "s\n";

		if (mReference) {
			if (mResult != mReference->mResult) {
				double diff = std::transform_reduce(std::begin(mResult), std::end(mResult), std::begin(mReference->mResult), double(0.0),
								    [](const double& x, const double& y) -> double
								    { return x + y;},
								    [](const std::complex<double>& x, const std::complex<double>&y) -> double
								    { return abs(y - x);}
								    );
				BOOST_CHECK_MESSAGE(diff <= tolerance, "Error: Results diff: " << diff);
			}

			BOOST_CHECK_MESSAGE(mElapsed <= mReference->mElapsed, "Error: " << mName << " too slow!");
		}
	}

	double elapsed() { return mElapsed; }

	Data& result() { return mResult; }
};

class MeasureDFT: public Measure {
public:
	MeasureDFT(const Data& in): Measure(in){ mName = "DFT";}
	void run_impl() override {
		mResult = dft(mIn);
	}
};

class MeasureFFT: public Measure {
public:
	MeasureFFT(const Data& in, Measure& reference): Measure(in) { mName = "FFT Cooley-Tukey"; mReference = &reference;}
	void run_impl() override {
		mResult = CooleyTukey::fft(mIn);
	}
};

class MeasureFFT_RR: public Measure {
	RIT::FFT mRR;
public:
	MeasureFFT_RR(const Data& in, Measure& reference): Measure(in), mRR(in.size()){ mName = "FFT RR"; mReference = &reference;}
	void run_impl() override {
		mResult = mRR(mIn);
	}
};

class MeasureFFT_RR_half: public Measure {
	RIT::FFT mRR;
public:
	MeasureFFT_RR_half(const Data& in, Measure& reference): Measure(in), mRR(in.size(), true){ mName = "FFT RR half"; mReference = &reference; }
	void run_impl() override {
		mResult = mRR(mIn);
	}
};

class MeasureFFT_RR_half_magnitudes: public Measure {
	RIT::FFT mRR;
public:
	MeasureFFT_RR_half_magnitudes(const Data& in, Measure& reference): Measure(in), mRR(in.size(), true){ mName = "FFT RR half magnitudes"; mReference = &reference; }
	void run_impl() override {
		mResult = mRR(mIn);
		RIT::magnitudes(mResult);
	}
};

class MeasureIFFT_RR: public Measure {
	RIT::IFFT mIFFT;
public:
	MeasureIFFT_RR(const Data& in): Measure(in), mIFFT(in.size()) { mName = "IFFT RR";}
	void run_impl() override {
		mResult = mIFFT(mIn);
	}
};

class MeasureAutoCorrelation_RR: public Measure {
	RIT::AutoCorrelation mAC;
public:
	MeasureAutoCorrelation_RR(const Data& in): Measure(in), mAC(in.size()) { mName = "AutoCorrelation RR";}
	void run_impl() override {
		mResult = mAC(mIn);
	}
};

class MeasureTuner_RR: public Measure {
	RIT::Tuner mTuner;
	RIT::Pitch mPitch;
public:
	MeasureTuner_RR(const Data& in): Measure(in), mTuner(in.size(), 44100) { mName = "Tuner RR";}
	void run_impl() override {
		mPitch = mTuner(mIn);
	}
};

BOOST_AUTO_TEST_CASE(performance)
{
	std::vector<std::complex<double>> v(4096, 0);

	generate(v);

	std::cout << "Timer precision: " << Timer::precision() << "s\n";

	MeasureDFT measureDFT(v);
	measureDFT.run();

	MeasureFFT measureFFT(v, measureDFT);
	measureFFT.run();

	MeasureFFT_RR measureFFT_RR(v, measureFFT);
	measureFFT_RR.run();

	MeasureFFT_RR_half measureFFT_RR_half(v, measureFFT_RR);
	measureFFT_RR_half.run();

	MeasureFFT_RR_half_magnitudes measureFFT_RR_half_magnitudes(v, measureDFT);
	measureFFT_RR_half_magnitudes.run();

	MeasureIFFT_RR measureIFFT_RR(v);
	measureIFFT_RR.run();

	MeasureAutoCorrelation_RR measureAutoCorrelation_RR(v);
	measureAutoCorrelation_RR.run();

	MeasureTuner_RR measureTuner_RR(v);
	measureTuner_RR.run();
}

// TODO:
// -0.5 <= deviation <= 0.5

BOOST_AUTO_TEST_CASE(fft_2)
{
        std::vector<std::complex<double>> v{{1, 0}, {0, 0}};
        RIT::FFT fft{2};
        auto fft_result = fft(v);

        BOOST_REQUIRE_EQUAL(fft_result.size(), 2);
        BOOST_REQUIRE_EQUAL(fft_result[0], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[1], (std::complex<double>{1.0, 0.0}));

        RIT::IFFT ifft{2};
        auto ifft_result = ifft(fft_result);

        BOOST_REQUIRE_EQUAL(ifft_result.size(), 2);
        BOOST_REQUIRE_EQUAL(ifft_result[0], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[1], (std::complex<double>{0.0, 0.0}));
}

BOOST_AUTO_TEST_CASE(fft_4a)
{
        std::vector<std::complex<double>> v{{1, 0}, {0, 0}, {0, 0}, {0, 0}};
        RIT::FFT fft{4};
        auto fft_result = fft(v);

        BOOST_REQUIRE_EQUAL(fft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(fft_result[0], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[1], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[2], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[3], (std::complex<double>{1.0, 0.0}));

        RIT::IFFT ifft{4};
        auto ifft_result = ifft(fft_result);

        BOOST_REQUIRE_EQUAL(ifft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(ifft_result[0], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[1], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[2], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[3], (std::complex<double>{0.0, 0.0}));
}

BOOST_AUTO_TEST_CASE(fft_4b)
{
        std::vector<std::complex<double>> v{{1, 0}, {-1, 0}, {1, 0}, {-1, 0}};
        RIT::FFT fft{4};
        auto fft_result = fft(v);

        BOOST_REQUIRE_EQUAL(fft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(fft_result[0], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[1], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[2], (std::complex<double>{4.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[3], (std::complex<double>{0.0, 0.0}));

        RIT::IFFT ifft{4};
        auto ifft_result = ifft(fft_result);

        BOOST_REQUIRE_EQUAL(ifft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(ifft_result[0], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[1], (std::complex<double>{-1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[2], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[3], (std::complex<double>{-1.0, 0.0}));
}

BOOST_AUTO_TEST_CASE(fft_4c)
{
        std::vector<std::complex<double>> v{{-1, 0}, {1, 0}, {-1, 0}, {1, 0}};
        RIT::FFT fft{4};
        auto fft_result = fft(v);

        BOOST_REQUIRE_EQUAL(fft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(fft_result[0], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[1], (std::complex<double>{0.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[2], (std::complex<double>{-4.0, 0.0}));
        BOOST_REQUIRE_EQUAL(fft_result[3], (std::complex<double>{0.0, 0.0}));

        RIT::IFFT ifft{4};
        auto ifft_result = ifft(fft_result);

        BOOST_REQUIRE_EQUAL(ifft_result.size(), 4);
        BOOST_REQUIRE_EQUAL(ifft_result[0], (std::complex<double>{-1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[1], (std::complex<double>{1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[2], (std::complex<double>{-1.0, 0.0}));
        BOOST_REQUIRE_EQUAL(ifft_result[3], (std::complex<double>{1.0, 0.0}));
}