mixing_graphs.jl

using Dictionaries
"""
A type that stores a pair of nodes representing an edge in the graph. 

We need to define a custom type for these so we can define a hash function on graph edges, in order to more efficiently use them in hashmaps (dictionaries)
"""
struct GraphEdge
    a::Int
    b::Int
end

"""
Define a hash on GraphEdge such that ``hash(a,b) = hash(b,a)`` (hash is commutative).

This is helpful because then we only need to store (a,b) in the graph edges weights dictionary, rather than both (a,b) and (b,a).
"""
function Base.hash(e::GraphEdge)
    return hash(minmax(e.a,e.b))
end

"""
Define symmetric edge equality, matches the hash function.
"""
function Base.isequal(e1::GraphEdge,e2::GraphEdge)
    return isequal(minmax(e1.a,e1.b),minmax(e2.a,e2.b))
end

"""
    sample_mixing_graph!(mixing_graph::Graph)


Resample all the weights in `mixing_graph`
"""
function sample_mixing_graph!(mixing_graph)
    mixing_edges = mixing_graph.mixing_edges
    # display(length.(mixing_edges.contact_array))
    # display(length.(mixing_edges.sample_cache))
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal     
        rand!(Random.default_rng(Threads.threadid()), mixing_edges.sampler_matrix[j,i],mixing_edges.sample_cache[j,i])
        for k in 1:length(mixing_edges.contact_array[j,i])
            edge_weight_k = mixing_edges.sample_cache[j,i][k]
            set!(mixing_edges.weights_dict, mixing_edges.contact_array[j,i][k], edge_weight_k)
        end
    end
end

"""
A type that defines the graph edges such that they can be easily resampled and then re-added to the graph. Created with `create_mixing_edges(..)`.

#Fields

    total_edges::Int

Total number of edges in the struct.    

    contact_array::Matrix{Tuple{Vector{Int},Vector{Int}}}

A matrix of vector pairs, such that one node is in the first vector and the other is in the second. The position of the vector pairs in the matrix corresponds to the distribution in sampler_matrix used to sample them. There is probably a better way to represent these, I did this so they could be sampled fast. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.

    sample_cache::Matrix{Vector{Int}}

The structure that pre-allocates the space for the weights of the edges in contact array. This is filled with weights and then the weights are added to weights_dict. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.

    weights_dict::Dictionary{GraphEdge,UInt8}

Stores the weights used in the graph, so they can be easily resampled.    

    sampler_matrix::M

This is the matrix of distributions, from which the edge weights are sampled. Specifically, weights for edges in `contact_array[i,j]` come from the distribution in `sampler_matrix[i,j]`, and are placed into `sample_cache[i,j]`.  We only use the upper triangle of this but Julia lacks a good Symmetric matrix type. See `sample_mixing_graph!`.
"""
struct MixingEdges{M}
    total_edges::Int
    contact_array::Matrix{Vector{GraphEdge}}
    sample_cache::Matrix{Vector{Int}}
    weights_dict::Dictionary{GraphEdge,UInt8}
    sampler_matrix::M
    function MixingEdges(total_edges,contact_array,sampler_matrix::Matrix{M}) where M<:Sampleable{Univariate, Discrete}
        sample_cache = map(v-> Vector{Int}(undef,length(v)),contact_array)
        weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = total_edges)
        new{typeof(sampler_matrix)}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix)
    end
end

"""
This function constructs the `MixingEdges` type for rest and WS graphs. Calls `generate_contact_vectors!(..)` to get the degree distributions for a bipartite graph for each pair of demographics, and then adds edges to MixingEdges.contact_array with the Chung-Lu approach. 
"""
function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
    contact_array =[Vector{GraphEdge}() for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
    tot = 0
    for i in 1:size(mixing_matrix)[1], j in 1:i  #diagonal
        num_degrees_ij = zeros(Int,length(demographic_index_vectors[i]))
        num_degrees_ji = zeros(Int,length(demographic_index_vectors[j]))
        generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],num_degrees_ij,num_degrees_ji)

        m = sum(num_degrees_ij)
        stubs_i = Vector{Int}(undef,m)
        stubs_j = Vector{Int}(undef,m)
        if m>0
            sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[i],Weights(num_degrees_ij./m),stubs_i)
            sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[j],Weights(num_degrees_ji./m),stubs_j)
            tot += m
        end    

        contact_array[j,i] = GraphEdge.(stubs_i,stubs_j)
    end
    return MixingEdges(tot,contact_array,weights_distribution_matrix)
end


"""
This function constructs the `MixingEdges` type for the home graphs. Simply adds edges to MixingEdges from the graph `g`, since we already create that in `generate_population`.
"""
function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix) 
    contact_array = [sizehint!(Vector{GraphEdge}(),ne(g)) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
    for e in edges(g)
        i = src(e)
        j = dst(e)
        push!(contact_array[Int(demographics[i]),Int(demographics[j])], GraphEdge(i,j))
    end
    sample_cache = map(v-> Vector{Int}(undef,length(v)),contact_array)
    return MixingEdges(ne(g),contact_array,weights_distribution_matrix)
end

"""

Stores the full time dependent mixing graph for the model. I think this might be a weird abstraction for this idea but it works fine.

#Fields 

These are references to the graphs that get resampled everyday.

    resampled_graphs::NTuple{N,G}

List of lists of graphs, one list for each day. 

    graph_list::Vector{Vector{G}}

"""
struct TimeDepMixingGraph{N,G,GNT}
    resampled_graphs::NTuple{N,GNT}
    graph_list::Vector{Vector{G}}
    function TimeDepMixingGraph(len,resampled_graphs::NTuple{N,GNT},base_graph_list::Vector{G}) where {GNT,G,N}
        return new{N,G,GNT}(
            resampled_graphs,
            [copy(base_graph_list) for i in 1:len]
        )
    end
end

"""
Creates the `TimeDepMixingGraph` for our specific model. 

Assumes the simulation begins on Thursday arbitrarily.
"""
function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
    home_static_edges = WeightedGraph(base_network,demographics,index_vectors,contact_time_distributions.hh) #network with households and LTC homes
   
    ws_static_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.daily,contact_time_distributions.ws)
    ws_weekly_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.twice_a_week,contact_time_distributions.ws)
    ws_daily_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.otherwise,contact_time_distributions.ws)
    
    rest_static_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.daily,contact_time_distributions.rest)
    rest_weekly_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.twice_a_week,contact_time_distributions.rest)
    rest_daily_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.otherwise,contact_time_distributions.rest)
   
    inf_network_list = [home_static_edges,rest_static_edges] 
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
    
    infected_mixing_graph = TimeDepMixingGraph(len,((ws_daily_edges,ws_matrix_tuple.daily),(rest_daily_edges,rest_matrix_tuple.daily)),inf_network_list)
    soc_mixing_graph = TimeDepMixingGraph(len,((ws_daily_edges,ws_matrix_tuple.daily),(rest_daily_edges,rest_matrix_tuple.daily)),soc_network_list)
    # display(infected_mixing_graph.graph_list)
    for (t,l) in enumerate(infected_mixing_graph.graph_list)
        day_of_week = mod(t,7)
        if !(day_of_week == 3 || day_of_week == 4) #simulation begins on thursday I guess
            push!(l, ws_static_edges)
        end
        if rand(Random.default_rng(Threads.threadid()))<5/7
            push!(l, ws_weekly_edges)
            push!(l, rest_weekly_edges)
        end
        push!(l,ws_daily_edges)
        push!(l,rest_daily_edges)
    end

    return infected_mixing_graph,soc_mixing_graph
end

"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
function remake!(time_dep_mixing_graph,demographic_index_vectors)
    for (weighted_graph,mixing_matrix) in time_dep_mixing_graph.resampled_graphs
        empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
        weighted_graph.g.ne = 0
        weighted_graph.mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weighted_graph.mixing_edges.sampler_matrix)
        graph_from_mixing_edges(weighted_graph.g,weighted_graph.mixing_edges)
    end
end

"""
Add the edges defined by MixingEdges to the actual graph G. Another big bottleneck since adjancency lists don't add edges super efficiently, and there are a ton of them. 
"""
function graph_from_mixing_edges(g,mixing_edges)
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal
        for k in 1:length(mixing_edges.contact_array[i,j])
            e = mixing_edges.contact_array[i,j][k]
            add_edge!(g,e.a,e.b)
        end
    end
end

"""
My own weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`. 

"""
mutable struct WeightedGraph{G,M} 
    g::G
    mixing_edges::M
    function WeightedGraph(demographics,demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
        mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
        g = Graph(length(demographics))
        graph_from_mixing_edges(g,mixing_edges)
        return new{typeof(g),typeof(mixing_edges)}(
            g,
            mixing_edges
        )
    end
    function WeightedGraph(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix)
        mixing_edges = create_mixing_edges(g, demographics, demographic_index_vectors,weights_distribution_matrix)
        weights_dict = Dict{GraphEdge,UInt8}()
        sample_cache = map(v-> Vector{Int}(undef,length(v)),mixing_edges.contact_array)
        return new{typeof(g),typeof(mixing_edges)}(
            g,
            mixing_edges,
        )
    end
end
neighbors(g::WeightedGraph,i) = neighbors(g.g,i)
get_weight(g::WeightedGraph,e) = g.mixing_edges.weights_dict[e]
function Base.show(io::IO, g::WeightedGraph) 
    print(io, "WG $(ne(g.g))")
end