## Copyright 2021 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.

# interpolation
@everywhere mutable struct Easyeq_approx
  bK::Float64
  bL::Float64
end

# compute energy
function easyeq_energy(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,approx)
  # compute gaussian quadrature weights
  weights=gausslegendre(order)

  # compute initial guess from previous rho
  (u,E,err)=easyeq_hatu(easyeq_init_u(a0,order,weights),order,(10.)^minlrho_init,v,maxiter,tolerance,weights,approx)
  for j in 2:nlrho_init
    rho_tmp=10^(minlrho_init+(log10(rho)-minlrho_init)*(j-1)/(nlrho_init-1))
    (u,E,err)=easyeq_hatu(u,order,rho_tmp,v,maxiter,tolerance,weights,approx)
  end

  # print energy
  @printf("% .15e % .15e\n",real(E),err)
end

# compute energy as a function of rho
function easyeq_energy_rho(rhos,order,a0,v,maxiter,tolerance,approx)
  # compute gaussian quadrature weights
  weights=gausslegendre(order)
  # init u
  u=easyeq_init_u(a0,order,weights)

  for j in 1:length(rhos)
    # compute u (init newton with previously computed u)
    (u,E,err)=easyeq_hatu(u,order,rhos[j],v,maxiter,tolerance,weights,approx)

    @printf("% .15e % .15e % .15e\n",rhos[j],real(E),err)

  end
end

# compute u(k)
function easyeq_uk(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,approx)
  weights=gausslegendre(order)

  # compute initial guess from previous rho
  (u,E,err)=easyeq_hatu(easyeq_init_u(a0,order,weights),order,(10.)^minlrho_init,v,maxiter,tolerance,weights,approx)
  for j in 2:nlrho_init
    rho_tmp=10^(minlrho_init+(log10(rho)-minlrho_init)*(j-1)/(nlrho_init-1))
    (u,E,err)=easyeq_hatu(u,order,rho_tmp,v,maxiter,tolerance,weights,approx)
  end

  for i in 1:order
    k=(1-weights[1][i])/(1+weights[1][i])
    @printf("% .15e % .15e\n",k,real(u[i]))
  end
end

# compute u(x)
function easyeq_ux(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,approx)
  weights=gausslegendre(order)

  # compute initial guess from previous rho
  (u,E,err)=easyeq_hatu(easyeq_init_u(a0,order,weights),order,(10.)^minlrho_init,v,maxiter,tolerance,weights,approx)
  for j in 2:nlrho_init
    rho_tmp=10^(minlrho_init+(log10(rho)-minlrho_init)*(j-1)/(nlrho_init-1))
    (u,E,err)=easyeq_hatu(u,order,rho_tmp,v,maxiter,tolerance,weights,approx)
  end

  for i in 1:nx
    x=xmin+(xmax-xmin)*i/nx
    @printf("% .15e % .15e\n",x,real(easyeq_u_x(x,u,weights)))
  end
end

# compute 2u(x)-rho u*u(x)
function easyeq_uux(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,approx)
  weights=gausslegendre(order)

  # compute initial guess from previous rho
  (u,E,err)=easyeq_hatu(easyeq_init_u(a0,order,weights),order,(10.)^minlrho_init,v,maxiter,tolerance,weights,approx)
  for j in 2:nlrho_init
    rho_tmp=10^(minlrho_init+(log10(rho)-minlrho_init)*(j-1)/(nlrho_init-1))
    (u,E,err)=easyeq_hatu(u,order,rho_tmp,v,maxiter,tolerance,weights,approx)
  end

  for i in 1:nx
    x=xmin+(xmax-xmin)*i/nx
    @printf("% .15e % .15e\n",x,real(easyeq_u_x(x,2*u-rho*u.*u,weights)))
  end
end

# condensate fraction
function easyeq_condensate_fraction(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,approx)
  # compute gaussian quadrature weights
  weights=gausslegendre(order)

  # compute initial guess from previous rho
  (u,E,err)=easyeq_hatu(easyeq_init_u(a0,order,weights),order,(10.)^minlrho_init,v,maxiter,tolerance,weights,approx)
  for j in 2:nlrho_init
    rho_tmp=10^(minlrho_init+(log10(rho)-minlrho_init)*(j-1)/(nlrho_init-1))
    (u,E,err)=easyeq_hatu(u,order,rho_tmp,v,maxiter,tolerance,weights,approx)
  end

  # compute eta
  eta=easyeq_eta(u,order,rho,v,maxiter,tolerance,weights,approx)

  # print energy
  @printf("% .15e % .15e\n",eta,err)
end

# condensate fraction as a function of rho
function easyeq_condensate_fraction_rho(rhos,order,a0,v,maxiter,tolerance,approx)
  weights=gausslegendre(order)
  # init u
  u=easyeq_init_u(a0,order,weights)

  for j in 1:length(rhos)
    # compute u (init newton with previously computed u)
    (u,E,err)=easyeq_hatu(u,order,rhos[j],v,maxiter,tolerance,weights,approx)

    # compute eta
    eta=easyeq_eta(u,order,rhos[j],v,maxiter,tolerance,weights,approx)

    @printf("% .15e % .15e % .15e\n",rhos[j],eta,err)
  end
end


# initialize u
@everywhere function easyeq_init_u(a0,order,weights)
  u=zeros(Float64,order)
  for j in 1:order
    # transformed k
    k=(1-weights[1][j])/(1+weights[1][j])
    u[j]=4*pi*a0/k^2
  end

  return u
end

# \hat u(k) computed using Newton
@everywhere function easyeq_hatu(u0,order,rho,v,maxiter,tolerance,weights,approx)
  # initialize V and Eta
  (V,V0)=easyeq_init_v(weights,v)
  (Eta,Eta0)=easyeq_init_H(weights,v)

  # init u
  u=rho*u0

  # iterate
  err=Inf
  for i in 1:maxiter-1
    new=u-inv(easyeq_dXi(u,V,V0,Eta,Eta0,weights,rho,approx))*easyeq_Xi(u,V,V0,Eta,Eta0,weights,rho,approx)

    err=norm(new-u)/norm(u)
    if(err<tolerance)
      u=new
      break
    end
    u=new
  end

  return (u/rho,easyeq_en(u,V0,Eta0,rho,weights)*rho/2,err)
end

# \Eta
@everywhere function easyeq_H(x,t,weights,v)
  return (x>t ? 2*t/x : 2)* integrate_legendre(y->2*pi*((x+t)*y+abs(x-t)*(1-y))*v((x+t)*y+abs(x-t)*(1-y)),0,1,weights)
end

# initialize V
@everywhere function easyeq_init_v(weights,v)
  order=length(weights[1])
  V=Array{Float64}(undef,order)
  V0=v(0)
  for i in 1:order
    k=(1-weights[1][i])/(1+weights[1][i])
    V[i]=v(k)
  end
  return(V,V0)
end

# initialize Eta
@everywhere function easyeq_init_H(weights,v)
  order=length(weights[1])
  Eta=Array{Array{Float64}}(undef,order)
  Eta0=Array{Float64}(undef,order)
  for i in 1:order
    k=(1-weights[1][i])/(1+weights[1][i])
    Eta[i]=Array{Float64}(undef,order)
    for j in 1:order
      y=(weights[1][j]+1)/2
      Eta[i][j]=easyeq_H(k,(1-y)/y,weights,v)
    end
    y=(weights[1][i]+1)/2
    Eta0[i]=easyeq_H(0,(1-y)/y,weights,v)
  end
  return(Eta,Eta0)
end

# Xi(u)
@everywhere function easyeq_Xi(u,V,V0,Eta,Eta0,weights,rho,approx)
  order=length(weights[1])

  # init
  out=zeros(Float64,order)

  # compute E before running the loop
  E=easyeq_en(u,V0,Eta0,rho,weights)

  for i in 1:order
    # k_i
    k=(1-weights[1][i])/(1+weights[1][i])
    # S_i
    S=V[i]-1/(rho*(2*pi)^3)*integrate_legendre_sampled(y->(1-y)/y^3,Eta[i].*u,0,1,weights)

    # A_K,i
    A=0.
    if approx.bK!=0.
      A+=approx.bK*S
    end
    if approx.bK!=1.
      A+=(1-approx.bK)*E
    end

    # T
    if approx.bK==1.
      T=1.
    else
      T=S/A
    end

    # B
    if approx.bK==approx.bL
      B=1.
    else
      B=(approx.bL*S+(1-approx.bL*E))/(approx.bK*S+(1-approx.bK*E))
    end

    # X_i
    X=k^2/(2*A*rho)

    # U_i
    out[i]=u[i]-T/(2*(X+1))*Phi(B*T/(X+1)^2)
  end

  return out
end

# derivative of Xi
@everywhere function easyeq_dXi(u,V,V0,Eta,Eta0,weights,rho,approx)
  order=length(weights[1])

  # init
  out=zeros(Float64,order,order)

  # compute E before the loop
  E=easyeq_en(u,V0,Eta0,rho,weights)

  for i in 1:order
    # k_i
    k=(1-weights[1][i])/(1+weights[1][i])
    # S_i
    S=V[i]-1/(rho*(2*pi)^3)*integrate_legendre_sampled(y->(1-y)/y^3,Eta[i].*u,0,1,weights)

    # A_K,i
    A=0.
    if approx.bK!=0.
      A+=approx.bK*S
    end
    if approx.bK!=1.
      A+=(1-approx.bK)*E
    end

    # T
    if approx.bK==1.
      T=1.
    else
      T=S/A
    end

    # B
    if approx.bK==approx.bL
      B=1.
    else
      B=(approx.bL*S+(1-approx.bL*E))/(approx.bK*S+(1-approx.bK*E))
    end

    # X_i
    X=k^2/(2*A*rho)

    for j in 1:order
      y=(weights[1][j]+1)/2
      dS=-1/rho*(1-y)*Eta[i][j]/(2*(2*pi)^3*y^3)*weights[2][j]
      dE=-1/rho*(1-y)*Eta0[j]/(2*(2*pi)^3*y^3)*weights[2][j]

      # dA
      dA=0.
      if approx.bK!=0.
	dA+=approx.bK*dS
      end
      if approx.bK!=1.
	dA+=(1-approx.bK)*dE
      end

      # dT
      if approx.bK==1.
	dT=0.
      else
	dT=(1-approx.bK)*(E*dS-S*dE)/A^2
      end

      # dB
      if approx.bK==approx.bL
	dB=0.
      else
	dB=(approx.bL*(1-approx.bK)-approx.bK*(1-approx.bL))*(E*dS-S*dE)/(approx.bK*S+(1-approx.bK*E))^2
      end

      dX=-k^2/(2*A^2*rho)*dA

      out[i,j]=(i==j ? 1 : 0)-(dT-T*dX/(X+1))/(2*(X+1))*Phi(B*T/(X+1)^2)-T/(2*(X+1)^3)*(B*dT+T*dB-2*B*T*dX/(X+1))*dPhi(B*T/(X+1)^2)
    end
  end

  return out
end

# derivative of Xi with respect to mu
@everywhere function easyeq_dXidmu(u,V,V0,Eta,Eta0,weights,rho,approx)
  order=length(weights[1])

  # init
  out=zeros(Float64,order)

  # compute E before running the loop
  E=easyeq_en(u,V0,Eta0,rho,weights)

  for i in 1:order
    # k_i
    k=(1-weights[1][i])/(1+weights[1][i])
    # S_i
    S=V[i]-1/(rho*(2*pi)^3)*integrate_legendre_sampled(y->(1-y)/y^3,Eta[i].*u,0,1,weights)

    # A_K,i
    A=0.
    if approx.bK!=0.
      A+=approx.bK*S
    end
    if approx.bK!=1.
      A+=(1-approx.bK)*E
    end

    # T
    if approx.bK==1.
      T=1.
    else
      T=S/A
    end

    # B
    if approx.bK==approx.bL
      B=1.
    else
      B=(approx.bL*S+(1-approx.bL*E))/(approx.bK*S+(1-approx.bK*E))
    end

    # X_i
    X=k^2/(2*A*rho)

    out[i]=T/(2*rho*A*(X+1)^2)*Phi(B*T/(X+1)^2)+B*T^2/(rho*A*(X+1)^4)*dPhi(B*T/(X+1)^2)
  end

  return out
end

# energy
@everywhere function easyeq_en(u,V0,Eta0,rho,weights)
  return V0-1/(rho*(2*pi)^3)*integrate_legendre_sampled(y->(1-y)/y^3,Eta0.*u,0,1,weights)
end

# condensate fraction
@everywhere function easyeq_eta(u,order,rho,v,maxiter,tolerance,weights,approx)
  (V,V0)=easyeq_init_v(weights,v)
  (Eta,Eta0)=easyeq_init_H(weights,v)

  du=-inv(easyeq_dXi(rho*u,V,V0,Eta,Eta0,weights,rho,approx))*easyeq_dXidmu(rho*u,V,V0,Eta,Eta0,weights,rho,approx)

  eta=-1/(2*(2*pi)^3)*integrate_legendre_sampled(y->(1-y)/y^3,Eta0.*du,0,1,weights)

  return eta
end

# inverse Fourier transform
@everywhere function easyeq_u_x(x,u,weights)
  order=length(weights[1])
  out=integrate_legendre_sampled(y->(1-y)/y^3*sin(x*(1-y)/y)/x/(2*pi^2),u,0,1,weights)
  return out
end