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

"Defines a nuclear system to be solved"
struct system
    Z::Int
    N::Int
    r_max::Float64
    divs::Int
end

"Get mass number of nucleus"
A(s::system) = s.Z + s.N

"Get r values in the mesh"
rs(s::system) = range(0, s.r_max, length=s.divs+1)

"Get Δr value for the mesh"
Δr(s::system) = s.r_max / s.divs

"Create an empty array for the size of the mesh"
zero_array(s::system) = zeros(1 + s.divs)

"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, initial_dens=nothing, initial_flds=(zeros(1 + s.divs) for _ in 1:4); monitor_print=true, monitor_plot=false)
    if isnothing(initial_dens)
        dens_guess = Woods_Saxon.(rs(s))
        ρ_sp = s.Z .* dens_guess
        ρ_vp = s.Z .* dens_guess
        ρ_sn = s.N .* dens_guess
        ρ_vn = s.N .* dens_guess
    else
        (ρ_sp, ρ_vp, ρ_sn, ρ_vn) = initial_dens
    end

    (Φ0s, W0s, B0s, A0s) = initial_flds

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

    E_total_previous = NaN

    while true
        @time "Meson fields" (Φ0s, W0s, B0s, A0s) = solveMesonWfs(ρ_sp, ρ_vp, ρ_sn, ρ_vn, s.r_max, s.divs, isnan(E_total_previous) ? 50 : 5; initial_sol = (Φ0s, W0s, B0s, A0s))

        S_interp = linear_interpolation(rs(s), Φ0s)
        V_interp = linear_interpolation(rs(s), W0s)
        R_interp = linear_interpolation(rs(s), B0s)
        A_interp = linear_interpolation(rs(s), A0s)

        # protons
        @time "Proton spectrum" (κs_p, Es_p) = findAllOrbitals(true, S_interp, V_interp, R_interp, A_interp, s.r_max)
        occs_p = fillNucleons(s.Z, κs_p, Es_p)
        @time "Proton densities" (ρ_sp, ρ_vp) = calculateNucleonDensity(κs_p, Es_p, occs_p, true, S_interp, V_interp, R_interp, A_interp, s.r_max, s.divs)

        # neutrons
        @time "Neutron spectrum" (κs_n, Es_n) = findAllOrbitals(false, S_interp, V_interp, R_interp, A_interp, s.r_max)
        occs_n = fillNucleons(s.N, κs_n, Es_n)
        @time "Neutron densities" (ρ_sn, ρ_vn) = calculateNucleonDensity(κs_n, Es_n, occs_n, false, S_interp, V_interp, R_interp, A_interp, s.r_max, s.divs)

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

        # total binding energy of nucleons
        E_total = sum(M_p .- Es_p) + sum(M_n .- Es_n)
        monitor_print && println("Total binding E per nucleon = $(E_total/A(s))")

        # check convergence
        abs(E_total_previous - E_total) < 0.1 && break
        E_total_previous = E_total
    end
end