/*************************    Program 5.3    ****************************/
/*                                                                      */
/************************************************************************/
/* Please Note:                                                         */
/*                                                                      */
/* (1) This computer program is written by Tao Pang in conjunction with */
/*     his book, "An Introduction to Computational Physics," published  */
/*     by Cambridge University Press in 1997.                           */
/*                                                                      */
/* (2) No warranties, express or implied, are made for this program.    */
/*                                                                      */
/************************************************************************/

#include <stdio.h>
#include <math.h>
#define NMAX1 128
#define NMAX2 128

void fft2d (fr,fi,n1,n2,m1,m2)
/* Subroutine for the two-dimensional fast Fourier transform with
   N=N1*N2 and N1=2**M1 and N2=2**M2.  Copyright (c) Tao Pang 1997. */
int n1,n2,m1,m2;
double fr[NMAX1][NMAX2],fi[NMAX1][NMAX2];
{
int i,j;
void fft();

/* Transformation on the second index */

for (i = 0; i < n1; ++i)
  {
  fft (fr[i][0],fi[i][0],n2,m2);
  }

/* Transformation on the first index */

for (j = 0; j < n2; ++j)
  {
  fft (fr[0][j],fi[0][j],n1,m1);
  }
}
 
void fft (ar,ai,n,m)
/* An example of the fast Fourier transform subroutine with
   N = 2**M.  AR and AI are the real and imaginary part of
   data in the input and corresponding Fourier coefficients
   in the output.  Copyright (c) Tao Pang 1997. */
int n,m;
double ar[],ai[];
{
int i,j,k,l,n1,n2,l1,l2;
double pi,a1,a2,q,v,u;

pi = 4*atan(1);
n2 = n/2;

n1 = pow(2,m);
if (n1 != n)
  {
  printf("Indices do not match\n");
  exit(1);
  }

/* Rearrange the data to the bit reversed order */

l = 0;
for (k = 0; k < n-1; ++k)
  {
  if (k < l)
    {
    a1 = ar[l];
    a2 = ai[l];
    ar[l] = ar[k];
    ar[k] = a1;
    ai[l] = ai[k];
    ai[k] = a2;
    }
  j = n2;
  while (j < l)
    {
    l = l-j;
    j = j/2;
    }
  l = l+j;
  }

/* Perform additions at all levels with reordered data */

l2 = 1;
for (l = 0; l < m; ++l)
  {
  q = 0;
  l1 = l2;
  l2 = 2*l1;
  for (k = 0; k < l1; ++k)
    {
    u =  cos(q);
    v = -sin(q);
    q =  q+pi/l1;
    for (j = k; j < n; j = j+l2)
      {
      i = j+l1;
      a1 = ar[i]*u-ai[i]*v;
      a2 = ar[i]*v+ai[i]*u;
      ar[i] = ar[j]-a1;
      ar[j] = ar[j]+a1;
      ai[i] = ai[j]-a2;
      ai[j] = ai[j]+a2;
      }
    }
  }
}