text:
#TODO: move shock sensor implementation to own file

#------------------------------------------------------------------------------
# Method from: Persson and Peraire, "Sub-Cell Shock Capturing for DG Methods"
#              AIAA 2006
# Some additional details are from Barter's Ph.D Thesis, "Shock Capturing
# with PDE-Based Artificial Viscosity for an Adaptive, Higher-Order DG
# Finite Element Method", MIT 2008.

"""
  Computes the shock sensor and the numerical viscoscity for the Persson
  and Perairi paper
"""
function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                          sensor::ShockSensorPP,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                          Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tmsh, Tdim}

  Tres = promote_type(Tsol, Tmsh)
  numNodesPerElement = size(q, 2)

  @unpack sensor up up_tilde up1_tilde s0 kappa e0
  #fill!(up_tilde, 0); fill!(up1_tilde, 0)

  # use density as the key variable
  for i=1:numNodesPerElement
    up[i] = q[1, i]
  end

  getFilteredSolutions(params, sensor.Vp, up, up_tilde, up1_tilde)
  #getFilteredSolution(params, sensor.Vp, up, up_tilde)
  #getFilteredSolution(params, sensor.Vp1, up, up1_tilde)

  # compute the inner product
  num = zero(Tres)
  den = zero(Tres)

  # use the filtered variables for (u, u) and ||u - u_tilde||.
  # This is a bit different than
  # finite element methods, where the original solution has a basis, and the
  # norm in any basis should be the same.  Here we use the filtered u rather
  # than the original because it is probably smoother.
  if sensor.use_filtered
    for i=1:numNodesPerElement
      fac = sbp.w[i]/jac[i]
      delta_u = up_tilde[i] - up1_tilde[i]

      num += delta_u*fac*delta_u
      den += up_tilde[i]*fac*up_tilde[i]
    end
  else
    # use the original u for (u, u) and ||u - u_tilde||.  For the diagonal E
    # operators, this means u - u_tilde contains the highest 2 degree
    # polynomials (because diagonal E operators have enough nodes for one
    # degree higher solution)
    for i=1:numNodesPerElement
      fac = sbp.w[i]/jac[i]
      delta_u = up[i] - up1_tilde[i]

      num += delta_u*fac*delta_u
      den += up[i]*fac*up[i]
    end
  end

  Se = num/den
  se = log10(Se)

  # should this be a separate function from computing Se?
  if se < s0 - kappa
    ee = Tres(0.0)
    is_nonzero = false
  elseif se > s0 - kappa && se < s0 + kappa
    ee = 0.5*e0*(1 + sinpi( (se - s0)/(2*kappa)))
    is_nonzero = true
  else
    ee = Tres(e0)
    is_nonzero = true
  end

  # scale by lambda_max * h/p to get subcell resolution
  lambda_max = zero(Tsol)
  h_avg = zero(Tmsh)
  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    lambda_max += getLambdaMax(params, q_i)
    h_avg += sbp.w[i]/jac[i]
  end

  lambda_max /= numNodesPerElement
  h_avg = h_avg^(1/Tdim)

  ee *= lambda_max*h_avg/sbp.degree

  # use elementwise constant
  fill!(Se_mat, Se)
  fill!(ee_mat, ee)

  return is_nonzero
end

"""
  Filter the solution for the shock sensor

  **Inputs**

   * params
   * vand: the [`VandermondeData`](@ref) to use
   * u_nodal: vector, length `numNodesPerElement` containing the solution at the
              nodes of the element

  **Inputs/Outputs**

   * u_filt: vector to be overwritten with the filtered solution, same length
             as `u_nodal`
"""
function getFilteredSolution(params::ParamType, vand::VandermondeData,
                             u_nodal::AbstractVector, u_filt::AbstractVector)

  # u_modal = Vpinv * u_nodal
  # u_filt = Vp * u_modal, thus this can be done in one step by:
  # u_filt = (Vp*Vpinv)*u_nodal
  smallmatvec!(vand.filt, u_nodal, u_filt)

  return nothing
end

"""
  Computed the two filtered solutions required by the shock sensor (well, only
  one of them is required, the other one may or may not be useful)

  **Inputs**

   * params
   * vand: `VandermondeData`
   * u_nodal: `numNodesPerElement` nodal solution

  **Inputs/Outputs**

   * u_filt1: The result of `u = Vp*pinv(Vp)*u_nodal`  This projects the
              solution onto the orthogonal basis and then projects it back.
              This may or may not be useful
   * u_filt2: The result of projecting `u` onto the orthgonal basis and then
              projecting all but the highest mode(s) back to a nodal back to
              the nodal basis
"""
function getFilteredSolutions(params::ParamType, vand::VandermondeData,
                              u_nodal::AbstractVector,
                              u_filt1::AbstractVector, u_filt2::AbstractVector)

  smallmatvec!(vand.filt, u_nodal, u_filt1)
  smallmatvec!(vand.filt1, u_nodal, u_filt2)

  return nothing
end

#------------------------------------------------------------------------------
# ShockSensorNone

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorNone,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                          Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tmsh}

  error("getShockSensor called for ShockSensorNone: did you forget to specify the shock capturing scheme?")
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorNone,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                         ) where {Tsol, Tmsh}

  error("isShockElement called for ShockSensorNone: did you forget to specify the shock capturing scheme?")
end

#------------------------------------------------------------------------------
# ShockSensorEverywhere

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorEverywhere{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  fill!(Se_mat, sensor.alpha)
  fill!(ee_mat, sensor.alpha)

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorEverywhere{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorVelocity

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorVelocity{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  numNodesPerElement = size(q, 2)
  dim = size(coords, 1)

  alpha = sensor.alpha
  for i=1:numNodesPerElement
    for d=1:dim
      Se_mat[d, i] = alpha*d*q[d+1, i]
      ee_mat[d, i] = alpha*d*q[d+1, i]
    end
  end

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorVelocity{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHIso

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHIso{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tres, Tmsh, Tdim}

  # compute | div(F) |
  # Do this in Cartesian coordinates because it makes the differentiation easier
  numDofPerNode, numNodesPerElement = size(q)
  res = sensor.res

  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)
  # Note: Hartmann's paper referrs to Jaffre's paper which says to use the
  #       maximum value of res over the element, but that is not differentiable,
  #       so take the L2 norm instead
  val = computeL2Norm(params, sbp, jac, res)
  h_avg = computeElementVolume(params, sbp, jac)

  h_fac = h_avg^((2 - sensor.beta)/Tdim)

  ee = Tres(sensor.C_eps*h_fac*val)
#=
  if ee < 1e-12  # make this zero so the shockmesh machinery knows to exclude
                 # this element, saving cost
    ee = Tres(0)
  end
=#

  fill!(Se_mat, val)
  fill!(ee_mat, ee)

  # calling | div(f) | as the shock sensor, which is somewhat arbitrary
  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHIso{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  return true
end

"""
  Computes del * F.  where del is the diverage in the xyz directions,
  for a single element.

  **Inputs**

   * params
   * sbp
   * data: `StrongFormData` to hold temporary arrays
   * q: solution on the entire element, `numDofPerNode` x `numNodesPerElement`
   * dxidx: scaled mapping jacobian for entire element, `dim` x `dim` x
            `numNodesPerElement`
   * jac: mapping jacobian determinant for entire element, `numNodesPerElement`

  **Inputs/Outputs**

   * res: array to overwrite with residual, same size as `q`.
"""
function computeStrongResidual(params::ParamType{Tdim}, sbp::AbstractOperator,
                               data::StrongFormData{Tsol, Tres},
                               q::AbstractMatrix,
                               dxidx::Abstract3DArray,
                               jac::AbstractVector,
                               res::AbstractMatrix) where {Tdim, Tsol, Tres}

  @unpack data flux nrm aux_vars work
  fill!(res, 0); fill!(nrm, 0)

  numDofPerNode, numNodesPerElement = size(q)

  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    for d=1:Tdim
      nrm[d] = 1
      flux_d = sview(flux, :, i, d)

      calcEulerFlux(params, q_i, aux_vars, nrm, flux_d)

      nrm[d] = 0
    end
  end

  applyDx(sbp, flux, dxidx, jac, work, res)

  return nothing
end

"""
  Computes the integral L2 norm for a single element

  **Inputs**

   * params
   * sbp
   * jac: mapping jacobian determinant for the element, `numNodesPerElement`
   * res: `numDofPerNode` x `numNodesPerElement` array to compute the norm
          of.
"""
function computeL2Norm(params::ParamType{Tdim}, sbp::AbstractOperator,
                       jac::AbstractVector, res::AbstractMatrix{Tres}
                      ) where {Tres, Tdim}

  numDofPerNode, numNodesPerElement = size(res)

  # compute norm
  val = zero(Tres)
  for i=1:numNodesPerElement
    fac = sbp.w[i]/jac[i]
    for j=1:numDofPerNode
      val += res[j, i]*fac*res[j, i]
    end
  end

  val = sqrt(val)

  return val
end

"""
  Computes volume of element by integrating the mapping jacobian determinant
"""
function computeElementVolume(params::ParamType{Tdim}, sbp::AbstractOperator,
                         jac::AbstractVector{Tmsh}) where {Tdim, Tmsh}

  # compute norm and mesh size h
  vol = zero(Tmsh)
  for i=1:length(jac)
    vol += sbp.w[i]/jac[i]
  end

  return vol
end

#------------------------------------------------------------------------------
# ShockSensorBO

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  lambda_max = zero(Tsol)
  h_avg = zero(Tmsh)
  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    lambda_max += getLambdaMax(params, q_i)
    h_avg += sbp.w[i]/jac[i]
  end

  lambda_max /= numNodesPerElement
  h_avg = h_avg^(1/Tdim)

  fill!(Se_mat, lambda_max)
  fill!(ee_mat, sensor.alpha*lambda_max*h_avg/sbp.degree)
#  println("element ", elnum, " SensorBO value = ", real(ee_mat[1, 1]))

  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHHO

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHHO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  @unpack sensor p_dot press_el press_dx work res Rp fp
  fill!(press_dx, 0); fill!(res, 0); fill!(Rp, 0)

  # compute Rm
  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)

  # compute Rp
  @simd for i=1:numNodesPerElement
    q_i = ro_sview(q, :, i)
    press = calcPressure_diff(params, q_i, p_dot)
    press_el[1, i] = press

    @simd for j=1:numDofPerNode
      Rp[i] += p_dot[j]*res[j, i]/press
    end
  end

  # compute pressure switch fp (excluding the h_k factor)
  applyDx(sbp, press_el, dxidx, jac, work, press_dx)
  @simd for i=1:numNodesPerElement
    val = zero(Tres)
    @simd for d=1:Tdim
      val += press_dx[1, i, d]*press_dx[1, i, d]
    end
    fp[i] = sqrt(val)/(press_el[1, i] + 1e-12)
  end

  @simd for i=1:numNodesPerElement
    @simd for d=1:Tdim
      h_fac = sensor.h_k_tilde[d, elnum]
      Se_mat[d, i] = h_fac*fp[i]*absvalue(Rp[i])
      ee_mat[d, i] = Se_mat[d, i]*h_fac*h_fac*sensor.C_eps
    end
  end

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHHOConst

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHHOConst{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  @unpack sensor p_dot press_el press_dx work res Rp fp epsilon
  fill!(press_dx, 0); fill!(res, 0); fill!(Rp, 0)

  # compute Rm
  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)

  # compute Rp
  @simd for i=1:numNodesPerElement
    q_i = ro_sview(q, :, i)
    press = calcPressure_diff(params, q_i, p_dot)
    press_el[1, i] = press

    @simd for j=1:numDofPerNode
      Rp[i] += p_dot[j]*res[j, i]/press
    end
  end

  # compute pressure switch fp (excluding the h_k factor)
  applyDx(sbp, press_el, dxidx, jac, work, press_dx)
  @simd for i=1:numNodesPerElement
    val = zero(Tres)
    @simd for d=1:Tdim
      val += press_dx[1, i, d]*press_dx[1, i, d]
    end
    fp[i] = sqrt(val)/(press_el[1, i] + 1e-12)
  end

  # compute L2 norm of fp * Rp
  @simd for i=1:numNodesPerElement
    epsilon[1, i] = fp[i]*Rp[i]
  end

  #val = computeL2Norm(params, sbp, jac, epsilon)
  val = zero(Tres)
  @simd for i=1:numNodesPerElement
    val += sbp.w[i]*absvalue(epsilon[1, i])
  end

  @simd for i=1:numNodesPerElement
    @simd for d=1:Tdim
      h_fac = sensor.h_k_tilde[d, elnum]
      Se_mat[d, i] = h_fac*val
      ee_mat[d, i] = Se_mat[d, i]*h_fac*h_fac*sensor.C_eps
    end
  end

  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::HHOSensors{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}
  return true
end

"""
  Compute the h_k tilde factors from the paper
"""
function calcAnisoFactors(mesh::AbstractMesh, sbp, opts,
                          hk_all::AbstractMatrix{T}) where {T}

  # compute h_k tilde = h_k/(p + 1), where h_k takes into account the
  # anisotropy of the element in directory k
  # compute the p_k first, then compute h_k, where p_k is the integral of the
  # face normal

  @assert size(hk_all) == (mesh.dim, mesh.numEl)

  fill!(hk_all, 0)

  # do all faces (interior, boundary, shared)
  @simd for i=1:mesh.numInterfaces
    iface_i = mesh.interfaces[i]
    nrm_i = ro_sview(mesh.nrm_face, :, :, i)

    @simd for j=1:mesh.numNodesPerFace
      @simd for k=1:mesh.dim
        fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
        hk_all[k, iface_i.elementL] += fac
        hk_all[k, iface_i.elementR] += fac
      end
    end
  end

  @simd for i=1:mesh.numBoundaryFaces
    bndry_i = mesh.bndryfaces[i]
    nrm_i = ro_sview(mesh.nrm_bndry, :, :, i)

    @simd for j=1:mesh.numNodesPerFace
      @simd for k=1:mesh.dim
        fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
        hk_all[k, bndry_i.element] += fac
      end
    end
  end

  @simd for peer=1:mesh.npeers
    @simd for i=1:length(mesh.shared_interfaces[peer])
      iface_i = mesh.shared_interfaces[peer]
      nrm_i = ro_sview(mesh.nrm_sharedface[peer], :, :, i)

      @simd for j=1:mesh.numNodesPerFace
        @simd for k=1:mesh.dim
          fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
          hk_all[k, iface_i.elementL] += fac
        end
      end
    end
  end

  # hk_all now contains the p_i in each direction
  # Now compute h_k tilde from p_i
  @simd for i=1:mesh.numEl

    jac_i = ro_sview(mesh.jac, :, i)
    vol = zero(T)
    @simd for j=1:mesh.numNodesPerElement
      vol += sbp.w[j]/jac_i[j]
    end

    h_k = vol^(1/mesh.dim)
    pk_prod = one(T)
    @simd for d=1:mesh.dim
      pk_prod *= hk_all[d, i]
    end
    pk_prod = pk_prod^(1/mesh.dim)

    @simd for d=1:mesh.dim
      hk_all[d, i] = h_k*pk_prod/(hk_all[d, i]*(sbp.degree + 1))
    end
  end

  return nothing
end

#------------------------------------------------------------------------------
# SensorHApprox

# even though this shared the name with all the other getShockSensor methods,
# it signature and outputs are different.  This is sort-of ok because this
# is a special shock sensor and should never be used as the main shock sensor.
function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHApprox{Tsol, Tres},
                        elnum::Integer, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numNodesPerElement = length(jac)

  shared_idx = elnum - first(sensor.shared_els) + 1  # index of shared element
  if (elnum in sensor.local_els) ||
     (elnum in sensor.shared_els && sensor.shared_isShocked[shared_idx])

    # compute element size h for both elements
    h = zero(Tmsh)
    @simd for i=1:numNodesPerElement
      h += sbp.w[i]/jac[i]
    end

    # compute estimates viscosity
    lambda_max = sensor.lambda_max
    eps_L = sensor.alpha*lambda_max*(h^(1/Tdim))/sbp.degree

#    println("element ", elnum, " HApprox value = ", real(eps_L), ", with h = ")
    return Tres(eps_L)
  else
    return zero(Tres)
  end

end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorHApprox,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                         ) where {Tsol, Tmsh}

  error("ShockSensorHApprox does not know which elements are shocked or not, don't ask it that.")
end

#------------------------------------------------------------------------------
# ShockSensorOddBO

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorOddBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  if elnum % 2 == 1
    getShockSensor(params, sbp, sensor.sensor, q, elnum, coords, dxidx, jac,
                   Se_mat, ee_mat)
  else
    fill!(Se_mat, 0)
    fill!(ee_mat, 0)
  end

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorOddBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return (elnum % 2 ) == 1
end


text:
# Copyright (c) 2021, R. Dong, C. Goodbarke, H. Harrington, G. Pogudin
# Copyright (c) 2020, A. Ovchinnikov, A. Pillay, G. Pogudin, T. Scanlon

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

function linear_compartment_model(graph, sinks)
    """
    Input:
        - graph - graph of the model network represented via adjacency lists
        - sinks - the indices of nodes having a sink
    Output: the corresponding ODE object where each parameter a_ij is replaced
            with a_ij + b_ij * x_0, where x_0 is a constant input encoded as a constant output
    """
    n = length(graph)
    x_vars_names = ["x$i" for i in 0:n]
    edges_vars_names = Array{String, 1}()
    for i in 1:n
        for j in graph[i]
            push!(edges_vars_names, "a_$(i)_$(j)")
            push!(edges_vars_names, "b_$(i)_$(j)")
        end
    end
    for s in sinks
        push!(edges_vars_names, "a_$(s)_0")
        push!(edges_vars_names, "b_$(s)_0")
    end
    R, vars = Nemo.PolynomialRing(Nemo.QQ, vcat(x_vars_names, edges_vars_names, ["y1", "y2"]))
    x_vars = vars[2:(n + 1)]
    x0 = vars[1]
    equations = Dict{fmpq_mpoly, Union{fmpq_mpoly, Generic.Frac{fmpq_mpoly}}}(x => R(0) for x in x_vars)
    equations[x0] = R(0)
    for i in 1:n
        for j in graph[i]
            rate = str_to_var("a_$(i)_$(j)", R) + str_to_var("b_$(i)_$(j)", R) * x0

            if i != j
                equations[x_vars[j]] += x_vars[i] * rate
                equations[x_vars[i]] -= x_vars[i] * rate
            else
                equations[x_vars[i]] -= x_vars[i] * rate
            end
        end
        if i in sinks
            rate = str_to_var("a_$(i)_0", R) + str_to_var("b_$(i)_0", R) * x0
            equations[x_vars[i]] += -x_vars[i] * rate
        end
    end
    return ODE{fmpq_mpoly}(equations, Dict(vars[end] => x_vars[1], vars[end - 1] => x0), Array{fmpq_mpoly, 1}())
end

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

function bicycle(n)
    """
    Generates a bidirected cycle of length n
    """
    graph = []
    for i in 1:n
        prev = (i == 1) ? n : (i - 1)
        next = (i == n) ? 1 : i + 1
        push!(graph, [prev, next])
    end
    return graph
end

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

function cycle(n)
    """
    Single directed cycle
    """
    graph = [[(i == n) ? 1 : (i + 1)] for i in 1:n]
    return graph
end

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

function catenary(n)
    """
    Bidirected chain from 1 to n
    """
    graph = [[] for i in 1:n]
    for i in 1:n
        if i != 1
            push!(graph[i], i - 1)
        end
        if i != n
            push!(graph[i], i + 1)
        end
    end
    return graph
end

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

function mammilary(n)
    """
    Bidirected 'star' with center at 1 and rays to 2, ..., n
    """
    graph = []
    push!(graph, [i for i in 2:n])
    for i in 2:n
        push!(graph, [1])
    end
    return graph
end

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

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

@testset "Assessing local identifiability (multiexperiment)" begin

    # checking bounds

    test_cases = [
        Dict(
            :name => "Cyclic",
            :graph => cycle,
            :bound => [3]
        ),
        Dict(
            :name => "Catenary",
            :graph => catenary,
            :bound => [4, 5]
        ),
        Dict(
            :name => "Mammilary",
            :graph => mammilary,
            :bound => [4, 5]
        )
    ]

    n_min = 3
    n_max = 8
    for case in test_cases
        for n in n_min:n_max
            model = linear_compartment_model(case[:graph](n), [1])
            println(case[:name] * ", n = $n")
            @time result = assess_local_identifiability(model, 0.97, :ME)
            correct = undef
            if n - n_min + 1 > length(case[:bound])
                correct = case[:bound][end]
            else
                correct = case[:bound][n - n_min + 1]
            end
            @test correct == result[2]
        end
    end

    # checking bounds and results

    test_cases = []

    ode = @ODEmodel(
        x0'(t) = -(a01 + a21) * x0(t) + a12 * x1(t),
        x1'(t) = a21 * x0(t) - a12 * x1(t),
        y(t) = x0(t)
    )
    funcs_to_test = [a01, a21, a12, a01 * a12, a01 + a12 + a21, (a01 + a12 + a21) // (a01 * a12)]
    correct = [false, false, false, true, true, true]
    push!(test_cases, Dict(
        :ode => ode,
        :funcs => funcs_to_test,
        :correct => (correct, 1)
    ))

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

    # Example 7.7 from https://arxiv.org/pdf/2011.10868.pdf
    ode = @ODEmodel(
        x0'(t) = 0,
        x1'(t) = x0(t) * x1(t) + mu1 * x0(t) + mu2,
        y(t) = x1(t)
    )
    funcs_to_test = [mu1, mu2]
    correct = [true, true]
    push!(test_cases, Dict(
        :ode => ode,
        :funcs => funcs_to_test,
        :correct => (correct, 2)
    ))

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

    # Example 7.8 from https://arxiv.org/pdf/2011.10868.pdf
    ode = @ODEmodel(
        S'(t) = -b * S(t) * I(t) / N,
        E'(t) = b * S(t) * I(t) / N - nu * E(t),
        I'(t) = nu * E(t) - a * I(t),
        c'(t) = 0,
        y1(t) = c(t) * I(t) + d * E(t),
        y2(t) = c(t),
        y3(t) = N
    )
    funcs_to_test = [b, nu, d, a]
    correct = [true, true, true, true]
        push!(test_cases, Dict(
        :ode => ode,
        :funcs => funcs_to_test,
        :correct => (correct, 1)
    ))

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

    for case in test_cases
        result = assess_local_identifiability(case[:ode], case[:funcs], 0.932, :ME)
        @test result == case[:correct]
    end

end

ext: jl


text:
# Testing censored distributions

module TestCensored

using Distributions, Test
using Distributions: Censored

function _as_mixture(d::Censored)
    d0 = d.uncensored
    dtrunc = if d0 isa DiscreteUniform || d0 isa Poisson
        truncated(
            d0,
            d.lower === nothing ? -Inf : floor(d.lower) + 1,
            d.upper === nothing ? Inf : ceil(d.upper) - 1,
        )
    elseif d0 isa ContinuousDistribution
        truncated(
            d0,
            d.lower === nothing ? -Inf : nextfloat(float(d.lower)),
            d.upper === nothing ? Inf : prevfloat(float(d.upper)),
        )
    else
        error("truncation to open interval not implemented for $d0")
    end
    prob_lower = d.lower === nothing ? 0 : cdf(d0, d.lower)
    prob_upper = if d.upper === nothing
        0
    elseif d0 isa ContinuousDistribution
        ccdf(d0, d.upper)
    else
        ccdf(d0, d.upper) + pdf(d0, d.upper)
    end
    prob_interval = 1 - (prob_lower + prob_upper)
    components = Distribution[dtrunc]
    probs = [prob_interval]
    if prob_lower > 0
        # workaround for MixtureModel currently not supporting mixtures of discrete and
        # continuous components
        push!(components, d0 isa DiscreteDistribution ? Dirac(d.lower) : Normal(d.lower, 0))
        push!(probs, prob_lower)
    end
    if prob_upper > 0
        push!(components, d0 isa DiscreteDistribution ? Dirac(d.upper) : Normal(d.upper, 0))
        push!(probs, prob_upper)
    end
    return MixtureModel(map(identity, components), probs)
end

@testset "censored" begin
    d0 = Normal(0, 1)
    @test_throws ArgumentError censored(d0, 1, -1)

    # bound argument constructors
    d = censored(d0, -1, 1.0)
    @test d isa Censored
    @test d.lower === -1.0
    @test d.upper === 1.0

    d = censored(d0, nothing, -1)
    @test d isa Censored
    @test d.lower === nothing
    @test d.upper == -1

    d = censored(d0, 1, nothing)
    @test d isa Censored
    @test d.upper === nothing
    @test d.lower == 1

    d = censored(d0, nothing, nothing)
    @test d === d0

    # bound keyword constructors
    d = censored(d0; lower=-2, upper=1.5)
    @test d isa Censored
    @test d.lower === -2.0
    @test d.upper === 1.5

    d = censored(d0; upper=true)
    @test d isa Censored
    @test d.lower === nothing
    @test d.upper === true

    d = censored(d0; lower=-3)
    @test d isa Censored
    @test d.upper === nothing
    @test d.lower === -3

    d = censored(d0)
    @test d === d0
end

@testset "Censored" begin
    @testset "basic" begin
        # check_args
        @test_throws ArgumentError Censored(Normal(0, 1), 2, 1)
        @test_throws ArgumentError Censored(Normal(0, 1), 2, 1; check_args=true)
        Censored(Normal(0, 1), 2, 1; check_args=false)
        Censored(Normal(0, 1), nothing, 1; check_args=true)
        Censored(Normal(0, 1), 2, nothing; check_args=true)

        d = Censored(Normal(0.0, 1.0), -1, 2)
        @test d isa Censored
        @test eltype(d) === Float64
        @test params(d) === (params(Normal(0.0, 1.0))..., -1, 2)
        @test partype(d) === Float64
        @test @inferred extrema(d) == (-1, 2)
        @test @inferred islowerbounded(d)
        @test @inferred isupperbounded(d)
        @test @inferred insupport(d, 0.1)
        @test insupport(d, -1)
        @test insupport(d, 2)
        @test !insupport(d, -1.1)
        @test !insupport(d, 2.1)
        @test sprint(show, "text/plain", d) == "Censored($(Normal(0.0, 1.0)); lower=-1, upper=2)"

        d = Censored(Cauchy(0, 1), nothing, 2)
        @test d isa Censored
        @test eltype(d) === Base.promote_type(eltype(Cauchy(0, 1)), Int)
        @test params(d) === (params(Cauchy(0, 1))..., nothing, 2)
        @test partype(d) === Float64
        @test extrema(d) == (-Inf, 2.0)
        @test @inferred !islowerbounded(d)
        @test @inferred isupperbounded(d)
        @test @inferred insupport(d, 0.1)
        @test insupport(d, -3)
        @test insupport(d, 2)
        @test !insupport(d, 2.1)
        @test sprint(show, "text/plain", d) == "Censored($(Cauchy(0.0, 1.0)); upper=2)"

        d = Censored(Gamma(1, 2), 2, nothing)
        @test d isa Censored
        @test eltype(d) === Base.promote_type(eltype(Gamma(1, 2)), Int)
        @test params(d) === (params(Gamma(1, 2))..., 2, nothing)
        @test partype(d) === Float64
        @test extrema(d) == (2.0, Inf)
        @test @inferred islowerbounded(d)
        @test @inferred !isupperbounded(d)
        @test @inferred insupport(d, 2.1)
        @test insupport(d, 2.0)
        @test !insupport(d, 1.9)
        @test sprint(show, "text/plain", d) == "Censored($(Gamma(1, 2)); lower=2)"

        d = Censored(Binomial(10, 0.2), -1.5, 9.5)
        @test extrema(d) === (0.0, 9.5)
        @test @inferred islowerbounded(d)
        @test @inferred isupperbounded(d)
        @test @inferred !insupport(d, -1.5)
        @test insupport(d, 0)
        @test insupport(d, 9.5)
        @test !insupport(d, 10)

        @test censored(Censored(Normal(), 1, nothing), nothing, 2) == Censored(Normal(), 1, 2)
        @test censored(Censored(Normal(), nothing, 1), -1, nothing) == Censored(Normal(), -1, 1)
        @test censored(Censored(Normal(), 1, 2), 1.5, 2.5) == Censored(Normal(), 1.5, 2.0)
        @test censored(Censored(Normal(), 1, 3), 1.5, 2.5) == Censored(Normal(), 1.5, 2.5)
        @test censored(Censored(Normal(), 1, 2), 0.5, 2.5) == Censored(Normal(), 1.0, 2.0)
        @test censored(Censored(Normal(), 1, 2), 0.5, 1.5) == Censored(Normal(), 1.0, 1.5)

        @test censored(Censored(Normal(), nothing, 1), nothing, 1) == Censored(Normal(), nothing, 1)
        @test censored(Censored(Normal(), nothing, 1), nothing, 2) == Censored(Normal(), nothing, 1)
        @test censored(Censored(Normal(), nothing, 1), nothing, 1.5) == Censored(Normal(), nothing, 1)
        @test censored(Censored(Normal(), nothing, 1.5), nothing, 1) == Censored(Normal(), nothing, 1)

        @test censored(Censored(Normal(), 1, nothing), 1, nothing) == Censored(Normal(), 1, nothing)
        @test censored(Censored(Normal(), 1, nothing), 2, nothing) == Censored(Normal(), 2, nothing)
        @test censored(Censored(Normal(), 1, nothing), 1.5, nothing) == Censored(Normal(), 1.5, nothing)
        @test censored(Censored(Normal(), 1.5, nothing), 1, nothing) == Censored(Normal(), 1.5, nothing)
    end

    @testset "Uniform" begin
        d0 = Uniform(0, 10)
        bounds = [
            (nothing, 8),
            (-Inf, 8),
            (nothing, Inf),
            (2, nothing),
            (2, Inf),
            (-Inf, nothing),
            (2, 8),
            (3.5, nothing),
            (3.5, Inf),
            (-Inf, Inf),
        ]
        @testset "lower = $(lower === nothing ? "nothing" : lower), upper = $(upper === nothing ? "nothing" : upper)" for (lower, upper) in bounds
            d = censored(d0, lower, upper)
            dmix = _as_mixture(d)
            l, u = extrema(d)
            if lower === nothing || !isfinite(lower)
                @test l == minimum(d0)
            else
                @test l == lower
            end
            if upper === nothing || !isfinite(upper)
                @test u == maximum(d0)
            else
                @test u == upper
            end
            @testset for f in [cdf, logcdf, ccdf, logccdf]
                @test @inferred(f(d, l)) ≈ f(dmix, l) atol=1e-8
                @test @inferred(f(d, l - 0.1)) ≈ f(dmix, l - 0.1) atol=1e-8
                @test @inferred(f(d, u)) ≈ f(dmix, u) atol=1e-8
                @test @inferred(f(d, u + 0.1)) ≈ f(dmix, u + 0.1) atol=1e-8
                @test @inferred(f(d, 5)) ≈ f(dmix, 5)
            end
            @testset for f in [mean, var]
                @test @inferred(f(d)) ≈ f(dmix)
            end
            @test @inferred(median(d)) ≈ clamp(median(d0), l, u)
            @inferred quantile(d, 0.5)
            @test quantile.(d, 0:0.01:1) ≈ clamp.(quantile.(d0, 0:0.01:1), l, u)
            # special-case pdf/logpdf/loglikelihood since when replacing Dirac(μ) with
            # Normal(μ, 0), they are infinite
            if lower === nothing || !isfinite(lower)
                @test @inferred(pdf(d, l)) ≈ pdf(d0, l)
                @test @inferred(logpdf(d, l)) ≈ logpdf(d0, l)
            else
                @test @inferred(pdf(d, l)) ≈ cdf(d0, l)
                @test @inferred(logpdf(d, l)) ≈ logcdf(d0, l)
            end
            if upper === nothing || !isfinite(upper)
                @test @inferred(pdf(d, u)) ≈ pdf(d0, u)
                @test @inferred(logpdf(d, u)) ≈ logpdf(d0, u)
            else
                @test @inferred(pdf(d, u)) ≈ ccdf(d0, u)
                @test @inferred(logpdf(d, u)) ≈ logccdf(d0, u)
            end
            # rand
            x = rand(d, 10_000)
            @test all(x -> insupport(d, x), x)
            # loglikelihood
            @test @inferred(loglikelihood(d, x)) ≈ sum(x -> logpdf(d, x), x)
            @test loglikelihood(d, [x; -1]) == -Inf
            # entropy
            @test @inferred(entropy(d)) ≈ mean(x -> -logpdf(d, x), x) atol = 1e-1
        end
    end

    @testset "Normal" begin
        d0 = Normal()
        bounds = [(nothing, 0.2), (-0.1, nothing), (-0.1, 0.2)]
        @testset "lower = $(lower === nothing ? "nothing" : lower), upper = $(upper === nothing ? "nothing" : upper)" for (lower, upper) in bounds
            d = censored(d0, lower, upper)
            dmix = _as_mixture(d)
            l, u = extrema(d)
            @testset for f in [cdf, logcdf, ccdf, logccdf]
                @test f(d, l) ≈ f(dmix, l) atol=1e-8
                @test f(d, l - 0.1) ≈ f(dmix, l - 0.1) atol=1e-8
                @test f(d, u) ≈ f(dmix, u) atol=1e-8
                @test f(d, u + 0.1) ≈ f(dmix, u + 0.1) atol=1e-8
                @test f(d, 5) ≈ f(dmix, 5)
            end
            @testset for f in [mean, var]
                @test f(d) ≈ f(dmix)
            end
            @test median(d) ≈ clamp(median(d0), l, u)
            @test quantile.(d, 0:0.01:1) ≈ clamp.(quantile.(d0, 0:0.01:1), l, u)
            # special-case pdf/logpdf/loglikelihood since when replacing Dirac(μ) with
            # Normal(μ, 0), they are infinite
            if lower === nothing
                @test pdf(d, l) ≈ pdf(d0, l)
                @test logpdf(d, l) ≈ logpdf(d0, l)
            else
                @test pdf(d, l) ≈ cdf(d0, l)
                @test logpdf(d, l) ≈ logcdf(d0, l)
            end
            if upper === nothing
                @test pdf(d, u) ≈ pdf(d0, u)
                @test logpdf(d, u) ≈ logpdf(d0, u)
            else
                @test pdf(d, u) ≈ ccdf(d0, u)
                @test logpdf(d, u) ≈ logccdf(d0, u)
            end
            # rand
            x = rand(d, 10_000)
            @test all(x -> insupport(d, x), x)
            # loglikelihood
            @test loglikelihood(d, x) ≈ sum(x -> logpdf(d, x), x)
            # entropy
            @test entropy(d) ≈ mean(x -> -logpdf(d, x), x) atol = 1e-1
        end
    end

    @testset "DiscreteUniform" begin
        d0 = DiscreteUniform(0, 10)
        bounds = [
            (nothing, 8),
            (-Inf, 8),
            (nothing, Inf),
            (2, nothing),
            (2, Inf),
            (-Inf, nothing),
            (2, 8),
            (3.5, nothing),
            (3.5, Inf),
            (-Inf, Inf),
        ]
        @testset "lower = $(lower === nothing ? "nothing" : lower), upper = $(upper === nothing ? "nothing" : upper)" for (lower, upper) in bounds
            d = censored(d0, lower, upper)
            dmix = _as_mixture(d)
            @test extrema(d) == extrema(dmix)
            l, u = extrema(d)
            @testset for f in [pdf, logpdf, cdf, logcdf, ccdf, logccdf]
                @test @inferred(f(d, l)) ≈ f(dmix, l) atol=1e-8
                @test @inferred(f(d, l - 0.1)) ≈ f(dmix, l - 0.1) atol=1e-8
                @test @inferred(f(d, u)) ≈ f(dmix, u) atol=1e-8
                @test @inferred(f(d, u + 0.1)) ≈ f(dmix, u + 0.1) atol=1e-8
                @test @inferred(f(d, 5)) ≈ f(dmix, 5)
            end
            @testset for f in [mean, var]
                @test @inferred(f(d)) ≈ f(dmix)
            end
            @test @inferred(median(d)) ≈ clamp(median(d0), l, u)
            @inferred quantile(d, 0.5)
            @test quantile.(d, 0:0.01:1) ≈ clamp.(quantile.(d0, 0:0.01:1), l, u)
            # rand
            x = rand(d, 10_000)
            @test all(x -> insupport(d, x), x)
            # loglikelihood
            @test @inferred(loglikelihood(d, x)) ≈ loglikelihood(dmix, x)
            # mean, std
            μ = @inferred mean(d)
            xall = unique(x)
            @test μ ≈ sum(x -> pdf(d, x) * x, xall)
            @test mean(x) ≈ μ atol = 1e-1
            v = @inferred var(d)
            @test v ≈ sum(x -> pdf(d, x) * abs2(x - μ), xall)
            @test std(x) ≈ sqrt(v) atol = 1e-1
            # entropy
            @test @inferred(entropy(d)) ≈ sum(x -> pdf(d, x) * -logpdf(d, x), xall)
        end
    end

    @testset "Poisson" begin
        d0 = Poisson(20)
        bounds = [(nothing, 12), (2, nothing), (2, 12), (8, nothing)]
        @testset "lower = $(lower === nothing ? "nothing" : lower), upper = $(upper === nothing ? "nothing" : upper)" for (lower, upper) in bounds
            d = censored(d0, lower, upper)
            dmix = _as_mixture(d)
            @test extrema(d) == extrema(dmix)
            l, u = extrema(d)
            @testset for f in [pdf, logpdf, cdf, logcdf, ccdf, logccdf]
                @test f(d, l) ≈ f(dmix, l) atol=1e-8
                @test f(d, l - 0.1) ≈ f(dmix, l - 0.1) atol=1e-8
                @test f(d, u) ≈ f(dmix, u) atol=1e-8
                @test f(d, u + 0.1) ≈ f(dmix, u + 0.1) atol=1e-8
                @test f(d, 5) ≈ f(dmix, 5)
            end
            @test median(d) ≈ clamp(median(d0), l, u)
            @test quantile.(d, 0:0.01:0.99) ≈ clamp.(quantile.(d0, 0:0.01:0.99), l, u)
            x = rand(d, 100)
            @test loglikelihood(d, x) ≈ loglikelihood(dmix, x)
            # rand
            x = rand(d, 10_000)
            @test all(x -> insupport(d, x), x)
            # mean, std
            @test mean(x) ≈ mean(x) atol = 1e-1
            @test std(x) ≈ std(x) atol = 1e-1
        end
    end

    @testset "mixed types are still type-inferrible" begin
        bounds = [(nothing, 8), (2, nothing), (2, 8)]
        @testset "lower = $(lower === nothing ? "nothing" : lower), upper = $(upper === nothing ? "nothing" : upper), uncensored partype=$T0, partype=$T" for (lower, upper) in bounds,
                T in (Int, Float32, Float64), T0 in (Int, Float32, Float64)
            d0 = Uniform(T0(0), T0(10))
            d = censored(d0, lower === nothing ? nothing : T(lower), upper === nothing ? nothing : T(upper))
            l, u = extrema(d)
            @testset for f in [pdf, logpdf, cdf, logcdf, ccdf, logccdf]
                @inferred f(d, 3)
                @inferred f(d, 4f0)
                @inferred f(d, 5.0)
            end
            @testset for f in [median, mean, var, entropy]
                @inferred f(d)
            end
            @inferred quantile(d, 0.3f0)
            @inferred quantile(d, 0.5)
            x = randn(Float32, 100)
            @inferred loglikelihood(d, x)
            x = randn(100)
            @inferred loglikelihood(d, x)
        end
    end
end

end # module


text:
###Projection methods aim at post-processing the
### saved scenario projections rather than
### directly operating on the saved model fits.

"""
    function create_NB_dist(μ,α)

Create a `NegativeBinomial` distribution with the correct `r,p` parameterisation from a `μ,α` parameterisation.
"""
function create_NB_dist(μ, α)
    μ += 0.001*(μ == 0.0)

    σ² = μ + α * μ^2
    p_negbin = 1 - (α * μ^2 / σ²)
    r_negbin = 1 / α
    d = NegativeBinomial(r_negbin, p_negbin)
    return d
end


text:
@testset "221.maximal-square.jl" begin
    @test 4 == maximal_square(hcat([['1', '0', '1', '0', '0'], ['1', '0', '1', '1', '1'],
                               ['1', '1', '1', '1', '1'], ['1', '0', '0', '1', '0']]...))
    @test maximal_square(hcat([['0','1'],['1','0']]...)) == 1
    @test maximal_square(hcat([['0']]...)) == 0
end


text:
using DataStructures, Logging

include("structs.jl")
include("gpu_batch.jl")
"""
#Ideally, syntax will eventually look something like this (for Collatz)
@init 1:1024
@branch x -> x == 1 :end
@out_function += 1 :add
@branch x -> x % 2 == 0 :even :odd
@function :even x -> x / 2
@function :odd x -> x * 3 + 1
@run
"""

#--- Program Building
"""
Return a `Program` object to run.
`values` - initial starting problem
`default_value` - effectively a null value for frames which aren't full
`frame_type` - either `Frame` for CPU running or `CuFrame` for GPU
`frame_size` - how many elements in frame
`output_type` - what type output array holds (it'll be same size as `values`)
"""
function initialize(values :: Array{T}, default_value :: T, frame_type, frame_size, output_type :: Type) where T
    @debug "Initializing..."
    F = frame_type{T, frame_size}
    pool = FramePool{T, F}(default_value)
    frame_dict = Dict(:root => VirtualStack{T, F}(Int(round(frame_size * 0.5))))
    output = Array{Int32}(undef, length(values)) .* 0
    program = Program{T, output_type, F}(frame_dict,
                      Dict{Symbol, Function}(),
                      Dict{Symbol, Function}(),
                      Dict{Symbol, Union{Symbol, Tuple{Symbol, Symbol}}}(),
                      pool, output)
    # TODO: implement constructor so scalar access can always be disallowed
    @debug "initializing root"
    for floor in 1:frame_size:length(values)

        ceil = min(frame_size + floor - 1, length(values))
        @debug "Fetching frame"
        frame = fetch!(pool)

        if F == Frame{T, frame_size}
            @debug "cpu frame"
            for (frame_index, value_index) in enumerate(floor:ceil)
                frame[frame_index] = Cell{T}(values[value_index], 0, value_index)
            end
        elseif F == CuFrame{T, frame_size}
            @debug "gpu frame"
            temp = collect(Cell{T}(values[i], 0, i) for i in floor:ceil)
            @debug "temp made"
            @time frame.cells[1:length(floor:ceil)] .= CuArray(temp)
        else
            return error("Unknown type")
        end
        frame.count = length(floor:ceil)
        push!(program, :root, frame)
    end

    @debug "ready!"
    return program
end

"""
Modify program by adding an edge in the process graph between `src` and `dest`
`dest` can be a single symbol or a tuple for branching
`value_function` can be a function or tuple of function and array of aux vars
"""
function add!(program :: Program{T, U, F}, src, dest=:root; value_function=nothing, dest_function=nothing) where {T, U, F}
    if ! isnothing(value_function)
        program.value_functions[src] = value_function
    end
    if ! isnothing(dest_function)
        program.dest_functions[src] = dest_function
    end

    if dest != nothing
        program.flow[src] = dest

        if !isa(dest, Tuple)
            dest = (dest,)
        end
        for d in dest
            if ! (d in keys(program.stacks))
                sz = size(program.pool)
                program.stacks[d] = VirtualStack{T, F}(floor(sz / 2))
            end
        end
    end
end

#---- Program Execution
"""
Get the name of the stack to run a frame from. If the program is done,
return :PROGRAM_TERMINATED
"""
function get_stack_name(program :: Program) :: Symbol
    for i in keys(program.stacks)
        if ready(program.stacks[i])
            return i
        end
    end
    for i in keys(program.stacks)
        if ! empty(program.stacks[i])
            return i
        end
    end
    return :PROGRAM_TERMINATED
end

function value_eval(program, dest, frame :: CuFrame, f_d :: Tuple{Function, Array})
    # break out function and domain
    f, domain = f_d
    for d in domain
        dest_frame = fetch!(program.pool)
        dest_frame.cells .= map(x -> f(x, d), frame.cells)
        dest_frame.count = sum(map(is_valid, frame.cells))
        push!(program, dest, dest_frame)
    end
    reallocate!(program.pool, frame)
end

function value_eval(program, dest, frame :: Frame, f_d :: Tuple{Function, Array})
    # break out function and domain
    f, domain = f_d
    for d in domain
        dest_frame = fetch!(program.pool)
        for i in 1:length(frame)
            push!(dest_frame, f(frame[i], d))
        end
        push!(program, dest, dest_frame)
    end
    reallocate!(program.pool, frame)
end

function value_eval(program, dest, frame :: AFrame, f :: Function)
    # This is a straight function evaluation, so the frame is moved from src
    # to dest without involving the pool
    # It's agnostic to the type of frame
    map!(f, frame.cells, frame.cells)
    push!(program, dest, frame)
end

function dest_eval(program, dest, frame :: AFrame, f)
    # This will copy off device if a CuArray, or do nothing if already Array
    cells = Array(frame.cells)
    for i in 1:length(frame)
        program.dest[cells[i].output_index] += cells[i] |> f
    end
end

function execute_batch(program, src, dest :: Symbol)
    frame = pop!(program, src)
    @debug "Frame $(frame.cells)"
    if src in keys(program.dest_functions)
        dest_eval(program, dest, frame, program.dest_functions[src])
    end
    if src in keys(program.value_functions)
        value_eval(program, dest, frame, program.value_functions[src])
    else
        push!(program, dest, frame)
    end
end

"""
For a cpu frame, compute `decider` for each element and move to `t_frame` if
true and to `f_frame` if false
"""
function split!(frame :: Frame, t_frame :: Frame, f_frame :: Frame, decider)
    for i in 1:length(frame)
        if frame[i] |> decider
            push!(t_frame, frame[i])
        else
            push!(f_frame, frame[i])
        end
    end
end

"""
For a GPU frame, compute `decider` for each element and move to `t_frame` if
true and to `f_frame` if false.

Does two boolean sorts: the first for whether a cell is a valid value, the
second `decider`. After those, the frame is all invalid values, followed by all
false, then all true. Once that is done, the false and trues are distributed to
destinations `t_frame` and `f_frame`.

Assumes that every thread of a SM can hold a Cell{T} and an Int32 worth of shmem
"""
function split!(frame :: CuFrame{T, L}, t_frame :: CuFrame, f_frame :: CuFrame, decider) where {T, L}
    N_t = 1024
    N_b = Int(L / N_t)
    flags = map(is_valid, frame.cells)
    indices = CUDA.zeros(Int32, L)
    batch_sums = CUDA.zeros(Int32, N_b)
    @cuda blocks=N_b threads=N_t index_sort_kernel(indices, flags, batch_sums)
    @cuda blocks=N_b threads=N_t index_move_kernel(frame.cells, indices)

    batch_valid_counts = Array(batch_sums)
    flags = map(decider, frame.cells)

    @cuda blocks=N_b threads=N_t index_sort_kernel(indices, flags, batch_sums)
    @cuda blocks=N_b threads=N_t index_move_kernel(frame.cells, indices)

    batch_true_counts = Array(batch_sums)
    t_floor = 0
    f_floor = 0
    for b in 1:N_b
        floor = (b - 1) * N_t
        valid_count = batch_valid_counts[b]
        invalid_count = N_t - valid_count
        true_count = batch_true_counts[b]
        false_count = valid_count - true_count
        if true_count > 0
            t_frame.cells[t_floor+1:t_floor+true_count] .= frame.cells[floor + (N_t - true_count + 1):floor + (N_t)]
            t_frame.count += true_count
        end
        if false_count > 0
            f_frame.cells[f_floor+1:f_floor+false_count] .= frame.cells[floor+invalid_count + 1: floor+invalid_count + false_count]
            f_frame.count += false_count
        end
        t_floor += true_count
        f_floor += false_count
    end
end

function execute_batch(program :: Program{T, U}, src, dest :: Tuple{Symbol, Symbol}) where {T, U}
    t_frame = fetch!(program.pool)
    f_frame = fetch!(program.pool)
    frame = pop!(program, src)
    #@debug "Frame count: $(length(frame))"
    split!(frame, t_frame, f_frame, program.value_functions[src])
    reallocate!(program.pool, frame)
    push!(program, dest[1], t_frame)
    push!(program, dest[2], f_frame)
end

function Base.run(program :: Program)
    stack_name = get_stack_name(program)
    while stack_name != :PROGRAM_TERMINATED
    #    @debug "Running $stack_name"
        execute_batch(program, stack_name, program.flow[stack_name])
        stack_name = get_stack_name(program)
    end
end


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

"""
    change_server_life_cycle_state(life_cycle, source_server_id)
    change_server_life_cycle_state(life_cycle, source_server_id, params::Dict{String,<:Any})

Allows the user to set the SourceServer.LifeCycle.state property for specific Source Server
IDs to one of the following: READY_FOR_TEST or READY_FOR_CUTOVER. This command only works
if the Source Server is already launchable (dataReplicationInfo.lagDuration is not null.)

# Arguments
- `life_cycle`: The request to change the source server migration lifecycle state.
- `source_server_id`: The request to change the source server migration lifecycle state by
  source server ID.

"""
change_server_life_cycle_state(lifeCycle, sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/ChangeServerLifeCycleState", Dict{String, Any}("lifeCycle"=>lifeCycle, "sourceServerID"=>sourceServerID); aws_config=aws_config)
change_server_life_cycle_state(lifeCycle, sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/ChangeServerLifeCycleState", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("lifeCycle"=>lifeCycle, "sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

"""
    create_replication_configuration_template(associate_default_security_group, bandwidth_throttling, create_public_ip, data_plane_routing, default_large_staging_disk_type, ebs_encryption, replication_server_instance_type, replication_servers_security_groups_ids, staging_area_subnet_id, staging_area_tags, use_dedicated_replication_server)
    create_replication_configuration_template(associate_default_security_group, bandwidth_throttling, create_public_ip, data_plane_routing, default_large_staging_disk_type, ebs_encryption, replication_server_instance_type, replication_servers_security_groups_ids, staging_area_subnet_id, staging_area_tags, use_dedicated_replication_server, params::Dict{String,<:Any})

Creates a new ReplicationConfigurationTemplate.

# Arguments
- `associate_default_security_group`: Request to associate the default Application
  Migration Service Security group with the Replication Settings template.
- `bandwidth_throttling`: Request to configure bandwidth throttling during Replication
  Settings template creation.
- `create_public_ip`: Request to create Public IP during Replication Settings template
  creation.
- `data_plane_routing`: Request to configure data plane routing during Replication Settings
  template creation.
- `default_large_staging_disk_type`: Request to configure the Staging Disk EBS volume type
  to \"gp2\" during Replication Settings template creation.
- `ebs_encryption`: Request to configure EBS enryption during Replication Settings template
  creation.
- `replication_server_instance_type`: Request to configure the Replication Server instance
  type during Replication Settings template creation.
- `replication_servers_security_groups_ids`: Request to configure the Replication Server
  Secuirity group ID during Replication Settings template creation.
- `staging_area_subnet_id`: Request to configure the Staging Area subnet ID during
  Replication Settings template creation.
- `staging_area_tags`: Request to configure Staiging Area tags during Replication Settings
  template creation.
- `use_dedicated_replication_server`: Request to use Dedicated Replication Servers during
  Replication Settings template creation.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"ebsEncryptionKeyArn"`: Request to configure an EBS enryption key during Replication
  Settings template creation.
- `"tags"`: Request to configure tags during Replication Settings template creation.
"""
create_replication_configuration_template(associateDefaultSecurityGroup, bandwidthThrottling, createPublicIP, dataPlaneRouting, defaultLargeStagingDiskType, ebsEncryption, replicationServerInstanceType, replicationServersSecurityGroupsIDs, stagingAreaSubnetId, stagingAreaTags, useDedicatedReplicationServer; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/CreateReplicationConfigurationTemplate", Dict{String, Any}("associateDefaultSecurityGroup"=>associateDefaultSecurityGroup, "bandwidthThrottling"=>bandwidthThrottling, "createPublicIP"=>createPublicIP, "dataPlaneRouting"=>dataPlaneRouting, "defaultLargeStagingDiskType"=>defaultLargeStagingDiskType, "ebsEncryption"=>ebsEncryption, "replicationServerInstanceType"=>replicationServerInstanceType, "replicationServersSecurityGroupsIDs"=>replicationServersSecurityGroupsIDs, "stagingAreaSubnetId"=>stagingAreaSubnetId, "stagingAreaTags"=>stagingAreaTags, "useDedicatedReplicationServer"=>useDedicatedReplicationServer); aws_config=aws_config)
create_replication_configuration_template(associateDefaultSecurityGroup, bandwidthThrottling, createPublicIP, dataPlaneRouting, defaultLargeStagingDiskType, ebsEncryption, replicationServerInstanceType, replicationServersSecurityGroupsIDs, stagingAreaSubnetId, stagingAreaTags, useDedicatedReplicationServer, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/CreateReplicationConfigurationTemplate", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("associateDefaultSecurityGroup"=>associateDefaultSecurityGroup, "bandwidthThrottling"=>bandwidthThrottling, "createPublicIP"=>createPublicIP, "dataPlaneRouting"=>dataPlaneRouting, "defaultLargeStagingDiskType"=>defaultLargeStagingDiskType, "ebsEncryption"=>ebsEncryption, "replicationServerInstanceType"=>replicationServerInstanceType, "replicationServersSecurityGroupsIDs"=>replicationServersSecurityGroupsIDs, "stagingAreaSubnetId"=>stagingAreaSubnetId, "stagingAreaTags"=>stagingAreaTags, "useDedicatedReplicationServer"=>useDedicatedReplicationServer), params)); aws_config=aws_config)

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

Deletes a single Job by ID.

# Arguments
- `job_id`: Request to delete Job from service by Job ID.

"""
delete_job(jobID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteJob", Dict{String, Any}("jobID"=>jobID); aws_config=aws_config)
delete_job(jobID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteJob", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("jobID"=>jobID), params)); aws_config=aws_config)

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

Deletes a single Replication Configuration Template by ID

# Arguments
- `replication_configuration_template_id`: Request to delete Replication Configuration
  Template from service by Replication Configuration Template ID.

"""
delete_replication_configuration_template(replicationConfigurationTemplateID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteReplicationConfigurationTemplate", Dict{String, Any}("replicationConfigurationTemplateID"=>replicationConfigurationTemplateID); aws_config=aws_config)
delete_replication_configuration_template(replicationConfigurationTemplateID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteReplicationConfigurationTemplate", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("replicationConfigurationTemplateID"=>replicationConfigurationTemplateID), params)); aws_config=aws_config)

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

Deletes a single source server by ID.

# Arguments
- `source_server_id`: Request to delete Source Server from service by Server ID.

"""
delete_source_server(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteSourceServer", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
delete_source_server(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DeleteSourceServer", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Retrieves detailed Job log with paging.

# Arguments
- `job_id`: Request to describe Job log job ID.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: Request to describe Job log item maximum results.
- `"nextToken"`: Request to describe Job log next token.
"""
describe_job_log_items(jobID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeJobLogItems", Dict{String, Any}("jobID"=>jobID); aws_config=aws_config)
describe_job_log_items(jobID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeJobLogItems", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("jobID"=>jobID), params)); aws_config=aws_config)

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

Returns a list of Jobs. Use the JobsID and fromDate and toData filters to limit which jobs
are returned. The response is sorted by creationDataTime - latest date first. Jobs are
normaly created by the StartTest, StartCutover, and TerminateTargetInstances APIs. Jobs are
also created by DiagnosticLaunch and TerminateDiagnosticInstances, which are APIs available
only to *Support* and only used in response to relevant support tickets.

# Arguments
- `filters`: Request to describe Job log filters.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: Request to describe Job log by max results.
- `"nextToken"`: Request to describe Job logby next token.
"""
describe_jobs(filters; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeJobs", Dict{String, Any}("filters"=>filters); aws_config=aws_config)
describe_jobs(filters, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeJobs", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("filters"=>filters), params)); aws_config=aws_config)

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

Lists all ReplicationConfigurationTemplates, filtered by Source Server IDs.

# Arguments
- `replication_configuration_template_ids`: Request to describe Replication Configuration
  template by template IDs.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: Request to describe Replication Configuration template by max results.
- `"nextToken"`: Request to describe Replication Configuration template by next token.
"""
describe_replication_configuration_templates(replicationConfigurationTemplateIDs; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeReplicationConfigurationTemplates", Dict{String, Any}("replicationConfigurationTemplateIDs"=>replicationConfigurationTemplateIDs); aws_config=aws_config)
describe_replication_configuration_templates(replicationConfigurationTemplateIDs, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeReplicationConfigurationTemplates", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("replicationConfigurationTemplateIDs"=>replicationConfigurationTemplateIDs), params)); aws_config=aws_config)

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

Retrieves all SourceServers or multiple SourceServers by ID.

# Arguments
- `filters`: Request to filter Source Servers list.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: Request to filter Source Servers list by maximum results.
- `"nextToken"`: Request to filter Source Servers list by next token.
"""
describe_source_servers(filters; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeSourceServers", Dict{String, Any}("filters"=>filters); aws_config=aws_config)
describe_source_servers(filters, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DescribeSourceServers", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("filters"=>filters), params)); aws_config=aws_config)

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

Disconnects specific Source Servers from Application Migration Service. Data replication is
stopped immediately. All AWS resources created by Application Migration Service for
enabling the replication of these source servers will be terminated / deleted within 90
minutes. Launched Test or Cutover instances will NOT be terminated. If the agent on the
source server has not been prevented from communciating with the Application Migration
Service service, then it will receive a command to uninstall itself (within approximately
10 minutes). The following properties of the SourceServer will be changed immediately:
dataReplicationInfo.dataReplicationState will be set to DISCONNECTED; The totalStorageBytes
property for each of dataReplicationInfo.replicatedDisks will be set to zero;
dataReplicationInfo.lagDuration and dataReplicationInfo.lagDurationwill be nullified.

# Arguments
- `source_server_id`: Request to disconnect Source Server from service by Server ID.

"""
disconnect_from_service(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DisconnectFromService", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
disconnect_from_service(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/DisconnectFromService", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Finalizes the cutover immediately for specific Source Servers. All AWS resources created by
Application Migration Service for enabling the replication of these source servers will be
terminated / deleted within 90 minutes. Launched Test or Cutover instances will NOT be
terminated. The AWS Replication Agent will receive a command to uninstall itself (within 10
minutes). The following properties of the SourceServer will be changed immediately:
dataReplicationInfo.dataReplicationState will be to DISCONNECTED; The
SourceServer.lifeCycle.state will be changed to CUTOVER; The totalStorageBytes property fo
each of dataReplicationInfo.replicatedDisks will be set to zero;
dataReplicationInfo.lagDuration and dataReplicationInfo.lagDurationwill be nullified.

# Arguments
- `source_server_id`: Request to finalize Cutover by Soure Server ID.

"""
finalize_cutover(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/FinalizeCutover", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
finalize_cutover(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/FinalizeCutover", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Lists all LaunchConfigurations available, filtered by Source Server IDs.

# Arguments
- `source_server_id`: Request to get Launch Configuration information by Source Server ID.

"""
get_launch_configuration(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/GetLaunchConfiguration", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
get_launch_configuration(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/GetLaunchConfiguration", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Lists all ReplicationConfigurations, filtered by Source Server ID.

# Arguments
- `source_server_id`: Request to get Replication Configuaration by Source Server ID.

"""
get_replication_configuration(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/GetReplicationConfiguration", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
get_replication_configuration(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/GetReplicationConfiguration", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Initialize Application Migration Service.

"""
initialize_service(; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/InitializeService"; aws_config=aws_config)
initialize_service(params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/InitializeService", params; aws_config=aws_config)

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

List all tags for your Application Migration Service resources.

# Arguments
- `resource_arn`: List tags for resource request by ARN.

"""
list_tags_for_resource(resourceArn; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("GET", "/tags/$(resourceArn)"; aws_config=aws_config)
list_tags_for_resource(resourceArn, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("GET", "/tags/$(resourceArn)", params; aws_config=aws_config)

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

Archives specific Source Servers by setting the SourceServer.isArchived property to true
for specified SourceServers by ID. This command only works for SourceServers with a
lifecycle.state which equals DISCONNECTED or CUTOVER.

# Arguments
- `source_server_id`: Mark as archived by Source Server ID.

"""
mark_as_archived(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/MarkAsArchived", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
mark_as_archived(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/MarkAsArchived", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Causes the data replication initiation sequence to begin immediately upon next Handshake
for specified SourceServer IDs, regardless of when the previous initiation started. This
command will not work if the SourceServer is not stalled or is in a DISCONNECTED or STOPPED
state.

# Arguments
- `source_server_id`: Retry data replication for Source Server ID.

"""
retry_data_replication(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/RetryDataReplication", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
retry_data_replication(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/RetryDataReplication", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Launches a Cutover Instance for specific Source Servers. This command starts a LAUNCH job
whose initiatedBy property is StartCutover and changes the SourceServer.lifeCycle.state
property to CUTTING_OVER.

# Arguments
- `source_server_ids`: Start Cutover by Source Server IDs.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tags"`: Start Cutover by Tags.
"""
start_cutover(sourceServerIDs; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/StartCutover", Dict{String, Any}("sourceServerIDs"=>sourceServerIDs); aws_config=aws_config)
start_cutover(sourceServerIDs, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/StartCutover", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerIDs"=>sourceServerIDs), params)); aws_config=aws_config)

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

Lauches a Test Instance for specific Source Servers. This command starts a LAUNCH job whose
initiatedBy property is StartTest and changes the SourceServer.lifeCycle.state property to
TESTING.

# Arguments
- `source_server_ids`: Start Test for Source Server IDs.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tags"`: Start Test by Tags.
"""
start_test(sourceServerIDs; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/StartTest", Dict{String, Any}("sourceServerIDs"=>sourceServerIDs); aws_config=aws_config)
start_test(sourceServerIDs, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/StartTest", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerIDs"=>sourceServerIDs), params)); aws_config=aws_config)

"""
    tag_resource(resource_arn, tags)
    tag_resource(resource_arn, tags, params::Dict{String,<:Any})

Adds or overwrites only the specified tags for the specified Application Migration Service
resource or resources. When you specify an existing tag key, the value is overwritten with
the new value. Each resource can have a maximum of 50 tags. Each tag consists of a key and
optional value.

# Arguments
- `resource_arn`: Tag resource by ARN.
- `tags`: Tag resource by Tags.

"""
tag_resource(resourceArn, tags; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/tags/$(resourceArn)", Dict{String, Any}("tags"=>tags); aws_config=aws_config)
tag_resource(resourceArn, tags, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/tags/$(resourceArn)", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("tags"=>tags), params)); aws_config=aws_config)

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

Starts a job that terminates specific launched EC2 Test and Cutover instances. This command
will not work for any Source Server with a lifecycle.state of TESTING, CUTTING_OVER, or
CUTOVER.

# Arguments
- `source_server_ids`: Terminate Target instance by Source Server IDs.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tags"`: Terminate Target instance by Tags.
"""
terminate_target_instances(sourceServerIDs; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/TerminateTargetInstances", Dict{String, Any}("sourceServerIDs"=>sourceServerIDs); aws_config=aws_config)
terminate_target_instances(sourceServerIDs, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/TerminateTargetInstances", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerIDs"=>sourceServerIDs), params)); aws_config=aws_config)

"""
    untag_resource(resource_arn, tag_keys)
    untag_resource(resource_arn, tag_keys, params::Dict{String,<:Any})

Deletes the specified set of tags from the specified set of Application Migration Service
resources.

# Arguments
- `resource_arn`: Untag resource by ARN.
- `tag_keys`: Untag resource by Keys.

"""
untag_resource(resourceArn, tagKeys; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("DELETE", "/tags/$(resourceArn)", Dict{String, Any}("tagKeys"=>tagKeys); aws_config=aws_config)
untag_resource(resourceArn, tagKeys, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("DELETE", "/tags/$(resourceArn)", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("tagKeys"=>tagKeys), params)); aws_config=aws_config)

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

Updates multiple LaunchConfigurations by Source Server ID.

# Arguments
- `source_server_id`: Update Launch configuration by Source Server ID request.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"copyPrivateIp"`: Update Launch configuration copy Private IP request.
- `"copyTags"`: Update Launch configuration copy Tags request.
- `"launchDisposition"`: Update Launch configuration launch disposition request.
- `"licensing"`: Update Launch configuration licensing request.
- `"name"`: Update Launch configuration name request.
- `"targetInstanceTypeRightSizingMethod"`: Update Launch configuration Target instance
  right sizing request.
"""
update_launch_configuration(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateLaunchConfiguration", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
update_launch_configuration(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateLaunchConfiguration", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Allows you to update multiple ReplicationConfigurations by Source Server ID.

# Arguments
- `source_server_id`: Update replication configuration Source Server ID request.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"associateDefaultSecurityGroup"`: Update replication configuration associate default
  Application Migration Service Security group request.
- `"bandwidthThrottling"`: Update replication configuration bandwidth throttling request.
- `"createPublicIP"`: Update replication configuration create Public IP request.
- `"dataPlaneRouting"`: Update replication configuration data plane routing request.
- `"defaultLargeStagingDiskType"`: Update replication configuration use default large
  Staging Disk type request.
- `"ebsEncryption"`: Update replication configuration EBS encryption request.
- `"ebsEncryptionKeyArn"`: Update replication configuration EBS encryption key ARN request.
- `"name"`: Update replication configuration name request.
- `"replicatedDisks"`: Update replication configuration replicated disks request.
- `"replicationServerInstanceType"`: Update replication configuration Replication Server
  instance type request.
- `"replicationServersSecurityGroupsIDs"`: Update replication configuration Replication
  Server Security Groups IDs request.
- `"stagingAreaSubnetId"`: Update replication configuration Staging Area subnet request.
- `"stagingAreaTags"`: Update replication configuration Staging Area Tags request.
- `"useDedicatedReplicationServer"`: Update replication configuration use dedicated
  Replication Server request.
"""
update_replication_configuration(sourceServerID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateReplicationConfiguration", Dict{String, Any}("sourceServerID"=>sourceServerID); aws_config=aws_config)
update_replication_configuration(sourceServerID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateReplicationConfiguration", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("sourceServerID"=>sourceServerID), params)); aws_config=aws_config)

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

Updates multiple ReplicationConfigurationTemplates by ID.

# Arguments
- `replication_configuration_template_id`: Update replication configuration template
  template ID request.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"arn"`: Update replication configuration template ARN request.
- `"associateDefaultSecurityGroup"`: Update replication configuration template associate
  default Application Migration Service Security group request.
- `"bandwidthThrottling"`: Update replication configuration template bandwidth throttling
  request.
- `"createPublicIP"`: Update replication configuration template create Public IP request.
- `"dataPlaneRouting"`: Update replication configuration template data plane routing
  request.
- `"defaultLargeStagingDiskType"`: Update replication configuration template use default
  large Staging Disk type request.
- `"ebsEncryption"`: Update replication configuration template EBS encryption request.
- `"ebsEncryptionKeyArn"`: Update replication configuration template EBS encryption key ARN
  request.
- `"replicationServerInstanceType"`: Update replication configuration template Replication
  Server instance type request.
- `"replicationServersSecurityGroupsIDs"`: Update replication configuration template
  Replication Server Security groups IDs request.
- `"stagingAreaSubnetId"`: Update replication configuration template Staging Area subnet ID
  request.
- `"stagingAreaTags"`: Update replication configuration template Staging Area Tags request.
- `"useDedicatedReplicationServer"`: Update replication configuration template use
  dedicated Replication Server request.
"""
update_replication_configuration_template(replicationConfigurationTemplateID; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateReplicationConfigurationTemplate", Dict{String, Any}("replicationConfigurationTemplateID"=>replicationConfigurationTemplateID); aws_config=aws_config)
update_replication_configuration_template(replicationConfigurationTemplateID, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = mgn("POST", "/UpdateReplicationConfigurationTemplate", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("replicationConfigurationTemplateID"=>replicationConfigurationTemplateID), params)); aws_config=aws_config)


text:
"""
    (data, lon, lat, time) = read_netcdf2d(file, varname="")

Read variable `varname` (optional) from netcdf `file`.
Returns data and coordinate and time axes.
"""
function  read_netcdf2d(file, varname="")

    nc = NetCDF.open(file, mode=NC_NOWRITE)
    if( haskey(nc.vars,"lon") )
        lat=ncread(file,"lat")
        lon=ncread(file,"lon")
    elseif (haskey(nc.vars,"longitude"))
        lat=ncread(file,"latitude")
        lon=ncread(file,"longitude")
    elseif (haskey(nc.vars,"x"))
        lat=ncread(file,"y")
        lon=ncread(file,"x")
    elseif (haskey(nc.vars,"X"))
        lat=ncread(file,"Y")
        lon=ncread(file,"X")
    else
        println("Input file does not contain lon or longitude dimensional variables")
        quit(1)
    end

    if(varname=="")
        for key in  keys(nc.vars)
            if(key!="lon" && key!="lat" && key!="time" && key!="longitude" &&
               key!="latitude" && key!="time_bnds" && key!="x" && key!="y" &&
               key!="X" && key!="Y" )
                varname = key
            end
        end
    end

    var = ncread(file, varname)

    var = float(var)
    if (haskey(nc.vars[varname].atts,"missing_value"))
        xmiss = nc.vars[varname].atts["missing_value"]
        xmiss = convert(typeof(var[1,1,1]),xmiss)
        var[var .== xmiss] .= NaN
    end

    return var, lon, lat, varname
end


text:
"""
    LTISDE

A lightweight wrapper around a `GP` `f` that tells this package to handle inference in `f`.
Can be constructed via the `to_sde` function.
"""
struct LTISDE{Tf<:GP{<:AbstractGPs.ZeroMean}, Tstorage<:StorageType} <: AbstractGP
    f::Tf
    storage::Tstorage
end

function to_sde(f::GP{<:AbstractGPs.ZeroMean}, storage_type=ArrayStorage(Float64))
    return LTISDE(f, storage_type)
end

storage_type(f::LTISDE) = f.storage

"""
    const FiniteLTISDE = FiniteGP{<:LTISDE}

A `FiniteLTISDE` is just a regular `FiniteGP` that happens to contain an `LTISDE`, as
opposed to any other `AbstractGP`.
"""
const FiniteLTISDE = FiniteGP{<:LTISDE}

# Deal with a bug in AbstractGPs.
function FiniteGP(f::LTISDE, x::AbstractVector{<:Real})
    return FiniteGP(f, x, convert(eltype(storage_type(f)), 1e-12))
end

# Implement the AbstractGP API.

AbstractGPs.mean(ft::FiniteLTISDE) = mean.(marginals(build_lgssm(ft)))

AbstractGPs.cov(ft::FiniteLTISDE) = cov(FiniteGP(ft.f.f, ft.x, ft.Σy))

function AbstractGPs.marginals(ft::FiniteLTISDE)
    return vcat(map(marginals, marginals(build_lgssm(ft)))...)
end

function AbstractGPs.rand(rng::AbstractRNG, ft::FiniteLTISDE)
    return destructure(rand(rng, build_lgssm(ft)))
end

AbstractGPs.rand(ft::FiniteLTISDE) = rand(Random.GLOBAL_RNG, ft)

function AbstractGPs.rand(rng::AbstractRNG, ft::FiniteLTISDE, N::Int)
    return hcat([rand(rng, ft) for _ in 1:N]...)
end

AbstractGPs.rand(ft::FiniteLTISDE, N::Int) = rand(Random.GLOBAL_RNG, ft, N)

AbstractGPs.logpdf(ft::FiniteLTISDE, y::AbstractVector{<:Real}) = _logpdf(ft, y)

function AbstractGPs.logpdf(ft::FiniteLTISDE, y::AbstractVector{<:Union{Missing, Real}})
    return _logpdf(ft, y)
end

function _logpdf(ft::FiniteLTISDE, y::AbstractVector{<:Union{Missing, Real}})
    model = build_lgssm(ft)
    return logpdf(model, restructure(y, model.emissions))
end

restructure(y::AbstractVector{<:Real}, ::StructArray{<:ScalarOutputLGC}) = y

destructure(y::AbstractVector{<:Real}) = y

# Converting GPs into LGSSMs.

function build_lgssm(ft::FiniteLTISDE)
    As, as, Qs, emission_proj, x0 = lgssm_components(ft.f.f.kernel, ft.x, ft.f.storage)
    return LGSSM(
        GaussMarkovModel(Forward(), As, as, Qs, x0),
        build_emissions(emission_proj, build_Σs(ft)),
    )
end

build_Σs(ft::FiniteLTISDE) = build_Σs(ft.x, ft.Σy)

build_Σs(::AbstractVector{<:Real}, Σ::Diagonal{<:Real}) = Σ.diag

function build_emissions(
    (Hs, hs)::Tuple{AbstractVector, AbstractVector}, Σs::AbstractVector,
)
    Hst = map(adjoint, Hs)
    return StructArray{get_type(Hst, hs, Σs)}((Hst, hs, Σs))
end

function get_type(Hs_prime, hs::AbstractVector{<:Real}, Σs)
    THs = eltype(Hs_prime)
    Ths = eltype(hs)
    TΣs = eltype(Σs)
    T = ScalarOutputLGC{THs, Ths, TΣs}
    return T
end

function get_type(Hs_prime, hs::AbstractVector{<:AbstractVector}, Σs)
    THs = eltype(Hs_prime)
    Ths = eltype(hs)
    TΣs = eltype(Σs)
    T = SmallOutputLGC{THs, Ths, TΣs}
    return T
end

@inline function Zygote.wrap_chainrules_output(x::NamedTuple)
    return map(Zygote.wrap_chainrules_output, x)
end

# Generic constructors for base kernels.

function lgssm_components(
    k::SimpleKernel, t::AbstractVector, storage::StorageType{T},
) where {T<:Real}

    # Compute stationary distribution and sde.
    x0 = stationary_distribution(k, storage)
    P = x0.P
    F, q, H = to_sde(k, storage)

    # Use stationary distribution + sde to compute finite-dimensional Gauss-Markov model.
    t = vcat([first(t) - 1], t)
    As = map(Δt -> time_exp(F, T(Δt)), diff(t))
    as = Fill(Zeros{T}(size(first(As), 1)), length(As))
    Qs = map(A -> Symmetric(P) - A * Symmetric(P) * A', As)
    Hs = Fill(H, length(As))
    hs = Fill(zero(T), length(As))
    emission_projections = (Hs, hs)

    return As, as, Qs, emission_projections, x0
end

function lgssm_components(
    k::SimpleKernel, t::Union{StepRangeLen, RegularSpacing}, storage_type::StorageType{T},
) where {T<:Real}

    # Compute stationary distribution and sde.
    x0 = stationary_distribution(k, storage_type)
    P = x0.P
    F, q, H = to_sde(k, storage_type)

    # Use stationary distribution + sde to compute finite-dimensional Gauss-Markov model.
    A = time_exp(F, T(step(t)))
    As = Fill(A, length(t))
    as = Fill(Zeros{T}(size(F, 1)), length(t))
    Q = Symmetric(P) - A * Symmetric(P) * A'
    Qs = Fill(Q, length(t))
    Hs = Fill(H, length(t))
    hs = Fill(zero(T), length(As))
    emission_projections = (Hs, hs)

    return As, as, Qs, emission_projections, x0
end

# Fallback definitions for most base kernels.
function to_sde(k::SimpleKernel, ::ArrayStorage{T}) where {T<:Real}
    F, q, H = to_sde(k, SArrayStorage(T))
    return collect(F), q, collect(H)
end

function stationary_distribution(k::SimpleKernel, ::ArrayStorage{T}) where {T<:Real}
    x = stationary_distribution(k, SArrayStorage(T))
    return Gaussian(collect(x.m), collect(x.P))
end

# Matern-1/2

function to_sde(k::Matern12Kernel, s::SArrayStorage{T}) where {T<:Real}
    F = SMatrix{1, 1, T}(-1)
    q = convert(T, 2)
    H = SVector{1, T}(1)
    return F, q, H
end

function stationary_distribution(k::Matern12Kernel, s::SArrayStorage{T}) where {T<:Real}
    return Gaussian(
        SVector{1, T}(0),
        SMatrix{1, 1, T}(1),
    )
end

Zygote.@adjoint function to_sde(k::Matern12Kernel, storage_type)
    return to_sde(k, storage_type), Δ->(nothing, nothing)
end

Zygote.@adjoint function stationary_distribution(k::Matern12Kernel, storage_type)
    return stationary_distribution(k, storage_type), Δ->(nothing, nothing)
end

# Matern - 3/2

function to_sde(k::Matern32Kernel, ::SArrayStorage{T}) where {T<:Real}
    λ = sqrt(3)
    F = SMatrix{2, 2, T}(0, -3, 1, -2λ)
    q = convert(T, 4 * λ^3)
    H = SVector{2, T}(1, 0)
    return F, q, H
end

function stationary_distribution(k::Matern32Kernel, ::SArrayStorage{T}) where {T<:Real}
    return Gaussian(
        SVector{2, T}(0, 0),
        SMatrix{2, 2, T}(1, 0, 0, 3),
    )
end

Zygote.@adjoint function to_sde(k::Matern32Kernel, storage_type)
    return to_sde(k, storage_type), Δ->(nothing, nothing)
end

Zygote.@adjoint function stationary_distribution(k::Matern32Kernel, storage_type)
    return stationary_distribution(k, storage_type), Δ->(nothing, nothing)
end

# Matern - 5/2

function to_sde(k::Matern52Kernel, ::SArrayStorage{T}) where {T<:Real}
    λ = sqrt(5)
    F = SMatrix{3, 3, T}(0, 0, -λ^3, 1, 0, -3λ^2, 0, 1, -3λ)
    q = convert(T, 8 * λ^5 / 3)
    H = SVector{3, T}(1, 0, 0)
    return F, q, H
end

function stationary_distribution(k::Matern52Kernel, ::SArrayStorage{T}) where {T<:Real}
    κ = 5 / 3
    m = SVector{3, T}(0, 0, 0)
    P = SMatrix{3, 3, T}(1, 0, -κ, 0, κ, 0, -κ, 0, 25)
    return Gaussian(m, P)
end

Zygote.@adjoint function to_sde(k::Matern52Kernel, storage_type)
    return to_sde(k, storage_type), Δ->(nothing, nothing)
end

Zygote.@adjoint function stationary_distribution(k::Matern52Kernel, storage_type)
    return stationary_distribution(k, storage_type), Δ->(nothing, nothing)
end

# Scaled

function lgssm_components(k::ScaledKernel, ts::AbstractVector, storage_type::StorageType)
    As, as, Qs, emission_proj, x0 = lgssm_components(k.kernel, ts, storage_type)
    σ = sqrt(convert(eltype(storage_type), only(k.σ²)))
    return As, as, Qs, _scale_emission_projections(emission_proj, σ), x0
end

function _scale_emission_projections((Hs, hs)::Tuple{AbstractVector, AbstractVector}, σ)
    return (map(H->σ * H, Hs), map(h->σ * h, hs))
end

function _scale_emission_projections((Cs, cs, Hs, hs), σ)
    return (Cs, cs, map(H->σ * H, Hs), map(h->σ * h, hs))
end

# Stretched

function lgssm_components(
    k::TransformedKernel{<:Kernel, <:ScaleTransform},
    ts::AbstractVector,
    storage_type::StorageType,
)
    return lgssm_components(k.kernel, apply_stretch(only(k.transform.s), ts), storage_type)
end

apply_stretch(a, ts::AbstractVector{<:Real}) = a * ts

apply_stretch(a, ts::StepRangeLen) = a * ts

apply_stretch(a, ts::RegularSpacing) = RegularSpacing(a * ts.t0, a * ts.Δt, ts.N)

# Sum

function lgssm_components(k::KernelSum, ts::AbstractVector, storage_type::StorageType)
    As_l, as_l, Qs_l, emission_proj_l, x0_l = lgssm_components(k.kernels[1], ts, storage_type)
    As_r, as_r, Qs_r, emission_proj_r, x0_r = lgssm_components(k.kernels[2], ts, storage_type)

    As = map(blk_diag, As_l, As_r)
    as = map(vcat, as_l, as_r)
    Qs = map(blk_diag, Qs_l, Qs_r)
    emission_projections = _sum_emission_projections(emission_proj_l, emission_proj_r)
    x0 = Gaussian(vcat(x0_l.m, x0_r.m), blk_diag(x0_l.P, x0_r.P))

    return As, as, Qs, emission_projections, x0
end

function _sum_emission_projections(
    (Hs_l, hs_l)::Tuple{AbstractVector, AbstractVector},
    (Hs_r, hs_r)::Tuple{AbstractVector, AbstractVector},
)
    return (map(vcat, Hs_l, Hs_r), hs_l + hs_r)
end

function _sum_emission_projections(
    (Cs_l, cs_l, Hs_l, hs_l)::Tuple{AbstractVector, AbstractVector, AbstractVector, AbstractVector},
    (Cs_r, cs_r, Hs_r, hs_r)::Tuple{AbstractVector, AbstractVector, AbstractVector, AbstractVector},
)
    Cs = map(vcat, Cs_l, Cs_r)
    cs = cs_l + cs_r
    Hs = map(blk_diag, Hs_l, Hs_r)
    hs = map(vcat, hs_l, hs_r)
    return (Cs, cs, Hs, hs)
end

Base.vcat(x::Zeros{T, 1}, y::Zeros{T, 1}) where {T} = Zeros{T}(length(x) + length(y))

function blk_diag(A::AbstractMatrix{T}, B::AbstractMatrix{T}) where {T}
    return hvcat(
        (2, 2),
        A, zeros(T, size(A, 1), size(B, 2)), zeros(T, size(B, 1), size(A, 2)), B,
    )
end

function blk_diag(A::SMatrix{DA, DA, T}, B::SMatrix{DB, DB, T}) where {DA, DB, T}
    zero_AB = zeros(SMatrix{DA, DB, T})
    zero_BA = zeros(SMatrix{DB, DA, T})
    return [[A zero_AB]; [zero_BA B]]
end

Zygote.@adjoint function blk_diag(A, B)
    function blk_diag_adjoint(Δ)
        ΔA = Δ[1:size(A, 1), 1:size(A, 2)]
        ΔB = Δ[size(A, 1)+1:end, size(A, 2)+1:end]
        return (ΔA, ΔB)
    end
    return blk_diag(A, B), blk_diag_adjoint
end


text:
module ChainComplexes

include("vector_spaces_F2.jl")
include("chain_complexes.jl")

end

max_stars_repo_name: osumray/ChainComplexes

max_issues_repo_name: osumray/ChainComplexes

max_forks_repo_name: osumray/ChainComplexes


text:
@testset "MaybeVector" begin
    a = MaybeVector{Int}()
    @test length(a) == 0
    @test_throws BoundsError a[1]

    b = MaybeVector(3)
    @test length(b) == 1
    @test b[1] == 3
    @test b[] == 3
    @test_throws BoundsError b[2]
end


text:
# graph represented by a light-weight edge list

mutable struct GenericEdgeList{V,E,VList,EList} <: AbstractGraph{V,E}
    is_directed::Bool
    vertices::VList
    edges::EList
end

@graph_implements GenericEdgeList vertex_list edge_list vertex_map edge_map

const SimpleEdgeList{E} = GenericEdgeList{Int,E,UnitRange{Int},Vector{E}}
const EdgeList{V,E} = GenericEdgeList{V,E,Vector{V},Vector{E}}

# construction

simple_edgelist(nv::Integer, edges::Vector{E}; is_directed::Bool=true) where {E} =
    SimpleEdgeList{E}(is_directed, intrange(nv), edges)

edgelist(vertices::Vector{V}, edges::Vector{E}; is_directed::Bool=true) where {V,E} =
    EdgeList{V,E}(is_directed, vertices, edges)

# required interface

is_directed(g::GenericEdgeList) = g.is_directed

num_vertices(g::GenericEdgeList) = length(g.vertices)
vertices(g::GenericEdgeList) = g.vertices

num_edges(g::GenericEdgeList) = length(g.edges)
edges(g::GenericEdgeList) = g.edges
edge_index(e, g::GenericEdgeList) = edge_index(e)

# mutation

add_vertex!(g::GenericEdgeList{V}, v::V) where {V} = (push!(g.vertices, v); v)
add_vertex!(g::GenericEdgeList{V}, x) where {V} = add_vertex!(g, make_vertex(g, x))

add_edge!(g::GenericEdgeList{V,E}, e::E) where {V,E} = (push!(g.edges, e); e)
add_edge!(g::GenericEdgeList{V,E}, u::V, v::V) where {V,E} = add_edge!(g, make_edge(g, u, v))


text:
using DrWatson
@quickactivate "OneVision"

using OneVision.Benchmarks

using DataFrames
import Plots
import AbstractPlotting
AbstractPlotting.__init__()
AbstractPlotting.inline!(false)
Plots.gr()

perf_path = datadir("results", "performance_exps.bson")
var_path = datadir("results", "variation_exps.bson")

# generate and save results
perf_tables = run_performance_exps(num_of_runs = 100)
wsave(perf_path, @dict perf_tables)

# load and display results
perf_tables = wload(perf_path)[:perf_tables]
show_performance_exps(perf_tables, datadir("processed", "performance_exps"))

# generate and save results
var_tables = run_variation_exps(num_of_runs = 10)
wsave(var_path, @dict var_tables)

# load and display results
var_tables = wload(var_path)[:var_tables]
show_variation_exps(var_tables, out_path = datadir("processed", "variation_exps"))

visualize_benchmark("2D Formation Switching", default_setting(), "OneVision")