# © 2022 Massachusetts Institute of Technology.  See LICENSE file for details.

"""
    GlobalParameters(βE::T, βI::T, γ::T, λ::T)
Structure for global scalar SEIR parameters.
"""
@with_kw mutable struct GlobalParameters{T<:Real} <: SEIRParameters{T}
    βE::T
    βI::T
    γ::T
    λ::T
end

GlobalParameters(βE::Real, βI::Real, γ::Real, λ::Real) = GlobalParameters(βE, βI, γ, λ)

"""
    CommunityParams(βE::Vector{T} βI::Vector{T}, γ::T, λ::T, c::Int) where T <: Real
Structure for individual community-based `SEIRParameters` constructor given vector β values and `c` communities. 

Can be constructed with scalar β values to set all communities to the same β values.
"""
@with_kw mutable struct CommunityParameters{T<:Real} <: SEIRParameters{T}
    βE::Vector{T}
    βI::Vector{T}
    γ::T
    λ::T
end

CommunityParameters(βE::Real, βI::Real, γ::Real, λ::Real, c::Int) = CommunityParameters(fill(βE, c), fill(βI, c), γ, λ)

"""
    CommunityPairParameters(βE::Matrix{T}, βI::Matrix{T}, γ::T, λ::T)
Structure for community-pair based SEIR parameters.

Can be constructed with scalar β values to set all community-pair β values.
"""
@with_kw mutable struct CommunityPairParameters{T<:Real} <: SEIRParameters{T}
    βE::Matrix{T}
    βI::Matrix{T}
    γ::T
    λ::T
end

CommunityPairParameters(βE::Real, βI::Real, γ::Real, λ::Real, c::Int) = CommunityPairParameters(fill(βE, (c, c)), fill(βI, (c, c)), γ, λ)


"""
    R_i(g, β_e, vs = []; γ = 0.2)
Compute expected number of infections generated by each vertex for specified β and γ values.
"""
function R_i(g, β_e, vs=[]; γ=0.2)
    R0 = zeros(length(vs))
    if isempty(vs)
        vs = vertices(g)
    end
    for (ind, i) in enumerate(vs)
        for j in neighbors(g, i)
            wij = get_prop(g, i, j, :weight)
            R0[ind] += ((1 - exp(-β_e * wij)) / (1 - ((exp(-β_e * wij)) * (1 - γ))))
        end
    end
    return R0
end

"""
    compute_R0(g, β_e; α=1.0, γ=0.2, filter_household = false, vweights = nothing)
Approximate the reproduction rate of a disease for a given graph `g`, infection parameters β and γ, and top proportion of nodes by degree α. `vweights` can be supplied if pre-computed.
"""
function compute_R0(g, β_e; α=1.0, γ=0.2, filter_household=false, vweights=nothing)
    if isnothing(vweights)
        vweights = [vertex_weight(g, v, filter_household) for v in vertices(g)]
    end
    ind = ceil(Int, α * nv(g))
    vs = partialsortperm(vweights, 1:ind, rev=true)
    r_i = R_i(g, β_e, vs; γ=γ)
    retval = mean(r_i)
    return retval
end

function betafromR0(g, R0tgt; α=1.0, γ=0.2, filter_household=false, range=[-5, 0])
    vweights = [vertex_weight(g, v, filter_household) for v in vertices(g)]
    fun(log10betaE) = compute_R0(g, 10.0^log10betaE, α=α, γ=γ, filter_household=filter_household, vweights=vweights) - R0tgt
    beta = 10.0 .^ find_zero(fun, (range[1], range[2]), Bisection(), xatol=1e-3)
    return beta
end

"""
    get_betas(g, R0; α=1.0, γ=0.2, λ=0.2, p = 1/5, filter_household = false, range = [-5,0])
Fit `β_e` and `β_i` parameters for a graph `g` with disease reproduction rate `R0`. `range` is the upper and lower bound as an order of magnitude for fitting. 
"""
function get_betas(g, R0; α=1.0, γ=0.2, λ=0.2, p=1 / 5, filter_household=false, range=[-5, 0])
    R0_E = R0
    R0_I = R0 * p
    beta_E = betafromR0(g, R0_E; α=α, γ=γ, filter_household=filter_household, range=range)
    beta_I = betafromR0(g, R0_I; α=α, γ=λ, filter_household=filter_household, range=range)
    return (beta_E, beta_I)
end
