solve.jl

#remove Base.@propagate_inbounds during tests if you get segfaults or change anything,
# as that might mean there is an indexing bug that is not being caught,
#and therefore accessing the wrong memory is causing background errors

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

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

Base.@propagate_inbounds @views function update_alert_durations!(t,modelsol) # Base.@propagate_inbounds 
    @unpack notification_parameter,notification_threshold = modelsol.params
    @unpack time_of_last_alert, app_user_index,inf_network,covid_alert_notifications,app_user = modelsol
    for (i,node) in enumerate(modelsol.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.g,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
            time_of_last_alert[i] = t
        end
    end
end

Base.@propagate_inbounds @views function update_infection_state!(t,modelsol)
    @unpack base_transmission_probability,immunization_loss_prob,recovery_rate,immunizing,immunization_begin_day = modelsol.params
    @unpack u_inf,u_vac,u_next_inf,u_next_vac,demographics,inf_network,status_totals, immunization_countdown = modelsol
    
    modelsol.daily_infected = 0
    modelsol.daily_immunizations_by_age.= 0
    function agent_transition!(node, from::AgentStatus,to::AgentStatus) 
        immunization_countdown[node] = -1 
        status_totals[Int(from)] -= 1
        status_totals[Int(to)] += 1
        u_next_inf[node] = to
    end
    u_next_inf .= u_inf
    for i in 1:modelsol.nodes
        agent_status = u_inf[i]
        is_vaccinator = u_vac[i]
        agent_demo = demographics[i]
        if agent_status == Susceptible
            if is_vaccinator && immunizing && immunization_countdown[i] == -1 && t> immunization_begin_day
                immunization_countdown[i] = 14
            else
                for mixing_graph in inf_network.graph_list[t]
                    for j in neighbors(mixing_graph.g,i)    
                        if u_inf[j] == Infected && u_next_inf[i] != Infected 
                            if rand(Random.default_rng(Threads.threadid())) < contact_weight(base_transmission_probability,get_weight(mixing_graph,GraphEdge(i,j)))
                                modelsol.daily_infected+=1
                                agent_transition!(i, Susceptible,Infected)
                            end
                        end
                    end
                end
            end
        elseif agent_status == Infected
            if rand(Random.default_rng(Threads.threadid())) < recovery_rate
               agent_transition!(i, Infected,Recovered)
            end
        elseif agent_status == Immunized
            if rand(Random.default_rng(Threads.threadid())) < immunization_loss_prob
               agent_transition!(i, Immunized,Susceptible)
            end
        end
        
        if immunization_countdown[i] == 0 
            modelsol.daily_immunizations_by_age[Int(agent_demo)] += 1
            agent_transition!(i, Susceptible,Immunized)
        elseif immunization_countdown[i]>0
            immunization_countdown[i] -= 1
        end
    end
end
Base.@propagate_inbounds @views function update_vaccination_opinion_state!(t,modelsol,total_infections)
    @unpack π_base, η,γ, κ, ω, ρ_y,ρ_m,ρ_o, ω_en,ρ_en,γ,β = modelsol.params
    @unpack demographics,time_of_last_alert, nodes, soc_network,u_vac,u_next_vac,app_user,app_user_list = modelsol
    app_user_pointer = 0
    for i in 1:nodes
        vac_payoff = 0
        soc_nbrs_vac = @MArray [0,0,0]
        soc_nbrs_nonvac = 0
        num_soc_nbrs = 0
        for sc_g in soc_network.graph_list[t]
            soc_nbrs = neighbors(sc_g.g,i)
            num_soc_nbrs += length(soc_nbrs)
            for nbr in soc_nbrs
                if u_vac[nbr]
                    soc_nbrs_vac[Int(demographics[nbr])] += 1
                else
                    soc_nbrs_nonvac += 1
                end
            end
        end
        ρ = @SVector [ρ_y,ρ_m,ρ_o]

        vac_payoff += π_base + dot(ρ,soc_nbrs_vac) + total_infections*ω +
            ifelse(num_soc_nbrs> 0, κ * ((sum(soc_nbrs_vac) - soc_nbrs_nonvac)/num_soc_nbrs),0)
        
        if app_user[i] && time_of_last_alert[app_user_list[i]]>=0
            vac_payoff += γ^(-1*(t - time_of_last_alert[app_user_list[i]]))* (η + dot(ρ_en,soc_nbrs_vac) + total_infections*ω_en) 
        end
        
        if u_vac[i]
            if rand(Random.default_rng(Threads.threadid())) < 1 - Φ(vac_payoff,β)
                u_next_vac[i] = !u_vac[i]
            else
                u_next_vac[i] = u_vac[i]
            end
        else
            if rand(Random.default_rng(Threads.threadid())) < Φ(vac_payoff,β)
                u_next_vac[i] = !u_vac[i]
            else
                u_next_vac[i] = u_vac[i]
            end
        end
    end

    modelsol.daily_vaccinators = count(==(true),u_vac) #could maybe make this more efficient
end


Base.@propagate_inbounds @views function update_vaccination_opinion_state_2!(t,modelsol,total_infections)
    @unpack π_base, η,γ, κ, ω, ρ_y,ρ_m,ρ_o, ω_en,ρ_en,γ,β = modelsol.params
    @unpack demographics,time_of_last_alert, nodes, soc_network,u_vac,u_next_vac,app_user,app_user_list = modelsol
    app_user_pointer = 0
    for i in 1:nodes
        vac_payoff = 0
        soc_nbrs = @MArray [0,0,0]
        num_soc_nbrs = 0
        for sc_g in soc_network.graph_list[t]
            soc_nbrs = neighbors(sc_g.g,i)
            num_soc_nbrs += length(soc_nbrs)
            for nbr in soc_nbrs
                soc_nbrs[Int(demographics[nbr])] += 1
            end
        end
        ρ = @SVector [ρ_y,ρ_m,ρ_o]
        # random_neighbour = 

        
            vac_payoff += π_base + dot(ρ,soc_nbrs_vac) + total_infections*ω +
                ifelse(num_soc_nbrs> 0, κ * ((sum(soc_nbrs_vac) - soc_nbrs_nonvac)/num_soc_nbrs),0)
            
            if app_user[i] && time_of_last_alert[app_user_list[i]]>=0
                vac_payoff += γ^(-1*(t - time_of_last_alert[app_user_list[i]]))* (η + dot(ρ_en,soc_nbrs_vac) + total_infections*ω_en) 
            end
        
        if u_vac[i]
            if rand(Random.default_rng(Threads.threadid())) < 1 - Φ(vac_payoff,β)
                u_next_vac[i] = !u_vac[i]
            else
                u_next_vac[i] = u_vac[i]
            end
        else
            if rand(Random.default_rng(Threads.threadid())) < Φ(vac_payoff,β)
                u_next_vac[i] = !u_vac[i]
            else
                u_next_vac[i] = u_vac[i]
            end
        end
    end

    modelsol.daily_vaccinators = count(==(true),u_vac) #could maybe make this more efficient
end


function agents_step!(t,modelsol)

    remake!(modelsol.inf_network,modelsol.index_vectors,modelsol.ws_matrix_tuple.daily)
    remake!(modelsol.soc_network,modelsol.index_vectors,modelsol.rest_matrix_tuple.daily)
    for network in modelsol.inf_network.graph_list[t]
        sample_mixing_graph!(network) #get new contact weights
    end
    if t == modelsol.params.infection_introduction_day
        init_indices = rand(Random.default_rng(Threads.threadid()), 1:modelsol.nodes, round(Int,modelsol.nodes*modelsol.params.I_0_fraction))
        modelsol.u_inf[init_indices] .= Infected
        modelsol.status_totals[Int(Infected)] += length(init_indices)
    end
    if t>modelsol.params.infection_introduction_day
        update_alert_durations!(t,modelsol)
    end
    update_vaccination_opinion_state!(t,modelsol,modelsol.status_totals[Int(Infected)])
    update_infection_state!(t,modelsol)

    modelsol.u_vac .= modelsol.u_next_vac
    modelsol.u_inf .= modelsol.u_next_inf
end



function solve!(modelsol,recordings...)

    for t in 1:modelsol.sim_length
        agents_step!(t,modelsol) 
        #advance agent states based on the new network
        for recording in recordings
            record!(t,modelsol,recording)
        end
    end
end