/*
Copyright 2017-2023 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 <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,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  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);

  // iterate
  time=starting_time;
  while(final_time<0 || time<final_time){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else if (algorithm==ALGORITHM_RK4) {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    } else if (algorithm==ALGORITHM_RKF45) {
      delta=ns_step_rkf45(u, adaptive_tolerance, adaptive_factor, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible);
    } else {
      fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
    }

    time+=delta;
    
    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");
    }
  }

  // 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,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  unsigned int algorithm,
  double print_freq,
  double 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;
  _Complex double* tmp4;
  _Complex double* tmp5;
  _Complex double* tmp6;
  _Complex double* tmp7;
  double time;
  double alpha, enstrophy;
  double prevtime;
  double avg_a,avg_en,avg_en_x_a;
  // 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);


  // init running average
  avg_a=0;
  avg_en=0;
  avg_en_x_a=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){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else if (algorithm==ALGORITHM_RK4) {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    } else if (algorithm==ALGORITHM_RKF45) {
      delta=ns_step_rkf45(u, adaptive_tolerance, adaptive_factor, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible);
    } else {
      fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
    }

    time+=delta;

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

    // add to running averages (estimate the total duration of interval as print_freq, will be adjusted later)
    avg_a+=alpha*(delta/print_freq);
    avg_en+=enstrophy*(delta/print_freq);
    avg_en_x_a+=enstrophy*alpha*(delta/print_freq);

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

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

      // print to stderr so user can follow along
      if(algorithm==ALGORITHM_RKF45){
	fprintf(stderr,"% .8e % .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",time, avg_a, avg_en, avg_en_x_a, alpha, enstrophy, alpha*enstrophy, delta);
	printf("% .15e % .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",time, avg_a, avg_en_x_a, avg_en, alpha, alpha*enstrophy, enstrophy, delta);
      } else {
	fprintf(stderr,"% .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",time, avg_a, avg_en, avg_en_x_a, alpha, enstrophy, alpha*enstrophy);
	printf("% .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",time, avg_a, avg_en_x_a, avg_en, alpha, alpha*enstrophy, enstrophy);
      }
      // print to stdout
    }

    // reset averages
    avg_a=0;
    avg_en=0;
    avg_en_x_a=0;
    prevtime=time;

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

  if(savefile!=NULL){
    fprintf(savefile,"# Continue computation with\n");

    // command to resume
    fprintf(savefile,"#! ");
    fprintf(savefile, cmd_string);
    // params
    // allocate buffer for params
    if(params_string!=NULL) {
      char* params=calloc(sizeof(char), strlen(params_string)+1);
      strcpy(params, params_string);
      remove_entry(params, "starting_time");
      remove_entry(params, "init");
      if(algorithm>ALGORITHM_ADAPTIVE_THRESHOLD){
	remove_entry(params, "delta");
      }
      fprintf(savefile," -p \"%s;init=file:%s", params, savefile_string);
      free(params);
      // write delta if adaptive, and not writing binary
      if(algorithm>ALGORITHM_ADAPTIVE_THRESHOLD && (savefile==stderr || savefile==stdout)){
	fprintf(savefile,";delta=%.15e", delta);
      }
      // write starting_time if not writing binary
      if(savefile==stderr || savefile==stdout){
	fprintf(savefile,";starting_time=%.15e", time);
      }
      fprintf(savefile,"\"");
    }
    fprintf(savefile," enstrophy\n");

    // save final u to savefile
    if(savefile==stderr || savefile==stdout){
      write_vec(u, K1, K2, savefile);
    } else {
      write_vec_bin(u, K1, K2, savefile);
      // last binary entry: starting time
      fwrite(&time, sizeof(double), 1, savefile);
      // extra binary data for adaptive algorithm
      if(algorithm>ALGORITHM_ADAPTIVE_THRESHOLD){
	fwrite(&delta, sizeof(double), 1, savefile);
      }
    }
  }

  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 starting_time,
  _Complex double* u0,
  _Complex double* g,
  bool irreversible,
  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);

  // iterate
  time=starting_time;
  while(final_time<0 || time<final_time){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else if (algorithm==ALGORITHM_RK4) {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    } else if (algorithm==ALGORITHM_RKF45) {
      delta=ns_step_rkf45(u, adaptive_tolerance, adaptive_factor, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, irreversible);
    } else {
      fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
    }

    time+=delta;
  }

  // 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(sizeof(_Complex double),K1*(2*K2+1)+K2);

  // allocate tmp vectors for computation
  if(algorithm==ALGORITHM_RK2){
    *tmp1=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp2=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
  } else if (algorithm==ALGORITHM_RK4){
    *tmp1=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp2=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp3=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
  } else if (algorithm==ALGORITHM_RKF45){
    *tmp1=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp2=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp3=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp4=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp5=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp6=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
    *tmp7=calloc(sizeof(_Complex double),K1*(2*K2+1)+K2);
  } 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){
    free(tmp1);
    free(tmp2);
    free(tmp3);
    free(tmp4);
    free(tmp5);
    free(tmp6);
    free(tmp7);
  } 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
// 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
){
  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)];
    }
  }

  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
){
  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)];
    }
  }

  return(0);
}

// next time step
// adaptive RK algorithm (Runge-Kutta-Fehlberg)
double ns_step_rkf45(
  _Complex double* u,
  double tolerance,
  double factor,
  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* k1,
  _Complex double* k2,
  _Complex double* k3,
  _Complex double* k4,
  _Complex double* k5,
  _Complex double* k6,
  _Complex double* tmp,
  bool irreversible
){
  int kx,ky;
  double err;

  // k1: u(t)
  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);
  
  // compute error
  err=0;
  for(kx=0;kx<=K1;kx++){
    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
      // difference between 5th order and 4th order
      err+=cabs(delta*(1./360*k1[klookup_sym(kx,ky,K2)]-128./4275*k3[klookup_sym(kx,ky,K2)]-2197./75240*k4[klookup_sym(kx,ky,K2)]+1./50*k5[klookup_sym(kx,ky,K2)]+2./55*k6[klookup_sym(kx,ky,K2)]));
    }
  }

  // new delta
  delta=factor*delta*pow(tolerance/err,0.2);
  // compare relative error with tolerance
  if(err<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)]+=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)]);
      }
    }
    return delta;
  }
  // error too big: repeat with smaller step
  else{
    return(ns_step_rkf45(u,tolerance,factor,K1,K2,N1,N2,nu,delta,L,g,fft1,fft2,ifft,k1,k2,k3,k4,k5,k6,tmp,irreversible));
  }
}

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

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