// FFT

#include "fft.h"

#include "util.h"

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


RIT::FFT::FFT(int size, bool halfOnly): mSize(size), order(size), expLUT(size/2), mFlagHalfOnly(halfOnly) {

	if (!is_power_of_two(size))
		throw std::invalid_argument("Size must be a power of two");

	// reorder LUT
	for (int i = 0; i < size; ++i) {
		order[i] = RIT::bitreverse(i, size);
	}

	// exp LUT
	for (int i = 0; i < size / 2; ++i) {
		expLUT[i] = exp(std::complex<double>(0,-2.*M_PI*i/size));
	}
}

std::vector<std::complex<double>> RIT::FFT::operator()(const std::vector<std::complex<double>> &v) const {
	if (v.size() != mSize)
		throw std::length_error("Bad input size");

	std::vector<std::complex<double>> result;

	reorder(v, result);
	if (mFlagHalfOnly) {
		fft_half(std::begin(result), mSize);
		result.resize(mSize/2);
	} else
		fft_recursive(std::begin(result), mSize);
	
	return result;
}

RIT::FFT& RIT::FFT::SetHalfOnly(bool enable) {
	mFlagHalfOnly = enable;
	return *this;
}

void RIT::FFT::reorder(const std::vector<std::complex<double>>& src, std::vector<std::complex<double>>& dst) const {
	int size = src.size();

	dst.resize(size);
	for (int i = 0; i < size; ++i) {
		dst[order[i]] = src[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 RIT::FFT::fft_recursive(std::vector<std::complex<double>>::iterator X, int N) const {
	if(N > 2) {
		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 = expLUT[k * mSize / N];
		X[k    ] = e + w * o;
		X[k+N/2] = e - w * o;
	}
}

// Same as fft_recursive, but results in only half the result due to symmetry in real input
void RIT::FFT::fft_half(std::vector<std::complex<double>>::iterator X, int N) const {
	if(N > 2) {
		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++) {
		X[k] += expLUT[k * mSize / N] * X[k+N/2];
	}
}

struct RIT::IFFT::Impl {
public:
	using transform_function = std::complex<double> (*)(const std::complex<double>&);

	Impl(int size): mFft(std::make_shared<RIT::FFT>(size)) {}
	Impl(std::shared_ptr<RIT::FFT> fft): mFft(fft) {}
	~Impl(){}
	std::shared_ptr<RIT::FFT> mFft;
};

RIT::IFFT::IFFT(int size): mImpl(std::make_unique<RIT::IFFT::Impl>(size))
{
}

RIT::IFFT::IFFT(std::shared_ptr<RIT::FFT> fft): mImpl(std::make_unique<RIT::IFFT::Impl>(fft))
{
}

RIT::IFFT::~IFFT()
{
}

std::vector<std::complex<double>> RIT::IFFT::operator()(const std::vector<std::complex<double>> &v) const
{
	std::vector<std::complex<double>> input(v.size());

	std::transform(std::begin(v), std::end(v), std::begin(input), static_cast<Impl::transform_function>(std::conj));

	auto result = (*mImpl->mFft)(input);
	
	const double factor = 1.0 / v.size();
	std::transform(std::begin(result), std::end(result), std::begin(result),
		       [=](const auto& x){return std::conj(x) * factor;});

	return result;
}