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

using Distributed
@everywhere using FastGaussQuadrature
@everywhere using LinearAlgebra
@everywhere using Polynomials
@everywhere using Printf
@everywhere using SpecialFunctions
include("chebyshev.jl")
include("integration.jl")
include("interpolation.jl")
include("tools.jl")
include("potentials.jl")
include("print.jl")
include("easyeq.jl")
include("simpleq-Kv.jl")
include("simpleq-iteration.jl")
include("simpleq-hardcore.jl")
include("anyeq.jl")
function main()
  ## defaults

  # values of rho,e, when needed
  rho=1e-6
  e=1e-4

  # incrementally initialize Newton algorithm
  nlrho_init=1

  # potential
  v=k->v_exp(k,1.)
  a0=a0_exp(1.)
  # arguments of the potential
  v_param_a=1.
  v_param_b=1.
  v_param_c=1.
  v_param_e=1.

  # plot range when plotting in rho
  minlrho=-6
  maxlrho=2
  nlrho=100
  rhos=Array{Float64}(undef,0)
  # plot range when plotting in e
  minle=-6
  maxle=2
  nle=100
  es=Array{Float64}(undef,0)
  # plot range when plotting in x
  xmin=0
  xmax=100
  nx=100

  # cutoffs
  tolerance=1e-11
  order=100
  maxiter=21

  # for anyeq
  # P
  P=11
  # N
  N=12
  # number of splines
  J=10

  # starting rho from which to incrementally initialize Newton algorithm
  # default must be set after reading rho, if not set explicitly
  minlrho_init=nothing

  # Hann window for Fourier transforms
  windowL=1e3

  # approximation for easyeq
  # bK,bL
  easyeq_simpleq_approx=Easyeq_approx(0.,0.)
  easyeq_medeq_approx=Easyeq_approx(1.,1.)
  easyeq_approx=easyeq_simpleq_approx

  # approximation for anyeq
  # aK,bK,gK,aL1,bL1,aL2,bL2,gL2,aL3,bL3,gl3
  anyeq_simpleq_approx=Anyeq_approx(0.,0.,1.,0.,0.,0.,0.,0.,0.,0.,0.)
  anyeq_bigeq_approx=Anyeq_approx(1.,1.,1.,1.,1.,1.,1.,1.,0.,0.,0.)
  anyeq_fulleq_approx=Anyeq_approx(1.,1.,1.,1.,1.,1.,1.,1.,1.,0.,1.)
  anyeq_medeq_approx=Anyeq_approx(0.,1.,1.,0.,1.,0.,0.,0.,0.,0.,0.)
  anyeq_compleq_approx=Anyeq_approx(1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.)
  anyeq_approx=anyeq_bigeq_approx

  # read cli arguments
  (params,potential,method,savefile,command)=read_args(ARGS)

  # read params
  if params!=""
    for param in split(params,";")
      terms=split(param,"=")
      if length(terms)!=2
	print(stderr,"error: could not read parameter '",param,"'.\n")
	exit(-1)
      end
      lhs=terms[1]
      rhs=terms[2]
      if lhs=="rho"
	rho=parse(Float64,rhs)
      elseif lhs=="minlrho_init"
	minlrho_init=parse(Float64,rhs)
      elseif lhs=="nlrho_init"
	nlrho_init=parse(Int64,rhs)
      elseif lhs=="e"
	e=parse(Float64,rhs)
      elseif lhs=="tolerance"
	tolerance=parse(Float64,rhs)
      elseif lhs=="order"
	order=parse(Int64,rhs)
      elseif lhs=="maxiter"
	maxiter=parse(Int64,rhs)
      elseif lhs=="rhos"
	rhos=parse_list(rhs)
      elseif lhs=="minlrho"
	minlrho=parse(Float64,rhs)
      elseif lhs=="maxlrho"
	maxlrho=parse(Float64,rhs)
      elseif lhs=="nlrho"
	nlrho=parse(Int64,rhs)
      elseif lhs=="es"
	es=parse_list(rhs)
      elseif lhs=="minle"
	minle=parse(Float64,rhs)
      elseif lhs=="maxle"
	maxle=parse(Float64,rhs)
      elseif lhs=="nle"
	nle=parse(Int64,rhs)
      elseif lhs=="xmin"
	xmin=parse(Float64,rhs)
      elseif lhs=="xmax"
	xmax=parse(Float64,rhs)
      elseif lhs=="nx"
	nx=parse(Int64,rhs)
      elseif lhs=="P"
	P=parse(Int64,rhs)
      elseif lhs=="N"
	N=parse(Int64,rhs)
      elseif lhs=="J"
	J=parse(Int64,rhs)
      elseif lhs=="window_L"
	windowL=parse(Float64,rhs)
      elseif lhs=="aK"
	anyeq_approx.aK=parse(Float64,rhs)
      elseif lhs=="bK"
	anyeq_approx.bK=parse(Float64,rhs)
	easyeq_approx.bK=parse(Float64,rhs)
      elseif lhs=="gK"
	anyeq_approx.gK=parse(Float64,rhs)
      elseif lhs=="aL1"
	anyeq_approx.aL1=parse(Float64,rhs)
      elseif lhs=="bL"
	easyeq_approx.bL=parse(Float64,rhs)
      elseif lhs=="bL1"
	anyeq_approx.bL1=parse(Float64,rhs)
      elseif lhs=="aL2"
	anyeq_approx.aL2=parse(Float64,rhs)
      elseif lhs=="bL2"
	anyeq_approx.bL2=parse(Float64,rhs)
      elseif lhs=="gL2"
	anyeq_approx.gL2=parse(Float64,rhs)
      elseif lhs=="aL3"
	anyeq_approx.aL3=parse(Float64,rhs)
      elseif lhs=="bL3"
	anyeq_approx.bL3=parse(Float64,rhs)
      elseif lhs=="gL3"
	anyeq_approx.gL3=parse(Float64,rhs)
      elseif lhs=="v_a"
	v_param_a=parse(Float64,rhs)
      elseif lhs=="v_b"
	v_param_b=parse(Float64,rhs)
      elseif lhs=="v_c"
	v_param_c=parse(Float64,rhs)
      elseif lhs=="v_e"
	v_param_e=parse(Float64,rhs)
      elseif lhs=="eq"
	if rhs=="simpleq"
	  easyeq_approx=easyeq_simpleq_approx
	  anyeq_approx=anyeq_simpleq_approx
	elseif rhs=="medeq"
	  easyeq_approx=easyeq_medeq_approx
	  anyeq_approx=anyeq_medeq_approx
	elseif rhs=="bigeq"
	  anyeq_approx=anyeq_bigeq_approx
	elseif rhs=="fulleq"
	  anyeq_approx=anyeq_fulleq_approx
	elseif rhs=="compleq"
	  anyeq_approx=anyeq_compleq_approx
	else
	    print(stderr,"error: unrecognized equation: ",rhs,"\n")
	    exit(-1)
	end
      else
	print(stderr,"unrecognized parameter '",lhs,"'.\n")
	exit(-1)
      end
    end
  end

  ## read potential
  if potential=="exp"
    v=k->v_exp(k,v_param_a)
    a0=a0_exp(v_param_a)
  elseif potential=="expcry"
    v=k->v_expcry(k,v_param_a,v_param_b)
    a0=a0_expcry(v_param_a,v_param_b)
  elseif potential=="npt"
    v=k->v_npt(k)
    a0=a0_npt()
  elseif potential=="alg"
    v=v_alg
    a0=a0_alg
  elseif potential=="algwell"
    v=v_algwell
    a0=a0_algwell
  elseif potential=="exact"
    v=k->v_exact(k,v_param_b,v_param_c,v_param_e)
    a0=a0_exact(v_param_b,v_param_c,v_param_e)
  elseif potential=="tent"
    v=k->v_tent(k,v_param_a,v_param_b)
    a0=a0_tent(v_param_a,v_param_b)
  else
    print(stderr,"unrecognized potential '",potential,"'.\n'")
    exit(-1)
  end

  ## set parameters
  # rhos
  if length(rhos)==0
    rhos=Array{Float64}(undef,nlrho)
    for j in 0:nlrho-1
      rhos[j+1]=(nlrho==1 ? 10^minlrho : 10^(minlrho+(maxlrho-minlrho)/(nlrho-1)*j))
    end
  end
  # es
  if length(es)==0
    es=Array{Float64}(undef,nle)
    for j in 0:nle-1
      es[j+1]=(nle==1 ? 10^minle : 10^(minle+(maxle-minle)/(nle-1)*j))
    end
  end

  # default minlrho_init
  if (minlrho_init==nothing)
    minlrho_init=log10(rho)
  end

  # splines
  taus=Array{Float64}(undef,J+1)
  for j in 0:J
    taus[j+1]=-1+2*j/J
  end

  ## run command
  if method=="easyeq"
    if command=="energy"
      easyeq_energy(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,easyeq_approx)
    # e(rho)
    elseif command=="energy_rho"
      easyeq_energy_rho(rhos,order,a0,v,maxiter,tolerance,easyeq_approx)
    # u(k)
    elseif command=="uk"
      easyeq_uk(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,easyeq_approx)
    # u(x)
    elseif command=="ux"
      easyeq_ux(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,easyeq_approx)
    elseif command=="uux"
      easyeq_uux(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,easyeq_approx)
    # condensate fraction
    elseif command=="condensate_fraction"
      easyeq_condensate_fraction(minlrho_init,nlrho_init,order,rho,a0,v,maxiter,tolerance,easyeq_approx)
    elseif command=="condensate_fraction_rho"
      easyeq_condensate_fraction_rho(rhos,order,a0,v,maxiter,tolerance,easyeq_approx)
    else
      print(stderr,"unrecognized command '",command,"'.\n")
      exit(-1)
    end
  elseif method=="simpleq-hardcore"
    if command=="energy_rho"
      simpleq_hardcore_energy_rho(rhos,taus,P,N,J,maxiter,tolerance)
    elseif command=="ux"
      simpleq_hardcore_ux(rho,taus,P,N,J,maxiter,tolerance)
    elseif command=="condensate_fraction_rho"
      simpleq_hardcore_condensate_fraction_rho(rhos,taus,P,N,J,maxiter,tolerance)
    end
  elseif method=="simpleq-iteration"
    # u_n(x) using iteration
    if command=="u"
      simpleq_iteration_ux(order,e,v,maxiter,xmin,xmax,nx)
    # rho(e) using iteration
    elseif command=="rho_e"
      simpleq_iteration_rho_e(es,order,v,maxiter)
    else
      print(stderr,"unrecognized command '",command,"'.\n")
      exit(-1)
    end
  elseif method=="anyeq"
    if command=="save_Abar"
      anyeq_save_Abar(taus,P,N,J,v,anyeq_approx)
    elseif command=="energy"
      anyeq_energy(rho,minlrho_init,nlrho_init,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    # e(rho)
    elseif command=="energy_rho"
      anyeq_energy_rho(rhos,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    elseif command=="energy_rho_init_prevrho"
      anyeq_energy_rho_init_prevrho(rhos,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    elseif command=="energy_rho_init_nextrho"
      anyeq_energy_rho_init_nextrho(rhos,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    # u(j)
    elseif command=="uk"
      anyeq_uk(minlrho_init,nlrho_init,taus,P,N,J,rho,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    # u(j)
    elseif command=="ux"
      anyeq_ux(minlrho_init,nlrho_init,taus,P,N,J,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,anyeq_approx,savefile)
    # condensate fraction
    elseif command=="condensate_fraction"
      anyeq_condensate_fraction(rho,minlrho_init,nlrho_init,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    elseif command=="condensate_fraction_rho"
      anyeq_condensate_fraction_rho(rhos,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    # momentum distribution
    elseif command=="momentum_distribution"
      anyeq_momentum_distribution(rho,minlrho_init,nlrho_init,taus,P,N,J,a0,v,maxiter,tolerance,anyeq_approx,savefile)
    elseif command=="2pt"
      anyeq_2pt_correlation(minlrho_init,nlrho_init,taus,P,N,J,windowL,rho,a0,v,maxiter,tolerance,xmin,xmax,nx,anyeq_approx,savefile)
    else
      print(stderr,"unrecognized command: '",command,"'\n")
      exit(-1)
    end
  elseif method=="simpleq-Kv"
    if command=="2pt"
      simpleq_Kv_2pt(minlrho_init,nlrho_init,taus,P,N,J,rho,a0,v,maxiter,tolerance,xmin,xmax,nx)
    elseif command=="condensate_fraction_rho"
      simpleq_Kv_condensate_fraction(rhos,taus,P,N,J,a0,v,maxiter,tolerance)
    elseif command=="Kv"
      simpleq_Kv_Kv(minlrho_init,nlrho_init,taus,P,N,J,rho,a0,v,maxiter,tolerance,xmin,xmax,nx)
    else
      print(stderr,"unrecognized command: '",command,"'\n")
      exit(-1)
    end
  else
    print(stderr,"unrecognized method: '",method,"'\n")
    exit(-1)
  end

end

# parse a comma separated list as an array of Float64
function parse_list(str)
  elems=split(str,",")
  out=Array{Float64}(undef,length(elems))
  for i in 1:length(elems)
    out[i]=parse(Float64,elems[i])
  end
  return out
end

# read cli arguments
function read_args(ARGS)
  # flag
  flag=""

  # output strings
  params=""
  # default potential
  potential="exp"
  # default method
  method="easyeq"
  
  savefile=""
  command=""

  # loop over arguments
  for arg in ARGS
    # argument that starts with a dash
    if arg[1]=='-'
      # go through characters after dash
      for char in arg[2:length(arg)]
	
	# set params
	if char=='p'
	  # raise flag
	  flag="params"
	elseif char=='U'
	  # raise flag
	  flag="potential"
	elseif char=='M'
	  # raise flag
	  flag="method"
	elseif char=='s'
	  # raise flag
	  flag="savefile"
	else
	  print_usage()
	  exit(-1)
	end
      end
    # arguments that do not start with a dash
    else
      if flag=="params"
	params=arg
      elseif flag=="potential"
	potential=arg
      elseif flag=="method"
	method=arg
      elseif flag=="savefile"
	savefile=arg
      else
	command=arg
      end
      # reset flag
      flag=""
    end
  end

  if command==""
    print_usage()
    exit(-1)
  end

  return (params,potential,method,savefile,command)
end

# usage
function print_usage()
  print(stderr,"usage:  simplesolv [-p params] [-U potential] [-M method] [-s savefile] <command>\n")
end

main()