using PolyLog, Plots
include("nucleons.jl")
include("mesons.jl")

"Total binding energy of the system"
total_E(s::system) = -(nucleon_E(s) + meson_E(s))

"Normalized Woods-Saxon form used for constructing an initial solution"
Woods_Saxon(r::Float64; R::Float64=7.0, a::Float64=0.5) = -1 / (8π * a^3 * reli3(-exp(R / a)) * (1 + exp((r - R) / a)))

"Run the full program by self-consistent solution of nucleon and meson densities"
function solve_system(s::system; reinitialize_densities=true, monitor_print=true, monitor_plot=false)
    if reinitialize_densities
        dens_guess = Woods_Saxon.(rs(s))
        @. s.ρ_sp = s.Z * dens_guess
        @. s.ρ_vp = s.Z * dens_guess
        @. s.ρ_sn = s.N * dens_guess
        @. s.ρ_vn = s.N * dens_guess
    end

    if monitor_plot
        p = plot(legends=false, size=(1024, 768), layout=(2, 4), title=["ρₛₚ" "ρᵥₚ" "ρₛₙ" "ρᵥₙ" "Φ₀" "W₀" "B₀" "A₀"])
    end

    previous_E_per_A = NaN

    while true
        # mesons
        @time "Meson fields" solveMesonFields!(s, isnan(previous_E_per_A) ? 50 : 15)

        # protons
        @time "Proton densities" (s.ρ_sp, s.ρ_vp) = solveNucleonDensity(true, s)

        # neutrons
        @time "Neutron densities" (s.ρ_sn, s.ρ_vn) = solveNucleonDensity(false, s)

        if monitor_plot
            for s in p.series_list
                s.plotattributes[:linecolor] = :gray
            end
            plot!(p, rs(s), hcat(s.ρ_sp, s.ρ_vp, s.ρ_sn, s.ρ_vn, s.Φ0, s.W0, s.B0, s.A0), linecolor=:red)
            display(p)
        end

        E_per_A = total_E(s) / A(s)
        monitor_print && println("Total binding E per nucleon = $E_per_A")

        # check convergence
        abs(previous_E_per_A - E_per_A) < 0.0001 && break
        previous_E_per_A = E_per_A
    end
end