#include <cmath>
#include <iostream>
#include <iomanip>

#include <eigen3/unsupported/Eigen/FFT>

#include "base/consts.h"
#include "base/utils.h"
#include "tbpm/dense.h"


double fermi_dirac2(double x) {
    double kb_evk = 8.617333262e-5;
    double temperature = 100.0;
    double beta = 1.0 / (kb_evk * temperature);
    double mu = 0.0;
    double rescale = 2.0;
    beta *= rescale;
    mu /= rescale;
    double tmp = std::exp(beta * (x - mu));
    return 1 / (1 + tmp);
}


void func_origin()
{
    // Grid data
    int num_data = 32768;
    double theta = tbplas::base::PI / num_data;

    // Perform consine transform
    Eigen::VectorXd cos_theta(num_data); // cos((n+0.5)*theta), n = 0, num_data - 1
    Eigen::VectorXd sin_theta(num_data); // sin((n+0.5)*theta), n = 0, num_data - 1
    Eigen::VectorXd cos_k_theta(num_data); // cos(k*(n+0.5)*theta), k = 0, num_data - 1
    Eigen::VectorXd sin_k_theta(num_data); // sin(k*(n+0.5)*theta), k = 0, num_data - 1
    Eigen::VectorXd fd(num_data);

    // Initialize essential arrays
#pragma omp parallel for simd
    for (int n = 0; n < num_data; ++n) {
        cos_theta[n] = std::cos(theta * (n + 0.5));
        sin_theta[n] = std::sin(theta * (n + 0.5));
        cos_k_theta[n] = 1.0;
        sin_k_theta[n] = 0.0;
        fd[n] = fermi_dirac2(cos_theta[n]);
    }

    tbplas::base::Timer timer;
    timer.tic("func_origin");
    // Perform cosine transformation
    Eigen::VectorXd fft_result(num_data);
    fft_result[0] = tbplas::tbpm::vec_dot(fd, cos_k_theta);
    double cos_next = 0.0;
    double sin_next = 0.0;
    for (int k = 1; k < num_data; ++k) {
#pragma omp parallel for simd private(cos_next, sin_next)
        for (int n = 0; n < num_data; ++n) {
            cos_next = cos_theta[n] * cos_k_theta[n] - sin_theta[n] * sin_k_theta[n];
            sin_next = sin_theta[n] * cos_k_theta[n] + cos_theta[n] * sin_k_theta[n];
            cos_k_theta[n] = cos_next;
            sin_k_theta[n] = sin_next;
        }
        fft_result[k] = tbplas::tbpm::vec_dot(fd, cos_k_theta);
    }
    tbplas::tbpm::vec_self_mul_sca(fft_result, 2.0 / num_data);
    timer.toc("func_origin");
    timer.report_total_time();
    std::cout.setf(std::ios::fixed);
    std::cout.precision(15);
    for (int k = 0; k < num_data; ++k) {
        std::cout << std::setw(24) << fft_result[k] << "\n";
    }
}

void func_fft() {
    // Grid data
    int num_data = 32768 * 2;
    double theta = tbplas::base::PI / num_data * 2;

    // Perform consine transform
    Eigen::VectorXd fd(num_data);

    // Initialize essential arrays
#pragma omp parallel for simd
    for (int n = 0; n < num_data; ++n) {
        fd[n] = fermi_dirac2(std::cos(theta * n));
    }

    tbplas::base::Timer timer;
    timer.tic("func_fft");
    Eigen::FFT<double> fft_plan;
    Eigen::VectorXcd temp(num_data);
    fft_plan.fwd(temp, fd);
    Eigen::VectorXd fft_result = temp.real() * 2.0 / num_data;
    timer.toc("func_fft");
    timer.report_total_time();
    std::cout.setf(std::ios::fixed);
    std::cout.precision(15);
    for (int k = 0; k < num_data; ++k) {
        std::cout << std::setw(24) << fft_result[k] << "\n";
    }
}

int main()
{
    //func_origin();
    func_fft();
}