//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Base/Math/FourierTransform.cpp
//! @brief     Implements class FourierTransform.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2015
//! @authors   Scientific Computing Group at MLZ Garching
//
//  ************************************************************************************************

#include "Base/Math/FourierTransform.h"
#include "Base/Util/Assert.h"
#include <algorithm>
#include <cmath>

namespace {

template <typename T> std::vector<T> fftshift_1d(const std::vector<T>& src, bool inverse)
{
    std::vector<T> result = src;

    // Centering FT around zero frequency
    int col_shift = (result.size() + 1) / 2;

    if (inverse)
        std::rotate(result.rbegin(), result.rbegin() + col_shift, result.rend());
    else
        std::rotate(result.begin(), result.begin() + col_shift, result.end());
    return result;
}

template <typename T> Field2D<T> fftshift_2d(const Field2D<T>& src, bool inverse)
{
    Field2D<T> result = src;

    // Centering FT around zero frequency
    int row_shift = (result.size() + 1) / 2;
    int col_shift = (result[0].size() + 1) / 2;

    // First shift the rows, then the cols
    if (inverse) {
        std::rotate(result.rbegin(), result.rbegin() + row_shift, result.rend());
        for (auto& row : result)
            std::rotate(row.rbegin(), row.rbegin() + col_shift, row.rend());
    } else {
        std::rotate(result.begin(), result.begin() + row_shift, result.end());
        for (auto& row : result)
            std::rotate(row.begin(), row.begin() + col_shift, row.end());
    }
    return result;
}

} // namespace

FourierTransform::FourierTransform() = default;

FourierTransform::Workspace::~Workspace()
{
    clear();
}

void FourierTransform::Workspace::clear()
{
    // rows (h) and columns (w)
    h = 0;
    w_real = 0;
    w_fftw = 0;

    if (arr_real != nullptr) {
        fftw_free((double*)arr_real);
        arr_real = nullptr;
    }

    if (arr_fftw != nullptr) {
        fftw_free((fftw_complex*)arr_fftw);
        arr_fftw = nullptr;
    }

    if (p_forw != nullptr) {
        fftw_destroy_plan(p_forw);
        p_forw = nullptr;
    }

    if (p_back != nullptr) {
        fftw_destroy_plan(p_back);
        p_back = nullptr;
    }
}

/* ************************************************************************* */
// Fourier Transform in 2D
/* ************************************************************************* */
complex2d_t FourierTransform::rfft(const double2d_t& src)
{
    // rows (h) and columns (w) of the input 'source'
    int h = src.size();
    int w_real = !src.empty() ? src[0].size() : 0;

    // initialization
    init_r2c(h, w_real);

    // Compute the forward FT
    fftw_forward_FT(src);

    // Pack the result as std vector of complex values
    return rfft2complex_vec();
}

double2d_t FourierTransform::irfft(const complex2d_t& src, int w_real)
{
    // rows (h) of the input 'source'
    int h = src.size();

    // initialization
    init_c2r(h, w_real);

    // Compute the unnormalized backward FT
    fftw_backward_FT(src);

    // Normalize aack the result as std vector of real values
    return irfft2double_vec();
}

double2d_t FourierTransform::ramplitude(const double2d_t& src)
{
    complex2d_t complex_result = rfft(src);
    return fft2amp(complex_result);
}

/* ************************************************************************* */
// Fourier Transform 2D shift - center array around zero frequency
/* ************************************************************************* */
double2d_t FourierTransform::fftshift(const double2d_t& src)
{
    return ::fftshift_2d(src, false);
}
complex2d_t FourierTransform::fftshift(const complex2d_t& src)
{
    return ::fftshift_2d(src, false);
}

double2d_t FourierTransform::ifftshift(const double2d_t& src)
{
    return ::fftshift_2d(src, true);
}
complex2d_t FourierTransform::ifftshift(const complex2d_t& src)
{
    return ::fftshift_2d(src, true);
}

/* ************************************************************************* */
// Fourier Transform in 1D
/* ************************************************************************* */
std::vector<complex_t> FourierTransform::rfft(const std::vector<double>& src)
{
    // we simply create 2d arrays with length of first dimension equal to 1, and call 2d FT
    double2d_t src2d{src};
    complex2d_t result2d = rfft(src2d);

    ASSERT(result2d.size() == 1);
    return result2d[0];
}

std::vector<double> FourierTransform::irfft(const std::vector<complex_t>& src, int w_src)
{
    // we simply create 2d arrays with length of first dimension equal to 1, and call 2d FT
    complex2d_t source2d{src};
    double2d_t result2d = irfft(source2d, w_src);

    ASSERT(result2d.size() == 1);
    return result2d[0];
}

std::vector<double> FourierTransform::ramplitude(const std::vector<double>& src)
{
    complex2d_t complex_result2d{rfft(src)};

    double2d_t result2d = fft2amp(complex_result2d);

    ASSERT(result2d.size() == 1);
    return result2d[0];
}

/* ************************************************************************* */
// Fourier Transform 1D shift - center around zero frequency
/* ************************************************************************* */
std::vector<double> FourierTransform::fftshift(const std::vector<double>& src)
{
    return ::fftshift_1d(src, false);
}
std::vector<complex_t> FourierTransform::fftshift(const std::vector<complex_t>& src)
{
    return ::fftshift_1d(src, false);
}

std::vector<double> FourierTransform::ifftshift(const std::vector<double>& src)
{
    return ::fftshift_1d(src, true);
}
std::vector<complex_t> FourierTransform::ifftshift(const std::vector<complex_t>& src)
{
    return ::fftshift_1d(src, true);
}

/* ************************************************************************************ */
// initialize input and output arrays in workspace for fast Fourier transformation (fftw)
/* ************************************************************************************ */

void FourierTransform::init(int h, int w_real)
{
    ASSERT(h);
    ASSERT(w_real);

    ws.clear();
    ws.h = h;
    ws.w_real = w_real;
    ws.w_fftw = w_real / 2 + 1;

    ws.arr_real = fftw_alloc_real(ws.h * ws.w_real);
    ws.arr_fftw = fftw_alloc_real(ws.h * ws.w_fftw * sizeof(fftw_complex));
}

void FourierTransform::init_r2c(int h, int w_real)
{
    init(h, w_real);

    // Initialization of the plan
    ws.p_forw = fftw_plan_dft_r2c_2d(ws.h, ws.w_real, ws.arr_real, (fftw_complex*)ws.arr_fftw,
                                     FFTW_ESTIMATE);
    ASSERT(ws.p_forw);
}

void FourierTransform::init_c2r(int h, int w_real)
{
    init(h, w_real);

    // Initialization of the plan
    ws.p_back = fftw_plan_dft_c2r_2d(ws.h, 2 * ws.w_fftw, (fftw_complex*)ws.arr_fftw, ws.arr_real,
                                     FFTW_ESTIMATE);
    ws.p_back = fftw_plan_dft_c2r_2d(ws.h, ws.w_real, (fftw_complex*)ws.arr_fftw, ws.arr_real,
                                     FFTW_ESTIMATE);
    ASSERT(ws.p_back);
}

/* ************************************************************************************ */
// convert output to standard containers
/* ************************************************************************************ */

complex2d_t FourierTransform::rfft2complex_vec() const
{
    ASSERT(ws.arr_fftw);

    complex2d_t result(ws.h, std::vector<complex_t>(ws.w_fftw));
    double* ptr = ws.arr_fftw;

    // Storing FT output
    for (int i = 0; i < ws.h; i++)
        for (int j = 0; j < ws.w_fftw; j++) {
            result[i][j] = complex_t(*ptr, *(ptr + 1));
            ptr += 2;
        }
    return result;
}

double2d_t FourierTransform::irfft2double_vec() const
{
    ASSERT(ws.arr_real);

    double* ptr = ws.arr_real;
    double2d_t result(ws.h, std::vector<double>(ws.w_real));

    double factor = ws.h * ws.w_real;

    // Storing FT output
    for (int i = 0; i < ws.h; i++)
        for (int j = 0; j < ws.w_real; j++) {
            result[i][j] = *ptr / factor;
            ptr++;
        }
    return result;
}

double2d_t FourierTransform::fft2amp(complex2d_t& source) const
{
    ASSERT(ws.arr_fftw);

    double2d_t result(ws.h, std::vector<double>(ws.w_real));

    // Storing FT output
    for (int i = 0; i < ws.h; i++) {
        int k = (i == 0) ? 0 : ws.h - i;
        for (int j = 0; j < ws.w_fftw; j++) {
            result[i][j] = std::hypot(source[i][j].real(), source[i][j].imag());
            if (j != 0)
                result[k][ws.w_real - j] = result[i][j];
        }
    }
    return result;
}

/* ************************************************************************* */
// initialize input and output arrays for fast Fourier transformation
/* ************************************************************************* */

void FourierTransform::fftw_forward_FT(const double2d_t& src) const
{
    // Initializing the content of input array to zero for all elements (Do we need this at all?)
    double* ptr_in_end = ws.arr_real + ws.h * ws.w_real;
    for (double* ptr = ws.arr_real; ptr != ptr_in_end; ++ptr)
        *ptr = 0.0;

    // Building the input signal for fftw algorithm
    for (int row = 0; row < ws.h; ++row)
        for (int col = 0; col < ws.w_real; ++col)
            ws.arr_real[row * ws.w_real + col] += src[row][col];

    // Computing the FFT with fftw plan
    fftw_execute(ws.p_forw);
}

void FourierTransform::fftw_backward_FT(const complex2d_t& src) const
{
    // Initializing the content of input array to zero for all elements (Do we need this at all?)
    double* ptr_in_end = ws.arr_fftw + 2 * ws.h * ws.w_fftw;
    for (double* ptr = ws.arr_fftw; ptr != ptr_in_end; ++ptr)
        *ptr = 0.0;

    // Building the fft input signal
    double* ptr = ws.arr_fftw;
    for (int row = 0; row < ws.h; ++row)
        for (int col = 0; col < ws.w_fftw; ++col) {
            *ptr += src[row][col].real();
            *(ptr + 1) += src[row][col].imag();
            ptr += 2;
        }

    // Computing the IFFT with fftw plan
    fftw_execute(ws.p_back);
}