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

// compute solution as a function of time
int uk(
  int K1,
  int K2,
  int N1,
  int N2,
  unsigned int nsteps,
  double nu,
  double delta,
  double L,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  unsigned int print_freq,
  unsigned int starting_time,
  unsigned int nthreads,
  FILE* savefile
){
  _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, nthreads);
  // copy initial condition
  copy_u(u, u0, K1, K2);

  // print column headers
  printf("#    1:i                    2:t ");
  t=3;
  for(kx=-K1;kx<=K1;kx++){
    for (ky=-K2;ky<=K2;ky++){
      printf("   %6d:(%4d,%4d)r ",t,kx,ky);
      t++;
      printf("   %6d:(%4d,%4d)i ",t,kx,ky);
      t++;
    }
  }

  // iterate
  for(t=starting_time;t<starting_time+nsteps;t++){
    ns_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    
    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__ getval_sym(u,kx,ky,K2), __imag__ getval_sym(u, kx,ky,K2));
	    
	  }
	  printf("% .15e % .15e ",__real__ getval_sym(u, kx,ky,K2), __imag__ getval_sym(u, kx,ky,K2));
	}
      }
      fprintf(stderr,"\n");
      printf("\n");
    }
  }

  // save final entry to savefile
  write_u(u, K1, K2, savefile);

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

// compute energy, enstrophy, alpha as a function of time in the I-NS equation
int eea(
  int K1,
  int K2,
  int N1,
  int N2,
  unsigned int nsteps,
  double nu,
  double delta,
  double L,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  unsigned int print_freq,
  unsigned int running_avg_window,
  unsigned int starting_time,
  unsigned int nthreads,
  FILE* savefile,
  // for interrupt recovery
  char* cmd_string,
  char* params_string,
  char* savefile_string
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  double alpha, energy, enstrophy;
  double avg_e,avg_a,avg_en;
  // quantities needed to compute averages
  double *save_print_e, *save_print_a, *save_print_en;
  double *save_print_short_e, *save_print_short_a, *save_print_short_en;
  // index
  unsigned int i;
  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, nthreads);
  // copy initial condition
  copy_u(u, u0, K1, K2);


  // init running average
  avg_e=0;
  avg_a=0;
  avg_en=0;

  // need this for running average
  unsigned int short_len=running_avg_window % print_freq;
  // number of full print intervals in running average window
  unsigned int nr_print_in_window = (running_avg_window-short_len)/print_freq;
  double avg_print_e=0;
  double avg_print_a=0;
  double avg_print_en=0;
  double avg_print_short_e=0;
  double avg_print_short_a=0;
  double avg_print_short_en=0;

  // save e a en for running average
  // only need to save if window is non zero
  if(running_avg_window!=0){
    save_print_e=calloc(sizeof(double), nr_print_in_window);
    save_print_a=calloc(sizeof(double), nr_print_in_window);
    save_print_en=calloc(sizeof(double), nr_print_in_window);
    save_print_short_e=calloc(sizeof(double), nr_print_in_window+1);
    save_print_short_a=calloc(sizeof(double), nr_print_in_window+1);
    save_print_short_en=calloc(sizeof(double), nr_print_in_window+1);
  }

  // iterate
  for(t=starting_time;t<starting_time+nsteps;t++){
    ns_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);

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

    energy=compute_energy(u, K1, K2);
    alpha=compute_alpha(u, K1, K2, N1, N2, g, L, ifft.fft);
    enstrophy=compute_enstrophy(u, K1, K2, L);
    
    // running average
    // if window=0, then take average over all times
    if(running_avg_window==0){
      if(t!=0){
	avg_e=avg_e-(avg_e-energy)/t;
	avg_a=avg_a-(avg_a-alpha)/t;
	avg_en=avg_en-(avg_en-enstrophy)/t;
      }
    }
    else{
      // reset print interval averages
      if(t % print_freq == 1){
	avg_print_e=0;
	avg_print_a=0;
	avg_print_en=0;
	avg_print_short_e=0;
	avg_print_short_a=0;
	avg_print_short_en=0;
      }

      // compute average over print interval
      avg_print_e=avg_print_e+energy/print_freq;
      avg_print_a=avg_print_a+alpha/print_freq;
      avg_print_en=avg_print_en+enstrophy/print_freq;

      // compute average over the last short_len elements of a print interval
      if (short_len != 0 && (t % print_freq > print_freq-short_len || t % print_freq == 0)){
	avg_print_short_e+=energy/short_len;
	avg_print_short_a+=alpha/short_len;
	avg_print_short_en+=enstrophy/short_len;
      }

      if(t % print_freq ==0 && t>starting_time){
	// compute averages
	if (t > running_avg_window + starting_time) {
	  avg_e=save_print_short_e[nr_print_in_window]*((double)short_len/running_avg_window);
	  avg_a=save_print_short_a[nr_print_in_window]*((double)short_len/running_avg_window);
	  avg_en=save_print_short_en[nr_print_in_window]*((double)short_len/running_avg_window);
	  for(i=0;i<nr_print_in_window;i++){
	    // prevent warnings: these will be initialized as long as t > running_avg_window
	    #pragma GCC diagnostic push
	    #pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
	    avg_e+=save_print_e[i]*((double)print_freq/running_avg_window);
	    avg_a+=save_print_a[i]*((double)print_freq/running_avg_window);
	    avg_en+=save_print_en[i]*((double)print_freq/running_avg_window);
	    #pragma GCC diagnostic pop
	  }
	}

	// memorize averages for future use
	// need to keep track of case ==0, otherwise nr_print_in_window-1 is not an unsigned int
	if(nr_print_in_window>0){
	  for(i=0;i<nr_print_in_window-1;i++){
	    save_print_e[i+1]=save_print_e[i];
	    save_print_a[i+1]=save_print_a[i];
	    save_print_en[i+1]=save_print_en[i];
	  }
	}
	for(i=0;i<nr_print_in_window;i++){
	  save_print_short_e[i+1]=save_print_short_e[i];
	  save_print_short_a[i+1]=save_print_short_a[i];
	  save_print_short_en[i+1]=save_print_short_en[i];
	}
	save_print_e[0]=avg_print_e;
	save_print_a[0]=avg_print_a;
	save_print_en[0]=avg_print_en;
	save_print_short_e[0]=avg_print_short_e;
	save_print_short_a[0]=avg_print_short_a;
	save_print_short_en[0]=avg_print_short_en;
      }
    }

    if(t>running_avg_window + starting_time && t%print_freq==0){
      fprintf(stderr,"%d % .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",t,t*delta, avg_a, avg_e, avg_en, alpha, energy, enstrophy);
      printf("%8d % .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",t,t*delta, avg_a, avg_e, avg_en, alpha, energy, enstrophy);
    }

    // catch abort signal
    if (g_abort){
      // print u to stderr if no savefile
      if (savefile==NULL){
	savefile=stderr;
      }
      fprintf(savefile,"# Interrupted computation. Resume with\n");
      // command to resume
      fprintf(savefile,"#! ");

      fprintf(savefile, cmd_string);

      // params
      // allocate buffer for params
      char* params=calloc(sizeof(char), strlen(params_string));
      strcpy(params, params_string);
      remove_entry(params, "starting_time");
      remove_entry(params, "init");
      remove_entry(params, "nsteps");
      fprintf(savefile," -p \"%s;starting_time=%u;nsteps=%u;init=file:%s\"", params, t+1, nsteps-t-1, savefile_string);
      free(params);

      fprintf(savefile," energy\n");

      break;
    }
  }

  // save final entry to savefile
  write_u(u, K1, K2, savefile);

  if(running_avg_window!=0){
    free(save_print_e);
    free(save_print_a);
    free(save_print_en);
    free(save_print_short_e);
    free(save_print_short_a);
    free(save_print_short_en);
  }

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

// compute solution as a function of time, but do not print anything (useful for debugging)
int quiet(
  int K1,
  int K2,
  int N1,
  int N2,
  unsigned int nsteps,
  double nu,
  double delta,
  double L,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  unsigned int nthreads,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  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, nthreads);
  // copy initial condition
  copy_u(u, u0, K1, K2);

  // iterate
  for(t=0;t<nsteps;t++){
    ns_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
  }

  // save final entry to savefile
  write_u(u, K1, K2, savefile);

  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,
  unsigned int nthreads
){
  // velocity field
  *u=calloc(sizeof(_Complex double),(K1+1)*(2*K2+1));

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

  // init threads
  fftw_init_threads();
  fftw_plan_with_nthreads(nthreads);

  // 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_threads();

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

  free(u);

  return 0;
}



// copy u0 to u
int copy_u(
  _Complex double* u,
  _Complex double* u0,
  int K1,
  int K2
){
  int i;

  for(i=0;i<(K1+1)*(2*K2+1);i++){
    u[i]=u0[i];
  }

  return 0;
}

// next time step
int ns_step(
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double delta,
  double L,
  _Complex double* g,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  _Complex double* tmp1,
  _Complex double* tmp2,
  _Complex double* tmp3,
  bool irreversible
){
  int kx,ky;

  // k1
  ns_rhs(tmp1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
  // add to output
  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta/6*tmp1[klookup_sym(kx,ky,K2)];
    }
  }

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

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

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

  return(0);
}

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

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

  if (irreversible) {
    alpha=nu;
  } else {
    alpha=compute_alpha(u,K1,K2,N1,N2,g,L,ifft.fft);
  }

  /*
  // compare convolution term (store result in fft1.fft)
  ns_T_nofft(fft1.fft, u, K1, K2, N1, N2);
  double cmp=0.;
  for(i=0;i<N1*N2;i++){
    cmp+=(ifft.fft[i]-fft1.fft[i])*(ifft.fft[i]-fft1.fft[i]);
  }
  printf("% .15e\n",sqrt(cmp));
  */

  for(i=0; i<(K1+1)*(2*K2+1); i++){
    out[i]=0;
  }
  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if(kx!=0 || ky!=0){
        out[klookup_sym(kx,ky,K2)]=-4*M_PI*M_PI/L/L*alpha*(kx*kx+ky*ky)*u[klookup_sym(kx,ky,K2)]+g[klookup_sym(kx,ky,K2)]+4*M_PI*M_PI/L/L/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)*getval_sym(u, kx,ky,K2)/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=ky*sqrt(kx*kx+ky*ky)*getval_sym(u, kx,ky,K2)/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)*getval_sym(u, kx,ky,K2)/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=kx*sqrt(kx*kx+ky*ky)*getval_sym(u, kx,ky,K2)/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);

  // enforce T(u,-k)=T(u,k)^*
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      ifft.fft[klookup(kx,ky,N1,N2)]=(ifft.fft[klookup(kx,ky,N1,N2)]+conj(ifft.fft[klookup(-kx,-ky,N1,N2)]))/2;
    }
  }

  return(0);
}

// convolution term in right side of convolution equation, computed without fourier transform
int ns_T_nofft(
  _Complex double* out,
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2
){
  int kx,ky;
  int px,py;
  int qx,qy;

  // loop over K's (needs N1>=4*K1+1 and N2>=4*K2+1)
  if (N1<4*K1+1 || N2<4*K2+1){
    fprintf(stderr,"error: N1 and N2 need t be >= 4*K1+1 and 4*K2+1 respectively\n");
    return(-1);
  }
  for(kx=-2*K1;kx<=2*K1;kx++){
    for(ky=-2*K2;ky<=2*K2;ky++){
      // init
      out[klookup(kx,ky,N1,N2)]=0.;

      for(px=-K1;px<=K1;px++){
	for(py=-K2;py<=K2;py++){
	  qx=kx-px;
	  qy=ky-py;

	  // cutoff in q
	  if(qx>=-K1 && qx<=K1 && qy>=-K2 && qy<=K2 && qx*qx+qy*qy>0 && px*px+py*py>0){
	    out[klookup(kx,ky,N1,N2)]+=(-qx*py+qy*px)*sqrt(qx*qx+qy*qy)/sqrt(px*px+py*py)*getval_sym(u, px,py,K2)*getval_sym(u, qx,qy,K2);
	  }
	}
      }
    }
  }

  // enforce T(u,-k)=T(u,k)^*
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      out[klookup(kx,ky,N1,N2)]=(out[klookup(kx,ky,N1,N2)]+conj(out[klookup(-kx,-ky,N1,N2)]))/2;
    }
  }

  return 0;
}

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

  num=0.;
  denom=0.;

  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      num+=(L*L/4/M_PI/M_PI*(kx*kx+ky*ky)*getval_sym(g, kx,ky,K2)+sqrt(kx*kx+ky*ky)*T[klookup(kx,ky,N1,N2)])*conj(getval_sym(u, kx,ky,K2));
      denom+=__real__ (kx*kx+ky*ky)*(kx*kx+ky*ky)*getval_sym(u, kx,ky,K2)*conj(getval_sym(u, kx,ky,K2));
    }
  }

  return __real__ num/denom;
}


/*
// compute alpha
double compute_alpha(
  _Complex double* u,
  int K1,
  int K2,
  _Complex double* g
){
  _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)*getval_sym(u, kx,ky,K2)*conj(getval_sym(u, kx,ky,K2))*(1+(ky!=0?kx*kx/ky/ky:0));
      num+=(kx*kx+ky*ky)*getval_sym(u, kx,ky,K2)*conj(g[getval(kx,ky,K1,K2)])*(1+(ky!=0?kx*kx/ky/ky:0));
    }
  }

  return __real__ (num/denom);
}*/

// compute energy
double compute_energy(
  _Complex double* u,
  int K1,
  int K2
){
  int kx,ky;
  double out=0.;
  for(kx=-K1;kx<=K1;kx++){
    for (ky=-K2;ky<=K2;ky++){
      out+=__real__ (getval_sym(u, kx,ky,K2)*conj(getval_sym(u, kx,ky,K2)));
    }
  }
  return out;
}

// compute enstrophy
double compute_enstrophy(
  _Complex double* u,
  int K1,
  int K2,
  double L
){
  int kx,ky;
  double out=0.;
  for(kx=-K1;kx<=K1;kx++){
    for (ky=-K2;ky<=K2;ky++){
      out+=__real__ (4*M_PI*M_PI/L/L*(kx*kx+ky*ky)*getval_sym(u, kx,ky,K2)*conj(getval_sym(u, kx,ky,K2)));
    }
  }
  return out;
}


// 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);
}

// get index for kx,ky in array of size K1,K2 in which only the terms with kx>=0 are stored
int klookup_sym(
  int kx,
  int ky,
  int K2
){
  if (kx<0) {
    fprintf(stderr, "bug!: attempting to access a symmetrized vector at kx<0\n Contact Ian at ian.jauslin@rutgers.edu!\n");
    exit(-1);
  }
  return kx*(2*K2+1) + (ky>=0 ? ky : (2*K2+1)+ky);
}

// get u_{kx,ky} from a vector u in which only the values for kx>=0 are stored, assuming u_{-k}=u_k^*
_Complex double getval_sym(
  _Complex double* u,
  int kx,
  int ky,
  int K2
){
  return (kx>=0 ? u[klookup_sym(kx,ky,K2)] : conj(u[klookup_sym(-kx,-ky,K2)]));
}