#ifndef FFTREAL_HH

#define FFTREAL_HH

// This code is based on:

//    fft-real

//    https://github.com/cyrilcode/fft-real

//    FFTReal is distributed under the terms of the Do What The Fuck You Want To Public License.

//

// This version is stripped down and specialized for our use case.

#include "cstd.hh"

#include "narrow.hh"

#include "xrange.hh"

#include <array>

#include <cassert>

#include <cmath>

#include <cstdint>

#include <numbers>

#include <span>

// Input mode for FFT optimization

enum class FFTInputMode {

	Full,           // All FFT_LEN input values provided

	                // Output: N complex values in FFTReal format (N/2+1 real, N/2-1 imag)

	HalfZeroPadded, // Only first FFT_LEN/2 values provided, second half is implicitly zero

	                // Example N=8: provide [a,b,c,d] → full array [a,b,c,d,0,0,0,0]

	                // Output: N complex values in FFTReal format (N/2+1 real, N/2-1 imag)

	Symmetric       // Only first FFT_LEN/2+1 values, mirrored to fill second half

	                // input[0] and input[N/2] are unique (endpoints)

	                // input[1..N/2-1] mirrored: input[N/2+i] = input[N/2-i]

	                // Example N=8: provide [a,b,c,d,e] → full array [a,b,c,d,e,d,c,b]

	                // Output: Only N/2+1 REAL values (imaginary parts are ~0 and skipped)

};

// Calculate the FFT for real input (as opposed to complex numbers input).

// NOTE: Uses exp(+2πi) sign convention (opposite from numpy/FFTW). Imaginary parts have opposite sign.

//

// The result however are still complex numbers, and arranged as follows:

//   output[0]         = Real(bin 0)

//   output[1]         = Real(bin 1)

//   ...

//   output[len/2 - 1] = Real(bin len/2 -1)

//   output[len/2    ] = Real(bin len/2)

//   output[len/2 + 1] = Imag(bin 1)

//   ...

//   output[len   - 2] = Imag(bin len/2 - 2)

//   output[len   - 1] = Imag(bin len/2 - 1)

// Note: There are two more real-parts than imaginary-parts. For the lowest bin

// (the DC-offset) and for the highest bin (at Nyquist frequency) the output is

// a pure real value (the imaginary part is zero).

template<unsigned FFT_LEN_L2, FFTInputMode MODE = FFTInputMode::Full>

class FFTReal

{

public:

	static_assert(FFT_LEN_L2 >=  4); // at least 16-points (allows optimization in final pass)

	static_assert(FFT_LEN_L2 <= 16); // at most 65536-points

	static constexpr unsigned FFT_LEN = 1 << FFT_LEN_L2;

	static constexpr unsigned INPUT_LEN =

		(MODE == FFTInputMode::HalfZeroPadded) ? (FFT_LEN / 2) :

		(MODE == FFTInputMode::Symmetric) ? (FFT_LEN / 2 + 1) :

		FFT_LEN;

	static constexpr unsigned OUTPUT_LEN =

		(MODE == FFTInputMode::Symmetric) ? (FFT_LEN / 2 + 1) :

		FFT_LEN;

	// Execute FFT and return meaningful output range

	// NOTE: The returned span aliases the 'output' parameter (first OUTPUT_LEN elements)

	//       Full 'output' buffer is still needed for intermediate calculations

	static std::span<float, OUTPUT_LEN> execute(std::span<const float, INPUT_LEN> input,

	                                            std::span<float, FFT_LEN> output,

	                                            std::span<float, FFT_LEN> tmpBuf)

	{

		pass<FFT_LEN_L2 - 1>(output, tmpBuf, input);

#ifndef NDEBUG

		// In debug mode, verify imaginary outputs are indeed ~0 for Symmetric mode

		if constexpr (MODE == FFTInputMode::Symmetric) {

			for (unsigned i = FFT_LEN / 2 + 1; i < FFT_LEN; ++i) {

				assert(std::abs(output[i]) < 1e-4f && "Symmetric mode: imaginary outputs should be ~0");

			}

		}

#endif

		return std::span<float, OUTPUT_LEN>(output.data(), OUTPUT_LEN);

	}

private:

	// Helper for dependent false in static_assert

	template<auto> static inline constexpr bool always_false = false;

	template<int PASS>

	static void pass(std::span<float, FFT_LEN> dst,

	                 std::span<float, FFT_LEN> src,

	                 std::span<const float, INPUT_LEN> input)

	{

		if constexpr (PASS == 1) {

			// Do 1st and 2nd pass at once.

			static constexpr auto len4 = FFT_LEN >> 2;

			if constexpr (MODE == FFTInputMode::Full) {

				// Full input path (general case)

				for (unsigned idx = 0; idx < FFT_LEN; idx += 4) {

					const unsigned ri_0 = bitRevBuf[idx >> 2];

					const unsigned ri_1 = ri_0 + 2 * len4; // bitRevBuf[idx + 1];

					const unsigned ri_2 = ri_0 + 1 * len4; // bitRevBuf[idx + 2];

					const unsigned ri_3 = ri_0 + 3 * len4; // bitRevBuf[idx + 3];

					dst[idx + 1] = input[ri_0] - input[ri_1];

					dst[idx + 3] = input[ri_2] - input[ri_3];

					const float sf_0 = input[ri_0] + input[ri_1];

					const float sf_2 = input[ri_2] + input[ri_3];

					dst[idx + 0] = sf_0 + sf_2;

					dst[idx + 2] = sf_0 - sf_2;

				}

			} else if constexpr (MODE == FFTInputMode::HalfZeroPadded) {

				// Optimized path: second (implied) half of input is zero

				// ri_1 and ri_3 always point to second half, so input[ri_1] == input[ri_3] == 0

				// This is a simplified version of the Full case above

				for (unsigned idx = 0; idx < FFT_LEN; idx += 4) {

					const unsigned ri_0 = bitRevBuf[idx >> 2];

					const unsigned ri_2 = ri_0 + 1 * len4; // bitRevBuf[idx + 2];

					dst[idx + 1] = input[ri_0];

					dst[idx + 3] = input[ri_2];

					dst[idx + 0] = input[ri_0] + input[ri_2];

					dst[idx + 2] = input[ri_0] - input[ri_2];

				}

			} else if constexpr (MODE == FFTInputMode::Symmetric) {

				// Symmetric input: input[0..N/2] provided, input[k] = input[N-k] for k > N/2

				// ri_0 and ri_2 always < N/2 (first half, use directly)

				// ri_1 and ri_3 always >= N/2 (second half, use mirrored index)

				for (unsigned idx = 0; idx < FFT_LEN; idx += 4) {

					const unsigned ri_0 = bitRevBuf[idx >> 2];

					const unsigned ri_1 = ri_0 + 2 * len4; // bitRevBuf[idx + 1];

					const unsigned ri_2 = ri_0 + 1 * len4; // bitRevBuf[idx + 2];

					const unsigned ri_3 = ri_0 + 3 * len4; // bitRevBuf[idx + 3];

					const float in_0 = input[ri_0];

					const float in_1 = input[FFT_LEN - ri_1];

					const float in_2 = input[ri_2];

					const float in_3 = input[FFT_LEN - ri_3];

					dst[idx + 1] = in_0 - in_1;

					dst[idx + 3] = in_2 - in_3;

					const float sf_0 = in_0 + in_1;

					const float sf_2 = in_2 + in_3;

					dst[idx + 0] = sf_0 + sf_2;

					dst[idx + 2] = sf_0 - sf_2;

				}

			} else {

				static_assert(always_false<MODE>, "Unsupported FFTInputMode");

			}

		} else if constexpr (PASS == 2) {

			// Start with 1st and 2nd passes (swap source and destination buffers).

			pass<1>(src, dst, input);

			// 3rd pass

			static constexpr float sqrt2_2 = std::numbers::sqrt2_v<float> * 0.5f;

			for (unsigned idx = 0; idx < FFT_LEN; idx += 8) {

				dst[idx + 0] = src[idx] + src[idx + 4];

				dst[idx + 4] = src[idx] - src[idx + 4];

				dst[idx + 2] = src[idx + 2];

				dst[idx + 6] = src[idx + 6];

				float v1 = (src[idx + 5] - src[idx + 7]) * sqrt2_2;

				dst[idx + 1] = src[idx + 1] + v1;

				dst[idx + 3] = src[idx + 1] - v1;

				float v2 = (src[idx + 5] + src[idx + 7]) * sqrt2_2;

				dst[idx + 5] = v2 + src[idx + 3];

				dst[idx + 7] = v2 - src[idx + 3];

			}

		} else {

			// First do the previous passes (swap source and destination buffers).

			pass<PASS - 1>(src, dst, input);

			static constexpr unsigned dist = 1 << (PASS - 1);

			static constexpr unsigned c1_r = 0;

			static constexpr unsigned c1_i = dist * 1;

			static constexpr unsigned c2_r = dist * 2;

			static constexpr unsigned c2_i = dist * 3;

			static constexpr unsigned cend = dist * 4;

			static constexpr unsigned table_step = COS_ARR_SIZE >> (PASS - 1);

			// For Symmetric input mode on final pass, skip imaginary outputs (they're ~0)

			// Final pass: dist = FFT_LEN/4, so c2_i and positions >= c2_r + 1 are imaginary

			// Only skip in release mode; in debug mode we compute them to validate they're ~0

			[[maybe_unused]] static constexpr bool isFinalPass = (PASS == FFT_LEN_L2 - 1);

#ifdef NDEBUG

			static constexpr bool skipImaginary = (MODE == FFTInputMode::Symmetric) && isFinalPass;

#else

			static constexpr bool skipImaginary = false;

#endif

			for (unsigned idx = 0; idx < FFT_LEN; idx += cend) {

				std::span<const float> sf = src.subspan(idx);

				std::span<      float> df = dst.subspan(idx);

				// Extreme coefficients are always real

				df[c1_r] = sf[c1_r] + sf[c2_r];

				df[c2_r] = sf[c1_r] - sf[c2_r];

				df[c1_i] = sf[c1_i];

				if constexpr (!skipImaginary) {

					df[c2_i] = sf[c2_i];  // Imaginary section

				}

				// Others are conjugate complex numbers

				for (unsigned i = 1; i < dist; ++ i) {

					const float c = cosBuf[        i  * table_step];

					const float s = cosBuf[(dist - i) * table_step];

					const float sf_r_i = sf[c1_r + i];

					const float sf_i_i = sf[c1_i + i];

					const float v1 = sf[c2_r + i] * c - sf[c2_i + i] * s;

					df[c1_r + i] = sf_r_i + v1;

					df[c2_r - i] = sf_r_i - v1;

					if constexpr (!skipImaginary) {

						const float v2 = sf[c2_r + i] * s + sf[c2_i + i] * c;

						df[c2_r + i] = v2 + sf_i_i;

						df[cend - i] = v2 - sf_i_i;

					}

				}

			}

		}

	}

	static constexpr auto bitRevBuf = []{

		constexpr int BR_ARR_SIZE = FFT_LEN / 4;

		std::array<uint16_t, BR_ARR_SIZE> result = {};

		for (auto cnt : xrange(narrow<unsigned>(result.size()))) {

			unsigned index = cnt << 2;

			unsigned res = 0;

			for (int bit_cnt = FFT_LEN_L2; bit_cnt > 0; --bit_cnt) {

				res <<= 1;

				res += (index & 1);

				index >>= 1;

			}

			result[cnt] = narrow<uint16_t>(res);

		}

		return result;

	}();

	static constexpr int COS_ARR_SIZE = FFT_LEN / 4;

	static constexpr auto cosBuf = []{

		std::array<float, COS_ARR_SIZE> result = {};

		const double mul = (0.5 * std::numbers::pi) / COS_ARR_SIZE;

		for (auto i : xrange(narrow<unsigned>(result.size()))) {

			result[i] = float(cstd::cos<4>(i * mul));

		}

		return result;

	}();

};

#endif