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

#include <pthread.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <complex.h>
#include "h.h"
#define cpus 2
// cpus is a flexible integer parameter > 0.
typedef struct{int first; int last;} blkmsg;
typedef unsigned char uchar;
typedef void * thnk(void *);
void fft(c *, const int);
#define wave 1
// wave is 0 or 1 depending on the test you want
static R scalex, scaley;
c * a; // 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) exit(printf("ERROR; return code from pthread_create() is %d\n", rc));
   pthread_attr_destroy(&attr);
   return zz;}
static int X, Y; static short LX, LY;
static void * doRows(void * ptm){blkmsg * x = ptm;
   if(1) 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[1<<LX];
      {int m=1<<LX; while(m--) ta[m] = a[n<<LX|m];}
      fft(ta, 1<<LX);
      {int m=1<<LX; while(m--) a[n<<LX|m] = scalex*ta[m];}}}  
   pthread_exit((void *) 0);}

static void * doCols(void * ptm){blkmsg * x = ptm;
   if(1) 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[1<<LY];
      {int m=1<<LY; while(m--) ta[m] = a[m<<LX|n];}
      fft(ta, 1<<LY);
      {int m=1<<LY; while(m--) a[m<<LX|n] = scaley*ta[m];}}}  
   pthread_exit((void *) 0);}

static void db(thnk fab, int sz) {
   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 R m2(c x){R r = creal(x), i = cimag(x); return r*r + i*i;}

void pp(int j, int k){c z = a[j<<LX|k]; printf(
         "%4d %4d %13.8f %13.8f %13.8f\n",
           j, k, creal(z), cimag(z), sqrt(m2(z)));}

static int fp(R th, char * m){R mx = 0, te = 0; int J=-1, K=-1;
   {int j = 1<<LY; while(j--) {int k=1<<LX; while(k--) 
     {R M2 = m2(a[j<<LX|k]); if(M2>mx) {J=j; K=k; mx=M2;} te += M2;}}}
    R my = sqrt(mx); int jm = J<<LX | K;
    printf("%s: th=%9.6f, peak at %d %d, max = %9.6f, te=%e\n",
       m, th, K, J, my, te); 
   if(0){int j = 1<<LY; R t = th*mx; int top=10;
      while(j--) {int k=1<<LX; while(k--) if(m2(a[j<<LX|k]) > t) {
      if(1) pp(j, k); if(!(top--)) return jm;}}}
  return jm;}

static void d2Dft(void){
  static int tcnt = 0;
  printf("Transform %d:\n", ++tcnt);
  db(doRows, 1<<LY); db(doCols, 1<<LX);}
void ext2D(c * b){c * ta = a; a = b; d2Dft(); a = ta;}


static int ri(FILE * in){int j=4;
  int X=0; while(j--) X = X<<8 | getc(in); return X;}
static void wi(FILE * out, int n){int y=4; while(y--) putc(n>>(8*y), out);}
static int li(int j){if((int)(0.5+cpus*((R)j/cpus)) != j) exit(printf("Oops!\n"));
  int k=1, z = j-1; while(z>>=1) ++k; return k;}

int main(int argc, char * * args){
  if(argc < 5) exit(printf("Usage:\n%s n infile outfile dotSize\n", args[0]));
  FILE * in = fopen(args[2], "r");
  if(!in) exit(printf("no such file: %s\n", args[2]));
  X = ri(in); Y = ri(in);
  printf("X=%d, Y=%d\n", X, Y);
  LX = li(X), LY = li(Y);
  int L2=LX+LY;
  scalex = 1/sqrt(1<<LX); scaley = 1/sqrt(1<<LY);
  uchar * gl(){long int * c = (long int *)malloc(1<<L2);
    int j=(1<<(L2-3)); while(j--) c[j] = 0; return (uchar *) c;}
  uchar *cr = gl(), *cg = gl(), *cb = gl();
  {int j; for(j=0; j<Y; ++j) {int k; for(k=0; k<X; ++k)
     {int n = j<<LX | k; cr[n]=getc(in); cg[n]=getc(in); cb[n]=getc(in);}}}
  fclose(in);
  a = (c *)malloc(sizeof(c)<<L2);
  if(1) // Choice of two demos
    { // transform each color separately.
    void dc(uchar * c){
        int j=1<<L2; while(j--) {a[j] = c[j];}
      if(0) {printf("avant "); pp(43, 46);}
      fp(.9, "frst"); d2Dft(); fp(.9, "mid");
      disp(argc-1, args+1, 1<<LX, 1<<LY);
      d2Dft(); R s = 255/sqrt(m2(a[fp(0.9, "after")]));
      if(0) {printf("après: "); pp(43, 46);}
      printf("scaled by %e\n", s);
      {int j = 1<<L2; while(j--) {c[j] = s*creal(a[j])+0.5;
         if(!(random()&0xffff)) pp(j>>LX, j&((1<<LX)-1));}}}
    {dc(cr); dc(cg); dc(cb);}}
  else { // average input colors to gray; transform and use color to report.
    {int j=1<<L2; while(j--) a[j] = cr[j]+cg[j]+cb[j];}
    R av = ({R s = 0; int j=X; while(j--) {int k=Y; while(k--){
    s += creal(a[k<<LX|j]);}} s;})/(X*Y);
    {int y; for(y=Y; y<1<<LY; ++y) {int x; for(x=0; x<1<<LX; ++x)
        a[y<<LX|y] = av;}}
    {int x; for(x=X; x<1<<LX; ++x) {int y; for(y=0; y<1<<LY; ++y)
        a[y<<LX|y] = av;}}
    if(0) d2Dft(); int m = fp(0.005, "color?"); R s = 22800/sqrt(m2(a[m]));
    c i = _Complex_I, vr = 1, vg = 0.5+i*sqrt(3)/2, vb = 0.5-i*sqrt(3)/2;
    printf("m = %d, s = %e+i*%e\n", m, creal(s), cimag(s));
    {int j=1<<L2; while(j--) {
      cr[j] = 128+s*creal(a[j]*vr);
      cg[j] = 128+s*creal(a[j]*vg);
      cb[j] = 128+s*creal(a[j]*vb);
      if(!(random()&0xfff)) printf("%d %d %d %d\n", j, cr[j], cg[j], cb[j]);
      }}}
  FILE * out = fopen(args[3], "w");
  wi(out, 1<<LX); wi(out, 1<<LY);
  {int j; for(j=0; j<(1<<L2); ++j) {
    putc(cr[j], out); putc(cg[j], out); putc(cb[j], out);}}
    fclose(out); return 0;}