solve.jl

function contact_weight(β, contact_time) 
    return 1 - (1-β)^contact_time
end

function Φ(payoff,ξ)
    return 1 / (1 + exp(-1*ξ*payoff))
end

Base.@propagate_inbounds @views function update_alert_durations!(t,modelsol) 
    @unpack notification_parameter,notification_threshold = modelsol.params
    @unpack time_of_last_alert, app_user_index,inf_network,covid_alert_notifications,app_user, output_data = modelsol
    
    for (i,node) in enumerate(app_user_index)
        for j in 2:14
            covid_alert_notifications[j-1,i] = covid_alert_notifications[j,i] #shift them all back  
        end
        total_weight_i = 0
        for mixing_graph in inf_network.graph_list[t]
            for j in neighbors(mixing_graph,node)
                if app_user[j]
                    total_weight_i+= get_weight(mixing_graph,GraphEdge(node,j))
                end
            end
        end
        coin_flip = 1 - (1 - notification_parameter)^total_weight_i
        r = rand(Random.default_rng(Threads.threadid()))
        if r < coin_flip
            covid_alert_notifications[end,i] = 1  #add the notifications for today
        else
            covid_alert_notifications[end,i] = 0
        end
        if sum(covid_alert_notifications[:,i])>=notification_threshold
            output_data.daily_total_notifications[t] += 1
            time_of_last_alert[i] = t
        end
    end
end

function agent_transition!(modelsol,agent, from::AgentStatus,to::AgentStatus) 
    @unpack u_next_inf,status_totals, immunization_countdown = modelsol
    immunization_countdown[agent] = -1 
    status_totals[Int(from)] -= 1
    status_totals[Int(to)] += 1
    u_next_inf[agent] = to
end

function infect_agent!(t, modelsol, agent,agent_inf_status,agent_demo)
    @unpack β_y,β_m,β_o,α_y,α_m,α_o = modelsol.params
    @unpack u_inf,inf_network, output_data = modelsol
 
    β_vec = @SVector [β_y,β_m,β_o]
    α_vec = @SVector [α_y,α_m,α_o]

    for mixing_graph in inf_network.graph_list[t]
        for neighbor in neighbors(mixing_graph,agent)    
            if u_inf[neighbor] == Infected
                infection_threshold = contact_weight(β_vec[Int(agent_demo)],get_weight(mixing_graph,GraphEdge(agent,neighbor)))

                if rand(Random.default_rng(Threads.threadid())) < infection_threshold
                    if agent_inf_status == Immunized
                        immunization_ineffective = rand(Random.default_rng(Threads.threadid())) < 1- α_vec[Int(agent_demo)]
                        if immunization_ineffective
                            agent_transition!(modelsol,agent, Immunized,Infected)
                            output_data.daily_cases_by_age[Int(agent_demo),t]+=1
                            return true
                        end
                    elseif agent_inf_status == Susceptible
                        output_data.daily_cases_by_age[Int(agent_demo),t]+=1
                        output_data.daily_unvac_cases_by_age[Int(agent_demo),t]+=1
                        agent_transition!(modelsol, agent, Susceptible,Infected)
                        return true
                    end
                end
            end
        end
    end
    return false
end


Base.@propagate_inbounds @views function update_infection_state!(t,modelsol; record_degrees = false)
    @unpack recovery_rate,immunizing,immunization_begin_day = modelsol.params
    @unpack u_inf,u_vac,u_next_inf,demographics,inf_network, immunization_countdown, output_data = modelsol
    
    u_next_inf .= u_inf
    
    for (agent,(agent_inf_status,agent_vac_status,agent_demo)) in enumerate(zip(u_inf,u_vac,demographics))
        if agent_inf_status == Susceptible && agent_vac_status && immunizing && immunization_countdown[agent] == -1 && t> immunization_begin_day
            immunization_countdown[agent] = 14
        end
        if  agent_inf_status == Susceptible || agent_inf_status == Immunized 
            infect_agent!(t, modelsol, agent,agent_inf_status,agent_demo)
        elseif agent_inf_status == Infected
            if rand(Random.default_rng(Threads.threadid())) < recovery_rate
               agent_transition!(modelsol, agent, Infected,Recovered)
            end
        end
        weighted_degree_of_i::Int = output_data.record_degrees_flag ? weighted_degree(t,agent,inf_network) : 0
        if immunization_countdown[agent] == 0 
            output_data.daily_immunized_by_age[Int(agent_demo),t] += 1
            fit!(output_data.avg_weighted_degree_of_vaccinators[Int(agent_demo)],weighted_degree_of_i)
            agent_transition!(modelsol, agent, Susceptible,Immunized)
        elseif immunization_countdown[agent]>0
            fit!(output_data.avg_weighted_degree[Int(agent_demo)],weighted_degree_of_i)
            immunization_countdown[agent] -= 1
        else
            fit!(output_data.avg_weighted_degree[Int(agent_demo)],weighted_degree_of_i)
        end
    end
end


Base.@propagate_inbounds @views function update_vaccination_opinion_state!(t,modelsol,total_infections)
    @unpack infection_introduction_day, π_base_y,π_base_m,π_base_o, η,Γ,ζ, ω, ω_en,ξ = modelsol.params
    @unpack demographics,time_of_last_alert, nodes, soc_network,u_vac,u_next_vac,app_user,app_user_list,output_data = modelsol
   
    for i in 1:nodes
        π_base = t<infection_introduction_day ? 
        (π_base_y,π_base_m,π_base_o) : 
        (π_base_y*ζ,π_base_m*ζ,π_base_o*ζ)
        random_soc_network = sample(Random.default_rng(Threads.threadid()), soc_network.graph_list[t])
        if !isempty(neighbors(random_soc_network,i))
                random_neighbour = sample(Random.default_rng(Threads.threadid()), neighbors(random_soc_network.g,i))
                app_vac_payoff = 0.0    
                if app_user[i] && time_of_last_alert[app_user_list[i]]>=0
                    output_data.daily_total_notified_agents[t] += 1
                    app_vac_payoff = Γ^((t - time_of_last_alert[app_user_list[i]])) * (η + total_infections*ω_en) 
                    # display(t - time_of_last_alert[app_user_list[i]])
                end
                if u_vac[random_neighbour] == u_vac[i]
                    vac_payoff = π_base[Int(demographics[i])] + total_infections*ω + app_vac_payoff
                    if u_vac[i]
                        # display(1 - Φ(vac_payoff,ξ))
                        if rand(Random.default_rng(Threads.threadid())) < 1 - Φ(vac_payoff,ξ)
                            u_next_vac[i] = false
                        end
                    else
                        # display( Φ(vac_payoff,ξ))
                        if rand(Random.default_rng(Threads.threadid())) < Φ(vac_payoff,ξ)
                            u_next_vac[i] = true
                        end
                    end
                end
        end
    end
    output_data.total_vaccinators[t] = count(==(true),u_vac) 
end

function weighted_degree(t,node,network::TimeDepMixingGraph)
    weighted_degree = 0
    for g in network.graph_list[t]
        for j in neighbors(g,node)
            weighted_degree += get_weight(g,GraphEdge(node,j))
        end
    end
    return weighted_degree
end

function sample_initial_nodes(nodes,graphs,I_0_fraction)
    weighted_degrees = zeros(nodes)
    for v in 1:nodes
        for g in graphs
            for w in neighbors(g,v)
                weighted_degrees[v] += get_weight(g,GraphEdge(v,w))
            end
        end
    end
    wv = Weights(weighted_degrees ./sum(weighted_degrees))
    num = round(Int,nodes*I_0_fraction)
    init_indices = sample(Random.default_rng(Threads.threadid()), 1:nodes,wv, num; replace = false)
    return init_indices
end    


function solve!(modelsol)
    init_indices = sample_initial_nodes(modelsol.nodes, modelsol.inf_network.graph_list[begin], modelsol.params.I_0_fraction)
    for t in 1:modelsol.sim_length  
        #this also resamples the soc network weights since they point to the same objects, but those are never used
        
        
        if t>1
            remake_all!(t,modelsol.inf_network,modelsol.index_vectors,modelsol.demographics)
        end

        if t>modelsol.params.infection_introduction_day
            if !isempty(init_indices)
                inf_index = pop!(init_indices)
                modelsol.u_inf[inf_index] = Infected
                modelsol.status_totals[Int(Infected)] += 1
            end
           
        end
        if t>modelsol.params.immunization_begin_day
            update_alert_durations!(t,modelsol)
        end
 
        update_infection_state!(t,modelsol; record_degrees = true)
        update_vaccination_opinion_state!(t,modelsol,modelsol.status_totals[Int(Infected)])
        #advance agent states based on the new network

        
        record!(t,modelsol)

        modelsol.u_vac .= modelsol.u_next_vac
        modelsol.u_inf .= modelsol.u_next_inf
        # display(mean.(modelsol.output_data.avg_weighted_degree))

    end
end