#include "navier-stokes.h"
#include "io.h"
#include <math.h>
#include <stdlib.h>

// compute solution as a function of time
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,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  unsigned int 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("   %6d:(%4d,%4d)r ",t,kx,ky);
      t++;
      printf("   %6d:(%4d,%4d)i ",t,kx,ky);
      t++;
    }
  }

  // iterate
  for(t=0;t<nsteps;t++){
    ins_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
    
    if(t%print_freq==0){
      fprintf(stderr,"%d % .8e ",t,t*delta);
      printf("%8d % .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__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
	    
	  }
	  printf("% .15e % .15e ",__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
	}
      }
      fprintf(stderr,"\n");
      printf("\n");
    }
  }

  // save final entry to savefile
  write_u(u, K1, K2, savefile);

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

// compute energy, enstrophy, alpha as a function of time in the I-NS equation
int eea(
  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,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  double alpha, energy, enstrophy;
  double avg_e,avg_a,avg_en;
  unsigned int 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_e=0;
  avg_a=0;
  avg_en=0;

  // iterate
  for(t=0;t<nsteps;t++){
    ins_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);

    energy=compute_energy(u, K1, K2);
    alpha=compute_alpha(u, K1, K2, g);
    enstrophy=compute_enstrophy(u, K1, K2, L);
    
    // running average
    if(t>0){
      avg_e=avg_e-(avg_e-energy)/t;
      avg_a=avg_a-(avg_a-alpha)/t;
      avg_en=avg_en-(avg_en-enstrophy)/t;
    }

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

  // save final entry to savefile
  write_u(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,
  unsigned int nsteps,
  double nu,
  double delta,
  double L,
  _Complex double* u0,
  _Complex double* g,
  unsigned int nthreads,
  FILE* savefile
){
  _Complex double* u;
  _Complex double* tmp1;
  _Complex double* tmp2;
  _Complex double* tmp3;
  unsigned int 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=0;t<nsteps;t++){
    ins_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
  }

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

  // allocate tmp vectors for computation
  *tmp1=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
  *tmp2=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
  *tmp3=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));

  // 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<(2*K1+1)*(2*K2+1);i++){
    u[i]=u0[i];
  }

  return 0;
}

// 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
){
  int kx,ky;

  // k1
  ins_rhs(tmp1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k1/2
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k2
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k2/2
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k3
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  // u+h*k3
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }
  // k4
  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft);
  // add to output
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u[klookup(kx,ky,2*K1+1,2*K2+1)]=tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
    }
  }

  return(0);
}

// 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
){
  int kx,ky;
  int i;

  // compute convolution term
  ns_T(u,K1,K2,N1,N2,fft1,fft2,ifft);


  /*
  // compare convolution term (store result in fft1.fft)
  ns_T_nofft(fft1.fft, u, K1, K2, N1, N2);
  double cmp=0.;
  for(i=0;i<N1*N2;i++){
    cmp+=(ifft.fft[i]-fft1.fft[i])*(ifft.fft[i]-fft1.fft[i]);
  }
  printf("% .15e\n",sqrt(cmp));
  */


  for(i=0; i<(2*K1+1)*(2*K2+1); i++){
    out[i]=0;
  }
  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if(kx!=0 || ky!=0){
	// enforce the reality of u by adding ifft.fft(k) and the conjugate of ifft.fft(-k)
        out[klookup(kx,ky,2*K1+1,2*K2+1)]=-4*M_PI*M_PI/L/L*nu*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]+g[klookup(kx,ky,2*K1+1,2*K2+1)]+4*M_PI*M_PI/L/L/sqrt(kx*kx+ky*ky)*(ifft.fft[klookup(kx,ky,N1,N2)]+conj(ifft.fft[klookup(-kx,-ky,N1,N2)]))/2;
      }
    }
  }

  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)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=ky*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/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)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/N1;
        fft2.fft[klookup(kx,ky,N1,N2)]=kx*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]/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>=4*K1+1 and N2>=4*K2+1)
  if (N1<4*K1+1 || N2<4*K2+1){
    fprintf(stderr,"error: N1 and N2 need t be >= 4*K1+1 and 4*K2+1 respectively\n");
    return(-1);
  }
  for(kx=-2*K1;kx<=2*K1;kx++){
    for(ky=-2*K2;ky<=2*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)*u[klookup(px,py,2*K1+1,2*K2+1)]*u[klookup(qx,qy,2*K1+1,2*K2+1)];
	  }
	}
      }
    }
  }

  return 0;
}

// compute alpha
double compute_alpha(
  _Complex double* u,
  int K1,
  int K2,
  _Complex double* g
){
  _Complex double num=0;
  _Complex double denom=0;
  int kx,ky;

  for(kx=-K1;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      denom+=(kx*kx+ky*ky)*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)])*(1+(ky!=0?kx*kx/ky/ky:0));
      num+=(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(g[klookup(kx,ky,2*K1+1,2*K2+1)])*(1+(ky!=0?kx*kx/ky/ky:0));
    }
  }

  return __real__ (num/denom);
}

// compute energy
double compute_energy(
  _Complex double* u,
  int K1,
  int K2
){
  int kx,ky;
  double out=0.;
  for(kx=-K1;kx<=K1;kx++){
    for (ky=-K2;ky<=K2;ky++){
      out+=__real__ (u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)]));
    }
  }
  return out;
}

// compute enstrophy
double compute_enstrophy(
  _Complex double* u,
  int K1,
  int K2,
  double L
){
  int kx,ky;
  double out=0.;
  for(kx=-K1;kx<=K1;kx++){
    for (ky=-K2;ky<=K2;ky++){
      out+=__real__ (4*M_PI*M_PI/L/L*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)]));
    }
  }
  return 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);
}