#include "fftwserver.h"
#include "FFTW/fftw.h"
#include "FFTW/rfftw.h"


#define MAXND_FFTW 5

class FFTWServerPlan {
public:
  FFTWServerPlan(int n, fftw_direction dir, bool fgreal=false);
  FFTWServerPlan(int nd, int * nxyz, fftw_direction dir, bool fgreal=false);
  ~FFTWServerPlan();
  void Recreate(int n);
  void Recreate(int nd, int * nxyz);

  int _n;  // Array dimension for 1-d arrays
  int _nd; // Nb of dimensions for n-d arrays
  int _nxyz[MAXND_FFTW]; // Array dimensions for n-d arrays
  fftw_direction _dir;

  fftw_plan p;
  rfftw_plan rp;
  fftwnd_plan pnd;
  rfftwnd_plan rpnd;
   
};

FFTWServerPlan::FFTWServerPlan(int n, fftw_direction dir, bool fgreal)
{
  if (n < 1) 
    throw ParmError("FFTWServerPlan: Array size <= 0 !");
  p = NULL;
  rp = NULL;
  pnd = NULL;
  rpnd = NULL;
  _nd = -1;
  for(int k=0; k<MAXND_FFTW; k++) _nxyz[k] = -10;
  _n = n;  
  _dir = dir;
  if (fgreal) rp = rfftw_create_plan(n, dir, FFTW_ESTIMATE);
  else p = fftw_create_plan(n, dir, FFTW_ESTIMATE);
}

FFTWServerPlan::FFTWServerPlan(int nd, int * nxyz, fftw_direction dir, bool fgreal)
{
  int k;
  if (nd > MAXND_FFTW)
    throw ParmError("FFTWServerPlan: Array rank (nd) > MAXND_FFTW !");
  p = NULL;
  rp = NULL;
  pnd = NULL;
  rpnd = NULL;
  _n = -10;
  _nd = nd;
  for(k=0; k<nd; k++) {
    if (nxyz[k] < 1) 
       throw ParmError("FFTWServerPlan: One of the Array size <= 0 !");   
    _nxyz[k] = nxyz[k];
  }
  for(k=nd; k<MAXND_FFTW; k++) _nxyz[k] = -10;
  _dir = dir;
  if (fgreal) rpnd = rfftwnd_create_plan(_nd, _nxyz, dir, FFTW_ESTIMATE);
  else pnd = fftwnd_create_plan(_nd, _nxyz, dir, FFTW_ESTIMATE);
}

FFTWServerPlan::~FFTWServerPlan()
{
  if (p) fftw_destroy_plan(p);
  if (rp) rfftw_destroy_plan(rp);
  if (pnd) fftwnd_destroy_plan(pnd);
  if (rpnd) rfftwnd_destroy_plan(rpnd);
}

void
FFTWServerPlan::Recreate(int n)
{
  if (n < 1) 
   throw ParmError("FFTWServerPlan::Recreate(n) n < 0 !");
  if (_nd > 0)  
   throw ParmError("FFTWServerPlan::Recreate(n) Multi-dimensional plan ! > 0 !");
  if (n == _n) return;
  _n = n;  
  if (p) {
    fftw_destroy_plan(p);
    p = fftw_create_plan(n, _dir, FFTW_ESTIMATE);
  }
  else {
    rfftw_destroy_plan(rp);
    rp = rfftw_create_plan(n, _dir, FFTW_ESTIMATE);
  }
}

void
FFTWServerPlan::Recreate(int nd, int * nxyz)
{
  if (_n > 0)
    throw ParmError("FFTWServerPlan::Recreate(nd, nxyz) 1-dimensional plan !");
  int k;
  if (nd == _nd) {
    bool samepl = true;
    for (int k=0; k<nd; k++) 
      if (nxyz[k] != _nxyz[k]) samepl = false;
    if (samepl) return;
  }
  if (nd > MAXND_FFTW)
    throw ParmError("FFTWServerPlan::Recreate(nd, nxyz) Array rank (nd) > MAXND_FFTW !");
  _nd = nd;
  for(k=0; k<nd; k++) {
    if (nxyz[k] < 1) 
       throw ParmError("FFTWServerPlan::Recreate(nd, nxyz) One of the Array size <= 0 !");   
    _nxyz[k] = nxyz[k];
  }
  for(k=nd; k<MAXND_FFTW; k++) _nxyz[k] = -10;
  if (pnd) {
    fftwnd_destroy_plan(pnd);
    pnd = fftwnd_create_plan(_nd, _nxyz, _dir, FFTW_ESTIMATE);
  }
  else {
    rfftwnd_destroy_plan(rpnd);
    rpnd = rfftwnd_create_plan(_nd, _nxyz, _dir, FFTW_ESTIMATE);
  }

}


/* --Methode-- */
FFTWServer::FFTWServer()
  : FFTServerInterface("FFTServer using FFTW package")
  , ckR8("FFTWServer - ", true, false)

{
  _p1df = NULL;
  _p1db = NULL;
  _pndf = NULL;
  _pndb = NULL;

  _p1drf = NULL;
  _p1drb = NULL;
  _pndrf = NULL;
  _pndrb = NULL;
}


/* --Methode-- */
FFTWServer::~FFTWServer()
{
  if (_p1df) delete _p1df ;
  if (_p1db) delete _p1db ;
  if (_pndf) delete _pndf ;
  if (_pndb) delete _pndb ;

  if (_p1drf) delete _p1drf ;
  if (_p1drb) delete _p1drb ;
  if (_pndrf) delete _pndrf ;
  if (_pndrb) delete _pndrb ;
}

/* --Methode-- */
FFTServerInterface * FFTWServer::Clone()
{
  return (new FFTWServer) ;
}

/* --Methode-- */
void 
FFTWServer::FFTForward(TArray< complex<r_8> > const & in, TArray< complex<r_8> > & out)
{
  int rank = ckR8.CheckResize(in, out);
  if (rank == 1) { // One dimensional transform 
    if (_p1df) _p1df->Recreate(in.Size());
    else _p1df = new FFTWServerPlan(in.Size(), FFTW_FORWARD, false);
    fftw_one(_p1df->p, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  }
  else {   // Multi dimensional 
    if (in.NbDimensions() > MAXND_FFTW) 
      throw ParmError("FFTWServer::FFTForward( complex<r_8>, complex<r_8> ) rank > MAXND_FFTW !"); 
    int sz[MAXND_FFTW];
    int k1 = 0;
    int k2 = 0;
    for(k1=in.NbDimensions()-1; k1>=0; k1--) {
      sz[k2] = in.Size(k1); k2++; 
    }
    if (_pndf) _pndf->Recreate(in.NbDimensions(), sz);
    else _pndf = new FFTWServerPlan(in.NbDimensions(), sz, FFTW_FORWARD, false);
    fftwnd_one(_pndf->pnd, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  }  
  if(this->getNormalize()) out=out/complex<r_8>((double)in.Size(),0.);
  return;
}

/* --Methode-- */
void FFTWServer::FFTBackward(TArray< complex<r_8> > const & in, TArray< complex<r_8> > & out)
{
  int rank = ckR8.CheckResize(in, out);
  if (rank == 1) { // One dimensional transform 
    if (_p1db) _p1db->Recreate(in.Size());
    else _p1db = new FFTWServerPlan(in.Size(), FFTW_BACKWARD, false);
    fftw_one(_p1db->p, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  }
  else {   // Multi dimensional 
    if (in.NbDimensions() > MAXND_FFTW) 
      throw ParmError("FFTWServer::FFTForward( complex<r_8>, complex<r_8> ) rank > MAXND_FFTW !"); 
    int sz[MAXND_FFTW];
    int k1 = 0;
    int k2 = 0;
    for(k1=in.NbDimensions()-1; k1>=0; k1--) {
      sz[k2] = in.Size(k1); k2++; 
    }
    if (_pndb) _pndb->Recreate(in.NbDimensions(), sz);
    else _pndb = new FFTWServerPlan(in.NbDimensions(), sz, FFTW_BACKWARD, false);
    fftwnd_one(_pndb->pnd, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  }

  return;
}


void FFTWServer::FFTForward(TArray< r_8 > const & in, TArray< complex<r_8> > & out)
{  
  int rank = ckR8.CheckResize(in, out);
  if (rank == 1) { // One dimensional transform 
    if (_p1drf) _p1drf->Recreate(in.Size());
    else _p1drf = new FFTWServerPlan(in.Size(), FFTW_REAL_TO_COMPLEX, true);
    TArray<r_8> outtemp; 
    outtemp.ReSize(in);
    rfftw_one(_p1drf->rp, (fftw_real *)(in.Data()) , (fftw_real *)(outtemp.Data()));
    ReShapetoCompl(outtemp, out);
  }
  else {   // Multi dimensional 
    if (in.NbDimensions() > MAXND_FFTW) 
      throw ParmError("FFTWServer::FFTForward( complex<r_8>, complex<r_8> ) rank > MAXND_FFTW !"); 
    int sz[MAXND_FFTW];
    int k1 = 0;
    int k2 = 0;
    for(k1=in.NbDimensions()-1; k1>=0; k1--) {
      sz[k2] = in.Size(k1); k2++; 
    }
    if (_pndrf) _pndrf->Recreate(in.NbDimensions(), sz);
    else _pndrf = new FFTWServerPlan(in.NbDimensions(), sz, FFTW_REAL_TO_COMPLEX, true);
    rfftwnd_one_real_to_complex(_pndrf->rpnd, (fftw_real *)(in.Data()) , 
				(fftw_complex *)(out.Data()) );
  }
  if(this->getNormalize()) out=out/complex<r_8>((double)in.Size(),0.);
  return;

}



void FFTWServer::FFTBackward(TArray< complex<r_8> > const & in, TArray< r_8 > & out)
{
  
  int rank = ckR8.CheckResize(in, out);
  if (rank == 1) { // One dimensional transform 
    TArray<r_8> intemp;
    intemp.ReSize(out);
    if (_p1drb) _p1drb->Recreate(out.Size());
    else _p1drb = new FFTWServerPlan(out.Size(), FFTW_COMPLEX_TO_REAL, true);

    ReShapetoReal(in, intemp);
    //    cerr << " DEBUG-FFTWServer::FFTBackward() in = \n" << in << endl;
    //    cerr << " DEBUG-FFTWServer::FFTBackward() intemp = \n" << intemp << endl;
    rfftw_one(_p1drb->rp, (fftw_real *)(intemp.Data()) , (fftw_real *)(out.Data()));

  }
  else {   // Multi dimensional 
    if (in.NbDimensions() > MAXND_FFTW) 
      throw ParmError("FFTWServer::FFTForward( complex<r_8>, complex<r_8> ) rank > MAXND_FFTW !"); 
    int sz[MAXND_FFTW];
    int k1 = 0;
    int k2 = 0;
    for(k1=in.NbDimensions()-1; k1>=0; k1--) {
      sz[k2] = in.Size(k1); k2++; 
    }
    if (_pndrb) _pndrb->Recreate(in.NbDimensions(), sz);
    else _pndrb = new FFTWServerPlan(in.NbDimensions(), sz, FFTW_COMPLEX_TO_REAL, true);

    rfftwnd_one_complex_to_real(_pndrb->rpnd, (fftw_complex *)(in.Data()) , (fftw_real *)(out.Data()) );
  }
  return;
}

/*
void FFTWServer::FFTForward(TArray< complex<r_8> > const & in, TArray< complex<r_8> > & out)
{
  out.ReSize(in.NRows(),in.NCols());

  if (_pndf) _pndf->Recreate( in.NRows(),in.NCols());
  else _pndf = new FFTWServerPlan( in.NCols(),in.NRows(), FFTW_FORWARD, false);
  
  fftwnd_one(_pndf->pnd, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  if(this->getNormalize()) out=out/complex<r_8>(pow(in.NRows()*in.NCols(),0.5),0.);  
}

void FFTWServer::FFTBackward(TArray< complex<r_8> > const & in, TArray< complex<r_8> > & out)
{
  if (_pndb) _pndb->Recreate(in.NCols(), in.NRows());
  else _pndb = new FFTWServerPlan(in.NCols(), in.NRows(), FFTW_BACKWARD, false);
  out.ReSize(in.NRows(), in.NCols());
  fftwnd_one(_pndb->pnd, (fftw_complex *)(in.Data()) , (fftw_complex *)(out.Data()) );
  if(this->getNormalize()) out=out/complex<r_8>(pow(in.NRows()*in.NCols(),0.5),0.);

}


void FFTWServer::FFTForward(TArray< r_8 > const & in, TArray< complex<r_8> > & out)
{

  TArray< r_8 > inNew(in.NCols(),in.NRows());
  for(int i=0; i<in.NRows(); i++)
    for(int j=0;j<in.NCols(); j++)
      inNew(j,i) = in(i,j);
  
  if (_pndrf) _pndrf->Recreate(inNew.NRows(),inNew.NCols());
  else _pndrf = new FFTWServerPlan(inNew.NRows(), inNew.NCols(),FFTW_REAL_TO_COMPLEX, true);
  //  rfftwnd_plan p;
  TArray< complex<r_8> > outTemp;
  outTemp.ReSize(in.NRows(),in.NCols());

  rfftwnd_one_real_to_complex(_pndrf->rpnd, (fftw_real *)(in.Data()) , (fftw_complex *)(out.Data()) );
}

void FFTWServer::FFTBackward(TArray< complex<r_8> > const & in, TArray< r_8 > & out)
{

  TArray< complex<r_8> > inNew(in.NCols(),in.NRows());
  for(int i=0; i<in.NRows(); i++)
    for(int j=0;j<in.NCols(); j++)
      inNew(j,i) = in(i,j);
  
  if (_pndrb) _pndrb->Recreate(inNew.NRows(),inNew.NCols());
  else _pndrb = new FFTWServerPlan(inNew.NRows(), inNew.NCols(),FFTW_COMPLEX_TO_REAL, true);
  //  rfftwnd_plan p;
  out.ReSize(in.NRows(),in.NCols());

  rfftwnd_one_complex_to_real(_pndrb->rpnd, (fftw_complex *)(in.Data()) , (fftw_real *)(out.Data()) );
  cout << " in the function !!!" << endl;
  if(this->getNormalize()) 
    {
      r_8 norm = (r_8)(in.NRows()*in.NCols());
      out=out/norm;
    }
}

*/


/* --Methode-- */
void FFTWServer::ReShapetoReal(TArray< complex<r_8> > const & ina, TArray< r_8 > & outa)
{
  TVector< complex<r_8> > in(ina);
  TVector< r_8> out(outa);
  int n = in.NElts();
  r_8 thr = FFTArrayChecker<r_8>::ZeroThreshold();
  sa_size_t ncs = ( (in(n-1).imag() < -thr) || (in(n-1).imag() > thr) ) ?
                    ncs = 2*n-1 : ncs = 2*n-2;

  if (out.NElts() != ncs) {
    cerr << " DEBUG-FFTWServer::ReShapetoReal() ncs= " << ncs 
	 << " out.NElts()= " << out.NElts() << endl;
    throw SzMismatchError("FFTWServer::ReShapetoReal() - Wrong output array size !");
  }
  //  if (ncs == 2*n-2)  {
  out(0) = in(0).real();
  for(int k=1; k<n; k++) {
      out(k) = in(k).real();
      out(ncs-k) = in(k).imag();
  }
  if (ncs == 2*n-2)
    out(ncs-1) = in(n-1).real();
  
  //  for(int k=0; k<out.NElts(); k++) cout << "ReShapetoReal out " << k << " " << out(k) << endl;
}


/* --Methode-- */
void FFTWServer::ReShapetoCompl(TArray< r_8 > const & ina, TArray< complex<r_8> > & outa)
{
  TVector< r_8> in(ina);
  TVector< complex<r_8> > out(outa);

  sa_size_t n = in.NElts();
  sa_size_t ncs = n/2+1;
  sa_size_t nc = (n%2 != 0) ? n/2+1 : n/2;
  if (out.NElts() != ncs)
    throw SzMismatchError("FFTWServer::ReShapetoCompl() - Wrong output array size !");

  out(0) = complex<r_8> (in(0),0.);
  for(int k=1; k<ncs; k++) 
    out(k) = complex<r_8>(in(k), in(n-k));
  if (n%2 == 0) out(ncs-1) = complex<r_8>(in(n/2), 0.);

}