#define VERSION "0.0"

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

// usage message
int print_usage();
// read command line arguments
int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nsteps, unsigned int* computation_nr);

// compute enstrophy as a function of time in the I-NS equation
int enstrophy(ns_params params, unsigned int Nsteps);


#define COMPUTATION_ENSTROPHY 1
int main (int argc, const char* argv[]){
  ns_params params;
  unsigned int nsteps;
  int ret;
  unsigned int computation_nr;

  // default computation: phase diagram
  computation_nr=COMPUTATION_ENSTROPHY;

  // read command line arguments
  ret=read_args(argc, argv, &params, &nsteps, &computation_nr);
  if(ret<0){
    return(-1);
  }
  if(ret>0){
    return(0);
  }

  // enstrophy
  if(computation_nr==COMPUTATION_ENSTROPHY){
    enstrophy(params, nsteps);
  }

  return(0);
}

// usage message
int print_usage(){
  fprintf(stderr, "usage:\n       nstrophy enstrophy [-h timestep] [-K modes] [-v] [-N nsteps]\n\n       nstrophy -V [-v]\n\n");
  return(0);
}

// read command line arguments
#define CP_FLAG_TIMESTEP 1
#define CP_FLAG_NSTEPS 2
#define CP_FLAG_MODES 2
#define CP_FLAG_NU 3
int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nsteps, unsigned int* computation_nr){
  int i;
  int ret;
  // temporary int
  int tmp_int;
  // temporary unsigned int
  unsigned int tmp_uint;
  // temporary double
  double tmp_double;
  // pointers
  char* ptr;
  // flag that indicates what argument is being read
  int flag=0;
  // print version and exit
  char Vflag=0;

  // defaults
  /*
  params->K=16;
  params->h=1e-3/(2*params->K+1);
  *nsteps=10000000;
  params->nu=1./1024/(2*params->K+1);
  */
  params->K=16;
  params->h=0.0001220703125;
  params->h=0.001953125;
  *nsteps=10000000;
  //params->nu=0.00048828125;
  params->nu=0.0078125;

  // loop over arguments
  for(i=1;i<argc;i++){
    // flag
    if(argv[i][0]=='-'){
      for(ptr=((char*)argv[i])+1;*ptr!='\0';ptr++){
	switch(*ptr){
	// timestep
	case 'h':
	  flag=CP_FLAG_TIMESTEP;
	  break;
	// nsteps
	case 'N':
	  flag=CP_FLAG_NSTEPS;
	  break;
	// modes
	case 'K':
	  flag=CP_FLAG_MODES;
	  break;
	// friction
	case 'n':
	  flag=CP_FLAG_NU;
	  break;
	// print version
	case 'V':
	  Vflag=1;
	  break;
	default:
	  fprintf(stderr, "unrecognized option '-%c'\n", *ptr);
	  print_usage();
	  return(-1);
	  break;
	}
      }
    }
    // timestep
    else if(flag==CP_FLAG_TIMESTEP){
      ret=sscanf(argv[i],"%lf",&tmp_double);
      if(ret!=1){
	fprintf(stderr, "error: '-h' should be followed by a double\n       got '%s'\n",argv[i]);
	return(-1);
      }
      params->h=tmp_double;
      flag=0;
    }
    // nsteps
    else if(flag==CP_FLAG_NSTEPS){
      ret=sscanf(argv[i],"%u",&tmp_uint);
      if(ret!=1){
	fprintf(stderr, "error: '-N' should be followed by an unsigned int\n       got '%s'\n",argv[i]);
	return(-1);
      }
      *nsteps=tmp_uint;
      flag=0;
    }
    // nsteps
    else if(flag==CP_FLAG_MODES){
      ret=sscanf(argv[i],"%d",&tmp_int);
      if(ret!=1){
	fprintf(stderr, "error: '-K' should be followed by an int\n       got '%s'\n",argv[i]);
	return(-1);
      }
      params->K=tmp_uint;
      flag=0;
    }
    // friction
    else if(flag==CP_FLAG_TIMESTEP){
      ret=sscanf(argv[i],"%lf",&tmp_double);
      if(ret!=1){
	fprintf(stderr, "error: '-n' should be followed by a double\n       got '%s'\n",argv[i]);
	return(-1);
      }
      params->nu=tmp_double;
      flag=0;
    }
    // computation to run
    else{
      if(strcmp(argv[i], "enstrophy")==0){
	*computation_nr=COMPUTATION_ENSTROPHY;
      }
      else{
	fprintf(stderr, "error: unrecognized computation: '%s'\n",argv[i]);
	print_usage();
	return(-1);
      }
      flag=0;
    }
  }

  // print version and exit
  if(Vflag==1){
    printf("nstrophy " VERSION "\n");
    return(1);
  }

  return(0);
}

// compute enstrophy as a function of time in the I-NS equation
int enstrophy(ns_params params, unsigned int Nsteps){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  _Complex double alpha;
  _Complex double avg;
  unsigned int t;
  int kx,ky;
  fft_vects fft_vects;
  double rescale;

  // sizes
  params.S=2*params.K+1;
  params.N=4*params.K+1;

  // velocity field
  u=calloc(sizeof(_Complex double),params.S*params.S);
  params.g=calloc(sizeof(_Complex double),params.S*params.S);
  // allocate tmp vectors for computation
  tmp1=calloc(sizeof(_Complex double),params.S*params.S);
  tmp2=calloc(sizeof(_Complex double),params.S*params.S);
  tmp3=calloc(sizeof(_Complex double),params.S*params.S);

  srand(17);

  // initial value
  for(ky=0;ky<=params.K;ky++){
    u[KLOOKUP(0,ky,params.S)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
  }
  for(kx=1;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      u[KLOOKUP(kx,ky,params.S)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
    }
  }
  for(ky=-params.K;ky<=-1;ky++){
    u[KLOOKUP(0,ky,params.S)]=conj(u[KLOOKUP(0,-ky,params.S)]);
  }
  for(kx=-params.K;kx<=-1;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      u[KLOOKUP(kx,ky,params.S)]=conj(u[KLOOKUP(-kx,-ky,params.S)]);
    }
  }
  rescale=0;
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      rescale=rescale+((__real__ u[KLOOKUP(kx,ky,params.S)])*(__real__ u[KLOOKUP(kx,ky,params.S)])+(__imag__ u[KLOOKUP(kx,ky,params.S)])*(__imag__ u[KLOOKUP(kx,ky,params.S)]))*(kx*kx+ky*ky);
    }
  }
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      u[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]*sqrt(155.1/rescale);
    }
  }

  /*
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      printf("%d %d % .8e % .8e\n",kx,ky, __real__ u[KLOOKUP(kx,ky,params.S)], __imag__ u[KLOOKUP(kx,ky,params.S)]);
    }
  }
  */


  // driving force
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      //params.g[KLOOKUP(kx,ky,params.S)]=sqrt(kx*kx*ky*ky)*exp(-(kx*kx+ky*ky));
      if(kx==2 && ky==-1){
	params.g[KLOOKUP(kx,ky,params.S)]=0.5+sqrt(3)/2*I;
      }
      else if(kx==-2 && ky==1){
	params.g[KLOOKUP(kx,ky,params.S)]=0.5-sqrt(3)/2*I;
      }
      else{
	params.g[KLOOKUP(kx,ky,params.S)]=0;
      }
    }
  }


  // prepare vectors for fft
  fft_vects.fft1=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
  fft_vects.fft1_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.fft1, fft_vects.fft1, FFTW_FORWARD, FFTW_MEASURE);
  fft_vects.fft2=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
  fft_vects.fft2_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.fft2, fft_vects.fft2, FFTW_FORWARD, FFTW_MEASURE);
  fft_vects.invfft=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
  fft_vects.invfft_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.invfft, fft_vects.invfft, FFTW_BACKWARD, FFTW_MEASURE);

  // init running average
  avg=0;

  // iterate
  for(t=0;t<Nsteps;t++){
    ins_step(u, params, fft_vects, tmp1, tmp2, tmp3);
    alpha=compute_alpha(u, params);

    /*
    // to avoid errors building up in imaginary part
    for(kx=-params.K;kx<=params.K;kx++){
      for(ky=-params.K;ky<=params.K;ky++){
	u[KLOOKUP(kx,ky,params.S)]=__real__ u[KLOOKUP(kx,ky,params.S)];
      }
    }
    */

    // running average
    if(t>0){
      avg=avg-(avg-alpha)/t;
    }

    if(t>0 && t%1000==0){
      fprintf(stderr,"%8d % .8e % .8e % .8e % .8e\n",t, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
      printf("%8d % .8e % .8e % .8e % .8e\n",t, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
    }
  }

  // free memory
  fftw_destroy_plan(fft_vects.fft1_plan);
  fftw_destroy_plan(fft_vects.fft2_plan);
  fftw_destroy_plan(fft_vects.invfft_plan);
  fftw_free(fft_vects.fft1);
  fftw_free(fft_vects.fft2);
  fftw_free(fft_vects.invfft);

  free(tmp3);
  free(tmp2);
  free(tmp1);
  free(params.g);
  free(u);

  return(0);
}