#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")
  ,  ckR4(true, false) , ckR8(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()) );
  }  
  // $CHECK$ Reza 9/2/2001 , Verifier normalisation
  if(this->getNormalize()) out=out/complex<r_8>(sqrt((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);
  }
  // $CHECK$ Reza 9/2/2001 , Verifier normalisation  
  if(this->getNormalize()) out=out/complex<r_8>(sqrt((double)in.Size()),0.);
  return;
}


void FFTWServer::FFTForward(TArray< r_8 > const & in, TArray< complex<r_8> > & out)
{  
  int rank = ckR8.CheckResize(in, out);
  TArray<r_8> outtemp(in, false); 

  if (rank == 1) { // One dimensional transform 
    if (_p1drf) _p1drf->Recreate(in.Size());
    else _p1drf = new FFTWServerPlan(in.Size(), FFTW_REAL_TO_COMPLEX, true);
    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()) );
  }
  // $CHECK$ Reza 9/2/2001 , Verifier normalisation  
  if(this->getNormalize()) out=out/complex<r_8>(sqrt((double)in.Size()),0.);
  return;

}



void FFTWServer::FFTBackward(TArray< complex<r_8> > const & in, TArray< r_8 > & out)
{
  throw ParmError("FFTWServer::FFTBackward(TArray< complex<r_8> > ... Not implemented ... !");
  /*
  int size;
  if(in(in.NElts()).imag()  == 0) { size = 2*in.NElts()-2;}
  else { size = 2*in.NElts()-1;}
  
  TArray< r_8 > inTemp(size);
  out.ReSize(size);

  if (_p1drb) _p1drb->Recreate(size);
  else _p1drb = new FFTWServerPlan(size, FFTW_COMPLEX_TO_REAL, true);

  ReShapetoReal(in, inTemp);
  rfftw_one(_p1drb->rp, (fftw_real *)(inTemp.Data()) , (fftw_real *)(out.Data()));
  if(this->getNormalize()) out=out/pow(size,0.5);
  */
}

/*
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 & in, TArray< r_8 >  & out)
{
  int N = in.Size();
  /*
  int i;
  if (in(in.Size()).imag() == 0) 
    {
      out(0) = in(0).real();
      for(i=1; i<in.NElts(); i++)
	{
	  out(i) = in(i).real();
	}
      for(i=1; i<in.NElts(); i++)
	{
	  out(i+in.NElts()-1) = in(in.NElts()-i-1).imag();
	}
    }
  else
    {
      out(0) = in(0).real();
      for(i=1; i<in.NElts(); i++)
	{
	  out(i) = in(i).real();
	}
      for(i=1; i<in.NElts(); i++)
	{
	  out(i+in.NElts()-1) = in(in.NElts()-i).imag();
	}
    }
  */
  out[0] = in[0].real();
  int k=0;
  for(int i=1; i<N; i++) {
    out[i] = in[i].real();
    out[N-i] = in[i].imag();
  }
  
  //  for(int k=0; k<out.NElts(); k++) cout << "ReShapetoReal out " << k << " " << out(k) << endl;
}


/* --Methode-- */
void FFTWServer::ReShapetoCompl(TArray< r_8 > const & in, TArray< complex<r_8> > & out)
{
  int N = in.Size();
  //  for(int k=0; k<in.NElts(); k++) cout << "ReShapetoCompl in " << k << " "  << in(k) << endl;
  out[0] = complex<r_8> (in[0],0.);
  for(int k=1; k<N+1/2; k++) 
    out[k] = complex<r_8>(in[k], in[N-k]);
  if (N%2 == 0) out[N/2] = complex<r_8>(in[N/2], 0.);
  //  for(int k=0; k<out.NElts(); k++) cout << "ReShapetoCompl out " << k << " "  << out(k) << endl;
}