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

/*!
   \defgroup IFFTW IFFTW module
   Module containing interface classes between Sophya objects
   and FFTW Fourier transform package (see http://www.fftw.org )
*/

/*! 
  \class SOPHYA::FFTWServer
  \ingroup IFFTW
  An implementation of FFTServerInterface based on FFTW, double
  precision arrays, using FFTW package, availabale from http://www.fftw.org.
  Refer to FFTServerInterface for details about FFTServer operations.

  \code
  #include "fftwserver.h"
  // ...
  TMatrix<r_8> in(24,32);
  TMatrix< complex<r_8> > out;
  in = RandomSequence();
  FFTWServer ffts;
  ffts.setNormalize(true);  // To have normalized transforms
  cout << " FFTServer info string= " << ffts.getInfo() << endl;
  cout << "in= " << in << endl;
  cout << " Calling ffts.FFTForward(in, out) : " << endl;
  ffts.FFTForward(in, out);
  cout << "out= " << out << endl;
  \endcode

*/

#define MAXND_FFTW 5

namespace SOPHYA {

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);

  static void FillSizes(const BaseArray & in, int * sz);

  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;
   
};

} // Fin du namespace

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);
    if (rp == NULL) 
      throw AllocationError("FFTWServerPlan: failed to create plan (rp) !");
  }
  else {
    p = fftw_create_plan(n, dir, FFTW_ESTIMATE);
    if (p == NULL) 
      throw AllocationError("FFTWServerPlan: failed to create plan (p) !");
  }
    
}

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);
    if (rpnd == NULL) 
      throw AllocationError("FFTWServerPlan: failed to create plan (rpnd) !");
  }
  else { 
    pnd = fftwnd_create_plan(_nd, _nxyz, dir, FFTW_ESTIMATE);
    if (pnd == NULL) 
      throw AllocationError("FFTWServerPlan: failed to create plan (pnd) !");
  }

}

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);
    if (p == NULL) 
      throw AllocationError("FFTWServerPlan::Recreate failed to create plan (p) !");
  }
  else {
    rfftw_destroy_plan(rp);
    rp = rfftw_create_plan(n, _dir, FFTW_ESTIMATE);
    if (rp == NULL) 
      throw AllocationError("FFTWServerPlan::Recreate failed to create plan (rp) !");
  }

}

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);
    if (pnd == NULL) 
      throw AllocationError("FFTWServerPlan::Recreate failed to create plan (pnd) !");
  }
  else {
    rfftwnd_destroy_plan(rpnd);
    rpnd = rfftwnd_create_plan(_nd, _nxyz, _dir, FFTW_ESTIMATE);
    if (rpnd == NULL) 
      throw AllocationError("FFTWServerPlan::Recreate failed to create plan (rpnd) !");
  }

}

void
FFTWServerPlan::FillSizes(const BaseArray & in, int * sz)
{
  int k1 = 0;
  int k2 = 0;
  for(k1=in.NbDimensions()-1; k1>=0; k1--) {
    sz[k2] = in.Size(k1); k2++; 
  }
}

/* --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];
    FFTWServerPlan::FillSizes(in, sz);
    //    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];
    FFTWServerPlan::FillSizes(in, sz);
    //    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];
    FFTWServerPlan::FillSizes(in, sz);
    //    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,
			     bool usoutsz)
{
  
  int rank = ckR8.CheckResize(in, out, usoutsz);
  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, usoutsz);
    //    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>, r_8 ) rank > MAXND_FFTW !"); 
    int sz[MAXND_FFTW];
    FFTWServerPlan::FillSizes(out, sz);
    //    int k1 = 0;
    //    int k2 = 0;
    //    for(k1=out.NbDimensions()-1; k1>=0; k1--) {
    //      sz[k2] = out.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;
}



/* --Methode-- */
void FFTWServer::ReShapetoReal(TArray< complex<r_8> > const & ina, TArray< r_8 > & outa, 
			       bool usz)
{
  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;
  if (usz) {
    if ( (out.NElts() != 2*n-1) && (out.NElts() != 2*n-2) )
      throw SzMismatchError("FFTWServer::ReShapetoReal(..., true) - Wrong output array size ");
    ncs = out.NElts();
  }
  else {
    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(n-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.);

}