agents.jl

#agent type definitions

"""
Define the possible agent demographic groups. Note that these can be converted to/from integers (starting at 1).
"""
@enum AgentDemographic begin
    Young = 1
    Adult
    Elderly
end

"""
Define the possible agent infection statuses. Note that these can be converted to/from integers (starting at 1).
"""
@enum AgentStatus begin
    Susceptible = 1
    Infected
    Recovered
    Immunized
end

"""

    complete_graph_from_households_composition(households_composition::AbstractVector{AbstractVector{Int}})

Generates a complete graph from a vector of household compositions. Each household composition is a 3 element vectors (one for each demographic group) of integers where each element describes the number of the corresponding demographic group present in that household.

This function wires together a graph such that each household is in a complete subgraph. It is much faster than the `mapreduce` solution used previously.
"""
function complete_graph_from_households_composition(households_composition)
    total_household_pop = sum(sum.(households_composition))
    network = SimpleGraph(total_household_pop)
    vertex_pointer = 1
    for household in households_composition
        num_vertices = sum(household)
        for v in vertex_pointer:(vertex_pointer + num_vertices - 1)
            for w in vertex_pointer:(vertex_pointer + num_vertices - 1)
                if v != w
                    add_edge!(network,v,w)
                end
            end
        end
        vertex_pointer+=num_vertices
    end
    return network
end


"""
    generate_population(num_households::Int)

Returns a namedtuple with fields `population_list`, `household_networks` and `index_vectors`.

population_list: Vector of AgentDemographic where each entry is the demographic of a person in the simulation. The length is equal to the total number of people in the model.

household_networks: SimpleGraph with all agents such that each household is in a complete subgraph.

index_vectors: Vector of 3 vectors, each of which contains the indexes of all the agents in a particular demographic. If `population_list` is a mapping from agents to demographics, these are mappings from demographics to agents.
"""
function generate_population(num_households)
    households_composition = sample_household_data(num_households)

    household_networks = complete_graph_from_households_composition(households_composition)
    
    #test the complete_graph_from_households_composition function against the mapreduce version to ensure it is correct.
    @c_assert all(adjacency_matrix(household_networks) .== adjacency_matrix( mapreduce(LightGraphs.complete_graph,blockdiag, sum.(households_composition); init = SimpleGraph())))

    households = map( l -> [fill(AgentDemographic(i),n) for (i,n) in enumerate(l)],households_composition) |>
    l -> reduce(vcat,l) |>
    l -> reduce(vcat,l)
    population_list = reduce(vcat,households)
    index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:(AgentDemographic.size-1)]

    
    return (;
        population_list,
        household_networks,
        index_vectors
    )
end


"""
Defines pretty printing methods so that `display(AgentStatus)` is more readable. 
"""
function Base.show(io::IO, status::AgentStatus) 
    if status == Susceptible
        print(io, "S")
    elseif status == Infected
        print(io, "I")
    elseif status == Recovered
        print(io, "R")
    elseif status == Immunized
        print(io, "Vac")
    end
end

"""
Defines pretty printing methods so that `display(AgentDemographic)` is more readable. 
"""
function Base.show(io::IO, status::AgentDemographic) 
    if status == Young
        print(io, "<25")
    elseif status == Adult
        print(io, "20-65")
    elseif status == Elderly
        print(io, ">65")
    end
end