// FFT

#include "fft.h"

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


namespace { // Helper functions
	bool is_power_of_two(unsigned int n) {
		return n != 0 && (n & (n - 1)) == 0;
	}

}

std::vector<double> RIT::magnitudes(std::vector<std::complex<double>>& v) {
	std::vector<double> result(v.size());
	std::transform(std::begin(v), std::end(v), std::begin(result), std::abs<double>);
	return result;
}

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] = bitreverse(i);
	}

	// 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) {
	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;
}

int RIT::FFT::bitreverse(int i) {
	int size{mSize};
	int result{0};

	while (size > 1) {
		result <<= 1;
		result |= i & 1;
		i >>= 1;
		size >>= 1;
	}

	return result;
}

void RIT::FFT::reorder(const std::vector<std::complex<double>>& src, std::vector<std::complex<double>>& dst) {
	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) {
	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) {
	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];
	}
}