#include "navier-stokes.h"
#include <math.h>
#include <stdlib.h>

#define M_PI 3.14159265358979323846

// compute solution as a function of time
int uk(
  int K1,
  int K2,
  int N1,
  int N2,
  unsigned int nsteps,
  double nu,
  double delta,
  _Complex double (*g)(int,int),
  unsigned int print_freq
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  unsigned int t;
  fft_vect fft1;
  fft_vect fft2;
  fft_vect ifft;
  int kx,ky;

  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
  ns_init_u(u, K1, K2);

  // iterate
  for(t=0;t<nsteps;t++){
    ins_step(u, K1, K2, N1, N2, nu, delta, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
    
    if(t%print_freq==0){
      fprintf(stderr,"%d % .8e ",t,t*delta);
      printf("%8d % .15e ",t,t*delta);

      for(kx=-K1;kx<=K1;kx++){
	for (ky=-K2;ky<=K2;ky++){
	  if (kx*kx+ky*ky<=1){
	    fprintf(stderr,"% .8e % .8e ",__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
	    
	  }
	  printf("% .15e % .15e ",__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
	}
      }
      fprintf(stderr,"\n");
      printf("\n");
    }
  }

  ns_free_tmps(u, tmp1, tmp2, tmp3, fft1, fft2, ifft);
  return(0);
}

// compute enstrophy as a function of time in the I-NS equation
int enstrophy(
  int K1,
  int K2,
  int N1,
  int N2,
  unsigned int nsteps,
  double nu,
  double delta,
  _Complex double (*g)(int,int),
  unsigned int print_freq
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  _Complex double alpha;
  _Complex double avg;
  unsigned int t;
  fft_vect fft1;
  fft_vect fft2;
  fft_vect ifft;

  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
  ns_init_u(u, K1, K2);


  // init running average
  avg=0;

  // iterate
  for(t=0;t<nsteps;t++){
    ins_step(u, K1, K2, N1, N2, nu, delta, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
    alpha=compute_alpha(u, K1, K2, g);
    
    // running average
    if(t>0){
      avg=avg-(avg-alpha)/t;
    }

    if(t>0 && t%print_freq==0){
      fprintf(stderr,"% .15e % .15e % .15e % .15e % .15e\n",t*delta, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
      printf("% .15e % .15e % .15e % .15e % .15e\n",t*delta, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
    }
  }

  ns_free_tmps(u, tmp1, tmp2, tmp3, fft1, fft2, ifft);
  return(0);
}


// initialize vectors for computation
int  ns_init_tmps(
  _Complex double ** u,
  _Complex double ** tmp1,
  _Complex double ** tmp2,
  _Complex double ** tmp3,
  fft_vect* fft1,
  fft_vect* fft2,
  fft_vect* ifft,
  int K1,
  int K2,
  int N1,
  int N2
){
  // velocity field
  *u=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));

  // allocate tmp vectors for computation
  *tmp1=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
  *tmp2=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
  *tmp3=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));

  // prepare vectors for fft
  fft1->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
  fft1->fft_plan=fftw_plan_dft_2d(N1,N2, fft1->fft, fft1->fft, FFTW_FORWARD, FFTW_MEASURE);
  fft2->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
  fft2->fft_plan=fftw_plan_dft_2d(N1,N2, fft2->fft, fft2->fft, FFTW_FORWARD, FFTW_MEASURE);
  ifft->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
  ifft->fft_plan=fftw_plan_dft_2d(N1,N2, ifft->fft, ifft->fft, FFTW_BACKWARD, FFTW_MEASURE);

  return 0;
}

// release vectors
int ns_free_tmps(
  _Complex double* u,
  _Complex double* tmp1,
  _Complex double* tmp2,
  _Complex double* tmp3,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft
){
  // free memory
  fftw_destroy_plan(fft1.fft_plan);
  fftw_destroy_plan(fft2.fft_plan);
  fftw_destroy_plan(ifft.fft_plan);
  fftw_free(fft1.fft);
  fftw_free(fft2.fft);
  fftw_free(ifft.fft);

  fftw_cleanup();

  free(tmp3);
  free(tmp2);
  free(tmp1);

  free(u);

  return 0;
}



// initial value
int ns_init_u(
  _Complex double* u,
  int K1,
  int K2
){
  int kx,ky;
  /*
  double rescale;

  srand(17);

  // random init (set half, then the other half are the conjugates)
  for(ky=0;ky<=K2;ky++){
    u[klookup(0,ky,2*K1+1,2*K2+1)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
  }
  for(kx=1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
    }
  }
  // conjugates
  for(ky=-K2;ky<=-1;ky++){
    u[klookup(0,ky,2*K1+1,2*K2+1)]=conj(u[klookup(0,-ky,2*K1+1,2*K2+1)]);
  }
  for(kx=-K1;kx<=-1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=conj(u[klookup(-kx,-ky,2*K1+1,2*K2+1)]);
    }
  }

  // rescale: 1/k decay
  rescale=0;
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      rescale=rescale+((__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)])*(__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)])+(__imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)])*(__imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]))*(kx*kx+ky*ky);
    }
  }
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]*sqrt(155.1/rescale);
    }
  }
  */


  /*
  // constant init
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=1.;
    }
  }
  */

  // exponentially decaying init
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=exp(-sqrt(kx*kx+ky*ky));
    }
  }

  return 0;
}


// next time step for Irreversible Navier-Stokes equation
int ins_step(
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double delta,
  _Complex double (*g)(int,int),
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  _Complex double* tmp1,
  _Complex double* tmp2,
  _Complex double* tmp3
){
  int kx,ky;

  // k1
  ins_rhs(tmp1, u, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k1/2
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k2
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k2/2
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k3
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k3
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k4
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  return(0);
}

// right side of Irreversible Navier-Stokes equation
int ins_rhs(
  _Complex double* out,
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  _Complex double (*g)(int,int),
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft
){
  int kx,ky;
  int i;

  // compute convolution term
  ns_T(u,K1,K2,N1,N2,fft1,fft2,ifft);

  for(i=0; i<(2*K1+1)*(2*K2+1); i++){
    out[i]=0;
  }
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if(kx!=0 || ky!=0){
        out[klookup(kx,ky,2*K1+1,2*K2+1)]=-4*M_PI*M_PI*nu*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]+(*g)(kx,ky)+4*M_PI*M_PI/sqrt(kx*kx+ky*ky)*ifft.fft[klookup(kx,ky,N1,N2)];
      }
    }
  }

  return(0);
}

// convolution term in right side of convolution equation
int ns_T(
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft
){
  int kx,ky;
  int i;

  // F(px/|p|*u)*F(qy*|q|*u)
  // init to 0
  for(i=0; i<N1*N2; i++){
    fft1.fft[i]=0;
    fft2.fft[i]=0;
    ifft.fft[i]=0;
  }
  // fill modes
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if(kx!=0 || ky!=0){
        fft1.fft[klookup(kx,ky,N1,N2)]=kx/sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=ky*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N2;
      }
    }
  }

  // fft
  fftw_execute(fft1.fft_plan);
  fftw_execute(fft2.fft_plan);
  // write to ifft
  for(i=0;i<N1*N2;i++){
    ifft.fft[i]=fft1.fft[i]*fft2.fft[i];
  }

  // F(py/|p|*u)*F(qx*|q|*u)
  // init to 0
  for(i=0; i<N1*N2; i++){
    fft1.fft[i]=0;
    fft2.fft[i]=0;
  }
  // fill modes
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if(kx!=0 || ky!=0){
        fft1.fft[klookup(kx,ky,N1,N2)]=ky/sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=kx*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N2;
      }
    }
  }

  // fft
  fftw_execute(fft1.fft_plan);
  fftw_execute(fft2.fft_plan);
  // write to ifft
  for(i=0;i<N1*N2;i++){
    ifft.fft[i]=ifft.fft[i]-fft1.fft[i]*fft2.fft[i];
  }

  // inverse fft
  fftw_execute(ifft.fft_plan);

  return(0);
}


// compute alpha
_Complex double compute_alpha(
  _Complex double* u,
  int K1,
  int K2,
  _Complex double (*g)(int,int)
){
  _Complex double num=0;
  _Complex double denom=0;
  int kx,ky;

  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      denom+=(kx*kx+ky*ky)*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)])*(1+(ky!=0?kx*kx/ky/ky:0));
      num+=(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj((*g)(kx,ky))*(1+(ky!=0?kx*kx/ky/ky:0));
    }
  }

  return(num/denom);
}


// get index for kx,ky in array of size S
int klookup(
  int kx,
  int ky,
  int S1,
  int S2
){
  return (kx>=0 ? kx : S1+kx)*S2 + (ky>=0 ? ky : S2+ky);
}