// This uses several cpus to do 2D square fft.

#include <pthread.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <complex.h>
#define cpus 1
// cpus is a flexible integer parameter > 0.
typedef struct{int first; int last; int j;} blkmsg;
typedef _Complex float c;
typedef double R;
typedef void * thnk(void *);
void fft(const int, c *);
#define sz 4096
// sz is a flexible parameter but must be a power of 2.
#define wave 1
// wave is 0 or 1 depending on the test you want
R scale;
int mc[cpus]; // .... count little fft's
c a[sz][sz]; // The data to be transformed.
static int min(int a, int b){return a<b?a:b;}
static pthread_t mt(void * msg, thnk rtn){
   pthread_attr_t attr;
   pthread_attr_init(&attr);
   pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
   {
      pthread_t zz;
      int rc = pthread_create(&zz, &attr, rtn, msg);
      if(rc) {printf(
            "ERROR; return code from pthread_create() is %d\n", rc);
         exit(-1);}
       pthread_attr_destroy(&attr);
       return zz;}}

static void * doRows(void * ptm){blkmsg * x = ptm;
   if(0) printf("Kilroy %d R %d thru %d\n", x-> j, x->first, x->last);
   {int l = x->last; for(int n = x->first; n<= l; ++n) {
      c ta[sz];
      for(int m=0; m<sz; ++m) ta[m] = a[n][m];
      fft(sz, ta); ++mc[x->j];
      for(int m=0; m<sz; ++m) a[n][m] = scale*ta[m];
      }}  
   pthread_exit((void *) 0);}

static void * doCols(void * ptm){blkmsg * x = ptm;
   if(0) printf("Kilroy R %d thru %d\n", x->first, x->last);
   {int l = x->last; for(int n = x->first; n<= l; ++n) {
      c ta[sz];
      for(int m=0; m<sz; ++m) ta[m] = a[m][n];
      fft(sz, ta); ++mc[x->j];
      for(int m=0; m<sz; ++m) a[m][n] = scale*ta[m];}}
   pthread_exit((void *) 0);}

static void db(thnk fab) {const int b = sz/cpus+(sz%cpus>0);
   blkmsg bm[cpus];
   pthread_t thrds[cpus];
   for(int n = 0; n < cpus; ++n) {bm[n] = (blkmsg){n*b, min((n+1)*b, sz)-1, n};      
          thrds[n] = mt(&bm[n], fab);}
   for(int n = 0; n < cpus; ++n) {int rc = pthread_join(thrds[n], NULL);
          if(rc) {printf("ERROR; return code from pthread_join() is %d\n", rc);
             exit(-1);}}}

static void d2Dft(void){
  static int tcnt = 0;
  printf("Transform %d:\n", ++tcnt);
  db(doRows); db(doCols);
  {int sm=0; 
  for(int j = 0; j<cpus; ++j) {printf("mc[%d]=%d\n", j, mc[j]);
    sm += mc[j];} printf("tot=%d\n", sm);}}

static R m2(c x){R r = creal(x), i = cimag(x); return r*r + i*i;}

static void pp(int j, int k){c z = a[j][k]; R m = m2(z); if(m>1.e-8) printf(
         "%4d %4d %13.8f %13.8f %13.8f\n",
           j, k, creal(z), cimag(z), sqrt(m));}

static void psam(int m){printf("sample %d\n", m);
   for(int j=81; j<85; ++j) for(int k=384; k<387; ++k) pp(j, k);}

static void fp(R th, char * m){R mx = 0, te = 0; int J=-1, K=-1;
   for(int j=0; j<sz; ++j) for(int k=0; k<sz; ++k)
      {R M2 = m2(a[j][k]); if(M2>mx) {J=j; K=k; mx=M2;} te += M2;}
   printf("%s: th=%9.6f, peak at %d %d, max = %9.6f, te=%e\n",
       m, th, J, K, sqrt(mx), te); 
   {R t = th*mx;
      for(int j=0; j<sz; ++j) for(int k=0; k<sz; ++k) if(m2(a[j][k]) > t) pp(j, k);}}

int main(){scale = 1./sqrt(sz); for(int j=0; j<cpus; ++j) mc[j]=0;
  if(sz&(sz-1)) exit(printf("fft works only for powers of 2\n"));
  printf(wave?"One wave\n":"One point\n");
  for(int j=0; j<sz; ++j) for(int k=0; k<sz; ++k)
            a[j][k] = wave?cexp((4*j+3*k)/25.*I) : (j==42)&(k==73);
  if(!wave) fp(0.999, "initial peak"); psam(0);
  d2Dft(); if(wave) fp(0.002, "once"); psam(1);
  d2Dft(); if(wave) fp(0.9999996, "twice"); else fp(0.999, "regenerated peak"); psam(2);
  d2Dft(); if(wave) fp(0.002, "thrice"); else psam(3);
  d2Dft(); fp(0.9999999, "fourth"); psam(4);
  return 0;}

/*
These results are for 2.7 GHz Intel Core i7
Four cores each doing two threads;
cpus=1; sz=4096
real	0m6.520s
user	0m6.456s
sys	0m0.051s

cpus=2; sz=4096
real	0m3.663s
user	0m6.587s
sys	0m0.044s

cpus=4; sz=4096
real	0m2.598s
user	0m8.129s
sys	0m0.054s

cpus=8; sz=4096
real	0m2.245s
user	0m12.361s
sys	0m0.057s

cpus=16; sz=4096
real	0m2.204s
user	0m12.465s
sys	0m0.055s
*/