/*
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.
*/

#ifndef NAVIERSTOKES_H
#define NAVIERSTOKES_H

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

#define M_PI 3.14159265358979323846

// abort signal
extern volatile bool g_abort;

// 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, uint64_t nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, bool irreversible, unsigned int algorithm, uint64_t print_freq, uint64_t starting_time, unsigned int nthreadsl, FILE* savefile);

// compute enstrophy and alpha
int enstrophy( int K1, int K2, int N1, int N2, uint64_t nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, bool irreversible, unsigned int algorithm, uint64_t print_freq, uint64_t starting_time, unsigned int nthreads, FILE* savefile, char* cmd_string, char* params_string, char* savefile_string);

// 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, uint64_t nsteps, double nu, double delta, double L, uint64_t starting_time, _Complex double* u0, _Complex double* g, bool irreversible, unsigned int algorithm, unsigned int nthreads, FILE* savefile);


// 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/reversible Navier-Stokes equation
// RK4
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);
// RK2
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);

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

// 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
double compute_alpha( _Complex double* u, int K1, int K2, _Complex double* g, double L);
// compute energy
double compute_energy( _Complex double* u, int K1, int K2);
// compute enstrophy
double compute_enstrophy( _Complex double* u, int K1, int K2, double L);

// get index for kx,ky in array of size S
int klookup( int kx, int ky, int S1, int S2);
// 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);
// 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);

#endif