using Interpolations, PolyLog, Plots

"Defines a nuclear system containing relevant parameters and meson/nucleon densities"
mutable struct system
    Z::Int
    N::Int

    r_max::Float64
    divs::Int

    Φ0::Vector{Float64}
    W0::Vector{Float64}
    B0::Vector{Float64}
    A0::Vector{Float64}

    ρ_sp::Vector{Float64}
    ρ_vp::Vector{Float64}
    ρ_sn::Vector{Float64}
    ρ_vn::Vector{Float64}

    "Initialize an unsolved system"
    system(Z, N, r_max, divs) = new(Z, N, r_max, divs, [zeros(1 + divs) for _ in 1:8]...)

    "Dummy struct to define the mesh"
    system(r_max, divs) = system(0, 0, r_max, divs)
end

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

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

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

"Get the number of protons or neutrons in the system"
Z_or_N(s::system, p::Bool)::Int = p ? s.Z : s.N

"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)))

"Create linear interpolations for the meson fields"
function fields_interp(s::system)
    S_interp = linear_interpolation(rs(s), s.Φ0)
    V_interp = linear_interpolation(rs(s), s.W0)
    R_interp = linear_interpolation(rs(s), s.B0)
    A_interp = linear_interpolation(rs(s), s.A0)
    return (S_interp, V_interp, R_interp, A_interp)
end

include("nucleons.jl")
include("mesons.jl")

"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

    E_total_previous = NaN

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

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

        # neutrons
        @time "Neutron spectrum" (κs_n, Es_n) = findAllOrbitals(false, s)
        occs_n = fillNucleons(s.N, κs_n, Es_n)
        @time "Neutron densities" (s.ρ_sn, s.ρ_vn) = calculateNucleonDensity(κs_n, Es_n, occs_n, 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

        # 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