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

# compute rho(e) using the iteration
function simpleq_iteration_rho_e(es,order,v,maxiter)
  for j in 1:length(es)
    (u,rho)=simpleq_iteration_hatun(es[j],order,v,maxiter)
    @printf("% .15e % .15e\n",es[j],real(rho[maxiter+1]))
  end
end

# compute u(x) using the iteration and print at every step
function simpleq_iteration_ux(order,e,v,maxiter,xmin,xmax,nx)
  (u,rho)=simpleq_iteration_hatun(e,order,v,maxiter)

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


# \hat u_n
function simpleq_iteration_hatun(e,order,v,maxiter)
  # gauss legendre weights
  weights=gausslegendre(order)

  # initialize V and Eta
  (V,V0)=easyeq_init_v(weights,v)
  (Eta,Eta0)=easyeq_init_H(weights,v)

  # init u and rho
  u=Array{Array{Float64}}(undef,maxiter+1)
  u[1]=zeros(Float64,order)
  rho=zeros(Float64,maxiter+1)

  # iterate
  for n in 1:maxiter
    u[n+1]=simpleq_iteration_A(e,weights,Eta)\simpleq_iteration_b(u[n],e,rho[n],V)
    rho[n+1]=simpleq_iteration_rhon(u[n+1],e,weights,V0,Eta0)
  end

  return (u,rho)
end

# A
function simpleq_iteration_A(e,weights,Eta)
  N=length(weights[1])
  out=zeros(Float64,N,N)
  for i in 1:N
    k=(1-weights[1][i])/(1+weights[1][i])
    out[i,i]=k^2+4*e
    for j in 1:N
      y=(weights[1][j]+1)/2
      out[i,j]+=weights[2][j]*(1-y)*Eta[i][j]/(2*(2*pi)^3*y^3)
    end
  end
  return out
end

# b
function simpleq_iteration_b(u,e,rho,V)
  out=zeros(Float64,length(V))
  for i in 1:length(V)
    out[i]=V[i]+2*e*rho*u[i]^2
  end
  return out
end

# rho_n
function simpleq_iteration_rhon(u,e,weights,V0,Eta0)
  S=V0
  for i in 1:length(weights[1])
    y=(weights[1][i]+1)/2
    S+=-weights[2][i]*(1-y)*u[i]*Eta0[i]/(2*(2*pi)^3*y^3)
  end
  return 2*e/S
end