text:
## BodyDict, JointDict
const BodyDict{V} = IndexDict{BodyID, Base.OneTo{BodyID}, V}
const BodyCacheDict{V} = CacheIndexDict{BodyID, Base.OneTo{BodyID}, V}
@propagate_inbounds Base.getindex(d::AbstractIndexDict{BodyID}, key::RigidBody) = d[BodyID(key)]
@propagate_inbounds Base.setindex!(d::AbstractIndexDict{BodyID}, value, key::RigidBody) = d[BodyID(key)] = value

const JointDict{V} = IndexDict{JointID, Base.OneTo{JointID}, V}
const JointCacheDict{V} = CacheIndexDict{JointID, Base.OneTo{JointID}, V}
@propagate_inbounds Base.getindex(d::AbstractIndexDict{JointID}, key::Joint) = d[JointID(key)]
@propagate_inbounds Base.setindex!(d::AbstractIndexDict{JointID}, value, key::Joint) = d[JointID(key)] = value

## SegmentedVector method overloads
@propagate_inbounds Base.getindex(v::SegmentedVector{JointID}, joint::Joint) = v[JointID(joint)]
@propagate_inbounds Base.view(v::SegmentedVector{JointID}, joint::Joint) = Base.view(v, JointID(joint))
function SegmentedVector(parent::AbstractVector{T}, joints::AbstractVector{<:Joint}, viewlengthfun) where T
    SegmentedVector{JointID, T, Base.OneTo{JointID}}(parent, joints, viewlengthfun)
end

"""
$(TYPEDEF)

A `MechanismState` stores state information for an entire `Mechanism`. It
contains the joint configuration and velocity vectors ``q`` and ``v``, and
a vector of additional states ``s``. In addition, it stores cache
variables that depend on ``q`` and ``v`` and are aimed at preventing double work.

Type parameters:
* `X`: the scalar type of the ``q``, ``v``, and ``s`` vectors.
* `M`: the scalar type of the `Mechanism`
* `C`: the scalar type of the cache variables (`== promote_type(X, M)`)
* `JointCollection`: the type of the `treejoints` and `nontreejoints` members (a `TypeSortedCollection` subtype)
"""
struct MechanismState{X, M, C, JointCollection}
    modcount::Int

    # mechanism layout
    mechanism::Mechanism{M}
    nonrootbodies::Vector{RigidBody{M}} # TODO: remove once https://github.com/JuliaLang/julia/issues/14955 is fixed
    treejoints::JointCollection
    nontreejoints::JointCollection
    jointids::Base.OneTo{JointID}
    treejointids::Base.OneTo{JointID}
    nontreejointids::UnitRange{JointID}
    predecessor_and_successor_ids::JointDict{Pair{BodyID, BodyID}}
    qranges::JointDict{UnitRange{Int}}
    vranges::JointDict{UnitRange{Int}}
    constraintranges::IndexDict{JointID, UnitRange{JointID}, UnitRange{Int}}
    support_set_masks::BodyDict{JointDict{Bool}} # TODO: use a Matrix-backed type
    constraint_jacobian_structure::JointDict{TreePath{RigidBody{M}, Joint{M}}} # TODO: use a Matrix-backed type
    q_index_to_joint_id::Vector{JointID}
    v_index_to_joint_id::Vector{JointID}

    # minimal representation of state
    q::SegmentedVector{JointID, X, Base.OneTo{JointID}, Vector{X}} # configurations
    v::SegmentedVector{JointID, X, Base.OneTo{JointID}, Vector{X}} # velocities
    s::Vector{X} # additional state

    # joint-related cache
    joint_transforms::JointCacheDict{Transform3D{C}}
    joint_twists::JointCacheDict{Twist{C}}
    joint_bias_accelerations::JointCacheDict{SpatialAcceleration{C}}
    tree_joint_transforms::VectorSegment{Transform3D{C}}
    non_tree_joint_transforms::VectorSegment{Transform3D{C}}

    #
    motion_subspaces::CacheElement{Vector{GeometricJacobian{SMatrix{3, 1, C, 3}}}} # TODO: use CacheIndexDict?
    constraint_wrench_subspaces::CacheElement{Vector{WrenchMatrix{SMatrix{3, 1, C, 3}}}} # TODO: use CacheIndexDict?

    # body-related cache
    transforms_to_root::BodyCacheDict{Transform3D{C}}
    twists_wrt_world::BodyCacheDict{Twist{C}}
    bias_accelerations_wrt_world::BodyCacheDict{SpatialAcceleration{C}}
    inertias::BodyCacheDict{SpatialInertia{C}}
    crb_inertias::BodyCacheDict{SpatialInertia{C}}
    contact_states::BodyCacheDict{Vector{Vector{DefaultSoftContactState{X}}}} # TODO: consider moving to separate type

    function MechanismState{X, M, C, JointCollection}(m::Mechanism{M}, q::Vector{X}, v::Vector{X}, s::Vector{X}) where {X, M, C, JointCollection}
        @assert length(q) == num_positions(m)
        @assert length(v) == num_velocities(m)
        @assert length(s) == num_additional_states(m)

        # mechanism layout
        nonrootbodies = collect(non_root_bodies(m))
        treejoints = JointCollection(tree_joints(m))
        nontreejoints = JointCollection(non_tree_joints(m))
        lastjointid = isempty(joints(m)) ? JointID(0) : JointID(last(joints(m)))
        lasttreejointid = isempty(tree_joints(m)) ? JointID(0) : JointID(last(tree_joints(m)))
        jointids = Base.OneTo(lastjointid)
        treejointids = Base.OneTo(lasttreejointid)
        nontreejointids = lasttreejointid + 1 : lastjointid
        predecessor_and_successor_ids = JointDict{Pair{BodyID, BodyID}}(
            JointID(j) => (BodyID(predecessor(j, m)) => BodyID(successor(j, m))) for j in joints(m))
        support_set = function (b::RigidBody)
            JointDict{Bool}(JointID(j) => (j ∈ tree_joints(m) ?  j ∈ path(m, b, root_body(m)) : false) for j in joints(m))
        end
        support_set_masks = BodyDict{JointDict{Bool}}(b => support_set(b) for b in bodies(m))
        constraint_jacobian_structure = JointDict{TreePath{RigidBody{M}, Joint{M}}}(
            JointID(j) => path(m, predecessor(j, m), successor(j, m)) for j in joints(m))
        qsegmented = SegmentedVector(q, tree_joints(m), num_positions)
        vsegmented = SegmentedVector(v, tree_joints(m), num_velocities)
        qranges = ranges(qsegmented)
        vranges = ranges(vsegmented)
        constraintranges = let start = 1
            rangevec = UnitRange{Int}[start : (start += num_constraints(j)) - 1 for j in non_tree_joints(m)]
            IndexDict(nontreejointids, rangevec)
        end
        q_index_to_joint_id = Vector{JointID}(undef, length(q))
        v_index_to_joint_id = Vector{JointID}(undef, length(v))
        for jointid in treejointids
            for qindex in qranges[jointid]
                q_index_to_joint_id[qindex] = jointid
            end
            for vindex in vranges[jointid]
                v_index_to_joint_id[vindex] = jointid
            end
        end

        # joint-related cache
        joint_transforms = JointCacheDict{Transform3D{C}}(jointids)
        joint_twists = JointCacheDict{Twist{C}}(treejointids)
        joint_bias_accelerations = JointCacheDict{SpatialAcceleration{C}}(treejointids)
        tree_joint_transforms = view(values(joint_transforms), 1 : Int(lasttreejointid))
        non_tree_joint_transforms = view(values(joint_transforms), Int(lasttreejointid) + 1 : Int(lastjointid))

        # motion subspaces, constraint wrench subspaces
        motion_subspaces = CacheElement(Vector{GeometricJacobian{SMatrix{3, 1, C, 3}}}(undef, num_velocities(m)))
        constraint_wrench_subspaces = CacheElement(Vector{WrenchMatrix{SMatrix{3, 1, C, 3}}}(undef, num_constraints(m)))

        # body-related cache
        bodyids = Base.OneTo(BodyID(last(bodies(m))))
        transforms_to_root = BodyCacheDict{Transform3D{C}}(bodyids)
        twists_wrt_world = BodyCacheDict{Twist{C}}(bodyids)
        bias_accelerations_wrt_world = BodyCacheDict{SpatialAcceleration{C}}(bodyids)
        inertias = BodyCacheDict{SpatialInertia{C}}(bodyids)
        crb_inertias = BodyCacheDict{SpatialInertia{C}}(bodyids)

        # contact. TODO: move out of MechanismState:
        contact_states = BodyCacheDict{Vector{Vector{DefaultSoftContactState{X}}}}(
            BodyID(b) => Vector{Vector{DefaultSoftContactState{X}}}() for b in bodies(m))
        startind = 1
        for body in bodies(m), point in contact_points(body)
            model = contact_model(point)
            n = num_states(model)
            push!(contact_states[body], collect(begin
                s_part = view(s, startind : startind + n - 1)
                contact_state = SoftContactState(model, s_part, root_frame(m))
                startind += n
                contact_state
            end for j = 1 : length(m.environment)))
        end

        # initialize state-independent cache data for root bodies:
        root = root_body(m)
        rootframe = default_frame(root)
        transforms_to_root[root] = one(Transform3D{C}, rootframe)
        twists_wrt_world[root] = zero(Twist{C}, rootframe, rootframe, rootframe)
        bias_accelerations_wrt_world[root] = zero(SpatialAcceleration{C}, rootframe, rootframe, rootframe)
        inertias[root] = zero(SpatialInertia{C}, rootframe)

        new{X, M, C, JointCollection}(
            modcount(m), m, nonrootbodies, treejoints, nontreejoints,
            jointids, treejointids, nontreejointids,
            predecessor_and_successor_ids, qranges, vranges, constraintranges, support_set_masks, constraint_jacobian_structure,
            q_index_to_joint_id, v_index_to_joint_id,
            qsegmented, vsegmented, s,
            joint_transforms, joint_twists, joint_bias_accelerations, tree_joint_transforms, non_tree_joint_transforms,
            motion_subspaces, constraint_wrench_subspaces,
            transforms_to_root, twists_wrt_world, bias_accelerations_wrt_world, inertias, crb_inertias,
            contact_states)
    end

    function MechanismState{X, M, C}(mechanism::Mechanism{M}, q::Vector{X}, v::Vector{X}, s::Vector{X}) where {X, M, C}
        JointCollection = typeof(TypeSortedCollection(joints(mechanism)))
        MechanismState{X, M, C, JointCollection}(mechanism, q, v, s)
    end

    function MechanismState{X, M}(mechanism::Mechanism{M}, q::Vector{X}, v::Vector{X}, s::Vector{X}) where {X, M}
        C = promote_type(X, M)
        MechanismState{X, M, C}(mechanism, q, v, s)
    end

    function MechanismState{X}(mechanism::Mechanism{M}) where {X, M}
        q = Vector{X}(undef, num_positions(mechanism))
        v = Vector{X}(undef, num_velocities(mechanism))
        s = Vector{X}(undef, num_additional_states(mechanism))
        state = MechanismState{X, M}(mechanism, q, v, s)
        zero!(state)
        state
    end
end

function MechanismState(mechanism::Mechanism{M}, q::Vector{X}, v::Vector{X}, s=zeros(X, num_additional_states(mechanism))) where {X, M}
    MechanismState{X, M}(mechanism, q, v, s)
end

MechanismState(mechanism::Mechanism{M}) where {M} = MechanismState{M}(mechanism)

Base.broadcastable(x::MechanismState) = Ref(x)

Base.show(io::IO, ::MechanismState{X, M, C}) where {X, M, C} = print(io, "MechanismState{$X, $M, $C, …}(…)")

modcount(state::MechanismState) = state.modcount
@propagate_inbounds predsucc(id::JointID, state::MechanismState) = state.predecessor_and_successor_ids[id]
@propagate_inbounds successorid(id::JointID, state::MechanismState) = last(state.predecessor_and_successor_ids[id])

"""
$(SIGNATURES)

Return the length of the joint configuration vector ``q``.
"""
num_positions(state::MechanismState) = length(state.q)

"""
$(SIGNATURES)

Return the length of the joint velocity vector ``v``.
"""
num_velocities(state::MechanismState) = length(state.v)

"""
$(SIGNATURES)

Return the length of the vector of additional states ``s`` (currently used
for stateful contact models).
"""
num_additional_states(state::MechanismState) = length(state.s)

state_vector_eltype(state::MechanismState{X, M, C}) where {X, M, C} = X
mechanism_eltype(state::MechanismState{X, M, C}) where {X, M, C} = M
cache_eltype(state::MechanismState{X, M, C}) where {X, M, C} = C

"""
$(SIGNATURES)

Return the part of the configuration vector ``q`` associated with `joint`.
"""
configuration(state::MechanismState, joint::Union{<:Joint, JointID}) = state.q[joint]

"""
$(SIGNATURES)

Return the part of the velocity vector ``v`` associated with `joint`.
"""
velocity(state::MechanismState, joint::Union{<:Joint, JointID}) = state.v[joint]

"""
$(SIGNATURES)

Invalidate all cache variables.
"""
function setdirty!(state::MechanismState)
    CustomCollections.setdirty!(state.joint_transforms)
    CustomCollections.setdirty!(state.joint_twists)
    CustomCollections.setdirty!(state.joint_bias_accelerations)
    CustomCollections.setdirty!(state.motion_subspaces)
    CustomCollections.setdirty!(state.constraint_wrench_subspaces)
    CustomCollections.setdirty!(state.transforms_to_root)
    CustomCollections.setdirty!(state.twists_wrt_world)
    CustomCollections.setdirty!(state.bias_accelerations_wrt_world)
    CustomCollections.setdirty!(state.inertias)
    CustomCollections.setdirty!(state.crb_inertias)
    nothing
end

"""
$(SIGNATURES)

Reset all contact state variables.
"""
function reset_contact_state!(state::MechanismState)
    for states_for_body in values(state.contact_states)
        for states_for_point in states_for_body
            for contact_state in states_for_point
                reset!(contact_state)
            end
        end
    end
end

"""
$(SIGNATURES)

'Zero' the configuration vector ``q``. Invalidates cache variables.

Note that when the `Mechanism` contains e.g. quaternion-parameterized joints,
``q`` may not actually be set to all zeros; the quaternion part of the
configuration vector would be set to identity. The contract is that each of the
joint transforms should be an identity transform.
"""
function zero_configuration!(state::MechanismState)
    foreach(state.treejoints, values(segments(state.q))) do joint, qjoint
        zero_configuration!(qjoint, joint)
    end
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Zero the velocity vector ``v``. Invalidates cache variables.
"""
function zero_velocity!(state::MechanismState)
    state.v .= 0
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Zero both the configuration and velocity. Invalidates cache variables.

See [`zero_configuration!`](@ref), [`zero_velocity!`](@ref).
"""
zero!(state::MechanismState) = (zero_configuration!(state); zero_velocity!(state))

"""
$(SIGNATURES)

Randomize the configuration vector ``q``. The distribution depends on
the particular joint types present in the associated `Mechanism`. The resulting
``q`` is guaranteed to be on the `Mechanism`'s configuration manifold.
Invalidates cache variables.
"""
function rand_configuration!(state::MechanismState)
    foreach(state.treejoints, values(segments(state.q))) do joint, qjoint
        rand_configuration!(qjoint, joint)
    end
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Randomize the velocity vector ``v``.
Invalidates cache variables.
"""
function rand_velocity!(state::MechanismState)
    rand!(state.v)
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Randomize both the configuration and velocity.
Invalidates cache variables.
"""
Random.rand!(state::MechanismState) = begin rand_configuration!(state); rand_velocity!(state) end

"""
$(SIGNATURES)

Return the configuration vector ``q``.

Note that this returns a reference to the underlying data in `state`. The user
is responsible for calling [`setdirty!`](@ref) after modifying this vector to
ensure that dependent cache variables are invalidated.
"""
configuration(state::MechanismState) = state.q

"""
$(SIGNATURES)

Return the velocity vector ``v``.

Note that this function returns a read-write reference to a field in `state`.
The user is responsible for calling [`setdirty!`](@ref) after modifying this
vector to ensure that dependent cache variables are invalidated.
"""
velocity(state::MechanismState) = state.v

"""
$(SIGNATURES)

Return the vector of additional states ``s``.
"""
additional_state(state::MechanismState) = state.s

for fun in (:num_velocities, :num_positions)
    @eval function $fun(p::TreePath{RigidBody{T}, <:Joint{T}} where {T})
        mapreduce($fun, +, p; init=0)
    end
end

"""
$(SIGNATURES)

Set the part of the configuration vector associated with `joint`.
Invalidates cache variables.
"""
function set_configuration!(state::MechanismState, joint::Joint, config)
    set_configuration!(configuration(state, joint), joint, config)
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Set the part of the velocity vector associated with `joint`.
Invalidates cache variables.
"""
function set_velocity!(state::MechanismState, joint::Joint, vel)
    set_velocity!(velocity(state, joint), joint, vel)
    reset_contact_state!(state)
    setdirty!(state)
end

"""
$(SIGNATURES)

Set the configuration vector ``q``. Invalidates cache variables.
"""
function set_configuration!(state::MechanismState, q::AbstractVector)
    copyto!(state.q, q)
    setdirty!(state)
end

"""
$(SIGNATURES)

Set the velocity vector ``v``. Invalidates cache variables.
"""
function set_velocity!(state::MechanismState, v::AbstractVector)
    copyto!(state.v, v)
    setdirty!(state)
end

"""
$(SIGNATURES)

Set the vector of additional states ``s``.
"""
function set_additional_state!(state::MechanismState, s::AbstractVector)
    copyto!(state.s, s)
    # note: setdirty! is currently not needed because no cache variables depend on s
end

"""
$(SIGNATURES)

Copy (minimal representation of) state `src` to state `dest`.
"""
function Base.copyto!(dest::MechanismState, src::MechanismState)
    dest.mechanism == src.mechanism || throw(ArgumentError("States are not associated with the same Mechanism."))
    @modcountcheck dest src
    copyto!(dest.q, src.q)
    copyto!(dest.v, src.v)
    copyto!(dest.s, src.s)
    setdirty!(dest)
    dest
end

"""
$(SIGNATURES)

Copy state information in vector `src` (ordered `[q; v; s]`) to state `dest`.
"""
function Base.copyto!(dest::MechanismState, src::AbstractVector)
    nq = num_positions(dest)
    nv = num_velocities(dest)
    ns = num_additional_states(dest)
    @boundscheck length(src) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(parent(dest.q), 1, src, 1, nq)
    @inbounds copyto!(parent(dest.v), 1, src, nq + 1, nv)
    @inbounds copyto!(dest.s, 1, src, nq + nv + 1, ns)
    setdirty!(dest)
    dest
end

"""
$(SIGNATURES)

Copy state information in state `dest` to vector `src` (ordered `[q; v; s]`).
"""
function Base.copyto!(dest::AbstractVector, src::MechanismState)
    nq = num_positions(src)
    nv = num_velocities(src)
    ns = num_additional_states(src)
    length(dest) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(dest, 1, src.q, 1, nq)
    @inbounds copyto!(dest, nq + 1, src.v, 1, nv)
    @inbounds copyto!(dest, nq + nv + 1, src.s, 1, ns)
    dest
end

"""
$(SIGNATURES)

Create a `Vector` that represents the same state as `state` (ordered `[q; v; s]`).
"""
function Base.Vector{T}(state::MechanismState) where T
    dest = Vector{T}(undef, num_positions(state) + num_velocities(state) + num_additional_states(state))
    copyto!(dest, state)
end

Base.Vector(state::MechanismState{X}) where {X} = Vector{X}(state)
Base.Array{T}(state::MechanismState) where {T} = Vector{T}(state)
Base.Array(state::MechanismState{X}) where {X} = Array{X}(state)
Base.convert(::Type{T}, state::MechanismState) where {T<:Array} = T(state)

"""
$(SIGNATURES)

Project the configuration vector ``q`` onto the configuration manifold.

For example:
* for a part of ``q`` corresponding to a revolute joint, this method is a no-op;
* for a part of ``q`` corresponding to a spherical joint that uses a unit quaternion
to parameterize the orientation of its successor with respect to its predecessor,
`normalize_configuration!` will renormalize the quaternion so that it is indeed
of unit length.

!!! warning
This method does not ensure that the configuration or velocity satisfy joint
configuration or velocity limits/bounds.
"""
function normalize_configuration!(state::MechanismState)
    foreach(state.treejoints, values(segments(state.q))) do joint, qjoint
        normalize_configuration!(qjoint, joint)
    end
end

"""
$(SIGNATURES)

Applies the principal_value functions from [Rotations.jl](https://github.com/FugroRoames/Rotations.jl/blob/d080990517f89b56c37962ad53a7fd24bd94b9f7/src/principal_value.jl)
to joint angles. This currently only applies to `SPQuatFloating` joints.

For example:
* for a part of ``q`` corresponding to a revolute joint, this method is a no-op;
* for a part of ``q`` corresponding to a `SPQuatFloating` joint this function applies
`principal_value the orientation.
"""
function principal_value!(state::MechanismState)
    foreach(state.treejoints, values(segments(state.q))) do joint, qjoint
        principal_value!(qjoint, joint)
    end
end

"""
$(SIGNATURES)

Return the range of indices into the joint configuration vector ``q`` corresponding to joint `joint`.
"""
@propagate_inbounds configuration_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.qranges[joint]

"""
$(SIGNATURES)

Return the range of indices into the joint velocity vector ``v`` corresponding to joint `joint`.
"""
@propagate_inbounds velocity_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.vranges[joint]

"""
$(SIGNATURES)

Return the `JointID` of the joint associated with the given index into the configuration vector ``q``.
"""
@propagate_inbounds configuration_index_to_joint_id(state::MechanismState, qindex::Integer) = state.q_index_to_joint_id[qindex]

"""
$(SIGNATURES)

Return the `JointID` of the joint associated with the given index into the velocity vector ``v``.
"""
@propagate_inbounds velocity_index_to_joint_id(state::MechanismState, qindex::Integer) = state.v_index_to_joint_id[qindex]

"""
$(SIGNATURES)

Return the range of row indices into the constraint Jacobian corresponding to joint `joint`.
"""
@propagate_inbounds constraint_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.constraintranges[joint]

"""
$(SIGNATURES)

Return whether `joint` supports `body`, i.e., `joint` is a tree joint on the path between `body` and the root.
"""
@propagate_inbounds function supports(joint::Union{<:Joint, JointID}, body::Union{<:RigidBody, BodyID}, state::MechanismState)
    state.support_set_masks[body][joint]
end

## Accessor functions for cached variables
"""
$(SIGNATURES)

Return the joint transform for the given joint, i.e. the transform from
`frame_after(joint)` to `frame_before(joint)`.
"""
@propagate_inbounds function joint_transform(state::MechanismState, joint::Union{<:Joint, JointID}, safe=true)
    safe && update_transforms!(state)
    state.joint_transforms[joint]
end

"""
$(SIGNATURES)
Return the joint twist for the given joint, i.e. the twist of
`frame_after(joint)` with respect to `frame_before(joint)`, expressed in the
root frame of the mechanism.
"""
@propagate_inbounds function twist(state::MechanismState, joint::Union{<:Joint, JointID}, safe=true)
    safe && update_joint_twists!(state)
    state.joint_twists[joint]
end

"""
$(SIGNATURES)

Return the bias acceleration across the given joint, i.e. the spatial acceleration
of `frame_after(joint)` with respect to `frame_before(joint)`, expressed in the
root frame of the mechanism when all joint accelerations are zero.
"""
@propagate_inbounds function bias_acceleration(state::MechanismState, joint::Union{<:Joint, JointID}, safe=true)
    safe && update_joint_bias_accelerations!(state)
    state.joint_bias_accelerations[joint]
end

"""
$(SIGNATURES)

Return the transform from `default_frame(body)` to the root frame of the
mechanism.
"""
@propagate_inbounds function transform_to_root(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    safe && update_transforms!(state)
    state.transforms_to_root[body]
end

"""
$(SIGNATURES)

Return the twist of `default_frame(body)` with respect to the root frame of the
mechanism, expressed in the root frame.
"""
@propagate_inbounds function twist_wrt_world(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    safe && update_twists_wrt_world!(state)
    state.twists_wrt_world[body]
end

"""
$(SIGNATURES)

Return the bias acceleration of the given body with respect to the world,
i.e. the spatial acceleration of `default_frame(body)` with respect to the
root frame of the mechanism, expressed in the root frame, when all joint
accelerations are zero.
"""
@propagate_inbounds function bias_acceleration(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    safe && update_bias_accelerations_wrt_world!(state)
    state.bias_accelerations_wrt_world[body]
end

"""
$(SIGNATURES)

Return the spatial inertia of `body` expressed in the root frame of the
mechanism.
"""
@propagate_inbounds function spatial_inertia(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    safe && update_spatial_inertias!(state)
    state.inertias[body]
end

"""
$(SIGNATURES)

Return the composite rigid body inertia `body` expressed in the root frame of the
mechanism.
"""
@propagate_inbounds function crb_inertia(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    safe && update_crb_inertias!(state)
    state.crb_inertias[body]
end

# Cache variable update functions
@inline update_transforms!(state::MechanismState) = isdirty(state.transforms_to_root) && _update_transforms!(state)
@noinline function _update_transforms!(state::MechanismState)
    @modcountcheck state state.mechanism

    # update tree joint transforms
    qs = values(segments(state.q))
    @inbounds map!(joint_transform, state.tree_joint_transforms, state.treejoints, qs)

    # update transforms to root
    transforms_to_root = state.transforms_to_root
    for joint in tree_joints(state.mechanism)
        jointid = JointID(joint)
        parentid, bodyid = predsucc(jointid, state)
        transforms_to_root[bodyid] = transforms_to_root[parentid] * joint_to_predecessor(joint) * state.joint_transforms[jointid]
    end
    state.transforms_to_root.dirty = false

    # update non-tree joint transforms
    if !isempty(state.nontreejointids)
        broadcast!(state.non_tree_joint_transforms, state, state.nontreejoints) do state, joint
            predid, succid = predsucc(JointID(joint), state)
            before_to_root = state.transforms_to_root[predid] * joint_to_predecessor(joint)
            after_to_root = state.transforms_to_root[succid] * joint_to_successor(joint)
            inv(before_to_root) * after_to_root
        end
    end
    state.joint_transforms.dirty = false
    nothing
end

@inline function update_joint_twists!(state::MechanismState)
    isdirty(state.joint_twists) && _update_joint_twists!(state)
    nothing
end
@noinline function _update_joint_twists!(state::MechanismState)
    qs = values(segments(state.q))
    vs = values(segments(state.v))
    @inbounds map!(joint_twist, values(state.joint_twists), state.treejoints, qs, vs)
    state.joint_twists.dirty = false
    nothing
end

@inline function update_joint_bias_accelerations!(state::MechanismState)
    isdirty(state.joint_bias_accelerations) && _update_joint_bias_accelerations!(state)
    nothing
end
@noinline function _update_joint_bias_accelerations!(state::MechanismState)
    qs = values(segments(state.q))
    vs = values(segments(state.v))
    @inbounds map!(bias_acceleration, values(state.joint_bias_accelerations), state.treejoints, qs, vs)
    state.joint_bias_accelerations.dirty = false
    nothing
end

@inline function update_motion_subspaces!(state::MechanismState)
    isdirty(state.motion_subspaces) && _update_motion_subspaces!(state)
    nothing
end
@inline function _motion_subspace(state::MechanismState, joint::Joint, qjoint::AbstractVector)
    jointid = JointID(joint)
    bodyid = successorid(jointid, state)
    transform(motion_subspace(joint, qjoint), transform_to_root(state, bodyid, false))
end
@noinline function _update_motion_subspaces!(state::MechanismState)
    update_transforms!(state)
    qs = values(segments(state.q))
    foreach_with_extra_args(state, state.treejoints, qs) do state, joint, qjoint
        S = _motion_subspace(state, joint, qjoint)
        @inbounds vrange = velocity_range(state, joint)
        @inbounds for col = Base.OneTo(size(S, 2))
            Scol = GeometricJacobian(S.body, S.base, S.frame, SMatrix{3, 1}(S.angular[:, col]), SMatrix{3, 1}(S.linear[:, col]))
            vindex = vrange[col]
            state.motion_subspaces.data[vindex] = Scol
        end
    end
    state.motion_subspaces.dirty = false
    nothing
end

@inline function update_twists_wrt_world!(state::MechanismState)
    isdirty(state.twists_wrt_world) && _update_twists_wrt_world!(state)
    nothing
end
@noinline function _update_twists_wrt_world!(state::MechanismState)
    update_transforms!(state)
    update_joint_twists!(state)
    twists = state.twists_wrt_world
    @inbounds for jointid in state.treejointids
        parentbodyid, bodyid = predsucc(jointid, state)
        jointtwist = state.joint_twists[jointid]
        twists[bodyid] = twists[parentbodyid] + transform(jointtwist, transform_to_root(state, bodyid, false))
    end
    state.twists_wrt_world.dirty = false
    nothing
end

@inline function update_constraint_wrench_subspaces!(state::MechanismState)
    isdirty(state.constraint_wrench_subspaces) && _update_constraint_wrench_subspaces!(state)
    nothing
end
@inline function _constraint_wrench_subspace(state::MechanismState, joint::Joint)
    jointid = JointID(joint)
    tf = state.joint_transforms[jointid]
    T = constraint_wrench_subspace(joint, tf)
    bodyid = successorid(jointid, state)
    toroot = state.transforms_to_root[bodyid] * joint_to_successor(joint)
    transform(T, toroot)
end

@noinline function _update_constraint_wrench_subspaces!(state::MechanismState)
    update_transforms!(state)
    foreach_with_extra_args(state, state.nontreejoints) do state, joint
        T = _constraint_wrench_subspace(state, joint)
        @inbounds crange = constraint_range(state, joint)
        @inbounds for col = Base.OneTo(size(T, 2))
            Tcol = WrenchMatrix(T.frame, SMatrix{3, 1}(T.angular[:, col]), SMatrix{3, 1}(T.linear[:, col]))
            cindex = crange[col]
            state.constraint_wrench_subspaces.data[cindex] = Tcol
        end
    end
    state.constraint_wrench_subspaces.dirty = false
    nothing
end

@inline function update_bias_accelerations_wrt_world!(state::MechanismState)
    isdirty(state.bias_accelerations_wrt_world) && _update_bias_accelerations_wrt_world!(state)
    nothing
end
@noinline function _update_bias_accelerations_wrt_world!(state::MechanismState) # TODO: make more efficient
    update_transforms!(state)
    update_twists_wrt_world!(state)
    update_joint_bias_accelerations!(state)
    biasaccels = state.bias_accelerations_wrt_world
    for jointid in state.treejointids
        parentbodyid, bodyid = predsucc(jointid, state)
        jointbias = bias_acceleration(state, jointid)
        # TODO: awkward way of doing this:
        toroot = transform_to_root(state, bodyid, false)
        twistwrtworld = transform(twist_wrt_world(state, bodyid), inv(toroot))
        jointtwist = twist(state, jointid)
        biasaccels[bodyid] = biasaccels[parentbodyid] + transform(jointbias, toroot, twistwrtworld, jointtwist)
    end
    state.bias_accelerations_wrt_world.dirty = false
    nothing
end

@inline function update_spatial_inertias!(state::MechanismState)
    isdirty(state.inertias) && _update_spatial_inertias!(state)
    nothing
end
@noinline function _update_spatial_inertias!(state::MechanismState)
    update_transforms!(state)
    mechanism = state.mechanism
    inertias = state.inertias
    @inbounds for body in state.nonrootbodies
        bodyid = BodyID(body)
        inertias[bodyid] = transform(spatial_inertia(body), transform_to_root(state, bodyid, false))
    end
    state.inertias.dirty = false
    nothing
end

@inline function update_crb_inertias!(state::MechanismState)
    isdirty(state.crb_inertias) && _update_crb_inertias!(state)
    nothing
end
@noinline function _update_crb_inertias!(state::MechanismState)
    update_spatial_inertias!(state)
    mechanism = state.mechanism
    crb_inertias = state.crb_inertias
    rootbodyid = BodyID(root_body(mechanism))
    crb_inertias[rootbodyid] = state.inertias[rootbodyid]
    for jointid in state.treejointids
        bodyid = successorid(jointid, state)
        crb_inertias[bodyid] = state.inertias[bodyid]
    end
    for jointid in reverse(state.treejointids)# TODO: reverse is kind of expensive
        parentbodyid, bodyid = predsucc(jointid, state)
        crb_inertias[parentbodyid] += crb_inertias[bodyid]
    end
    state.crb_inertias.dirty = false
    nothing
end

contact_states(state::MechanismState, body::Union{<:RigidBody, BodyID}) = state.contact_states[body]

@propagate_inbounds function newton_euler(state::MechanismState, body::Union{<:RigidBody, BodyID}, accel::SpatialAcceleration, safe=true)
    inertia = spatial_inertia(state, body, safe)
    twist = twist_wrt_world(state, body, safe)
    newton_euler(inertia, accel, twist)
end

@propagate_inbounds function momentum(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    spatial_inertia(state, body, safe) * twist_wrt_world(state, body, safe)
end

@propagate_inbounds function momentum_rate_bias(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    newton_euler(state, body, bias_acceleration(state, body, safe), safe)
end

@propagate_inbounds function kinetic_energy(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    kinetic_energy(spatial_inertia(state, body, safe), twist_wrt_world(state, body, safe))
end

"""
$(SIGNATURES)

Return the gravitational potential energy in the given state, computed as the
negation of the dot product of the gravitational force and the center
of mass expressed in the `Mechanism`'s root frame.
"""
@propagate_inbounds function gravitational_potential_energy(state::MechanismState, body::Union{<:RigidBody, BodyID}, safe=true)
    inertia = spatial_inertia(body)
    m = inertia.mass
    m > 0 || return zero(cache_eltype(state))
    com = transform_to_root(state, body, safe) * center_of_mass(inertia)
    -m * dot(state.mechanism.gravitational_acceleration, FreeVector3D(com))
end

function configuration_derivative!(q̇::SegmentedVector{JointID}, state::MechanismState)
    q̇s = values(segments(q̇))
    qs = values(segments(state.q))
    vs = values(segments(state.v))
    foreach((joint, q̇, q, v) -> velocity_to_configuration_derivative!(q̇, joint, q, v), state.treejoints, q̇s, qs, vs)
end

function configuration_derivative_to_velocity_adjoint!(
        fq::SegmentedVector{JointID}, state::MechanismState, fv::SegmentedVector{JointID})
    fqs = values(segments(fq))
    qs = values(segments(state.q))
    fvs = values(segments(fv))
    foreach((joint, fq, q, fv) -> configuration_derivative_to_velocity_adjoint!(fq, joint, q, fv), state.treejoints, fqs, qs, fvs)
end

function configuration_derivative(state::MechanismState{X}) where {X}
    ret = SegmentedVector(Vector{X}(undef, num_positions(state.mechanism)), tree_joints(state.mechanism), num_positions)
    configuration_derivative!(ret, state)
    ret
end

function velocity_to_configuration_derivative_jacobian!(J::SegmentedBlockDiagonalMatrix, state::MechanismState)
    qs = values(segments(state.q))
    foreach(state.treejoints, qs, CustomCollections.blocks(J)) do joint, qjoint, block
        copyto!(block, velocity_to_configuration_derivative_jacobian(joint, qjoint))
    end
    nothing
end

function velocity_to_configuration_derivative_jacobian(state::MechanismState{X, M, C}) where {X, M, C}
    q_ranges = values(ranges(configuration(state)))
    v_ranges = values(ranges(velocity(state)))
    J = SegmentedBlockDiagonalMatrix(zeros(C, num_positions(state), num_velocities(state)), zip(q_ranges, v_ranges))
    velocity_to_configuration_derivative_jacobian!(J, state)
    J
end

function configuration_derivative_to_velocity_jacobian!(J::SegmentedBlockDiagonalMatrix, state::MechanismState)
    qs = values(segments(state.q))
    foreach(state.treejoints, qs, CustomCollections.blocks(J)) do joint, qjoint, block
        copyto!(block, configuration_derivative_to_velocity_jacobian(joint, qjoint))
    end
    nothing
end

function configuration_derivative_to_velocity_jacobian(state::MechanismState{X, M, C}) where {X, M, C}
    q_ranges = values(ranges(configuration(state)))
    v_ranges = values(ranges(velocity(state)))
    J = SegmentedBlockDiagonalMatrix(zeros(C, num_velocities(state), num_positions(state)), zip(v_ranges, q_ranges))
    configuration_derivative_to_velocity_jacobian!(J, state)
    J
end

function transform_to_root(state::MechanismState, frame::CartesianFrame3D, safe=true)
    body = body_fixed_frame_to_body(state.mechanism, frame) # FIXME: expensive
    tf = transform_to_root(state, body, safe)
    if tf.from != frame
        tf = tf * body_fixed_frame_definition(state.mechanism, frame) # TODO: consider caching
    end
    tf
end

@propagate_inbounds function statesum(fun::F, state::MechanismState, itr, init, safe=true) where F
    ret = init
    for x in itr
        ret += fun(state, x, safe)
    end
    ret
end

@propagate_inbounds function momentum(state::MechanismState, body_itr)
    T = cache_eltype(state)
    update_twists_wrt_world!(state)
    update_spatial_inertias!(state)
    statesum(momentum, state, body_itr, zero(Momentum{T}, root_frame(state.mechanism)), false)
end

@propagate_inbounds function momentum_rate_bias(state::MechanismState, body_itr)
    T = cache_eltype(state)
    update_bias_accelerations_wrt_world!(state)
    update_spatial_inertias!(state)
    statesum(momentum_rate_bias, state, body_itr, zero(Wrench{T}, root_frame(state.mechanism)), false)
end

@propagate_inbounds function kinetic_energy(state::MechanismState, body_itr)
    T = cache_eltype(state)
    update_twists_wrt_world!(state)
    update_spatial_inertias!(state)
    statesum(kinetic_energy, state, body_itr, zero(T), false)
end

@propagate_inbounds function gravitational_potential_energy(state::MechanismState, body_itr)
    T = cache_eltype(state)
    update_transforms!(state)
    statesum(gravitational_potential_energy, state, body_itr, zero(T), false)
end

for fun in (:momentum, :momentum_rate_bias, :kinetic_energy, :gravitational_potential_energy)
    @eval $fun(state::MechanismState) = @inbounds return $fun(state, state.nonrootbodies)
end

"""
$(SIGNATURES)

Return the homogeneous transform from `from` to `to`.
"""
function relative_transform(state::MechanismState, from::CartesianFrame3D, to::CartesianFrame3D)
    update_transforms!(state)
    inv(transform_to_root(state, to, false)) * transform_to_root(state, from, false)
end

"""
$(SIGNATURES)

Return the twist of `body` with respect to `base`, expressed in the
`Mechanism`'s root frame.
"""
function relative_twist(state::MechanismState, body::Union{<:RigidBody, BodyID}, base::Union{<:RigidBody, BodyID})
    update_twists_wrt_world!(state)
    -twist_wrt_world(state, base, false) + twist_wrt_world(state, body, false)
end

"""
$(SIGNATURES)

Return the twist of `body_frame` with respect to `base_frame`, expressed in the
`Mechanism`'s root frame.
"""
function relative_twist(state::MechanismState, body_frame::CartesianFrame3D, base_frame::CartesianFrame3D)
    body = body_fixed_frame_to_body(state.mechanism, body_frame)
    base = body_fixed_frame_to_body(state.mechanism, base_frame)
    twist = relative_twist(state, body, base)
    Twist(body_frame, base_frame, twist.frame, angular(twist), linear(twist))
end

for VectorType in (:Point3D, :FreeVector3D, :Twist, :Momentum, :Wrench)
    @eval begin
        function transform(state::MechanismState, v::$VectorType, to::CartesianFrame3D)
            # TODO: consider transforming in steps, so that computing the relative transform is not necessary
            transform(v, relative_transform(state, v.frame, to))
        end
    end
end

function transform(state::MechanismState, accel::SpatialAcceleration, to::CartesianFrame3D)
    old_to_root = transform_to_root(state, accel.frame)
    root_to_old = inv(old_to_root)
    twist_of_body_wrt_base = transform(relative_twist(state, accel.body, accel.base), root_to_old)
    twist_of_old_wrt_new = transform(relative_twist(state, accel.frame, to), root_to_old)
    old_to_new = inv(transform_to_root(state, to)) * old_to_root
    transform(accel, old_to_new, twist_of_old_wrt_new, twist_of_body_wrt_base)
end

"""
$(SIGNATURES)

Compute local coordinates ``\\phi`` centered around (global) configuration vector
``q_0``, as well as their time derivatives ``\\dot{\\phi}``.
""" # TODO: refer to the method that takes a joint once it's moved to its own Joints module
function local_coordinates!(
        ϕ::SegmentedVector{JointID}, ϕd::SegmentedVector{JointID}, state::MechanismState, q0::SegmentedVector{JointID})
    ϕs = values(segments(ϕ))
    ϕds = values(segments(ϕd))
    q0s = values(segments(q0))
    qs = values(segments(state.q))
    vs = values(segments(state.v))
    foreach((joint, ϕ, ϕ̇, q0, q, v) -> local_coordinates!(ϕ, ϕ̇, joint, q0, q, v), state.treejoints, ϕs, ϕds, q0s, qs, vs)
end

"""
$(SIGNATURES)

Convert local coordinates ``\\phi`` centered around ``q_0`` to (global)
configuration vector ``q``.
""" # TODO: refer to the method that takes a joint once it's moved to its own Joints module
function global_coordinates!(state::MechanismState, q0::SegmentedVector{JointID}, ϕ::SegmentedVector{JointID})
    qs = values(segments(state.q))
    q0s = values(segments(q0))
    ϕs = values(segments(ϕ))
    foreach((joint, q, q0, ϕ) -> global_coordinates!(q, joint, q0, ϕ), state.treejoints, qs, q0s, ϕs)
end


text:
using StructuralEstimation
# Set RNG seed
srand(124)

# State space
n = 500
o_max = 450
Xval = linspace(0, o_max, n)
p = [0.0937; 0.4475; 0.4459; 0.0127; 0.0002];

F1 = zeros(n, n)
offset = -1
for i = 1:n, j = 1:length(p)
    i+offset+j > n && continue
    F1[i,i+offset+j] = i+offset+j == n ? sum(p[j:end]) : p[j]
end

# We can handle sparse matrices
sparse_F = [sparse(F1), sparse(F1)[ones(Int64, n), :]]
F = [F1, F1[ones(Int64, n), :]]

S = State(Xval, F)
sparse_S = State(Xval, sparse_F)
bdisc = 0.9999
Z1 = [zeros(n) -0.001*Xval]
Z2 = [-ones(n) zeros(n)]

U = LinearUtility([Z1, Z2], bdisc, [11.;2.5])
V, iters = solve!(U,S)
T = 100
N = 100

D = simulate(U, S, 1, T, N)

#ESTIMATE TRANSITIONS

U.Θ*=0.
rust_nfxp = fit_nfxp(U, S, D)
rust_npl = fit_npl(U, S, D)

U = LinearUtility([Z1, Z2], 0.9999, [11., 3.]*0)
@time rust_nfxp = fit_nfxp(U, S, D; verbose = false);
U = LinearUtility([Z1, Z2], 0.9999, [11., 3.]*0)
@time rust_npl = fit_npl(U, S, D)

ext: jl


text:
using ComplexPhasePortrait, Colors
using Test

nx = 1000
x = range(-1, stop=1, length=nx)
Z = x' .+ reverse(x)*im

f = z -> (z - 0.5im)^2 * (z + 0.5+0.5im)/z
fz = f.(Z)

img = portrait(fz)
@test img[1,5] ≈ RGB{Float64}(0.11686143572621054,1.0,0.0)

img = portrait(fz, ctype="nist")
@test img[1,5] ≈ RGB{Float64}(0.4916573971078975,1.0,0.0)

img = portrait(fz, PTstepphase)
@test img[1,5] ≈ RGB{Float64}(0.09621043816546014,0.8232864637301505,0.0)

img = portrait(fz, PTstepmod)
@test img[1,5] ≈ RGB{Float64}(0.10275080341985139,0.8792533035498697,0.0)

img = portrait(fz, PTcgrid)
@test img[1,5] ≈ RGB{Float64}(0.08803468544171197,0.7533253797083627,0.0)


text:
abstract type AbstractCosmology end
abstract type AbstractCosmologicalGrid end
abstract type AbstractBackgroundQuantities end
abstract type BoltzmannSolverParams end
abstract type AbstractDensity end
abstract type AbstractConvolvedDensity end
abstract type AbstractInstrumentResponse end
abstract type AbstractWeightFunction end
abstract type AbstractSourceFunction end
abstract type AbstractTransferFunction end
abstract type AbstractPowerSpectrum end
abstract type AbstractCℓ end
abstract type Abstract∂Cℓ end
abstract type AbstractBias end
abstract type LensingEfficiencyMethod end
abstract type AbstractIntrinsicAlignment end
abstract type AbstractFFTLog end
abstract type AbstractCovariance end
abstract type AbstractFisher end

"""
    w0waCDMCosmology(w0::Float64 = -1, wa::Float64 = 0, ΩM::Float64 = 0.32,
    ΩB::Float64  = 0.05, ΩDE::Float64 = 0.68, Ωk::Float64  = 0.,
    Ωr::Float64  = 0., ns::Float64  = 0.96, Mν::Float64  = 0.06,
    σ8::Float64  = 0.816, H0::Float64  = 67.)

This struct contains the value of the cosmological parameters for ``w_0 w_a``CDM cosmologies:
- ``w_0`` and ``w_a``, the parameters in the [CPL parameterization](https://arxiv.org/abs/astro-ph/0208512)

- ``\\Omega_M``, ``\\Omega_B``, ``\\Omega_{DE}``, ``\\Omega_R``, and ``\\Omega_k`` the density parameters for respectively  matter, baryons, Dark Energy, radiation, curvature

- ``n_s``, the scalar spectral index

- ``M_\\nu``, the sum of the neutrino mass eigenstates in eV

- ``\\sigma_8``, the amplitude of the scalar fluctuations

- ``H_0``, the value of the Hubble paramater
"""
@kwdef mutable struct w0waCDMCosmology <: AbstractCosmology
    w0::Float64  = -1.
    wa::Float64  = 0.
    Mν::Float64  = 0.06 #neutrino mass in eV
    H0::Float64  = 67.
    ΩM::Float64  = 0.32
    ΩB::Float64  = 0.05
    ΩDE::Float64 = 0.68
    Ωk::Float64  = 0.
    Ωr::Float64  = 0.
    ns::Float64  = 0.96
    σ8::Float64  = 0.816
end

@kwdef mutable struct Flatw0waCDMCosmology <: AbstractCosmology
    w0::Float64  = -1.
    wa::Float64  = 0.
    Mν::Float64  = 0.06 #neutrino mass in eV
    H0::Float64  = 67.
    ΩM::Float64  = 0.32
    ΩB::Float64  = 0.05
    ns::Float64  = 0.96
    σ8::Float64  = 0.816
end

"""
    CosmologicalGrid(ZArray::Vector{Float64} = Array(LinRange(0.001, 2.5, 300)),
    KArray::Vector{Float64} = LogSpaced(1e-5, 50., 1000))

This struct contains the value of the Cosmological Grid, both in ``k`` and ``z``.
"""
@kwdef mutable struct CosmologicalGrid <: AbstractCosmologicalGrid
    ZArray::Vector{Float64} = Array(LinRange(0.001, 2.5, 300))
    KArray::Vector{Float64} = LogSpaced(1e-5, 50., 1000)
    ℓBinCenters::Vector{Float64} = LinRange(10., 3000., 2991)
    ℓBinWidths::Vector{Float64} = LinRange(10., 3000., 2991)
    KLimberArray::AbstractArray{Float64, 2} = zeros(length(ℓBinCenters),
    length(ZArray))
    KBeyondLimberArray::AbstractArray{Float64, 2} = zeros(100, 1000)
end

"""
    BackgroundQuantities(HZArray::Vector{Float64},
    rZArray::Vector{Float64})

This struct contains the arrays with the values of the Hubble parameter ``H(z)``
and the comoving distance ``r(z)``.
"""
@kwdef mutable struct BackgroundQuantities <: AbstractBackgroundQuantities
    HZArray::Vector{Float64} = zeros(500)
    rZArray::Vector{Float64} = zeros(500)
    DZArray::Vector{Float64} = zeros(500)
end

"""
    classyParams(classyParamsDict::Dict)

This struct contains the dictionary with the classy parameters. For a detalied
explanation of the parameters, please refer to the
[CLASS website](http://class-code.net/)
"""
@kwdef mutable struct classyParams <: BoltzmannSolverParams
    classyParamsDict::Dict = Dict("output" => "mPk",
        "non linear"=> "halofit",
        "Omega_b"=> 0.05,
        "Omega_cdm"=> 0.2735,
        "N_ur"=> 2.0328,
        "h"=> 0.67,
        "sigma8" => 0.816,
        "n_s" => 0.96,
        "m_ncdm" => 0.06,
        "P_k_max_1/Mpc" => 50,
        "z_max_pk" =>  5,
        "use_ppf" =>  "yes",
        "w0_fld" =>  -1.,
        "Omega_k" =>  0,
        "Omega_fld" =>  0.68,
        "wa_fld" =>  0.,
        "cs2_fld" =>  1.,
        "N_ncdm" =>  1,
        "tau_reio" =>  0.058)
end

"""
    AnalitycalDensity(z0::Float64 = 0.9/sqrt(2.), zmin::Float64 = 0.001,
    zmax::Float64 = 2.5, surfacedensity::Float64 = 30.,
    normalization::Float64 = 1.)

This struct contains the parameters of the source galaxy density as given by the
[official Euclid forecast](https://arxiv.org/abs/1910.09273), whose expression
is given by:
```math
n(z)\\propto\\left(\\frac{z}{z_0}\\right)^2
\\exp{\\left(-\\left(\\frac{z}{z_0}\\right)^{-3/2}\\right)}
```
The parameters contained in this struct are
- ``z_{min}`` and ``z_{max}``, the minimum and the maximum redshift considered

- ``z_0``, the parameter present in the galaxy distribution

- surfacedensity , the value of the galaxy source density integrated between ``z_{min}`` and ``z_{max}``
- normalization, the value of parameter which multiplies the source dennsity in order to match the correct surface density
"""
@kwdef mutable struct AnalitycalDensity <: AbstractDensity
    Z0::Float64 = 0.9/sqrt(2.)
    ZMin::Float64 = 0.001
    ZMax::Float64 = 4.0
    SurfaceDensity::Float64 = 30.
    Normalization::Float64 = 1.
end

"""
    ConvolvedDensity(AnalitycalDensity::AnalitycalDensity = AnalitycalDensity(),
    InstrumentResponse::InstrumentResponse = InstrumentResponse()
    ZBinArray::Vector{Float64} = Array([0.001, 0.418, 0.560, 0.678, 0.789, 0.900, 1.019, 1.155, 1.324, 1.576, 2.50])
    DensityNormalizationArray::Vector{Float64} = ones(length(ZBinArray)-1)
    DensityGridArray::AbstractArray{Float64, 2} = ones(length(ZBinArray)-1, 300))

In order to take into account the error in the redshift measurement, the
source density is convolved with the [`InstrumentResponse`](@ref),
according to [the following equation](https://arxiv.org/abs/1910.09273)
```math
n_{i}(z)=\\frac{\\int_{z_{i}^{-}}^{z_{i}^{+}}
\\mathrm{d} z_{\\mathrm{p}} n(z) p \\left(z_{\\mathrm{p}}
\\mid z\\right)}{\\int_{z_{\\min }}^{z_{\\max }} \\mathrm{d} z
\\int_{z_{i}^{-}}^{z_{i}^{+}} \\mathrm{d} z_{\\mathrm{p}} n(z) p
\\left(z_{\\mathrm{p}} \\mid z\\right)}
```
"""
@kwdef mutable struct ConvolvedDensity <: AbstractConvolvedDensity
    ZBinArray::Vector{Float64} = Array([0.001, 0.418, 0.560, 0.678, 0.789,
    0.900, 1.019, 1.155, 1.324, 1.576, 2.50])
    DensityNormalizationArray::Vector{Float64} = ones(length(ZBinArray)-1)
    DensityGridArray::AbstractArray{Float64, 2} = ones(length(ZBinArray)-1, 300)
    ShiftArray::Vector{Float64} = zeros(length(ZBinArray)-1)
    SurfaceDensityArray = ones(10)
end

"""
    InstrumentResponse(cb::Float64 = 1.0, zb::Float64 = 0.0,
    σb::Float64 = 0.05, co::Float64 = 1.0, zo::Float64 = 0.1,
    σo::Float64 = 0.05, fout::Float64 = 0.1)

When we measure the redshift of a galaxy with redshit ``z``, we will measure a
redshift ``z_p`` with a probability given by
[the following expression](https://arxiv.org/abs/1910.09273):
```math
p(z_p|z)  = \\frac{1-f_{out}}{\\sqrt{2 \\pi} \\sigma_{b}(1+z)} \\exp \\left(
-\\frac{1}{2}\\left(\\frac{z-c_{b} z_{b}-z_{b}}{\\sigma_{b}(1+z)}\\right)^{2}
\\right) + \\frac{f_{out}}{\\sqrt{2 \\pi} \\sigma_{\\mathrm{o}}(1+z)} \\exp
\\left(-\\frac{1}{2}\\left(\\frac{z-c_{o} z_{p}-z_{o}}{\\sigma_{o}(1+z)}
\\right)^{2}\\right)
```
This struct contains all these parameters.
"""
@kwdef struct InstrumentResponse <: AbstractInstrumentResponse
    cb::Float64   = 1.0
    zb::Float64   = 0.0
    σb::Float64   = 0.05
    co::Float64   = 1.0
    zo::Float64   = 0.1
    σo::Float64   = 0.05
    fout::Float64 = 0.1
end

"""
    GCWeightFunction()

This struct contains the array with the Galaxy Bias and Galaxy Clustering Weight Function
values for all tomographic bins and redshift values in the [`CosmologicalGrid`](@ref).
"""
@kwdef mutable struct GCWeightFunction <: AbstractWeightFunction
    WeightFunctionArray::AbstractArray{Float64, 2} = zeros(10, 500)
    BiasArray::AbstractArray{Float64, 2} = zeros(size(WeightFunctionArray))
    BiasKind::AbstractBias = PiecewiseBias()
end

@kwdef mutable struct PiecewiseBias <: AbstractBias
    BiasMultiplier::AbstractArray{Float64, 1} = ones(10)
end

"""
    EuclidBias()

This struct contains the parameter for the bias model measured by
[the Euclid Collaboration](https://arxiv.org/abs/2005.00055):

```math
b(z)= A +\\frac{B}{1+\\exp \\left( \\left(D-z \\right)C   \\right)}
```
"""
@kwdef mutable struct EuclidBias <: AbstractBias
    A::Float64 = 1.0
    B::Float64 = 2.5
    C::Float64 = 2.8
    D::Float64 = 1.6
end

struct AbsentIA <: AbstractIntrinsicAlignment end

@kwdef mutable struct ExtendedNLIA <: AbstractIntrinsicAlignment
    𝓐IA::Float64 = 1.72
    βIA::Float64 = 2.17
    𝓒IA::Float64 = 0.0134
    ηIA::Float64 = -0.41
end

"""
    WLWeightFunction()

This struct contains the array with the Lensing Efficiency and Weak Lensing
Weight Function values for all tomographic bins and redshift values in the
[`CosmologicalGrid`](@ref)
"""
@kwdef mutable struct WLWeightFunction <: AbstractWeightFunction
    WeightFunctionArray::AbstractArray{Float64, 2} = zeros(10, 500)
    LensingEfficiencyArray::AbstractArray{Float64, 2} = zeros(10, 500)
    IntrinsicAlignmentArray::AbstractArray{Float64, 2} = zeros(10, 500)
    IntrinsicAlignmentModel::AbstractIntrinsicAlignment = ExtendedNLIA()
end

@kwdef mutable struct LensingSourceFunction <: AbstractSourceFunction
    SourceFunctionArray::AbstractArray{Float64, 2} = zeros(10, 500)
    LensingEfficiencyArray::AbstractArray{Float64, 2} = zeros(10, 500)
end

"""
    PowerSpectrum()

This struct contains the array with the Linear and Nonlinear Power Spectrum
evaluated on the ``k-z`` grid and the interpolated Nonlinear Power Spectrum on
Limber ``k-z`` grid.
"""
@kwdef mutable struct PowerSpectrum <: AbstractPowerSpectrum
    PowerSpectrumLinArray::AbstractArray{Float64, 2} = zeros(1000, 300)
    PowerSpectrumNonlinArray::AbstractArray{Float64, 2} = zeros(1000, 300)
    InterpolatedPowerSpectrum::AbstractArray{Float64, 2} = zeros(2991, 300)
    GrowthFactor::AbstractArray{Float64, 1} = zeros(
    length(PowerSpectrumLinArray[:,1]))
    InterpolatedPowerSpectrumBeyondLimber::AbstractArray{Float64, 2} =
    zeros(100, 1000)
end

"""
    Cℓ(CℓArray::AbstractArray{Float64, 3})

This struct contains the array with the Angular Coefficients.
"""
@kwdef mutable struct Cℓ <: AbstractCℓ
    CℓArray::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
end

"""
    ∂Cℓ(∂CℓArray::AbstractArray{Float64, 3})

This struct contains the array with the derivatives of the Angular Coefficients.
"""
@kwdef mutable struct ∂Cℓ <: Abstract∂Cℓ
    ∂CℓArray::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
end

"""
    Fisherαβ()

This struct contains the array with the Fisher Matrix.
"""
@kwdef mutable struct Fisherαβ <: AbstractFisher
    FisherMatrix::AbstractArray{Float64, 2} = zeros(8,8)
    CorrelationMatrix::AbstractArray{Float64, 2} = zeros(8,8)
    FisherDict::Dict = Dict()
    ParametersList::Vector{String} = []
    SelectedParametersList::Vector{String} = []
    MarginalizedErrors::Dict = Dict()
end

"""
    aₗₘCovariance()

This struct contains the Covariance in the field perspective, i.e. when the observables are the
``a_{\\ell m}``'s.
"""
@kwdef mutable struct aₗₘCovariance  <: AbstractCovariance
    Covariance::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
    Cℓ::AbstractCℓ = Cℓ()
    Noise::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
    Covariance⁻¹::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
end

"""
    CℓCovariance()

This struct contains the Covariance in the estimator perspective, i.e. when the observables
are the ``C_{\\ell}``'s.
"""
@kwdef mutable struct CℓCovariance  <: AbstractCovariance
    Covariance::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
    Covariance⁻¹::AbstractArray{Float64, 3} = zeros(2991, 10, 10)
end

"""
    InterpolationMethod

This type is used to specify the method to interpolate the Power Spectrum;
actually are included:

- Dierckx.jl, the Julia wrapper for the Fortran library Dierckx

- GriddedLinear, from Interpolations.jl

- BSpliceCubic (recommended, for its speed and accuracy), from Interpolations.jl
"""
abstract type InterpolationMethod end

struct RectBivSplineDierckx <: InterpolationMethod end
struct GriddedLinear <: InterpolationMethod end
struct BSplineCubic <: InterpolationMethod end

struct StandardLensingEfficiency <: LensingEfficiencyMethod end
struct CustomLensingEfficiency <: LensingEfficiencyMethod end

@kwdef mutable struct FFTLog <: AbstractFFTLog
    XArray::Vector{Float64}
    DLnX::Float64 = log(XArray[2]/XArray[1])
    FXArray::Vector{Float64}
    OriginalLenght::Int64 = length(XArray)
    ν::Float64 = 1.01
    NExtrapLow::Int64 = 0
    NExtrapHigh::Int64 = 0
    CWindowWidth::Float64 = 0.25
    NPad::Int64 = 500
    N::Int64 = OriginalLenght+NExtrapHigh+NExtrapLow+2*NPad
    M::Vector{Float64} = zeros(N)
    CM::Vector{ComplexF64} = zeros(N)
    ηM::Vector{Float64} = zeros(N)
end

"""
    κTransferFunction()

This struct contains the array with the Lensing Efficiency and Weak Lensing
Weight Function values for all tomographic bins and redshift values in the
[`CosmologicalGrid`](@ref)
"""
@kwdef mutable struct κTransferFunction <: AbstractTransferFunction
    LensingSourceFunction::LensingSourceFunction = LensingSourceFunction
    TransferFunctionArray::AbstractArray{Float64, 3} = zeros(10, 100, 1000)
end


text:
abstract masks_Forward_Backward

function BackwardD2( r2::Array{Float64} , M :: masks_Forward_Backward;)
  error("not serving for a pure masks_Forward_Backward")
end

function LBackwardD2( r2::Array{Float64,1} , M :: masks_Forward_Backward;)
  error("not serving for a pure masks_Forward_Backward")
end

function LBackwardD2( r2::Array{Float64,2} , M :: masks_Forward_Backward;)
  error("not serving for a pure masks_Forward_Backward")
end

export masks_Forward_Backward
export ForwardD2,BackwardD2,LBackwardD2;

abstract Met_DA_Functor{N}
abstract REG_Functor <: Met_DA_Functor{2}
export Met_DA_Functor
export REG_Functor

type _AR2_REG <: REG_Functor end

type _EOF_Functor <: Met_DA_Functor{1} end
type _SVD_Functor <: Met_DA_Functor{2} end
type _COV_Functor <: Met_DA_Functor{2} end

EOFF = _EOF_Functor(); SVDF = _SVD_Functor();
COVF = _COV_Functor();
export EOFF, _EOF_Functor, SVDF, _SVD_Functor;
export COVF, _COV_Functor;

type _Mean_Functor <: Met_DA_Functor{1} end
type _Sum_Functor <: Met_DA_Functor{1} end
SUMF = _Sum_Functor(); MEANF = _Mean_Functor()
export SUMF, MEANF
function Fevaluate{N}(EE :: Met_DA_Functor{N} ,x)
  error("No function for pure")
end

function DFevaluate{N}(EE :: Met_DA_Functor{N} ,x)
  error("No function for pure")
end

function Fevaluate( EE :: _EOF_Functor, xarr:: Array{Float64,2}; nsv = 3)
  xarr = RemoveTr(xarr)
  xarr = broadcast(-, xarr,mean(xarr,2))
  #= (S,V,D) = svds(xarr,nsv=nsv) =#
  (S,V,D) = svd(xarr);
  S = S[:,1:nsv]; V = V[1:nsv]; D= D[:,1:nsv];
  PC = S' * xarr
  ratio = V.^2/sum(xarr.^2)
  return(S, V, PC, ratio)
end

function Fevaluate( FF :: _SVD_Functor, xarr :: Array{Float64,2},
  yarr :: Array{Float64,2}; nsv = 3)
  xarr = RemoveTr(xarr)
  yarr = RemoveTr(yarr)
  xarr = broadcast(-, xarr,mean(xarr,2))
  yarr = broadcast(-, yarr,mean(yarr,2))
  covxy = xarr*yarr'
  #= (Sx,V,Sy) = svds(covxy,nsv=nsv) =#
  (Sx,V,Sy) = svd(covxy);
  Sx = Sx[:,1:nsv]; V = V[1:nsv]; Sy = Sy[:,1:nsv];
  PCx = Sx'*xarr
  PCy = Sy'*yarr
  ratio = V.^2/sum(covxy.^2)
  return (Sx, Sy, V, ratio, PCx, PCy)
end
export Fevaluate

function Fevaluate( CV :: _COV_Functor, xarr :: Array{Float64,2}, yarr :: Array{Float64,2};
                   meanidx :: Int64 = size(xarrr,2))
  return xarr * yarr'/ meanidx
end

include("NaN_Forward_and_Backward.jl")
include("Weighted_Forward_Backward.jl")


text:
"""
    MPS{L,T<:Number}

Matrix product state for `L` sites with numerical entries of type `T`.

The tensors are stored with the indices `[a, m, b]` corresponding to

       m
       |
    a -#- b

where the physical index (`m`) is between the bond indices (`a`, `b`). The end
tensors also have this structure, so they must be capped explicitly.
"""
struct MPS{L,T<:Number}
    "Tensors making up the MPS."
    tnsrs::Vector{Array{T,3}}
end

"""
    MPS(vs::Vector{Vector})

Create an `MPS` representing a product state (all bonds have dimension 1),
where each site is described by the corresponding element of `vs`.
"""
function MPS(vs::Vector{Vector{T}}) where {T}
    L = length(vs)

    tnsrs = Vector{Array{T,3}}(undef, L)
    for i in 1:L
        tnsrs[i] = reshape(copy(vs[i]), 1, :, 1)
    end

    MPS{L,T}(tnsrs)
end

"""
    MPS(v::Vector, L)

Create an `MPS` for `L` sites representing a uniform product state (all bonds
have dimension 1), where each site is described by `v`.
"""
MPS(v::Vector, L) = MPS([v for _ in 1:L])


text:
using GeometricAlgebra
const ga = GeometricAlgebra

@basis "++++"
include("definitions.jl")
include("../operators.jl")


text:
using Test

using FFTW

using Fourier

function test_core(ary; forward)
    if forward
        truth = fftshift(fft(ifftshift(ary)))
        candidate = mdft(ary, size(ary))
        @test candidate ≈ truth
        return
    end
    truth = fftshift(ifft(ifftshift(ary)))
    candidate = imdft(ary, size(ary)) ./ length(ary)  # match ifft convention
    @test candidate ≈ truth
end

@testset "regression against FFTW (even square array)" begin
    ary_sizes = [2, 4, 8, 16, 32, 64, 128, 256, 512, 100] # includes a non power of two
    for sz in ary_sizes
        a = rand(sz, sz)
        test_core(a, forward = true)
        test_core(a, forward = false)
    end
end

@testset "regression against FFTW (odd square array)" begin
    ary_sizes = [3, 5, 11, 19, 33, 51, 101, 201, 707, 555]
    for sz in ary_sizes
        a = rand(sz, sz)
        test_core(a, forward = true)
        test_core(a, forward = false)
    end
end

@testset "regression against FFTW (nonsquare even array)" begin
    ary_sizes =
        [(2, 4), (4, 8), (4, 2), (8, 4), (100, 200), (200, 100), (256, 512), (512, 2)]
    for sz in ary_sizes
        a = rand(sz...)
        test_core(a, forward = true)
        test_core(a, forward = false)
    end
end

@testset "regression against FFTW (nonsquare odd array)" begin
    ary_sizes = [(3, 5), (5, 7), (7, 9), (7, 101), (101, 7), (3, 9), (101, 201), (301, 201)]
    for sz in ary_sizes
        a = rand(sz...)
        test_core(a, forward = true)
        test_core(a, forward = false)
    end
end

@testset "regression against FFTW (zoom, Q=2)" begin
    b = rand(2, 2)
    a = zeros(4, 4)
    a[2:3, 2:3] = b
    # the comparison is slightly subtle
    # mdft does not require padding to achieve
    # sampling, so fft uses a (padded) while mdft uses b (unpadded)
    truth = fftshift(fft(ifftshift(a)))
    candidate = mdft(b, 4, Q = 2)
    @test candidate ≈ truth
    truth = fftshift(ifft(ifftshift(a)))
    candidate = imdft(b, 4, Q = 2) ./ length(truth)
    @test candidate ≈ truth
end

@testset "shift works as expected" begin
    data = ones(100, 100)  # uniform plane has all energy in zero DFT bin
    # => bump at (50,50) before shift

    # X shifts
    shifted_fwd = abs.(mdft(data, size(data), shift = (50, 0)))
    @test argmax(shifted_fwd)[2] == 1
    shifted_fwd = abs.(mdft(data, size(data), shift = (-49, 0)))
    @test argmax(shifted_fwd)[2] == 100

    shifted_fwd = abs.(imdft(data, size(data), shift = (50, 0)))
    @test argmax(shifted_fwd)[2] == 1
    shifted_fwd = abs.(imdft(data, size(data), shift = (-49, 0)))
    @test argmax(shifted_fwd)[2] == 100

    # Y shifts
    shifted_fwd = abs.(mdft(data, size(data), shift = (0, 50)))
    @test argmax(shifted_fwd)[1] == 1
    shifted_fwd = abs.(mdft(data, size(data), shift = (0, -49)))
    @test argmax(shifted_fwd)[1] == 100

    shifted_fwd = abs.(imdft(data, size(data), shift = (0, 50)))
    @test argmax(shifted_fwd)[1] == 1
    shifted_fwd = abs.(imdft(data, size(data), shift = (0, -49)))
    @test argmax(shifted_fwd)[1] == 100
end

max_stars_repo_name: JuliaOptics/Fourier

max_issues_repo_name: JuliaOptics/Fourier

max_forks_repo_name: JuliaOptics/Fourier


text:
using FlightMechanicsSimulator
using FlightMechanicsUtils
using OrdinaryDiffEq
using SimulationLogs
using Test

x_stev = [
    502.0 * FT2M,
    0.2392628,
    0.0005061803,
    1.366289,
    0.05000808,
    0.2340769,
    -0.01499617,
    0.2933811,
    0.06084932,
    0.0 * FT2M,
    0.0 * FT2M,
    0.0 * FT2M,
    64.12363,
]
controls_stev = [0.8349601, -1.481766, 0.09553108, -0.4118124]
xcg = 0.35
dssd, outputs = f(
    time,
    SixDOFAeroEuler(x_stev),
    controls_stev,
    F16(F16Stevens.MASS, F16Stevens.INERTIA, xcg),
    F16StevensAtmosphere(x_stev[12]),
    LHDownGravity(FlightMechanicsSimulator.F16Stevens.GD*FT2M),
)

x_dot = get_xdot(dssd)

# RETRIM to refine flying condition
dssd, controls_trim, outputs_trim, cost = trim(
    SixDOFAeroEuler(x_stev),
    controls_stev,
    F16(F16Stevens.MASS, F16Stevens.INERTIA, xcg),
    F16StevensAtmosphere(x_stev[12]),
    LHDownGravity(FlightMechanicsSimulator.F16Stevens.GD*FT2M),
    0.0,
    0.3,
)

x_trim = get_x(dssd)
x_dot_trim = get_xdot(dssd)

dt = 0.01  # s
t0 = 0.0  # s
t1 = 180.0  # s

x = x_trim
# Transform to Input for simulate
controls = ConstantInput.(controls_trim)

results = FlightMechanicsSimulator._simulate(
    t0,
    t1,
    SixDOFAeroEuler(x),
    controls,
    F16(F16Stevens.MASS, F16Stevens.INERTIA, xcg),
    F16StevensAtmosphere,
    LHDownGravity(FlightMechanicsSimulator.F16Stevens.GD*FT2M);
    solver=RK4(),
    solve_args=Dict(:reltol=>1e-10, :saveat=>dt),
    )

out = get_log(results)

@testset "Logging" begin
    for (i, name) in enumerate(get_x_names(dssd)[1:9])
        @test results[i,:] == getproperty(out, name)
    end
end

out.Fx


text:
# capitalize.jl - basic operations to chain price change arrays. Operaciones
# básicas para encadenar variaciones intermensuales en índices de precios.

"""
    capitalize(v::AbstractVector, base_index::Real = 100)
    capitalize(vmat::AbstractMatrix, base_index::Real = 100)
    capitalize(vmat::AbstractMatrix, base_index::AbstractVector)

Función para encadenar un vector o matriz con variaciones intermensuales de
índices de precios `v` o `vmat` para conformar un índice de precios cuyo valor
base sea `base_index`.
"""
function capitalize end

"""
    capitalize!(idx:: AbstractVector, v::AbstractVector, base_index::Real)
    capitalize!(vmat::AbstractMatrix, base_index = 100)

Función para encadenar un vector o matriz con variaciones intermensuales de
índices de precios `v` o `vmat` para conformar un índice de precios cuyo valor
base sea `base_index` y sea almacenado en `idx` o en la propia matriz `vmat`.
"""
function capitalize! end

function capitalize(v::AbstractVector, base_index::Real = 100)
    idx = similar(v)
    capitalize!(idx, v, base_index)
    idx
end

function capitalize(vmat::AbstractMatrix, base_index::Real = 100)
    c = size(vmat, 2)
    idxmat = similar(vmat)
    for j in 1:c
        vcol = @view vmat[:, j]
        idxcol = @view idxmat[:, j]
        capitalize!(idxcol, vcol, base_index)
    end
    idxmat
end

function capitalize(vmat::AbstractMatrix, base_index::AbstractVector)
    c = size(vmat, 2)
    idxmat = similar(vmat)
    for j in 1:c
        vcol = @view vmat[:, j]
        idxcol = @view idxmat[:, j]
        capitalize!(idxcol, vcol, base_index[j])
    end
    idxmat
end

## Version in-place

function capitalize!(idx:: AbstractVector, v::AbstractVector, base_index::Real)
    l = length(v)
    idx[1] = base_index * (1 + v[1]/100)
    for i in 2:l
        @inbounds idx[i] = idx[i-1] * (1 + v[i]/100)
    end
end

capitalize!(v::AbstractVector, base_index::Real) = capitalize!(v, v, base_index)

function capitalize!(vmat::AbstractMatrix, base_index = 100)
    r = size(vmat, 1)
    @views @. vmat[1, :] = base_index * (1 + vmat[1, :]/100)
    for i in 2:r
        @views @. vmat[i, :] = vmat[i-1, :] * (1 + vmat[i, :]/100)
    end
end

## Versiones para tipos contenedores

function _offset_back(dates)
    start = first(dates) - Month(1)
    start:Month(1):last(dates)
end

"""
    capitalize(base::VarCPIBase)

Esto devuelve una nueva instancia (copia) de tipo `IndexCPIBase` de un objeto
`VarCPIBase`.
"""
function capitalize(base::VarCPIBase)
    IndexCPIBase(base)
end

## Otras versiones (descontinuado)

#=
"""
    capitalize_addbase(vmat::AbstractMatrix, base_index = 100)

Function to chain a matrix of price changes with an index starting with `base_index`. This function adds the base index as the first row of the returned matrix.
"""
=#
function capitalize_addbase(vmat::AbstractMatrix, base_index = 100)
    r, c = size(vmat)
    idxmat = zeros(eltype(vmat), r+1, c)
    idxmat[1, :] .= base_index
    for i in 1:r
        @views @. idxmat[i+1, :] = idxmat[i, :] * (1 + vmat[i, :]/100)
    end
    idxmat
end


text:
export vertices,
    transitions,
    minimize_automaton!

#################
### Overloads ###
#################

#= WordVertex - Returns whether the child s exists for n =#
_childexists(n::WordVertex{T}, s::Symbol) where {T <: Integer}  = !isnothing(getfield(n, s))

#= Automaton - Overloads of Base functions =#
Base.getindex(g::Automaton{T}, w::T) where {T <: Integer}                       = g.data[w]
Base.setindex!(g::Automaton{T}, v::WordVertex{T}, w::T) where {T <: Integer}    = (g.data[w] = v)
Base.length(g::Automaton{T}) where {T <: Integer}                               = length(g.data)
Base.haskey(g::Automaton{T}, w::T) where {T <: Integer}                         = haskey(g.data, w)

######################
### Util Functions ###
######################

"""
    vertices(automaton::Automaton)

Returns all the vertices in `automaton`.
"""
function vertices(automaton::Automaton{T}) where {T <: Integer}
    return values(automaton.data)
end # function

"""
    transitions(automaton::Automaton)

Returns all the transitions in `automaton` in the form of a set of pairs where each pair consists of `(v1, v2, c12)`, i.e., the tail of the transition `v1`, the head of the transition `v2`, and the label of the transition `c12`.
"""
function transitions(automaton::Automaton{T}) where {T <: Integer}

    trans = Set{Tuple{T, T, UInt8}}();
    sizehint!(trans, 2*length(vertices(automaton)))

    for vertex in vertices(automaton)
        push!(trans, (vertex.w, vertex.hit.w, Hit))
        if _childexists(vertex, :miss)
            push!(trans, (vertex.w, vertex.miss.w, Miss))
        end # if
    end # for

    return trans
end # function

function _vertices_equivalent(v1::WordVertex{T}, v2::WordVertex{T}) where {T <: Integer}
    if v1 === v2 # If they are the same vertex
        return false
    end

    if _childexists(v1, :miss) && _childexists(v2, :miss)
        return v1.miss.w == v2.miss.w && v1.hit.w == v2.hit.w
    elseif !_childexists(v1, :miss) && !_childexists(v2, :miss)
        return v1.hit.w == v2.hit.w
    else
        return false
    end #if
end # function

"""
    minimize_automaton!(automaton::Automaton)

Minimises the automaton representation of a set of weakly-hard constraints.
"""
function minimize_automaton!(automaton::Automaton{T}) where {T <: Integer}
    changed = true
    while changed
        changed = false

        for v1 in vertices(automaton)
            v1_identical = Vector{WordVertex{T}}()
            for v2 in vertices(automaton)
                if _vertices_equivalent(v1, v2)
                    changed = true
                    push!(v1_identical, v2)
                end
            end # for

            if !isempty(v1_identical)
                push!(v1_identical, v1)
                # find shortest word to replace v1_identical with
                v1_new = reduce(v1_identical) do x, y
                (x.w < y.w) ? x : y end
                for v3 in vertices(automaton)
                    if v3 in v1_identical && !(v3 === v1_new)
                        delete!(automaton.data, v3.w)
                    else
                        if _childexists(v3, :miss) && v3.miss in v1_identical
                            v3.miss = v1_new
                        end
                        if v3.hit in v1_identical
                            v3.hit = v1_new
                        end
                    end # if
                end # for
            end # if
        end # for
    end # while
    return automaton
end # function

###########################
### Auxiliary Functions ###
###########################

#= Find minimal sized Integer to use for automaton construction =#
function _word_type(Lambda::Set{T}, k::S) where {T <: Constraint, S <: Integer}
    n = 0
    for lambda in Lambda
        if typeof(lambda) == RowHitConstraint && lambda.x > n
            n = lambda.x
        end # if
    end # for

    if k + 2*n < 64 - 2
        return Int64
    else
        return BigInt
    end # if
end # function

_word_type(lambda::T, k::S) where {T <: Constraint, S <: Integer} = _word_type(Set{T}([lambda]), k)

#= Find minimal size of automaton (specifically used by sizehint for the dictionary in the automaton) =#
_minimal_size(lambda::AnyHitConstraint)  = binomial(lambda.k, lambda.x)
_minimal_size(lambda::AnyMissConstraint) = binomial(lambda.k, lambda.k - lambda.x)
_minimal_size(lambda::RowMissConstraint) = lambda.x + 1
function _minimal_size(lambda::RowHitConstraint)
    if lambda.x > lambda.k
        return 0
    elseif lambda.k < 2*lambda.x
        return 1
    elseif lambda.k == 2*lambda.x
        return lambda.x + 1
    elseif lambda.k == 2*lambda.x + 1
        return lambda.x + 2
    elseif lambda.k < 3*lambda.x
        return 2*_minimal_size(RowHitConstraint(lambda.x, lambda.k - 1)) - _minimal_size(RowHitConstraint(lambda.x, lambda.k - 2)) + 1
    else
        return _minimal_size(RowHitConstraint(lambda.x, lambda.k - 1)) + lambda.x
    end # if
end # function


text:
using SpecialFunctions, Roots
using Parameters, Distributions
using Plots, Random
const sf = SpecialFunctions

export PDFParameters
export plot_input_pdfs, int_xtotx, xtotx
export get_input_pdf_func
export input_pdf_map

function get_dirichlet_samples(λ_u::Float64, K_u::Float64, λ_d::Float64,
                               K_d::Float64, seed::Integer, weights::Vector{Float64})

    A_u = 2 / sf.beta(λ_u, K_u + 1)
    A_d = 1 / sf.beta(λ_d, K_d + 1)

    I_u = A_u * sf.beta(λ_u+1, K_u+1)
    I_d = A_d * sf.beta(λ_d+1, K_d+1)

    remaining = 1 - (I_u + I_d)

    Random.seed!(seed);
    dirichlet = Dirichlet(weights)

    return rand(dirichlet) * remaining
end

"""
    PDFParameters

Parameters of the input PDFs.
"""
@with_kw struct PDFParameters
    λ_u::Float64
    K_u::Float64
    λ_d::Float64
    K_d::Float64
    λ_g1::Float64
    λ_g2::Float64
    K_g::Float64
    λ_q::Float64
    seed::Integer = 0
    weights::Vector{Float64}
    θ::Vector{Float64} = get_dirichlet_samples(λ_u, K_u, λ_d, K_d, seed, weights)
end

"""
    x_uv_x(x, λ_u, K_u)

Momentum density of u valence component.
Beta function
    A_u x^λ_u (1-x)^K_u
A_u is set by λ_u and K_u.
"""
function x_uv_x(x::Float64, λ_u::Float64, K_u::Float64)

    A_u = 2 / sf.beta(λ_u, K_u + 1)

    return A_u * x^λ_u * (1 - x)^K_u

end

"""
    x_dv_x(x, λ_d, K_d)

Momentum density of d valence component.
Beta function
    A_d x^λ_d (1 - x)^K_d
A_d is set by λ_d and K_d.
"""
function x_dv_x(x::Float64, λ_d::Float64, K_d::Float64)

    A_d = 1 / sf.beta(λ_d, K_d + 1)

    return A_d * x^λ_d * (1 - x)^K_d

end

"""
        x_g_x(x, λ_g1, λ_g2, K_g, w1, w2)

Momentum density of gluon component.
    A_g1 x^λ_g1 (1 - x)^K_g + A_g2 x^λ_g2
Amplitudes are set by weights `w1` and `w2`.
"""
function x_g_x(x::Float64, λ_g1::Float64, λ_g2::Float64, K_g::Float64,
               w1::Float64, w2::Float64)

    A_g1 = w1 / sf.beta(λ_g1 + 1, K_g + 1)

    x_g1_x = A_g1 * x^λ_g1 * (1 - x)^K_g

    A_g2 = w2 / sf.beta(λ_g2 + 1, 5 + 1)

    x_g2_x = A_g2 * x^λ_g2 * (1 - x)^5

    return x_g1_x + x_g2_x

end

"""
    x_q_x(x, λ_q, w)

Momentum density of non-valence
quark component.
    A_q x^λ_q
Amplitude is set by weight `w`
"""
function x_q_x(x::Float64, λ_q::Float64, w::Float64)

    A_q = (w / 2) / sf.beta(λ_q + 1, 5 + 1)

    return A_q * x^λ_q * (1 - x)^5

end

"""
    x_total_x(x, λ_u, K_u, λ_d, K_d, λ_g1, λ_g2, K_g, λ_q, θ)

Total momentum density.
"""
function xtotx(x::Float64, λ_u::Float64, K_u::Float64, λ_d::Float64, K_d::Float64,
               λ_g1::Float64, λ_g2::Float64, K_g::Float64, λ_q::Float64, θ::Array{Float64})

    xuvx = x_uv_x(x, λ_u, K_u)

    xdvx = x_dv_x(x, λ_d, K_d)

    xgx = x_g_x(x, λ_g1, λ_g2, K_g, θ[1], θ[2])

    xubarx = x_q_x(x, λ_q, θ[3])

    xdbarx = x_q_x(x, λ_q, θ[4])

    xsx = x_q_x(x, λ_q, θ[5])

    xcx = x_q_x(x, λ_q, θ[6])

    xbx = x_q_x(x, λ_q, θ[7])

    return xuvx + xdvx + xgx + xubarx + xdbarx + xsx + xcx + xbx

end

"""
    x_total_x(x, hyper_params)

Total momentum density.
"""
function xtotx(x::Float64, hyper_params::PDFParameters)

    hp = hyper_params

    return xtotx(x, hp.λ_u, hp.K_u, hp.λ_d, hp.K_d, hp.λ_g1,
                 hp.λ_g2, hp.K_g, hp.λ_q, hp.θ)
end

"""
    int_xtotx(λ_u, K_u, λ_d, K_d, λ_g1, λ_g2, K_g, λ_q, θ)

Total integrated momentum density. Should equal 1.
"""
function int_xtotx(λ_u::Float64, K_u::Float64, λ_d::Float64, K_d::Float64,
                   λ_g1::Float64, λ_g2::Float64, K_g::Float64, λ_q::Float64, θ::Array{Float64})

    A_u = 2 / sf.beta(λ_u, K_u + 1)
    A_d = 1 / sf.beta(λ_d, K_d + 1)

    A_g1 = θ[1] / sf.beta(λ_g1 + 1, K_g + 1)
    A_g2 = θ[2] / sf.beta(λ_g2 + 1, 5 + 1)

    A_ubar = (θ[3] / 2) / sf.beta(λ_q + 1, 5 + 1)

    A_dbar = (θ[4] / 2) / sf.beta(λ_q + 1, 5 + 1)

    A_s = (θ[5] / 2) / sf.beta(λ_q + 1, 5 + 1)

    A_c = (θ[6] / 2) / sf.beta(λ_q + 1, 5 + 1)

    A_b = (θ[7] / 2) / sf.beta(λ_q + 1, 5 + 1)

    # Integrate
    I_u = A_u * sf.beta(λ_u+1, K_u+1)

    I_d = A_d * sf.beta(λ_d+1, K_d+1)

    I_g = A_g1 * sf.beta(λ_g1 + 1, K_g + 1) + A_g2 * sf.beta(λ_g2 + 1, 5 + 1)

    I_q = 2 * (A_ubar + A_dbar + A_s + A_c + A_b) * sf.beta(λ_q + 1, 5 + 1)

    I_tot = I_u + I_d + I_g + I_q

    I_tot
end

"""
    int_xtotx(hyper_params)

Total integrated momentum density. Should equal 1.
"""
function int_xtotx(hyper_params::PDFParameters)

    hp = hyper_params

    result = int_xtotx(hp.λ_u, hp.K_u, hp.λ_d, hp.K_d,
                       hp.λ_g1, hp.λ_g2, hp.K_g, hp.λ_q, hp.θ)

    return result
end

"""
    plot_input_pdfs(hyper_params, xmin, xmax, nx)

Plot the input PDFs defined by hyper_params over
the given x range.
"""
function plot_input_pdfs(hyper_params::PDFParameters, xmin::Float64=1.0e-2,
                         xmax::Float64=1.0, nx::Integer=1000)

    x_grid = range(xmin, stop=xmax, length=nx)
    hp = hyper_params

    p = plot(x_grid, [x_uv_x(x, hp.λ_u, hp.K_u) for x in x_grid], label="x uv(x)", lw=3)
    p = plot!(x_grid, [x_dv_x(x, hp.λ_d, hp.K_d) for x in x_grid], label="x dv(x)", lw=3)
    p = plot!(x_grid, [x_g_x(x, hp.λ_g1, hp.λ_g2, hp.K_g, hp.θ[1], hp.θ[2])
                   for x in x_grid], label="x g(x)", lw=3)

    p = plot!(x_grid, [x_q_x(x, hp.λ_q, hp.θ[3]) for x in x_grid], label="x ubar(x)", lw=3)
    p = plot!(x_grid, [x_q_x(x, hp.λ_q, hp.θ[4]) for x in x_grid], label="x dbar(x)", lw=3)
    p = plot!(x_grid, [x_q_x(x, hp.λ_q, hp.θ[5]) for x in x_grid], label="x s(x)", lw=3)
    p = plot!(x_grid, [x_q_x(x, hp.λ_q, hp.θ[6]) for x in x_grid], label="x c(x)", lw=3)
    p = plot!(x_grid, [x_q_x(x, hp.λ_q, hp.θ[7]) for x in x_grid], label="x b(x)", lw=3)

    p = plot!(xlabel="x")
    p = plot!(xaxis=:log, yaxis=:log, legend=:outertopright)
    p = ylims!(1e-8, 30)

    return p
end

function get_input_pdf_func(hyper_params::PDFParameters)

    hp = hyper_params

    func = function _input_pdfs(i, x)::Float64
        i = i[]
        x = x[]

        f = 0.0

        # gluon
        if (i == 0)
            f = x_g_x(x, hp.λ_g1, hp.λ_g2, hp.K_g, hp.θ[1], hp.θ[2])
        end

        # u valence
        if (i == 1)
            f = x_uv_x(x, hp.λ_u, hp.K_u)
        end

        # d valence
        if (i == 2)
            f = x_dv_x(x, hp.λ_d, hp.K_d)
        end

        # ubar
        if (i == 3)
            f = x_q_x(x, hp.λ_q, hp.θ[3])
        end

        # dbar
        if (i == 4)
            f = x_q_x(x, hp.λ_q, hp.θ[4])
        end

        # s and sbar
        if (i == 5) || (i == 6)
            f = x_q_x(x, hp.λ_q, hp.θ[5])
        end

        # c and cbar
        if (i == 7) || (i == 8)
            f = x_q_x(x, hp.λ_q, hp.θ[6])
        end

        # d and dbar
        if (i == 9) || (i == 10)
            f = x_q_x(x, hp.λ_q, hp.θ[7])
        end

        return f
    end

    return func
end

#                         tb  bb  cb  sb  ub  db  g   d   u   s   c   b   t
input_pdf_map = Float64.([0., 0., 0., 0.,-1., 0., 0., 0., 1., 0., 0., 0., 0., # 1 # U valence
                          0., 0., 0., 0., 0.,-1., 0., 1., 0., 0., 0., 0., 0., # 2 # D valence
                          0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., # 3 # u sea
                          0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., # 4 # d sea
                          0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., # 5 # s
                          0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., # 6 # sbar
                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., # 7 # c
                          0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., # 8 # cbar
                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., # 9 # b
                          0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., # 10 # bbar
                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., # 11
                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]); # 12


text:
@testset "polblogs" begin

#io = ZipFile.Reader("gml/polblogs.zip").files[1]
#println("polblogs")
#io = open("gml/polblogs.gml")
#parse_dict(io; debug=false)
#io = open("gml/polblogs.gml")
#@time parse_dict(io; debug=false)

#println("polblogs id")
#io = open("gml/polblogs.gml")
#parse_id_dict(io; debug=false)
#io = open("gml/polblogs.gml")
#@time parse_id_dict(io; debug=false)

s = open(s -> read(s, String), "gml/polblogs.gml")
x = parse_dict(s; debug=false)
@test length(x[:graph][1][:edge]) == 19090
@time parse_raw(s; debug=false)

end


text:
"""
    drawtrappedpointsR2(f [; hasescaped, maxiterations])

Return the drawing of the trapped points set of a function
    \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$,
    in a rectangular region in \$\\mathbb{R}^2\$,
    using the escape time of iterations algorithm.

The trapped points set of \$f\$ is defined as
\$\\mathcal{K}(f)=\\{(x,y)\\in\\mathbb{R}^2\\,|\\,|f^n(x,y)|\\nrightarrow\\infty\\,n\\rightarrow\\infty\\}\$

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `hasescaped::Function`: A boolean function to check if the iterations has escaped.
- `maxiterations::Integer`: Maximum number of iterations to check.
"""
function drawtrappedpointsR2(f::Function;
    hasescaped::Function=(x::Real,y::Real) -> x*x+y*y > 4, maxiterations::Int=100)

    # Veryfying if graphics backend supports functions
    SDDGraphics.supported(:drawpixel)

    # Verifying functions
    @assert typeof(f(1., 1.)) <: Tuple{Real,Real}
    @assert typeof(hasescaped(1., 1.)) <: Bool

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(maxiterations)

    @sweeprectregion SDDGraphics.xlims() SDDGraphics.ylims() SDDGraphics.canvassize() begin
        xn,yn = x,y
        escapetime = maxiterations
        for n in 1:maxiterations
            if hasescaped(xn,yn)
                escapetime = n
                break
            end # if hasescaped
            xn,yn = f(xn,yn)
        end # for n maxiterations
        SDDGraphics.color(escapetime)
        SDDGraphics.drawpixel(i,j)
    end # Implemented algorithm

    SDDGraphics.drawing()

#= # Classical code
    width, height = SDDGraphics.canvassize()
    xmin, xmax = xlims
    ymin, ymax = ylims
    Δx, Δy = (xmax-xmin)/width, (ymax-ymin)/height

    x, y = 0., 0.

    y = ymax
    for j in 1:height

        x = xmin
        for i in 1:width

            xn, yn = x, y
            escapetime = 1.
            for n in 0:maxiterations

                if hasescaped(xn,yn)
                    escapetime = n/maxiterations
                    break
                end # if hasescaped

                xn,yn = f(xn,yn)

            end # for n maxiterations

            SDDGraphics.drawpixel(j,i,escapetime)

            x += Δx

        end # for i width

        y -= Δy

    end # for j height

    SDDGraphics.plot()
=#
end # function drawtrappedpointsR2

"""
    drawtrappedpointsC(f [; hasescaped, maxiterations])

Return the drawing of the trapped points set of a function
    \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$,
    in a rectangular region in \$\\mathbb{C}\$,
    using the escape time of iterations algorithm.

The trapped points set of \$f\$ is defined as
\$\\mathcal{K}(f)=\\{z\\in\\mathbb{C}\\,|\\,|f^n(z)|\\nrightarrow\\infty\\,n\\rightarrow\\infty\\}\$

#### Arguments
- `f::Function`: A function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$.
- `hasescaped::Function`: A boolean function to check if the iterations has escaped.
- `maxiterations::Integer`: Maximum number of iterations to check.
"""
function drawtrappedpointsC(f::Function;
    hasescaped::Function=z::Number -> abs2(z) > 4, maxiterations::Int=100)

    # Veryfying if graphics backend supports functions
    SDDGraphics.supported(:drawpixel)

    # Verifying functions
    @assert typeof(f(1.0im)) <: Number
    @assert typeof(hasescaped(1.0im)) <: Bool

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(maxiterations)

    @sweeprectregion SDDGraphics.xlims() SDDGraphics.ylims() SDDGraphics.canvassize() begin
        zn = complex(x,y)
        escapetime = maxiterations
        for n in 1:maxiterations
            if hasescaped(zn)
                escapetime = n
                break
            end # if hasescaped
            zn = f(zn)
        end # for n maxiterations
        SDDGraphics.color(escapetime)
        SDDGraphics.drawpixel(i,j)
    end # Implemented algorithm

    SDDGraphics.drawing()

end # function drawtrappedpointsC