#include <cassert>
#include <concepts>
#include <format>
#include <iostream>
#include <math.h>
#include <numbers>
#include <ranges>
#include <vector>


/*
 *
 * Helper classes and functions
 *
 */


template <typename value_t>
class Complex {
public:
    constexpr Complex() noexcept : mRe(), mIm() {
    }

    constexpr Complex(const value_t& re, const value_t& im) noexcept
        : mRe(re), mIm(im) {
    }

    constexpr Complex operator-() const {
        return {-mRe, -mIm};
    }

    constexpr Complex operator+(const Complex& other) const {
        return {mRe + other.mRe, mIm + other.mIm};
    }

    constexpr Complex operator-(const Complex& other) const {
        return *this + (-other);
    }

    constexpr Complex operator*(const Complex& other) const {
        return {mRe * other.mRe - mIm * other.mIm,
                mIm * other.mRe + mRe * other.mIm};
    }

    operator std::string() {
        return std::format("{:.2f} + j*{:.2f}", mRe, mIm);
    }

    constexpr value_t& real() {
        return mRe;
    }

    constexpr value_t& imag() {
        return mIm;
    }

private:
    value_t mRe;
    value_t mIm;
};

using ComplexFloat = Complex<float>;


template <typename value_t>
constexpr void print(std::vector<Complex<value_t>> vec) {
    for (auto& elem : vec) {
        std::cout << "\t" << elem.real() << "; " << elem.imag() << std::endl;
    }
}


/*
 *
 * FFT implementation
 *
 */


template <std::integral value_t>
constexpr value_t bit_reverse(value_t n, std::size_t b = sizeof(value_t) * 8) {
    assert(b <= std::numeric_limits<value_t>::digits);

    value_t rv = 0;

    for (size_t i = 0; i < b; ++i, n >>= 1) {
        rv = (rv << 1) | (n & 0x01);
    }

    return rv;
}

constexpr std::vector<ComplexFloat>
fft_reorder(const std::vector<ComplexFloat>& x) {
    std::vector<ComplexFloat> result;

    for (int n = 0; n < x.size(); ++n) {
        const unsigned rev_n = bit_reverse(n, log2(x.size()));
        result.push_back(x[rev_n]);
    }

    return result;
}

//// TODO: Precompute Factors
//// TODO: Avoid unnecessary initialization cost of second buffer
//// TODO: Stream- or Vector-Hardware
// constexpr std::vector<ComplexFloat> fft(const std::vector<ComplexFloat>& x) {
//     const unsigned N0 = x.size();
//
//     std::array<std::vector<ComplexFloat>, 2> buffers;
//     buffers[0] = fft_reorder(x);
//     buffers[1].resize(N0);
//
//     for (int i_iter = 1; i_iter <= log2(N0); ++i_iter) {
//         std::vector<ComplexFloat>& input  = buffers[(i_iter - 1) % 2];
//         std::vector<ComplexFloat>& output = buffers[i_iter % 2];
//
//         const unsigned N            = std::pow(2, i_iter);
//         const unsigned num_sub_ffts = N0 / N;
//
//         for (int i_sub_fft = 0; i_sub_fft < num_sub_ffts; ++i_sub_fft) {
//             for (int k = 0; k < N; ++k) {
//                 float factor_re = std::cos((2 * std::numbers::pi * k) / N);
//                 float factor_im = -std::sin((2 * std::numbers::pi * k) / N);
//                 ComplexFloat factor = {factor_re, factor_im};
//
//                 const unsigned index1 = i_sub_fft * N + k % (N / 2);
//                 const unsigned index2 = i_sub_fft * N + N / 2 + k % (N / 2);
//
//                 output[i_sub_fft * N + k] =
//                         input[index1] + factor * input[index2];
//             }
//         }
//     }
//
//     return buffers[static_cast<int>(log2(N0)) % 2];
// }


std::vector<ComplexFloat> precompute_factors(unsigned N0) {
    std::vector<ComplexFloat> result;

    for (int N = 2; N <= N0; N *= 2) {
        for (int k = 0; k < N / 2; k++) {
            float factor_re = std::cos((2 * std::numbers::pi * k) / N);
            float factor_im = -std::sin((2 * std::numbers::pi * k) / N);

            result.push_back({factor_re, factor_im});
        }
    }

    return result;
}

constexpr unsigned get_factor_index(unsigned N, unsigned k) {
    return (N - 2) / 2 + k;
}


const std::vector<ComplexFloat> factors = precompute_factors(8);


// TODO: Stream- or Vector-Hardware
std::vector<ComplexFloat> fft(const std::vector<ComplexFloat>& input) {
    const unsigned N0 = input.size();

    std::vector<ComplexFloat> x = fft_reorder(input);

    for (int N = 2; N <= N0; N *= 2) {
        for (int i_sub_fft = 0; i_sub_fft < N0; i_sub_fft += N) {
            for (int k = 0; k < N / 2; k++) {
                const ComplexFloat& factor = factors[get_factor_index(N, k)];

                ComplexFloat z           = factor * x[i_sub_fft + k + N / 2];
                x[i_sub_fft + k + N / 2] = x[i_sub_fft + k] - z;
                x[i_sub_fft + k]         = x[i_sub_fft + k] + z;
            }
        }
    }

    return x;
}


/*
 *
 * Usage
 *
 */


int main() {
    std::vector<ComplexFloat> x = {{1, 0}, {2, 0}, {3, 0}, {4, 0},
                                   {5, 0}, {6, 0}, {7, 0}, {8, 0}};
    //        std::vector<ComplexFloat> x = {{1, 0}, {2, 0}, {3, 0}, {4, 0}};

    const auto& X = fft(x);

    //    int N0 = 8;
    //    for (int N = 2; N <= N0; N *= 2) {
    //        for (int k = 0; k < N / 2; k++) {
    //            std::cout << get_factor_index(N, k) << std::endl;
    //        }
    //    }

    std::cout << "================" << std::endl;
    std::cout << "result: " << std::endl;
    print(X);

    return 0;
}