/*
  in-place FFT and inverse FFT, with bit-reversed result locations

  R. Perry, June 2017

  Algorithm 4.2 from "Inside the FFT Black Box: Serial and Parallel
  Fast Fourier Transform Algorithms", Eleanor Chu and Alan George,
  CRC Press, 1999

  requires len(X) = power of 2
*/

#include "Python.h"

#include <stdlib.h> // for realloc() and qsort()
#include <complex.h>

// CMPLX macro may not be available in glib
// https://github.com/JuliaLang/openlibm/issues/15
// https://gcc.gnu.org/onlinedocs/gcc-4.7.2/gcc/Other-Builtins.html
//
#ifndef CMPLX
#define CMPLX(x,y) __builtin_complex((x),(y))
#endif

// for index and value list when finding largest absolute result values
//
typedef struct
{
  size_t i; double v;
}
ivlist;

// comparsion function for qsort()
//
static int compare( const void *v1, const void *v2)
{
  const ivlist *iv1 = v1, *iv2 = v2;

  if( iv1->v < iv2->v) return -1;
  else if( iv1->v > iv2->v) return +1;
  else return 0;
}

//------------------------------------------------------------------------------

static PyObject *fft( PyObject *self, PyObject *args, PyObject *keywords)
{
  PyObject *X, *P = NULL;  int inv = 0, pre = 1;  size_t nmax = 0;

  static char *keywordlist[] = { "X", "inv", "pre", "nmax", NULL };

  if( !PyArg_ParseTupleAndKeywords( args, keywords, "O!|ppO!", keywordlist,
      &PyList_Type, &X, &inv, &pre, &PyLong_Type, &P))
    return NULL;

  if( P) nmax = PyLong_AsSize_t( P); // might overflow

  if( PyErr_Occurred()) return NULL;

#if 0 // debug
  printf( "fft: inv = %i, pre = %i, nmax = %llu\n", inv, pre, (unsigned long long) nmax);
#endif

 // check the first X value
 //
 // should check all of the X values, but skipping that for speed
 // 
  if( !PyComplex_Check( PyList_GET_ITEM( X, 0)))
  {
    PyErr_SetString( PyExc_TypeError, "fft argument X must be list of complex");
    return NULL;
  }

 // check that size is a power of 2
 //
  size_t N = PyList_Size( X), t = N >> 1, n = 0;

  while( t > 0) { t >>= 1; ++n; }

  if( N != (1U << n))
  {
    PyErr_SetString( PyExc_RuntimeError, "len(X) must be a power of 2");
    return NULL;
  }

 // precompute W to improve speed -- but requires extra memory
 //
  double pi = 4*atan(1.0);  double complex c = -2*I*pi/N;  if( inv) c = conj(c);

  static size_t Wpre_N = 0; static int Wpre_inv = 0;

  static double complex *Wpre = NULL;

  if( pre)
  {
    if( !Wpre || Wpre_N != N) // create or resize
    {
      double complex *Wpre_new = realloc( Wpre, N*sizeof(double complex));

      if( !Wpre_new)
      {
        printf( "fft: WARNING: realloc %zu bytes failed.\n", N*sizeof(double complex));
        pre = 0;
      }
      else
      {
        #if 0 // debug
          printf( "fft: computing Wpre\n");
        #endif

        Wpre = Wpre_new;  Wpre_N = N;  Wpre_inv = inv;

        for( size_t wi = 0; wi < N; ++wi) Wpre[wi] = cexp(c*wi);
      }
    }
    else if( Wpre_inv != inv) // replace with conjugate values
    {
      #if 0 // debug
        printf( "fft: conjugate Wpre\n");
      #endif

      for( size_t wi = 0; wi < N; ++wi) Wpre[wi] = conj(Wpre[wi]);

      Wpre_inv = inv;
    }
  }

 // perform FFT
 //
  size_t ProblemSize = N, NumOfProblems = 1;

  while( ProblemSize > 1)
  {
    size_t HalfSize = ProblemSize/2;

    for( size_t K = 0; K < NumOfProblems; ++K)
    {
      size_t JFirst = K*ProblemSize, JLast = JFirst + HalfSize, wi = 0;

      for( size_t J = JFirst; J < JLast; ++J)
      {
        PyComplexObject *XJ  = (PyComplexObject *) PyList_GET_ITEM( X, J);
        PyComplexObject *XJH = (PyComplexObject *) PyList_GET_ITEM( X, J+HalfSize);
        double complex W = pre ? Wpre[wi] : cexp(c*wi);
        double complex t = CMPLX( XJ->cval.real, XJ->cval.imag),
	 xjh = CMPLX( XJH->cval.real, XJH->cval.imag), xj = t + xjh;
        XJ->cval.real = creal(xj); XJ->cval.imag = cimag(xj);
        xjh = W*(t - xjh);
        XJH->cval.real = creal(xjh); XJH->cval.imag = cimag(xjh);
        wi += NumOfProblems;
      }
    }
    NumOfProblems *= 2;
    ProblemSize = HalfSize;
  }

  if( nmax == 0) return Py_None;

 // create list of indices of the nmax largest absolute values
 //
  if( nmax > N) nmax = N;

  ivlist iv[nmax];

  PyObject *plist = PyList_New(nmax);  if( !plist) return NULL;

  for( size_t k = 0; k < nmax; ++k)
  {
    iv[k].i = k;
    PyComplexObject *Xk  = (PyComplexObject *) PyList_GET_ITEM( X, k);
    iv[k].v = Xk->cval.real*Xk->cval.real + Xk->cval.imag*Xk->cval.imag;
  }

 // put list in increasing order, and then keep it that way
 //
  qsort( iv, nmax, sizeof(ivlist), compare);

 // check the rest of the result values
 //
  for( size_t k = nmax; k < N; ++k)
  {
    PyComplexObject *Xk  = (PyComplexObject *) PyList_GET_ITEM( X, k);
    double v = Xk->cval.real*Xk->cval.real + Xk->cval.imag*Xk->cval.imag;

    if( v > iv[0].v) // if greater than current min value, replace it
    {
      iv[0].i = k; iv[0].v = v;
      for( size_t j = 0; j < nmax-1; ++j) // resort
        if( iv[j].v <= iv[j+1].v) break;
        else { ivlist t = iv[j]; iv[j] = iv[j+1]; iv[j+1] = t; }
    }
  }

  for( size_t k = 0; k < nmax; ++k)
  {
    PyList_SET_ITEM( plist, k, PyLong_FromSize_t( iv[nmax-k-1].i));
  }

  return plist;
}

// Notes:
//
// https://docs.python.org/3/c-api/complex.html
//
//  typedef struct {
//    double real;
//    double imag;
//  } Py_complex;
//
// https://docs.python.org/3/c-api/memory.html#overview
// 
//  void* PyMem_RawMalloc(size_t n)
//  void* PyMem_RawCalloc(size_t nelem, size_t elsize)
//  void* PyMem_RawRealloc(void *p, size_t n)
//  void PyMem_RawFree(void *p)
//
// Include/complexobject.h:
//
//   #ifndef Py_LIMITED_API
//   typedef struct {
//       PyObject_HEAD
//       Py_complex cval;
//   } PyComplexObject;
//   #endif
//
// Objects/complexobject.c:
// 
//   Py_complex
//   PyComplex_AsCComplex(PyObject *op)
//   {
//       Py_complex cv;
//       PyObject *newop = NULL;
//  
//       assert(op);
//       /* If op is already of type PyComplex_Type, return its value */
//       if (PyComplex_Check(op)) {
//           return ((PyComplexObject *)op)->cval;
//       }
//   ... 

//------------------------------------------------------------------------------

PyDoc_STRVAR( module_doc, "in-place FFT");

static PyMethodDef fft_methods[] =
{
  // The cast of fft is necessary since PyCFunction values only take two
  // PyObject* parameters, but fft() takes three because of the keyword list.
  //
  { "fft", (PyCFunction) fft, METH_VARARGS | METH_KEYWORDS,
	PyDoc_STR( "fft(X,inv=False,pre=True,nmax=0)\n\n\
Overwrites X with FFT, with bit-reversed result locations.\n\n\
X must be a list of complex values, and len(X) must be a power of 2.\n\n\
If inv is true then the inverse FFT is computed.\n\n\
If pre is true then the FFT W factors are precomputed and saved.\n\
This can improve speed but requires extra memory equivalent to that used by X.\n\n\
Returns indices of the nmax largest result absolute values."
)},

  { NULL, NULL}
};

static struct PyModuleDef fft_module =
{
  PyModuleDef_HEAD_INIT,
  "fft",
  module_doc,
  0,
  fft_methods,
  NULL,
  NULL,
  NULL,
  NULL
};

PyMODINIT_FUNC PyInit_fft( void)
{
  return PyModule_Create( &fft_module);
}