const noalloc_doc = """This method does its computation in place, performing no dynamic memory allocation."""

"""
$(SIGNATURES)

Return the mass of a subtree of a `Mechanism`, rooted at `base` (including
the mass of `base`).
"""
function subtree_mass(base::RigidBody{T}, mechanism::Mechanism{T}) where {T}
    result = isroot(base, mechanism) ? zero(T) : spatial_inertia(base).mass
    for joint in joints_to_children(base, mechanism)
        result += subtree_mass(successor(joint, mechanism), mechanism)
    end
    result
end

"""
$(SIGNATURES)

Return the total mass of the `Mechanism`.
"""
mass(m::Mechanism) = subtree_mass(root_body(m), m)

"""
$(SIGNATURES)

Compute the center of mass of an iterable subset of a `Mechanism`'s bodies in
the given state. Ignores the root body of the mechanism.
"""
function center_of_mass(state::MechanismState, itr)
    update_transforms!(state)
    T = cache_eltype(state)
    mechanism = state.mechanism
    frame = root_frame(mechanism)
    com = Point3D(frame, zero(SVector{3, T}))
    mass = zero(T)
    for body in itr
        if body.inertia !== nothing
            inertia = spatial_inertia(body)
            if inertia.mass > 0
                bodycom = transform_to_root(state, body, false) * center_of_mass(inertia)
                com += inertia.mass * FreeVector3D(bodycom)
                mass += inertia.mass
            end
        end
    end
    com /= mass
    com
end

"""
$(SIGNATURES)

Compute the center of mass of the whole `Mechanism` in the given state.
"""
center_of_mass(state::MechanismState) = center_of_mass(state, bodies(state.mechanism))

const geometric_jacobian_doc = """Compute a geometric Jacobian (also known as a
basic, or spatial Jacobian) associated with a directed path in the `Mechanism`'s
spanning tree, (a collection of `Joint`s and traversal directions) in the given state.

A geometric Jacobian maps the `Mechanism`'s joint velocity vector ``v``
to the twist of the target of the joint path (the body succeeding the last joint
in the path) with respect to the source of the joint path (the body preceding the
first joint in the path).

See also [`path`](@ref), [`GeometricJacobian`](@ref), [`Twist`](@ref).
"""

"""
$(SIGNATURES)

$geometric_jacobian_doc

`transformfun` is a callable that may be used to transform the individual motion
subspaces of each of the joints to the frame in which `out` is expressed.

$noalloc_doc
"""
function geometric_jacobian!(jac::GeometricJacobian, state::MechanismState, path::TreePath, transformfun)
    @boundscheck num_velocities(state) == size(jac, 2) || throw(DimensionMismatch())
    @framecheck jac.body default_frame(target(path))
    @framecheck jac.base default_frame(source(path))
    update_motion_subspaces!(state)
    fill!(jac.angular, 0)
    fill!(jac.linear, 0)
    for i in eachindex(path.edges) # TODO: just use @inbounds here; currently messes with frame check in set_col!
        @inbounds joint = path.edges[i]
        vrange = velocity_range(state, joint)
        @inbounds direction = directions(path)[i]
        for col in eachindex(vrange)
            @inbounds vindex = vrange[col]
            @inbounds Scol = transformfun(state.motion_subspaces.data[vindex])
            direction == PathDirections.up && (Scol = -Scol)
            set_col!(jac, vindex, Scol)
        end
    end
    jac
end

"""
$(SIGNATURES)

$geometric_jacobian_doc

`root_to_desired` is the transform from the `Mechanism`'s root frame to the frame
in which `out` is expressed.

$noalloc_doc
"""
function geometric_jacobian!(out::GeometricJacobian, state::MechanismState, path::TreePath, root_to_desired::Transform3D)
    geometric_jacobian!(out, state, path, S -> transform(S, root_to_desired))
end

"""
$(SIGNATURES)

$geometric_jacobian_doc

See [`geometric_jacobian!(out, state, path, root_to_desired)`](@ref). Uses `state`
to compute the transform from the `Mechanism`'s root frame to the frame in which
`out` is expressed.

$noalloc_doc
"""
function geometric_jacobian!(out::GeometricJacobian, state::MechanismState, path::TreePath)
    if out.frame == root_frame(state.mechanism)
        geometric_jacobian!(out, state, path, identity)
    else
        geometric_jacobian!(out, state, path, inv(transform_to_root(state, out.frame)))
    end
end

"""
$(SIGNATURES)

$geometric_jacobian_doc

The Jacobian is computed in the `Mechanism`'s root frame.

See [`geometric_jacobian!(out, state, path)`](@ref).
"""
function geometric_jacobian(state::MechanismState{X, M, C}, path::TreePath{RigidBody{M}, Joint{M}}) where {X, M, C}
    nv = num_velocities(state)
    angular = Matrix{C}(undef, 3, nv)
    linear = Matrix{C}(undef, 3, nv)
    bodyframe = default_frame(target(path))
    baseframe = default_frame(source(path))
    jac = GeometricJacobian(bodyframe, baseframe, root_frame(state.mechanism), angular, linear)
    geometric_jacobian!(jac, state, path, identity)
end

const point_jacobian_doc = """
Compute the Jacobian mapping the `Mechanism`'s joint velocity vector ``v`` to
the velocity of a point fixed to the target of the joint path (the body
succeeding the last joint in the path) with respect to the source of the joint
path (the body preceding the first joint in the path).
"""

"""
$(SIGNATURES)

$point_jacobian_doc

$noalloc_doc
"""
function _point_jacobian!(Jp::PointJacobian, state::MechanismState, path::TreePath,
                         point::Point3D, transformfun)
    @framecheck Jp.frame point.frame
    update_motion_subspaces!(state)
    fill!(Jp.J, 0)
    p̂ = Spatial.hat(point.v)
    @inbounds for i in eachindex(path.edges)
        joint = path.edges[i]
        vrange = velocity_range(state, joint)
        direction = path.directions[i]
        for col in eachindex(vrange)
            vindex = vrange[col]
            Scol = transformfun(state.motion_subspaces.data[vindex])
            direction == PathDirections.up && (Scol = -Scol)
            newcol =  -p̂ * angular(Scol) + linear(Scol)
            for row in 1:3
                Jp.J[row, vindex] = newcol[row]
            end
        end
    end
    Jp
end

"""
$(SIGNATURES)

$point_jacobian_doc

Uses `state` to compute the transform from the `Mechanism`'s root frame to the
frame in which `out` is expressed if necessary.

$noalloc_doc
"""
function point_jacobian!(out::PointJacobian, state::MechanismState, path::TreePath, point::Point3D)
    if out.frame == root_frame(state.mechanism)
        _point_jacobian!(out, state, path, point, identity)
    else
        root_to_desired = inv(transform_to_root(state, out.frame))
        let root_to_desired = root_to_desired
            _point_jacobian!(out, state, path, point, S -> transform(S, root_to_desired))
        end
    end
end

"""
$(SIGNATURES)

$point_jacobian_doc
"""
function point_jacobian(state::MechanismState{X, M, C},
                        path::TreePath{RigidBody{M}, Joint{M}},
                        point::Point3D) where {X, M, C}
    Jp = PointJacobian(point.frame, Matrix{C}(undef, 3, num_velocities(state)))
    point_jacobian!(Jp, state, path, point)
    Jp
end

const mass_matrix_doc = """Compute the joint-space mass matrix
(also known as the inertia matrix) of the `Mechanism` in the given state, i.e.,
the matrix ``M(q)`` in the unconstrained joint-space equations of motion
```math
M(q) \\dot{v} + c(q, v, w_\\text{ext}) = \\tau
```

This method implements the composite rigid body algorithm.
"""

"""
$(SIGNATURES)

$mass_matrix_doc

$noalloc_doc

The `out` argument must be an ``n_v \\times n_v`` lower triangular
`Symmetric` matrix, where ``n_v`` is the dimension of the `Mechanism`'s joint
velocity vector ``v``.
"""
function mass_matrix!(M::Symmetric, state::MechanismState)
    nv = num_velocities(state)
    @boundscheck size(M, 1) == nv || throw(DimensionMismatch("mass matrix has wrong size"))
    @boundscheck M.uplo == 'L' || throw(ArgumentError(("expected a lower triangular symmetric matrix")))
    update_motion_subspaces!(state)
    update_crb_inertias!(state)
    motion_subspaces = state.motion_subspaces.data
    @inbounds for i in Base.OneTo(nv)
        jointi = velocity_index_to_joint_id(state, i)
        bodyi = successorid(jointi, state)
        Ici = crb_inertia(state, bodyi, false)
        Si = motion_subspaces[i]
        Fi = Ici * Si
        for j in Base.OneTo(i)
            jointj = velocity_index_to_joint_id(state, j)
            M.data[i, j] = if supports(jointj, bodyi, state)
                Sj = motion_subspaces[j]
                (transpose(Fi) * Sj)[1]
            else
                zero(eltype(M))
            end
        end
    end
    M
end

mass_matrix!(result::DynamicsResult, state::MechanismState) = mass_matrix!(result.massmatrix, state)

"""
$(SIGNATURES)

$mass_matrix_doc
"""
function mass_matrix(state::MechanismState{X, M, C}) where {X, M, C}
    nv = num_velocities(state)
    ret = Symmetric(Matrix{C}(undef, nv, nv), :L)
    mass_matrix!(ret, state)
    ret
end

const momentum_matrix_doc = """Compute the momentum matrix ``A(q)`` of the
`Mechanism` in the given state.

The momentum matrix maps the `Mechanism`'s joint velocity vector ``v`` to
its total momentum.

See also [`MomentumMatrix`](@ref).
"""

const momentum_matrix!_doc = """$momentum_matrix_doc
The `out` argument must be a mutable `MomentumMatrix` with as many columns as
the dimension of the `Mechanism`'s joint velocity vector ``v``.
"""

"""
$(SIGNATURES)

$momentum_matrix!_doc

`transformfun` is a callable that may be used to transform the individual
momentum matrix blocks associated with each of the joints to the frame in which
`out` is expressed.

$noalloc_doc
"""
function momentum_matrix!(mat::MomentumMatrix, state::MechanismState, transformfun)
    nv = num_velocities(state)
    @boundscheck size(mat, 2) == nv || throw(DimensionMismatch())
    update_motion_subspaces!(state)
    update_crb_inertias!(state)
    @inbounds for i in 1 : nv
        jointi = velocity_index_to_joint_id(state, i)
        bodyi = successorid(jointi, state)
        Ici = crb_inertia(state, bodyi)
        Si = state.motion_subspaces.data[i]
        Fi = transformfun(Ici * Si)
        set_col!(mat, i, Fi)
    end
    mat
end

"""
$(SIGNATURES)

$momentum_matrix!_doc

`root_to_desired` is the transform from the `Mechanism`'s root frame to the frame
in which `out` is expressed.

$noalloc_doc
"""
function momentum_matrix!(mat::MomentumMatrix, state::MechanismState, root_to_desired::Transform3D)
    momentum_matrix!(mat, state, F -> transform(F, root_to_desired))
end

"""
$(SIGNATURES)

$momentum_matrix!_doc

See [`momentum_matrix!(out, state, root_to_desired)`](@ref). Uses `state`
to compute the transform from the `Mechanism`'s root frame to the frame in which
`out` is expressed.

$noalloc_doc
"""
function momentum_matrix!(out::MomentumMatrix, state::MechanismState)
    if out.frame == root_frame(state.mechanism)
        momentum_matrix!(out, state, identity)
    else
        momentum_matrix!(out, state, inv(transform_to_root(state, out.frame)))
    end
end

"""
$(SIGNATURES)

$momentum_matrix_doc

See [`momentum_matrix!(out, state)`](@ref).
"""
function momentum_matrix(state::MechanismState)
    ncols = num_velocities(state)
    T = cache_eltype(state)
    ret = MomentumMatrix(root_frame(state.mechanism), Matrix{T}(undef, 3, ncols), Matrix{T}(undef, 3, ncols))
    momentum_matrix!(ret, state, identity)
    ret
end

function bias_accelerations!(out::AbstractDict{BodyID, SpatialAcceleration{T}}, state::MechanismState{X, M}) where {T, X, M}
    update_bias_accelerations_wrt_world!(state)
    gravitybias = convert(SpatialAcceleration{T}, -gravitational_spatial_acceleration(state.mechanism))
    for jointid in state.treejointids
        bodyid = successorid(jointid, state)
        out[bodyid] = gravitybias + bias_acceleration(state, bodyid, false)
    end
    nothing
end

function spatial_accelerations!(out::AbstractDict{BodyID, SpatialAcceleration{T}}, state::MechanismState, v̇::SegmentedVector{JointID}) where T
    update_transforms!(state)
    update_twists_wrt_world!(state)

    # Compute joint accelerations
    joints = state.treejoints
    qs = values(segments(state.q))
    vs = values(segments(state.v))
    v̇s = values(segments(v̇))
    foreach_with_extra_args(state, out, joints, qs, vs, v̇s) do state, accels, joint, qjoint, vjoint, v̇joint
        bodyid = successorid(JointID(joint), state)
        accels[bodyid] = joint_spatial_acceleration(joint, qjoint, vjoint, v̇joint)
    end

    # Recursive propagation
    # TODO: manual frame changes. Find a way to not avoid the frame checks here.
    mechanism = state.mechanism
    root = root_body(mechanism)
    out[root] = convert(SpatialAcceleration{T}, -gravitational_spatial_acceleration(mechanism))
    @inbounds for jointid in state.treejointids
        parentbodyid, bodyid = predsucc(jointid, state)
        toroot = transform_to_root(state, bodyid, false)
        parenttwist = twist_wrt_world(state, parentbodyid, false)
        bodytwist = twist_wrt_world(state, bodyid, false)
        jointaccel = out[bodyid]
        jointaccel = SpatialAcceleration(jointaccel.body, parenttwist.body, toroot.to,
            Spatial.transform_spatial_motion(jointaccel.angular, jointaccel.linear, rotation(toroot), translation(toroot))...)
        out[bodyid] = out[parentbodyid] + (-bodytwist) × parenttwist + jointaccel
    end
    nothing
end

spatial_accelerations!(result::DynamicsResult, state::MechanismState) = spatial_accelerations!(result.accelerations, state, result.v̇)

function relative_acceleration(accels::AbstractDict{BodyID, SpatialAcceleration{T}}, body::RigidBody{M}, base::RigidBody{M}) where {T, M}
    -accels[BodyID(base)] + accels[BodyID(body)]
end

# TODO: ensure that accelerations are up-to-date
relative_acceleration(result::DynamicsResult, body::RigidBody, base::RigidBody) = relative_acceleration(result.accelerations, body, base)

function newton_euler!(
        out::AbstractDict{BodyID, Wrench{T}}, state::MechanismState{X, M},
        accelerations::AbstractDict{BodyID, SpatialAcceleration{T}},
        externalwrenches::AbstractDict{BodyID, Wrench{W}}) where {T, X, M, W}
    update_twists_wrt_world!(state)
    update_spatial_inertias!(state)
    for jointid in state.treejointids
        bodyid = successorid(jointid, state)
        wrench = newton_euler(state, bodyid, accelerations[bodyid], false)
        out[bodyid] = haskey(externalwrenches, bodyid) ? wrench - externalwrenches[bodyid] : wrench
    end
end

function joint_wrenches_and_torques!(
        torquesout::SegmentedVector{JointID},
        net_wrenches_in_joint_wrenches_out::AbstractDict{BodyID, <:Wrench}, # TODO: consider having a separate Associative{Joint{M}, Wrench{T}} for joint wrenches
        state::MechanismState)
    update_motion_subspaces!(state)
    # Note: pass in net wrenches as wrenches argument. wrenches argument is modified to be joint wrenches
    @boundscheck length(torquesout) == num_velocities(state) || error("length of torque vector is wrong")
    wrenches = net_wrenches_in_joint_wrenches_out
    for jointid in reverse(state.treejointids)
        parentbodyid, bodyid = predsucc(jointid, state)
        jointwrench = wrenches[bodyid]
        wrenches[parentbodyid] += jointwrench # action = -reaction
        for vindex in velocity_range(state, jointid)
            Scol = state.motion_subspaces.data[vindex]
            torquesout[vindex] = (transpose(Scol) * jointwrench)[1]
        end
    end
end

const dynamics_bias_doc = """Compute the 'dynamics bias term', i.e. the term
```math
c(q, v, w_\\text{ext})
```
in the unconstrained joint-space equations of motion
```math
M(q) \\dot{v} + c(q, v, w_\\text{ext}) = \\tau
```
given joint configuration vector ``q``, joint velocity vector ``v``,
joint acceleration vector ``\\dot{v}`` and (optionally) external
wrenches ``w_\\text{ext}``.

The `externalwrenches` argument can be used to specify additional
wrenches that act on the `Mechanism`'s bodies.
"""

"""
$(SIGNATURES)

$dynamics_bias_doc

$noalloc_doc
"""
function dynamics_bias!(
        torques::SegmentedVector{JointID},
        biasaccelerations::AbstractDict{BodyID, <:SpatialAcceleration},
        wrenches::AbstractDict{BodyID, <:Wrench},
        state::MechanismState{X},
        externalwrenches::AbstractDict{BodyID, <:Wrench} = NullDict{BodyID, Wrench{X}}()) where X
    bias_accelerations!(biasaccelerations, state)
    newton_euler!(wrenches, state, biasaccelerations, externalwrenches)
    joint_wrenches_and_torques!(torques, wrenches, state)
    torques
end

function dynamics_bias!(result::DynamicsResult, state::MechanismState)
    dynamics_bias!(result.dynamicsbias, result.accelerations, result.jointwrenches, state, result.totalwrenches)
end

"""
$(SIGNATURES)

$dynamics_bias_doc
"""
function dynamics_bias(
        state::MechanismState{X, M},
        externalwrenches::AbstractDict{BodyID, Wrench{W}} = NullDict{BodyID, Wrench{X}}()) where {X, M, W}
    T = promote_type(X, M, W)
    mechanism = state.mechanism
    torques = similar(velocity(state), T)
    rootframe = root_frame(mechanism)
    jointwrenches = BodyDict(BodyID(b) => zero(Wrench{T}, rootframe) for b in bodies(mechanism))
    accelerations = BodyDict(BodyID(b) => zero(SpatialAcceleration{T}, rootframe, rootframe, rootframe) for b in bodies(mechanism))
    dynamics_bias!(torques, accelerations, jointwrenches, state, externalwrenches)
    torques
end

const inverse_dynamics_doc = """Do inverse dynamics, i.e. compute ``\\tau``
in the unconstrained joint-space equations of motion
```math
M(q) \\dot{v} + c(q, v, w_\\text{ext}) = \\tau
```
given joint configuration vector ``q``, joint velocity vector ``v``,
joint acceleration vector ``\\dot{v}`` and (optionally) external
wrenches ``w_\\text{ext}``.

The `externalwrenches` argument can be used to specify additional
wrenches that act on the `Mechanism`'s bodies.

This method implements the recursive Newton-Euler algorithm.

Currently doesn't support `Mechanism`s with cycles.
"""

"""
$(SIGNATURES)

$inverse_dynamics_doc

$noalloc_doc
"""
function inverse_dynamics!(
        torquesout::SegmentedVector{JointID},
        jointwrenchesout::AbstractDict{BodyID, Wrench{T}},
        accelerations::AbstractDict{BodyID, SpatialAcceleration{T}},
        state::MechanismState,
        v̇::SegmentedVector{JointID},
        externalwrenches::AbstractDict{BodyID, <:Wrench} = NullDict{BodyID, Wrench{T}}()) where T
    @boundscheck length(tree_joints(state.mechanism)) == length(joints(state.mechanism)) || error("This method can currently only handle tree Mechanisms.")
    spatial_accelerations!(accelerations, state, v̇)
    newton_euler!(jointwrenchesout, state, accelerations, externalwrenches)
    joint_wrenches_and_torques!(torquesout, jointwrenchesout, state)
end

"""
$(SIGNATURES)

$inverse_dynamics_doc
"""
function inverse_dynamics(
        state::MechanismState{X, M},
        v̇::SegmentedVector{JointID, V},
        externalwrenches::AbstractDict{BodyID, Wrench{W}} = NullDict{BodyID, Wrench{X}}()) where {X, M, V, W}
    T = promote_type(X, M, V, W)
    mechanism = state.mechanism
    torques = similar(velocity(state), T)
    rootframe = root_frame(mechanism)
    jointwrenches = BodyDict(BodyID(b) => zero(Wrench{T}, rootframe) for b in bodies(mechanism))
    accelerations = BodyDict(BodyID(b) => zero(SpatialAcceleration{T}, rootframe, rootframe, rootframe) for b in bodies(mechanism))
    inverse_dynamics!(torques, jointwrenches, accelerations, state, v̇, externalwrenches)
    torques
end

function constraint_jacobian!(jac::AbstractMatrix, rowranges, state::MechanismState)
    # TODO: traversing jac in the wrong order
    update_motion_subspaces!(state)
    update_constraint_wrench_subspaces!(state)
    fill!(jac, 0)
    @inbounds for nontreejointid in state.nontreejointids
        path = state.constraint_jacobian_structure[nontreejointid]
        for cindex in constraint_range(state, nontreejointid)
            Tcol = state.constraint_wrench_subspaces.data[cindex]
            for i in eachindex(path.edges)
                treejoint = path.edges[i]
                vrange = velocity_range(state, treejoint)
                direction = directions(path)[i]
                sign = ifelse(direction == PathDirections.up, -1, 1)
                for col in eachindex(vrange)
                    vindex = vrange[col]
                    Scol = state.motion_subspaces.data[vindex]
                    jacelement = flipsign((transpose(Tcol) * Scol)[1], sign)
                    jac[cindex, vindex] = jacelement
                end
            end
        end
    end
    jac
end

function constraint_jacobian!(result::DynamicsResult, state::MechanismState)
    constraint_jacobian!(result.constraintjacobian, result.constraintrowranges, state)
end

"""
$(SIGNATURES)

Return the default Baumgarte constraint stabilization gains. These gains result in
critical damping, and correspond to ``T_{stab} = 0.1`` in Featherstone (2008), section 8.3.
"""
function default_constraint_stabilization_gains(scalar_type::Type{T}) where T
    ConstDict{JointID}(SE3PDGains(PDGains(T(100), T(20)), PDGains(T(100), T(20))))
end

const stabilization_gains_doc = """
The `stabilization_gains` keyword argument can be used to set PD gains for Baumgarte
stabilization, which can be used to prevent separation of non-tree (loop) joints.
See Featherstone (2008), section 8.3 for more information. There are several options
for specifying gains:

* `nothing` can be used to completely disable Baumgarte stabilization.
* Gains can be specifed on a per-joint basis using any `AbstractDict{JointID, <:RigidBodyDynamics.PDControl.SE3PDGains}`,
  which maps the `JointID` for the non-tree joints of the mechanism to the gains for that joint.
* As a special case of the second option, the same gains can be used for all joints
  by passing in a `RigidBodyDynamics.CustomCollections.ConstDict{JointID}`.

The [`default_constraint_stabilization_gains`](@ref) function is called to produce the default gains,
which use the last option.
"""

function constraint_bias!(bias::SegmentedVector, state::MechanismState{X};
        stabilization_gains::Union{AbstractDict{JointID, <:SE3PDGains}, Nothing}=default_constraint_stabilization_gains(X)) where X
    update_transforms!(state)
    update_twists_wrt_world!(state)
    update_bias_accelerations_wrt_world!(state)
    update_constraint_wrench_subspaces!(state)
    constraint_wrench_subspaces = state.constraint_wrench_subspaces.data
    for nontreejointid in state.nontreejointids
        predid, succid = predsucc(nontreejointid, state)

        predtwist = twist_wrt_world(state, predid, false)
        succtwist = twist_wrt_world(state, succid, false)
        crossterm = succtwist × predtwist

        succbias = bias_acceleration(state, succid, false)
        predbias = bias_acceleration(state, predid, false)
        jointbias = -predbias + succbias

        biasaccel = crossterm + jointbias # what's inside parentheses in 8.47 in Featherstone

        if stabilization_gains !== nothing
            # TODO: make this nicer (less manual frame juggling and no manual transformations)
            jointtransform = joint_transform(state, nontreejointid, false)
            jointtwist = -predtwist + succtwist
            jointtwist = Twist(jointtransform.from, jointtransform.to, jointtwist.frame, jointtwist.angular, jointtwist.linear) # make frames line up
            successor_to_root = transform_to_root(state, jointtransform.from) # TODO: expensive
            jointtwist = transform(jointtwist, inv(successor_to_root)) # twist needs to be in successor frame for pd method
            stabaccel = pd(stabilization_gains[nontreejointid], jointtransform, jointtwist, SE3PDMethod{:Linearized}()) # in successor frame
            stabaccel = SpatialAcceleration(stabaccel.body, stabaccel.base, biasaccel.frame,
                Spatial.transform_spatial_motion(stabaccel.angular, stabaccel.linear,
                rotation(successor_to_root), translation(successor_to_root))...) # back to world frame. TODO: ugly way to do this
            stabaccel = SpatialAcceleration(biasaccel.body, biasaccel.body, stabaccel.frame, stabaccel.angular, stabaccel.linear) # make frames line up
            @inbounds biasaccel = biasaccel + -stabaccel
        end

        for cindex in constraint_range(state, nontreejointid)
            Tcol = constraint_wrench_subspaces[cindex]
            # TODO: make nicer:
            @framecheck Tcol.frame biasaccel.frame
            bias[cindex] = (transpose(Tcol.angular) * biasaccel.angular + transpose(Tcol.linear) * biasaccel.linear)[1]
        end
    end
    bias
end

function constraint_bias!(result::DynamicsResult, state::MechanismState{X};
        stabilization_gains::Union{AbstractDict{JointID, <:SE3PDGains}, Nothing}=default_constraint_stabilization_gains(X)) where X
    constraint_bias!(result.constraintbias, state; stabilization_gains=stabilization_gains)
end

function contact_dynamics!(result::DynamicsResult{T, M}, state::MechanismState{X, M, C}) where {X, M, C, T}
    update_transforms!(state)
    update_twists_wrt_world!(state)
    mechanism = state.mechanism
    root = root_body(mechanism)
    frame = default_frame(root)
    for body in bodies(mechanism)
        bodyid = BodyID(body)
        wrench = zero(Wrench{T}, frame)
        points = contact_points(body)
        if !isempty(points)
            # TODO: AABB
            body_to_root = transform_to_root(state, bodyid, false)
            twist = twist_wrt_world(state, bodyid, false)
            states_for_body = contact_states(state, bodyid)
            state_derivs_for_body = contact_state_derivatives(result, bodyid)
            for i = 1 : length(points)
                @inbounds c = points[i]
                point = body_to_root * location(c)
                velocity = point_velocity(twist, point)
                states_for_point = states_for_body[i]
                state_derivs_for_point = state_derivs_for_body[i]
                for j = 1 : length(mechanism.environment.halfspaces)
                    primitive = mechanism.environment.halfspaces[j]
                    contact_state = states_for_point[j]
                    contact_state_deriv = state_derivs_for_point[j]
                    model = contact_model(c)
                    # TODO: would be good to move this to Contact module
                    # arguments: model, state, state_deriv, point, velocity, primitive
                    if point_inside(primitive, point)
                        separation, normal = Contact.detect_contact(primitive, point)
                        force = Contact.contact_dynamics!(contact_state_deriv, contact_state,
                            model, -separation, velocity, normal)
                        wrench += Wrench(point, force)
                    else
                        Contact.reset!(contact_state)
                        Contact.zero!(contact_state_deriv)
                    end
                end
            end
        end
        set_contact_wrench!(result, body, wrench)
    end
end

function dynamics_solve!(result::DynamicsResult, τ::AbstractVector)
    # version for general scalar types
    # TODO: make more efficient
    M = result.massmatrix
    c = parent(result.dynamicsbias)
    v̇ = parent(result.v̇)

    K = result.constraintjacobian
    k = result.constraintbias
    λ = result.λ

    nv = size(M, 1)
    nl = size(K, 1)
    G = [Matrix(M) K'; # TODO: Matrix because of https://github.com/JuliaLang/julia/issues/21332
         K zeros(nl, nl)]
    r = [τ - c; -k]
    v̇λ = G \ r
    v̇ .= view(v̇λ, 1 : nv)
    λ .= view(v̇λ, nv + 1 : nv + nl)
    nothing
end

function dynamics_solve!(result::DynamicsResult{T, S}, τ::AbstractVector{T}) where {S, T<:LinearAlgebra.BlasReal}
    # optimized version for BLAS floats
    M = result.massmatrix
    c = parent(result.dynamicsbias)
    v̇ = parent(result.v̇)

    K = result.constraintjacobian
    k = result.constraintbias
    λ = result.λ

    L = result.L
    A = result.A
    z = result.z
    Y = result.Y

    L .= M.data
    uplo = M.uplo
    LinearAlgebra.LAPACK.potrf!(uplo, L) # L <- Cholesky decomposition of M; M == L Lᵀ (note: Featherstone, page 151 uses M == Lᵀ L instead)
    τbiased = v̇
    τbiased .= τ .- c

    if size(K, 1) > 0
        # Loops.
        # Basic math:
        # DAE:
        #   [M Kᵀ] [v̇] = [τ - c]
        #   [K 0 ] [λ]  = [-k]
        # Solve for v̇:
        #   v̇ = M⁻¹ (τ - c - Kᵀ λ)
        # Plug into λ equation:
        #   K M⁻¹ (τ - c - Kᵀ λ) = -k
        #   K M⁻¹ Kᵀ λ = K M⁻¹(τ - c) + k
        # which can be solved for λ, after which λ can be used to solve for v̇.

        # Implementation loosely follows page 151 of Featherstone, Rigid Body Dynamics Algorithms.
        # It is slightly different because Featherstone uses M = Lᵀ L, whereas BLAS uses M = L Lᵀ.
        # M = L Lᵀ (Cholesky decomposition)
        # Y = K L⁻ᵀ, so that Y Yᵀ = K L⁻ᵀ L⁻¹ Kᵀ = K M⁻¹ Kᵀ = A
        # z = L⁻¹ (τ - c)
        # b = Y z + k
        # solve A λ = b for λ
        # solve M v̇  = τ - c for v̇

        # Compute Y = K L⁻ᵀ
        Y .= K
        LinearAlgebra.BLAS.trsm!('R', uplo, 'T', 'N', one(T), L, Y)

        # Compute z = L⁻¹ (τ - c)
        z .= τbiased
        LinearAlgebra.BLAS.trsv!(uplo, 'N', 'N', L, z) # z <- L⁻¹ (τ - c)

        # Compute A = Y Yᵀ == K * M⁻¹ * Kᵀ
        LinearAlgebra.BLAS.gemm!('N', 'T', one(T), Y, Y, zero(T), A) # A <- K * M⁻¹ * Kᵀ

        # Compute b = Y z + k
        b = λ
        b .= k
        LinearAlgebra.BLAS.gemv!('N', one(T), Y, z, one(T), b) # b <- Y z + k

        # Compute λ = A⁻¹ b == (K * M⁻¹ * Kᵀ)⁻¹ * (K * M⁻¹ * (τ - c) + k)
        # LinearAlgebra.LAPACK.posv!(uplo, A, b) # NOTE: doesn't work in general because A is only guaranteed to be positive semidefinite
        singular_value_zero_tolerance = 1e-10 # TODO: more principled choice
        # TODO: https://github.com/JuliaLang/julia/issues/22242
        b[:], _ = LinearAlgebra.LAPACK.gelsy!(A, b, singular_value_zero_tolerance) # b == λ <- (K * M⁻¹ * Kᵀ)⁻¹ * (K * M⁻¹ * (τ - c) + k)

        # Update τbiased: subtract Kᵀ * λ
        LinearAlgebra.BLAS.gemv!('T', -one(T), K, λ, one(T), τbiased) # τbiased <- τ - c - Kᵀ * λ

        # Solve for v̇ = M⁻¹ * (τ - c - Kᵀ * λ)
        LinearAlgebra.LAPACK.potrs!(uplo, L, τbiased) # τbiased ==v̇ <- M⁻¹ * (τ - c - Kᵀ * λ)
    else
        # No loops.
        LinearAlgebra.LAPACK.potrs!(uplo, L, τbiased) # τbiased == v̇ <- M⁻¹ * (τ - c)
    end
    nothing
end

"""
$(SIGNATURES)

Compute the joint acceleration vector ``\\dot{v}`` and Lagrange multipliers
``\\lambda`` that satisfy the joint-space equations of motion
```math
M(q) \\dot{v} + c(q, v, w_\\text{ext}) = \\tau - K(q)^{T} \\lambda
```
and the constraint equations
```math
K(q) \\dot{v} = -k
```
given joint configuration vector ``q``, joint velocity vector ``v``, and
(optionally) joint torques ``\\tau`` and external wrenches ``w_\\text{ext}``.

The `externalwrenches` argument can be used to specify additional
wrenches that act on the `Mechanism`'s bodies.

$stabilization_gains_doc
"""
function dynamics!(result::DynamicsResult, state::MechanismState{X},
        torques::AbstractVector = ConstVector(zero(X), num_velocities(state)),
        externalwrenches::AbstractDict{BodyID, <:Wrench} = NullDict{BodyID, Wrench{X}}();
        stabilization_gains=default_constraint_stabilization_gains(X)) where X
    configuration_derivative!(result.q̇, state)
    contact_dynamics!(result, state)
    for jointid in state.treejointids
        bodyid = successorid(jointid, state)
        contactwrench = result.contactwrenches[bodyid]
        result.totalwrenches[bodyid] = haskey(externalwrenches, bodyid) ? externalwrenches[bodyid] + contactwrench : contactwrench
    end
    dynamics_bias!(result, state)
    mass_matrix!(result, state)
    if has_loops(state.mechanism)
        constraint_jacobian!(result, state)
        constraint_bias!(result, state; stabilization_gains=stabilization_gains)
    end
    dynamics_solve!(result, parent(torques))
    nothing
end

"""
$(SIGNATURES)

Convenience function for use with standard ODE integrators that takes a
`Vector` argument
```math
x = \\left(\\begin{array}{c}
q\\\\
v
\\end{array}\\right)
```
and returns a `Vector` ``\\dot{x}``.
"""
function dynamics!(ẋ::StridedVector{X},
        result::DynamicsResult, state::MechanismState{X}, x::AbstractVector{X},
        torques::AbstractVector = ConstVector(zero(X), num_velocities(state)),
        externalwrenches::AbstractDict{BodyID, <:Wrench} = NullDict{BodyID, Wrench{X}}();
        stabilization_gains=default_constraint_stabilization_gains(X)) where X
    copyto!(state, x)
    dynamics!(result, state, torques, externalwrenches; stabilization_gains=stabilization_gains)
    copyto!(ẋ, result)
    ẋ
end


#!/usr/bin/env julia

using LaTeXStrings
import PyPlot.plt
using NPZ

img = npzread("img.npy")
tau = npzread("tau.npy")
mask = npzread("mask.npy")
mask_height = npzread("mask_height.npy")

tau8 = tau .> 8.0

nx, ny, nv = size(img)

# Estimate how many pixels show emission above some threshold
max_pix = maximum(img)
println("Max pix value ", max_pix)
above = (img .> (1e-5 * max_pix))
println("Total pixels traced ", sum(above))
println("Fraction of pixes unmasked ", sum(mask)/(nx * ny * nv))
println("Fraction of pixes unmasked-height ", sum(mask_height)/(nx * ny * nv))
println("Fraction of pixels brighter than max * 1e-3, ", sum(above)/(nx * ny * nv))

println("Total pixels mask ", sum(mask))
println("Total pixels mask height ", sum(mask_height))

fig, ax = plt[:subplots](nrows=5, ncols=nv, figsize=(1.5 * nv, 8.0))
for k=1:nv
    ax[1,k][:imshow](img[:,:,k], origin="lower", interpolation="none")
    ax[2,k][:imshow](tau[:,:,k], origin="lower", interpolation="none")
    ax[3,k][:imshow](tau8[:,:,k], origin="lower", interpolation="none")
    ax[4,k][:imshow](mask[:,:,k], origin="lower", interpolation="none")
    ax[5,k][:imshow](mask_height[:,:,k], origin="lower", interpolation="none")
end
fig[:subplots_adjust](left=0.03, right=0.97, wspace=0.01)
plt[:savefig]("chmaps.png")


module TimeCorrelationFunctions

using NQCDynamics: AbstractSimulation, Estimators
using NQCDynamics.NonadiabaticDistributions: Diabatic, Adiabatic
using NQCDynamics.NQCModels: nstates

export PopulationCorrelationFunction

"""
    TimeCorrelationFunction

Abstract type for defining time correlation functions
"""
abstract type TimeCorrelationFunction end

function (correlation::TimeCorrelationFunction)(sol, _)
    sim = correlation.sim
    template = correlation_template(correlation)
    out = [similar(template) for _ in 1:length(sol.u)]

    normalisation = evaluate_normalisation(sim, correlation)
    initial_value = evaluate_initial_value(sim, correlation, first(sol.u))

    for i in 1:length(sol.u)
        final_value = evaluate_final_value(sim, correlation, sol.u[i])
        correlate!(out[i], normalisation, initial_value, final_value, correlation)
    end

    return (out, false)
end

evaluate_normalisation(sim::AbstractSimulation, correlation::TimeCorrelationFunction) = 1
evaluate_initial_value(sim::AbstractSimulation, correlation::TimeCorrelationFunction, u) =
    evaluate_final_value(sim, correlation, u)

function correlate!(out, normalisation, initial_value, final_value, ::TimeCorrelationFunction)
    @. out = normalisation * initial_value * final_value
end

function correlation_template(::TimeCorrelationFunction)
    error("Implement this for your correlation function type.")
end

"""
    PopulationCorrelationFunction{T,S<:AbstractSimulation} <: TimeCorrelationFunction

Output type for computing the population correlation function.
The `statetype` determines the population type (diabatic or adiabatic).
`sim` must also be provided to access the parameters to compute the population.
"""
struct PopulationCorrelationFunction{T,S<:AbstractSimulation} <: TimeCorrelationFunction
    sim::S
    statetype::T
end

function correlation_template(correlation::PopulationCorrelationFunction)
    zeros(nstates(correlation.sim), nstates(correlation.sim))
end

function evaluate_final_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Diabatic}, u)
    Estimators.diabatic_population(sim, u)
end

function evaluate_initial_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Diabatic}, u)
    Estimators.initial_diabatic_population(sim, u)
end

function evaluate_final_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Adiabatic}, u)
    Estimators.adiabatic_population(sim, u)
end

function evaluate_initial_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Adiabatic}, u)
    Estimators.initial_adiabatic_population(sim, u)
end

function correlate!(out, normalisation::Real, initial_value::Vector, final_value::Vector, ::PopulationCorrelationFunction)
    out .= normalisation .* initial_value * final_value'
end

end # module


module LITS

####################################### Structs Exports ####################################
#AVR Exports
export AVR
export AVRFixed
export AVRSimple
export AVRTypeI
export AVRTypeII
export AVRTypeIIManual

#Machine Exports
export Machine
export BaseMachine
export OneDOneQMachine
export MarconatoMachine
export SimpleMarconatoMachine
export AndersonFouadMachine
export SimpleAFMachine
export FullMachine
export SimpleFullMachine

#PSS Exports
export PSS
export PSSFixed
export PSSFixed

#Shaft Exports
export SingleMass
export FiveMassShaft

#TG Exports
export TurbineGov
export TGFixed
export TGTypeI
export TGTypeII

# Converter Exports
export Converter
export AvgCnvFixedDC

# DC Source Exports
export DCSource
export FixedDCSource

# Filter Exports
export Filter
export LCLFilter

# FrequencyEstimator Exports
export FrequencyEstimator
export PLL

# Outer Control Exports
export OuterControl
export VirtualInertiaQdroop
export VirtualInertia
export ReactivePowerDroop

# VSControl Export
export VSControl
export CombinedVIwithVZ

# DynBranches Export
export DynLine

# Sources Exports
export StaticSource

# Base Exports
export DynGenerator
export DynInverter
export DynamicSystem
export DynamicSimulation
export run_simulation!
export get_state_series
export get_voltagemag_series
export add_device!
export add_devices!
export add_network!

####################################### Package Imports ####################################
import DiffEqBase
import SparseArrays: SparseMatrixCSC
import LinearAlgebra: BLAS
import Base.to_index
import NLsolve
import PowerSystems
const PSY = PowerSystems

#Structs for General Devices and System
include("utils/util_macros.jl")
include("structs/base.jl")
include("structs/perturbations.jl")
include("structs/simulation.jl")

#Structs for Dynamic Devices
include("structs/common.jl")
include("structs/sources.jl")
include("structs/generator.jl")
include("structs/inverter.jl")
include("structs/dyn_branches.jl")

#Common Models
include("models/branch_model.jl")
include("models/device_model.jl")
include("models/kcl.jl")
include("models/dynline_model.jl")
include("models/ref_transformations.jl")

#Generator Component Models
include("models/generator_models/machine_models.jl")
include("models/generator_models/pss_models.jl")
include("models/generator_models/avr_models.jl")
include("models/generator_models/tg_models.jl")
include("models/generator_models/shaft_models.jl")

#Inverter Component Models
include("models/inverter_models/DCside_models.jl")
include("models/inverter_models/filter_models.jl")
include("models/inverter_models/frequency_estimator_models.jl")
include("models/inverter_models/outer_control_models.jl")
include("models/inverter_models/voltage_source_control_models.jl")
include("models/inverter_models/converter_models.jl")

#Injection Models
include("models/load_models.jl")
include("models/source_models.jl")

#System Model
include("models/system_model.jl")

#Perturbations
include("perturbations/common.jl")
include("perturbations/ThreePhaseFault.jl")
include("perturbations/PowerStepChange.jl")

#Utils
include("utils/plot_utils.jl")
include("utils/print.jl")

end # module


using MacroTools

function Base.convert(::Type{Expression}, ex::Expr)
    ex.head === :if && (ex = Expr(:call, ifelse, ex.args...))
    ex.head === :call || throw(ArgumentError("internal representation does not support non-call Expr"))

    op = eval(ex.args[1])  # HACK
    args = convert.(Expression, ex.args[2:end])

    return Operation(op, args)
end
Base.convert(::Type{Expression}, x::Expression) = x
Base.convert(::Type{Expression}, x::Number) = Constant(x)

function build_expr(head::Symbol, args)
    ex = Expr(head)
    append!(ex.args, args)
    ex
end

# used in parsing
isblock(x) = length(x) == 1 && x[1] isa Expr && x[1].head == :block
function flatten_expr!(x)
    isb = isblock(x)
    if isb
        x = MacroTools.striplines(x[1])
        filter!(z->z isa Symbol || z.head != :line, x.args)
        x = (x.args...,)
    end
    x
end

function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr, expression = Val{true}; constructor=nothing)
    _vs = map(x-> x isa Operation ? x.op : x, vs)
    _ps = map(x-> x isa Operation ? x.op : x, ps)
    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(_ps)]
    (ls, rs) = zip(var_pairs..., param_pairs...)

    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))

    fname = gensym(:ModelingToolkitFunction)

    X = gensym(:MTIIPVar)
    ip_sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
    ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))

    sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
    let_expr = Expr(:let, var_eqs, sys_expr)

    fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))

    oop_ex = :(
        ($(fargs.args...),) -> begin
            X = $let_expr
            T = promote_type(map(typeof,X)...)
            convert.(T,X)
            construct = $(constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, T, $(size(rhss)...)); vec(du) .= x; du)) : constructor)
            construct(X)
        end
    )

    iip_ex = :(
        ($X,$(fargs.args...)) -> begin
            $ip_let_expr
            nothing
        end
    )

    if expression == Val{true}
        return oop_ex, iip_ex
    else
        return GeneralizedGenerated.mk_function(@__MODULE__,oop_ex), GeneralizedGenerated.mk_function(@__MODULE__,iip_ex)
    end
end

is_constant(::Constant) = true
is_constant(::Any) = false

is_operation(::Operation) = true
is_operation(::Any) = false

is_derivative(O::Operation) = isa(O.op, Differential)
is_derivative(::Any) = false

Base.occursin(t::Expression) = Base.Fix1(occursin, t)
Base.occursin(t::Expression, x::Operation ) = isequal(x, t) || any(occursin(t), x.args)
Base.occursin(t::Expression, x::Expression) = isequal(x, t)

clean(x::Variable) = x
clean(O::Operation) = isa(O.op, Variable) ? O.op : throw(ArgumentError("invalid variable: $(O.op)"))

vars(exprs) = foldl(vars!, exprs; init = Set{Variable}())
function vars!(vars, O)
    isa(O, Operation) || return vars
    for arg ∈ O.args
        if isa(arg, Operation)
            isa(arg.op, Variable) && push!(vars, arg.op)
            vars!(vars, arg)
        end
    end

    return vars
end


function optimize(je_mmle_model::JointEstimationMMLEModel)
    local parameters = je_mmle_model.parameters
    local latents = je_mmle_model.latents
    local dist = je_mmle_model.dist
    local n_index = je_mmle_model.n_index
    local i_index = je_mmle_model.i_index
    local responses_per_item = je_mmle_model.responses_per_item
    local responses_per_examinee = je_mmle_model.responses_per_examinee

    I = size(parameters, 1)
    N = size(latents, 1)

    batch_den = max(1,(nprocs()-1))
    batch_size_N = Int(ceil(N/batch_den))
    batch_size_I = Int(ceil(I/batch_den))

    iter = 1

    # starting values
    before_time = time()
    stop = false
    old_likelihood = 0.0
    old_pars = hcat(_get_parameters_vals.(parameters)...)
    start_time = time()

    latents = pmap(
        (l, idx, r) ->
            l.assessed || _update_posterior(l, parameters[idx], r),
        latents,
        i_index,
        responses_per_examinee,
        batch_size = batch_size_N,
        distributed = true
    )
    # Threads.@threads for n in 1:N
    #     latents[n] = latents[n].assessed || _update_posterior(latents[n], parameters[i_index[n]], responses_per_examinee[n])
    # end
    # latents = @sync @distributed (vcat) for n in 1:N
    #     if latents[n].assessed
    #         latents[n]
    #     else
    #         _update_posterior(latents[n], parameters[i_index[n]], responses_per_examinee[n])
    #     end
    # end
    while !stop
        # calibrate items

        # parameters = @sync @distributed (vcat) for i in 1:I
        #     if parameters[i].calibrated
        #         parameters[i]
        #     else
        #         m_step(parameters[i], latents[n_index[i]], responses_per_item[i], je_mmle_model.int_opt_settings)
        #     end
        # end
        # Threads.@threads for i in 1:I
        #     parameters[i] = parameters[i].calibrated || _calibrate_item_mmle(parameters[i], latents[n_index[i]], responses_per_item[i], je_mmle_model.int_opt_settings)
        # end
        parameters = pmap(
             (p, idx, r) ->
                 p.calibrated || m_step(p, latents[idx], r, je_mmle_model.int_opt_settings),
             parameters,
             n_index,
             responses_per_item,
             batch_size = batch_size_I,
             distributed = true
        )

        # #update posteriors

        # latents = @sync @distributed (vcat) for n in 1:N
        #     if latents[n].assessed
        #         latents[n]
        #     else
        #         _update_posterior(latents[n], parameters[i_index[n]], responses_per_examinee[n])
        #     end
        # end
        # Threads.@threads for n in 1:N
        #     latents[n] = latents[n].assessed || _update_posterior(latents[n], parameters[i_index[n]], responses_per_examinee[n])
        # end

        latents = pmap(
            (l, idx, r) ->
                l.assessed || _update_posterior(l, parameters[idx], r),
            latents,
            i_index,
            responses_per_examinee,
            batch_size = batch_size_N,
            distributed = true
        )
        #rescale dist
        if je_mmle_model.rescale_latent
            dist = _rescale(
                dist,
                latents;
                metric = je_mmle_model.metric
            )
            for n = 1:N
                latents[n].prior = dist
            end
        end

        if any([
            check_iter(iter; max_iter = Int64(je_mmle_model.ext_opt_settings[1]), verbosity = Int(je_mmle_model.ext_opt_settings[5])),
            check_time(start_time; max_time = Int64(je_mmle_model.ext_opt_settings[2]), verbosity = Int(je_mmle_model.ext_opt_settings[5])),
            check_f_tol_rel!(
                latents,
                old_likelihood;
                f_tol_rel = je_mmle_model.ext_opt_settings[3],
                verbosity = Int(je_mmle_model.ext_opt_settings[5])
                ),
            check_x_tol_rel!(
                parameters,
                old_pars;
                x_tol_rel = je_mmle_model.ext_opt_settings[4],
                verbosity = Int(je_mmle_model.ext_opt_settings[5])
            )]
            )
            stop = true
        end
        iter += 1
    end

    return parameters::Vector{<:AbstractParameters}, latents::Vector{<:AbstractLatent}, dist::Distributions.DiscreteUnivariateDistribution
end


## sample.jl : collection of methods to add samples to an Estimator
#
# Methods to add samples to the Estimator.
#
# This file is part of MultilevelEstimators.jl - A Julia toolbox for
# Multilevel Monte Carlo Methods (c) Pieterjan Robbe, 2019

function update_samples(estimator::Estimator, n_opt)
    for τ in keys(n_opt)
        n_due = n_opt[τ] - nb_of_samples(estimator, τ)
        n_due > 0 && sample!(estimator, τ, n_due)
    end
end

function update_samples(estimator::Estimator{<:U}, n_opt)
    for τ in keys(n_opt)
        n_opt[τ] > 0 && sample!(estimator, τ, n_opt[τ])
    end
end

function sample!(estimator::Estimator{I}, index::Index, n::Integer) where I
    warm_up = !has_samples_at_index(estimator, index)
    estimator[:verbose] && print_sample!_header(estimator, index, n, warm_up)

    Istart = I <: U ? sum(values(nb_of_samples(estimator))) + 1 : nb_of_samples(estimator, index) + 1 # count zero level if U
    Iend = Istart + n - 1
    t = @elapsed parallel_sample!(estimator, index, Istart, Iend)

    add_to_nb_of_samples(estimator, index, n)
    add_to_total_work(estimator, index, n)
    add_to_total_time(estimator, index, t)

    estimator[:verbose] && print_sample!_footer()
end

#
# MC
#
function parallel_sample!(estimator::Estimator{<:AbstractIndexSet, <:MC}, index::Index, Istart::Integer, Iend::Integer)

    m = estimator[:nb_of_uncertainties](index)
    f(i) = estimator.sample_function(index, transform.(view(distributions(estimator), 1:m), rand(m)))
    all_workers = workers()
    worker_idcs = 1:min(estimator[:nb_of_workers](index), nworkers())
    pool = CachingPool(all_workers[worker_idcs])
    #batch_size = ceil(Int, (Iend - Istart + 1)/length(worker_idcs))
    batch_size = 1
    retry_delays = ExponentialBackOff(n = 3)
    S = pmap(f, pool, Istart:Iend, batch_size = batch_size, retry_delays = retry_delays)
    clear!(pool)

    s_diff = first.(S)
    s = last.(S)

    append_samples!(estimator, index, s_diff, s)
end

function append_samples!(estimator::Estimator{<:AbstractIndexSet, <:MC}, index::Index, s_diff, s)
    for n_qoi in 1:estimator[:nb_of_qoi]
        append_samples_diff!(estimator, n_qoi, index, getindex.(s_diff, n_qoi))
        append_samples!(estimator, n_qoi, index, getindex.(s, n_qoi))
    end
end

function append_samples!(estimator::Estimator{<:U, <:MC}, index::Index, s_diff, s)
    for idx in zero(index):index
        for n_qoi in 1:estimator[:nb_of_qoi]
            append_samples_diff!(estimator, n_qoi, idx, getindex.(getindex.(s_diff, Ref(idx + one(idx))), n_qoi))
            append_samples!(estimator, n_qoi, idx, getindex.(getindex.(s, Ref(idx + one(idx))), n_qoi))
            idx == zero(idx) && append_to_accumulator!(estimator, n_qoi, zeros(length(s_diff)))
            add_to_accumulator!(estimator, n_qoi, 1 / Prob(estimator, idx) * getindex.(getindex.(s_diff, Ref(idx + one(idx))), n_qoi))
        end
    end
end

#
# QMC
#
function parallel_sample!(estimator::Estimator{I, <:QMC}, index::Index, Istart::Integer, Iend::Integer) where I<:AbstractIndexSet

    m = estimator[:nb_of_uncertainties](index)
    x(i) = begin
        _index = I <: U ? zero(index) : index # take points from level 0 if unbiased
        point = get_point(generator(estimator, index, i[1]), i[2] - 1)
        length(point) < m && append!(point, rand(m - length(point)))
        view(point, 1:m)
    end
    f(i) = estimator.sample_function(index, transform.(view(distributions(estimator), 1:m), x(i)))
    all_workers = workers()
    worker_idcs = 1:min(estimator[:nb_of_workers](index), nworkers())
    pool = CachingPool(all_workers[worker_idcs])
    #batch_size = ceil(Int, (Iend - Istart + 1)/length(worker_idcs))
    batch_size = 1
    retry_delays = ExponentialBackOff(n = 3)
    S = pmap(f, pool, Iterators.product(1:estimator[:nb_of_shifts](index), Istart:Iend), batch_size = batch_size, retry_delays = retry_delays)
    clear!(pool)

    s_diff = first.(S)
    s = last.(S)

    append_samples!(estimator, index, s_diff, s)
end

function append_samples!(estimator::Estimator{<:AbstractIndexSet, <:QMC}, index::Index, s_diff, s)
    for n_shift in 1:estimator[:nb_of_shifts](index)
        for n_qoi in 1:estimator[:nb_of_qoi]
            append_samples_diff!(estimator, n_qoi, n_shift, index, getindex.(view(s_diff, n_shift, :), n_qoi))
            append_samples!(estimator, n_qoi, n_shift, index, getindex.(view(s, n_shift, :), n_qoi))
        end
    end
end

function append_samples!(estimator::Estimator{<:U, <:QMC}, index::Index, s_diff, s)
    for idx in zero(index):index
        for n_shift in 1:estimator[:nb_of_shifts](index)
            for n_qoi in 1:estimator[:nb_of_qoi]
                append_samples_diff!(estimator, n_qoi, n_shift, idx, getindex.(getindex.(view(s_diff, n_shift, :), Ref(idx + one(idx))), n_qoi))
                append_samples!(estimator, n_qoi, n_shift, idx, getindex.(getindex.(view(s, n_shift, :), Ref(idx + one(idx))), n_qoi))
                idx == zero(idx) && append_to_accumulator!(estimator, n_qoi, n_shift, zeros(size(s_diff, 2)))
                add_to_accumulator!(estimator, n_qoi, n_shift, 1 / Prob(estimator, idx) * getindex.(getindex.(view(s_diff, n_shift, :), Ref(idx + one(idx))), n_qoi))
            end
        end
    end
end


@testset "nl_convert" begin
    @testset "check special conversion cases" begin
        special_cases = [:cbrt, :abs2, :inv, :log2, :log1p, :exp2, :expm1, :sec,
                         :csc, :cot, :sind, :cosd, :tand, :asind, :acosd, :atand,
                         :secd, :cscd, :cotd, :sech, :csch, :coth, :asech, :acsch]
        for func in special_cases
            x = rand()
            expr = Expr(:call, func, x)
            @test isapprox(eval(JAMSDWriter.convert_formula(expr)), eval(expr), atol=1e-6)
        end
        # These functions need input >1
        for func in [:acoth, :asec, :acsc, :acot, :asecd, :acscd, :acotd]
            x = rand() + 1
            expr = Expr(:call, func, x)
            @test isapprox(eval(JAMSDWriter.convert_formula(expr)), eval(expr), atol=1e-6)
        end
    end

    @testset "check numeric values" begin
        x = rand()
        @test JAMSDWriter.convert_formula(:($x)) == :($x)
        x = -rand()
        @test JAMSDWriter.convert_formula(:($x)) == :($x)
    end

    @testset "check unary, binary and n-ary plus" begin
        expr = :(+(1))
        @test JAMSDWriter.convert_formula(expr) == :(1)
        expr = :(1 + 2)
        @test JAMSDWriter.convert_formula(expr) == :(1 + 2)
        expr = :(1 + 2 + 3)
        @test JAMSDWriter.convert_formula(expr) == :(sum(1, 2, 3))
    end

    @testset "check unary, binary and n-ary minus" begin
        expr = :(- x)
        @test JAMSDWriter.convert_formula(expr) == :(neg(x))
        expr = :(x - y)
        @test JAMSDWriter.convert_formula(expr) == :(x - y)
        expr = :(x - y - z)
        @test JAMSDWriter.convert_formula(expr) == :((x - y) - z)
    end

    @testset "check n-ary multiplication" begin
        expr = :(x * y * z)
        @test JAMSDWriter.convert_formula(expr) == :(x * (y * z))
    end

    @testset "check comparison expansion" begin
        expr = :(1 < 2 < 3)
        @test JAMSDWriter.convert_formula(expr) == :(1 < 2 && 2 < 3)
        expr = :(1 < 2 < 3 < 4)
        @test JAMSDWriter.convert_formula(expr) == :((1 < 2 && 2 < 3) && 3 < 4)
    end
end


module RGFPrecompTest2
    using RuntimeGeneratedFunctions
    RuntimeGeneratedFunctions.init(@__MODULE__)

    y_in_RGFPrecompTest2 = 2

    # Simulates a helper function which generates an RGF, but caches it in a
    # different module.
    function generate_rgf(cache_module)
        context_module = @__MODULE__
        RuntimeGeneratedFunction(cache_module, @__MODULE__, :((x)->y_in_RGFPrecompTest2+x))
    end
end


const t=2.1

@inline initial_condition_generator(rng,n,model) = (X=[0.0],φ=[1])

const tau = fixed_time(t)

const x = range(0,10,length=10_001)


#### ExportFuncs

# Provides a set of export functions that operate on StateVec.

using Printf, DelimitedFiles
export export_state, import_state!

include("ExportFuncsPrimitive.jl")

"""
    export_state(sv::StateVec, grid::Grid, dir, filename)

Export StateVec to a human-readable file `filename` in directory `dir`.
"""
function export_state(sv::StateVec, grid::Grid, dir, filename)
  domain = over_elems(grid)
  vn = var_names(sv)
  FT = eltype(grid)

  headers = ["z", [var_string(sv, ϕ, i) for ϕ in vn for i in over_sub_domains(sv, ϕ)]...]

  data = zeros(FT,length(domain), length(headers)-1)
  for ϕ in vn, i in over_sub_domains(sv, ϕ), k in domain
    i_eff = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
    data[k, i_eff] = sv[ϕ, k, i]
  end

  z = grid.zc[domain]
  data_all = hcat(z, data)
  mkpath(dir)

  file = joinpath(dir, filename)
  max_len = 18 # max_len = max(length.(headers)...) # Must match fmt in export_data
  headers = [repeat(" ", max_len-length(x))*x for x in headers]
  export_data(data_all, file, headers)
end

"""
    import_state!(sv::StateVec, grid::Grid, dir, filename)

Import StateVec from a human-readable file `filename` in directory `dir`.
"""
function import_state!(sv::StateVec, grid::Grid, dir, filename)
  domain = over_elems(grid)
  vn = var_names(sv)
  FT = eltype(grid)

  headers = ["z", [var_string(sv, ϕ, i) for ϕ in vn for i in over_sub_domains(sv, ϕ)]...]
  data_all = zeros(FT, length(domain), length(headers))
  mkpath(dir)

  file = joinpath(dir, filename)
  max_len = 18 # max_len = max(length.(headers)...) # Must match fmt in export_data
  headers = [repeat(" ", max_len-length(x))*x for x in headers]
  import_data!(data_all, file, headers)

  data = data_all[:,2:end] # remove z
  for ϕ in vn, i in over_sub_domains(sv, ϕ), k in domain
    i_eff = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
    sv[ϕ, k, i] = data[k, i_eff]
  end

end


function MultivariateFunctionMappingAdapter(arg0::MultivariateFunction, arg1::Vector{jdouble}, arg2::Vector{jdouble})
    return MultivariateFunctionMappingAdapter((MultivariateFunction, Vector{jdouble}, Vector{jdouble}), arg0, arg1, arg2)
end

function bounded_to_unbounded(obj::MultivariateFunctionMappingAdapter, arg0::Vector{jdouble})
    return jcall(obj, "boundedToUnbounded", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function unbounded_to_bounded(obj::MultivariateFunctionMappingAdapter, arg0::Vector{jdouble})
    return jcall(obj, "unboundedToBounded", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function value(obj::MultivariateFunctionMappingAdapter, arg0::Vector{jdouble})
    return jcall(obj, "value", jdouble, (Vector{jdouble},), arg0)
end


import ..UncertainValues: CertainValue
using RecipesBase

@recipe f(::Type{CertainValue{T}}, x::CertainValue{T}) where {T} = [x.value]

@recipe function f(certainvals::Vector{CertainValue})
    @series begin
        [val.value for val in certainvals]
    end
end

@recipe function f(certainvals::Vararg{CertainValue,N}) where {N}
    @series begin
        [val.value for val in certainvals]
    end
end


include("aes-code.jl")

global const BLOCK_BYTES = (Nb * WORDLENGTH)

# Carries out AES in ECB mode on the given blocks using the given key and
# returns the result.
# ECB mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, ECB mode is carried out in decryption
# direction.
# Both the blocks and key must be represented as strings consisting of
# only hexadecimal characters.
# The returned result is a string consisting of only hexadecimal
# characters.
function AESECB(blocks::String, key::String, encrypt::Bool)
	bytes2hex(AESECB(hex2bytes(blocks), hex2bytes(key), encrypt))
end

# Carries out AES in ECB mode on the given blocks using the given key.
# CBC mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, ECB mode is carried out in decryption
# direction.
function AESECB(blocks::Array{UInt8, 1}, key::Array{UInt8, 1}, encrypt::Bool)
	noBlocks = paddedCheck(blocks, key)
	o = Array{UInt8}(undef, length(blocks))

	for i=1:noBlocks
		indices = blockIndices(blocks, i)
		o[indices] = encrypt ? AESEncrypt(blocks[indices], key) : AESDecrypt(blocks[indices], key)
	end

	return o
end

# Carries out AES in CBC mode on the given blocks using the given key and
# returns the result.
# CBC mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, CBC mode is carried out in decryption
# direction.
# Both the blocks and key must be represented as strings consisting of
# only hexadecimal characters.
# The returned result is a string consisting of only hexadecimal
# characters.
function AESCBC(blocks::String, key::String,
	iv::String, encrypt::Bool)
	bytes2hex(AESCBC(hex2bytes(blocks), hex2bytes(key), hex2bytes(iv),
	encrypt))
end

# Carries out AES in CBC mode on the given blocks using the given key.
# CBC mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, CBC mode is carried out in decryption
# direction.
function AESCBC(blocks::Array{UInt8, 1}, key::Array{UInt8, 1},
	iv::Array{UInt8, 1}, encrypt::Bool)
	noBlocks = paddedCheck(blocks, key)
	if length(iv) != BLOCK_BYTES
		error("IV does not have 16 bytes!")
	end
	o = Array{UInt8}(undef, length(blocks))
	prev = iv

	for i=1:noBlocks
		indices = blockIndices(blocks, i)
		if encrypt
			curr = AESEncrypt(xor.(prev , blocks[indices]), key)
			o[indices] = curr
			prev = curr
		else
			# decrypt
			curr = xor.(AESDecrypt(blocks[indices], key) , prev)
			o[indices] = curr
			prev = blocks[indices]
		end
	end

	return o
end

# Carries out AES in CFB mode on the given blocks using the given key and
# returns the result.
# CFB mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, CFB mode is carried out in decryption
# direction.
# Both the blocks and key must be represented as strings consisting of
# only hexadecimal characters.
# The returned result is a string consisting of only hexadecimal
# characters.
function AESCFB(blocks::String, key::String,
	iv::String, encrypt::Bool)
	bytes2hex(AESCFB(hex2bytes(blocks), hex2bytes(key), hex2bytes(iv),
	encrypt))
end

# Carries out AES in CFB mode on the given blocks using the given key.
# CFB mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, CFB mode is carried out in decryption
# direction.
function AESCFB(blocks::Array{UInt8, 1}, key::Array{UInt8, 1},
	iv::Array{UInt8, 1}, encrypt::Bool)
	noBlocks = keyStreamCheck(blocks, key, iv)
	o = Array{UInt8}(undef, length(blocks))
	curr = iv

	for i=1:noBlocks
		indices = blockIndices(blocks, i)
		o[indices] = xor.(AESEncrypt(curr, key)[1:length(indices)] , blocks[indices])
		if encrypt
			curr = o[indices]
		else
			# decrypt
			curr = blocks[indices]
		end
	end

	return o
end

# Carries out AES in OFB mode on the given blocks using the given key and
# returns the result.
# Both the blocks and key must be represented as strings consisting of
# only hexadecimal characters.
# The returned result is a string consisting of only hexadecimal
# characters.
function AESOFB(blocks::String, key::String, iv::String)
	bytes2hex(AESOFB(hex2bytes(blocks), hex2bytes(key), hex2bytes(iv)))
end

# Carries out AES in OFB mode on the given blocks using the given key.
function AESOFB(blocks::Array{UInt8, 1}, key::Array{UInt8, 1},
	iv::Array{UInt8, 1})
	noBlocks = keyStreamCheck(blocks, key, iv)
	o = Array{UInt8}(undef, length(blocks))
	prev = iv

	for i=1:noBlocks
		indices = blockIndices(blocks, i)
		eb = AESEncrypt(prev, key)
		o[indices] = xor.(eb[1:length(indices)] , blocks[indices])
		prev = eb
	end

	return o
end

# Carries out AES in CTR mode on the given blocks using the given key and
# returns the result.
# Both the blocks and key must be represented as strings consisting of
# only hexadecimal characters.
# The returned result is a string consisting of only hexadecimal
# characters. Treats the low eight bytes of the iv array as the little endian
# counter.
function AESCTR(blocks::String, key::String, iv::String)
	bytes2hex(AESCTR(hex2bytes(blocks), hex2bytes(key), hex2bytes(iv)))
end

# Carries out AES in CTR mode on the given blocks using the given key.
# CTR mode is carried out in encryption direction if the parameter
# encrypt is true. Otherwise, CTR mode is carried out in decryption
# direction. Treats the least-significant eight bytes of the iv array
# as the big endian counter.
function AESCTR(blocks::Array{UInt8, 1}, key::Array{UInt8, 1},
	iv::Array{UInt8, 1})
	noBlocks = keyStreamCheck(blocks, key, iv)
	curr = copy(iv)
	o = Array{UInt8}(undef, length(blocks))

	for i=1:noBlocks
		indices = blockIndices(blocks, i)
		eb = AESEncrypt(curr, key)
		o[indices] = xor.(eb[1:length(indices)] , blocks[indices])
		for bi=reverse((length(curr) - 7):length(curr))
			tmp = curr[bi]
			curr[bi] = UInt8(mod(Int(tmp) + 1, 256))
			if curr[bi] > tmp
				# no byte overflow
				break
			end
		end
	end

	return o
end

# Checks whether the parameters are OK for ECB/CBC mode.
# Returns the number of blocks including the last block whose length must
# be as long as the standard block length (16 bytes).
function paddedCheck(blocks::Array{UInt8, 1}, key::Array{UInt8, 1})
	noBlocks = div(length(blocks), BLOCK_BYTES)
	if (noBlocks < 1) || ((noBlocks * BLOCK_BYTES) != length(blocks))
		error("No blocks or length of blocks is not a multplile of ",
		"16!")
	end
	# Check if key is OK
	AESParameters(key)
	return noBlocks
end

# Checks whether the parameters are OK for CFB/OFB/CTR mode.
# Returns the number of blocks including the last block whose length may
# be less than standard block length (16 bytes).
function keyStreamCheck(blocks::Array{UInt8, 1}, key::Array{UInt8, 1},
	iv::Array{UInt8, 1})
	if length(blocks) < 1
		error("blocks must contain at least one byte!")
	end
	noBlocks = div(length(blocks) + BLOCK_BYTES - 1, BLOCK_BYTES)
	# Check if key is OK
	AESParameters(key)
	if length(iv) != BLOCK_BYTES
		error("iv does not have 16 bytes!")
	end
	return noBlocks
end

# Returns the indices of the bytes of a block given a one-based
# block number. The last block may be incomplete.
function blockIndices(blocks::Array{UInt8, 1}, blockNumber::Int)
	@assert(blockNumber >= 1)
	((blockNumber - 1) * BLOCK_BYTES + 1):(min(blockNumber * BLOCK_BYTES, length(blocks)))
end


ℓ = multivariate_normal(ones(1))
κ = GaussianKineticEnergy(1)
Q = evaluate_ℓ(ℓ, [1.0])

reporters_1 = [
    NoProgressReport(),
    ProgressMeterReport(),
]

reporters_2 = [
    LogProgressReport(),
]

for reporter in deepcopy(vcat(reporters_1, reporters_2))
    results = mcmc_with_warmup(RNG, ℓ, 10_000; reporter = reporter)
end

for reporter in deepcopy(vcat(reporters_1, reporters_2))
    DynamicHMC.report(reporter, "")

    mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)
    DynamicHMC.report(mcmc_reporter_1, "")
    DynamicHMC.report(mcmc_reporter_1, 1)

    mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)
    DynamicHMC.report(mcmc_reporter_2, "")
    DynamicHMC.report(mcmc_reporter_2, 1)
end

for reporter in deepcopy(reporters_1)
    DynamicHMC.report(reporter, "")
    DynamicHMC.report(reporter, 1)

    mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)
    DynamicHMC.report(mcmc_reporter_1, "")
    DynamicHMC.report(mcmc_reporter_1, 1)

    mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)
    DynamicHMC.report(mcmc_reporter_2, "")
    DynamicHMC.report(mcmc_reporter_2, 1)
end


using FTCTests  # reexport FaulTolerantControl
using Transducers
using Random
using Test

"""
This is used for 2nd-year report.
"""
function run_multiple_sim_2nd_year_report(N=1;
        manoeuvres=[:hovering, :forward],
        methods=[:adaptive, :adaptive2optim],
        t0=0.0, tf=20.0,
        h_threshold=5.0,  # m (nothing: no constraint)
        actual_time_limit=60.0,  # s
        N_thread=Threads.nthreads(),
        collector=Transducers.tcollect, will_plot=false, seed=2021,
    )
    fault_time = 0.0
    _fault_list_single = [[
                          [LoE(fault_time, index, effectiveness)]
                          for index in 1:6
                          for effectiveness in [0.9, 0.5, 0.1, 0.0]
                         ]...]
    _fault_list_double = [[
                          [LoE(fault_time, index, effectiveness),
                           LoE(fault_time, index2, effectiveness)]
                          for index in 1:5
                          for index2 in index+1:6
                          for effectiveness in [0.9, 0.5, 0.1, 0.0]
                         ]...]
    _fault_list_single_failure_single_fault = [[
                                               [LoE(fault_time, index, 0.0),
                                                LoE(fault_time, index2, effectiveness)]
                                               for index in 1:6
                                               for index2 in filter(x -> x != index, 1:6)
                                               for effectiveness in [0.9, 0.5, 0.1]
                                              ]...]
    _fault_list = [_fault_list_single..., _fault_list_double..., _fault_list_single_failure_single_fault...]
    _faults = 1:N |> Map(i -> rand(_fault_list)) |> collect  # randomly sampled N faults
    θs = [[0, 0, -10.0]]  # Control points of Bezier curve; constant position tracking
    θs_array = 1:N |> Map(i -> θs) |> collect
    run_multiple_sim(N,
                     manoeuvres,
                     methods,
                     _faults,
                     θs_array,
                     t0, tf,
                     h_threshold, actual_time_limit, N_thread, collector,
                     will_plot, seed
                    )
end


# RolloutSimulator
# maintained by @zsunberg

"""
    RolloutSimulator(rng, max_steps)
    RolloutSimulator(; <keyword arguments>)

A fast simulator that just returns the reward

The simulation will be terminated when either
1) a terminal state is reached (as determined by `isterminal()` or
2) the discount factor is as small as `eps` or
3) max_steps have been executed

# Keyword arguments:
- rng: A random number generator to use.
- eps: A small number; if γᵗ where γ is the discount factor and t is the time step becomes smaller than this, the simulation will be terminated.
- max_steps: The maximum number of steps to simulate.

# Usage (optional arguments in brackets):

    ro = RolloutSimulator()
    history = simulate(ro, pomdp, policy, [updater [, init_belief [, init_state]]])

See also: [`HistoryRecorder`](@ref), [`run_parallel`](@ref)
"""
struct RolloutSimulator{RNG<:AbstractRNG} <: Simulator
    rng::RNG

    # optional: if these are null, they will be ignored
    max_steps::Union{Nothing,Int}
    eps::Union{Nothing,Float64}
end

RolloutSimulator(rng::AbstractRNG, d::Int=typemax(Int)) = RolloutSimulator(rng, d, nothing)
function RolloutSimulator(;rng=Random.GLOBAL_RNG,
                           eps=nothing,
                           max_steps=nothing)
    return RolloutSimulator{typeof(rng)}(rng, max_steps, eps)
end

@POMDP_require simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy) begin
    @req updater(::typeof(policy))
    bu = updater(policy)
    @subreq simulate(sim, pomdp, policy, bu)
end

@POMDP_require simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy, bu::Updater) begin
    @req initialstate_distribution(::typeof(pomdp))
    dist = initialstate_distribution(pomdp)
    @subreq simulate(sim, pomdp, policy, bu, dist)
end

function simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy, bu::Updater=updater(policy))
    dist = initialstate_distribution(pomdp)
    return @inferred simulate(sim, pomdp, policy, bu, dist)
end

@POMDP_require simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy, updater::Updater, initial_belief) begin
    @req rand(::typeof(sim.rng), ::typeof(initial_belief))
    @subreq simulate(sim, pomdp, policy, updater, initial_belief, s)
end

function simulate(sim::RolloutSimulator, pomdp::POMDP{S}, policy::Policy, updater::Updater, initial_belief) where {S}
    s = rand(sim.rng, initial_belief)::S
    return @inferred simulate(sim, pomdp, policy, updater, initial_belief, s)
end

@POMDP_require simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy, updater::Updater, initial_belief, s) begin
    P = typeof(pomdp)
    S = statetype(P)
    A = actiontype(P)
    O = obstype(P)
    @req initialize_belief(::typeof(updater), ::typeof(initial_belief))
    @req isterminal(::P, ::S)
    @req discount(::P)
    @req generate_sor(::P, ::S, ::A, ::typeof(sim.rng))
    b = initialize_belief(updater, initial_belief)
    @req action(::typeof(policy), ::typeof(b))
    @req update(::typeof(updater), ::typeof(b), ::A, ::O)
end

function simulate(sim::RolloutSimulator, pomdp::POMDP, policy::Policy, updater::Updater, initial_belief, s)
    if isterminal(pomdp, s)
        return 0.0
    end

    if sim.eps == nothing
        eps = 0.0
    else
        eps = sim.eps
    end

    if sim.max_steps == nothing
        max_steps = typemax(Int)
    else
        max_steps = sim.max_steps
    end

    disc = 1.0
    r_total = 0.0

    b = initialize_belief(updater, initial_belief)

    step = 1

    while disc > eps && step <= max_steps
        a = action(policy, b)

        sp, o, r = generate_sor(pomdp, s, a, sim.rng)

        r_total += disc*r

        s = sp

        bp = update(updater, b, a, o)
        b = bp

        disc *= discount(pomdp)
        step += 1

        if isterminal(pomdp, s, a, o)
            break
        end
    end

    return r_total
end

@POMDP_require simulate(sim::RolloutSimulator, mdp::MDP, policy::Policy) begin
    istate = initialstate(mdp, sim.rng)
    @subreq simulate(sim, mdp, policy, istate)
end

@POMDP_require simulate(sim::RolloutSimulator, mdp::MDP, policy::Policy, initialstate) begin
    P = typeof(mdp)
    S = typeof(initialstate)
    A = actiontype(mdp)
    @req isterminal(::P, ::S)
    @req action(::typeof(policy), ::S)
    @req generate_sr(::P, ::S, ::A, ::typeof(sim.rng))
    @req discount(::P)
end

function simulate(sim::RolloutSimulator, mdp::MDP, policy::Policy)
    istate = initialstate(mdp, sim.rng)
    @inferred simulate(sim, mdp, policy, istate)
end

# This function is also defined in Entropy for an InformationRewardBeliefMDP
function simulate(sim::RolloutSimulator, mdp::Union{MDP{S}, POMDP{S}}, policy::Policy, initialstate::S) where {S}
    if isterminal(mdp, initialstate)
        return 0.0
    end

    if sim.eps == nothing
        eps = 0.0
    else
        eps = sim.eps
    end

    if sim.max_steps == nothing
        max_steps = typemax(Int)
    else
        max_steps = sim.max_steps
    end

    s = initialstate

    disc = 1.0
    r_total = 0.0
    step = 1

    while disc > eps && step <= max_steps
        a = action(policy, s)

        sp, r = generate_sr(mdp, s, a, sim.rng)

        r_total += disc*r

        s = sp

        disc *= discount(mdp)
        step += 1

        if isterminal(mdp, s, a)
            break
        end
    end

    return r_total
end


using Base.Threads
function blockranges(nblocks, total_len)
    rem = total_len % nblocks
    main_len = div(total_len, nblocks)

    starts=Int[1]
    ends=Int[]
    for ii in 1:nblocks
        len = main_len
        if rem>0
            len+=1
            rem-=1
        end
        push!(ends, starts[end]+len-1)
        push!(starts, ends[end] + 1)
    end
    @assert ends[end] == total_len
    starts[1:end-1], ends
end

function sumby_multi_rg(by::AbstractVector{T},  val::AbstractVector{S})::Dict{T,S} where {T,S <: Number}
    (starts, ends) = blockranges(nthreads(), length(by))
    res = Dict{T,S}[]
    @threads for (ss, ee) in collect(zip(starts, ends))
        # @inbounds byv = @view by[ss:ee]
        # @inbounds valv = @view val[ss:ee]
        @inbounds push!(res, sumby_radixgroup!(by[ss:ee], val[ss:ee]))
    end

    szero  = zero(S)
    @inbounds res_fnl = res[1]
    for res1 in res[2:end]
        for k in keys(res1)
            res_fnl[k] = get(res_fnl,k,szero) + res1[k]
        end
    end
    return res_fnl
end

function sumby_multi_rs(by::AbstractVector{T},  val::AbstractVector{S})::Dict{T,S} where {T,S <: Number}
    (starts, ends) = blockranges(nthreads(), length(by))
    res = Dict{T,S}[]
    @threads for (ss, ee) in collect(zip(starts, ends))
        # @inbounds byv = @view by[ss:ee]
        # @inbounds valv = @view val[ss:ee]
        @inbounds push!(res, sumby_radixsort!(by[ss:ee], val[ss:ee]))
    end

    szero  = zero(S)
    @inbounds res_fnl = res[1]
    for res1 in res[2:end]
        for k in keys(res1)
            res_fnl[k] = get(res_fnl,k,szero) + res1[k]
        end
    end
    return res_fnl
end

function sumby_multi_van(by::AbstractVector{T},  val::AbstractVector{S})::Dict{T,S} where {T,S <: Number}
    (starts, ends) = blockranges(nthreads(), length(by))
    res = Dict{T,S}[]
    @threads for (ss, ee) in collect(zip(starts, ends))
        # @inbounds byv = @view by[ss:ee]
        # @inbounds valv = @view val[ss:ee]
        @inbounds push!(res, sumby!(by[ss:ee], val[ss:ee]))
    end

    szero  = zero(S)
    @inbounds res_fnl = res[1]
    for res1 in res[2:end]
        for k in keys(res1)
            res_fnl[k] = get(res_fnl,k,szero) + res1[k]
        end
    end
    return res_fnl
end


""" Schedules

This module will contain the tooling which converts diagrams to DWDs. This will
also include optimization tooling relating to DWD schedules (compressing matrix
operations, pre-computing constant values, moving addition to reduce linear).
"""
module Schedules
using Catlab.Syntax
using Catlab.WiringDiagrams.DirectedWiringDiagrams
using Catlab.Theories
using Catlab.CategoricalAlgebra
using Catlab.Programs.DiagrammaticPrograms: NamedGraph

using ..Diagrams

export diag2dwd

sp_otimes(expr) = expr isa ObExpr{:otimes} ? expr.args : [expr]
n_args(expr) = length(sp_otimes(expr))

isa_proj(hom) = startswith(string(hom), "proj")
proj_ind(hom) = unsub(split(string(hom), "_")[1][5:end])

super2int = Dict(
  '⁰'=>'0', '¹'=>'1', '²'=>'2', '³'=>'3', '⁴'=>'4', '⁵'=>'5',
  '⁶'=>'6', '⁷'=>'7', '⁸'=>'8', '⁹'=>'9'
)
sub2int = Dict(
  '₀'=>'0', '₁'=>'1', '₂'=>'2', '₃'=>'3', '₄'=>'4', '₅'=>'5',
  '₆'=>'6', '₇'=>'7', '₈'=>'8', '₉'=>'9'
)

unsub(str) = begin
  parse(Int64, join([sub2int[a] for a in str]))
end

function eval_deps!(dwd, graph, el, w2b, el2p, obs; in_els = Dict{Int, Int}(), boundaries = [])
  if el in keys(el2p)
    return el2p[el]
  end

  inputs = incident(graph, el, :tgt)

  labels = graph[:ename]
  el_types = sp_otimes(obs[el])
  outp = if length(inputs) > 1
    sort!(inputs, by = i->proj_ind(labels[i]))
    vcat(map(w -> eval_deps!(dwd, graph, graph[w, :src], w2b, el2p, obs; in_els = in_els), inputs)...)
  elseif length(inputs) == 0
    el ∈ keys(in_els) || error("Element $el has no dependencies, but is not defined in `in_els`")
    [(input_id(dwd), in_els[el])]
  else
    in_ps = eval_deps!(dwd, graph, graph[inputs[1], :src], w2b, el2p, obs; in_els = in_els)
    if isa_proj(graph[inputs[1], :ename])
      [in_ps[proj_ind(graph[inputs[1], :ename])]]
    else
      wires = map(enumerate(in_ps)) do (ip, op)
        Wire(port_value(dwd, Port(op[1], OutputPort, op[2])), op, (w2b[inputs[1]], ip))
      end
      add_wires!(dwd, wires)
      map(1:length(output_ports(dwd, w2b[inputs[1]]))) do p
        (w2b[inputs[1]], p)
      end
    end
  end

  if el in keys(el2p)
    return el2p[el]
  end

  # Evaluate Boundary Conditions
  bcs = incident(graph, el, :src)
  filter!(a->"$(Symbol(graph[a, :ename]))"[1] == '∂', bcs)
  for bc in bcs
    wires = map(enumerate(outp)) do (ip, op)
      Wire(port_value(dwd, Port(op[1], OutputPort, op[2])), op, (w2b[bc], ip))
    end
    add_wires!(dwd, wires)
    outp .= map(1:length(output_ports(dwd, w2b[bc]))) do p
      (w2b[bc], p)
    end
  end
  el2p[el] = outp
  outp
end

name(a::HomExpr) = head(a) == :generator ? args(a)[1] : head(a)

function diag2dwd(diagram; clean = false, calc_states = [], out_vars=[], in_vars=[])
  homs = Vector{Any}(copy(diagram.hom_map))
  obs = Vector{Any}(copy(diagram.ob_map))
  graph = NamedGraph{Any, Any}()
  copy_parts!(graph, dom(diagram).graph)
  graph[:ename] .= homs

  # Expand homs which are composition
  for h in parts(graph, :E)
    h_name = graph[h, :ename]
    elsrc, eltgt = graph[h, :src], graph[h, :tgt]
    if h_name isa HomExpr{:compose}
      args = h_name.args
      rem_part!(graph, :E, h)
      verts = add_parts!(graph, :V, length(args) - 1, vname = [Symbol("anon_", gensym()) for i in 1:(length(args)-1)])
      append!(obs, codom.(args)[1:(end-1)])
      add_parts!(graph, :E, length(args), ename = args,
                          src = vcat([elsrc], verts),
                          tgt = vcat(verts, [eltgt]))
    end
  end

  flipped = fill(false, ne(graph))
  # Flip projections for graph eval
  for (i, h) in enumerate(homs)
    if isa_proj(h) && !flipped[i]
      iw = copy(incident(graph, graph[i, :src], :tgt))
      ow = copy(incident(graph, graph[i, :src], :src))
      if isempty(iw)
        for w in ow
          if isa_proj(homs[w]) && !flipped[w]
            wsrc = graph[w, :src]
            graph[w, :src] = graph[w, :tgt]
            graph[w, :tgt] = wsrc
            flipped[w] = true
          end
        end
      end
    end
  end

  pres = presentation(codom(diagram))

  # Decide on State and Parameter variables
  time_arrs = findall(h-> h isa HomExpr{:∂ₜ}, graph[:ename])
  params = findall(e -> isempty(incident(graph, e, :tgt)), parts(graph, :V))
  state_vals = graph[time_arrs, :src]

  # Mix State and Parameters into single in_els (need to ensure alignment between input and outputs
  in_els = unique(vcat(state_vals, params))
  out_els = copy(graph[time_arrs, :tgt])
  out_names = [Symbol(:∂ₜ, graph[v, :vname]) for v in state_vals]

  # TODO: Establish some kind of consistency expectations here
  if !isempty(in_vars) && !isempty(out_vars)
    in_els = [incident(graph, a, :vname) for a in in_vars]
    out_els = [incident(graph, a, :vname) for a in out_vars]
    out_names = graph[out_els, :vname]
  elseif !isempty(in_vars)
    # This ensures that the output values are in-order with input values
    new_inds = [findfirst(v -> graph[v, :vname] == a, in_els) for a in in_vars]
    non_param_inds = filter(i -> in_els[i] in state_vals , new_inds)
    in_els .= in_els[new_inds]
    out_els .= out_els[non_param_inds]
    out_names .= out_names[non_param_inds]
  elseif !isempty(out_vars)
    out_els = [incident(graph, a, :vname) for a in out_vars]
    out_names = graph[out_els, out_els]
  end
  if !isempty(calc_states)
    calc_inds = findall(v->graph[v, :vname] ∈ calc_states, state_vals)
    out_els = out_els[calc_inds]
    out_names = out_names[calc_inds]
  end

  in_types = [Dict(:name => graph[v, :vname],
                   :type => obs[v])
              for v in in_els]
  out_types = [Dict(:name => out_names[v],
                    :type => obs[out_els[v]])
               for v in 1:length(out_els)]

  # FIXME: Hacky solution to ensuring time_arrs aren't included later
  # need more thoughtful approach to sorting through edges here
  rem_parts!(graph, :E, time_arrs)

  dwd = WiringDiagram(collect(in_types), collect(out_types))

  # Add all necessary boxes to the DWD based on arrows in the diagram
  w2b = map(parts(graph, :E)) do a
    w_type = graph[a, :ename]
    # Fix this if-case. This was meant for times when multiple arguments would
    # go into a single box
    if isa_proj(w_type)
      nothing
    elseif w_type isa HomExpr
      src_el = graph[a, :src]
      tgt_el = graph[a, :tgt]
      in_ports = sp_otimes(obs[src_el])
      out_ports = sp_otimes(obs[tgt_el])
      add_box!(dwd, Box(name(w_type), [Dict(:type=>ip) for ip in in_ports],
                        [Dict(:type=>op) for op in out_ports]))
    end
  end

  # Mapping from elements (vertices in `graph`) to ports in the DWD
  # This ensures that elements which have already been computed are
  # not recomputed
  el2p = Dict{Int, Vector}()
  in_ps = vcat(map(out_els) do el
    eval_deps!(dwd, graph, el, w2b, el2p, obs; in_els=Dict(in_els[i] => i for i in 1:length(in_els)))
  end...)

  wires = map(enumerate(in_ps)) do (ip, op)
    Wire(out_types[ip], op, (output_id(dwd), ip))
  end
  add_wires!(dwd, wires)

  # Remove any boxes without any connecting wires
  if clean
    to_rem = Vector{Int64}()
    for b in 1:nparts(dwd.diagram, :Box)
      if isempty(vcat(incident(dwd.diagram, b, [:tgt, :in_port_box]),
                      incident(dwd.diagram, b, [:in_tgt, :in_port_box]))) &&
         isempty(vcat(incident(dwd.diagram, b, [:src, :out_port_box]),
                      incident(dwd.diagram, b, [:out_src, :out_port_box])))
         push!(to_rem, b)
      end
    end
    rem_boxes!(dwd, to_rem)
  end

  dwd
end

end


"""

Run Stan's diagnose binary on a model.

$(SIGNATURES)

### Required arguments
```julia
* `model`                    : SampleModel
```
"""
function diagnose(model::SampleModel)
  #local csvfile
  n_chains = model.n_chains[1]

  samplefiles = String[]
  for i in 1:n_chains
    push!(samplefiles, "$(model.output_base)_chain_$(i).csv")
  end
  try
    pstring = joinpath("$(model.cmdstan_home)", "bin", "diagnose")
    if Sys.iswindows()
      pstring = pstring * ".exe"
    end
    cmd = `$(pstring) $(StanBase.par(samplefiles))`
    resfile = open(cmd; read=true);
    print(read(resfile, String))
  catch e
    println(e)
  end

  return
end

export
    diagnose


draw(PNG("round-round.png", 140, 100),
     Axes(border=0.2),
     Paper("lightgrey"),
     Pen(0.3, cap="round", join="round")) do
    move(0, 0)
    line(0.5, 0)
    line(0.5, 0.5)
    line(1, 1)
end


@testset "evaluation.jl" begin

    @testset "SameInstancesEvaluator constructor" begin
        generator = SeaPearl.HomogenousGraphColoringGenerator(10, 0.1)

        my_heuristic(x::SeaPearl.IntVar; cpmodel=nothing) = minimum(x.domain)
        valueheuristic = SeaPearl.BasicHeuristic(my_heuristic)
        valueSelectionArray = [valueheuristic]
        evaluator = SeaPearl.SameInstancesEvaluator(valueSelectionArray,generator; seed=nothing, evalFreq = 50, nbInstances = 2)

        @test evaluator.nbInstances == 2
        @test evaluator.evalFreq == 50
        @test length(evaluator.instances) == 2
        @test length(values(evaluator.instances[1].variables)) == 11
        @test length(values(evaluator.instances[2].variables)) == 11
    end

    @testset "evaluate(::SameInstancesEvaluator)" begin

        generator = SeaPearl.HomogenousGraphColoringGenerator(10, 0.1)

        my_heuristic(x::SeaPearl.IntVar; cpmodel=nothing) = minimum(x.domain)
        valueheuristic = SeaPearl.BasicHeuristic(my_heuristic)
        valueSelectionArray = [valueheuristic]
        evaluator = SeaPearl.SameInstancesEvaluator(valueSelectionArray,generator; seed=1, evalFreq = 50, nbInstances = 2)

        variableheuristic = SeaPearl.MinDomainVariableSelection{false}()

        SeaPearl.evaluate(evaluator, variableheuristic, SeaPearl.DFSearch())
        @test evaluator.metrics[1,1].nodeVisited[1] == [12, 23]
    end
end


abstract type Galaxy end

struct Spiral <: Galaxy
	name::Symbol
	diskHeight::Number
	radius::Number
	units::Quantity
end


"""
    filter(Y, P)

"Filter" data `Y` with iterator `P`
calling (handling of `nothing`s omitted)
```
(t, y), state = iterate(Y, state)
(s => x) = dyniterate(P, s => x, (observation = t => y,))
```

"""
struct Filtered{T,S} <: DynamicIterator
    Y::S
    P::T
end
#filter(Y, P, ::Nothing) = Filtered(Y, P)
filter(Y, P) = Filtered(Y, P)

function dyniterate(M::Filtered, start::Start)
    v, q = @returnnothing dyniterate(M.Y, nothing)
    u, p = @returnnothing dyniterate(M.P, Condition(start, 0.0), v)
    u, (q, p)
end

function dyniterate(M::Filtered, (q, p)::Tuple)
    v, q = @returnnothing iterate(M.Y, q)
    u, p = @returnnothing dyniterate(M.P, p, v)
    u, (q, p)
end

#=
struct ControlledFilter{R,T,S} <: DynamicIterator
    C::R
    Y::S
    P::T
end
filter(Y, P; control = nothing) = filter(Y, P, control)
filter(Y, P, C) = ControlledFilter(C, Y, P)

function dyniterate(M::ControlledFilter, ::Nothing, (value,)::Value)
    w, r = @returnnothing dyniterate(M.C, nothing, (control = value,))
    v, q = @returnnothing iterate(M.Y)
    u, p = @returnnothing dyniterate(M.P, nothing, (value = value, control = w, observation = v))
    u, (u, r, q, p)
end

function iterate(M::ControlledFilter, (u, r, q, p)::Tuple)
    w, r = @returnnothing iterate(M.C, r, (control = u,))
    v, q = @returnnothing iterate(M.Y, q)
    u, p = @returnnothing dyniterate(M.P, p, (control = w, observation = v))
    u, (u, r, q, p)
end
=#


@testset "Tests: types" begin
    gate = QCO.GateData("H_1", 2)
    @test gate.isreal
    gate = QCO.GateData("S_3", 3)
    @test !gate.isreal

    # Add more tests for U3 and R gates.
end


__precompile__(true)

module UnitfulPlasmas

  import Unitful
  using Unitful: @unit

  import Unitful.DefaultSymbols

  @unit n20 "n20" PlasmaDensity 1e20/Unitful.m^3 false

  # Some gymnastics required here because if we precompile, we cannot add to
  # Unitful.basefactors at compile time and expect the changes to persist to runtime.

  const localunits = Unitful.basefactors

  function __init__()
      merge!(Unitful.basefactors, localunits)
      Unitful.register(UnitfulPlasmas)
  end

end


#Heatmap plots

import .LibCImPlot: PlotHeatmap

function PlotHeatmap(x::AbstractArray{T}, rows, cols, scale_min = 0.0, scale_max = 1.0;
                 label_id::String = "", label_fmt="%.1f", bounds_min = LibCImPlot.ImPlotPoint(0.0,0.0),
                 bounds_max = LibCImPlot.ImPlotPoint(1.0,1.0)) where {T<:ImPlotData}

    LibCImPlot.PlotHeatmap(label_id, x, rows, cols, scale_min, scale_max, label_fmt, bounds_min, bounds_max)
end


module Simhashes

    include("simhash.jl")
    include("simhash_index.jl")

    export Simhash, SimhashIndex
    export distance, get_near_dups, delete_from_index, add_to_index

end # module


@testset "Composition" begin
    encodings = (ProjectiveTransforms((5, 5)), ImagePreprocessing(), OneHot())
    blocks = (Image{2}(), Label(1:10))
    obs = (rand(RGB{N0f8}, 10, 10), 7)
    testencoding(encodings, blocks, obs)
end