/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * 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++)
    {
        // 这里我们把2.0改成了1.0
        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++)
    {
        // 这里我们删除了 *0.5
        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;
}

#else

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 n = in_real.size();
    int s = n/2 + 1; // out size
    int m = log2(n);

    double *ar, *ai;
    ar = new double [n];
    ai = new double [n];
    for (size_t i = 0; i < n; i++)
    {
        ar[i] = in_real[i];
        ai[i] = 0.0;
    }

    int i,j,k,k1,l1,l2,ip;
	double *ir,*ii,*i1,*i2;
	double f1,t1,t2,v,u,w1,w2,d,vu,pi1;
	
	f1=(double)1;
	j=0;
	for(i=0,ir=ar,ii=ai;i<n;i++,ir++,ii++)
    {	if(i<j)
        {	i1=ar+j; i2=ai+j;
            t1=*i1; t2=*i2;	*i1=*ir; *i2=*ii; *ir=t1; *ii=t2;
        }
        k1=n/2;
        while((j>=k1)&&(k1!=0))
        {	j=j-k1; k1=k1/2;
        }
        j=j+k1;
    }

	for(j=1;j<m+1;j++)
	{	l1=2;
		for (i=1;i<j;i++) l1=2*l1;
		l2=l1/2; pi1=(double)GCTL_Pi/(double)l2; v=(double)1; u=(double)0;
		w1=(double)cos(pi1); w2=(double)sin(pi1);
		for(k=1;k<l2+1;k++)
		{	d=f1*u;
			for(i=k-1;i<n-l1+k;i+=l1)
			{	ip=i+l2; i1=ar+ip; i2=ai+ip; ir=ar+i; ii=ai+i;
				t1=*i1*v+*i2*d; t2=*i2*v-*i1*d;
				*i1=*ir-t1; *i2=*ii-t2; *ir=*ir+t1; *ii=*ii+t2;
			}
			vu=v*w1-u*w2; u=v*w2+u*w1; v=vu;
	}	}

	for(i=0,ir=ar,ii=ai;i<n;i++,ir++,ii++)
    {
        *ir/=n; *ii/=n;
    }

    out_spectrum.resize(s);
    for (int i = 0; i < s; i++)
    {
        out_spectrum[i].real(ar[i]);
        out_spectrum[i].imag(ai[i]);
    }

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

    delete[] ar;
    delete[] ai;
    return;
}

void gctl::dft_c2r_1d(const _1cd_array &in_spectrum, _1d_array &out_real)
{
    return;
}

void gctl::dft2d(const _2d_matrix &in_arr, _1cd_array &out_arr)
{
    return;
}

void gctl::idft2d(const _1cd_array &in_arr, int m, int n, _2d_matrix &out_arr)
{
    return;
}

void gctl::dct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
    return;
}

void gctl::idct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
    return;
}

#endif // GCTL_FFTW3