using FastGaussQuadrature

"Tabulates a nucleon spectrum (protons or neutrons) containing κ and occupancy"
struct spectrum
    κ::Vector{Int}
    E::Vector{Float64}
    occ::Vector{Int}
end

"Initializes an unfilled spectrum"
unfilled_spectrum() = spectrum(Int[], Float64[], Int[])

"Defines a mesh"
struct mesh
    r_max::Float64
    r::Vector{Float64}
    w::Vector{Float64}
end

"Create a uniform mesh"
function uniform_mesh(r_max::Float64, divs::Int)::mesh
    r = range(0, r_max, length=divs+1) |> collect
    w = fill(r_max / divs, divs+1)
    return mesh(r_max, r, w)
end

"Create a Gauss-Legendre mesh"
function gauleg_mesh(r_max::Float64, n::Int)::mesh
    a, b = (0.0, r_max)
    nodes, weights = gausslegendre(n)
    transformed_nodes = (b - a) / 2 .* nodes .+ (b + a) / 2
    transformed_weights = (b - a) / 2 .* weights
    return mesh(r_max, transformed_nodes, transformed_weights)
end

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

    r_mesh::mesh

    p_spectrum::spectrum
    n_spectrum::spectrum

    Φ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 with a uniform r-mesh"
    system(Z::Int, N::Int, r_max::Float64, divs::Int) = new(Z, N, uniform_mesh(r_max, divs), unfilled_spectrum(), unfilled_spectrum(), [zeros(1 + divs) for _ in 1:8]...)

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

    "Initialize an unsolved system with a Gauss-Legendre r-mesh"
    system(Z::Int, N::Int, mesh_size::Int, r_max::Float64) = new(Z, N, gauleg_mesh(r_max, mesh_size), unfilled_spectrum(), unfilled_spectrum(), [zeros(mesh_size) for _ in 1:8]...)
end

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

"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(length(s.r_mesh.r))