// @(#)root/fft:$Id$
// Author: Anna Kreshuk   07/4/2006

/*************************************************************************
 * Copyright (C) 1995-2006, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

////////////////////////////////////////////////////////////////////////////////
/// \class TFFTReal
/// One of the interface classes to the FFTW package, can be used directly
/// or via the TVirtualFFT class. Only the basic interface of FFTW is implemented.
///
/// Computes transforms called r2r in FFTW manual:
/// - transforms of real input and output in "halfcomplex" format i.e.
///   real and imaginary parts for a transform of size n stored as
///   (r0, r1, r2, ..., rn/2, i(n+1)/2-1, ..., i2, i1)
/// - discrete Hartley transform
/// - sine and cosine transforms (DCT-I,II,III,IV and DST-I,II,III,IV)
/// For the detailed information on the computed
/// transforms please refer to the FFTW manual, chapter "What FFTW really computes".
///
/// How to use it:
/// 1. Create an instance of TFFTReal - this will allocate input and output
///    arrays (unless an in-place transform is specified)
/// 2. Run the Init() function with the desired flags and settings (see function
///    comments for possible kind parameters)
/// 3. Set the data (via SetPoints()or SetPoint() functions)
/// 4. Run the Transform() function
/// 5. Get the output (via GetPoints() or GetPoint() functions)
/// 6. Repeat steps 3)-5) as needed
///
/// For a transform of the same size, but of different kind (or with different flags),
/// rerun the Init() function and continue with steps 3)-5)
///
/// NOTE:
///       1. running Init() function will overwrite the input array! Don't set any data
///          before running the Init() function!
///       2. FFTW computes unnormalized transform, so doing a transform followed by
///          its inverse will lead to the original array scaled BY:
///          - transform size (N) for R2HC, HC2R, DHT transforms
///          - 2*(N-1) for DCT-I (REDFT00)
///          - 2*(N+1) for DST-I (RODFT00)
///          - 2*N for the remaining transforms
///
/// Transform inverses:
/// - R2HC<-->HC2R
/// - DHT<-->DHT
/// - DCT-I<-->DCT-I
/// - DCT-II<-->DCT-III
/// - DCT-IV<-->DCT-IV
/// - DST-I<-->DST-I
/// - DST-II<-->DST-III
/// - DST-IV<-->DST-IV
///
////////////////////////////////////////////////////////////////////////////////

#include "TFFTReal.h"
#include "fftw3.h"

ClassImp(TFFTReal);

////////////////////////////////////////////////////////////////////////////////
///default

TFFTReal::TFFTReal()
{
   fIn    = 0;
   fOut   = 0;
   fPlan  = 0;
   fN     = 0;
   fKind  = 0;
   fNdim = 0;
   fTotalSize = 0;
}

////////////////////////////////////////////////////////////////////////////////
///For 1d transforms
///n here is the physical size of the transform (see FFTW manual for more details)

TFFTReal::TFFTReal(Int_t n, Bool_t inPlace)
{
   fIn = fftw_malloc(sizeof(Double_t)*n);
   if (inPlace) fOut = 0;
   else fOut = fftw_malloc(sizeof(Double_t)*n);

   fPlan = 0;
   fNdim = 1;
   fN = new Int_t[1];
   fN[0] = n;
   fKind = 0;
   fTotalSize = n;
}

////////////////////////////////////////////////////////////////////////////////
///For multidimensional transforms
///1st parameter is the # of dimensions,
///2nd is the sizes (physical) of the transform in each dimension

TFFTReal::TFFTReal(Int_t ndim, Int_t *n, Bool_t inPlace)
{
   fTotalSize = 1;
   fNdim = ndim;
   fN = new Int_t[ndim];
   fKind = 0;
   fPlan = 0;
   for (Int_t i=0; i<ndim; i++){
      fTotalSize*=n[i];
      fN[i] = n[i];
   }
   fIn = fftw_malloc(sizeof(Double_t)*fTotalSize);
   if (!inPlace)
      fOut = fftw_malloc(sizeof(Double_t)*fTotalSize);
   else
      fOut = 0;
}

////////////////////////////////////////////////////////////////////////////////
///clean-up

TFFTReal::~TFFTReal()
{
   fftw_destroy_plan((fftw_plan)fPlan);
   fPlan = 0;
   fftw_free(fIn);
   fIn = 0;
   if (fOut){
      fftw_free(fOut);
   }
   fOut = 0;
   if (fN)
      delete [] fN;
   fN = 0;
   if (fKind)
      fftw_free((fftw_r2r_kind*)fKind);
   fKind = 0;
}

////////////////////////////////////////////////////////////////////////////////
///Creates the fftw-plan
///
///NOTE:  input and output arrays are overwritten during initialisation,
///       so don't set any points, before running this function!!!!!
///
/// #### 1st parameter:
///    Possible flag_options:
///
/// - "ES" (from "estimate") - no time in preparing the transform, but probably sub-optimal
///       performance
/// - "M" (from "measure") - some time spend in finding the optimal way to do the transform
/// - "P" (from "patient") - more time spend in finding the optimal way to do the transform
/// - "EX" (from "exhaustive") - the most optimal way is found
///
///  This option should be chosen depending on how many transforms of the same size and
///  type are going to be done. Planning is only done once, for the first transform of this
///  size and type.
///
/// #### 2nd parameter:
///    is dummy and doesn't need to be specified
///
/// #### 3rd parameter:
///    transform kind for each dimension
///     4 different kinds of sine and cosine transforms are available
///  -   DCT-I   - kind=0
///  -   DCT-II  - kind=1
///  -   DCT-III - kind=2
///  -   DCT-IV  - kind=3
///  -   DST-I   - kind=4
///  -   DST-II  - kind=5
///  -   DSTIII  - kind=6
///  -   DSTIV   - kind=7

void TFFTReal::Init( Option_t* flags,Int_t /*sign*/, const Int_t *kind)
{
   if (fPlan)
      fftw_destroy_plan((fftw_plan)fPlan);
   fPlan = 0;

   if (!fKind)
      fKind = (fftw_r2r_kind*)fftw_malloc(sizeof(fftw_r2r_kind)*fNdim);

   if (MapOptions(kind)){
      if (fOut)
         fPlan = (void*)fftw_plan_r2r(fNdim, fN, (Double_t*)fIn, (Double_t*)fOut, (fftw_r2r_kind*)fKind, MapFlag(flags));
      else
         fPlan = (void*)fftw_plan_r2r(fNdim, fN, (Double_t*)fIn, (Double_t*)fIn, (fftw_r2r_kind*)fKind, MapFlag(flags));
      fFlags = flags;
   }
}

////////////////////////////////////////////////////////////////////////////////
///Computes the transform, specified in Init() function

void TFFTReal::Transform()
{
   if (fPlan)
      fftw_execute((fftw_plan)fPlan);
   else {
      Error("Transform", "transform hasn't been initialised");
      return;
   }
}

////////////////////////////////////////////////////////////////////////////////
///Returns the type of the transform

Option_t *TFFTReal::GetType() const
{
   if (!fKind) {
      Error("GetType", "Type not defined yet (kind not set)");
      return "";
   }
   if (((fftw_r2r_kind*)fKind)[0]==FFTW_R2HC) return "R2HC";
   if (((fftw_r2r_kind*)fKind)[0]==FFTW_HC2R) return "HC2R";
   if (((fftw_r2r_kind*)fKind)[0]==FFTW_DHT) return "DHT";
   else return "R2R";
}

////////////////////////////////////////////////////////////////////////////////
///Copies the output (or input) points into the provided array, that should
///be big enough

void TFFTReal::GetPoints(Double_t *data, Bool_t fromInput) const
{
   const Double_t * array = GetPointsReal(fromInput);
   if (!array) return;
   std::copy(array, array+fTotalSize, data);
}

////////////////////////////////////////////////////////////////////////////////
///For 1d tranforms. Returns point #ipoint

Double_t TFFTReal::GetPointReal(Int_t ipoint, Bool_t fromInput) const
{
   if (ipoint<0 || ipoint>fTotalSize){
      Error("GetPointReal", "No such point");
      return 0;
   }
   const Double_t * array = GetPointsReal(fromInput);
   return ( array ) ? array[ipoint] : 0;
}

////////////////////////////////////////////////////////////////////////////////
///For multidim.transforms. Returns point #ipoint

Double_t TFFTReal::GetPointReal(const Int_t *ipoint, Bool_t fromInput) const
{
   Int_t ireal = ipoint[0];
   for (Int_t i=0; i<fNdim-1; i++)
      ireal=fN[i+1]*ireal + ipoint[i+1];

   const Double_t * array = GetPointsReal(fromInput);
   return ( array ) ? array[ireal] : 0;
}

////////////////////////////////////////////////////////////////////////////////
///Only for input of HC2R and output of R2HC

void TFFTReal::GetPointComplex(Int_t ipoint, Double_t &re, Double_t &im, Bool_t fromInput) const
{
   const Double_t * array = GetPointsReal(fromInput);
   if (!array) return;
   if ( ( ((fftw_r2r_kind*)fKind)[0]==FFTW_R2HC && !fromInput ) ||
        ( ((fftw_r2r_kind*)fKind)[0]==FFTW_HC2R &&  fromInput ) )
   {
      if (ipoint<fN[0]/2+1){
         re = array[ipoint];
         im = array[fN[0]-ipoint];
      } else {
         re = array[fN[0]-ipoint];
         im = -array[ipoint];
      }
      if ((fN[0]%2)==0 && ipoint==fN[0]/2) im = 0;
   }
}
////////////////////////////////////////////////////////////////////////////////
///Only for input of HC2R and output of R2HC and for 1d

void TFFTReal::GetPointComplex(const Int_t *ipoint, Double_t &re, Double_t &im, Bool_t fromInput) const
{
   GetPointComplex(ipoint[0], re, im, fromInput);
}

////////////////////////////////////////////////////////////////////////////////
///Returns the output (or input) array
/// we have 4 different cases:
///
/// - fromInput = false; fOut = !NULL (transformed is not in place) : return fOut
///  - fromInput = false; fOut = NULL (transformed is in place) : return fIn
/// - fromInput = true; fOut = !NULL :   return fIn
/// - fromInput = true; fOut = NULL return an error since input array is overwritten

Double_t* TFFTReal::GetPointsReal(Bool_t fromInput) const
{

   if (!fromInput && fOut)
      return (Double_t*)fOut;
   else if (fromInput && !fOut) {
      Error("GetPointsReal","Input array was destroyed");
      return 0;
   }
   //R__ASSERT(fIn);
   return (Double_t*)fIn;
}

////////////////////////////////////////////////////////////////////////////////

void TFFTReal::SetPoint(Int_t ipoint, Double_t re, Double_t im)
{
   if (ipoint<0 || ipoint>fTotalSize){
      Error("SetPoint", "illegal point index");
      return;
   }
   if (((fftw_r2r_kind*)fKind)[0]==FFTW_HC2R){
      if ((fN[0]%2)==0 && ipoint==fN[0]/2)
         ((Double_t*)fIn)[ipoint] = re;
      else {
         ((Double_t*)fIn)[ipoint] = re;
         ((Double_t*)fIn)[fN[0]-ipoint]=im;
      }
   }
   else
      ((Double_t*)fIn)[ipoint]=re;
}

////////////////////////////////////////////////////////////////////////////////
///Since multidimensional R2HC and HC2R transforms are not supported,
///third parameter is dummy

void TFFTReal::SetPoint(const Int_t *ipoint, Double_t re, Double_t /*im*/)
{
   Int_t ireal = ipoint[0];
   for (Int_t i=0; i<fNdim-1; i++)
      ireal=fN[i+1]*ireal + ipoint[i+1];
   if (ireal < 0 || ireal >fTotalSize){
      Error("SetPoint", "illegal point index");
      return;
   }
   ((Double_t*)fIn)[ireal]=re;
}

////////////////////////////////////////////////////////////////////////////////
///Sets all points

void TFFTReal::SetPoints(const Double_t *data)
{
   for (Int_t i=0; i<fTotalSize; i++)
      ((Double_t*)fIn)[i] = data[i];
}

////////////////////////////////////////////////////////////////////////////////
///transfers the r2r_kind parameters to fftw type

Int_t TFFTReal::MapOptions(const Int_t *kind)
{
   if (kind[0] == 10){
      if (fNdim>1){
         Error("Init", "Multidimensional R2HC transforms are not supported, use R2C interface instead");
         return 0;
      }
      ((fftw_r2r_kind*)fKind)[0] = FFTW_R2HC;
   }
   else if (kind[0] == 11) {
      if (fNdim>1){
         Error("Init", "Multidimensional HC2R transforms are not supported, use C2R interface instead");
         return 0;
      }
      ((fftw_r2r_kind*)fKind)[0] = FFTW_HC2R;
   }
   else if (kind[0] == 12) {
      for (Int_t i=0; i<fNdim; i++)
         ((fftw_r2r_kind*)fKind)[i] = FFTW_DHT;
   }
   else {
      for (Int_t i=0; i<fNdim; i++){
         switch (kind[i]) {
         case 0: ((fftw_r2r_kind*)fKind)[i] = FFTW_REDFT00;  break;
         case 1: ((fftw_r2r_kind*)fKind)[i] = FFTW_REDFT01;  break;
         case 2: ((fftw_r2r_kind*)fKind)[i] = FFTW_REDFT10;  break;
         case 3: ((fftw_r2r_kind*)fKind)[i] = FFTW_REDFT11;  break;
         case 4: ((fftw_r2r_kind*)fKind)[i] = FFTW_RODFT00;  break;
         case 5: ((fftw_r2r_kind*)fKind)[i] = FFTW_RODFT01;  break;
         case 6: ((fftw_r2r_kind*)fKind)[i] = FFTW_RODFT10;  break;
         case 7: ((fftw_r2r_kind*)fKind)[i] = FFTW_RODFT11;  break;
         default:
         ((fftw_r2r_kind*)fKind)[i] = FFTW_R2HC; break;
         }
      }
   }
   return 1;
}

////////////////////////////////////////////////////////////////////////////////
///allowed options:
/// - "ES" - FFTW_ESTIMATE
/// - "M" - FFTW_MEASURE
/// - "P" - FFTW_PATIENT
/// - "EX" - FFTW_EXHAUSTIVE

UInt_t TFFTReal::MapFlag(Option_t *flag)
{
   TString opt = flag;
   opt.ToUpper();
   if (opt.Contains("ES"))
      return FFTW_ESTIMATE;
   if (opt.Contains("M"))
      return FFTW_MEASURE;
   if (opt.Contains("P"))
      return FFTW_PATIENT;
   if (opt.Contains("EX"))
      return FFTW_EXHAUSTIVE;
   return FFTW_ESTIMATE;
}