/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2023  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "fft.h"

#ifdef GCTL_FFTW3

void gctl::dft_r2c_1d(const _1d_array &in_real, _1cd_array &out_spectrum, double sampling, _1d_array *freq_ptr)
{
    if (in_real.empty())
    {
        throw std::invalid_argument("[GCTL] Invalid input array size for gctl::dft_r2c_1d(...)");
    }

    int m = in_real.size();
    int n = m/2 + 1; // out size
    double *in = fftw_alloc_real(m);
    fftw_complex *out= fftw_alloc_complex(n);

    fftw_plan p = fftw_plan_dft_r2c_1d(m, in, out, FFTW_ESTIMATE);

    for (int i = 0; i < m; i++)
    {
        in[i] = in_real[i];
    }

    fftw_execute(p);
    fftw_destroy_plan(p);

    out_spectrum.resize(n);
    for (int i = 0; i < n; i++)
    {
        out_spectrum[i].real(out[i][0]*2.0/m);
        out_spectrum[i].imag(out[i][1]*2.0/m);
    }

    if (freq_ptr != nullptr)
    {
        freq_ptr->resize(n);
        for (size_t i = 0; i < n; i++)
        {
            freq_ptr->at(i) = 0.5*sampling*i/n;
        }
    }

    fftw_free(in); 
    fftw_free(out);
    return;
}

void gctl::dft_c2r_1d(const _1cd_array &in_spectrum, _1d_array &out_real)
{
    if (in_spectrum.empty())
    {
        throw std::invalid_argument("[GCTL] Invalid input array size for gctl::dft_c2r_1d(...)");
    }

    int m = in_spectrum.size();
    int n = 2*(m-1);
    fftw_complex *in = fftw_alloc_complex(m);
    double *out= fftw_alloc_real(n);

    fftw_plan p = fftw_plan_dft_c2r_1d(n, in, out, FFTW_ESTIMATE);

    for (int i = 0; i < m; i++)
    {
        in[i][0] = in_spectrum[i].real()*0.5*n;
        in[i][1] = in_spectrum[i].imag()*0.5*n;
    }

    fftw_execute(p);
    fftw_destroy_plan(p);

    out_real.resize(n);
    for (int i = 0; i < n; i++)
    {
        out_real[i] = out[i]/n;
    }
    
    fftw_free(in); 
    fftw_free(out);
    return;
}

void gctl::dft2d(const _2d_matrix &in_arr, _1cd_array &out_arr)
{
    if (in_arr.empty())
    {
        throw length_error("The input array is empty. From gctl::dft2d(...)");
    }

    int m = in_arr.row_size();
    int n = in_arr.col_size();
    fftw_complex *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
    fftw_complex *out= (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);

    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < n; j++)
        {
            in[i*n+j][0] = in_arr[i][j];
            in[i*n+j][1] = 0.0;
        }
    }

    fftw_plan p = fftw_plan_dft_2d(m, n, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
    fftw_execute(p);

    out_arr.resize(m*n);
    for (int i = 0; i < m*n; i++)
    {
        out_arr[i].real(out[i][0]);
        out_arr[i].imag(out[i][1]);
    }

    fftw_destroy_plan(p);
    fftw_free(in); 
    fftw_free(out);
    return;
}

void gctl::idft2d(const _1cd_array &in_arr, int m, int n, _2d_matrix &out_arr)
{
    if (in_arr.empty())
    {
        throw length_error("The input array is empty. From gctl::idft2d(...)");
    }

    if (m <= 0 || n <= 0 || m*n != in_arr.size())
    {
        throw invalid_argument("Invalid parameter. From gctl::idft2d(...)");
    }

    fftw_complex *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
    fftw_complex *out= (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);

    for (int i = 0; i < m*n; i++)
    {
        in[i][0] = in_arr[i].real();
        in[i][1] = in_arr[i].imag();
    }

    fftw_plan p = fftw_plan_dft_2d(m, n, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
    fftw_execute(p);

    out_arr.resize(m, n);
    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < n; j++)
        {
            out_arr[i][j] = out[i*n+j][0]/(m*n);
        }
    }

    fftw_destroy_plan(p);
    fftw_free(in); 
    fftw_free(out);
    return;
}

void gctl::dct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
    if (in_arr.empty())
    {
        throw length_error("The input array is empty. From gctl::dct1d(...)");
    }

    int m = in_arr.size();
    out_arr.resize(m);

    fftw_plan p = fftw_plan_r2r_1d(m, in_arr.get(), out_arr.get(), 
        FFTW_REDFT10, FFTW_ESTIMATE);
    fftw_execute(p);

    fftw_destroy_plan(p);
    return;
}

void gctl::idct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
   if (in_arr.empty())
    {
        throw length_error("The input array is empty. From gctl::idct1d(...)");
    }

    int m = in_arr.size();
    out_arr.resize(m);

    fftw_plan p = fftw_plan_r2r_1d(m, in_arr.get(), out_arr.get(), 
        FFTW_REDFT01, FFTW_ESTIMATE);
    fftw_execute(p);

    for (int i = 0; i < m; i++)
    {
        out_arr[i] = out_arr[i]*0.5/m;
    }

    fftw_destroy_plan(p);
    return;
}

#endif // GCTL_FFTW3