text:
# types needed for different shock capturing schemes (basically different
# diffusion terms)

#------------------------------------------------------------------------------
# Projection-based shock capturing

mutable struct ProjectionShockCapturing{Tsol, Tres} <: AbstractVolumeShockCapturing
  filt::Matrix{Float64}  # the filter operator

  # storage
  w::Matrix{Tsol}
  t1::Matrix{Tres}
  t2::Matrix{Tres}
  Se_jac::Array{Tres, 4}
  ee_jac::Array{Tres, 4}
  A0inv::Matrix{Tsol}

  sensor::AbstractShockSensor

  function ProjectionShockCapturing{Tsol, Tres}(mesh::AbstractMesh, sbp::AbstractOperator, eqn, opts, sensor::AbstractShockSensor) where {Tsol, Tres}

    numDofPerNode = mesh.numDofPerNode
    filt = zeros(Float64, sbp.numnodes, sbp.numnodes)
    getFilterOperator!(sbp, filt)
    println("max entry in filt = ", maximum(abs.(filt)))

    w = zeros(Tsol, numDofPerNode, sbp.numnodes)
    t1 = zeros(Tres, numDofPerNode, sbp.numnodes)
    t2 = zeros(Tres, numDofPerNode, sbp.numnodes)

    Se_jac = zeros(Tres, mesh.dim, numDofPerNode, mesh.numNodesPerElement,
                          mesh.numNodesPerElement)
    ee_jac = zeros(Tres, mesh.dim, numDofPerNode, mesh.numNodesPerElement,
                          mesh.numNodesPerElement)

    A0inv = zeros(Tsol, numDofPerNode, numDofPerNode)

    return new(filt, w, t1, t2, Se_jac, ee_jac, A0inv, sensor)
  end
end

#------------------------------------------------------------------------------
# LDG shock capturing

"""
  Shock capturing using the entropy-stable varient of the Local Discontinuous
  Galerkin method.  [`allocateArrays`](@ref) must be called after the
  shock mesh is known before this type is usable.

  The sizes of these arrays may be larger than required.

  **Fields**

   * w_el: `numDofPerNode` x `numNodesPerElement` x `shockmesh.numEl` array
           for storing entropy variables
   * q_j: `numDofPerNode` x `numNodesPerElement` x `dim` x `shockmesh.numEl`
          array for storing theta_j and q_j
   * entropy_vars: an [`AbstractVariables`](@ref) specifying which entropy
                   variables to use
   * `flux`: an [`AbstractLDGFlux`](@ref)
   * diffusion: an [`AbstractDiffusion`](@ref).
"""
mutable struct LDGShockCapturing{Tsol, Tres} <: AbstractFaceShockCapturing
  # Note: the variable names are from Chen and Shu's "Entropy Stable High Order
  #       DG Methods" paper
  w_el::Array{Tsol, 3}  # entropy variables for all elements in elnums_all
  q_j::Array{Tres, 4}  # auxiliary equation solution for elements in
                       # elnums_all (this gets used for both theta and q)
  entropy_vars::AbstractVariables  # convert to entropy variables
  flux::AbstractLDGFlux
  diffusion::AbstractDiffusion
  sensor::AbstractShockSensor
  function LDGShockCapturing{Tsol, Tres}(mesh::AbstractMesh, sbp,
                      eqn, opts, sensor::AbstractShockSensor) where {Tsol, Tres}
    # we don't know the right sizes yet, so just make them zero size
    w_el = Array{Tsol}(0, 0, 0)
    q_j = Array{Tsol}(0, 0, 0, 0)

    # default values
    entropy_vars = getSCVariables(opts)
    flux = LDG_ESFlux()
    diffusion = ShockDiffusion{Tres}(mesh, sensor)

    return new(w_el, q_j, entropy_vars, flux, diffusion, sensor)
  end
end

"""
  This function allocates the arrays of [`LDGShockCapturing`](@ref) that
  depend on the shock mesh.  The arrays are only re-allocated if they are
  too small.

  **Inputs**

   * capture: [`LDGShockCapturing`](@ref)
   * mesh
   * shockmesh: a `ShockedElements` object, fully initialized
"""
function allocateArrays(capture::LDGShockCapturing{Tsol, Tres}, mesh::AbstractMesh,
                        shockmesh::ShockedElements) where {Tsol, Tres}

  # can't resize multi-dimension arrays, so reallocate
  if size(capture.q_j, 4) < shockmesh.numEl
    capture.q_j = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement,
                              mesh.dim, shockmesh.numEl)
    # the final dimension must be numEl, not numShock, because q_j is zero
    # for the neighboring elements later
  end

  if size(capture.w_el, 3) < shockmesh.numEl
    capture.w_el = Array{Tsol}(mesh.numDofPerNode, mesh.numNodesPerElement,
                               shockmesh.numEl)
  end

  return nothing
end

"""
  Entropy stable LDG flux
"""
struct LDG_ESFlux  <: AbstractLDGFlux
  alpha::Float64
  beta::Float64  # Chen and Shu say this can be an arbitrary vector, but my
                 # analysis says all entries must the the same, so it acts
                 # like a scalar
  function LDG_ESFlux()
    alpha = 1
    beta = 1

    return new(alpha, beta)
  end
end

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

"""
  Shock capturing type that errors out.  Used when shock capturing is not
  supposed to be added.
"""
mutable struct ErrorShockCapturing{Tsol, Tres} <: AbstractShockCapturing

  sensor::AbstractShockSensor
  diffusion::AbstractDiffusion
  function ErrorShockCapturing{Tsol, Tres}(mesh::AbstractMesh, sbp::AbstractSBP,
                                           eqn, opts,
                                           sensor::AbstractShockSensor) where {Tsol, Tres}

    diffusion = ShockDiffusion{Tres}(mesh, sensor)
    return new(sensor, diffusion)
  end
end

"""
  Parallel communications for alpha
"""
struct AlphaComm
  tag::Cint
  send_req::Vector{MPI.Request}
  recv_req::Vector{MPI.Request}
  recv_waited::Vector{Bool}
  send_bufs::Vector{Vector{UInt8}}  # make the buffers UInt8 so (hopefully)
  recv_bufs::Vector{Vector{UInt8}}  # the send will be done eagerly
  comm::MPI.Comm
  peer_parts::Vector{Int}  # the list of all processes that might be
                           # communicated with (same order as mesh.peer_parts)
  peer_parts_red::Vector{Cint}  # list of peer parts that will actually be
                                # communicated with
  function AlphaComm(comm::MPI.Comm, peer_parts::Vector{Int})

    tag = getNextTag(TagManager)
    send_req = Vector{MPI.Request}(0)
    recv_req = Vector{MPI.Request}(0)
    recv_waited = Vector{Bool}(0)
    send_bufs = Vector{Vector{UInt8}}(0)
    recv_bufs = Vector{Vector{UInt8}}(0)
    peer_parts_red = Vector{Cint}(0)

    return new(tag, send_req, recv_req, recv_waited, send_bufs, recv_bufs,
               comm, peer_parts, peer_parts_red)
  end
end

function allocateArrays(comm::AlphaComm, shockmesh::ShockedElements)

  # make sure previous communications have finished before reallocating arrays
  MPI.Waitall!(comm.send_req)
  assertReceivesWaited(comm)

  resize!(comm.send_bufs, shockmesh.npeers)
  resize!(comm.recv_bufs, shockmesh.npeers)
  resize!(comm.send_req, shockmesh.npeers)
  resize!(comm.recv_req, shockmesh.npeers)
  resize!(comm.recv_waited, shockmesh.npeers)
  resize!(comm.peer_parts_red, shockmesh.npeers)

  # extract needed MPI ranks
  for i=1:shockmesh.npeers
    comm.peer_parts_red[i] = comm.peer_parts[shockmesh.peer_indices[i]]
    comm.recv_bufs[i] = Vector{UInt8}(shockmesh.numSharedInterfaces[i])
    comm.send_bufs[i] = Vector{UInt8}(shockmesh.numSharedInterfaces[i])
  end

  return nothing
end

import MPI: Waitall!, Waitany!, Isend, Irecv!

"""
  Wrapper around MPI.Waitall! for the receive requests
"""
function Waitall!(comm::AlphaComm)

  MPI.Waitall!(comm.recv_req)
end

"""
  Wrapper around MPI.Waitany! for the receive requests.
  Returns only the index, not the Status object.
"""
function Waitany!(comm::AlphaComm)

  idx, stat = MPI.Waitany!(comm.recv_req)
  comm.recv_req[idx] = MPI.REQUEST_NULL
  comm.recv_waited[idx] = true

  return idx
end

import Utils: assertReceivesWaited

"""
  Asserts all receives have been waited on
"""
function assertReceivesWaited(comm::AlphaComm)

  if length(comm.recv_waited) > 0
    for i=1:length(comm.recv_waited)
      @assert comm.recv_waited[i]
    end
  end

  return nothing
end

"""
  Starts communication for a given peer index (note: this is the index in
  the range 1:shockmesh.npeer)

  **Inputs**

   * comm: `AlphaComm` object
   * peeridx: peer index
"""
function Isend(comm::AlphaComm, peeridx::Integer)

  req = MPI.Isend(comm.send_bufs[peeridx], comm.peer_parts_red[peeridx],
                  comm.tag, comm.comm)
  comm.send_req[peeridx] = req

  return nothing
end

"""
  Counterpart to Isend, starts the receive for a given peer index.  Same
  arguments as Isend.
"""
function Irecv!(comm::AlphaComm, peeridx::Integer)

  req = MPI.Irecv!(comm.recv_bufs[peeridx], comm.peer_parts_red[peeridx],
                   comm.tag, comm.comm)
  comm.recv_waited[peeridx] = false
  comm.recv_req[peeridx] = req

  return nothing
end

"""
  Computes any diffusion scheme from the SBPParabolic paper.

  Note that the sizes for the arrays are the *minimum* sizes.  In practice
  they may be larger.  Do not call `size()` on these arrays, use the appropriate
  integer field of `mesh` or `shockmesh`.

  **Fields**

   * w_el: entropy variables at each volume node of each element,
           `numDofPerNode` x `numNodesPerElement` x `shockmesh.numEl`
   * grad_w: stores Lambda * D * q at volume nodes, numDofPerNode x
             numNodesPerElement x dim x shockmesh.numEl.  Should be zero
             for all neighboring elements (see `ShockedElement`) because
             the viscoscity is zero there, but storing values in this array
             significantly simplifies the code
  * entropy_vars: an [`AbstractVariables`](@ref) specifying which entropy
                   variables to use.
  * diffusion: an [`AbstractDiffusion`](@ref) object that specifies the
               diffusion tensor. Defaults to `ShockDiffusion`
  * penalty: an [`AbstractDiffusionPenalty`](@ref) object that specifies which
             scheme to use.  Defaults to opts["DiffusionPenalty"]
  * alpha: a 2 x `shockmesh.numInterfaces` array containing alpha_gk and
           alpha_gn for each interface.

  Note that `convert_entropy`, `diffusion` and `penalty` are abstract types
  and should only be accessed through a function barrier in performance
  critical functions.  The user is allowed to change these fields at any
  time to customize the behavior of the scheme.

  Note that this function is not fully constructed by the constructor,
  see the `allocateArrays` function.

"""
mutable struct SBPParabolicSC{Tsol, Tres} <: AbstractFaceShockCapturing
  w_el::Array{Tsol, 3}
  grad_w::Array{Tres, 4}
  entropy_vars::AbstractVariables
  diffusion::AbstractDiffusion
  penalty::AbstractDiffusionPenalty
  sensor::AbstractShockSensor
  alpha::Array{Float64, 2}  # 2 x numInterfaces
  alpha_b::Array{Float64, 1}  # numBoundaryFaces
  alpha_parallel::Array{Array{Float64, 2}, 1}  # one array for each peer

  #------------------
  # temporary arrays

  # getFaceVariables
  w_faceL::Matrix{Tsol}
  w_faceR::Matrix{Tsol}
  grad_faceL::Matrix{Tres}
  grad_faceR::Matrix{Tres}

  # getFaceVariables_diff
  wL_dot::Array{Tsol, 3}
  wR_dot::Array{Tsol, 3}
  Dx::Array{Float64, 3}
  t1::Array{Tres, 3}
  t1_dot::Array{Tres, 5}
  t2L_dot::Array{Tres, 5}
  t2R_dot::Array{Tres, 5}
  t3L_dot::Array{Tres, 5}
  t3R_dot::Array{Tres, 5}

  # applyDgkTranspose
  temp1L::Matrix{Tres}
  temp1R::Matrix{Tres}
  temp2L::Array{Tres, 3}
  temp2R::Array{Tres, 3}
  temp3L::Array{Tres, 3}
  temp3R::Array{Tres, 3}
  work::Array{Tres, 3}

  # applyDgkTranspose_diff
  t3L_dotL::Array{Tres, 5}
  t3L_dotR::Array{Tres, 5}
  t3R_dotL::Array{Tres, 5}
  t3R_dotR::Array{Tres, 5}

  t4L_dotL::Array{Tres, 5}
  t4L_dotR::Array{Tres, 5}
  t4R_dotL::Array{Tres, 5}
  t4R_dotR::Array{Tres, 5}

  t5L_dotL::Array{Tres, 5}
  t5L_dotR::Array{Tres, 5}
  t5R_dotL::Array{Tres, 5}
  t5R_dotR::Array{Tres, 5}

  # computeBoundaryTerm_diff
  w_dot::Array{Tsol, 3}
  t2_dot::Array{Tres, 5}
  t3_dot::Array{Tres, 5}
  t4_dot::Array{Tres, 4}

  alpha_comm::AlphaComm # MPI communications for alpha

  function SBPParabolicSC{Tsol, Tres}(mesh::AbstractMesh, sbp::AbstractOperator,
                                      eqn, opts, sensor::AbstractShockSensor) where {Tsol, Tres}

    @unpack mesh numDofPerNode numNodesPerFace dim numNodesPerElement

    # we don't know the right sizes yet, so make them zero size
    w_el = Array{Tsol}(0, 0, 0)
    grad_w = Array{Tres}(0, 0, 0, 0)

    # default values
    entropy_vars = getSCVariables(opts)
    diffusion = ShockDiffusion{Tres}(mesh, sensor)
    penalty = getDiffusionPenalty(mesh, sbp, eqn, opts)
    alpha = zeros(Float64, 0, 0)
    alpha_b = zeros(Float64, 0)
    alpha_parallel = Array{Array{Float64, 2}}(0)

    # temporary arrays

    # getFaceVariables
    w_faceL = zeros(Tsol, numDofPerNode, numNodesPerFace)
    w_faceR = zeros(Tsol, numDofPerNode, numNodesPerFace)
    grad_faceL = zeros(Tres, numDofPerNode, numNodesPerFace)
    grad_faceR = zeros(Tres, numDofPerNode, numNodesPerFace)

    # getFaceVariables_diff
    wL_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)
    wR_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)
    Dx = zeros(numNodesPerElement, numNodesPerElement, dim)
    t1 = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    t1_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t2L_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t2R_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t3L_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                          numNodesPerElement)
    t3R_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                          numNodesPerElement)

    # applyDgk transpose
    temp1L = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp1R = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp2L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp2R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    work = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)

    # applyDgk_transpose_diff
    t3L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)

    t4L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    t5L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    # computeBoundaryTerm_diff
    w_dot = zeros(Tsol, mesh.numDofPerNode, mesh.numDofPerNode,
                        mesh.numNodesPerElement)
    #t1 = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
    #t1_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
    #                     mesh.numNodesPerElement, mesh.numNodesPerElement)
    t2_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerElement, mesh.numNodesPerElement)

    t3_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    t4_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    alpha_comm = AlphaComm(mesh.comm, mesh.peer_parts)

    return new(w_el, grad_w, entropy_vars, diffusion, penalty, sensor, alpha,
               alpha_b, alpha_parallel, w_faceL, w_faceR, grad_faceL,
               grad_faceR,
               wL_dot, wR_dot, Dx, t1, t1_dot, t2L_dot, t2R_dot, t3L_dot,
               t3R_dot,
               temp1L, temp1R, temp2L, temp2R, temp3L, temp3R, work,
               t3L_dotL, t3L_dotR, t3R_dotL, t3R_dotR,
               t4L_dotL, t4L_dotR, t4R_dotL, t4R_dotR,
               t5L_dotL, t5L_dotR, t5R_dotL, t5R_dotR,
               w_dot, t2_dot, t3_dot, t4_dot,
               alpha_comm)
  end
end

"""
  Sets up the shock capturing object for use, using the fully initialized
  `shockmesh`.  Also does some other misc. setup stuff that needs `shockmesh`.
"""
function allocateArrays(capture::SBPParabolicSC{Tsol, Tres}, mesh::AbstractMesh,
                        shockmesh::ShockedElements) where {Tsol, Tres}

  # can't resize multi-dimension arrays, so reallocate
  if size(capture.grad_w, 4) < shockmesh.numEl
    capture.grad_w = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement,
                                 mesh.dim, shockmesh.numEl)
  end

  if size(capture.w_el, 3) < shockmesh.numEl
    capture.w_el = Array{Tsol}(mesh.numDofPerNode, mesh.numNodesPerElement,
                               shockmesh.numEl)
  end

  allocateArrays(capture.alpha_comm, shockmesh)
  computeAlpha(capture, mesh, shockmesh)

  return nothing
end

#------------------------------------------------------------------------------
# SBPParabolicReducedSC

"""
  Computes a reduced form of the scheme from the SBPParabolic paper.

  Note that the sizes for the arrays are the *minimum* sizes.  In practice
  they may be larger.  Do not call `size()` on these arrays, use the appropriate
  integer field of `mesh` or `shockmesh`.

  **Fields**

   * w_el: entropy variables at each volume node of each element,
           `numDofPerNode` x `numNodesPerElement` x `shockmesh.numEl`
   * entropy_vars: an [`AbstractVariables`](@ref) specifying which entropy
                   variables to use.
   * diffusion: an [`AbstractDiffusion`](@ref) object that specifies the
               diffusion tensor. Defaults to `ShockDiffusion`
   * penalty: an [`AbstractDiffusionPenalty`](@ref) object that specifies which
             scheme to use.  Defaults to opts["DiffusionPenalty"]
   * alpha: a 2 x `shockmesh.numInterfaces` array containing alpha_gk and
           alpha_gn for each interface.

  Note that `convert_entropy`, `diffusion` and `penalty` are abstract types
  and should only be accessed through a function barrier in performance
  critical functions.  The user is allowed to change these fields at any
  time to customize the behavior of the scheme.

  Note that this function is not fully constructed by the constructor,
  see the `allocateArrays` function.

"""
mutable struct SBPParabolicReducedSC{Tsol, Tres} <: AbstractFaceShockCapturing
  w_el::Array{Tsol, 3}
  entropy_vars::AbstractVariables
  diffusion::AbstractDiffusion
  penalty::AbstractDiffusionPenalty
  sensor::AbstractShockSensor
  sensor_const::ShockSensorHApprox{Tsol, Tres}  # shock sensor used for
                                                   # face terms
  #------------------
  # temporary arrays

  # getFaceVariables
  w_faceL::Matrix{Tsol}
  w_faceR::Matrix{Tsol}

  # getFaceVariables_diff
  wL_dot::Array{Tsol, 3}
  wR_dot::Array{Tsol, 3}

  # applyDgkTranspose
  temp1L::Matrix{Tres}
  temp1R::Matrix{Tres}
  temp2L::Array{Tres, 3}
  temp2R::Array{Tres, 3}
  temp3L::Array{Tres, 3}
  temp3R::Array{Tres, 3}
  work::Array{Tres, 3}

  # applyDgkTranspose_diff
  Dx::Array{Float64, 3}
  t3L_dotL::Array{Tres, 5}
  t3L_dotR::Array{Tres, 5}
  t3R_dotL::Array{Tres, 5}
  t3R_dotR::Array{Tres, 5}

  t4L_dotL::Array{Tres, 5}
  t4L_dotR::Array{Tres, 5}
  t4R_dotL::Array{Tres, 5}
  t4R_dotR::Array{Tres, 5}

  t5L_dotL::Array{Tres, 5}
  t5L_dotR::Array{Tres, 5}
  t5R_dotL::Array{Tres, 5}
  t5R_dotR::Array{Tres, 5}

  # computeBoundaryTerm_diff
  w_dot::Array{Tsol, 3}
  t2_dot::Array{Tres, 5}
  t3_dot::Array{Tres, 5}
  t4_dot::Array{Tres, 4}

  function SBPParabolicReducedSC{Tsol, Tres}(mesh::AbstractMesh, sbp::AbstractOperator,
                                      eqn, opts, sensor::AbstractShockSensor) where {Tsol, Tres}

    @unpack mesh numDofPerNode numNodesPerFace dim numNodesPerElement

    # we don't know the right sizes yet, so make them zero size
    w_el = Array{Tsol}(0, 0, 0)

    # default values
    entropy_vars = getSCVariables(opts)
    diffusion = ShockDiffusion{Tres}(mesh, sensor)
    penalty = getDiffusionPenalty(mesh, sbp, eqn, opts)
    sensor_const = ShockSensorHApprox{Tsol, Tres}(mesh, sbp, opts)

    # temporary arrays

    # getFaceVariables
    w_faceL = zeros(Tsol, numDofPerNode, numNodesPerFace)
    w_faceR = zeros(Tsol, numDofPerNode, numNodesPerFace)

    # getFaceVariables_diff
    wL_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)
    wR_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)

    # applyDgk transpose
    temp1L = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp1R = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp2L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp2R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    work = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)

    # applyDgk_transpose_diff
    Dx = zeros(numNodesPerElement, numNodesPerElement, dim)
    t3L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)

    t4L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    t5L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    # computeBoundaryTerm_diff
    w_dot = zeros(Tsol, mesh.numDofPerNode, mesh.numDofPerNode,
                        mesh.numNodesPerElement)
    #t1 = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
    #t1_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
    #                     mesh.numNodesPerElement, mesh.numNodesPerElement)
    t2_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerElement, mesh.numNodesPerElement)

    t3_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    t4_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    return new(w_el, entropy_vars, diffusion, penalty, sensor, sensor_const,
               w_faceL, w_faceR,
               wL_dot, wR_dot,
               temp1L, temp1R, temp2L, temp2R, temp3L, temp3R, work,
               Dx,
               t3L_dotL, t3L_dotR, t3R_dotL, t3R_dotR,
               t4L_dotL, t4L_dotR, t4R_dotL, t4R_dotR,
               t5L_dotL, t5L_dotR, t5R_dotL, t5R_dotR,
               w_dot, t2_dot, t3_dot, t4_dot)
  end
end

"""
  Sets up the shock capturing object for use, using the fully initialized
  `shockmesh`.  Also does some other misc. setup stuff that needs `shockmesh`.
"""
function allocateArrays(capture::SBPParabolicReducedSC{Tsol, Tres},
                        mesh::AbstractMesh,
                        shockmesh::ShockedElements) where {Tsol, Tres}

  if size(capture.w_el, 3) < shockmesh.numEl
    capture.w_el = Array{Tsol}(mesh.numDofPerNode, mesh.numNodesPerElement,
                               shockmesh.numEl)
  end

  return nothing
end

function getSparsityPattern(capture::SBPParabolicReducedSC)

  return INVISCID
end

#------------------------------------------------------------------------------
# SBPParabolicReduced2

mutable struct SBPParabolicReduced2SC{Tsol, Tres} <: AbstractFaceShockCapturing
  w_el::Array{Tsol, 3}
  grad_w::Array{Tres, 4}
  entropy_vars::AbstractVariables
  diffusion::AbstractDiffusion
  penalty::AbstractDiffusionPenalty
  sensor::AbstractShockSensor
  alpha::Array{Float64, 2}  # 2 x numInterfaces
  alpha_b::Array{Float64, 1}  # numBoundaryFaces
  alpha_parallel::Array{Array{Float64, 2}, 1}  # one array for each peer

  #------------------
  # temporary arrays

  # getFaceVariables
  w_faceL::Matrix{Tsol}
  w_faceR::Matrix{Tsol}
  grad_faceL::Matrix{Tres}
  grad_faceR::Matrix{Tres}

  # getFaceVariables_diff
  wL_dot::Array{Tsol, 3}
  wR_dot::Array{Tsol, 3}
  Dx::Array{Float64, 3}
  t1::Array{Tres, 3}
  t1_dot::Array{Tres, 5}
  t2L_dot::Array{Tres, 5}
  t2R_dot::Array{Tres, 5}
  t3L_dot::Array{Tres, 5}
  t3R_dot::Array{Tres, 5}

  # applyDgkTranspose
  temp1L::Matrix{Tres}
  temp1R::Matrix{Tres}
  temp2L::Array{Tres, 3}
  temp2R::Array{Tres, 3}
  temp3L::Array{Tres, 3}
  temp3R::Array{Tres, 3}
  work::Array{Tres, 3}

  # applyDgkTranspose_diff
  t3L_dotL::Array{Tres, 5}
  t3L_dotR::Array{Tres, 5}
  t3R_dotL::Array{Tres, 5}
  t3R_dotR::Array{Tres, 5}

  t4L_dotL::Array{Tres, 5}
  t4L_dotR::Array{Tres, 5}
  t4R_dotL::Array{Tres, 5}
  t4R_dotR::Array{Tres, 5}

  t5L_dotL::Array{Tres, 5}
  t5L_dotR::Array{Tres, 5}
  t5R_dotL::Array{Tres, 5}
  t5R_dotR::Array{Tres, 5}

  # computeBoundaryTerm_diff
  w_dot::Array{Tsol, 3}
  t2_dot::Array{Tres, 5}
  t3_dot::Array{Tres, 5}
  t4_dot::Array{Tres, 4}

  alpha_comm::AlphaComm # MPI communications for alpha

  function SBPParabolicReduced2SC{Tsol, Tres}(mesh::AbstractMesh, sbp::AbstractOperator,
                                      eqn, opts, sensor::AbstractShockSensor) where {Tsol, Tres}

    @unpack mesh numDofPerNode numNodesPerFace dim numNodesPerElement

    # we don't know the right sizes yet, so make them zero size
    w_el = Array{Tsol}(0, 0, 0)
    grad_w = Array{Tres}(0, 0, 0, 0)

    # default values
    entropy_vars = getSCVariables(opts)
    diffusion = ShockDiffusion{Tres}(mesh, sensor)
    penalty = getDiffusionPenalty(mesh, sbp, eqn, opts)
    alpha = zeros(Float64, 0, 0)
    alpha_b = zeros(Float64, 0)
    alpha_parallel = Array{Array{Float64, 2}}(0)

    # temporary arrays

    # getFaceVariables
    w_faceL = zeros(Tsol, numDofPerNode, numNodesPerFace)
    w_faceR = zeros(Tsol, numDofPerNode, numNodesPerFace)
    grad_faceL = zeros(Tres, numDofPerNode, numNodesPerFace)
    grad_faceR = zeros(Tres, numDofPerNode, numNodesPerFace)

    # getFaceVariables_diff
    wL_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)
    wR_dot = zeros(Tsol, numDofPerNode, numDofPerNode, numNodesPerElement)
    Dx = zeros(numNodesPerElement, numNodesPerElement, dim)
    t1 = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    t1_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t2L_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t2R_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                         numNodesPerElement)
    t3L_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                          numNodesPerElement)
    t3R_dot = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                          numNodesPerElement)

    # applyDgk transpose
    temp1L = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp1R = zeros(Tres, numDofPerNode, numNodesPerFace)
    temp2L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp2R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3L = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    temp3R = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)
    work = zeros(Tres, numDofPerNode, numNodesPerElement, mesh.dim)

    # applyDgk_transpose_diff
    t3L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)
    t3R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                           numNodesPerElement)

    t4L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t4R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    t5L_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5L_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)
    t5R_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                           numNodesPerElement)

    # computeBoundaryTerm_diff
    w_dot = zeros(Tsol, mesh.numDofPerNode, mesh.numDofPerNode,
                        mesh.numNodesPerElement)
    #t1 = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
    #t1_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
    #                     mesh.numNodesPerElement, mesh.numNodesPerElement)
    t2_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerElement, mesh.numNodesPerElement)

    t3_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode, mesh.dim,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    t4_dot = zeros(Tres, mesh.numDofPerNode, mesh.numDofPerNode,
                         mesh.numNodesPerFace, mesh.numNodesPerElement)

    alpha_comm = AlphaComm(mesh.comm, mesh.peer_parts)

    return new(w_el, grad_w, entropy_vars, diffusion, penalty, sensor, alpha,
               alpha_b, alpha_parallel, w_faceL, w_faceR, grad_faceL,
               grad_faceR,
               wL_dot, wR_dot, Dx, t1, t1_dot, t2L_dot, t2R_dot, t3L_dot,
               t3R_dot,
               temp1L, temp1R, temp2L, temp2R, temp3L, temp3R, work,
               t3L_dotL, t3L_dotR, t3R_dotL, t3R_dotR,
               t4L_dotL, t4L_dotR, t4R_dotL, t4R_dotR,
               t5L_dotL, t5L_dotR, t5R_dotL, t5R_dotR,
               w_dot, t2_dot, t3_dot, t4_dot,
               alpha_comm)
  end
end

"""
  Sets up the shock capturing object for use, using the fully initialized
  `shockmesh`.  Also does some other misc. setup stuff that needs `shockmesh`.
"""
function allocateArrays(capture::SBPParabolicReduced2SC{Tsol, Tres}, mesh::AbstractMesh,
                        shockmesh::ShockedElements) where {Tsol, Tres}

  # can't resize multi-dimension arrays, so reallocate
  if size(capture.grad_w, 4) < shockmesh.numEl
    capture.grad_w = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement,
                                 mesh.dim, shockmesh.numEl)
  end

  if size(capture.w_el, 3) < shockmesh.numEl
    capture.w_el = Array{Tsol}(mesh.numDofPerNode, mesh.numNodesPerElement,
                               shockmesh.numEl)
  end

  return nothing
end

#------------------------------------------------------------------------------
# VolumeShockCapturing

mutable struct VolumeShockCapturing{Tsol, Tres} <: AbstractVolumeShockCapturing
  entropy_vars::AbstractVariables
  diffusion::AbstractDiffusion
  sensor::AbstractShockSensor

  function VolumeShockCapturing{Tsol, Tres}(mesh::AbstractMesh,
                                sbp::AbstractOperator, eqn,  opts,
                                sensor::AbstractShockSensor) where {Tsol, Tres}
    entropy_vars = getSCVariables(opts)
    diffusion = ShockDiffusion{Tres}(mesh, sensor)
    return new(entropy_vars, diffusion, sensor)
  end
end

"""
  Get the entropy variables specified by the options dictionary for use in the
  shock capturing dissipation term.
"""
function getSCVariables(opts)

  name = opts["shock_capturing_variables"]
  if name == "IR"
    return IRVariables()
  elseif name == "conservative"
    return ConservativeVariables()
  else
    error("unrecognized variables for shock capturing dissipation: $name")
  end
end

#------------------------------------------------------------------------------
# Creating shock capturing

global const ShockCapturingDict = Dict{String, Type{T} where T <: AbstractShockCapturing}(
"ShockCapturingNone" => ErrorShockCapturing,
"ElementProjection" => ProjectionShockCapturing,
"LDG" => LDGShockCapturing,
"SBPParabolic" => SBPParabolicSC,
"SBPParabolicReduced" => SBPParabolicReducedSC,
"SBPParabolicReduced2" => SBPParabolicReduced2SC,
"Volume" => VolumeShockCapturing,
)

"""
  This function populates the `eqn.shock_capturing` field of the mesh with the
  shock capturing object for the scheme.  The shock capturing scheme is
  determined by opts["shock_capturing_name"]
"""
function getShockCapturing(mesh, sbp, eqn::EulerData{Tsol, Tres}, opts,
                           sensor::AbstractShockSensor) where {Tsol, Tres}

  name = opts["shock_capturing_name"]
  obj = ShockCapturingDict[name]{Tsol, Tres}(mesh, sbp, eqn, opts, sensor)
  assertArraysUnique(obj)
  eqn.shock_capturing = obj

  return nothing
end


text:
module QuantumEvolution
using Random, StatsBase, LinearAlgebra

struct QEState{vP,vC,vO}
    P::Vector{vP}
    C::Vector{vC}
    eqP::Vector{vO}
    eqC::Vector{vC}
    N::Int
    K::Int
    iter::Int
    fnevals::Int
    improv::Int
end

evals(state::QEState) = state.fnevals

struct QEProblem{tF,tS,tP,T}
    f::tF
    s::tS
    projector::tP
    beta::T
    max_iters::Int
end

function opt_state(problem::QEProblem, N, K; kw...)
    P = [problem.projector(problem.s(i)) for i = 1:N]
    C = [problem.f(P[i]) for i = 1:N]
    for i in eachindex(P)
        isfinite(C[i]) && continue
        error("The cost function has produced non-finite values for the initial point: $(P[i])")
    end
    order = sortperm(C)
    eqP = [P[i] * 1 for i in order[1:K]]
    eqC = [C[i] * 1 for i in order[1:K]]
    (eqP[1], eqC[1], QEState(P[order], C[order], eqP, eqC, N, K, 1, N, 1))
end

function opt_state(problem::QEProblem, S::QEState; rng, verbose, kw...)
    P, C, eqP, eqC, N, K, iters, fnevals, improv =
        S.P, S.C, S.eqP, S.eqC, S.N, S.K, S.iter, S.fnevals, S.improv
    cost, beta, max_iters, projector =
        problem.f, problem.beta, problem.max_iters, problem.projector
    eqPav = sum(x -> x / K, eqP)
    dsol = size(eqPav)
    tr = iters / max_iters
    t = (1 - tr)^(2tr)
    worst = C[end]
    best = C[1]

    for i = ceil(Int, (N - K + 1) * (1 - t) + t):N
        @label regen
        eqPr = (K + 1) * rand(rng) < 1 ? eqPav : rand(rng, eqP)
        log_u = log(rand(rng)) / log(2)
        v = sign.(rand(rng, dsol...) .- 0.5)
        nP = projector(@.(eqPr + beta * abs(P[i] - eqPr) * log_u * v))
        nC = cost(nP)
        fnevals += 1
        isfinite(nC) || begin
            verbose && println("Failed function evaluation at $nP")
            @goto regen
        end
        if nC <= worst
            push!(P, nP)
            push!(C, nC)
            nC < best ? improv += 1 : nothing
        end
    end
    iters += 1
    #efficiency = (improv/iters)/(fnevals/(N*0.497595+K*(1-0.497595)))
    #@show efficiency
    order = sortperm(C)[1:N]
    eqP = [P[i] * 1 for i in order[1:K]]
    eqC = [C[i] * 1 for i in order[1:K]]
    (eqP[1], eqC[1], QEState(P[order], C[order], eqP, eqC, N, K, iters, fnevals, improv))
end

"""
Implementation of `Quantum PSO` with some modifications

## API

```julia
function optimize(
    f,
    s,
    N;
    K = 4,
    beta = 0.95,
    max_iters = 100,
    min_cost = -Inf,
    verbose = false,
    rng = Random.GLOBAL_RNG,
    callback = state -> nothing,
    projector = identity
)
```

- `f` : cost function to minimize, whose argument is either a scalar or a vector, must returns a scalar value.
- `s` : function whose input is the particle number and output is a random initial point to be ranked by `f`.
- `N` : number of particles to use, choose a number greater than `d+4` where `d` is the number of dimensions.
- `K` : number of particles used as best individuals to propose new particles.
- `beta` : dilation/contraction factor for proposal distribution.
- `max_iters` : maximum number of iterations.
- `min_cost` : threshold for stopping the algorithm whenever the cost of the current best solution is below `min_cost`
- `verbose` : enables verbosity
- `rng` : random number generator
- `callback` : callback function to access the internal state of the optimizer, if it returns `true` the optimizer is stopped
- `projector` : function whose role is projecting parameters back into their constrained space

## Usage example:

```julia
using QuantumEvolution
rosenbrock2d(x) = (x[1] - 1)^2 + 100*(x[2] - x[1]^2)^2
initpoint(i) = randn(2)
optimize(rosenbrock2d, initpoint, 30)
```

```
(x = [0.9999998740708989, 0.9999997219238367], fx = 8.459636979174772e-14, f_nevals = 1661)
```
as expected the global optimum has been found.
"""
function optimize(
    f,
    s,
    N;
    K = 4,
    beta = 0.95,
    max_iters = 100,
    min_cost = -Inf,
    verbose = false,
    rng = Random.GLOBAL_RNG,
    callback = state -> nothing,
    projector = identity,
)
    N < 2K && error("N must be at least $(2K)")
    problem = QEProblem(f, s, projector, beta, max_iters)
    best, cost, state = opt_state(problem, N, K)
    callback(state) === true && @goto ret
    for i = 1:(max_iters)
        best, cost, state = opt_state(problem, state; rng, verbose)
        verbose &&
            println(i, " ", cost, " ", state.improv, " ", state.improv / evals(state))
        callback(state) === true && break
        cost < min_cost && break
    end
    @label ret
    (x = best, fx = cost, f_nevals = evals(state))
end

export optimize
end

ext: jl


text:
"""
```
get_jstep(m, n_sim)
```

Retrieve `forecast_jstep` setting (thinning step size for forecast
step) from `m.settings`. If `n_sim ==  1`, returns 1.
"""
function get_jstep(m::AbstractDSGEModel, n_sim::Int)
    if n_sim == 1
        jstep = 1
    else
        jstep = get_setting(m, :forecast_jstep)
    end
end

"""
```
n_forecast_draws(m, input_type)
```

Returns the number of forecast draws in the file
`get_forecast_input_file(m, input_type)`.
"""
function n_forecast_draws(m::AbstractDSGEModel, input_type::Symbol)
    if input_type in [:mean, :mode, :init]
        return 1
    elseif input_type in [:full, :subset]
        input_file = get_forecast_input_file(m, input_type)
        draws = h5open(input_file, "r") do file
            if get_setting(m, :sampling_method) == :MH
                dataset = HDF5.o_open(file, "mhparams")
            elseif get_setting(m, :sampling_method) == :SMC
                dataset = HDF5.o_open(file, "smcparams")
            else
                throw("Invalid :sampling_method setting specification.")
            end
            size(dataset)[1]
        end
        return draws
    elseif input_type == :prior || input_type == :mode_draw_shocks
        return 5000
    else
        throw(ArgumentError("Invalid input_type: $(input_type)"))
    end
end

"""
```
forecast_block_inds(m, input_type; subset_inds = 1:0)
```

Returns a pair of `Vector{AbstractRange{Int64}}`s, `block_inds` and `block_inds_thin`,
each of length equal to the number of forecast blocks. `block_inds[i]` is the
range of indices for block `i` before thinning by `jstep` and
`block_inds_thin[i]` is the range after thinning.
"""
function forecast_block_inds(m::AbstractDSGEModel, input_type::Symbol; subset_inds::AbstractRange{Int64} = 1:0, params::Vector{Float64} = Vector{Float64}(undef, 0))

    if !isempty(params)
        ndraws    = 1000
        jstep     = get_jstep(m, ndraws)
        start_ind = 1
        end_ind   = ndraws
    else
        if input_type == :full || input_type == :prior || input_type == :init_draw_shocks || input_type == :mode_draw_shocks
            regime_switching = (haskey(m.settings, :regime_switching) && haskey(m.settings, :regime_switching_ndraws)) ?
                get_setting(m, :regime_switching) : false
            set_ndraw = haskey(m.settings, :forecast_ndraws) ? get_setting(m, :forecast_ndraws) : 0
            if set_ndraw > 0
                ndraws = set_ndraw
            else
                ndraws = regime_switching ? get_setting(m, :regime_switching_ndraws) : n_forecast_draws(m, :full)
            end
            jstep     = get_jstep(m, ndraws)
            start_ind = 1
            end_ind   = ndraws
        elseif input_type == :subset
            ndraws    = length(subset_inds)
            jstep     = get_jstep(m, ndraws)
            start_ind = first(subset_inds)
            end_ind   = last(subset_inds)
        else
            throw(ArgumentError("Cannot call forecast_block_inds with input_type = $input_type."))
        end
    end
    all_inds = start_ind:jstep:end_ind
    end_ind_thin = convert(Int64, floor(end_ind / jstep))

    # Make sure block_size is a multiple of jstep
    block_size = forecast_block_size(m)
    if end_ind_thin < block_size
        if input_type != :subset
            error("Number of draw ($(end_ind_thin)) divided by jstep ($(jstep)) is less than forecast block size.")
        else
            error("Last index of subset of draws ($(end_ind_thin)) divided by jstep ($(jstep)) is less than forecast block size.")
        end
    end
    if block_size % jstep != 0
        error("forecast_block_size(m) must be a multiple of jstep = $jstep")
    end
    block_size_thin = convert(Int64, block_size / jstep)
    nblocks = convert(Int64, ceil(ndraws / block_size))

    # Fill in draw indices for each block
    block_inds      = Vector{AbstractRange{Int64}}(undef, nblocks)
    block_inds_thin = Vector{AbstractRange{Int64}}(undef, nblocks)
    current_draw      = start_ind - 1
    current_draw_thin = start_ind - 1
    for i = 1:(nblocks-1)
        block_inds[i]      = (current_draw+jstep):jstep:(current_draw+block_size)
        block_inds_thin[i] = (current_draw_thin+1):(current_draw_thin+block_size_thin)
        current_draw      += block_size
        current_draw_thin += block_size_thin
    end
    block_inds[end]      = (current_draw+jstep):jstep:end_ind
    block_inds_thin[end] = (current_draw_thin+1):end_ind_thin
    return block_inds, block_inds_thin
end

"""
```
add_requisite_output_vars(output_vars::Vector{Symbol}; bdd_fcast = true)
```

Based on the given `output_vars`, this function determines which
additional `output_vars` must be computed and stored for future
plotting.

Specifically, when plotting a shock decomposition, the trend and
deterministic trend series are required (the trend is subtracted from
the value of each shock, and the deterministic trend represents
deviations from steady-state that would realize even in the absence of
shocks). For example, if `output_vars` contains `shockdecobs`, the
variables `dettrendobs` and `trendobs` will be added to `output_vars`.

Note that this case is distinct from a case in which computing a
different product is required to compute the desired `output_var`. For
example, smoothed historical states (`histstates`) must be computed in
order to compute a shock decomposition for a state variable, but need
not be saved to produce plots later on. Therefore, `histstates` is not
added to `output_vars` when calling
`add_requisite_output_vars([shockdecstates])`.

The option `bdd_fcast` can allow us to avoid adding `bdd` forecast output vars
if these are not wanted by the user.
"""
function add_requisite_output_vars(output_vars::Vector{Symbol}; bdd_fcast::Bool = true)
    # Add :bddforecast<class> if :forecast<class> is in output_vars
    forecast_outputs = Base.filter(output -> get_product(output) in [:forecast, :forecastut, :forecast4q, :forecastlvl],
                                   output_vars)
    if !isempty(forecast_outputs) && bdd_fcast
        bdd_vars = [Symbol("bdd$(var)") for var in forecast_outputs]
        output_vars = unique(vcat(output_vars, bdd_vars))
    end

    # Add :trend<class> and :dettrend<class> if :shockdec<class> is in output_vars
    shockdec_outputs = Base.filter(output -> get_product(output) == :shockdec, output_vars)
    if !isempty(shockdec_outputs)
        classes = [get_class(output) for output in shockdec_outputs]
        dettrend_vars = [Symbol("dettrend$c") for c in classes]
        trend_vars = [Symbol("trend$c") for c in classes]
        output_vars = unique(vcat(output_vars, dettrend_vars, trend_vars))
    end

    return output_vars
end

"""
```
remove_meansbands_only_output_vars(output_vars)
```
"""
function remove_meansbands_only_output_vars(output_vars::Vector{Symbol})
    # All the <product>ut<class> and <product>4q<class> variables are computed
    # during compute_meansbands
    meansbands_only_products = [:histut, :hist4q, :forecastut, :forecast4q,
                                :bddforecastut, :bddforecast4q, :histlvl, :forecastlvl,
                                :bddforecastlvl]

    Base.filter(var -> !(get_product(var) in meansbands_only_products), output_vars)
end

"""
```
transplant_history(history, last_hist_period)
```

Remove the smoothed states, shocks, or pseudo-observables corresponding to
conditional data periods. This is necessary because when we forecast with
conditional data, we smooth beyond the last historical period.
"""
function transplant_history(history::Matrix{T}, last_hist_period::Int) where {T<:AbstractFloat}

    if isempty(history)
        return history
    else
        return history[:, 1:last_hist_period]
    end
end

"""
```
transplant_forecast(history, forecast, last_hist_period)
```

Transplant the smoothed states, shocks, or pseudo-observables corresponding to
conditional data periods from the history to the forecast.
"""
function transplant_forecast(history::Matrix{T}, forecast::Matrix{T},
                             last_hist_period::Int) where {T<:AbstractFloat}

    ncondperiods = size(history, 2) - last_hist_period
    cond_range   = (last_hist_period + 1):(last_hist_period + ncondperiods)
    condhist     = history[:, cond_range]

    return hcat(condhist, forecast)
end

"""
```
transplant_forecast_observables(histstates, forecastobs, system, last_hist_period)
```

Transplant the observables implied by `histstates` corresponding to
conditional data periods from the history to the forecast.

This exists as a separate function from `transplant_forecast` because we don't
usually map the smoothed historical states into observables, since they would
just give us back the data. However, in the conditional data periods, we only
have data for a subset of observables, so we need to get the remaining
observables by mapping the smoothed states.
"""
function transplant_forecast_observables(histstates::Matrix{T}, forecastobs::Matrix{T},
                                         system::System{T},
                                         last_hist_period::Int) where {T<:AbstractFloat}

    nvars        = size(forecastobs, 1)
    ncondperiods = size(histstates, 2) - last_hist_period
    cond_range   = (last_hist_period + 1):(last_hist_period + ncondperiods)

    condstates   = histstates[:, cond_range]
    condobs      = system[:ZZ]*condstates .+ system[:DD]

    return hcat(condobs, forecastobs)
end

"""
```
standardize_shocks(shocks, QQ)
```

Normalize shocks by their standard deviations. Shocks with zero standard
deviation will be set to zero.
"""
function standardize_shocks(shocks::AbstractMatrix{T}, QQ::AbstractMatrix{T}) where {T<:AbstractFloat}
    stdshocks = shocks ./ sqrt.(diag(QQ))

    zeroed_shocks = findall(diag(QQ) .== 0)
    stdshocks[zeroed_shocks, :] .= 0

    return stdshocks
end

"""
```
standardize_shocks(shocks, QQs, start_date, end_date)
```

Normalize shocks by their standard deviations when there is regime switching.
Shocks with zero standard deviation will be set to zero.
"""
function standardize_shocks(shocks::Matrix{T}, QQs::Vector{Matrix{T}},
                            regime_inds::Vector{UnitRange{Int}}) where {T<:AbstractFloat}

    stdshocks = similar(shocks)
    for(QQ, inds) in zip(QQs, regime_inds)
        inshocks = @view shocks[:, inds]
        stdshocks[:, inds] = standardize_shocks(inshocks, QQ)
    end

    return stdshocks
end

"""
```
assemble_block_outputs(dicts)
```

Given a vector `dicts` of forecast output dictionaries, concatenate each output
along the draw dimension and return a new dictionary of the concatenated
outputs.
"""
function assemble_block_outputs(dicts::Vector{Dict{Symbol, Array{Float64}}}; show_failed_percent::Bool = false)
    out = Dict{Symbol, Array{Float64}}()
    if !isempty(dicts)
        _populate_empty_dictionaries!(dicts; show_failed_percent = show_failed_percent)
        for var in keys(dicts[1])
            outputs  = map(dict -> reshape(dict[var], (1, size(dict[var])...)), dicts)
            out[var] = cat(outputs..., dims=1)
        end
    end
    return out
end

"""
```
_populate_empty_dictionaries!(dicts)
```

helps `assemble_block_outputs` handle empty dictionaries by populating
them with NaNs.
"""
function _populate_empty_dictionaries!(dicts::Vector{Dict{Symbol, Array{Float64}}}; show_failed_percent::Bool = false)
    check = isempty.(dicts)
    if any(check)
        empty_dicts  = findall(check)
        if length(empty_dicts) == length(dicts)
            error("All dictionaries in dicts are empty. No valid forecast output was computed by forecast_one_draw")
        else
            ref_dict = dicts[findfirst(.!check)]
            for i in empty_dicts
                for (k, v) in ref_dict
                    dicts[i][k] = fill(NaN, size(v))
                end
            end
            if show_failed_percent
                nan_percent = round(100. * count(check) / length(dicts), digits = 2)
                println("The percentage of failed forecasts is $(nan_percent)%")
            end
        end
    end

    ct = 0
    for dict in dicts
        for (k, v) in dict
            if any(isnan.(v))
                ct += 1
                break
            end
        end
    end
    if show_failed_percent
        nan_percent = round(100. * ct / length(dicts), digits = 2)
        println("The percentage of failed forecasts is $(nan_percent)%")
    end

end

"""
```
get_forecast_output_dims(m, input_type, output_var)
```

Returns the dimension of the forecast output specified by `input_type` and
`output_var`.
"""
function get_forecast_output_dims(m::AbstractDSGEModel, input_type::Symbol, output_var::Symbol;
                                  subset_inds::AbstractRange{Int64} = 1:0)
    prod  = get_product(output_var)
    class = get_class(output_var)

    ndraws = if input_type in [:mode, :mean, :init]
        1
    elseif input_type in [:full, :subset, :prior, :init_draw_shocks, :mode_draw_shocks]
        _, block_inds_thin = forecast_block_inds(m, input_type; subset_inds = subset_inds)
        sum(map(length, block_inds_thin))
    end

    nvars = if class == :states
        n_states_augmented(m)
    elseif class == :obs
        n_observables(m)
    elseif class == :pseudo
        n_pseudo_observables(m)
    elseif class in [:shocks, :stdshocks]
        n_shocks_exogenous(m)
    end

    nperiods = if prod == :hist
        n_mainsample_periods(m)
    elseif prod in [:forecast, :bddforecast]
        forecast_horizons(m)
    elseif prod in [:shockdec, :dettrend]
        n_shockdec_periods(m)
    elseif prod == :irf
        impulse_response_horizons(m)
    end

    if prod == :trend
        regime_switching = haskey(m.settings, :regime_switching) ? get_setting(m, :regime_switching) : false
        if regime_switching
            return (ndraws, nvars, get_setting(m, :n_regimes))
        else
            return (ndraws, nvars)
        end
    elseif prod in [:hist, :forecast, :bddforecast, :dettrend]
        return (ndraws, nvars, nperiods)
    elseif prod in [:shockdec, :irf]
        nshocks = n_shocks_exogenous(m)
        return (ndraws, nvars, nperiods, nshocks)
    end
end


text:
############################################################ Sentido 1: México-Acapulco #############################################################

function Sentido1()
   Sentido_1 = Carretera2D(2, 3640+6)
    return Sentido_1, Sentido_1.carretera
end

Secciones1 = [150, 200, 400, 600, 750, 900, 1100, 1400, 1700, 1850, 2100, 2300, 2350, 2600, 2900,
            3100, 3200, 3600 ]

          # = (x, lramp)

Secciones1_CR = [150, 200, 320, 390, 540, 570, 810, 850, 980, 1030, 1400, 1530, 1730, 1810, 2190, 2230,
                 2300, 2345, 2450, 2515, 2940, 2970, 3140, 3175, 3600]

S1_P_camion_s1 = [0.1, 0.09, 0.07, 0.05, 0.03, 0.01, 0.01, 0.008, 0.006]

S1_P_camion_s2 = [0.11, 0.08, 0.06, 0.04, 0.02, 0.009, 0.007, 0.005, 0.003]

S1_in_rampa = [(178, 21), (371, 38), (557, 3), (838, 30), (1015, 31), (1512, 100), (1755, 10), (1791, 10), (2215, 15),
               (2321, 1), (2500, 20), (2958, 6), (3158, 3)]

S1_out_rampa = [(176, 10), (257, 11), (472, 20), (800, 22), (957, 15), (1414, 50), (1628, 15), (1729, 10), (1755, 10),
                (2158, 7), (2321, 1), (2358, 20), (2928, 6), (3115, 10), (3557, 13)]

S1_in_rampa_ampli = [(10, 20), (2500, 20), (2958, 6), (3158, 3)]

P_camion_S1 = [
                  (0.19, 0.02, 0.04),
                  (0.25, 0.05, 0.5),
                  (0.23, 0.07, 0.07),
                  (0.18, 0.09, 0.12),
                  (0.16, 0.05, 0.13),
                  (0.12, 0.11, 0.18),
                  (0.1, 0.14, 0.2),
                  (0.08, 0.16, 0.21),
                  (0.05, 0.19, 0.23)
              ]

############################################################## Sentido 2: Acapulco-México ##########################################################

function Sentido2()
    Sentido_2 = Carretera2D(3, 3640+6)
    for i = 1:3440 # 3685 = 3863 - 177 - 1
        Obstaculo!(Sentido_2.carretera[1].carretera[i])
    end
    return Sentido_2, Sentido_2.carretera
end

Secciones2 = [100, 250, 500, 600, 700, 850, 1000, 1300, 1550, 1600, 1700, 2000, 2200,
            2300, 2500, 2700, 2800, 3000, 3300, 3400, 3500, 3600]

          # = (x, lramp)

Secciones2_CR = [190, 230, 500, 580, 790, 825, 900, 950, 1450, 1520, 1565, 1600, 1650, 1700, 2120, 2165, 2450, 2500,
                 2750, 2800, 2890, 2930, 3300, 3400, 3550, 3600]

S2_P_camion_s1 = [0.1, 0.09, 0.07, 0.05, 0.03, 0.01, 0.009, 0.007, 0.006]

S2_P_camion_s2 = [0.1, 0.08, 0.06, 0.04, 0.02, 0.009, 0.007, 0.005, 0.003]

S2_in_rampa = [(214, 8), (562, 35), (810, 5), (920, 4), (1500, 20), (1585, 4), (1686, 5), (2140, 4),  (2479, 10), (2785, 6),
               (2915, 5), (3336, 8), (3378, 6), (3572, 13)]

S2_out_rampa = [(107, 12), (528, 6), (760, 6), (900, 4), (1357, 20), (1571, 4), (1629, 4), (2057, 6), (2107, 4), (2342, 7),
                  (2757, 7), (2871, 22), (3314, 4), (3542, 5)]

S2_in_rampa_ampli =  [(810, 5), (920, 4), (1500, 20), (3600, 20)]

P_camion_S2 = [
                  (0.15, 0.04, 0.04),
                  (0.19, 0.04, 0.02),
                  (0.14, 0.02, 0.04),
                  (0.13, 0.01, 0.05),
                  (0.09, 0.04, 0.07),
                  (0.07, 0.06, 0.12),
                  (0.07, 0.07, 0.15),
                  (0.05, 0.09, 0.14),
                  (0.04, 0.08, 0.22)
              ]


text:
@testset "BBC & SID exp" begin
    α = big(1.1)

    @testset "BBC exp" begin
        x = 1.1
        p0 = 1 # Order of approx for 0 matrix multiplies
        (graph, cref) = graph_bbc_exp(0)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p0 + 1

        x = 1.1
        p1 = 2 # Order of approx for 1 matrix multiplies
        (graph, cref) = graph_bbc_exp(1)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p1 + 1

        x = 1.1
        p2 = 4 # Order of approx for 2 matrix multiplies
        (graph, cref) = graph_bbc_exp(2, T = BigFloat)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p2 + 1

        x = 1.1
        p3 = 8 # Order of approx for 3 matrix multiplies
        (graph, cref) = graph_bbc_exp(3)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p3 + 1

        x = 2.0
        p4bbc = 12 # Order of approx for 4 matrix multiplies
        (graph, cref) = graph_bbc_exp(4)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p4bbc + 1

        x = 2.3
        p5bbc = 18 # Order of approx for 5 matrix multiplies
        (graph, cref) = graph_bbc_exp(5)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p5bbc + 1
    end

    @testset "SID exp" begin
        x = 2.0
        p0sid = 1 # Order of approx for 0 matrix multiplies
        (graph, cref) = graph_sid_exp(0)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p0sid + 1

        x = 2.0
        p1sid = 2 # Order of approx for 1 matrix multiplies
        (graph, cref) = graph_sid_exp(1)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p1sid + 1

        x = 2.0
        p2sid = 4 # Order of approx for 2 matrix multiplies
        (graph, cref) = graph_sid_exp(2)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p2sid + 1

        x = 2.0
        p3sid = 6 # Order of approx for 3 matrix multiplies
        (graph, cref) = graph_sid_exp(3)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p3sid + 1

        x = 2.0
        p4sid = 15 # Order of approx for 4 matrix multiplies
        (graph, cref) = graph_sid_exp(4)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p4sid + 1

        x = 3.2
        p5sid = 21 # Order of approx for 5 matrix multiplies
        (graph, cref) = graph_sid_exp(5)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p5sid + 1

        x = 3.5
        p6sid = 24 # Order of approx for 6 matrix multiplies
        (graph, cref) = graph_sid_exp(6)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p6sid + 1

        x = 5.2
        p7sid = 30 # Order of approx for 7 matrix multiplies
        (graph, cref) = graph_sid_exp(7)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p7sid + 1

        # Test scaling and squaring
        x = 6.5
        (graph, cref) = graph_sid_exp(8)
        @test abs(eval_graph(big(graph), big(x)) - exp(big(x))) < 5e2 * eps()
    end
end


text:
module Dist

import Distributions
const Djl = Distributions
using DocStringExtensions: SIGNATURES

import ..Inference: defalg

import ..Omega
import ..Omega:lift
import ..Prim
import ..Util: *ₛ
import ..NonDet: RandVar, elemtype
import ..Causal: ReplaceRandVar
import ..Inference: FailUnsat
using Spec

# Modules with dist props
import Statistics: mean, var
import Base: minimum, maximum

# Dist Ops
export  succprob,
        support,
        failprob,
        maximum,
        minimum,
        islowerbounded,
        isupperbounded,
        isbounded,
        std,
        mean,
        median,
        mode,
        modes,
        skewness,
        kurtosis,
        isplatykurtic,
        ismesokurtic,
        isleptokurtic,
        entropy,
        prob

        # Dist Algs
export samplemean,
       sampleprob

        # Lifted distributional functions
export  succprobᵣ,
        supportᵣ,
        failprobᵣ,
        maximumᵣ,
        minimumᵣ,
        islowerboundedᵣ,
        isupperboundedᵣ,
        isboundedᵣ,
        stdᵣ,
        medianᵣ,
        modeᵣ,
        modesᵣ,
        meanᵣ,
        probᵣ,
        skewnessᵣ,
        kurtosisᵣ,
        isplatykurticᵣ,
        ismesokurticᵣ,
        isleptokurticᵣ,
        entropyᵣ,
        meanᵣ,
        varᵣ,
        samplemeanᵣ,
        sampleprobᵣ

include("interface.jl")     # Distributional operators
include("distalgs.jl")      # Algorithms that compute distritbuional operators
include("djl.jl")           # Distributions.jl interop
include("liftdistop.jl")    # Lifting for higher order
include("defaults.jl")      # Distributions.jl interop

end


text:
#=
Project trees are usually made up of multiple files. An AbstractServer holds the AbstractFiles that represent this tree. FileServer is the basic implementation and assumes files are available and readable from disc. (LanguageServer illustrates another implementaiton). The accompanying functions summarised below are required for making an alternative implementation.

Interface spec.
AbstractServer :-> (has/canload/load/set/get)file, getsymbols, getsymbolextends
AbstractFile :-> (get/set)path, (get/set)root, (get/set)cst, semantic_pass, (get/set)server
=#
abstract type AbstractServer end
abstract type AbstractFile end

mutable struct File
    path::String
    source::String
    cst::EXPR
    root::Union{Nothing,File}
    server
end

mutable struct FileServer <: AbstractServer
    files::Dict{String,File}
    roots::Set{File}
    workspacepackages::Dict{String,File} # list of files that may represent within-workspace packages
    external_env::ExternalEnv
end
FileServer() = FileServer(Dict{String,File}(), Set{File}(), Dict{String,File}(), ExternalEnv(Dict{Symbol,SymbolServer.ModuleStore}(:Base => SymbolServer.stdlibs[:Base], :Core => SymbolServer.stdlibs[:Core]), SymbolServer.collect_extended_methods(SymbolServer.stdlibs), Symbol[]))

hasfile(server::FileServer, path::String) = haskey(server.files, path)
canloadfile(server, path) = isfile(path)
function setfile(server::FileServer, path::String, file::File)
    server.files[path] = file
end
getfile(server::FileServer, path::String) = server.files[path]
function loadfile(server::FileServer, path::String)
    source = read(path, String)
    cst = CSTParser.parse(source, true)
    f = File(path, source, cst, nothing, server)
    setroot(f, f)
    setfile(server, path, f)
    return getfile(server, path)
end

getsymbols(env::ExternalEnv) = env.symbols
getsymbols(state::State) = getsymbols(state.env)

getsymbolextendeds(env::ExternalEnv) = env.extended_methods
getsymbolextendeds(state::State) = getsymbolextendeds(state.env)

"""
    getenv(file::File, server::FileServer)

Get the relevant `ExternalEnv` for a given file.
"""
function getenv(file::File, server::FileServer)
    # For FileServer this approach is equivalent to the previous behaviour. Other AbstractServers
    # (e.g. LanguageServerInstance) can use this function to associate different files (or trees of
    # files) with different environments.
    server.external_env
end

getpath(file::File) = file.path

getroot(file::File) = file.root
function setroot(file::File, root::File)
    file.root = root
    return file
end

getcst(file::File) = file.cst
function setcst(file::File, cst::EXPR)
    file.cst = cst
    return file
end

getserver(file::File) = file.server
function setserver(file::File, server::FileServer)
    file.server = server
    return file
end

function Base.display(f::File)
    println(f.path)
end

function Base.display(s::FileServer)
    n = length(s.files)
    println(n, "-file Server")
    cnt = 0
    for p in keys(s.files)
        cnt += 1
        println(" ", p)
        cnt > 10 && break
    end
end


text:
"""
    PrecipFormation

Storage for tendencies due to precipitation formation

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct PrecipFormation{FT}
    θ_liq_ice_tendency::FT
    qt_tendency::FT
    qr_tendency::FT
    qs_tendency::FT
end

"""
    EntrDetr

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct EntrDetr{FT}
    "Dynamical entrainment"
    ε_dyn::FT
    "Dynamical detrainment"
    δ_dyn::FT
    "Turbulent entrainment"
    ε_turb::FT
    "Horizontal eddy-diffusivity"
    K_ε::FT
end

"""
    GeneralizedEntr

A general set of variables entrainment might depend on.

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct GeneralizedEntr{FT}
    "updraft condensate (liquid water + ice)"
    q_cond_up::FT
    "environment condensate (liquid water + ice)"
    q_cond_en::FT
    "updraft vertical velocity"
    w_up::FT
    "environment vertical velocity"
    w_en::FT
    "updraft buoyancy"
    b_up::FT
    "environment buoyancy"
    b_en::FT
    "environment tke"
    tke::FT
    "updraft momentum divergence"
    dMdz::FT
    "updraft momentum"
    M::FT
    "updraft area fraction"
    a_up::FT
    "environment area fraction"
    a_en::FT
    "pressure plume spacing"
    R_up::FT
    "updraft relative humidity"
    RH_up::FT
    "environment relative humidity"
    RH_en::FT
    "maximum updraft area"
    max_area::FT
    "vertical coordinate"
    zc_i::FT
    "Model time step"
    Δt::FT
end

Base.eltype(::GeneralizedEntr{FT}) where {FT} = FT

struct MDEntr end  # existing model
struct NNEntr end

"""
    GradBuoy

Environmental buoyancy gradients.

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct GradBuoy{FT}
    "environmental vertical buoyancy gradient"
    ∂b∂z::FT
    "vertical buoyancy gradient in the unsaturated part of the environment"
    ∂b∂z_unsat::FT
    "vertical buoyancy gradient in the saturated part of the environment"
    ∂b∂z_sat::FT
end

abstract type AbstractEnvBuoyGradClosure end
struct BuoyGradMean <: AbstractEnvBuoyGradClosure end
struct BuoyGradQuadratures <: AbstractEnvBuoyGradClosure end

"""
    EnvBuoyGrad

Variables used in the environmental buoyancy gradient computation.

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct EnvBuoyGrad{FT, EBC <: AbstractEnvBuoyGradClosure}
    "temperature in the saturated part"
    t_sat::FT
    "vapor specific humidity  in the saturated part"
    qv_sat::FT
    "total specific humidity in the saturated part"
    qt_sat::FT
    "potential temperature in the saturated part"
    θ_sat::FT
    "liquid ice potential temperature in the saturated part"
    θ_liq_ice_sat::FT
    "virtual potential temperature gradient in the non saturated part"
    ∂θv∂z_unsat::FT
    "total specific humidity gradient in the saturated part"
    ∂qt∂z_sat::FT
    "liquid ice potential temperature gradient in the saturated part"
    ∂θl∂z_sat::FT
    "reference pressure"
    p0::FT
    "cloud fraction"
    en_cld_frac::FT
    "specific volume"
    alpha0::FT
end
function EnvBuoyGrad(::EBG; t_sat::FT, bg_kwargs...) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
    return EnvBuoyGrad{FT, EBG}(; t_sat, bg_kwargs...)
end

"""
    MixLen

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct MixLen{FT}
    "minimum length number"
    min_len_ind::Int
    "mixing length"
    mixing_length::FT
    "length ratio"
    ml_ratio::FT
end

"""
    MinDisspLen

Minimum dissipation model

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct MinDisspLen{FT}
    "height"
    z::FT
    "obukhov length"
    obukhov_length::FT
    "surface TKE values"
    tke_surf::FT
    "u star - surface velocity scale"
    ustar::FT
    "turbulent Prandtl number"
    Pr::FT
    "reference pressure"
    p0::FT
    "vertical buoyancy gradient struct"
    ∇b::GradBuoy{FT}
    "env shear"
    Shear²::FT
    "environment turbulent kinetic energy"
    tke::FT
    "Updraft tke source"
    b_exch::FT
end

Base.@kwdef mutable struct PrecipVariables
    precipitation_model::String = "default_precipitation_model"
    mean_rwp::Float64 = 0
    mean_swp::Float64 = 0
    cutoff_precipitation_rate::Float64 = 0
end
function PrecipVariables(namelist, grid::Grid)

    precipitation_model = parse_namelist(
        namelist,
        "microphysics",
        "precipitation_model";
        default = "None",
        valid_options = ["None", "cutoff", "clima_1m"],
    )

    if !(precipitation_model in ["None", "cutoff", "clima_1m"])
        error("precipitation model not recognized")
    end

    return PrecipVariables(; precipitation_model)
end

mutable struct UpdraftVariables{A1}
    n_updrafts::Int
    cloud_base::A1
    cloud_top::A1
    cloud_cover::A1
    updraft_top::A1
    lwp::Float64
    iwp::Float64
    function UpdraftVariables(nu, namelist, grid::Grid)
        n_updrafts = nu

        # cloud and precipitation diagnostics for output
        cloud_base = zeros(nu)
        cloud_top = zeros(nu)
        cloud_cover = zeros(nu)
        updraft_top = zeros(nu)

        lwp = 0.0
        iwp = 0.0

        A1 = typeof(cloud_base)
        return new{A1}(n_updrafts, cloud_base, cloud_top, cloud_cover, updraft_top, lwp, iwp)
    end
end

Base.@kwdef mutable struct GridMeanVariables{PS}
    param_set::PS
    lwp::Float64
    iwp::Float64
    cloud_base::Float64
    cloud_top::Float64
    cloud_cover::Float64
    EnvThermo_scheme::String
end
function GridMeanVariables(namelist, grid::Grid, param_set::PS) where {PS}
    lwp = 0.0
    iwp = 0.0

    cloud_base = 0.0
    cloud_top = 0.0
    cloud_cover = 0.0

    EnvThermo_scheme = parse_namelist(namelist, "thermodynamics", "sgs"; default = "mean")

    return GridMeanVariables(; param_set, lwp, iwp, cloud_base, cloud_top, cloud_cover, EnvThermo_scheme)
end

Base.@kwdef mutable struct EnvironmentVariables
    cloud_base::Float64 = 0
    cloud_top::Float64 = 0
    cloud_cover::Float64 = 0
    lwp::Float64 = 0
    iwp::Float64 = 0
    EnvThermo_scheme::String = "default_EnvThermo_scheme"
end
function EnvironmentVariables(namelist, grid::Grid)
    # TODO: EnvThermo_scheme is repeated in GridMeanVariables
    EnvThermo_scheme = parse_namelist(namelist, "thermodynamics", "sgs"; default = "mean")
    return EnvironmentVariables(; EnvThermo_scheme)
end

struct EnvironmentThermodynamics
    quadrature_order::Int
    quadrature_type::String
    function EnvironmentThermodynamics(namelist, grid::Grid)
        quadrature_order = parse_namelist(namelist, "thermodynamics", "quadrature_order"; default = 3)
        quadrature_type = parse_namelist(namelist, "thermodynamics", "quadrature_type"; default = "gaussian")
        return new(quadrature_order, quadrature_type)
    end
end

# Stochastic entrainment/detrainment closures:
struct NoneClosureType end
struct LogNormalClosureType end
struct SDEClosureType end

# Stochastic differential equation memory
Base.@kwdef mutable struct sde_struct{T}
    u0::Float64
    dt::Float64
end

# SurfaceMoninObukhovDry:
#     Needed for dry cases (qt=0). They have to be initialized with nonzero qtg for the
#     reference profiles. This surface subroutine sets the latent heat flux to zero
#     to prevent errors due to nonzero qtg in vars such as the obukhov_length.
# SurfaceSullivanPatton
#     Not fully implemented yet. Maybe not needed - Ignacio
struct SurfaceNone end
struct SurfaceFixedFlux end
struct SurfaceFixedCoeffs end
struct SurfaceMoninObukhov end
struct SurfaceMoninObukhovDry end
struct SurfaceSullivanPatton end

Base.@kwdef mutable struct SurfaceBase{T}
    zrough::Float64 = 0
    interactive_zrough::Bool = false
    Tsurface::Float64 = 0
    qsurface::Float64 = 0
    shf::Float64 = 0
    lhf::Float64 = 0
    cm::Float64 = 0
    ch::Float64 = 0
    cq::Float64 = 0
    bflux::Float64 = 0
    windspeed::Float64 = 0
    ustar::Float64 = 0
    rho_qtflux::Float64 = 0
    rho_hflux::Float64 = 0
    rho_uflux::Float64 = 0
    rho_vflux::Float64 = 0
    obukhov_length::Float64 = 0
    Ri_bulk_crit::Float64 = 0
    ustar_fixed::Bool = false
    ref_params::NamedTuple = NamedTuple()
end

function SurfaceBase(::Type{T}; namelist::Dict, ref_params) where {T}
    Ri_bulk_crit = namelist["turbulence"]["Ri_bulk_crit"]
    return SurfaceBase{T}(; Ri_bulk_crit, ref_params)
end

struct ForcingBaseType end
struct ForcingNone end
struct ForcingStandard end
struct ForcingDYCOMS_RF01 end
struct ForcingLES end

struct RadiationBaseType end
struct RadiationNone end
struct RadiationStandard end
struct RadiationDYCOMS_RF01 end
struct RadiationLES end

Base.@kwdef mutable struct LESData
    "Start time index of LES"
    imin::Int = 0
    "End time index of LES"
    imax::Int = 0
    "Path to LES stats file used to drive SCM"
    les_filename::String = nothing
    "Drive SCM with LES data from t = [end - t_interval_from_end_s, end]"
    t_interval_from_end_s::Float64 = 6 * 3600.0
    "Length of time to average over for SCM initialization"
    initial_condition_averaging_window_s::Float64 = 3600.0
end

"""
    ForcingBase

LES-driven forcing

$(DocStringExtensions.FIELDS)
"""
Base.@kwdef mutable struct ForcingBase{T}
    "Boolean specifying whether Coriolis forcing is applied"
    apply_coriolis::Bool = false
    "Boolean specifying whether subsidence forcing is applied"
    apply_subsidence::Bool = false
    "Coriolis parameter"
    coriolis_param::Float64 = 0
    "Momentum relaxation timescale"
    nudge_tau::Float64 = 0.0
end

force_type(::ForcingBase{T}) where {T} = T

Base.@kwdef mutable struct RadiationBase{T}
    divergence::Float64 = 0
    alpha_z::Float64 = 0
    kappa::Float64 = 0
    F0::Float64 = 0
    F1::Float64 = 0
end

rad_type(::RadiationBase{T}) where {T} = T

Base.@kwdef mutable struct CasesBase{T}
    case::T
    casename::String = "default_casename"
    inversion_option::String = "default_inversion_option"
    les_filename::String = "None"
    Sur::SurfaceBase
    Fo::ForcingBase
    Rad::RadiationBase
    rad_time::StepRangeLen = range(10, 360; length = 36) .* 60
    rad::AbstractMatrix{Float64} = zeros(1, 1)
    lhf0::Float64 = 0
    shf0::Float64 = 0
    LESDat::Union{LESData, Nothing} = nothing
end

CasesBase(case::T; kwargs...) where {T} = CasesBase{T}(; case = case, casename = string(nameof(T)), kwargs...)

mutable struct EDMF_PrognosticTKE{N_up, A1, EBGC, EC, SDES, UPVAR}
    Ri_bulk_crit::Float64
    zi::Float64
    n_updrafts::Int
    asp_label::String
    extrapolate_buoyancy::Bool
    surface_area::Float64
    max_area::Float64
    aspect_ratio::Float64
    tke_ed_coeff::Float64
    tke_diss_coeff::Float64
    static_stab_coeff::Float64
    minimum_area::Float64
    Precip::PrecipVariables
    UpdVar::UPVAR
    EnvVar::EnvironmentVariables
    EnvThermo::EnvironmentThermodynamics
    area_surface_bc::A1
    w_surface_bc::A1
    h_surface_bc::A1
    qt_surface_bc::A1
    pressure_plume_spacing::A1
    prandtl_number::Float64
    wstar::Float64
    entr_surface_bc::Float64
    detr_surface_bc::Float64
    dt_max::Float64
    sde_model::SDES
    bg_closure::EBGC
    entr_closure::EC
    function EDMF_PrognosticTKE(namelist, grid::Grid, param_set::PS) where {PS}
        # get values from namelist
        prandtl_number = namelist["turbulence"]["prandtl_number_0"]
        Ri_bulk_crit = namelist["turbulence"]["Ri_bulk_crit"]
        zi = 0.0

        # Set the number of updrafts (1)
        n_updrafts = parse_namelist(namelist, "turbulence", "EDMF_PrognosticTKE", "updraft_number"; default = 1)

        pressure_func_drag_str = parse_namelist(
            namelist,
            "turbulence",
            "EDMF_PrognosticTKE",
            "pressure_closure_drag";
            default = "normalmode",
            valid_options = ["normalmode", "normalmode_signdf"],
        )

        asp_label = parse_namelist(
            namelist,
            "turbulence",
            "EDMF_PrognosticTKE",
            "pressure_closure_asp_label";
            default = "const",
        )

        extrapolate_buoyancy =
            parse_namelist(namelist, "turbulence", "EDMF_PrognosticTKE", "extrapolate_buoyancy"; default = true)

        # Get values from namelist
        # set defaults at some point?
        surface_area = namelist["turbulence"]["EDMF_PrognosticTKE"]["surface_area"]
        max_area = namelist["turbulence"]["EDMF_PrognosticTKE"]["max_area"]
        # entrainment parameters
        # pressure parameters
        aspect_ratio = namelist["turbulence"]["EDMF_PrognosticTKE"]["aspect_ratio"]

        # mixing length parameters
        tke_ed_coeff = namelist["turbulence"]["EDMF_PrognosticTKE"]["tke_ed_coeff"]
        tke_diss_coeff = namelist["turbulence"]["EDMF_PrognosticTKE"]["tke_diss_coeff"]
        static_stab_coeff = namelist["turbulence"]["EDMF_PrognosticTKE"]["static_stab_coeff"]
        # Need to code up as namelist option?
        minimum_area = 1e-5

        # Create the class for precipitation
        Precip = PrecipVariables(namelist, grid)

        # Create the updraft variable class (major diagnostic and prognostic variables)
        UpdVar = UpdraftVariables(n_updrafts, namelist, grid)

        # Create the environment variable class (major diagnostic and prognostic variables)
        EnvVar = EnvironmentVariables(namelist, grid)
        # Create the class for environment thermodynamics
        EnvThermo = EnvironmentThermodynamics(namelist, grid)

        # Near-surface BC of updraft area fraction
        area_surface_bc = zeros(n_updrafts)
        w_surface_bc = zeros(n_updrafts)
        h_surface_bc = zeros(n_updrafts)
        qt_surface_bc = zeros(n_updrafts)
        pressure_plume_spacing = zeros(n_updrafts)

        # Initialize SDE parameters
        dt = parse_namelist(namelist, "time_stepping", "dt_min"; default = 1.0)
        closure = parse_namelist(
            namelist,
            "turbulence",
            "EDMF_PrognosticTKE",
            "stochastic",
            "closure";
            default = "none",
            valid_options = ["none", "lognormal", "sde"],
        )
        closure_type = if closure == "none"
            NoneClosureType
        elseif closure == "lognormal"
            LogNormalClosureType
        elseif closure == "sde"
            SDEClosureType
        else
            error("Something went wrong. Invalid stochastic closure type '$closure'")
        end
        sde_model = sde_struct{closure_type}(u0 = 1, dt = dt)

        bg_type = parse_namelist(
            namelist,
            "turbulence",
            "EDMF_PrognosticTKE",
            "env_buoy_grad";
            default = "mean",
            valid_options = ["mean", "quadratures"],
        )
        bg_closure = if bg_type == "mean"
            BuoyGradMean()
        elseif bg_type == "quadratures"
            BuoyGradQuadratures()
        else
            error("Something went wrong. Invalid environmental buoyancy gradient closure type '$bg_type'")
        end

        # entr closure
        entr_type = parse_namelist(
            namelist,
            "turbulence",
            "EDMF_PrognosticTKE",
            "entrainment";
            default = "moisture_deficit",
            valid_options = ["moisture_deficit", "NN"],
        )
        entr_closure = if entr_type == "moisture_deficit"
            MDEntr()
        elseif entr_type == "NN"
            NNEntr()
        else
            error("Something went wrong. Invalid entrainment type '$entr_type'")
        end
        EC = typeof(entr_closure)

        wstar = 0
        entr_surface_bc = 0
        detr_surface_bc = 0
        dt_max = 0
        A1 = typeof(area_surface_bc)
        EBGC = typeof(bg_closure)
        SDES = typeof(sde_model)
        UPVAR = typeof(UpdVar)
        return new{n_updrafts, A1, EBGC, EC, SDES, UPVAR}(
            Ri_bulk_crit,
            zi,
            n_updrafts,
            asp_label,
            extrapolate_buoyancy,
            surface_area,
            max_area,
            aspect_ratio,
            tke_ed_coeff,
            tke_diss_coeff,
            static_stab_coeff,
            minimum_area,
            Precip,
            UpdVar,
            EnvVar,
            EnvThermo,
            area_surface_bc,
            w_surface_bc,
            h_surface_bc,
            qt_surface_bc,
            pressure_plume_spacing,
            prandtl_number,
            wstar,
            entr_surface_bc,
            detr_surface_bc,
            dt_max,
            sde_model,
            bg_closure,
            entr_closure,
        )
    end
end
parameter_set(obj) = obj.param_set
n_updrafts(edmf::EDMF_PrognosticTKE{N_up}) where {N_up} = N_up

struct State{P, A, T}
    prog::P
    aux::A
    tendencies::T
end