//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Base/Math/FourierTransform.h
//! @brief     Defines class Math::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
//
//  ************************************************************************************************

#ifdef SWIG
#error no need to expose this header to Swig
#endif // SWIG
#ifndef BORNAGAIN_BASE_MATH_FOURIERTRANSFORM_H
#define BORNAGAIN_BASE_MATH_FOURIERTRANSFORM_H

#include "Base/Type/Field2D.h"
#include <fftw3.h>

//! Fourier transform of vectors (in 1D or 2D) using Fast Fourier Transform (fftw package).
class FourierTransform {
public:
    FourierTransform();

    // forward transform
    std::vector<complex_t> rfft(const std::vector<double>& src);
    complex2d_t rfft(const double2d_t& src);

    // forward transform with real-valued amplitudes in output
    std::vector<double> ramplitude(const std::vector<double>& src);
    double2d_t ramplitude(const double2d_t& src);

    // backward transform
    std::vector<double> irfft(const std::vector<complex_t>& src, int w_src);
    double2d_t irfft(const complex2d_t& src, int w_real);

    // Frequency shifts are useful for providing different representations of spectrum: with lowest
    // frequency at edges of array (default output of fft) or with lowest frequency in center of
    // array

    // shift low frequency from corners to center
    static std::vector<double> fftshift(const std::vector<double>& src);
    static std::vector<complex_t> fftshift(const std::vector<complex_t>& src);
    static double2d_t fftshift(const double2d_t& src);
    static complex2d_t fftshift(const complex2d_t& src);

    // shift low frequency from center to corners
    static std::vector<double> ifftshift(const std::vector<double>& src);
    static std::vector<complex_t> ifftshift(const std::vector<complex_t>& src);
    static double2d_t ifftshift(const double2d_t& src);
    static complex2d_t ifftshift(const complex2d_t& src);

private:
    // prepare arrays and workspaces for FT
    void init(int h, int w_real);
    void init_r2c(int h, int w_real);
    void init_c2r(int h, int w_real);

    // convert FT output to the 2D vector
    complex2d_t rfft2complex_vec() const;
    double2d_t irfft2double_vec() const;

    // convert complex FT to the real amplitudes
    double2d_t fft2amp(complex2d_t& source) const;

    // compute forward/backward FT of source using Fast Fourier transformation from fftw
    void fftw_forward_FT(const double2d_t& src) const;
    void fftw_backward_FT(const complex2d_t& src) const;

    //! Workspace for Fourier Transform.

    //! Workspace contains input, intermediate and output
    //! arrays to run FT via fft; 'real' is our signal
    //! Output arrays are allocated via fftw3 allocation for maximum performance.
    class Workspace {
    public:
        Workspace() = default;
        ~Workspace();
        void clear();
        friend class FourierTransform;

    private:
        //! Here, h = height (# rows), w = width (# columns)
        int h{0};      // height of input AND output arrays in 2D
        int w_real{0}; // size of input (r2c) or output (c2r) real-valued array in 2D
        int w_fftw{0}; // size of output (r2c) or input (c2r) 'complex-valued FT' array in 2D

        double* arr_real{nullptr}; // pointer to input/output real-valued array
        double* arr_fftw{nullptr}; // pointer to input/output 'complex-valued FT' array

        fftw_plan p_forw{nullptr};
        fftw_plan p_back{nullptr};
    };

    //! input and output data for fftw3
    Workspace ws;
};

#endif // BORNAGAIN_BASE_MATH_FOURIERTRANSFORM_H
