// 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 2
// cpus is a flexible integer parameter > 0.
typedef struct{int first; int last;} blkmsg;
typedef _Complex float c;
typedef double R;
typedef void * thnk(void *);
void fft(c *, const int);
#define sz 2048
// 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;
c a[sz][sz]; // The data to be transformed.
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 R %d thru %d\n", x->first, x->last);
   {int l = x->last, n; for(n = x->first; n<= l; ++n) {
      c ta[sz];
      {int m=sz; while(m--) ta[m] = a[n][m];}
      fft(ta, sz);
      {int m=sz; while(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, n; for(n = x->first; n<= l; ++n) {
      c ta[sz];
      {int m=sz; while(m--) ta[m] = a[m][n];}
      fft(ta, sz);
      {int m=sz; while(m--) a[m][n] = scale*ta[m];}}}  
   pthread_exit((void *) 0);}

static void db(thnk fab) {
   int n; R rps = (R)sz/cpus;
   blkmsg bm[cpus];
   pthread_t thrds[cpus];
      for(n = 0; n < cpus; ++n) {bm[n] = (blkmsg){n*rps, (n+1)*rps-0.5};
          thrds[n] = mt(&bm[n], fab);}
      for(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);}

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]; printf(
         "%4d %4d %13.8f %13.8f %13.8f\n",
           j, k, creal(z), cimag(z), sqrt(m2(z)));}

static void psam(char * m){printf("sample %s\n", m);
  int j, k; for(j=81; j<85; ++j) for(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;
   {int j = sz; while(j--) {int k=sz; while(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); 
   {int j = sz; R t = th*mx;
      while(j--) {int k=sz; while(k--) if(m2(a[j][k]) > t) pp(j, k);}}}
   
int main(){scale = 1./sqrt(sz);
  if((int)(0.5+cpus*((R)sz/cpus)) != sz) exit(printf("Oops!\n"));
  if(sz&(sz-1)) exit(printf("fft works only for powers of 2\n"));
  {int j = sz; while(j--) {int k=sz; while(k--) a[j][k] =
        wave?cexp((4*j+3*k)/25.*I) : (j==42)&(k==73);}}
  psam("initial");
  if(!wave) fp(0.999, "initial peak");
  d2Dft();
  if(wave) fp(0.01, "once");
  psam("one trnsf");
  d2Dft();
  if(wave) fp(0.9999993, "twice");
  else fp(0.999, "regenerated peak");
  psam("two trnsf");
  d2Dft();
  if(wave) fp(0.002, "thrice");
  else psam("two trnsf");
  d2Dft();
  fp(0.999999, "fourth");
   if(!wave)  psam("4 trnsf");
  return 0;}