#include <cmath>
#include <iomanip>
#include <complex>
#include <iostream>

#include <fftw3.h>

static const int nx = 8;
static const int ny = 8;
static const int nyk = ny / 2 + 1;

static constexpr int width = 8;

using namespace std;

// Access the element at column x, row y in an array.  xy are cartesian
// coordinates.  
int xy(int x, int y, int number_of_cols) { return x + y * number_of_cols; }
void print(double* m, int width, int height)
{
  std::cout << std::fixed << std::setprecision(2);
  for (int y = 0; y < height; ++y)
  {
    for (int x = 0; x < width; ++x)
      std::cout << std::setw(6) << m[xy(x, y, width)] << ' ';
    std::cout << '\n';
  }

  std::cout << '\n';
}

void print(fftw_complex* m, int width, int height)
{
  std::cout << std::fixed << std::setprecision(2);
  for (int y = 0; y < height; ++y)
  {
    for (int x = 0; x < width; ++x)
      std::cout << std::setw(6) << std::complex{m[xy(x, y, width)][0], m[xy(x, y, width)][1]} <<  ' ';
    std::cout << '\n';
  }

  std::cout << '\n';
}

int main(){	
  double* uFun        = fftw_alloc_real(nx * (ny + 1));
  fftw_complex* uFunk = fftw_alloc_complex(nx * (ny + 1));
  double* Out         = fftw_alloc_real(nx * (ny + 1));
  
  // This plan computes independent one-dimensional DFTs of each row of the
  // array.
  static fftw_plan r2c_plan = fftw_plan_many_dft_r2c(
    1,                          // int rank,
    &nx,                        // const int *n,
    ny + 1,                     // int howmany,
    uFun,                       // double *in,
    nullptr,                    // const int *inembed,
    1,                          // int istride, // adjacent elements within a row are adjacent.
    nx,                         // int idist,   // adjacent elements within a column are separate
    uFunk,                      // fftw_complex *out,
    nullptr,                    // const int *onembed,
    1,                          // int ostride, 
    nx,                         // int odist,
    FFTW_ESTIMATE               // unsigned flags
    );
  
  static fftw_plan c2r_plan = fftw_plan_many_dft_c2r(
    1,                          // int rank,
    &nx,                        // const int *n,
    ny + 1,                     // int howmany,
    uFunk,                      // double *in,
    nullptr,                    // const int *inembed,
    1,                          // int istride, // adjacent elements within a row are adjacent.
    nx,                         // int idist,   // adjacent elements within a column are separate
    Out,                        // fftw_complex *out,
    nullptr,                    // const int *onembed,
    1,                          // int ostride, 
    nx,                         // int odist,
    FFTW_ESTIMATE               // unsigned flags
    );

  for(int y = 0; y < ny + 1; y++) 
    for(int x = 0; x < nx; x++)
    {
      double const XXxy = +std::sin(x * 2.0 * 3.14159 / nx);
      double const YYxy = -std::cos(y * 3.14159 / ny) - 0.5;
      uFun[xy(x, y, nx)] = YYxy * YYxy * XXxy;
    }
  
  print(uFun, nx, ny+1);
  fftw_execute_dft_r2c(r2c_plan, uFun, uFunk);
  print(uFunk, nx, ny+1);
  fftw_execute_dft_c2r(c2r_plan, uFunk, Out);

  // renormalize:
  for(int y = 0; y < ny + 1; y++) 
    for(int x = 0; x < nx; x++)
      Out[xy(x, y, nx)] /= nx;
  
  print(Out, nx, ny+1);
  
  // Good practice: free stuff in reverse order of initialization
  fftw_free(Out);
  fftw_free(uFunk);
  fftw_free(uFun);
}

#include <cmath>
#include <iomanip>
#include <complex>
#include <cstring>
#include <iostream>

#include <fftw3.h>

static const int nx = 8;
static const int ny = 8;
static const int ncomp = 2;
static const int nyk = ny/2 + 1;
static const int nxk = nx/2 + 1;
using namespace std;

// NOTE(mbozzi): Non-static member data was removed from these classes because
// it is only used in the constructor.
class MyFftwClass {
public:
  MyFftwClass(void);

  // ~MyFftwClass(void);
  void r2cexecute(double rArr[], double cArr[][ncomp]);

private:
  static fftw_plan s_plan; // <-- shared by all instances
  // double *m_buffer_in;
  // fftw_complex *m_buffer_out;
};

class MyiFftwClass { //inverse
public:
  MyiFftwClass(void);
  // ~MyiFftwClass(void);
  void c2rexecute(double cArr[][ncomp],double rArr[]);
private:
  static fftw_plan s_plan1; // <-- shared by all instances
  // fftw_complex *m_buffer_in1;
  // double *m_buffer_out1;
};

MyFftwClass r2c;
MyiFftwClass c2r;

// Assumes row-major storage where i = row index, j = column index.
int ij(int i, int j, int number_of_cols) { return j + i * number_of_cols; }

void print(double* m, int width, int height)
{
  std::cout << std::fixed << std::setprecision(2);
  for (int i = 0; i < height; ++i)
  {
    for (int j = 0; j < width; ++j)
      std::cout << std::setw(6) << m[ij(i, j, width)] << ' ';
    std::cout << '\n';
  }

  std::cout << '\n';
}

void print(fftw_complex* m, int width, int height)
{
  std::cout << std::fixed << std::setprecision(2);
  for (int i = 0; i < height; ++i)
  {
    for (int j = 0; j < width; ++j)
      std::cout << std::setw(6) << std::complex{m[ij(i, j, width)][0], m[ij(i, j, width)][1]} <<  ' ';
    std::cout << '\n';
  }

  std::cout << '\n';
}

int main()
{
  double *uFun;
  uFun = (double*) fftw_malloc((nx*(ny+1))*sizeof(double));

  // for(int i = 0; i< nx; i++){
  //   for(int j = 0; j< ny+1; j++){
  //     uFun[j + (ny+1)*i] =//some real function
  //       }
  // }
  // NOTE(mbozzi): Using the same test function as earlier
  for(int i = 0; i < ny + 1; i++)
    for(int j = 0; j < nx; j++)
    {
      double const XX = +std::sin(j * 2.0 * 3.14159 / nx);
      double const YY = -std::cos(i * 3.14159 / ny) - 0.5;
      uFun[ij(i, j, nx)] = YY * YY * XX;
    }

  print(uFun, nx, ny + 1);

  // fftw_complex *uFunk;
  // uFunk = (fftw_complex*) fftw_malloc((nx*nyk) * sizeof(fftw_complex));
  // memset(uFunk, 42, (nx*nyk)* sizeof(fftw_complex));
  fftw_complex *uFunk;
  uFunk = (fftw_complex*) fftw_malloc((nxk * (ny + 1)) * sizeof(fftw_complex));
  memset(uFunk, 42, (nxk*(ny + 1))* sizeof(fftw_complex));

  r2c.r2cexecute(uFun,uFunk); //take forward FFT

  print(uFunk, nxk, ny + 1);

  double *Out;
  Out = (double*) fftw_malloc((nx*(ny+1))*sizeof(double));
  memset(Out, 42, (nx*(ny+1)) * sizeof(double));

  c2r.c2rexecute(uFunk,Out);

  print(Out, nx, ny + 1);
}

MyFftwClass::MyFftwClass(void)
{
  // m_buffer_in = fftw_alloc_real((nx*ny));
  // m_buffer_out = fftw_alloc_complex((nx*nyk));
  double* buffer_in = fftw_alloc_real(nx * (ny + 1));
  fftw_complex* buffer_out = fftw_alloc_complex(nxk * (ny + 1));
  if (!(buffer_in && buffer_out))
  {
    // throw new std::runtime_error("Failed to allocate memory!");
    throw std::runtime_error("Failed to allocate memory!");
  }

  if (!s_plan)
  {
    // s_plan = fftw_plan_many_dft_r2c(1,&nx, (ny+1),m_buffer_in,nullptr,(ny+1),1,m_buffer_out,nullptr, (ny+1),1, FFTW_MEASURE);
    s_plan = fftw_plan_many_dft_r2c(
      1,           // int rank,
      &nx,         // const int *n,
      ny + 1,      // int howmany,
      buffer_in,   // double *in,
      nullptr,     // const int *inembed,
      1,           // int istride, // adjacent elements within a row are adjacent.
      nx,          // int idist,   // adjacent elements within a column are separate
      buffer_out,  // fftw_complex *out,
      nullptr,     // const int *onembed,
      1,           // int ostride,
      nxk,         // int odist,
      FFTW_MEASURE // unsigned flags
      );

    if (!s_plan)
    {
      // throw new std::runtime_error("Failed to create plan!");
      throw std::runtime_error("Failed to create plan!");
    }
  }

  fftw_free(buffer_out);
  fftw_free(buffer_in);
}

// MyFftwClass::~MyFftwClass(void)
// {
//   fftw_free(m_buffer_in);
//   fftw_free(m_buffer_out);
// }

// void MyFftwClass::r2cexecute(double rArr[], double cArr[][ncomp]) MyFftwClass::execute(double *const in, fftw_complex *const out)
// {
//     memcpy(m_buffer_in, rArr,  sizeof(double) * (nx*ny));
//     fftw_execute_dft_r2c(s_plan, m_buffer_in, m_buffer_out);
//     memcpy(cArr, m_buffer_out, sizeof(fftw_complex) * (nx*nyk));
// }
void MyFftwClass::r2cexecute(double *const rArr, fftw_complex *const cArr)
{
  // NOTE(mbozzi): Use cArr anr rArr directly, instead of using
  // m_buffer_in/m_buffer_out followed by copying their contents.
  fftw_execute_dft_r2c(s_plan, rArr, cArr);
}


MyiFftwClass::MyiFftwClass(void)
{
  // m_buffer_in1 = fftw_alloc_complex((nx*nyk));
  // m_buffer_out1 = fftw_alloc_real((nx*ny));
  fftw_complex* buffer_in1 = fftw_alloc_complex(nxk * (ny + 1));
  double* buffer_out1 = fftw_alloc_real(nx*(ny + 1));

  if (!(buffer_in1 && buffer_out1))
  {
    // NOTE(mbozzi): prefer to "throw" instead of "throw new".  It's a
    // particularly bad idea in the case of out-of-memory error, because "new"
    // itself allocates memory.
    //
    // If you did "throw new", and care about error recovery, you ought to catch
    // the pointer and delete it:
    //   try { throw new std::runtime_error(""); } catch(std::runtime_error* pe) { delete pe; }
    // Else the new exception object will not be freed (a small memory leak).

    // NOTE(mbozzi): Consider calling std::abort instead of throwing exceptions.

    // throw new std::runtime_error("Failed to allocate memory!");
    throw std::runtime_error("Failed to allocate memory!");
  }

  if (!s_plan1)
  {
    // s_plan1 = fftw_plan_many_dft_c2r(1,&nx, (nyk),m_buffer_in1,nullptr,(nyk),1,m_buffer_out1,nullptr, (ny+1),1, FFTW_MEASURE);
    s_plan1 = fftw_plan_many_dft_c2r(
      1,           // int rank,
      &nx,         // const int *n,
      ny + 1,      // int howmany,
      buffer_in1,  // double *in,
      nullptr,     // const int *inembed,
      1,           // int istride, // adjacent elements within a row are adjacent.
      nxk,         // int idist,   // adjacent elements within a column are separate
      buffer_out1, // fftw_complex *out,
      nullptr,     // const int *onembed,
      1,           // int ostride,
      nx,          // int odist,
      FFTW_MEASURE // unsigned flags
      );
    if (!s_plan1)
    {
      // throw new std::runtime_error("Failed to create plan!");
      throw std::runtime_error("Failed to create plan!");
    }
  }

  fftw_free(buffer_out1);
  fftw_free(buffer_in1);
}

// MyiFftwClass::~MyiFftwClass(void)
// {
//   fftw_free(m_buffer_in1);
//   fftw_free(m_buffer_out1);
// }

//void MyiFftwClass::c2rexecute(double cArr[][ncomp],double rArr[]) void MyFftwClass::execute(double *const in, fftw_complex *const out)
// {
//   memcpy(m_buffer_in1,cArr, sizeof(fftw_complex) * (nx*nyk));
//   fftw_execute_dft_c2r(s_plan1, m_buffer_in1, m_buffer_out1); //instead of fftw_excute(plan)
//   memcpy(rArr, m_buffer_out1,sizeof(double) * (nx*ny));
//   //renormalize
//   for(int i = 0; i < nx; i++){
//     for( int j = 0; j < ny+1; j++){
//       //rArr[j + (ny+1)*i] = rArr[j + (ny+1)*i] * 1.0 / (nx);
//     }
//   }
// }
void MyiFftwClass::c2rexecute(double cArr[][ncomp],double rArr[])
{
  fftw_execute_dft_c2r(s_plan1, cArr, rArr); // instead of fftw_excute(plan)
  // renormalize
  for (int i = 0; i < ny + 1; ++i)
    for (int j = 0; j < nx; ++j)
      rArr[ij(i, j, nx)] /= nx;
}

fftw_plan MyiFftwClass::s_plan1 = NULL;
fftw_plan MyFftwClass::s_plan = NULL;