using Distributions
using Random

function sample_infection_time(;omicron=false)
    """
    Sample infection times for num individuals based on the generation 
    interval distribution, Gamma(shape_gen, scale_gen). 
    """

    (shape_gen, scale_gen) = (2.75, 1.00)
    (shape_gen_omicron, scale_gen_omicron) = (1.58, 1.32)
    
    shape = (1-omicron)*shape_gen + omicron*shape_gen_omicron
    scale = (1-omicron)*scale_gen + omicron*scale_gen_omicron
    
    infection_time = ceil(Int, rand(Gamma(shape, scale)))
    
    return infection_time
end

function sample_onset_time(;omicron=false)
    """
    Sample incubation times for num individuals based on incubation period 
    distribution, Gamma(shape_inc, scale_inc). 
    """
    
    (shape_inc, scale_inc) = (5.807, 0.948)
    (shape_inc_omicron, scale_inc_omicron) = (3.33, 1.34)
    
    shape = (1-omicron)*shape_inc + omicron*shape_inc_omicron
    scale = (1-omicron)*scale_inc + omicron*scale_inc_omicron
    
    onset_time = ceil(Int, rand(Gamma(shape, scale)))
    
    return onset_time
end


function set_simulation_constants(state)
    """
    Contains the assumptions for simulation parameters. This includes all the dynamical constants: 
        - k = heterogeneity parameter
        - p_symp = probability of symptoms 
        - γ = relative infectiousness of asymptomatic individuals 
        - p_symp_given_detect = probability of symptoms given detection
        - p_asymp_given_detect = probability of being asymptomatic given detection
        - consistency_multiplier = chosen such that sim_cases < consistency_multiplier*actual cases 
            results in cases being injected into the simulation. This is used to account 
            for superspreading events after periods of low incidence. 
    These values are stored in sim_constants which is a dictionary indexed by the parameter name and 
    ultimately stored on the stack in a SimulationParameters object. 
    """

    # overdispersion parameter 
    k = 0.15
    # assumptions surrouding the probability of symptomatic, 
    # relative infectiousness γ and the ratio of Reff (α's) 
    p_symp = 0.4
    
    p_detect_given_symp_dict = Dict{String, Float64}(
        "NSW" => 0.7,
        "QLD" => 0.7,
        "SA" => 0.7,
        "TAS" => 0.7,
        "WA" => 0.7,
        "ACT" => 0.7,
        "NT" => 0.7,
        "VIC" => 0.7,
    )
    
    p_detect_given_asymp_dict = Dict{String, Float64}(
        "NSW" => 0.467,
        "QLD" => 0.467,
        "SA" => 0.467,
        "TAS" => 0.467,
        "WA" => 0.467,
        "ACT" => 0.467,
        "NT" => 0.467,
        "VIC" => 0.467,
    )
    
    γ = 0.5     # relative infectiousness of asymptomatic
    # solve the system α_s*ps + α_a(1-ps) = 1 with α_a = γ*α_s
    p_detect_given_symp = p_detect_given_symp_dict[state]
    p_detect_given_asymp = p_detect_given_asymp_dict[state]
    # prob of detection
    p_detect = p_symp*p_detect_given_symp + (1-p_symp)*p_detect_given_asymp
    # prob symptomatic given detect
    p_symp_given_detect = p_detect_given_symp*p_symp/p_detect 
    
    α_s = 1/(p_symp + γ*(1-p_symp))
    α_a = γ * α_s
    consistency_multiplier = 5.0

    # prob of detecting an international import 
    qi = 0.98
    # prior parametes for the import model
    prior_alpha = 0.5
    prior_beta = 0.2
    # ema smoothing factor 
    ϕ = 0.1

    simulation_constants = Constants(
        k,
        p_symp,
        p_detect_given_symp,
        p_detect_given_asymp,
        p_detect, 
        p_symp_given_detect,
        γ,
        α_s,
        α_a,
        qi,
        prior_alpha,
        prior_beta,
        ϕ,
        consistency_multiplier,
    )

    # mapping between types 
    individual_type_map = IndividualTypeMap(1, 2, 3)

    return (simulation_constants, individual_type_map)
end
