text:
"""
    DifferentialOperator(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition}, method::DifferenceMethod)

Returns a discretized differential operator of `length(x̄)` by `length(x̄)` matrix
under mixed boundary conditions from `bc` using finite difference method specified by `method`.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> DifferentialOperator(x̄, (Reflecting(), Reflecting()), BackwardFirstDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
  0.0   0.0    ⋅    ⋅
 -1.0   1.0   0.0   ⋅
   ⋅   -1.0   1.0  0.0
   ⋅     ⋅   -1.0  1.0

julia> DifferentialOperator(x̄, (Reflecting(), Reflecting()), ForwardFirstDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   1.0    ⋅    ⋅
  0.0  -1.0   1.0   ⋅
   ⋅    0.0  -1.0  1.0
   ⋅     ⋅    0.0  0.0

julia> DifferentialOperator(x̄, (Reflecting(), Reflecting()), CentralSecondDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   1.0    ⋅     ⋅
  1.0  -2.0   1.0    ⋅
   ⋅    1.0  -2.0   1.0
   ⋅     ⋅    1.0  -1.0

julia> x̄ = 0:5
0:5

julia> DifferentialOperator(x̄, (Mixed(ξ = 1.0), Mixed(ξ = 1.0)), BackwardFirstDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   0.0    ⋅    ⋅
 -1.0   1.0   0.0   ⋅
   ⋅   -1.0   1.0  0.0
   ⋅     ⋅   -1.0  1.0

julia> DifferentialOperator(x̄, (Mixed(ξ = 1.0), Mixed(ξ = 1.0)), ForwardFirstDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   1.0    ⋅     ⋅
  0.0  -1.0   1.0    ⋅
   ⋅    0.0  -1.0   1.0
   ⋅     ⋅    0.0  -1.0

julia> DifferentialOperator(x̄, (Mixed(ξ = 1.0), Mixed(ξ = 1.0)), CentralSecondDifference())
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 0.0   1.0    ⋅     ⋅
 1.0  -2.0   1.0    ⋅
  ⋅    1.0  -2.0   1.0
  ⋅     ⋅    1.0  -2.0
```
"""
function DifferentialOperator(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition}, method::BackwardFirstDifference)
    # reflecting bcs are special cases of mixed bcs with ξ = 0
    if (typeof(bc[1]) <: Reflecting)
        return DifferentialOperator(x̄, (Mixed(ξ = 0), bc[2]), method)
    end
    if (typeof(bc[2]) <: Reflecting)
        return DifferentialOperator(x̄, (bc[1], Mixed(ξ = 0)), method)
    end

    # setup for operator
    M = length(x̄) - 2
    d = diff(x̄)
    Δ₋⁻¹ = 1 ./ d[1:end-1] # (1 ./ Δ₋)
    T = eltype(Δ₋⁻¹)

    # construct the operator
    L = Tridiagonal(-Δ₋⁻¹[2:M], Δ₋⁻¹, zeros(T, M-1))

    # setup for boundary conditions
    Δ_1m = x̄[2] - x̄[1]

    # apply boundary conditions
    # (under homogeneous absorbing bc on lb, the first column in invariant)
    if !(typeof(bc[1]) <: Absorbing)
        ξ_lb = bc[1].ξ
        L[1,1] += (bc[1].direction == :backward) ? (-1/Δ_1m - ξ_lb) : 1/(-1+ξ_lb*Δ_1m)/Δ_1m
    end

    return L
end

function DifferentialOperator(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition}, method::ForwardFirstDifference)
    # reflecting bcs are special cases of mixed bcs with ξ = 0
    if (typeof(bc[1]) <: Reflecting)
        return DifferentialOperator(x̄, (Mixed(ξ = 0), bc[2]), method)
    end
    if (typeof(bc[2]) <: Reflecting)
        return DifferentialOperator(x̄, (bc[1], Mixed(ξ = 0)), method)
    end

    # setup for operator
    M = length(x̄) - 2
    d = diff(x̄)
    Δ₊⁻¹ = 1 ./ d[2:end] # (1 ./ Δ₊), extracting elements on the interior
    T = eltype(Δ₊⁻¹)

    # construct the operator
    L = Tridiagonal(zeros(T, M-1), -Δ₊⁻¹, Δ₊⁻¹[1:M-1])

    # setup for boundary conditions
    Δ_Mp = x̄[end] - x̄[end-1]

    # apply boundary conditions
    # (under homogeneous absorbing bc on ub, the last column in invariant)
    if !(typeof(bc[2]) <: Absorbing)
        ξ_ub = bc[2].ξ
        L[end,end] += (bc[2].direction == :forward) ? (1/Δ_Mp - ξ_ub) : 1/(1+ξ_ub*Δ_Mp)/Δ_Mp
    end

    return L
end

function DifferentialOperator(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition}, method::CentralSecondDifference)
    # reflecting bcs are special cases of mixed bcs with ξ = 0
    if (typeof(bc[1]) <: Reflecting)
        return DifferentialOperator(x̄, (Mixed(ξ = 0), bc[2]), method)
    end
    if (typeof(bc[2]) <: Reflecting)
        return DifferentialOperator(x̄, (bc[1], Mixed(ξ = 0)), method)
    end

    # setup for operators
    M = length(x̄) - 2
    d = diff(x̄)
    Δ₋⁻¹ = 1 ./ d[1:end-1] # 1 ./ Δ₋
    Δ₊⁻¹ = 1 ./ d[2:end] # 1 ./ Δ₊
    Δ⁻¹ = 1 ./ (d[1:end-1] + d[2:end]) # 1 ./ (Δ₋ + Δ₊)

    # construct the operator
    L = 2*Tridiagonal((Δ⁻¹.*Δ₋⁻¹)[2:M], -Δ₋⁻¹ .* Δ₊⁻¹, (Δ⁻¹.*Δ₊⁻¹)[1:M-1])

    # setup for boundary conditions
    Δ_1p = x̄[3] - x̄[2]
    Δ_1m = x̄[2] - x̄[1]
    Δ_Mp = x̄[end] - x̄[end-1]
    Δ_Mm = x̄[end-1] - x̄[end-2]

    # apply boundary conditions
    # (under homogeneous absorbing bc on lb, the first column in invariant)
    if !(typeof(bc[1]) <: Absorbing)
        ξ_lb = bc[1].ξ
        Ξ_1p = L[1,1] - 2/((-1+ξ_lb*Δ_1m)*(Δ_1p+Δ_1m)*(Δ_1m))
        Ξ_1m = 2*(-1/(Δ_1p*Δ_1m) + (1+ξ_lb*Δ_1m)/(Δ_1p+Δ_1m)/Δ_1m)
        L[1,1] = (bc[1].direction == :backward) ? Ξ_1m : Ξ_1p
    end
    # (under homogeneous absorbing bc on ub, the last column in invariant)
    if !(typeof(bc[2]) <: Absorbing)
        ξ_ub = bc[2].ξ
        Ξ_Mm = L[end,end] + 2/((1+ξ_ub*Δ_Mp)*(Δ_Mp+Δ_Mm)*(Δ_Mp))
        Ξ_Mp = 2*(-1/(Δ_Mp*Δ_Mm) - (-1+ξ_ub*Δ_Mp)/(Δ_Mp+Δ_Mm)/Δ_Mp)
        L[end,end] = (bc[2].direction == :forward) ? Ξ_Mp : Ξ_Mm
    end
    return L
end

# Convenience calls
"""
    L₁₋bc(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition})

Returns a discretized first-order differential operator of `length(x̄)` by `length(x̄)` matrix
using backward difference under boundary conditions specified by `bc`.

The first element of `bc` is applied to the lower bound, and second element of `bc` to the upper.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> L₁₋bc(x̄, (Reflecting(), Reflecting()))
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
  0.0   0.0    ⋅    ⋅
 -1.0   1.0   0.0   ⋅
   ⋅   -1.0   1.0  0.0
   ⋅     ⋅   -1.0  1.0
```
"""
L₁₋bc(x̄, bc) = DifferentialOperator(x̄, bc, BackwardFirstDifference())

"""
    L₁₊bc(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition})

Returns a discretized first-order differential operator of `length(x̄)` by `length(x̄)` matrix
using forward difference under boundary conditions specified by `bc`.

The first element of `bc` is applied to the lower bound, and second element of `bc` to the upper.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> L₁₊bc(x̄, (Reflecting(), Reflecting()))
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   1.0    ⋅    ⋅
  0.0  -1.0   1.0   ⋅
   ⋅    0.0  -1.0  1.0
   ⋅     ⋅    0.0  0.0
```
"""
L₁₊bc(x̄, bc) = DifferentialOperator(x̄, bc, ForwardFirstDifference())

"""
    L₂bc(x̄, bc::Tuple{BoundaryCondition, BoundaryCondition})

Returns a discretized second-order differential operator of `length(x̄)` by `length(x̄)` matrix
using central difference under boundary conditions specified by `bc`.

The first element of `bc` is applied to the lower bound, and second element of `bc` to the upper.
# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> L₂bc(x̄, (Reflecting(), Reflecting()))
4×4 LinearAlgebra.Tridiagonal{Float64,Array{Float64,1}}:
 -1.0   1.0    ⋅     ⋅
  1.0  -2.0   1.0    ⋅
   ⋅    1.0  -2.0   1.0
   ⋅     ⋅    1.0  -1.0
```
"""
L₂bc(x̄, bc) = DifferentialOperator(x̄, bc, CentralSecondDifference())

"""
    L₁₋(x̄)

Returns a discretized first-order differential operator of `length(x̄)` by `length(x̄) + 2` matrix
using backward difference under no boundary condition.

The first and last columns are applied to the ghost nodes just before `x̄[1]` and `x̄[end]` respectively.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 1:3
1:3

julia> Array(L₁₋(x̄))
1×3 Array{Float64,2}:
 -1.0  1.0  0.0
```
"""
L₁₋(x̄) = ExtensionDifferentialOperator(x̄, BackwardFirstDifference())

"""
    L₁₊(x̄)

Returns a discretized first-order differential operator of `length(x̄)` by `length(x̄) + 2` matrix using
forward difference under no boundary condition.

The first and last columns are applied to the ghost nodes just before `x̄[1]` and `x̄[end]` respectively.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> Array(L₁₊(x̄))
4×6 Array{Float64,2}:
 0.0  -1.0   1.0   0.0   0.0  0.0
 0.0   0.0  -1.0   1.0   0.0  0.0
 0.0   0.0   0.0  -1.0   1.0  0.0
 0.0   0.0   0.0   0.0  -1.0  1.0
```
"""
L₁₊(x̄) = ExtensionDifferentialOperator(x̄, ForwardFirstDifference())

"""
    L₂(x̄)

Returns a discretized second-order differential operator of `length(x̄)` by `length(x̄) + 2` matrix
using central difference under no boundary condition.

The first and last columns are applied to the ghost nodes just before `x̄[1]` and `x̄[end]` respectively.

# Examples
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> Array(L₂(x̄))
4×6 Array{Float64,2}:
 1.0  -2.0   1.0   0.0   0.0  0.0
 0.0   1.0  -2.0   1.0   0.0  0.0
 0.0   0.0   1.0  -2.0   1.0  0.0
 0.0   0.0   0.0   1.0  -2.0  1.0
```
"""
L₂(x̄)  = ExtensionDifferentialOperator(x̄, CentralSecondDifference())

"""
    interiornodes(x̄)

Returns an interior grid of length `length(x̄)-2` given extended grid `x̄`.
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> interiornodes(x̄)
1:4

julia> x̄ = [1.0; 1.5; 1.7]
3-element Array{Float64,1}:
 1.0
 1.5
 1.7

julia> interiornodes(x̄)
1-element Array{Float64,1}:
 1.5
```
"""
interiornodes(x̄) = x̄[2:end-1]

"""
    interiornodes(x̄, bc)

Returns an interior grid corresponding to the boundary condition `bc` given extended grid `x̄`.
```jldoctest; setup = :(using SimpleDifferentialOperators)
julia> x̄ = 0:5
0:5

julia> interiornodes(x̄, (Reflecting(), Reflecting()))
1:4

julia> x̄ = [1.0; 1.5; 1.7]
3-element Array{Float64,1}:
 1.0
 1.5
 1.7

julia> interiornodes(x̄, (Mixed(ξ = 1.0), Mixed(ξ = 1.0)))
1-element Array{Float64,1}:
 1.5
```
"""
interiornodes(x̄, bc) = interiornodes(x̄)


text:
# the show function
function show(io::IO, v::Node)
    print(io, "Node($(v.key))")
end

function show(io::IO, v::AttributeNode)
    print(io, "AttributeNode($(v.key))")
end

function show(io::IO, v::TimeNode)
    print(io, "TimeNode($(v.key), $(v.timestamp))")
end

function show(io::IO, e::Edge)
    print(io, "$(e.source)->$(e.target)")
end

function show(io::IO, e::TimeEdge)
    print(io, "$(e.source)->$(e.target) at time $(e.timestamp)")
end

function show(io::IO, e::WeightedTimeEdge)
    print(io, "$(e.source)-$(e.weight)->$(e.target) at time $(e.timestamp)")
end

function show{V, E}(io::IO, g::DiGraph{V, E})
    print(io, "DiGraph $(num_nodes(g)) nodes, $(num_edges(g)) edges")
end

function show{V, E, T, KV}(io::IO, g::EvolvingGraph{V, E, T, KV})
    title = is_directed(g) ? "Directed EvolvingGraph" : "Undirected EvolvingGraph"
    print(io, "$(title) $(num_nodes(g)) nodes, $(num_edges(g)) static edges, $(num_timestamps(g)) timestamps")
end

function show(io::IO, g::IntAdjacencyList)
    title = is_directed(g) ? "Directed IntAdjacencyList" : "Undirected IntAdjacencyList"
    print(io, "$(title) ($(num_nodes(g)) nodes, $(num_edges(g)) static edges, $(num_timestamps(g)) timestamps)")
end

function show(io::IO, g::MatrixList)
    print(io, "MatrixList ($(num_matrices(g)) matrices)")
end

function show(io::IO, p::AbstractPath)
    result = ""
    for i in 1:length(p)
        if i == length(p)
            result = string(result, p.nodes[i])
        else
            result = string(result, p.nodes[i], "->")
        end
    end
    print(io, "$(result)")
end

ext: jl


text:
# Compression
# #+HTML: Source
# #+HTML:

# [[file:~/.julia/dev/MatrixProductStates/README.org::*Compression][Compression:1]]
function compress(ψ::MPS{L, T}, to_the::Right; Dcut::Int=typemax(Int)) where {L, T}
    tensors = Array{T, 3}[]

    B = ψ[1]
    d = length(B[1, 1, :])

    @cast Bm[(σ¹, a⁰), a¹] |= B[a⁰, a¹, σ¹]
    U, S, V = psvd(Bm, rank=Dcut)
    #S = S/√sum(S .^ 2)

    @cast A[a⁰, a¹, σ¹] |= U[(σ¹, a⁰), a¹] (σ¹:d)
    push!(tensors, A)

    for i ∈ 2:L
        B = ψ[i]
        d = length(B[1, 1, :])

        @tensor M[aⁱ⁻¹, aⁱ, σⁱ] := (Diagonal(S)*V')[aⁱ⁻¹, aⁱ⁻¹′] * B[aⁱ⁻¹′, aⁱ, σⁱ]
        @cast   Mm[(σⁱ, aⁱ⁻¹), aⁱ] |= M[aⁱ⁻¹, aⁱ, σⁱ]

        U, S, V = psvd(Mm, rank=Dcut)
        #S = S/√sum(S .^ 2)

        @cast A[aⁱ⁻¹, aⁱ, σⁱ] |= U[(σⁱ, aⁱ⁻¹), aⁱ] (σⁱ:d)
        push!(tensors, A)
    end
    MPS{L, T}(tensors), Left()
end

leftcanonical(ψ) = compress(ψ, right)[1]

function compress(ψ::MPS{L, T}, to_the::Left; Dcut::Int=typemax(Int)) where {L, T}
    tensors = Array{T, 3}[]

    A = ψ[L]
    d = length(A[1, 1, :])
    @cast Am[aᴸ⁻¹, (σᴸ, aᴸ)] |= A[aᴸ⁻¹, aᴸ, σᴸ]

    U, S, V = psvd(Am, rank=Dcut)
    #S = S/√sum(S .^ 2)

    @cast B[aᴸ⁻¹, aᴸ, σᴸ] |= V'[aᴸ⁻¹, (σᴸ, aᴸ)] (σᴸ:d)
    push!(tensors, B)

    for i ∈ (L-1):-1:1
        A = ψ[i]
        d = length(A[1, 1, :])
        @tensor M[aⁱ⁻¹, aⁱ, σⁱ]    := A[aⁱ⁻¹, aⁱ′, σⁱ] * (U * Diagonal(S))[aⁱ′, aⁱ]
        @cast   Mm[aⁱ⁻¹, (σⁱ, aⁱ)] |= M[aⁱ⁻¹, aⁱ, σⁱ]

        U, S, V = psvd(Mm, rank=Dcut)
        #S = S/√sum(S .^ 2)

        @cast B[aⁱ⁻¹, aⁱ, σⁱ] |= V'[aⁱ⁻¹, (σⁱ, aⁱ)] (σⁱ:d)
        push!(tensors, B)
    end
    MPS{L, T}(reverse(tensors)), Right()
end

rightcanonical(ψ) = compress(ψ, left)[1]

compress(ψ; Dcut) = compress(ψ, left, Dcut=Dcut)[1]
# Compression:1 ends here


text:
using VSFlow
using Test

@testset "VSFlow.jl" begin
    @test 1+1 == 2
    # Write your tests here.
end


text:
include("generic.jl")
include("discrete.jl")
include("continuous.jl")
include("HopfieldQAP.jl")


text:
@with_kw mutable struct Sampler{P, V, Pol <: Policy, T1 <: AbstractSpace, T2 <: AbstractSpace}
    mdp::P
    agent::PolicyParams{Pol, T2} # The agent
    adversary = nothing # The adversary
    s::V = rand(initialstate(mdp)) # Current State
    S::T1 = state_space(initial_observation(mdp, s)) # State space
    svec::AbstractArray = tovec(initial_observation(mdp, s), S) # Current observation
    max_steps::Int = 100
    required_columns::Array{Symbol} = !isnothing(adversary) ? [:x] : []
    γ::Float32 = discount(mdp)
    λ::Float32 = NaN32
    episode_length::Int64 = 0
    episode_checker::Function = (data, start, stop) -> true

    # Parameters for cost constraints
    Vc::Union{ContinuousNetwork, Nothing} = nothing
    λcost::Float32 = NaN32
end

Sampler(mdp, π::T; kwargs...) where {T <: Policy} = Sampler(;mdp=mdp, agent=PolicyParams(π), kwargs...)
Sampler(mdp, agent::T; kwargs...) where {T <: PolicyParams} = Sampler(;mdp=mdp, agent=agent, kwargs...)

# Construct a vector of samplers from a vector of mdps
Sampler(mdps::AbstractVector, π::T; kwargs...) where {T <: Policy} = [Sampler(mdps[i], agent; kwargs...) for i in 1:length(mdps)]
Sampler(mdps::AbstractVector, agent::T; kwargs...) where {T <: PolicyParams} = [Sampler(mdps[i], agent; kwargs...) for i in 1:length(mdps)]

function reset_sampler!(sampler::Sampler)
    sampler.s = rand(initialstate(sampler.mdp))
    sampler.svec = tovec(initial_observation(sampler.mdp, sampler.s), sampler.S)
    sampler.episode_length = 0
end

function initial_observation(mdp, s)
    if mdp isa POMDP
        return convert_o(AbstractArray, rand(initialobs(mdp, s)), mdp)
    else
        return convert_s(AbstractArray, s, mdp)
    end
end

function terminate_episode!(sampler::Sampler, data, j)
    haskey(data, :episode_end) && (data[:episode_end][1,j] = true)
    ep = j - sampler.episode_length + 1 : j
    haskey(data, :advantage) && fill_gae!(data, ep, sampler.agent.π, sampler.λ, sampler.γ)
    haskey(data, :return) && fill_returns!(data, ep, sampler.γ)

    # Dealing with cost constraints
    haskey(data, :cost_advantage) && fill_gae!(data, ep, sampler.Vc, sampler.λ, sampler.γ, source=:cost, target=:cost_advantage)
    haskey(data, :cost_return) && fill_returns!(data, ep, sampler.γ, source=:cost, target=:cost_return)

    reset_sampler!(sampler)
end

function step!(data, j::Int, sampler::Sampler; explore=false, i=0)
    a, logprob = explore ? exploration(sampler.agent.π_explore, sampler.svec, π_on=sampler.agent.π, i=i) : (action(sampler.agent.π, sampler.svec), NaN)
    (a isa AbstractArray || a isa Tuple) && length(a) == 1 && (a = a[1])

    # This implements the ability to get cost information from safety gym
    info = Dict()
    kwargs = haskey(data, :cost) ? (info=info,) : ()

    args = (a,)
    if !isnothing(sampler.adversary)
        x, xlogprob = explore ? exploration(sampler.adversary.π_explore, sampler.svec, π_on=sampler.adversary.π, i=i) : (action(sampler.adversary.π, sampler.svec), NaN)
        (x isa AbstractArray || x isa Tuple) && length(x) == 1 && (x = x[1]) # disturbances always come out as an array
        data[:x][:, j:j] .= tovec(x, sampler.adversary.space)
        haskey(data, :xlogprob) && (data[:xlogprob][:, j] .= xlogprob)
        args = (a, x)
    end

    if sampler.mdp isa POMDP
        sp, o, r = gen(sampler.mdp, sampler.s, args...; kwargs...)
        spvec = convert_o(AbstractArray, o, sampler.mdp)
    else
        sp, r = gen(sampler.mdp, sampler.s, args...; kwargs...)
        spvec = convert_s(AbstractArray, sp, sampler.mdp)
    end
    spvec = tovec(spvec, sampler.S)
    done = isterminal(sampler.mdp, sp)

    # Save the tuple
    bslice(data[:s], j:j) .= sampler.svec
    data[:a][:, j:j] .= tovec(a, sampler.agent.space)
    bslice(data[:sp], j:j) .= spvec
    data[:r][1, j] = r
    data[:done][1, j] = done

    # Handle optional data storage
    haskey(data, :logprob) && (data[:logprob][:, j] .= logprob)
    haskey(data, :t) && (data[:t][1, j] = sampler.episode_length + 1)
    haskey(data, :i) && (data[:i][1, j] = i+1)
    haskey(data, :cost) && (data[:cost][1,j] = info["cost"])
    haskey(data, :fail) && (data[:fail][1,j] = extra_functions["isfailure"](sampler.mdp, sp)) #TODO Changed this to "s" instead of "sp" for the continuum world

    # Cut the episode short if needed
    sampler.episode_length += 1
    if done || sampler.episode_length >= sampler.max_steps
        terminate_episode!(sampler, data, j)
    else
        sampler.s = sp
        sampler.svec = spvec
    end
end

function steps!(sampler::Sampler; Nsteps = 1, explore=false, i=0, reset=false, return_episodes=false, return_at_episode_end=false)
    data = mdp_data(sampler.S, sampler.agent.space, Nsteps, return_episodes ? [sampler.required_columns..., :episode_end] : sampler.required_columns)
    for j=1:Nsteps
        step!(data, j, sampler, explore=explore, i=i + (j-1))
        if return_at_episode_end && sampler.episode_length == 0
            trim!(data, j)
            break
        end
    end
    reset && terminate_episode!(sampler, data, Nsteps)
    return_episodes ? (data, episodes(data)) : data
end

function steps!(samplers::Vector{T}; Nsteps = 1, explore = false, i = 0, reset = false, return_episodes = false) where {T<:Sampler}
    data = mdp_data(samplers[1].S, samplers[1].agent.space, Nsteps*length(samplers), return_episodes ? [samplers[1].required_columns..., :episode_end] : samplers[1].required_columns)
    j = 1
    for s=1:Nsteps
        for sampler in samplers
            step!(data, j, sampler, explore = explore, i = i + (j-1))
            j += 1
        end
    end
    reset && terminate_episode!(sampler, data, Nsteps)
    return_episodes ? (data, episodes(data)) : data
end

function episodes!(sampler::Sampler; Neps=1, explore=false, i=0, return_episodes=false)
    reset_sampler!(sampler)
    data = mdp_data(sampler.S, sampler.agent.space,
     Neps*sampler.max_steps, sampler.required_columns)
    episode_starts = zeros(Int, Neps)
    episode_ends = zeros(Int, Neps)
    j, k = 0, 1
    while k <= Neps
        episode_starts[k] = j+1
        while true
            j = j+1
            step!(data, j, sampler, explore=explore, i=i + (i-1))
            if sampler.episode_length == 0
                episode_ends[k] = j
                sampler.episode_checker(data, episode_starts[k], j) ? (k = k+1) : (j = episode_starts[k])
                break
            end
        end
    end
    trim!(data, j)
    return_episodes ? (data, zip(episode_starts, episode_ends)) : data
end

function fillto!(b::ExperienceBuffer, s::Union{Sampler, Vector{T}}, N::Int; i=1, explore=false) where {T <: Sampler}
    Nfill = max(0, N - length(b))
    Nfill > 0 && push!(b, steps!(s, i=i, Nsteps=Nfill, explore=explore))
    Nfill
end

## metric

# Recover multiple metrics from a single sampler
metrics_by_key(data, start, stop; keys) = [sum(data[key][1,start:stop]) for key in keys]

function metrics_by_key(s::Sampler; keys, Neps=100, kwargs...)
    data = episodes!(s, Neps=Neps; kwargs...)
    [sum(data[key]) / Neps for key in keys]
end

# recover a single metric
metric_by_key(data, start, stop; key) = metrics_by_key(data, start, stop; keys=[key])[1]

metric_by_key(s::Sampler; key, Neps=100, kwargs...) = metrics_by_key(s; keys=[key], Neps=Neps, kwargs...)[1]

# Get the undiscounted return
undiscounted_return(data, start, stop) = metric_by_key(data, start, stop; key=:r)
undiscounted_return(s::Sampler; Neps=100, kwargs...) = metric_by_key(s; Neps=Neps, key=:r, kwargs...)

## Discounted returns
function discounted_return(data, start, stop, γ)
    r = 0f0
    for i in reverse(start:stop)
        r = data[:r][1, i] + γ*r
    end
    r
end

function discounted_return(s::Sampler; Neps=100, kwargs...)
    data, episodes = episodes!(s, Neps=Neps, return_episodes=true; kwargs...)
    mean([discounted_return(data, e..., discount(s.mdp)) for e in episodes])
end

## Failures
failure(data, start, stop; threshold = 0.) = undiscounted_return(data, start, stop) < threshold

function failure(s::Sampler; threshold=0., Neps=100, kwargs...)
    data, episodes = episodes!(s, Neps = Neps, return_episodes = true; kwargs...)
    mean([failure(data, e..., threshold = threshold) for e in episodes])
end

## Generalized Advantage Estimation
function fill_gae!(d::ExperienceBuffer, V, λ::Float32, γ::Float32)
    eps = episodes(d)
    for ep in eps
        fill_gae!(d, ep, V, λ, γ)
    end
end

function fill_gae!(d, episode_range, V, λ::Float32, γ::Float32; source = :r, target = :advantage)
    A, c = 0f0, λ*γ
    nd = ndims(d[:s])
    for i in reverse(episode_range)
        Vsp = value(V, bslice(d[:sp], i:i))
        Vs = value(V, bslice(d[:s], i:i))
        @assert length(Vs) == 1
        A = c*A + d[source][1,i] + (1.f0 - d[:done][1,i])*γ*Vsp[1] - Vs[1]
        @assert !isnan(A)
        d[target][:, i] .= A
    end
end

function fill_returns!(data, episode_range, γ::Float32; source=:r, target=:return)
    r = 0f0
    for i in reverse(episode_range)
        r = data[source][1, i] + γ*r
        data[target][:, i] .= r
    end
end

# Utils
function trim!(data::Dict{Symbol, Array}, N)
    for k in keys(data)
        data[k] = bslice(data[k], 1:N)
    end
    data
end


text: Pkg.checkout("XGBoost")


text:
include("SumSpace.jl")
include("ArraySpace.jl")
include("ProductSpaceOperators.jl")
include("SubSpace.jl")
include("QuotientSpace.jl")

⊕(A::Space,B::Space) = domainscompatible(A,B) ? SumSpace(A,B) : PiecewiseSpace(A,B)
⊕(f::Fun,g::Fun) = Fun(space(f) ⊕ space(g), interlace(coefficients(f),coefficients(g)))

⊕(f::Fun,g::Fun,h::Fun...) = ⊕((f ⊕ g), h...)

+(A::Space,B::Space) = A ⊕ B

#split the cfs into component spaces
function coefficients(cfs::AbstractVector,A::SumSpace,B::SumSpace)
    if spacescompatible(A,B)
        cfs
    else
        mapreduce(f->Fun(f,B),+,components(Fun(A,cfs))).coefficients
    end
end
function coefficients(cfs::AbstractVector,A::PiecewiseSpace,B::PiecewiseSpace)
    if spacescompatible(A,B)
        cfs
    else
        mapreduce(f->Fun(f,B),+,components(Fun(A,cfs))).coefficients
    end
end

# spread a single space into a sum space by placing
# its coefficients depending on k
function interlacewithzeros(cfs::AbstractVector,k,it)
    n = length(cfs)

    ret = Array{eltype(cfs)}(undef, 0)
    n == 0 && return ret

    for (K,j) in it
        if K == k
            push!(ret,cfs[j])
            n == j && break
        else
            push!(ret,0)
        end
    end

    ret
end

interlacewithzeros(cfs::AbstractVector,k,B::DirectSumSpace) = interlacewithzeros(cfs,k,interlacer(B))

function sumspacecoefficients(cfsin::AbstractVector,A::Space,B::SumSpace)
    m=length(components(B))

    for k=1:m
        if isconvertible(A,component(B,k))
            cfs = coefficients(cfsin,A,component(B,k))
            return interlacewithzeros(cfs,k,B)
        end
    end

    defaultcoefficients(cfsin,A,B)
end

function sumspacecoefficients(cfsin::AbstractVector,A::Space,B::PiecewiseSpace)
    m=length(components(B))

    for k=1:m
        if domain(component(B,k)) == domain(A) && isconvertible(A,component(B,k))
            cfs = coefficients(cfsin,A,component(B,k))
            return interlacewithzeros(cfs,k,B)
        end
    end

    defaultcoefficients(cfsin,A,B)
end

for TYP in (:SumSpace,:PiecewiseSpace), ATYP in (:ConstantSpace,:(ConstantSpace{<:Domain{<:Number}}),:Space)
    @eval coefficients(cfsin::AbstractVector,A::$ATYP,B::$TYP) = sumspacecoefficients(cfsin,A,B)
end

## LowRank Constructors

# convert a vector of functionals and an operator to a LowRnakPertOperator
# the rangespace is a DirectSumSpace specified by ST of the input rangespaces
# the default is a  ArraySpace, but support is there for PiecewiseSpace
# for bcs
for TYP in (:PiecewiseSpace,:VectorSpace)
    @eval function LowRankPertOperator(A::AbstractVector{OT},::Type{$TYP}) where OT<:Operator
        A=promotedomainspace(A)
        for k=1:length(A)-1
            @assert isafunctional(A[k])
        end
        @assert isbanded(A[end])
        L=LowRankOperator(A[1:end-1],$TYP)
        # add zero functionals to shift down
        BB=[fill(ZeroOperator(domainspace(BB),ConstantSpace()),length(A)-1);A[end]]
        S=InterlaceOperator(BB,domainspace(BB),$TYP(map(rangespace,A)))
        L+S
    end
end

LowRankPertOperator(A::AbstractVector{OT}) where {OT<:Operator}=LowRankPertOperator(A,VectorSpace)

function LowRankOperator(Bin::AbstractVector{FT},::Type{PiecewiseSpace}) where FT<:Operator
    B=promotedomainspace(Bin)
    rsp=PiecewiseSpace(map(rangespace,B))
    LowRankOperator(
        VFun{typeof(rsp),Float64}[Fun(rsp,[zeros(k-1);1]) for k=1:length(B)],
        B)
end

function LowRankOperator(Bin::AbstractVector{FT},::Type{VectorSpace}) where FT<:Operator
    B=promotedomainspace(Bin)
    rsp=Space([map(rangespace,B);ZeroSpace()])  #TODO: Why the hack?
    LowRankOperator(
        VFun{typeof(rsp),Float64}[Fun(rsp,[zeros(k-1);1]) for k=1:length(B)],
        B)
end

LowRankOperator(Bin::AbstractVector{FT}) where {FT<:Operator} = LowRankOperator(Bin,VectorSpace)


text:
model,obj = TrajectoryOptimization.Dynamics.dubinscar
opts = TrajectoryOptimization.SolverOptions()
opts.verbose = false
solver = TrajectoryOptimization.Solver(model,obj,dt=0.1,opts=opts)
solver_sqrt = TrajectoryOptimization.Solver(model,obj,dt=0.1,opts=opts)
solver_sqrt.opts.square_root = true
U0 = ones(solver.model.m,solver.N)
results = init_results(solver,Array{Float64}(undef,0,0),U0)
results_sqrt = init_results(solver_sqrt,Array{Float64}(undef,0,0),U0)
results.X[1] = solver.obj.x0
copyto!(results.U, ones(model.m,solver.N-1))
results_sqrt.X[1] = solver_sqrt.obj.x0
copyto!(results_sqrt.U, ones(model.m,solver_sqrt.N-1))

TrajectoryOptimization.rollout!(results,solver)
TrajectoryOptimization.rollout!(results_sqrt,solver_sqrt)
TrajectoryOptimization.update_jacobians!(results,solver)
TrajectoryOptimization.update_jacobians!(results_sqrt,solver_sqrt)

ΔV = TrajectoryOptimization.backwardpass!(results,solver)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_sqrt,solver_sqrt)

@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results.K,results_sqrt.K))
@test isapprox(to_array(results.d),to_array(results_sqrt.d))
S_sqrt = [results_sqrt.S[k]'*results_sqrt.S[k] for k = 1:solver.N]
@test isapprox(to_array(results.S),to_array(S_sqrt))
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt.bp)
results_sqrt.bp
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt)
@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# backward pass square root for constrained solve
u_min = -10
u_max = 10
obj_c = TrajectoryOptimization.ConstrainedObjective(obj,u_min=u_min,u_max=u_max) # constrained objective
opts_con = TrajectoryOptimization.SolverOptions()
opts_con_sqrt = TrajectoryOptimization.SolverOptions()
opts_con_sqrt.square_root = true
solver_con = TrajectoryOptimization.Solver(model,obj_c,dt=0.1,opts=opts_con)
solver_con_sqrt = TrajectoryOptimization.Solver(model,obj_c,dt=0.1,opts=opts_con_sqrt)
U0 = ones(solver_con.model.m,solver_con.N)
results_con = init_results(solver_con,Array{Float64}(undef,0,0),U0)
results_con_sqrt = init_results(solver_con_sqrt,Array{Float64}(undef,0,0),U0)
TrajectoryOptimization.update_jacobians!(results_con, solver_con)
TrajectoryOptimization.update_jacobians!(results_con_sqrt, solver_con_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results_con, solver_con)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_con_sqrt, solver_con_sqrt)

@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results_con.K,results_con_sqrt.K))
@test all(isapprox.(results_con.d,results_con_sqrt.d))
@test all(isapprox.(results_con.s,results_con_sqrt.s))
S_con_sqrt = [results_con_sqrt.S[k]'*results_con_sqrt.S[k] for k = 1:solver_con.N]
@test isapprox(to_array(results_con.S),to_array(S_con_sqrt))
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results_con.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results_con)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_con_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_con_sqrt)
@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# square root + infeasible
solver_con_sqrt.state.infeasible = true
U0 = ones(solver_con.model.m,solver_con.N)
X0 = line_trajectory(solver_con)

results_con = init_results(solver_con,X0,U0)
results_con_sqrt = init_results(solver_con_sqrt,X0,U0)
TrajectoryOptimization.update_jacobians!(results_con, solver_con)
TrajectoryOptimization.update_jacobians!(results_con_sqrt, solver_con_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results_con, solver_con)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_con_sqrt, solver_con_sqrt)

@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results_con.K,results_con_sqrt.K))
@test all(isapprox.(results_con.d,results_con_sqrt.d))
@test all(isapprox.(results_con.s,results_con_sqrt.s))
S_con_sqrt = [results_con_sqrt.S[k]'*results_con_sqrt.S[k] for k = 1:solver_con.N]
@test isapprox(to_array(results_con.S),to_array(S_con_sqrt))
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results_con.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results_con)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_con_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_con_sqrt)
@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# Minimum Time
N = solver_con.N
obj_mintime = update_objective(obj_c,tf=:min)
opts_mintime = TrajectoryOptimization.SolverOptions()
opts_mintime_sqrt = TrajectoryOptimization.SolverOptions()
opts_mintime_sqrt.square_root = true
solver_mintime = TrajectoryOptimization.Solver(model,obj_mintime,N=N,opts=opts_mintime)
solver_mintime_sqrt = TrajectoryOptimization.Solver(model,obj_mintime,N=N,opts=opts_mintime_sqrt)
U0 = ones(solver_mintime.model.m,solver_mintime.N)
results_mintime = init_results(solver_mintime,Array{Float64}(undef,0,0),U0)
results_mintime_sqrt = init_results(solver_mintime_sqrt,Array{Float64}(undef,0,0),U0)
TrajectoryOptimization.update_jacobians!(results_mintime, solver_mintime)
TrajectoryOptimization.update_jacobians!(results_mintime_sqrt, solver_mintime_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results_mintime, solver_mintime)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_mintime_sqrt, solver_mintime_sqrt)

@test all(isapprox.(results_mintime.K,results_mintime_sqrt.K))
@test all(isapprox.(results_mintime.d,results_mintime_sqrt.d))
@test all(isapprox.(results_mintime.s,results_mintime_sqrt.s))
S_mintime_sqrt = [results_mintime_sqrt.S[k]'*results_mintime_sqrt.S[k] for k = 1:solver_mintime.N]
@test isapprox(to_array(results_mintime.S),to_array(S_mintime_sqrt))
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_sqrt)
@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# Minimum Time + Infeasible
solver_mintime.state.infeasible = true
solver_mintime_sqrt.state.infeasible = true
results_mintime_inf = init_results(solver_mintime,X0,U0)
results_mintime_inf_sqrt = init_results(solver_mintime_sqrt,X0,U0)
TrajectoryOptimization.update_jacobians!(results_mintime_inf, solver_mintime)
TrajectoryOptimization.update_jacobians!(results_mintime_inf_sqrt, solver_mintime_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results_mintime_inf, solver_mintime)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_mintime_inf_sqrt, solver_mintime_sqrt)

@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results_mintime_inf.K,results_mintime_inf_sqrt.K))
@test all(isapprox.(results_mintime_inf.d,results_mintime_inf_sqrt.d))
@test all(isapprox.(results_mintime_inf.s,results_mintime_inf_sqrt.s))
S_mintime_inf_sqrt = [results_mintime_inf_sqrt.S[k]'*results_mintime_inf_sqrt.S[k] for k = 1:N]
@test isapprox(to_array(results_mintime_inf.S),to_array(S_mintime_inf_sqrt))
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_inf.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_inf)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_inf_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_mintime_inf_sqrt)

@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# Terminal cost Hessian .= 0
obj_ = copy(obj)
obj_.cost.Qf .= 0.
obj_.cost.qf .= 0.

opts = TrajectoryOptimization.SolverOptions()
solver = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt.opts.square_root = true
U0 = ones(solver.model.m,solver.N)
results = init_results(solver,Array{Float64}(undef,0,0),U0)
results_sqrt = init_results(solver_sqrt,Array{Float64}(undef,0,0),U0)
results.X[1] = solver.obj.x0
copyto!(results.U, ones(model.m,solver.N-1))
results_sqrt.X[1] = solver.obj.x0
copyto!(results_sqrt.U, ones(model.m,solver.N-1))

TrajectoryOptimization.rollout!(results,solver)
TrajectoryOptimization.rollout!(results_sqrt,solver_sqrt)
TrajectoryOptimization.update_jacobians!(results,solver)
TrajectoryOptimization.update_jacobians!(results_sqrt,solver_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results,solver)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_sqrt,solver_sqrt)

# test that results from square root backward pass are the same as backward pass
@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results.K,results_sqrt.K))
@test isapprox(to_array(results.d),to_array(results_sqrt.d))
S_sqrt = [results_sqrt.S[k]'*results_sqrt.S[k] for k = 1:N]
@test isapprox(to_array(results.S),to_array(S_sqrt))

max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt)

@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# Stage cost Q .= 0
obj_ = copy(obj)
obj_.cost.Q .= 0.
obj_.cost.q .= 0.
opts = TrajectoryOptimization.SolverOptions()
solver = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt.opts.square_root = true
U0 = ones(solver.model.m,solver.N)
results = init_results(solver,Array{Float64}(undef,0,0),U0)
results_sqrt = init_results(solver_sqrt,Array{Float64}(undef,0,0),U0)
results.X[1] = solver.obj.x0
copyto!(results.U, ones(model.m,solver.N-1))
results_sqrt.X[1] = solver.obj.x0
copyto!(results_sqrt.U, ones(model.m,solver.N-1))

TrajectoryOptimization.rollout!(results,solver)
TrajectoryOptimization.rollout!(results_sqrt,solver_sqrt)
TrajectoryOptimization.update_jacobians!(results,solver)
TrajectoryOptimization.update_jacobians!(results_sqrt,solver_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results,solver)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_sqrt,solver_sqrt)

# test that results from square root backward pass are the same as backward pass
@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results.K,results_sqrt.K))
@test isapprox(to_array(results.d),to_array(results_sqrt.d))
S_sqrt = [results_sqrt.S[k]'*results_sqrt.S[k] for k = 1:N]
@test isapprox(to_array(results.S),to_array(S_sqrt))

max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt)

@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

# Stage cost Q .= 0, Terminal Cost Qf .= 0
obj_ = copy(obj)
obj_.cost.Q .= 0.
obj_.cost.q .= 0.
obj_.cost.Qf .= 0.
obj_.cost.qf .= 0.
opts = TrajectoryOptimization.SolverOptions()
solver = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt = TrajectoryOptimization.Solver(model,obj_,dt=0.1,opts=opts)
solver_sqrt.opts.square_root = true
solver_sqrt.opts.bp_sqrt_inv_type = :pseudo
U0 = ones(solver.model.m,solver.N)
results = init_results(solver,Array{Float64}(undef,0,0),U0)
results_sqrt = init_results(solver_sqrt,Array{Float64}(undef,0,0),U0)
results.X[1] = solver.obj.x0
copyto!(results.U, ones(model.m,solver.N-1))
results_sqrt.X[1] = solver.obj.x0
copyto!(results_sqrt.U, ones(model.m,solver.N-1))

TrajectoryOptimization.rollout!(results,solver)
TrajectoryOptimization.rollout!(results_sqrt,solver_sqrt)
TrajectoryOptimization.update_jacobians!(results,solver)
TrajectoryOptimization.update_jacobians!(results_sqrt,solver_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results,solver)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_sqrt,solver_sqrt)

# test that results from square root backward pass are the same as backward pass
@test isapprox(ΔV,ΔV_sqrt)
@test norm(to_array(results.K) .- to_array(results_sqrt.K)) < 1e-12
@test norm(to_array(results.d) .- to_array(results_sqrt.d)) < 1e-12
S_sqrt = [results_sqrt.S[k]'*results_sqrt.S[k] for k = 1:N]
@test norm(to_array(results.S) .- to_array(S_sqrt)) < 1e-12
#
max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt)

# NOTE: If you have not state stage costs or state terminal cost the square root method may not be better numerically conditioned
@test max_cn_Quu_sqrt <= max_cn_Quu
@test max_cn_S_sqrt <= max_cn_S # this test fails but that may be ok

# w/ regularization
opts = TrajectoryOptimization.SolverOptions()
opts.bp_reg_initial = 10.0
solver = TrajectoryOptimization.Solver(model,obj,dt=0.1,opts=opts)
solver_sqrt = TrajectoryOptimization.Solver(model,obj,dt=0.1,opts=opts)
solver_sqrt.opts.square_root = true
U0 = ones(solver.model.m,solver.N)
results = init_results(solver,Array{Float64}(undef,0,0),U0)
results_sqrt = init_results(solver_sqrt,Array{Float64}(undef,0,0),U0)
results.X[1] = solver.obj.x0
copyto!(results.U, ones(model.m,solver.N-1))
results_sqrt.X[1] = solver.obj.x0
copyto!(results_sqrt.U, ones(model.m,solver.N-1))

TrajectoryOptimization.rollout!(results,solver)
TrajectoryOptimization.rollout!(results_sqrt,solver_sqrt)
TrajectoryOptimization.update_jacobians!(results,solver)
TrajectoryOptimization.update_jacobians!(results_sqrt,solver_sqrt)
ΔV = TrajectoryOptimization.backwardpass!(results,solver)
ΔV_sqrt = TrajectoryOptimization.backwardpass!(results_sqrt,solver_sqrt)

# test that results from square root backward pass are the same as backward pass
@test isapprox(ΔV,ΔV_sqrt)
@test all(isapprox.(results.K,results_sqrt.K))
@test isapprox(to_array(results.d),to_array(results_sqrt.d))
S_sqrt = [results_sqrt.S[k]'*results_sqrt.S[k] for k = 1:N]
@test isapprox(to_array(results.S),to_array(S_sqrt))

max_cn_Quu = TrajectoryOptimization.backwardpass_max_condition_number(results.bp)
max_cn_S = TrajectoryOptimization.backwardpass_max_condition_number(results)
max_cn_Quu_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt.bp)
max_cn_S_sqrt = TrajectoryOptimization.backwardpass_max_condition_number(results_sqrt)

@test max_cn_Quu_sqrt < max_cn_Quu
@test max_cn_S_sqrt < max_cn_S

## Simple pendulum
u_bound = 2.
model, obj = TrajectoryOptimization.Dynamics.pendulum
obj_c = Dynamics.pendulum_constrained[2]
opts = TrajectoryOptimization.SolverOptions()
opts.verbose = false

# unconstrained
solver = TrajectoryOptimization.Solver(model,obj,dt=0.1,opts=opts)
U = zeros(solver.model.m, solver.N-1)
results, = TrajectoryOptimization.solve(solver,U)
solver.opts.square_root = false
results_sqrt, = TrajectoryOptimization.solve(solver,U)
@test norm(results.X[end]-results_sqrt.X[end]) < 1e-3

# constrained
solver = TrajectoryOptimization.Solver(model,obj_c,dt=0.1,opts=opts)
results_c, = TrajectoryOptimization.solve(solver, U)
solver.opts.square_root = true
results_c_sqrt, = TrajectoryOptimization.solve(solver, U)
max_c = TrajectoryOptimization.max_violation(results_c_sqrt)
@test norm(results_c.X[end]-results_c_sqrt.X[end]) < 1e-3
@test max_c < 1e-3


text:
module SIforSFS

export compute_solution_path, truncated_interval, parametric_SFS_SI, parametric_SFS_CI, SelectionEvent, Active, ActiveSign, ActiveOrder, ActiveSignOrder

using InvertedIndices
using LinearAlgebra
using IntervalSets
using Distributions
include("sfs.jl")
using .SFS
include("../selective_p.jl")
using .MultiTruncatedDistributions
include("../ci.jl")
using .ConfidenceInterval

const ε = 1e-4

function parametric_representation(𝜼, 𝐲, 𝚺)
    z = 𝜼' * 𝐲
    σ²_z = 𝜼' * 𝚺 * 𝜼
    𝐛 = 𝚺 * 𝜼 / σ²_z
    𝐚 = 𝐲 - 𝐛 * 𝜼' * 𝐲
    @assert 𝐚 + 𝐛 * z ≈ 𝐲
    return (z, σ²_z, 𝐚, 𝐛)
end

function polytope_t(𝐗, 𝐚, 𝐛, 𝐴, 𝐬, t)
    L, U = -Inf, +Inf
    for i = 1:t
        𝐗_𝐴ᵢ = 𝐗[:, 𝐴[1:(i - 1)]]
        𝐏 = I - 𝐗_𝐴ᵢ / (𝐗_𝐴ᵢ' * 𝐗_𝐴ᵢ) * 𝐗_𝐴ᵢ'
        for j′ in collect(1:size(𝐗, 2))[Not(𝐴[1:i])]
            𝐱_j′ = 𝐗[:, j′] # remain feature
            𝐱_jᵢ = 𝐗[:, 𝐴[i]] # i-th selected feature
            a₊, b₊ = (+𝐱_j′ - 𝐬[i] * 𝐱_jᵢ)' * 𝐏 * [𝐚 𝐛] # selection event: a₊ + b₊*z ≤ 0
            a₋, b₋ = (-𝐱_j′ - 𝐬[i] * 𝐱_jᵢ)' * 𝐏 * [𝐚 𝐛] # selection event: a₋ + b₋*z ≤ 0
            if b₊ < 0
                L = max(L, -a₊ / b₊)
            elseif b₊ > 0
                U = min(U, -a₊ / b₊)
            end
            if b₋ < 0
                L = max(L, -a₋ / b₋)
            elseif b₋ > 0
                U = min(U, -a₋ / b₋)
            end
        end
    end
    return L..U
end

function polytope(𝐗, 𝐲_z, 𝐚, 𝐛, k; until=k)
    𝐴_z, 𝐬_z = sfs(𝐗, 𝐲_z, k)
    until != k && (𝐴_z = 𝐴_z[1:until]; 𝐬_z = 𝐬_z[1:until])
    I_z = polytope_t(𝐗, 𝐚, 𝐛, 𝐴_z, 𝐬_z, until)
    return (I_z, 𝐴_z, 𝐬_z)
end

function compute_solution_path(𝐗, 𝐚, 𝐛, k, z_min, z_max; until=k)
    z = z_min
    𝐳 = float([z_min]) # transition point list
    𝑨 = Vector{Vector{Int}}(undef, 0) # active set list
    𝐒 = Vector{Vector{Float64}}(undef, 0) # sign vector list
    while z < z_max
        𝐲_z = 𝐚 + 𝐛 * z
        I_z, 𝐴_z, 𝐬_z = polytope(𝐗, 𝐲_z, 𝐚, 𝐛, k; until=until)
        z = I_z.right + ε
        z < z_max ? push!(𝐳, z - ε) : push!(𝐳, z_max)
        push!(𝑨, 𝐴_z)
        push!(𝐒, 𝐬_z)
    end
    return (𝐳, 𝑨, 𝐒)
end

function parametric_SFS_SI(𝐗, 𝐲_obs, k, 𝚺; selection_event=Active())
    𝐴_obs, 𝐬_obs = sfs(𝐗, 𝐲_obs, k)
    𝐗_obs = 𝐗[:, 𝐴_obs]
    𝜼 = 𝐗_obs / (𝐗_obs' * 𝐗_obs)
    p_values = Vector{Float64}(undef, length(𝐴_obs)) # selective p-value list
    for j = 1:length(𝐴_obs)
        # parametric representation for 𝐲(z)
        z_obs, σ²_z, 𝐚, 𝐛 = parametric_representation(𝜼[:, j], 𝐲_obs, 𝚺)

        # compute Z = {z | selection_event}
        Z = truncated_interval(selection_event, 𝐴_obs, 𝐬_obs, 𝐗, 𝐚, 𝐛, k, -abs(z_obs) - 10 * √(σ²_z), abs(z_obs) + 10 * √(σ²_z))

        𝜋_j = cdf(TruncatedDistribution(Normal(0, √(σ²_z)), Z), z_obs)
        p_values[j] = 2 * min(𝜋_j, 1 - 𝜋_j)
    end
    return (𝐴_obs, p_values)
end

function parametric_SFS_CI(𝐗, 𝐲_obs, k, 𝚺; selection_event=Active(), α=0.05)
    𝐴_obs, 𝐬_obs = sfs(𝐗, 𝐲_obs, k)
    𝐗_obs = 𝐗[:, 𝐴_obs]
    𝜼 = 𝐗_obs / (𝐗_obs' * 𝐗_obs)
    CIs = Vector{ClosedInterval{Float64}}(undef, length(𝐴_obs)) # selective CI list
    for j = 1:length(𝐴_obs)
        # parametric representation for 𝐲(z)
        z_obs, σ²_z, 𝐚, 𝐛 = parametric_representation(𝜼[:, j], 𝐲_obs, 𝚺)

        # compute Z = {z | selection_event}
        Z = truncated_interval(selection_event, 𝐴_obs, 𝐬_obs, 𝐗, 𝐚, 𝐛, k, -abs(z_obs) - 10 * √(σ²_z), abs(z_obs) + 10 * √(σ²_z))

        CIs[j] = confidence_interval(z_obs, √(σ²_z), Z; α=α)
    end
    return (𝐴_obs, CIs)
end

abstract type SelectionEvent end
struct Active <: SelectionEvent end
struct ActiveSign <: SelectionEvent end
struct ActiveOrder <: SelectionEvent end
struct ActiveSignOrder <: SelectionEvent end

function truncated_interval(::Active, 𝐴_obs, 𝐬_obs, params...; kwargs...)
    𝐳, 𝑨, _ = compute_solution_path(params...; kwargs...)
    Z = Array{ClosedInterval{eltype(𝐳)}}(undef, 0) # truncated intevals
    for i = 1:length(𝑨)
        Set(𝐴_obs) == Set(𝑨[i]) && push!(Z, (𝐳[i]..(𝐳[i + 1]))) # selective event: 𝐴(z) = 𝐴_obs
    end
    return Z
end
function truncated_interval(::ActiveSign, 𝐴_obs, 𝐬_obs, params...; kwargs...)
    𝐳, 𝑨, 𝐒 = compute_solution_path(params...; kwargs...)
    Z = Array{ClosedInterval{eltype(𝐳)}}(undef, 0) # truncated intevals
    for i = 1:length(𝑨)
        Set(𝐴_obs .* 𝐬_obs) == Set(𝑨[i] .* 𝐒[i]) && push!(Z, (𝐳[i]..(𝐳[i + 1]))) # selective event: 𝐴(z) = 𝐴_obs and 𝐬(z) = 𝐬_obs
    end
    return Z
end
function truncated_interval(::ActiveOrder, 𝐴_obs, 𝐬_obs, params...; kwargs...)
    𝐳, 𝑨, _ = compute_solution_path(params...; kwargs...)
    Z = Array{ClosedInterval{eltype(𝐳)}}(undef, 0) # truncated intevals
    for i = 1:length(𝑨)
        𝐴_obs == 𝑨[i] && push!(Z, (𝐳[i]..(𝐳[i + 1]))) # selective event: 𝐴(z) = 𝐴_obs and selection order
    end
    return Z
end
function truncated_interval(::ActiveSignOrder, 𝐴_obs, 𝐬_obs, params...; kwargs...)
    𝐳, 𝑨, 𝐒 = compute_solution_path(params...; kwargs...)
    Z = Array{ClosedInterval{eltype(𝐳)}}(undef, 0) # truncated intevals
    for i = 1:length(𝑨)
        𝐴_obs == 𝑨[i] && 𝐬_obs == 𝐒[i] && push!(Z, (𝐳[i]..(𝐳[i + 1]))) # selective event: 𝐴(z) = 𝐴_obs and 𝐬(z) = 𝐬_obs and selection order
    end
    return Z
end

end # module SIforSFS


text:
## p-value adjustment methods ##

"""
    adjust(PValues, <:PValueAdjustment)
    adjust(PValues, Int, <:PValueAdjustment)

Adjustment of p-values

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BenjaminiHochberg())
4-element Array{Float64,1}:
 0.004
 0.02
 0.039999999999999994
 0.5

julia> adjust(pvals, 6, BenjaminiHochberg()) # 4 out of 6 p-values
4-element Array{Float64,1}:
 0.006
 0.03
 0.06
 0.75

julia> adjust(pvals, BarberCandes())
4-element Array{Float64,1}:
 0.3333333333333333
 0.3333333333333333
 0.3333333333333333
 1.0
```

# See also

`PValueAdjustment`s:

[`Bonferroni`](@ref)
[`BenjaminiHochberg`](@ref)
[`BenjaminiHochbergAdaptive`](@ref)
[`BenjaminiYekutieli`](@ref)
[`BenjaminiLiu`](@ref)
[`Hochberg`](@ref)
[`Holm`](@ref)
[`Hommel`](@ref)
[`Sidak`](@ref)
[`ForwardStop`](@ref)
[`BarberCandes`](@ref)
"""
function adjust end

# promotion from float vectors to PValues type

function adjust(pValues::Vector{T}, method::M) where {T <: AbstractFloat,M <: PValueAdjustment}
    return adjust(PValues(pValues), method)
end

function adjust(pValues::Vector{T}, n::Integer, method::M) where {T <: AbstractFloat,M <: PValueAdjustment}
    return adjust(PValues(pValues), n, method)
end

# Bonferroni

"""
Bonferroni adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, Bonferroni())
4-element Array{Float64,1}:
 0.004
 0.04
 0.12
 1.0

julia> adjust(pvals, 6, Bonferroni())
4-element Array{Float64,1}:
 0.006
 0.06
 0.18
 1.0
```

# References

Bonferroni, C.E. (1936). Teoria statistica delle classi e calcolo delle
probabilita (Libreria internazionale Seeber).
"""
struct Bonferroni <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::Bonferroni) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::Bonferroni) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    pAdjusted = clamp.(pValues * n, 0, 1)
    return pAdjusted
end

# Benjamini-Hochberg

"""
Benjamini-Hochberg adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BenjaminiHochberg())
4-element Array{Float64,1}:
 0.004
 0.02
 0.039999999999999994
 0.5

julia> adjust(pvals, 6, BenjaminiHochberg())
4-element Array{Float64,1}:
 0.006
 0.03
 0.06
 0.75
```

# References

Benjamini, Y., and Hochberg, Y. (1995). Controlling the False Discovery Rate: A
Practical and Powerful Approach to Multiple Testing. Journal of the Royal
Statistical Society. Series B (Methodological) 57, 289–300.
"""
struct BenjaminiHochberg <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::BenjaminiHochberg) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::BenjaminiHochberg) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if k <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = pValues[sortedOrder]
    pAdjusted .*= n ./ (1:k)
    stepup!(pAdjusted)
    pAdjusted[sortedOrder] = clamp.(pAdjusted, 0, 1)
    return pAdjusted
end

# Benjamini-Hochberg Adaptive

"""
Adaptive Benjamini-Hochberg adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BenjaminiHochbergAdaptive(Oracle(0.5))) # known π₀ of 0.5
4-element Array{Float64,1}:
 0.002
 0.01
 0.019999999999999997
 0.25

julia> adjust(pvals, BenjaminiHochbergAdaptive(StoreyBootstrap())) # π₀ estimator
4-element Array{Float64,1}:
 0.0
 0.0
 0.0
 0.0

julia> adjust(pvals, 6, BenjaminiHochbergAdaptive(StoreyBootstrap()))
4-element Array{Float64,1}:
 0.0
 0.0
 0.0
 0.0
```

# References

Benjamini, Y., Krieger, A. M. & Yekutieli, D. (2006). Adaptive linear step-up
procedures that control the false discovery rate. Biometrika 93, 491–507.
"""
struct BenjaminiHochbergAdaptive <: PValueAdjustment
    pi0estimator::Pi0Estimator
end

## default to BenjaminiHochberg
BenjaminiHochbergAdaptive(π0::AbstractFloat) = BenjaminiHochbergAdaptive(Oracle(π0))

BenjaminiHochbergAdaptive() = BenjaminiHochbergAdaptive(1.0)

adjust(pValues::PValues{T}, method::BenjaminiHochbergAdaptive) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::BenjaminiHochbergAdaptive) where T <: AbstractFloat
    π0 = estimate(pValues, method.pi0estimator)
    pAdjusted = adjust(pValues, n, BenjaminiHochberg()) * π0
    return pAdjusted
end

# Benjamini-Yekutieli

"""
Benjamini-Yekutieli adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BenjaminiYekutieli())
4-element Array{Float64,1}:
 0.00833333333333334
 0.0416666666666667
 0.0833333333333334
 1.0

julia> adjust(pvals, 6, BenjaminiYekutieli())
4-element Array{Float64,1}:
 0.01470000000000001
 0.07350000000000005
 0.14700000000000008
 1.0
```

# References

Benjamini, Y., and Yekutieli, D. (2001). The Control of the False Discovery Rate
in Multiple Testing under Dependency. The Annals of Statistics 29, 1165–1188.
"""
struct BenjaminiYekutieli <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::BenjaminiYekutieli) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::BenjaminiYekutieli) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if k <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = pValues[sortedOrder]
    pAdjusted .*= harmonic_number(n) .* n ./ (1:k)
    stepup!(pAdjusted)
    pAdjusted[sortedOrder] = clamp.(pAdjusted, 0, 1)
    return pAdjusted
end

# Benjamini-Liu

"""
Benjamini-Liu adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BenjaminiLiu())
4-element Array{Float64,1}:
 0.003994003998999962
 0.022275750000000066
 0.02955000000000002
 0.125

julia> adjust(pvals, 6, BenjaminiLiu())
4-element Array{Float64,1}:
 0.0059850199850060015
 0.04084162508333341
 0.07647146000000005
 0.4375
```

# References

Benjamini, Y., and Liu, W. (1999). A step-down multiple hypotheses testing
procedure that controls the false discovery rate under independence. Journal of
Statistical Planning and Inference 82, 163–170.
"""
struct BenjaminiLiu <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::BenjaminiLiu) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::BenjaminiLiu) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if n <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = pValues[sortedOrder]
    # a bit more involved because cutoffs at significance α have the form:
    # P_(i) <= 1- [1 - min(1, m/(m-i+1)α)]^{1/(m-i+1)}
    s = n .- (1:k) .+ 1
    pAdjusted = (1 .- (1 .- pAdjusted).^s) .* s ./ n
    stepdown!(pAdjusted)
    pAdjusted[sortedOrder] = clamp.(pAdjusted, 0, 1)
    return pAdjusted
end

# Hochberg

"""
Hochberg adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, Hochberg())
4-element Array{Float64,1}:
 0.004
 0.03
 0.06
 0.5

julia> adjust(pvals, 6, Hochberg())
4-element Array{Float64,1}:
 0.006
 0.05
 0.12
 1.0
```

# References

Hochberg, Y. (1988). A sharper Bonferroni procedure for multiple tests of
significance. Biometrika 75, 800–802.
"""
struct Hochberg <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::Hochberg) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::Hochberg) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if k <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = pValues[sortedOrder]
    pAdjusted .*= (n .- (1:k) .+ 1)
    stepup!(pAdjusted)
    pAdjusted[sortedOrder] = clamp.(pAdjusted, 0, 1)
    return pAdjusted
end

# Holm

"""
Holm adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, Holm())
4-element Array{Float64,1}:
 0.004
 0.03
 0.06
 0.5

julia> adjust(pvals, 6, Holm())
4-element Array{Float64,1}:
 0.006
 0.05
 0.12
 1.0
```

# References

Holm, S. (1979). A Simple Sequentially Rejective Multiple Test Procedure.
Scandinavian Journal of Statistics 6, 65–70.
"""
struct Holm <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::Holm) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::Holm) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if n <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = pValues[sortedOrder]
    pAdjusted .*= (n .- (1:k) .+ 1)
    stepdown!(pAdjusted)
    pAdjusted[sortedOrder] = clamp.(pAdjusted, 0, 1)
    return pAdjusted
end

# Hommel

"""
Hommel adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, Hommel())
4-element Array{Float64,1}:
 0.004
 0.03
 0.06
 0.5

julia> adjust(pvals, 6, Hommel())
4-element Array{Float64,1}:
 0.006
 0.05
 0.12
 1.0
```

# References

Hommel, G. (1988). A stagewise rejective multiple test procedure based on a
modified Bonferroni test. Biometrika 75, 383–386.
"""
struct Hommel <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::Hommel) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::Hommel) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    if k <= 1
        return pValues
    end
    sortedOrder = sortperm(pValues)
    pAdjusted = vcat(pValues[sortedOrder], fill(one(T), n - k))
    lower = n * minimum(pAdjusted ./ (1:n))
    q = fill(lower, n)
    pa = fill(lower, n)
    for j in (n - 1):-1:2
        idx_left = 1:(n - j + 1)
        idx_right = (n - j + 2):n
        q_right = minimum(view(pAdjusted, idx_right) ./ (2:j))
        q[idx_left] .= j .* min.(view(pAdjusted, idx_left), q_right)
        q[idx_right] .= q[n - j + 1]
        pa .= max.(pa, q)
    end
    pa = pa[1:k]
    pa[sortedOrder] = pa
    pAdjusted = max.(pa, pValues)
    return pAdjusted
end

# Sidak

"""
Šidák adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, Sidak())
4-element Array{Float64,1}:
 0.003994003998999962
 0.039403990000000055
 0.11470719000000007
 0.9375

julia> adjust(pvals, 6, Sidak())
4-element Array{Float64,1}:
 0.0059850199850060015
 0.058519850599
 0.1670279950710002
 0.984375
```

# References

Šidák, Z. (1967). Rectangular Confidence Regions for the Means of Multivariate
Normal Distributions. Journal of the American Statistical Association 62,
626–633.
"""
struct Sidak <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::Sidak) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::Sidak) where T <: AbstractFloat
    check_number_tests(length(pValues), n)
    pAdjusted = clamp.(1 .- (1 .- pValues).^n, 0, 1)
    return pAdjusted
end

# Forward Stop

"""
Forward-Stop adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, ForwardStop())
4-element Array{Float64,1}:
 0.0010005003335835344
 0.005525418093542492
 0.013836681223931188
 0.1836643060579347

julia> adjust(pvals, 6, ForwardStop())
4-element Array{Float64,1}:
 0.0010005003335835344
 0.005525418093542492
 0.013836681223931188
 0.1836643060579347
```

# References

G’Sell, M.G., Wager, S., Chouldechova, A., and Tibshirani, R. (2016). Sequential
selection procedures and false discovery rate control. J. R. Stat. Soc. B 78,
423–444.
"""
struct ForwardStop <: PValueAdjustment
end

adjust(pValues::PValues{T}, method::ForwardStop) where T <: AbstractFloat = adjust(pValues, length(pValues), method)

function adjust(pValues::PValues{T}, n::Integer, method::ForwardStop) where T <: AbstractFloat
    k = length(pValues)
    check_number_tests(k, n)
    sortedOrder = sortperm(pValues)
    logsums = -cumsum(log.(1 .- pValues[sortedOrder]))
    logsums ./= (1:k)
    stepup!(logsums)
    logsums[sortedOrder] = clamp.(logsums, 0, 1)
    return logsums
end

# Barber-Candès

"""
Barber-Candès adjustment

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> adjust(pvals, BarberCandes())
4-element Array{Float64,1}:
 0.3333333333333333
 0.3333333333333333
 0.3333333333333333
 1.0
```

# References

Barber, R.F., and Candès, E.J. (2015). Controlling the false discovery rate via
knockoffs. Ann. Statist. 43, 2055–2085.

Arias-Castro, E., and Chen, S. (2017). Distribution-free multiple testing.
Electron. J. Statist. 11, 1983–2001.
"""
struct BarberCandes <: PValueAdjustment
end

function adjust(pValues::PValues{T}, method::BarberCandes) where T <: AbstractFloat
    n = length(pValues)
    # special cases unlike other p-adjust methods
    if n <= 1
    return fill(1.0, size(pValues))
    end
    if maximum(pValues) < 0.5
        return fill(1 / n, size(pValues))
    end

    sorted_indexes = sortperm(pValues)
    estimated_fdrs = pValues[sorted_indexes]

    Rt = 1 # current number of discoveries
    Vt = 1 # estimated false discoveries at t

    left_pv = estimated_fdrs[1]
    right_pv = estimated_fdrs[n]

    while left_pv < 0.5
        while (1 - right_pv <= left_pv) && (right_pv >= 0.5) && (Vt + Rt <= n)
            estimated_fdrs[n - Vt + 1] = 1.0
            right_pv = estimated_fdrs[n - Vt]
            Vt += 1
        end
        estimated_fdrs[Rt] = Vt / Rt
        Rt += 1
        left_pv = estimated_fdrs[Rt]
    end

    while (right_pv >= 0.5) && (Vt + Rt <= n + 1)
      estimated_fdrs[n - Vt + 1] = 1.0
      right_pv = (Vt + Rt <= n) ? estimated_fdrs[n - Vt] : 0.0
      Vt += 1
    end

    stepup!(estimated_fdrs)
    pAdjusted = clamp.(estimated_fdrs, 0, 1)
    pAdjusted[sorted_indexes] = pAdjusted
    return pAdjusted
end

## internal ##

# step-up / step-down

function stepup!(values::AbstractVector{T}) where T <: AbstractFloat
    accumulate!(min, values, reverse!(values))
    reverse!(values)
    return values
end

function stepdown!(values::AbstractVector{T}) where T <: AbstractFloat
    accumulate!(max, values, values)
    return values
end

function check_number_tests(k::Integer, n::Integer)
    if k > n
        msg = "Number of total tests ($n) is smaller than number of p-values ($k)"
        throw(ArgumentError(msg))
    end
end

harmonic_number(n::Integer) =  digamma(n + 1) + MathConstants.γ


text:
using Asap
using Test

@testset "Asap.jl" begin
    # 2D truss test: Example 3.8 from Kassimali "Matrix Analysis of Structures 2e"
    # in kips, ft

    #nodes
    n1 = Node([0., 0.] .* 12, :truss, :fixed)
    n2 = Node([12., 0.] .* 12, :truss, :fixed)
    n3 = Node([24., 0.] .* 12, :truss, :fixed)
    n4 = Node([12., 16.] .* 12, :truss, :free)

    nodes = [n1, n2, n3, n4]

    # elements
    E = 29e3 #ksi

    e1 = Element(nodes, [1,4], E, 8.)
    e2 = Element(nodes, [2,4], E, 6.)
    e3 = Element(nodes, [3,4], E, 8.)

    elements = [e1, e2, e3]

    # load
    l1 = Load(nodes, n4.position, [150., -300.])

    loads = [l1]

    # assemble
    ex38 = Structure(nodes, elements, loads)

    # Analyze
    analyze(ex38)

    # value in textbook
    d_textbook = [0.21552, -0.13995]

    @test round.(d_textbook, sigdigits = 4) ≈ round.(n4.disp, sigdigits = 4)

    #2D frame test: Example 6.6
    # in kips, in

    # nodes
    n1 = Node([0., 0.] .* 12, :frame, :fixed)
    n2 = Node([10., 20.] .* 12, :frame, :free)
    n3 = Node([30., 20.] .* 12, :frame, :fixed)

    nodes = [n1, n2, n3]

    # elements
    E = 29e3 #ksi
    A = 11.8 # in^2
    I = 310. # in^4

    e1 = Element(nodes, [1,2], E, A, I)
    e2 = Element(nodes, [2,3], E, A, I)

    elements = [e1, e2]

    # loads
    l1 = Load(nodes, n2.position, [0., -60., -750.])

    loads = [l1]

    # assembly + analysis
    ex66 = Structure(nodes, elements, loads)
    analyze(ex66)

    d_textbook = [0.021302, -0.06732, -0.0025499]

    @test round.(d_textbook, sigdigits = 4) ≈ round.(n2.disp, sigdigits = 4)

    # 3D truss test: Example 8.1
    # in kips, in

    # nodes
    n1 = Node([0., 0., 0.] .* 12, :truss, :fixed)
    n2 = Node([18., 0., 0.] .* 12, :truss, :fixed)
    n3 = Node([12., 16., 0.] .* 12, :truss, :fixed)
    n4 = Node([-6., 16., 0.] .* 12, :truss, :fixed)
    n5 = Node([6., 8., 24.] .* 12, :truss, :free)

    nodes = [n1, n2, n3, n4, n5]

    # elements
    E = 10e3 #ksi
    A = 8.4 #in^2

    e1 = Element(nodes, [1,5], E, A)
    e2 = Element(nodes, [2,5], E, A)
    e3 = Element(nodes, [3,5], E, A)
    e4 = Element(nodes, [4,5], E, A)

    elements = [e1, e2, e3, e4]

    # loads
    l1 = Load(nodes, n5.position, [0., 50., -100.])

    loads = [l1]

    ex81 = Structure(nodes, elements, loads)
    analyze(ex81)

    d_textbook = [.10913, .57202, -.12104]

    @test round.(d_textbook, sigdigits = 4) ≈ round.(n5.disp, sigdigits = 4)

    # 3D frame test: Example 8.3 Ferreira "MATLAB Codes for Finite Element Analysis"
    # in kN, m

    # nodes
    n1 = Node([0., 0., 0.], :frame, :free)
    n2 = Node([3., 0., 0.], :frame, :fixed)
    n3 = Node([0., 0., -3.], :frame, :fixed)
    n4 = Node([0., -4., 0.], :frame, :fixed)

    nodes = [n1, n2, n3, n4]

    # elements
    E = 210e6
    A = 0.02
    Iy = 10e-5
    Iz = 20e-5
    J = 5e-5
    G = 84e6

    e1 = Element(nodes, [1,2], E, A, G, Iz, Iy, J)
    e2 = Element(nodes, [1,3], E, A, G, Iz, Iy, J)
    e3 = Element(nodes, [1,4], E, A, G, Iz, Iy, J)

    elements = [e1, e2, e3]

    # loads
    l = Load(nodes, n1.position, [-10., 0., 20., 0., 0., 0.])

    loads = [l]

    ex83 = Structure(nodes, elements, loads)
    analyze(ex83)

    d_textbook = [-7.05e-6, -7e-8, 1.418e-5, 1.45e-6, 1.75e-6, 1.14e-6]

    @test round.(d_textbook, digits = 8) ≈ round.(n1.disp, digits = 8)
end

max_stars_repo_name: keithjlee/Asap

max_issues_repo_name: keithjlee/Asap

max_forks_repo_name: keithjlee/Asap


text:
module BangBang

# Use README as the docstring of the module:
@doc read(joinpath(dirname(@__DIR__), "README.md"), String) BangBang

export
    @!,
    append!!,
    delete!!,
    empty!!,
    lmul!!,
    materialize!!,
    mul!!,
    pop!!,
    popfirst!!,
    push!!,
    pushfirst!!,
    rmul!!,
    setindex!!,
    setproperty!!,
    splice!!

using Base.Broadcast: materialize!
using InitialValues
using LinearAlgebra
using Requires
using ZygoteRules: @adjoint

include("NoBang/NoBang.jl")
using .NoBang: ImmutableContainer

include("core.jl")
include("base.jl")
include("linearalgebra.jl")
include("initials.jl")
include("macro.jl")
include("zygote.jl")

function __init__()
    @require StaticArrays="90137ffa-7385-5640-81b9-e52037218182" begin
        include("staticarrays.jl")
    end
end

end # module


text:
function pvdwrite(filename, cgnodes, cgelcon, cgcells, celltype, scalars, vectors, fields, dt, shape)

if ENDIAN_BOM==0x04030201
    byte_order = "LittleEndian";
elseif ENDIAN_BOM==0x01020304
    byte_order = "BigEndian";
else
    error("Endian is not valid");
end

pvdfile = filename * ".pvd";
fid = open(pvdfile, "w");
mystr = "<?xml version=\"1.0\"?>\n";
mystr = mystr * "<VTKFile type=\"Collection\" version=\"0.1\"\n";
mystr = mystr * "         byte_order=\"" * byte_order * "\"\n";
mystr = mystr * "         compressor=\"vtkZLibDataCompressor\">\n";
mystr = mystr * "  <Collection>\n";

nt = size(fields,4);
for i = 1:nt
    vtufile = filename * string(i);

    ind = findlast("/", vtufile);
    outfile = vtufile[(ind[end]+1):end] * ".vtu";

    tm = shape*reshape(fields[:,:,:,i],(size(shape,2), size(fields,2)*size(fields,3)));
    tm = reshape(tm,(size(shape,1), size(fields,2), size(fields,3)));

    vtuwrite(vtufile, cgnodes, cgelcon, cgcells, celltype, scalars, vectors, tm);
    mystr = mystr * "    <DataSet timestep=\"" * string(i*dt) * "\" group=\"\" part=\"0\"\n";
    mystr = mystr * "             file=\"" * outfile * "\"/>\n"
end
write(fid,mystr);

write(fid,"  </Collection>\n");
write(fid,"</VTKFile>\n");
close(fid);

end

max_stars_repo_name: glwagner/Exasim

max_issues_repo_name: glwagner/Exasim

max_forks_repo_name: glwagner/Exasim


text:
using Plots
x=-3:0.1:3
f(x)=x^2
y=f.(x)

gr()
plot(x,y,label="line")
scatter!(x,y,label="points")

max_stars_repo_name: stepinski/julia-in

max_issues_repo_name: stepinski/julia-init

max_forks_repo_name: stepinski/julia-init


text:
"""
    barplot(x, y; kwargs...)

Plots a barplot; `y` defines the height.  `x` and `y` should be 1 dimensional.

## Attributes
$(ATTRIBUTES)
"""
@recipe(BarPlot, x, y) do scene
    Attributes(;
        fillto = 0.0,
        color = theme(scene, :color),
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        marker = Rect,
        strokewidth = 0,
        strokecolor = :white,
        width = automatic,
        direction = :y,
        visible = theme(scene, :visible),
    )
end

conversion_trait(::Type{<: BarPlot}) = PointBased()

function bar_rectangle(xy, width, fillto)
    x, y = xy
    # y could be smaller than fillto...
    ymin = min(fillto, y)
    ymax = max(fillto, y)
    w = abs(width)
    return FRect(x - (w / 2f0), ymin, w, ymax - ymin)
end

flip(r::Rect2D) = Rect2D(reverse(origin(r)), reverse(widths(r)))

function AbstractPlotting.plot!(p::BarPlot)

    in_y_direction = lift(p.direction) do dir
        if dir == :y
            true
        elseif dir == :x
            false
        else
            error("Invalid direction $dir. Options are :x and :y.")
        end
    end

    bars = lift(p[1], p.fillto, p.width, in_y_direction) do xy, fillto, width, in_y_direction
        # compute half-width of bars
        if width === automatic
            # times 0.8 for default gap
            width = mean(diff(first.(xy))) * 0.8
            width = ifelse(isfinite(width), width, 0.8)
        end

        rects = bar_rectangle.(xy, width, fillto)
        return in_y_direction ? rects : flip.(rects)
    end

    poly!(
        p, bars, color = p.color, colormap = p.colormap, colorrange = p.colorrange,
        strokewidth = p.strokewidth, strokecolor = p.strokecolor, visible = p.visible
    )
end


text:
myFunction (a{T};
            x::Float64=sin(1))::Bool where {T} = false


text:
add_one(args...) = 1.0

function test_forcing(fld)
    kwarg = Dict(Symbol(:F, fld) => add_one)
    forcing = Forcing(; kwarg...)
    f = getfield(forcing, fld)
    f() == 1.0
end

function time_step_with_forcing_functions(arch)
    @inline Fu(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.u[i, j, k] / 60, 0)
    @inline Fv(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.v[i, j, k] / 60, 0)
    @inline Fw(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.w[i, j, k] / 60, 0)

    forcing = Forcing(Fu=Fu, Fv=Fv, Fw=Fw)

    model = BasicModel(N=(16, 16, 16), L=(1, 1, 1), architecture=arch, forcing=forcing)
    time_step!(model, 1, 1)
    return true
end

function time_step_with_forcing_functions_params(arch)
    @inline Fu(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.u[i, j, k] / params.τ, 0)
    @inline Fv(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.v[i, j, k] / params.τ, 0)
    @inline Fw(i, j, k, grid, time, U, Φ, params) = @inbounds ifelse(k == grid.Nz, -U.w[i, j, k] / params.τ, 0)

    forcing = Forcing(Fu=Fu, Fv=Fv, Fw=Fw)

    model = BasicModel(N=(16, 16, 16), L=(1, 1, 1), architecture=arch, forcing=forcing, parameters=(τ=60,))
    time_step!(model, 1, 1)
    return true
end

function time_step_with_forcing_functions_sin_exp(arch)
    @inline Fu(i, j, k, grid, time, U, Φ, params) = @inbounds sin(grid.xC[i])
    @inline FT(i, j, k, grid, time, U, Φ, params) = @inbounds exp(-Φ.T[i, j, k])

    forcing = Forcing(Fu=Fu, FT=FT)

    model = BasicModel(N=(16, 16, 16), L=(1, 1, 1), architecture=arch, forcing=forcing)
    time_step!(model, 1, 1)
    return true
end

@testset "Forcings" begin
    println("Testing forcings...")

    @testset "Forcing function initialization" begin
        println("  Testing forcing function initialization...")
        for fld in (:u, :v, :w, :T, :S)
            @test test_forcing(fld)
        end
    end

    for arch in archs
        @testset "Forcing function time stepping [$(typeof(arch))]" begin
            println("  Testing forcing function time stepping [$(typeof(arch))]...")
            @test time_step_with_forcing_functions(arch)
            @test time_step_with_forcing_functions_params(arch)
            @test time_step_with_forcing_functions_sin_exp(arch)
        end
    end
end


text:
using BenchmarkTools, Plots, SNOPT7

T = Float64

# options
max_con_viol = 1.0e-8
verbose=false

opts_ilqr = iLQRSolverOptions{T}(verbose=verbose,live_plotting=:off)

opts_al = AugmentedLagrangianSolverOptions{T}(verbose=false,
    opts_uncon=opts_ilqr,
    cost_tolerance=1.0e-4,
    cost_tolerance_intermediate=1.0e-3,
    constraint_tolerance=max_con_viol,
    penalty_scaling=50.,
    penalty_initial=1.)

opts_pn = ProjectedNewtonSolverOptions{T}(verbose=verbose,
    feasibility_tolerance=max_con_viol)

opts_altro = ALTROSolverOptions{T}(verbose=verbose,
    opts_al=opts_al,
    opts_pn=opts_pn,
    projected_newton=true,
    projected_newton_tolerance=1.0e-3);

opts_ipopt = DIRCOLSolverOptions{T}(verbose=verbose,
    nlp=Ipopt.Optimizer(),
    feasibility_tolerance=max_con_viol)

opts_snopt = DIRCOLSolverOptions{T}(verbose=verbose,
    nlp=SNOPT7.Optimizer(),
    feasibility_tolerance=max_con_viol)

# ALTRO w/ Newton
prob_altro = copy(Problems.cartpole)
@time p1, s1 = solve(prob_altro, opts_altro)
@benchmark p1, s1 = solve($prob_altro, $opts_altro)
max_violation_direct(p1)
plot(p1.X,title="Cartpole state (ALTRO)")
plot(p1.U,title="Cartpole control (ALTRO)")

# DIRCOL w/ Ipopt
prob_ipopt = copy(Problems.cartpole)
rollout!(prob_ipopt)
prob_ipopt = update_problem(prob_ipopt,model=Dynamics.cartpole) # get continuous time model
@time p2, s2 = solve(prob_ipopt, opts_ipopt)
@benchmark p2, s2 = solve($prob_ipopt, $opts_ipopt)
push!(s2.stats[:iter_time],s2.stats[:time])
push!(s2.stats[:c_max], max_violation_direct(p2))
max_violation_direct(p2)
plot(p2.X,title="Cartpole state (Ipopt)")
plot(p2.U,title="Cartpole control (Ipopt)")

# DIRCOL w/ SNOPT
prob_snopt = copy(Problems.cartpole)
rollout!(prob_snopt)
prob_snopt = update_problem(prob_snopt,model=Dynamics.cartpole) # get continuous time model
@time p3, s3 = solve(prob_snopt, opts_snopt)
@benchmark p3, s3 = solve($prob_snopt, $opts_snopt)
max_violation_direct(p3)
s3.stats[:c_max][end]
plot(p3.X,title="Cartpole state (SNOPT)")
plot(p3.U,title="Cartpole control (SNOPT)")

# AL-iLQR
prob_alilqr = copy(Problems.cartpole)
opts_al.constraint_tolerance = max_con_viol
@time p4, s4 = solve(prob_alilqr, opts_al)
@benchmark p4, s4 = solve($prob_alilqr, $opts_al)
max_violation_direct(p4)
plot(p4.X,title="Cartpole state (ALTRO)")
plot(p4.U,title="Cartpole control (ALTRO)")

# Convergence plot
t_pn = s1.stats[:time_al]
t_span_al = range(0,stop=s1.stats[:time_al],length=s1.solver_al.stats[:iterations])
t_span_pn = range(t_pn,stop=s1.stats[:time],length=s1.solver_pn.stats[:iterations]+1)
t_span = [t_span_al;t_span_pn[2:end]]
c_span = [s1.solver_al.stats[:c_max]...,s1.solver_pn.stats[:c_max]...]

al_iter = s1.solver_al.stats[:iterations]
t_span2 = [t_span[1:al_iter]; t_pn .+ collect(range(0,stop=s4.stats[:time]-t_pn, length=s4.stats[:iterations]-al_iter))]
c_span2 = s4.stats[:c_max]

s_range = [1,(2:4:length(s3.stats[:iter_time])-1)...,length(s3.stats[:iter_time])]
i_range = [1,(2:4:length(s2.stats[:iter_time])-1)...,length(s2.stats[:iter_time])]
# p = plot(t_pn*ones(100),range(1.0e-10,stop=10.0,length=100),color=:red,linestyle=:dash,label="Projected Newton",width=2)
# p = plot!(s3.stats[:iter_time][s_range],s3.stats[:c_max][s_range],marker=:circle,yscale=:log10,ylim=[1.0e-10,10.0],color=:green,label="SNOPT")
# p = plot!(s2.stats[:iter_time][i_range],s2.stats[:c_max][i_range],marker=:circle,yscale=:log10,ylim=[1.0e-10,10.0],color=:blue,label="Ipopt")
# p = plot!(t_span,c_span,title="Cartpole c_max",xlabel="time (s)",marker=:circle,color=:orange,width=2,yscale=:log10,ylim=[1.0e-10,10.0],label="ALTRO")

###############################################
#            Create PGF Plot                  #
###############################################
include("vars.jl")
lwidth = 1.5
w1 = PGF.Plots.Linear(s3.stats[:iter_time][s_range],s3.stats[:c_max][s_range],mark="none", legendentry="SNOPT",style="line width=$lwidth, color=$col_snopt");
w2 = PGF.Plots.Linear(s2.stats[:iter_time][i_range],s2.stats[:c_max][i_range], mark="none",legendentry="Ipopt",style="line width=$lwidth, color=$col_ipopt");
w3 = PGF.Plots.Linear(t_span, c_span,  mark="none", legendentry="ALTRO",  style="line width=$lwidth, color=$col_altro");
w4 = PGF.Plots.Linear(t_span2,c_span2, mark="none", legendentry="AL-iLQR",style="line width=$lwidth, color=$col_alilqr");
w5 = PGF.Plots.Linear(t_pn*ones(100),range(1.0e-9,stop=1e5,length=100),legendentry="Projected Newton",mark="none",style="thick, color=red, dashed");

a = Axis([w1;w2;w4;w3;w5],
    xmin=0., ymin=1e-9, xmax=s3.stats[:iter_time][end], ymax=s3.stats[:c_max][1],
    legendPos="north east",
    ymode="log",
    hideAxis=false,
    xlabel="time (s)",
    ylabel="maximum constraint violation",
    style="grid=both")

save(joinpath(paper,"cartpole_c_max.tikz"), a, include_preamble=false)

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
#	GridWhale module
#	Copyright (c) 2020 Kronosaur Productions, LLC. All Rights Reserved.
#
#	This file provides functions for interacting with the Canvas port. It is
#	part of the GridWhale module.

module Canvas

function beginpath()
	ccall("CanvasBeginPath", Cvoid, ())
end

function closepath()
	ccall("CanvasClosePath", Cvoid, ())
end

function fill()
	ccall("CanvasFill", Cvoid, ())
end

function lineto(x::Real, y::Real)
	ccall("CanvasLineTo", Cvoid, (Cdouble, Cdouble), x, y)
end

function moveto(x::Real, y::Real)
	ccall("CanvasMoveTo", Cvoid, (Cdouble, Cdouble), x, y)
end

function stroke()
	ccall("CanvasStroke", Cvoid, ())
end

function setfillstyle(style::AbstractString)
	ccall("CanvasFillStyle", Cvoid, (Cstring, ), style)
end

function setstrokestyle(style::AbstractString)
	ccall("CanvasStrokeStyle", Cvoid, (Cstring, ), style)
end

end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

### basic definitions (types, aliases, constructors, abstractarray interface, sundry similar)

# note that QuasiAdjoint and QuasiTranspose must be able to wrap not only vectors and matrices
# but also factorizations, rotations, and other linear algebra objects, including
# user-defined such objects. so do not restrict the wrapped type.
struct QuasiAdjoint{T,S} <: AbstractQuasiMatrix{T}
    parent::S
end
struct QuasiTranspose{T,S} <: AbstractQuasiMatrix{T}
    parent::S
end

# basic outer constructors
QuasiAdjoint(A) = QuasiAdjoint{Base.promote_op(adjoint,eltype(A)),typeof(A)}(A)
QuasiTranspose(A) = QuasiTranspose{Base.promote_op(transpose,eltype(A)),typeof(A)}(A)

Base.dataids(A::Union{QuasiAdjoint, QuasiTranspose}) = Base.dataids(A.parent)
Base.unaliascopy(A::Union{QuasiAdjoint,QuasiTranspose}) = typeof(A)(Base.unaliascopy(A.parent))

# wrapping lowercase quasi-constructors
"""
    adjoint(A)

Lazy adjoint (conjugate transposition) (also postfix `'`).
Note that `adjoint` is applied recursively to elements.

This operation is intended for linear algebra usage - for general data manipulation see
[`permutedims`](@ref Base.permutedims).

# Examples
```jldoctest
julia> A = [3+2im 9+2im; 8+7im  4+6im]
2×2 Array{Complex{Int64},2}:
 3+2im  9+2im
 8+7im  4+6im

julia> adjoint(A)
2×2 QuasiAdjoint{Complex{Int64},Array{Complex{Int64},2}}:
 3-2im  8-7im
 9-2im  4-6im
```
"""
adjoint(A::AbstractQuasiVecOrMat) = QuasiAdjoint(A)

"""
    transpose(A)

Lazy transpose. Mutating the returned object should appropriately mutate `A`. Often,
but not always, yields `QuasiTranspose(A)`, where `QuasiTranspose` is a lazy transpose wrapper. Note
that this operation is recursive.

This operation is intended for linear algebra usage - for general data manipulation see
[`permutedims`](@ref Base.permutedims), which is non-recursive.

# Examples
```jldoctest
julia> A = [3+2im 9+2im; 8+7im  4+6im]
2×2 Array{Complex{Int64},2}:
 3+2im  9+2im
 8+7im  4+6im

julia> transpose(A)
2×2 QuasiTranspose{Complex{Int64},Array{Complex{Int64},2}}:
 3+2im  8+7im
 9+2im  4+6im
```
"""
transpose(A::AbstractQuasiVecOrMat) = QuasiTranspose(A)

# unwrapping lowercase quasi-constructors
adjoint(A::QuasiAdjoint) = A.parent
transpose(A::QuasiTranspose) = A.parent
adjoint(A::QuasiTranspose{<:Real}) = A.parent
transpose(A::QuasiAdjoint{<:Real}) = A.parent

# some aliases for internal convenience use
const AdjOrTrans{T,S} = Union{QuasiAdjoint{T,S},QuasiTranspose{T,S}} where {T,S}
const QuasiAdjointAbsVec{T} = QuasiAdjoint{T,<:AbstractQuasiVector}
const QuasiTransposeAbsVec{T} = QuasiTranspose{T,<:AbstractQuasiVector}
const AdjOrTransAbsVec{T} = AdjOrTrans{T,<:AbstractQuasiVector}
const AdjOrTransAbsMat{T} = AdjOrTrans{T,<:AbstractQuasiMatrix}

# for internal use below
wrapperop(A::QuasiAdjoint) = adjoint
wrapperop(A::QuasiTranspose) = transpose

# AbstractQuasiArray interface, basic definitions
length(A::AdjOrTrans) = length(A.parent)
size(v::AdjOrTransAbsVec) = (1, length(v.parent))
size(A::AdjOrTransAbsMat) = reverse(size(A.parent))
axes(v::AdjOrTransAbsVec) = (Base.OneTo(1), axes(v.parent)...)
axes(A::AdjOrTransAbsMat) = reverse(axes(A.parent))
IndexStyle(::Type{<:AdjOrTransAbsVec}) = IndexCartesian()
IndexStyle(::Type{<:AdjOrTransAbsMat}) = IndexCartesian()
@propagate_inbounds function getindex(A::AdjOrTransAbsVec, i::Number, j::Number)
    @boundscheck i == 1 || throw(BoundsError(A,i,j))
    wrapperop(A)(A.parent[j])
end
@propagate_inbounds getindex(A::AdjOrTransAbsMat, i::Number, j::Number) = wrapperop(A)(A.parent[j, i])
@propagate_inbounds function setindex!(A::AdjOrTransAbsVec, x, i::Number, j::Number)
    @boundscheck i == 1 || throw(BoundsError(A,i,j))
    (setindex!(A.parent, wrapperop(A)(x), j); A)
end
@propagate_inbounds setindex!(A::AdjOrTransAbsMat, x, i::Number, j::Number) = (setindex!(A.parent, wrapperop(A)(x), j, i); A)
# AbstractQuasiArray interface, additional definitions to retain wrapper over vectors where appropriate
@propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, is::AbstractArray{<:Number}) = wrapperop(v)(v.parent[is])
@propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, ::Colon) = wrapperop(v)(v.parent[:])

# conversion of underlying storage
convert(::Type{QuasiAdjoint{T,S}}, A::QuasiAdjoint) where {T,S} = QuasiAdjoint{T,S}(convert(S, A.parent))
convert(::Type{QuasiTranspose{T,S}}, A::QuasiTranspose) where {T,S} = QuasiTranspose{T,S}(convert(S, A.parent))

# for vectors, the semantics of the wrapped and unwrapped types differ
# so attempt to maintain both the parent and wrapper type insofar as possible
similar(A::AdjOrTransAbsVec) = wrapperop(A)(similar(A.parent))
similar(A::AdjOrTransAbsVec, ::Type{T}) where {T} = wrapperop(A)(similar(A.parent, Base.promote_op(wrapperop(A), T)))
# for matrices, the semantics of the wrapped and unwrapped types are generally the same
# and as you are allocating with similar anyway, you might as well get something unwrapped
similar(A::AdjOrTrans) = similar(A.parent, eltype(A), axes(A))
similar(A::AdjOrTrans, ::Type{T}) where {T} = similar(A.parent, T, axes(A))
similar(A::AdjOrTrans, ::Type{T}, dims::Dims{N}) where {T,N} = similar(A.parent, T, dims)

# sundry basic definitions
parent(A::AdjOrTrans) = A.parent
vec(v::AdjOrTransAbsVec) = v.parent

cmp(A::AdjOrTransAbsVec, B::AdjOrTransAbsVec) = cmp(parent(A), parent(B))
isless(A::AdjOrTransAbsVec, B::AdjOrTransAbsVec) = isless(parent(A), parent(B))

### concatenation
# preserve QuasiAdjoint/QuasiTranspose wrapper around vectors
# to retain the associated semantics post-concatenation
hcat(avs::Union{Number,QuasiAdjointAbsVec}...) = _adjoint_hcat(avs...)
hcat(tvs::Union{Number,QuasiTransposeAbsVec}...) = _transpose_hcat(tvs...)
_adjoint_hcat(avs::Union{Number,QuasiAdjointAbsVec}...) = adjoint(vcat(map(adjoint, avs)...))
_transpose_hcat(tvs::Union{Number,QuasiTransposeAbsVec}...) = transpose(vcat(map(transpose, tvs)...))
typed_hcat(::Type{T}, avs::Union{Number,QuasiAdjointAbsVec}...) where {T} = adjoint(typed_vcat(T, map(adjoint, avs)...))
typed_hcat(::Type{T}, tvs::Union{Number,QuasiTransposeAbsVec}...) where {T} = transpose(typed_vcat(T, map(transpose, tvs)...))
# otherwise-redundant definitions necessary to prevent hitting the concat methods in sparse/sparsevector.jl
hcat(avs::QuasiAdjoint{<:Any,<:Vector}...) = _adjoint_hcat(avs...)
hcat(tvs::QuasiTranspose{<:Any,<:Vector}...) = _transpose_hcat(tvs...)
hcat(avs::QuasiAdjoint{T,Vector{T}}...) where {T} = _adjoint_hcat(avs...)
hcat(tvs::QuasiTranspose{T,Vector{T}}...) where {T} = _transpose_hcat(tvs...)
# TODO unify and allow mixed combinations

### higher order functions
# preserve QuasiAdjoint/QuasiTranspose wrapper around vectors
# to retain the associated semantics post-map/broadcast
#
# note that the caller's operation f operates in the domain of the wrapped vectors' entries.
# hence the adjoint->f->adjoint shenanigans applied to the parent vectors' entries.
map(f, avs::QuasiAdjointAbsVec...) = adjoint(map((xs...) -> adjoint(f(adjoint.(xs)...)), parent.(avs)...))
map(f, tvs::QuasiTransposeAbsVec...) = transpose(map((xs...) -> transpose(f(transpose.(xs)...)), parent.(tvs)...))

### linear algebra

(-)(A::QuasiAdjoint)   = QuasiAdjoint(  -A.parent)
(-)(A::QuasiTranspose) = QuasiTranspose(-A.parent)

## multiplication *

# QuasiAdjoint/QuasiTranspose-vector * vector

# allow re-expansion using Legendre
_dot(ax, a, b) = Base.invoke(dot, NTuple{2,Any}, a, b)

function dot(a::AbstractQuasiArray, b::AbstractQuasiArray)
    ax = axes(a,1)
    ax == axes(b,1) || throw(DimensionMismatch())
    _dot(ax, a, b)
end

*(u::QuasiAdjointAbsVec, v::AbstractQuasiVector) = dot(u.parent, v)
*(u::QuasiTransposeAbsVec{T}, v::AbstractQuasiVector{T}) where {T<:Real} = dot(u.parent, v)
function *(u::QuasiTransposeAbsVec, v::AbstractQuasiVector)
    @boundscheck axes(u,2) == axes(v,1) || throw(DimensionMismatch())
    return sum(@inbounds(u.parent[k]*v[k]) for k in axes(v,1))
end
# vector * QuasiAdjoint/QuasiTranspose-vector
*(u::AbstractQuasiVector, v::AdjOrTransAbsVec) = broadcast(*, u, v)
# QuasiAdjoint/QuasiTranspose-vector * QuasiAdjoint/QuasiTranspose-vector
# (necessary for disambiguation with fallback methods in linalg/matmul)
*(u::QuasiAdjointAbsVec, v::QuasiAdjointAbsVec) = throw(MethodError(*, (u, v)))
*(u::QuasiTransposeAbsVec, v::QuasiTransposeAbsVec) = throw(MethodError(*, (u, v)))

# AdjOrTransAbsVec{<:Any,<:AdjOrTransAbsVec} is a lazy conj vectors
# We need to expand the combinations to avoid ambiguities
(*)(u::QuasiTransposeAbsVec, v::QuasiAdjointAbsVec{<:Any,<:QuasiTransposeAbsVec}) =
    sum(uu*vv for (uu, vv) in zip(u, v))
(*)(u::QuasiAdjointAbsVec,   v::QuasiAdjointAbsVec{<:Any,<:QuasiTransposeAbsVec}) =
    sum(uu*vv for (uu, vv) in zip(u, v))
(*)(u::QuasiTransposeAbsVec, v::QuasiTransposeAbsVec{<:Any,<:QuasiAdjointAbsVec}) =
    sum(uu*vv for (uu, vv) in zip(u, v))
(*)(u::QuasiAdjointAbsVec,   v::QuasiTransposeAbsVec{<:Any,<:QuasiAdjointAbsVec}) =
    sum(uu*vv for (uu, vv) in zip(u, v))

## pseudoinversion
pinv(v::QuasiAdjointAbsVec, tol::Number = 0) = pinv(v.parent, tol).parent
pinv(v::QuasiTransposeAbsVec, tol::Number = 0) = pinv(conj(v.parent)).parent

## left-division \
\(u::AdjOrTransAbsVec, v::AdjOrTransAbsVec) = pinv(u) * v

## right-division \
/(u::QuasiAdjointAbsVec, A::AbstractQuasiMatrix) = adjoint(adjoint(A) \ u.parent)
/(u::QuasiTransposeAbsVec, A::AbstractQuasiMatrix) = transpose(transpose(A) \ u.parent)
/(u::QuasiAdjointAbsVec, A::QuasiTranspose{<:Any,<:AbstractQuasiMatrix}) = adjoint(conj(A.parent) \ u.parent) # technically should be adjoint(copy(adjoint(copy(A))) \ u.parent)
/(u::QuasiTransposeAbsVec, A::QuasiAdjoint{<:Any,<:AbstractQuasiMatrix}) = transpose(conj(A.parent) \ u.parent) # technically should be transpose(copy(transpose(copy(A))) \ u.parent)

function adjoint(M::Mul)
    Mul(reverse(adjoint.(M.args))...)
end

==(A::QuasiAdjoint, B::QuasiAdjoint) = parent(A) == parent(B)

adjointlayout(::Type, ::QuasiArrayLayout) = QuasiArrayLayout()
transposelayout(::QuasiArrayLayout) = QuasiArrayLayout()

MemoryLayout(::Type{QuasiTranspose{T,P}}) where {T,P} = transposelayout(MemoryLayout(P))
MemoryLayout(::Type{QuasiAdjoint{T,P}}) where {T,P} = adjointlayout(T, MemoryLayout(P))

call(::ApplyLayout{typeof(*)}, V::QuasiAdjoint) = *
call(::ApplyLayout{typeof(*)}, V::QuasiTranspose) = *
arguments(LAY::ApplyLayout{typeof(*)}, V::QuasiAdjoint) = reverse(adjoint.(arguments(LAY, V')))
arguments(LAY::ApplyLayout{typeof(*)}, V::QuasiTranspose) = reverse(transpose.(arguments(LAY, V')))

# This is used in ContinuumArrays.jl to ensure x' is lazy
BroadcastStyle(::Type{<:QuasiAdjoint{<:Any,<:Inclusion}}) = LazyQuasiArrayStyle{2}()
BroadcastStyle(::Type{<:QuasiTranspose{<:Any,<:Inclusion}}) = LazyQuasiArrayStyle{2}()

###
# adjoint concat support
###

arguments(::ApplyLayout{typeof(vcat)}, A::QuasiAdjoint) = map(adjoint, arguments(ApplyLayout{typeof(hcat)}(), parent(A)))
arguments(::ApplyLayout{typeof(hcat)}, A::QuasiAdjoint) = map(adjoint, arguments(ApplyLayout{typeof(vcat)}(), parent(A)))
arguments(::ApplyLayout{typeof(vcat)}, A::QuasiTranspose) = map(transpose, arguments(ApplyLayout{typeof(hcat)}(), parent(A)))
arguments(::ApplyLayout{typeof(hcat)}, A::QuasiTranspose) = map(transpose, arguments(ApplyLayout{typeof(vcat)}(), parent(A)))

copy(M::Mul{ApplyLayout{typeof(vcat)},QuasiArrayLayout}) = vcat((arguments(vcat, M.A) .* Ref(M.B))...)


text:
#-----------------------------------------------------------------------# utils

_repr(o::Nothing) = ""
_repr(o) = repr("text/plain", o)

#-----------------------------------------------------------------------# renderers
"""
    render(o::CodeBlock, r::Val{:renderer}; kw...)::String

Render a [`CodeBlock`](@ref) for the given `:renderer`, which can be used via

    ```julia; renderer;
    ...
    ```

`kw...` always contains `builddir`.
"""
function render(o::CodeBlock, r::Val{T}; kw...) where {T}
    @warn "Reproducible does not recognize the renderer `$T`.  Code block not eval-ed."
    block(codestring(o))
end

#-----------------------------------------------------------------------# run
render(o::CodeBlock, r::Val{:run}; kw...) = block(codestring(o), "julia")

#-----------------------------------------------------------------------# hide
render(o::CodeBlock, r::Val{:hide}; kw...) = ""

#-----------------------------------------------------------------------# block
function render(o::CodeBlock, r::Val{:block}; transform=identity, kw...)
    block(codestring(o)) * "\n" * block(transform(o.rows[end].repldisplay), "julia")
end

#-----------------------------------------------------------------------# repl
function render(o::CodeBlock, r::Val{:repl}; transform=identity, kw...)
    s = ""
    for row in o.rows
        s *= "julia> " * strip(row.input)
        s *= '\n' * strip(row.repldisplay)
        s *= "\n\n"
    end
    block(s, "julia")
end

#-----------------------------------------------------------------------# docstring
function render(o::CodeBlock, r::Val{:docstring}; kw...)
    Markdown.plain(Docs.doc(output(o)))
end

#-----------------------------------------------------------------------# test utils
function render(o::CodeBlock, r::Val{:rendertest}; renderer::Symbol=:block)
    """
    **Render Test for $renderer**:

    ````
    $(block(codestring(o), "julia; $renderer"))````

    $(render(o, Val(renderer)))
    """
end

function render(o::CodeBlock, r::Val{:testall}; kw...)
    render(o, Val(:rendertest); renderer = :block) *
        render(o, Val(:rendertest); renderer = :repl) *
        render(o, Val(:rendertest); renderer = :run)
end


text:
using SequentialMonteCarlo
using RNGPool
import SMCExamples.SMCSampler.defaultSMCSampler
using MonteCarloMarkovKernels
using StaticArrays
using Plots
Plots.gr()

setRNGs(0)

model, ltarget = defaultSMCSampler()

smcio = SMCIO{model.particle, model.pScratch}(1024*1024, model.maxn,
  Threads.nthreads(), false)
smc!(model, smcio)

xs = (p->p.x[1]).(smcio.zetas)
ys = (p->p.x[2]).(smcio.zetas)

## bimodal target needs a kde bandwidth adjustment since its variance is large
# x, y, f = MonteCarloMarkovKernels.kde(xs, ys, 0.005)
x, y, f = MonteCarloMarkovKernels.kde(xs, ys, 1.0, 0.001)
contour(x, y, f)
contour!(x,y, (x,y) -> exp(ltarget((SVector{2,Float64}(x,y)))))


text:
# # Twisted beam: Export solid model to Abaqus

# Source code: [`twisted_beam-export-to-abaqus_tut.jl`](twisted_beam-export-to-abaqus_tut.jl)

# ## Description

# In this example  we show how to export a model  to the finite element software Abaqus.

# The model is solved also in the example `twisted_beam_algo.jl`.  Here we export the model for execution in Abaqus.

# The task begins with defining the input parameters, creating the mesh, identifying the nodes  to which essential boundary conditions are to be applied,  and extracting from the boundary the surface finite elements to which the traction loading at the end of the beam is to be applied.

# This is the finite element toolkit itself.
using FinEtools

# The linear stress analysis application is implemented in this package.
using FinEtoolsDeforLinear

# We will also use specifically some functions from these modules.
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.MeshExportModule

# Define some parameters, in consistent units.
E = 0.29e8;
nu = 0.22;
W = 1.1;
L = 12.;
t =  0.32;
nl = 2; nt = 1; nw = 1; ref = 2;
p =   1/W/t;
#  Loading in the Z direction. Reference (publication by Harder): 5.424e-3.
loadv = [0;0;p]; dir = 3; uex = 0.005424534868469;
#   Loading in the Y direction. Reference (Harder): 1.754e-3. And comment the line below to obtain a solution in the other direction.
#loadv = [0;p;0]; dir = 2; uex = 0.001753248285256;
tolerance  = t/1000;

fens,fes  = H8block(L,W,t, nl*ref,nw*ref,nt*ref)

# Reshape the rectangular block into a twisted beam shape.
for i = 1:count(fens)
    let
        a = fens.xyz[i,1]/L*(pi/2); y = fens.xyz[i,2]-(W/2); z = fens.xyz[i,3]-(t/2);
        fens.xyz[i,:] = [fens.xyz[i,1],y*cos(a)-z*sin(a),y*sin(a)+z*cos(a)];
    end
end

# Clamped face of the beam: select all the nodes in this cross-section.
l1  = selectnode(fens; box = [0 0 -100*W 100*W -100*W 100*W], inflate  =  tolerance)

# Traction on the opposite face
boundaryfes  =   meshboundary(fes);
Toplist   = selectelem(fens,boundaryfes, box =  [L L -100*W 100*W -100*W 100*W], inflate =   tolerance);

# The tutorial proper begins here. We create the Abaqus exporter and start writing the .inp file.
AE = AbaqusExporter("twisted_beam");
HEADING(AE, "Twisted beam example");

# The  part definition is trivial: all will be defined rather for the instance of the part.
PART(AE, "part1");
END_PART(AE);

# The assembly will consist  of a single instance (of the empty part defined above).  The node set will be defined for the instance itself.
ASSEMBLY(AE, "ASSEM1");
INSTANCE(AE, "INSTNC1", "PART1");
NODE(AE, fens.xyz);

# We export the finite elements themselves.  Note that the elements  need to have  distinct numbers.  We start numbering the hexahedra at 1. The definition of the element creates simultaneously an element set  which is used below in the section assignment (and the definition of the load).
ELEMENT(AE, "c3d8rh", "AllElements", 1, connasarray(fes))

# The traction is applied to surface elements.  Because the elements in the Abaqus model need to have unique numbers, we need to start from an integer  which is  the number of the solid elements plus one.
ELEMENT(AE, "SFM3D4", "TractionElements", 1+count(fes), connasarray(subset(boundaryfes,Toplist)))

# The nodes in the clamped cross-section are going to be grouped in the node set `l1`.
NSET_NSET(AE, "l1", l1)

# We define a coordinate system  (orientation of the material  coordinate system), in this example it is the global Cartesian coordinate system. The sections are defined for the solid elements of the interior and the surface elements to which the traction is applied, and the assignment to the  elements is by element set (`AllElements` and `TractionElements`). Note that for the solid section we also define reference  to hourglass control named `Hourglassctl`.
ORIENTATION(AE, "GlobalOrientation", vec([1. 0 0]), vec([0 1. 0]));
SOLID_SECTION(AE, "elasticity", "GlobalOrientation", "AllElements", "Hourglassctl");
SURFACE_SECTION(AE, "TractionElements")

# This concludes the definition  of the instance  and of the assembly.
END_INSTANCE(AE);
END_ASSEMBLY(AE);

# This is the definition of the isotropic elastic material.
MATERIAL(AE, "elasticity")
ELASTIC(AE, E, nu)

# The element properties for the interior hexahedra are controlled by the section-control.  In this case we are selecting enhanced hourglass stabilization (much preferable to the default  stiffness stabilization).
SECTION_CONTROLS(AE, "Hourglassctl", "HOURGLASS=ENHANCED")

# The static perturbation  analysis step is defined  next.
STEP_PERTURBATION_STATIC(AE)

# The boundary conditions are applied directly to the node set `l1`.   Since the node set is defined for the instance, we need to refer to it by the qualified name `ASSEM1.INSTNC1.l1`.
BOUNDARY(AE, "ASSEM1.INSTNC1.l1", 1)
BOUNDARY(AE, "ASSEM1.INSTNC1.l1", 2)
BOUNDARY(AE, "ASSEM1.INSTNC1.l1", 3)

# The traction is applied to the surface  quadrilateral elements exported above.
DLOAD(AE, "ASSEM1.INSTNC1.TractionElements", vec(loadv))

# Now we have defined  the analysis step and the definition of the model can be concluded.
END_STEP(AE)
close(AE)

# As quick check, here is the contents of the  exported model file:
@show readlines("twisted_beam.inp")

# What remains is to load the model into Abaqus and execute it as a job.  Alternatively Abaqus can be called on the input file to carry out the analysis at the command line as
# ```
# abaqus job=twisted_beam.inp
# ```
# The output database `twisted_beam.odb` can then be loaded for postprocessing, for instance from the command line as
# ```
# abaqus viewer database=twisted_beam.odb
# ```

nothing