# solve Navier-Stokes
function navier_stokes_uk(
  nsteps::Int64,
  nu::Float64,
  g::Function,
  N1::Int64,
  N2::Int64,
  K1::Int64,
  K2::Int64,
  delta::Float64,
  print_freq::Int64
)
  # init
  u=ns_init(K1,K2)

  for i in 1:nsteps
    u=ns_RK4(u,nu,g,N1,N2,K1,K2,delta)

    if i % print_freq==0
      @printf("% .15e ",delta*i)
      print_u_inline(u,N1,N2,K1,K2)

      # print to stderr
      @printf(stderr,"% .8e ",delta*i)
      for k1 in -K1:K1
	for k2 in -K2:K2
	  index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
	  if (k1^2+k2^2<=1)
	    @printf(stderr,"% .8e % .8e ",real(u[index]),imag(u[index]))
	  end
	end
      end
      @printf(stderr,"\n")
    end
  end

end

# compute the enstrophy
function navier_stokes_alpha(
  nsteps::Int64,
  nu::Float64,
  g::Function,
  N1::Int64,
  N2::Int64,
  K1::Int64,
  K2::Int64,
  delta::Float64,
  print_freq::Int64
)
  # init
  u=ns_init(K1,K2)

  for i in 1:nsteps
    u=ns_RK4(u,nu,g,N1,N2,K1,K2,delta)
    if i % print_freq==0
      alpha=ns_alpha(u,g,K1,K2)
      @printf("% .15e % .15e % .15e\n",i*delta,real(alpha),imag(alpha))
      @printf(stderr,"% .15e % .15e % .15e\n",i*delta,real(alpha),imag(alpha))
    end
  end
end

# initialize u
function ns_init(
  K1::Int64,
  K2::Int64
)
  u=zeros(Complex{Float64},(2*K1+1)*(2*K2+1))
  for k1 in -K1:K1
    for k2 in -K2:K2
      index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
      u[index]=exp(-sqrt(k1^2+k2^2))
    end
  end

  return u
end

# Navier-Stokes equation (right side)
function ns_rhs(
  u::Array{Complex{Float64},1},
  nu::Float64,
  g::Function,
  N1::Int64,
  N2::Int64,
  K1::Int64,
  K2::Int64
)
  # compute the T term
  out=ns_T(u,N1,N2,K1,K2)

  # add the rest of the terms
  for k1 in -K1:K1
    for k2 in -K2:K2
      index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
      out[index]=-(4*pi^2*nu*(k1^2+k2^2))*u[index]+g(k1,k2)+(k1^2+k2^2>0 ? 4*pi^2/sqrt(k1^2+k2^2)*out[index] : 0)
    end
  end

  return out
end

# compute k\in K from n\in N
function momentum_from_indices(
  i::Int64,
  K::Int64,
  N::Int64,
)
  if i<=K
    k=i
  elseif i>=N-K
    k=i-N
  else
    k=i
  end
  
  return(k)
end
# inverse
function indices_from_momentum(
  k::Int64,
  N::Int64
)
  if k>=0
    return k
  else
    return N+k
  end
end

# T term in Navier-Stokes
function ns_T(
  u::Array{Complex{Float64},1},
  N1::Int64,
  N2::Int64,
  K1::Int64,
  K2::Int64
)
  Fx1=zeros(Complex{Float64},N1,N2)
  Fy1=zeros(Complex{Float64},N1,N2)
  Fx2=zeros(Complex{Float64},N1,N2)
  Fy2=zeros(Complex{Float64},N1,N2)

  for k1 in -K1:K1
    for k2 in -K2:K2
      i1=indices_from_momentum(k1,N1)+1
      i2=indices_from_momentum(k2,N2)+1
      index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
      # no need to divide by N: ifft does the division already
      Fx1[i1,i2]=(k1==0 ? 0. : k1/sqrt(k1^2+k2^2)*u[index])
      Fx2[i1,i2]=(k1==0 ? 0. : k1*sqrt(k1^2+k2^2)*u[index])
      Fy1[i1,i2]=(k2==0 ? 0. : k2/sqrt(k1^2+k2^2)*u[index])
      Fy2[i1,i2]=(k2==0 ? 0. : k2*sqrt(k1^2+k2^2)*u[index])
    end
  end

  fft!(Fx1)
  fft!(Fx2)
  fft!(Fy1)
  fft!(Fy2)
  Fx1=Fx1.*Fy2.-Fx2.*Fy1
  ifft!(Fx1)

  out=zeros(Complex{Float64},(2*K1+1)*(2*K2+1))
  for k1 in -K1:K1
    for k2 in -K2:K2
      i1=indices_from_momentum(k1,N1)+1
      i2=indices_from_momentum(k2,N2)+1
      index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
      out[index]=Fx1[i1,i2]
    end
  end

  return out
end

# Runge-Kutta step to solve Navier-Stokes
function ns_RK4(
  u::Array{Complex{Float64},1},
  nu::Float64,
  g::Function,
  N1::Int64,
  N2::Int64,
  K1::Int64,
  K2::Int64,
  delta::Float64
)
  k1=delta*ns_rhs(u,nu,g,N1,N2,K1,K2)
  k2=delta*ns_rhs(u.+(k1./2),nu,g,N1,N2,K1,K2)
  k3=delta*ns_rhs(u.+(k2./2),nu,g,N1,N2,K1,K2)
  k4=delta*ns_rhs(u.+k3,nu,g,N1,N2,K1,K2)

  return u.+(k1+2 .*k2+2 .*k3.+k4)./6
end

# enstrophy
function ns_alpha(
  u::Array{Complex{Float64},1},
  g::Function,
  K1::Int64,
  K2::Int64
)
  denom=0. *1im
  num=0. *1im
  for k1 in -K1:K1
    for k2 in -K2:K2
      index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
      denom+=(k1^2+k2^2)^2*u[index]*conj(u[index])*(1. +(k2!=0 ? k1^2/k2^2 : 0.))
      num+=(k1^2+k2^2)*u[index]*conj(g(k1,k2))*(1. +(k2!=0 ? k1^2/k2^2 : 0.))
    end
  end

  return num/denom
end