text:
module CompactBasisHelFEM
using ContinuumArrays: SimplifyStyle
using CompactBases: CompactBases, Basis, @materialize, LazyArrays, ContinuumArrays,
    QuasiAdjoint, QuasiDiagonal, Derivative, BroadcastQuasiArray
import CompactBases: locs
using HelFEM: HelFEM

# This comes from EllipsisNotation via CompactBases, but can't be imported with using
const .. = CompactBases.:(..)

struct HelFEMBasis <: Basis{Float64}
    b :: HelFEM.FEMBasis
end

function Base.axes(b::HelFEMBasis)
    rmin, rmax = first(HelFEM.boundaries(b.b)), last(HelFEM.boundaries(b.b))
    (CompactBases.Inclusion(rmin..rmax), Base.OneTo(length(b.b)))
end

function Base.getindex(b::HelFEMBasis, r::Number, i)
    r ∈ axes(b,1) && i ∈ axes(b,2) || throw(BoundsError(b, [r,i]))
    b.b([r])[i]
end

function Base.getindex(b::HelFEMBasis, r::Vector, i)
    all(∈(axes(b,1)), r) && (i isa Colon || all(∈(axes(b,2)),i)) || throw(BoundsError(b, [r,i]))
    b.b(r)[:,i]
end

Base.getindex(b::HelFEMBasis, r::AbstractVector, i) = b[collect(r), i]

# Overlap matrix
@materialize function *(Ac::QuasiAdjoint{<:Any,<:HelFEMBasis}, B::HelFEMBasis)
    SimplifyStyle
    T -> begin
        A = parent(Ac)
        @assert A == B
        Matrix{T}(undef, length(A.b), length(B.b))
    end
    dest::Matrix{T} -> begin
        A = parent(Ac)
        dest .= HelFEM.radial_integral(A.b, r -> 1.0)
    end
end

# Potential
@materialize function *(Ac::QuasiAdjoint{<:Any,<:HelFEMBasis}, D::QuasiDiagonal, B::HelFEMBasis)
    SimplifyStyle
    T -> begin
        A = parent(Ac)
        @assert A == B
        Matrix{T}(undef, length(A.b), length(B.b))
    end
    dest::Matrix{T} -> begin
        A = parent(Ac)
        dest .= HelFEM.radial_integral(A.b, r -> first(getindex.(Ref(D.diag), [r])))
    end
end

@materialize function *(Ac::QuasiAdjoint{<:Any,<:HelFEMBasis}, D::Derivative, B::HelFEMBasis)
    SimplifyStyle
    T -> begin
        A = parent(Ac)
        @assert A == B
        Matrix{T}(undef, length(A.b), length(B.b))
    end
    dest::Matrix{T} -> begin
        A = parent(Ac)
        dest .= HelFEM.radial_integral(A.b, r -> 1.0, rderivative = true)
    end
end

@materialize function *(Ac::QuasiAdjoint{<:Any,<:HelFEMBasis}, Dc::QuasiAdjoint{<:Any,<:Derivative}, D::Derivative, B::HelFEMBasis)
    SimplifyStyle
    T -> begin
        A = parent(Ac)
        @assert A == B
        Matrix{T}(undef, length(A.b), length(B.b))
    end
    dest::Matrix{T} -> begin
        A = parent(Ac)
        dest .= HelFEM.radial_integral(A.b, r -> 1.0, lderivative = true, rderivative = true)
    end
end

function Sinvh(A::HelFEMBasis)
    S = HelFEM.radial_integral(A.b, r -> 1.0)
    return real(inv(sqrt(Hermitian(S))))
end

# Interpolating functions over HelFEMBasis
locs(B::HelFEMBasis) = HelFEM.controlpoints(B.b)

function Base.:(\ )(B::HelFEMBasis, f::BroadcastQuasiArray)
    @assert B.b.pb.primbas == 4 # only works for LIPs at the moment
    axes(f,1) == axes(B,1) || throw(DimensionMismatch("Function on $(axes(f,1).domain) cannot be interpolated over basis on $(axes(B,1).domain)"))
    cs = locs(B)
    collect(f[cs[2:end-1]])
end

# TODO: Is this sensible?
CompactBases.assert_compatible_bases(A::HelFEMBasis, B::HelFEMBasis) = A == B
Base.:(==)(A::HelFEMBasis, B::HelFEMBasis) = A.b == B.b

end


text:
@doc raw"""
Nonlocal term coming from norm-conserving pseudopotentials in Kleinmann-Bylander form.
``\text{Energy} = \sum_a \sum_{ij} \sum_{n} f_n <ψ_n|p_{ai}> D_{ij} <p_{aj}|ψ_n>.``
"""
struct AtomicNonlocal end
function (::AtomicNonlocal)(basis::PlaneWaveBasis)
    # keep only pseudopotential atoms
    atoms = [elem.psp => positions
             for (elem, positions) in basis.model.atoms
             if elem isa ElementPsp]

    isempty(atoms) && return NoopTerm(basis)
    ops = map(basis.kpoints) do kpt
        P = build_projection_vectors_(basis, atoms, kpt)
        D = build_projection_coefficients_(basis, atoms)
        NonlocalOperator(basis, kpt, P, D)
    end
    TermAtomicNonlocal(basis, ops)
end

struct TermAtomicNonlocal <: Term
    basis::PlaneWaveBasis
    ops::Vector{NonlocalOperator}
end

function ene_ops(term::TermAtomicNonlocal, ψ, occ; kwargs...)
    basis = term.basis
    T = eltype(basis)
    ψ === nothing && return (E=T(Inf), ops=term.ops)

    E = zero(T)
    for (ik, k) in enumerate(basis.kpoints)
        for iband = 1:size(ψ[1], 2)
            ψnk = @views ψ[ik][:, iband]
            Pψnk = term.ops[ik].P' * ψnk
            E += basis.kweights[ik] * occ[ik][iband] * real(dot(Pψnk, term.ops[ik].D * Pψnk))
        end
    end

    (E=E, ops=term.ops)
end

@timing "forces_nonlocal" function forces(term::TermAtomicNonlocal, ψ, occ; kwargs...)
    basis = term.basis
    T = eltype(basis)
    atoms = basis.model.atoms
    unit_cell_volume = basis.model.unit_cell_volume

    # early return if no pseudopotential atoms
    any(attype isa ElementPsp for (attype, positions) in atoms) || return nothing

    # energy terms are of the form <psi, P C P' psi>, where P(G) = form_factor(G) * structure_factor(G)
    forces = [zeros(Vec3{T}, length(positions)) for (el, positions) in atoms]
    for (iel, (el, positions)) in enumerate(atoms)
        el isa ElementPsp || continue

        C = build_projection_coefficients_(el.psp)
        # TODO optimize: switch this loop and the kpoint loop
        for (ir, r) in enumerate(positions)
            fr = zeros(T, 3)
            for α = 1:3
                tot_red_kpt_number = sum([length(symops) for symops in basis.ksymops])
                ind_red = 1
                for (ik, kpt_irred) in enumerate(basis.kpoints)
                    # Here we need to do an explicit loop over
                    # symmetries, because the atom displacement might break them
                    for isym in 1:length(basis.ksymops[ik])
                        (S, τ) = basis.ksymops[ik][isym]
                        Skpoint, ψSk = apply_ksymop((S, τ), basis, kpt_irred, ψ[ik])
                        Skcoord = Skpoint.coordinate
                        # energy terms are of the form <psi, P C P' psi>,
                        # where P(G) = form_factor(G) * structure_factor(G)
                        qs = [basis.model.recip_lattice * (Skcoord + G)
                              for G in G_vectors(Skpoint)]
                        form_factors = build_form_factors(el.psp, qs)
                        structure_factors = [cis(-2T(π) * dot(Skcoord + G, r))
                                             for G in G_vectors(Skpoint)]
                        P = structure_factors .* form_factors ./ sqrt(unit_cell_volume)
                        dPdR = [-2T(π)*im*(Skcoord + G)[α] for G in G_vectors(Skpoint)] .* P

                        # TODO BLASify this further
                        for iband = 1:size(ψ[ik], 2)
                            ψnSk = @view ψSk[:, iband]
                            fr[α] -= (occ[ik][iband] / tot_red_kpt_number
                                      * real(  dot(ψnSk, P * C * dPdR' * ψnSk)
                                             + dot(ψnSk, dPdR * C * P' * ψnSk)))
                        end
                        ind_red += 1
                    end
                end
            end
            forces[iel][ir] += fr
        end
    end
    forces
end

# Count the number of projection vectors implied by the potential array
# generated by term_nonlocal
function count_n_proj_(psp::PspHgh)
    psp.lmax < 0 ? 0 : sum(size(psp.h[l + 1], 1) * (2l + 1) for l in 0:psp.lmax)::Int
end
function count_n_proj_(atoms)
    sum(count_n_proj_(psp)*length(positions) for (psp, positions) in atoms)::Int
end

# Build projection coefficients for a atoms array generated by term_nonlocal
# The ordering of the projector indices is (A,l,m,i), where A is running over all
# atoms, l, m are AM quantum numbers and i is running over all projectors for a
# given l. The matrix is block-diagonal with non-zeros only if A, l and m agree.
function build_projection_coefficients_(basis::PlaneWaveBasis{T}, atoms) where {T}
    # TODO In the current version the proj_coeffs still has a lot of zeros.
    #      One could improve this by storing the blocks as a list or in a
    #      BlockDiagonal data structure
    n_proj = count_n_proj_(atoms)
    proj_coeffs = zeros(T, n_proj, n_proj)

    count = 0
    for (psp, positions) in atoms, r in positions
        n_proj_psp = count_n_proj_(psp)
        proj_coeffs[count+1:count+n_proj_psp, count+1:count+n_proj_psp] = build_projection_coefficients_(psp)
        count += n_proj_psp
    end # psp, r
    @assert count == n_proj

    proj_coeffs
end

# Builds the projection coefficient matrix for a single atom
# The ordering of the projector indices is (l,m,i), where l, m are the
# AM quantum numbers and i is running over all projectors for a given l.
# The matrix is block-diagonal with non-zeros only if l and m agree.
function build_projection_coefficients_(psp::PspHgh)
    n_proj = count_n_proj_(psp)
    proj_coeffs = zeros(n_proj, n_proj)
    count = 0
    for l in 0:psp.lmax, m in -l:l
        n_proj_l = size(psp.h[l + 1], 1)  # Number of i's
        range = count .+ (1:n_proj_l)
        proj_coeffs[range, range] = psp.h[l + 1]
        count += n_proj_l
    end # l, m
    proj_coeffs
end

"""
Build projection vectors for a atoms array generated by term_nonlocal

H_at  = sum_ij Cij |pi> <pj|
H_per = sum_R sum_ij Cij |pi(x-R)> <pj(x-R)|
      = sum_R sum_ij Cij |pi(x-R)> <pj(x-R)|

<e_kG'|H_per|e_kG> = ...
                   = 1/Ω sum_ij Cij pihat(k+G') pjhat(k+G)^*

where pihat(q) = ∫_R^3 pi(r) e^{-iqr} dr

We store 1/√Ω pihat(k+G) in proj_vectors.
"""
function build_projection_vectors_(basis::PlaneWaveBasis{T}, atoms, kpt::Kpoint) where {T}
    n_proj = count_n_proj_(atoms)
    model = basis.model

    proj_vectors = zeros(Complex{T}, length(G_vectors(kpt)), n_proj)
    qs = [model.recip_lattice * (kpt.coordinate + G) for G in G_vectors(kpt)]

    # Compute the columns of proj_vectors = 1/√Ω pihat(k+G)
    # Since the pi are translates of each others, pihat(k+G) decouples as
    # pihat(q) = ∫ p(r-R) e^{-iqr} dr = e^{-iqR} phat(q).
    # The first term is the structure factor, the second the form factor.
    offset = 0 # offset into proj_vectors
    for (psp, positions) in atoms
        # Compute position-independent form factors
        form_factors = build_form_factors(psp, qs)

        # Combine with structure factors
        for r in positions
            # k+G in this formula can also be G, this only changes an unimportant phase factor
            structure_factors = [cis(-2T(π)*dot(kpt.coordinate + G, r)) for G in G_vectors(kpt)]
            for iproj = 1:count_n_proj_(psp)
                @views proj_vectors[:, offset+iproj] .= (structure_factors
                                                        .* form_factors[:, iproj]
                                                        ./ sqrt(model.unit_cell_volume))
            end
            offset += count_n_proj_(psp)
        end
    end
    @assert offset == n_proj
    proj_vectors
end

"""
Build form factors (Fourier transforms of projectors) for an atom centered at 0.
"""
function build_form_factors(psp, qs)
    qnorms = norm.(qs)
    T = real(eltype(qnorms))
    # Compute position-independent form factors
    form_factors = zeros(Complex{T}, length(qs), count_n_proj_(psp))
    count = 1
    for l in 0:psp.lmax, m in -l:l
        prefac_lm = im^l .* ylm_real.(l, m, qs)
        n_proj_l = size(psp.h[l + 1], 1)

        for iproj in 1:n_proj_l
            radial_il = eval_psp_projection_radial.(psp, iproj, l, qnorms)
            form_factors[:, count] = prefac_lm .* radial_il
            count += 1
        end
    end
    @assert count == count_n_proj_(psp) + 1
    form_factors
end

ext: jl


text:
"""
A Field with this form
y = x⋅G + y0
"""
struct AffineMap{D1,D2,T,L} <:Field
  gradient::TensorValue{D1,D2,T,L}
  origin::Point{D2,T}
  function AffineMap(
    gradient::TensorValue{D1,D2,T,L},
    origin::Point{D2,T}) where {D1,D2,T,L}

    new{D1,D2,T,L}(gradient,origin)
  end
end

@inline affine_map(gradient,origin) = AffineMap(gradient,origin)

@inline function evaluate!(cache,f::AffineMap,x::Point)
  G = f.gradient
  y0 = f.origin
  x⋅G + y0
end

function return_cache(f::AffineMap,x::AbstractVector{<:Point})
  T = return_type(f,testitem(x))
  y = similar(x,T,size(x))
  CachedArray(y)
end

@inline function evaluate!(cache,f::AffineMap,x::AbstractVector{<:Point})
  setsize!(cache,size(x))
  y = cache.array
  G = f.gradient
  y0 = f.origin
  for i in eachindex(x)
    xi = x[i]
    yi = xi⋅G + y0
    y[i] = yi
  end
  y
end

function gradient(h::AffineMap)
  ConstantField(h.gradient)
end

function push_∇∇(∇∇a::Field,ϕ::AffineMap)
  # Assuming ϕ is affine map
  Jt = ∇(ϕ)
  Jt_inv = pinvJt(Jt)
  Operation(push_∇∇)(∇∇a, Jt_inv)
end

function push_∇∇(∇∇a::Number,Jt_inv::MultiValue{Tuple{D,D}} where D)
  Jt_inv⋅Jt_inv⋅∇∇a
end

function push_∇∇(∇∇a::Number,Jt_inv::MultiValue{Tuple{D1,D2}} where {D1,D2})
  Jt_inv⋅∇∇a⋅transpose(Jt_inv)
end

function lazy_map(
  k::Broadcasting{typeof(push_∇∇)},
  cell_∇∇a::AbstractArray,
  cell_map::AbstractArray{<:AffineMap})
  cell_Jt = lazy_map(∇,cell_map)
  cell_invJt = lazy_map(Operation(pinvJt),cell_Jt)
  lazy_map(Broadcasting(Operation(push_∇∇)),cell_∇∇a,cell_invJt)
end

function lazy_map(
  k::Broadcasting{typeof(push_∇∇)},
  cell_∇∇at::LazyArray{<:Fill{typeof(transpose)}},
  cell_map::AbstractArray{<:AffineMap})
  cell_∇∇a = cell_∇∇at.args[1]
  cell_∇∇b = lazy_map(k,cell_∇∇a,cell_map)
  cell_∇∇bt = lazy_map(transpose,cell_∇∇b)
  cell_∇∇bt
end

function inverse_map(f::AffineMap)
  Jt = f.gradient
  y0 = f.origin
  invJt = inv(Jt)
  x0 = -y0⋅invJt
  AffineMap(invJt,x0)
end

function lazy_map(::typeof(∇),a::LazyArray{<:Fill{typeof(affine_map)}})
  gradients = a.args[1]
  lazy_map(constant_field,gradients)
end


text:
module PatientAllocation

using JuMP
using Gurobi

using LinearAlgebra
using MathOptInterface
using Distributions
using Memoize
using Statistics

##############################################
############# Standard Model #################
##############################################

function patient_redistribution(
		capacity::Array{<:Real},
		initial_patients::Array{<:Real,1},
		discharged_patients::Array{<:Real,2},
		admitted_patients::Array{<:Real,2},
		adj_matrix::BitArray{2}, los::Union{<:Distribution,Array{<:Real,1},Int};

		capacity_cushion::Real=-1, capacity_weights::Array{<:Real,1}=Int[],
		no_artificial_overflow::Bool=false, no_worse_overflow::Bool=false,
		sent_penalty::Real=0, smoothness_penalty::Real=0,

		sendreceive_gap::Int=0, min_send_amt::Real=0,
		balancing_thresh::Real=1.0, balancing_penalty::Real=0,
		severity_weighting::Bool=false, setup_cost::Real=0,

		verbose::Bool=false,
	)

	###############
	#### Setup ####
	###############

	if ndims(capacity) == 1
		capacity = reshape(capacity, (:,1))
	end

	N, T = size(admitted_patients)
	C = size(capacity, 2)
	check_sizes(initial_patients, discharged_patients, admitted_patients, capacity)

	if 0 < capacity_cushion < 1
		capacity = capacity .* (1.0 - capacity_cushion)
	end

	if isempty(capacity_weights)
		capacity_weights = ones(Int, C)
	end

	L = discretize_los(los, T)

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	###############
	## Variables ##
	###############

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, overflow[1:N,1:T,1:C] >= 0)

	#################
	## Expressions ##
	#################

	# expressions for the number of active patients
	@expression(model, active_patients[i=1:N,t=1:T],
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (
			admitted_patients[i,t₁]
			- sum(sent[i,:,t₁])
			+ sum(sent[:,i,t₁])
		) for t₁ in 1:t)
		+ sum(sent[i,:,t])
	)
	active_null = compute_active_null(initial_patients, discharged_patients, admitted_patients, L)

	# objective function
	objective = @expression(model, dot(capacity_weights, sum(overflow, dims=(1,2))))

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [i=1:N,t=1:T], active_patients[i,t] >= 0)

	# only send new patients
	@constraint(model, [t=1:T], sum(sent[:,:,t], dims=2) .<= admitted_patients[:,t])

	# objective constraint
	@constraint(model, [i=1:N,t=1:T,c=1:C], overflow[i,t,c] >= active_patients[i,t] - capacity[i,c])

	################################
	## Optional Constraints/Costs ##
	################################

	enforce_adj!(model, sent, adj_matrix)
	enforce_no_artificial_overflow!(model, no_artificial_overflow, active_patients, active_null, capacity)
	enforce_no_worse_overflow!(model, no_worse_overflow, active_patients, active_null, capacity)
	enforce_minsendamt!(model, sent, min_send_amt)
	enforce_sendreceivegap!(model, sent, sendreceive_gap)

	add_sent_penalty!(model, sent, objective, sent_penalty)
	add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	add_setup_cost!(model, sent, objective, setup_cost)
	add_loadbalancing_penalty!(model, sent, objective, balancing_penalty, balancing_thresh, active_patients, capacity)
	add_severity_weighting!(model, sent, objective, severity_weighting, overflow, active_null, capacity)

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
############### Load Balance #################
##############################################

function patient_loadbalance(
		capacity::Array{<:Real},
		initial_patients::Array{<:Real,1},
		discharged_patients::Array{<:Real,2},
		admitted_patients::Array{<:Real,2},
		adj_matrix::BitArray{2}, los::Union{<:Distribution,Array{<:Real,1},Int};

		capacity_cushion::Real=-1, capacity_weights::Array{<:Real,1}=Int[],
		no_artificial_overflow::Bool=false, no_worse_overflow::Bool=false,
		sent_penalty::Real=0, smoothness_penalty::Real=0,

		sendreceive_gap::Int=0, min_send_amt::Real=0,
		setup_cost::Real=0,

		verbose::Bool=false,
	)

	###############
	#### Setup ####
	###############

	if ndims(capacity) == 1
		capacity = reshape(capacity, (:,1))
	end

	N, T = size(admitted_patients)
	C = size(capacity, 2)
	check_sizes(initial_patients, discharged_patients, admitted_patients, capacity)

	if 0 < capacity_cushion < 1
		capacity = capacity .* (1.0 - capacity_cushion)
	end

	if isempty(capacity_weights)
		capacity_weights = ones(Int, C)
	end

	L = discretize_los(los, T)

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	###############
	## Variables ##
	###############

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, load_objective[1:N,1:T,1:C] >= 0)

	#################
	## Expressions ##
	#################

	# expressions for the number of active patients
	@expression(model, active_patients[i=1:N,t=1:T],
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (
			admitted_patients[i,t₁]
			- sum(sent[i,:,t₁])
			+ sum(sent[:,i,t₁])
		) for t₁ in 1:t)
		+ sum(sent[i,:,t])
	)
	active_null = compute_active_null(initial_patients, discharged_patients, admitted_patients, L)

	# expression for the patient load
	@expression(model, load[i=1:N,t=1:T,c=1:C], active_patients[i,t] / capacity[i,c])

	# objective function
	objective = @expression(model, dot(capacity_weights, sum(load_objective, dims=(1,2))))

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [i=1:N,t=1:T], active_patients[i,t] >= 0)

	# only send new patients
	@constraint(model, [t=1:T], sum(sent[:,:,t], dims=2) .<= admitted_patients[:,t])

	# objective constraint
	@constraint(model, [i=1:N,t=1:T,c=1:C],  (load[i,t,c] - mean(load[:,t,c])) <= load_objective[i,t,c])
	@constraint(model, [i=1:N,t=1:T,c=1:C], -(load[i,t,c] - mean(load[:,t,c])) <= load_objective[i,t,c])

	################################
	## Optional Constraints/Costs ##
	################################

	enforce_adj!(model, sent, adj_matrix)
	enforce_no_artificial_overflow!(model, no_artificial_overflow, active_patients, active_null, capacity)
	enforce_no_worse_overflow!(model, no_worse_overflow, active_patients, active_null, capacity)
	enforce_minsendamt!(model, sent, min_send_amt)
	enforce_sendreceivegap!(model, sent, sendreceive_gap)

	add_sent_penalty!(model, sent, objective, sent_penalty)
	add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	add_setup_cost!(model, sent, objective, setup_cost)

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
############ Surge-Level Model ###############
##############################################

function patient_surgelevel_redistribution(
		capacity::Array{<:Real,2},
		initial_patients::Array{<:Real,1},
		discharged_patients::Array{<:Real,2},
		admitted_patients::Array{<:Real,2},
		adj_matrix::BitArray{2}, los::Union{<:Distribution,Array{<:Real,1},Int};

		capacity_cushion::Real=-1, capacity_weights::Array{<:Real,1}=Int[],
		no_artificial_overflow::Bool=false, no_worse_overflow::Bool=false,
		sent_penalty::Real=0, smoothness_penalty::Real=0,

		sendreceive_gap::Int=0, min_send_amt::Real=0,
		balancing_thresh::Real=1.0, balancing_penalty::Real=0,
		setup_cost::Real=0,

		mipgap::Real=0.1, timelimit::Real=240, M::Real=1e6,

		verbose::Bool=false,
	)

	###############
	#### Setup ####
	###############

	N, T = size(admitted_patients)
	C = size(capacity, 2)
	check_sizes(initial_patients, discharged_patients, admitted_patients, capacity)

	if 0 < capacity_cushion < 1
		capacity = capacity .* (1.0 - capacity_cushion)
	end

	if isempty(capacity_weights)
		capacity_weights = ones(Int, C)
	end

	L = discretize_los(los, T)

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	set_optimizer_attribute(model, "MIPGap", mipgap)
	set_optimizer_attribute(model, "TimeLimit", timelimit)

	###############
	## Variables ##
	###############

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, surgelevel[1:N,1:T,1:C], Bin)

	#################
	## Expressions ##
	#################

	# expressions for the number of active patients
	@expression(model, active_patients[i=1:N,t=1:T],
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (
			admitted_patients[i,t₁]
			- sum(sent[i,:,t₁])
			+ sum(sent[:,i,t₁])
		) for t₁ in 1:t)
		+ sum(sent[i,:,t])
	)
	active_null = compute_active_null(initial_patients, discharged_patients, admitted_patients, L)

	# expression for the patient overflow
	@expression(model, overflow[i=1:N,t=1:T,c=1:C], active_patients[i,t] - capacity[i,c])

	# weights
	weights = hcat(capacity[:,1], diff(capacity, dims=2))
	weights = weights .* permutedims(repeat(capacity_weights, 1, N), (2,1))

	# objective function
	objective = @expression(model, sum(sum(surgelevel, dims=2)[:,1,:] .* weights))

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [i=1:N,t=1:T], active_patients[i,t] >= 0)

	# only send new patients
	@constraint(model, [t=1:T], sum(sent[:,:,t], dims=2) .<= admitted_patients[:,t])

	# set surgelevel
	for i in 1:N, t in 1:T, c in 1:C
		@constraint(model, M * surgelevel[i,t,c] >= overflow[i,t,c])
		@constraint(model, M * (1 - surgelevel[i,t,c]) >= -overflow[i,t,c])
	end

	################################
	## Optional Constraints/Costs ##
	################################

	enforce_adj!(model, sent, adj_matrix)
	enforce_no_artificial_overflow!(model, no_artificial_overflow, active_patients, active_null, capacity)
	enforce_no_worse_overflow!(model, no_worse_overflow, active_patients, active_null, capacity)
	enforce_minsendamt!(model, sent, min_send_amt)
	enforce_sendreceivegap!(model, sent, sendreceive_gap)

	add_sent_penalty!(model, sent, objective, sent_penalty)
	add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	add_setup_cost!(model, sent, objective, setup_cost)
	add_loadbalancing_penalty!(model, sent, objective, balancing_penalty, balancing_thresh, active_patients, capacity)

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
####### Surge-Level Load Balancing ###########
##############################################

function patient_surgelevel_loadbalance(
		capacity::Array{<:Real},
		initial_patients::Array{<:Real,1},
		discharged_patients::Array{<:Real,2},
		admitted_patients::Array{<:Real,2},
		adj_matrix::BitArray{2}, los::Union{<:Distribution,Array{<:Real,1},Int};

		capacity_cushion::Real=-1,
		no_artificial_overflow::Bool=false, no_worse_overflow::Bool=false,
		sent_penalty::Real=0, smoothness_penalty::Real=0,

		sendreceive_gap::Int=0, min_send_amt::Real=0,
		setup_cost::Real=0,

		mipgap::Real=0.1, timelimit::Real=240, M::Real=1e6,
		verbose::Bool=false,
	)

	###############
	#### Setup ####
	###############

	N, T = size(admitted_patients)
	C = size(capacity, 2)
	check_sizes(initial_patients, discharged_patients, admitted_patients, capacity)

	if 0 < capacity_cushion < 1
		capacity = capacity .* (1.0 - capacity_cushion)
	end

	L = discretize_los(los, T)

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	set_optimizer_attribute(model, "MIPGap", mipgap)
	set_optimizer_attribute(model, "TimeLimit", timelimit)

	###############
	## Variables ##
	###############

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, surgelevel[1:N,1:T,1:C], Bin)
	@variable(model, load_obj[1:N,1:T] >= 0)

	#################
	## Expressions ##
	#################

	# expressions for the number of active patients
	@expression(model, active_patients[i=1:N,t=1:T],
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (
			admitted_patients[i,t₁]
			- sum(sent[i,:,t₁])
			+ sum(sent[:,i,t₁])
		) for t₁ in 1:t)
		+ sum(sent[i,:,t])
	)
	active_null = compute_active_null(initial_patients, discharged_patients, admitted_patients, L)

	# expression for the patient overflow
	@expression(model, overflow[i=1:N,t=1:T,c=1:C], active_patients[i,t] - capacity[i,c])

	# objective function
	objective = @expression(model, sum(load_obj))

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [i=1:N,t=1:T], active_patients[i,t] >= 0)

	# only send new patients
	@constraint(model, [t=1:T], sum(sent[:,:,t], dims=2) .<= admitted_patients[:,t])

	# set surge-level
	for i in 1:N, t in 1:T, c in 1:C
		@constraint(model, M * surgelevel[i,t,c] >= overflow[i,t,c])
		@constraint(model, M * (1 - surgelevel[i,t,c]) >= -overflow[i,t,c])
	end

	# load balancing objective
	@constraint(model, [i=1:N,t=1:T],  (sum(surgelevel[i,t,:]) - mean(sum(surgelevel[:,t,:], dims=2))) <= load_obj[i,t])
	@constraint(model, [i=1:N,t=1:T], -(sum(surgelevel[i,t,:]) - mean(sum(surgelevel[:,t,:], dims=2))) <= load_obj[i,t])

	################################
	## Optional Constraints/Costs ##
	################################

	enforce_adj!(model, sent, adj_matrix)
	enforce_no_artificial_overflow!(model, no_artificial_overflow, active_patients, active_null, capacity)
	enforce_no_worse_overflow!(model, no_worse_overflow, active_patients, active_null, capacity)
	enforce_minsendamt!(model, sent, min_send_amt)
	enforce_sendreceivegap!(model, sent, sendreceive_gap)

	add_sent_penalty!(model, sent, objective, sent_penalty)
	add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	add_setup_cost!(model, sent, objective, setup_cost)

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
############### Block Model ##################
##############################################

function patient_block_allocation(
		beds::Array{<:Real,2},
		initial::Array{<:Real,2},
		discharged::Array{<:Real,3},
		admitted::Array{<:Real,3},
		los::Array{<:Any,1},
		adj_matrix::BitArray{2},
		transfer_graph::BitArray{2},
		bed_types::Array{Int,1};
		sendreceive_switch_time::Int=0,
		min_send_amt::Real=0,
		smoothness_penalty::Real=0,
		setup_cost::Real=0,
		sent_penalty::Real=0,
		balancing_thresh::Real=1.0,
		balancing_penalty::Real=0,
		severity_weighting::Bool=false,
		no_artificial_overflow::Bool=false,
		no_worse_overflow::Bool=false,
		capacity_cushion::Real=0.0,
		verbose::Bool=false,
)
	B = size(beds,1)
	G, N, T = size(admitted)

	@assert size(beds) == (B,N)
	@assert size(initial) == (G,N)
	@assert size(discharged) == (G,N,T)
	@assert size(adj_matrix) == (N,N)
	@assert size(transfer_graph) == (G,G)
	@assert size(los) == (G,)
	@assert size(bed_types) == (G,)

	###############
	#### Setup ####
	###############

	F = zeros(Float64, G, T)
	f = zeros(Float64, G, T)
	for g in 1:G
		if isa(los[g], Int)
			F[g,:] = vcat(ones(Int, los[g]), zeros(Int, T-los[g]))
			f[g,los[g]+1] = 1.0
		elseif isa(los[g], Distribution)
			F[g,:] = 1.0 .- cdf.(los[g], 0:T-1)
			f[g,:] = pdf.(los[g], 0:T-1)
		else
			error("Invalid length of stay distribution")
		end
		f[g,:] = f[g,:] / sum(f[g,:])
	end

	if capacity_cushion > 0.0
		beds = beds .* (1.0 - capacity_cushion)
	end

	bed_groups = [sort(findall(x -> x == b, bed_types)) for b in 1:B]

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	###############
	## Variables ##
	###############

	@variable(model, sent[1:G,1:N,1:N,1:T] >= 0)
	@variable(model, obj_dummy[1:B,1:N,1:T] >= 0)

	#################
	## Expressions ##
	#################

	@memoize q(g,i,t) = @expression(model,
		admitted[g,i,t] + tr(g,i,t) + sum(sent[g,j,i,t] - sent[g,i,j,t] for j in 1:N)
	)
	@memoize l(g,i,t) = @expression(model,
		discharged[g,i,t] + sum(f[g,t-t₁+1] * q(g,i,t₁) for t₁ in 1:t)
	)
	@memoize tr(g,i,t) = @expression(model,
		any(transfer_graph[:,g]) ?
			sum(l(g₁,i,t) for g₁ in findall(transfer_graph[:,g]))
			: 0
	)
	@expression(model, α[g=1:G,i=1:N,t=1:T],
		initial[g,i] + sum(sent[g,i,:,t]) + sum((q(g,i,t₁) - l(g,i,t₁)) for t₁ in 1:t)
	)

	@memoize q_null(g,i,t) = (
		admitted[g,i,t] + tr_null(g,i,t)
	)
	@memoize l_null(g,i,t) = (
		discharged[g,i,t] + sum(f[g,t-t₁+1] * q_null(g,i,t₁) for t₁ in 1:t)
	)
	@memoize tr_null(g,i,t) = (
		any(transfer_graph[:,g]) ?
			sum(l_null(g₁,i,t) for g₁ in findall(transfer_graph[:,g]))
			: 0
	)
	α_null = [(
		initial[g,i] + sum((q_null(g,i,t₁) - l_null(g,i,t₁)) for t₁ in 1:t)
		) for g in 1:G, i in 1:N, t in 1:T
	]

	# expression for the patient overflow
	@expression(model, overflow[b=1:B,i=1:N,t=1:T],
		sum(α[g,i,t] for g in bed_groups[b]) - beds[b,i]
	)

	objective = @expression(model, sum(obj_dummy))

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [g=1:G,i=1:N,t=1:T], α[g,i,t] >= 0)

	# only send new patients
	@constraint(model, [g=1:G,i=1:N,t=1:T], sum(sent[g,i,:,t]) <= admitted[g,i,t])

	# only send patients between connected locations
	for i in 1:N, j in 1:N
		if ~adj_matrix[i,j]
			for g in 1:G, t in 1:T
				fix(sent[g,i,j,t], 0, force=true)
			end
		end
	end

	# only transfer patients within groups with incoming patients
	for g in findall(sum(admitted, dims=[2,3])[:] .== 0)
		for i in 1:N, j in 1:N, t in 1:T
			fix(sent[g,i,j,t], 0, force=true)
		end
	end

	# objective constraint
	@constraint(model, [b=1:B,i=1:N,t=1:T], obj_dummy[b,i,t] >= overflow[b,i,t])

	##########################
	## Optional Constraints ##
	##########################

	# enforce minimum transfer amount if enabled
	if min_send_amt > 0
		semi_cont_set = MOI.Semicontinuous(Float64(min_send_amt), Inf)
		for g in 1:G, i in 1:N, j in 1:N, t in 1:T
			if !is_fixed(sent[g,i,j,t])
				delete_lower_bound(sent[g,i,j,t])
				@constraint(model, sent[g,i,j,t] in semi_cont_set)
			end
		end
	end

	# weight objective per-location by max load
	if severity_weighting
		load_null = [(
			sum(α_null[g,i,t] for g in bed_groups[b]) / beds[b,i]
			) for b in bed_types, i in 1:N, t in 1:T
		]
		max_load_null = maximum(load_null, dims=3)[:,:,1]
		severity_weight = [max_load_null[b,i] > 1.0 ? 0.0 : 9.0 for b in 1:B, i in 1:N]

		add_to_expression!(objective, sum(sum(obj_dummy, dims=3) .* severity_weight))
	end

	if no_artificial_overflow
		for b in bed_types, i in 1:N, t in 1:T
			if sum(α_null[g,i,t] for g in bed_groups[b]) < beds[b,i]
				@constraint(model, sum(α[g,i,t] for g in bed_groups[b]) <= beds[b,i])
			end
		end
	end

	if no_worse_overflow
		for b in bed_types, i in 1:N, t in 1:T
			if sum(α_null[g,i,t] for g in bed_groups[b]) > beds[i]
				@constraint(model, sum(α[g,i,t] for g in bed_groups[b]) <= sum(α_null[g,i,t] for g in bed_groups[b]))
			end
		end
	end

	# penalize total sent if enabled
	if sent_penalty > 0
		add_to_expression!(objective, sent_penalty*sum(sent))
	end

	# penalize non-smoothness in sent patients if enabled
	if smoothness_penalty > 0
		@variable(model, smoothness_dummy[i=1:N,j=1:N,t=1:T-1] >= 0)
		@constraint(model, [t=1:T-1],  (sum(sent[:,:,:,t], dims=1) - sum(sent[:,:,:,t+1], dims=1))[1,:,:,:] .<= smoothness_dummy[:,:,t])
		@constraint(model, [t=1:T-1], -(sum(sent[:,:,:,t], dims=1) - sum(sent[:,:,:,t+1], dims=1))[1,:,:,:] .<= smoothness_dummy[:,:,t])

		add_to_expression!(objective, smoothness_penalty * sum(smoothness_dummy))
		add_to_expression!(objective, smoothness_penalty * sum(sent[:,:,:,1]))
	end

	# add setup costs if enabled
	if setup_cost > 0
		@variable(model, setup_dummy[i=1:N,j=i+1:N], Bin)
		@constraint(model, [i=1:N,j=i+1:N], [1-setup_dummy[i,j], sum(sent[:,i,j,:])+sum(sent[:,j,i,:])] in MOI.SOS1([1.0, 1.0]))
		add_to_expression!(objective, setup_cost*sum(setup_dummy))
	end

	# enforce a minimum time between sending and receiving
	if sendreceive_switch_time > 0
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,:,i,t]), sum(sent[:,i,:,t:min(t+sendreceive_switch_time,T)])] in MOI.SOS1([1.0, 1.0])
		)
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,:,i,t:min(t+sendreceive_switch_time,T)]), sum(sent[:,i,:,t])] in MOI.SOS1([1.0, 1.0])
		)
	end

	# load balancing
	if balancing_penalty > 0
		@variable(model, balancing_dummy[bed_types,1:N,1:T] >= 0)
		@constraint(model, [b in bed_types,i=1:N,t=1:T],
			balancing_dummy[b,i,t] >= (sum(α[g,i,t] for g in bed_groups[b]) / beds[b,i]) - balancing_thresh)
		add_to_expression!(objective, balancing_penalty * sum(balancing_dummy))
	end

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
############### Robust Model #################
##############################################

function patient_redistribution_robust(
		capacity::Array{<:Real},
		initial_patients::Array{<:Real,1},
		discharged_patients::Array{<:Real,2},
		admitted_patients_nominal::Array{<:Real,2},
		admitted_patients_lb::Array{<:Real,2},
		admitted_patients_ub::Array{<:Real,2},
		adj_matrix::BitArray{2}, los::Union{<:Distribution,Array{<:Real,1},Int};

		Γ::Int=3,
		capacity_cushion::Real=-1, capacity_weights::Array{<:Real,1}=Int[],
		no_artificial_overflow::Bool=false, no_worse_overflow::Bool=false,
		sent_penalty::Real=0, smoothness_penalty::Real=0,
		sendreceive_gap::Int=0, min_send_amt::Real=0, setup_cost::Real=0,

		verbose::Bool=false,
	)

	###############
	#### Setup ####
	###############

	if ndims(capacity) == 1
		capacity = reshape(capacity, (:,1))
	end

	N, T = size(admitted_patients_nominal)
	C = size(capacity, 2)
	check_sizes(initial_patients, discharged_patients, admitted_patients_nominal, capacity)

	if 0 < capacity_cushion < 1
		capacity = capacity .* (1.0 - capacity_cushion)
	end

	if isempty(capacity_weights)
		capacity_weights = ones(Int, C)
	end

	L = discretize_los(los, T)

	###############
	#### Model ####
	###############

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	###############
	## Variables ##
	###############

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, overflow[1:N,1:T,1:C] >= 0)

	#################
	## Expressions ##
	#################

	# objective function
	objective = @expression(model, dot(capacity_weights, sum(overflow, dims=(1,2))))

	active_null_ub = [(
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (t₁ > (t-Γ) ? admitted_patients_ub[i,t₁] : admitted_patients_nominal[i,t₁]) for t₁ in 1:t)
	) for i in 1:N, t in 1:T]

	active_null_lb = [(
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (t₁ > (t-Γ) ? admitted_patients_lb[i,t₁] : admitted_patients_nominal[i,t₁]) for t₁ in 1:t)
	) for i in 1:N, t in 1:T]

	######################
	## Hard Constraints ##
	######################

	# ensure the number of active patients is non-negative
	@constraint(model, [i=1:N,t=1:T], 0 <=
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (
			(t₁ > (t-Γ) ? admitted_patients_ub[i,t₁] : admitted_patients_nominal[i,t₁])
			- sum(sent[i,:,t₁])
			+ sum(sent[:,i,t₁])
		) for t₁ in 1:t)
		+ sum(sent[i,:,t])
	)

	# only send new patients
	@constraint(model, [i=1:N,t=1:T], sum(sent[i,:,t]) <= admitted_patients_lb[i,t])

	# objective constraint
	@constraint(model, [i=1:N,t=1:T,c=1:C], overflow[i,t,c] >=
		initial_patients[i]
		- sum(discharged_patients[i,1:t])
		+ sum(L[t-t₁+1] * (t₁ > (t-Γ) ? admitted_patients_ub[i,t₁] : admitted_patients_nominal[i,t₁]) for t₁ in 1:t)
		- sum(L[t-t₁+1] * sum(sent[i,:,t₁]) for t₁ in 1:t)
		+ sum(L[t-t₁+1] * sum(sent[:,i,t₁]) for t₁ in 1:t)
		+ sum(sent[i,:,t])
		- capacity[i,c]
	)

	##########################
	## Optional Constraints ##
	##########################

	enforce_adj!(model, sent, adj_matrix)
	enforce_minsendamt!(model, sent, min_send_amt)
	enforce_sendreceivegap!(model, sent, sendreceive_gap)

	add_sent_penalty!(model, sent, objective, sent_penalty)
	add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	add_setup_cost!(model, sent, objective, setup_cost)

	if no_artificial_overflow
		for i in 1:N, t in 1:T
			if active_null_ub[i,t] < capacity[i,end]
				@constraint(model, capacity[i,end] >=
					initial_patients[i]
					- sum(discharged_patients[i,1:t])
					+ sum(L[t-t₁+1] * (t₁ > (t-Γ) ? admitted_patients_ub[i,t₁] : admitted_patients_nominal[i,t₁]) for t₁ in 1:t)
					- sum(L[t-t₁+1] * sum(sent[i,:,t₁]) for t₁ in 1:t)
					+ sum(L[t-t₁+1] * sum(sent[:,i,t₁]) for t₁ in 1:t)
					+ sum(sent[i,:,t])
				)
			end
		end
	end

	if no_worse_overflow
		for i in 1:N, t in 1:T
			if active_null_lb[i,t] > capacity[i,end]
				@constraint(model, 0 >=
					- sum(L[t-t₁+1] * sum(sent[i,:,t₁]) for t₁ in 1:t)
					+ sum(L[t-t₁+1] * sum(sent[:,i,t₁]) for t₁ in 1:t)
					+ sum(sent[i,:,t])
				)
			end
		end
	end

	###############
	#### Solve ####
	###############

	@objective(model, Min, objective)
	optimize!(model)

	return model
end

##############################################
############# Helper Functions ###############
##############################################

function discretize_los(los, T)
	L = nothing
	if isa(los, Int)
		L = vcat(ones(Int, los), zeros(Int, T-los))
	elseif isa(los, Array{<:Real,1})
		if length(los) >= T
			L = los
		else
			L = vcat(los, zeros(Float64, T-length(los)))
		end
	elseif isa(los, Distribution)
		L = 1.0 .- cdf.(los, 0:T)
	else
		error("Invalid length of stay distribution")
	end
	return L
end

function compute_active_null(initial_patients, discharged_patients, admitted_patients, L)
	N, T = size(admitted_patients)
	active_null = [(
			initial_patients[i]
			- sum(discharged_patients[i,1:t])
			+ sum(L[t-t₁+1] * admitted_patients[i,t₁] for t₁ in 1:t)
		) for i in 1:N, t in 1:T
	]
	return active_null
end

function check_sizes(initial_patients, discharged_patients, admitted_patients, beds)
	N, T = size(admitted_patients)
	@assert(size(initial_patients) == (N,))
	@assert(size(discharged_patients) == (N, T))
	@assert(size(beds, 1) == N)
	return
end

##############################################
########### Optional Constraints #############
##############################################

# only send patients between connected locations
function enforce_adj!(model, sent, adj_matrix)
	N, _, T = size(sent)
	@assert(size(adj_matrix) == (N,N))
	for i in 1:N, j in 1:N
		if ~adj_matrix[i,j]
			for t in 1:T
				fix(sent[i,j,t], 0, force=true)
			end
		end
	end
	return
end

function enforce_no_artificial_overflow!(model, no_artificial_overflow, active_patients, active_null, capacity)
	if no_artificial_overflow
		N, T = size(active_null)
		for i in 1:N, t in 1:T
			if active_null[i,t] <= capacity[i,end]
				@constraint(model, active_patients[i,t] <= capacity[i,end])
			end
		end
	end
	return
end

function enforce_no_worse_overflow!(model, no_worse_overflow, active_patients, active_null, capacity)
	if no_worse_overflow
		N, T = size(active_null)
		for i in 1:N, t in 1:T
			if active_null[i,t] >= capacity[i,end]
				@constraint(model, active_patients[i,t] <= active_null[i,t])
			end
		end
	end
	return
end

# enforce minimum transfer amount if enabled
function enforce_minsendamt!(model, sent, min_send_amt)
	if min_send_amt > 0
		N, _, T = size(sent)
		semi_cont_set = MOI.Semicontinuous(Float64(min_send_amt), Inf)
		for i in 1:N, j in 1:N, t in 1:T
			if !is_fixed(sent[i,j,t])
				delete_lower_bound(sent[i,j,t])
				@constraint(model, sent[i,j,t] in semi_cont_set)
			end
		end
	end
	return
end

# enforce a minimum time between sending and receiving
function enforce_sendreceivegap!(model, sent, sendreceive_gap)
	if sendreceive_gap > 0
		N, _, T = size(sent)
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,i,t]), sum(sent[i,:,t:min(t+sendreceive_gap,T)])] in MOI.SOS1([1.0, 1.0])
		)
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,i,t:min(t+sendreceive_gap,T)]), sum(sent[i,:,t])] in MOI.SOS1([1.0, 1.0])
		)
	end
	return
end

##############################################
############ Optional Penalties ##############
##############################################

# penalize total sent if enabled
function add_sent_penalty!(model, sent, objective, sent_penalty)
	if sent_penalty > 0
		add_to_expression!(objective, sent_penalty*sum(sent))
	end
	return
end

# penalize non-smoothness in sent patients if enabled
function add_smoothness_penalty!(model, sent, objective, smoothness_penalty)
	if smoothness_penalty > 0
		N, _, T = size(sent)

		@variable(model, smoothness_dummy[i=1:N,j=1:N,t=1:T-1] >= 0)
		@constraint(model, [t=1:T-1],  (sent[:,:,t] - sent[:,:,t+1]) .<= smoothness_dummy[:,:,t])
		@constraint(model, [t=1:T-1], -(sent[:,:,t] - sent[:,:,t+1]) .<= smoothness_dummy[:,:,t])

		add_to_expression!(objective, smoothness_penalty * sum(smoothness_dummy))
		add_to_expression!(objective, smoothness_penalty * sum(sent[:,:,1]))
	end
	return
end

# add setup costs if enabled
function add_setup_cost!(model, sent, objective, setup_cost)
	if setup_cost > 0
		N, _, T = size(sent)
		@variable(model, setup_dummy[i=1:N,j=i+1:N], Bin)
		@constraint(model, [i=1:N,j=i+1:N], [1-setup_dummy[i,j], sum(sent[i,j,:])+sum(sent[j,i,:])] in MOI.SOS1([1.0, 1.0]))
		add_to_expression!(objective, setup_cost*sum(setup_dummy))
	end
	return
end

# load balancing penalty
function add_loadbalancing_penalty!(model, sent, objective, balancing_penalty, balancing_thresh, active_patients, capacity)
	if balancing_penalty > 0
		N, _, T = size(sent)
		@variable(model, balancing_dummy[1:N,1:T] >= 0)
		@constraint(model, [i=1:N,t=1:T], balancing_dummy[i,t] >= (active_patients[i,t] / capacity[i,1]) - balancing_thresh)
		add_to_expression!(objective, balancing_penalty * sum(balancing_dummy))
	end
	return
end

# weight objective per-location by max load
function add_severity_weighting!(model, sent, objective, severity_weighting, overflow, active_null, capacity)
	if severity_weighting
		N, _, T = size(sent)
		max_load_null = [maximum(active_null[i,:] / capacity[i,1]) for i in 1:N]
		severity_weight = [max_load_null[i] > 1.0 ? 0.0 : 9.0 for i in 1:N]
		add_to_expression!(objective, dot(sum(overflow, dims=2), severity_weight))
	end
	return
end

end;

max_stars_repo_name: flixpar/covid-resource-allocation

max_issues_repo_name: flixpar/covid-resource-allocation

max_forks_repo_name: flixpar/covid-resource-allocation


text:
using RationalApproximations
using PyPlot

S = range(-1, 1, length=100)
F = exp.(S)
max_m = 5
tol = 1e-6
r_v1, err = aaa_v1(F, S, max_m, tol)

figure(1)
t = range(-1, 1, length=201)
y = exp.(t)
for m = 2:4
    semilogy(t, abs.(y-r_v1[m].(t)), label="m=$m")
end
legend(loc="lower left")
grid(true)
title("AAA Version 1: absolute errors")

r_v2, err = aaa_v2(F, S, max_m, tol)

figure(2)
for m = 2:4
    semilogy(t, abs.(y-r_v2[m].(t)), label="m=$m")
end
legend(loc="lower left")
grid(true)
title("AAA Version 2: absolute errors")

figure(3)
plot(t, y-r_v1[4].(t), label="V1")
plot(t, y-r_v2[4].(t), label="V2")
legend()
grid(true)
title("Comparison of errors when m=4")

max_stars_repo_name: billmclean/RationalApproximations

max_issues_repo_name: billmclean/RationalApproximations

max_forks_repo_name: billmclean/RationalApproximations


text:
struct EPMat{T<:AbstractFloat} <: AbstractEPMat
    αStS::SparseMatrixCSC{T,Int64} # A container of KK = βS'S (Its diag will change)
    αStS_diag::AbstractArray{T} # Store the original diag of KK = βS'S
    αStb::AbstractArray{T}
    Σ::Matrix{T} # ΣQ Enforced to be dense
    v::AbstractArray{T}
    lb::AbstractArray{T}
    ub::AbstractArray{T}
end

function EPMat(S::AbstractArray{T}, b::AbstractArray{T},
        lb::AbstractArray{T}, ub::AbstractArray{T}, alpha::T) where T <: Real
    isinf(alpha) && error("For Inf alpha use 'EPMatT0'")
    M, N = size(S)
    αStS = sparse(alpha * S' * S)
    return EPMat(
        #= αStS      =# αStS,
        #= αStS_diag =# diag(αStS),
        #= αStb      =# alpha * S' * b,
        #= Σ         =# zeros(T,N,N),
        #= v         =# zeros(T,N),
        #= lb        =# lb,
        #= ub        =# ub
    )
end


text:
module vehicle_platoon_5
using ReachabilityAnalysis
A = [0 1 0 0 0 0 0 0 0 0 0 0 0 0 0;
    0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0;
    1.7152555329 3.9705119979 -4.3600526739 -0.9999330812 -1.5731541104 0.2669165553 -0.2215507198 -0.4303855023 0.0669078193 -0.0881500219 -0.1881468451 0.0322187056 -0.0343095071 -0.0767587194 0.0226660281;
    0 0 0 0 1 0 0 0 0 0 0 0 0 0 0;
    0 0 1 0 0 -1 0 0 0 0 0 0 0 0 0;
    0.7153224517 2.3973578876 0.2669165553 1.4937048131 3.5401264957 -4.2931448546 -1.0880831031 -1.7613009555 0.2991352608 -0.2558602268 -0.5071442217 0.0895738474 -0.0881500219 -0.1881468451 0.0548847337;
    0 0 0 0 0 0 0 1 0 0 0 0 0 0 0;
    0 0 0 0 0 1 0 0 -1 0 0 0 0 0 0;
    0.493771732 1.9669723853 0.0669078193 0.6271724298 2.2092110425 0.2991352608 1.4593953061 3.4633677762 -4.2704788265 -1.0880831031 -1.7613009555 0.3218012889 -0.2215507198 -0.4303855023 0.121792553;
    0 0 0 0 0 0 0 0 0 0 1 0 0 0 0;
    0 0 0 0 0 0 0 0 1 0 0 -1 0 0 0;
    0.40562171 1.7788255402 0.0322187056 0.4594622249 1.8902136659 0.0895738474 0.6271724298 2.2092110425 0.3218012889 1.4937048131 3.5401264957 -4.2382601209 -0.9999330812 -1.5731541104 0.3887091083;
    0 0 0 0 0 0 0 0 0 0 0 0 0 1 0;
    0 0 0 0 0 0 0 0 0 0 0 1 0 0 -1;
    0.371312203 1.7020668208 0.0226660281 0.40562171 1.7788255402 0.0548847337 0.493771732 1.9669723853 0.121792553 0.7153224517 2.3973578876 0.3887091083 1.7152555329 3.9705119979 -3.9713435656;
]

function model(X0)
    S = @system(x' = Ax)
    return IVP(S, X0)
end
end # module

max_stars_repo_name: mforets/ReachabilityLibrary

max_issues_repo_name: mforets/ReachabilityLibrary

max_forks_repo_name: mforets/ReachabilityLibrary


text:
using DelimitedFiles

x = readdlm("data_1.dat", '\n', Int, '\n')
diff = x[2:length(x)] - x[1:length(x)-1]
nincrease = length(diff[diff.>0])

print(nincrease)


text:
# Model hyperparameters

"""
RAPIDS Docs for PCA:
    https://docs.rapids.ai/api/cuml/stable/api.html#principal-component-analysis

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = PCA(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)
inverse_transform(mach, X)

println(mach.fitresult.components_)
```
"""
MLJModelInterface.@mlj_model mutable struct PCA <: MMI.Unsupervised
    handle = nothing
    copy::Bool = false # we are passing a numpy array, so modifying the data does not matter
    iterated_power::Int = 15::(_ > 0)
    n_components = nothing
    random_state = nothing
    svd_solver::String = "full"::(_ in ("auto", "full", "jacobi"))
    tol::Float64 = 1e-7::(_ > 0)
    verbose::Bool = false
end

"""
RAPIDS Docs for IncrementalPCA:
    https://docs.rapids.ai/api/cuml/stable/api.html#incremental-pca

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = IncrementalPCA(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)

println(mach.fitresult.components_)
```
"""
MLJModelInterface.@mlj_model mutable struct IncrementalPCA <: MMI.Unsupervised
    handle = nothing
    copy::Bool = false # we are passing a numpy array, so modifying the data does not matter
    whiten::Bool = false
    n_components = nothing
    batch_size = nothing
    verbose::Bool = false
end

"""
RAPIDS Docs for TruncatedSVD:
    https://docs.rapids.ai/api/cuml/stable/api.html#truncated-svd

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = TruncatedSVD(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)

println(mach.fitresult.components_)
```
"""
MLJModelInterface.@mlj_model mutable struct TruncatedSVD <: MMI.Unsupervised
    handle = nothing
    n_components = nothing
    n_iter::Int = 15::(_ > 0)
    random_state = nothing
    tol::Float64 = 1e-7::(_ > 0)
    verbose::Bool = false
end

"""
RAPIDS Docs for UMAP:
    https://docs.rapids.ai/api/cuml/stable/api.html#umap

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = UMAP(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)

println(mach.fitresult.embedding_)
```
"""
MLJModelInterface.@mlj_model mutable struct UMAP <: MMI.Unsupervised
    handle = nothing
    n_components = nothing
    n_epochs::Int = 15::(_ > 0)
    learning_rate::Float64 = 1.0::(_ > 0)
    init::String = "spectral"::(_ in ("random", "spectral"))
    min_dist::Float64 = 0.1::(_ > 0)
    spread::Float64 = 1.0::(_ > 0)
    set_op_mix_ratio::Float64 = 1.0::(_ >= 0 && _ <= 1)
    local_connectivity::Int = 1::(_ > 0)
    repulsion_strength::Float64 = 1.0::(_ >= 0)
    negative_sample_rate::Int = 5::(_ >= 0)
    transform_queue_size::Float64 = 4.0::(_ >= 0)
    a = nothing
    b = nothing
    hash_input::Bool = false
    random_state = nothing
    #optim_batch_size = nothing # TODO: This is in the cuML docs, but not an actual argument
    callback = nothing
    verbose::Bool = false
end

"""
RAPIDS Docs for GaussianRandomProjection:
    https://docs.rapids.ai/api/cuml/stable/api.html#random-projection

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = GaussianRandomProjection(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)
```
"""
MLJModelInterface.@mlj_model mutable struct GaussianRandomProjection <: MMI.Unsupervised
    handle = nothing
    n_components = nothing
    eps::Float64 = 0.1::(_ > 0)
    random_state = nothing
    verbose::Bool = false
end

"""
RAPIDS Docs for TSNE:
    https://docs.rapids.ai/api/cuml/stable/api.html#tsne

Example:
```
using RAPIDS
using MLJ

X = rand(100, 5)

model = TSNE(n_components=2)
mach = machine(model, X)
fit!(mach)
X_trans = transform(mach, X)

println(mach.fitresult.kl_divergence_)
```
"""
MLJModelInterface.@mlj_model mutable struct TSNE <: MMI.Unsupervised
    handle = nothing
    n_components::Int = 2::(_ > 0)
    perplexity::Float64 = 30.0::(_ > 0)
    early_exaggeration::Float64 = 12.0::(_ > 0)
    late_exaggeration::Float64 = 1.0::(_ > 0)
    learning_rate::Float64 = 200.0::(_ > 0)
    n_iter::Int = 1000::(_ > 0)
    n_iter_without_progress::Int = 300::(_ > 0)
    min_grad_norm::Float64 = 1e-7::(_ > 0)
    metric::String = "euclidean"::(_ in ("euclidean", ))
    init::String = "random"::(_ in ("random", ))
    method::String = "fft"::(_ in ("barnes_hut", "exact", "fft"))
    angle::Float64 = 0.5::(_ >= 0 && _ <= 1)
    learning_rate_method = "adaptive"
    n_neighbors::Int = 90::(_ > 0)
    perplexity_max_iter::Int = 100::(_ >= 0)
    exaggeration_iter::Int = 250::(_ >= 0)
    pre_momentum::Float64 = 0.5::(_ >= 0)
    post_momentum::Float64 = 0.8::(_ >= 0)
    square_distances::Bool = false
    random_state = nothing
    verbose::Bool = false
end

# Multiple dispatch for initializing models
model_init(mlj_model::PCA) = cuml.PCA(; mlj_to_kwargs(mlj_model)...)
model_init(mlj_model::IncrementalPCA) = cuml.IncrementalPCA(; mlj_to_kwargs(mlj_model)...)
model_init(mlj_model::TruncatedSVD) = cuml.TruncatedSVD(; mlj_to_kwargs(mlj_model)...)
model_init(mlj_model::UMAP) = cuml.UMAP(; mlj_to_kwargs(mlj_model)...)
model_init(mlj_model::GaussianRandomProjection) = cuml.GaussianRandomProjection(; mlj_to_kwargs(mlj_model)...)
model_init(mlj_model::TSNE) = cuml.TSNE(; mlj_to_kwargs(mlj_model)...)

const CUML_DIMENSIONALITY_REDUCTION = Union{PCA,
                                            IncrementalPCA,
                                            TruncatedSVD,
                                            UMAP,
                                            GaussianRandomProjection,
                                            TSNE}

# add metadata
MMI.load_path(::Type{<:PCA}) = "$PKG.PCA"
MMI.load_path(::Type{<:IncrementalPCA}) = "$PKG.IncrementalPCA"
MMI.load_path(::Type{<:TruncatedSVD}) = "$PKG.TruncatedSVD"
MMI.load_path(::Type{<:UMAP}) = "$PKG.UMAP"
MMI.load_path(::Type{<:HDBSCAN}) = "$PKG.HDBSCAN"
MMI.load_path(::Type{<:GaussianRandomProjection}) = "$PKG.GaussianRandomProjection"

MMI.input_scitype(::Type{<:CUML_DIMENSIONALITY_REDUCTION}) = Union{AbstractMatrix, Table(Continuous)}

MMI.docstring(::Type{<:PCA}) = "cuML's PCA: https://docs.rapids.ai/api/cuml/stable/api.html#principal-component-analysis"
MMI.docstring(::Type{<:IncrementalPCA}) = "cuML's IncrementalPCA: https://docs.rapids.ai/api/cuml/stable/api.html#incremental-pca"
MMI.docstring(::Type{<:TruncatedSVD}) = "cuML's TruncatedSVD: https://docs.rapids.ai/api/cuml/stable/api.html#truncated-svd"
MMI.docstring(::Type{<:UMAP}) = "cuML's UMAP: https://docs.rapids.ai/api/cuml/stable/api.html#umap"
MMI.docstring(::Type{<:GaussianRandomProjection}) = "cuML's GaussianRandomProjection: https://docs.rapids.ai/api/cuml/stable/api.html#random-projection"
MMI.docstring(::Type{<:TSNE}) = "cuML's TSNE: https://docs.rapids.ai/api/cuml/stable/api.html#tsne"

function MMI.fit(mlj_model::CUML_DIMENSIONALITY_REDUCTION, verbosity, X, w=nothing)
    # fit model
    model = model_init(mlj_model)
    model.fit(prepare_input(X))
    fitresult = model

    # save result
    cache = nothing
    report = ()
    return (fitresult, cache, report)
end

# transform methods
function MMI.transform(mlj_model::Union{PCA, IncrementalPCA,
                                        UMAP, TruncatedSVD,
                                        GaussianRandomProjection},
                        fitresult, Xnew)
    model  = fitresult
    py_preds = model.transform(prepare_input(Xnew))
    preds = pyconvert(Array, py_preds)

    return preds
end

# todo figure out how to handle TSNE's lack of transform
# right now we have to refit the model when transforming
function MMI.transform(mlj_model::TSNE, fitresult, Xnew)
    model  = fitresult
    py_preds = model.fit_transform(prepare_input(Xnew))
    preds = pyconvert(Array, py_preds)

    return preds
end

function MMI.inverse_transform(mlj_model::Union{PCA, TruncatedSVD}, fitresult, Xnew)
    model = fitresult
    py_preds = model.inverse_transform(prepare_input(Xnew))
    preds = pyconvert(Array, py_preds)

    return preds
end

# Clustering metadata
MMI.metadata_pkg.((PCA, IncrementalPCA, TruncatedSVD, UMAP, GaussianRandomProjection, TSNE),
    name = "cuML Dimensionality Reduction and Manifold Learning Methods",
    uuid = "2764e59e-7dd7-4b2d-a28d-ce06411bac13", # see your Project.toml
    url  = "https://github.com/tylerjthomas9/RAPIDS.jl",  # URL to your package repo
    julia = false,          # is it written entirely in Julia?
    license = "MIT",        # your package license
    is_wrapper = true,      # does it wrap around some other package?
)


text:
#############################################
# Van der Pol Oscillator Controlled by RFPT #
#############################################
Adaptive=1
######################
# Control Parameters #
######################
K=1e5
B=-1
A=1e-5

#########
# Time  #
#########
δt=1e-3
LONG=Int(2e4)

########################################
# Kinematic Block Parameter (2nd order)#
########################################
Λ=3

####################################
# Parameters for Nominal Trajectory#
####################################

ω=0.5
Ampl=2
##########################
# Exact Model Parameters #
##########################

μₑ=0.4
ωₑ=0.46
αₑ=1
λₑ=0.1
mₑ=1

################################
# Approximate Model Parameters #
################################

μₐ=0.5
ωₐ=0.42
αₐ=0.9
λₐ=0.09
mₐ=0.8

############################
# Functions for simulation #
############################

# function nom_traj(t)
#   #egys=[1,-1.5,2]
#   qN=Ampl*sin(ω*δt*t)
#   qN_p=Ampl*ω*cos(ω*δt*t)
#   qN_pp=-Ampl*ω^2*sin(ω*δt*t)
# 	#println(qN[t,:])
#   return [qN;qN_p;qN_pp]
# end

#######################
# The Kinematic Block #
#######################

# function Kinematic_Block(t,err_int)
# 	#Compute the Error
# 	err=(qN-q)#'
# 	#Compute the 1st time derivative of the Error
# 	err_p=(qN_p-q_p)#'
# 	#Compute te integral of the Error (with Euler formula)
#    #println(err)
# 	err_i=err_int+err*δt
#   #println(q,q_p,q_pp)
# 	#Compute the Desired acceleration (2nd time derivative)
# 	#out=qN_pp[t]'+3*Λ^2*err+3*Λ*err_p+Λ^3*err_i
#   out=qN_pp+3*Λ^2*err+3*Λ*err_p+Λ^3*err_i #Has to be colum vector
#   #println(out)
# 	#println(err_i)
#    #println(typeof(err_i))
# 	return [out err_i]
# end

###############
# Exact Model #
###############
function Exact(u,q,q_p)
  q_pp=(u+μₑ*(1-q^2)*q_p-ωₑ^2*q-αₑ*q^3-λₑ^5)/mₑ
  return q_pp
end

################
# Approx Model #
################

function Approx(q,q_p,q_pp)
  u=mₐ*q_pp-μₐ*(1-q^2)*q_p+ωₐ^2*q+αₐ*q^3+λₐ^5
  return u
end

function sigmoid(x)
	s=x/(1+abs(x))
  return s
end

function G(past_input,past_response,xDnow)#(f,x,xᵈ)
  out=(K+past_input) * (1 + B * sigmoid(A * (past_response - xDnow))) - K
  return out #q_ppReq
end

##############################
# Define arrays for Plotting #
##############################
#time
time_mem=zeros(LONG)
#nominal trajectory
#qN_mem=zeros(LONG)
#qN_p_mem=zeros(LONG)
#qN_pp_mem=zeros(LONG)

#q_ppDes_mem=zeros(LONG)
#x=0
#f=0
#q_ppDef_mem=zeros(LONG)
#Q_mem=zeros(LONG)
#err_int=0
#qN=0
#qN_p=0
#qN_pp=0
#q=0
#q_p=0
#q_pp=0
#q_mem=zeros(LONG)
##################
# The simulation #
##################
q_mem=zeros(LONG) #A megvalósult pálya
q_p_mem=zeros(LONG)

qN_mem=zeros(LONG) #A nominális pálya
qN_p_mem=zeros(LONG)

u_mem=zeros(LONG) #A szanályozó jel elmentése

qN_pp_mem=zeros(LONG) #A nominális gyorsulások értéke
q_pp_Des_mem=zeros(LONG)#A PID korrekciós adatok
q_pp_Req_mem=zeros(LONG) #Az Adaptívan torzított jelre
q_pp_mem=zeros(LONG) #Az Adaptívan torzított jelre

h_int=0 #A követési hiba integrálja

past_input=0
past_response=0

for t=1:LONG-1
  #store time for plotting
  time_mem[t]=t*δt
  #define the nominal trajectory
  qN_mem[t]=Ampl*sin(ω*δt*t)
  qN_p_mem[t]=ω*Ampl*cos(ω*δt*t)
  qN_pp=-ω^2*Ampl*sin(ω*δt*t)
  qN_pp_mem[t]=qN_pp
  #the kinematic
  q_ppDes=Λ^3*h_int+3*Λ^2*(qN_mem[t]-q_mem[t])+3*Λ*(qN_p_mem[t]-q_p_mem[t])+qN_pp
  q_pp_Des_mem[t]=q_ppDes
  if Adaptive==1 && t>10
    q_ppReq=G(past_input,past_response,q_ppDes)
  else
    q_ppReq=q_ppDes
  end
  q_pp_Req_mem[t]=q_ppReq
  u=Approx(q_mem[t],q_p_mem[t],q_ppReq)
  u_mem[t]=u
  #println(Q)
  past_input=q_ppReq

  q_pp=Exact(u,q_mem[t],q_p_mem[t])
  #println(x)
  #println(x)
  q_pp_mem[t]=q_pp

  past_response=q_pp
  q_p_mem[t+1]=q_p_mem[t]+δt*q_pp
  q_mem[t+1]=q_mem[t]+δt*q_p_mem[t]
  h_int=h_int+δt*(qN_mem[t]-q_mem[t])

end #for t=1:LONG-1

using(PyPlot)
figure(1)
grid("on")
title("Tracking Error vs Time")
xlabel("Time [s]")
ylabel("Tracking Error")
plot(time_mem[11:LONG-1],qN_mem[11:LONG-1]-q_mem[11:LONG-1])

figure(2)
grid("on")
title("Nominal and Realized Trajectory")
xlabel("Time [s]")
ylabel("Nominal and Realized Trajectory")
plot(time_mem[3:LONG-1],qN_mem[3:LONG-1],color="red")
plot(time_mem[3:LONG-1],q_mem[3:LONG-1],color="green")

max_stars_repo_path: examples/RFPT_SISO/van der Pol Oscillator/van_der_Pol_RFPT.jl

max_stars_repo_name: krisztiankosi/Control

max_issues_repo_name: krisztiankosi/Control

max_forks_repo_name: krisztiankosi/Control


text:
using ExtensibleMCMCPlots
using Test

@testset "ExtensibleMCMCPlots.jl" begin
    # Write your own tests here.
end


text:
@testset "static assignment to/from array" begin
    submap = StaticChoiceMap((a=1., b=[2., 2.5]),NamedTuple())
    outer = StaticChoiceMap((c=3.,), (d=submap, e=submap))

    arr = to_array(outer, Float64)
    @test to_array(outer, Float64) == Float64[3.0, 1.0, 2.0, 2.5, 1.0, 2.0, 2.5]

    choices = from_array(outer, Float64[1, 2, 3, 4, 5, 6, 7])
    @test choices[:c] == 1.0
    @test choices[:d => :a] == 2.0
    @test choices[:d => :b] == [3.0, 4.0]
    @test choices[:e => :a] == 5.0
    @test choices[:e => :b] == [6.0, 7.0]
    @test length(collect(get_submaps_shallow(choices))) == 2
    @test length(collect(get_values_shallow(choices))) == 1
    submap1 = get_submap(choices, :d)
    @test length(collect(get_values_shallow(submap1))) == 2
    @test length(collect(get_submaps_shallow(submap1))) == 0
    submap2 = get_submap(choices, :e)
    @test length(collect(get_values_shallow(submap2))) == 2
    @test length(collect(get_submaps_shallow(submap2))) == 0
end

@testset "dynamic assignment to/from array" begin
    outer = choicemap()
    set_value!(outer, :c, 3.)
    submap = choicemap()
    set_value!(submap, :a, 1.)
    set_value!(submap, :b, [2., 2.5])
    set_submap!(outer, :d, submap)
    set_submap!(outer, :e, submap)

    arr = to_array(outer, Float64)
    @test to_array(outer, Float64) == Float64[3.0, 1.0, 2.0, 2.5, 1.0, 2.0, 2.5]

    choices = from_array(outer, Float64[1, 2, 3, 4, 5, 6, 7])
    @test choices[:c] == 1.0
    @test choices[:d => :a] == 2.0
    @test choices[:d => :b] == [3.0, 4.0]
    @test choices[:e => :a] == 5.0
    @test choices[:e => :b] == [6.0, 7.0]
    @test length(collect(get_submaps_shallow(choices))) == 2
    @test length(collect(get_values_shallow(choices))) == 1
    submap1 = get_submap(choices, :d)
    @test length(collect(get_values_shallow(submap1))) == 2
    @test length(collect(get_submaps_shallow(submap1))) == 0
    submap2 = get_submap(choices, :e)
    @test length(collect(get_values_shallow(submap2))) == 2
    @test length(collect(get_submaps_shallow(submap2))) == 0
end

@testset "dynamic assignment copy constructor" begin
    other = choicemap()
    other[:x] = 1
    other[:y] = 2
    other_nested = choicemap()
    other_nested[:z] = 3
    other_nested[:w] = 4
    set_submap!(other, :u, other_nested)
    choices = DynamicChoiceMap(other)
    @test choices[:x] == 1
    @test choices[:y] == 2
    @test choices[:u => :z] == 3
    @test choices[:u => :w] == 4
end

@testset "internal vector assignment to/from array" begin
    inner = choicemap()
    set_value!(inner, :a, 1.)
    set_value!(inner, :b, 2.)
    outer = vectorize_internal([inner, inner, inner])

    arr = to_array(outer, Float64)
    @test to_array(outer, Float64) == Float64[1, 2, 1, 2, 1, 2]

    choices = from_array(outer, Float64[1, 2, 3, 4, 5, 6])
    @test choices[1 => :a] == 1.0
    @test choices[1 => :b] == 2.0
    @test choices[2 => :a] == 3.0
    @test choices[2 => :b] == 4.0
    @test choices[3 => :a] == 5.0
    @test choices[3 => :b] == 6.0
    @test length(collect(get_submaps_shallow(choices))) == 3
end

@testset "dynamic assignment merge" begin
    submap = choicemap()
    set_value!(submap, :x, 1)
    choices1 = choicemap()
    set_value!(choices1, :a, 1.)
    set_value!(choices1, :b, 2.)
    set_submap!(choices1, :c, submap)
    set_submap!(choices1, :shared, submap)
    choices2 = choicemap()
    set_value!(choices2, :d, 3.)
    set_submap!(choices2, :e, submap)
    set_submap!(choices2, :f, submap)
    submap2 = choicemap()
    set_value!(submap2, :y, 4.)
    set_submap!(choices2, :shared, submap2)
    choices = merge(choices1, choices2)
    @test choices[:a] == 1.
    @test choices[:b] == 2.
    @test choices[:d] == 3.
    @test choices[:c => :x] == 1
    @test choices[:e => :x] == 1
    @test choices[:f => :x] == 1
    @test choices[:shared => :x] == 1
    @test choices[:shared => :y] == 4.
    @test length(collect(get_submaps_shallow(choices))) == 4
    @test length(collect(get_values_shallow(choices))) == 3
end

@testset "dynamic assignment variadic merge" begin
    choices1 = choicemap((:a, 1))
    choices2 = choicemap((:b, 2))
    choices3 = choicemap((:c, 3))
    choices_all = choicemap((:a, 1), (:b, 2), (:c, 3))
    @test merge(choices1) == choices1
    @test merge(choices1, choices2, choices3) == choices_all
end

@testset "static assignment merge" begin
    submap = choicemap()
    set_value!(submap, :x, 1)
    submap2 = choicemap()
    set_value!(submap2, :y, 4.)
    choices1 = StaticChoiceMap((a=1., b=2.), (c=submap, shared=submap))
    choices2 = StaticChoiceMap((d=3.,), (e=submap, f=submap, shared=submap2))
    choices = merge(choices1, choices2)
    @test choices[:a] == 1.
    @test choices[:b] == 2.
    @test choices[:d] == 3.
    @test choices[:c => :x] == 1
    @test choices[:e => :x] == 1
    @test choices[:f => :x] == 1
    @test choices[:shared => :x] == 1
    @test choices[:shared => :y] == 4.
    @test length(collect(get_submaps_shallow(choices))) == 4
    @test length(collect(get_values_shallow(choices))) == 3
end

@testset "static assignment variadic merge" begin
    choices1 = StaticChoiceMap((a=1,), NamedTuple())
    choices2 = StaticChoiceMap((b=2,), NamedTuple())
    choices3 = StaticChoiceMap((c=3,), NamedTuple())
    choices_all = StaticChoiceMap((a=1, b=2, c=3), NamedTuple())
    @test merge(choices1) == choices1
    @test merge(choices1, choices2, choices3) == choices_all
end

@testset "static assignment errors" begin

    # get_choices on an address that contains a value throws a KeyError
    choices = StaticChoiceMap((x=1,), NamedTuple())
    threw = false
    try get_submap(choices, :x) catch KeyError threw = true end
    @test threw

    # static_get_submap on an address that contains a value throws a KeyError
    choices = StaticChoiceMap((x=1,), NamedTuple())
    threw = false
    try static_get_submap(choices, Val(:x)) catch KeyError threw = true end
    @test threw

    # get_choices on an address whose prefix contains a value throws a KeyError
    choices = StaticChoiceMap((x=1,), NamedTuple())
    threw = false
    try get_submap(choices, :x => :y) catch KeyError threw = true end
    @test threw

    # static_get_choices on an address whose prefix contains a value throws a KeyError
    choices = StaticChoiceMap((x=1,), NamedTuple())
    threw = false
    try static_get_submap(choices, Val(:x)) catch KeyError threw = true end
    @test threw

    # get_choices on an address that contains nothing gives empty assignment
    choices = StaticChoiceMap(NamedTuple(), NamedTuple())
    @test isempty(get_submap(choices, :x))
    @test isempty(get_submap(choices, :x => :y))

    # static_get_choices on an address that contains nothing throws a KeyError
    choices = StaticChoiceMap(NamedTuple(), NamedTuple())
    threw = false
    try static_get_submap(choices, Val(:x)) catch KeyError threw = true end
    @test threw

    # get_value on an address that contains a submap throws a KeyError
    submap = choicemap()
    submap[:y] = 1
    choices = StaticChoiceMap(NamedTuple(), (x=submap,))
    threw = false
    try get_value(choices, :x) catch KeyError threw = true end
    @test threw

    # static_get_value on an address that contains a submap throws a KeyError
    submap = choicemap()
    submap[:y] = 1
    choices = StaticChoiceMap(NamedTuple(), (x=submap,))
    threw = false
    try static_get_value(choices, Val(:x)) catch KeyError threw = true end
    @test threw

    # get_value on an address that contains nothing throws a KeyError
    choices = StaticChoiceMap(NamedTuple(), NamedTuple())
    threw = false
    try get_value(choices, :x) catch KeyError threw = true end
    @test threw
    threw = false
    try get_value(choices, :x => :y) catch KeyError threw = true end
    @test threw

    # static_get_value on an address that contains nothing throws a KeyError
    choices = StaticChoiceMap(NamedTuple(), NamedTuple())
    threw = false
    try static_get_value(choices, Val(:x)) catch KeyError threw = true end
    @test threw
end

@testset "dynamic assignment errors" begin

    # get_choices on an address that contains a value throws a KeyError
    choices = choicemap()
    choices[:x] = 1
    threw = false
    try get_submap(choices, :x) catch KeyError threw = true end
    @test threw

    # get_choices on an address whose prefix contains a value throws a KeyError
    choices = choicemap()
    choices[:x] = 1
    threw = false
    try get_submap(choices, :x => :y) catch KeyError threw = true end
    @test threw

    # get_choices on an address that contains nothing gives empty assignment
    choices = choicemap()
    @test isempty(get_submap(choices, :x))
    @test isempty(get_submap(choices, :x => :y))

    # get_value on an address that contains a submap throws a KeyError
    choices = choicemap()
    choices[:x => :y] = 1
    threw = false
    try get_value(choices, :x) catch KeyError threw = true end
    @test threw

    # get_value on an address that contains nothing throws a KeyError
    choices = choicemap()
    threw = false
    try get_value(choices, :x) catch KeyError threw = true end
    @test threw
    threw = false
    try get_value(choices, :x => :y) catch KeyError threw = true end
    @test threw
end

@testset "dynamic assignment overwrite" begin

    # overwrite value with a value
    choices = choicemap()
    choices[:x] = 1
    choices[:x] = 2
    @test choices[:x] == 2

    # overwrite value with a submap
    choices = choicemap()
    choices[:x] = 1
    @test has_value(choices, :x)
    @test !has_submap(choices, :x)
    submap = choicemap(); submap[:y] = 2
    set_submap!(choices, :x, submap)
    @test !has_value(choices, :x)
    @test !isempty(get_submap(choices, :x))
    @test has_submap(choices, :x)

    # overwrite subassignment with a value
    choices = choicemap()
    choices[:x => :y] = 1
    choices[:x] = 2
    threw = false
    @test !has_submap(choices, :x)
    try get_submap(choices, :x) catch KeyError threw = true end
    @test threw
    @test choices[:x] == 2

    # overwrite subassignment with a subassignment
    choices = choicemap()
    choices[:x => :y] = 1
    @test has_submap(choices, :x)
    @test !has_submap(choices, :x => :y)
    submap = choicemap(); submap[:z] = 2
    set_submap!(choices, :x,  submap)
    @test !isempty(get_submap(choices, :x))
    @test !has_value(choices, :x => :y)
    @test choices[:x => :z] == 2

    # illegal set value under existing value
    choices = choicemap()
    choices[:x] = 1
    threw = false
    try set_value!(choices, :x => :y, 2) catch KeyError threw = true end
    @test threw

    # illegal set submap under existing value
    choices = choicemap()
    choices[:x] = 1
    submap = choicemap(); choices[:z] = 2
    threw = false
    try set_submap!(choices, :x => :y, submap) catch KeyError threw = true end
    @test threw
end

@testset "dynamic assignment constructor" begin

    choices = choicemap((:x, 1), (:y => :a, 2))
    @test choices[:x] == 1
    @test choices[:y => :a] == 2
end

@testset "choice map equality" begin

    a = choicemap((:a, 1), (:b => :c => 2), (:d => :e => :f => 3))
    b = choicemap((:a, 1), (:d => :e => :f => 3))
    c = choicemap((:a, 1), (:b => :c => 2), (:d => :e => :f => 4))

    @test a == a
    @test b == b
    @test a != b
    @test b != a
    @test a != c
    @test c != a

end

@testset "choice map nested view" begin
    c = choicemap((:a, 1),
                  (:b => :c, 2))
    cv = nested_view(c)
    @test c[:a] == cv[:a]
    @test c[:b => :c] == cv[:b][:c]
    # Base.length
    @test length(cv) == 2
    @test length(cv[:b]) == 1
    # Base.keys
    @test Set(keys(cv)) == Set([:a, :b])
    @test Set(keys(cv[:b])) == Set([:c])
    # Base.:(==)
    c1 = choicemap((:a, 1),
                   (:b => :c, 2))
    c2 = choicemap((:a, 4),
                   (:b => :c, 2))
    @test nested_view(c1) == cv
    @test nested_view(c2) != cv
end

@testset "filtering choicemaps with selections" begin

    c = choicemap((:a, 1), (:b, 2))

    filtered = get_selected(c, select(:a))
    @test filtered[:a] == 1
    @test !has_value(filtered, :b)

    c = choicemap((:x => :y, 1), (:x => :z, 2))

    filtered = get_selected(c, select(:x))
    @test filtered[:x => :y] == 1
    @test filtered[:x => :z] == 2

    filtered = get_selected(c, select(:x => :y))
    @test filtered[:x => :y] == 1
    @test !has_value(filtered, :x => :z)
end

@testset "invalid choice map constructor" begin
    threw = false
    try c = choicemap((:a, 1, :b, 2)) catch Exception threw = true end
    @test threw
end


text:
#=
capacitated facility location problem with polynomial flows over continuous time

time is t ∈ [0, 1]
P are polynomials in t of max degree deg ≥ 1, and P₊ ⊂ P are nonnegative on [0, 1]

sets:
i ∈ 1..n   facilities
j ∈ 1..m   customers

parameters:
fᵢ  ∈ ℝ₊  fixed cost of opening i
cᵢⱼ ∈ ℝ₊  cost (per unit of flow) to ship from i to j
dⱼ  ∈ P₊  demand at j
uᵢ  ∈ P₊  maximum output of i

variables:
xᵢ  ∈ {0,1}  whether i is opened
yᵢⱼ ∈ P₊     flow from i to j

objective:
min ∑ᵢⱼ cᵢⱼ ∫ yᵢⱼ + ∑ᵢ fᵢ xᵢ   minimize total cost

constraints:
∑ⱼ yᵢⱼ ≤ uᵢ xᵢ, ∀i   limit production at i
∑ᵢ yᵢⱼ ≥ dⱼ, ∀j      meet demand at j
=#

struct PolyFacilityLocation <: ExampleInstance
    n::Int
    m::Int
    halfdeg::Int
    use_nat::Bool # use WSOS cone formulation, else SDP formulation
end

function build(inst::PolyFacilityLocation)
    # setup polynomial interpolation
    dom = PolyUtils.BoxDomain{Float64}([0.0], [1.0])
    (U, _, Ps, _, w) = PolyUtils.interpolate(dom, inst.halfdeg, get_quadr = true)

    # generate parameter values
    (n, m) = (inst.n, inst.m)
    f = rand(n)
    c = rand(n, m)
    d = make_nonneg_polys(m, Ps)
    u = make_nonneg_polys(n, Ps)

    # make feasible by ensuring sum(u) > 2 * sum(d)
    excess = find_poly_min(sum(u) - 2 * sum(d), Ps)
    incr = 0.1 - excess / n
    for i in 1:n
        u[i] .+= incr + 5 * rand() + 1
    end
    excess = find_poly_min(sum(u) - 2 * sum(d), Ps)
    @assert excess > 0

    # build model
    model = JuMP.Model()
    JuMP.@variable(model, x[1:n], Bin)
    JuMP.@variable(model, y[1:n, 1:m, 1:U])

    JuMP.@objective(
        model,
        Min,
        sum(JuMP.dot(c[i, j] * w, y[i, j, :]) for j in 1:m, i in 1:n) + JuMP.dot(f, x)
    )

    add_wsos(aff) = (inst.use_nat ? add_wsos_nat : add_wsos_ext)(Ps, aff, model)

    for i in 1:n, j in 1:m
        add_wsos(y[i, j, :])
    end

    JuMP.@expression(model, prod[i in 1:n], u[i] * x[i] - sum(y[i, j, :] for j in 1:m))
    for i in 1:n
        add_wsos(prod[i])
    end

    JuMP.@expression(model, dem[j in 1:m], sum(y[i, j, :] for i in 1:n) - d[j])
    for j in 1:m
        add_wsos(dem[j])
    end

    # save for use in tests
    model.ext[:x] = x
    model.ext[:y] = y
    model.ext[:prod] = prod
    model.ext[:dem] = dem
    model.ext[:Ps] = Ps

    return model
end

function test_extra(inst::PolyFacilityLocation, model::JuMP.Model)
    stat = JuMP.termination_status(model)
    @test stat in OPT_OR_LIMIT
    (stat == MOI.OPTIMAL) || return

    tol = eps()^0.2
    x = JuMP.value.(model.ext[:x])
    @test x ≈ round.(Int, x) atol = tol rtol = tol
    @test all(-tol .<= x .<= 1 + tol)
    # check WSOS feasibility
    Ps = model.ext[:Ps]
    y = JuMP.value.(model.ext[:y])
    @test all(check_nonneg(y[i, j, :] .+ tol, Ps) for i in 1:(inst.n), j in 1:(inst.m))
    @test all(check_nonneg(JuMP.value.(p_i) .+ tol, Ps) for p_i in model.ext[:prod])
    @test all(check_nonneg(JuMP.value.(d_j) .+ tol, Ps) for d_j in model.ext[:dem])
    return
end


text:
using Makie
using GeoInterface, GeoJSON

states = download("https://raw.githubusercontent.com/PublicaMundi/MappingAPI/master/data/geojson/us-states.json")

states_geo = GeoJSON.parse(read(states, String))

features = states_geo.features

props = getproperty.(features, :properties)

densities = getindex.(props, "density")

cr = let ds = sort(densities); (ds[1], ds[end-1]); end

Base.convert(poly::GeoInterface.Polygon, ::Type{Point2f0}) = Point2f0.(poly.coordinates[1])

using PlotUtils

viridis = cgrad(:viridis; scale = :log10)

convert_arguments(::Poly, poly::Polygon) = Point2f0.(poly.coordinates[1])

# convert_arguments(::Poly, mp::MultiPolygon) =

function plo(sc::Scene, poly::Polygon; kwargs...)
    poly!(sc, convert(poly, Point2f0); kwargs...)
end

function plo(sc, poly::MultiPolygon; kwargs...)
    # map(x -> poly!(sc, Point2f0.(x[1]); kwargs...), poly.coordinates)
    data = map(x -> Point2f0.(x[1]), poly.coordinates)
    poly!(sc, data; kwargs...)
    sc
end

Point2{T}(a::Vector{T}) where T = Point2(a[1], a[2])

toPointVecs(poly::Polygon) = Vector{Point2}[Point2.(poly.coordinates[1])]
toPointVecs(mp::MultiPolygon) = map(x -> Point2.(x[1]), mp.coordinates)

sc = Scene()

normalize(val::Real; min = cr[1], max = cr[2]) = (val - min) / (max - min)

for (i, feature) in enumerate(features)
    println(i)
    try
        plo(sc, feature.geometry, color = viridis[normalize(feature.properties["density"])])
    catch err
        @warn("Feature had dimension zero!")
        print(err)
        pop!(sc.plots)
    end
    # display(sc)
end

save("usa-pop-density-log.png", sc)

###############

_getsize(::Polygon) = 1
_getsize(mp::MultiPolygon) = length(mp.coordinates)
_getsize(feature::Feature) = _getsize(f.geometry)
_getsize(fc::FeatureCollection) = sum(_getsize.(fc.features))

cs = Vector{Point2f0}[]
props = Dict[]

for feature in statesgj.features
    coords = toPointVecs(feature.geometry)
    l = length(coords)
    for coord in coords
        push!(cs, coord)
        push!(props, feature.properties)
    end
end

zerodiff(pt::GeometryTypes.Point{2, T}) where {T} = pt[1] != zero(T) && pt[2] != zero(T)

function imprecise(arr)
    dif = Point2{Float32}.(arr) |> diff
    return !any(zerodiff.(dif))
end

inds = findall(imprecise, cs)

if inds != nothing
    @warn("Imprecise points!")
end

deleteat!(cs, inds)
deleteat!(props, inds)


text:
@service S3
struct MinioConfig <: AbstractAWSConfig
    endpoint::String
    region::String
    creds
end
struct SimpleCredentials
    access_key_id::String
    secret_key::String
    token::String
end

AWS.region(c::MinioConfig) = c.region
AWS.credentials(c::MinioConfig) = c.creds
AWS.check_credentials(c::SimpleCredentials) = c

function AWS.generate_service_url(aws::MinioConfig, service::String, resource::String)
    service == "s3" || throw(ArgumentError("Can only handle s3 requests to Minio"))
    return string(aws.endpoint, resource)
end

AWS.global_aws_config(
    MinioConfig(
        "http://127.0.0.1:9000",
        ENV["MINIO_REGION_NAME"],
        SimpleCredentials(ENV["MINIO_ACCESS_KEY"], ENV["MINIO_SECRET_KEY"], ""),
    ),
)
datadir = joinpath(dirname(pathof(AWS)), "..", "data")

# Create bucket and place objects in it
S3.create_bucket("anewbucket")
try
    S3.put_object("anewbucket", "myobject", Dict("body" => "Hi from Minio"))
    S3.put_object("anewbucket", "empty")
    S3.put_object("anewbucket", "foo/bar", Dict("body" => "a nested object"))
    S3.put_object("anewbucket", "foo/baz", Dict("body" => "a secondnested object"))

    # Test retrieving an object
    @test String(S3.get_object("anewbucket", "myobject")) == "Hi from Minio"

    # Test retrieving an object into a stream target
    mktemp() do f, io
        S3.get_object("anewbucket", "myobject", Dict("response_stream" => io))
        flush(io)
        @test read(f, String) == "Hi from Minio"
    end

    # Test listing
    objs = S3.list_objects_v2("anewbucket")
    @test length(objs["Contents"]) == 4

    # Test api version 2 of list-objects
    objs_truncated = S3.list_objects_v2("anewbucket", Dict("max-keys" => 2))
    @test length(objs_truncated["Contents"]) == 2
    @test objs_truncated["IsTruncated"] == "true"
    @test haskey(objs_truncated, "NextContinuationToken")

    # Test listing with prefixes
    objs_prefix = S3.list_objects_v2("anewbucket", Dict("prefix" => "", "delimiter" => "/"))
    @test length(objs_prefix["Contents"]) == 2
    @test sort(getindex.(objs_prefix["Contents"], "Key")) == ["empty", "myobject"]
    @test objs_prefix["CommonPrefixes"]["Prefix"] == "foo/"

    # Duplicated testset from "test/issues.jl". Useful for testing outside the CI. Ideally,
    # the tests should be revised such that local testing works without having to duplicate
    # testsets.

    @testset "issue 466" begin
        file_name = "hang.txt"

        try
            S3.put_object("anewbucket", file_name)

            # Note: Using `eof` for these tests can hang when using an unclosed `Base.BufferStream`

            stream = S3.get_object("anewbucket", file_name, Dict("return_stream" => true))
            if AWS.DEFAULT_BACKEND[] isa AWS.HTTPBackend
                @test !isopen(stream)
            else
                @test isopen(stream)
            end

            stream = Base.BufferStream()
            S3.get_object("anewbucket", file_name, Dict("response_stream" => stream))
            if AWS.DEFAULT_BACKEND[] isa AWS.HTTPBackend
                @test !isopen(stream)
            else
                @test_broken isopen(stream)
            end

            stream = Base.BufferStream()
            S3.get_object(
                "anewbucket",
                file_name,
                Dict("response_stream" => stream, "return_stream" => true),
            )
            if AWS.DEFAULT_BACKEND[] isa AWS.HTTPBackend
                @test !isopen(stream)
            else
                @test isopen(stream)
            end
        finally
            S3.delete_object("anewbucket", file_name)
        end
    end

    @testset "issue 474" begin
        body = "foo\0bar"
        expected = Vector{UInt8}(body)
        file_name = "null.txt"
        bucket_name = "anewbucket"

        try
            S3.put_object(bucket_name, file_name, Dict("body" => body))

            raw = S3.get_object(bucket_name, file_name, Dict("return_raw" => true))
            @test raw isa Vector{UInt8}
            @test raw == expected

            stream = S3.get_object(bucket_name, file_name, Dict("return_stream" => true))
            if AWS.DEFAULT_BACKEND[] isa AWS.HTTPBackend
                @test stream isa Base.BufferStream
                @test !isopen(stream)

                if !isopen(stream)
                    @test read(stream) == expected
                end
            else
                @test stream isa IOBuffer
                @test isopen(stream)
                seekstart(stream)
                @test read(stream) == expected
            end
        finally
            S3.delete_object(bucket_name, file_name)
        end
    end

finally
    # Delete all objects and the bucket
    objs = S3.list_objects_v2("anewbucket")

    for obj in objs["Contents"]
        S3.delete_object("anewbucket", obj["Key"])
    end

    S3.delete_bucket("anewbucket")
end

@test isempty(filter(!isequal(".minio.sys"), readdir(datadir)))


text:
# ECOS orders differently than MOI the second and third dimension of the exponential cone

struct PermutedExponentialCone <: MOI.AbstractVectorSet end
Base.copy(set::PermutedExponentialCone) = set

const MOIB = MOI.Bridges

struct PermutedExponentialBridge{T,F} <: MOIB.Constraint.SetMapBridge{
    T,
    PermutedExponentialCone,
    MOI.ExponentialCone,
    F,
    F,
}
    constraint::MOI.ConstraintIndex{F,PermutedExponentialCone}
end

function MOIB.Constraint.concrete_bridge_type(
    ::Type{<:PermutedExponentialBridge{T}},
    F::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.ExponentialCone},
) where {T}
    return PermutedExponentialBridge{T,F}
end

function MOIB.map_set(
    ::Type{<:PermutedExponentialBridge},
    ::MOI.ExponentialCone,
)
    return PermutedExponentialCone()
end

function MOIB.inverse_map_set(
    ::Type{<:PermutedExponentialBridge},
    ::PermutedExponentialCone,
)
    return MOI.ExponentialCone()
end

function MOIB.map_function(
    ::Type{<:PermutedExponentialBridge{T}},
    func,
) where {T}
    scalars = MOIU.eachscalar(func)
    return scalars[[1, 3, 2]]
end

# Swapping indices 2 <-> 3 is an involution (it is its own inverse)

function MOIB.inverse_map_function(BT::Type{<:PermutedExponentialBridge}, func)
    return MOIB.map_function(BT, func)
end

# It is also symmetric (it is its own transpose)

function MOIB.adjoint_map_function(BT::Type{<:PermutedExponentialBridge}, func)
    return MOIB.map_function(BT, func)
end

function MOIB.inverse_adjoint_map_function(
    BT::Type{<:PermutedExponentialBridge},
    func,
)
    return MOIB.map_function(BT, func)
end


text:
using KoalaLow
using Base.Test


text:
using LinearAlgebra

function assembleDiffusionTerm(gamma, phi, gradPhi)
    fluxCf = [0.0 for _ in 1:mesh.nFaces]
    fluxFf = [0.0 for _ in 1:mesh.nFaces]
    fluxVf = [0.0 for _ in 1:mesh.nFaces]

    faces = [mesh.faces[iFace] for iFace in 1:mesh.nFaces]
    geoDiff = [faces[iFace].geoDiff for iFace in 1:mesh.nFaces]
    Tf = [faces[iFace].Tf for iFace in 1:mesh.nFaces]
    gammaFaces = gamma.faceValues
    phiFaces = phi.faceValues
    gradPhiFaces = gradPhi.faceValues

    # Assemble internal faces
    faceRange = 1:mesh.nInternalFaces

    fluxCf[faceRange] = gammaFaces[faceRange] .* geoDiff[faceRange]
    fluxFf[faceRange] = -fluxCf[faceRange]
    fluxVf[faceRange] = [dot(gammaFaces[iFace] * gradPhiFaces[iFace], Tf[iFace]) for iFace in faceRange]

    # Assemble boundary faces
    for boundary in mesh.boundaries
        faceRange = boundary.startFace:boundary.startFace + boundary.nFaces - 1

        if boundary.bcType == "fixedValue"
            fluxCf[faceRange] = gammaFaces[faceRange] .* geoDiff[faceRange]
            fluxFf[faceRange] .= 0.0
            fluxVf[faceRange] = -fluxCf[faceRange] .* phiFaces[faceRange] +
                [dot(gammaFaces[iFace] * gradPhiFaces[iFace], Tf[iFace]) for iFace in faceRange]
        elseif boundary.bcType == "zeroGradient"
            fluxCf[faceRange] .= 0.0
            fluxFf[faceRange] .= 0.0
            fluxVf[faceRange] = [dot(gammaFaces[iFace] * gradPhiFaces[iFace], Tf[iFace]) for iFace in faceRange]
        end
    end

    return fluxCf, fluxFf, fluxVf
end


text:
using IVPTestSuite
using Base.Test

include("innards.jl")
include("testcases.jl")


text:
"""
    Imputor

An imputor stores information about imputing values in `AbstractArray`s and `DataFrame`s.
New imputation methods are expected to sutype `Imputor` and, at minimum,
implement the `impute!{T<:Any}(imp::<MyImputor>, ctx::Context, data::AbstractArray{T, 1})`
method.
"""

abstract type Imputor end

"""
    impute!(imp::Imputor, data::Dataset, limit::Float64=0.1)

Creates a `Context` using information about `data`. These include

1. missing data function which defaults to `missing`

2. number of elements: `*(size(data)...)`

# Arguments
* `imp::Imputor`: the Imputor method to use
* `data::Dataset`: the data to impute
* `limit::Float64: missing data ratio limit/threshold (default: 0.1)`

# Return
* `Dataset`: the input `data` with values imputed.
"""
function impute!(imp::Imputor, data::Dataset, limit::Float64=0.1)
    ctx = Context(*(size(data)...), 0, limit, ismissing)
    return impute!(imp, ctx, data)
end

"""
    impute!(imp::Imputor, ctx::Context, data::AbstractMatrix)

Imputes the data in a matrix by imputing the values 1 column at a time;
if this is not the desired behaviour custom imputor methods should overload this method.

# Arguments
* `imp::Imputor`: the Imputor method to use
* `ctx::Context`: the contextual information for missing data
* `data::AbstractMatrix`: the data to impute

# Returns
* `AbstractMatrix`: the input `data` with values imputed
"""
function impute!(imp::Imputor, ctx::Context, data::AbstractMatrix)
    for i in 1:size(data, 2)
        impute!(imp, ctx, view(data, :, i))
    end
    return data
end

"""
    impute!(imp::Imputor, ctx::Context, data::DataFrame)

Imputes the data in a DataFrame by imputing the values 1 column at a time;
if this is not the desired behaviour custom imputor methods should overload this method.

# Arguments
* `imp::Imputor`: the Imputor method to use
* `ctx::Context`: the contextual information for missing data
* `data::DataFrame`: the data to impute

# Returns
* `DataFrame`: the input `data` with values imputed
"""
function impute!(imp::Imputor, ctx::Context, data::DataFrame)
    colwise(data) do c
        impute!(imp, ctx, c)
    end

    return data
end

for file in ("drop.jl", "locf.jl", "nocb.jl", "interp.jl", "fill.jl", "chain.jl")
    include(joinpath("imputors", file))
end


text:
using Vlasiator, PyPlot
using Documenter

makedocs(;
    modules=[Vlasiator],
    authors="Hongyang Zhou <hongyang.zhou@helsinki.fi> and contributors",
    repo="https://github.com/henry2004y/Vlasiator.jl/blob/{commit}{path}#L{line}",
    sitename="Vlasiator.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://henry2004y.github.io/Vlasiator.jl",
        assets=String[],
        sidebar_sitename=false
    ),
    pages=[
        "Home" => "index.md",
        "User Guide" => "manual.md",
        "Gallery" => "gallery.md",
        "API Reference" => "internal.md",
        "Contributing" => "contributing.md",
        "Log" => "log.md",
    ],
)

deploydocs(;
    repo="github.com/henry2004y/Vlasiator.jl",
)


text:
# This file is auto-generated by AWSMetadata.jl
using AWS
using AWS.AWSServices: qldb_session
using AWS.Compat
using AWS.UUIDs

"""
    send_command()
    send_command(params::Dict{String,<:Any})

Sends a command to an Amazon QLDB ledger.  Instead of interacting directly with this API,
we recommend using the QLDB driver or the QLDB shell to execute data transactions on a
ledger.   If you are working with an AWS SDK, use the QLDB driver. The driver provides a
high-level abstraction layer above this QLDB Session data plane and manages SendCommand API
calls for you. For information and a list of supported programming languages, see Getting
started with the driver in the Amazon QLDB Developer Guide.   If you are working with the
AWS Command Line Interface (AWS CLI), use the QLDB shell. The shell is a command line
interface that uses the QLDB driver to interact with a ledger. For information, see
Accessing Amazon QLDB using the QLDB shell.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"AbortTransaction"`: Command to abort the current transaction.
- `"CommitTransaction"`: Command to commit the specified transaction.
- `"EndSession"`: Command to end the current session.
- `"ExecuteStatement"`: Command to execute a statement in the specified transaction.
- `"FetchPage"`: Command to fetch a page.
- `"SessionToken"`: Specifies the session token for the current command. A session token is
  constant throughout the life of the session. To obtain a session token, run the
  StartSession command. This SessionToken is required for every subsequent command that is
  issued during the current session.
- `"StartSession"`: Command to start a new session. A session token is obtained as part of
  the response.
- `"StartTransaction"`: Command to start a new transaction.
"""
function send_command(; aws_config::AbstractAWSConfig=global_aws_config())
    return qldb_session(
        "SendCommand"; aws_config=aws_config, feature_set=SERVICE_FEATURE_SET
    )
end
function send_command(
    params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()
)
    return qldb_session(
        "SendCommand", params; aws_config=aws_config, feature_set=SERVICE_FEATURE_SET
    )
end


text:
"""
```
augment_states{T<:AbstractFloat}(m::AbstractModel, TTT::Matrix{T}, RRR::Matrix{T}, CCC::Matrix{T})
```

### Arguments

-`m`: the m object
-`TTT`, `RRR`, and `CCC`: matrices of the state transition equation

### Return values

- `TTT_aug`, `RRR_aug`, and `CCC_aug`: extend the corresponding input matrices to include
  jobservables which are growth rates.

### Description

Some observables in the model are growth rates, which are calculated as a linear combination
of a present and lagged state (which is not yet accounted for in the `TTT`, `RRR`,and `CCC`
matrices). To improve the performance of `gensys`, these additional states are added after
the model is solved. `augment_states` assigns an index to each lagged state, and extends the
input `TTT`, `RRR`, and `CCC` matrices to accommodate the additional states and capture the
lagged state value in the current state vector. `RRR` and `CCC` are mostly augmented with
zeros.

The diagram below shows how `TTT` is extended to `TTT_aug`.

                TTT_aug
     (m.endogenous_states_additional
                  x
     m.endogenous_states_additional)
     _________________________________
    |                     |           |
    |          TTT        | endog_    |
    | (endogenous_states  | states_   |
    |          x          | augmented |
    |  endogenous_states) |           |
    |_____________________|           |
    |                                 |
    |    endogenous_states_augmented  |
    |_________________________________|

"""
function augment_states(m::RealBondMkup, TTT::Matrix{T}, TTT_jump::Matrix{T}, RRR::Matrix{T}, CCC::Vector{T}, GDPeqn::AbstractMatrix{T}) where T<:AbstractFloat
    endo     = m.endogenous_states
    #endo_unnorm = m.endogenous_states_unnormalized
    endo_new = m.model_states_augmented
    exo      = m.exogenous_shocks

    _n_model_states     = n_model_states(m)
    _n_states           = n_backward_looking_states(m)
    _n_jumps            = n_jumps(m)
    _n_observables      = n_observables(m)
    _n_shocks_exogenous = n_shocks_exogenous(m)

    @assert (_n_model_states, _n_model_states) == size(TTT)
    @assert (_n_model_states, _n_shocks_exogenous)  == size(RRR)
    @assert (_n_model_states,)        == size(CCC)

    # Initialize augmented matrices
    n_states_add = length(endo_new)
    TTT_aug = zeros(_n_model_states + n_states_add, _n_model_states + n_states_add)
    TTT_aug[1:_n_model_states, 1:_n_model_states] = TTT
    RRR_aug = [RRR; zeros(n_states_add, _n_shocks_exogenous)]
    CCC_aug = [CCC; zeros(n_states_add)]

    ### TTT modifications

    # Track Lags
              # Aug[:y_t1]    States[:y_t]
     TTT_aug[endo_new[:y_t1], 1:_n_states] = GDPeqn

    return TTT_aug, RRR_aug, CCC_aug
end