text:
derivative_order(coefficients::AbstractDerivativeCoefficients) = coefficients.derivative_order
accuracy_order(coefficients::AbstractDerivativeCoefficients) = coefficients.accuracy_order
Base.eltype(coefficients::AbstractDerivativeCoefficients{T}) where {T} = T
LinearAlgebra.issymmetric(coefficients::AbstractDerivativeCoefficients) = coefficients.symmetric

derivative_order(D::AbstractDerivativeOperator) = derivative_order(D.coefficients)
accuracy_order(D::AbstractDerivativeOperator) = accuracy_order(D.coefficients)
Base.eltype(D::AbstractDerivativeOperator{T}) where {T} = T
LinearAlgebra.issymmetric(D::AbstractDerivativeOperator) = issymmetric(D.coefficients)
function Base.size(D::AbstractDerivativeOperator)
    N = length(grid(D))
    (N, N)
end
function Base.size(D::AbstractDerivativeOperator, i::Int)
    if i <= 0
        error("arraysize: dimension out of range")
    elseif i <= 2
        length(grid(D))
    else
        1
    end
end

@inline grid(D::AbstractDerivativeOperator) = D.grid
@inline grid(D::AbstractPeriodicDerivativeOperator) = D.grid_compute

Base.@propagate_inbounds function mul!(dest, D::AbstractDerivativeOperator, u)
    mul!(dest, D, u, one(eltype(dest)))
end

@noinline function *(D::AbstractDerivativeOperator, u)
    @boundscheck begin
        @argcheck size(D,1) == size(D,2) == length(u) DimensionMismatch
    end
    T = promote_type(eltype(D), eltype(u))
    dest = similar(u, T); fill!(dest, zero(T))
    @inbounds mul!(dest, D, u)
    dest
end

function Base.Matrix(D::AbstractDerivativeOperator{T}) where {T}
    v = Array{T}(undef, size(D, 2)...)
    fill!(v, T(0))
    A = Array{T}(undef, size(D)...)
    for i in 1:size(D,2)
        v[i] = T(1)
        mul!(view(A,:,i), D, v)
        v[i] = T(0)
    end
    A
end

function SparseArrays.sparse(D::AbstractDerivativeOperator{T}) where {T}
    M, N = size(D)
    rowind = Vector{Int}()
    nzval = Vector{T}()
    colptr = Vector{Int}(undef, N+1)
    v = fill(zero(T), N)
    dest = Array{T}(undef, M)

    for i = 1:N
        v[i] = T(1)
        mul!(dest, D, v)
        js = findall(!iszero, dest)
        colptr[i] = length(nzval)+1
        if length(js) > 0
            append!(rowind, js)
            append!(nzval, dest[js])
        end
        v[i] = T(0)
    end
    colptr[N+1] = length(nzval)+1

    return SparseMatrixCSC(M, N, colptr, rowind, nzval)
end

"""
    compute_coefficients(u, D::AbstractDerivativeOperator)

Compute the nodal values of the function `u` at the grid associated to the
derivative operator `D`.
"""
function compute_coefficients(u, D::AbstractDerivativeOperator)
    x = grid(D)
    xmin = first(x)
    xmax = last(x)
    uval = Array{typeof(u((xmin+xmax)/2))}(undef, length(x))
    compute_coefficients!(uval, u, D)
end

"""
    compute_coefficients!(uval::AbstractVector, u, D::AbstractDerivativeOperator)

Compute the nodal values of the function `u` at the grid associated to the
derivative operator `D` and stores the result in `uval`.
"""
function compute_coefficients!(uval::AbstractVector, u, D::AbstractDerivativeOperator)
    uval .= u.(grid(D))
end

"""
    compute_coefficients(u, D::AbstractPeriodicDerivativeOperator)

Compute the nodal values of the function `u` at the grid associated to the
derivative operator `D`.
"""
function compute_coefficients(u, D::AbstractPeriodicDerivativeOperator)
    x = D.grid_compute
    xmin = first(x)
    xmax = last(x)
    uval = Array{typeof(u((xmin+xmax)/2))}(undef, length(x))
    compute_coefficients!(uval, u, D)
end

"""
    compute_coefficients!(uval::AbstractVector, u, D::AbstractPeriodicDerivativeOperator)

Compute the nodal values of the function `u` at the grid associated to the
derivative operator `D` and stores the result in `uval`.
"""
function compute_coefficients!(uval::AbstractVector, u, D::AbstractPeriodicDerivativeOperator)
    uval .= u.(D.grid_compute)
end

"""
    evaluate_coefficients(u, D::AbstractDerivativeOperator)

Evaluates the nodal coefficients `u` at a grid associated to the derivative
operator `D`.
Returns `xplot, uplot`, where `xplot` contains the nodes and `uplot` the
corresponding values of `u`.
"""
function evaluate_coefficients(u, D::AbstractDerivativeOperator)
    x = grid(D)
    xplot = Array{eltype(x)}(undef, length(x))
    uplot = Array{eltype(u)}(undef, length(x))

    evaluate_coefficients!(xplot, uplot, u, D)
end

"""
    evaluate_coefficients!(xplot, uplot, u, D::AbstractDerivativeOperator)

Evaluates the nodal coefficients `u` at a grid associated to the derivative
operator `D` and stores the result in `xplot, uplot`.
Returns `xplot, uplot`, where `xplot` contains the nodes and `uplot` the
corresponding values of `u`.
"""
function evaluate_coefficients!(xplot, uplot, u, D::AbstractDerivativeOperator)
    @argcheck length(uplot) == length(xplot)
    @argcheck length(uplot) == length(grid(D))

    xplot .= grid(D)
    uplot .= u

    xplot, uplot
end

"""
    evaluate_coefficients(u, D::AbstractPeriodicDerivativeOperator)

Evaluates the nodal coefficients `u` at a grid including both endpoints
associated to the derivative periodic operator `D`.
Returns `xplot, uplot`, where `xplot` contains the equally spaced nodes and
`uplot` the corresponding values of `u`.
"""
function evaluate_coefficients(u, D::AbstractPeriodicDerivativeOperator)
    x = D.grid_evaluate
    xplot = Array{eltype(x)}(undef, length(x))
    uplot = Array{eltype(u)}(undef, length(x))

    evaluate_coefficients!(xplot, uplot, u, D)
end

"""
    evaluate_coefficients!(xplot, uplot, u, D::AbstractPeriodicDerivativeOperator)

Evaluates the nodal coefficients `u` at a grid including both endpoints
associated to the derivative periodic operator `D` and stores the result in
`xplot, uplot`.
Returns `xplot, uplot`, where `xplot` contains the equally spaced nodes and
`uplot` the corresponding values of `u`.
"""
function evaluate_coefficients!(xplot, uplot, u, D::AbstractPeriodicDerivativeOperator)
    @argcheck length(uplot) == length(xplot)
    @argcheck length(uplot) == length(D.grid_evaluate)

    xplot .= D.grid_evaluate
    uplot[1:end-1] = u
    uplot[end] = uplot[1]

    xplot, uplot
end

"""
    integrate(func, u, D::AbstractPeriodicDerivativeOperator)

Map the function `func` to the coefficients `u` and integrate with respect to
the quadrature rule associated with the derivative operator `D`.
"""
function integrate(func, u::AbstractVector, D::AbstractPeriodicDerivativeOperator)
    @boundscheck begin
        length(u) == length(grid(D))
    end
    @unpack Δx = D

    res = sum(func, u)

    Δx * res
end

"""
    SumOfDerivativeOperators

Sum several derivative operators lazily.
"""
struct SumOfDerivativeOperators{T,N,Operators<:Tuple{Vararg{AbstractDerivativeOperator{T},N}}} <: AbstractDerivativeOperator{T}
    operators::Operators

    function SumOfDerivativeOperators(operators::Operators) where {T,N,Operators<:Tuple{Vararg{AbstractDerivativeOperator{T},N}}}
        @argcheck all(i->size(operators[i]) == size(first(operators)), eachindex(operators)) DimensionMismatch
        @argcheck all(i->grid(operators[i]) == grid(first(operators)), eachindex(operators)) ArgumentError
        new{T,N,Operators}(operators)
    end
end

SumOfDerivativeOperators(ops...) = SumOfDerivativeOperators(ops)

Base.size(sum::SumOfDerivativeOperators) = size(first(sum.operators))
Base.size(sum::SumOfDerivativeOperators, i::Int) = size(first(sum.operators), i)
function Base.length(::Type{SumOfDerivativeOperators{T,N,Operators}}) where {T,N,Operators}
    N
end
grid(sum::SumOfDerivativeOperators) = grid(first(sum.operators))

function Base.:+(D1::AbstractDerivativeOperator, D2::AbstractDerivativeOperator)
    SumOfDerivativeOperators(D1, D2)
end

function Base.:+(sum::SumOfDerivativeOperators, D::AbstractDerivativeOperator)
    SumOfDerivativeOperators(sum.operators..., D)
end

function Base.:+(D::AbstractDerivativeOperator, sum::SumOfDerivativeOperators)
    SumOfDerivativeOperators(D, sum.operators...)
end

function Base.:+(sum1::SumOfDerivativeOperators, sum2::SumOfDerivativeOperators)
    SumOfDerivativeOperators(sum1.operators..., sum2.operators...)
end

function Base.show(io::IO, sum::SumOfDerivativeOperators)
    print(io, "Sum of operators:\n")
    for D in sum.operators
        print(io, D)
    end
end

@unroll function mul!(dest::AbstractVector, sum::SumOfDerivativeOperators, u::AbstractVector,
                      α, β)
    @unpack operators = sum
    @boundscheck begin
        @argcheck size(first(operators), 2) == length(u) DimensionMismatch
        @argcheck size(first(operators), 1) == length(dest) DimensionMismatch
    end

    @inbounds mul!(dest, operators[1], u, α, β)
    @unroll for i in 1:length(sum)
        if i != 1
            @inbounds mul!(dest, operators[i], u, α, one(β))
        end
    end

    nothing
end

@unroll function mul!(dest::AbstractVector, sum::SumOfDerivativeOperators, u::AbstractVector,
                      α)
    @unpack operators = sum
    @boundscheck begin
        @argcheck size(first(operators), 2) == length(u) DimensionMismatch
        @argcheck size(first(operators), 1) == length(dest) DimensionMismatch
    end

    @inbounds mul!(dest, operators[1], u, α)
    @unroll for i in 1:length(sum)
        if i != 1
            @inbounds mul!(dest, operators[i], u, α, one(α))
        end
    end

    nothing
end


text:
"""
Create normal budget distribution for each product with mean given by
the number of resources in the product.
"""
function normal_budgets_per_resource(P::AbstractArray{<:Product}, resource_budget_μ ::Float64, σ::Float64)
    P_sizes = Float64[sum(p) for p in P]
    budget_means = P_sizes .* resource_budget_μ
    [Normal(μ, σ) for μ in budget_means]
end

ext: jl


text:
include("types/__init__.jl")
include("printing/__init__.jl")
include("corpus/__init__.jl")
include("utils/__init__.jl")

P = symbol("P", Statements)
Q = symbol("Q", Statements)
R = symbol("R", Statements)

x = symbol("x", Statements)
A = symbol("A", Sets)
B = symbol("B", Sets)
C = symbol("C", Sets)


text:
# this file has to be run before the start of the framework
# in case the framework hasn't been used yet. It installs all
# necessary packages for the framework.
import Pkg;
Pkg.add("CSV")
Pkg.add("DataFrames")
Pkg.add("JuMP")
Pkg.add("Shapefile")
Pkg.add("Plots")
Pkg.add("StatsPlots")
Pkg.add("Distributions")


text:
using Tables
using Formatting
using FactCheck

# FactCheck.setstyle(:compact)
# write your own tests here

box1 = Box{Int}(Point(0,0), Point(10,10))
box2 = Box{Int}(Point(1,1), Point(2,2))
box3 = Box{Int}(Point(1,1), Point(12,12))
box4 = Box{Int}(Point(11,11), Point(12,12))
box5 = Box{Int}(Point(5,10), Point(12,12))
box6 = Box{Int}(Point(-1,-2), Point(13,14))

r1 = Range(box1, 1,1)
r2 = Range(box2, 2,2)
r3 = Range(box3, 3,3)
r4 = Range(box4, 4,4)
r5 = Range(box5, 5,5)
r6 = Range(box6, 6,6)

include("box.jl")
include("range.jl")

include("general/examples.jl")
include("html/html.jl")


text:
"""
    logseriescdf(p)

Returns the vector of samples from the discrete cdf of logarithmic distribution
"""
function logseriescdf(p)
  T = eltype(p)
  cdfs = T.([0.])
  for i in 1:100000000
    @inbounds push!(cdfs, cdfs[i]-(p^i)/(i*log(1-p)))
    # one can change atol here
    if isapprox(cdfs[i],  1.)
      return cdfs
    end
  end
  cdfs
end

"""
    levyel(θ, u1, u2)

u1 and u2 are supposed to be random numbers form [0,1]
An element from Levy stable distribution with parameters α = 1/θ, β = 1,
γ = (cos(pi/(2*θ)))^θ and δ = 0.
"""
function levyel(θ, u1, u2)
  ϕ = pi*u1-pi/2
  v = quantile.(Exponential(1.), u2)
  γ = (cos(pi/(2*θ)))^θ
  v = ((cos(pi/(2*θ)+(1/θ-1)*ϕ))/v)^(θ-1)
  γ*v*sin(1/θ*(pi/2+ϕ))*(cos(pi/(2*θ))*cos(ϕ))^(-θ)
end

"""
    tiltedlevygen(V0, α; rng::AbstractRNG)

Returns a sample from the expotencialy tilted Levy stable pdf
f(x; V0, α) = exp(-V0^α) g(x; α)/exp(-V0), where g(x; α) is a stable Levy pdf
with parameters α = 1/θ, β = 1, γ = (cos(pi/(2*θ)))^θ and δ = 0.
"""
function tiltedlevygen(V0, α; rng)
  T = eltype(V0)
  x = levyel(α, rand(rng, T), rand(rng, T))
  u = rand(rng, T)
  while exp(-V0^α*x)/(1500*exp(-V0)) < u
    x = levyel(α, rand(rng, T), rand(rng, T))
    u = rand(rng, T)
  end
  return x.*V0.^α
end

"""
    Ginv(y, α)

Returns Real, helper for the joe/frank nested copula generator
"""
Ginv(y, α) = ((1-y)*SpecialFunctions.gamma(1-α))^(-1/α)

"""
    InvlaJ(n::Int, α)

Returns the inverse Laplace transform of generator of Joe nested copula
parametered by α
"""
function InvlaJ(n, α)
  T = eltype(α)
  return 1-1/(n*beta(T(n), 1-α))
end

"""
  sampleInvlaJ(α, v)

Returns Int, a sample of inverce laplacea transform of generator of Joe nested copula,
given a parameter α and a random numver v ∈ [0,1]
"""
function sampleInvlaJ(α, v)
  if v <= α
    return 1
  else
    G = Ginv(v, α)
    if G > 2^62
      return 2^62
    else
      return (InvlaJ(floor(Int, G), α) < v) ? ceil(Int, G) : floor(Int, G)
    end
  end
end

"""
  elInvlaF(θ₁, θ₀, logseriescdf; rng::AbstractRNG)

Returns Int, a single sample of the inverse Laplace transform of the generator
of nested Frank copula
"""
function elInvlaF(θ₁, θ₀, logseriescdf; rng)
  T = eltype(θ₁)
  c1 = 1-exp(-θ₁)
  α = θ₀/θ₁
  if θ₀ <= 1
    v = rand(rng, T)
    X = findlast(logseriescdf .< rand(rng, T))[1]
    while v > 1/((X-α)*beta(X, 1-α))
      v = rand(rng, T)
      X = findlast(logseriescdf .< rand(rng, T))[1]
    end
    return X
  else
    v = rand(rng, T)
    X = sampleInvlaJ(α, rand(rng, T))
    while v > c1^(X-1)
      X = sampleInvlaJ(α, rand(rng, T))
      v = rand(rng, T)
    end
    return X
  end
end

"""
  nestedfrankgen(θ₁, θ₀, V0::Int, logseriescdf; rng::AbstractRNG)

Return int, sampled from the Inverse Laplace trensform of nested
Frank copula given parametes θ₁ θ₀ (child and parent)
and an element of V0 - vector of samples of invlaplace of the parents copula
"""
function nestedfrankgen(θ₁, θ₀, V0, logseriescdf; rng)
  T = eltype(θ₁)
  return sum([elInvlaF(θ₁, θ₀, logseriescdf; rng = rng) for j in 1:V0])
end


text:
using CUDA, FixedEffectModels, CategoricalArrays, CSV, DataFrames, Test, LinearAlgebra
df = DataFrame(CSV.File(joinpath(dirname(pathof(FixedEffectModels)), "../dataset/Cigar.csv")))
df.StateC = categorical(df.State)
df.YearC = categorical(df.Year)

##############################################################################
##
## coefficients
##
##############################################################################

# simple
m = @formula Sales ~ Price
x = reg(df, m)
@test coef(x) ≈ [139.73446,-0.22974] atol = 1e-4
m = @formula Sales ~ Price
x = reg(df, m, weights = :Pop)
@test coef(x) ≈ [137.72495428982756,-0.23738] atol = 1e-4

df.SalesInt = round.(Int64, df.Sales)
m = @formula SalesInt ~ Price
x = reg(df, m)
@test coef(x) ≈ [139.72674,-0.2296683205] atol = 1e-4

# absorb
m = @formula Sales ~ Price + fe(State)
x = reg(df, m)
@test coef(x) ≈ [-0.20984] atol = 1e-4
@test x.iterations == 1

m = @formula Sales ~ Price + fe(State) + fe(Year)
x = reg(df, m)
@test coef(x)  ≈ [-1.08471] atol = 1e-4
m = @formula Sales ~ Price + fe(State) + fe(State)*Year
x = reg(df, m)
@test coef(x) ≈  [-0.53470, 0.0] atol = 1e-4
m = @formula Sales ~ Price + fe(State)*Year
x = reg(df, m)
@test coef(x) ≈  [-0.53470, 0.0] atol = 1e-4

#@test isempty(coef(reg(df, @formula(Sales ~ 0), @fe(State*Price))))
df.mState = div.(df.State, 10)
m = @formula Sales ~ Price + fe(mState)&fe(Year)
x = reg(df, m)
@test coef(x)  ≈  [-1.44255] atol = 1e-4

m = @formula Sales ~ Price + fe(State)&Year
x = reg(df, m)
@test coef(x) ≈ [13.993028174622104,-0.5804357763515606] atol = 1e-4
m = @formula Sales ~ Price + Year&fe(State)
x = reg(df, m)
@test coef(x) ≈ [13.993028174622104,-0.5804357763515606] atol = 1e-4
m = @formula Sales ~ 1 + Year&fe(State)
x = reg(df, m)
@test coef(x) ≈ [174.4084407796102] atol = 1e-4

m = @formula Sales ~ Price + fe(State)&Year + fe(Year)&State
x = reg(df, m)
@test coef(x) ≈ [51.2359,- 0.5797] atol = 1e-4
m = @formula Sales ~ Price + NDI + fe(State)&Year + fe(Year)&State
x = reg(df, m)
@test coef(x) ≈ [-46.4464,-0.2546, -0.005563] atol = 1e-4
m = @formula Sales ~ 0 + Price + NDI + fe(State)&Year + fe(Year)&State
x = reg(df, m)
@test coef(x) ≈  [-0.21226562244177932,-0.004775616634862829] atol = 1e-4

# recheck these two below
m = @formula Sales ~ Pimin + Price&NDI&fe(State)
x = reg(df, m)
@test coef(x)  ≈  [122.98713, 0.30933] atol = 1e-4
# SSR does not work well here
m = @formula Sales ~ Pimin + (Price&NDI)*fe(State)
x = reg(df, m)
@test coef(x) ≈   [0.421406, 0.0] atol = 1e-4

# only one intercept
m = @formula Sales ~ 1 + fe(State) + fe(Year)
x = reg(df, m)

# TO DO: REPORT INTERCEPT IN CASE OF FIXED EFFFECTS, LIKE STATA
df.id3  = categorical(mod.(1:size(df, 1), Ref(3)))
df.id4  = categorical(div.(1:size(df, 1), Ref(10)))

m = @formula Sales ~ id3
x = reg(df, m)
@test length(coef(x)) == 3

m = @formula Sales ~ 0 + id3
x = reg(df, m)
@test length(coef(x)) == 3

# with fixed effects it's like implicit intercept
m = @formula Sales ~  id3 + fe(id4)
x = reg(df, m)
@test length(coef(x)) == 2

m = @formula Sales ~ Year + fe(State)
x = reg(df, m)

m = @formula Sales ~ id3&Price
x = reg(df, m)
@test length(coef(x)) == 4
m = @formula Sales ~ id3&Price + Price
x = reg(df, m)
@test length(coef(x)) == 4

m = @formula Sales ~ Year + fe(State)
x = reg(df, m)

m = @formula Sales ~ Year&Price + fe(State)
x = reg(df, m)

# absorb + weights
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, weights = :Pop)
@test coef(x) ≈ [- 0.21741] atol = 1e-4
m = @formula Sales ~ Price + fe(State) + fe(Year)
x = reg(df, m, weights = :Pop)
@test coef(x) ≈ [- 0.88794] atol = 1e-3
m = @formula Sales ~ Price + fe(State) + fe(State)&Year
x = reg(df, m, weights = :Pop)
@test coef(x) ≈ [- 0.461085492] atol = 1e-4

# iv
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m)
@test coef(x) ≈ [138.19479,- 0.20733] atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin)
x = reg(df, m)
@test coef(x) ≈  [137.45096,0.00516,- 0.76276] atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin + Pop)
x = reg(df, m)
@test coef(x) ≈ [137.57335,0.00534,- 0.78365] atol = 1e-4
## multiple endogeneous variables
m = @formula Sales ~ (Price + NDI ~ Pimin + Pop)
x = reg(df, m)
@test coef(x) ≈ [139.544, .8001, -.00937] atol = 1e-4
m = @formula Sales ~ 1 + (Price + NDI ~ Pimin + Pop)
x = reg(df, m)
@test coef(x) ≈ [139.544, .8001, -.00937] atol = 1e-4
result = [196.576, 0.00490989, -2.94019, -3.00686, -2.94903, -2.80183, -2.74789, -2.66682, -2.63855, -2.52394, -2.34751, -2.19241, -2.18707, -2.09244, -1.9691, -1.80463, -1.81865, -1.70428, -1.72925, -1.68501, -1.66007, -1.56102, -1.43582, -1.36812, -1.33677, -1.30426, -1.28094, -1.25175, -1.21438, -1.16668, -1.13033, -1.03782]

m = @formula Sales ~ NDI + (Price&YearC ~ Pimin&YearC)
x = reg(df, m)
@test coef(x) ≈ result atol = 1e-4

# iv + weight
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m,  weights = :Pop)
@test coef(x) ≈ [137.03637,- 0.22802] atol = 1e-4

# iv + weight + absorb
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m)
@test coef(x) ≈ [-0.20284] atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m,  weights = :Pop)
@test coef(x) ≈ [-0.20995] atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin)
x = reg(df, m)
@test coef(x) ≈  [137.45096580480387,0.005169677634275297,-0.7627670265757879] atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin) + fe(State)
x = reg(df, m)
@test coef(x)  ≈  [0.0011021722526916768,-0.3216374943695231] atol = 1e-4

# non high dimensional factors
m = @formula Sales ~ Price + YearC
x = reg(df, m)
m = @formula Sales ~ YearC + fe(State)
x = reg(df, m)
m = @formula Sales ~ Price + YearC + fe(State)
x = reg(df, m)
@test coef(x)[1] ≈ -1.08471 atol = 1e-4
m = @formula Sales ~ Price + YearC + fe(State)
x = reg(df, m,  weights = :Pop)
@test coef(x)[1] ≈ -0.88794 atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin) + YearC + fe(State)
x = reg(df, m)
@test coef(x)[1] ≈ -0.00525 atol = 1e-4

##############################################################################
##
## Programming
##
##############################################################################
reg(df, term(:Sales) ~ term(:NDI) + fe(:State) + fe(:Year), Vcov.cluster(:State))
@test fe(:State) + fe(:Year) === reduce(+, fe.([:State, :Year])) === fe(term(:State)) + fe(term(:Year))

##############################################################################
##
## Functions
##
##############################################################################

# function
m = @formula Sales ~ log(Price)
x = reg(df, m)
@test coef(x)[1] ≈184.98520688 atol = 1e-4

# function defined in user space
mylog(x) = log(x)
m = @formula Sales ~ mylog(Price)
x = reg(df, m)

# Function returning Inf
df.Price_zero = copy(df.Price)
df.Price_zero[1] = 0.0
m = @formula Sales ~ log(Price_zero)
@test_throws  "Some observations for the regressor are infinite" reg(df, m)
##############################################################################
##
## collinearity
## add more tests
##
##############################################################################
# ols
df.Price2 = df.Price
m = @formula Sales ~ Price + Price2
x = reg(df, m)
@test coef(x) ≈ [139.7344639806166,-0.22974688593485126,0.0] atol = 1e-4

## iv
df.NDI2 = df.NDI
m = @formula Sales ~ NDI2 + NDI + (Price ~ Pimin)
x = reg(df, m)
@test iszero(coef(x)[2]) || iszero(coef(x)[3])

## endogeneous variables collinear with instruments are reclassified
df.zPimin = df.Pimin
m = @formula Sales ~ zPimin + (Price ~ NDI + Pimin)
x = reg(df, m)

m2 = @formula Sales ~ (zPimin + Price ~ NDI + Pimin)
NDI = reg(df, m2)
@test coefnames(x) == coefnames(NDI)
@test coef(x) ≈ coef(NDI)
@test vcov(x) ≈ vcov(NDI)

# catch when IV underidentified
@test_throws "Model not identified. There must be at least as many ivs as endogeneneous variables" reg(df, @formula(Sales ~ Price + (NDI + Pop ~ NDI)))
@test_throws "Model not identified. There must be at least as many ivs as endogeneneous variables" reg(df, @formula(Sales ~ Price + (Pop ~ Price)))

# catch when IV underidentified
@test_throws "Model not identified. There must be at least as many ivs as endogeneneous variables" reg(df, @formula(Sales ~ Price + (NDI + Pop ~ NDI)))

# Make sure all coefficients are estimated
p = [100.0, -40.0, 30.0, 20.0]
df_r = DataFrame(y = p, x = p.^4)
result = reg(df_r, @formula(y ~ x))
@test sum(abs.(coef(result)) .> 0)  == 2

##############################################################################
##
## std errors
##
##############################################################################

# Simple - matches stata
m = @formula Sales ~ Price
x = reg(df, m)
@test stderror(x) ≈ [1.521269, 0.0188963] atol = 1e-6
# Stata ivreg - ivreg2 discrepancy - matches with df_add=-2
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m)
@test stderror(x) ≈ [1.53661, 0.01915] atol = 1e-4
# Stata areg
m = @formula Sales ~ Price + fe(State)
x = reg(df, m)
@test stderror(x) ≈  [0.0098003] atol = 1e-7

# White
# Stata reg - matches stata
m = @formula Sales ~ Price
x = reg(df, m, Vcov.robust())
@test stderror(x) ≈ [1.686791, 0.0167042] atol = 1e-6
# Stata ivreg - ivreg2 discrepancy - matches with df_add=-2
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m, Vcov.robust())
@test stderror(x) ≈ [1.63305, 0.01674] atol = 1e-4
# Stata areg - matches areg
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, Vcov.robust())
@test stderror(x) ≈ [0.0110005] atol = 1e-7

# Clustering models
# cluster - matches stata
m = @formula Sales ~ Price
x = reg(df, m, Vcov.cluster(:State))
@test stderror(x)[2] ≈ 0.0379228 atol = 1e-7
# cluster with fe - matches areg & reghdfe
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, Vcov.cluster(:Year))
@test stderror(x) ≈ [0.0220563] atol = 1e-7
# stata reghxe - matches reghdfe (not areg)
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, Vcov.cluster(:State))
@test stderror(x) ≈ [0.0357498] atol = 1e-7
# iv + fe + cluster - matches ivreghdfe
m = @formula Sales ~ NDI + (Price ~Pimin) + fe(State)
x = reg(df, m, Vcov.cluster(:State))
@test stderror(x) ≈ [0.0019704, 0.1893396] atol = 1e-7
# iv + fe + cluster + weights - matches ivreghdfe
m = @formula Sales ~ NDI + (Price ~ Pimin) + fe(State)
x = reg(df, m,  Vcov.cluster(:State), weights = :Pop)
@test stderror(x) ≈ [0.000759, 0.070836] atol = 1e-6
# iv + fe + cluster + weights - matches ivreghdfe
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m,  Vcov.cluster(:State), weights = :Pop)
@test stderror(x) ≈ [0.0337439] atol = 1e-7
# multiway clustering - matches reghdfe
m = @formula Sales ~ Price
x = reg(df, m, Vcov.cluster(:State, :Year))
@test stderror(x) ≈ [6.196362, 0.0403469] atol = 1e-6
# multiway clustering - matches reghdfe
m = @formula Sales ~ Price
x = reg(df, m, Vcov.cluster(:State, :Year))
@test stderror(x) ≈ [6.196362, 0.0403469] atol = 1e-6
# fe + multiway clustering - matches reghdfe
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, Vcov.cluster(:State, :Year))
@test stderror(x) ≈ [0.0405335] atol = 1e-7
m = @formula Sales ~ Price + fe(State)
x = reg(df, m, Vcov.cluster(:State, :Year))
@test stderror(x) ≈ [0.0405335] atol = 1e-7
# fe + clustering on interactions - matches reghdfe
#m = @formula Sales ~ Price + fe(State) vcov = cluster(State&id2)
#x = reg(df, m)
#@test stderror(x) ≈ [0.0110005] atol = 1e-7
# fe partially nested in interaction clusters - matches reghdfe
#m=@formula Sales ~ Price + fe(State) vcov=cluster(State&id2)
#x = reg(df, m)
#@test stderror(x) ≈ [0.0110005] atol = 1e-7
# regressor partially nested in interaction clusters - matches reghdfe
#m=@formula Sales ~ Price + StateC vcov=cluster(State&id2)
#x = reg(df, m)
#@test stderror(x)[1:2] ≈ [3.032187, 0.0110005] atol=1e-5

#check palue printed is correct
#m = @formula Sales ~ Price + fe(State)
#x = reg(df, m, Vcov.cluster(:State))
#@test coeftable(x).mat[1, 4] ≈ 4.872723900371927e-7 atol = 1e-7

##############################################################################
##
## subset
##
##############################################################################
m = @formula Sales ~ Price + StateC
x0 = reg(df[df.State .<= 30, :], m)

m = @formula Sales ~ Price + StateC
Price = reg(df, m, subset = df.State .<= 30)
@test length(Price.esample) == size(df, 1)
@test coef(x0) ≈ coef(Price) atol = 1e-4
@test vcov(x0) ≈ vcov(Price) atol = 1e-4

df.State_missing = ifelse.(df.State .<= 30, df.State, missing)
df.StateC_missing = categorical(df.State_missing)
m = @formula Sales ~ Price + StateC_missing
NDI = reg(df, m)
@test length(NDI.esample) == size(df, 1)
@test coef(x0) ≈ coef(NDI) atol = 1e-4
@test vcov(x0) ≈ vcov(NDI) atol = 1e-2

# missing weights
df.Price_missing = ifelse.(df.State .<= 30, df.Price, missing)
m = @formula Sales ~ NDI + StateC
x = reg(df, m, weights = :Price_missing)
@test length(x.esample) == size(df, 1)

# missing interaction
m = @formula Sales ~ NDI + fe(State)&Price_missing
x = reg(df, m)
@test nobs(x) == count(.!ismissing.(df.Price_missing))

m = @formula Sales ~ Price + fe(State)
x3 = reg(df, m, subset = df.State .<= 30)
@test length(x3.esample) == size(df, 1)
@test coef(x0)[2] ≈ coef(x3)[1] atol = 1e-4

m = @formula Sales ~ Price + fe(State_missing)
x4 = reg(df, m)
@test coef(x0)[2] ≈ coef(x4)[1] atol = 1e-4

# categorical variable as fixed effects
m = @formula Sales ~ Price + fe(State)
x5 = reg(df, m, subset = df.State .>= 30)

#Error reported by Erik
m = @formula Sales ~ Pimin + CPI
x = reg(df, m, Vcov.cluster(:State), subset = df.State .>= 30)
@test diag(x.vcov) ≈ [130.7464887, 0.0257875, 0.0383939] atol = 1e-4

##############################################################################
##
## R2
##
##############################################################################
m = @formula Sales ~ Price
x = reg(df, m)
@test r2(x) ≈ 0.0969 atol = 1e-4
@test adjr2(x) ≈ 0.09622618 atol = 1e-4

##############################################################################
##
## F Stat
##
##############################################################################
m = @formula Sales ~ Price
x = reg(df, m)
@test x.F  ≈ 147.82425 atol = 1e-4
m = @formula Sales ~ Price + fe(State)
x = reg(df, m)
@test x.F  ≈ 458.45825 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m)
@test x.F  ≈ 117.17329 atol = 1e-4
m = @formula Sales ~ Price + Pop
x = reg(df, m)
@test x.F ≈ 79.091576 atol = 1e-4
m = @formula Sales ~ Price
x = reg(df, m, Vcov.cluster(:State))
@test x.F    ≈ 36.70275 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m, Vcov.cluster(:State))
@test x.F  ≈ 39.67227 atol = 1e-4
#  xtivreg2
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m)
@test x.F  ≈ 422.46444 atol = 1e-4

# p value
m = @formula Pop ~ Pimin
x = reg(df, m)
@test x.p ≈ 0.003998718283554043 atol = 1e-4
m = @formula Pop ~ Pimin + Price
x = reg(df, m)
@test x.p ≈ 3.369423076033613e-14 atol = 1e-4

# Fstat  https://github.com/FixedEffects/FixedEffectModels.jl/issues/150
df_example = DataFrame(CSV.File(joinpath(dirname(pathof(FixedEffectModels)), "../dataset/Ftest.csv")))
x = reg(df_example, @formula(Y ~ fe(id) + (X1 + X2 ~ Z1 + Z2) ), Vcov.robust() )
@test x.F_kp ≈ 14.5348449357 atol = 1e-4

##############################################################################
##
## F_kp r_kp statistics for IV. Difference degrees of freedom.
##
##############################################################################
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m)
@test x.F_kp ≈ 52248.79247 atol = 1e-4
m = @formula Sales ~ NDI + (Price ~ Pimin)
x = reg(df, m)
@test x.F_kp ≈ 5159.812208193612 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin + CPI)
x = reg(df, m)
@test x.F_kp   ≈ 27011.44080 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m)
@test x.F_kp ≈ 100927.75710 atol = 1e-4

# exactly same with ivreg2
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m, Vcov.robust())
@test x.F_kp ≈ 23160.06350 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin) + fe(State)
x = reg(df, m, Vcov.robust())
@test x.F_kp  ≈ 37662.82808 atol = 1e-4
m = @formula Sales ~ CPI + (Price ~ Pimin)
x = reg(df, m, Vcov.robust())
@test x.F_kp ≈ 2093.46609 atol = 1e-4
m = @formula Sales ~ (Price ~ Pimin + CPI)
x = reg(df, m, Vcov.robust())
@test x.F_kp  ≈ 16418.21196 atol = 1e-4

# like in ivreg2 but += 5 difference. combination iv difference, degrees of freedom difference?
# iv + cluster - F_kp varies from ivreg2 about 7e-4 (SEs off by more than 2 df)
m = @formula Sales ~ (Price ~ Pimin)
x = reg(df, m, Vcov.cluster(:State))
@test x.F_kp     ≈ 7249.88606 atol = 1e-4
# iv + cluster - F_kp varies from ivreg2 about 5e-6 (SEs off by more than 2 df)
m = @formula Sales ~ CPI + (Price ~ Pimin)
x = reg(df, m, Vcov.cluster(:State))
@test x.F_kp   ≈ 538.40393 atol = 1e-4
# Modified test values below after multiway clustering update
# iv + 2way clustering - F_kp matches ivreg2 (SEs match with df_add=-2)
m = @formula Sales ~ CPI + (Price ~ Pimin)
x = reg(df, m, Vcov.cluster(:State, :Year))
@test x.F_kp   ≈ 421.9651 atol = 1e-4
# multivariate iv + clustering - F_kp varies from ivreg2 about 3 (SEs off by more than 2 df)
m = @formula Sales ~ (Price ~ Pimin + CPI)
x = reg(df, m, Vcov.cluster(:State))
@test x.F_kp      ≈ 4080.66081 atol = 1e-4
# multivariate iv + multiway clustering - F_kp varies from ivreg2 about 2 (SEs off by more than 2 df)
m = @formula Sales ~ (Price ~ Pimin + CPI)
x = reg(df, m, Vcov.cluster(:State, :Year))
@test x.F_kp  ≈ 2873.1405 atol = 1e-4

##############################################################################
##
## Test singleton
##
##
##############################################################################
df.n = max.(1:size(df, 1), 60)
df.pn = categorical(df.n)
m = @formula Sales ~ Price + fe(pn)
x = reg(df, m, Vcov.cluster(:State))
@test x.nobs == 60

m = @formula Sales ~ Price + fe(pn)
x = reg(df, m, Vcov.cluster(:State), drop_singletons = false)
@test x.nobs == 1380

##############################################################################
##
## Test unbalanced panel
##
## corresponds to abdata in Stata, for instance reghxe wage emp [w=indoutpt], a(id year)
##
##############################################################################
df = DataFrame(CSV.File(joinpath(dirname(pathof(FixedEffectModels)), "../dataset/EmplUK.csv")))

m = @formula Wage ~ Emp + fe(Firm)
x = reg(df, m)
@test coef(x) ≈ [- 0.11981270017206136] atol = 1e-4
m = @formula Wage ~ Emp + fe(Firm)&Year
x = reg(df, m)
@test coef(x)  ≈ [-315.0000747500431,- 0.07633636891202833] atol = 1e-4
m = @formula Wage ~ Emp + Year&fe(Firm)
x = reg(df, m)
@test coef(x) ≈  [-315.0000747500431,- 0.07633636891202833] atol = 1e-4
m = @formula Wage ~ 1 + Year&fe(Firm)
x = reg(df, m)
@test coef(x) ≈  [- 356.40430526316396] atol = 1e-4
m = @formula Wage ~ Emp + fe(Firm)
x = reg(df, m, weights = :Output)
@test coef(x) ≈ [- 0.11514363590574725] atol = 1e-4

# absorb + weights
m = @formula Wage ~ Emp + fe(Firm) + fe(Year)
x = reg(df, m)
@test coef(x)  ≈  [- 0.04683333721137311] atol = 1e-4
m = @formula Wage ~ Emp + fe(Firm) + fe(Year)
x = reg(df, m, weights = :Output)
@test coef(x) ≈  [- 0.043475472188120416] atol = 1e-3

## the last two ones test an ill conditioned model matrix
# SSR does not work well here
m = @formula Wage ~ Emp + fe(Firm) + fe(Firm)&Year
x = reg(df, m)
@test coef(x)  ≈   [- 0.122354] atol = 1e-4
@test x.iterations <= 30

# SSR does not work well here
m = @formula Wage ~ Emp + fe(Firm) + fe(Firm)&Year
x = reg(df, m, weights = :Output)
@test coef(x) ≈ [- 0.11752306001586807] atol = 1e-4
@test x.iterations <= 50

methods_vec = [:cpu]
if FixedEffectModels.FixedEffects.has_CUDA()
	push!(methods_vec, :gpu)
end
for method in methods_vec
	# same thing with float32 precision
	local m = @formula Wage ~ Emp + fe(Firm)
	local x = reg(df, m, method = method, double_precision = false)
	@test coef(x) ≈ [- 0.11981270017206136] rtol = 1e-4
	local m = @formula Wage ~ Emp + fe(Firm)&Year
	local x = reg(df, m, method = method, double_precision = false)
	@test coef(x)  ≈ [-315.0000747500431,- 0.07633636891202833] rtol = 1e-4
	local m = @formula Wage ~ Emp + Year&fe(Firm)
	local x = reg(df, m, method = method, double_precision = false)
	@test coef(x) ≈  [-315.0000747500431,- 0.07633636891202833] rtol = 1e-4
	local m = @formula Wage ~ 1 + Year&fe(Firm)
	local x = reg(df, m, method = method, double_precision = false)
	@test coef(x) ≈  [- 356.40430526316396] rtol = 1e-4
	local m = @formula Wage ~ Emp + fe(Firm)
	local x = reg(df, m, weights = :Output, method = method, double_precision = false)
	@test coef(x) ≈ [- 0.11514363590574725] rtol = 1e-4
	local m = @formula Wage ~ Emp + fe(Firm) + fe(Year)
	local x = reg(df, m, method = method, double_precision = false)
	@test coef(x)  ≈  [- 0.04683333721137311] rtol = 1e-4
	local m = @formula Wage ~ Emp + fe(Firm) + fe(Year)
	local x = reg(df, m, weights = :Output, method = method, double_precision = false)
	@test coef(x) ≈  [- 0.043475472188120416] atol = 1e-3
end

# add tests with missing fixed effects
df.Firm_missing = ifelse.(df.Firm .<= 30, missing, df.Firm)

## test with missing fixed effects
m = @formula Wage ~ Emp + fe(Firm_missing)
x = reg(df, m)
@test coef(x) ≈ [-.1093657] atol = 1e-4
@test stderror(x) ≈  [.032949 ] atol = 1e-4
@test r2(x) ≈ 0.8703 atol = 1e-2
@test adjr2(x) ≈ 0.8502 atol = 1e-2
@test x.nobs == 821

## test with missing interaction
df.Year2 = df.Year .>= 1980
m = @formula Wage ~ Emp + fe(Firm_missing) & fe(Year2)
x = reg(df, m)
@test coef(x) ≈ [-0.100863] atol = 1e-4
@test stderror(x) ≈  [0.04149] atol = 1e-4
@test x.nobs == 821

##############################################################################
##
## Missing dependent / independent variables
##
##############################################################################

# Missing
df1 = DataFrame(a=[1.0, 2.0, 3.0, 4.0], b=[5.0, 7.0, 11.0, 13.0])
df2 = DataFrame(a=[1.0, missing, 3.0, 4.0], b=[5.0, 7.0, 11.0, 13.0])
x = reg(df1, @formula(a ~ b))
@test coef(x) ≈ [ -0.6500000000000004, 0.35000000000000003] atol = 1e-4
x = reg(df1, @formula(b ~ a))
@test coef(x) ≈ [2.0, 2.8] atol = 1e-4
x = reg(df2,@formula(a ~ b))
@test coef(x) ≈ [ -0.8653846153846163, 0.3653846153846155] atol = 1e-4
x = reg(df2,@formula(b ~ a))
@test coef(x) ≈ [ 2.4285714285714253, 2.7142857142857157] atol = 1e-4

# Works with Integers
df1 = DataFrame(a=[1, 2, 3, 4], b=[5, 7, 11, 13], c = categorical([1, 1, 2, 2]))
x = reg(df1, @formula(a ~ b))
@test coef(x) ≈ [-0.65, 0.35] atol = 1e-4
x = reg(df1, @formula(a ~ b + fe(c)))
@test coef(x) ≈ [0.5] atol = 1e-4


text:
type Point
    x::Float64
    y::Float64
    z::Float64
end

typeof(Point) #> DataType

names(Point) #> 3-element Array{Symbol,1}: :x :y :z

orig = Point(0, 0, 0) #> Point(0.0,0.0,0.0)
p1 = Point(2, 4, 1.3) #> Point(2.0,4.0,1.3)
names(p1) #> 3-element Array{Symbol,1}: :x :y :z

methods(Point)
# 2 methods for generic function "Point":
# Point(x::Float64,y::Float64,z::Float64)
# Point(x,y,z)

typeof(p1) #> Point (constructor with 1 method)
subtypes(Point) #> 0-element Array{Any,1}

p1.y #> 4.0
p1.z = 3.14 #> 3.14
p1 #> Point(2.0,4.0,3.14)
# p1.x = "A" #> ERROR: `convert` has no method matching ...

immutable Vector3D
    x::Float64
    y::Float64
    z::Float64
end
typeof(Vector3D) #> DataType
p = Vector3D(1, 2, 3) #> Vector3D(1.0,2.0,3.0)
# p.y = 5 #> ERROR: type Vector3D is immutable

5 == 5 #> true
5 == 5.0 #> true
is(5, 5) #> true
is(5, 5.0) #> false
q = Vector3D(4.0, 3.14, 2.71)
r = Vector3D(4.0, 3.14, 2.71)
is(q, r) #> false

typealias Point3D Point
p31 = Point3D(1, 2, 3) #> Point(1.0,2.0,3.0)

# constructors, subtyping and multiple dispatch
abstract Employee
# Employee() #> type cannot be constructed

type Developer <: Employee
    name::String
    iq
    favorite_lang
end

#type MobileDeveloper <: Developer
#    platform
#end
# ERROR: invalid subtyping in definition of MobileDeveloper

# outer constructor:
Developer(name, iq) = Developer(name, iq, "Java")
#> Developer (constructor with 3 methods)

type Manager
    name::String
    iq
    department::String
end

devel1 = Developer("Bob", 110) #> Developer("Bob",110,"Java")
devel2 = Developer("William", 145, "Julia")
#> Developer("William",145,"Julia")
man1 = Manager("Julia", 120, "ICT") #> Manager("Julia",120,"ICT")

cleverness(emp::Employee) = emp.iq
cleverness(devel1) #> 110
# cleverness(man1) #> ERROR: `cleverness` has no method matching cleverness(::Manager)

function cleverer(m::Manager, e::Employee)
    println("The manager $(m.name) is cleverer!")
end

cleverer(man1, devel1) #> The manager Julia is cleverer!
cleverer(man1, devel2) #> The manager Julia is cleverer!

function cleverer(d::Developer, e::Employee)
    println("The developer $(d.name) is cleverer!")
end

cleverer(devel1, devel2) #> The developer Bob is cleverer!
# cleverer(devel1, man1) #> ERROR: `cleverer` has no method matching cleverer(::Developer,::Manager)

function cleverer(e::Employee, d::Developer)
    if e.iq <= d.iq
        println("The developer $(d.name) is cleverer!")
    else
        println("The employee $(e.name) is cleverer!")
    end
end
#>
#Warning: New definition
#    cleverer(Employee,Developer) at none:2
#is ambiguous with:
#    cleverer(Developer,Employee) at none:2.
#To fix, define
#    cleverer(Developer,Developer)
#before the new definition.

function cleverer(d1::Developer, d2::Developer)
    if d1.iq <= d2.iq
        println("The developer $(d2.name) is cleverer!")
    else
        println("The developer $(d1.name) is cleverer!")
    end
end

cleverer(devel1, devel2) #> The developer William is cleverer!
cleverer(devel2, devel1) #> The developer William is cleverer!

max_stars_repo_path: Module 1/Chapter06/user_defined.jl

max_issues_repo_path: Module 1/Chapter06/user_defined.jl

max_forks_repo_path: Module 1/Chapter06/user_defined.jl


text:
"""
```
`function solve_adaptive_ϕ(cloud::Cloud, proposed_fixed_schedule::Vector{Float64},
                           i::Int64, j::Int64, ϕ_prop::Float64, ϕ_n1::Float64,
                           tempering_target::Float64, resampled_last_period::Bool)`
```
Solves for next Φ. Returns ϕ_n, resampled_last_period, j, ϕ_prop.
"""
function solve_adaptive_ϕ(cloud::Cloud, proposed_fixed_schedule::Vector{Float64},
                          i::Int64, j::Int64, ϕ_prop::Float64, ϕ_n1::Float64,
                          tempering_target::Float64, resampled_last_period::Bool)
    n_Φ = length(proposed_fixed_schedule)

    if resampled_last_period
        # The ESS_bar is reset to target an evenly weighted particle population
        ESS_bar = tempering_target * length(cloud)
        resampled_last_period = false
    else
        ESS_bar = tempering_target*cloud.ESS[i-1]
    end

    # Setting up the optimal ϕ solving function for endogenizing the tempering schedule
    optimal_ϕ_function(ϕ) = compute_ESS(get_loglh(cloud), get_weights(cloud), ϕ, ϕ_n1,
                                        old_loglh = get_old_loglh(cloud)) - ESS_bar

    # Find ϕ_prop s.t. optimal ϕ_n lies between ϕ_n1 and ϕ_prop --
    # do so by iterating through proposed_fixed_schedule and finding the first
    # ϕ_prop s.t. the ESS falls by more than the targeted amount, ESS_bar
    while optimal_ϕ_function(ϕ_prop) >= 0 && j <= n_Φ
        ϕ_prop = proposed_fixed_schedule[j]
        j += 1
    end

    # Note: optimal_ϕ_function(ϕ_n1) > 0, since ESS_{t-1} always positive.
    # When ϕ_prop != 1.0, there are still ϕ increments strictly below 1 that
    # give the optimal ϕ step, ϕ_n.

    # When ϕ_prop == 1.0 but optimal_ϕ_function(ϕ_prop) < 0, there still exists
    # an optimal ϕ step, ϕ_n, that does not equal 1.
    # Thus the interval [optimal_ϕ_function(ϕ_n1), optimal_ϕ_function(ϕ_prop)] always
    # contains a 0 by construction.

    # Modification makes it such that ϕ_n is the minimum of ϕ_prop (the fixed schedule)
    # at a given stage or the root-solved ϕ such that the ESS drops by the target amount.
    # Thus the adaptive ϕ_schedule is strictly bounded above by the fixed schedule
    # i.e. the adaptive ϕ schedule should not outpace the fixed schedule at the end
    # (when the fixed schedule tends to drop by less than 5% per iteration)
    if ϕ_prop != 1. || optimal_ϕ_function(ϕ_prop) < 0
        ϕ_n = fzero(optimal_ϕ_function, [ϕ_n1, ϕ_prop], xtol = 0.)
        push!(cloud.tempering_schedule, ϕ_n)
    else
        ϕ_n = 1.
        push!(cloud.tempering_schedule, ϕ_n)
    end
    return ϕ_n, resampled_last_period, j, ϕ_prop
end

"""
```
`mvnormal_mixture_draw(θ_old::Vector{T}, d_prop::Distribution;
                       c::T = 1.0, α::T = 1.0) where T<:AbstractFloat`
```

Create a `DegenerateMvNormal` distribution object, `d`, from a parameter vector, `p`, and a
standard deviation matrix (obtained from SVD), `σ`.

Generate a draw from the mixture distribution of:
1. A `DegenerateMvNormal` centered at θ_old with the standard deviation matrix `σ`, scaled by `cc^2` and with mixture proportion `α`.
2. A `DegenerateMvNormal` centered at the same mean, but with a standard deviation matrix of the diagonal entries of `σ` scaled by `cc^2` with mixture proportion `(1 - α)/2`.
3. A `DegenerateMvNormal`  with the same standard deviation matrix `σ` but centered at the new proposed mean, `θ_prop`, scaled by `cc^2`, and with mixture proportion `(1 - α)/2`.

If no `θ_prop` is given, but an `α` is specified, then the mixture will consist of `α` of
the standard distribution and `(1 - α)` of the diagonalized distribution.

### Arguments
- `θ_old::Vector{T}`: The mean of the desired distribution
- `σ::Matrix{T}`: The standard deviation matrix of the desired distribution

### Keyword Arguments
- `cc::T`: The standard deviation matrix scaling factor
- `α::T`: The mixing proportion
- `θ_prop::Vector{T}`: The proposed parameter vector to be used as part of the mixture distribution, set by default to be the weighted mean of the particles, prior to mutation.

### Outputs
- `θ_new::Vector{T}`: The draw from the mixture distribution to be used as the MH proposed step
"""
function mvnormal_mixture_draw(θ_old::Vector{T}, d_prop::Distribution;
                               c::T = 1.0, α::T = 1.0) where T<:AbstractFloat
    @assert 0 <= α <= 1
    d_bar = MvNormal(d_prop.μ, c^2 * d_prop.Σ)

    # Create mixture distribution conditional on the previous parameter value, θ_old
    d_old      = MvNormal(θ_old, c^2 * d_prop.Σ)
    d_diag_old = MvNormal(θ_old, Diagonal(diag(c^2 * d_prop.Σ)))
    d_mix_old  = MixtureModel(MvNormal[d_old, d_diag_old, d_bar], [α, (1 - α)/2, (1 - α)/2])

    θ_new = rand(d_mix_old)

    return θ_new
end

get_cov(d::MvNormal)::Matrix = d.Σ.mat
get_cov(d::DegenerateMvNormal)::Matrix = d.σ

"""
```
compute_proposal_densities(para_draw::Vector{T}, para_subset::Vector{T},
                           d_subset::Distribution;
                           α::T = 1.0, c::T = 1.0, catch_near_zeros::Bool = false) where {T<:AbstractFloat}
```
Called in Metropolis-Hastings mutation step. After you have generated proposal draw
ϑ_b from the mixture distrubtion, computes the density of the proposal distribution
for computation of acceptance probability.

### Inputs
- `para_draw::Vector{T}`: ϑ_b
- `para_subset::Vector{T}`: θ^i_{n,b,m-1}
- `d_subset::Distribution`: MvNormal(θ*_b, Σ*_b)

### Optional Inputs
- `α::T`
- `c::T`

### Outputs
- `q0::T`: q(ϑ_b | θ^i_{n,b,m-1}, θ^i_{n,-b,m}, θ*_b, Σ*_b)
- `q1::T`: q(θ^i_{n,b,m-1} | ϑ_b, θ^i_{n,b,m-1}, θ^i_{n,-b,m}, θ*_b, Σ*_b)
"""
function compute_proposal_densities(para_draw::Vector{T}, para_subset::Vector{T},
                                    d_subset::Distribution;
                                    α::T = 1.0, c::T = 1.0,
                                    catch_near_zeros::Bool = false,
                                    tol::Float64 = 1e-6) where {T<:AbstractFloat}
    d_Σ = get_cov(d_subset)

    q0 = α * exp(logpdf(DegenerateMvNormal(para_draw,   c^2 * d_Σ), para_subset))
    q1 = α * exp(logpdf(DegenerateMvNormal(para_subset, c^2 * d_Σ), para_draw))

    ind_pdf = 1.0

    if catch_near_zeros
        # If approximately near zero and negative, assume that the entries are actually supposed to be zero
        near_zero_inds = findall(-tol .<= d_Σ .< 0.)
        d_Σ[near_zero_inds] .= 0.
    end
    for i = 1:length(para_subset)
        Σ_ii    = sqrt(d_Σ[i,i])
        zstat   = (para_subset[i] - para_draw[i]) / Σ_ii
        ind_pdf = ind_pdf / (Σ_ii * sqrt(2.0 * π)) * exp(-0.5 * zstat^2)
    end

    q0 += (1.0-α)/2.0 * ind_pdf
    q1 += (1.0-α)/2.0 * ind_pdf

    q0 += (1.0-α)/2.0 * exp(logpdf(DegenerateMvNormal(d_subset.μ, c^2 * d_Σ), para_subset))
    q1 += (1.0-α)/2.0 * exp(logpdf(DegenerateMvNormal(d_subset.μ, c^2 * d_Σ), para_draw))

    q0 = log(q0)
    q1 = log(q1)

    if (q0 == Inf && q1 == Inf)
        q0 = 0.0
    end
    return q0, q1
end

"""
```
function `compute_ESS(loglh::Vector{T}, current_weights::Vector{T}, ϕ_n::T, ϕ_n1::T;
                     old_loglh::Vector{T} = zeros(length(loglh))) where {T<:AbstractFloat}`
```
Compute ESS given log likelihood, current weights, ϕ_n, ϕ_{n-1}, and old log likelihood.
"""
function compute_ESS(loglh::Vector{T}, current_weights::Vector{T}, ϕ_n::T, ϕ_n1::T;
                     old_loglh::Vector{T} = zeros(length(loglh))) where T<:AbstractFloat
    N            = length(loglh)
    inc_weights  = exp.((ϕ_n1 - ϕ_n) * old_loglh + (ϕ_n - ϕ_n1) * loglh)
    new_weights  = current_weights .* inc_weights
    norm_weights = N * new_weights / sum(new_weights) # Normalize to N
    ESS          = N^2 / sum(norm_weights .^ 2)       # Transform back for ESS
    return ESS
end

function generate_param_blocks(n_params::Int64, n_blocks::Int64)
    if n_blocks == 1
        return [collect(1:n_params)]
    end

    rand_inds = shuffle(1:n_params)

    subset_length     = cld(n_params, n_blocks) # ceiling division
    last_block_length = n_params - subset_length*(n_blocks - 1)

    blocks_free = Vector{Vector{Int64}}(undef, n_blocks)
    for i in 1:n_blocks
        if i < n_blocks
            blocks_free[i] = rand_inds[((i-1)*subset_length + 1):(i*subset_length)]
        else
            # To account for the fact that the last block may be smaller than the others
            blocks_free[i] = rand_inds[end-last_block_length+1:end]
        end
    end
    blocks_free = [sort(p_block) for p_block in blocks_free]
    return blocks_free
end

"""
```
`generate_free_blocks(n_free_para::Int64, n_blocks::Int64)`
```

Return a Vector{Vector{Int64}} where each internal Vector{Int64} contains a subset of the range
1:n_free_para of randomly permuted indices. This is used to index out random blocks of free
parameters from the covariance matrix for the mutation step.
"""
function generate_free_blocks(n_free_para::Int64, n_blocks::Int64)
    rand_inds = shuffle(1:n_free_para)

    subset_length     = cld(n_free_para, n_blocks) # ceiling division
    last_block_length = n_free_para - subset_length*(n_blocks - 1)

    blocks_free = Vector{Vector{Int64}}(undef, n_blocks)
    for i in 1:n_blocks
        if i < n_blocks
            blocks_free[i] = rand_inds[((i-1)*subset_length + 1):(i*subset_length)]
        else
            # To account for the fact that the last block may be smaller than the others
            blocks_free[i] = rand_inds[end-last_block_length+1:end]
        end
    end
    return blocks_free
end

"""
```
`generate_all_blocks(blocks_free::Vector{Vector{Int64}}, free_para_inds::Vector{Int64})`
```

Return a Vector{Vector{Int64}} where each internal Vector{Int64} contains indices
corresponding to those in `blocks_free` but mapping to `1:n_para` (as opposed to
`1:n_free_para`). These blocks are used to reconstruct the particle vector by
inserting the mutated free parameters into the size `n_para,` particle vector,
which also contains fixed parameters.
"""
function generate_all_blocks(blocks_free::Vector{Vector{Int64}}, free_para_inds::Vector{Int64})
    n_free_para = length(free_para_inds)
    ind_mappings = Dict{Int64, Int64}()

    for (k, v) in zip(1:n_free_para, free_para_inds)
        ind_mappings[k] = v
    end

    blocks_all = similar(blocks_free)
    for (i, block) in enumerate(blocks_free)
        blocks_all[i] = similar(block)
        for (j, b) in enumerate(block)
            blocks_all[i][j] = ind_mappings[b]
        end
    end
    return blocks_all
end

"""
```
check_nan_ess(cloud::Cloud, stage::Int64, incremental_weights::Vector{T},
              normalized_weights::Vector{T}, savepath::String, debug_assertion::Bool) where {T <: AbstractFloat}
```
checks whether the calculated ESS is a `NaN` and determines what error output we should give in the
assertion error.
"""
function check_nan_ess(cloud::Cloud, stage::Int64, incremental_weights::Vector{T},
                       normalized_weights::Vector{T}, savepath::String, debug_assertion::Bool) where {T <: AbstractFloat}
    if isnan(cloud.ESS[stage])
        inf_loglh = any(x -> isinf(x), incremental_weights)
        nan_loglh = any(x -> isnan(x), incremental_weights)
        zero_norm_wts = sum(normalized_weights .^2) <= eps()
        nan_norm_wts = isnan(sum(normalized_weights .^2))
        assert_str = "No particles have non-zero weight."
        if inf_loglh
            assert_str *= " Some particles have approximately infinite log-likelihoods."
        end
        if nan_loglh
            assert_str *= " Some particles have approximately NaN log-likelihoods."
        end
        if zero_norm_wts
            assert_str *= " The squared sum of the normalized weights is at machine-error."
        end
        if nan_norm_wts
            assert_str *= " The squared sum of the normalized weights is returning a NaN."
            if nan_norm_wts && any(x -> isnan(x), normalized_weights)
                assert_str *= " Part of the reason is that one of the normalized weights is a NaN"
            end
        end
        if debug_assertion
            jldopen(replace(savepath, ".jld2" => "_debug_assertion.jld2"), true, true, true, IOStream) do file
                write(file, "incremental_weights", incremental_weights)
                write(file, "normalized_weights", normalized_weights)
                write(file, "cloud", cloud)
            end
        end

        @assert false assert_str
    end

    nothing
end


text:
"""
$(TYPEDEF)

The `QuaternionSpherical` joint type allows rotation in any direction. It is an
implementation of a ball-and-socket joint.

The 4-dimensional configuration vector ``q`` associated with a `QuaternionSpherical` joint
is the unit quaternion that describes the orientation of the frame after the joint
with respect to the frame before the joint. In other words, it is the quaternion that
can be used to rotate vectors from the frame after the joint to the frame before the
joint.

The 3-dimensional velocity vector ``v`` associated with a `QuaternionSpherical` joint is
the angular velocity of the frame after the joint with respect to the frame before
the joint, expressed in the frame after the joint (body frame).
"""
struct QuaternionSpherical{T} <: JointType{T} end

Base.show(io::IO, jt::QuaternionSpherical) = print(io, "Quaternion spherical joint")
Random.rand(::Type{QuaternionSpherical{T}}) where {T} = QuaternionSpherical{T}()
num_positions(::Type{<:QuaternionSpherical}) = 4
num_velocities(::Type{<:QuaternionSpherical}) = 3
has_fixed_subspaces(jt::QuaternionSpherical) = true
isfloating(::Type{<:QuaternionSpherical}) = false

@propagate_inbounds function rotation(jt::QuaternionSpherical, q::AbstractVector, normalize::Bool = true)
    Quat(q[1], q[2], q[3], q[4], normalize)
end

@propagate_inbounds function set_rotation!(q::AbstractVector, jt::QuaternionSpherical, rot::Rotation{3, T}) where {T}
    quat = convert(Quat{T}, rot)
    q[1] = quat.w
    q[2] = quat.x
    q[3] = quat.y
    q[4] = quat.z
    nothing
end

@propagate_inbounds function set_configuration!(q::AbstractVector, joint::Joint{<:Any, <:QuaternionSpherical}, rot::Rotation{3})
    @boundscheck check_num_positions(joint, q)
    @inbounds set_rotation!(q, joint_type(joint), rot)
    q
end

@propagate_inbounds function joint_transform(jt::QuaternionSpherical, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D, q::AbstractVector)
    quat = rotation(jt, q, false)
    Transform3D(frame_after, frame_before, quat)
end

@inline function motion_subspace(jt::QuaternionSpherical{T}, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector{X}) where {T, X}
    S = promote_type(T, X)
    angular = one(SMatrix{3, 3, S})
    linear = zero(SMatrix{3, 3, S})
    GeometricJacobian(frame_after, frame_before, frame_after, angular, linear)
end

@inline function constraint_wrench_subspace(jt::QuaternionSpherical{T}, joint_transform::Transform3D{X}) where {T, X}
    S = promote_type(T, X)
    angular = zero(SMatrix{3, 3, S})
    linear = one(SMatrix{3, 3, S})
    WrenchMatrix(joint_transform.from, angular, linear)
end

@inline function bias_acceleration(jt::QuaternionSpherical{T}, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector{X}, v::AbstractVector{X}) where {T, X}
    S = promote_type(T, X)
    zero(SpatialAcceleration{S}, frame_after, frame_before, frame_after)
end

@propagate_inbounds function configuration_derivative_to_velocity!(v::AbstractVector, jt::QuaternionSpherical, q::AbstractVector, q̇::AbstractVector)
    quat = rotation(jt, q, false)
    quatdot = SVector(q̇[1], q̇[2], q̇[3], q̇[4])
    v .= angular_velocity_in_body(quat, quatdot)
    nothing
end

@propagate_inbounds function configuration_derivative_to_velocity_adjoint!(fq, jt::QuaternionSpherical, q::AbstractVector, fv)
    quatnorm = sqrt(q[1]^2 + q[2]^2 + q[3]^2 + q[4]^2) # TODO: make this nicer
    quat = Quat(q[1] / quatnorm, q[2] / quatnorm, q[3] / quatnorm, q[4] / quatnorm, false)
    fq .= (velocity_jacobian(angular_velocity_in_body, quat)' * fv) ./ quatnorm
    nothing
end

@propagate_inbounds function velocity_to_configuration_derivative!(q̇::AbstractVector, jt::QuaternionSpherical, q::AbstractVector, v::AbstractVector)
    quat = rotation(jt, q, false)
    q̇ .= quaternion_derivative(quat, v)
    nothing
end

@propagate_inbounds function velocity_to_configuration_derivative_jacobian(jt::QuaternionSpherical, q::AbstractVector)
    quat = rotation(jt, q, false)
    velocity_jacobian(quaternion_derivative, quat)
end

@propagate_inbounds function configuration_derivative_to_velocity_jacobian(jt::QuaternionSpherical, q::AbstractVector)
    quat = rotation(jt, q, false)
    velocity_jacobian(angular_velocity_in_body, quat)
end

@propagate_inbounds function zero_configuration!(q::AbstractVector, jt::QuaternionSpherical)
    T = eltype(q)
    set_rotation!(q, jt, one(Quat{T}))
    nothing
end

@propagate_inbounds function rand_configuration!(q::AbstractVector, jt::QuaternionSpherical)
    T = eltype(q)
    set_rotation!(q, jt, rand(Quat{T}))
    nothing
end

@propagate_inbounds function joint_twist(jt::QuaternionSpherical{T}, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector{X}, v::AbstractVector{X}) where {T, X}
    S = promote_type(T, X)
    angular = SVector{3, S}(v)
    linear = zero(SVector{3, S})
    Twist(frame_after, frame_before, frame_after, angular, linear)
end

@propagate_inbounds function joint_spatial_acceleration(jt::QuaternionSpherical{T}, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector{X}, v::AbstractVector{X}, vd::AbstractVector{XD}) where {T, X, XD}
    S = promote_type(T, X, XD)
    angular = SVector{3, S}(vd)
    linear = zero(SVector{3, S})
    SpatialAcceleration(frame_after, frame_before, frame_after, angular, linear)
end

@propagate_inbounds function joint_torque!(τ::AbstractVector, jt::QuaternionSpherical, q::AbstractVector, joint_wrench::Wrench)
    τ .= angular(joint_wrench)
    nothing
end

# uses exponential coordinates centered around q0
@propagate_inbounds function local_coordinates!(ϕ::AbstractVector, ϕ̇::AbstractVector,
        jt::QuaternionSpherical, q0::AbstractVector, q::AbstractVector, v::AbstractVector)
    quat = inv(rotation(jt, q0, false)) * rotation(jt, q, false)
    rv = RodriguesVec(quat)
    ϕstatic = SVector(rv.sx, rv.sy, rv.sz)
    ϕ .= ϕstatic
    ϕ̇ .= rotation_vector_rate(ϕstatic, v)
    nothing
end

@propagate_inbounds function global_coordinates!(q::AbstractVector, jt::QuaternionSpherical, q0::AbstractVector, ϕ::AbstractVector)
    quat0 = rotation(jt, q0, false)
    quat = quat0 * Quat(RodriguesVec(ϕ[1], ϕ[2], ϕ[3]))
    set_rotation!(q, jt, quat)
    nothing
end

@propagate_inbounds normalize_configuration!(q::AbstractVector, jt::QuaternionSpherical) = set_rotation!(q, jt, rotation(jt, q, true))

@propagate_inbounds function is_configuration_normalized(jt::QuaternionSpherical, q::AbstractVector, rtol, atol)
    isapprox(quatnorm(rotation(jt, q, false)), one(eltype(q)); rtol = rtol, atol = atol)
end

@propagate_inbounds function principal_value!(q::AbstractVector, jt::QuaternionSpherical)
    quat = rotation(jt, q, false)
    set_rotation!(q, jt, principal_value(quat))
    nothing
end


text:
abstract type Graph end

graph(g::Graph) = g
labels(g::Graph; kw...) = [] #error("labels() not defined for $x")
edgestyles(g::Graph; kw...) = Dict()

makedot(g::Graph) = begin
    node(i, l) = """$i [label="$l"]\n"""
    N = [node(i, l) for (i, l) in enumerate(labels(g))]

    edge(a, b) = """$a -> $b [style="$(get(ES, (a, b), ""))"]\n"""
    ES = edgestyles(g)
    E = [edge(e.src, e.dst) for e in edges(graph(g))]

    """
    digraph {
    ratio=0.5
    ranksep=0.2
    node[
        width=0
        height=0
        margin=0.03
        shape=plaintext
        fontname=Courier
        fontsize=9
    ]
    edge [
        arrowsize=0.2
        penwidth=0.2
    ]
    $(N...)
    $(E...)
    }
    """
end

writedot(g::Graph) = writedot(tempname(), g)
writedot(name::AbstractString, g::Graph) = begin
    !endswith(name, ".dot") && (name *= ".dot")
    write(name, makedot(g))
    name
end

#TODO: wait until Graphviz_jll adds Windows support
import Conda
writeimage(name::AbstractString, g::Graph; format=nothing) = begin
    ext = splitext(name)[2]
    if isnothing(format)
        format = ext[2:end]
        isempty(format) && error("format unspecified")
    else
        format = string(format)
        ext != format && (name *= "."*format)
    end
    dot = writedot(g)
    let exe = joinpath(Conda.bin_dir(:cropbox), "dot"),
        cmd = `$exe -T$format $dot -o $name`
        success(cmd) || error("cannot execute: $cmd")
    end
    name
end

Base.show(io::IO, ::MIME"image/svg+xml", g::Graph) = begin
    f = writeimage(tempname(), g; format=:svg)
    s = read(f, String)
    print(io, s)
end


text:
@testset "Reverse" begin
    @test reverse(A₀) == A₀
    @test reverse(A₁) == A₁
    @test reverse(A₂) == -A₂
    @test reverse(A₃) == -A₃
    @test reverse(A₄) == A₄
    @test reverse(A * B) == reverse(B) * reverse(A)
end


text:
#-------------------------------------------------------------------------------------------------
#--
#-- Ususal libraries
#--
using CSV, Dates
using DataFrames, DataFramesMeta

#-------------------------------------------------------------------------------------------------
#--
#-- DATA
#--
# Update latest data from Neherlab
function updateData()

    ### COUNTRY CODES
    filename = "https://github.com/neherlab/covid19_scenarios_data/raw/master/country_codes.csv"
    download(filename, "data/country_codes.csv")

    ### LIST OF CASES
    filename = "https://opendata.ecdc.europa.eu/covid19/casedistribution/csv"
    download(filename, "data/ecdc.csv")

    ### ICU beds
    filename = "https://github.com/neherlab/covid19_scenarios_data/raw/master/hospital-data/ICU_capacity.tsv"
    download(filename, "data/ICU_capacity.tsv")

    ### Hospital beds
    filename = "https://github.com/neherlab/covid19_scenarios_data/raw/master/hospital-data/hospital_capacity.csv"
    download(filename, "data/hospital_capacity.csv")
end

function loadData()
    ### Population Data
    popData = CSV.read("data/populationData.tsv", delim = "\t")

    ### COUNTRY CODES
    country_codes = select(CSV.read("data/country_codes.csv"), :name, Symbol("alpha-3"))

    ### LIST OF CASES
    ecdc = CSV.read("data/ecdc.csv")

    ### ICU beds
    ICU_capacity = select(CSV.read("data/ICU_capacity.tsv", delim = "\t"), :country, :CriticalCare)

    ### Hospital beds
    hospital_capacity = select(
        CSV.read("data/hospital_capacity.csv", types = Dict(:COUNTRY => String), limit = 1267),
        :COUNTRY,
        :YEAR,
        :VALUE,
    )
    hospital_capacity = @where(hospital_capacity, Not(ismissing.(:COUNTRY)))

    ### Age pyramids
    age_distribution = CSV.read("data/country_age_distribution.csv")

    ### Last saved parameters
    optimisedParameters = try
        CSV.read("data/allCountryParameters.csv")
    catch
        nothing
    end

    return Dict(
        :codes => country_codes,
        :popData => popData,
        :cases => ecdc,
        :ICU => ICU_capacity,
        :Beds => hospital_capacity,
        :age_distribution => age_distribution,
        :optimisedParameters => optimisedParameters
    )
end

# Let's get the Neherlab repo's data (the files can be updated with `updateData()`).
COUNTRY_DB = loadData()

#-------------------------------------------------------------------------------------------------
#--
#-- Create a dictionary populated with the information of a given country
#--
function populateCountryData(country, hemisphere; useOptimised = true)

    # This is the start of a country specific data structure. A dictionary is good enough for that purpose.
    countryData = Dict()

    countryData = Dict([(:name, country), (:hemisphere, hemisphere)])

    countryData[:peak] = peakDate[countryData[:hemisphere]]
    countryData[:ϵ] = ϵ[countryData[:hemisphere]]
    countryData[:mitigation] = DEFAULT_MITIGATION

    # Cases from the ECDC database
    caseDB = COUNTRY_DB[:cases]

    # select the right country
    ecdc = @where(caseDB, country .== :countriesAndTerritories)

    # Add a date comlumn in the Date type and convert to days
    ecdc[!, :time] = Date.(ecdc.year, ecdc.month, ecdc.day)
    ecdc[!, :t] = date2days.(ecdc.time)

    # and order all dates in ascending order
    ecdc = @orderby(ecdc, :time)

    # Convert cases and deaths from daily deaths to cumulative amounts
    ecdc.cases = cumsum(ecdc.cases)
    ecdc.deaths = cumsum(ecdc.deaths)

    # Only retain entries where there is a strictily positive number of deaths
    ecdc = @where(ecdc, :deaths .> 0)

    # Extract country-specific information
    countryShort = ecdc[1, :countryterritoryCode]

    population = ecdc[1, :popData2018]

    ICU_capacity = try
        @where(COUNTRY_DB[:popData], occursin.(country, :name))[!, :ICUBeds][1]
    catch
        population / 10.0
    end
    ICU_capacity = convert(Float64, ICU_capacity)

    hospital_capacity = try
        @where(COUNTRY_DB[:popData], occursin.(country, :name))[!, :hospitalBeds][1]
    catch
        population / 100.0
    end
    hospital_capacity = convert(Float64, hospital_capacity)

    age_distribution = @where(COUNTRY_DB[:age_distribution], occursin.(country, :_key))[1, 2:10]

    countryData[:population] = population
    countryData[:country_code] = countryShort
    countryData[:cases] = ecdc
    countryData[:ICU_capacity] = ICU_capacity
    countryData[:hospital_capacity] = hospital_capacity
    countryData[:age_distribution] = age_distribution

    countryData[:lossFunction] = p -> singleCountryLoss(country, p)

    if useOptimised == true
        p =  @where(COUNTRY_DB[:optimisedParameters], occursin.(country, :country))[:, 3:end]
        if nrow(p) == 1
            countryData[:params] =  convert(Array, p)
        end
    else
        countryData[:params] = COUNTRY_INIT
        countryData[:range] = COUNTRY_RANGE
    end

    return countryData
end

# Update the record of cases for countries after an update
function updateCountryCases(country)

    # Cases from the ECDC database
    caseDB = COUNTRY_DB[:cases]

    # select the right country
    ecdc = @where(caseDB, country .== :countriesAndTerritories)

    # Add a date comlumn in the Date type and convert to days
    ecdc[!, :time] = Date.(ecdc.year, ecdc.month, ecdc.day)
    ecdc[!, :t] = date2days.(ecdc.time)

    # and order all dates in ascending order
    ecdc = @orderby(ecdc, :time)

    # Convert cases and deaths from daily deaths to cumulative amounts
    ecdc.cases = cumsum(ecdc.cases)
    ecdc.deaths = cumsum(ecdc.deaths)

    # Only retain entries where there is a strictily positive number of deaths
    ecdc = @where(ecdc, :deaths .> 0)

    global countryData[country][:cases] = ecdc
end

#--------------------------------------------------------------------------------------------------
#--
#-- Find the date at which cases exceed 5 for a given country
#--
function approximateModelStartRange(country::String)
    # Estimates when deaths number is about DEATH_AT_MODEL_START
    above = @where(countryData[country][:cases], :deaths .> DEATH_AT_MODEL_START)

    # Take first date
    if nrow(above) > 0
        first_date_above = first(above[:, :time])
    else
        first_date_above = BASE_DATE
    end

    # Convert to day
    above_day = date2days(first_date_above)

    return(above_day - 0.0, above_day + 50.0)
end


text:
module WebAssembly

# All of the WABt utilities are accessed through the following module:
module WABT
include("../deps/wabt_deps.jl")
end

module Binaryen
include("../deps/binaryen_deps.jl")
end

function __init__()
    WABT.check_deps()
    Binaryen.check_deps()
end

include("wasm.jl")
include("ir.jl")
include("passes.jl")
include("io.jl")
include("parser.jl")
include("interpret.jl")

module Instructions

using ..WebAssembly: Instruction, Const, Nop, Local, SetLocal, Op, Select, Convert,
  Block, If, Loop, Branch, Call, Return, Unreachable, nop, unreachable

export Instruction, Const, Nop, Local, SetLocal, Op, Select, Convert,
  Block, If, Loop, Branch, Call, Return, Unreachable, Label, Goto, nop, unreachable

end

end # module


text:
# Sorting a DVector using samplesort

function sample_n_setup_ref(d::DVector, sample_size; kwargs...)
    lp = localpart(d)
    llp = length(lp)
    np = length(procs(d))
    sample_size = llp > sample_size ? sample_size : llp
    sorted = sort(lp; kwargs...)
    sample = sorted[collect(1:div(llp,sample_size):llp)]
    ref = RemoteChannel(()->Channel(np+1))             # To collect parts to be sorted locally later.
                                                       # First element is the locally sorted vector
    put!(ref, sorted)
    return (sample, ref)
end

function scatter_n_sort_localparts{T}(d, myidx, refs::Array{RemoteChannel}, boundaries::Array{T}; by = identity, kwargs...)
    if d==nothing
        sorted = take!(refs[myidx])  # First entry in the remote channel is sorted localpart
    else
        sorted = sort(localpart(d); by = by, kwargs...)
    end

    # send respective parts to correct workers, iterate over sorted array
    p_sorted = 1
    for (i,r) in enumerate(refs)
        p_till = length(sorted)+1

        # calculate range to send to refs[i]
        ctr=1
        for x in sorted[p_sorted:end]
            if by(x) > by(boundaries[i+1])
                p_till = p_sorted+ctr-1
                break
            else
                ctr += 1
            end
        end

        if p_till == p_sorted
            @async put!(r, Array(T,0))
        else
            v = sorted[p_sorted:p_till-1]
            @async put!(r, v)
        end

        p_sorted = p_till
    end

    # wait to receive all of my parts from all other workers
    lp_sorting=T[]
    for _ in refs
        v = take!(refs[myidx])
        append!(lp_sorting, v)
    end

    sorted_ref=RemoteChannel()
    put!(sorted_ref, sort!(lp_sorting; by = by, kwargs...))
    return (sorted_ref, length(lp_sorting))
end

function compute_boundaries{T}(d::DVector{T}; kwargs...)
    pids = procs(d)
    np = length(pids)
    sample_sz_on_wrkr = 512

    results = asyncmap(p -> remotecall_fetch(sample_n_setup_ref, p, d, sample_sz_on_wrkr; kwargs...), pids)

    samples = Array(T,0)
    for x in results
        append!(samples, x[1])
    end
    sort!(samples; kwargs...)
    samples[1] = typemin(T)

    refs=RemoteChannel[x[2] for x in results]

    boundaries = samples[[1+(x-1)*div(length(samples), np) for x in 1:np]]
    push!(boundaries, typemax(T))

    return (boundaries, refs)
end

"""
    sort(d::DVector; sample=true, kwargs...) -> DVector

Sorts and returns a new distributed vector.

The sorted vector may not have the same distribution as the original.

Keyword argument `sample` can take values:

- `true`: A sample of max size 512 is first taken from all nodes. This is used to balance the distribution of the sorted array on participating workers. Default is `true`.

- `false`: No sampling is done. Assumes a uniform distribution between min(d) and max(d)

- 2-element tuple of the form `(min, max)`: No sampling is done. Assumes a uniform distribution between specified min and max values

- Array{T}: The passed array is assumed to be a sample of the distribution and is used to balance the sorted distribution.

Keyword argument `alg` takes the same options `Base.sort`
"""
function Base.sort{T}(d::DVector{T}; sample=true, kwargs...)
    pids = procs(d)
    np = length(pids)

    # Only `alg` and `sample` are supported as keyword arguments
    if length(filter(x->!(x in (:alg, :by)), [x[1] for x in kwargs])) > 0
        throw(ArgumentError("Only `alg`, `by` and `sample` are supported as keyword arguments"))
    end

    if sample==true
        boundaries, refs = compute_boundaries(d; kwargs...)
        presorted=true

    elseif sample==false
        # Assume an uniform distribution between min and max values
        minmax=asyncmap(p->remotecall_fetch(d->(minimum(localpart(d)), maximum(localpart(d))), p, d), pids)
        min_d = minimum(T[x[1] for x in minmax])
        max_d = maximum(T[x[2] for x in minmax])

        return sort(d; sample=(min_d,max_d), kwargs...)

    elseif isa(sample, Tuple)
        # Assume an uniform distribution between min and max values in the tuple
        lb=sample[1]
        ub=sample[2]

        @assert lb<=ub

        s = Array(T, np)
        part = abs(ub - lb)/np
        (isnan(part) || isinf(part)) && throw(ArgumentError("lower and upper bounds must not be infinities"))

        for n in 1:np
            v = lb + (n-1)*part
            if T <: Integer
                s[n] = round(v)
            else
                s[n] = v
            end
        end
        return sort(d; sample=s, kwargs...)

    elseif isa(sample, Array)
        # Provided array is used as a sample
        samples = sort(copy(sample))
        samples[1] = typemin(T)
        boundaries = samples[[1+(x-1)*div(length(samples), np) for x in 1:np]]
        push!(boundaries, typemax(T))
        presorted=false

        refs=RemoteChannel[RemoteChannel(p) for p in procs(d)]
    else
        throw(ArgumentError("keyword arg `sample` must be Boolean, Tuple(Min,Max) or an actual sample of data : " * string(sample)))
    end

    local_sort_results = Array(Tuple, np)

    Base.asyncmap!((i,p) -> remotecall_fetch(
            scatter_n_sort_localparts, p, presorted ? nothing : d, i, refs, boundaries; kwargs...),
                                    local_sort_results, 1:np, pids)

    # Construct a new DArray from the sorted refs. Remove parts with 0-length since
    # the DArray constructor_from_refs does not yet support it. This implies that
    # the participating workers for the sorted darray may be different from the original
    # for highly non-uniform distributions.
    local_sorted_refs = RemoteChannel[x[1] for x in filter(x->x[2]>0, local_sort_results)]
    return DArray(local_sorted_refs)
end


text:
abstract type AbstractContinuation <: Function end

struct RationalContinuation{T}
    pmax::T
    kmax::Int
    b::T
end
function Base.getproperty(c::RationalContinuation, f::Symbol)
    f === :length && return length(c)
    return getfield(c, f)
end
function RationalContinuation(xmin, steps::Int)
    pmax = (1 - xmin)/xmin
    kmax = steps
    b = 1/steps * (1 - sqrt(xmin)) / (1 + sqrt(xmin))
    return RationalContinuation(pmax, kmax, b)
end
function Base.iterate(s::RationalContinuation, k = 1)
    (k > s.kmax + 1) && return nothing
    b = s.b
    return s(k), k + 1
end
Base.length(s::RationalContinuation) = s.kmax + 1
function (s::RationalContinuation{T})(k) where T
    k = clamp(k, 1, s.kmax + 1)
    b = s.b
    return 4*b*(k-1) / (1 - b*(k-1))^2
end

struct FixedContinuation{T} <: AbstractContinuation
    param::T
    length::Int
end
function Base.iterate(s::FixedContinuation, x=1)
    (x > s.length) && return nothing
    s.param, x+1
end
Base.length(s::FixedContinuation) = s.length
(s::FixedContinuation{T})(x...) where T = s.param

"""
p(x) = 1/(a + b*ℯ^(-c*x))
"""
struct SigmoidContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
SigmoidContinuation(; kwargs...) = SigmoidContinuation{Float64}(; kwargs...)
function SigmoidContinuation{T}(; c::T = T(0.1), start::T = T(1), finish::T = T(5), steps::Int = 30, min::T = -Inf) where T
    a = 1 - T(finish-start)/finish/(ℯ^(-c) - ℯ^(-steps*c)) * ℯ^(-c)
    b = T(finish-start)/finish/(ℯ^(-c) - ℯ^(-steps*c))
    SigmoidContinuation(a, b, c, steps, min)
end
function Base.iterate(s::SigmoidContinuation, x=1)
    (x > s.length) && return nothing
    max(1 / (s.a + s.b*ℯ^(-s.c*x)), s.min), x+1
end
Base.length(s::SigmoidContinuation) = s.length
(s::SigmoidContinuation{T})(x) where T = max(1 / (s.a + s.b * ℯ^(-s.c * T(x))), s.min)

"""
p(x) = a(x-b*steps)^3 + c
"""
struct CubicSplineContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
CubicSplineContinuation(; kwargs...) = CubicSplineContinuation{Float64}(; kwargs...)
function CubicSplineContinuation{T}(;b::T=T(0.5), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf) where T
    a = (finish - start) / (steps^3 * (1 - b)^3 - (1 - b * steps)^3)
    c = start - a * (1 - b * steps)^3
    CubicSplineContinuation(a, b, c, steps, min)
end

function Base.iterate(s::CubicSplineContinuation, x=1)
    x > s.length && return nothing
    max(s.a*(x-s.b*s.length)^3 + s.c, s.min), x+1
end
Base.length(s::CubicSplineContinuation) = s.length
(s::CubicSplineContinuation{T})(x) where T = max(s.a*(T(x)-s.b*s.length)^3 + s.c, s.min)

"""
p(x) = a*x^b + c
"""
struct PowerContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
PowerContinuation(; kwargs...) = PowerContinuation{Float64}(; kwargs...)
function PowerContinuation{T}(; b=T(2), start=T(1), finish=T(5), steps::Int=30, min = -T(Inf)) where T
    a = (finish - start) / max(T(1), steps^b - 1)
    c = start - a
    PowerContinuation(a, b, c, steps, min)
end
function Base.iterate(s::PowerContinuation, x=1)
    x > s.length && return nothing
    max(s.a * x^s.b + s.c, s.min), x+1
end
Base.length(s::PowerContinuation) = s.length
(s::PowerContinuation{T})(x) where T = max(s.a * T(x)^s.b + s.c, s.min)

"""
p(x) = a*ℯ^(b*x) + c
"""
struct ExponentialContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
ExponentialContinuation(; kwargs...) = ExponentialContinuation{Float64}(; kwargs...)
function ExponentialContinuation{T}(; b::T = T(0.1), start::T = T(1), finish::T = T(5), steps::Int = 30, min::T = -Inf) where T
    a = (finish - start) / max(T(1), ℯ^(b * steps) - ℯ^(b))
    c = start - a * ℯ^(b)
    ExponentialContinuation(a, b, c, steps, min)
end
function Base.iterate(s::ExponentialContinuation, x=1)
    x > s.length && return nothing
    max(s.a*ℯ^(s.b*x) + s.c, s.min), x+1
end
Base.length(s::ExponentialContinuation) = s.length
(s::ExponentialContinuation{T})(x) where T = max(s.a*ℯ^(s.b*T(x)) + s.c, s.min)

"""
p(x) = a*log(b*x) + c
"""
struct LogarithmicContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
LogarithmicContinuation(; kwargs...) = LogarithmicContinuation{Float64}(; kwargs...)
function LogarithmicContinuation{T}(; b::T = T(1), start::T = T(1), finish::T = T(5), steps::Int = 30, min::T = -Inf) where T
    a = (finish - start) / max(T(1), log(b * steps) - log(b))
    c = start - a * log(b)
    LogarithmicContinuation(a, b, c, steps, min)
end
function Base.iterate(s::LogarithmicContinuation, x=1)
    x > s.length && return nothing
    max(s.a*log(s.b*x) + s.c, s.min), x+1
end
Base.length(s::LogarithmicContinuation) = s.length
(s::LogarithmicContinuation{T})(x) where T = max(s.a * log(s.b * T(x)) + s.c, s.min)

Continuation(::PowerPenalty; kwargs...) = Continuation(PowerPenalty; kwargs...)
Continuation(::RationalPenalty; kwargs...) = Continuation(RationalPenalty; kwargs...)
function Continuation(::Type{<:PowerPenalty}; steps, pmax = 5.0, kwargs...)
    return PowerContinuation{Float64}(b = 1.0, start = 1.0, steps = steps + 1, finish = pmax)
end
function Continuation(::Type{<:RationalPenalty}; steps, xmin = 0.001, kwargs...)
    return RationalContinuation(xmin, steps)
end


text:
using Distributions, Random

function simulatePoisson(;Nd=1,Ns=10,a0=1.0,a1=.5,a0_sig=.3,kwargs...)
    N = Nd*Ns
    y = fill(0,N)
    x = fill(0.0,N)
    idx = similar(y)
    i = 0
    for s in 1:Ns
        a0s = rand(Normal(0,a0_sig))
        logpop = rand(Normal(9,1.5))
        λ = exp(a0 + a0s + a1*logpop)
        for nd in 1:Nd
            i+=1
            x[i] = logpop
            idx[i] = s
            y[i] = rand(Poisson(λ))
        end
    end
    return (y=y,x=x,idx=idx,N=N,Ns=Ns)
 end


text:
## sets module
## TODO: SymSets class so that we can better manage dispatch?

"""
`SymPy` has the [sets](http://docs.sympy.org/latest/modules/sets.html) module that provides:

* `FiniteSet` to specify a finite set of values
* `Interval` to specify an interval
* `Union` to specify a union of finite sets and/or intervals.

The `FiniteSet` and `Interval` constructors are provided, unions are formed from these.

The `S` module provides some built-in sets: `S.Reals`, `S.Integers`, `S.Naturals`, `S.Naturals0`, and `S.UniversalSet`, `S.EmptySet`.

Various methods for sets and intervals are provided, though not all are defined for unions:

* `powerset`
* `union`, `intersection`, `contains`, `boundary`
* `complement`, `is_disjoint`, `is_subset`, `is_superset`, `is_proper_subset`
* `sup`, `inf`, `measure`
* `is_Union`, `is_Interval`, `is_FiniteSet`
* `as_relational`, `left_open`, `right_open`, `is_left_unbounded`, `is_right_unbounded`

Examples:
```
a,b = Interval(0,1), Interval(1//2,2, false, true)  # [0,1], [1/2, 2)
union(a, b)     # [0,2), also  a ∪ b
intersect(a,b)  # [1/2,1],  also a ∩ b
contains(a, 3/2), contains(b, 3/2) # false, true
inf(a), sup(b)  # 1, 2
right_open(b)   # true
complement(b, S.Reals)  # (-∞, 1/2) ∪ [2, ∞)
complement(b, a)  # [0, 1/2)  (This is a \ b, but that notation is not provided)

x = symbols("x")
as_relational(a, x)  #  0 ≤ x ∧ x ≤ 1

s = FiniteSet(1,2)
powerset(s) # {∅, {1}, {2}, {1, 2}}
elements(s)   # collect(args(s))
```

SymPy `FiniteSet`s are not iterable, but can be collected with `collect(args(u))`.

SymPy `Union`s are not iterable, but can be collected after calling `args` via `collect(args(s))`.

"""
sypmy_sets = nothing

module S
using SymPy
function init_set()
    S = sympy[:S]
    global Reals = S[:Reals]
    global UniversalSet = S[:UniversalSet]
    global Naturals = S[:Naturals]
    global Naturals0 = S[:Naturals0]
    global Integers = S[:Integers]
    global EmptySet = S[:EmptySet]
end
end

export S, FiniteSet, Interval, ProductSet, ConditionSet
export powerset, imageset
export contains,  boundary, sup, measure
export is_disjoint
export is_proper_subset, is_subset, is_superset
export Interval
export as_relational
export is_left_unbounded, is_right_unbounded, left_open, right_open
export is_FiniteSet, is_Union, is_Interval

"Convert a SymPy finite set to a Set of Sym objects"
function Base.convert(::Type{Set}, s::Sym)
    is_FiniteSet(s) || throw(ArgumentError("`s` must be a finite set"))
    s1 = Set(Any[u for u in s.x])
    for el in copy(s1)
        if !isa(el, Set) && !(el.x[:is_Symbol]) && is_FiniteSet(s) && length(el.x) > 1
            ## replace python FiniteSet with julian Set
            setdiff!(s1, Set([el]))
            push!(s1, Set(convert(Set, el)))
        end
    end
    s1
end

## to convert a Union to an array of intervals
function _union_to_intervals(s::Sym)
    is_Union(s) || throw(ArgumentError("`s` must be a Union of Intervals"))
    collect(args(s))
end

"Find power set of set"
powerset(s::Sym) = s[:powerset]()

"Does I contain x?"
Base.contains(I::Sym, x) = (I[:contains](x) == Sym(true))
Base.in(x::Number, I::Sym) = contains(I, x)

"Elements of finite set"
elements(x::Sym) = PyCall.pyeval("[i for i in s]", s=project(x))
export elements

VERSION < v"0.5.0-" && eval(Expr(:import, :Base, :complement))
"Complement of set within the universe"
complement(I::Sym, U::Sym=S.Reals) = I[:complement](U)
export complement

"boundary, returnsa set"
boundary(I::Sym) = I.x[:boundary]

VERSION < v"0.5.0-" && eval(Expr(:import, :Base, :inf))
"Infinum of I"
inf(I::Sym) = I.x[:inf]
export inf

"Intersection of two intervals"
Base.intersect(I::Sym, J::Sym) = I[:intersect](J)

"Are `I` and `J` disjoint?"
is_disjoint(I::Sym, J::Sym) = I[:is_disjoint](J) == Sym(true)

"Is `J` a proper subset of `I`?"
is_proper_subset(I::Sym, J::Sym) = I[:is_proper_subset](J) == Sym(true)

"Is `J` a  subset of `I`?"
is_subset(I::Sym, J::Sym) = I[:is_subset](J) == Sym(true)

"Is `I` a  superset of `J`?"
is_superset(I::Sym, J::Sym) = I[:is_superset](J) == Sym(true)

"Lebesgue mesuare of an interval"
measure(I::Sym) = I.x[:measure]

"Supremum of a set"
sup(I::Sym) = I.x[:sup]

"Union of two intervals"
Base.union(I::Sym, J::Sym) = I[:union](J)

"Symmetric difference of two intervals, in one or other, but not both: `(I ∪ J) \ (I  ∩ J)"
Base.symdiff(I::Sym, J::Sym) = complement(intersection(I,J), union(I,J))

"rexpress I in terms of relations involving variable `x`"
as_relational(I::Sym, x::Sym) = I[:as_relational](x)

# not implemented: `end` => `sup`, `left` =>

"Looks like (-oo, a)?"
is_left_unbounded(I::Sym) = I.x[:is_left_unbounded] == Sym(true)

"Looks like (a, oo)?"
is_right_unbounded(I::Sym) = I.x[:is_right_unbounded] == Sym(true)

"Looks like (a, b...?"
left_open(I::Sym) = I.x[:left_open] == Sym(true)

"Looks like ..., b)?"
right_open(I::Sym) = I.x[:right_open] == Sym(true)

"Is `I` a finite set?"
is_FiniteSet(I::Sym) = I.x[:is_FiniteSet]

"Is `I` an interval?"
is_Interval(I::Sym) = I.x[:is_Interval]

"Is `I` a union?"
is_Union(I::Sym) = I.x[:is_Union]

function init_sets()
    S.init_set()

"""
    `FiniteSet(1,2,3)`, `FiniteSet(1:10...)`: create a finite set

See also `Interval` to create subsets of the real line.
"""
    "FiniteSet: http://docs.sympy.org/latest/modules/sets.html"
    global FiniteSet(args...) = sympy_meth(:FiniteSet, args...)

    "ProductSet: http://docs.sympy.org/latest/modules/sets.html"
    global ProductSet(args...) = sympy_meth(:ProductSet, args...)

    """
    Means to filter a set to pull out elements by a condition.

    ConditionSet:  A set of elements which satisfies a given condition.

    `ConditionSet(x, condition, S) = {x | condition(x) == true for x in S}`

    (Introduced in SymPy 1.0)
    """
    global ConditionSet(var::Sym, pred::Sym, S) = sympy_meth(:ConditionSet, var, pred, S)

    """
    Represent a ComplexRegion.
    working?
```
    I, J = Interval(0,1), Interval(0,2)
    R = ComplexRegion(I * J)
    1 in R        # true
    2 + 1im in R  # false
```
"""
    ComplexRegion(IJ::Sym; kwargs...) = sympy_meth(:ComplexRegion, IJ; kwargs...)

    "imageset: http://docs.sympy.org/latest/modules/sets.html"
    global imageset(fn::Function, args...) = begin
        x = Sym("x")
        imageset(x, fn(x), args...)
    end
    global imageset(args...) = sympy_meth(:imageset, args...)

    ## Interval

"""
Create an interval object

```
Interval(0,1) # [0,1]
Interval(0,1,true, false) # (0,1]
```

"""
    global Interval(l,r,left_open=false, right_open=false) = sympy_meth(:Interval, Sym(l), Sym(r), left_open ,right_open)

end


text:
module Helpers
export default_mean_f, default_add_x, default_residual

"""
Default mean function.
"""
function default_mean_f(
    Σ::AbstractArray{Float64, 2}, m_w::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 1}
    Σ' * m_w
end

"""
Default addition function.
"""
function default_add_x(
    y::AbstractArray{Float64, 1}, x::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 1}
    y + x
end

"""
Default residual function.
"""
function default_residual(
    y::AbstractArray{Float64, 1}, x::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 1}
    y - x
end

function default_residual(
    y::AbstractArray{Float64, 2}, x::AbstractArray{Float64, 1},
)::AbstractArray{Float64, 2}
    y .- x'
end

end

max_stars_repo_name: pxl-th/KalmanFilters

max_issues_repo_name: pxl-th/KalmanFilters

max_forks_repo_name: pxl-th/KalmanFilters


text:
## Helper functions to deal with multivariate parameters

# check_samples(samples, names) returns a widden matrix of scalar samples, and an associative array of parameters names and positions in the matrix
function check_samples(samples::TV, ::Nothing) where TV <: AbstractMatrix
    names = [Symbol("arg$i") for i = 1:size(samples, 1)]
    return check_samples(samples, names)
end

function check_samples(samples::TV, names::Tuple) where {T <: Real, TV <: AbstractMatrix{T}}
    return samples, NamedTuple{names}(eachindex(names))
end

function check_samples(samples, names)
    nn = length(names)
    nn == size(samples, 1) || throw(DimensionMismatch("$nn names for $(size(samples, 1)) parameters. Make sure the samples dimensions are nparams x nsamples"))
    ns = size(samples,2)
    nd = 0
    positions = Vector(undef, nn)
    for ni = 1:nn
        l = length(samples[ni, 1])
        positions[ni] = l > 1 ? (nd+1:nd+l) : nd+1
        nd += l
    end
    samples_wide = zeros(nd, ns)
    for i = 1:ns
        samples_wide[:,i] .= vcat(widden.(samples[:,i])...)
    end
    return samples_wide, NamedTuple{Tuple(names)}(positions)
end

widden(x::Real) = x
widden(x::AbstractArray) = vec(x)

# format_samples(samples, name_map) returns a compact matrix of samples of various dimensions, given the map name_map
function format_samples(samples::TV, name_map::NamedTuple) where {T <: Real, TV <: AbstractMatrix{T}}
    nd = length(name_map)
    ns = size(samples, 2)
    samples_comp = Matrix(undef, nd, ns)
    for i = 1:ns
        for j = 1:nd
            samples_comp[j,i] = samples[name_map[j],i]
        end
    end
    return samples_comp
end


text:
# Kernel definition is taken from here: https://en.wikipedia.org/wiki/Kernel_(statistics)#Kernel_functions_in_common_use

export Epanechnikov_kernel, Uniform_kernel, Triangular_kernel, Quartic_Kernel, Triweight_Kernel, Tricube_Kernel, Gaussian_Kernel, Cosine_Kernel, Logistic_Kernel, Sigmoid_Kernel, Silverman_Kernel

function Epanechnikov_kernel(t)
    if abs(t) > 1
        return 0
    else
        return 0.75*(1-t^2)
    end
end

function Uniform_kernel(t)
    if abs(t) > 1
        return 0
    else
        return 0.5
    end
end

function Triangular_kernel(t)
    if abs(t) > 1
        return 0
    else
        return (1-abs(t))
    end
end

function Quartic_Kernel(t)
  if abs(t)>0
    return 0
  else
    return (15*(1-t^2)^2)/16
  end
end

function Triweight_Kernel(t)
  if abs(t)>0
    return 0
  else
    return (35*(1-t^2)^3)/32
  end
end

function Tricube_Kernel(t)
  if abs(t)>0
    return 0
  else
    return (70*(1-abs(t)^3)^3)/80
  end
end

function Gaussian_Kernel(t)
  exp(-0.5*t^2)/(sqrt(2*π))
end

function Cosine_Kernel(t)
  if abs(t)>0
    return 0
  else
    return (π*cos(π*t/2))/4
  end
end

function Logistic_Kernel(t)
  1/(exp(t)+2+exp(-t))
end

function Sigmoid_Kernel(t)
  2/(π*(exp(t)+exp(-t)))
end

function Silverman_Kernel(t)
  sin(abs(t)/2+π/4)*0.5*exp(-abs(t)/sqrt(2))
end


text:
using BenchmarkTools
using TaylorIntegration
using DiffEqCallbacks
using ParameterizedFunctions
using RecursiveArrayTools
using Parameters
using Plots: plot, plot!, scatter, scatter!, histogram, histogram!, savefig,
    pgfplots, contour
using Colors

pgfplots()

function prob_setup(g, t; rescaling=false, Ttr=1e3, τ=5.)
    p0, q0 = ics_separate(g, E, PoincareRand(n=5), 0.15)
    z0 = ics(g)

    p = PhysicalParameters(B=0.15)
    if !rescaling
        prob1 = DynamicalODEProblem(ṗ, q̇, p0[1], q0[1], t, p)
        prob2 = ODEProblem(ż, z0[1], t, p)
    else
        alg = TimeRescaling(T=t, Ttr=Ttr, τ=τ)
        prob1 = λproblem(p0[1], q0[1], alg, params=p)
        prob2 = λproblem(z0[1], alg, params=p)
    end
    @unpack A,B,D = p
    prob3 = ODEProblem(h_eqs, Array(z0[1]), t, (A,B,D))

    return prob1, prob2, prob3
end

h_eqs = @ode_def HamiltEqs begin
  dp₀ = -A * q₀ - 3 * B / √2 * (q₂^2 - q₀^2) - D * q₀ * (q₀^2 + q₂^2)
  dp₂ = -q₂ * (A + 3 * √2 * B * q₀ + D * (q₀^2 + q₂^2))
  dq₀ = A * p₀
  dq₂ = A * p₂
end A B D

function g(resid, u)
    resid[1] = H((u[1],u[2]),(u[3],u[4]), PhysicalParameters(B=0.15)) - E
    resid[2:4] .= 0
end

const cb = ManifoldProjection(g, nlopts=Dict(:ftol=>1e-13))
const E = 0.1

function timings(g, t; rescaling=false)
    p1, p2, p3 = prob_setup(g, t, rescaling=rescaling)
    suite = BenchmarkGroup()

    suite["ODEProblem"] = BenchmarkGroup()
    suite["ODEProblem"]["Vern9"] = @benchmarkable solve($p2, Vern9(), abstol=1e-14, reltol=1e-14)
    if !rescaling
        suite["ODEProblem"]["Vern9+ManifoldProjection"] = @benchmarkable solve($p3, Vern9(), callback=cb, abstol=1e-14, reltol=1e-14)
        suite["ODEProblem"]["TaylorMethod"] = @benchmarkable solve($p3, TaylorMethod(50), abstol=1e-20)
    end

    suite["DynamicalODEProblem"] = BenchmarkGroup()
    suite["DynamicalODEProblem"]["DPRKN12"] =  @benchmarkable solve($p1, DPRKN12(), abstol=1e-14, reltol=1e-14)
    suite["DynamicalODEProblem"]["KahanLi8"] =  @benchmarkable solve($p1, KahanLi8(), dt=1e-2)
    suite["DynamicalODEProblem"]["SofSpa10"] =  @benchmarkable solve($p1, SofSpa10(), dt=1e-2)

    tune!(suite)
    run(suite)
end

energy_err(sol,offset=0) = DiffEqArray(map(
    i->H((sol[1,i], sol[2,i]),
         (sol[3+offset,i], sol[4+offset,i]),
         PhysicalParameters(B=0.15)
        ) - E,
    axes(sol,2)), sol.t)

function E_conservation(g, t; rescaling=false, short=true)
    subsampling = rescaling ? 10 : 1
    if short
        Ttr = 0.
        saveat = 0.1
        subsampling = rescaling ? subsampling ÷ 10 : subsampling
    else
        Ttr = 1e3
        saveat = 500.
    end
    p = prob_setup(g, t, rescaling=rescaling, Ttr=Ttr)

    suite = Dict{String,NamedTuple}()
    GC.gc()
    sol, t = @timed solve(p[2], Vern9(), abstol=1e-14, reltol=1e-14)
    if short
        suite["Vern9"] = (t=t, err=energy_err(sol))
    else
        suite["Vern9"] = (t=t, err=energy_err(sol)[1:1000:end])
    end
    if !rescaling
        GC.gc()
        sol, t = @timed solve(p[3], Vern9(), callback=cb, abstol=1e-14, reltol=1e-14)
        if short
            suite["Vern9+ManifoldProjection"] = (t=t, err=energy_err(sol))
            GC.gc()
            sol, t = @timed solve(p[3], TaylorMethod(50), abstol=1e-20)
            suite["TaylorMethod"] = (t=t, err=energy_err(sol))
        else
            # forced subsampling
            suite["Vern9+ManifoldProjection"] = (t=t, err=energy_err(sol)[1:1000:end])
        end
    end
    offset = rescaling ? 2 : 0
    GC.gc()
    sol, t = @timed solve(p[1], DPRKN12(), abstol=1e-14, reltol=1e-14)
    if short
        suite["DPRKN12"] = (t=t, err=energy_err(sol, offset)[1:subsampling:end])
    else
        # forced subsampling
        suite["DPRKN12"] = (t=t, err=energy_err(sol, offset)[1:500:end][1:subsampling:end])
    end
    GC.gc()
    sol, t = @timed solve(p[1], KahanLi8(), dt=1e-2, saveat=saveat)
    suite["KahanLi8"] = (t=t, err=energy_err(sol, offset)[1:subsampling:end])
    GC.gc()
    sol, t = @timed solve(p[1], SofSpa10(), dt=1e-2, saveat=saveat)
    suite["SofSpa10"] = (t=t, err=energy_err(sol, offset)[1:subsampling:end])

    return suite
end

function collect_results(g, t; rescaling=false, short=true)
    E_results = E_conservation(g, t, rescaling=rescaling, short=short)
    if short
        t_results = timings(g, t, rescaling=rescaling)
        integ1 = keys(t_results["ODEProblem"])
        integ2 = keys(t_results["DynamicalODEProblem"])
        solvers = Iterators.flatten((integ1,integ2))

        # times are in nanoseconds
        ts1 = [minimum(t_results["ODEProblem"][i].times) for i in integ1]
        ts2 = [minimum(t_results["DynamicalODEProblem"][i].times) for i in integ2]
        ts = Iterators.flatten((ts1,ts2))
    else
        solvers = keys(E_results)
        ts = [E_results[i].t for i in solvers]
    end
    E_err = [E_results[i].err for i in solvers]

    return solvers, ts, E_err
end

function short_benchmark(g; rescaling=false)
    solvers, ts, E_errs = collect_results(g, 100., rescaling=rescaling, short=true)
    p1 = plot(background_color=bg,
              xlabel="t",
              ylabel="Energy error",
              framestyle=:grid,
              tex_output_standalone=true
        )
    for (E_err,solver) in zip(E_errs, solvers)
        plot!(p1, E_err, label=solver)
    end
    p2 = plot(background_color=bg,
              xlabel="Solver",
              ylabel="Computation time (ms)",
              framestyle=:grid,
              tex_output_standalone=true,
              legend=false
        )
    scatter!(p2, string.(solvers), ts./1e6, m=10, color=ac, markerstrokealpha=0,
        markerstrokewidth=0)

    return p1, p2
end

function long_benchmark(g; rescaling=false)
    solvers, ts, E_errs = collect_results(g, 1e5, rescaling=rescaling, short=false)
    p1 = plot(background_color=bg,
              xlabel="t",
              ylabel="Energy error",
              framestyle=:grid,
              tex_output_standalone=true
        )
    for (E_err,solver) in zip(E_errs, solvers)
        plot!(p1, E_err, label=solver)
    end
    p2 = plot(background_color=bg,
              xlabel="Solver",
              ylabel="Computation time (s)",
              framestyle=:grid,
              tex_output_standalone=true,
              legend=false
        )
    scatter!(p2, string.(solvers), ts, m=10, color=ac, markerstrokewidth=0,
        markerstrokealpha=0)

    return p1, p2
end


text:
## This is to test the first derivative of sin
using MAT,Plots,LinearAlgebra
const USE_GPU=false;  # Use GPU? If this is set false, then no GPU needs to be available
using ParallelStencil
using ParallelStencil.FiniteDifferences1D
@static if USE_GPU
    @init_parallel_stencil(CUDA,Float64,3);
else
    @init_parallel_stencil(Threads,Float64,3);
end
dx=.001;
x=0:dx:6*pi;
nx=length(x);
y=cos.(x);
## analytical expression of d^2y/dx^2
d2y_dx2=-cos.(x);
## d^2y/dx^2 2nd-order accuracy
@parallel function Dx_inn_2nd_order_accuracy(in,out,dx)
    @inn(out)=@d_1(in)/dx;
    return nothing
end

@parallel_indices (iy) function u_1_plus_2nd_order_accuracy(in,out)
out[2:end-1]=in[:];
return nothing
end

@parallel_indices (iy) function u_1_minus_2nd_order_accuracy(in,out)
out[1:end-2]=in[:];
return nothing
end

dtt=@zeros(nx-2,);
tt=@zeros(nx,);
d2y_dx2_2nd_order_accuracy=@zeros(nx,);

@parallel Dx_inn_2nd_order_accuracy(y,dtt,dx);
@parallel (1:1) u_1_minus_2nd_order_accuracy(dtt,tt);
@parallel Dx_inn_2nd_order_accuracy(tt,dtt,dx);
@parallel (1:1) u_1_plus_2nd_order_accuracy(dtt,d2y_dx2_2nd_order_accuracy);
## d^2y/dx^2 8th-order accuracy
@parallel function Dx_inn_8th_order_accuracy(in,out,dx)
    @inn(out)=@d_8(in)/dx;
    return nothing
end

@parallel_indices (iy) function u_1_plus_8th_order_accuracy(in,out)
out[4:end-2]=in[:];
return nothing
end

@parallel_indices (iy) function u_1_minus_8th_order_accuracy(in,out)
out[3:end-3]=in[:];
return nothing
end

dtt=@zeros(nx-5,);
tt=@zeros(nx,);
d2y_dx2_8th_order_accuracy=@zeros(nx,);

@parallel Dx_inn_8th_order_accuracy(y,dtt,dx);
@parallel (1:1) u_1_minus_8th_order_accuracy(dtt,tt);
@parallel Dx_inn_8th_order_accuracy(tt,dtt,dx);
@parallel (1:1) u_1_plus_8th_order_accuracy(dtt,d2y_dx2_8th_order_accuracy);
## d^2y/dx^2 12th-order accuracy
@parallel function Dx_inn_12th_order_accuracy(in,out,dx)
    @inn(out)=@d_12(in)/dx;
    return nothing
end

@parallel_indices (iy) function u_1_plus_12th_order_accuracy(in,out)
out[6:end-4]=in[:];
return nothing
end

@parallel_indices (iy) function u_1_minus_12th_order_accuracy(in,out)
out[5:end-5]=in[:];
return nothing
end

dtt=@zeros(nx-9,);
tt=@zeros(nx,);
d2y_dx2_12th_order_accuracy=@zeros(nx,);

@parallel Dx_inn_12th_order_accuracy(y,dtt,dx);
@parallel (1:1) u_1_minus_12th_order_accuracy(dtt,tt);
@parallel Dx_inn_12th_order_accuracy(tt,dtt,dx);
@parallel (1:1) u_1_plus_12th_order_accuracy(dtt,d2y_dx2_12th_order_accuracy);
## plot derivatives
IND=30:nx-30;

plot(x[IND],d2y_dx2[IND],label="analytical",color="red",title="d^2y/dx^2",lw=3);
plot!(x[IND],d2y_dx2_2nd_order_accuracy[IND],label="2nd-order",color="blue",lw=3)
plot!(x[IND],d2y_dx2_8th_order_accuracy[IND],label="8th-order",color="green",lw=3)
plot!(x[IND],d2y_dx2_12th_order_accuracy[IND],label="12th-order",color="yellow",lw=3)
## plot error
plot(x[IND],d2y_dx2_2nd_order_accuracy[IND]-d2y_dx2[IND],label="2nd-order",color="blue",lw=3);
plot!(x[IND],d2y_dx2_8th_order_accuracy[IND]-d2y_dx2[IND],label="8th-order",color="green",lw=3);
plot!(x[IND],d2y_dx2_12th_order_accuracy[IND]-d2y_dx2[IND],label="12th-order",color="yellow",lw=3)
## norm error
e1=norm(d2y_dx2_2nd_order_accuracy[IND]-d2y_dx2[IND],2)
e2=norm(d2y_dx2_8th_order_accuracy[IND]-d2y_dx2[IND],2)
e3=norm(d2y_dx2_12th_order_accuracy[IND]-d2y_dx2[IND],2)