/*
Copyright 2017-2025 Ian Jauslin

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include "constants.cpp"
#include "io.h"
#include "navier-stokes.h"
#include "complex_tools.h"
#include <math.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

// compute solution as a function of time
int uk(
  int K1,
  int K2,
  int N1,
  int N2,
  double final_time,
  double nu,
  double delta,
  double L,
  double adaptive_tolerance,
  double adaptive_factor,
  double max_delta,
  unsigned int adaptive_cost,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  unsigned int algorithm,
  double print_freq,
  double starting_time,
  unsigned int nthreads,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  _Complex double* tmp4;
  _Complex double* tmp5;
  _Complex double* tmp6;
  _Complex double* tmp7;
  double time;
  fft_vect fft1;
  fft_vect fft2;
  fft_vect ifft;
  int kx,ky;

  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &tmp4, &tmp5, &tmp6, &tmp7, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads, algorithm);
  // copy initial condition
  copy_u(u, u0, K1, K2);

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

  // period
  // add 0.1 to ensure proper rounding
  uint64_t n=(uint64_t)((starting_time-fmod(starting_time, print_freq))/print_freq+0.1);

  // store step (useful for adaptive step methods)
  double step=delta;
  double next_step=step;

  // iterate
  time=starting_time;
  while(final_time<0 || time<final_time){
    // step
    ns_step(algorithm, u, K1, K2, N1, N2, nu, &step, &next_step, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, L, g, time, starting_time, fft1, fft2, ifft, &tmp1, &tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible, keep_en_cst, target_en);

    time+=step;
    step=next_step;
    
    if(time>(n+1)*print_freq){
      n++;

      fprintf(stderr,"% .8e ",time);
      printf("% .15e ",time);

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

  // TODO: update handling of savefile
  // save final entry to savefile
  write_vec_bin(u, K1, K2, savefile);

  ns_free_tmps(u, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, fft1, fft2, ifft, algorithm);
  return(0);
}

// compute enstrophy, alpha as a function of time
int enstrophy(
  int K1,
  int K2,
  int N1,
  int N2,
  double final_time,
  double nu,
  double delta,
  double L,
  double adaptive_tolerance,
  double adaptive_factor,
  double max_delta,
  unsigned int adaptive_cost,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  unsigned int algorithm,
  double print_freq,
  double starting_time,
  bool print_alpha,
  unsigned int nthreads,
  FILE* savefile,
  FILE* utfile,
  // for interrupt recovery
  const char* cmd_string,
  const char* params_string,
  const char* savefile_string,
  const char* utfile_string
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  _Complex double* tmp4;
  _Complex double* tmp5;
  _Complex double* tmp6;
  _Complex double* tmp7;
  double time;
  double alpha, enstrophy, energy;
  double prevtime;
  double avg_a,avg_enstrophy,avg_energy;
  // index
  fft_vect fft1;
  fft_vect fft2;
  fft_vect ifft;

  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &tmp4, &tmp5, &tmp6, &tmp7, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads, algorithm);
  // copy initial condition
  copy_u(u, u0, K1, K2);

  // store step (useful for adaptive step methods
  double step=delta;
  double next_step=step;

  // init running average
  avg_a=0;
  avg_enstrophy=0;
  avg_energy=0;
  prevtime=starting_time;

  // period
  // add 0.1 to ensure proper rounding
  uint64_t n=(uint64_t)((starting_time-fmod(starting_time, print_freq))/print_freq+0.1);

  // iterate
  time=starting_time;
  while(final_time<0 || time<final_time){
    ns_step(algorithm, u, K1, K2, N1, N2, nu, &step, &next_step, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, L, g, time, starting_time, fft1, fft2, ifft, &tmp1, &tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible, keep_en_cst, target_en);

    time+=step;

    alpha=compute_alpha(u, K1, K2, g, L);
    enstrophy=compute_enstrophy(u, K1, K2, L);
    energy=compute_energy(u, K1, K2);

    // check that anything is nan
    if(isnan(alpha) || isnan(enstrophy)){
      fprintf(stderr,"nan encountered, aborting.\n");
      break;
    }

    // add to running averages (estimate the total duration of interval as print_freq, will be adjusted later)
    avg_a+=alpha*(step/print_freq);
    avg_enstrophy+=enstrophy*(step/print_freq);
    avg_energy+=energy*(step/print_freq);

    if(time>(n+1)*print_freq){
      n++;

      // adjust duration of interval to its actual value
      avg_a*=print_freq/(time-prevtime);
      avg_enstrophy*=print_freq/(time-prevtime);
      avg_energy*=print_freq/(time-prevtime);

      // print to stderr so user can follow along
      if(algorithm>ALGORITHM_ADAPTIVE_THRESHOLD){
	fprintf(stderr,"% .8e % .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",time, avg_a, avg_energy, avg_enstrophy, alpha, energy, enstrophy, step);
	printf("% .15e % .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",time, avg_a, avg_energy, avg_enstrophy, alpha, energy, enstrophy, step);
      } else {
	fprintf(stderr,"% .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",time, avg_a, avg_energy, avg_enstrophy, alpha, energy, enstrophy);
	printf("% .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",time, avg_a, avg_energy, avg_enstrophy, alpha, energy, enstrophy);
      }

      // reset averages
      avg_a=0;
      avg_enstrophy=0;
      avg_energy=0;
      prevtime=time;
    }

    // print alpha when it gets negative
    if(print_alpha && alpha<0){
      fprintf(stderr,"## negative alpha: % .8e at time % .8e\n",alpha, time);
      printf("## negative alpha: % .15e at time % .15e\n",alpha, time);
    }

    // get ready for next step
    step=next_step;

    // catch abort signal
    if (g_abort){
      // print u to stderr if no savefile
      if (savefile==NULL){
	savefile=stderr;
      }
      break;
    }
  }

  if(savefile!=NULL){
    save_state(u, savefile, K1, K2, cmd_string, params_string, savefile_string, utfile_string, utfile, COMMAND_ENSTROPHY, algorithm, step, time, nthreads);
  }

  // save final u to utfile in txt format
  if(utfile!=NULL){
    write_vec(u, K1, K2, utfile);
  }

  ns_free_tmps(u, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, fft1, fft2, ifft, algorithm);
  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,
  double final_time,
  double nu,
  double delta,
  double L,
  double adaptive_tolerance,
  double adaptive_factor,
  double max_delta,
  unsigned int adaptive_cost,
  double starting_time,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  unsigned int algorithm,
  unsigned int nthreads,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  _Complex double* tmp4;
  _Complex double* tmp5;
  _Complex double* tmp6;
  _Complex double* tmp7;
  double time;
  fft_vect fft1;
  fft_vect fft2;
  fft_vect ifft;

  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &tmp4, &tmp5, &tmp6, &tmp7, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads, algorithm);
  // copy initial condition
  copy_u(u, u0, K1, K2);

  // store delta (useful for adaptive step algorithms)
  double step=delta;
  double next_step=step;


  // iterate
  time=starting_time;
  while(final_time<0 || time<final_time){
    ns_step(algorithm, u, K1, K2, N1, N2, nu, &step, &next_step, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, L, g, time, starting_time, fft1, fft2, ifft, &tmp1, &tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible, keep_en_cst, target_en);

    time+=step;
    step=next_step;
  }

  // TODO: update handling of savefile
  // save final entry to savefile
  write_vec(u, K1, K2, savefile);

  ns_free_tmps(u, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, fft1, fft2, ifft, algorithm);
  return(0);
}


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

  // allocate tmp vectors for computation
  if(algorithm==ALGORITHM_RK2){
    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
  } else if (algorithm==ALGORITHM_RK4){
    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
  } else if (algorithm==ALGORITHM_RKF45 || algorithm==ALGORITHM_RKDP54){
    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp4=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp5=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp6=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp7=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
  } else if (algorithm==ALGORITHM_RKBS32){
    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp4=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
    *tmp5=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
  } else {
    fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
  };

  // 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,
  _Complex double* tmp4,
  _Complex double* tmp5,
  _Complex double* tmp6,
  _Complex double* tmp7,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  unsigned int algorithm
){
  // 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();

  if(algorithm==ALGORITHM_RK2){
    free(tmp1);
    free(tmp2);
  } else if (algorithm==ALGORITHM_RK4){
    free(tmp1);
    free(tmp2);
    free(tmp3);
  } else if (algorithm==ALGORITHM_RKF45 || algorithm==ALGORITHM_RKDP54){
    free(tmp1);
    free(tmp2);
    free(tmp3);
    free(tmp4);
    free(tmp5);
    free(tmp6);
    free(tmp7);
  } else if (algorithm==ALGORITHM_RKBS32){
    free(tmp1);
    free(tmp2);
    free(tmp3);
    free(tmp4);
    free(tmp5);
  } else {
    fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
  };

  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*(2*K2+1)+K2;i++){
    u[i]=u0[i];
  }

  return 0;
}

// next time step
int ns_step(
  unsigned int algorithm,
  _Complex double* u,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double* delta,
  double* next_delta,
  double adaptive_tolerance,
  double adaptive_factor,
  double max_delta,
  unsigned int adaptive_cost,
  double L,
  _Complex double* g,
  double time,
  double starting_time,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  // the pointers tmp1 and tmp2 will be exchanged at the end of the routine for first-same-as-last RK algorithms
  _Complex double** tmp1,
  _Complex double** tmp2,
  _Complex double* tmp3,
  _Complex double* tmp4,
  _Complex double* tmp5,
  _Complex double* tmp6,
  _Complex double* tmp7,
  bool irreversible,
  bool keep_en_cst,
  double target_en
){
  if(algorithm==ALGORITHM_RK2){
    ns_step_rk2(u, K1, K2, N1, N2, nu, *delta, L, g, fft1, fft2, ifft, *tmp1, *tmp2, irreversible, keep_en_cst, target_en);
  } else if (algorithm==ALGORITHM_RK4) {
    ns_step_rk4(u, K1, K2, N1, N2, nu, *delta, L, g, fft1, fft2, ifft, *tmp1, *tmp2, tmp3, irreversible, keep_en_cst, target_en);
  } else if (algorithm==ALGORITHM_RKF45) {
    ns_step_rkf45(u, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, K1, K2, N1, N2, nu, delta, next_delta, L, g, fft1, fft2, ifft, *tmp1, *tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible, keep_en_cst, target_en, true);
  } else if (algorithm==ALGORITHM_RKDP54) {
    // only compute k1 on the first step
    // first-same-as-last with 2-nd argument
    ns_step_rkdp54(u, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, K1, K2, N1, N2, nu, delta, next_delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible, keep_en_cst, target_en, time==starting_time);
  } else if (algorithm==ALGORITHM_RKBS32) {
    // only compute k1 on the first step
    // first-same-as-last with 4-th argument
    ns_step_rkbs32(u, adaptive_tolerance, adaptive_factor, max_delta, adaptive_cost, K1, K2, N1, N2, nu, delta, next_delta, L, g, fft1, fft2, ifft, tmp1, tmp3, tmp4, tmp2, tmp5, irreversible, keep_en_cst, target_en, time==starting_time);
  } else {
    fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers.edu\n",algorithm);
  }

  return (0);
}

// TODO: do not need to use klookup in any of the rk routines
// RK 4 algorithm
int ns_step_rk4(
  _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,
  bool keep_en_cst,
  double target_en
){
  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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      u[klookup_sym(kx,ky,K2)]=tmp3[klookup_sym(kx,ky,K2)]+delta/6*tmp1[klookup_sym(kx,ky,K2)];
    }
  }

  // keep enstrophy constant
  if(keep_en_cst){
    double en=compute_enstrophy(u, K1, K2, L);
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
      }
    }
  }

  return(0);
}

// RK 2 algorithm
int ns_step_rk2(
  _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,
  bool irreversible,
  bool keep_en_cst,
  double target_en
){
  int kx,ky;

  // k1
  ns_rhs(tmp1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // u+h*k1/2
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);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=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      u[klookup_sym(kx,ky,K2)]+=delta*tmp1[klookup_sym(kx,ky,K2)];
    }
  }

  // keep enstrophy constant
  if(keep_en_cst){
    double en=compute_enstrophy(u, K1, K2, L);
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
      }
    }
  }

  return(0);
}

// next time step
// adaptive RK algorithm (Runge-Kutta-Fehlberg)
int ns_step_rkf45(
  _Complex double* u,
  double tolerance,
  double factor,
  double max_delta,
  unsigned int adaptive_cost,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double* delta,
  double* next_delta,
  double L,
  _Complex double* g,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  _Complex double* k1,
  _Complex double* k2,
  _Complex double* k3,
  _Complex double* k4,
  _Complex double* k5,
  _Complex double* k6,
  _Complex double* tmp,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  // whether to compute k1 or leave it as is
  bool compute_k1
){
  int kx,ky;
  double cost;

  // k1: u(t)
  if(compute_k1){
    ns_rhs(k1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
  }

  // k2 : u(t+1/4*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/4*k1[klookup_sym(kx,ky,K2)];
    }
  }
  ns_rhs(k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k3 : u(t+3/8*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./32*k1[klookup_sym(kx,ky,K2)]+9./32*k2[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k4 : u(t+12./13*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(1932./2197*k1[klookup_sym(kx,ky,K2)]-7200./2197*k2[klookup_sym(kx,ky,K2)]+7296./2197*k3[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k5 : u(t+1*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(439./216*k1[klookup_sym(kx,ky,K2)]-8*k2[klookup_sym(kx,ky,K2)]+3680./513*k3[klookup_sym(kx,ky,K2)]-845./4104*k4[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k5, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k6 : u(t+1./2*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(-8./27*k1[klookup_sym(kx,ky,K2)]+2*k2[klookup_sym(kx,ky,K2)]-3544./2565*k3[klookup_sym(kx,ky,K2)]+1859./4104*k4[klookup_sym(kx,ky,K2)]-11./40*k5[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k6, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
  
  // next step
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      // u
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(25./216*k1[klookup_sym(kx,ky,K2)]+1408./2565*k3[klookup_sym(kx,ky,K2)]+2197./4104*k4[klookup_sym(kx,ky,K2)]-1./5*k5[klookup_sym(kx,ky,K2)]);
      // U: save to k6, which is no longer needed
      k6[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(16./135*k1[klookup_sym(kx,ky,K2)]+6656./12825*k3[klookup_sym(kx,ky,K2)]+28561./56430*k4[klookup_sym(kx,ky,K2)]-9./50*k5[klookup_sym(kx,ky,K2)]+2./55*k6[klookup_sym(kx,ky,K2)]);
    }
  }

  // cost function
  cost=ns_adaptive_cost(tmp,k6,adaptive_cost,K1,K2,g,L);

  // compare relative error with tolerance
  if(cost<tolerance){
    // copy to output
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
      }
    }
    // next delta to use in future steps (do not exceed max_delta)
    *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,0.2));

    // keep enstrophy constant
    if(keep_en_cst){
      double en=compute_enstrophy(u, K1, K2, L);
      for(kx=0;kx<=K1;kx++){
	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
	}
      }
    }
  }
  // error too big: repeat with smaller step
  else{
    *delta=factor*(*delta)*pow(tolerance/cost,0.2);
    // do not recompute k1
    ns_step_rkf45(u,tolerance,factor,max_delta,adaptive_cost,K1,K2,N1,N2,nu,delta,next_delta,L,g,fft1,fft2,ifft,k1,k2,k3,k4,k5,k6,tmp,irreversible,keep_en_cst,target_en,false);
  }

  return 0;
}


// next time step
// adaptive RK algorithm (Runge-Kutta-Bogacki-Shampine method)
int ns_step_rkbs32(
  _Complex double* u,
  double tolerance,
  double factor,
  double max_delta,
  unsigned int adaptive_cost,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double* delta,
  double* next_delta,
  double L,
  _Complex double* g,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  // the pointers k1 and k4 will be exchanged at the end of the routine
  _Complex double** k1,
  _Complex double* k2,
  _Complex double* k3,
  _Complex double** k4,
  _Complex double* tmp,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  // whether to compute k1
  bool compute_k1
){
  int kx,ky;
  double cost;

  // k1: u(t)
  // only compute it if it is the first step (otherwise, it has already been computed due to the First Same As Last property)
  if(compute_k1){
    ns_rhs(*k1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
  }

  // k2 : u(t+1/4*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/2*(*k1)[klookup_sym(kx,ky,K2)];
    }
  }
  ns_rhs(k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k3 : u(t+3/4*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./4*k2[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k4 : u(t+delta)
  // tmp computed here is the next step
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(2./9*(*k1)[klookup_sym(kx,ky,K2)]+1./3*k2[klookup_sym(kx,ky,K2)]+4./9*k3[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(*k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // next step
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      // U: store in k3, which is no longer needed
      k3[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(7./24*(*k1)[klookup_sym(kx,ky,K2)]+1./4*k2[klookup_sym(kx,ky,K2)]+1./3*k3[klookup_sym(kx,ky,K2)]+1./8*(*k4)[klookup_sym(kx,ky,K2)]);
    }
  }
  
  // compute cost
  cost=ns_adaptive_cost(tmp, k3, adaptive_cost, K1, K2, g, L);

  // compare cost with tolerance
  if(cost<tolerance){
    // add to output
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
      }
    }
    // next delta to use in future steps (do not exceed max_delta)
    *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,1./3));

    // k1 in the next step will be this k4 (first same as last)
    tmp=*k1;
    *k1=*k4;
    *k4=tmp;

    // keep enstrophy constant
    if(keep_en_cst){
      double en=compute_enstrophy(u, K1, K2, L);
      for(kx=0;kx<=K1;kx++){
	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
	}
      }
    }
  }
  // error too big: repeat with smaller step
  else{
    *delta=factor*(*delta)*pow(tolerance/cost,1./3);
    // this will reuse the same k1 without re-computing it
    ns_step_rkbs32(u,tolerance,factor,max_delta,adaptive_cost,K1,K2,N1,N2,nu,delta,next_delta,L,g,fft1,fft2,ifft,k1,k2,k3,k4,tmp,irreversible,keep_en_cst,target_en,false);
  }

  return 0;
}

// next time step
// adaptive RK algorithm (Runge-Kutta-Dormand-Prince method)
int ns_step_rkdp54(
  _Complex double* u,
  double tolerance,
  double factor,
  double max_delta,
  unsigned int adaptive_cost,
  int K1,
  int K2,
  int N1,
  int N2,
  double nu,
  double* delta,
  double* next_delta,
  double L,
  _Complex double* g,
  fft_vect fft1,
  fft_vect fft2,
  fft_vect ifft,
  // the pointers k1 and k2 will be exchanged at the end of the routine
  _Complex double** k1,
  _Complex double** k2,
  _Complex double* k3,
  _Complex double* k4,
  _Complex double* k5,
  _Complex double* k6,
  _Complex double* tmp,
  bool irreversible,
  bool keep_en_cst,
  double target_en,
  // whether to compute k1
  bool compute_k1
){
  int kx,ky;
  double cost;

  // k1: u(t)
  // only compute it if it is the first step (otherwise, it has already been computed due to the First Same As Last property)
  if(compute_k1){
    ns_rhs(*k1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
  }

  // k2 : u(t+1/5*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/5*(*k1)[klookup_sym(kx,ky,K2)];
    }
  }
  ns_rhs(*k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k3 : u(t+3/10*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./40*(*k1)[klookup_sym(kx,ky,K2)]+9./40*(*k2)[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k4 : u(t+4/5*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(44./45*(*k1)[klookup_sym(kx,ky,K2)]-56./15*(*k2)[klookup_sym(kx,ky,K2)]+32./9*k3[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k5 : u(t+8/9*delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(19372./6561*(*k1)[klookup_sym(kx,ky,K2)]-25360./2187*(*k2)[klookup_sym(kx,ky,K2)]+64448./6561*k3[klookup_sym(kx,ky,K2)]-212./729*k4[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k5, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k6 : u(t+delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(9017./3168*(*k1)[klookup_sym(kx,ky,K2)]-355./33*(*k2)[klookup_sym(kx,ky,K2)]+46732./5247*k3[klookup_sym(kx,ky,K2)]+49./176*k4[klookup_sym(kx,ky,K2)]-5103./18656*k5[klookup_sym(kx,ky,K2)]);
    }
  }
  ns_rhs(k6, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  // k7 : u(t+delta)
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      // tmp computed here is the step
      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(35./384*(*k1)[klookup_sym(kx,ky,K2)]+500./1113*k3[klookup_sym(kx,ky,K2)]+125./192*k4[klookup_sym(kx,ky,K2)]-2187./6784*k5[klookup_sym(kx,ky,K2)]+11./84*k6[klookup_sym(kx,ky,K2)]);
    }
  }
  // store in k2, which is not needed anymore
  ns_rhs(*k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);

  //next step
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      // U: store in k6, which is not needed anymore
      k6[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(5179./57600*(*k1)[klookup_sym(kx,ky,K2)]+7571./16695*k3[klookup_sym(kx,ky,K2)]+393./640*k4[klookup_sym(kx,ky,K2)]-92097./339200*k5[klookup_sym(kx,ky,K2)]+187./2100*k6[klookup_sym(kx,ky,K2)]+1./40*(*k2)[klookup_sym(kx,ky,K2)]);
    }
  }
  
  // compute cost
  cost=ns_adaptive_cost(tmp, k6, adaptive_cost, K1, K2, g, L);

  // compare relative error with tolerance
  if(cost<tolerance){
    // add to output
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
      }
    }
    // next delta to use in future steps (do not exceed max_delta)
    *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,0.2));

    // k1 in the next step will be this k7 (first same as last)
    tmp=*k1;
    *k1=*k2;
    *k2=tmp;

    // keep enstrophy constant
    if(keep_en_cst){
      double en=compute_enstrophy(u, K1, K2, L);
      for(kx=0;kx<=K1;kx++){
	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
	}
      }
    }
  }
  // error too big: repeat with smaller step
  else{
    *delta=factor*(*delta)*pow(tolerance/cost,0.2);
    // this will reuse the same k1 without re-computing it
    ns_step_rkdp54(u,tolerance,factor,max_delta,adaptive_cost,K1,K2,N1,N2,nu,delta,next_delta,L,g,fft1,fft2,ifft,k1,k2,k3,k4,k5,k6,tmp,irreversible,keep_en_cst,target_en,false);
  }

  return 0;
}


// compute error for adaptive step methods
double ns_adaptive_cost(
  _Complex double* u,
  _Complex double* U,
  unsigned int adaptive_cost,
  int K1,
  int K2,
  _Complex double* g,
  double L
){
  int kx,ky;

  if(adaptive_cost==COST_L1){
    double cost=0;
    double relative=0;
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	cost+=cabs(u[klookup_sym(kx,ky,K2)]-U[klookup_sym(kx,ky,K2)]);
	relative+=cabs(u[klookup_sym(kx,ky,K2)]);
      }
    }
    return cost/relative;
  }
  else if(adaptive_cost==COST_k3){
    double cost=0;
    double relative=0;
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	cost+=cabs(u[klookup_sym(kx,ky,K2)]-U[klookup_sym(kx,ky,K2)])/pow(kx*kx+ky*ky,3);
	relative+=cabs(u[klookup_sym(kx,ky,K2)])/pow(kx*kx+ky*ky,3);
      }
    }
    return cost/relative;
  }
  else if(adaptive_cost==COST_k32){
    double cost=0;
    double relative=0;
    for(kx=0;kx<=K1;kx++){
      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
	cost+=cabs(u[klookup_sym(kx,ky,K2)]-U[klookup_sym(kx,ky,K2)])/pow(kx*kx+ky*ky,1.5);
	relative+=cabs(u[klookup_sym(kx,ky,K2)])/pow(kx*kx+ky*ky,1.5);
      }
    }
    return cost/relative;
  }
  else if(adaptive_cost==COST_ENSTROPHY){
    double enu=compute_enstrophy(u,K1,K2,L);
    return fabs(enu-compute_enstrophy(U,K1,K2,L))/enu;
  }
  else if(adaptive_cost==COST_ALPHA){
    double alu=compute_alpha(u,K1,K2,g,L);
    return fabs((alu-compute_alpha(U,K1,K2,g,L))/alu);
  }
  else{
    fprintf(stderr,"bug: unknown norm: %u, contact ian.jauslin@rutgers.edu\n", adaptive_cost);
    exit(-1);
  }

  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,g,L);
  }

  for(i=0; i<K1*(2*K2+1)+K2; i++){
    out[i]=0;
  }
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      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);

  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>=2*K1+1 and N2>=2*K2+1)
  if (N1<2*K1+1 || N2<2*K2+1){
    fprintf(stderr,"error: N1 and N2 need t be >= 2*K1+1 and 2*K2+1 respectively\n");
    return(-1);
  }
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=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);
	  }
	}
      }
    }
  }

  return 0;
}


/*
// Jacobian of rhs
int ns_D_rhs(
  _Complex double* out,
  _Complex double* u,
  int K1,
  int K2,
  double nu,
  double L,
  bool irreversible
){
  int kx,ky,lx,ly;
  int i;
  double alpha;

  if (irreversible) {
    alpha=nu;
  } else {
    // TODO
    alpha=0;
  }

  for(i=0; i<MATSIZE*MATSIZE; i++){
    out[i]=0;
  }
  // diagonal term
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      out[klookup_sym(kx,ky,K2)*MATSIZE+klookup_sym(kx,ky,K2)]=-4*M_PI*M_PI/L/L*alpha*(kx*kx+ky*ky);
    }
  }
  // convolution term
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      for(lx=0;lx<=K1;lx++){
	for(ly=(lx>0 ? -K2 : 1);ly<=K2;ly++){
	  // no contribution from k=l
	  if (kx!=lx && ky!=ly){
	    // use row-major ordering
	    out[klookup_sym(kx,ky,K2)*MATSIZE+klookup_sym(lx,ly,K2)]=4*M_PI*M_PI/L/L*ns_d_T(u,kx,ky,lx,ly,K2);
	  }
	}
      }
    }
  }

  return(0);
}

// derivative of T with respect to u_k
_Complex double ns_d_T(
  _Complex double* u,
  int kx,
  int ky,
  int lx,
  int ly,
  int K2
){
  int qx=kx-lx;
  int qy=ky-ly;
  // trivial case
  if (qx*qx+qy*qy==0){
    return 0.;
  }
  if (qx>0 || (qx==0 && qy>=1)){
    return (ky*lx-kx*ly)*(qx*qx+qy*qy-lx*lx-ly*ly)/sqrt(lx*lx+ly*ly)/sqrt(qx*qx+qy*qy)*u[klookup_sym(qx,qy,K2)];
  }else{
    return (ky*lx-kx*ly)*(qx*qx+qy*qy-lx*lx-ly*ly)/sqrt(lx*lx+ly*ly)/sqrt(qx*qx+qy*qy)*u[klookup_sym(-qx,-qy,K2)];
  }
}
*/


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

  num=0.;
  denom=0.;

  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      num+=L*L/4/M_PI/M_PI*(kx*kx+ky*ky)*getval_sym(g, kx,ky,K2)*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 energy
double compute_energy(
  _Complex double* u,
  int K1,
  int K2
){
  int kx,ky;
  double out=0.;
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      out+=__real__ (getval_sym(u, kx,ky,K2)*conj(getval_sym(u, kx,ky,K2)));
    }
  }
  return 2*out;
}

// compute enstrophy
double compute_enstrophy(
  _Complex double* u,
  int K1,
  int K2,
  double L
){
  int kx,ky;
  double out=0.;
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);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 2*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 and if kx=0, ky>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);
  }
  if (kx==0) {
    if (ky<=0){
      fprintf(stderr, "bug!: attempting to access a symmetrized vector at kx=0 and ky<=0\n Contact Ian at ian.jauslin@rutgers.edu!\n");
      exit(-1);
    }
    return ky-1;
  }
  return K2+(kx-1)*(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
){
  if(kx>0 || (kx==0 && ky>0)){
    return u[klookup_sym(kx,ky,K2)];
  }
  else if(kx==0 && ky==0){
    return 0;
  } else {
    return conj(u[klookup_sym(-kx,-ky,K2)]);
  }
}