/*
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 "statistics.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,
  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
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  uint64_t 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("   %6lu:(%4d,%4d)r ",t,kx,ky);
      t++;
      printf("   %6lu:(%4d,%4d)i ",t,kx,ky);
      t++;
    }
  }

  // iterate
  for(t=starting_time;nsteps==0 || t<starting_time+nsteps;t++){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    }
    
    if(t%print_freq==0){
      fprintf(stderr,"%lu % .8e ",t,t*delta);
      printf("%8lu % .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_vec_bin(u, K1, K2, savefile);

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

// compute enstrophy, alpha as a function of time
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,
  // 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, enstrophy;
  double avg_a,avg_en,avg_en_x_a;
  // index
  uint64_t 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_a=0;
  avg_en=0;
  avg_en_x_a=0;

  // special first case when starting_time is not a multiple of print_freq
  uint64_t first_box = print_freq - (starting_time % print_freq);

  // iterate
  for(t=starting_time;nsteps==0 || t<starting_time+nsteps;t++){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    }

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

    avg_a=average_step(alpha, avg_a, t, starting_time, print_freq, first_box);
    avg_en=average_step(enstrophy, avg_en, t, starting_time, print_freq, first_box);
    avg_en_x_a=average_step(enstrophy*alpha, avg_en_x_a, t, starting_time, print_freq, first_box);

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

    // 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");
      remove_entry(params, "nsteps");
      fprintf(savefile," -p \"%s;starting_time=%lu;nsteps=%lu;init=file:%s\"", params, t+1, (nsteps+starting_time < t+1 ? 0 : nsteps+starting_time-t-1), savefile_string);
      free(params);
    }
    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);
    }
  }

  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,
  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
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  uint64_t 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=starting_time;nsteps==0 || t<starting_time+nsteps;t++){
    if(algorithm==ALGORITHM_RK2){
      ns_step_rk2(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, irreversible);
    } else {
      ns_step_rk4(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3, irreversible);
    }
  }

  // save final entry to savefile
  write_vec(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*(2*K2+1)+K2);

  // allocate tmp vectors for computation
  *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);

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

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