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

#define M_PI 3.14159265358979323846

// next time step for Irreversible Navier-Stokes equation
int ins_step(_Complex double* u, ns_params params, fft_vects vects, _Complex double* tmp1, _Complex double* tmp2, _Complex double* tmp3){
  int kx,ky;

  // k1
  ins_rhs(tmp1, u, params, vects);
  // add to output
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp3[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/6*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }

  // u+h*k1/2
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/2*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }
  // k2
  ins_rhs(tmp1, tmp2, params, vects);
  // add to output
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp3[KLOOKUP(kx,ky,params.S)]+=params.h/3*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }

  // u+h*k2/2
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/2*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }
  // k3
  ins_rhs(tmp1, tmp2, params, vects);
  // add to output
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp3[KLOOKUP(kx,ky,params.S)]+=params.h/3*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }

  // u+h*k3
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }
  // k4
  ins_rhs(tmp1, tmp2, params, vects);
  // add to output
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      u[KLOOKUP(kx,ky,params.S)]=tmp3[KLOOKUP(kx,ky,params.S)]+params.h/6*tmp1[KLOOKUP(kx,ky,params.S)];
    }
  }

  return(0);
}

// right side of Irreversible Navier-Stokes equation
int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vects vects){
  int kx,ky;

  // F(px/|p|*u)*F(qy*|q|*u)
  // init to 0
  for(kx=0; kx<params.N*params.N; kx++){
    vects.fft1[kx]=0;
    vects.fft2[kx]=0;
    vects.invfft[kx]=0;
  }
  // fill modes
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      if(kx!=0 || ky!=0){
        vects.fft1[KLOOKUP(kx,ky,params.N)]=kx/sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
        vects.fft2[KLOOKUP(kx,ky,params.N)]=ky*sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
      }
    }
  }

  // fft
  fftw_execute(vects.fft1_plan);
  fftw_execute(vects.fft2_plan);
  // write to invfft
  for(kx=-2*params.K;kx<=2*params.K;kx++){
    for(ky=-2*params.K;ky<=2*params.K;ky++){
      vects.invfft[KLOOKUP(kx,ky,params.N)]=vects.fft1[KLOOKUP(kx,ky,params.N)]*vects.fft2[KLOOKUP(kx,ky,params.N)];
    }
  }

  // F(py/|p|*u)*F(qx*|q|*u)
  // init to 0
  for(kx=0; kx<params.N*params.N; kx++){
    vects.fft1[kx]=0;
    vects.fft2[kx]=0;
  }
  // fill modes
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      if(kx!=0 || ky!=0){
        vects.fft1[KLOOKUP(kx,ky,params.N)]=ky/sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
        vects.fft2[KLOOKUP(kx,ky,params.N)]=kx*sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
      }
    }
  }

  // fft
  fftw_execute(vects.fft1_plan);
  fftw_execute(vects.fft2_plan);
  // write to invfft
  for(kx=-2*params.K;kx<=2*params.K;kx++){
    for(ky=-2*params.K;ky<=2*params.K;ky++){
      vects.invfft[KLOOKUP(kx,ky,params.N)]=vects.invfft[KLOOKUP(kx,ky,params.N)]-vects.fft1[KLOOKUP(kx,ky,params.N)]*vects.fft2[KLOOKUP(kx,ky,params.N)];
    }
  }

  // inverse fft
  fftw_execute(vects.invfft_plan);
  
  // write out
  for(kx=0; kx<params.S*params.S; kx++){
    out[kx]=0;
  }
  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      if(kx!=0 || ky!=0){
        out[KLOOKUP(kx,ky,params.S)]=-4*M_PI*M_PI*params.nu*(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]+params.g[KLOOKUP(kx,ky,params.S)]+4*M_PI*M_PI/sqrt(kx*kx+ky*ky)*vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N;
      }
    }
  }

  return(0);
}


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

  for(kx=-params.K;kx<=params.K;kx++){
    for(ky=-params.K;ky<=params.K;ky++){
      denom+=(kx*kx+ky*ky)*(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]*conj(u[KLOOKUP(kx,ky,params.S)])*(1+(ky!=0?kx*kx/ky/ky:0));
      num+=(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]*conj(params.g[KLOOKUP(kx,ky,params.S)])*(1+(ky!=0?kx*kx/ky/ky:0));
    }
  }

  return(num/denom);
}