text:
using ..CovidSurvey: RandomWalk
@model function model_cases(
    θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    predict,
	vectorize = true,
	::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
) where {TV, V}

	@unpack num_observations, num_total_days, num_regions, num_rt_steps, invlink, hospitmodel, casemodel, seromodel, num_covariates, include_early_dynamic = θ
    # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
    num_time_steps = predict ? num_total_days : num_observations

	############# 2.) time varying reproduction number

	# κ          ~ truncated(Normal(0, 1.), 0, Inf)
	# R0s        ~ filldist(truncated(Normal(2., κ), 1., 5.), num_regions)
	# σ_rt       ~ truncated(Normal(0.15, .05), 0, .25)
	R0         ~ truncated(Normal(3.0, 0.5), 1., 4.5)
	R0s        ~ filldist(truncated(Normal(R0, 0.2), 1., 4.5), num_regions)
	σ_rt       ~ truncated(Normal(0.15*2, .05*2), 0, .25*3)

	init_rw = if include_early_dynamic
			R0s
		else
			R1  ~ truncated(Normal(0.8, 0.1), 0., 1.2)
			R1s ~ filldist(truncated(Normal(R1, .1), 0., 1.2), num_regions)
	end

	latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
	latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
	for m in 1:num_regions
		latent_Rts_z[m] ~ RandomWalk(num_rt_steps[m])
		latent_Rts[m] = latent_Rts_z[m] * σ_rt .+ invlink(init_rw[m])
	end

	Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	if num_covariates > 0
		grouped_effect ~ filldist( Laplace(0,0.2), num_covariates)
		effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
		effects = [ effects_z[:,m] .* effect_std .+ grouped_effect for m in 1:num_regions]

		# grouped_effect ~ filldist( Exponential(0.2), num_covariates)
		# effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		# pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
		# pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
		# effects ~ filldist( pooled_effects_distr, num_regions )

		if θ.semiparametric
			semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		else
			mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		end
	else
		random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
	end

	############ 3.) infection dynamics
	τ  ~ Exponential(100/5) # `Exponential` has inverse parameterization of the one in Stan
	T  = typeof(τ)
	ys ~ filldist(truncated(Normal(τ, 30),T(0),T(500)), num_regions)

	newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	cumulative_infecteds = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	# effective_Rts        = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]

	infections!(newly_infecteds, cumulative_infecteds, θ, τ, ys, Rts)
	# infections!(newly_infecteds, cumulative_infecteds, effective_Rts, θ, τ, ys, Rts)
	########### 4.) derive observables
    μ_i2h ~ truncated(Normal(9.5, .5), 7, 15)
	μ_i2c ~ truncated(Normal(2., .5), 0, 5)
	ihr   ~ truncated(Normal(1.8/100,0.5/100), 1/100, 5/100)
	iar   ~ Beta(1,10)
	ϕ_h   ~ truncated(Normal(50, 10), 20, Inf)
	ϕ_c   ~ truncated(InverseGamma2(90, .5), 10, Inf)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	weekdayeffect = TV(undef, 7)
	weekdayeffect!(weekdayeffect, weekdayeffect_simplex)

	### holiday effect
	holidayeffect ~ Beta(1,1)

	########### 3.) observation model
	expected_daily_hospits = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	expected_daily_cases   = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]

	hospit_observation_model = SimpleObsModel(hospitmodel, μ_i2h, ihr, ϕ_h, expected_daily_hospits, vectorize)
	# cases_observation_model = SimpleObsModel(casemodel, μ_i2c, iar, ϕ_c, expected_daily_cases)
	cases_observation_model = WeekdayHolidayObsModel(
        casemodel, μ_i2c, iar, ϕ_c, expected_daily_cases, weekdayeffect, holidayeffect, vectorize
    )
	sero_observation_model   = SimpleSeroObsModel(seromodel, cumulative_infecteds)

	expected!(hospit_observation_model, newly_infecteds)
	expected!(cases_observation_model, newly_infecteds)

	########### 4.) compare model to observations
	## 4.1) observe hospitalizations

	ℓ  = zero(V)
	ℓ += logpdf(hospit_observation_model, θ.hospits)
	ℓ += logpdf(cases_observation_model, θ.cases)
	ℓ += logpdf(sero_observation_model, θ.sero)
	Turing.@addlogprob! ℓ

	return (
		# newly_infecteds = newly_infecteds,
		expected_daily_hospits = expected_daily_hospits,
		expected_daily_cases = expected_daily_cases,
		expected_seropos = cumulative_infecteds,
		Rts = Rts,
		# effective_Rt = effective_Rts
	)
end

# ============================================================
@model function model_hospit(
    θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
	::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
    predict = false,
	vectorize = true,
	context   = :inference
) where {TV, V}

	@unpack num_observations, num_total_days, num_regions, num_rt_steps, hospitmodel, invlink, seromodel, num_covariates, rwscale, include_early_dynamic = θ
    # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
    num_time_steps = predict ? num_total_days : num_observations

	############# 2.) time varying reproduction number

	R0_std     ~ truncated(Normal(.5, .2), 0.01, Inf)
	R0         ~ truncated(Normal(3.5, 1.), 1., 4.5)
	R0s        ~ filldist(truncated(Normal(R0, R0_std), 1., 4.5), num_regions)
	# R0         ~ truncated(Normal(3.0, 0.5), 0., 6.)
	# R0s        ~ filldist(truncated(Normal(R0, 0.2), 0., 6.), num_regions) #CovidSurvey.CustomMvNormal1(R0, 0.2, 1., 4.5, num_regions) #filldist(truncated(Normal(R0, 0.2), 1., 4.5), num_regions)
	σ_rt       ~ truncated(Normal(0.3*rwscale, .02*rwscale), 0, .5*rwscale) #N(0.32,0.5)

	init_rw = if include_early_dynamic
            R0s
		else
			# R1  ~ truncated(Normal(0.8, 0.1), 0., 2.)
			# R1s ~ filldist(truncated(Normal(R1, 0.1), 0., 2.), num_regions)

			R1_std    ~ truncated(Normal(0.1, .05), 0.05, Inf)
			R1s       ~ filldist(truncated(Normal(0.8, R1_std), 0., 1.5), num_regions)
			R1s
	end

	latent_Rts = if include_early_dynamic
		latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
		latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
		for m in 1:num_regions
			latent_Rts_z[m] ~ CovidSurvey.RandomWalk(num_rt_steps[m])
			@. latent_Rts[m] = latent_Rts_z[m] * σ_rt + invlink(init_rw[m])
		end
		latent_Rts
	else
		n = first( num_rt_steps )
		latent_Rts_z ~ filldist( CovidSurvey.RandomWalk( n ), num_regions)
		CovidSurvey.rescale( latent_Rts_z, σ_rt, invlink.(init_rw) )
	end
	# latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
	# latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
	# for m in 1:num_regions
	# 	latent_Rts_z[m] ~ RandomWalk(num_rt_steps[m])
	# 	latent_Rts[m] = latent_Rts_z[m] * σ_rt .+ invlink(init_rw[m])
	# end

	Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	if num_covariates > 0
		grouped_effects ~ filldist( Laplace(0,0.2), num_covariates) #Exponential(0.2)
		effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
		effects = [ effects_z[:,m] .* effect_std .+ grouped_effects for m in 1:num_regions]

		# grouped_effect ~ filldist( Exponential(0.2), num_covariates)
		# effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		# pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
		# pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
		# effects ~ filldist( pooled_effects_distr, num_regions )

		if θ.semiparametric
			CovidSurvey.semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		else
			CovidSurvey.mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		end
	else
		CovidSurvey.random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
	end

	############ 3.) infection dynamics
	τ  ~ Exponential(20) # `Exponential` has inverse parameterization of the one in Stan
	# ys ~ filldist(truncated(Normal(τ, 30),T(1),T(500)), num_regions)
	ys ~ filldist(Exponential(τ), num_regions)

	newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	cumulative_infecteds = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	# effective_Rts        = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]

	CovidSurvey.infections!(newly_infecteds, cumulative_infecteds, θ, ys, Rts)
	# infections!(newly_infecteds, cumulative_infecteds, effective_Rts, θ, τ, ys, Rts)
	########### 4.) derive observables
	μ_i2h ~ truncated(Normal(11.7, .5), 8, 14)
	σ_i2h ~ truncated(Normal(6.5,0.5), 1, 14)
	ihr   ~ truncated(Normal(1.8/100,0.5/100), 0.5/100, 5/100) #GammaMeanCv(0.025, 0.1)
	ϕ_h   ~ truncated(Normal(50, 10), 20, Inf)

	expected_daily_hospit    = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	hospit_observation_model = CovidSurvey.SimpleObsModel2(hospitmodel, μ_i2h, σ_i2h, ihr, ϕ_h, expected_daily_hospit, vectorize)
	# hospit_observation_model = CovidSurvey.SimpleObsModel(hospitmodel, μ_i2h, ihr, ϕ_h, expected_daily_hospit, vectorize)
	sero_observation_model   = CovidSurvey.SimpleSeroObsModel(seromodel, cumulative_infecteds)
	CovidSurvey.expected!(hospit_observation_model, newly_infecteds)

    ########### 4.) compare model to observations
	## 4.1) observe hospitalizations

	if context == :inference
		ℓ  = zero(V)
		ℓ += CovidSurvey.logpdf(hospit_observation_model, θ.hospits)
		ℓ += CovidSurvey.logpdf(sero_observation_model, θ.sero)
		Turing.@addlogprob! ℓ
		return (
			newly_infecteds = newly_infecteds,
			expected_daily_hospits = expected_daily_hospit,
			expected_seropos = cumulative_infecteds,
			Rts = Rts,
		)
	else
		if num_covariates > 0
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_hospits = expected_daily_hospit,
				expected_seropos = cumulative_infecteds,
				Rts = Rts,
				predicted_daily_hospits = CovidSurvey.prediction(hospit_observation_model), #effects = effects,
				grouped_effects = [grouped_effects for m in 1:num_regions],
				effects = effects
			)
		else
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_hospits = expected_daily_hospit,
				expected_seropos = cumulative_infecteds,
				Rts = Rts,
				predicted_daily_hospits = CovidSurvey.prediction(hospit_observation_model) #effects = effects,
				#effective_Rt = effective_Rts
			)
		end
	end
end

# @model function model_hospit(
#     θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
# 	::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
#     predict = false,
# 	vectorize = true,
# 	context = :inference #:prediction, :loglikelihood
# ) where {TV, V}
#
# 	@unpack num_observations, num_total_days, num_regions, num_rt_steps, invlink, hospitmodel, seromodel, num_covariates, include_early_dynamic = θ
#     # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
#     num_time_steps = predict ? num_total_days : num_observations
#
# 	############# 2.) time varying reproduction number
#
# 	# κ          ~ truncated(Normal(0, 1.), 0, Inf)
# 	# R0s        ~ filldist(truncated(Normal(2., κ), 1., 5.), num_regions)
# 	# σ_rt       ~ truncated(Normal(0.15, .05), 0, .25)
# 	# R0         ~ truncated(Normal(3.0, 0.5), 1., 4.5)
# 	# R0_std     ~ GammaMeanCv(0.2, 0.5)
# 	# R0s        ~ filldist( truncated( Normal( R0, R0_std ), V(1.), V(4.5)), num_regions)
# 	R0s        ~ filldist(truncated(Normal(3., 1.), 1., 5.), num_regions)
# 	σ_rt       ~ truncated(Normal(0.15*2, .05*2), 0, .25*3)
#
# 	init_rw = if include_early_dynamic
# 			R0s
# 		else
# 			R1s ~ filldist(truncated(Normal(.8, .1), .5, 1.1), num_regions)
# 	end
#
# 	latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
# 	latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
# 	for m in 1:num_regions
# 		latent_Rts_z[m] ~ RandomWalk(num_rt_steps[m])
# 		latent_Rts[m] = latent_Rts_z[m] * σ_rt .+ invlink(init_rw[m])
# 	end
#
# 	Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	if num_covariates > 0
# 		grouped_effects ~ filldist( Laplace(0,0.2), num_covariates)
# 		effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
# 		effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
# 		effects = [ effects_z[:,m] .* effect_std .+ grouped_effects for m in 1:num_regions]
#
# 		# grouped_effect ~ filldist( Exponential(0.2), num_covariates)
# 		# effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
# 		# pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
# 		# pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
# 		# effects ~ filldist( pooled_effects_distr, num_regions )
#
# 		if θ.semiparametric
# 			semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
# 		else
# 			mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
# 		end
# 	else
# 		random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
# 	end
#
#
# 	############ 3.) infection dynamics
# 	# y  ~ InverseGamma2(120, 0.2) # `Exponential` has inverse parameterization of the one in Stan
# 	# y_std   ~ InverseGamma2(10, 0.5)
# 	# ys ~ filldist(truncated(Normal(y, 10), 1., 500), num_regions)
# 	#ys      = y_z * y_std .+ y_mean
# 	τ  ~ Exponential(30.) # `Exponential` has inverse parameterization of the one in Stan
# 	T  = typeof(τ)
# 	ys ~ filldist(truncated(Exponential(τ),T(0),T(1000)), num_regions)
#
#
# 	newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	cumulative_infecteds = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	# effective_Rts        = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
#
# 	infections!(newly_infecteds, cumulative_infecteds, θ, ys, Rts)
# 	# infections!(newly_infecteds, cumulative_infecteds, effective_Rts, θ, τ, ys, Rts)
# 	########### 4.) derive observables
#     μ_i2h ~ truncated(Normal(11., .5), 8, 14)
# 	σ_i2h ~ truncated(Normal(5.41,1), 1, 14)
# 	ihr   ~ truncated(Normal(1.8/100,0.5/100), 1/100, 5/100)
# 	ϕ_h   ~ truncated(Normal(50, 10), 20, Inf)
#
# 	expected_daily_hospits   = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	hospit_observation_model = SimpleObsModel2(hospitmodel, μ_i2h, σ_i2h, ihr, ϕ_h, expected_daily_hospits, vectorize)
# 	sero_observation_model   = SimpleSeroObsModel(seromodel, cumulative_infecteds)
#
# 	expected!(hospit_observation_model, newly_infecteds)
#
# 	########### 4.) compare model to observations
# 	## 4.1) observe hospitalizations
# 	ℓ  = zero(V)
# 	ℓ += logpdf(hospit_observation_model, θ.hospits)
# 	ℓ += logpdf(sero_observation_model, θ.sero)
# 	Turing.@addlogprob! ℓ
#
# 	return (
# 		newly_infecteds = newly_infecteds,
# 		expected_daily_hospits = expected_daily_hospits,
# 		expected_seropos = cumulative_infecteds,
# 		Rts = Rts,
# 		#predicted_daily_hospits = predicted_daily_hospits,
# 		# effective_Rt = effective_Rts
# 	)
# 	# if context == :inference
# 	# 	ℓ  = zero(V)
# 	# 	ℓ += logpdf(hospit_observation_model, θ.hospits)
# 	# 	ℓ += logpdf(sero_observation_model, θ.sero)
# 	# 	Turing.@addlogprob! ℓ
# 	#
# 	# elseif context == :loglikelihood
# 	# 	hospit_loglikelihoods = loglikelihoods(hospit_observation_model, θ.hospits)
# 	# 	sero_loglikelihoods = loglikelihoods(sero_observation_model, θ.sero)
# 	# 	return (;
# 	# 		hospit_loglikelihoods,
# 	# 		sero_loglikelihoods
# 	# 	)
# 	#
# 	# else #prediction
# 	# 	predicted_daily_hospits = prediction(hospit_observation_model)
# 	#
# 	# 	if num_covariates == 0
# 	# 		return (
# 	# 			newly_infecteds = newly_infecteds,
# 	# 			expected_daily_hospits = expected_daily_hospits,
# 	# 			expected_seropos = cumulative_infecteds,
# 	# 			Rts = Rts,
# 	# 			predicted_daily_hospits = predicted_daily_hospits,
# 	# 			# effective_Rt = effective_Rts
# 	# 		)
# 	# 	else
# 	# 		return (
# 	# 			newly_infecteds = newly_infecteds,
# 	# 			expected_daily_hospits = expected_daily_hospits,
# 	# 			expected_seropos = cumulative_infecteds,
# 	# 			Rts = Rts,
# 	# 			predicted_daily_hospits = predicted_daily_hospits,
# 	# 			grouped_effects = [grouped_effects for m in 1:num_regions],
# 	# 			effects = effects
# 	# 			#effective_Rt = effective_Rts
# 	# 		)
# 	# 	end
# 	# end
# end
#
#
# @model function model_hospit2(
#     θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
# 	hospits,
# 	::Type{TV} = Vector{Float64},
# 	::Type{V}  = Float64;
#     predict = false,
# 	vectorize = true,
# 	context = :inference #:prediction, :loglikelihood
# ) where {TV, V}
#
# 	@unpack num_observations, num_total_days, num_regions, num_rt_steps, invlink, hospitmodel, seromodel, num_covariates, include_early_dynamic = θ
#     # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
#     num_time_steps = predict ? num_total_days : num_observations
#
# 	############# 2.) time varying reproduction number
#
# 	# κ          ~ truncated(Normal(0, 1.), 0, Inf)
# 	# R0s        ~ filldist(truncated(Normal(2., κ), 1., 5.), num_regions)
# 	# σ_rt       ~ truncated(Normal(0.15, .05), 0, .25)
# 	R0         ~ truncated(Normal(3.0, 0.5), 1., 4.5)
# 	R0s        ~ filldist(truncated(Normal(R0, 0.2), 1., 4.5), num_regions)
# 	σ_rt       ~ truncated(Normal(0.15*2, .05*2), 0, .25*3)
#
# 	init_rw = if include_early_dynamic
# 			R0s
# 		else
# 			R1s ~ filldist(truncated(Normal(.8, .1), .5, 1.1), num_regions)
# 	end
#
# 	latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
# 	latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
# 	for m in 1:num_regions
# 		latent_Rts_z[m] ~ RandomWalk(num_rt_steps[m])
# 		latent_Rts[m] = latent_Rts_z[m] * σ_rt .+ invlink(init_rw[m])
# 	end
#
# 	Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	if num_covariates > 0
# 		grouped_effects ~ filldist( Laplace(0,0.2), num_covariates)
# 		effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
# 		effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
# 		effects = [ effects_z[:,m] .* effect_std .+ grouped_effects for m in 1:num_regions]
#
# 		# grouped_effect ~ filldist( Exponential(0.2), num_covariates)
# 		# effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
# 		# pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
# 		# pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
# 		# effects ~ filldist( pooled_effects_distr, num_regions )
#
# 		if θ.semiparametric
# 			semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
# 		else
# 			mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
# 		end
# 	else
# 		random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
# 	end
#
#
# 	############ 3.) infection dynamics
# 	τ  ~ Exponential(100 / 5) # `Exponential` has inverse parameterization of the one in Stan
# 	T  = typeof(τ)
# 	ys ~ filldist(truncated(Normal(τ, 30),T(0),T(500)), num_regions)
#
# 	newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	cumulative_infecteds = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	# effective_Rts        = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
#
# 	infections!(newly_infecteds, cumulative_infecteds, θ, τ, ys, Rts)
# 	# infections!(newly_infecteds, cumulative_infecteds, effective_Rts, θ, τ, ys, Rts)
# 	########### 4.) derive observables
#     μ_i2h ~ truncated(Normal(11., .5), 8, 14)
# 	σ_i2h ~ truncated(Normal(5.41,1), 1, 14)
# 	ihr   ~ truncated(Normal(1.8/100,0.5/100), 1/100, 5/100)
# 	ϕ_h   ~ truncated(Normal(50, 10), 20, Inf)
#
# 	expected_daily_hospits   = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
# 	hospit_observation_model = SimpleObsModel2(hospitmodel, μ_i2h, σ_i2h, ihr, ϕ_h, expected_daily_hospits, vectorize)
# 	sero_observation_model   = SimpleSeroObsModel(seromodel, cumulative_infecteds)
#
# 	expected!(hospit_observation_model, newly_infecteds)
#
# 	########### 4.) compare model to observations
# 	## 4.1) observe hospitalizations
#
# 	for m in 1:num_regions
# 		start = hospitmodel.starts[m]
# 		stop  = hospitmodel.stops[m]
# 		ts    = start:stop
# 		μs = expected_daily_hospits[m][ts]
# 		ys = hospits[m][ts]
#
# 		_out_of_bounds(μs, populations[m]) && (return T(Inf))
# 		if vectorize
# 			ys ~ arraydist(NegativeBinomial2.(μs, Ref(ϕ_h)))
# 		else
# 			for j in eachindex(ys)
# 				ys[j] ~ NegativeBinomial2(μs[j], ϕ_h)
# 			end
# 		end
# 	end
#
# 	if num_covariates == 0
# 		return (
# 			newly_infecteds = newly_infecteds,
# 			expected_daily_hospits = expected_daily_hospits,
# 			expected_seropos = cumulative_infecteds,
# 			Rts = Rts,
# 			# effective_Rt = effective_Rts
# 		)
# 	else
# 		return (
# 			newly_infecteds = newly_infecteds,
# 			expected_daily_hospits = expected_daily_hospits,
# 			expected_seropos = cumulative_infecteds,
# 			Rts = Rts,
# 			#grouped_effects = grouped_effects,
# 			#effects = effects
# 			# effective_Rt = effective_Rts
# 		)
# 	end
# end
# ============================================================
@model function model_deaths(
    θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
	::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
    predict = false,
	vectorize = true,
	context   = :inference
) where {TV, V}

	@unpack num_observations, num_total_days, num_regions, num_rt_steps, deathmodel, invlink, seromodel, num_covariates, rwscale, include_early_dynamic = θ
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
    num_time_steps = predict ? num_total_days : num_observations

	############# 2.) time varying reproduction number

	R0_std     ~ truncated(Normal(.5, .2), 0.01, Inf)
	R0         ~ truncated(Normal(3.5, 1.), 1., 4.5)
	R0s        ~ filldist(truncated(Normal(R0, R0_std), 1., 4.5), num_regions)
	# R0         ~ truncated(Normal(3.0, 0.5), 0., 6.)
	# R0s        ~ filldist(truncated(Normal(R0, 0.2), 0., 6.), num_regions) #CovidSurvey.CustomMvNormal1(R0, 0.2, 1., 4.5, num_regions) #filldist(truncated(Normal(R0, 0.2), 1., 4.5), num_regions)
	σ_rt       ~ truncated(Normal(0.3*rwscale, .02*rwscale), 0, .5*rwscale) #N(0.32,0.5)

	init_rw = if include_early_dynamic
            R0s
		else
			# R1  ~ truncated(Normal(0.8, 0.1), 0., 2.)
			# R1s ~ filldist(truncated(Normal(R1, 0.1), 0., 2.), num_regions)

			R1_std    ~ truncated(Normal(0.1, .05), 0.05, Inf)
			R1s       ~ filldist(truncated(Normal(0.8, R1_std), 0., 1.5), num_regions)
			R1s
	end

	latent_Rts = if include_early_dynamic
		latent_Rts_z = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
		latent_Rts = [TV(undef, num_rt_steps[m]) for m in 1:num_regions]
		for m in 1:num_regions
			latent_Rts_z[m] ~ CovidSurvey.RandomWalk(num_rt_steps[m])
			@. latent_Rts[m] = latent_Rts_z[m] * σ_rt + invlink(init_rw[m])
		end
		latent_Rts
	else
		n = first( num_rt_steps )
		latent_Rts_z ~ filldist( CovidSurvey.RandomWalk( n ), num_regions)
		CovidSurvey.rescale( latent_Rts_z, σ_rt, invlink.(init_rw) )
	end

	Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	if num_covariates > 0
		grouped_effects ~ filldist( Laplace(0,0.2), num_covariates) #Exponential(0.2)
		effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
		effects = [ effects_z[:,m] .* effect_std .+ grouped_effects for m in 1:num_regions]

		# grouped_effect ~ filldist( Exponential(0.2), num_covariates)
		# effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
		# pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
		# pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
		# effects ~ filldist( pooled_effects_distr, num_regions )

		if θ.semiparametric
			CovidSurvey.semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		else
			CovidSurvey.mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
		end
	else
		CovidSurvey.random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
	end

	############ 3.) infection dynamics
	τ  ~ Exponential(20) # `Exponential` has inverse parameterization of the one in Stan
	# ys ~ filldist(truncated(Normal(τ, 30),T(1),T(500)), num_regions)
	ys ~ filldist(Exponential(τ), num_regions)

	newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	cumulative_infecteds = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	# effective_Rts        = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]

	CovidSurvey.infections!(newly_infecteds, cumulative_infecteds, θ, ys, Rts)

	########### 4.) derive observables
	μ_i2d ~ truncated(Normal(21.82, 1.01), 18, 24)
	σ_i2h ~ truncated(Normal(4.26, 5.18), 1, 14)
	idr   ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ϕ_d   ~ truncated(Normal(50, 10), 20, Inf)

	expected_daily_deaths    = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
	deaths_observation_model = CovidSurvey.SimpleObsModel2(deathmodel, μ_i2d, σ_i2h, idr, ϕ_d, expected_daily_deaths, vectorize)
	sero_observation_model   = CovidSurvey.SimpleSeroObsModel(seromodel, cumulative_infecteds)
	CovidSurvey.expected!(deaths_observation_model, newly_infecteds)

	########### 4.) compare model to observations
	## 4.1) observe hospitalizations

	if context == :inference
		ℓ  = zero(V)
		ℓ += logpdf(deaths_observation_model, θ.deaths)
		ℓ += logpdf(sero_observation_model, θ.sero)
		Turing.@addlogprob! ℓ

		return (
			newly_infecteds = newly_infecteds,
			expected_daily_deaths = expected_daily_deaths,
			expected_seropos = cumulative_infecteds,
			Rts = Rts,
		)
	else
		if num_covariates > 0
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_deaths = expected_daily_deaths,
				expected_seropos = cumulative_infecteds,
				Rts = Rts,
				predicted_daily_deaths = prediction(deaths_observation_model), #effects = effects,
				grouped_effects = [grouped_effects for m in 1:num_regions],
				effects = effects
			)
		else
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_deaths = expected_daily_deaths,
				expected_seropos = cumulative_infecteds,
				Rts = Rts,
				predicted_daily_deaths = prediction(deaths_observation_model), #effects = effects,
			)
		end
	end
end

@model function model_international_deaths(
    θ,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
	::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	predict = false,
	vectorize = true,
	context   = :inference
) where {TV, V}

    @unpack num_observations, num_total_days, num_regions, num_rt_steps, deathmodel, init_infected, num_covariates, invlink, rwscale, include_early_dynamic = θ
      # If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
    num_time_steps = num_observations

    ############# 2.) time varying reproduction number

    R0_std     ~ truncated(Normal(.5, .2), 0.01, Inf)
    R0s        ~ filldist(truncated(Normal(0.8, R0_std), 0., 4.5), num_regions)
    σ_rt       ~ truncated(Normal(0.3*rwscale, .02*rwscale), 0, .5*rwscale)

    n = first( num_rt_steps )
    latent_Rts_z ~ filldist( CovidSurvey.RandomWalk( n ), num_regions)
    latent_Rts = CovidSurvey.rescale( latent_Rts_z, σ_rt, invlink.(R0s) )

    Rts = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
    if num_covariates > 0
        grouped_effect ~ filldist( Laplace(0,0.2), num_covariates)
        effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
        effects_z ~ filldist( MvNormal( num_covariates, 1.), num_regions)
        effects = [ effects_z[:,m] .* effect_std .+ grouped_effect for m in 1:num_regions]

        # grouped_effect ~ filldist( Exponential(0.2), num_covariates)
        # effect_std ~ filldist( GammaMeanCv(0.1, 0.5), num_covariates)
        # pooled_effects_distr_raw = Normal.( grouped_effect, effect_std )
        # pooled_effects_distr = arraydist( truncated.( pooled_effects_distr_raw, zero(V), V(Inf) ) )
        # effects ~ filldist( pooled_effects_distr, num_regions )

    	if θ.semiparametric
    	    CovidSurvey.semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
    	else
            CovidSurvey.mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
    	end
    else
        CovidSurvey.random_walks!(Rts, θ, predict, latent_Rts, R0s, σ_rt)
    end

    ############ 3.) infection dynamics
    # τ  ~ Exponential(100.) # `Exponential` has inverse parameterization of the one in Stan
    ys ~ arraydist(Exponential.(init_infected))
    # ys = fill(100, num_regions)

    newly_infecteds      = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
    CovidSurvey.infections!(newly_infecteds, θ, ys, Rts)

    ########### 4.) derive observables
    μ_i2d ~ truncated(Normal(21.82, 1.01), 18, 24)
    σ_i2d ~ truncated(Normal(4.26, 5.18), 1, 14)
    # idr   = 0.8 / 100
    idr   ~ truncated(Normal(8/1000, 1/1000), 1/1000, 15/1000)
    ϕ_d   ~ truncated(Normal(50, 10), 20, Inf)

    ### day of week effect
    weekdayeffect_simplex ~ filldist(Dirichlet([7,7,7,7,7,3,1]), num_regions)
    weekdayeffect = TV[TV(undef, 7) for m in 1:num_regions]
    CovidSurvey.weekdayeffect!(weekdayeffect, weekdayeffect_simplex)

    ### holiday effect
    holidayeffect ~ filldist(Beta(1,1), num_regions)

    expected_daily_deaths    = TV[TV(undef, num_time_steps[m]) for m in 1:num_regions]
    deaths_observation_model = CovidSurvey.WeekdayHolidayObsModel2(
        deathmodel, μ_i2d, σ_i2d, idr, ϕ_d, expected_daily_deaths, weekdayeffect, holidayeffect, vectorize
    )

    # deaths_observation_model = CovidSurvey.SimpleObsModel2(deathmodel, μ_i2d, σ_i2d, idr, ϕ_d, expected_daily_deaths, vectorize)
    CovidSurvey.expected!(deaths_observation_model, newly_infecteds)

    ########### 4.) compare model to observations
    ## 4.1) observe deaths

    if context == :inference
		ℓ = CovidSurvey.logpdf(deaths_observation_model, θ.deaths)
		Turing.@addlogprob! ℓ
		return (
			newly_infecteds = newly_infecteds,
			expected_daily_deaths = expected_daily_deaths,
			Rts = Rts,
		)
	else
		if num_covariates > 0
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_deaths = expected_daily_deaths,
				Rts = Rts,
				predicted_daily_deaths = CovidSurvey.prediction(deaths_observation_model), #effects = effects,
				grouped_effects = [grouped_effects for m in 1:num_regions],
				effects = effects
			)
		else
			return (
				newly_infecteds = newly_infecteds,
				expected_daily_deaths = expected_daily_deaths,
				Rts = Rts,
				predicted_daily_deaths = CovidSurvey.prediction(deaths_observation_model) #effects = effects,
				#effective_Rt = effective_Rts
			)
		end
	end
end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
"""
    LP_norm(in1,in2,p,imax,type)

Calculate robust inner product in the Lp norm space: Find the parameter alpha which minimize this cost function ||in1-alpha*in2||_p

# Arguments
- `in1`: vector 1
- `in2`: vector 2
- `p`: Lp norm space
- `type`: 1 output=alpha, 2 output=inner product between in1 and in2 in lp norm space
"""
function LP_norm(in1,in2,p,imax,type)
    N=length(in1);
    if N==0
        alpha=0;
    else
        alpha=0;

        for i=1:imax
            w=1 ./((abs.(in2-alpha*in1)).^(2-p).+0.000000001);
            alpha=sum(in1.*w.*in2)/sum(in1.*w.*in1);
        end
    end
    if type==1
        out=alpha;
    end
    if type==2
        out=sum(alpha.*in1);
    end
    return out
end

ext: jl

max_stars_repo_name: Ji-seismic/SeisTest

max_issues_repo_name: Ji-seismic/SeisTest

max_forks_repo_name: Ji-seismic/SeisTest


text:
#=
Given an undirected graph represented as an adjacency matrix and an integer k, write a function to determine whether each vertex in the graph can be colored such that no two adjacent vertices share the same color using at most k colors.
=#

graph = [0 1 1 0;1 0 0 0;1 0 0 1;0 0 1 0]
colors = [1, 3,3,3]

function iscolorvalid(graph::Array{Int, 2}, colors::Array{Int, 1})
    vertex_index = length(colors)
    neighbors = findall(isequal(1), graph[vertex_index,:])
    colorvalid = true
    for neighbor in neighbors
        if neighbor <= vertex_index && colors[neighbor] == colors[vertex_index]
            colorvalid = false
        end
    end
    return colorvalid
end

function colorable(graph::Array{Int,2}, k::Int)
    is_colorable = true
    colors = Array{Int,1}()
    for vertex in 1:size(graph, 1)
        vertex_colored = false
        color = 1
        while color <= k && !vertex_colored
            push!(colors, color)
            if !iscolorvalid(graph, colors) && color == k
                is_colorable = false
                color = color + 1
            elseif !iscolorvalid(graph, colors) && color < k
                pop!(colors)
                color = color + 1
            else
                vertex_colored = true
            end
        end
    end
    return is_colorable
end

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
# instantiate project
#import Pkg
#Pkg.activate(@__DIR__)
#Pkg.instantiate()

const TARGET_FOLDER = "results"
const RESULTS_FILE = "results.csv"

function main()
    if !isdir(TARGET_FOLDER)
        mkdir(TARGET_FOLDER)
    end
    global io = open(joinpath(TARGET_FOLDER, RESULTS_FILE), "w")

    println("Running evaluation...")

    println(">>> Running Logistic model")
    include("evaluation/Logistic/Logistic.jl")

    println(">>> Running SEIR model")
    include("evaluation/SEIR/SEIR.jl")

    println(">>> Running Burgers model")
    include("evaluation/Burgers/Burgers.jl")

    println("Finished running evaluation. Results stored in $TARGET_FOLDER")
    close(io)
    nothing
end

main()


text:
import Base.==

#############
## Feature ##
#############

# features are strings of the following forms: "R=f", "L=f", "f"
# where f is any string of lowercase or capital letters of the standard alphabet
Feature = String
is_feature(s :: String) = match(r"(R=|L=|\+|-)?[a-zA-Z]+", s).offset == 1
is_selector(f :: Feature) = length(f) > 2 && f[2] == '='
selects_right(f :: Feature) = is_selector(f) && f[1] == 'R'
selects_left(f :: Feature) = is_selector(f) && f[1] == 'L'
is_selectee(f :: Feature) = !is_selector(f)
name(f :: Feature) = if is_selector(f)
                        f[3:end]
                     else f[1:end]
                     end

"extracts the category (selectee) feature of a feature vector"
function catfeature(fs :: Vector{Feature})
    if isempty(fs)
        error("no category feature found.")
    elseif is_selectee(fs[1])
        fs[1]
    else
        catfeature(fs[2:end])
    end
end

#################
## LexicalItem ##
#################

type LexicalItem{T}
    phon :: T
    features :: Vector{Feature}
    id :: String # this is just the string representation of the LI
end

#function LexicalItem(s :: AbstractString)
#    splititem = map(String, map(strip, split(s, "::")))
#    (phon, features) = (String(splititem[1]), map(String, split(splititem[2])))
#    if !all(map(is_feature, features))
#        error(s, "contains a non feature. Check formatting.")
#    end
#    LexicalItem{String}(phon, map(Feature, features), String(s))
#end

function LexicalItem(s :: AbstractString)
    splititem = map(String, map(strip, split(strip(s), "::")))
    (phon, features) = (String(splititem[1]), map(String, split(splititem[2])))
    if !all(map(is_feature, features))
        error(s, "contains a non feature. Check formatting.")
    end
    features = map(Feature, features)
    id = ""
    for f in features
        id = string(id, " ", f)
    end
    id = string(phon, " ::", id)
    LexicalItem{String}(phon, features, id)
end

function LexicalItem{T}(phon :: T, features :: Vector{Feature})
    id = ""
    for f in features
        id = string(id, " ", f)
    end
    id = string(phon, " ::", id)
    LexicalItem(phon, features, id)
end

phon{T}(lex :: LexicalItem{T}) = lex.phon
features{T}(lex :: LexicalItem{T}) = lex.features
id{T}(lex :: LexicalItem{T}) = lex.id

function show{T}(io :: IO, lex :: LexicalItem{T})
    print(io, "$(lex.phon) ::")
    map(x -> print(io, " $x"), lex.features)
end
show(lex :: LexicalItem) = show(STDOUT, lex)
==(l1 :: LexicalItem, l2 :: LexicalItem) = id(l1) == id(l2)
isequal(l1 :: LexicalItem, l2 :: LexicalItem) = isequal(id(l1), id(l2))

#########################################
## helper methods for minimalist rules ##
#########################################

match(f :: Vector{Feature}, g :: Vector{Feature}) =
    name(f[1]) == name(g[1]) && (
        (selects_right(f[1]) && is_selectee(g[1])) ||
        (selects_left(g[1]) && is_selectee(f[1])))

is_finalfeat(f :: Vector{Feature}) = length(f) == 1

adjacent(r :: Tuple{Int, Int}, s :: Tuple{Int, Int}) =  r[1] <= s[2] && r[2] == s[1]

######################
## minimalist rules ##
######################

function merge(f :: Vector{Vector{Feature}}, g :: Vector{Vector{Feature}})
    @assert match(f[1], g[1])
    if selects_right(f[1][1]) && is_selectee(g[1][1]) && is_finalfeat(g[1])
        merge_right(f, g)
    elseif selects_left(g[1][1]) && is_selectee(f[1][1]) && is_finalfeat(f[1])
        merge_left(f, g)
    elseif selects_right(f[1][1]) && is_selectee(g[1][1])
        merge_move_right(f, g)
    elseif selects_left(g[1][1]) && is_selectee(f[1][1])
        merge_move_left(f, g)
    end
end

function merge_right(f :: Vector{Vector{Feature}}, g :: Vector{Vector{Feature}})
    [[[f[1][2:end]]; f[2:end]]; g[2:end]]
end

function merge_left(f :: Vector{Vector{Feature}}, g :: Vector{Vector{Feature}})
    [[[g[1][2:end]]; g[2:end]]; f[2:end]]
end

#######################
## MinimalistGrammar ##
#######################

MinimalistCategoryID = String
MinimalistRule = Tuple{MinimalistCategoryID,  Tuple{MinimalistCategoryID, MinimalistCategoryID}}
MinimalistState = State{MinimalistCategoryID,  MinimalistRule}
id(feats :: Vector{Feature}) = join(feats, ' ') # convert features to category id
features(id :: String) = map(String, split(id, ' ')) # convert category id to features

type MinimalistGrammar{T, S, Cond}
    lexicon :: Dict{T, Vector{LexicalItem{T}}}
    start_symbols :: Vector{Vector{Feature}}
    lexitem_cond :: Cond
    startstate :: MinimalistState
end
show(mg :: MinimalistGrammar) = string(("$(id(li)) ; $(exp(prob(mg, li)))\n"
                                        for (w,lexitems) in mg.lexicon
                                        for li in lexitems)...)
show_counts(mg :: MinimalistGrammar) =
    string(("$(id(li)) ; $(mg.lexitem_cond.dists[catfeature(features(li))].counts[id(li)])\n"
            for (w,lexitems) in mg.lexicon
            for li in lexitems)...)

function getlexicalitem(lexid :: String, g :: MinimalistGrammar)
    li = LexicalItem(lexid)
    @assert haskey(g.lexicon, li.phon) "lexicon has no lexical items for $(li.phon)"
    lis = g.lexicon[li.phon]
    try
        filter(li -> id(li) == lexid, lis)[1]
    catch
        error("lexicon has no lexical items matching $lexid")
    end
end

function construct_minimalist_states!(lexicon_features :: Vector{Vector{Feature}}, startstate :: MinimalistState)
    i = 1; while i <= length(lexicon_features)
        fs = lexicon_features[i]
        if is_selector(fs[1])
            cats = if selects_right(fs[1])
                [[fs], [[name(fs[1])]]]
            else
                [[[name(fs[1])]], [fs]]
            end
            merged = merge(cats...)

            cats_ids = Tuple(map(x -> join(x[1], ' '), cats))
            merged_id = join(merged[1], ' ')
            rule = (merged_id, cats_ids)
            add_rule!(startstate, rule, merged_id, cats_ids)

            if is_selector(merged[1][1])
                push!(lexicon_features, merged[1])
            end
        end
        #construct_minimalist_states!(new_lexicon_features, startstate)
        i+= 1
    end
end

"Given a lexicon as a Vector, constructs a MinimalistGrammar with a uniformly initialized Dirichlet distribution, using the DirCatVar type."
function MinimalistGrammar{T}(lexicon :: Vector{LexicalItem{T}}, start_symbols=[["c"]])
    lexitem_cond = SimpleCond(Dict(
        cat => DirCatVar(map(id, filter(li -> catfeature(features(li)) == cat, lexicon)))
        for cat in unique([catfeature(features(li)) for li in lexicon])))
    startstate = State(MinimalistCategoryID,  MinimalistRule)
    construct_minimalist_states!(unique(map(features, lexicon)), startstate)
    lex_dict = Dict(word => filter(li -> phon(li) == word, lexicon)
                    for word in unique(map(phon, lexicon)))
    MinimalistGrammar{T,LogProb,SimpleCond}(lex_dict, start_symbols, lexitem_cond, startstate)
end

"Given a lexicon as a Dict, constructs a MinimalistGrammar with a uniformly initialized Dirichlet distribution, using the DirCatVar type."
function MinimalistGrammar{T}(lexicon :: Dict{T, Vector{LexicalItem{T}}}, start_symbols=[["c"]])
    lex_list = vcat(values(lexicon)...)
    lexitem_cond = SimpleCond(Dict(
        cat => DirCatVar(map(id, filter(li -> catfeature(features(li)) == cat, lex_list)))
        for cat in unique([catfeature(features(li)) for li in lex_list])))
    MinimalistGrammar{T, LogProb, SimpleCond}(lexicon, start_symbols, lexitem_cond)
end

"Given a lexicon and a conditional distritbution, constructs a MinimalistGrammar."
function MinimalistGrammar{T,Cond}(lexicon :: Dict{T, Vector{LexicalItem{T}}}, lexitem_cond :: Cond, start_symbols=[["c"]], Prob=LogProb)
    lex_list = vcat(values(lexicon)...)
    startstate = State(MinimalistCategoryID,  MinimalistRule)
    construct_minimalist_states!(unique(map(features, lex_list)), startstate)
    MinimalistGrammar{T,Prob,Cond}(lexicon, start_symbols, lexitem_cond, startstate)
end

"""
Given a string containing lines of the form
    word :: feat1 feat2 ... featN ; prob
construct a MinimalistGrammar with a distribution populated with the specified values.
Default distribution is dirichlet distribution.
"""
function MinimalistGrammar(s :: AbstractString, start_symbols=[["c"]], Dist=DirCatVar{String,Float64}, logprob=true)
    pairs = map(x -> split(strip(x), ";"), split(strip(s), "\n"))
    lexicon_probs = [(LexicalItem(li), float(parse(p))) for (li,p) in pairs]
    if logprob
        lexitem_cond = SimpleCond(Dict(
            cat => Dist(Dict(
                id(lexitem) => prob
                for (lexitem,prob) in filter(
                    pair -> catfeature(features(pair[1])) == cat,
                    lexicon_probs)))
            for cat in unique([catfeature(features(li)) for (li,p) in lexicon_probs])))
    else
        lexitem_cond = SimpleCond(Dict(
            cat => Dist(Dict(
                id(lexitem) => LogProb(prob)
                for (lexitem,prob) in filter(
                    pair -> catfeature(features(pair[1])) == cat,
                    lexicon_probs)))
            for cat in unique([catfeature(features(li)) for (li,p) in lexicon_probs])))
    end
    startstate = State(MinimalistCategoryID, MinimalistRule)
    lexicon = map(first, lexicon_probs)
    construct_minimalist_states!(unique(map(features, lexicon)), startstate)
    lex_dict = Dict(word => filter(li -> phon(li) == word, lexicon)
                    for word in unique(map(phon, lexicon)))
    MinimalistGrammar{String, LogProb, SimpleCond}(lex_dict, start_symbols, lexitem_cond, startstate)
end

lexicon{T,S,C}(g :: MinimalistGrammar{T,S,C}) = g.lexicon
start_symbols{T,S,C}(g :: MinimalistGrammar{T,S,C}) = g.start_symbols
lexitem_cond{T,S,C}(g :: MinimalistGrammar{T,S,C}) = g.lexitem_cond
startstate{T,S,C}(g :: MinimalistGrammar{T,S,C}) = g.startstate

"""
This function is needed for calculate_outside_probs_and_expected_rule_usages!
Gives the probability that grammar g assigns to the rule.
"""
prob(g :: MinimalistGrammar, lexitem :: LexicalItem) =
    logscore(g.lexitem_cond.dists[catfeature(features(lexitem))], id(lexitem))
function prob(g :: MinimalistGrammar, cat_id :: MinimalistCategoryID, rule)
    if typeof(rule) == terminal_rule_type(g)
        word = rule[2]
        lexitem = filter(li -> features(li) == features(cat_id), lexicon(g)[word])[1]
        logscore(g.lexitem_cond.dists[catfeature(features(cat_id))], id(lexitem))
    elseif typeof(rule) == category_rule_type(g)
        LogProb(1)
    end
end

# type interface functions. needed by the parser.
terminal_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) = T
category_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) = MinimalistCategoryID
category_rule_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) = MinimalistRule
terminal_rule_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) =
Tuple{MinimalistCategoryID, T}
score_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) = Score
state_type{T, Score, Cond}(g :: MinimalistGrammar{T, Score, Cond}) = MinimalistState

function unkify!(g :: MinimalistGrammar)
    feature_seqs = unique(Tuple(features(li)) for (word, lexitems) in lexicon(g) for li in lexitems)
    unks = [LexicalItem("UNK", [f for f in fs]) for fs in feature_seqs]
    g.lexicon["UNK"] = unks
    try
        for unkli in unks
            g.lexitem_cond.dists[catfeature(features(unkli))].counts[id(unkli)] = 1
        end
    catch y
        error("unkify! only works with MGs specified with Dirichlet distributions. Error thrown:")
        throw(y)
    end
end

##############################
## parser-grammar interface ##
##############################

is_final{T,Sc}(g :: MinimalistGrammar{T,Sc}, s :: MinimalistState) = isfinal(s)

is_possible_transition{T,Sc}(g :: MinimalistGrammar{T,Sc}, s :: MinimalistState, c :: MinimalistCategoryID) = !isfinal(s) && haskey(s.trans, c)

transition{T,Sc}(g :: MinimalistGrammar{T,Sc}, s :: MinimalistState, c :: MinimalistCategoryID) = s.trans[c]

startsymbols{T,Sc}(g :: MinimalistGrammar{T,Sc}) = map(x -> join(x, ' '), start_symbols(g))

completions{T,Sc}(g :: MinimalistGrammar{T,Sc}, s :: MinimalistState) = [
    (c, r, prob(g, c, r)) for (c, r) in s.comp]

function completions{T,Sc}(g :: MinimalistGrammar{T,Sc}, t :: T)
    comps = Vector{Tuple{category_type(g), terminal_rule_type(g), score_type(g)}}()
    for li in lexicon(g)[t]
        cat = id(features(li))
        rule = (cat, t)
        push!(comps, (cat, rule, prob(g, cat, rule)))
    end
    comps
end

#############################
# Converting from CFG to MG #
#############################

# FIXME these might not work! they haven't been updated since I overhauled the
# code in this file for MGs and LexicalItems. -chris

"""
Removes unary rules from lists of CFG rules in place, giving a weakly-equivalent
grammar.
"""
function remove_unaries!{C,T}(category_rules::Vector{ContextFreeRule{C,C}}, terminal_rules::Vector{ContextFreeRule{C,T}})
    i = 1
    while i <= length(category_rules)
        if length(rhs(category_rules[i])) == 1
            unary_rule = splice!(category_rules, i)
            f(x::ContextFreeRule{C,C}) = lhs(x) == rhs(unary_rule)[1]
            for rules in [category_rules, terminal_rules]
                I = find(f, rules)
                for j in I
                    old_rule = rules[j]
                    new_rule = ContextFreeRule(lhs(unary_rule), rhs(old_rule))
                    rules[j] = new_rule
                end
            end
        else i += 1
        end
    end
end

"""
Convert a CFG encoded as lists of rules into an MG.
"""
# FIXME: For now this assigns a score (probability) of 1 to each lexical item
# So the resulting grammar may not have valid probability distributions
function cfg2mg{C<:AbstractString, T<:AbstractString, D<:AbstractString}(
        category_rules::Vector{ContextFreeRule{C,C}},
        terminal_rules::Vector{ContextFreeRule{C,T}},
        startsymbols::Vector{D},
        Score=LogProb)

    if !all(length(rhs(rule)) > 1 for rule in category_rules)
        error("Category rules cannot be unary!")
    end
    if !all(length(rhs(rule)) == 1 for rule in terminal_rules)
        error("Terminal rules must be unary!")
    end

    terminals = [[rhs(rule)[1] for rule in terminal_rules]; ""]
    terminal_map = Dict{String, Set{Vector{Feature}}}(t => Set() for t in terminals)
    complex_map = Dict{String, Set{Vector{Feature}}}(t => Set() for t in terminals)

    for rule in terminal_rules
        t = rhs(rule)[1]
        cat = lhs(rule)
        push!(terminal_map[t], [cat])
    end

    for rule in category_rules
        # find categories in the RHS of the rule which are mapped to by terminals
        f(x) = [x] in union(values(terminal_map)...)
        I = find(f, rhs(rule))
        # create a complex category for the first matching terminal found
        if !isempty(I)
            i = I[1] # let the first matching category in the RHS be the head
            head_cat = rhs(rule)[i]
            for t in keys(filter((x,y)->[head_cat] in y, terminal_map))
                feats = [["L=$B" for B in rhs(rule)[i-1:-1:1]];
                    ["R=$B" for B in rhs(rule)[i+1:end]];
                    ["$(lhs(rule))"]]
                    push!(complex_map[t], feats)
            end
            # if no matching terminals found, create a complex category for the empty string
        else
            feats = [["R=$B" for B in rhs(rule)]; ["$(lhs(rule))"]]
            push!(complex_map[""], feats)
        end
    end

    map = merge(union, terminal_map, complex_map)
    lexicon = [LexicalItem(t,fs,Score(1)) for (t,featsset) in map for fs in featsset]
    return MinimalistGrammar(lexicon, [[x] for x in startsymbols], 1)
end

########

# helpers for getting information from trees

getrule(node :: TreeNode) = node.data[2]

function getlexid(l :: TreeNode)
    r = getrule(l)
    "$(r[2]) :: $(r[1])"
end

"Give list of lexical item ids used in tree in order, preserves duplicates."
lexids(t :: TreeNode) = map(getlexid, leafs(t))

"Probability of a MG derivation tree"
function prob(g :: MinimalistGrammar, t :: TreeNode)
    f(r) = prob(g, r[1], r)
    prod(map(f, map(getrule, leafs(t))))
end


text:
using JLD2, FileIO, NetCDF
using Elefridge

path = "/network/aopp/chaos/pred/kloewer/esowc/"
member1files = filter(x->endswith(x,".nc"),readdir(joinpath(path,"member1")))
steps = [parse(Int,split(split(file,"step")[2],".")[1]) for file in member1files]
sort!(steps)

# parameter
Nens = 25
Nbits = 32
Nlon = 1800
Nlat = 901
Nvert = 91
Ntsteps = length(steps)

# preallocate/load bitwise information content
BI = zeros(Float64,Ntsteps,Nbits)
try
    BI[:,:] = load(joinpath(path,"analysis/bitinformation_$(Nens)members.jld2"),"BI")
    println("Result array loaded from file.")
catch
    println("Result array preallocated.")
end

temp = fill(0f0,Nens,Nlon,Nlat,Nvert)

for t in 1:Ntsteps
    println("Time step $(steps[t])h")

    # read the data
    print("Ensemble member ")
    for ie in 1:Nens
        print("$ie,")
        ncfile = NetCDF.open(joinpath(path,"member$ie","ensemble.t.member$ie.step$(steps[t]).ll.nc"))
        temp[ie,:,:,:] = ncfile.vars["t"][1:Nlon,1:Nlat,:]
    end
    println("All loaded.")

    # calculate bitwise information
    BI[t,:] = bitinformation(temp)
    println("Information calculated.")
    save(joinpath(path,"analysis/bitinformation_$(Nens)members.jld2"),"BI",BI)
    println("---")
end


text:
function binding_sites(ρ_P::Vector{Float64}, R_E::Vector{Float64})
    σ_EP = 0.1
    #ρ_P  = 1.35     # g/cm^3
    #R_E  = 3.29e-7  # cm
    R_P  = 20e-3
    N_A  = 6.022e23
    α_kin = zeros(size(ρ_P,1), size(R_E,1))
    for i in 1:size(ρ_P,1)
        for j in 1:size(R_E,1)
            α_kin[i,j] = @. 3*σ_EP/(pi*ρ_P[i]*R_P*R_E[j]^2*N_A)
        end
    end
    return α_kin
end

function polymer_affinity(R_E::Vector{Float64}, R_P::Vector{Float64}, V_E::Vector{Float64}, D_E::Vector{Float64})
    σ_EP = 0.1
    N_A  = 6.022e23
    p_inv = zeros(size(R_P,1), size(R_E,1))
    K_EP_0 = zeros(size(R_P,1), size(R_E,1))
    for i in 1:size(R_P,1)
        for j in 1:size(R_E,1)
            p_inv[i,j] = @. (4*σ_EP*R_P[i] + pi*R_E[j])/(4*σ_EP*R_P[i])
            K_EP_0[i,j] = @. σ_EP*R_P[i]*V_E[j]/(pi*D_E[j]*R_E[j]^2*N_A)*p_inv[i,j]
        end
    end
    return K_EP_0
end

function max_hydrolyis_rate(R_E::Vector{Float64}, R_P::Vector{Float64}, V_E::Vector{Float64})
    σ_EP = 0.1
    N_A  = 6.022e23
    ρ_P  = 1.35
    k_max_RP = @. 432*V_E*σ_EP/(12*N_A*π*R_P*R_E^2*ρ_P)
end

max_stars_repo_name: giannamars/DEBMicroTrait

max_issues_repo_name: giannamars/DEBMicroTrait

max_forks_repo_name: giannamars/DEBMicroTrait


text:
using Documenter, Stheno

makedocs(
	modules = [Stheno],
    format = Documenter.HTML(),
    sitename = "Stheno.jl",
    pages = [
        "Home" => "index.md",
        "Getting Started" => "getting_started.md",
        "Internals" => "internals.md",
        "Input Types" => "input_types.md",
    ],
)

deploydocs(repo="github.com/willtebbutt/Stheno.jl.git")


text:
"""
Counts the Differential operators for given variable x. This is used to determine
the order of a PDE.
"""
function count_differentials(term, x::Symbolics.Symbolic)
    S = Symbolics
    SU = SymbolicUtils
    if !S.istree(term)
        return 0
    else
        op = SU.operation(term)
        count_children = sum(map(arg -> count_differentials(arg, x), SU.arguments(term)))
        if op isa Differential && isequal(op.x, x)
            return 1 + count_children
        end
        return count_children
    end
end

"""
return list of differential orders in the equation
"""
function differential_order(eq, x::Symbolics.Symbolic)
    S = Symbolics
    SU = SymbolicUtils
    orders = Set{Int}()
    if S.istree(eq)
        op = SU.operation(eq)
        if op isa Differential
            push!(orders, count_differentials(eq, x))
        else
            for o in map(ch -> differential_order(ch, x), SU.arguments(eq))
                union!(orders, o)
            end
        end
    end
    return filter(!iszero, orders)
end

"""
find all the dependent variables given by depvar_ops in an expression
"""
function get_depvars(eq,depvar_ops)
    S = Symbolics
    SU = SymbolicUtils
    depvars = Set()
    if eq isa Num
       eq = eq.val
    end
    if S.istree(eq)
        if eq isa Term && any(u->isequal(operation(eq),u),depvar_ops)
              push!(depvars, eq)
        else
            for o in map(x->get_depvars(x,depvar_ops), SU.arguments(eq))
                union!(depvars, o)
            end
        end
    end
    return depvars
end

@inline function get_all_depvars(pdesys, depvar_ops)
    pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
    return collect(mapreduce(x->get_depvars(x.lhs,depvar_ops), union, pdeeqs) ∪ mapreduce(x->get_depvars(x.rhs,depvar_ops), union, pdeeqs))
end

"""
A function that creates a tuple of CartesianIndices of unit length and `N` dimensions, one pointing along each dimension.
"""
function unitindices(N::Int) #create unit CartesianIndex for each dimension
    null = zeros(Int, N)
    if N == 0
        return CartesianIndex()
    else
        return map(1:N) do i
            unit_i = copy(null)
            unit_i[i] = 1
            CartesianIndex(Tuple(unit_i))
        end
    end
end

@inline function unitindex(N, j)
    N == 0  && return CartesianIndex()
    null = zeros(Int, N)
    null[j] = 1
    CartesianIndex(Tuple(null))
end

function _split_terms(term)
    S = Symbolics
    SU = SymbolicUtils
    # TODO: Update this to be exclusive of derivatives and depvars rather than inclusive of +-/*
    if S.istree(term) && ((operation(term) == +) | (operation(term) == -) | (operation(term) == *) | (operation(term) == /))
        return mapreduce(_split_terms, vcat, SU.arguments(term))
    else
        return [term]
    end
end

function split_terms(eq::Equation)
    lhs = _split_terms(eq.lhs)
    rhs = _split_terms(eq.rhs)
    return vcat(lhs,rhs)
end

# Additional handling to get around limitations in rules
# Splits out derivatives from containing math expressions for ingestion by the rules
function _split_terms(term, x̄)
    S = Symbolics
    SU = SymbolicUtils
    st(t) = _split_terms(t, x̄)
    # TODO: Update this to handle more ops e.g. exp sin tanh etc.
    # TODO: Handle cases where two nonlinear laplacians are multiplied together
    if S.istree(term)
        # Additional handling for upwinding
        if (operation(term) == *)
            args = SU.arguments(term)
            for (i,arg) in enumerate(args)
                # Incase of upwinding, we need to keep the original term
                if S.istree(arg) && operation(arg) isa Differential
                    # Flatten the arguments of the differential to make nonlinear laplacian work in more cases
                    try
                        args[i] = operation(arg)(flatten_division.(SU.arguments(arg))...)
                    catch e
                        println("Argument to derivative in $term is not a dependant variable, is trivially differentiable or is otherwise not differentiable.")
                        throw(e)
                    end
                    return [*(flatten_division.(args)...)]
                end
            end
            return mapreduce(st, vcat, SU.arguments(term))
        elseif (operation(term) == /)
            args = SU.arguments(term)
            # Incase of upwinding or spherical, we need to keep the original term
            if S.istree(args[1])
                if args[1] isa Differential
                    try
                        args[1] = operation(arg)(flatten_division.(SU.arguments(arg))...)
                    catch e
                        println("Argument to derivative in $term is not a dependant variable, is trivially differentiable or is otherwise not differentiable.")
                        throw(e)
                    end
                    return [/(flatten_division.(args)...)]
                # Handle with care so that spherical still works
                elseif operation(args[1]) == *
                    subargs = SU.arguments(args[1])
                    # look for a differential in the arguments
                    for (i,arg) in enumerate(subargs)
                        if S.istree(arg) && operation(arg) isa Differential
                            # Flatten the arguments of the differential to make nonlinear laplacian/spherical work in more cases
                            try
                                subargs[i] = operation(arg)(flatten_division.(SU.arguments(arg))...)
                                args[1] = operation(args[1])(flatten_division.(subargs)...)
                            catch e
                                println("Argument to derivative in $term is not a dependant variable, is trivially differentiable or is otherwise not differentiable.")
                                throw(e)
                            end
                            return [/(flatten_division.(args)...)]
                        end
                    end
                end
            end
            # Basecase for division
            return vcat(st(args[1]), st(args[2]))
        elseif (operation(term) == +) | (operation(term) == -)
            return mapreduce(st, vcat, SU.arguments(term))
        elseif (operation(term) isa Differential)
            return [operation(term)(flatten_division.(SU.arguments(term))...)]
        else
            return [term]
        end
    else
        return [term]
    end
end

function split_terms(eq::Equation, x̄)
    lhs = _split_terms(eq.lhs, x̄)
    rhs = _split_terms(eq.rhs, x̄)
    return flatten_division.(vcat(lhs,rhs))
end

function split_additive_terms(eq)
    # Calling the methods from symbolicutils matches the expressions
    rhs_arg = istree(eq.rhs) && (SymbolicUtils.operation(eq.rhs) == +) ? SymbolicUtils.arguments(eq.rhs) : [eq.rhs]
    lhs_arg = istree(eq.lhs) && (SymbolicUtils.operation(eq.lhs) == +) ? SymbolicUtils.arguments(eq.lhs) : [eq.lhs]

    return vcat(lhs_arg,rhs_arg)
end

# Filthy hack to get around limitations in rules and avoid simplification to a dividing expression
@inline function flatten_division(term)
    #=rules = [@rule(/(~a, ~b) => *(~a, b^(-1.0))),
             @rule(/(*(~~a), ~b) => *(~a..., b^(-1.0))),
             @rule(/(~a, *(~~b)) => *(~a, *(~b...)^(-1.0))),
             @rule(/(*(~~a), *(~~b)) => *(~a..., *(~b...)^(-1.0)))]
    for r in rules
        if r(term) !== nothing
            return r(term)
        end
    end
    =#return term
end

@inline clip(II::CartesianIndex{M}, j, N) where M = II[j] > N ? II - unitindices(M)[j] : II

subsmatch(expr, rule) = isequal(substitute(expr, rule), expr) ? false : true

#substitute(eq::Equation, rules) = substitute(eq.lhs, rules) ~ substitute(eq.rhs, rules)

remove(args, t) = filter(x -> t === nothing || !isequal(x, t.val), args)

half_range(x) = -div(x,2):div(x,2)

@inline function _wrapperiodic(I, N, j, l)
    I1 = unitindex(N, j)
    # -1 because of the relation u[1] ~ u[end]
    if I[j] <= 1
        I = I + I1*(l  - 1)
    elseif I[j] > l
        I = I - I1*(l - 1)
    end
    return I
end

@inline function wrapperiodic(I, s, ::Val{true}, u, jx)
    j, x = jx
    return _wrapperiodic(I, ndims(u,s), j, length(s,x))
end

@inline function wrapperiodic(I, s, ::Val{false}, u, jx)
    return I
end


text:
if write_slices
    simulation.output_writers[:surface] = surface_output_writer
    simulation.output_writers[:middepth] = middepth_output_writer
    simulation.output_writers[:zonal] = zonal_output_writer
    simulation.output_writers[:meridional] = meridional_output_writer
end
if write_zonal
    simulation.output_writers[:statistics] = zonal_statistics
end
simulation.output_writers[:checkpoint] = checkpointer

max_stars_repo_name: sandreza/Mesoscale

max_issues_repo_name: sandreza/Mesoscale

max_forks_repo_name: sandreza/Mesoscale


text:
# push!(LOAD_PATH, abspath(@__DIR__,".."))
using LinearAlgebra,Statistics,StatsBase
using Plots,NamedColors ; theme(:dark)
using ProgressMeter
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using Distributions
using Random ; Random.seed!(0)
using FileIO
##

const Ne = 500
const Ni = 0
const Ntot = Ne+Ni

const Nas = 3
const p_as = 0.2

const Δt_as = 0.2
const Δt_blank = 0.1
const t_as_delay = 0.3

const Ttot = 10.0

# input to all neurons
const lowrate = 1E3
# additional input to assembly neurons
const highrate = 5E3
##

as_idxs = map(_->findall(rand(Ne) .< p_as),1:Nas)

inputfun=S.pattern_functor(Δt_as,Ttot,lowrate,highrate,Ntot,
  as_idxs;t_pattern_delay=t_as_delay,Δt_pattern_blank=Δt_blank)
inputfun_upper = S.pattern_functor_upperlimit(lowrate,highrate,Ntot,as_idxs)

## let's visualize it !
_  = let times = range(-1,Ttot+3;length=500)
  ys = hcat([inputfun.(t,1:Ne) for t in times]...)
  heatmap(times,1:Ntot,ys)
end

idxs_sort = S.order_by_pattern_idxs(as_idxs,Ne)

_  = let times = range(-1,Ttot+3;length=500)
  ys = hcat([inputfun.(t,1:Ne) for t in times]...)
  heatmap(times,1:Ntot,ys[idxs_sort,:])
end

##
# let's say kernel and v_reversal are same as E population
# excitatory neurons are Auguste's
const taueplus = 6E-3
const taueminus = 1E-3
const eifslope = 2.0
const τe = 20E-3
const Cap = 300.0
const vthexp = -50.0
const v_rev_e = 0.0
const vth_e = 20.0
const v_reset_e = -60.0
const v_rest_e = -60.0
const τrefr  = 2E-3
nt_e = let sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SpikeGenEIF(vthexp,eifslope)
  S.NTLIFConductance(sker,sgen,τe,Cap,
   vth_e,v_reset_e,v_rest_e,τrefr,v_rev_e)
end
ps_e = S.PSLIFConductance(nt_e,Ne)

const w_in = 3.0
nt_in = let sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SGPoissonFExact(inputfun,inputfun_upper)
  S.NTInputConductance(sgen,sker,v_rev_e)
end
ps_in = S.PSInputPoissonConductanceExact(nt_in,w_in,Ntot)

## now define E to E connections
# with triplets plasticity rule!
const jee = 6.0 # this should scale as sqrt(Ne), or something like that
const pee = 0.2
wmat_start = S.sparse_constant_wmat(Ne,Ne,pee,jee;no_autapses=true)
# now the plasticity
# parameters from Augiste's code
triplets_plasticity = let (n_post,n_pre) = size(wmat_start),
  τplus = 17E-3 # 16.8
  τminus = 34E-3 # 33.7
  τx = 100E-3 # 101
  τy = 120E-3 # 125
  plast_eps = 5E-5 # 1E-3
  A2plus = 7.5E-7*plast_eps
  A3plus = 9.3*plast_eps
  A2minus = 7.0*plast_eps
  A3minus = 2.3E-1*plast_eps
  S.PlasticityTriplets(τplus,τminus,τx,τy,A2plus,A3plus,
    A2minus,A3minus,n_post,n_pre)
end

conn_e_e = S.ConnGeneralIF2(copy(wmat_start),triplets_plasticity)

## Now add stabilization to E neurons

ps_istab = let v_rev_i = v_reset_e
  Aloc = 100.0 # 500
  Aglo = 3000.0 # 2000
  τloc = 50E-3
  τglo = 8E-3
  S.PSConductanceInputInhibitionStabilization(v_rev_i,Aglo,Aloc,τglo,τloc,Ne)
end

##
pop_e = S.Population(ps_e,
  (S.FakeConnection(),ps_in),
  (S.FakeConnection(),ps_istab),
  (conn_e_e,ps_e))

## network
const dt = 0.1E-3
ntw = S.RecurrentNetwork(dt,pop_e)

# record spiketimes and internal potential
krec = 1
n_e_rec = 3
t_wup = 0.0
rec_state_e = S.RecStateNow(ps_e,krec,dt,Ttot;idx_save=collect(1:n_e_rec),t_warmup=t_wup)
rec_spikes_e = S.RecSpikes(ps_e,500.0,Ttot;idx_save=collect(1:Ne),t_warmup=t_wup)

## Run

times = (0:ntw.dt:Ttot)
nt = length(times)
# clean up
S.reset!.([rec_state_e,rec_spikes_e])
S.reset!.([ps_e,ps_in,ps_istab])
# initial conditions
ps_e.state_now .= v_reset_e #rand(Uniform(v_reset_e,v_reset_e+0.3),Ne)

@time begin
  @showprogress 5.0 "network simulation " for (k,t) in enumerate(times)
    rec_state_e(t,k,ntw)
    rec_spikes_e(t,k,ntw)
    S.dynamics_step!(t,ntw)
  end
end

S.add_fake_spikes!(1.0vth_e,rec_state_e,rec_spikes_e)

## visualize activity

_ = let neu = 1,
  plt=plot()
  plot!(plt,rec_state_e.times,rec_state_e.state_now[neu,:];linewidth=2,leg=false)
end

##
idxs_sort = S.order_by_pattern_idxs(as_idxs,Ne) |> reverse
_ = let Traster = Ttot
  tmp = S.raster_png(1E-3,rec_spikes_e;spike_height=3,Ttot=Traster,Nneurons=Ne,
   reorder=idxs_sort)
  save("/tmp/tmp.png",tmp)
end

##
_ = let (_times,binned) =S.binned_spikecount(0.1,rec_spikes_e)
  heatmap(_times,1:size(binned,1),binned[idxs_sort,:])
end

# order E neurons by assembly

##
w_start = copy(wmat_start)
w_end = copy(conn_e_e.weights)
w_diff = w_end .- w_start

##
histogram(nonzeros(w_start))
histogram(nonzeros(w_end))
histogram(nonzeros(w_diff))

heatmap(Matrix(w_end)[idxs_sort,idxs_sort])
heatmap(Matrix(w_diff)[idxs_sort,idxs_sort])

function plus_minus_rescale(mat::Array{R}) where R
  down,up = extrema(mat)
  @assert (down < 0.0) && (up > 0.0)
  pluspart = mat ./ up
  minuspart = mat ./ abs(down)
  ret = similar(mat)
  for i in eachindex(mat)
    if mat[i] > 0
      ret[i] = pluspart[i]
    else
      ret[i] = minuspart[i]
    end
  end
  return ret
end

theme(:sand)
heatmap(plus_minus_rescale(Matrix(w_diff))[idxs_sort,idxs_sort];
  color=:seismic,ratio=1,ylabel="postsynapic",xlabel="presynaptic",
  colorbar_title="normalized weight change",xlims=(0,Ne))
savefig("/tmp/w.svg")


text:
# function that handles the dispatch of cars
function incident_dispatch!(sim_data::DataFrame,
                            incidents::DataFrame,
                            ressource_flow::Array{Int64,3},
                            drivingtime::Array{Float64,2},
                            traffic::Array{Float64,2},
                            locations::Vector{Int64},
                            current_case::Int64,
                            now::Int64,
                            district::Int64,
                            k::Int64,
                            dispatch::String,
                            fastest_time::Int64)
#  create the neccessary variables
    inc::Int64           = incidents[current_case,:location] # incident_location
    dc::Int64            = locations[district]               # location district center

    # create a vector with all necessary timestamps
    # timestamps[1]: minute of arrival at incident
    # timestamps[2]: minute of departure at incident
    # timestamps[3]: minute of arrival at district center
    timestamps = Vector{Int64}(undef,3)

# calculate the minute the car will arrive at the incident location
    if fastest_time == 0
        timestamps[1] = now + dispatch_drivingtime(drivingtime[dc,inc],traffic,now)
    else
        timestamps[1] = now + fastest_time
    end

# calculate the minute of arrival at the incident and the minute of departure from the incident
    if sim_data[current_case,:cars_dispatched] == 0
        timestamps[2] = timestamps[1] + incidents[current_case,:length]
        timestamps[3] = timestamps[2] + dispatch_drivingtime(drivingtime[inc,dc],traffic,timestamps[2])
        ressource_flow[timestamps[3],district,7] += incidents[current_case,:backlog]
    else
        timestamps[2] = min(timestamps[1] + incidents[current_case,:length], timestamps[1] +
                        incidents[current_case,:length] - sim_data[current_case,:dispatch_minute_first] + now)
        timestamps[3] = timestamps[2] + dispatch_drivingtime(drivingtime[inc,dc],traffic,timestamps[2])
    end

# calculate the number of cars to dispatch to the incident
    # cars[1]:  number of cars dispatched from location currently free
    # cars[2]:  number of cars dispatched from location currently working on backlog
    # cars[3]:  number of cars dispatched currently on patrol
    cars = Vector{Int64}(undef,3) .= 0

    # k = 1: only one car will be dispatched
    # k = 2: the number of dispatched cars dependends on the availability
    if dispatch == "own_department"
        if ressource_flow[now,district,1] > 0
            if k == 1
                cars[1] = 1
            else
                cars[1] = min(sim_data[current_case,:cars_missing], ressource_flow[now,district,1])
            end
        else
            cars[2] = 1
        end
    elseif dispatch == "own_patrol"
        cars[3] = 1
    elseif dispatch == "exchange_department"
        cars[1] = 1
    end

    # calculate the numer of dispatched cars
    total = sum(cars)

    # assign the calculated capacities into the ressource flow matrix
    ressource_flow[now,district,1]             -= cars[1]
    ressource_flow[now,district,5]             -= cars[2]
    ressource_flow[now,district,6]             -= cars[3]
    ressource_flow[now,district,2]             += total
    ressource_flow[timestamps[1],district,2]   -= total
    ressource_flow[timestamps[1],district,3]   += total
    ressource_flow[timestamps[2],district,3]   -= total
    ressource_flow[timestamps[2],district,4]   += total
    ressource_flow[timestamps[3],district,4]   -= total
    ressource_flow[timestamps[3],district,1]   += total

    # write the resulting times to the sim_data DataFrame
    if sim_data[current_case,:cars_dispatched] == 0
        sim_data[current_case,:location_dispatched_first] = locations[district]
        sim_data[current_case,:dispatch_minute_first]     = now
        sim_data[current_case,:arrival_minute_first]      = timestamps[1]
    end
    sim_data[current_case,:arrival_minute_first]    = min(sim_data[current_case,:arrival_minute_first],timestamps[1])
    sim_data[current_case,:cars_dispatched]         += total
    sim_data[current_case,:cars_missing]            -= total
    if sim_data[current_case,:cars_missing] == 0
        sim_data[current_case,:dispatch_minute_all] = now
        sim_data[current_case,:arrival_minute_all]  = timestamps[1]
    end
    if  dispatch == "own_department" || dispatch == "own_patrol"
        sim_data[current_case,:cars_location_responsible] += total
    end
end


text:
# Neural stick-breaking processe

struct MeanFieldSBC <: StochasticLayer
    kuma
    gumbel
end

function MeanFieldSBC(i_dim::Int, k_init::Int=0; kwargs...)
    kuma = KumaraswamyLazyDense(i_dim, k_init; kwargs...)
    gumbel = GumbelBernoulliLogitLazyDense(i_dim, k_init; kwargs...)
    return MeanFieldSBC(kuma, gumbel)
end

function (enc::MeanFieldSBC)(i, k::Int; lowerbound=(FT == Float64 ? FT(0.1) : FT(0.2)))
    dist_kuma = enc.kuma(i, k)
    dist_nu = BatchKumaraswamy(dist_kuma.a .+ lowerbound, dist_kuma.b .+ lowerbound)
    dist_gumbel = enc.gumbel(i, k)
    return dist_nu, dist_gumbel
end

struct StructuredSBC <: StochasticLayer
    a
    b
    gumbel
    α
    β
end

"""
Helper function to initialise StructuredSBC with an IBP.
"""
function StructuredSBC(ibp::IBP{T}, i_dim::Int, k_init::Int=0; kwargs...) where {T<:AbstractFloat}
    return StructuredSBC(i_dim, k_init, ibp.α, one(T); kwargs...)
end

function StructuredSBC(i_dim::Int, k_init::Int=0, α::AbstractFloat=one(FT), β::AbstractFloat=one(FT); kwargs...)
    a = Knet.Param(AT(zeros(FT, k_init, 1) .+ invsoftplus(α)))
    b = Knet.Param(AT(zeros(FT, k_init, 1) .+ invsoftplus(β)))
    gumbel = GumbelBernoulliLogitLazyDense(i_dim, k_init; kwargs...)
    return StructuredSBC(a, b, gumbel, α, β)
end

function (enc::StructuredSBC)(i, k::Int=size(enc.a, 1); α::AbstractFloat=enc.α, β::AbstractFloat=enc.β, lowerbound=(FT == Float64 ? FT(0.1) : FT(0.2)))
    if k > size(enc.a, 1)
        k_diff = k - size(enc.a, 1)
        enc.a.value = vcat([enc.a.value, AT(zeros(FT, k_diff, 1) .+ invsoftplus(α))]...)
        enc.b.value = vcat([enc.b.value, AT(zeros(FT, k_diff, 1) .+ invsoftplus(β))]...)
    end
    dist_nu = BatchKumaraswamy(softplus.(enc.a[1:k,:]) .+ lowerbound,
                               softplus.(enc.b[1:k,:]) .+ lowerbound)
    dist_gumbel = enc.gumbel(i, k)
    return dist_nu, dist_gumbel
end


text:
# a lightweight CUDA array type for testing purposes

## ctor & finalizer

mutable struct CuTestArray{T,N}
    buf::Mem.DeviceBuffer
    shape::NTuple{N,Int}
    function CuTestArray{T,N}(shape::NTuple{N,Int}) where {T,N}
        len = prod(shape)
        buf = Mem.alloc(Mem.Device, len*sizeof(T))

        obj = new{T,N}(buf, shape)
        finalizer(unsafe_free!, obj)
        return obj
    end
end

function unsafe_free!(a::CuTestArray)
    CUDAdrv.isvalid(a.buf.ctx) && Mem.free(a.buf)
end

Base.cconvert(::Type{<:CuPtr}, x::CuTestArray) = x.buf

Base.length(xs::CuTestArray) = prod(xs.shape)

## memory copy operations

function CuTestArray(src::Array{T,N}) where {T,N}
    dst = CuTestArray{T,N}(size(src))
    Mem.copy!(dst.buf, pointer(src), length(src) * sizeof(T))
    return dst
end

function Base.Array(src::CuTestArray{T,N}) where {T,N}
    dst = Array{T,N}(undef, src.shape)
    Mem.copy!(pointer(dst), src.buf, length(src) * sizeof(T))
    return dst
end

## conversions

using Adapt
function Adapt.adapt_storage(::CUDAnative.Adaptor, a::CuTestArray{T,N}) where {T,N}
    ptr = convert(CuPtr{T}, a.buf)
    devptr = CUDAnative.DevicePtr{T,AS.Global}(ptr)
    CuDeviceArray{T,N,AS.Global}(a.shape, devptr)
end


text:
using TrajectoryOptimization: generate_general_constraint_jacobian, is_inplace_function, count_inplace_output, stage_cost
using Test

""" Simple Pendulum """
model, = TrajectoryOptimization.Dynamics.pendulum
n = model.n
m = model.m
x0 = [0; 0.];
xf = [pi; 0]; # (ie, swing up)
u0 = [1]
Q = 1e-3*Diagonal(I,n)
Qf = 100. *Diagonal(I,n)
R = 1e-3*Diagonal(I,m)
tf = 5.

costfun = LQRCost(Q,R,Qf,xf)
@test_nowarn UnconstrainedObjective(costfun, tf, x0, xf)
obj_uncon = UnconstrainedObjective(costfun, tf, x0, xf)
@test obj_uncon.tf == tf

# Test minimum time constructor
obj_uncon = UnconstrainedObjective(costfun, :min, x0, xf)
@test obj_uncon.tf == 0

# Try invalid inputs
tf = -1.
@test_throws ArgumentError("tf must be non-negative") UnconstrainedObjective(costfun, tf, x0, xf)
tf = 1.; R_ = Diagonal(I,m)*-1.
@test_throws ArgumentError("R must be positive definite") LQRObjective(Q, R_, Qf, tf, x0, xf)
@test_throws ArgumentError(":min is the only recognized Symbol for the final time") obj_uncon = UnconstrainedObjective(costfun, :max, x0, xf)

function myfun(c,x,u)
    c[1:2] = x + u
end
c = zeros(2)
myfun(c,ones(2),ones(2))
is_inplace_function(myfun,ones(2),ones(2))
count_inplace_output(myfun,ones(2),ones(2))

### Constraints ###
# Test defaults
obj = ConstrainedObjective(costfun,tf,x0,xf)
@test obj.u_min == [-Inf]
@test obj.u_max == [Inf]
@test obj.x_min == -[Inf,Inf]
@test obj.x_max == [Inf,Inf]
@test isa(obj.cI(ones(0),x0,u0),Nothing)
@test isa(obj.cE(ones(0),x0,u0),Nothing)

@test obj.p == 0
@test obj.use_xf_equality_constraint == true
@test obj.p_N == 2

# Use scalar control constraints
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-1,u_max=1)
@test obj.p == 2
@test obj.p_N == 2

# Single-sided
obj = ConstrainedObjective(costfun,tf,x0,xf,u_max=1)
@test obj.p == 1
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=1, u_max=Inf)
@test obj.p == 1

# Error testing
@test_throws ArgumentError ConstrainedObjective(costfun,tf,x0,xf,u_min=1, u_max=-1)
@test_throws DimensionMismatch ConstrainedObjective(costfun,tf,x0,xf,u_min=[1], u_max=[1,2])

# State constraints
obj = ConstrainedObjective(costfun,tf,x0,xf,x_min=-[1,2], x_max=[1,2],use_xf_equality_constraint=false)
@test obj.p == 4
@test obj.pI == 4
@test obj.pI_N == 4

@test_throws DimensionMismatch ConstrainedObjective(costfun,tf,x0,xf,x_min=-[Inf,2,3,4], x_max=[1,Inf,3,Inf])
obj = ConstrainedObjective(costfun,tf,x0,xf,x_min=-[Inf,4], x_max=[3,Inf],use_xf_equality_constraint=false)
@test obj.p == 2
@test obj.pI == 2
@test obj.pI_N == 2

# Scalar to array constraint
obj = ConstrainedObjective(costfun,tf,x0,xf,x_min=-4, x_max=4)
@test obj.p == 4
@test obj.x_max == [4,4]

# Custom constraints
function cI(cdot,x,u)
    cdot[1] = x[2]+u[1]-2
end
obj = ConstrainedObjective(costfun,tf,x0,xf,cI=cI)
@test obj.p == 1
@test obj.pI == 1
obj = ConstrainedObjective(costfun,tf,x0,xf,cE=cI)
@test obj.p == 1
@test obj.pI == 0

# Construct from unconstrained
obj = ConstrainedObjective(obj_uncon)
@test obj.u_min == [-Inf]
@test obj.u_max == [Inf]
@test obj.x_min == -[Inf,Inf]
@test obj.x_max == [Inf,Inf]
@test isa(obj.cI(ones(0),x0,u0),Nothing)
@test isa(obj.cE(ones(0),x0,u0),Nothing)
@test obj.p == 0
@test obj.use_xf_equality_constraint == true
@test obj.p_N == 2

obj = ConstrainedObjective(obj_uncon, u_min=-1)
@test obj.p == 1

# Update objectve
obj = update_objective(obj, u_max=2, x_max = 4, cE=cI)
@test obj.p == 5

# Minimum time
c = 0.1
obj = ConstrainedObjective(costfun,tf,x0,xf, x_max=2)
@test obj.p == 2
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2)
@test obj.p == 1
@test obj.tf == tf
tf_ = :min
obj = ConstrainedObjective(costfun,tf_,x0,xf,u_min=-2)
@test obj.tf == 0
@test obj.p == 1
obj = ConstrainedObjective(costfun,tf_,x0,xf,u_min=-2)
@test obj.tf == 0

# Test constraint function
function cI!(cres,x,u)
    cres[1] = x[1]*x[2] + u[1]
    cres[2] = u[1]*x[2] + 3x[1]
end
function cE!(cres,x,u)
    cres[1] = x[1]^2
end
function cI!(cres,x)
    cres[1] = x[1]^2
    cres[2] = x[2]^2
end
function cE!(cres,x)
    cres[1] = x[1] + x[2] - 5
end
c = zeros(2)
cI!(c,ones(2))
is_inplace_function(cI!,ones(n))
is_inplace_function(cE!,ones(n))
count_inplace_output(cE!,ones(n))

obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cI=cI!, cE=cE!, cI_N=cI!, cE_N=cE!, use_xf_equality_constraint=false)
@test obj.p_N == 7
@test obj.pI_N == 6
@test obj.p == 9
@test obj.pI_custom == 2
@test obj.pE_custom == 1
@test obj.pI_N_custom == 2
@test obj.pE_N_custom == 1
@test_throws ArgumentError ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cI=cI!, cE=cE!, cI_N=cI!, cE_N=cE!)
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cI=cI!, cE=cE!)
@test obj.p_N == length(x0)
@test obj.pI_N == 0
@test obj.p == 9
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cE=cE!)
@test obj.p_N == length(x0)
@test obj.p == 7
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cE_N=cE!, use_xf_equality_constraint=false)
@test obj.p_N == 5
@test obj.p == 6

@test get_sizes(obj) == (2,1)
obj = ConstrainedObjective(costfun,tf,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cI=cI!, cE=cE!)
c,c_jacob,c_labels = TrajectoryOptimization.generate_constraint_functions(obj)
cres = zeros(9)
x = [1,5]
u = [1]
cans = [0,-3,-3,1,-4,-8,6,8,1]
c(cres,x,u)
@test cres == cans

cresN = zeros(2)
c(cresN,x)
cresN
cansN = [x[1] - obj.xf[1]; x[2] - obj.xf[2]]
@test cresN == cansN
# test with minimum time
obj = ConstrainedObjective(costfun,:min,x0,xf,u_min=-2,u_max=1,x_min=-3,x_max=4, cI=cI!, cE=cE!)
@test obj.p == 9
c, c_jacob = TrajectoryOptimization.generate_constraint_functions(obj,max_dt=1.0,min_dt=0.01)
cres = zeros(12)
u_dt = [u; 0.1]
cans = [0,-0.9,-3,sqrt(0.01)-0.1,-3,1,-4,-8,6,8,1,0.1-0.9]
c(cres,[x; 0.9],u_dt)
@test cres == cans

# Change upper bound on dt
c, c_jacob = TrajectoryOptimization.generate_constraint_functions(obj,max_dt=10.)
cres = zeros(12)
cans = [0,.1-sqrt(10), -3,0  ,-3,1,-4,-8,  6,8,1,0.1-0.9]
c(cres,[x;0.9],u_dt)
@test cres == cans

# use infeasible start
u_inf = [u_dt; -1; -1]
cans = [0,.1-sqrt(10), -3,0  ,-3,1,-4,-8,  6,8,1, -1,-1, 0.1-0.9]
cres = zeros(14)
c(cres,[x;0.9],u_inf)
cres
@test cres == cans

# Verify jacobians
∇x_cI(x,u) = [x[2] x[1]; 3 u[1]]
∇u_cI(x,u) = [1; x[2]]
∇x_cE(x,u) = [2x[1] 0]
∇u_cE(x,u) = [0,]
cx = [zeros(2(m+1),n+1) ;
      Matrix(I,n,n) zeros(n);
     -Matrix(I,n,n) zeros(n);
      ∇x_cI(x,u) zeros(2);
      ∇x_cE(x,u) zeros(1);
      zeros(n,n) zeros(n);
      zeros(1,n) -1.]
cu = [Matrix(I,m+1,m+1)     zeros(m+1,n);
     -Matrix(I,m+1,m+1)     zeros(m+1,n);
      zeros(2n,m+1)         zeros(2n,n);
      ∇u_cI(x,u) zeros(2,1) zeros(2,n);
      ∇u_cE(x,u) zeros(1,1) zeros(1,n);
      zeros(n,m+1)       Matrix(I,n,n);
      zeros(1,m) 1.0 zeros(1,n)]
cx_res,cu_res = zero(cx),zero(cu)
c_jacob(cx_res,cu_res,[x; 0.9],u_inf)
@test cx_res == cx
@test cu_res == cu

### GENERAL CONSTRAINTS JACOBIAN ###
n, m = 3,2
function cE(cdot,x,u)
     cdot[1:2] = 2x[1:2]+u
     cdot[3] = x'x + 5
end
pE = 3

x = [1;2;3]
u = [1;-1]
cdot = zeros(3)
cE(cdot,x,u)
@test cdot == [3; 3; 19]

# Jacobians
jac_cE = generate_general_constraint_jacobian(cE,pE,n,m)
jac_x(x,u) = [2 0 0;
              0 2 0;
              2x']
jac_u(x,u) = [1 0;
              0 1;
              0 0]
A = zeros(3,3)
B = zeros(3,2)
jac_cE(A,B,x,u)
@test A == jac_x(x,u)
@test B == jac_u(x,u)

# Add terminal
function _cE(xdot,x)
    xdot[1:2] = [cos(x[1]) + x[2]*x[3]; x[1]*x[2]^2]
end

pE_N = 2

jac_xN(x) = [-sin(x[1]) x[3] x[2]; x[2]^2 2x[1]*x[2] 0]

jac_cE = generate_general_constraint_jacobian(cE,pE,n,m)

_jac_cE = generate_general_constraint_jacobian(_cE,pE_N,n)
tmp1 = zeros(pE_N,n)
_jac_cE(tmp1,x)
@test tmp1 == jac_xN(x)

A1 = zeros(3,3)
B1 = zeros(3,2)
jac_cE(A1,B1,x,u)
@test A1 == jac_x(x,u)
@test B1 == jac_u(x,u)

# COST FUNCTION TESTS
using TrajectoryOptimization: taylor_expansion, stage_cost
n,m = 3,2
Q = Diagonal([1.,2,3])
R = Diagonal([4.,5])
Qf = Diagonal(ones(n)*10)
xf = Vector{Float64}(1:n)
x0 = zeros(n)

x = ones(n)
u = ones(m)*2

LinQuad = LQRCost(Q,R,Qf,xf)
J = stage_cost(LinQuad,x,u)
@test taylor_expansion(LinQuad,x,u) == (Q,R,zeros(m,n),Q*(x-xf),R*u)
@test taylor_expansion(LinQuad,x) == (Qf,Qf*(x-xf))

# Generic Cost Function
my_stage_cost(x,u) = x[1]^2 + 2*x[2]*x[3] + x[3] + 3*u[1] + u[2]^2 + u[2]*u[1] + x[2]*u[1] + log(x[3]) + sin(u[2])
my_final_cost(x) = (x[1] - 1)^2 + x[1]*x[2]
qfun(x,u) = [2*x[1], 2*x[3] + u[1], 1 + 2*x[2] + 1/x[3]]
rfun(x,u) = [3 + x[2] + u[2], 2*u[2] + u[1] + cos(u[2])]
Qfun(x,u) = [2 0 0;
          0 0 2;
          0 2 -1/x[3]^2]
Rfun(x,u) = [0 1; 1 2-sin(u[2])]
Hfun(x,u) = [0 0; 1 0; 0 0]'
qffun(x) = [2(x[1] -1) + x[2], x[1], 0]
Qffun(x) = [2x[1] 1 0; 1 0 0; 0 0 0]

mycost = GenericCost(my_stage_cost,my_final_cost,n,m)
stage_expansion = (Qfun(x,u), Rfun(x,u), Hfun(x,u), qfun(x,u), rfun(x,u))
@test taylor_expansion(mycost,x,u) == stage_expansion
@test taylor_expansion(mycost,x) == (Qffun(x), qffun(x))

mygrad(x,u) = qfun(x,u), rfun(x,u)
myhess(x,u) = Qfun(x,u), Rfun(x,u), Hfun(x,u)
myexpansion(x,u) = Qfun(x,u), Rfun(x,u), Hfun(x,u), qfun(x,u), rfun(x,u)
mygrad(x) = qffun(x)
myhess(x) = Qffun(x)
myexpansion(x) = Qffun(x), qffun(x)

mycost2 = GenericCost(my_stage_cost,my_final_cost,myexpansion,n,m)
@test taylor_expansion(mycost2,x,u) == stage_expansion
@test taylor_expansion(mycost2,x) == (Qffun(x), qffun(x))

mycost2 = GenericCost(my_stage_cost,my_final_cost,mygrad,myhess,n,m)
@test taylor_expansion(mycost2,x,u) == stage_expansion
@test taylor_expansion(mycost2,x) == (Qffun(x), qffun(x))

# Unconstrained Objective
costfun = LinQuad
obj = UnconstrainedObjective(costfun,:min,x0,xf)
@test obj.tf == 0
obj = UnconstrainedObjective(costfun,tf,x0,xf)
@test_throws ArgumentError UnconstrainedObjective(costfun,tf,x0,u)
UnconstrainedObjective(costfun,tf,x0,Float64[])
@test stage_cost(obj.cost,x,u) == J
@test stage_cost(obj,x,u) == J


text:
function Pool(
            ctx::Context,
            user::String,
            password::String,
            connect_string::String,
            common_params::CommonCreateParams,
            pool_create_params::OraPoolCreateParams
        )

    ora_common_params = OraCommonCreateParams(ctx, common_params)
    dpi_pool_handle_ref = Ref{Ptr{Cvoid}}()
    result = dpiPool_create(ctx.handle, user, password, connect_string, Ref(ora_common_params), Ref(pool_create_params), dpi_pool_handle_ref)
    error_check(ctx, result)

    # retrieves pool name from OraPoolCreateParams
    @assert pool_create_params.out_pool_name != C_NULL "Failed to create a Pool: no pool name was returned."
    pool_name = unsafe_string(pool_create_params.out_pool_name, pool_create_params.out_pool_name_length)

    return Pool(ctx, dpi_pool_handle_ref[], pool_name)
end

function Pool(
            ctx::Context,
            user::String,
            password::String,
            connect_string::String;
            encoding::AbstractString=DEFAULT_CONNECTION_ENCODING,
            nencoding::AbstractString=DEFAULT_CONNECTION_NENCODING,
            create_mode::Union{Nothing, OraCreateMode}=nothing,
            edition::Union{Nothing, String}=nothing,
            driver_name::Union{Nothing, String}=nothing,
            enable_soda_metadata_cache::Bool=false,
            stmt_cache_size::Integer=0,
            min_sessions::Union{Nothing, Integer}=nothing,
            max_sessions::Union{Nothing, Integer}=nothing,
            session_increment::Union{Nothing, Integer}=nothing,
            ping_interval::Union{Nothing, Integer}=nothing,
            ping_timeout::Union{Nothing, Integer}=nothing,
            homogeneous::Union{Nothing, Bool}=nothing,
            external_auth::Union{Nothing, Bool}=nothing,
            get_mode::Union{Nothing, OraPoolGetMode}=nothing,
            timeout::Union{Nothing, Integer}=nothing,
            wait_timeout::Union{Nothing, Integer}=nothing,
            max_lifetime_session::Union{Nothing, Integer}=nothing
        )

    function EmptyOraPoolCreateParams(ctx::Context)
        pool_create_params_ref = Ref{OraPoolCreateParams}(OraPoolCreateParams(UInt32(1), UInt32(1), UInt32(0), Int32(60), Int32(5000), Int32(1), Int32(0), ORA_MODE_POOL_GET_NOWAIT, C_NULL, UInt32(0), UInt32(0), UInt32(0), UInt32(0), C_NULL, 0))
        result = dpiContext_initPoolCreateParams(ctx.handle, pool_create_params_ref)
        error_check(ctx, result)
        return pool_create_params_ref[]
    end

    pool_create_params = EmptyOraPoolCreateParams(ctx)

    # parse optional pool parameters
    @parse_opt_field_param(pool_create_params, min_sessions, UInt32)
    @parse_opt_field_param(pool_create_params, max_sessions, UInt32)
    @parse_opt_field_param(pool_create_params, session_increment, UInt32)
    @parse_opt_field_param(pool_create_params, ping_interval, Int32)
    @parse_opt_field_param(pool_create_params, ping_timeout, Int32)
    @parse_opt_field_param(pool_create_params, homogeneous, Int32)
    @parse_opt_field_param(pool_create_params, external_auth, Int32)
    @parse_opt_field_param(pool_create_params, get_mode)
    @parse_opt_field_param(pool_create_params, timeout, UInt32)
    @parse_opt_field_param(pool_create_params, wait_timeout, UInt32)
    @parse_opt_field_param(pool_create_params, max_lifetime_session, UInt32)

    common_params = CommonCreateParams(create_mode, encoding, nencoding, edition, driver_name, enable_soda_metadata_cache, stmt_cache_size)
    return Pool(ctx, user, password, connect_string, common_params, pool_create_params)
end

function Pool(
            user::String,
            password::String,
            connect_string::String;
            encoding::AbstractString=DEFAULT_CONNECTION_ENCODING,
            nencoding::AbstractString=DEFAULT_CONNECTION_NENCODING,
            create_mode::Union{Nothing, OraCreateMode}=nothing,
            edition::Union{Nothing, String}=nothing,
            driver_name::Union{Nothing, String}=nothing,
            enable_soda_metadata_cache::Bool=false,
            stmt_cache_size::Integer=0,
            min_sessions::Union{Nothing, Integer}=nothing,
            max_sessions::Union{Nothing, Integer}=nothing,
            session_increment::Union{Nothing, Integer}=nothing,
            ping_interval::Union{Nothing, Integer}=nothing,
            ping_timeout::Union{Nothing, Integer}=nothing,
            homogeneous::Union{Nothing, Bool}=nothing,
            external_auth::Union{Nothing, Bool}=nothing,
            get_mode::Union{Nothing, OraPoolGetMode}=nothing,
            timeout::Union{Nothing, Integer}=nothing,
            wait_timeout::Union{Nothing, Integer}=nothing,
            max_lifetime_session::Union{Nothing, Integer}=nothing
        )
    return Pool(Context(), user, password, connect_string,
            encoding=encoding,
            nencoding=nencoding,
            create_mode=create_mode,
            edition=edition,
            driver_name=driver_name,
            enable_soda_metadata_cache=enable_soda_metadata_cache,
            stmt_cache_size=stmt_cache_size,
            min_sessions=min_sessions,
            max_sessions=max_sessions,
            session_increment=session_increment,
            ping_interval=ping_interval,
            ping_timeout=ping_timeout,
            homogeneous=homogeneous,
            external_auth=external_auth,
            get_mode=get_mode,
            timeout=timeout,
            wait_timeout=wait_timeout,
            max_lifetime_session=max_lifetime_session
        )
end

function pool_get_mode(pool::Pool) :: OraPoolGetMode
    pool_get_mode_ref = Ref{OraPoolGetMode}()
    result = dpiPool_getGetMode(pool.handle, pool_get_mode_ref)
    error_check(context(pool), result)
    return pool_get_mode_ref[]
end

function pool_get_open_count(pool::Pool) :: UInt32
    result = Ref{UInt32}()
    err = dpiPool_getOpenCount(pool.handle, result)
    error_check(context(pool), err)
    return result[]
end

function pool_get_busy_count(pool::Pool) :: UInt32
    result = Ref{UInt32}()
    err = dpiPool_getBusyCount(pool.handle, result)
    error_check(context(pool), err)
    return result[]
end

function Connection(pool::Pool; auth_mode::OraAuthMode=ORA_MODE_AUTH_DEFAULT)
    conn_create_params = ConnCreateParams(auth_mode, pool)
    ora_conn_create_params = OraConnCreateParams(context(pool), conn_create_params)
    connection_handle_ref = Ref{Ptr{Cvoid}}()
    result = dpiPool_acquireConnection(pool.handle, Ref(ora_conn_create_params), connection_handle_ref)
    error_check(context(pool), result)
    return Connection(context(pool), connection_handle_ref[], pool)
end

function close(pool::Pool; close_mode::OraPoolCloseMode=ORA_MODE_POOL_CLOSE_DEFAULT)
    result = dpiPool_close(pool.handle, close_mode)
    error_check(context(pool), result)
    nothing
end