#include <stdlib.h>
#include <math.h>
#include <time.h>               // for RNG seed
#include <fftw3.h>              // for Fourier Transforms
#include <gsl/gsl_rng.h>        // for random number generation

#define M_PI 3.14159265358979323846264338327

void example_FFTW_inplace();
void example_FFTW();
void example_FFTW_GURU();
void example_FFTW_2D();

void initialiseData(double *, int, int, int);
void printDataReal(double * , int , int , int );
void printDataFourier(double * , int , int , int );

int seed;

enum {FALSE, TRUE};


int main(void)
{
    seed = (int) time(NULL);

    example_FFTW_2D();

    return 0;
}





void example_FFTW_2D()
{

    printf ("\n\n");
    printf ("==============================================================\n");
    printf ("Example of taking Fourier transforms with FFTW (out of place):\n");
    printf ("==============================================================\n");


    // array dimensions
    int Nx = 8; int Ny = 8; int Nz = 3;

    const int dims = Nx * Ny * Nz;

    // allocate memory
    double *input = fftw_alloc_real(dims);              // input data (pre Fourier transform)
    double *input2 = fftw_alloc_real(dims);             // store copy of input data so we can compute accuracy


    const int outdims = Nx * (Ny/2 + 1) * Nz;
    fftw_complex *output = fftw_alloc_complex(outdims);    // we want to perform the transform out of place (so seperate array for output)

#define input(i,j,k) ( input[(k) + Nz * ((j) + Ny * (i))] )
#define input2(i,j,k) ( input2[(k) + Nz * ((j) + Ny * (i))] )


    ////////////////////////////////////////////////////
    // setup "plans" for forward and backward transforms
    const int rank = 2; const int howmany = Nz;
    const int istride = Nz; const int ostride = Nz;
    const int idist = 1; const int odist = 1;
    int n[] = {Nx, Ny};
    int *inembed = NULL, *onembed = NULL;


    fftw_plan fp = fftw_plan_many_dft_r2c(rank, n, howmany,
            input, inembed, istride, idist,
            output, onembed, ostride, odist,
            FFTW_PATIENT);

    fftw_plan bp = fftw_plan_many_dft_c2r(rank, n, howmany,
            output, onembed, ostride, odist,
            input, inembed, istride, idist,
            FFTW_PATIENT);



    ///////////////////////////////////////////
    // initialise the data (important that this
    // is done after initialisation)
    initialiseData(input, Nx, Ny, Nz);
    initialiseData(input2, Nx, Ny, Nz);

    // and print for record
    printDataReal(input, Nx, Ny, Nz);


    /////////////////////////////
    // take the forward transform
    fftw_execute(fp); printf ("FFT complete!\n");

    // and print for record
    printDataFourier((double *) output, Nx, Ny, Nz);




    /////////////////////////////
    // take the backward transform
    fftw_execute(bp); printf ("IFFT complete!\n");

    double maxError = 0.0;

    for (int i = 0; i < Nx; ++i)
        for (int j = 0; j < Ny; ++j)
            for (int d = 0; d < Nz; ++d)
            {
                // remember that FFTW doesn't normalise
                input(i,j,d) /= (double) (Nx*Ny);

                const double error = fabs(input2(i,j,d) - input(i,j,d));
                if (error > maxError) maxError = error;
            }

    printDataReal(input, Nx, Ny, Nz);

    printf ("maximum error after 1 cycle: %.4g\n", maxError);

    fftw_free(input);
    fftw_free(input2);
    fftw_free(output);


    // clean up FTTW stuff
    fftw_destroy_plan(fp);
    fftw_destroy_plan(bp);

    fftw_cleanup();
}





void initialiseData(double * input, int Nx, int Ny, int Nz)
{
    // initialise random number generator with Mersenne Twister
    gsl_rng * rng = gsl_rng_alloc ( gsl_rng_mt19937 );
    // change this to try a different seed
    gsl_rng_set(rng, seed);


    // setup input data
    for (int k = 0; k < Nz; ++k )
        for (int i = 0; i < Nx; ++i)
            for (int j = 0; j < Ny; ++j)
            {

                if (k == 0)
                    input(i,j,k) = gsl_rng_uniform(rng);
                else
                {
                    const double x = (double) i / (double) Nx;
                    const double y = (double) j / (double) Ny;

                    input(i,j,k) = cos(2* M_PI * y);
                    input(i,j,k) += cos(4* M_PI * x);

                    input(i,j,k) *= (k + 1);
                }

            }

    gsl_rng_free(rng);
}

void printDataReal(double * dataToPrint, int Nx, int Ny, int Nz)
{

    printf ("Data in real space\n");
#define dataToPrint(i,j,k) ( dataToPrint[(k) + Nz * ((j) + Ny * (i))] )

    // print input data
    for (int k = 0; k < Nz; ++k )
    {
        printf ("\nk = %d\n", k);

        for (int i = 0; i < Nx; ++i)
        {
            for (int j = 0; j < Ny; ++j)
                printf ("i[%3d][%d][%d] = %.3f\t", i, j, k, dataToPrint(i,j,k));

            printf ("\n");
        }
    }

    printf ("\n\n\n");
}


void printDataFourier(double * dataToPrint, int Nx, int Ny, int Nz)
{

    printf ("Data in Fourier space\n");

#define dataToPrintRe(qx,qy,k) ( dataToPrint[2*((k) + Nz * ((qy) + (Ny/2 + 1) * (qx)))    ] )
#define dataToPrintIm(qx,qy,k) ( dataToPrint[2*((k) + Nz * ((qy) + (Ny/2 + 1) * (qx))) + 1] )

    // print input data
    for (int k = 0; k < Nz; ++k )
    {
        printf ("\nk = %d\n", k);

        for (int qx = 0; qx <= Nx/2; ++qx)
        {
            for (int qy = 0; qy <= Ny/2; ++qy)
                printf ("o[%3d][%d][%d] = %.3f + %.3fi\t", qx, qy, k, dataToPrintRe(qx, qy, k), dataToPrintIm(qx, qy, k));

            printf ("\n");
        }
    }

    printf ("\n\n\n");
}