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