text:
const ShortTuple = Union{
    Tuple{},
    Tuple{A} where {A},
    Tuple{A, B} where {A, B},
    Tuple{A, B, C} where {A, B, C},
    Tuple{A, B, C, D} where {A, B, C, D},
    Tuple{A, B, C, D, E} where {A, B, C, D, E},
    Tuple{A, B, C, D, E, F} where {A, B, C, D, E, F},
    Tuple{A, B, C, D, E, F, G} where {A, B, C, D, E, F, G},
    Tuple{A, B, C, D, E, F, G, H} where {A, B, C, D, E, F, G, H},
    Tuple{A, B, C, D, E, F, G, H, I} where {A, B, C, D, E, F, G, H, I},
    Tuple{A, B, C, D, E, F, G, H, I, J} where {A, B, C, D, E, F, G, H, I, J},
    Tuple{A, B, C, D, E, F, G, H, I, J, K} where {A, B, C, D, E, F, G, H, I, J, K},
    Tuple{A, B, C, D, E, F, G, H, I, J, K, L} where {A, B, C, D, E, F, G, H, I, J, K, L},
    Tuple{A, B, C, D, E, F, G, H, I, J, K, L, M} where {A, B, C, D, E, F, G, H, I, J, K, L, M},
    Tuple{A, B, C, D, E, F, G, H, I, J, K, L, M, N} where {A, B, C, D, E, F, G, H, I, J, K, L, M, N},
    Tuple{A, B, C, D, E, F, G, H, I, J, K, L, M, N, O} where {A, B, C, D, E, F, G, H, I, J, K, L, M, N, O},
    Tuple{A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P} where {A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P}
}

export fill_tuple
"""
```jldoctest
julia> using Keys

julia> fill_tuple((1, 'a', 1.0), "a")
("a", "a", "a")
```
"""
fill_tuple(t::ShortTuple, value) = ntuple(
    let value = value
        x -> value
    end,
    length(t)
)

reduce_unrolled(reducer, ::Tuple{}) = Base._empty_reduce_error()
reduce_unrolled(reducer, args::ShortTuple) = reduce_unrolled(
    reducer,
    args[1],
    tail(args)
)
reduce_unrolled(reducer, default, args::Tuple{}) = default

export reduce_unrolled
"""
```jldoctest
julia> using Keys

julia> reduce_unrolled(&, (true, false, true))
false
```
"""
reduce_unrolled(reducer, default, args::ShortTuple) =
    reducer(default, reduce_unrolled(reducer, args[1], tail(args)))

export getindex_unrolled
getindex_unrolled(into::Tuple{}, switch::Tuple{}) = ()
getindex_unrolled(into::Tuple{}, switch::ShortTuple) = ()
getindex_unrolled(into::ShortTuple, switch::Tuple{}) = ()
"""
```jldoctest
julia> using Keys

julia> getindex_unrolled((1, "a", 1.0), (true, false, true))
(1, 1.0)
```
"""
function getindex_unrolled(into::ShortTuple, switch::ShortTuple)
    next = getindex_unrolled(tail(into), tail(switch))
    if Bool(switch[1])
        (into[1], next...)
    else
        next
    end
end

export setindex_unrolled
"""
```jldoctest
julia> using Keys

julia> setindex_unrolled(
            (1, "a", 1.0),
            ('a', 1//1),
            (True(), False(), True())
        )
('a', "a", 1//1)
```
"""
setindex_unrolled(::Tuple{}, ::Tuple{}, ::Tuple{}) = ()
setindex_unrolled(::Tuple{}, ::Tuple{}, switch::ShortTuple) = ()
setindex_unrolled(::Tuple{}, new::ShortTuple, ::Tuple{}) = ()
setindex_unrolled(::Tuple{}, new::ShortTuple, switch::ShortTuple) = ()
setindex_unrolled(old::ShortTuple, ::Tuple{}, ::Tuple{}) = old
setindex_unrolled(old::ShortTuple, ::Tuple{}, switch::ShortTuple) = old
setindex_unrolled(old::ShortTuple, new::ShortTuple, ::Tuple{}) = ()
function setindex_unrolled(old::ShortTuple, new::ShortTuple, switch::ShortTuple)
    first_tuple, tail_tuple =
        if Bool(switch[1])
            (new[1], tail(new))
        else
            (old[1], new)
        end
    first_tuple, setindex_unrolled(tail(old), tail_tuple, tail(switch))...
end

export find_unrolled
find_unrolled(t::ShortTuple) = find_unrolled(t, 1)
find_unrolled(t::Tuple{}, n::Integer) = ()

"""
```jldoctest
julia> using Keys

julia> find_unrolled((true, false, true))
(1, 3)
```
"""
function find_unrolled(t::ShortTuple, n::Integer)
    next = find_unrolled(tail(t), n + 1)
    if Bool(t[1])
        (n, next...)
    else
        next
    end
end

export flatten_unrolled
"""
```jldoctest
julia> using Keys

julia> flatten_unrolled(((1, 2.0), ("c", 4//4)))
(1, 2.0, "c", 1//1)
```
"""
flatten_unrolled(x::ShortTuple) = x[1]..., flatten_unrolled(tail(x))...
flatten_unrolled(::Tuple{}) = ()

export product_unrolled
"""
```jldoctest
julia> using Keys

julia> product_unrolled((1, 2.0), ("c", 4//4))
((1, "c"), (2.0, "c"), (1, 1//1), (2.0, 1//1))
```
"""
product_unrolled(x::ShortTuple, y::ShortTuple) = flatten_unrolled(map(
    let x = x
        y1 ->
            map(
                let y1 = y1
                    x1 -> (x1, y1)
                end,
                x
            )
    end,
    y
))

export filter_unrolled
"""
```jldoctest
julia> using Keys

julia> filter_unrolled(identity, (True(), False()))
(True(),)
```
"""
filter_unrolled(f, x::ShortTuple) = getindex_unrolled(x, map(f, x))


text:
Base.reverse(r::StepSRange) = srange(last(r), step=-step(r), stop=first(r))
Base.reverse(r::StepMRange) = mrange(last(r), step=-step(r), stop=first(r))
Base.reverse(r::AbstractLinRange) = similar_type(r)(last(r), first(r), length(r))

function Base.reverse(r::AbstractStepRangeLen)
    # If `r` is empty, `length(r) - r.offset + 1 will be nonpositive hence
    # invalid. As `reverse(r)` is also empty, any offset would work so we keep
    # `r.offset`
    offset = isempty(r) ? _offset(r) : length(r) - _offset(r) + 1
    return similar_type(r)(_ref(r), -step_hp(r), length(r), offset)
end

function Base.reverse!(r::StepMRange)
    setfield!(r, :step, -step(r))
    f = first(r)
    l = last(r)
    setfield!(r, :stop, f)
    setfield!(r, :start, l)
    return r
end

function Base.reverse!(r::StepMRangeLen)
    # If `r` is empty, `length(r) - r.offset + 1 will be nonpositive hence
    # invalid. As `reverse(r)` is also empty, any offset would work so we keep
    # `r.offset`
    setfield!(r, :offset, isempty(r) ? _offset(r) : length(r) - _offset(r) + 1)
    setfield!(r, :step, -step_hp(r))
    return r
end

function Base.reverse!(r::LinMRange)
    f = first(r)
    l = last(r)
    setfield!(r, :stop, f)
    setfield!(r, :start, l)
    return r
end

function Base.iterate(r::Union{AbstractLinRange,AbstractStepRangeLen}, i::Int=1)
    Base.@_inline_meta
    length(r) < i && return nothing
    unsafe_getindex(r, i), i + 1
end

Base.isempty(r::Union{AbstractLinRange,AbstractStepRangeLen}) = length(r) == 0

###
### ==(r1, r2)
###
Base.:(==)(r::OneToRange, s::OneToRange) = last(r) == last(s)

function Base.:(==)(r::StepMRangeLen{T,R,S}, s::StepMRangeLen{T,R,S}) where {T,R,S}
    return (first(r) == first(s)) & (length(r) == length(s)) & (last(r) == last(s))
end

function Base.:(==)(r::StepSRangeLen{T,R,S}, s::StepSRangeLen{T,R,S}) where {T,R,S}
    return (first(r) == first(s)) & (length(r) == length(s)) & (last(r) == last(s))
end

function Base.:(==)(r::AbstractLinRange{T}, s::AbstractLinRange{T}) where {T}
    return (first(r) == first(s)) & (length(r) == length(s)) & (last(r) == last(s))
end

###
### +(r1, r2)
###
Base.:(+)(r1::Union{AbstractStepRangeLen,AbstractLinRange}, r2::AbstractRange) = +(promote(r1, r2)...)
Base.:(+)(r1::AbstractRange, r2::Union{AbstractStepRangeLen,AbstractLinRange}) =  +(promote(r1, r2)...)
Base.:(+)(r1::Union{AbstractStepRangeLen,AbstractLinRange}, r2::Union{AbstractStepRangeLen,AbstractLinRange}) = +(promote(r1, r2)...)

function Base.:(+)(r1::StepMRangeLen{T,S}, r2::StepMRangeLen{T,S}) where {T,S}
    len = length(r1)
    (len == length(r2) ||
        throw(DimensionMismatch("argument dimensions must match")))
    return StepMRangeLen(first(r1)+first(r2), step(r1)+step(r2), len)
end

function Base.:(+)(r1::StepMRangeLen{T,TwicePrecision{T}}, r2::StepMRangeLen{T,TwicePrecision{T}}) where {T}
    len = length(r1)
    (len == length(r2) || throw(DimensionMismatch("argument dimensions must match")))
    if r1.offset == r2.offset
        imid = r1.offset
        ref = r1.ref + r2.ref
    else
        imid = round(Int, (r1.offset + r2.offset)/2)
        ref1mid = _getindex_hiprec(r1, imid)
        ref2mid = _getindex_hiprec(r2, imid)
        ref = ref1mid + ref2mid
    end
    step = twiceprecision(r1.step + r2.step, nbitslen(T, len, imid))
    return StepMRangeLen{T,typeof(ref),typeof(step)}(ref, step, len, imid)
end

function Base.:(+)(r1::StepSRangeLen{T,TwicePrecision{T}}, r2::StepSRangeLen{T,TwicePrecision{T}}) where {T}
    len = length(r1)
    (len == length(r2) ||
        throw(DimensionMismatch("argument dimensions must match")))
    if _offset(r1) == _offset(r2)
        imid = _offset(r1)
        ref = _ref(r1) + _ref(r2)
    else
        imid = round(Int, (_offset(r1)+_offset(r2))/2)
        ref1mid = _getindex_hiprec(r1, imid)
        ref2mid = _getindex_hiprec(r2, imid)
        ref = ref1mid + ref2mid
    end
    step = twiceprecision(r1.step + r2.step, nbitslen(T, len, imid))
    return StepSRangeLen{T,typeof(ref),typeof(step)}(ref, step, len, imid)
end

#_add(r1::StepSRangeLen{T,R,S}, r2::Union{OneToSRange,UnitSRange,StepSRange,LinSRange}) where {T,R,S} = +(r1, StepSRangeLen{T,R,S}(r2))
#_add(r2::Union{OneToSRange,UnitSRange,StepSRange,LinSRange}, r1::StepSRangeLen{T,R,S}) where {T,R,S} = +(r1, StepSRangeLen{T,R,S}(r2))

function Base.sum(r::AbstractStepRangeLen)
    l = length(r)
    # Compute the contribution of step over all indices.
    # Indexes on opposite side of r.offset contribute with opposite sign,
    #    r.step * (sum(1:np) - sum(1:nn))
    np, nn = l - r.offset, r.offset - 1  # positive, negative
    # To prevent overflow in sum(1:n), multiply its factors by the step
    sp, sn = sumpair(np), sumpair(nn)
    tp = Base._tp_prod(r.step, sp[1], sp[2])
    tn = Base._tp_prod(r.step, sn[1], sn[2])
    s_hi, s_lo = add12(tp.hi, -tn.hi)
    s_lo += tp.lo - tn.lo
    # Add in contributions of ref
    ref = r.ref * l
    sm_hi, sm_lo = add12(s_hi, ref.hi)
    return add12(sm_hi, sm_lo + ref.lo)[1]
end

###
### -(r1, r2)
###
Base.:(-)(r1::Union{AbstractStepRangeLen,AbstractLinRange}, r2::AbstractRange) = -(promote(r1, r2)...)
Base.:(-)(r1::AbstractRange, r2::Union{AbstractStepRangeLen,AbstractLinRange}) =  -(promote(r1, r2)...)
Base.:(-)(r1::Union{AbstractStepRangeLen,AbstractLinRange}, r2::Union{AbstractStepRangeLen,AbstractLinRange}) = -(promote(r1, r2)...)

function Base.:(-)(r::AbstractStepRangeLen)
    return similar_type(r)(-_ref(r), -step(r), length(r), _offset(r))
end
Base.:(-)(r1::StepMRangeLen{T,R,S}, r2::StepMRangeLen{T,R,S}) where {T,R,S} = +(r1, -r2)
Base.:(-)(r1::StepSRangeLen{T,R,S}, r2::StepSRangeLen{T,R,S}) where {T,R,S} = +(r1, -r2)

###
### *(r1, r2)
###
Base.:(*)(r::AbstractStepRangeLen{T,TwicePrecision{T}}, x::Real) where {T<:Real} = x*r
function Base.:(*)(x::Real, r::StepMRangeLen{T,TwicePrecision{T}}) where {T<:Real}
    return StepMRangeLen(x * _ref(r), Base.twiceprecision(x * step(r), nbitslen(r)), length(r), _offset(r))
end
function Base.:(*)(x::Real, r::StepSRangeLen{T,TwicePrecision{T}}) where {T<:Real}
    return StepSRangeLen(x * _ref(r), Base.twiceprecision(x * step(r), nbitslen(r)), length(r), _offset(r))
end

###
### /(r1, r2)
###
function Base.:(/)(r::StepMRangeLen{T,TwicePrecision{T}}, x::Real) where {T<:Real}
    return StepMRangeLen(_ref(r)/x, Base.twiceprecision(step(r)/x, Base.nbitslen(r)), length(r), _offset(r))
end
function Base.:(/)(r::StepSRangeLen{T,TwicePrecision{T}}, x::Real) where {T<:Real}
    return StepSRangeLen(_ref(r)/x, Base.twiceprecision(step(r)/x, Base.nbitslen(r)), length(r), _offset(r))
end

for (frange,R) in ((mrange, :StepMRange), (srange, :StepSRange))
    for f in (:+, :-)
        @eval begin
            Base.$(f)(r1::$R, r2::OrdinalRange) = $(f)(promote(r1, r2)...)
            Base.$(f)(r1::OrdinalRange, r2::$R) =  $(f)(promote(r1, r2)...)
            function Base.$(f)(r1::$R, r2::$R)
                r1l = length(r1)
                (r1l == length(r2) ||
                 throw(DimensionMismatch("argument dimensions must match: length of r1 is $r1l, length of r2 is $(length(r2))")))
                $(frange)($f(first(r1), first(r2)), step=$f(step(r1), step(r2)), length=r1l)
            end

        end
    end
end

#= TODO
float(r::StepRange) = float(r.start):float(r.step):float(last(r))
float(r::UnitRange) = float(r.start):float(last(r))
float(r::StepRangeLen{T}) where {T} =
    StepRangeLen{typeof(float(T(r.ref)))}(float(r.ref), float(r.step), length(r), r.offset)
function float(r::LinRange)
    LinRange(float(r.start), float(r.stop), length(r))
end
=#

ext: jl


text:
const NEW_APP_PATH = joinpath("files", "new_app")

const GENIE_CONTROLLER_FILE_POSTFIX             = "Controller.jl"

const ROUTES_FILE_NAME                          = "routes.jl"
const ASSETS_FINGERPRINT_INITIALIZER_FILE_NAME  = "assets_fingerprint.jl"
const SEARCHLIGHT_INITIALIZER_FILE_NAME         = "searchlight.jl"
const SECRETS_FILE_NAME                         = "secrets.jl"
const BOOTSTRAP_FILE_NAME                       = "bootstrap.jl"
const ENV_FILE_NAME                             = "env.jl"
const GENIE_FILE_NAME                           = "genie.jl"
const GLOBAL_ENV_FILE_NAME                      = "global.jl"
const TESTS_FILE_NAME                           = "runtests.jl"

const PARAMS_REQUEST_KEY    = :REQUEST
const PARAMS_RESPONSE_KEY   = :RESPONSE
const PARAMS_SESSION_KEY    = :SESSION
const PARAMS_FLASH_KEY      = :FLASH
const PARAMS_POST_KEY       = :POST
const PARAMS_GET_KEY        = :GET
const PARAMS_WS_CLIENT      = :WS_CLIENT
const PARAMS_JSON_PAYLOAD   = :JSON_PAYLOAD
const PARAMS_RAW_PAYLOAD    = :RAW_PAYLOAD
const PARAMS_FILES          = :FILES
const PARAMS_ROUTE_KEY      = :ROUTE
const PARAMS_CHANNELS_KEY   = :CHANNEL

SECRET_TOKEN = ""
ASSET_FINGERPRINT = ""


text:
#----------------------------------------------------------------------------
# load data from Imperial College
struct Data
    turing_data
    cases
    deaths
    hospit
    num_observations
    num_total
    dates
    pop
    predictors
    seroprev
end

dan2eng = Dict(
    "Dato"          => "date",
    "Total"         => "total",
    "Hovedstaden"   => "capital",
    "Sj\xe6lland"   => "zealand",
    "Sjælland"      => "zealand",
    "Syddanmark"    => "south",
    "Midtjylland"   => "central",
    "Nordjylland"   => "north",
    "Ukendt Region" => "unknown"
)

function missing2zero(xs)
    i = findfirst(!ismissing, xs)
    xs[1:i] .= 0
    convert(Vector{Int64}, xs)
end
# weekly_average(xs) = [mean(xs[i:(i+6)]) for i in 1:div(length(xs),7)]
# weekly_average(deaths)

function padlast(xs, l)
    n  = length(xs)
    i  = findlast(x -> typeof(x) <: Real, xs) #avoid catching a missing value
    ys = Vector{Union{Float64, Missing}}(undef, l)
    ys[1:n] = xs
    ys[n+1:end] .= xs[i]
    ys
end

function padmissing(xs, l)
    n  = length(xs)
    ys = Vector{Union{eltype(xs), Missing}}(missing, l)
    ys[1:n] = xs
    ys
end

function cum2counts(x)
    x = x |> missing2zero |> diff
    pushfirst!(x, 0)
    x[x.<0] .= 0 # ignore data inconsistencies
    x
end

# plotnpi(x) = plot(dk.date, x; ticks=:native)
# dk.cancel_events|>plotnpi
# dk.internal_movement_restrictions|>plotnpi
# dk.gatherings_restrictions|>plotnpi
# dk.stay_home_restrictions|>plotnpi
# dk.transport_closing|>plotnpi
# dk.workplace_closing|>plotnpi
# dk.school_closing|>plotnpi
# # dk.stringency_index|>plotnpi

const dformat1 = DateFormat("d u Y")
const dformat2 = DateFormat("d-u-Y")
function parsetime(ds)
    ds = replace(ds, "Mrz"=>"Mar")
    ds = replace(ds, "Dez"=>"Dec")
    ds = replace(ds, "."=>"")
    d = try
        Date(ds, dformat1)
    catch
        try
            Date(ds, dformat2)
        catch
            @error "Not a valid time format: $ds"
        end
    end
    d + Year(2000)
end

function znorm!(xs)
    xs .-= mean(skipmissing(xs))
    xs ./= std(skipmissing(xs))
    nothing
end

function minmaxnorm!(xs)
    xs .-= minimum(skipmissing(xs))
    xs ./= maximum(skipmissing(xs))
    nothing
end

function surveynorm!(xs)
    i = findfirst(x->!ismissing(x), xs)
    xs ./= abs(xs[i])
    xs .-= xs[i]
    nothing
end

function preproc(xs::AbstractVector, num_obs, Ntot; normalize=nothing)
    mask = 1:num_obs
    xs = padlast(xs[mask], Ntot)
    !isnothing(normalize) && normalize(xs)
    xs
end

function select_predictors(dk, predictors, num_obs, Ntot, surveyvars)
    predictors = String.(predictors)

    #preprocess mobility data
    dk[!,:nonresidential_mobility] = preproc(
        (
            dk.retail_and_recreation_percent_change_from_baseline +
            dk.grocery_and_pharmacy_percent_change_from_baseline +
            dk.workplaces_percent_change_from_baseline
        ) / 3 / 100 , num_obs, Ntot; normalize=nothing)
    # begin
    #     Mavg = (
    #     preproc(dk.retail_and_recreation_percent_change_from_baseline, num_obs, Ntot) +
    #     preproc(dk.grocery_and_pharmacy_percent_change_from_baseline, num_obs, Ntot) +
    #     preproc(dk.workplaces_percent_change_from_baseline, num_obs, Ntot)
    #     ) ./ -3
    # end
    dk[!,:residential_mobility] = preproc(dk.residential_percent_change_from_baseline, num_obs, Ntot; normalize=nothing) ./ 100

    #preprocess survey data
    for k in surveyvars
        dk[!,k] = preproc(dk[:,k], num_obs, Ntot; normalize=surveynorm!)
    end

    dk[!, :lockdown] .= 0.
    t = findfirst(==(Date("2020-03-14")), dk.date) # lockdown date in DK: 11. (partial), 17. (full)
    dk[t:end, :lockdown] .+= 1
    # select predictors
    miss = .!( in.(predictors, Ref(names(dk)) ) )
    any(miss) && (@error "predictor $(predictors[miss]) not found. Select either nonresidential_mobility, residential_mobility or from survey: $(surveyvars)")
    Array( dk[!,predictors] )
end

rolling(xs) = begin
    n  = length(xs)
    ys = Vector{Union{Missing, Float64}}(missing, n)
    ys[1:n-6] = [maximum(xs[i:i+6]) for i in 1:n-6]
    ys
end

function gamma(μ, cv)
    α = 1/cv^2
    θ = μ/α
    Gamma(α, θ)
end

function inf2hosp(nbins=50, ihr = 1)
    N = 10_000_000
    # SSI: ihr ∈ 0.05%-0.5%
    # derived from Verity et al. See parameters.jl for details
    # incubatine time ~ gamma(5.1, .86) see Flaxman et al
    # onset to hospitalization see Zhang et al. Lancet (2020)
    x = rand(gamma(5.1, .86), N) .+ rand(gamma(2.6, .92), N)
    h = fit(Histogram, x, vcat([0], 1.5:1:(nbins + .5)))
    h.weights / N * ihr
end

function inf2death(nbins = 75, ifr = 1)
    N = 10_000_000
    # ifr = 89 / 100_000 for age < 70 see Erikstrup et al. in Clin. Inf. Dis.
    # report34: onset to death = 19.8, Cv = .85
    r = vcat([0], 1.5:1:(nbins + .5))
    x = rand(gamma(5.1, .86), N) .+ rand(gamma(19.8, .85), N) #rand(gamma(17.8, .45), N)
    h = fit(Histogram, x, r)
    h.weights / N * ifr
end

function seropositive(nbins)
    N = 10_000_000
    # For details see Imperial Report34
    rbins = 1.5:1:(nbins + .5)
    inf2symp = rand(gamma(5.1, .86), N)
    inf2conv = inf2symp .+ rand(Exponential(13.3), N)
    inf2rev = inf2conv .+ rand(Weibull(3.67, 132.88), N)
    seroconverted = map( ecdf(symp2conv), rbins)
    notseroreverted = 1 .- map( ecdf(symp2rev), rbins)
    serconverted .* notseroreverted
end

function inf2case(nbins = 30)
    # ifr = 89 / 100_000 for age < 70 see Erikstrup et al. in Clin. Inf. Dis.
    r = vcat([0], 1.5:1:(nbins + .5))
    diff(cdf.(gamma(5.1, .86), r))
end

function serialinterval(nbins = 30)
    r = vcat([0], 1.5:1:(nbins + .5))
    diff(cdf.(gamma(6.5, .62), r))
end

# function inf2serocon(nbins)
#     N = 10_000_000
#     # onset to seroconversion time from Imperial Report 34
#     x = rand(gamma(5.1, .86), N) .+ rand(Exponential(18.3), N)
#     h = fit(Histogram, x, vcat([0], 1.5:1:150.5))
#     h.weights / N * ihr
# end

function aggregateregions!(df, mun2reg)
  regions = values(mun2reg) |> unique
  for r in regions
    df[!, r] .= 0
  end
  for m in keys(mun2reg)
    df[!, mun2reg[m]] += df[!,m]
  end
end

# function getcases()
#     municipalities = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/municipalities.csv"); delim="\t", )|>DataFrame
#     mun2reg = Dict(municipalities.municipality[i] => municipalities.region[i]
#                     for i in 1:size(municipalities,1))
#     cases    = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Municipality_cases_time_series.csv") )|>DataFrame
#     cases = select(cases, municipalities.municipality, :date_sample => :date)
#     cases[!, :total] = sum(select(cases, municipalities.municipality)|>eachcol)
#     aggregateregions!(cases, mun2reg)
#     cases
# end

function insert_missing_dates!(df)
    ts = df.date[1] : Day(1) : df.date[end]
    missing_date = .! in.(ts, Ref(df.date))
    if any(missing_date)
        for t in ts[missing_date]
            insertrow!(df, t)
        end
    end
end

function replace_missing!(df, v=0)
    @assert !any( ismissing.( df.date ) )
    for c in eachcol(df)
        replace!(c, missing=>v)
    end
    disallowmissing!(df)
end

function process_ssi(df)
    rename!(df, 1=>:region, 2=>:date, 3=>:counts)
    df.region = replace(df.region, dan2eng...)
    dropmissing!(df, :region)
    df = unstack(df, :date, :region, :counts)
    insert_missing_dates!(df)
    replace_missing!(df)
    DataFrames.transform(df, AsTable(Not(:date)) => sum => :country)
end

# function getdeaths()
#     deaths   = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Deaths_over_time.csv"))|>DataFrame
#     deaths = deaths[1:end-1,:]
#     rename!(deaths, "Dato"=>"date", "Antal_døde"=>"total")
#     deaths.date = Date.(deaths.date, Ref(dateformat"y-m-d"))
#     deaths
# end
#
# function gethospit()
#     hosp     = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Newly_admitted_over_time.csv"))|>DataFrame
#     dan2eng = Dict(
#         "Dato"          => "date",
#         "Total"         => "total",
#         "Hovedstaden"   => "capital",
#         "Sjælland"      => "zealand",
#         "Syddanmark"    => "south",
#         "Midtjylland"   => "central",
#         "Nordjylland"   => "north",
#         "Ukendt Region" => "unknown"
#     )
#     rename!(hosp, dan2eng)
#     hosp
# end

function startdate(deaths)
    i = findfirst(>=(10), skipmissing(deaths.country))
    deaths.date[i] - Day(30)
end

enddate(df1, df2) = min(df1.date[end], df2.date[end])
enddate(df, args...) = min(df.date[end], enddate(args...))

function padzeros(xs::AbstractVector, padfirst=0, total=length(xs)+padfirst)
    n  = length(xs)
    ys = zeros(eltype(xs), total)
    ys[padfirst+1:n+padfirst] = xs
    ys
end

function insertrow!(df, d::Date, v=0)
    i = findfirst(==(d-Day(1)), df.date)
    v = [eltype(c) == Date ? d : v for c in eachcol(df)]
    insert!.(eachcol(df), i+1, v)
    nothing
end

function padzeros(df::DataFrame, s::Date)
    npad = (df.date[1] - s).value
    dfpad = DataFrame( :date => s:Day(1):df.date[end] )
    for n in names(df)
        n == "date" && continue
        dfpad[!,n] = padzeros(df[!,n], npad)
    end
    dfpad
end

function limit(df::DataFrame, s::Date, e::Date, pad=true)
    pad && df.date[1] > s && ( df = padzeros(df, s) )
    is = findfirst(==(s), df.date)
    ie = findlast(==(e), df.date)
    df[is:ie, :]
end

function parseweek(s)
    _,w = split(s)
    w = parse(Int, w)
    y = w > 40 ? Date("2020") : Date("2021")
    y + Week(w) - Day(3) #center at day 4
end

# ≈ 15days delay from infection to antibody production
# see Ferguson Report and Zhao 2020: "Antibody responses to SARS-CoV-2 in patients of novel coronavirus disease 2019"
function readseroprev(dates, delay=15)
    df0 = CSV.File(projectdir("data/seroprevalence/", "seroprev_early.csv"), delim="\t")|>DataFrame
    # df1 = CSV.File(projectdir("data/seroprevalence/", "seroprev.csv"), delim="\t")|>DataFrame
    # df1[!,:date] = df1.week .|> parseweek
    # select!(df1, Not(:week))
    # df = vcat(df0, df1)
    df = df0
    df = df[ dates[1] .<= df.date .<= dates[end], :]
    df.date .-= Day(delay)
    parsepercent!(df)
    df.CI = df.CI .|> parseCI
    # df[!,:std_u] = [(CIs[i][2] - df.country[i]) / 1.96 for i in 1:length(CIs)]
    # estimates of std are not symmetric with respect to lower / upper CI bound
    # we choose the larger estimate, i.e. the lower bound
    turing_data = (
        std  = [(df.country[i] - df.CI[i][1]) / 1.96 for i in 1:length(df.CI)],
        mean = df.country,
        idx  = [findfirst(==(d), dates) for d in df.date]
    )
    df, turing_data
end

function parsepercent(x::String)
    x = strip(x, '%')
    x = replace(x, "," => ".")
    parse(Float64, x) / 100
end

function parseCI(x)
    x = strip(x, '%')
    x = replace(x, "," => ".")
    lq, uq = split(x, '-')
    parse(Float64, lq) / 100, parse(Float64, uq) / 100
end

function parsepercent!(df)
    regions = ["capital", "zealand", "south", "central", "north"]
    df.country = parsepercent.(df.country)
    for r in regions
        df[!,r] = parsepercent.(df[:,r])
    end
end

function findindices(dates, ds)
    is = Int[]
    for d in ds
        i = findfirst(==(d), dates)
        push!(is, i)
    end
    is
end

function readsurvey(fname)
    s = CSV.File(fname)|>DataFrame
    s.date = s.date .|> parsetime
    dropmissing!(s)
    for n in names(s)
        c = s[:,n]
        if eltype(c) == String
            c = replace.(c, ","=>".")
            s[!,n] = parse.(Ref(Float64), c)
        end
    end
    s
end

function readsurveys_raw()
    s1 = normpath( homedir(), "data/covidsurvey/contacts.csv" )  |> readsurvey #projectdir("data", "survey", "timeseries_colleagues_strangers.csv") |> readsurvey
    s1 = select(s1, Not([:Rt, :uncertainty_lower, :uncertainty_upper]))
    s2 = normpath( homedir(), "data/covidsurvey/predictors.csv" )  |> readsurvey #projectdir("data", "survey", "predictors.csv") |> readsurvey
    s = innerjoin(s1, s2, on=:date)
    s[3:end,:] # skip first for data quality
end
function readsurveys()
    fname = normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" )
    df = CSV.File(fname)|>DataFrame
    return df
end

country(df) = df.country

function get_iar_idx(num_tot, num_obs, start_idx, iar_step)
    iar_idx = zeros(Int64, num_tot)
    iar_idx[1:start_idx-1] .= 1
    idx = 2
    for (j,i) in enumerate(start_idx:num_obs)
         iar_idx[i] = idx
         j % iar_step == 0 && (idx += 1)
    end
    if num_obs < num_tot
        iar_idx[num_obs+1:num_tot] .= iar_idx[num_obs]
    end
    iar_idx
end

function load_data(
    observations_end=nothing,
    predictors=nothing,
    cases_start=nothing,
    addhospital=false,
    addseroprev=false,
    addtests=false;
    update=false
)
    iar_step = 7

    cases = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/08_bekraeftede_tilfaelde_pr_dag_pr_regions.csv" ) ) |>
        DataFrame |> process_ssi

    deaths = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/07_antal_doede_pr_dag_pr_region.csv" ) ) |>
        DataFrame |> process_ssi

    hospit = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/06_nye_indlaeggelser_pr_region_pr_dag.csv" ) ) |>
        DataFrame |> process_ssi

    s      = startdate(deaths)
    e      = enddate(deaths, hospit, cases)
    dates  = collect( s : Day(1) : e )
    deaths = limit(deaths, s, e) |> country
    hospit = limit(hospit, s, e) |> country
    cases  = limit(cases, s, e)  |> country
    num_tot= length(dates)

    dk = begin
        update && run(`Rscript $(datadir())/COVID19.R`)
        dk = load(datadir("COVID19", "DK.rds"))
        limit(dk, s, e, false)
    end

    survey = readsurveys()
    surveyvars = filter(x->x!=="date", names(survey))
    dk = leftjoin(dk, survey, on = :date)

    num_obs = if isnothing(observations_end)
        num_tot
    else
        findfirst(==(Date.(observations_end)), dates)
    end

    lockdown = findfirst(==(Date("2020-03-18")), dates)

    num_case_obs, cases_start_idx = if isnothing(cases_start)
        0, nothing
    else
        cases_start_idx = findfirst(==(Date.(cases_start)), dates)
        num_obs - cases_start_idx + 1, cases_start_idx
    end

    iar_idx = get_iar_idx(num_tot, num_obs, cases_start_idx, iar_step)
    num_iar_steps = length(unique(iar_idx))

    turing_data = OrderedDict(
        :num_impute                  => 6,
        :num_total_days              => num_tot,
        :cases                       => cases[1:num_obs], #padmissing(cases[mask], num_tot),
        :deaths                      => deaths[1:num_obs], #padmissing(deaths[mask], num_tot),
        :π                           => padzeros(inf2death(90), 0, num_tot),#[1:100], #length = 100
        :π2                          => padzeros(inf2case(30), 0, num_tot),#[1:30], #WARNING US-NY DATA
        :epidemic_start              => 31,
        :population                  => dk.population[1],
        :serial_intervals            => padzeros(serialinterval(30), 0, num_tot), #data_usa.turing_data.serial_intervals[1:50],
        :num_iar_steps               => num_iar_steps,
        :iar_idx                     => iar_idx,
        :lockdown                    => lockdown,
        :num_case_obs                => num_case_obs,
        )

    if !isnothing(predictors)
        preds = select_predictors(dk, predictors, num_obs, num_tot, surveyvars)
        covariates_start = findfirst(x->!ismissing(x), preds[:,end])
        preds = convert(Array{Float64,2}, preds[covariates_start:end,:] )
        turing_data[:covariates] = preds
        turing_data[:covariates_start] = covariates_start
    end

    if addhospital
        turing_data[:πh] = padzeros(inf2hosp(50, 1), 0,  num_tot) #length 50
        turing_data[:hospit] = hospit[1:num_obs]
    end

    seroprev, turing_seroprev = readseroprev(dates[1:num_obs])
    if addseroprev
        #turing_data[:πs] = padzeros(inf2serocon(50, 1), 0,  num_tot)
        turing_data[:seroprev_mean] = turing_seroprev.mean * dk.population[1]
        turing_data[:seroprev_std]  = turing_seroprev.std * dk.population[1]
        turing_data[:seroprev_idx]  = turing_seroprev.idx
    end

    turing_data = (;turing_data...)
    Data(
        turing_data,
        cases,
        deaths,
        hospit,
        num_obs, #num_observations
        num_tot,
        dates,
        dk.population,
        predictors,
        seroprev,
    )
end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
# AuxDataDict
# ===========
#
# Auxiliary data dictionary for BAM.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

# This type is not designed for very large dictionaries: time complexities in lookup
# and update operations are O(N). The memory layout is depicted as follows:
#
#   data: |tag (2 bytes)|value type (1 byte)|value (variable)|tag|...
#
# Note that tag = 0xffff indicates the value must be skipped because it is outdated.
immutable AuxDataDict <: Associative{String,Any}
    data::Vector{UInt8}
end

function AuxDataDict()
    return AuxDataDict(UInt8[])
end

function AuxDataDict{K<:AbstractString}(pairs::Pair{K}...)
    dict = AuxDataDict()
    for (tag, val) in pairs
        dict[tag] = val
    end
    return dict
end

function Base.getindex(dict::AuxDataDict, tag::AbstractString)
    checkkeytag(tag)
    return getvalue(dict.data, 1, UInt8(tag[1]), UInt8(tag[2]))
end

function Base.setindex!(dict::AuxDataDict, val, tag::AbstractString)
    checkkeytag(tag)
    setvalue!(dict.data, 1, val, UInt8(tag[1]), UInt8(tag[2]))
    return dict
end

function Base.delete!(dict::AuxDataDict, tag::AbstractString)
    checkkeytag(tag)
    deletevalue!(dict.data, 1, UInt8(tag[1]), UInt8(tag[2]))
    return dict
end

function Base.length(dict::AuxDataDict)
    data = dict.data
    p = 1
    len = 0
    while p ≤ length(data)
        if !(data[p] == data[p+1] == 0xff)
            len += 1
        end
        p = next_tag_position(data, p)
    end
    return len
end

function Base.start(dict::AuxDataDict)
    return 1
end

function Base.done(dict::AuxDataDict, pos)
    return pos > length(dict.data)
end

function Base.next(dict::AuxDataDict, pos)
    data = dict.data
    @label doit
    t1 = data[pos]
    t2 = data[pos+1]
    pos, typ = loadauxtype(data, pos + 2)
    pos, value = loadauxvalue(data, pos, typ)
    if t1 == t2 == 0xff
        @goto doit
    end
    return Pair{String,Any}(String([t1, t2]), value), pos
end

# Internals
# ---------

function getvalue(data, pos, t1, t2)
    pos = findtag(data, pos, t1, t2)
    if pos == 0
        throw(KeyError(String([t1, t2])))
    end
    pos, T = loadauxtype(data, pos + 2)
    _, val = loadauxvalue(data, pos, T)
    return val
end

function setvalue!(data, pos, val, t1, t2)
    pos = findtag(data, pos, t1, t2)
    if pos > 0
        # cancel tag
        data[pos] = data[pos+1] = 0xff
    end

    # TODO: in-place update if possible
    pos = length(data)
    if isa(val, AbstractString)
        resize!(data, length(data) + 2 + 1 + sizeof(val) + 1)
    else
        resize!(data, length(data) + 2 + 1 + sizeof(val))
    end
    data[pos+1] = t1
    data[pos+2] = t2
    storeauxtype!(data, pos + 3, typeof(val))
    storeauxvalue!(data, pos + 4, val)
    return data
end

function deletevalue!(data, pos, t1, t2)
    pos = findtag(data, pos, t1, t2)
    if pos > 0
        # cancel tag
        data[pos] = data[pos+1] = 0xff
    end
    return data
end

function findtag(data, pos, t1, t2)
    while pos ≤ length(data) && !(data[pos] == t1 && data[pos+1] == t2)
        pos = next_tag_position(data, pos)
    end
    if pos > length(data)
        # not found
        return 0
    else
        return pos
    end
end

function loadauxtype(data, p)
    t = data[p]
    if t == UInt8('B')
        return p + 2, Vector{auxtype[data[p+1]]}
    else
        return p + 1, auxtype[t]
    end
end

function loadauxvalue{T}(data, p, ::Type{T})
    return p + sizeof(T), unsafe_load(Ptr{T}(pointer(data, p)))
end

function loadauxvalue(data, p, ::Type{Char})
    return p + 1, Char(unsafe_load(pointer(data, p)))
end

function loadauxvalue{T}(data, p, ::Type{Vector{T}})
    n = unsafe_load(Ptr{Int32}(pointer(data, p)))
    p += 4
    xs = Array(T, n)
    unsafe_copy!(pointer(xs), Ptr{T}(pointer(data, p)), n)
    return p + n * sizeof(T), xs
end

function loadauxvalue(data, p, ::Type{String})
    dataptr = pointer(data, p)
    endptr = ccall(:memchr, Ptr{Void}, (Ptr{Void}, Cint, Csize_t),
                   dataptr, '\0', length(data) - p + 1)
    q::Int = p + (endptr - dataptr) - 1
    return q + 2, String(data[p:q])
end

function storeauxtype!(data, p, T)
    if T <: AbstractVector
        data[p] = UInt8('B')
    else
        data[p] = UInt8(auxtypechar[T])
    end
    return data
end

function storeauxvalue!(data, p, val)
    if isa(val, AbstractString)
        n = length(val)
        ccall(:memcpy, Ptr{Void}, (Ptr{Void}, Ptr{Void}, Csize_t), pointer(data, p), pointer(val), n)
        data[p + n] = 0x00
    elseif isa(val, AbstractVector)
        n = length(val)
        storeauxtype!(data, p, eltype(val))
        unsafe_store!(Ptr{UInt32}(pointer(data, p + 1)), n)
        ccall(:memcpy, Ptr{Void}, (Ptr{Void}, Ptr{Void}, Csize_t), pointer(data, p + 5), pointer(val), n)
    else
        unsafe_store!(Ptr{eltype(val)}(pointer(data, p)), val)
    end
    return data
end

# Find the starting position of a next tag in `data` after `p`.
# `(data[p], data[p+1])` is supposed to be a current tag.
function next_tag_position(data, p)
    typ = Char(data[p+2])
    p += 3
    if typ == 'A'
        p += 1
    elseif typ == 'c' || typ == 'C'
        p += 1
    elseif typ == 's' || typ == 'S'
        p += 2
    elseif typ == 'i' || typ == 'I'
        p += 4
    elseif typ == 'f'
        p += 4
    elseif typ == 'd'
        p += 8
    elseif typ == 'Z' || typ == 'H'
        while data[p] != 0x00  # NULL-terminalted string
            p += 1
        end
        p += 1
    elseif typ == 'B'
        eltyp = Char(data[p])
        elsize = eltyp == 'c' || eltyp == 'C'                 ? 1 :
                 eltyp == 's' || eltyp == 'S'                 ? 2 :
                 eltyp == 'i' || eltye == 'I' || eltyp == 'f' ? 4 :
                 error("unrecognized auxiliary type: ", eltyp)
        p += 1
        n = unsafe_load(Ptr{Int32}(pointer(data, p)))
        p += elsize * n
    else
        error("unrecognized auxiliary type: ", typ)
    end
    return p
end


text:
# create optimizer type
abstract type Optimizer end

# include optimizers
include("optimizers/adam.jl")


text:
using ITensors,
      Random,
      Test

Random.seed!(12345)

include("2d_classical_ising.jl")

function ctmrg(T::ITensor,
               Clu::ITensor,
               Al::ITensor;
               χmax::Int,nsteps::Int)
  Clu = addtags(Clu,"orig","link")
  Al = addtags(Al,"orig","link")
  for i = 1:nsteps
    ## Get the grown corner transfer matrix (CTM)
    Au = replacetags(replacetags(replacetags(Al,"down","left","link"),"up","right","link"),"left","up","site")
    Clu⁽¹⁾ = Clu*Al*Au*T

    ## Diagonalize the grown CTM
    ld = findindex(Clu⁽¹⁾, "link,down")
    sd = findindex(Clu⁽¹⁾, "site,down")
    lr = findindex(Clu⁽¹⁾, "link,right")
    sr = findindex(Clu⁽¹⁾, "site,right")
    Ud,Cdr = eigen(Clu⁽¹⁾, (ld,sd), (lr,sr);
                   maxdim=χmax,
                   lefttags="link,down,renorm",
                   righttags="link,right,renorm")

    ## The renormalized CTM is the diagonal matrix of eigenvalues
    Clu = replacetags(replacetags(replacetags(Cdr,"renorm","orig"),"down","up"),"right","left")
    Clu = Clu/norm(Clu)

    ## Calculate the renormalized half row transfer matrix (HRTM)
    Al = Al*replacetags(Ud,"down","up")*T*Ud
    Al = replacetags(replacetags(Al,"renorm","orig"),"right","left","site")
    Al = Al/norm(Al)
  end
  Clu = removetags(Clu,"orig")
  Al = removetags(Al,"orig")
  return Clu,Al
end

@testset "ctmrg" begin
  # Make Ising model MPO
  β = 1.1*βc
  d = 2
  s = Index(d,"site")
  sl = addtags(s,"left")
  sr = addtags(s,"right")
  su = addtags(s,"up")
  sd = addtags(s,"down")

  T = ising_mpo((sl,sr),(su,sd),β)

  χ0 = 1
  l = Index(χ0,"link")
  ll = addtags(l,"left")
  lu = addtags(l,"up")
  ld = addtags(l,"down")

  Clu = randomITensor(lu,ll)
  Al = randomITensor(lu,ld,sl)

  Clu = 0.5*(Clu+swaptags(Clu,"up","left"))
  Al = 0.5*(Al+swaptags(Al,"up","down"))

  Clu,Al = ctmrg(T,Clu,Al;χmax=30,nsteps=2000)

  # Normalize corner matrix
  trC⁴ = Clu*replacetags(Clu,"0","1","up")*
         replacetags(Clu,"0","1")*replacetags(Clu,"0","1","left")
  Clu = Clu/scalar(trC⁴)^(1/4)

  # Normalize MPS tensor
  trA² = Clu*replacetags(Clu,"0","1","up")*Al*
         replacetags(Al,"0","1","link")*
         replacetags(replacetags(Clu,"up,0","down,1"),"0","1","left")*
         replacetags(replacetags(Clu,"0","1","left"),"up","down")
  Al = Al/sqrt(scalar(trA²))

  ## Get environment tensors for a single site measurement
  Ar = replacetags(replacetags(Al,"left","right","site"),"0","1","link")
  Au = replacetags(replacetags(replacetags(Al,"left","up","site"),"down","left","link"),"up","right","link")
  Ad = replacetags(replacetags(Au,"up","down","site"),"0","1","link")
  Cld = replacetags(replacetags(Clu,"up","down"),"0","1","left")
  Cru = replacetags(replacetags(Clu,"left","right"),"0","1","up")
  Crd = replacetags(replacetags(Cru,"0","1","right"),"up","down")

  ## Calculate partition function per site
  κ = scalar(Clu*Al*Cld*Au*T*Ad*Cru*Ar*Crd)
  @test κ≈exp(-β*ising_free_energy(β))

  ## Calculate magnetization
  Tsz = ising_mpo((sl,sr),(su,sd),β;sz=true)
  m = scalar(Clu*Al*Cld*Au*Tsz*Ad*Cru*Ar*Crd)/κ
  @test abs(m)≈ising_magnetization(β)
end


text:
function chebyshev_polynomial_derivative(order::S,x::T) where {T<:AbstractFloat,S<:Integer}

  polynomial = ones(T,1,order+1)
  poly_deriv = zeros(T,1,order+1)

  for i = 2:order+1
    if i == 2
      polynomial[i] = x
      poly_deriv[i] = 1.0
    else
      polynomial[i] = 2*x*polynomial[i-1]-polynomial[i-2]
      poly_deriv[i] = ((i-1)*polynomial[i-1]-(i-1)*x*polynomial[i])/(1-x^2)
    end
  end

  return poly_deriv

end

function chebyshev_polynomial_derivative(order::S,x::Array{T,1}) where {T<:AbstractFloat,S<:Integer}

  polynomial = ones(T,length(x),order+1)
  poly_deriv = zeros(T,length(x),order+1)

  for i = 1:length(x)
    poly_deriv[i,:] = chebyshev_polynomial_derivative(order,x[i])
  end

#=

  for i = 2:order+1
	for j = 1:length(x)
      if i == 2
        polynomial[j,i] = x[j]
        poly_deriv[j,i] = 1.0
      else
        polynomial[j,i] = 2*x[j]*polynomial[j,i-1]-polynomial[j,i-2]
        poly_deriv[j,i] = ((i-1)*polynomial[j,i-1]-(i-1)*x[j]*polynomial[j,i])/(1-x[j]^2)
      end
    end
  end

=#

  return poly_deriv

end

max_issues_repo_name: RJDennis/ChebyshevApprox

max_forks_repo_name: RJDennis/ChebyshevApprox


text:
############################################################################################
# Import External Packages
using Test
using Random: Random, AbstractRNG, seed!
using UnPack: UnPack, @unpack, @pack!
using Distributions

############################################################################################
# Import Baytes Packages
using
    BaytesCore,
    ModelWrappers,
    BaytesMCMC,
    BaytesFilters,
    BaytesPMCMC,
    BaytesSMC

############################################################################################
# Include Files
include("testhelper/TestHelper.jl")

############################################################################################
# Run Tests
@testset "All tests" begin
    include("test-utility.jl")
    include("test-construction.jl")
end


text:
tests = ["categorical",
		 "binomial",
		 "poisson",
		 "exponential",
		 "gamma"]

println("Run tests for Sampling:")

for t in tests
    tfile = string(t, ".jl")
    println(" * $(tfile) ...")
    include(tfile)
end


text:
import JSOSolvers: lbfgs, trunk, tron
for fun in (:lbfgs, :trunk, :tron)
  premeth = Meta.parse("JSOSolvers." * string(fun))
  @eval begin
    @doc """
        `$($fun)(stp::NLPStopping; subsolver_verbose::Int = 0, kwargs...)`

    Stopping-version of the $($fun) function from JSOSolvers.jl.
    This function calls `fill_in!` (doesn't update hessian) and `stop!` after $($fun) call,
    if the problem is a success (`:first_order` or `:acceptable`) and `fill_in_on_success` is true or if it failed and `fill_in_on_failure` is true.
    The keyword arguments are passed to the $($fun) call.
    """
    function $premeth(
      stp::NLPStopping;
      subsolver_verbose::Int = 0,
      fill_in_on_success = true,
      fill_in_on_failure = true,
      kwargs...,
    )
      max_ev = if :neval_obj in keys(stp.meta.max_cntrs)
        stp.meta.max_cntrs[:neval_obj]
      else
        typemax(Int)
      end

      stp.meta.start_time = time()

      nlp = stp.pb
      T = eltype(nlp.meta.x0)
      stats = $fun(
        nlp;
        verbose = subsolver_verbose,
        atol = T(stp.meta.atol),
        rtol = T(stp.meta.rtol),
        x = stp.current_state.x,
        # max_iter = stp.meta.max_iter,
        max_time = stp.meta.max_time,
        max_eval = max_ev,
        kwargs...,
      )

      stp = stats_status_to_meta!(stp, stats)

      if status(stp) == :Unknown
        @warn "Error in StoppingInterface statuses: return status is $(stats.status)"
        @show stats.solver_specific
      end

      stp.meta.nb_of_stop = stats.iter
      Stopping._update_time!(stp.current_state, time()) # stats.elapsed_time

      success = stats.status ∈ [:first_order, :acceptable]
      if (success && fill_in_on_success) || (!success && fill_in_on_failure)
        x = stats.solution
        # Not mandatory, but in case some entries of the State are used to stop
        fill_in!(stp, x, Hx = stp.current_state.Hx)

        stop!(stp)
      end

      return stp
    end
  end
end


text:
using Rubin
using Tests
using Elliptic
using HypergeometricFunctions
using Polylogarithms
using SpecialFunctions

@test integrate(x^2*SpecialFunctions.zeta(2,a+b*x), x) == :(2*b^-3*SpecialFunctions.polygamma(-2,x)+b^-1*x^2*SpecialFunctions.polygamma(0,x)+-2*x*b^-2*loggamma(a+b*x))
@test integrate(x*SpecialFunctions.zeta(2,a+b*x), x) == :(-1*b^-2*loggamma(a+b*x)+x*b^-1*SpecialFunctions.polygamma(0,x))
@test integrate(SpecialFunctions.zeta(2,a+b*x), x) == :(b^-1*SpecialFunctions.polygamma(0,x))
@test integrate(x^-2*SpecialFunctions.zeta(2,a+b*x)+-1*b*x^-1*SpecialFunctions.polygamma(2,x), x) == :(-1*x^-1*SpecialFunctions.polygamma(1,x))
@test integrate(x*SpecialFunctions.zeta(s,a+b*x), x) == :(x*b^-1*(1+-1s)^-1*SpecialFunctions.zeta(-1+s,a+b*x)+-1*b^-2*(1+-1s)^-1*(2+-1s)^-1*SpecialFunctions.zeta(-2+s,a+b*x))
@test integrate(SpecialFunctions.zeta(s,a+b*x), x) == :(b^-1*(1+-1s)^-1*SpecialFunctions.zeta(-1+s,a+b*x))
@test integrate(x^-2*SpecialFunctions.zeta(s,a+b*x)+b*s*x^-1*SpecialFunctions.zeta(1+s,a+b*x), x) == :(-1*x^-1*SpecialFunctions.zeta(s,a+b*x))

max_stars_repo_path: test/8 Special functions/8.7 Zeta function.jl


text:
# This file contains functions for the Cost General Direct Approach DEA model
"""
    CostGDADEAModel
An data structure representing a Cost General Direct Approach  DEA model.
"""
struct CostGDADEAModel <: AbstractCostDEAModel
    n::Int64
    m::Int64
    s::Int64
    measure::Symbol
    rts::Symbol
    monetary::Bool
    dmunames::Union{Vector{String},Nothing}
    eff::Vector
    lambda::SparseMatrixCSC{Float64, Int64}
    techeff::Vector
    alloceff::Vector
    normalization::Vector
    Xtarget::Matrix
    Ytarget::Matrix
end

"""
    deacostgda(X, Y, W, measure)
Compute cost efficiency using data envelopment analysis General Direct Approach model for
inputs `X`, outputs `Y`, price of inputs `W`, and efficiency `measure`.

# Measure specification:
- `:ERG`: Enhanced Russell Graph (or Slack Based Measure (SBM)).

# Optional Arguments
- `rts=:VRS`: choose between constant returns to scale `:CRS` or variable returns to scale `:VRS`.
- `atol=1e-6`: tolerance for DMU to be considered efficient.
- `names`: a vector of strings with the names of the decision making units.

# Examples
```jldoctest
julia> X = [2 2; 1 4; 4 1; 4 3; 5 5; 6 1; 2 5; 1.6 8];

julia> Y = [1; 1; 1; 1; 1; 1; 1; 1];

julia> W = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

julia> deacostgda(X, Y, W, :ERG)
General Direct Approach Cost DEA Model
DMUs = 8; Inputs = 2; Outputs = 1
Returns to Scale = VRS
Associated efficiency measure = ERG
──────────────────────────────────
       Cost  Technical  Allocative
──────────────────────────────────
1  0.0        0.0        0.0
2  0.5        0.0        0.5
3  0.5        0.0        0.5
4  0.416667   0.416667   0.0
5  0.6        0.6        0.0
6  0.25       0.166667   0.0833333
7  0.525      0.35       0.175
8  0.532609   0.4375     0.0951087
──────────────────────────────────
```
"""
function deacostgda(X::Union{Matrix,Vector}, Y::Union{Matrix,Vector},
    W::Union{Matrix,Vector},
    measure::Symbol;
    rts::Symbol = :VRS, monetary::Bool = false, atol::Float64 = 1e-6,
    names::Union{Vector{String},Nothing} = nothing,
    optimizer::Union{DEAOptimizer,Nothing} = nothing)::CostGDADEAModel

    # Check parameters
    nx, m = size(X, 1), size(X, 2)
    ny, s = size(Y, 1), size(Y, 2)

    nw, mw = size(W, 1), size(W, 2)

    if nx != ny
        throw(DimensionMismatch("number of rows in X and Y ($nx, $ny) are not equal"));
    end
    if nw != nx
        throw(DimensionMismatch("number of rows in W and X ($nw, $nx) are not equal"));
    end
    if mw != m
        throw(DimensionMismatch("number of columns in W and X ($mw, $m) are not equal"));
    end

    n = nx

    # Get efficiency measure model
    if measure == :ERG
        measuremodel = deacostrussell(X, Y, W, rts = rts, optimizer = optimizer)
        measuretechnicalmodel = dearussell(X, Y, orient = :Input, rts = rts, slack = false, optimizer = optimizer)
    else
        throw(ArgumentError("Invalid efficiency `measure`"));
    end

    normmeasure = normfactor(measuremodel)
    cimeasure = efficiency(measuremodel)  .* normmeasure
    timeasure = efficiency(measuremodel, :Technical)
    clambdaeff = peersmatrix(measuremodel)
    Xtarget = targets(measuremodel, :X)
    Ytarget = targets(measuremodel, :Y)
    Xtargetti = targets(measuretechnicalmodel, :X)
    Ytargetti = targets(measuretechnicalmodel, :Y)

    # Calculate slack as the difference between target and observed
    slackXmeasure = X .- Xtargetti

    # Get profit inefficiency of projection (target)
    if measure == :ERG
        targetmodel = deacostrussell(Xtargetti, Ytargetti, W, rts = rts, optimizer = optimizer)
    end

    citarget = efficiency(targetmodel) .* normfactor(targetmodel)

    monetaryti = sum(W .* slackXmeasure, dims = 2)

    # Normalization factor (equation 13.13)
    normalization = zeros(n)
    for i=1:n
        if abs(timeasure[i] ) <= atol
            # For technical efficient DMU's
            normalization[i] = normmeasure[i]
        else
            # For inefficient DMU's
            normalization[i] = monetaryti[i] ./ timeasure[i]
        end
    end

    if monetary
        cefficiency = cimeasure
        techefficiency = timeasure .* normalization
        allocefficiency = citarget
    else
        cefficiency = cimeasure ./ normalization
        techefficiency = timeasure
        allocefficiency = citarget ./ normalization
    end

    return CostGDADEAModel(n, m, s, measure, rts, monetary, names, cefficiency, clambdaeff, techefficiency, allocefficiency, normalization, Xtarget, Ytarget)
end

function Base.show(io::IO, x::CostGDADEAModel)
    compact = get(io, :compact, false)

    n = nobs(x)
    m = ninputs(x)
    s = noutputs(x)
    dmunames = names(x)

    eff = efficiency(x)
    techeff = efficiency(x, :Technical)
    alloceff = efficiency(x, :Allocative)

    if !compact
        print(io, "General Direct Approach Cost DEA Model \n")
        print(io, "DMUs = ", n)
        print(io, "; Inputs = ", m)
        print(io, "; Outputs = ", s)
        print(io, "\n")
        print(io, "Returns to Scale = ", string(x.rts))
        print(io, "\n")
        print(io, "Associated efficiency measure = ", string(x.measure))
        print(io, "\n")
        show(io, CoefTable(hcat(eff, techeff, alloceff), ["Cost", "Technical", "Allocative"], dmunames))
    end

end


text:
### Kinetic matrix elements

function kinetic(a::PGBF,b::PGBF)
    return a.norm*b.norm*kinetic(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
                                b.expn,b.x,b.y,b.z,b.I,b.J,b.K)
end

function kinetic(aexpn::Float64,ax::Float64,ay::Float64,az::Float64,aI::Int64,aJ::Int64,aK::Int64,
            bexpn::Float64,bx::Float64,by::Float64,bz::Float64,bI::Int64,bJ::Int64,bK::Int64)
    overlap0 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)
    overlapx1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI+2,bJ,bK)
    overlapy1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ+2,bK)
    overlapz1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK+2)
    overlapx2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI-2,bJ,bK)
    overlapy2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ-2,bK)
    overlapz2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK-2)
    term0 = bexpn*(2*(bI+bJ+bK)+3)*overlap0
    term1 = -2*(bexpn^2)*(overlapx1+overlapy1+overlapz1)
    term2 = -0.5*(bI*(bI-1)*overlapx2+bJ*(bJ-1)*overlapy2+bK*(bK-1)*overlapz2)
    return term0+term1+term2
end

kinetic(a::CGBF,b::CGBF) = contract(kinetic,a,b)

function test_kinetic()
    @testset "test_kinetic" begin
    s = pgbf(1.0)
    c = cgbf(0.0,0.0,0.0)
    push!(c,1,1)
    @test amplitude(c,0,0,0) ≈ 0.71270547
    @test kinetic(1.,0.,0.,0.,0,0,0,1.,0.,0.,0.,0,0,0) ≈ 2.9530518648229536
    @test kinetic(s,s) ≈ 1.5
    @test kinetic(c,c) ≈ 1.5
    end #testset
end


text:
module PhysicalOceanography

export f, β

Ω = 7.2921 * 10^(-5) #rad/s
R = 6.3781 * 10^(6) #m

"""
    f(φ)

Coriolis Parameter
"""
f(φ) = 2Ω * sind(φ)

"""
    β(φ)

Rossby Parameter
"""
β(φ) = 2Ω / R * cosd(φ)

end


text:
using Test
@isdefined(TestHelpers) || include("TestHelpers.jl")

TestHelpers.@timed_testset "utils/" begin
    include("utils/import_datasets.jl")
    include("utils/import_example_nets.jl")
end


text:
# Train an normalizing flow for seismic denosing
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia train_warm-start_unsupervised_hint.jl
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using JLD
using JLD2
using Random
using Statistics
using ArgParse
using Printf
using Flux: gpu, cpu
using Flux.Data: DataLoader
using Flux.Optimise: Optimiser, ExpDecay, update!, ADAM
using JOLI4Flux

# Random seed
Random.seed!(0)

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 15
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-4
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 5
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 4
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 12
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 1f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "warm-start-unsupervised-hint-seismic-denoising-256x256-0.3CS"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

# Loading the experiment—only network weights and training loss
pretrained_args = Dict(
    "max_epoch" => 50,
    "epoch"     => 50,
    "lr"        => 1f-3,
    "lr_step"   => 1,
    "batchsize" => 4,
    "n_hidden"  => 64,
    "depth"     => 12,
    "sigma"     => 2f-2,
    "sim_name"  => "seismic-denoising-256x256-0.3CS"
)
Params, fval, exp_path = load_experiment(pretrained_args; return_path=true)
A = wload(exp_path)["A"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "seismic_samples_256_by_256_num_10k.jld")
label_path = datadir("training-data", "seismic_samples_256_by_256_num_10k_labels.jld")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/vapl62yhh8fxgwy/'
        'seismic_samples_256_by_256_num_10k.jld -q -O $data_path`)
end
if isfile(label_path) == false
    run(`wget https://www.dropbox.com/s/blxkh6tdszudlcq/'
        'seismic_samples_256_by_256_num_10k_labels.jld -q -O $label_path`)
end

# Load seismic images and create one observed data
JLD.jldopen(data_path, "r") do file
    X_orig = read(file, "X")

    labels = JLD.jldopen(label_path, "r")["labels"][:]
    idx = findall(x -> x == -1, labels)

    X_orig = X_orig[:, :, :, idx]
    global nx, ny, nc, nsamples = size(X_orig)

    global nx = Int(nx/2)
    global ny = Int(ny/2)
    global nc = Int(nc*4)

    # Whiten data
    X_orig = whiten_input(X_orig)

    idx_fixed = 733
    global X_true = X_orig[:, :, :, idx_fixed:idx_fixed]
end

# Define prior and warmstart network
CH = NetworkConditionalHINT(nc, n_hidden, depth)
CH_prior = NetworkConditionalHINT(nc, n_hidden, depth)

put_params!(CH, convert(Array{Any,1}, Params))
put_params!(CH_prior, convert(Array{Any,1}, copy(Params)))

# Observed data
Y_obs = A*X_true + sigma*randn(Float32, size(X_true))

X_true = wavelet_squeeze(X_true)
Y_obs = wavelet_squeeze(Y_obs)

CH = CH |> gpu
Y_obs = Y_obs |> gpu

# Corresponding data latent variable
Zy_obs = CH.forward_Y(Y_obs)

# reversed network
CH_prior = CH_prior |> gpu
CHrev = reverse(CH)

# Training
# Batch extractor
ntrain = 1000
Z_train = randn(Float32, nx, ny, nc, ntrain) |> gpu

num_batches = cld(ntrain, batchsize)
train_loader = DataLoader(Z_train, batchsize=batchsize, shuffle=true)

# Optimizer
opt = Optimiser(ExpDecay(lr, .9f0, num_batches*lr_step, 0f0), ADAM(lr))

# Training log keeper
fval = zeros(Float32, num_batches*max_epoch)

for epoch=1:max_epoch
    for (itr, Zx) in enumerate(train_loader)

        fval[(epoch-1)*num_batches + itr] = loss_unsupervised(
            CHrev, Y_obs, Zx, Zy_obs, A, sigma, CH_prior
        )

        @printf(
            "Epoch: [%d/%d] | Itr: [%d/%d] | Average log-likelihood: %4.2f\r",
            epoch, max_epoch, itr, num_batches, fval[(epoch-1)*num_batches + itr]
        )

        # Update params
        for p in get_params(CHrev)
            update!(opt, p.data, p.grad)
        end
        clear_grad!(CHrev)
    end

    # Saving parameters and logs
    Params = get_params(CHrev) |> cpu
    save_dict = @strdict epoch max_epoch lr lr_step batchsize n_hidden depth sigma sim_name Array(Y_obs) Array(Zy_obs) X_true Array(Z_train) A pretrained_args Params fval
    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits=6)),
        save_dict;
        safe=true
    )

end


text:
module tmodel

# truncated power law model

export write_grid, write_model, write_lambda, write_dust, Parameters, Grid

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using ..constants

# Write the wavelength sampling file. Only run on setup
function write_lambda(lams::Array{Float64, 1}, basedir::AbstractString)
    fcam = open(basedir * "camera_wavelength_micron.inp", "w")
    nlam = length(lams)
    @printf(fcam, "%d\n", nlam)
    for lam in lams
        @printf(fcam, "%.9e\n", lam) # [microns]
    end

    close(fcam)
end

const eqmirror = true # mirror the grid about the z=0 midplane ?
# if we decide to mirror, then ncells = 1/2 of the true value

# Define a grid object which stores all of these variables
# This will not change for the duration of the run
immutable Grid
    nr::Int
    ntheta::Int
    nphi::Int
    ncells::Int
    # cell edges
    Rs::Vector{Float64}
    Thetas::Vector{Float64}
    Phis::Vector{Float64}
    # cell centers
    rs::Vector{Float64}
    thetas::Vector{Float64}
    phis::Vector{Float64}
end

function Grid(nr::Int, ntheta::Int, r_in::Real, r_out::Real, eqmirror::Bool)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # Inner extent of disk
    r_out = convert(Float64, r_out) * AU # Outer extent of disk

    #Define the cell *walls*
    Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced

    if eqmirror
        ped = 0.1
        #Thetas = linspace(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = linspace(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end])
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

#This function only needs to be run once, upon setup.
function write_grid(basedir::AbstractString, grid::Grid)
    #amr_grid.inp
    f = open(basedir * "amr_grid.inp", "w")

    #Write the header
    @printf(f, "%d\n", 1) #iformat
    @printf(f, "%d\n", 0) #regular grid (no AMR or Oct-tree)
    @printf(f, "%d\n", 100) #spherical coordiantes
    @printf(f, "%d\n", 0) #gridinfo (none needed for now)
    #incl_r incl_phi incl_z #use this axis?
    @printf(f, "%d %d %d \n", 1, 1, 0) # 2D axisymmetric
    #n_r    n_phi   n_z #number of cells in this dimension
    @printf(f, "%d %d %d \n", grid.nr, grid.ntheta, grid.nphi)

    for R in grid.Rs
        @printf(f, "%.9e\n", R)
    end

    for Theta in grid.Thetas
        @printf(f, "%.9e\n", Theta)
    end

    for Phi in grid.Phis
        @printf(f, "%.9e\n", Phi)
    end

    close(f)
end

#Define the 7 parameters listed in Rosenfeld et al., then write functions to go from those parameters to dust density, velocity, microturbulence.

#Let's try defining a parameters type, an object of which gets passed around.

type Parameters
    M_star::Float64 # [M_sun] stellar mass
    r_in::Float64 # [AU] hard inner radius of the disk
    r_out::Float64 # [AU] radius at which density is depleted by delta
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface temperature gradient exponent
    M_gas::Float64 # [M_Sun] disk mass of gas
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

# Typical values, from Rosenfeld et al. 2012, Table 1 for V4046 Sgr
M_star = 1.75 # [M_sun] stellar mass
r_c =  45. # [AU] characteristic radius
T_10 =  115. # [K] temperature at 10 AU
q = 0.63 # temperature gradient exponent
gamma = 1.0 # surface temperature gradient exponent
M_CO = 0.933 # [M_earth] disk mass of CO
ksi = 0.14e5 # [cm s^{-1}] microturbulence
dpc = 73.
incl = 33.5 # [degrees] inclination
PA = 73.
vel = 0.0
mu_RA = 0.0
mu_DEC = 0.0

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
function velocity{T}(r::T, M_star::Float64)
    sqrt(G * M_star ./ r)
end
velocity{T}(r::T, pars::Parameters) = velocity(r, pars.M_star * M_sun)

function temperature{T}(r::T, T_10::Float64, q::Float64)
    T_10 * (r ./ (10. * AU)).^(-q)
end
temperature{T}(r::T, pars::Parameters) = temperature(r, pars.T_10, pars.q)

function Hp{T}(r::T, M_star::Float64, T_10::Float64, q::Float64)
    temp = temperature(r, T_10, q)
    sqrt(kB * temp .* r.^3./(mu_gas * m_H * G * M_star))
end
Hp{T}(r::T,  pars::Parameters) = Hp(r, pars.M_star * M_sun, pars.T_10, pars.q)

# Calculate the gas surface density
function Sigma(r::Float64, pars::Parameters)
    r_c = 1. * AU # [cm]
    r_in = pars.r_in * AU
    r_out = pars.r_out * AU

    if r > r_in && r < r_out
        Sigma_c = pars.M_gas * M_sun * (2 - pars.gamma) / (2 * pi * r_c^2 * ((r_out/r_c)^(2 - pars.gamma) - (r_in/r_c)^(2 - pars.gamma)))
        return Sigma_c * (r/r_c)^(-pars.gamma)
    else
        return 0.0
    end
end

# Delivers a gas density in g/cm^3
function rho_gas(r::Float64, z::Float64, pars::Parameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    return S/(sqrt(2. * pi) * H) * exp(-0.5 * (z/H)^2)
end

# Now, replace these functions to simply multiply rho_gas by X_12CO/m_12CO, or X_13CO/m_13CO, etc.

n_12CO(r::Float64, z::Float64, pars::Parameters) = number_densities["12CO"] * rho_gas(r, z, pars)

n_13CO(r::Float64, z::Float64, pars::Parameters) = number_densities["13CO"] * rho_gas(r, z, pars)

n_C18O(r::Float64, z::Float64, pars::Parameters) = number_densities["C18O"] * rho_gas(r, z, pars)

# It is realistic to include freezout of CO onto dust grains.
# This is the amount by which the number density of the CO is reduced (X_freeze) relative to
# the nominal value.
function X_freeze(temp::Float64, pars::Parameters)
    # If it's cooler than the freezout temperature, reduce the number density by the given factor
    if temp <= pars.T_freeze
        return pars.X_freeze
    # Otherwise, just keep it as is.
    else
        return 1.0
    end
end

function rho_dust(r::Float64, z::Float64, pars::Parameters)
    nCO = n_CO(r, z, pars) # number of CO molecules per cm^3

    # Convert from nCO to nH2
    nH2 = nCO / 7.e-5 # number density ratio

    # Convert from nH2 (assuming nH2 ~ nGas ) to mGas
    mGas = mu_gas * amu * nH2 # [g]

    # Convert from mGas to mDust using Gas/Dust ratio of 100
    mDust = mGas * 0.01 # [g]

    return mDust
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Float64)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::Parameters) = microturbulence(pars.ksi)

function write_model(pars::Parameters, basedir::AbstractString, grid::Grid, species::AbstractString)

    funcs = Dict([("12CO", n_12CO), ("13CO", n_13CO), ("C18O", n_C18O)])
    n_CO = funcs[species]

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", n_CO(r_cyl, z, pars))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, pars))
                @printf(ftemp, "%.9e\n", temperature(r_cyl, pars))
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_dust(pars::Parameters, basedir::AbstractString, grid::Grid)
    fdens = open(basedir * "dust_density.inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)
    @printf(fdens, "%d\n", 1) # number of dust species

    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", rho_dust(r_cyl, z, pars))
            end
        end
    end

    close(fdens)
end

end

max_stars_repo_name: elnjensen/DiskJockey

max_issues_repo_name: elnjensen/DiskJockey

max_forks_repo_name: elnjensen/DiskJockey


text:
export sos, chordal_sos, putinar, chordal_putinar

export csp_graph, chordal_csp_graph

const CEG = ChordalExtensionGraph

function csp_graph(poly::MP.APL, ::FullSpace)
    G = CEG.LabelledGraph{MP.variable_union_type(poly)}()
    for mono in MP.monomials(poly)
        CEG.add_clique!(G, MP.effective_variables(mono))
    end
    return G
end

function csp_graph(poly::MP.APL, domain::AbstractAlgebraicSet)
    G = csp_graph(poly, FullSpace())
    for p in equalities(domain)
        CEG.add_clique!(G, MP.effective_variables(p))
    end
    return G
end

function csp_graph(poly::MP.APL, domain::BasicSemialgebraicSet)
    G = csp_graph(poly, domain.V)
    for p in inequalities(domain)
        CEG.add_clique!(G, MP.effective_variables(p))
    end
    return G
end

function chordal_csp_graph(poly::MP.APL, domain::AbstractBasicSemialgebraicSet)
    H, cliques = CEG.chordal_extension(csp_graph(poly, domain), CEG.GreedyFillIn())
    for clique in cliques
        sort!(clique, rev=true)
        unique!(clique)
    end
    return H, cliques

end

struct ChordalPutinar{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: AbstractPreorderCertificate
    cone::CT
    basis::Type{BT}
    maxdegree::Int
end

get(certificate::ChordalPutinar, ::Cone) = certificate.cone

struct ChordalDomain{S, V}
    domain::S
    cliques::Vector{V}
end

function get(::ChordalPutinar, ::PreprocessedDomain, domain::BasicSemialgebraicSet, p)
    H, cliques = chordal_csp_graph(p, domain)
    return ChordalDomain(domain, cliques)
end

function get(::ChordalPutinar, index::PreorderIndices, domain::ChordalDomain)
    return map(PreorderIndex, eachindex(domain.domain.p))
end

function get(certificate::ChordalPutinar, ::MultiplierBasis, index::PreorderIndex, domain::ChordalDomain)
    q = domain.domain.p[index.value]
    maxdegree_s2 = certificate.maxdegree - MP.maxdegree(q)
    # If maxdegree_s2 is odd, `div(maxdegree_s2, 2)` would make s^2 have degree up to maxdegree_s2-1
    # for this reason, we take `div(maxdegree_s2 + 1, 2)` so that s^2 have degree up to maxdegree_s2+1
    return [maxdegree_gram_basis(certificate.basis, clique, maxdegree_s2 + 1) for clique in domain.cliques if variables(q) ⊆ clique]
end
function get(::Type{ChordalPutinar{CT, BT}}, ::MultiplierBasisType) where {CT, BT}
    return Vector{BT}
end

function get(::ChordalPutinar, ::Generator, index::PreorderIndex, domain::ChordalDomain)
    return domain.domain.p[index.value]
end

get(certificate::ChordalPutinar, ::IdealCertificate) = ChordalIdeal(certificate.cone, certificate.basis, certificate.maxdegree)
get(::Type{<:ChordalPutinar{CT, BT}}, ::IdealCertificate) where {CT, BT} = ChordalIdeal{CT, BT}

SumOfSquares.matrix_cone_type(::Type{<:ChordalPutinar{CT}}) where {CT} = SumOfSquares.matrix_cone_type(CT)

struct ChordalIdeal{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: SimpleIdealCertificate{CT, BT}
    cone::CT
    basis::Type{BT}
    maxdegree::Int
end
function get(certificate::ChordalIdeal, ::GramBasis, poly)
    H, cliques = chordal_csp_graph(poly, FullSpace())
    return map(cliques) do clique
        return maxdegree_gram_basis(certificate.basis, clique, certificate.maxdegree)
    end
end
function get(::Type{ChordalIdeal{CT, BT}}, ::GramBasisType) where {CT, BT}
    return Vector{BT}
end


text:
using BinDeps

# Find the directory with installed Julia libraries
julialibpath = dirname(Libdl.dlpath(Libdl.dlopen("libgmp")))

# Set link flags for Autotools external packages
# This is only needed for the ecm package
ENV["LDFLAGS"] = "-L$julialibpath"
# config.log shows that BinDeps also set the following as well.
# We have copied gmp.h from Julia source tree to this location.
# Maybe Julia should also copy it to the installation tree.
ENV["CPPFLAGS"] = "-I../../usr/include"

# -Wl... makes the libecm search for libgmp in the Julia installation rather than the system.
# (The  -lgmp should not be neccessary, but it is.
# The Makefile writes -lgmp multiple times in link commands. -lm occurs five times per link command.)
ENV["LIBS"] = "-lgmp -Wl,-rpath -Wl,$julialibpath"

@BinDeps.setup

gmpecm = library_dependency("gmpecm", aliases = ["libecm"])
primesieve = library_dependency("primesieve", aliases = ["libprimesieve"])
primecount = library_dependency("primecount", aliases = ["libprimecount"], depends = [primesieve])
cprimecount = library_dependency("cprimecount", aliases = ["libcprimecount"], depends = [primecount])
smsieve = library_dependency("smsieve", aliases = ["libsmsieve"], depends = [gmpecm])

provides(Sources, URI("http://dl.bintray.com/kimwalisch/primesieve/primesieve-5.4.1.tar.gz"), primesieve)
provides(Sources, URI("http://dl.bintray.com/kimwalisch/primecount/primecount-1.4.tar.gz"), primecount)
provides(Sources, URI("https://gforge.inria.fr/frs/download.php/file/32159/ecm-6.4.4.tar.gz"), gmpecm)
# Getting zip- or tarball from github with a predictable name is mysterious to me.
# But, pushing tags allows downloading this way...
provides(Sources, URI("https://github.com/jlapeyre/msieve-shared/archive/v0.0.3.tar.gz"), smsieve,unpacked_dir="msieve-shared-0.0.3")

# The Autotools BuildProcess will try to download the source using the data above.
# It would not be hard to modify BinDeps to allow skipping the download.
provides(BuildProcess, Autotools(libtarget = ".libs/libecm."*BinDeps.shlib_ext, configure_options =
                                 String["--enable-shared", "--enable-openmp","--with-gmp-lib=$julialibpath"]), gmpecm)
provides(BuildProcess, Autotools(libtarget = ".libs/libprimesieve."*BinDeps.shlib_ext), primesieve)
provides(BuildProcess, Autotools(libtarget = ".libs/libprimecount."*BinDeps.shlib_ext), primecount)

# BinDeps.depsdir(cprimecount) is /pathto/PackageName/deps
const cpcsrcdir = joinpath(BinDeps.depsdir(cprimecount),"src","cprimecount")
# This source for this library will not be downloaded; it is in this repo.
# BuildProcess and SimpleBuild may use the information in provides(Sources...
# Here, we need no such information, and so omit GetSources.
provides(SimpleBuild,
    (@build_steps begin
        ChangeDirectory(cpcsrcdir)
        `make`
        `make install`
    end),cprimecount, os = :Unix)

# libsmsieve.so needs to know the location of libecm.so . Both are in deps/usr/lib .
# Since this is our own library, the hardcoded rpath is specified in the Makefile rather than
# passed via ENV as above.
smsrcdir = joinpath(BinDeps.depsdir(cprimecount),"src","msieve-shared-0.0.3")
provides(SimpleBuild,
         (@build_steps begin
             GetSources(smsieve)
             @build_steps begin
                 ChangeDirectory(smsrcdir)
                 `cp ../localmsieve/Makefile ../localmsieve/msieveshared.c .`
                 `make ECM=1  msieveshared`
                 `cp libsmsieve.so ../../usr/lib/libsmsieve.so`
             end
         end),smsieve, os = :Unix)

@BinDeps.install Dict([(:gmpecm, :gmpecm),(:primecount, :primecount), (:primesieve, :primesieve),
                       (:cprimecount, :cprimecount), (:smsieve, :smsieve) ])


text:
# Internal file
# Here we write some Julia code!

# Note that for writing package code, you should try to keep your functions "loose"
# output_string(s::Symbol) = string(s)

# This definition is normally a bad idea because the functions will work on
# many other types! Note that the `::Symbol` only does dispatch and does not
# improve performance! For details on this, see
# http://www.stochasticlifestyle.com/7-julia-gotchas-handle/

output_string(s::Symbol) = string(s)

# On the otherhand, packages should strive to achieve good performance and generality.
# This is usually done through good usage of types. In Julia, you must strictly
# define the types on the fields for your types in order to get good performance.

type BadPerformance
  A
end

# will get bad performance. Instead, if you know the type that A should be:

type GoodPerformance
  A::Float64
end

# But again, libraries should strive to be type-generic. If our algorithm should
# work on any floating point number, you should do:

type GoodPerformanceAnyFloat{T<:AbstractFloat}
  A::T
end

# Note that this is not the same as

type BadPerformanceAnyFloat
  A::AbstractFloat
end

# where this way uses an abstract type on the field, and thus has bad performance.
# Note that using a type parameter gives good performance, even if it's un-restricted:

type GoodPerformanceAny{T}
  A::T
end


text:
"""  Functor to estimate RV based on velocity at minimum of CCF.
Warning:  Since a discrete grid of velocities are evaluated, this should only be used in limited situations (e.g., an initial guess). """
struct MeasureRvFromMinCCF <: AbstractMeasureRvFromCCF
end

function (mrv::MeasureRvFromMinCCF)(vels::A1, ccf::A2 ) where {T1<:Real, A1<:AbstractArray{T1,1}, T2<:Real, A2<:AbstractArray{T2,1} }
    idx_min_ccf = findmin(ccf)[2]
    vels_min_ccf = vels[idx_min_ccf]
    Δv = idx_min_ccf<length(vels) ? vels[idx_min_ccf+1] - vels[idx_min_ccf] : vels[idx_min_ccf] - vels[idx_min_ccf-1]
    return (rv = vels_min_ccf, σ_rv = Δv)
end

function (mrv::MeasureRvFromMinCCF)(vels::A1, ccf::A2, ccf_var::A3 ) where {T1<:Real, A1<:AbstractArray{T1,1}, T2<:Real, A2<:AbstractArray{T2,1}, T3<:Real, A3<:AbstractArray{T3,1} }
    return mrv(vels,ccf)
end

max_stars_repo_name: RvSpectML/EchelleCCFs

max_issues_repo_name: RvSpectML/EchelleCCFs

max_forks_repo_name: RvSpectML/EchelleCCFs


text:
using GeoTables
using Tables
using Meshes
using Test, Random, Plots
using ReferenceTests, ImageIO

# workaround GR warnings
ENV["GKSwstype"] = "100"

# environment settings
isCI = "CI" ∈ keys(ENV)
islinux = Sys.islinux()
visualtests = !isCI || (isCI && islinux)
datadir = joinpath(@__DIR__,"data")

@testset "GeoTables.jl" begin
  table = GeoTables.load(joinpath(datadir,"test1.gpkg"))
  @test table.id isa Vector{String}
  @test table.x isa Vector{Float64}
  @test table.y isa Vector{Float64}
  @test table.z isa Vector{Float64}
  @test table.geometry isa Vector{<:Point}

  table = GeoTables.load(joinpath(datadir,"test2.gpkg"))
  row = first(Tables.rows(table))
  @test row.municipio  == "Fortaleza"
  @test row.ano_cria   == 1725.0
  @test row.regiao     == "Grande Fortaleza"
  @test row.area_km2   == 312.21031
  @test row.codigo_ibg == "2304400"
  @test row.geometry isa Multi

  if visualtests
    zone = GeoTables.load(joinpath(datadir,"zone.shp"))
    path = GeoTables.load(joinpath(datadir,"path.shp"))

    Random.seed!(123)
    p = plot(size=(600,400))
    plot!(domain(zone), fill=true, color=:gray)
    plot!(domain(path), fill=true, color=:gray90)
    @test_reference "data/zonepath.png" p
  end
end


text:
module VolumeRHS

using BenchmarkTools
using CUDA
using StableRNGs
using StaticArrays

function loopinfo(name, expr, nodes...)
    if expr.head != :for
        error("Syntax error: pragma $name needs a for loop")
    end
    push!(expr.args[2].args, Expr(:loopinfo, nodes...))
    return expr
end

macro unroll(expr)
    expr = loopinfo("@unroll", expr, (Symbol("llvm.loop.unroll.full"),))
    return esc(expr)
end

# HACK: module-local versions of core arithmetic; needed to get FMA
for (jlf, f) in zip((:+, :*, :-), (:add, :mul, :sub))
    for (T, llvmT) in ((Float32, "float"), (Float64, "double"))
        ir = """
            %x = f$f contract nsz $llvmT %0, %1
            ret $llvmT %x
        """
        @eval begin
            # the @pure is necessary so that we can constant propagate.
            @inline Base.@pure function $jlf(a::$T, b::$T)
                Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
            end
        end
    end
    @eval function $jlf(args...)
        Base.$jlf(args...)
    end
end

let (jlf, f) = (:div_arcp, :div)
    for (T, llvmT) in ((Float32, "float"), (Float64, "double"))
        ir = """
            %x = f$f fast $llvmT %0, %1
            ret $llvmT %x
        """
        @eval begin
            # the @pure is necessary so that we can constant propagate.
            @inline Base.@pure function $jlf(a::$T, b::$T)
                Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
            end
        end
    end
    @eval function $jlf(args...)
        Base.$jlf(args...)
    end
end
rcp(x) = div_arcp(one(x), x) # still leads to rcp.rn which is also a function call

# div_fast(x::Float32, y::Float32) = ccall("extern __nv_fast_fdividef", llvmcall, Cfloat, (Cfloat, Cfloat), x, y)
# rcp(x) = div_fast(one(x), x)

# note the order of the fields below is also assumed in the code.
const _nstate = 5
const _ρ, _U, _V, _W, _E = 1:_nstate
const stateid = (ρ = _ρ, U = _U, V = _V, W = _W, E = _E)

const _nvgeo = 14
const _ξx, _ηx, _ζx, _ξy, _ηy, _ζy, _ξz, _ηz, _ζz, _MJ, _MJI,
_x, _y, _z = 1:_nvgeo
const vgeoid = (ξx = _ξx, ηx = _ηx, ζx = _ζx,
                ξy = _ξy, ηy = _ηy, ζy = _ζy,
                ξz = _ξz, ηz = _ηz, ζz = _ζz,
                MJ = _MJ, MJI = _MJI,
                x = _x,   y = _y,   z = _z)

const N = 4
const nmoist = 0
const ntrace = 0

Base.@irrational grav 9.81 BigFloat(9.81)
Base.@irrational gdm1 0.4 BigFloat(0.4)

function volumerhs!(rhs, Q, vgeo, gravity, D, nelem)
    Q = Base.Experimental.Const(Q)
    vgeo = Base.Experimental.Const(vgeo)
    D = Base.Experimental.Const(D)

    nvar = _nstate + nmoist + ntrace

    Nq = N + 1

    s_D = CuStaticSharedArray(eltype(D), (Nq, Nq))
    s_F = CuStaticSharedArray(eltype(Q), (Nq, Nq, _nstate))
    s_G = CuStaticSharedArray(eltype(Q), (Nq, Nq, _nstate))

    r_rhsρ = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsU = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsV = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsW = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsE = MArray{Tuple{Nq}, eltype(rhs)}(undef)

    e = blockIdx().x
    j = threadIdx().y
    i = threadIdx().x

    @inbounds begin
        for k in 1:Nq
            r_rhsρ[k] = zero(eltype(rhs))
            r_rhsU[k] = zero(eltype(rhs))
            r_rhsV[k] = zero(eltype(rhs))
            r_rhsW[k] = zero(eltype(rhs))
            r_rhsE[k] = zero(eltype(rhs))
        end

        # fetch D into shared
        s_D[i, j] = D[i, j]
        @unroll for k in 1:Nq
            sync_threads()

            # Load values will need into registers
            MJ = vgeo[i, j, k, _MJ, e]
            ξx, ξy, ξz = vgeo[i,j,k,_ξx,e], vgeo[i,j,k,_ξy,e], vgeo[i,j,k,_ξz,e]
            ηx, ηy, ηz = vgeo[i,j,k,_ηx,e], vgeo[i,j,k,_ηy,e], vgeo[i,j,k,_ηz,e]
            ζx, ζy, ζz = vgeo[i,j,k,_ζx,e], vgeo[i,j,k,_ζy,e], vgeo[i,j,k,_ζz,e]
            z = vgeo[i,j,k,_z,e]

            U, V, W = Q[i, j, k, _U, e], Q[i, j, k, _V, e], Q[i, j, k, _W, e]
            ρ, E = Q[i, j, k, _ρ, e], Q[i, j, k, _E, e]

            # GPU performance trick
            # Allow optimizations to use the reciprocal of an argument rather than perform division.
            # IEEE floating-point division is implemented as a function call
            ρinv = rcp(ρ)
            ρ2inv = rcp(2ρ)
            # ρ2inv = 0.5f0 * pinv

            P = gdm1*(E - (U^2 + V^2 + W^2)*ρ2inv - ρ*gravity*z)

            fluxρ_x = U
            fluxU_x = ρinv * U * U + P
            fluxV_x = ρinv * U * V
            fluxW_x = ρinv * U * W
            fluxE_x = ρinv * U * (E + P)

            fluxρ_y = V
            fluxU_y = ρinv * V * U
            fluxV_y = ρinv * V * V + P
            fluxW_y = ρinv * V * W
            fluxE_y = ρinv * V * (E + P)

            fluxρ_z = W
            fluxU_z = ρinv * W * U
            fluxV_z = ρinv * W * V
            fluxW_z = ρinv * W * W + P
            fluxE_z = ρinv * W * (E + P)

            s_F[i, j,  _ρ] = MJ * (ξx * fluxρ_x + ξy * fluxρ_y + ξz * fluxρ_z)
            s_F[i, j,  _U] = MJ * (ξx * fluxU_x + ξy * fluxU_y + ξz * fluxU_z)
            s_F[i, j,  _V] = MJ * (ξx * fluxV_x + ξy * fluxV_y + ξz * fluxV_z)
            s_F[i, j,  _W] = MJ * (ξx * fluxW_x + ξy * fluxW_y + ξz * fluxW_z)
            s_F[i, j,  _E] = MJ * (ξx * fluxE_x + ξy * fluxE_y + ξz * fluxE_z)

            s_G[i, j,  _ρ] = MJ * (ηx * fluxρ_x + ηy * fluxρ_y + ηz * fluxρ_z)
            s_G[i, j,  _U] = MJ * (ηx * fluxU_x + ηy * fluxU_y + ηz * fluxU_z)
            s_G[i, j,  _V] = MJ * (ηx * fluxV_x + ηy * fluxV_y + ηz * fluxV_z)
            s_G[i, j,  _W] = MJ * (ηx * fluxW_x + ηy * fluxW_y + ηz * fluxW_z)
            s_G[i, j,  _E] = MJ * (ηx * fluxE_x + ηy * fluxE_y + ηz * fluxE_z)

            r_Hρ = MJ * (ζx * fluxρ_x + ζy * fluxρ_y + ζz * fluxρ_z)
            r_HU = MJ * (ζx * fluxU_x + ζy * fluxU_y + ζz * fluxU_z)
            r_HV = MJ * (ζx * fluxV_x + ζy * fluxV_y + ζz * fluxV_z)
            r_HW = MJ * (ζx * fluxW_x + ζy * fluxW_y + ζz * fluxW_z)
            r_HE = MJ * (ζx * fluxE_x + ζy * fluxE_y + ζz * fluxE_z)

            # one shared access per 10 flops
            for n = 1:Nq
                Dkn = s_D[k, n]

                r_rhsρ[n] += Dkn * r_Hρ
                r_rhsU[n] += Dkn * r_HU
                r_rhsV[n] += Dkn * r_HV
                r_rhsW[n] += Dkn * r_HW
                r_rhsE[n] += Dkn * r_HE
            end

            r_rhsW[k] -= MJ * ρ * gravity

            sync_threads()

            # loop of ξ-grid lines
            @unroll for n = 1:Nq
                Dni = s_D[n, i]
                Dnj = s_D[n, j]

                r_rhsρ[k] += Dni * s_F[n, j, _ρ]
                r_rhsρ[k] += Dnj * s_G[i, n, _ρ]

                r_rhsU[k] += Dni * s_F[n, j, _U]
                r_rhsU[k] += Dnj * s_G[i, n, _U]

                r_rhsV[k] += Dni * s_F[n, j, _V]
                r_rhsV[k] += Dnj * s_G[i, n, _V]

                r_rhsW[k] += Dni * s_F[n, j, _W]
                r_rhsW[k] += Dnj * s_G[i, n, _W]

                r_rhsE[k] += Dni * s_F[n, j, _E]
                r_rhsE[k] += Dnj * s_G[i, n, _E]
            end
        end # k

        @unroll for k in 1:Nq
            MJI = vgeo[i, j, k, _MJI, e]

            # Updates are a performance bottleneck
            # primary source of stall_long_sb
            rhs[i, j, k, _U, e] += MJI*r_rhsU[k]
            rhs[i, j, k, _V, e] += MJI*r_rhsV[k]
            rhs[i, j, k, _W, e] += MJI*r_rhsW[k]
            rhs[i, j, k, _ρ, e] += MJI*r_rhsρ[k]
            rhs[i, j, k, _E, e] += MJI*r_rhsE[k]
        end
    end
    return
end

function main()
    DFloat = Float32
    nelem = 240_000

    rng = StableRNG(123)

    Nq = N + 1
    nvar = _nstate + nmoist + ntrace

    Q = 1 .+ CuArray(rand(rng, DFloat, Nq, Nq, Nq, nvar, nelem))
    Q[:, :, :, _E, :] .+= 20
    vgeo = CuArray(rand(rng, DFloat, Nq, Nq, Nq, _nvgeo, nelem))

    # make sure the entries of the mass matrix satisfy the inverse relation
    vgeo[:, :, :, _MJ, :] .+= 3
    vgeo[:, :, :, _MJI, :] .= 1 ./ vgeo[:, :, :, _MJ, :]

    D = CuArray(rand(rng, DFloat, Nq, Nq))

    rhs = CuArray(zeros(DFloat, Nq, Nq, Nq, nvar, nelem))

    threads=(N+1, N+1)

    kernel = @cuda launch=false volumerhs!(rhs, Q, vgeo, DFloat(grav), D, nelem)
    # XXX: should we print these for all kernels? maybe upload them to Codespeed?
    @info """volumerhs! details:
              - $(CUDA.registers(kernel)) registers, max $(CUDA.maxthreads(kernel)) threads
              - $(Base.format_bytes(CUDA.memory(kernel).local)) local memory,
                $(Base.format_bytes(CUDA.memory(kernel).shared)) shared memory,
                $(Base.format_bytes(CUDA.memory(kernel).constant)) constant memory"""
    results = @benchmark begin
        CUDA.@sync $kernel($rhs, $Q, $vgeo, $(DFloat(grav)), $D, $nelem;
                           threads=$threads, blocks=$nelem)
    end

    # BenchmarkTools captures inputs, JuliaCI/BenchmarkTools.jl#127, so forcibly free them
    CUDA.unsafe_free!(rhs)
    CUDA.unsafe_free!(Q)
    CUDA.unsafe_free!(vgeo)
    CUDA.unsafe_free!(D)

    results
end

end

VolumeRHS.main()


text:
"""
Download ConceptNetNumberbatch embeddings given a `url` and saves them
to a file pointed to by `localfile`.
"""
function download_embeddings(;url=CONCEPTNET_EN_LINK,
                             localfile=abspath("./_conceptnet_/" *
                                               split(url,"/")[end]))
    directory = join(split(localfile, "/")[1:end-1], "/")
    !isempty(directory) && !isdir(directory) && mkpath(directory)
    @info "Download ConceptNetNumberbatch to $localfile..."
    if !isfile(localfile)
        download(url, localfile)
        if isfile(localfile) return localfile end
    else
        @warn "$localfile already exists. Will not download."
        return localfile
    end
end

"""
Load the embeddings given a valid ConceptNetNumberbatch `filepath`,
lading at most `max_vocab_size` embeddings if no specific `keep_words` are
specified, filtering on `languages`.
"""
function load_embeddings(filepath::AbstractString;
                         max_vocab_size::Union{Nothing,Int}=nothing,
                         keep_words=String[],
                         languages::Union{Nothing, Languages.Language,
                                          Vector{<:Languages.Language},
                                          Symbol, Vector{Symbol}}=nothing,
                         data_type::Type{E}=Float64) where E<:Real
    if languages isa Nothing
        languages = unique(collect(values(LANGUAGES)))
    elseif languages isa Symbol
        languages = LANGUAGES[languages]
    elseif languages isa Vector{Symbol}
        languages = [LANGUAGES[lang] for lang in languages]
    end

    if any(endswith.(filepath, [".gz", ".gzip"]))
        conceptnet = _load_gz_embeddings(filepath,
                                         GzipDecompressor(),
                                         max_vocab_size,
                                         keep_words,
                                         languages=languages,
                                         data_type=data_type)
    elseif any(endswith.(filepath, [".h5", ".hdf5"]))
        conceptnet = _load_hdf5_embeddings(filepath,
                                           max_vocab_size,
                                           keep_words,
                                           languages=languages,
                                           data_type=data_type)
    else
        conceptnet = _load_gz_embeddings(filepath,
                                         Noop(),
                                         max_vocab_size,
                                         keep_words,
                                         languages=languages,
                                         data_type=data_type)
    end
    return conceptnet
end

"""
Load the ConceptNetNumberbatch embeddings from a .gz or uncompressed file.
"""
function _load_gz_embeddings(filepath::S1,
                             decompressor::TranscodingStreams.Codec,
                             max_vocab_size::Union{Nothing,Int},
                             keep_words::Vector{S2};
                             languages::Union{Nothing, Languages.Language,
                                              Vector{<:Languages.Language}
                                             }=nothing,
                         data_type::Type{E}=Float64) where
        {E<:Real, S1<:AbstractString, S2<:AbstractString}
    local lang_embs, _length::Int, _width::Int, type_lang, fuzzy_words
    type_word = String
    open(filepath, "r") do fid
        cfid = TranscodingStream(decompressor, fid)
        _length, _width = parse.(Int64, split(readline(cfid)))
        vocab_size = _get_vocab_size(_length,
                                     max_vocab_size,
                                     keep_words)
        lang_embs, languages, type_lang, english_only =
            process_language_argument(languages, type_word, data_type)
		fuzzy_words = Dict{type_lang, Vector{type_word}}()
        no_custom_words = length(keep_words)==0
        lang = :en
        cnt = 0
        for (idx, line) in enumerate(eachline(cfid))
            word, _ = _parseline(line, data_type, word_only=true)
            if !english_only
                _, _, _lang, word = split(word,"/")
                lang = Symbol(_lang)
            end
            if word in keep_words || no_custom_words
                if lang in keys(LANGUAGES) && LANGUAGES[lang] in languages  # use only languages mapped in LANGUAGES
                    _llang = LANGUAGES[lang]
                    if !haskey(lang_embs, _llang)
                        push!(lang_embs, _llang=>Dict{type_word,
                                                      Vector{data_type}}())
                        push!(fuzzy_words, _llang=>type_word[])
                    end
                    _, embedding = _parseline(line, data_type, word_only=false)
                    occursin("#", word) && push!(fuzzy_words[_llang], word)
                    push!(lang_embs[_llang], word=>embedding)
                    cnt+=1
                    if cnt > vocab_size-1
                        break
                    end
                end
            end
        end
        close(cfid)
    end
    return ConceptNet{type_lang, type_word, data_type}(lang_embs, _width, fuzzy_words)
end

"""
Load the ConceptNetNumberbatch embeddings from a HDF5 file.
"""
function _load_hdf5_embeddings(filepath::S1,
                               max_vocab_size::Union{Nothing,Int},
                               keep_words::Vector{S2};
                               languages::Union{Nothing, Languages.Language,
                                    Vector{<:Languages.Language}}=nothing,
                               data_type::Type{E}=Int8) where
        {S1<:AbstractString, S2<:AbstractString, E<:Real}
    local fuzzy_words
    type_word = String
    payload = h5open(read, filepath)["mat"]
    words = map(payload["axis1"]) do val
        _, _, lang, word = split(val, "/")
        return Symbol(lang), word
    end
    embeddings = payload["block0_values"]
    vocab_size = _get_vocab_size(length(words),
                                 max_vocab_size,
                                 keep_words)
    lang_embs, languages, type_lang, _ =
        process_language_argument(languages, type_word, E)
	fuzzy_words = Dict{type_lang, Vector{type_word}}()
    no_custom_words = length(keep_words)==0
    cnt = 0
    for (idx, (lang, word)) in enumerate(words)
        if word in keep_words || no_custom_words
            if haskey(LANGUAGES, lang) && LANGUAGES[lang] in languages  # use only languages mapped in LANGUAGES
                _llang = LANGUAGES[lang]
                if !haskey(lang_embs, _llang)
                    push!(lang_embs, _llang=>Dict{type_word, Vector{E}}())
                    push!(fuzzy_words, _llang=>type_word[])
                end
                occursin("#", word) && push!(fuzzy_words[_llang], word)
                push!(lang_embs[_llang], word=>embeddings[:,idx])
                cnt+=1
                if cnt > vocab_size-1
                    break
                end
            end
        end
    end
    return ConceptNet{type_lang, type_word, E}(lang_embs, size(embeddings,1), fuzzy_words)
end

# Function that returns some needed structures based on the languages provided
# Returns:
#   - a dictionary to store the embeddings
#   - a vector of Languages.Language (used to check whether to load embedding or not
#     while parsing)
#   - the type of the language
#   - a flag specifying whether only English is used or not
function process_language_argument(languages::Nothing,
                                   type_word::T1,
                                   type_data::T2) where {T1, T2}
    return Dict{Languages.Language, Dict{type_word, Vector{type_data}}}(),
           collect(language for language in LANGUAGES),
           Languages.Language, false
end

function process_language_argument(languages::Languages.English,
                                   type_word::T1,
                                   type_data::T2) where {T1, T2}
    return Dict{Languages.English, Dict{type_word, Vector{type_data}}}(), [languages],
           Languages.English, true
end

function process_language_argument(languages::L,
                                   type_word::T1,
                                   type_data::T2) where {L<:Languages.Language, T1, T2}
    return Dict{L, Dict{type_word, Vector{type_data}}}(), [languages], L, false
end

function process_language_argument(languages::Vector{L},
                                   type_word::T1,
                                   type_data::T2) where {L<:Languages.Language, T1, T2}
    if length(languages) == 1
        return process_language_argument(languages[1], type_word, type_data)
    else
        return Dict{L, Dict{type_word, Vector{type_data}}}(), languages, L, false
    end
end

"""
Calculate how many embeddings to retreive.
"""
function _get_vocab_size(real_vocab_size,
                         max_vocab_size=nothing,
                         keep_words=String[])

    # The real dataset cannot contain negative samples
    real_vocab_size = max(0, real_vocab_size)
    # If no maximum number of words is specified,
    # maximum size is the actual size
    if max_vocab_size == nothing
        max_vocab_size = real_vocab_size
    end
    # The maximum has to be at most the real size
    max_vocab_size = min(real_vocab_size, max_vocab_size)
    # The maximum cannot be more than the number of custom words
    # if there are custom words
    n_custom_words = length(keep_words)
    if n_custom_words > 0
        max_vocab_size = min(max_vocab_size, n_custom_words)
    end
    return max_vocab_size
end

"""
Parse a line of text from a ConceptNetNumberbatch delimited file.
"""
function _parseline(buf, data_type::Type{E}; word_only=false) where E<:Real
    bufvec = split(buf, " ")
    word = string(popfirst!(bufvec))
    if word_only
        return word, E[]
    else
        embedding = parse.(E, bufvec)
        return word, embedding
    end
end