text:
using ..CovidSurvey: RandomWalk

struct WeekdayHolidayObsModel{T,M,A,P,E,W,H} <: ObservationsModel
    θ::T
    μ::M
    α::A
    ϕ::P
	expected::E
    weekdayeffect::W
    holidayeffect::H
end

struct SimpleObsModel{T,M,A,P,E} <: ObservationsModel
    θ::T
    μ::M
    α::A
    ϕ::P
	expected::E
end

struct SimpleSeroObsModel2{T,E}# <: ObservationsModel
    θ::T
	# σ::S
	expected::E
end

# ============================================================================
# random walk model
function random_walks!(Rt, θ, predict, latent_Rt, R0, σ_rt)
	@unpack lockdown_index, num_observations, link = θ #rt_step_indices

	Rt[1:lockdown_index] .= R0
	Rt[lockdown_index+1:num_observations] = link.(latent_Rt) #rt_step_indices
	if predict
		@unpack num_total_days, invlink = θ
		Rt[num_observations+1:num_total_days] .= Rt[num_observations]
		# n  = num_total_days - num_observations
		# n < 1 && return nothing
		# rw = RandomWalk(n, σ_rt, invlink(Rt[num_observations]))
		# Rt[num_observations+1:end] = link.( rand( rw ) )
	end
	return nothing
end

# ============================================================================
# semi-parametric model
function semiparametric!(Rt, θ, predict, latent_Rt, R0, σ_rt, α)
	@unpack lockdown_index, num_observations, link, covariates, covariates_start, num_total_days = θ #rt_step_indices

	Rt[1:lockdown_index] .= R0
	Rt[lockdown_index+1:covariates_start-1] = link.(latent_Rt)
	Rt[covariates_start:num_observations]   = link.(
        latent_Rt[end] .+ covariates * α
    )
	if predict
		Rt[num_observations+1:num_total_days] .= Rt[num_observations]
		# @unpack num_total_days, invlink = θ
		# n  = num_total_days - num_observations
		# n < 1 && return nothing
		# rw = RandomWalk(n, σ_rt, invlink(Rt[num_observations]))
		# Rt[num_observations+1:end] = link.( rand( rw ) )
	end
	return nothing
end
# ============================================================================
# semi-parametric model
function mixed!(Rt, θ, predict, latent_Rt, R0, σ_rt, α)
	@unpack lockdown_index, num_observations, link, covariates, covariates_start, num_total_days = θ #rt_step_indices
	predictor = covariates * α

	Rt[1:lockdown_index] .= R0

	for (j,i) in enumerate(lockdown_index+1:covariates_start-1)
		Rt[i] = link(latent_Rt[j])
	end
	n = covariates_start - lockdown_index - 1
	for (j,i) in enumerate(covariates_start:num_observations)
		Rt[i] = link(latent_Rt[j+n] + predictor[j])
	end
	if predict
		Rt[num_observations+1:num_total_days] .= Rt[num_observations]
		# @unpack num_total_days, invlink = θ
		# n  = num_total_days - num_observations
		# n < 1 && return nothing
		# rw = RandomWalk(n, σ_rt, invlink(Rt[num_observations]))
		# Rt[num_observations+1:end] = link.( rand( rw ) )
	end
	return nothing
end

# ============================================================================
# time varying ascertainment rate

function iar!(iar, θ, predict, latent_iar, iar0, σ_iar)
	@unpack iar_start_idx, num_observations, casemodel = θ
	iar[1:iar_start_idx-1] .= iar0
	iar[iar_start_idx:num_observations] = logistic.( latent_iar )

	if predict
		@unpack num_total_days = θ
		n  = num_total_days - num_observations
		rw = RandomWalk(n, σ_iar, logit(iar[num_observations]))
		iar[num_observations+1:end] = logistic.( rand( rw ) )
	end
	return nothing
end
# ============================================================================
# infection model
function infections!(newly_infected, cumulative_infected, effective_Rt, θ, y, Rt)
	initepidemic!(newly_infected, cumulative_infected, effective_Rt, θ, Rt, y)
	runepidemic!(newly_infected, cumulative_infected, effective_Rt, θ, Rt)
end

function initepidemic!(newly_infected, cumulative_infected, effective_Rt, θ, Rt, y)
	@unpack population, num_impute = θ

	newly_infected[1]      = y
	cumulative_infected[1] = zero( eltype(cumulative_infected) )
	effective_Rt[1]        = Rt[1]

	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
		St = max(population - cumulative_infected[t], 0) / population
		effective_Rt[t] = St * Rt[t]
	end
end

function runepidemic!(newly_infected, cumulative_infected, effective_Rt, θ, Rt)
	@unpack population, num_impute, serial_interval, num_si = θ

	num_time_step = length(newly_infected)

	for t = (num_impute + 1):num_time_step
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * serial_interval[t - τ] for τ = (t - 1):-1:max(t-num_si,1))

		effective_Rt[t] = St * Rt[t]
		# number of new infections (unobserved)
		newly_infected[t] = effective_Rt[t] * effectively_infectious
	end
	return nothing
end

# ============================================================================
# observations
function inf2obs(param, μ)
	@unpack delay_dispersion, delay_length, delay_dist = param
	i2o = pdf.( Ref(delay_dist(μ, delay_dispersion)), 1:delay_length )
	i2o /= sum(i2o)
	return i2o
end

infectious(t, newly_infected, i2o, stop) =
    sum(newly_infected[τ] * i2o[t - τ] for τ = (t - 1):-1:stop)

_expected!(expected, t, α, args...) =
	@inbounds expected[t] = α * infectious(t, args...)

function expected!(obsmodel::SimpleObsModel, newly_infected)
	@unpack θ, μ, α, expected = obsmodel
	@unpack delay_length = θ

	i2o = inf2obs(θ, μ)
	num_time_step = length(expected)

	@inbounds expected[1] = 1e-15 * newly_infected[1]
	for t = 2:delay_length+1
		ᾱ = selectvalue(α, t)
		_expected!(expected, t, ᾱ, newly_infected, i2o, 1)
	end
	for t = delay_length+2:num_time_step
		ᾱ = selectvalue(α, t)
		_expected!(expected, t, ᾱ, newly_infected, i2o, t-delay_length)
	end
	return nothing
end

selectvalue(x::Number, i) = x
selectvalue(x::Vector, i) = x[i]

function expected!(obsmodel::WeekdayHolidayObsModel, newly_infected)
	@unpack θ, μ, α, expected, holidayeffect, weekdayeffect = obsmodel
	@unpack delay_length, holiday, weekday = θ

	i2o = inf2obs(θ, μ)
	num_time_step = length(expected)

	@inbounds expected[1] = 1e-15 * newly_infected[1]
	for t = 2:delay_length+1
		ᾱ = selectvalue(α, t)
		weekday_holiday_effect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		_expected!(expected, t, ᾱ * weekday_holiday_effect, newly_infected, i2o, 1)
	end
	for t = delay_length+2:num_time_step
		ᾱ = selectvalue(α, t)
		weekday_holiday_effect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		_expected!(expected, t, ᾱ * weekday_holiday_effect, newly_infected, i2o, t-delay_length)
	end
	return nothing
end

_out_of_bounds(x, upper_lim) = !all(@. 0 < x < upper_lim )

function Turing.logpdf(obsmodel::ObservationsModel, observed)
	@unpack θ, ϕ, expected = obsmodel
	@unpack population, start, stop = θ

	T     = eltype(expected)
	ts    = start:stop
	μs = expected[ts]
	ys = observed[ts]
	# _out_of_bounds(μs, population) && (return T(Inf))
	dist  = arraydist(NegativeBinomial2.(μs, ϕ))
	return logpdf(dist, ys)
end

_unrealistic(ŷ, y) = any( ŷ .> 20*y )
function Turing.logpdf(obsmodel::SimpleSeroObsModel2, observed)
	@unpack θ, expected = obsmodel
	@unpack index, delay, std, population = θ

	T    = eltype(expected)
	μ    = expected[index .- delay] / population
	dist = MvNormal(μ, std) # select only early observations
	# _unrealistic(μ, observed) && (@error "unrealistic"; return T(Inf))
	return logpdf(dist, observed)
end

# =============================================================================

function weekdayeffect!(weekdayeffect, weekdayeffect_simplex)
	num_days = 7
	for i in eachindex(weekdayeffect)
		@inbounds weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end
end

# struct WeekDayEffect{Tα} <: ContinuousMultivariateDistribution
# 	α::Tα
# end
# WeekDayEffect() = WeekDayEffect(ones(Int64, 7))
#
# Distributions.rand(rng::AbstractRNG, d::WeekDayEffect{Tα}) where {Tα} = begin
# 	weekdayeffect_simplex = rand( Dirichlet(d.α) )
# 	weekdayeffect_simplex .* length(d.α)
# end
#
# Distributions.logpdf(d::WeekDayEffect{Tα}, x::AbstractVector{T}) where {Tα,T} =
#     logpdf( Dirichlet( d.α ), x ./ length(d.α) )# + logpdf( Normal( zero(Ts), x[1] ) )
#
# Bijectors.bijector(d::RandomWalk) = Bijectors.Scale Identity{1}()
## =============================================================================
#                      include
# ==============================================================================
# include("models/parametric_model.jl")
include("models/random-walk-models.jl")
# ## =============================================================================
# #                      model_constant_iar_contacts
# # ==============================================================================
# @model function model_constant_iar_contacts(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar       ~ Beta(1,10) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# 	α         ~ filldist(Beta(1, 5), num_covariates)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_time_steps-lockdown)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
# 	latent_Rt[num_Rt_steps+1:end] .= latent_Rt[num_Rt_steps]
# 	latent_Rt[covariates_start-lockdown:end] .= latent_Rt[covariates_start-lockdown:end] .+ covariates[1:num_time_steps-covariates_start+1, :] * α #
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_time_steps] = link.(latent_Rt)
#
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar * expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:num_obs
# 	# μs  = expected_daily_deaths[ts]
# 	# σ2s = @. μs + μs^2 / ϕ
#     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# 	μs  = expected_daily_hospit[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
# ## =============================================================================
# #                      model_constant_iar
# # ==============================================================================
# @model function model_constant_iar(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar       ~ Beta(1,10) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute * y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_Rt_steps)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
#
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_obs] = link.(latent_Rt)
# 	if predict #if prediction
# 		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
# 	end
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar * expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start: num_obs
# 	# μs  = expected_daily_deaths[ts]
# 	# σ2s = @. μs + μs^2 / ϕ
#     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# 	μs  = expected_daily_hospit[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
# ## =============================================================================
# #                      model_contacts_v3
# # ==============================================================================
# @model function model_contacts_v3(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar       ~ Beta(1,10) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# 	α         ~ filldist(Beta(1, 5), num_covariates)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_time_steps-lockdown)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
# 	latent_Rt[num_Rt_steps+1:end] .= latent_Rt[num_Rt_steps]
# 	latent_Rt[covariates_start-lockdown:end] .= latent_Rt[covariates_start-lockdown:end] .+ covariates[1:num_time_steps-covariates_start+1, :] * α #
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_time_steps] = link.(latent_Rt)
#
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar * expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:(num_obs-num_case_obs) #DEBUG JUST FOR TESTING: REPLACE WITH num_obs !
# 	μs  = expected_daily_deaths[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
# 	(ϕ <= 0) && (@error "ϕ <= 0"; Turing.@addlogprob! -Inf; return)
# 	deaths[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ))
# 	# μs  = expected_daily_hospit[ts]
# 	# !all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	# (ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	# hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean[1:4] ~ MvNormal(expected_seropositive[seroprev_idx[1:4]], seroprev_σ[1:4])
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
# ## =============================================================================
# #                      model_contacts_v4
# # ==============================================================================
# @model function model_contacts_v4(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar0      ~ Beta(10,100) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	σ_iar     ~ truncated(Normal(0.05, .03), 0, .15)
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# 	α         ~ filldist(Beta(1, 5), num_covariates)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_time_steps-lockdown)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
# 	latent_Rt[num_Rt_steps+1:end] .= latent_Rt[num_Rt_steps]
# 	latent_Rt[covariates_start-lockdown:end] .= latent_Rt[covariates_start-lockdown:end] .+ covariates[1:num_time_steps-covariates_start+1, :] * α #
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_time_steps] = link.(latent_Rt)
#
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
# 	latent_iar = TV(undef, num_iar_steps)
# 	latent_iar[1] = logit(iar0)
# 	for i in 2:num_iar_steps
# 		latent_iar[i] ~ Normal(latent_iar[i-1], σ_iar)
# 	end
# 	iar = logistic.( latent_iar[iar_index[1:num_time_steps]] )
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar .* expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:(num_obs-num_case_obs) #DEBUG JUST FOR TESTING: REPLACE WITH num_obs !
# 	μs  = expected_daily_deaths[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
# 	(ϕ <= 0) && (@error "ϕ <= 0"; Turing.@addlogprob! -Inf; return)
# 	deaths[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ))
# 	# μs  = expected_daily_hospit[ts]
# 	# !all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	# (ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	# hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
#
# ## =============================================================================
# #                      model_contacts_fitted
# # ==============================================================================
# @model function model_contacts_fitted(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
#     covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	priors,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_covariates = size(covariates, 2)
# 	num_Rt_steps = num_obs - lockdown
#     # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     y         ~ truncated( priors[:y], 0, 1000)
#     ϕ2        ~ truncated( priors[:ϕ2], 0, Inf)
#     ϕ3        ~ truncated( priors[:ϕ3], 0, Inf)
# 	R0        ~ truncated( priors[:R0], 2, 4.5)
# 	R1        ~ truncated( priors[:R1], .5, 1.1)
#     ifr       ~ truncated( priors[:ifr], 3/1000, 10/1000)
# 	ihr       ~ truncated( priors[:ihr],.1/100,5/100)
#     iar0      ~ truncated( priors[:iar0], 0, 1)
# 	α         ~ filldist(Beta(1, 100), num_covariates)
# 	# α         = zeros(V, num_covariates)
# 	seroprev_σ ~ arraydist( priors[:seroprev_σ] )
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
# 	latent_Rt             = TV(undef, num_Rt_steps)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
# 	latent_iar             = TV(undef, num_iar_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	# hidden markov model for Rt
# 	latent_Rt[1] = invlink(R1)
# 	latent_Rt[2:num_Rt_steps] ~ arraydist(priors[:logRt])
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_obs] = link.(latent_Rt)
# 	if predict #if prediction
# 		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
# 	end
#
# 	# Adjusts for portion of pop that are susceptible
# 	@. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	any( infected .< 0 ) && (i = findfirst(<(0), infected); @error "infected<0 at $(i)"; Turing.@addlogprob! -Inf; return)
# 	any( infected .> population ) && (i = findfirst(>(population), infected); @error "infected>population at $(i)"; Turing.@addlogprob! -Inf; return)
# 	any( .! isfinite.(infected) ) && (i = findfirst(x->!isfinite(x), infected); @error "infected is not finite at $(i)"; Turing.@addlogprob! -Inf; return)
#
# 	latent_iar[1] = logit(iar0)
# 	latent_iar[2:num_iar_steps] ~ arraydist( priors[:latent_iar] )
# 	iar = logistic.( latent_iar[ iar_index[1:num_time_steps] ] )
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1] = infected[1]
#     for t = 2:num_time_steps
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
#     end
# 	observed_daily_cases = iar .* expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:num_obs
# 	μs  = expected_daily_hospit[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && (@error "observed_daily_cases"; Turing.@addlogprob! -Inf; return)
# 		(ϕ3 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
# 		cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
# 	end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		cumulative_cases      = cumcases,
# 		iar = iar
#     )
# end
# # ## =============================================================================
# # #                      model_contacts
# # # ==============================================================================
# # @model function model_contacts(
# #     num_impute,        # [Int] num. of days for which to impute infections
# #     num_total_days,    # [Int] days of observed data + num. of days to forecast
# #     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
# #     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
# #     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     epidemic_start,    # [AbstractVector{<:Int}]
# #     population,        # [AbstractVector{<:Real}]
# #     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# # 	num_iar_steps,         # [Int] number of weeks (21)
# #     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
# #     lockdown,         # [Int] number of weeks (21)
# #     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# # 	covariates,
# # 	covariates_start,
# # 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     hospit,
# # 	seroprev_mean,
# #     seroprev_std,
# #     seroprev_idx,
# #     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
# #     ::Type{TV} = Vector{Float64},
# # 	::Type{V}  = Float64;
# # 	invlink = log,#KLogit(V(3))
# # 	link    = exp,#KLogistic(V(3))
# # ) where {TV, V}
# #     num_obs = length(cases)
# # 	num_covariates = size(covariates, 2)
# # 	num_Rt_steps = num_obs - lockdown
# # 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
# #     num_time_steps = predict ? num_total_days : num_obs
# #
# #     # Latent variables
# #     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# # 	T = typeof(τ)
# #     y         ~ truncated(Exponential(τ),T(0),T(1000))
# # 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
# #     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
# #     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# # 	R0        ~ truncated(Normal(3.6, .8), 2, 5.)
# # 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
# #     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# # 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
# #     iar0      ~ Beta(10,100) #Beta(20,25)
# # 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# # 	σ_iar     ~ InverseGamma2(0.1, .2)
# # 	# σ_rt      = .05
# # 	# σ_iar     = .05
# # 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.001)))
# # 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# # 	α         ~ filldist(Beta(1, 100), num_covariates)
# #
# # 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# # 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# # 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
# #
# #     infected              = TV(undef, num_time_steps)
# #     expected_daily_cases  = TV(undef, num_time_steps)
# #     expected_daily_hospit = TV(undef, num_time_steps)
# #     cumcases              = TV(undef, num_time_steps)
# #     expected_daily_deaths = TV(undef, num_time_steps)
# # 	latent_Rt             = TV(undef, num_Rt_steps)
# # 	Rt                    = TV(undef, num_time_steps)
# #     Rt_adj                = TV(undef, num_time_steps)
# #
# #     # Imputation of `num_impute` days
# #     infected[1:num_impute] .= y
# #     cumcases[1] = zero(V)
# #     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
# #
# # 	latent_Rt[1] = invlink(R1)
# # 	for t in 2:num_Rt_steps
# # 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# # 	end
# # 	latent_Rt[covariates_start-lockdown:num_Rt_steps] .= latent_Rt[covariates_start-lockdown:num_Rt_steps] .+ covariates[1:num_obs-covariates_start+1, :] * α
# #
# # 	Rt[1:lockdown] .= R0
# # 	Rt[lockdown+1:num_obs] = link.(latent_Rt)
# # 	if predict #if prediction
# # 		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
# # 	end
# # 	# effect1 = 1 .+ α[1] * covariates[1:num_obs-covariates_start+1,1]
# # 	# effect2 = 1 .+ α[2] * covariates[1:num_obs-covariates_start+1,2]
# # 	# effect3 = 1 .+ α[3] * covariates[1:num_obs-covariates_start+1,3]
# # 	# effect4 = 1 .+ α[4] * covariates[1:num_obs-covariates_start+1,4]
# # 	# @. Rt[covariates_start:num_obs] *= effect1 + effect2 + effect3 + effect4
# #
# # 	# Adjusts for portion of pop that are susceptible
# #     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
# #     for t = (num_impute + 1):num_time_steps
# #         # Update cumulative infected
# #         cumcases[t] = cumcases[t - 1] + infected[t - 1]
# #         # Adjusts for portion of pop that are susceptible
# #         susceptible = (max(population - cumcases[t], zero(V)) / population)
# #         Rt_adj[t] = susceptible * Rt[t]
# #         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = 1:t-1) #max(t-50, 1)
# #     end
# # 	any( infected .< 0 ) && (i = findfirst(<(0), infected); @error "infected<0 at $(i)"; Turing.@addlogprob! -Inf; return)
# # 	any( .! isfinite.(infected) ) && (i = findfirst(x->!isfinite(x), infected); @error "infected is not finite at $(i)"; Turing.@addlogprob! -Inf; return)
# # 	any( infected .> population ) && (i = findfirst(>(population), infected); @error "infected>population at $(i)"; Turing.@addlogprob! -Inf; return)
# # 	#
# # 	latent_iar = TV(undef, num_iar_steps)
# # 	latent_iar[1] = logit(iar0)
# # 	for i in 2:num_iar_steps
# # 		latent_iar[i] ~ Normal(latent_iar[i-1], σ_iar)
# # 	end
# # 	iar = logistic.( latent_iar[iar_index[1:num_time_steps]] )
# #
# #     expected_daily_deaths[1] = 1e-15 * infected[1]
# # 	expected_daily_hospit[1] = 1e-15 * infected[1]
# # 	expected_daily_cases[1] = infected[1]
# #     for t = 2:num_time_steps
# # 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = 1:t-1) #max(t-100, 1)
# #         expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = 1:t-1) #max(t-30, 1)
# #         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = 1:t-1) #max(t-50, 1)
# #     end
# # 	observed_daily_cases = iar .* expected_daily_cases
# # 	# Extract time-steps for which we have observations
# #
# # 	ts = epidemic_start:num_obs
# # 	# μs  = expected_daily_deaths[ts]
# # 	# σ2s = @. μs + μs^2 / ϕ
# #     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# # 	μs  = expected_daily_hospit[ts]
# # 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# # 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# # 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
# #     if num_case_obs > 0
# # 		ts = (num_obs-num_case_obs):num_obs
# # 		μs  = observed_daily_cases[ts]
# # 		!all( 0 .< μs .< population ) && (@error "observed_daily_cases"; Turing.@addlogprob! -Inf; return)
# # 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
# #         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
# #     end
# #
# #     # seroprevalence study
# # 	seroprev_mean ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
# # 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
# #
# #     return (
# # 		observed_daily_cases  = observed_daily_cases,
# #         expected_daily_deaths = expected_daily_deaths,
# #         expected_daily_hospit = expected_daily_hospit,
# #         Rt = Rt,
# # 		cumulative_cases      = cumcases,
# # 		iar = iar
# #     )
# # end
# ## =============================================================================
# #                      model_contacts_v2
# # ==============================================================================
# @model function model_contacts_v2(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar0      ~ Beta(10,100) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	σ_iar     ~ truncated(Normal(0.05, .03), 0, .15)
# 	# σ_rt      = .05
# 	# σ_iar     = .05
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# 	α         ~ filldist(Beta(1, 5), num_covariates)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_Rt_steps)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
# 	latent_Rt[covariates_start-lockdown:num_Rt_steps] .= latent_Rt[covariates_start-lockdown:num_Rt_steps] .+ covariates[1:num_obs-covariates_start+1, :] * α #
#
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_obs] = link.(latent_Rt)
# 	if predict #if prediction
# 		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
# 	end
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
# 	#
# 	latent_iar = TV(undef, num_iar_steps)
# 	latent_iar[1] = logit(iar0)
# 	for i in 2:num_iar_steps
# 		latent_iar[i] ~ Normal(latent_iar[i-1], σ_iar)
# 	end
# 	iar = logistic.( latent_iar[iar_index[1:num_time_steps]] )
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar .* expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:num_obs
# 	# μs  = expected_daily_deaths[ts]
# 	# σ2s = @. μs + μs^2 / ϕ
#     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# 	μs  = expected_daily_hospit[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
# ## =============================================================================
# #                      model_contacts_shifted
# # ==============================================================================
# @model function model_contacts_shifted(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
#     τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# 	T = typeof(τ)
#     y         ~ truncated(Exponential(τ),T(0),T(1000))
# 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
#     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
#     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# 	R0        ~ truncated(Normal(3.6, .8), 2, 4.5)
# 	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
#     ifr       ~ truncated(Normal(6/1000, 3/1000), 3/1000, 10/1000)
# 	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
#     iar0      ~ Beta(10,100) #Beta(20,25)
# 	σ_rt      ~ truncated(Normal(0.05, .03), 0, .15)
# 	σ_iar     ~ truncated(Normal(0.05, .03), 0, .15)
# 	# σ_rt      = .05
# 	# σ_iar     = .05
# 	seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.1)))
# 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# 	α         ~ filldist(Beta(1, 5), num_covariates)
#
# 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_time_steps-lockdown)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	latent_Rt[1] = invlink(R1)
# 	for t in 2:num_Rt_steps
# 		latent_Rt[t] ~ Normal(latent_Rt[t-1], σ_rt)
# 	end
# 	latent_Rt[num_Rt_steps+1:end] .= latent_Rt[num_Rt_steps]
# 	latent_Rt[covariates_start-lockdown:end] .= latent_Rt[covariates_start-lockdown:end] .+ covariates[1:num_time_steps-covariates_start+1, :] * α #
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:num_time_steps] = link.(latent_Rt)
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
# 	!all( 0 .< infected .< population ) && (@error "infected<0"; Turing.@addlogprob! -Inf; return)
#
# 	#
# 	latent_iar = TV(undef, num_iar_steps)
# 	latent_iar[1] = logit(iar0)
# 	for i in 2:num_iar_steps
# 		latent_iar[i] ~ Normal(latent_iar[i-1], σ_iar)
# 	end
# 	iar = logistic.( latent_iar[iar_index[1:num_time_steps]] )
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar .* expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:num_obs
# 	# μs  = expected_daily_deaths[ts]
# 	# σ2s = @. μs + μs^2 / ϕ
#     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# 	μs  = expected_daily_hospit[ts]
# 	!all( 0 .< μs .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
# 	(ϕ3 <= 0) && (@error "ϕ3 <= 0"; Turing.@addlogprob! -Inf; return)
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 		!all( 0 .< μs .< population ) && ( @error "observed_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
# 		(ϕ2 <= 0) && (@error "ϕ2 <= 0"; Turing.@addlogprob! -Inf; return)
#         cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
#
# ## =============================================================================
# #                      model_contacts_v2
# # ==============================================================================
# @model function model_contacts_test(
#     num_impute,        # [Int] num. of days for which to impute infections
#     num_total_days,    # [Int] days of observed data + num. of days to forecast
#     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
#     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
#     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     epidemic_start,    # [AbstractVector{<:Int}]
#     population,        # [AbstractVector{<:Real}]
#     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# 	num_iar_steps,         # [Int] number of weeks (21)
#     iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
#     lockdown,         # [Int] number of weeks (21)
#     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# 	covariates,
# 	covariates_start,
# 	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
#     hospit,
# 	seroprev_mean,
#     seroprev_std,
#     seroprev_idx,
# 	π4,
# 	priors,
#     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
#     ::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
# 	invlink = log,#KLogistic(V(3))
# 	link    = exp,#KLogit(V(k))
# ) where {TV, V}
# 	num_si = length(serial_intervals)
# 	num_π = length(π)
# 	num_π2 = length(π2)
# 	num_π3 = length(π3)
# 	num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
#     num_obs = length(cases)
# 	num_Rt_steps = num_obs - lockdown
# 	num_covariates = size(covariates, 2)
# 	num_latent_Rt = covariates_start-lockdown-1
# 	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
#     num_time_steps = predict ? num_total_days : num_obs
#
#     # Latent variables
# 	y         ~ truncated( priors[:y], 0, 1000)
#     ϕ2        ~ truncated( priors[:ϕ2], 0, Inf)
#     ϕ3        ~ truncated( priors[:ϕ3], 0, Inf)
# 	R0        ~ truncated( priors[:R0], 2, 4.5)
# 	R1        ~ truncated( priors[:R1], .5, 1.1)
#     ifr       ~ truncated( priors[:ifr], 3/1000, 10/1000)
# 	ihr       ~ truncated( priors[:ihr],.1/100,5/100)
#     iar0      ~ truncated( priors[:iar0], 0, 1)
# 	seroprev_σ~ arraydist( priors[:seroprev_σ] )
#
#     infected              = TV(undef, num_time_steps)
#     expected_daily_cases  = TV(undef, num_time_steps)
#     expected_daily_hospit = TV(undef, num_time_steps)
# 	expected_seropositive = TV(undef, num_time_steps)
#     cumcases              = TV(undef, num_time_steps)
#     expected_daily_deaths = TV(undef, num_time_steps)
#     latent_Rt             = TV(undef, num_latent_Rt)
# 	Rt                    = TV(undef, num_time_steps)
#     Rt_adj                = TV(undef, num_time_steps)
#
#     # Imputation of `num_impute` days
#     infected[1:num_impute] .= y
#     cumcases[1] = zero(V)
#     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
#
# 	num_latent_Rt = covariates_start-lockdown-1
# 	latent_Rt[1] = invlink(R1)
# 	latent_Rt[2:num_latent_Rt] ~ arraydist(priors[:logRt][1:num_latent_Rt-1])
# 	Rt[1:lockdown] .= R0
# 	Rt[lockdown+1:lockdown+num_latent_Rt] = link.(latent_Rt)
# 	Rt[lockdown+num_latent_Rt:num_obs] ~ arraydist(priors[:Rt])
# 	if predict #if prediction
# 		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
# 	end
#
# 	# Adjusts for portion of pop that are susceptible
#     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * Rt[1:num_impute]
#     for t = (num_impute + 1):num_time_steps
#         # Update cumulative infected
#         cumcases[t] = cumcases[t - 1] + infected[t - 1]
#         # Adjusts for portion of pop that are susceptible
#         susceptible = (max(population - cumcases[t], zero(V)) / population)
#         Rt_adj[t] = susceptible * Rt[t]
#         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
#     end
#
# 	latent_iar[1] = logit(iar0)
# 	latent_iar[2:num_iar_steps] ~ arraydist( priors[:latent_iar] )
# 	iar = logistic.( latent_iar[ iar_index[1:num_time_steps] ] )
#
#     expected_daily_deaths[1] = 1e-15 * infected[1]
# 	expected_daily_hospit[1] = 1e-15 * infected[1]
# 	expected_daily_cases[1]  = 1e-15 * infected[1]
# 	expected_seropositive[1] = 1e-15 * infected[1]
# 	for t = 2:num_time_steps
# 		expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# 		expected_daily_deaths[t] = ifr * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
#         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# 		expected_seropositive[t] = sum( infected[τ] * π4[t - τ] for τ = 1:t-1 )
#     end
# 	observed_daily_cases = iar .* expected_daily_cases
# 	# Extract time-steps for which we have observations
#
# 	ts = epidemic_start:num_obs
# 	μs  = expected_daily_hospit[ts]
# 	hospit[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ3))
#     if num_case_obs > 0
# 		ts = (num_obs-num_case_obs):num_obs
# 		μs  = observed_daily_cases[ts]
# 	    cases[ts] ~ arraydist(NegativeBinomial2.(μs, ϕ2))
#     end
#
#     # seroprevalence study
# 	# seroprev_mean ~ MvNormal(expected_seropositive[seroprev_idx], seroprev_σ)
# 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
#
#     return (
# 		observed_daily_cases  = observed_daily_cases,
#         expected_daily_deaths = expected_daily_deaths,
#         expected_daily_hospit = expected_daily_hospit,
#         Rt = Rt,
# 		expected_seropositive = expected_seropositive,
# 		iar = iar,
#     )
# end
# #---------------------------------------------------------------------------------------
# # GP for Rt
#
# # struct MvNormalGP{T<:AbstractGPs.FiniteGP} <: AbstractMvNormal
# #     f::T
# # end
# # AbstractGPs.rand(rng::Random.AbstractRNG, f::MvNormalGP) = rand(rng, f.f)
# # AbstractGPs.logpdf(f::MvNormalGP, y::AbstractVector{<:Real}) = logpdf(f.f, y)
# #
# # kernel(s::Real, l::Real) =
# #     s^2 * KernelFunctions.transform(Matern32Kernel(), 1/l)
# #
# # @model function model_gp(
# #     num_impute,        # [Int] num. of days for which to impute infections
# #     num_total_days,    # [Int] days of observed data + num. of days to forecast
# #     cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
# #     deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
# #     π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     epidemic_start,    # [AbstractVector{<:Int}]
# #     population,        # [AbstractVector{<:Real}]
# #     serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
# #     lockdown,         # [Int] number of weeks (21)
# #     num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
# #     π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
# #     hospit,
# # 	seroprev_mean,
# #     seroprev_std,
# #     seroprev_idx,
# #     predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
# #     ::Type{TV} = Vector{Float64},
# # 	::Type{V}  = Float64
# # ) where {TV, V}
# # 	num_si = length(serial_intervals)
# # 	num_π = length(π)
# # 	num_π2 = length(π2)
# # 	num_π3 = length(π3)
# # 	num_obs = length(cases)
# # 	num_Rt_steps = num_obs - lockdown
# #     num_obs = length(cases)
# #     # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
# #     num_time_steps = predict ? num_total_days : num_obs
# #
# #     # Latent variables
# #     # τ         ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
# # 	# T = typeof(τ)
# #     # y         ~ truncated(Exponential(τ),T(0),T(1000))
# # 	y         ~ truncated(Exponential(100),0.,1000.)
# # 	ϕ         ~ truncated(Normal(0, 5), 0, Inf)
# #     ϕ2        ~ truncated(Normal(25, 5), 10, 40)
# #     ϕ3        ~ truncated(Normal(25, 5), 10, 40)
# # 	R0        ~ truncated(Normal(3.28, .4), 2, 4)
# # 	R1        ~ truncated(Normal(.8, .2), .5, 1.2)
# #     ifr_noise ~ truncated(Normal(1., 0.2), 0, Inf)
# # 	ihr       ~ truncated(Normal(2/100,2/100),.1/100,6/100)
# #     iar0      ~ Beta(10,40) #Beta(20,25)
# # 	# σ_rt      ~ InverseGamma2(0.2, .3)
# # 	σ_rt      = .1
# # 	σ_iar     = .05
# # 	# seroprev_σ ~ arraydist(InverseGamma2.(seroprev_std, Ref(.01)))
# # 	# seroprev_μ ~ MvNormal(seroprev_mean, seroprev_mean*.01)
# #
# # 	!isfinite(R0) && (@error "R0"; Turing.@addlogprob! -Inf; return) #DEBUG
# # 	population < num_impute*y && (@error "population < num_impute*y"; Turing.@addlogprob! -Inf; return) #DEBUG
# # 	y<0 && (@error "y<0"; Turing.@addlogprob! -Inf; return) #DEBUG
# #
# #     infected              = TV(undef, num_time_steps)
# #     expected_daily_cases  = TV(undef, num_time_steps)
# #     expected_daily_hospit = TV(undef, num_time_steps)
# #     cumcases              = TV(undef, num_time_steps)
# #     expected_daily_deaths = TV(undef, num_time_steps)
# #     logRt                 = TV(undef, num_time_steps)
# #     Rt_adj                = TV(undef, num_time_steps)
# # 	latent_iar             = TV(undef, num_time_steps)
# #
# #     # Imputation of `num_impute` days
# #     infected[1:num_impute] .= y
# #     cumcases[1] = zero(V)
# #     cumcases[2:num_impute] .= cumsum(infected[1:num_impute - 1])
# #
# # 	gp = GP(kernel(.3, 20.))
# # 	f = to_sde(gp, SArrayStorage(Float64))
# # 	X = RegularSpacing(0.0, 1., num_obs - lockdown)
# #     slow_Rt_effect ~ MvNormalGP( f(X) )
# #
# # 	logRt[1:lockdown] .= log(R0)
# # 	logRt[lockdown+1:num_obs] = slow_Rt_effect .- slow_Rt_effect[1] .+ log(R1)
# # 	if predict #if prediction
# # 		logRt[num_obs+1:num_total_days] .= logRt[num_obs]
# # 	end
# #
# # 	# Adjusts for portion of pop that are susceptible
# #     @. Rt_adj[1:num_impute] = (max(population - cumcases[1:num_impute], zero(V)) / population) * exp(logRt[1:num_impute])
# #     for t = (num_impute + 1):num_time_steps
# #         # Update cumulative infected
# #         cumcases[t] = cumcases[t - 1] + infected[t - 1]
# #         # Adjusts for portion of pop that are susceptible
# #         susceptible = (max(population - cumcases[t], zero(V)) / population)
# #         Rt_adj[t] = susceptible * exp(logRt[t])
# #         infected[t] = Rt_adj[t] * sum(infected[τ] * serial_intervals[t - τ] for τ = (t - 1):-1:max(t-num_si,1)) #max(t-50, 1)
# #     end
# # 	any( infected .< 0 ) && (i = findfirst(<(0), infected); @error "infected<0 at $(i)"; Turing.@addlogprob! -Inf; return)
# # 	any( .! isfinite.(infected) ) && (i = findfirst(x->!isfinite(x), infected); @error "infected is not finite at $(i)"; Turing.@addlogprob! -Inf; return)
# # 	any( infected .> population ) && (i = findfirst(>(population), infected); @error "infected>population at $(i)"; Turing.@addlogprob! -Inf; return)
# #
# # 	cases_start = num_obs - num_case_obs + 1
# # 	latent_iar[1:cases_start-1] .= logit(iar0)
# # 	if num_case_obs > 0
# # 		gp = GP(kernel(.3, 20.))
# # 		f = to_sde(gp, SArrayStorage(Float64))
# # 		X = RegularSpacing(0.0, 1., num_case_obs)
# # 		slow_IAR_effect ~ MvNormalGP( f(X) )
# # 		latent_iar[cases_start:num_obs] = slow_IAR_effect .- slow_IAR_effect[1] .+ logit(iar0)
# # 	end
# # 	if predict #if prediction
# # 		latent_iar[num_obs+1:num_total_days] .= latent_iar[num_obs]
# # 	end
# # 	iar = logistic.(latent_iar)
# #
# #     expected_daily_deaths[1] = 1e-15 * infected[1]
# # 	expected_daily_hospit[1] = 1e-15 * infected[1]
# # 	expected_daily_cases[1] = infected[1]
# #     for t = 2:num_time_steps
# # 		expected_daily_deaths[t] = ifr_noise * sum(infected[τ] * π[t - τ] for τ = (t - 1):-1:max(t-num_π,1)) #max(t-100, 1)
# #         expected_daily_cases[t]  = sum(infected[τ] * π2[t - τ] for τ = (t - 1):-1:max(t-num_π2,1)) #max(t-30, 1)
# #         expected_daily_hospit[t] = ihr * sum(infected[τ] * π3[t - τ] for τ = (t - 1):-1:max(t-num_π3,1)) #max(t-50, 1)
# #     end
# # 	observed_daily_cases = iar .* expected_daily_cases
# # 	# Extract time-steps for which we have observations
# #
# # 	ts = epidemic_start:num_obs
# # 	# μs  = expected_daily_deaths[ts]
# # 	# σ2s = @. μs + μs^2 / ϕ
# #     # deaths[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# # 	μs  = expected_daily_hospit[ts]
# # 	σ2s = @. μs + μs^2 / ϕ3
# # 	hospit[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# #     if num_case_obs > 0
# # 		ts = (num_obs-num_case_obs):num_obs
# # 		μs  = observed_daily_cases[ts]
# # 		σ2s = @. μs + μs^2 / ϕ2
# #         cases[ts] ~ arraydist(NegativeBinomial3.(μs, σ2s))
# #     end
# #
# #     # seroprevalence study
# # 	seroprev_mean ~ MvNormal(cumcases[seroprev_idx], seroprev_std)
# # 	# seroprev_μ ~ MvNormal(cumcases[seroprev_idx], seroprev_σ)
# #
# #     return (
# # 		observed_daily_cases  = observed_daily_cases,
# #         expected_daily_deaths = expected_daily_deaths,
# #         expected_daily_hospit = expected_daily_hospit,
# #         Rt = exp.(logRt),
# # 		cumulative_cases      = cumcases,
# # 		iar = iar
# #     )
# # end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
export Vec, Vec1, Vec2, Vec3, Vec1f, Vec2f, Vec3f
export ⊙, ⊘, inv, norm²

"""
    Vec{Dim,T}

A vector in `Dim`-dimensional space with coordinates of type `T`.

A vector can be obtained by subtracting two [`Point`](@ref) objects.

## Examples

```julia

A = Point(0.0, 0.0)
B = Point(1.0, 0.0)
v = B - A
```
### Notes

- A `Vec` is an `SVector` from StaticArrays.jl
- Type aliases are `Vec1`, `Vec2`, `Vec3`, `Vec1f`, `Vec2f`, `Vec3f`
"""
const Vec = SVector

# type aliases for convenience
const Vec1  = SVector{1,Float64}
const Vec2  = SVector{2,Float64}
const Vec3  = SVector{3,Float64}
const Vec1f = SVector{1,Float32}
const Vec2f = SVector{2,Float32}
const Vec3f = SVector{3,Float32}

"""
    ⊙(A::AbstractArray, B::AbstractArray)

Hadamard product (element-wise multiplication) of two equal size arrays.
"""
⊙(A::AbstractArray, B::AbstractArray) = map(*, A, B)

"""
    ⊘(a::AbstractArray, b::AbstractArray)

Hadamard division (element-wise division) of two equal size arrays.
"""
⊘(A::AbstractArray, B::AbstractArray) = map(/, A, B)

"""
    inv(v::AbstractVector)

Return the Samelson inverse of a.

The returned vector, a⁻¹, satisfies a⁻¹ ⋅ a = 1.
"""
Base.inv(a::AbstractVector) = inv(a ⋅ a) * a'

"""
    norm²(v::AbstractVector)

Return the squared 2-norm of the vector.
"""
norm²(v::AbstractVector) = v ⋅ v

ext: jl

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# Initialize data structures in element container
function init_elements!(elements, mesh::StructuredMesh{1}, basis::LobattoLegendreBasis)
  @unpack node_coordinates, left_neighbors,
          jacobian_matrix, contravariant_vectors, inverse_jacobian = elements

  # Calculate node coordinates, Jacobian matrix, and inverse Jacobian determinant
  for cell_x in 1:size(mesh, 1)
    calc_node_coordinates!(node_coordinates, cell_x, mesh.mapping, mesh, basis)

    calc_jacobian_matrix!(jacobian_matrix, cell_x, node_coordinates, basis)

    calc_inverse_jacobian!(inverse_jacobian, cell_x, jacobian_matrix)
  end

  # Contravariant vectors don't make sense in 1D, they would be identical to inverse_jacobian
  fill!(contravariant_vectors, NaN)

  initialize_neighbor_connectivity!(left_neighbors, mesh)

  return nothing
end

# Calculate physical coordinates to which every node of the reference element is mapped
# `mesh.mapping` is passed as an additional argument for type stability (function barrier)
function calc_node_coordinates!(node_coordinates, cell_x, mapping, mesh::StructuredMesh{1},
                                basis::LobattoLegendreBasis)
  @unpack nodes = basis

  # Get cell length in reference mesh
  dx = 2 / size(mesh, 1)

  # Calculate node coordinates of reference mesh
  cell_x_offset = -1 + (cell_x-1) * dx + dx/2

  for i in eachnode(basis)
    # node_coordinates are the mapped reference node_coordinates
    node_coordinates[1, i, cell_x] = mapping(cell_x_offset + dx/2 * nodes[i])[1]
  end
end

# Calculate Jacobian matrix of the mapping from the reference element to the element in the physical domain
function calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates::AbstractArray{<:Any,3},
                               basis::LobattoLegendreBasis)
  @views mul!(jacobian_matrix[1, 1, :, element], basis.derivative_matrix, node_coordinates[1, :, element]) # x_ξ

  return jacobian_matrix
end

# Calculate inverse Jacobian (determinant of Jacobian matrix of the mapping) in each node
function calc_inverse_jacobian!(inverse_jacobian::AbstractArray{<:Any, 2}, element, jacobian_matrix)
  @views inverse_jacobian[:, element] .= inv.(jacobian_matrix[1, 1, :, element])

  return inverse_jacobian
end

# Save id of left neighbor of every element
function initialize_neighbor_connectivity!(left_neighbors, mesh::StructuredMesh{1})
  # Neighbors in x-direction
  # Inner elements
  for cell_x in 2:size(mesh, 1)
    left_neighbors[1, cell_x] = cell_x - 1
  end

  if isperiodic(mesh)
    # Periodic boundary
    left_neighbors[1, 1] = size(mesh, 1)
  else
    # Use boundary conditions
    left_neighbors[1, 1] = 0
  end

  return left_neighbors
end

end # @muladd


text:
module GPUContagion

using RandomNumbers.Xorshifts
using CUDA
using LightGraphs
using UnPack
using Test
using GPUArrays
using FLoops
using Random123
using BenchmarkTools
@enum AgentStatus begin
    Susceptible = 1
    Infected
    Recovered
end

# const Susceptible, Infected, Recovered = 1,2,3
const RNG = Xoroshiro128Star(1)
const nodes = 10_000

function main()
    g = BitArray(adjacency_matrix(erdos_renyi(nodes,0.005)))
    u_0 = get_u_0(nodes)
    p = get_params()
    steps = 1000
    solution = vcat([u_0], [fill(Susceptible,length(u_0)) for i in 1:steps])
    sol_cu = cu.(solution)
    cu_graph = cu(g)
    @btime solve($solution,$p,$steps,$g)
    @btime CUDA.@sync solve_cuda($sol_cu,$p,$steps,$cu_graph)
    # display(solution)
    # solve_cuda(solution,p,100,g)
end
function get_u_0(nodes)
        u_0 = fill(Susceptible,nodes)
        init_indices = rand(RNG, 1:nodes, 5)
        u_0[init_indices] .= Infected #start with five infected
        return u_0
end

function get_params()
    params = (
        p = 0.005,
        recovery_rate = 0.01,
    )
    return params
end
function solve(solution,params,steps,g)
    for t in 1:steps
        agents_step!(t,solution[t+1],solution[t],g,params)
    end

    return solution
end
function solve_cuda(sol_cu,params,steps,cu_graph)
    for t in 1:steps
        agents_step_gpu!(t,sol_cu[t+1],sol_cu[t],cu_graph,params)
    end
    return sol_cu
end
function agents_step!(t,u_next,u,g,params)
    @unpack p, recovery_rate  = params
    u_next .= u#keep state the same if nothing else happens

    for i in 1:length(u)
        agent_status = u[i]
        if agent_status == Susceptible
            for j in 1:length(u)
                if g[i,j] == 1 && u[j] == Infected && rand(RNG) < p
                    u_next[i] = Infected
                end
            end
        elseif agent_status == Infected
            if rand(RNG) < recovery_rate
                u_next[i] = Recovered
            end
        end
    end
end

function agents_step_gpu!(t,u_next,u,graph,params)
    u_next .= u
    function kernel(state, _, (t,u_next,u,g,params,randstate))
        @unpack p, recovery_rate  = params
        i = linear_index(state)
        if i <= length(u)
            agent_status = u[i]
            if agent_status == Susceptible
                for j in 1:length(u)
                    if g[i,j] == 1 && u[j] == Infected && GPUArrays.gpu_rand(Float64, state, randstate) < p
                        u_next[i] = Infected
                    end
                end
            elseif agent_status == Infected
                if GPUArrays.gpu_rand(Float64, state, randstate) < recovery_rate
                    u_next[i] = Recovered
                end
            end
        end
        return nothing
    end
    gpu_call(kernel, u, (t,u_next,u,graph,params, GPUArrays.default_rng(typeof(u)).state))
end

function generate_contact_vectors!(ij_dist,ji_dist,i_to_j_contacts, j_to_i_contacts)
    rand!(RNG,ij_dist,i_to_j_contacts)
    rand!(RNG,ji_dist,j_to_i_contacts)
    l_i = length(i_to_j_contacts)
    l_j = length(j_to_i_contacts)
    contacts_sums = sum(i_to_j_contacts) - sum(j_to_i_contacts)
    sample_list_length = max(l_i,l_j) #better heuristic for this based on stddev of dist?
    index_list_i = sample(RNG,1:l_i,sample_list_length)
    index_list_j = sample(RNG,1:l_j,sample_list_length)
    sample_list_i = rand(RNG,ij_dist,sample_list_length)
    sample_list_j = rand(RNG,ji_dist,sample_list_length)
    for k = 1:sample_list_length
        if (contacts_sums != 0)
            i_index = index_list_i[k]
            j_index = index_list_j[k]
            contacts_sums +=  j_to_i_contacts[j_index] - i_to_j_contacts[i_index]

            i_to_j_contacts[i_index] = sample_list_i[k]
            j_to_i_contacts[j_index] = sample_list_j[k]
            contacts_sums += i_to_j_contacts[i_index] -  j_to_i_contacts[j_index]
        else
            break
        end
    end
    while contacts_sums != 0
        i_index = sample(RNG,1:l_i)
        j_index = sample(RNG,1:l_j)
        contacts_sums +=  j_to_i_contacts[j_index] - i_to_j_contacts[i_index]
        i_to_j_contacts[i_index] = rand(RNG,ij_dist)
        j_to_i_contacts[j_index] = rand(RNG,ji_dist)
        contacts_sums += i_to_j_contacts[i_index] -  j_to_i_contacts[j_index]
    end
    return nothing
end
function gpu_contact_vectors!(ij_dist,ji_dist,i_to_j_contacts, j_to_i_contacts)
    rand!(RNG,ij_dist,i_to_j_contacts)
    rand!(RNG,ji_dist,j_to_i_contacts)
    l_i = length(i_to_j_contacts)
    l_j = length(j_to_i_contacts)
    contacts_sums = sum(i_to_j_contacts) - sum(j_to_i_contacts)
    sample_list_length = max(l_i,l_j) #better heuristic for this based on stddev of dist?
    index_list_i = sample(RNG,1:l_i,sample_list_length)
    index_list_j = sample(RNG,1:l_j,sample_list_length)
    sample_list_i = rand(RNG,ij_dist,sample_list_length)
    sample_list_j = rand(RNG,ji_dist,sample_list_length)
    for k = 1:sample_list_length
        if (contacts_sums != 0)
            i_index = index_list_i[k]
            j_index = index_list_j[k]
            contacts_sums +=  j_to_i_contacts[j_index] - i_to_j_contacts[i_index]

            i_to_j_contacts[i_index] = sample_list_i[k]
            j_to_i_contacts[j_index] = sample_list_j[k]
            contacts_sums += i_to_j_contacts[i_index] -  j_to_i_contacts[j_index]
        else
            break
        end
    end
    while contacts_sums != 0
        i_index = sample(RNG,1:l_i)
        j_index = sample(RNG,1:l_j)
        contacts_sums +=  j_to_i_contacts[j_index] - i_to_j_contacts[i_index]
        i_to_j_contacts[i_index] = rand(RNG,ij_dist)
        j_to_i_contacts[j_index] = rand(RNG,ji_dist)
        contacts_sums += i_to_j_contacts[i_index] -  j_to_i_contacts[j_index]
    end
end
function Base.show(io::IO, status::AgentStatus)
    if status == Susceptible
        print(io, "S")
    elseif status == Infected
        print(io, "I")
    elseif status == Recovered
        print(io, "R")
    elseif status == Immunized
        print(io, "Vac")
    end
end

end


text:
export iauPvstar
"""
Convert star position+velocity vector to catalog coordinates.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  support function.

Given (Note 1):
   pv     double[2][3]   pv-vector (au, au/day)

Returned (Note 2):
   ra     double         right ascension (radians)
   dec    double         declination (radians)
   pmr    double         RA proper motion (radians/year)
   pmd    double         Dec proper motion (radians/year)
   px     double         parallax (arcsec)
   rv     double         radial velocity (km/s, positive = receding)

Returned (function value):
         int            status:
                           0 = OK
                           -1 = superluminal speed (Note 5)
                           -2 = null position vector

Notes:

   1. The specified pv-vector is the coordinate direction (and its rate
   of change) for the date at which the light leaving the star
   reached the solar-system barycenter.

   2. The star data returned by this function are "observables" for an
   imaginary observer at the solar-system barycenter.  Proper motion
   and radial velocity are, strictly, in terms of barycentric
   coordinate time, TCB.  For most practical applications, it is
   permissible to neglect the distinction between TCB and ordinary
   "proper" time on Earth (TT/TAI).  The result will, as a rule, be
   limited by the intrinsic accuracy of the proper-motion and
   radial-velocity data;  moreover, the supplied pv-vector is likely
   to be merely an intermediate result (for example generated by the
   function iauStarpv), so that a change of time unit will cancel
   out overall.

   In accordance with normal star-catalog conventions, the object's
   right ascension and declination are freed from the effects of
   secular aberration.  The frame, which is aligned to the catalog
   equator and equinox, is Lorentzian and centered on the SSB.

   Summarizing, the specified pv-vector is for most stars almost
   identical to the result of applying the standard geometrical
   "space motion" transformation to the catalog data.  The
   differences, which are the subject of the Stumpff paper cited
   below, are:

   (i) In stars with significant radial velocity and proper motion,
   the constantly changing light-time distorts the apparent proper
   motion.  Note that this is a classical, not a relativistic,
   effect.

   (ii) The transformation complies with special relativity.

   3. Care is needed with units.  The star coordinates are in radians
   and the proper motions in radians per Julian year, but the
   parallax is in arcseconds; the radial velocity is in km/s, but
   the pv-vector result is in au and au/day.

   4. The proper motions are the rate of change of the right ascension
   and declination at the catalog epoch and are in radians per Julian
   year.  The RA proper motion is in terms of coordinate angle, not
   true angle, and will thus be numerically larger at high
   declinations.

   5. Straight-line motion at constant speed in the inertial frame is
   assumed.  If the speed is greater than or equal to the speed of
   light, the function aborts with an error status.

   6. The inverse transformation is performed by the function iauStarpv.

Called:
   iauPn        decompose p-vector into modulus and direction
   iauPdp       scalar product of two p-vectors
   iauSxp       multiply p-vector by scalar
   iauPmp       p-vector minus p-vector
   iauPm        modulus of p-vector
   iauPpp       p-vector plus p-vector
   iauPv2s      pv-vector to spherical
   iauAnp       normalize angle into range 0 to 2pi

Reference:

   Stumpff, P., 1985, Astron.Astrophys. 144, 232-240.

This revision:  2017 March 16

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# int iauPvstar(double pv[2][3], double *ra, double *dec,
#               double *pmr, double *pmd, double *px, double *rv)

function iauPvstar(pv::AbstractMatrix{<:Real})
   # Preallocate return values
   ref_ra  = Ref{Float64}(0.0)
   ref_dec = Ref{Float64}(0.0)
   ref_pmr = Ref{Float64}(0.0)
   ref_pmd = Ref{Float64}(0.0)
   ref_px  = Ref{Float64}(0.0)
   ref_rv  = Ref{Float64}(0.0)

   status = ccall((:iauPvstar, libsofa_c), Cint,
         (Ptr{Cdouble},
         Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble},
         Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}),
         convert(Array{Float64, 2}, pv'),
         ref_ra, ref_dec, ref_pmr, ref_pmd, ref_px, ref_rv)

   return status, ref_ra[], ref_dec[], ref_pmr[], ref_pmd[], ref_px[], ref_rv[]
end


text:
A(K1, K2) = dot(K1, K2) / sqrt(norm(K1) * norm(K2))

function alignment(nnpool::NNPool, data::Array{Float64, 2}, K_full::Array{Float64, 2}, nn_stability; adjusted)
    idx = collect(keys(nnpool.nn_pools[:Lin]) ∪ keys(nnpool.nn_pools[:Lout]));

    y = zeros(size(data, 2));
    y[idx] .= [get(nnpool.nn_pools[:Lin], id, 0) / (get(nnpool.nn_pools[:Lin], id, 0) + get(nnpool.nn_pools[:Lout], id, 0)) for id in idx]
    y[idx] = ifelse.(y[idx] .> 0.5, 1, -1)

    idx_in = findall(y .== 1)
    idx_out = findall(y .== -1)

    K_opt = y * y'

    kdtree = KDTree(data; leafsize = nn_stability)
    nn = knn(kdtree, data, nn_stability, true)[1]
    reverse_nn = KernelLearning.reverseKNN(nn)
    symmetric_nn = [n ∩ r for (n,r) in zip(nn, reverse_nn)];

    mask = falses(size(K_opt))

    for (_, id) in enumerate(findall(nnpool.pools .!= :U))
        if y[id] == 1
            nn_in = nn[id][1:end] ∩ idx_in
            nn_out = nn[id][1:end] ∩ idx_out
        else
            nn_in = setdiff(nn[id][1:end], reverse_nn[id][1:end]) ∩ idx_in
            nn_out = symmetric_nn[id][1:end] ∩ idx_out
        end

        if length(nn_in) > 0
            mask[id, nn_in[1:min(end, nn_stability)]] .= true
            mask[nn_in[1:min(end, nn_stability)], id] .= true
        end

        if length(nn_out) > 0
            mask[id, nn_out[1:min(end, nn_stability)]] .= true
            mask[nn_out[1:min(end, nn_stability)], id] .= true
        end
    end

    if adjusted
        K_full = centerkernelmatrix(K_full)
        K_opt = centerkernelmatrix(K_opt)
    end

    return A(K_full[mask], K_opt[mask])
end

function alignment(nnpool::NNPool, data::Array{Float64, 2}, gamma::Float64, nn_stability; adjusted=true)
    K_full = kernelmatrix(Val(:col), GaussianKernel(gamma), data)
    return alignment(nnpool, data, K_full, nn_stability; adjusted=adjusted)
end

function alignment_score(nnpool::NNPool, data::Array{Float64, 2}, K_full::Array{Float64, 2}, query_history, k; adjusted=true, n_sample = 100)
    candidates = [x for x ∈ 1:size(K_full, 1) if x ∉ query_history]
    scores = zeros(length(candidates))
    A = alignment(nnpool, data, K_full, k; adjusted=adjusted)

    if length(candidates) > n_sample
        candidates_subset = StatsBase.sample(eachindex(candidates), n_sample, replace = false)
    else
        candidates_subset = eachindex(candidates)
    end

    for i in candidates_subset
        tmp_in = deepcopy(nnpool)
        add_to_pools!(tmp_in, candidates[i], :inlier, k)
        A_in = alignment(tmp_in, data, K_full, k; adjusted=adjusted)

        tmp_out = deepcopy(nnpool)
        add_to_pools!(tmp_out, candidates[i], :outlier, k)
        A_out = alignment(tmp_out, data, K_full, k; adjusted=adjusted)

        scores[i] = min(abs(A - A_in), abs(A - A_out))
    end
    return (candidates, scores)
end

function next_query(nnpool::NNPool, data::Array{Float64, 2}, K_full::Array{Float64, 2}, query_history, nn_stability; adjusted)
    @assert length(query_history) < size(data,2)
    candidates, scores = alignment_score(nnpool, data, K_full, query_history, nn_stability; adjusted=adjusted)
    candidates[sortperm(scores, rev=true)][1]
end

function find_gamma(nnpool::NNPool, data::Array{Float64, 2}, nn_stability; adjusted=true, min_γ=0.1, max_γ=30.0)
    neg_alignment(gamma) = -alignment(nnpool, data, gamma, nn_stability; adjusted=adjusted)
    res = Optim.optimize(neg_alignment, min_γ, max_γ, abs_tol=0.1)
    return res.minimizer
end

max_stars_repo_name: holtri/lama

max_issues_repo_name: holtri/lama

max_forks_repo_name: holtri/lama


text:
@node typeof(*) Deterministic [ out, A, in ]

# By default multiplication node uses TinyCorrection() strategy for precision matrix on `in` edge to ensure precision is always invertible
default_meta(::typeof(*)) = TinyCorrection()