#ifndef NAVIERSTOKES_H
#define NAVIERSTOKES_H

#include <complex.h>
#include <fftw3.h>

#define M_PI 3.14159265358979323846

// arrays on which the ffts are performed
typedef struct fft_vects {
  fftw_complex* fft;
  fftw_plan fft_plan;
} fft_vect;

// compute u_k
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, unsigned int print_freq, unsigned int nthreads);

// compute the energy as a function of time
int energy( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int print_freq, unsigned int nthreads);

// compute alpha
int alpha( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int print_freq, unsigned int nthreads);

// 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, unsigned int nthreads);


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

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

// 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, double L, _Complex double* g, fft_vect fft1, fft_vect fft2,fft_vect ifft, _Complex double* tmp1, _Complex double *tmp2, _Complex double *tmp3);

// 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, double L, _Complex double* g, fft_vect fft1, fft_vect fft2, fft_vect ifft);

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

// convolution term in right side of equation (computed without fft)
int ns_T_nofft( _Complex double* out, _Complex double* u, int K1, int K2, int N1, int N2);

// compute alpha
_Complex double compute_alpha( _Complex double* u, int K1, int K2, _Complex double* g);

// get index for kx,ky in array of size S
int klookup( int kx, int ky, int S1, int S2);

#endif