text:
# New simplification rules
exponent_rules = [
    @rule(exp(log(~x)) => ~x),
    @acrule(exp(~x + ~y) => exp(~x) * exp(~y)),
    @acrule(exp(~x * ~y) => exp(~y)^~x),
]
rw_exp = RestartedChain(exponent_rules)
rw_chain = RestartedChain([SymbolicUtils.default_simplifier(),rw_exp])
rewriter = Postwalk(rw_chain)

# Constitutive relations
cr(n::Component) = n.equations
function cr(m::EqualEffort)
    if isempty(freeports(m))
        return Vector{Equation}([])
    end

    N = numports(m)
    @variables E[1:N](t) F[1:N](t)

    flow_constraint = [0 ~ sum(collect(F))]
    effort_constraints = [0 ~ E[1] - e for e in E[2:end]]
    return vcat(flow_constraint,effort_constraints)
end
function cr(m::EqualFlow)
    if isempty(freeports(m))
        return Vector{Equation}([])
    end
    N = numports(m)
    @variables E[1:N](t) F[1:N](t)

    weighted_e = m.weights.*collect(E)
    weighted_f = m.weights.*collect(F)

    effort_constraint = [0 ~ sum(weighted_e)]
    flow_constraints = [0 ~ weighted_f[1] - f for f in weighted_f[2:end]]
    return vcat(effort_constraint,flow_constraints)
end
function cr(m::BondGraph)
    return simplify.(ModelingToolkit.equations(de_system(m)),rewriter=rewriter)
end

@connector function MTKPort(;name)
    vars = @variables E(t) F(t)
    ODESystem(Equation[], t, vars, []; name=name)
end

ModelingToolkit.connect(::Type{MTKPort}, p1, p2) =
    [0 ~ p1.F + p2.F, 0 ~ p1.E - p2.E]

function ModelingToolkit.ODESystem(n::AbstractNode)
    N = numports(n)
    ps = [MTKPort(name=Symbol("p$i")) for i in 1:N]

    @variables E[1:N](t) F[1:N](t)
    e_sub_rules = Dict(E[i] => ps[i].E for i in 1:N)
    f_sub_rules = Dict(F[i] => ps[i].F for i in 1:N)
    sub_rules = merge(e_sub_rules,f_sub_rules)
    eqs = [substitute(eq,sub_rules) for eq in cr(n)]

    sys = ODESystem(eqs, t, state_vars(n), params(n);
                    name=n.name, defaults=default_value(n))
    return compose(sys, ps...)
end
function ModelingToolkit.ODESystem(m::BondGraph; simplify_eqs=true)
    subsystems = OrderedDict(c => ODESystem(c) for c in components(m))

    # Add constraints from bonds
    connections = Equation[]
    for (s,d) in bonds(m)
        src_port = getproperty(subsystems[s.node],Symbol("p$(s.index)"))
        dst_port = getproperty(subsystems[d.node],Symbol("p$(d.index)"))
        append!(connections,connect(src_port, dst_port))
    end

    @named _model = ODESystem(connections,t; name=m.name)
    model = compose(_model,collect(values(subsystems)))

    if simplify_eqs
        simplified_model = structural_simplify(model)
        simplified_eqs = simplify.(ModelingToolkit.equations(simplified_model); expand=true, rewriter=rewriter)
        # '@set!' allows mutation of an immutable field
        @set! simplified_model.eqs = simplified_eqs
        return simplified_model
    else
        return initialize_system_structure(model)
    end
end

equations(m::Model;simplify_eqs=true) =
    ModelingToolkit.equations(ODESystem(m;simplify_eqs))

function simulate(m::BondGraph, tspan; u0=[], pmap=[], probtype::Symbol=:ODE, kwargs...)

    sys = ODESystem(m)
    flag_ODE = !any([isequal(eq.lhs,0) for eq in ModelingToolkit.equations(sys)])
    if probtype==:ODE || flag_ODE
        prob = ODEProblem(sys, u0, tspan, pmap)
    else
        prob = ODAEProblem(sys, u0, tspan, pmap)
    end
    return solve(prob; kwargs...)
end

==(sym1::T, sym2::T) where T <: SymbolicUtils.Symbolic{Real} = (simplify(sym1) - simplify(sym2)) == 0

max_stars_repo_name: jedforrest/BondGraphs

max_issues_repo_name: jedforrest/BondGraphs

max_forks_repo_name: jedforrest/BondGraphs


text:
#=
    ----------------------------------------------------------------------------
    Implements the main functionality of the package: Markov chain Monte Carlo
    sampler, which allows for exploration of a product space between a path
    space (on which the diffusion of interest is defined) and a parameter space.
    The main function is `mcmc(...)` and it is supported by two main types of
    routines: `impute!(...)` and `update_param(...)`.
    ----------------------------------------------------------------------------
=#

import ForwardDiff: value # currently not needed, but will be

#===============================================================================
                        Routines for imputations
===============================================================================#

"""
    update!(updt::Imputation{NoBlocking}, ws::Workspace, θ, ll,
                 step, ::Any, headstart=false)

Imputation step of the MCMC scheme (without blocking).
...
# Arguments
- `updt`: [Imputation](@ref)
- `ws`: workspace containing most containers and definitions of the chain
- `ll`: log-likelihood of the old (previously accepted) diffusion path
- `θ`: current parameter
- `headstart`: flag for whether to 'ease into' fpt conditions
...
"""
function update!(updt::Imputation{NoBlocking}, ws::Workspace{OS}, θ, ll,
                 step, ::Any, headstart=false) where OS
    ρ = updt.ρs[1]
    # sample proposal starting point
    zᵒ, yᵒ = proposal_start_pt(ws, ws.P[1], ρ)

    # sample proposal path
    m = length(ws.WWᵒ)
    success = sample_segments!(1:m, ws, yᵒ, ρ, Val{headstart}())
    llᵒ = success ? (logpdf(ws.x0_prior, yᵒ) +
                     path_log_likhd(OS(), ws.XXᵒ, ws.P, 1:m, ws.fpt) +
                     lobslikelihood(ws.P[1], yᵒ)) : -Inf

    print_info(step, value(ll), value(llᵒ), "impute")

    if accept_sample(llᵒ-ll, step.verbose)
        swap!(ws.XX, ws.XXᵒ, ws.WW, ws.WWᵒ, 1:m)
        set!(ws.z, zᵒ)
        return llᵒ, true, θ
    else
        return ll, false, θ
    end
end

"""
    proposal_start_pt(ws::Workspace, P, ρ)

Sample a new proposal starting point using a random walk with pCN step
...
# Arguments
- `ws`: Workspace, contains previously accepted `z` and the definition of a prior
- `P`: diffusion law according to which auxiliary posterior is to be computed
- `ρ`: memory parameter of the preconditioned Crank Nicolson scheme
...
"""
function proposal_start_pt(ws::Workspace, P, ρ)
    x0_prior, z = ws.x0_prior, ws.z.val
    zᵒ = rand(x0_prior, z, ρ)
    yᵒ = start_pt(zᵒ, x0_prior, P)
    zᵒ, yᵒ
end

"""
    sample_segments!(iRange, ws::Workspace{OS}, y, ρ,
                     headstart::Val{false}=Val{false}()
                     ) where OS

Sample paths segments in index range `iRange` using preconditioned
Crank-Nicolson scheme
...
# Arguments
- `iRange`: range of indices of the segments that need to be sampled
- `ws`: Workspace, containing all relevant containers of the Markov chain
- `y`: starting point
- `ρ`: memory parameter of the preconditioned Crank-Nicolson scheme
- `headstart`: whether to ease into first-passage time sampling
...
"""
function sample_segments!(iRange, ws::Workspace, y, ρ,
                          headstart=Val{false}())
    _sample_segments!(iRange, ws, y, ρ, headstart)
end

function _sample_segments!(iRange, ws::Workspace{OS}, y, ρ,
                          headstart::Val{false}) where OS
    for i in iRange
        success, y = sample_segment!(i, ws, y, ρ)
        !success && return false
    end
    true
end

function _sample_segments!(iRange, ws::Workspace{OS}, y, ρ,
                          headstart::Val{true}) where OS
    for i in iRange
        success, _ = sample_segment!(i, ws, y, ρ)
        while !success && !checkFpt(OS(), ws.XXᵒ[i], ws.fpt[i])
            success, _ = sample_segment!(i, ws, y, ρ)
        end
        y = ws.XXᵒ[i].yy[end]
    end
    true
end

"""
    sample_segment!(i, ws, y, ρ)

Sample the `i`th path segment using preconditioned Crank-Nicolson scheme
...
# Arguments
- `i`: index of the segment to be sampled
- `ws`: Workspace, containing all relevant containers of the Markov chain
- `y`: starting point
- `ρ`: memory parameter of the preconditioned Crank-Nicolson scheme
...
"""
function sample_segment!(i, ws, y, ρ)
    sample!(ws.WWᵒ[i], ws.Wnr)
    crank_nicolson!(ws.WWᵒ[i].yy, ws.WW[i].yy, ρ)
    success, _ = solve!(EulerMaruyamaBounded(), ws.XXᵒ[i], y, ws.WWᵒ[i], ws.P[i]) # always according to trgt law
    success, ws.XXᵒ[i].yy[end]
end

"""
    crank_nicolson!(yᵒ, y, ρ)

Preconditioned Crank-Nicolson update with memory parameter `ρ`, previous vector
`y` and new vector `yᵒ`
"""
crank_nicolson!(yᵒ, y, ρ) = (yᵒ .= √(1-ρ)*yᵒ + √(ρ)*y)

"""
    path_log_likhd(::OS, XX, P, iRange, fpt; skipFPT=false
                   ) where OS <: AbstractObsScheme

Compute likelihood for path `XX` to be observed under `P`. Only segments with
index numbers in `iRange` are considered. `fpt` contains relevant info about
checks regarding adherence to first passage time pattern. `skipFPT` if set to
`true` can skip the step of checking adherence to fpt pattern (used for
conjugate updates, or any updates that keep `XX` unchanged)
"""
function path_log_likhd(::OS, XX, P, iRange, fpt; skipFPT=false
                        ) where OS <: AbstractObsScheme
    ll = 0.0
    for i in iRange
        ll += llikelihood(LeftRule(), XX[i], P[i])
    end
    !skipFPT && (ll = check_full_path_fpt(OS(), XX, iRange, fpt) ? ll : -Inf)
    !skipFPT && (ll += check_domain_adherence(P, XX, iRange) ? 0.0 : -Inf)
    ll
end

"""
    print_info(step, ll, llᵒ, msg="update")

Print information to the console about current likelihood values

...
# Arguments
- `verbose`: flag for whether to print anything at all
- `it`: iteration of the Markov chain
- `ll`: likelihood of the previous, accepted sample
- `llᵒ`: likelihood of the proposal sample
- `msg`: message to start with
...
"""
function print_info(step, ll, llᵒ, msg="update")
    step.verbose && print(msg, ": ", step.iter, " ll ", round(ll, digits=3), " ",
                     round(llᵒ, digits=3), " diff_ll: ", round(llᵒ-ll,digits=3))
end

"""
    accept_sample(log_threshold, verbose=false)

Make a random MCMC decision for whether to accept a sample or reject it.
"""
function accept_sample(log_threshold, verbose=false)
    if rand(Exponential(1.0)) > -log_threshold # Reject if NaN
        verbose && print("\t ✓\n")
        return true
    else
        verbose && print("\t .\n")
        return false
    end
end

"""
    swap!(A, Aᵒ, iRange)

Swap contents between containers A & Aᵒ in the index range iRange
"""
function swap!(A, Aᵒ, iRange)
    for i in iRange
        A[i], Aᵒ[i] = Aᵒ[i], A[i]
    end
end

"""
    swap!(A, Aᵒ, B, Bᵒ, iRange)

Swap contents between containers A & Aᵒ in the index range iRange, do the same
for containers B & Bᵒ
"""
function swap!(A, Aᵒ, B, Bᵒ, iRange)
    swap!(A, Aᵒ, iRange)
    swap!(B, Bᵒ, iRange)
end

"""
    impute!(ws::Workspace{OS,ChequeredBlocking}, ll, verbose=false, it=NaN,
            headstart=false) where OS

Imputation step of the MCMC scheme (without blocking).
...
# Arguments
- `ws`: workspace containing most containers and definitions of the chain
- `ll`: log-likelihood of the old (previously accepted) diffusion path
- `verbose`: whether to print update info while sampling
- `it`: iteration index of the MCMC algorithm
- `headstart`: flag for whether to 'ease into' fpt conditions
...
"""
function update!(updt::Imputation{<:Block}, ws::Workspace{OS}, θ, ll, step,
                 ::Any, headstart=false) where OS
    P, XXᵒ, XX = ws.P, ws.XXᵒ, ws.XX

    recompute_accepted_law!(updt, ws)
    # proposal starting point:
    z_prop, y_prop = proposal_start_pt(ws, P[1], updt.ρs[1])

    ll_total = 0.0
    acc = fill(false, length(updt.blocking.blocks))
    for (block_idx, block) in enumerate(updt.blocking.blocks)
        block_flag = Val{block[1]}()

        # previously accepted starting point
        y = XX[block[1]].yy[1]

        # proposal starting point for the block (can be non-y only for the first block)
        yᵒ = choose_start_pt(block_flag, y, y_prop)

        # sample path in block
        success = sample_segments!(block, ws, yᵒ, updt.ρs[block_idx])
        set_end_pt_manually!(block_idx, block, ws, length(updt.blocking.blocks))

        # starting point, path and observations contribution
        llᵒ = success ? (start_pt_log_pdf(block_flag, ws.x0_prior, yᵒ) +
                         path_log_likhd(OS(), XXᵒ, P, block, ws.fpt) +
                         lobslikelihood(P[block[1]], yᵒ)) : -Inf

        ll_prev = start_pt_log_pdf(block_flag, ws.x0_prior, y)
        ll_prev += path_log_likhd(OS(), XX, P, block, ws.fpt; skipFPT=true)
        ll_prev += lobslikelihood(P[block[1]], y)

        print_info(step, value(ll_prev), value(llᵒ), "impute")
        if accept_sample(llᵒ-ll_prev, step.verbose)
            swap!(XX, XXᵒ, ws.WW, ws.WWᵒ, block)
            set_z!(block_flag, ws, z_prop)
            acc[block_idx] = true
            ll_total += llᵒ
        else
            ll_total += ll_prev
        end
    end
    return ll_total, acc, θ
end

"""
    recompute_accepted_law!(ws::Workspace)

Recompute the (H, Hν, c) triplet as well as the noise that corresponds to law
`ws.P` that is obtained after switching blocks
"""
function recompute_accepted_law!(updt, ws::Workspace)
    solve_back_rec!(updt.blocking.blocks, updt.solver, ws.P)         # compute (H, Hν, c) for given blocks
    noise_from_path!(updt.blocking.blocks, ws.XX, ws.WW, ws.P) # find noise WW that generates XX under 𝔅.P

    # compute white noise generating starting point under 𝔅
    z = inv_start_pt(ws.XX[1].yy[1], ws.x0_prior, ws.P[1])
    set!(ws.z, z)
end

"""
    solve_back_rec!(::NoBlocking, ws::Workspace{OS,B,ST}, P) where {OS,B,ST}

Solve backward recursion to find H, Hν and c, which together define r̃(t,x)
and p̃(x, 𝓓) under the auxiliary law, when no blocking is done
"""
function solve_back_rec!(::NoBlocking, solver, P)
    m = length(P)
    gpupdate!(P[m]; solver=solver)
    for i in (m-1):-1:1
        gpupdate!(P[i], P[i+1].H[1], P[i+1].Hν[1], P[i+1].c[1]; solver=solver)
    end
end

"""
    solve_back_rec!(ws::Workspace{OS,B,ST}, P) where {OS,B,ST}

Solve backward recursion to find H, Hν and c, which together define r̃(t,x)
and p̃(x, 𝓓) under the auxiliary law, when blocking is done
"""
function solve_back_rec!(blocks::Vector, solver, P)
    for block in reverse(blocks)
        gpupdate!(P[block[end]]; solver=solver)
        for i in reverse(block[1:end-1])
            gpupdate!(P[i], P[i+1].H[1], P[i+1].Hν[1], P[i+1].c[1];
                      solver=solver)
        end
    end
end

"""
    noise_from_path!(ws::Workspace, XX, WW, P)

Compute driving Wiener noise `WW` from path `XX` drawn under law `P`
"""
function noise_from_path!(blocks, XX, WW, P)
    for block in blocks
        for i in block
            inv_solve!(EulerMaruyamaBounded(), XX[i], WW[i], P[i])
        end
    end
end

choose_start_pt(::Val{1}, y, yᵒ) = copy(yᵒ)
choose_start_pt(::Any, y, yᵒ) = copy(y)

"""
    set_end_pt_manually!(block_idx, block, ws::Workspace)

Manually set the end-point of the proposal path under blocking so that it agrees
with the end-point of the previously accepted path. If it is the last block,
then do nothing
"""
function set_end_pt_manually!(block_idx, block, ws::Workspace, block_len)
    if block_idx < block_len
        ws.XXᵒ[block[end]].yy[end] = ws.XX[block[end]].yy[end]
    end
end

"""
    start_pt_log_pdf(::Val{1}, yPr::StartingPtPrior, y)

Compute the log-likelihood contribution of the starting point for a given prior
under a blocking scheme (intended to be used with a first block only)
"""
start_pt_log_pdf(::Val{1}, x0_prior::StartingPtPrior, y) = logpdf(x0_prior, y)

"""
    start_pt_log_pdf(::Any, yPr::StartingPtPrior, y)

Default contribution to log-likelihood from the startin point under blocking
"""
start_pt_log_pdf(::Any, x0_prior::StartingPtPrior, y) = 0.0

set_z!(::Val{1}, ws::Workspace, zᵒ) = set!(ws.z, zᵒ)

set_z!(::Any, ::Workspace, ::Any) = nothing

#===============================================================================
                Parameter update routines via Metropolis-Hastings
===============================================================================#

"""
    update_param!(pu::ParamUpdtDefn{MetropolisHastingsUpdt,UpdtIdx}, θ,
                  ws::Workspace{OS,NoBlocking}, ll, verbose=false, it=NaN
                  ) where {UpdtIdx,OS}
Update parameters
...
# Arguments
- `pu`: determines transition kernel, priors and which parameters to update
- `θ`: current value of the parameter
- `ws`: workspace with all necessary containers needed by the Markov chain
- `ll`: likelihood of the old (previously accepted) parametrisation
- `verbose`: whether to print updates info while sampling
- `it`: iteration index of the MCMC algorithm
...
"""
function update!(pu::ParamUpdate{MetropolisHastingsUpdt}, ws::Workspace{OS}, θ,
                 ll, step, blocking::NoBlocking) where OS
    θᵒ = rand(pu.t_kernel, θ, pu.updt_coord)               # sample new parameter
    (logpdf(pu.priors, θᵒ) === -Inf) && (return ll, false, θ)

    WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
    m = length(WW)

    update_laws!(Pᵒ, θᵒ)
    pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, Pᵒ) # compute (H, Hν, c)

    # find white noise which for a given θᵒ gives a correct starting point
    y = XX[1].yy[1]
    zᵒ = inv_start_pt(y, ws.x0_prior, Pᵒ[1])

    success = find_path_from_wiener!(XXᵒ, y, WW, Pᵒ, 1:m)

    llᵒ = success ? (logpdf(ws.x0_prior, y) +
                     path_log_likhd(OS(), XXᵒ, Pᵒ, 1:m, fpt) +
                     lobslikelihood(Pᵒ[1], y)) : -Inf

    print_info(step, value(ll), value(llᵒ))

    llr = ( llᵒ - ll + prior_kernel_contrib(pu, θ, θᵒ))

    # Accept / reject
    if accept_sample(llr, step.verbose)
        swap!(XX, XXᵒ, P, Pᵒ, 1:m)
        set!(ws.z, zᵒ)
        return llᵒ, true, θᵒ
    else
        return ll, false, θ
    end
end

"""
    update_laws!(Ps, θᵒ)

Set new parameter `θᵒ` for the laws in vector `Ps`
"""
function update_laws!(Ps, θᵒ)
    m = length(Ps)
    for i in 1:m
        Ps[i] = GuidPropBridge(Ps[i], θᵒ)
    end
end

"""
    find_path_from_wiener!(XX, y, WW, P, iRange)

Find path `XX` (that starts from `y`) that is generated under law `P` from the
Wiener process `WW`. Only segments with indices in range `iRange` are considered
"""
function find_path_from_wiener!(XX, y, WW, P, iRange)
    for i in iRange
        success, _ = solve!(EulerMaruyamaBounded(), XX[i], y, WW[i], P[i])
        !success && return false
        y = XX[i].yy[end]
    end
    true
end

"""
    prior_kernel_contrib(t_kern, priors, θ, θᵒ)

Contribution to the log-likelihood ratio from transition kernel `t_kern` and
`priors`.
"""
function prior_kernel_contrib(pu::T, θ, θᵒ) where {T <: MCMCUpdate}
    c = pu.updt_coord
    llr = ( logpdf(pu.t_kernel, c, θᵒ, θ) - logpdf(pu.t_kernel, c, θ, θᵒ)
            + logpdf(pu.priors, θᵒ) - logpdf(pu.priors, θ) )
    llr
end

#NOTE blocking and no-blocking param update should be joined into one function
"""
    update_param!(pu::ParamUpdtDefn{MetropolisHastingsUpdt,UpdtIdx}, θ,
                  ws::Workspace{OS,B,ST}, ll, verbose=false, it=NaN
                  ) where {UpdtIdx,OS,B,ST}
Update parameters
"""
function update!(pu::ParamUpdate{MetropolisHastingsUpdt}, ws::Workspace{OS}, θ,
                 ll, step, blocking::Block) where OS
    WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
    m = length(P)
    θᵒ = rand(pu.t_kernel, θ, pu.updt_coord)
    (logpdf(pu.priors, θᵒ) === -Inf) ? (return ll, false, θ) : nothing
    update_laws!(Pᵒ, θᵒ)                   # update law `Pᵒ` accordingly
    solve_back_rec!(blocking.blocks, pu.aux.solver, Pᵒ)   # compute (H, Hν, c)

    y = XX[1].yy[1]
    zᵒ = inv_start_pt(y, ws.x0_prior, Pᵒ[1])

    llᵒ = logpdf(ws.x0_prior, y)
    for (block_idx, block) in enumerate(blocking.blocks)
        y = XX[block[1]].yy[1]
        success = find_path_from_wiener!(XXᵒ, y, WW, Pᵒ, block)
        !success && accept_sample(-Inf, step.verbose)
        !success && return ll, false, θ

        set_end_pt_manually!(block_idx, block, ws, length(blocking.blocks))

        # Compute log-likelihood ratio
        llᵒ += path_log_likhd(OS(), XXᵒ, Pᵒ, block, ws.fpt)
        llᵒ += lobslikelihood(Pᵒ[block[1]], y)
    end
    print_info(step, value(ll), value(llᵒ))

    llr = ( llᵒ - ll + prior_kernel_contrib(pu, θ, θᵒ))

    # Accept / reject
    if accept_sample(llr, step.verbose)
        swap!(XX, XXᵒ, P, Pᵒ, 1:m)
        set!(ws.z, zᵒ)
        return llᵒ, true, θᵒ
    else
        return ll, false, θ
    end
end

#===============================================================================
                Parameter update routines via conjugate updates
===============================================================================#

"""
    update_param!(pu::ParamUpdtDefn{ConjugateUpdt,UpdtIdx}, θ,
                  ws::Workspace{OS,NoBlocking,ST}, ll, verbose=false, it=NaN
                  ) where {UpdtIdx,OS,ST}
Update parameters using conjugate draws
"""
function update!(pu::ParamUpdate{ConjugateUpdt}, ws::Workspace{OS}, θ, ll,
                 step, blocking::NoBlocking) where OS
    WW, P, XX, fpt = ws.WW, ws.P, ws.XX, ws.fpt
    m = length(WW)
    θᵒ = conjugate_draw(θ, XX, P[1].Target, pu.priors[1], pu.updt_coord)   # sample new parameter

    update_laws!(P, θᵒ)
    pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, P) # compute (H, Hν, c)

    for i in 1:m    # compute wiener path WW that generates XX
        inv_solve!(EulerMaruyamaBounded(), XX[i], WW[i], P[i])
    end
    # compute white noise that generates starting point
    y = XX[1].yy[1]
    z = inv_start_pt(y, ws.x0_prior, P[1])

    llᵒ = logpdf(ws.x0_prior, y)
    llᵒ += path_log_likhd(OS(), XX, P, 1:m, fpt; skipFPT=true)
    llᵒ += lobslikelihood(P[1], y)
    print_info(step, value(ll), value(llᵒ))
    set!(ws.z, z)
    return llᵒ, true, θᵒ
end

#NOTE blocking and no-blocking param conjugate update should be joined into one function
"""
    update_param!(pu::ParamUpdtDefn{ConjugateUpdt,UpdtIdx}, θ,
                  ws::Workspace{OS,B}, ll, verbose=false, it=NaN
                  ) where {UpdtIdx,OS,B}
Update parameters
see the definition of  update_param!(…, ::MetropolisHastingsUpdt, …) for the
explanation of the arguments.
"""
function update!(pu::ParamUpdate{ConjugateUpdt}, ws::Workspace{OS}, θ, ll,
                 step, blocking::Block) where OS
    WW, P, XX, fpt = ws.WW, ws.P, ws.XX, ws.fpt
    m = length(WW)
    θᵒ = conjugate_draw(θ, XX, P[1].Target, pu.priors[1], pu.updt_coord)   # sample new parameter

    update_laws!(P, θᵒ)
    pu.aux.recompute_ODEs && solve_back_rec!(blocking.blocks, pu.aux.solver, P)
    for i in 1:m    # compute wiener path WW that generates XX
        inv_solve!(EulerMaruyamaBounded(), XX[i], WW[i], P[i])
    end
    # compute white noise that generates starting point
    y = XX[1].yy[1]
    z = inv_start_pt(y, ws.x0_prior, P[1])

    llᵒ = logpdf(ws.x0_prior, y)
    for block in blocking.blocks
        llᵒ += path_log_likhd(OS(), XX, P, block, ws.fpt; skipFPT=true)
        llᵒ += lobslikelihood(P[block[1]], XX[block[1]].yy[1])
    end
    print_info(step, value(ll), value(llᵒ))
    set!(ws.z, z)
    return llᵒ, true, θᵒ
end

ext: jl


text:
Base.@kwdef struct ScreeningParams
  start_time::Float64 = 0.0
  precision::Float64 = 0.8
  period::Float64 = 7.0
  lower_bound_age::Int64 = 8
  upper_bound_age::Int64 = 16
end

function screening!(state::AbstractSimState, params::AbstractSimParams, event::Event)
  if params.screening_params != nothing
    for id in 1:numindividuals(state)
      health = MocosSim.health(state, id)
      if health == Healthy || health == Recovered || health == Incubating
        continue
      end

      age = params.ages[id]
      if age < params.screening_params.lower_bound_age || age > params.screening_params.upper_bound_age
        continue
      end

      if rand(state.rng) >= params.screening_params.precision
        continue
      end
      push!(
        state.queue,
        Event(
          Val(DetectionEvent),
          time(event),
          id,
          OutsideQuarantineDetection),
          immediate=true)
    end
  end
end

function add_screening!(state::AbstractSimState, params::AbstractSimParams, time_limit::TimePoint=typemax(TimePoint))
  for screening_time in params.screening_params.start_time:params.screening_params.period:time_limit
    event = Event(Val(ScreeningEvent), screening_time)
    push!(state.queue, event)
  end
end


text:
"""
    EarlyFilteredLogger(filter, logger)

Wraps `logger` in an filter that runs before the log message is created.

For contrast see the [`ActiveFilteredLogger`](@ref) which has full control,
but runs after the log message content is computed.
In most circumstances this is fine, but if your log messages are expensive
to create (e.g. they include summary statistics), then the `EarlyFilteredLogger`
is going to be better.

The `filter` should be a function that returns a boolean.
`true` if the message should be logged and `false` if not.
As input it will be given a named tuple with the following fields:
`(level, _module, group, id)`
See [`LoggingExtras.shouldlog_args`](@ref) for more information on what each is.
"""
struct EarlyFilteredLogger{T <: AbstractLogger, F} <: AbstractLogger
    filter::F
    logger::T
end

function handle_message(logger::EarlyFilteredLogger, args...; kwargs...)
    if comp_handle_message_check(logger.logger, args...; kwargs...)
        return handle_message(logger.logger, args...; kwargs...)
    end
end

function shouldlog(logger::EarlyFilteredLogger, args...)
    log_args = shouldlog_args(args...)
    comp_shouldlog(logger.logger, args...) && logger.filter(log_args)
end

min_enabled_level(logger::EarlyFilteredLogger) = min_enabled_level(logger.logger)
catch_exceptions(logger::EarlyFilteredLogger) = catch_exceptions(logger.logger)

"""
    shouldlog_args

This returns a NamedTuple containing all the arguments the logger gives
to `shouldlog`
It is passed to the early logger filter.
These argument come from the logging macro (@info`, `@warn` etc).

  * `level::LogLevel` Warn, Info, etc,
  * `_module::Module` can be used to specify a different originating module from
    the source location of the message.
  * `group::Symbol` can be used to override the message group (this is
    normally derived from the base name of the source file).
  * `id::Symbol` can be used to override the automatically generated unique
    message identifier.  This is useful if you need to very closely associate
    messages generated on different source lines.
"""
function shouldlog_args(fieldvals...)
    fieldnames = (:level, :_module, :group, :id)
    return NamedTuple{fieldnames, typeof(fieldvals)}(fieldvals)
end


text:
mutable struct EKF
    # Measurement times
    ts::Vector{Float64}
    imuΔt::Float64
    gpsΔt::Float64

    # Measurement residuals
    rs::Matrix{Float64}

    # State estimates and covariance matrix diags
    txp::Vector{Float64} # Epochs accosiated with entries in xhats and Ps
    ixp::Int # Index associated with latest entry
    xhats::Matrix{Float64}
    es::Matrix{Float64}
    Ps::Matrix{Float64}

    # Preallocate matricies
    P⁺::Matrix{Float64}
    P⁻::Matrix{Float64}
    H::Matrix{Float64}
    HPH::Matrix{Float64}
    PH::Matrix{Float64}
    KH::Matrix{Float64}
    F::Matrix{Float64}
    Q::Diagonal{Float64, Vector{Float64}}
    y::Vector{Float64}
    dy::Vector{Float64}

    # Simulators
    gpsSim::GPSSim
    imuSim::IMUSim
    scSim::SpacecraftSim

    # Measurement covariance terms
    σ2GPS::Float64
    σ2Acc::Float64

    # Steps to save per propagation
    stp::Int64

    # Step Index
    k::Int64

    # Update flag
    updated::Bool

    # Finished flag
    finished::Bool

    # Luna perturbation flag
    lunaPerts::Bool
end

function EKF(xhat0, P0, Q, σ2GPS, σ2Acc, ts, gpsΔt, gpsSim, imuSim, scSim; lunaPerts = false, steps2save = 5)
    # Instantiate preallocated matriciesAdd
    ixp     = 1
    rs      = zeros(length(ts), 32 + 3)
    txp     = zeros((steps2save + 1)*length(ts) + 1); txp[1] = scSim.t0
    xhats   = zeros((steps2save + 1)*length(ts) + 1, 7); xhats[1, :] .= xhat0
    es      = zeros((steps2save + 1)*length(ts) + 1, 7);
    Ps      = zeros((steps2save + 1)*length(ts) + 1, 7); Ps[1, :] .= diag(P0)
    P⁻      = zeros(7,7)
    P⁺      = deepcopy(P0)
    H       = zeros(32 + 3, 7)
    HPH     = zeros(32 + 3, 32 + 3)
    PH      = zeros(7, 32 + 3)
    KH      = zeros(7,7)
    F       = zeros(7,7)
    y       = zeros(7 + 49)
    dy      = deepcopy(y)

    # Compute IMU Δt from ts vector
    imuΔt   = ts[2] - ts[1]

    # Compute initial estimates error
    (ytrue, u) = GetStateAndControl(scSim, scSim.t0)
    @views es[1,:] .= xhat0 .- ytrue[1:7]

    # Instantiate EKF
    EKF(ts,imuΔt,gpsΔt,rs,txp,ixp,xhats,es,Ps,P⁺,P⁻,H,HPH,PH,KH,F,Q,y,dy,gpsSim,imuSim,scSim,σ2GPS,σ2Acc,steps2save,0,true,false,lunaPerts)
end

function propagate!(ekf::EKF)
    if ekf.updated == false
        throw(ErrorException("Cannot propagate filter before updating."))
    else
        ekf.k += 1
        ekf.updated = false
    end

    # Get time
    if ekf.k == 1
        t0 = ekf.scSim.t0
        tf = ekf.ts[1]
    else
        t0 = ekf.ts[ekf.k - 1]
        tf = ekf.ts[ekf.k]
    end

    # Generate vector of times to save at
    saveat = range(t0; length = ekf.stp, stop = tf)

    # Construct integration vector
    if ekf.k == 1
        @views ekf.y[1:7]   .= ekf.xhats[1, :]
    else
        @views ekf.y[1:7]   .= ekf.xhats[ekf.ixp, :]
    end
    pVec = reshape(ekf.P⁺, (49, 1))
    for i in 1:49
        ekf.y[7 + i] = pVec[i]
    end

    # Create ODE Problem
    if ekf.lunaPerts == false
        prob = ODEProblem(ekfNoLunarPertEOM!, ekf.y, (t0, tf), ekf)
    else
        prob = ODEProblem(ekfWithLunarPertEOM!, ekf.y, (t0, tf), ekf)
    end

    # Solve ODE
    sol = solve(prob, Vern7(), reltol=1e-10, abstol=1e-10, saveat=saveat)

    # Set P⁻
    ekf.P⁻ .= reshape(sol.u[end][8:end], (7, 7))

    # Update storage matricies
    start = ekf.ixp + 1
    stop  = ekf.ixp + ekf.stp
    step  = 0
    for i in start:stop
        step += 1

        # Time
        ekf.txp[i] = sol.t[step]

        # State
        @views ekf.xhats[i,:] .= sol.u[step][1:7]

        # Error
        (ytrue, u) = GetStateAndControl(ekf.scSim, ekf.txp[i])
        @views ekf.es[i,:] .= ekf.xhats[i,:] .- ytrue[1:7]

        # Covariance
        @views ekf.Ps[i,:] .= diag(reshape(sol.u[step][8:end], (7, 7)))
        ekf.ixp = i
    end

    return nothing
end

function updateGPS!(ekf::EKF)
    if ekf.updated == true
        throw(ErrorException("Cannot update before propagating."))
    else
        ekf.updated = true
    end

    # Get current time
    t = ekf.ts[ekf.k]

    # Get state estimate
    @views xhat = ekf.xhats[ekf.ixp, :]

    # Get true state
    (ytrue, u) = GetStateAndControl(ekf.scSim, t)

    # Generate gps measurements
    @views measGPS = gpsMeasurement(ekf.gpsSim, ytrue[1:3], t)

    # Get satelite ids
    @views sats = measGPS[2:2:end]
    numSats = length(sats)

    # Get expected gps measurements
    @views (expMeasGPS, ps) = gpsMeasurement(ekf.gpsSim, xhat[1:3], sats, t; type = :expected)

    # Compute residuals and measurement partials
    numRejected = 0
    rejectTol   = 10.0
    for i in 1:numSats
        r = measGPS[3 + 2*(i - 1)] - expMeasGPS[3 + 2*(i - 1)]
        if abs(r) < rejectTol
            # Compute residual while catching bad measurements
            @views ekf.rs[ekf.k, i - numRejected] = r

            # Compute measurement Jacobian
            ekf.H .*= 0.0
            @views ekf.H[i - numRejected, 1:3] .= -ps[i,1:3]
            @views ekf.H[i - numRejected, 1:3] ./= norm(ps[i,1:3])
        else
            numRejected += 1
        end
    end
    numSats -= numRejected

    # Compute gain
    numMeas = numSats
    @views mul!(ekf.PH[:,1:numMeas], ekf.P⁻, transpose(ekf.H[1:numMeas, :]))
    @views mul!(ekf.HPH[1:numMeas,1:numMeas], ekf.H[1:numMeas, :], ekf.PH[:,1:numMeas])
    for i in 1:numSats; ekf.HPH[i,i] += ekf.σ2GPS; end
    @views HPHpRFact = factorize(ekf.HPH[1:numMeas,1:numMeas])
    @views rdiv!(ekf.PH[:,1:numMeas], HPHpRFact) # Stores K in PH matrix

    # Compute state update
    ekf.ixp += 1
    @views mul!(ekf.xhats[ekf.ixp, :], ekf.PH[:,1:numMeas], ekf.rs[ekf.k, 1:numMeas])
    ekf.xhats[ekf.ixp, :] .+= ekf.xhats[ekf.ixp - 1, :]

    # Set time in storage vector
    ekf.txp[ekf.ixp] = t

    # Compute error
    (ytrue, u) = GetStateAndControl(ekf.scSim, ekf.txp[ekf.ixp])
    @views ekf.es[ekf.ixp, :] .= ekf.xhats[ekf.ixp, :] .- ytrue[1:7]

    # Compute covariance update
    @views mul!(ekf.KH, ekf.PH[:,1:numMeas], ekf.H[1:numMeas,:], -1.0, 0.0)
    for i in 1:7; ekf.KH[i,i] += 1.0; end
    mul!(ekf.P⁺, ekf.KH, ekf.P⁻)

    ekf.Ps[ekf.ixp, :] .= diag(ekf.P⁺)

    return nothing
end

function updateIMU!(ekf::EKF)
    if ekf.updated == true
        throw(ErrorException("Cannot update before propagating."))
    else
        ekf.updated = true
    end

    # Get current time
    t = ekf.ts[ekf.k]

    # Get state estimate
    @views xhat = ekf.xhats[ekf.ixp, :]

    # Generate accelerometer measurements
    (dy, _us) = GetStateAndControl(ekf.scSim, t; derivs = true)
    @views measAcc = ComputeMeasurement(ekf.imuSim, dy[4:6])

    # Compute expected accelerometer measurement
    @views ekf.y[1:7] .= xhat
    ekf.y[8:end] .= 0.0
    for i in 1:7; ekf.y[7 + 7*(i - 1) + i] = 1.0; end
    if ekf.lunaPerts == false
        ekfNoLunarPertEOM!(ekf.dy, ekf.y, ekf, t)
    else
        ekfWithLunarPertEOM!(ekf.dy, ekf.y, ekf, t)
    end
    @views expMeasAcc = ComputeMeasurement(ekf.imuSim, ekf.dy[4:6]; type = :expected)

    # Compute accelerometer residuals and partials
    ekf.rs[ekf.k, 1:3]   .= measAcc .- expMeasAcc
    @views ekf.H[1:3, 1:3] .= ekf.F[4:6, 1:3]
    @views ekf.H[1:3, 7] .= ekf.F[4:6, 7]

    # Compute gain
    numMeas = 3
    @views mul!(ekf.PH[:,1:numMeas], ekf.P⁻, transpose(ekf.H[1:numMeas, :]))
    @views mul!(ekf.HPH[1:numMeas,1:numMeas], ekf.H[1:numMeas, :], ekf.PH[:,1:numMeas])
    for i in 1:3; ekf.HPH[i,i] += ekf.σ2Acc; end
    @views HPHpRFact = factorize(ekf.HPH[1:numMeas,1:numMeas])
    @views rdiv!(ekf.PH[:,1:numMeas], HPHpRFact) # Stores K in PH matrix

    # Compute state update
    ekf.ixp += 1
    @views mul!(ekf.xhats[ekf.ixp, :], ekf.PH[:,1:numMeas], ekf.rs[ekf.k, 1:numMeas])
    ekf.xhats[ekf.ixp, :] .+= ekf.xhats[ekf.ixp - 1, :]

    # Set time in storage vector
    ekf.txp[ekf.ixp] = t

    # Compute error
    (ytrue, u) = GetStateAndControl(ekf.scSim, ekf.txp[ekf.ixp])
    @views ekf.es[ekf.ixp, :] .= ekf.xhats[ekf.ixp, :] .- ytrue[1:7]

    # Compute covariance update
    @views mul!(ekf.KH, ekf.PH[:,1:numMeas], ekf.H[1:numMeas,:], -1.0, 0.0)
    for i in 1:7; ekf.KH[i,i] += 1.0; end
    mul!(ekf.P⁺, ekf.KH, ekf.P⁻)

    ekf.Ps[ekf.ixp, :] .= diag(ekf.P⁺)

    return nothing
end

function updateGPSIMU!(ekf::EKF)
    if ekf.updated == true
        throw(ErrorException("Cannot update before propagating."))
    else
        ekf.updated = true
    end

    # Get current time
    t = ekf.ts[ekf.k]

    # Get state estimate
    @views xhat = ekf.xhats[ekf.ixp, :]

    # Get true state
    (ytrue, u) = GetStateAndControl(ekf.scSim, t)

    # Generate gps measurements
    @views measGPS = gpsMeasurement(ekf.gpsSim, ytrue[1:3], t)

    # Get satelite ids
    @views sats = measGPS[2:2:end]
    numSats = length(sats)

    # Get expected gps measurements
    @views (expMeasGPS, ps) = gpsMeasurement(ekf.gpsSim, xhat[1:3], sats, t; type = :expected)

    # Compute residuals and measurement partials
    numRejected = 0
    rejectTol   = 1.0
    for i in 1:numSats
        r = measGPS[3 + 2*(i - 1)] - expMeasGPS[3 + 2*(i - 1)]
        if abs(r) < rejectTol
            # Compute residual while catching bad measurements
            @views ekf.rs[ekf.k, i - numRejected] = r

            # Compute measurement Jacobian
            ekf.H .*= 0.0
            @views ekf.H[i - numRejected, 1:3] .= -ps[i,1:3]
            @views ekf.H[i - numRejected, 1:3] ./= norm(ps[i,1:3])
        else
            numRejected += 1
        end
    end
    numSats -= numRejected

    # Generate accelerometer measurements
    (dy, _us) = GetStateAndControl(ekf.scSim, t; derivs = true)
    @views measAcc = ComputeMeasurement(ekf.imuSim, dy[4:6])

    # Compute expected accelerometer measurement
    @views ekf.y[1:7] .= xhat
    ekf.y[8:end] .= 0.0
    for i in 1:7; ekf.y[7 + 7*(i - 1) + i] = 1.0; end
    if ekf.lunaPerts == false
        ekfNoLunarPertEOM!(ekf.dy, ekf.y, ekf, t)
    else
        ekfWithLunarPertEOM!(ekf.dy, ekf.y, ekf, t)
    end
    @views expMeasAcc = ComputeMeasurement(ekf.imuSim, ekf.dy[4:6]; type = :expected)

    # Compute accelerometer residuals and partials
    ekf.rs[ekf.k, numSats + 1:numSats + 3]   .= measAcc .- expMeasAcc
    @views ekf.H[numSats + 1:numSats + 3, 1:3] .= ekf.F[4:6, 1:3]
    @views ekf.H[numSats + 1:numSats + 3, 7] .= ekf.F[4:6, 7]

    # Compute gain
    numMeas = numSats + 3
    @views mul!(ekf.PH[:,1:numMeas], ekf.P⁻, transpose(ekf.H[1:numMeas, :]))
    @views mul!(ekf.HPH[1:numMeas,1:numMeas], ekf.H[1:numMeas, :], ekf.PH[:,1:numMeas])
    for i in 1:numSats; ekf.HPH[i,i] += ekf.σ2GPS; end
    for i in numSats + 1:numMeas; ekf.HPH[i,i] += ekf.σ2Acc; end
    @views HPHpRFact = factorize(ekf.HPH[1:numMeas,1:numMeas])
    @views rdiv!(ekf.PH[:,1:numMeas], HPHpRFact) # Stores K in PH matrix

    # Compute state update
    ekf.ixp += 1
    @views mul!(ekf.xhats[ekf.ixp, :], ekf.PH[:,1:numMeas], ekf.rs[ekf.k, 1:numMeas])
    ekf.xhats[ekf.ixp, :] .+= ekf.xhats[ekf.ixp - 1, :]

    # Set time in storage vector
    ekf.txp[ekf.ixp] = t

    # Compute error
    (ytrue, u) = GetStateAndControl(ekf.scSim, ekf.txp[ekf.ixp])
    @views ekf.es[ekf.ixp, :] .= ekf.xhats[ekf.ixp, :] .- ytrue[1:7]

    # Compute covariance update
    @views mul!(ekf.KH, ekf.PH[:,1:numMeas], ekf.H[1:numMeas,:], -1.0, 0.0)
    for i in 1:7; ekf.KH[i,i] += 1.0; end
    mul!(ekf.P⁺, ekf.KH, ekf.P⁻)

    ekf.Ps[ekf.ixp, :] .= diag(ekf.P⁺)

    return nothing
end

function runFilter!(ekf::EKF)
    for i in 1:length(ekf.ts)
        propagate!(ekf)
        if i*ekf.imuΔt % ekf.gpsΔt == 0
            updateGPSIMU!(ekf)
        else
            updateIMU!(ekf)
        end
    end
end

function ekfNoLunarPertEOM!(dy, y, ekf::EKF, t)
    # Zero F
    ekf.F .= 0.0

    # Earth gravitational parameter
    μ = 3.986004418e5 # [km^3/s^2]

    # Thrust magnitued
    T = 10e-3 # [kN]

    # Get control
    u = GetControl(ekf.scSim, t)

    # States
    @views r        = norm(y[1:3])
    @views dy[1:3] .= y[4:6]
    @views dy[4:6] .= -(μ/r^3).*y[1:3] .+ (T*u[1] / y[7]).*u[2:4]
    dy[7]           = -u[1]*T / ekf.scSim.ps.sp.c

    # Dynamics Jacobian
    # dr/dv
    ekf.F[1,4] = 1.0
    ekf.F[2,5] = 1.0
    ekf.F[3,6] = 1.0

    # dv/dr
    @views mul!(ekf.F[4:6, 1:3], y[1:3], transpose(y[1:3]))
    @views ekf.F[4:6, 1:3] .*= (3.0 / r^2)
    for i in 1:3; ekf.F[3 + i, i] -= 1.0; end
    @views ekf.F[4:6, 1:3] .*= (μ/r^3)

    # dv/dm
    @views ekf.F[4:6, 7] .= -(u[1]*T / y[7]^2) .* u[2:4]

    # Covariance Dynamics
    @views P   = reshape(y[8:end], (7, 7))
    @views dP  = reshape(dy[8:end], (7, 7))
    mul!(dP, ekf.F, P)
    mul!(dP, P, transpose(ekf.F), 1.0, 1.0)
    dP .+= ekf.Q
end

function ekfWithLunarPertEOM!(dy, y, ekf::EKF, t)
    # Zero F
    ekf.F .= 0.0

    # Earth gravitational parameter
    μ = 5.9742e24*6.673e-20

    # Luna gravitational parameter
    μl = 7.3483e22*6.673e-20

    # Thrust magnitued
    T = 10e-3 # [kN]

    # Get control
    (u, rLuna) = GetControlAndLunaPos(ekf.scSim, t)

    # Luna gravitational perturbation
    rls         = @SVector [rLuna[1] - y[1], rLuna[2] - y[2], rLuna[3] - y[3]]
    μldRls3     = μl / (norm(rls)^3)
    μldRluna3   = μl / (norm(rLuna)^3)
    al          = @SVector [rls[1]*μldRls3 - rLuna[1]*μldRluna3,
                            rls[2]*μldRls3 - rLuna[2]*μldRluna3,
                            rls[3]*μldRls3 - rLuna[3]*μldRluna3]
    # States
    @views r        = norm(y[1:3])
    @views dy[1:3] .= y[4:6]
    @views dy[4:6] .= -(μ/r^3).*y[1:3] .+ al .+ (T*u[1] / y[7]).*u[2:4]
    dy[7]           = -u[1]*T / ekf.scSim.ps.sp.c

    # Dynamics Jacobian
    # dr/dv
    ekf.F[1,4] = 1.0
    ekf.F[2,5] = 1.0
    ekf.F[3,6] = 1.0

    # dv/dr keplarian
    @views mul!(ekf.F[4:6, 1:3], y[1:3], transpose(y[1:3]))
    @views ekf.F[4:6, 1:3] .*= (3.0 / r^2)
    for i in 1:3; ekf.F[3 + i, i] -= 1.0; end
    @views ekf.F[4:6, 1:3] .*= (μ/r^3)

    # dv/dr luna perturbation (stealing memory from ekf.H)
    @views mul!(ekf.H[1:3, 1:3], rls, transpose(rls))
    @views ekf.H[1:3, 1:3] .*= (3.0 / norm(rls)^2)
    for i in 1:3; ekf.H[i,i] -= 1.0; end
    @views ekf.H[1:3, 1:3] .*= (μl / norm(rls)^3)

    # Combine dv/dr partials
    @views ekf.F[4:6, 1:3] .+= ekf.H[1:3, 1:3]

    # dv/dm
    @views ekf.F[4:6, 7] .= -(u[1]*T / y[7]^2) .* u[2:4]

    # Covariance Dynamics
    @views P   = reshape(y[8:end], (7, 7))
    @views dP  = reshape(dy[8:end], (7, 7))
    mul!(dP, ekf.F, P)
    mul!(dP, P, transpose(ekf.F), 1.0, 1.0)
    dP .+= ekf.Q
end

function plot(ekf::EKF, xtrue, n)
    ts   = ekf.txp[1:n]
    xhat = ekf.xhats[1:n, :]
    xt   = xtrue[1:n, :]
    es   = ekf.es[1:n, :]
    σ311 = 3*sqrt.(ekf.Ps[1:n,1])
    σ322 = 3*sqrt.(ekf.Ps[1:n,2])
    σ333 = 3*sqrt.(ekf.Ps[1:n,3])
    σ344 = 3*sqrt.(ekf.Ps[1:n,4])
    σ355 = 3*sqrt.(ekf.Ps[1:n,5])
    σ366 = 3*sqrt.(ekf.Ps[1:n,6])
    σ377 = 3*sqrt.(ekf.Ps[1:n,7])

    mat"""
    ts      = $ts;
    ts      = (ts - $(ekf.scSim.t0))/86400;
    xhat    = $xhat;
    xt      = $xt;
    es      = $es;
    s311    = $σ311;
    s322    = $σ322;
    s333    = $σ333;
    s344    = $σ344;
    s355    = $σ355;
    s366    = $σ366;
    s377    = $σ377;

    figure()
    subplot(3,1,1)
    plot(ts, es(:,1), 'k')
    hold on
    plot(ts, s311, 'r')
    plot(ts, -s311, 'r')
    legend('Estimation Error', 'EKF \$3\\sigma\$', 'Interpreter', 'latex')
    grid on

    subplot(3,1,2)
    plot(ts, es(:,1), 'k')
    hold on
    plot(ts, s322, 'r')
    plot(ts, -s322, 'r')
    grid on

    subplot(3,1,3)
    plot(ts, es(:,3), 'k')
    hold on
    plot(ts, s333, 'r')
    plot(ts, -s333, 'r')
    xlabel("Time, days")
    grid on

    figure()
    subplot(3,1,1)
    plot(ts, es(:,4), 'k')
    hold on
    plot(ts, s344, 'r')
    plot(ts, -s344, 'r')

    subplot(3,1,2)
    plot(ts, es(:,5), 'k')
    hold on
    plot(ts, s355, 'r')
    plot(ts, -s355, 'r')

    subplot(3,1,3)
    plot(ts, es(:,6), 'k')
    hold on
    plot(ts, s366, 'r')
    plot(ts, -s366, 'r')

    figure()
    plot(ts, es(:,7), 'k')
    hold on
    plot(ts, s377, 'r')
    plot(ts, -s377, 'r')

    figure()
    plot3(xhat(:,1),xhat(:,2), xhat(:,3), 'r')
    hold on
    plot3(xt(:,1),xt(:,2),xt(:,3), 'b')
    grid on
    axis equal
    """
end

function plotMC(ekf::EKF, em, ev, n)
    ts      = ekf.txp[1:n]
    xhat    = ekf.xhats[1:n, :]
    es      = ekf.es[1:n, :]
    σ311    = 3*sqrt.(ekf.Ps[1:n,1])
    σ311mc  = 3*sqrt.(ev[:,1])
    σ322    = 3*sqrt.(ekf.Ps[1:n,2])
    σ322mc  = 3*sqrt.(ev[:,2])
    σ333    = 3*sqrt.(ekf.Ps[1:n,3])
    σ333mc  = 3*sqrt.(ev[:,3])
    σ344    = 3*sqrt.(ekf.Ps[1:n,4])
    σ344mc  = 3*sqrt.(ev[:,4])
    σ355    = 3*sqrt.(ekf.Ps[1:n,5])
    σ355mc  = 3*sqrt.(ev[:,5])
    σ366    = 3*sqrt.(ekf.Ps[1:n,6])
    σ366mc  = 3*sqrt.(ev[:,6])
    σ377    = 3*sqrt.(ekf.Ps[1:n,7])
    σ377mc  = 3*sqrt.(ev[:,7])

    mat"""
    ts      = $ts;
    ts      = (ts - $(ekf.scSim.t0))/86400;
    xhat    = $xhat;
    es      = $es;
    s311    = $σ311;
    s311mc  = $σ311mc;
    s322    = $σ322;
    s322mc  = $σ322mc;
    s333    = $σ333;
    s333mc  = $σ333mc;
    s344    = $σ344;
    s344mc  = $σ344mc;
    s355    = $σ355;
    s355mc  = $σ355mc;
    s366    = $σ366;
    s366mc  = $σ366mc;
    s377    = $σ377;
    s377mc  = $σ377mc;
    em      = $em;

    figure()
    subplot(3,1,1)
    plot(ts, es(:,1).*1000, 'k')
    hold on
    plot(ts, s311.*1000, 'r')
    plot(ts, -s311.*1000, 'r')
    plot(ts, s311mc.*1000, '--r')
    plot(ts, -s311mc.*1000, '--r')
    legend('Estimation Error', 'EKF \$3\\sigma\$', 'MC \$3\\sigma\$','Interpreter', 'latex')
    ylabel('\$e_{r_x}\$, m', 'Interpreter', 'latex')

%    ylim([-1000, 1000])
    grid on

    subplot(3,1,2)
    plot(ts, es(:,1).*1000, 'k')
    hold on
    plot(ts, s322.*1000, 'r')
    plot(ts, -s322.*1000, 'r')
    plot(ts, s322mc.*1000, '--r')
    plot(ts, -s322mc.*1000, '--r')
    ylabel('\$e_{r_y}\$, m', 'Interpreter', 'latex')

%    ylim([-1000, 1000])
    grid on

    subplot(3,1,3)
    plot(ts, es(:,3).*1000, 'k')
    hold on
    plot(ts, s333.*1000, 'r')
    plot(ts, -s333.*1000, 'r')
    plot(ts, s333mc.*1000, '--r')
    plot(ts, -s333mc.*1000, '--r')
    xlabel('Time, days')
    ylabel('\$e_{r_z}\$, m', 'Interpreter', 'latex')

    grid on

    figure()
    subplot(3,1,1)
    plot(ts, es(:,4).*1000, 'k')
    hold on
    plot(ts, s344.*1000, 'r')
    plot(ts, -s344.*1000, 'r')
    plot(ts, s344mc.*1000, '--r')
    plot(ts, -s344mc.*1000, '--r')
    legend('Estimation Error', 'UKF \$3\\sigma\$', 'MC \$3\\sigma\$','Interpreter', 'latex')
    ylabel('\$e_{v_x}\$, m/s', 'Interpreter', 'latex')

    ylim([-1, 1])
    grid on

    subplot(3,1,2)
    plot(ts, es(:,5).*1000, 'k')
    hold on
    plot(ts, s355.*1000, 'r')
    plot(ts, -s355.*1000, 'r')
    plot(ts, s355mc.*1000, '--r')
    plot(ts, -s355mc.*1000, '--r')
    ylabel('\$e_{v_y}\$, m/s', 'Interpreter', 'latex')

    ylim([-1.5, 1.5])
    grid on

    subplot(3,1,3)
    plot(ts, es(:,6).*1000, 'k')
    hold on
    plot(ts, s366.*1000, 'r')
    plot(ts, -s366.*1000, 'r')
    plot(ts, s366mc.*1000, '--r')
    plot(ts, -s366mc.*1000, '--r')
    xlabel('Time, days')
    ylabel('\$e_{v_z}\$, m/s', 'Interpreter', 'latex')

    ylim([-1, 1])
    grid on

    figure()
    plot(ts, es(:,7), 'k')
    hold on
    plot(ts, s377, 'r')
    plot(ts, -s377, 'r')
    plot(ts, s377mc, '--r')
    plot(ts, -s377mc, '--r')
    xlabel('Time, days')
    ylabel('\$e_m\$, kg', 'Interpreter', 'latex')
    legend('Estimation Error', 'UKF \$3\\sigma\$', 'MC \$3\\sigma\$','Interpreter', 'latex')

    grid on

    figure()
    histogram(em(:,1).*1000, 20)
    hold on
    xline(mean(em(:,1))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{r_x}\$, m', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,2).*1000, 20)
    hold on
    xline(mean(em(:,2))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{r_y}\$, m', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,3).*1000, 20)
    hold on
    xline(mean(em(:,3))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{r_z}\$, m', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,4).*1000, 20)
    hold on
    xline(mean(em(:,4))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{v_x}\$, m/s', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,5).*1000, 20)
    hold on
    xline(mean(em(:,5))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{v_y}\$, m/s', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,6).*1000, 30)
    hold on
    xline(mean(em(:,6))*1000, '--k', 'LineWidth', 3)
    xlabel('\$e_{v_z}\$, m/s', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    figure()
    histogram(em(:,7), 20)
    hold on
    xline(mean(em(:,7)), '--k', 'LineWidth', 3)
    xlabel('\$e_m\$, m/s', 'Interpreter', 'latex')
    ylabel('# Occurances')
    grid on

    """
end

function saveMCmat(ekf::EKF, em, ev, n)
    ts      = ekf.txp[1:n]
    xhat    = ekf.xhats[1:n, :]
    es      = ekf.es[1:n, :]
    σ311    = 3*sqrt.(ekf.Ps[1:n,1])
    σ311mc  = 3*sqrt.(ev[:,1])
    σ322    = 3*sqrt.(ekf.Ps[1:n,2])
    σ322mc  = 3*sqrt.(ev[:,2])
    σ333    = 3*sqrt.(ekf.Ps[1:n,3])
    σ333mc  = 3*sqrt.(ev[:,3])
    σ344    = 3*sqrt.(ekf.Ps[1:n,4])
    σ344mc  = 3*sqrt.(ev[:,4])
    σ355    = 3*sqrt.(ekf.Ps[1:n,5])
    σ355mc  = 3*sqrt.(ev[:,5])
    σ366    = 3*sqrt.(ekf.Ps[1:n,6])
    σ366mc  = 3*sqrt.(ev[:,6])
    σ377    = 3*sqrt.(ekf.Ps[1:n,7])
    σ377mc  = 3*sqrt.(ev[:,7])

    fname   = datadir("MATLABdata","ekf.mat")

    mat"""
    ts      = $ts;
    ts      = (ts - $(ekf.scSim.t0))/86400;
    xhat    = $xhat;
    es      = $es;
    s311    = $σ311;
    s311mc  = $σ311mc;
    s322    = $σ322;
    s322mc  = $σ322mc;
    s333    = $σ333;
    s333mc  = $σ333mc;
    s344    = $σ344;
    s344mc  = $σ344mc;
    s355    = $σ355;
    s355mc  = $σ355mc;
    s366    = $σ366;
    s366mc  = $σ366mc;
    s377    = $σ377;
    s377mc  = $σ377mc;
    em      = $em;

    save($fname)
    """

    return nothing
end

max_stars_repo_name: GrantHecht/OptimalEstimationProject


text:
module TestICModel

using Test
using MLJIteration
using MLJBase
using IterationControl
using ..DummyModel

X, y = make_dummy(N=20)

@testset "IterationControl overloadings" begin

    model=DummyIterativeModel(n=10)

    # evaluate after 10 iterations by hand:
    mach = machine(model, X, y)
    e = evaluate!(mach,
                 resampling=Holdout(fraction_train=0.5),
                  measure=l2).measurement[1]

    losses = training_losses(mach)

    # evaluate after 10 iterations using IterationControl, and in two stages:
    resampling_machine =
        machine(Resampler(model=model,
                          resampling=Holdout(fraction_train=0.5),
                          measure=l2),
                X,
                y)

    ic_model = MLJIteration.ICModel(resampling_machine, :n, 0)
    IterationControl.train!(ic_model, 4)
    IterationControl.train!(ic_model, 6)

    # compare:
    @test IterationControl.loss(ic_model) ≈ e

    # test `Δi` is remembered:
    @test ic_model.Δi[] == 6

    # test training_losses:
    @test IterationControl.training_losses(ic_model) == losses[end - 5:end]

    # test expose:
    @test IterationControl.expose(ic_model) ==
        fitted_params(ic_model.mach).machine
    @test IterationControl.expose(MLJIteration.ICModel(mach, :n, 0)) == mach

    # test methods in case applying when `resampling == nothing`:
    model=DummyIterativeModel(n=0)
    mach = machine(model, X, y)
    ic_model2 = MLJIteration.ICModel(mach, :n, 0)
    IterationControl.train!(ic_model2, 10)
    losses2 = training_losses(mach)
    @test losses2 ≈ IterationControl.training_losses(ic_model2)
    @test losses2[end] == IterationControl.loss(ic_model2)
end

mutable struct FooBar <: MLJBase.Deterministic
    n::Int
end
iteration_parameter(::Type{<:FooBar}) = :n
MLJBase.fit(::FooBar, ::Any, data...) = nothing, nothing, nothing
MLJBase.predict(::FooBar, ::Any, Xnew) = ones(length(Xnew))

@testset "IterationControl.training_losses: error for unsupported models" begin
    resampling_machine =
        machine(Resampler(model=FooBar(10),
                          resampling=Holdout(fraction_train=0.5),
                          measure=l2),
                X,
                y)
    fit!(resampling_machine, verbosity=0)
    ic_model = MLJIteration.ICModel(resampling_machine, :n, 0)
    @test IterationControl.training_losses(ic_model) === nothing

    mach = machine(FooBar(10), X, y)
    fit!(mach, verbosity=0)
    ic_model = MLJIteration.ICModel(mach, :n, 0)
    @test IterationControl.training_losses(ic_model) === nothing
    @test IterationControl.loss(ic_model) === nothing
end

@testset "ICModel interface for users" begin
    model=DummyIterativeModel(n=0)
    resampling_machine =
        machine(Resampler(model=model,
                          resampling=Holdout(fraction_train=0.5),
                          measure=l2),
                X,
                y)

    ic_model = MLJIteration.ICModel(resampling_machine, :n, 0)
    IterationControl.train!(ic_model, 4)
    IterationControl.train!(ic_model, 6)
    @test ic_model.machine == fitted_params(resampling_machine).machine
    @test ic_model.machine.model == DummyIterativeModel(n=10)
    @test ic_model.n_cycles == 2
    @test ic_model.n_iterations == 10
    @test ic_model.Δiterations == 6
    @test ic_model.loss == IterationControl.loss(ic_model)
    @test ic_model.training_losses == IterationControl.training_losses(ic_model)
end

end

true


text:
#=

  Euler #30 in Julia.

  Problem 30
  """
  Surprisingly there are only three numbers that can be written
  as the sum of fourth powers of their digits:

     1634 = 1^(4) + 6^(4) + 3^(4) + 4^(4)
     8208 = 8^(4) + 2^(4) + 0^(4) + 8^(4)
     9474 = 9^(4) + 4^(4) + 7^(4) + 4^(4)

  As 1 = 1^(4) is not a sum it is not included.

  The sum of these numbers is 1634 + 8208 + 9474 = 19316.

  Find the sum of all the numbers that can be written as the sum of
  fifth powers of their digits.
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

# 0.07169755ss
function euler30a()
    t = 0
    m = 5
    for n in 10:6*(9^5)
        # nn = (split(string(n),"").|>i->parse(Int,i)^m)|>sum
        # nn = (digits(n).|>i->i^m)|>sum
        nn = sum(digits(n).|>i->i^m)
        if n === nn
            t += n
        end
    end
    return t
end

# 0.15361390s
function euler30b()
    m = 5;
    # s = (10:6*9^5).|>n-> n*(n==(split(string(n),"").|>i->parse(Int,i)^m)|>sum)
    s = (10:6*9^5).|>n-> n*(n==(digits(n).|>i->i^m)|>sum)
    return s # |>sum
end

function euler30c()
    m = 5;
    sum([n*(n==(digits(n).|>i->i^m)) for n in 10:6*9^5])
end

run_euler(euler30a)
# run_euler(euler30b)
# run_euler(euler30c)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
using DrWatson
@quickactivate
using ArgParse
using GenerativeAD
using StatsBase: fit!, predict, sample
using BSON

s = ArgParseSettings()
@add_arg_table! s begin
	"max_seed"
		default = 1
		arg_type = Int
		help = "maximum number of seeds to run through"
	"dataset"
		default = "iris"
		arg_type = String
		help = "dataset"
    "contamination"
    	arg_type = Float64
    	help = "contamination rate of training data"
    	default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, max_seed, contamination = parsed_args

modelname = "RealNVP_005_010"

function sample_params()
	parameter_rng = (
		nflows 		= 2 .^ (1:3),
		hdim 		= 2 .^(4:10),
		nlayers 	= 2:3,
		lr 			= [1f-4],
		batchsize 	= 2 .^ (5:7),
		act_loc		= ["relu", "tanh"],
		act_scl		= ["relu", "tanh"],
		bn 			= [true, false],
		wreg 		= [0.0f0, 1f-5, 1f-6],
		init_seed 	= 1:Int(1e8),
		init_I 		= [true, false],
		tanhscaling = [true, false]
	)

	(;zip(keys(parameter_rng), map(x->sample(x, 1)[1], parameter_rng))...)
end

function fit(data, parameters)
	model = GenerativeAD.Models.RealNVPFlow(;idim=size(data[1][1], 1), parameters...)

	try
		global info, fit_t, _, _, _ = @timed fit!(model, data; max_train_time=82800/max_seed,
										patience=200, check_interval=10, quantile=(0.05, 0.1),parameters...)
	catch e
		@info "Failed training due to \n$e"
		return (fit_t = NaN, history=nothing, npars=nothing, model=nothing), []
	end

	training_info = (
		fit_t = fit_t,
		history = info.history,
		niter = info.niter,
		npars = info.npars,
		model = info.model
		)

	training_info, [(x -> predict(info.model, x), parameters)]
end

try_counter = 0
max_tries = 10*max_seed
cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
while try_counter < max_tries
    parameters = sample_params()

    for seed in 1:max_seed
		savepath = datadir("experiments/tabular$cont_string/$(modelname)/$(dataset)/seed=$(seed)")
		mkpath(savepath)

		# get data
		data = GenerativeAD.load_data(dataset, seed=seed, contamination=contamination)
		edited_parameters = GenerativeAD.edit_params(data, parameters)

		if GenerativeAD.check_params(savepath, edited_parameters)
			@info "Started training $(modelname)$(edited_parameters) on $(dataset):$(seed)"
			@info "Train/valdiation/test splits: $(size(data[1][1], 2)) | $(size(data[2][1], 2)) | $(size(data[2][1], 2))"
			@info "Number of features: $(size(data[1][1], 1))"

			training_info, results = fit(data, edited_parameters)

			if training_info.model != nothing
				tagsave(joinpath(savepath, savename("model", edited_parameters, "bson", digits=5)),
						Dict("model"=>training_info.model,
							"fit_t"=>training_info.fit_t,
							"history"=>training_info.history,
							"parameters"=>edited_parameters
							), safe = true)
				training_info = merge(training_info, (model = nothing,))
			end
			save_entries = merge(training_info, (modelname = modelname, seed = seed, dataset = dataset, contamination = contamination))

			for result in results
				GenerativeAD.experiment(result..., data, savepath; save_entries...)
			end
			global try_counter = max_tries + 1
		else
			@info "Model already present, trying new hyperparameters..."
			global try_counter += 1
		end
	end
end
(try_counter == max_tries) ? (@info "Reached $(max_tries) tries, giving up.") : nothing


text:
using SunlightAPIs
using Base.Test

function sunlight_test(fn, args...; auth = "", named_args...)
    print("Testing $fn ... ")

    # test different ways of passing in api key
    auth_first = fn(auth, args...)
    auth_named = fn(args...; auth = auth)
    @test auth_first != nothing
    @test auth_named != nothing
    @test auth_first == auth_named

    # test each of the named args
    for arg in named_args
        w_arg = fn(auth, args...; [arg]...)
        @test w_arg != nothing
        @test w_arg != auth_first || isempty(w_arg) # if thre's no data we can use to check, let's assume the named arg is working
    end

    # test a silly named arg that should error
    @test_throws ErrorException fn(auth, args...; this_arg_doesnt_work = "I hope...")

    println("Success!")
end

# -------

function runtests(api_key; pol_id = Pol("597e02e7d1b04d83976913da1b8e2998"),
                           ind_id = Ind("a03df5d9b20e467fa0ceaefa94c4491e"),
                           org_id = Org("52a1620b2ff543ebb74718fbff742529"),
                           industry_id = "165d820dd48441e1befdc47f3fa3d236",
                           limit = 1,
                           cycle = 2012)

    sunlight_test(top_contributors, pol_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(top_industries, pol_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(unknown_industries, pol_id; cycle = cycle, auth = api_key)
    sunlight_test(top_sectors, pol_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(local_breakdown, pol_id; cycle = cycle, auth = api_key)
    sunlight_test(contributor_breakdown, pol_id; cycle = cycle, auth = api_key)
    sunlight_test(fec_summary, pol_id; auth = api_key)
    sunlight_test(fec_independent_expenditures, pol_id; auth = api_key)
    sunlight_test(top_recipient_organizations, ind_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(top_recipient_politicians, ind_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(party_breakdown, ind_id; cycle = cycle, auth = api_key)
    sunlight_test(lobbying_registrants, ind_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(lobbying_clients, ind_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(lobbying_issues, ind_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(top_recipients, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(pac_recipients, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(party_breakdown, org_id; cycle = cycle, auth = api_key)
    sunlight_test(state_federal_breakdown, org_id; cycle = cycle, auth = api_key)
    sunlight_test(lobbing_registrants, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(lobbying_issues, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(bills, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(lobbyists, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(registrant_clients, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(registrant_issues, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(registrant_bills, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(registrant_lobbyists, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(mentions_in_regulations, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(regulatory_comment_submissions, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(faca_memberships, org_id; limit = limit, cycle = cycle, auth = api_key)
    sunlight_test(fec_summary, org_id; auth = api_key)
    sunlight_test(top_industry_organizations, industry_id; limit = limit, cycle = cycle, auth = api_key)

end

# -------
# Run these tests automatically when an api-key is passed as the first arg on startup

!isempty(ARGS) && runtests(ARGS[1])


text:
# ideal gas constant, from unitful
default_units(property::typeof(molar_mass)) = u"kg/mol"
default_units(property::typeof(mol_helmholtz)) = u"J/mol"
default_units(property::typeof(mass_helmholtz)) = u"J/kg"
default_units(property::typeof(total_helmholtz)) = u"J"

default_units(property::typeof(mol_gibbs)) = u"J/mol"
default_units(property::typeof(mass_gibbs)) = u"J/kg"
default_units(property::typeof(total_gibbs)) = u"J"

default_units(property::typeof(mol_internal_energy)) = u"J/mol"
default_units(property::typeof(mass_internal_energy)) = u"J/kg"
default_units(property::typeof(total_internal_energy)) = u"J"

default_units(property::typeof(mol_enthalpy)) = u"J/mol"
default_units(property::typeof(mass_enthalpy)) = u"J/kg"
default_units(property::typeof(total_enthalpy)) = u"J"

default_units(property::typeof(mol_entropy)) = u"J/(mol*K)"
default_units(property::typeof(mass_entropy)) = u"J/(kg*K)"
default_units(property::typeof(total_entropy)) = u"J/(K)"

default_units(property::typeof(mol_volume)) = u"(m^3)/mol"
default_units(property::typeof(mass_volume)) = u"(m^3)/kg"
default_units(property::typeof(total_volume)) = u"(m^3)"

default_units(property::typeof(mol_density)) = u"mol/(m^3)"
default_units(property::typeof(mass_density)) = u"kg/(m^3)"

default_units(property::typeof(moles)) = u"mol"
default_units(property::typeof(mass)) = u"kg"
default_units(property::typeof(temperature)) = u"K"
default_units(property::typeof(pressure)) = u"Pa"

default_units(property::typeof(mol_fraction)) = Unitful.NoUnits
default_units(property::typeof(mol_number)) = u"mol"
default_units(property::typeof(mass_fraction)) = Unitful.NoUnits
default_units(property::typeof(mass_number)) = u"kg"

function to_spec end

to_spec(property::typeof(mol_helmholtz)) = Helmholtz{MOLAR}()
to_spec(property::typeof(mass_helmholtz)) = Helmholtz{MASS}()
to_spec(property::typeof(total_helmholtz)) = Helmholtz{TOTAL}()

to_spec(property::typeof(mol_gibbs)) = Gibbs{MOLAR}()
to_spec(property::typeof(mass_gibbs)) =Gibbs{MASS}()
to_spec(property::typeof(total_gibbs)) = Gibbs{TOTAL}()

to_spec(property::typeof(mol_internal_energy)) =InternalEnergy{MOLAR}()
to_spec(property::typeof(mass_internal_energy)) = InternalEnergy{MASS}()
to_spec(property::typeof(total_internal_energy)) = InternalEnergy{TOTAL}()

to_spec(property::typeof(mol_enthalpy)) = Enthalpy{MOLAR}()
to_spec(property::typeof(mass_enthalpy)) = Enthalpy{MASS}()
to_spec(property::typeof(total_enthalpy)) = Enthalpy{TOTAL}()

to_spec(property::typeof(mol_entropy)) = Entropy{MOLAR}()
to_spec(property::typeof(mass_entropy)) = Entropy{MASS}()
to_spec(property::typeof(total_entropy)) = Entropy{TOTAL}()

to_spec(property::typeof(mol_volume)) = VolumeAmount{MOLAR,VOLUME}()
to_spec(property::typeof(mass_volume)) = VolumeAmount{MASS,VOLUME}()
to_spec(property::typeof(total_volume)) = VolumeAmount{TOTAL,VOLUME}()

to_spec(property::typeof(mol_density)) = VolumeAmount{MOLAR,VOLUME}()
to_spec(property::typeof(mass_density)) = VolumeAmount{MASS,DENSITY}()

to_spec(property::typeof(moles)) = MaterialAmount{MOLAR}()
to_spec(property::typeof(mass)) = MaterialAmount{MASS}()
to_spec(property::typeof(temperature)) = Temperature()
to_spec(property::typeof(pressure)) = Pressure()

to_spec(property::typeof(mol_fraction)) = MaterialCompounds{MOLAR,FRACTION}()
to_spec(property::typeof(mol_number)) = MaterialCompounds{MOLAR,TOTAL_AMOUNT}()
to_spec(property::typeof(mass_fraction)) = MaterialCompounds{MASS,FRACTION}()
to_spec(property::typeof(mass_number)) = MaterialCompounds{MASS,TOTAL_AMOUNT}()

#to_spec(property::typeof(sound_speed)) = u"m/s"
#to_spec(property::typeof(compressibility_factor)) = Unitful.NoUnits

#to_spec(property::typeof(molecular_weight)) = u"g/mol"
#to_spec(property::typeof(compounds_number)) = Unitful.NoUnits
#to_spec(property::typeof(covolumes)) = u"(m^3)/mol"

==#

"""
    @to_units f(x)

return `f(x)` with Unitful units.
the measurement unit is determined by `default_units(a::typeof(f))`

#example:
```julia-repl
julia> critical_pressure(water)
2.2064e7 Pa
```

"""
macro to_units(fx)
    f = fx.args[1]
    if length(fx.args) <=3
        return quote
            @__MODULE__().default_units($f)*$fx

        end |> esc
    else
        return quote
            $fx*$(fx.args[4])
        end |> esc
    end
end

macro to_spec(fx)
    f = fx.args[1]
    model = fx.args[2]
    specs = fx.args[3]
    if length(fx.args) <=3
        return quote
            begin
                pp =  @__MODULE__()
                if typeof($model) === FromSpecs
                    pp.get_spec(pp.Types.Temperature(),$specs) #change here
                end
            end
        end |> esc
    else
        return quote
            begin
                pp =  @__MODULE__()
                if typeof($model) === FromSpecs
                    pp.get_spec(pp.Types.Temperature(),$specs) #change here
                else
                    val = pp.@to_units $fx
                    pp.spec(Temperature(),val)
                end
            end
        end |> esc
    end
end

#generic transfommation to get spec from defined functions
function spec(f::F,val,normalize_units=true) where F<:Function
    return spec(to_spec(f)val,normalize_units)
end

#directly calling from constructed specs:

(x::Helmholtz{MOLAR})(args...)=mol_helmholtz(args...)
(x::Helmholtz{MASS})(args...)=mass_helmholtz(args...)
(x::Helmholtz{TOTAL})(args...)=total_helmholtz(args...)

(x::Gibbs{MOLAR})(args...)=mol_helmholtz(args...)
(x::Gibbs{MASS})(args...)=mass_helmholtz(args...)
(x::Gibbs{TOTAL})(args...)=total_helmholtz(args...)

(x::InternalEnergy{MOLAR})(args...)=mol_internal_energy(args...)
(x::InternalEnergy{MASS})(args...)=mass_internal_energy(args...)
(x::InternalEnergy{TOTAL})(args...)=total_internal_energy(args...)

(x::Enthalpy{MOLAR})(args...)=mol_enthalpy(args...)
(x::Enthalpy{MASS})(args...)=mass_enthalpy(args...)
(x::Enthalpy{TOTAL})(args...)=total_enthalpy(args...)

(x::Entropy{MOLAR})(args...)=mol_entropy(args...)
(x::Entropy{MASS})(args...)=mol_entropy(args...)
(x::Entropy{TOTAL})(args...)=mol_entropy(args...)

(x::VolumeAmount{MOLAR,VOLUME})(args...)=mol_volume(args...)
(x::VolumeAmount{MASS,VOLUME})(args...)=mass_volume(args...)
(x::VolumeAmount{TOTAL,VOLUME})(args...)=total_volume(args...)

(x::VolumeAmount{MOLAR,DENSITY})(args...)=mol_density(args...)
(x::VolumeAmount{MASS,DENSITY})(args...)=mass_density(args...)

(x::Temperature)(args...)=temperature(args...)
(x::Pressure)(args...)=pressure(args...)

(x::MaterialAmount{MOLAR})(args...)=moles(args...)
(x::MaterialAmount{MASS})(args...)=mass(args...)

(x::MaterialCompounds{MOLAR,FRACTION})(args...)=mol_fraction(args...)
(x::MaterialCompounds{MOLAR,TOTAL_AMOUNT})(args...)=mol_number(args...)
(x::MaterialCompounds{MASS,FRACTION})(args...)=mass_fraction(args...)
(x::MaterialCompounds{MASS,TOTAL_AMOUNT})(args...)=mass_number(args...)
(x::SingleComponent)(args...) = nothing

_units_tuple(st) = (default_units(typeof(st[1])),
default_units(typeof(st[2])),
default_units(typeof(st[3])))

function properties(st::T,
        	state::ThermodynamicState,
            mw = nothing,
            units=_units_tuple(st)) where T <: Tuple{AbstractSpec,AbstractSpec,AbstractSpec}

    f(val,unit) = val(FromState(),state,unit,mw)

    return (f(st[1],units[1]),f(st[2],units[2]),f(st[3],units[3]))
end


text:
using Test
using Pkg;
using REPLHistory

history(12)
using REPLHistory
ab = "ααα"
∇ϕ(x) = grad(x)
ab = "2343511ω" * "ααα"
history()
print(history())
using REPLHistory
print(history())
ab = "ααα"
∇ϕ(x) = grad(x)
ab = "ααα"
print(history())


text:
using Documenter, MonteCarlo

makedocs(
    modules = [MonteCarlo],
    doctest = false,
    sitename = "MonteCarlo.jl",
    pages = [
        "Introduction" => "index.md",
        "DQMC" => [
            "Introduction" => "DQMC/Introduction.md",
            "Model" => "DQMC/model.md",
            "Lattices" => "DQMC/lattice.md",
            "DQMC" => "DQMC/dqmc.md",
            "Configuration Recorder" => "DQMC/recorder.md",
            "Update Scheduler" => "DQMC/scheduler.md",
            "Measurements" => "DQMC/measurements.md",
        ],
        "Examples" => [
            "examples/triangular_Hubbard.md",
            "examples/ALF1.md",
            "examples/HofmannBergChowdhury.md"
        ],
        "Legacy" => [
            "Introduction" => "legacy/index.md",
            "Manual" => [
                "Getting started" => "legacy/manual/gettingstarted.md",
                "Showcase" => "legacy/manual/showcase.md",
            ],
            "Physical models" => [
                "Ising model" => "legacy/models/ising.md",
                "Attractive Hubbard model" => "legacy/models/hubbardattractive.md",
            ],
            "Monte Carlo flavors" => [
                "MC" => "legacy/flavors/mc.md",
                "DQMC" => "legacy/flavors/dqmc.md",
            ],
            "Lattices" => "legacy/lattices.md",
            "Customize" => "legacy/customize.md",
            "Interfaces" => [
                "MC" => "legacy/interfaces/MC.md",
                "DQMC" => "legacy/interfaces/DQMC.md",
            ],
            "General exports" => [
                "General" => "legacy/methods/general.md"
            ],
        ]
    ],
    # assets = ["assets/custom.css", "assets/custom.js"]
)

deploydocs(
    repo = "github.com/carstenbauer/MonteCarlo.jl.git",
    push_preview = true,
    # target = "site",
)


text:
using CSV
using DataFrames

df = DataFrame(CSV.File("files/mintsUTDData.csv"))
print(typeof(df.dateTime[1]))


text:
@testset "peirls phases" begin
    println("Checking Peirls phases")
    init!(mc)
    # check peirls phases result
    @test isapprox(filter(x->!isnan(x), mc.l.peirls[1,1]), filter(x->!isnan(x), load("data/O3.jld", "peirls11")))
    @test isapprox(filter(x->!isnan(x), mc.l.peirls[1,2]), filter(x->!isnan(x), load("data/O3.jld", "peirls12")))
    @test isapprox(filter(x->!isnan(x), mc.l.peirls[2,1]), filter(x->!isnan(x), load("data/O3.jld", "peirls21")))
    @test isapprox(filter(x->!isnan(x), mc.l.peirls[2,2]), filter(x->!isnan(x), load("data/O3.jld", "peirls22")))

    p = mc.p
    l = mc.l
    L = mc.l.L
    phis = Matrix{Array{Float64}}(undef,2,2) # colidx = flavor, rowidx = spin
    for f in 1:2
      for s in 1:2
        phis[s,f] = permutedims(reshape(l.peirls[s,f], (L,L,L,L)), [2,1,4,3])
        B = zeros(2,2) # colidx = flavor, rowidx = spin up,down
        B[1,1] = B[2,2] = 2 * pi / l.sites
        B[1,2] = B[2,1] = - 2 * pi / l.sites

        # Flux check
        fluxes_nn = zeros(ComplexF64, L, L)
        fluxes_nnn_1 = zeros(ComplexF64, L, L)
        fluxes_nnn_2 = zeros(ComplexF64, L, L)
        fluxes_nnn_3 = zeros(ComplexF64, L, L)
        fluxes_nnn_4 = zeros(ComplexF64, L, L)
        fluxes_nnnn = zeros(ComplexF64, L, L)
        fluxes_mixed_1 = zeros(ComplexF64, L, L)
        fluxes_mixed_2 = zeros(ComplexF64, L, L)

        for x in 1:L
            for y in 1:L
                xp = mod1((x + 1),L)
                yp = mod1((y + 1),L)
                ym = mod1((y - 1),L)
                xm = mod1((x - 1),L)
                xp2 = mod1((x + 2),L)
                yp2 = mod1((y + 2),L)

                #### Nearest neighbor square
                # up, right, down, left
                flux = phis[s,f][x, y, xp, y] + phis[s,f][xp, y, xp, yp] + phis[s,f][xp, yp, x,yp] + phis[s,f][x,yp, x,y]
                fluxes_nn[x,y] = exp(im * flux)

                #### Next nearest neighbor triangles
                # up, right, down-left - yoni
                flux = phis[s,f][x, y, xp, yp] + phis[s,f][xp, yp, x, yp] + phis[s,f][x, yp, x,y]
                fluxes_nnn_1[x,y] = exp(im * flux)

                # down, left, up-right
                flux = phis[s,f][x, y, xm, ym] + phis[s,f][xm, ym, x, ym] + phis[s,f][x, ym, x,y]
                fluxes_nnn_2[x,y] = exp(im * flux)

                # left, up, down-right - yoni
                flux = phis[s,f][x, y, xm, yp] + phis[s,f][xm, yp, xm, y] + phis[s,f][xm, y, x, y]
                fluxes_nnn_3[x,y] = exp(im * flux)

                # right, down, up-left
                flux = phis[s,f][x, y, xp, ym] + phis[s,f][xp, ym, xp, y] + phis[s,f][xp, y, x,y]
                fluxes_nnn_4[x,y] = exp(im * flux)

                #### Next next nearest neighbor square
                # 2up, 2right, 2down, 2left
                flux = phis[s,f][x, y, xp2, y] + phis[s,f][xp2, y, xp2, yp2] + phis[s,f][xp2, yp2, x,yp2] + phis[s,f][x,yp2, x,y]
                fluxes_nnnn[x,y] = exp(im * flux)

                #### Other
                # up-right, down-right, 2left
                flux = phis[s,f][x, y, xp2, y] + phis[s,f][xp2, y, xp, yp] + phis[s,f][xp, yp, x,y]
                fluxes_mixed_1[x,y] = exp(im * flux)

                # up-left, up-right, 2down
                flux = phis[s,f][x, y, x, yp2] + phis[s,f][x, yp2, xm, yp] + phis[s,f][xm, yp, x,y]
                fluxes_mixed_2[x,y] = exp(im * flux)
            end
        end
        println(s," ",f)
        @test all( isapprox.(fluxes_nn, exp(im * B[s,f])) )

        if mc.p.Nhoppings != "none" && mc.p.NNhoppings != "none"
          @test all( isapprox.(fluxes_nnn_1, exp(im * B[s,f] / 2)) )
          @test all( isapprox.(fluxes_nnn_2, exp(im * B[s,f] / 2)) )
          @test all( isapprox.(fluxes_nnn_3, exp(im * B[s,f] / 2 )) )
          @test all( isapprox.(fluxes_nnn_4, exp(im * B[s,f] / 2)) )
          @test all( isapprox.(fluxes_nnnn, exp(im * 4 * B[s,f])) )
          @test all( isapprox.(fluxes_mixed_1, exp(im *  B[s,f])) )
          @test all( isapprox.(fluxes_mixed_2, exp(im *  B[s,f])) )
        end
      end
    end
end

max_stars_repo_name: crstnbr/dqmc

max_issues_repo_name: crstnbr/dqmc

max_forks_repo_name: crstnbr/dqmc


text:
# - - parameters - - - - - - - - - - -
# Physical setting
system["L"] = (4, 4, 4, 4)
system["β"] = 1.9
system["NC"] = 2
system["Nthermalization"] = 0
system["Nsteps"] = 5
system["initial"] = "cold"
system["initialtrj"] = 1
system["update_method"] = "HMC"
system["Nwing"] = 1

# Physical setting(fermions)
system["quench"] = true
system["Dirac_operator"] = nothing
wilson["Clover_coefficient"] = 0
wilson["r"] = 1
wilson["hop"] = 0
staggered["Nf"] = 0
staggered["mass"] = 0
system["BoundaryCondition"] = [1, 1, 1, -1]

# System Control
system["log_dir"] = "./logs"
system["logfile"] = "HMC_L04040404_beta5.7_quenched_su2.txt"
system["saveU_dir"] = ""
system["saveU_format"] = nothing
system["saveU_every"] = 1
system["verboselevel"] = 2
system["randomseed"] = 111
measurement["measurement_basedir"] = "./measurements"
measurement["measurement_dir"] = "HMC_L04040404_beta5.7_quenched_su2"

# HMC related
md["Δτ"] = 0.06666666666666667
md["SextonWeingargten"] = false
md["N_SextonWeingargten"] = 2
md["MDsteps"] = 15
cg["eps"] = 1.0e-19
cg["MaxCGstep"] = 3000

# Action parameter for SLMC
actions["use_autogeneratedstaples"] = false
actions["couplingcoeff"] = Any[]
actions["couplinglist"] = Any[]
actions["coupling_loops"] = nothing

# Measurement set
measurement["measurement_methods"] = Dict[
  Dict{Any,Any}("methodname" => "Chiral_condensate",
    "eps" => 1.0e-19,
    "fermiontype" => "Staggered",
    "Nf" => 4,
    "mass" => 0.5,
    "measure_every" => 5,
    "MaxCGstep" => 3000
  ),
  Dict{Any,Any}("methodname" => "Polyakov_loop",
    "measure_every" => 1
  ),
  Dict{Any,Any}("methodname" => "Topological_charge",
    "Nflowsteps" => 4,
    "numflow" => 10,
    "measure_every" => 10,
    "eps_flow" => 0.01
  ),
  Dict{Any,Any}("methodname" => "Pion_correlator",
    "eps" => 1.0e-19,
    "fermiontype" => "Wilson",
    "r" => 1,
    "measure_every" => 5,
    "hop" => 0.141139,
    "MaxCGstep" => 3000
  ),
  Dict{Any,Any}("methodname" => "Plaquette",
    "measure_every" => 1
  )
]

# - - - - - - - - - - - - - - - - - - -


text:
# discovering binary CUDA dependencies

using CompilerSupportLibraries_jll
using LazyArtifacts
import Libdl

const dependency_lock = ReentrantLock()

# lazily initialize a Ref containing a path to a library.
# the arguments to this macro is the name of the ref, an expression to populate it
# (possibly returning `nothing` if the library wasn't found), and an optional initialization
# hook to be executed after successfully discovering the library and setting the ref.
macro initialize_ref(ref, ex, hook=:())
    quote
        ref = $(esc(ref))

        # test and test-and-set
        if !isassigned(ref)
            Base.@lock dependency_lock begin
                if !isassigned(ref)
                    val = $(esc(ex))
                    if val === nothing && !(eltype(ref) <: Union{Nothing,<:Any})
                        error($"Could not find a required library")
                    end
                    ref[] = val
                    if val !== nothing
                        $(esc(hook))
                    end
                end
            end
        end

        ref[]
    end
end

#
# CUDA toolkit
#

export toolkit

abstract type AbstractToolkit end

struct ArtifactToolkit <: AbstractToolkit
    release::VersionNumber
    artifact::String
end

struct LocalToolkit <: AbstractToolkit
    release::VersionNumber  # approximate, from the CUDA runtime library
    dirs::Vector{String}
end

const __toolkit = Ref{AbstractToolkit}()

function toolkit()
    @initialize_ref __toolkit begin
        toolkit = nothing

        # load the driver library before anything else, because we may want to load a
        # forward-compatible alternative instead, and dlopening libcudart as part of
        # the discovery process below will irreversibly load the system version.
        CUDA.libcuda()

        # CI runs in a well-defined environment, so prefer a local CUDA installation there
        if getenv("CI", false) && !haskey(ENV, "JULIA_CUDA_USE_BINARYBUILDER")
            toolkit = find_local_cuda()
        end

        if toolkit === nothing && getenv("JULIA_CUDA_USE_BINARYBUILDER", true)
            toolkit = find_artifact_cuda()
        end

        # if the user didn't specifically request an artifact version, look for a local installation
        if toolkit === nothing && !haskey(ENV, "JULIA_CUDA_VERSION")
            toolkit = find_local_cuda()
        end

        if toolkit === nothing
            error("Could not find a suitable CUDA installation")
        end

        toolkit
    end CUDA.__init_toolkit__()
    __toolkit[]::Union{ArtifactToolkit,LocalToolkit}
end

# workaround @artifact_str eagerness on unsupported platforms by passing a variable
function cuda_artifact(id, cuda::VersionNumber)
    platform = Base.BinaryPlatforms.HostPlatform()
    platform.tags["cuda"] = "$(cuda.major).$(cuda.minor)"

    dir = try
        @artifact_str(id, platform)
    catch ex
        @debug "Could not load artifact '$id' for CUDA $(cuda.release)" exception=(ex,catch_backtrace())
        return nothing
    end

    # sometimes artifact downloads fail (e.g. JuliaGPU/CUDA.jl#1003)
    if isempty(readdir(dir))
        error("""The artifact at $dir is empty.
                 This is probably caused by a failed download. Remove the directory and try again.""")
    end

    return dir
end

# NOTE: we don't use autogenerated JLLs, because we have multiple artifacts and need to
#       decide at run time (i.e. not via package dependencies) which one to use.
const cuda_toolkits = [
    (release=v"11.5", preferred=true),
    (release=v"11.4", preferred=true),
    (release=v"11.3", preferred=true),
    (release=v"11.2", preferred=true),
    (release=v"11.1", preferred=true),
    (release=v"11.0", preferred=true),
    (release=v"10.2", preferred=true),
    (release=v"10.1", preferred=true),
    (release=v"10.0", preferred=true),
    (release=v"9.0",  preferred=true),
    (release=v"9.2",  preferred=true),
]

function find_artifact_cuda()
    @debug "Trying to use artifacts..."

    # select compatible artifacts
    if haskey(ENV, "JULIA_CUDA_VERSION")
        wanted = VersionNumber(ENV["JULIA_CUDA_VERSION"])   # misnomer: actually the release
        @debug "Selecting artifacts based on requested $wanted"
        candidate_toolkits = filter(cuda_toolkits) do toolkit
            toolkit.release == wanted
        end
        isempty(candidate_toolkits) && @debug "Requested CUDA release $wanted is not provided by any artifact"
    else
        driver_release = CUDA.release()
        @debug "Selecting artifacts based on driver compatibility $driver_release"
        candidate_toolkits = filter(cuda_toolkits) do toolkit
            toolkit.preferred &&
                (toolkit.release <= driver_release ||
                 # CUDA 11: Enhanced Compatibility (aka. semver)
                 (driver_release >= v"11" &&
                  toolkit.release.major <= driver_release.major))
        end
        isempty(candidate_toolkits) && @debug "CUDA driver compatibility $driver_release is not compatible with any artifact"
    end

    # download and install
    artifact = nothing
    for cuda in sort(candidate_toolkits; rev=true, by=toolkit->toolkit.release)
        dir = cuda_artifact("CUDA", cuda.release)
        if dir !== nothing
            artifact = (release=cuda.release, dir)
            break
        end
    end
    if artifact === nothing
        @debug "Could not find a compatible artifact."
        return nothing
    end

    @debug "Using CUDA $(artifact.release) from an artifact at $(artifact.dir)"
    return ArtifactToolkit(artifact.release, artifact.dir)
end

function find_local_cuda()
    @debug "Trying to use local installation..."

    dirs = find_toolkit()

    let path = find_cuda_binary(dirs, "nvdisasm")
        if path === nothing
            @debug "Could not find nvdisasm"
            return nothing
        end
        __nvdisasm[] = path
    end

    cudart_versions = cuda_library_versions("cudart")
    let path = find_cuda_library(dirs, "cudart", cudart_versions)
        if path === nothing
            @debug "Could not find the CUDA runtime library"
            return nothing
        end
        __libcudart[] = path
    end
    release = CUDA.runtime_version()

    @debug "Found local CUDA $(release) at $(join(dirs, ", "))"
    return LocalToolkit(release, dirs)
end

## properties

export toolkit_origin, toolkit_release

"""
    toolkit_origin()

Returns the origin of the CUDA toolkit in use (either :artifact, or :local).
"""
toolkit_origin() = toolkit_origin(toolkit())::Symbol
toolkit_origin(::ArtifactToolkit) = :artifact
toolkit_origin(::LocalToolkit) = :local

"""
    toolkit_release()

Returns the CUDA release version of the toolkit in use.
"""
toolkit_release() = toolkit().release::VersionNumber

## binaries

export ptxas, nvlink, nvdisasm, compute_sanitizer, has_compute_sanitizer

# pxtas: used for compiling PTX to SASS
const __ptxas = Ref{String}()
function ptxas()
    @initialize_ref __ptxas begin
        find_binary(toolkit(), "ptxas")
    end
end

# nvlink: used for linking additional libraries
const __nvlink = Ref{String}()
function nvlink()
    @initialize_ref __nvlink begin
        find_binary(toolkit(), "nvlink")
    end
end

# nvdisasm: used for reflection (decompiling SASS code)
const __nvdisasm = Ref{String}()
function nvdisasm()
    @initialize_ref __nvdisasm begin
        find_binary(toolkit(), "nvdisasm")
    end
end

# compute-santizer: used by the test suite
const __compute_sanitizer = Ref{Union{Nothing,String}}()
function compute_sanitizer(; throw_error::Bool=true)
    path = @initialize_ref __compute_sanitizer begin
        if toolkit_release() < v"11.0"
            nothing
        else
            find_binary(toolkit(), "compute-sanitizer"; optional=true)
        end
    end
    if path === nothing && throw_error
        error("This functionality is unavailabe as compute_sanitizer is missing.")
    end
    path
end
has_compute_sanitizer() = compute_sanitizer(throw_error=false) !== nothing

function artifact_binary(artifact_dir, name)
    path = joinpath(artifact_dir, "bin", Sys.iswindows() ? "$name.exe" : name)
    if !ispath(path)
        error("""Could not find binary '$name' in $artifact_dir!
                 This is a bug; please file an issue with a verbose directory listing of $artifact_dir
                 If this directory is empty, delete it and try again.""")
    end
    return path
end

function find_binary(cuda::ArtifactToolkit, name; optional=false)
    # NOTE: optional is ignored, and we'll error if not found
    artifact_binary(cuda.artifact, name)
end

function find_binary(cuda::LocalToolkit, name; optional=false)
    path = find_cuda_binary(cuda.dirs, name)
    if path !== nothing
        return path
    else
        optional || error("Could not find binary '$name' in your local CUDA installation.")
        return nothing
    end
end

## libraries

# XXX: we don't correctly model the dependencies of these libraries, and hack around them
#      by loading libcublasLt before libcublas, or CUDNN's sublibraries before libcudnn.
#      this is necessary (even if the dependent libraries are in the same directory)
#      to avoid a local toolkit from messing with our artifacts (JuliaGPU/CUDA.jl#609).

export libcudart, libcublas, libcusparse, libcufft, libcurand, libcusolver,
       libcusolvermg, has_cusolvermg, libcupti, has_cupti, libnvtx, has_nvtx

const __libcudart = Ref{String}()
function libcudart()
    @initialize_ref __libcudart begin
        find_library(toolkit(), "cudart")
    end
end

const __libcublaslt = Ref{String}()
function libcublaslt()
    @initialize_ref __libcublaslt begin
        if toolkit_release() < v"10.1"
            nothing
        else
            find_library(toolkit(), "cublasLt")
        end
    end
end

const __libcublas = Ref{String}()
function libcublas()
    @initialize_ref __libcublas begin
        libcublaslt()
        find_library(toolkit(), "cublas")
    end CUDA.CUBLAS.__runtime_init__()
end

const __libcusparse = Ref{String}()
function libcusparse()
    @initialize_ref __libcusparse begin
        find_library(toolkit(), "cusparse")
    end
end

const __libcufft = Ref{String}()
function libcufft()
    @initialize_ref __libcufft begin
        find_library(toolkit(), "cufft")
    end
end

const __libcurand = Ref{String}()
function libcurand()
    @initialize_ref __libcurand begin
        find_library(toolkit(), "curand")
    end
end

const __libcusolver = Ref{String}()
function libcusolver()
    @initialize_ref __libcusolver begin
        find_library(toolkit(), "cusolver")
    end
end

const __libcusolverMg = Ref{Union{String,Nothing}}()
function libcusolvermg(; throw_error::Bool=true)
     path = @initialize_ref __libcusolverMg begin
        if toolkit_release() < v"10.1"
            nothing
        else
            find_library(toolkit(), "cusolverMg"; optional=true)
        end
    end
    if path === nothing && throw_error
        error("This functionality is unavailabe as cuSolverMg is missing.")
    end
    path
end
has_cusolvermg() = libcusolvermg(throw_error=false) !== nothing

const __libcupti = Ref{Union{String,Nothing}}()
function libcupti(; throw_error::Bool=true)
    path = @initialize_ref __libcupti begin
        find_library(toolkit(), "cupti"; optional=true)
    end
    if path === nothing && throw_error
        error("This functionality is unavailabe as CUPTI is missing.")
    end
    path
end
has_cupti() = libcupti(throw_error=false) !== nothing

const __libnvtx = Ref{Union{String,Nothing}}()
function libnvtx(; throw_error::Bool=true)
    path = @initialize_ref __libnvtx begin
        find_library(toolkit(), "nvToolsExt"; optional=true)
    end
    if path === nothing && throw_error
        error("This functionality is unavailabe as NVTX is missing.")
    end
    path
end
has_nvtx() = libnvtx(throw_error=false) !== nothing

function artifact_library(artifact, name, versions=[])
    # XXX: we don't want to consider multiple library names based on all candidate versions,
    #      since all that is known when building the artifact, but not saved anywhere.
    dir = joinpath(artifact, Sys.iswindows() ? "bin" : "lib")
    all_names = library_names(name, versions)
    for name in all_names
        path = joinpath(dir, name)
        ispath(path) && return path
    end

    # we should _always_ find libraries in artifacts
    error("""Could not find library '$name' in $artifact
             This is a bug; please file an issue with a verbose directory listing of $dir
             If this directory is empty, delete it and try again.""")
end

function find_library(cuda::ArtifactToolkit, name; optional=false)
    # NOTE: optional is ignored, and we'll error if not found
    versions = cuda_library_versions(name)
    path = artifact_library(cuda.artifact, name, versions)
    Libdl.dlopen(path)
    return path
end

function find_library(cuda::LocalToolkit, name; optional=false)
    versions = cuda_library_versions(name)
    path = find_cuda_library(cuda.dirs, name, versions)
    if path !== nothing
        Libdl.dlopen(path)
        return path
    else
        optional || error("Could not find library '$name' in your local CUDA installation.")
        return nothing
    end
end

## other

export libdevice, libcudadevrt

const __libdevice = Ref{String}()
function libdevice()
    @initialize_ref __libdevice begin
        find_libdevice(toolkit())
    end
end

function artifact_file(artifact_dir, name)
    path = joinpath(artifact_dir, name)
    if !ispath(path)
        error("""Could not find '$name' in $artifact_dir!
                 This is a bug; please file an issue with a verbose directory listing of $artifact_dir
                 If this directory is empty, delete it and try again.""")
    end
    return path
end

function find_libdevice(cuda::ArtifactToolkit)
    path = artifact_file(cuda.artifact, joinpath("share", "libdevice", "libdevice.10.bc"))
    if isfile(path)
        return path
    else
    end
end

function find_libdevice(cuda::LocalToolkit)
    path = find_libdevice(cuda.dirs)
    if path !== nothing
        return path
    else
        error("Could not find libdevice in your local CUDA installation.")
    end
end

const __libcudadevrt = Ref{String}()
function libcudadevrt()
    @initialize_ref __libcudadevrt begin
        find_libcudadevrt(toolkit())
    end
end

function artifact_static_library(artifact_dir, name)
    path = joinpath(artifact_dir, "lib", Sys.iswindows() ? "$name.lib" : "lib$name.a")
    if !ispath(path)
        error("""Could not find static library '$name' in $artifact_dir
                 This is a bug; please file an issue with a verbose directory listing of $artifact_dir
                 If this directory is empty, delete it and try again.""")
    end
    return path
end

find_libcudadevrt(cuda::ArtifactToolkit) = artifact_static_library(cuda.artifact, "cudadevrt")

function find_libcudadevrt(cuda::LocalToolkit)
    path = find_libcudadevrt(cuda.dirs)
    if path !== nothing
        return path
    else
        error("Could not find libcudadevrt in your local CUDA installation.")
    end
end

#
# CUDNN
#

export libcudnn, has_cudnn

const __libcudnn = Ref{Union{String,Nothing}}()
function libcudnn(; throw_error::Bool=true)
    path = @initialize_ref __libcudnn begin
        # CUDNN depends on CUBLAS
        libcublas()

        find_cudnn(toolkit(), v"8")
    end CUDA.CUDNN.__runtime_init__()
    if path === nothing && throw_error
        error("This functionality is unavailabe as CUDNN is missing.")
    end
    path
end
has_cudnn() = libcudnn(throw_error=false) !== nothing

function find_cudnn(cuda::ArtifactToolkit, version)
    artifact_dir = cuda_artifact("CUDNN", cuda.release)
    if artifact_dir === nothing
        return nothing
    end
    path = artifact_library(artifact_dir, "cudnn", [version])

    # HACK: eagerly open CUDNN sublibraries to avoid dlopen discoverability issues
    for sublibrary in ("ops_infer", "ops_train",
                       "cnn_infer", "cnn_train",
                       "adv_infer", "adv_train")
        sublibrary_path = artifact_library(artifact_dir, "cudnn_$(sublibrary)", [version])
        Libdl.dlopen(sublibrary_path)
    end

    @debug "Using CUDNN from an artifact at $(artifact_dir)"
    Libdl.dlopen(path)
    return path
end

function find_cudnn(cuda::LocalToolkit, version)
    path = find_library("cudnn", [version]; locations=cuda.dirs)
    if path === nothing
        return nothing
    end

    # HACK: eagerly open CUDNN sublibraries to avoid dlopen discoverability issues
    for sublibrary in ("ops_infer", "ops_train",
                       "cnn_infer", "cnn_train",
                       "adv_infer", "adv_train")
        sublibrary_path = find_library("cudnn_$(sublibrary)", [version]; locations=cuda.dirs)
        sublibrary_path === nothing && error("Could not find local CUDNN sublibrary $sublibrary")
        Libdl.dlopen(sublibrary_path)
    end

    @debug "Using local CUDNN at $(path)"
    Libdl.dlopen(path)
    return path
end

#
# CUTENSOR
#

export libcutensor, has_cutensor

const __libcutensor = Ref{Union{String,Nothing}}()
function libcutensor(; throw_error::Bool=true)
    path = @initialize_ref __libcutensor begin
        # CUTENSOR depends on CUBLAS
        libcublas()

        find_cutensor(toolkit(), v"1")
    end
    if path === nothing && throw_error
        error("This functionality is unavailabe as CUTENSOR is missing.")
    end
    path
end
has_cutensor() = libcutensor(throw_error=false) !== nothing

function find_cutensor(cuda::ArtifactToolkit, version)
    artifact_dir = cuda_artifact("CUTENSOR", cuda.release)
    if artifact_dir === nothing
        return nothing
    end
    path = artifact_library(artifact_dir, "cutensor", [version])

    @debug "Using CUTENSOR from an artifact at $(artifact_dir)"
    Libdl.dlopen(path)
    return path
end

function find_cutensor(cuda::LocalToolkit, version)
    path = find_library("cutensor", [version]; locations=cuda.dirs)
    if path === nothing
        return nothing
    end

    @debug "Using local CUTENSOR at $(path)"
    Libdl.dlopen(path)
    return path
end