text:
export fillflows!, makepdct, check_flowgrad, update_payoffs!, fduψ, fillflows_grad!

# ----------------------------------- wrappers to fill a vector -------------------------------------------

function duθ!(du::AbstractVector{T}, FF::Type, wp::AbstractUnitProblem, θ::AbstractVector{T}, σ::T, z::Tuple, stinfo::Tuple, typestuff::Real...) where {T}
    K = length(θ)
    length(du) == K  || throw(DimensionMismatch())
    stinfo[end]
    @inbounds for k = 1:K
        du[k] = flowdθ(FF, wp, θ, σ, z, stinfo[end-1], k, stinfo[end], typestuff...)
    end
end

# ------------------------ check flow grad --------------

function check_flowgrad(FF::Type, θ::AbstractVector{T}, σ::T, zspace::Tuple, ψspace::AbstractRange, wp::AbstractUnitProblem, itype::Real...) where {T}
    K = length(θ)
    dx = Vector{T}(undef, K)
    fdx = similar(dx)

    zpdct = Base.product(zspace...)
    pdct = Base.product(zpdct, ψspace, 0:_dmax(wp))

    for zψd in pdct
        z, ψ, d = zψd

        for sidx ∈ 1:ShaleDrillingModel._nS(wp)
            u(θ) = flow(FF, wp, sidx, θ, σ, z, ψ, d, itype...)
            Calculus.finite_difference!(u, θ, fdx, :central)
            @inbounds for k ∈ 1:K
                dx[k] = flowdθ(FF, wp, sidx, θ, σ, z, ψ, k, d, itype...)
            end
            if !(fdx ≈ dx)
                @warn "FF = $FF. Bad θ diff at sidx=$sidx, (z,ψ,d)=$zψd. du=$dx and fd = $dxfd"
                return false
            end
        end

        # check σ
        for sidx ∈ 1:min(2,end_ex0(wp))
            dσ = flowdσ(FF, wp, sidx, θ, σ, z, ψ, d, itype...)
            fdσ = Calculus.derivative((σh::Real) -> flow(FF, wp, sidx, θ, σh, z, ψ, d, itype...), σ, :central)
            if !(dσ ≈ fdσ) && !isapprox(dσ, fdσ, atol= 1e-7)
                @warn "Bad σ diff at sidx=$sidx with (z,ψ,d)=$zψd. duσ = $dσ and fd = $fdσ"
                return false
            end
        end

        # check ψ
        for sidx in 1:min(4,end_ex0(wp))
            z,ψ,d = zψd
            dψ = flowdψ(FF, wp, sidx, θ, σ, z, ψ, d, itype...)
            fdψ = Calculus.derivative((h::Real) -> flow(FF, wp, sidx, θ, σ, z, ψ+h, d, itype...), 0.0, :central)
            if !(dψ ≈ fdψ) && !isapprox(dψ, fdψ, atol=1e-7)
                @warn "Bad ψ diff at sidx=$sidx with (z,ψ,d)=$(zψd). duψ = $dψ and fdψ = $fdψ"
                return false
            end
        end
    end
    return true
end

check_flowgrad(θ::AbstractVector, σ::Real, p::dcdp_primitives{FF}, itype::Real...) where {FF} = check_flowgrad(FF, θ, σ, p.zspace, p.ψspace, p.wp, itype...)
check_flowgrad(θ::AbstractVector,          p::dcdp_primitives{FF}, itype::Real...) where {FF} = check_flowgrad(FF, _θt(θ, p), _σv(θ), p, itype...)

# ------------------------ fill flows --------------

function fillflows!(X::AbstractArray, f::Function, p::dcdp_primitives{FF}, θ::AbstractVector, σ::Real, sidx::Integer, itype::Real...) where {FF}
    zpdct = Base.product(p.zspace...)
    idxd = actionspace(p.wp, sidx)
    pdct = Base.product(zpdct, p.ψspace, idxd)
    length(X) == length(pdct) || throw(DimensionMismatch("state[$sidx] is $(state(p.wp,sidx)). size(X) = $(size(X)) != size(pdct) = $(size(pdct))"))
    @inbounds for (i,zψd,) in enumerate(pdct)
        X[i] = f(FF, p.wp, sidx, θ, σ, zψd..., itype...)
    end
end

function fillflows_grad!(X::AbstractArray, f::Function, p::dcdp_primitives{FF}, θ::AbstractVector, σ::Real, sidx::Integer, itype::Real...) where {FF}
    zpdct = Base.product(p.zspace...)
    idxd = 0:size(X, ndims(X))-1
    pdct = Base.product(zpdct, p.ψspace, Base.OneTo(length(θ)), idxd)
    length(X) == length(pdct) || throw(DimensionMismatch("state[$sidx] is $(state(p.wp,sidx)). size(X) = $(size(X)) != size(pdct) = $(size(pdct))"))
    @inbounds for (i,zψkd,) in enumerate(pdct)
        X[i] = f(FF, p.wp, sidx, θ, σ, zψkd..., itype...)
    end
end


text:
## object methods

Base.show(io::IO, o::Object) = println(io, "Object of type $(typeof(o))")

## Use o[:symbol] for properties (will also use o[:integer] for containers and children)
## We make it so o[:symbol] calls getSymbol (if present), else it looks in o.attrs[:symbol]
## Similarly, o[:symbol] = value will look for setSymbol (if present). This makes it easy to
## overload the assignment.
function getindex(o::Object, i::Symbol)
    getProp(o, i)
end

function getProp(o::Object, i::Symbol)
    ## do we have a function?
    aname = string(i)
    aname = uppercase(aname[1:1]) * aname[2:end]
    fname = "get" * aname
    try
        JGUI.(symbol(fname))(o)
    catch e
        if haskey(o.attrs, i)
            o.attrs[i]
        else
            nothing
        end
    end
end

## By default, we look for a method getPropname, if that is defined we use it
## otherwise we look for the key in attrs Dict.
function getProp(o::Object, i::Symbol, default)
    val = getProp(o, i)
    isa(val, Nothing) ? default : val
end

function setindex!(o::Object, value, i::Symbol)
    setProp(o, i, value)
end

function setProp(o::Object, i::Symbol, value)
    aname = string(i)
    aname = uppercase(aname[1:1]) * aname[2:end]
    fname = "set" * aname
    try
        JGUI.(symbol(fname))(o, value)
    catch e
        println(e)
        o.attrs[i] = value
    end
    o
end

## A means to document inheritance of properties
## SOmething like `list_props(::@PROP("WidgetType")) = {:prop => "description"}`
## is there a cleaner way? This is cribbed from MIME
## This allows us to dispatch on a DataType
immutable PROP{atype} end
PROP(s) = PROP{symbol(s)}()
macro PROP(s)
    quote
        PROP{symbol($s)}
    end
end

## list properties
immutable ObjectProperties
    x::Dict
end
function Base.show(io::IO, p::ObjectProperties)
    for (k, v) in p.x
        println(k)
        [println(io, "\t$k => $v") for (k,v) in v]
    end
end

## return available properties for a widget
function properties(o::Object)
    d = Dict()
    s = typeof(o)
    while s != Any
        try
            d[string(s)] = list_props(PROP(string(s)))
        catch e
            nothing
        end
        s = super(s)
    end
    ObjectProperties(d)
end

ext: jl


text:
struct DecoupledControlParams{T}
    V_min::T
    V_max::T

    k_V::T
    k_s::T

    δ̇_max::T

    Q_Δψ::T
    Q_e::T
    W_β::T
    W_r::T

    R_δ::T
    R_Δδ::T
end
function DecoupledControlParams(;V_min=1.0,
                                 V_max=15.0,
                                 k_V=10/4/100,
                                 k_s=10/4/10000,
                                 δ̇_max=0.344,
                                 Q_Δψ=1/(10*π/180)^2,
                                 Q_e=1.0,
                                 W_β=50/(10*π/180),
                                 W_r=50.0,
                                 R_δ=0.0,
                                 R_Δδ=0.01/(10*π/180)^2)
    DecoupledControlParams(V_min, V_max, k_V, k_s, δ̇_max, Q_Δψ, Q_e, W_β, W_r, R_δ, R_Δδ)
end

function DecoupledTrajectoryTrackingMPC(vehicle::Dict{Symbol,T}, trajectory::TrajectoryTube{T}; control_params=DecoupledControlParams(),
                                        N_short=10, N_long=20, dt_short=T(0.01), dt_long=T(0.2), use_correction_step=true) where {T}
    dynamics = VehicleModel(vehicle)
    current_state = zeros(BicycleState{T})
    current_control = zeros(BicycleControl{T})
    time_steps = MPCTimeSteps(N_short, N_long, dt_short, dt_long, use_correction_step)

    N = 1 + N_short + N_long
    qs = rand(LateralTrackingBicycleState{T}, N)    # not zeros so that construct_lateral_tracking_QP below works
    us = rand(BicycleControl2{T}, N)
    ps = rand(LateralTrackingBicycleParams{T}, N)
    tracking_dynamics = VehicleModel(vehicle, LateralTrackingBicycleModel(vehicle))
    model, variables, parameters = construct_lateral_tracking_QP(tracking_dynamics, control_params, time_steps, qs, us, ps)
    TrajectoryTrackingMPC(vehicle, trajectory, dynamics, control_params,
                          current_state, current_control, 0, NaN,
                          time_steps,
                          qs, us, ps,
                          tracking_dynamics, model, variables, parameters)
end

function compute_linearization_nodes!(mpc::TrajectoryTrackingMPC{T},
                                      qs::Vector{LateralTrackingBicycleState{T}},
                                      us::Vector{BicycleControl2{T}},
                                      ps::Vector{LateralTrackingBicycleParams{T}}) where {T}
    traj = mpc.trajectory
    X  = mpc.tracking_dynamics
    U  = mpc.control_params
    q0 = mpc.current_state
    u0 = mpc.current_control
    TS = mpc.time_steps
    N_short, N_long = TS.N_short, TS.N_long
    ts, dt = TS.ts, TS.dt

    s, e0, t0 = path_coordinates(traj, q0)

    V = hypot(q0.Ux, q0.Uy)
    β0 = atan(q0.Uy, q0.Ux)
    r0 = q0.r
    δ0 = u0.δ
    Fyf0, Fyr0 = lateral_tire_forces(X.bicycle_model, q0, u0)
    for i in 1:N_short+N_long+1
        τ = (i == N_short+N_long+1 ? dt[i-1] : dt[i])
        tj = traj[s]
        κ = tj.κ
        A_des = tj.A + U.k_V*(tj.V - V)/τ + (isnan(mpc.time_offset) ? 0 : U.k_s*(traj(ts[i]).s - s)/τ/τ)   # alternatively, - U.k_t*(tj.t - ts[i])/τ/τ
        A_des = min(max(A_des, (U.V_min - V)/τ), (U.V_max - V)/τ)
        if i == 1
            q = LateralTrackingBicycleState(q0.Uy, q0.r, adiff(q0.ψ, tj.ψ), e0)    # to match paper should be (..., 0.0, 0.0)
            u = BicycleControl2(u0)
            p = LateralTrackingBicycleParams(q0.Ux, κ, T(0), T(0))
            q̇ = mpc.dynamics(q0, BicycleControl2(u0), LocalRoadGeometry(tj))
            A = (q̇[4] - q0.r*q0.Uy)*cos(β0) + (q̇[5] + q0.r*q0.Ux)*sin(β0)
        elseif i <= N_short+1
            q = LateralTrackingBicycleState(q0.Uy, q0.r, adiff(q0.ψ, tj.ψ), e0)    # to match paper should be (..., 0.0, 0.0)
            est = steady_state_estimates(X, V, A_des, κ, num_iters=1, r=r0, β0=β0, δ0=δ0, Fyf0=Fyf0)
            u = BicycleControl2(est.δ, est.Fxf + est.Fxr)
            p = LateralTrackingBicycleParams(est.Ux, κ, T(0), T(0))
            A = est.A
        else
            est = steady_state_estimates(X, V, A_des, κ)
            q = LateralTrackingBicycleState(est.Uy, est.r, -est.β, T(0))
            u = BicycleControl2(est.δ, est.Fxf + est.Fxr)
            p = LateralTrackingBicycleParams(est.Ux, κ, T(0), T(0))
            A = est.A
        end
        qs[i] = q
        us[i] = u
        ps[i] = p
        i == N_short+N_long+1 && break
        V = V + A*τ
        s = s + V*τ + A*τ*τ/2
    end
end

struct LateralTrackingQPParams{T}
    Q_Δψ  ::Parameter{Diagonal{T,Vector{T}},typeof(identity),true}
    Q_e   ::Parameter{Diagonal{T,Vector{T}},typeof(identity),true}
    R_δ   ::Parameter{Diagonal{T,Vector{T}},typeof(identity),true}
    R_Δδ  ::Parameter{Diagonal{T,Vector{T}},typeof(identity),true}
    W_β   ::Parameter{Vector{T},typeof(identity),true}
    W_r   ::Parameter{Vector{T},typeof(identity),true}
    q_curr::Parameter{Vector{T},typeof(identity),true}
    δ_curr::Parameter{Vector{T},typeof(identity),true}
    A     ::Vector{Parameter{Matrix{T},typeof(identity),true}}
    B     ::Vector{Parameter{Matrix{T},typeof(identity),true}}
    B0    ::Vector{Parameter{Matrix{T},typeof(identity),true}}
    Bf    ::Vector{Parameter{Matrix{T},typeof(identity),true}}
    c     ::Vector{Parameter{Vector{T},typeof(identity),true}}
    H     ::Vector{Parameter{Matrix{T},typeof(identity),true}}
    G     ::Vector{Parameter{Vector{T},typeof(identity),true}}
    δ_min ::Vector{Parameter{Vector{T},typeof(identity),true}}
    δ_max ::Vector{Parameter{Vector{T},typeof(identity),true}}
    Δδ_min::Vector{Parameter{Vector{T},typeof(identity),true}}
    Δδ_max::Vector{Parameter{Vector{T},typeof(identity),true}}
end

struct LateralTrackingQPVariables
    q::Matrix{Variable}
    δ::Matrix{Variable}
    σ::Matrix{Variable}
end

function construct_lateral_tracking_QP(dynamics::VehicleModel{T}, control_params, time_steps, qs, us, ps) where {T}
    N_short, N_long, dt = time_steps.N_short, time_steps.N_long, time_steps.dt

    optimizer = OSQPOptimizer()
    MOI.set!(optimizer, OSQPSettings.Verbose(), false)
    MOI.set!(optimizer, OSQPSettings.WarmStart(), true)
    m = Model{T}(optimizer)

    Q_Δψ   = Parameter(Diagonal(control_params.Q_Δψ .* dt), m)
    Q_e    = Parameter(Diagonal(control_params.Q_e  .* dt), m)
    R_δ    = Parameter(Diagonal(control_params.R_δ  .* dt), m)
    R_Δδ   = Parameter(Diagonal(control_params.R_Δδ ./ dt), m)
    W_β    = Parameter(control_params.W_β .* dt, m)
    W_r    = Parameter(control_params.W_r .* dt, m)
    q_curr = Parameter(Array(qs[1]), m)
    δ_curr = Parameter([us[1].δ], m)
    A      = Parameter{Matrix{T},typeof(identity),true}[]
    B      = Parameter{Matrix{T},typeof(identity),true}[]
    B0     = Parameter{Matrix{T},typeof(identity),true}[]
    Bf     = Parameter{Matrix{T},typeof(identity),true}[]
    c      = Parameter{Vector{T},typeof(identity),true}[]
    H      = Parameter{Matrix{T},typeof(identity),true}[]
    G      = Parameter{Vector{T},typeof(identity),true}[]
    δ_min  = Parameter{Vector{T},typeof(identity),true}[]
    δ_max  = Parameter{Vector{T},typeof(identity),true}[]
    Δδ_min = Parameter{Vector{T},typeof(identity),true}[]
    Δδ_max = Parameter{Vector{T},typeof(identity),true}[]

    q  = [Variable(m) for i in 1:4, t in 1:N_short+N_long+1]
    δ  = [Variable(m) for i in 1:1, t in 1:N_short+N_long+1]
    σ  = [Variable(m) for i in 1:2, t in 1:N_short+N_long]
    Δδ = [Variable(m) for t in 1:N_short+N_long]
    @constraint(m, vec(σ) >= fill(T(0), 2*(N_short+N_long)))
    @constraint(m, diff(δ[1,:]) == Δδ)

    @constraint(m, q[:,1] == q_curr)
    @constraint(m, δ[:,1] == δ_curr)
    for t in 1:N_short
        ZOHt = linearize(linearize(dynamics, qs[t], [us[t]; ps[t]]),
                         qs[t], StepControl(dt[t], [us[t]; ps[t]]), keep_control_dims=SVector(1))
        # ZOHt = linearize(dynamics, qs[t], StepControl(dt[t], [us[t]; ps[t]]), keep_control_dims=SVector(1))
        At = push!(A, Parameter(Array(ZOHt.A), m))[end]
        Bt = push!(B, Parameter(Array(ZOHt.B), m))[end]
        ct = push!(c, Parameter(Array(ZOHt.c), m))[end]
        @constraint(m, At*q[:,t] + Bt*δ[:,t] + ct == q[:,t+1])
    end

    for t in N_short+1:N_short+N_long
        FOHt = linearize(linearize(dynamics, qs[t], [us[t]; ps[t]]),
                         qs[t], RampControl(dt[t], [us[t]; ps[t]], [us[t+1]; ps[t+1]]), keep_control_dims=SVector(1))
        # FOHt = linearize(dynamics, qs[t], RampControl(dt[t], [us[t]; ps[t]], [us[t+1]; ps[t+1]]), keep_control_dims=SVector(1))
        At  = push!(A,  Parameter(Array(FOHt.A),  m))[end]
        B0t = push!(B0, Parameter(Array(FOHt.B0), m))[end]
        Bft = push!(Bf, Parameter(Array(FOHt.Bf), m))[end]
        ct  = push!(c,  Parameter(Array(FOHt.c),  m))[end]
        @constraint(m, At*q[:,t] + B0t*δ[:,t] + Bft*δ[:,t+1] + ct == q[:,t+1])
    end

    δ_hwmax, δ̇_max = dynamics.control_limits.δ_max, control_params.δ̇_max
    for t in 1:N_short+N_long
        Uxt = ps[t+1][1]
        Fxft, Fxrt = longitudinal_tire_forces(dynamics.longitudinal_params, us[t+1][2])
        δ_mint, δ_maxt, Ht, Gt = stable_limits(dynamics.bicycle_model, Uxt, Fxft, Fxrt)
        Ht = push!(H, Parameter(Array(Ht), m))[end]
        Gt = push!(G, Parameter(Array(Gt), m))[end]
        δ_mint = push!(δ_min, Parameter([max(δ_mint, -δ_hwmax)], m))[end]
        δ_maxt = push!(δ_max, Parameter([min(δ_maxt, δ_hwmax)], m))[end]
        Uy_r = q[1:2,t+1]
        σt   = [σ[1,t],σ[1,t],σ[2,t],σ[2,t]]
        Δδt  = [Δδ[t]]
        Δδ_mint = push!(Δδ_min, Parameter([-δ̇_max*dt[t]], m))[end]
        Δδ_maxt = push!(Δδ_max, Parameter([δ̇_max*dt[t]],  m))[end]
        @constraint(m, δ[:,t+1] <= δ_maxt)
        @constraint(m, δ[:,t+1] >= δ_mint)
        @constraint(m, Ht*Uy_r - Gt <= σt)
        @constraint(m, Δδt <= Δδ_maxt)
        @constraint(m, Δδt >= Δδ_mint)
    end

    Δψ   = q[3,2:end]
    e    = q[4,2:end]
    δ⁺   = δ[1,2:end]
    σ1   = σ[1,:]
    σ2   = σ[2,:]
    obj  = @expression transpose(Δψ)*Q_Δψ*Δψ +
                       transpose(e)*Q_e*e +
                       transpose(δ⁺)*R_δ*δ⁺ +
                       transpose(Δδ)*R_Δδ*Δδ +
                       W_β⋅σ1 + W_r⋅σ2
    @objective(m, Minimize, obj)
    m, LateralTrackingQPVariables(q, δ, σ), LateralTrackingQPParams(Q_Δψ, Q_e, R_δ, R_Δδ, W_β, W_r, q_curr, δ_curr,
                                                                    A, B, B0, Bf, c, H, G, δ_min, δ_max, Δδ_min, Δδ_max)
end

function update_QP!(mpc::TrajectoryTrackingMPC, QPP::LateralTrackingQPParams)
    dynamics = mpc.tracking_dynamics
    control_params  = mpc.control_params
    time_steps = mpc.time_steps
    N_short, N_long, dt = time_steps.N_short, time_steps.N_long, time_steps.dt
    qs, us, ps = mpc.qs, mpc.us, mpc.ps

    QPP.Q_Δψ().diag .= control_params.Q_Δψ .* dt
    QPP.Q_e().diag  .= control_params.Q_e  .* dt
    QPP.R_δ().diag  .= control_params.R_δ  .* dt
    QPP.R_Δδ().diag .= control_params.R_Δδ ./ dt
    QPP.W_β()    .= control_params.W_β .* dt
    QPP.W_r()    .= control_params.W_r .* dt
    QPP.q_curr() .= qs[1]
    QPP.δ_curr() .= us[1].δ

    for t in 1:N_short
        ZOHt = linearize(linearize(dynamics, qs[t], [us[t]; ps[t]]),
                         qs[t], StepControl(dt[t], [us[t]; ps[t]]), keep_control_dims=SVector(1))
        # ZOHt = linearize(dynamics, qs[t], StepControl(dt[t], [us[t]; ps[t]]), keep_control_dims=SVector(1))
        QPP.A[t]() .= ZOHt.A
        QPP.B[t]() .= ZOHt.B
        QPP.c[t]() .= ZOHt.c
    end
    for t in N_short+1:N_short+N_long
        FOHt = linearize(linearize(dynamics, qs[t], [us[t]; ps[t]]),
                         qs[t], RampControl(dt[t], [us[t]; ps[t]], [us[t+1]; ps[t+1]]), keep_control_dims=SVector(1))
        # FOHt = linearize(dynamics, qs[t], RampControl(dt[t], [us[t]; ps[t]], [us[t+1]; ps[t+1]]), keep_control_dims=SVector(1))
        QPP.A[t]() .= FOHt.A
        QPP.B0[t-N_short]() .= FOHt.B0
        QPP.Bf[t-N_short]() .= FOHt.Bf
        QPP.c[t]() .= FOHt.c
    end
    δ_max, δ̇_max = dynamics.control_limits.δ_max, control_params.δ̇_max
    for t in 1:N_short+N_long
        Uxt = ps[t+1][1]
        Fxft, Fxrt = longitudinal_tire_forces(dynamics.longitudinal_params, us[t+1][2])
        δ_mint, δ_maxt, Ht, Gt = stable_limits(dynamics.bicycle_model, Uxt, Fxft, Fxrt)
        QPP.H[t]() .= Ht
        QPP.G[t]() .= Gt
        QPP.δ_min[t]() .= max(δ_mint, -δ_max)
        QPP.δ_max[t]() .= min(δ_maxt, δ_max)
        QPP.Δδ_min[t]() .= -δ̇_max*dt[t]
        QPP.Δδ_max[t]() .= δ̇_max*dt[t]
    end
end

function get_next_control(mpc::TrajectoryTrackingMPC, variables::LateralTrackingQPVariables)
    BicycleControl(mpc.dynamics.longitudinal_params,
                   BicycleControl2(value(mpc.model, variables.δ[2]), mpc.us[2].Fx))
end

max_stars_repo_name: Michelle-NYX/CollisionAvoidance

max_issues_repo_name: Michelle-NYX/CollisionAvoidance

max_forks_repo_name: Michelle-NYX/CollisionAvoidance


text:
#///////////////////////////////////////
#// File Name: RiskSensitiveSAC.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2020/12/04
#// Description: Julia package for Risk Sensitive Sequential Action
#// Control, with public version of Trajectron++
#///////////////////////////////////////

module RiskSensitiveSAC

import Base: ==, isequal, isapprox
import PlotUtils: plot_color
import RobotOS: Time, Duration
import StatsBase: fit, Histogram
import StatsFuns: logsumexp

using CUDA
using DataStructures
using Distributions
using ForwardDiff
using LinearAlgebra
using Plots
pyplot();
using Printf
using ProgressMeter
using PyCall
using Random

# Modules for Buffered Input Cell and Related Functions (and constrained SAC optimizations)
import Convex: Variable, norm, quadform, minimize, dot, solve!
using GeometricalPredicates
using LazySets
using Polyhedra
using SCS
using ECOS
using VoronoiCells

function __init__()
    @info "Number of Julia Thread(s): $(Threads.nthreads())";
    @info "CUDA Device: $(name(CuDevice(0)))"
    @info "Python executable used by PyCall: $(pyimport("sys").executable)";

    # Append the code directory to python path
    # (see https://github.com/JuliaPy/PyCall.jl#troubleshooting)
    file_path = @__FILE__;
    pushfirst!(PyVector(pyimport("sys")."path"),
               py"'/'.join($file_path.split('/')[:-1]) + '/../Trajectron-plus-plus/trajectron'");

    # Define python functions for Trajectron bridge
    py"""
    import os
    import time
    import json
    import torch
    import dill
    import random
    import pathlib
    import pandas as pd
    # import evaluation
    import numpy as np
    import visualization as vis
    from easydict import EasyDict
    from model.online.online_trajectron import OnlineTrajectron
    from model.model_registrar import ModelRegistrar
    from environment import Environment, Scene, Node

    # For CrowdNav Baseline
    import configparser
    import crowd_nav
    import gym
    from crowd_nav.policy.policy_factory import policy_factory
    from crowd_sim.envs.utils.robot import Robot
    from crowd_sim.envs.utils.state import ObservableState

    def load_hyperparams(model_dir, conf_file):
        '''
        param model_dir: directory where the learned model resides
        param conf_file: config file name
        '''
        # Load hyperparameters from json
        config_file = os.path.join(model_dir, conf_file)
        if not os.path.exists(config_file):
            raise ValueError('Config json not found!')
        with open(config_file, 'r') as conf_json:
            hyperparams = json.load(conf_json)
        return hyperparams

    def load_online_scene(data_dir, data_name, maximum_history_length, state_def,
                         scene_idx, incl_robot_node, init_timestep=2,
                         ado_id_removed=None):

        '''
        param data_dir: directory to look in for evaluation data
        param data_name: name of the evaluation data
        param maximum_history_length: maximum history length considered in the model
        param state_def: state definition
        param scene_idx: scene_idx (starting from 0)
        param incl_robot_node: whether to include the robot node in the scene
        param init_timestep: init_timestep (starting form 0). You need to have at least acceleration, so you want 2 timesteps of prior data, e.g. [0, 1],
                             so that you can immediately start incremental inference from the 3rd timestep onwards.
        param ado_id_removed: ado_id in string to be removed from online_scene. Note this ado will not be removed from eval_scene
        '''
        log = []
        eval_data_path = os.path.join(data_dir, data_name)
        with open(eval_data_path, 'rb') as f:
            eval_env = dill.load(f, encoding='latin1')
        log.append('Loaded evaluation data from %s' % (eval_data_path,))
        if eval_env.robot_type is None and incl_robot_node:
             eval_env.robot_type = eval_env.NodeType[0]  # TODO: Make more general, allow the user to specify?

             x = [np.nan]
             y = [np.nan]
             vx = [np.nan]
             vy = [np.nan]
             ax = [np.nan]
             ay = [np.nan]
             data_dict = {('position', 'x'): x,
                          ('position', 'y'): y,
                          ('velocity', 'x'): vx,
                          ('velocity', 'y'): vy,
                          ('acceleration', 'x'): ax,
                          ('acceleration', 'y'): ay}
             data_columns = pd.MultiIndex.from_product([['position', 'velocity', 'acceleration'], ['x', 'y']])
             node_data = pd.DataFrame(data_dict, columns=data_columns)
             robot_node = Node(eval_env.robot_type, 'ROBOT', node_data, is_robot=True, first_timestep=init_timestep)
             for scene in eval_env.scenes:
                 #scene.add_robot_from_nodes(eval_env.robot_type)
                 scene.robot = robot_node
                 scene.nodes.append(scene.robot)

        assert 0 <= scene_idx < len(eval_env.scenes), "scene_idx must be between 0 and {}!".format(len(eval_env.scenes) - 1)
        eval_scene = eval_env.scenes[scene_idx]

        online_scene = Scene(timesteps=init_timestep,
                             map=eval_scene.map,
                             dt=eval_scene.dt)
        online_scene.nodes = eval_scene.get_nodes_clipped_at_time(timesteps=np.arange(init_timestep + 1 - maximum_history_length,
                                                                                init_timestep),
                                                                  state=state_def)
        online_scene.robot = eval_scene.robot
        node_to_remove = None
        if ado_id_removed is not None:
            for node in online_scene.nodes:
                if str(node) == ado_id_removed:
                    node_to_remove = node
        if node_to_remove is not None:
            online_scene.nodes.remove(node_to_remove)

        return eval_env, eval_scene, online_scene, log

    def initialize_model(args, model_dir, conf_file):
        '''
        param args: easydict object for parsing arguments
        param model_dir: directory where the learned model resides
        param confi_file: config file name
        '''
        log = []
        # Load hyperparameters from json
        hyperparams = load_hyperparams(model_dir, conf_file)

        if args.seed is not None:
            random.seed(args.seed)
            np.random.seed(args.seed)
            torch.manual_seed(args.seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(args.seed)

        # Add hyperparams from arguments
        hyperparams['dynamic_edges'] = args.dynamic_edges
        hyperparams['edge_state_combine_method'] = args.edge_state_combine_method
        hyperparams['edge_influence_combine_method'] = args.edge_influence_combine_method
        hyperparams['edge_addition_filter'] = args.edge_addition_filter
        hyperparams['edge_removal_filter'] = args.edge_removal_filter
        hyperparams['offline_scene_graph'] = args.offline_scene_graph
        hyperparams['incl_robot_node'] = args.incl_robot_node
        hyperparams['edge_encoding'] = not args.no_edge_encoding
        hyperparams['use_map_encoding'] = args.map_encoding

        iter_num = 100
        model_registrar = ModelRegistrar(model_dir, args.eval_device)
        model_registrar.load_models(iter_num=iter_num)
        log.append('Loaded Trajectron model from %s' % (os.path.join(model_dir, 'model_registrar-{}.pt'.format(iter_num))))
        trajectron = OnlineTrajectron(model_registrar, hyperparams,
                                  args.eval_device)
        return trajectron, hyperparams, log

    def configure_rl_robot(model_dir, env_config_path, policy_config_path,
                           policy_name='sarl'):
        env_config_file = os.path.join(model_dir, os.path.basename(env_config_path))
        policy_config_file = os.path.join(model_dir, os.path.basename(policy_config_path))
        model_weights = os.path.join(model_dir, 'rl_model.pth')
        device = torch.device('cpu')

        # configure policy
        policy = policy_factory[policy_name]()
        policy_config = configparser.RawConfigParser()
        policy_config.read(policy_config_file)
        policy.configure(policy_config)
        if policy.trainable:
            policy.get_model().load_state_dict(torch.load(model_weights))
        phase = 'test'
        policy.set_phase(phase)
        policy.set_device(device)
        if hasattr(policy, 'query_env'):
            policy.query_env = False

        # configure environment
        env_config = configparser.RawConfigParser()
        env_config.read(env_config_file)

        robot = Robot(env_config, 'robot')
        robot.set_policy(policy)

        # additional setup (from crowd_sim.CrowdSim.reset)
        time_step = env_config.getfloat('env', 'time_step')
        robot.time_step = time_step
        robot.policy.time_step = time_step
        if robot.sensor == 'RGB':
            raise NotImplementedError

        return robot
    """
end

# Type Definition
export
    Parameter,
    State,
    WayPoint2D,
    Trajectory2D,
    get_position,
    RobotState,
    get_velocity,
    WorldState
include("type_definition.jl")

# State Transition
export
    transition,
    transition_jacobian,
    transition_control_coeff
    #sample_future_positions
include("state_transition.jl")

# Cost
export
    instant_position_cost,
    instant_control_cost,
    instant_collision_cost,
    terminal_position_cost,
    terminal_collision_cost,
    instant_position_cost_gradient,
    instant_collision_cost_gradient,
    terminal_position_cost_gradient,
    terminal_collision_cost_gradient,
    CostParameter
include("cost.jl")

# Cost (CUDA kernels)
export
    kernel_instant_position_cost!,
    kernel_instant_control_cost!,
    kernel_instant_collision_cost!,
    kernel_terminal_position_cost!,
    kernel_terminal_collision_cost!,
    kernel_instant_position_cost_gradient!,
    kernel_instant_collision_cost_gradient!,
    kernel_terminal_position_cost_gradient!,
    kernel_terminal_collision_cost_gradient!
include("cost_gpu.jl")

# Scene Loader
export
    fetch_ado_positions!,
    convert_nodes_to_str,
    SceneLoader,
    TrajectronSceneParameter,
    TrajectronSceneLoader,
    reduce_to_positions,
    get_trajectory_for_ado,
    SyntheticSceneParameter,
    SyntheticSceneLoader
include("scene_loader.jl")

# Predictor
export
    initialize_scene_graph!,
    sample_future_ado_positions!,
    Predictor,
    TrajectronPredictorParameter,
    TrajectronPredictor,
    OraclePredictorParameter,
    OraclePredictor,
    GaussianPredictorParameter,
    GaussianPredictor
include("predictor.jl")

# Forward-Backward Simulation
export
    get_measurement_schedule,
    get_target_pos_array,
    process_u_arrays,
    simulate_forward,
    sample_ado_positions,
    process_ap_dict,
    compute_costs,
    integrate_costs,
    choose_best_nominal_control,
    compute_cost_gradients,
    sum_cost_gradients,
    simulate_backward,
    simulate,
    # evaluate_risk,
    SimulationParameter,
    SimulationCostResult,
    SimulationCostGradResult,
    SimulationResult
include("forward_backward_simulation.jl")

# Risk Sensitive SAC Controller
export
    control!,
    adjust_old_prediction!,
    schedule_prediction!,
    schedule_control_update!,
    convert_to_schedule,
    get_nominal_u_arrays,
    get_robot_present_and_future,
    get_control_coeffs,
    solve_multi_qcqp,
    get_control_schedule,
    determine_control_time,
    sac_control_update,
    ControlParameter,
    ControlSchedule,
    RSSACController
include("rs_sac_controller.jl")

# Buffered Input Cells and Related Functions (benchmark)
export
    vertex_direction,
    sort_vertices,
    scale_and_shift,
    find_extrema,
    retract_cell,
    get_input_positions,
    compute_buffered_voronoi_cells,
    compute_buffered_input_cell,
    project,
    lq_tracking,
    BICControlParameter,
    BICController,
    bic_control_update,
    bic_parameter_setup
include("bic.jl")

# CrowdNav RL Controller (benchmark)
export
    CrowdNavControlParameter,
    CrowdNavController,
    get_action!,
    crowdnav_control_update!,
    control!,
    schedule_control_update!
include("crowd_nav_controller.jl")

# Evaluation Functions
export
    evaluate,
    EvaluationResult,
    BICEvaluationResult
include("evaluation.jl")

# Helper Functions
export
    get_nominal_trajectory,
    init_condition_setup,
    controller_setup,
    display_log,
    visualize!,
    make_gif,
    fetch_stats_filtered,
    plot_histogram
include("utils.jl")

end # module


text:
# Structs and Functions
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#chap16
#include("ch15module.jl")
include("AllnixThinkJulia.jl")
import .AllnixThinkJulia
ns = AllnixThinkJulia

# functional programming style

# Time
time = ns.MyTime(11, 59, 30)
# Exercise 16-1
ns.printtime(time)
# Exercise 16-2
t1 = ns.MyTime(9, 30, 15)
t2 = ns.MyTime(8, 45, 30)
ns.isafter(t2, t1)

# Pure Functions
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_pure_functions

# prototype and patch

# A pure function
start = ns.MyTime(9, 45, 0)
duration = ns.MyTime(1, 35, 0)
done = ns.addtime(start, duration)

ns.printtime(done)

# Modifiers
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#modifiers

# Function that modifies its parameter

# Exercise 16-3: See ch16module.jl
time = ns.MyTime(9, 45, 0)
ns.increment!(time, 130)
ns.printtime(time)

# Exercise 16-4: See ch16module.jl
time = ns.MyTime(9, 45, 0)
time = ns.increment(time, 130)
ns.printtime(time)

# Prototyping Versus Planning
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#prototyping_versus_planning

# prototype and patch: could be messy
# designed development: more insight into the problem then do it
time = ns.MyTime(9, 45, 0)
totalsecond = ns.timetosecond(time)
time = ns.secondtotime(totalsecond)

# Exercise 16-5: See ch16module.jl


text:
using ImageIO
using Test
using ColorTypes
using FixedPointNumbers
using ImageCore
using Logging
using Random

#logger = ConsoleLogger(stdout, Logging.Debug)
logger = ConsoleLogger(stdout, Logging.Info)
global_logger(logger)

tmpdir = joinpath(@__DIR__,"temp")
@testset "ImageIO.jl" begin
    # Write your own tests here.

    @testset "libtiff" begin
        @test ImageIO.ModTIFF.version() == "LIBTIFF, Version 4.0.10\nCopyright (c) 1988-1996 Sam Leffler\nCopyright (c) 1991-1996 Silicon Graphics, Inc."
    end

    @testset "libpng" begin
        isdir(tmpdir) && rm(tmpdir, recursive = true)
        mkdir(tmpdir)

        img = rand(Bool, 5, 5, 5, 5)
        filepath = joinpath(tmpdir, "5x5x5x5.png")
        @test_throws ErrorException ModPNG.writeimage(filepath, img)

        @testset "Binary Image" begin
            a = rand(Bool, 11, 10)
            filepath = joinpath(tmpdir, "binary1.png")
            ModPNG.writeimage(filepath, a)
            b1 = ModPNG.readimage(filepath)
            @test b1 == convert(Array{Gray{N0f8}}, a)

            a = bitrand(5,5)
            filepath = joinpath(tmpdir, "binary2.png")
            ModPNG.writeimage(filepath, a)
            b2 = ModPNG.readimage(filepath)
            @test b2 == convert(Array{Gray{N0f8}}, a)

            a = colorview(Gray, a)
            filepath = joinpath(tmpdir, "binary3.png")
            ModPNG.writeimage(filepath, a)
            b3 = ModPNG.readimage(filepath)
            @test b3 == convert(Array{Gray{N0f8}}, a)
        end

        @testset "Gray image" begin
            gray = vcat(fill(Gray(1.0), 10, 10), fill(Gray(0.5), 10, 10), fill(Gray(0.0), 10, 10))
            filepath = joinpath(tmpdir, "gray1.png")
            ModPNG.writeimage(filepath, gray)
            g1 = ModPNG.readimage(filepath)
            @test g1 == convert(Array{Gray{N0f8}}, gray)

            gray = rand(Gray{N0f8}, 5, 5)
            filepath = joinpath(tmpdir, "gray2.png")
            ModPNG.writeimage(filepath, gray)
            g2 = ModPNG.readimage(filepath)
            @test g2 == gray
        end

        @testset "Color - RGB" begin
            #rgb8 = rand(RGB{N0f8}, 10, 5)
            rgb8 = reshape(range(RGB{N0f8}(1,0,0),RGB{N0f8}(0,1,1), length=10*5), 10, 5)
            filepath = joinpath(tmpdir, "rgb_n0f8.png")
            ModPNG.writeimage(filepath, rgb8)
            r1 = ModPNG.readimage(filepath)
            @test r1 == rgb8

            #rgb16 = rand(RGB{N0f16}, 10, 5)
            rgb16 = reshape(range(RGB{N0f16}(1,0,0),RGB{N0f16}(0,1,1), length=10*5), 10, 5)
            filepath = joinpath(tmpdir, "rgb_n0f16.png")
            ModPNG.writeimage(filepath, rgb16)
            r2 = ModPNG.readimage(filepath)
            ModPNG.writeimage(joinpath(tmpdir, "rgb_n0f16_resave.png"), r2)
            @test r2 == rgb16
        end

        @testset "Alpha" begin
            # RGBA
            r = RGBA(1.0,0.0,0.0, 0.2)
            g = RGBA(0.0,1.0,0.0, 0.8)
            b = RGBA(0.0,0.0,1.0, 1.0)
            rgba1 = vcat(fill(r, 50,100), fill(g, 50,100), fill(b, 50,100))
            filepath = joinpath(tmpdir, "rgba1.png")
            ModPNG.writeimage(filepath, rgba1)
            r1 = ModPNG.readimage(filepath)
            @test r1 == rgba1

            # GrayA
            r = GrayA(1.0, 0.25)
            g = GrayA(0.5, 0.5)
            b = GrayA(0.0, 0.75)
            graya = vcat(fill(r, 50,100), fill(g, 50,100), fill(b, 50,100))
            filepath = joinpath(tmpdir, "graya1.png")
            ModPNG.writeimage(filepath, graya)
            g1 = ModPNG.readimage(filepath)
            @test g1 == convert(Array{GrayA{N0f8}}, graya)
        end
        # TODO implement palette
    end

    @testset "libjpeg" begin
        @test unsafe_string(ImageIO.ModJPEG.tjGetErrorStr()) == "No error"
    end

end

# try
#     rm(tmpdir, recursive = true)
# catch
#     @error "Unable to remove temp directory at: $(tmpdir)"
# end


text:
using ConservationLawsParticles
using SpecialFunctions, RecursiveArrayTools, DifferentialEquations, Plots

# model
V(t, x) = 0.0 # 2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = 0.0 # x^2/2 - abs(x)
W′(t, x) = 0.0 # x - sign(x)
mob(ρ) = 0.0 # max((ρ-1)^2 * (2ρ+1), 0)
smodel = DiffusiveSampledModel((V,), ((W′,),), (mob,), (Diffusion(1),))
imodel = DiffusiveIntegratedModel((V,), ((W,),), (mob,), (Diffusion(1),))

# initial condition
n = 100
x0 = ArrayPartition(vcat(
    range(-2, -1; length=n+1),
    range(0, 2; length=n+1)
    ))

boxsol(t,x;a,b) = (erf((x-a)/2sqrt(t)) - erf((x-b)/2sqrt(t))) / (2*(b-a))
refsol(t, x) = boxsol(t,x; a=-2,b=-1) * n/(2n+1) +
    boxsol(t,x; a=0,b=2) * n/(2n+1) +
    boxsol(t,x; a=-1,b=0) * 1/(2n-1)

# time span
tspan = (0., 2.0)

# ODE system for the particles
sprob = ODEProblem(velocities_diff!, x0, tspan, smodel)
# iprob = ODEProblem(velocities_diff!, x0, tspan, imodel)

abstol = reltol = 2e-8

@time ssol = solve(sprob, BS5(); abstol=abstol, reltol=reltol);
# @time isol = solve(iprob, BS5(); abstol=abstol, reltol=reltol);

# plot the particle trajectories
plot(legend=false, title="Trajectories", xlabel="time", ylabel="position")
plot!(ssol, vars=1:5:(2n+2); color=:blue)
savefig("diffusion2.png")

# plot an animation of the density
anim = @animate for t in range(tspan...; length=101)
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-7,7), ylims=(0,1/sqrt(pi)))
    plot!(x -> refsol(t, x); color=:blue, label="analytic")
    densityplot!(ssol(t); color=:red, label="particles")
end
gif(anim, "diffusion2.gif")


text:
function findRequiredIdxs(A::LagrangeInterpolation, t)
  idxs = sortperm(A.t,by=x->abs(t-x))
  idxs[1:(A.n+1)]
end

function spline_coefficients(n, d, k, u::Number)
  N = zeros(n)
  if u == k[1]
    N[1] = one(u)
  elseif u == k[end]
    N[end] = one(u)
  else
    i = findfirst(x->x>u,k) - 1
    N[i] = one(u)
    for deg = 1:d
      N[i-deg] = (k[i+1]-u)/(k[i+1]-k[i-deg+1]) * N[i-deg+1]
      for j = (i-deg+1):(i-1)
        N[j] = (u-k[j])/(k[j+deg]-k[j]) * N[j] + (k[j+deg+1]-u)/(k[j+deg+1]-k[j+1]) * N[j+1]
      end
      N[i] = (u-k[i])/(k[i+deg]-k[i]) * N[i]
    end
  end
  N
end

function spline_coefficients(n, d, k, u::AbstractVector)
  N = zeros(eltype(u),n,n)
  for i = 1:n
    N[i,:] .= spline_coefficients(n,d,k,u[i])
  end
  N
end

function weibull_fun(x,p)
  ans = copy(x)
  for i = 1:length(x)
    if x[i] >= 0
      ans[i] = p[1] * (1 - exp(-1 * (x[i]/p[2])^p[3]))
    else
      ans[i] = 0.0
    end
  end
  ans
end

# helper functions to get t_max and c_max (OneCompartmentPK)
function OneCompartmentPK_tmax(A)
  p = A.c_f.param
  2.303 * log(10.0,p[1]/p[4]) / (p[1] - p[4])
end

function OneCompartmentPK_cmax(A)
  p = A.c_f.param
  t = OneCompartmentPK_tmax(A)
  p[1] * p[2] * p[3] * (exp(-p[4]*t) - exp(-p[1]*t)) / (p[5] * (p[1] - p[4]))
end

# helper function for data manipulation
function munge_data(u::AbstractVector, t::AbstractVector)
  Tu = Base.nonmissingtype(eltype(u))
  Tt = Base.nonmissingtype(eltype(t))
  newu = Tu[]
  newt = Tt[]
  @assert length(t) == length(u)
  @inbounds for i in eachindex(t)
    ui = u[i]
    ti = t[i]
    if !ismissing(ui) && !ismissing(ti)
      push!(newu, ui)
      push!(newt, ti)
    end
  end
  return newu, newt
end

function munge_data(U::StridedMatrix, t::AbstractVector)
  TU = Base.nonmissingtype(eltype(U))
  Tt = Base.nonmissingtype(eltype(t))
  newUs = [TU[] for i in 1:size(U, 1)]
  newt  = Tt[]
  @assert length(t) == size(U,2)
  @inbounds for j in eachindex(t)
    tj = t[j]
    if ismissing(tj) || any(ismissing, view(U, :, j))
      continue
    end

    push!(newt, tj)
    for i in 1:length(newUs)
      push!(newUs[i], U[i,j])
    end
  end

  return vcat(adjoint.(newUs)...), newt
end


text:
"""
    upside_potential_ratio(returns, threshold; method=:partial)

The Upside Potential Ratio is a risk-adjusted performance measure similarly to the Sharpe Ratio and the Sortino Ratio. This ratio considers only upside returns (above `threshold`) in the numerator, and only downside returns (below `threshold`) in the denominator (see `downside_deviation`).

# Arguments
- `returns`:     Vector of asset returns.
- `threshold`:   Scalar value or vector denoting the threshold returns.
- `method`:      One of `:full` (default) or `:partial`. Indicates whether to use the number of all returns (`:full`), or only the number of returns above the threshold (`:partial`) in the denominator.

# Sources
- Plantinga, A., van der Meer, R. and Sortino, F. (2001). The Impact of Downside Risk on Risk-Adjusted Performance of Mutual Funds in the Euronext Markets.
"""
function upside_potential_ratio(returns, threshold; method::Symbol=:partial)
    if method == :full
        n = length(returns)
    elseif method == :partial
        n = count(returns .> threshold)
    else
        throw(ArgumentError("Passed method parameter '$(method)' is invalid, must be one of :full, :partial."))
    end

    dd = downside_deviation(returns, threshold; method=method)
    excess = returns .- threshold
    (sum(map(x -> max(0.0, x), excess)) / n) / dd
end


text:
export Conv
using Knet.Ops21: conv, zeropad, symmetric
import CUDA

using CUDA.CUDNN:
    cudnnConvolutionForward,
    cudnnConvolutionForward!,
    cudnnConvolutionBackwardFilter,
    cudnnConvolutionBackwardData,
    cudnnGetConvolutionNdForwardOutputDim,
    cudnnSetConvolutionMathType,
    cudnnSetConvolutionReorderType,
    cudnnSetConvolutionGroupCount,
    cudnnFindConvolutionForwardAlgorithmEx,
        cudnnConvolutionFwdAlgoPerf_t,
    cudnnFindConvolutionBackwardFilterAlgorithmEx,
        cudnnConvolutionBwdFilterAlgoPerf_t,
    cudnnFindConvolutionBackwardDataAlgorithmEx,
        cudnnConvolutionBwdDataAlgoPerf_t,
    cudnnConvolutionDescriptor,
        cudnnConvolutionDescriptor_t,
        cudnnCreateConvolutionDescriptor,
        cudnnSetConvolutionNdDescriptor,
        cudnnDestroyConvolutionDescriptor,
    cudnnConvolutionMode_t,
        CUDNN_CONVOLUTION,       # 0
        CUDNN_CROSS_CORRELATION, # 1
    cudnnActivationMode_t,
        CUDNN_ACTIVATION_SIGMOID,      # 0
        CUDNN_ACTIVATION_RELU,         # 1
        CUDNN_ACTIVATION_TANH,         # 2
        CUDNN_ACTIVATION_CLIPPED_RELU, # 3
        CUDNN_ACTIVATION_ELU,          # 4
        CUDNN_ACTIVATION_IDENTITY,     # 5
    cudnnNanPropagation_t,
        CUDNN_NOT_PROPAGATE_NAN, # 0
        CUDNN_PROPAGATE_NAN,     # 1
    cudnnMathType_t,
        CUDNN_DEFAULT_MATH,                    # 0
        CUDNN_TENSOR_OP_MATH,                  # 1
        CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION, # 2
        CUDNN_FMA_MATH,                        # 3
    cudnnReorderType_t,
        CUDNN_DEFAULT_REORDER, # 0
        CUDNN_NO_REORDER,      # 1
    cudnnConvolutionFwdAlgo_t,
        CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM,         # 0
        CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, # 1
        CUDNN_CONVOLUTION_FWD_ALGO_GEMM,                  # 2
        CUDNN_CONVOLUTION_FWD_ALGO_DIRECT,                # 3
        CUDNN_CONVOLUTION_FWD_ALGO_FFT,                   # 4
        CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING,            # 5
        CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD,              # 6
        CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED,     # 7
        CUDNN_CONVOLUTION_FWD_ALGO_COUNT,                 # 8
    cudnnConvolutionBwdFilterAlgo_t,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0,                 # 0, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1,                 # 1,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT,               # 2,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_3,                 # 3, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD,          # 4, /* not implemented */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED, # 5,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING,        # 6,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT,             # 7
    cudnnConvolutionBwdDataAlgo_t,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_0,                 # 0, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_1,                 # 1,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT,               # 2,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING,        # 3,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD,          # 4,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED, # 5,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT,             # 6
    cudnnTensorFormat_t,
        CUDNN_TENSOR_NCHW,        # 0, /* row major (wStride = 1, hStride = w) */
        CUDNN_TENSOR_NHWC,        # 1, /* feature maps interleaved ( cStride = 1 )*/
        CUDNN_TENSOR_NCHW_VECT_C, # 2, /* each image point is vector of element of C, vector length in data type */
    cudnnDataType,
    convdims,
    math_mode,
    handle

"""
    c = Conv(wdims...; winit, binit, kwargs...)
    c = Conv(w; bias, kwargs...)
    y = c(x, [z=nothing])

Return a layer that can perform convolution and optionally scaling, bias/residual addition,
normalization and/or activation. The constructor can take the convolution weight tensor `w`
or its dimensions `wdims`. If `wdims` is used a weight tensor will be initialized using the
`winit` function which defaults to uniform random in `±√(6/(fanin+fanout))` and a bias will
be initialized using the `binit` function if provided. No bias will be used unless either
`bias` or `binit` is specified.  The layer can be called with a single input `x`, or two
inputs `x,z` where `z` has the same size as the output `y`. The computed result is:

    y = activation(normalization(alpha * conv(w,x) + beta * z .+ bias))

For tensor sizes and keyword arguments with their defaults see `@doc Knet.Ops21.conv`.
"""
mutable struct Conv
    w  # No type here, we do not want to restrict the type of array
    wdims::Dims
    winit

    bias  # Can be nothing
    bdims::Dims
    binit

    # y; dw; dbias; dx; dz; These are overwritten if multiple calls to same layer.

    padding::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}},Tuple{Vararg{NTuple{2,<:Integer}}}}
    stride::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}}
    dilation::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}}
    groups::Integer
    crosscorrelation::Bool
    channelmajor::Bool
    activation
    normalization
    convDesc::Union{Nothing,cudnnConvolutionDescriptor}

    alpha::Real
    beta::Real
end

# If the user provides size, we can't immediately initialize w because we do not know the
# array type until we see the first input
function Conv(
    wdims::Integer...;
    w = nothing,
    bias = nothing,
    crosscorrelation::Bool = true,
    channelmajor::Bool = false, # CUDNN_TENSOR_NHWC if true.
    winit = 𝑼(√(6/(fanin(wdims; channelmajor)+fanout(wdims; channelmajor)))),
    binit = nothing,
    activation = nothing,
    normalization = nothing,
    padding::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}},Tuple{Vararg{NTuple{2,<:Integer}}}} = 0,
    stride::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}} = 1,
    dilation::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}} = 1,
    groups::Integer = 1,
    alpha::Real = 1,
    beta::Real = 0,
)
    wdims = Int.(wdims)
    dw = Ref{Any}(nothing)
    dbias = Ref{Any}(nothing)
    ndims = length(wdims)
    cdim = channelmajor ? 1 : ndims-1
    bdims = ntuple(i->(i===cdim ? wdims[end] : 1), ndims)
    convDesc = nothing          # To be initialized at first call
    Conv(w, wdims, winit, bias, bdims, binit, padding, stride, dilation, groups, crosscorrelation, channelmajor, activation, normalization, convDesc, alpha, beta)
end

Conv(w; o...) = Conv(size(w)...; w, o...)

# Some of the initialization can only be done at the first call when we know the type of x
# TODO: are we sure about this? should we use atype() instead?
function initconv(c::Conv, x, z, padding)
    issimilar(u,v,s=size(v))=(typeof(value(u)) === typeof(value(v)) && size(u) === s)
    if c.w === nothing
        c.w = Param(oftype(value(x), c.winit(c.wdims...)))
    end
    @assert issimilar(c.w, x, c.wdims)  "w=$(summary(c.w)) x=$(summary(x)) wdims=$(c.wdims)"
    @assert (c.channelmajor ? size(x,1) === size(c.w,1)*c.groups : size(x)[end-1] === size(c.w)[end-1]*c.groups)  "x=$(size(x)) w=$(size(c.w)) grp=$(c.groups)"
    if c.bias === nothing && c.binit !== nothing
        c.bias = Param(oftype(value(x), c.binit(c.bdims...)))
    end
    if c.bias === nothing && ((c.activation === relu && c.normalization ∈ (nothing, identity)) || (z !== nothing && c.beta != 0))
        # will call cudnnConvolutionBiasActivationForward, must have bias
        c.bias = fill!(similar(c.w, c.bdims), 0)
    end
    @assert c.bias === nothing || issimilar(c.bias, c.w, c.bdims)
    if CUDA.functional() && (c.convDesc === nothing || c.convDesc.ptr === C_NULL)
        mode = c.crosscorrelation ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION
        format = c.channelmajor ? CUDNN_TENSOR_NHWC : CUDNN_TENSOR_NCHW
        reorderType, mathType = CUDNN_DEFAULT_REORDER, math_mode()
        c.convDesc = cudnnConvolutionDescriptor(convdims(padding,size(x),format), convdims(c.stride,size(x),format), convdims(c.dilation,size(x),format), mode, cudnnDataType(eltype(x)), mathType, reorderType, Cint(c.groups))
    end
end

function (c::Conv)(x, z=nothing)
    if symmetric(c.padding)
        padding = c.padding
    else
        x = zeropad(x, c.padding)
        padding = 0
    end
    initconv(c, x, z, padding)
    conv(c.w, x; z, c.bias, c.activation, c.normalization, c.alpha, c.beta, c.channelmajor,
         c.convDesc, c.crosscorrelation, c.dilation, c.groups, padding, c.stride)
end


text:
export SimpleProblem,
        solve,
        value_of_information

"""
    struct SimpleProblem
        bn::BayesianNetwork
        chance_vars::Vector{Variable}
        decision_vars::Vector{Variable}
        utility_vars::Vector{Variable}
        utilities::Dict{Symbol, Vector{Float64}}
    end

Data structure representing a decision network.

# Fields

    bn::BayesianNetwork

    chance_vars::Vector{Variable}

    decision_vars::Vector{Variable}

    utility_vars::Vector{Variable}

    utilities::Dict{Symbol, Vector{Float64}}

"""
struct SimpleProblem
    bn::BayesianNetwork
    chance_vars::Vector{Variable}
    decision_vars::Vector{Variable}
    utility_vars::Vector{Variable}
    utilities::Dict{Symbol, Vector{Float64}}
end

"""
    function solve(𝒫::SimpleProblem, evidence, M)

Given a decision network `𝒫` and observed variable assignment `evidence`, determine the optimal action variable assignment and associated expected utility, using inference method M.

# Arguments

    𝒫::SimpleProblem
        Decision network
    evidence
        Vector of NamedTuples mapping Symbol name of Variable to assignment of observed values of condition vars.
    M
        Method to use for inference over utility vars given evidence
"""
function solve(𝒫::SimpleProblem, evidence, M)
    query = [var.name for var in 𝒫.utility_vars]
    U(a) = sum(𝒫.utilities[uname][a[uname]] for uname in query)
    best = (a=nothing, u=-Inf)
    for assignment in assignments(𝒫.decision_vars)
        evidence = merge(evidence, assignment)
        φ = infer(M, 𝒫.bn, query, evidence)
        u = sum(p*U(a) for (a, p) in φ.table)
        if u > best.u
            best = (a=assignment, u=u)
        end
    end
    return best
end

"""
    function value_of_information(𝒫, query, evidence, M)

Determine the Value of Information (VOI) of `query` variables, given `evidence` assignment of (specific) values to already observed variables, using inference method `M`.
"""
function value_of_information(𝒫, query, evidence, M)

    φ = infer(M, 𝒫.bn, query, evidence)
    voi = -solve(𝒫, evidence, M).u
    query_vars = filter(v->v.name ∈ query, 𝒫.chance_vars)

    for o′ in assignments(query_vars)
        oo′ = merge(evidence, o′)
        p = φ.table[o′]
        voi += p*solve(𝒫, oo′, M).u
    end

    return voi
end


text:
using LinearAlgebra
using SparseArrays
using Plots

"
This code runs a convergence test for finite difference methods applied to Laplace's equation.
The error in different norms is plotted, along with O(h) and O(h^2) slopes for reference.
"

f(x) = pi^2*sin(pi*x)
uexact(x) = sin(pi*x)

mvec = 2 .^ (2:9)
hvec = zeros(length(mvec))
err_inf = zeros(length(mvec))
err_1 = zeros(length(mvec))
err_2 = zeros(length(mvec))
for (i,m) in enumerate(mvec)
    x = LinRange(-1,1,m+2) # add 2 endpoints
    xint = x[2:end-1]
    h = x[2]-x[1]
    A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))
    b = f.(xint)
    u = A\b

    hvec[i] = h
    err_inf[i] = maximum(@. abs(uexact(xint) - u))
    err_1[i] = norm((@. abs(uexact(xint) - u)),1) # 1-norm
    err_2[i] = norm(@. abs(uexact(xint) - u)) # 2-norm
end

plot(hvec,err_inf,marker=:dot,label="Max error")
plot!(hvec,err_1,marker=:dot,label="1-norm error")
plot!(hvec,err_2,marker=:dot,label="2-norm error")
plot!(hvec,hvec*1.5,linestyle=:dash,label="O(h)")
plot!(hvec,hvec.^2,linestyle=:dash,label="O(h^2)")
plot!(xaxis=:log,yaxis=:log,legend=:bottomright)


text:
export HH3DHyperSingularFDBIO

abstract type Helmholtz3DOp <: MaxwellOperator3D end
abstract type Helmholtz3DOpReg <: MaxwellOperator3DReg end

"""
```
∫_Γ dx ∫_Γ dy \left(α G g(x) n_x ⋅ n_y f(y) + β G \mbox{curl} g(x) ⋅ \mbox{curl} f(y) \right)
```

with ``G(x,y) = \frac{e^{-γ |x-y|}}{4 π |x-y|}``
"""
struct HH3DHyperSingularFDBIO{T,K} <: Helmholtz3DOp
    "coefficient of the weakly singular term"
    alpha::T
    "coefficient of the hyper singular term"
    beta::T
    "`im*κ` with `κ` the wave number"
    gamma::K
end

HH3DHyperSingularFDBIO(gamma) = HH3DHyperSingularFDBIO(gamma^2, one(gamma), gamma)
scalartype(op::HH3DHyperSingularFDBIO) = promote_type(typeof(op.alpha), typeof(op.beta), typeof(op.gamma))

"""
```math
a(u,v) = α ∬_{Γ×Γ} u(x) G_{γ}(|x-y|) v(y)
```

with ``G_{γ}(r) = \frac{e^{-γr}}{4πr}``.
"""
struct HH3DSingleLayerFDBIO{T,K} <: Helmholtz3DOp
    alpha::T
    gamma::K
end

struct HH3DSingleLayerReg{T,K} <: Helmholtz3DOpReg
    alpha::T
    gamma::K
end

struct HH3DSingleLayerSng{T,K} <: Helmholtz3DOp
    alpha::T
    gamma::K
end

struct HH3DDoubleLayer{T,K} <: Helmholtz3DOp
    alpha::T
    gamma::K
end

struct HH3DDoubleLayerTransposed{T,K} <: Helmholtz3DOp
    alpha::T
    gamma::K
end

function quaddata(op::Helmholtz3DOp, test_refspace::LagrangeRefSpace,
        trial_refspace::LagrangeRefSpace, test_elements, trial_elements)

    test_eval(x)  = test_refspace(x,  Val{:withcurl})
    trial_eval(x) = trial_refspace(x, Val{:withcurl})

    # The combinations of rules (6,7) and (5,7 are) BAAAADDDD
    # they result in many near singularity evaluations with any
    # resemblence of accuracy going down the drain! Simply don't!
    # (same for (5,7) btw...).
    # test_qp = quadpoints(test_eval,  test_elements,  (6,))
    # bssi_qp = quadpoints(trial_eval, trial_elements, (7,))

    test_qp = quadpoints(test_eval,  test_elements,  (4,))
    bsis_qp = quadpoints(trial_eval, trial_elements, (7,))

    return test_qp, bsis_qp
end

function quadrule(op::HH3DSingleLayerFDBIO, test_refspace, trial_refspace, i,
        test_element, j, trial_element, quadrature_data)

    tol, hits = sqrt(eps(eltype(eltype(test_element.vertices)))), 0
    for t in test_element.vertices
        for s in trial_element.vertices
            norm(t-s) < tol && (hits +=1; break)
    end end

    test_quadpoints  = quadrature_data[1]
    trial_quadpoints = quadrature_data[2]

    hits != 0 && return WiltonSEStrategy(
        test_quadpoints[1,i],
        DoubleQuadStrategy(
            test_quadpoints[1,i],
            trial_quadpoints[1,j]))

    return DoubleQuadStrategy(
        quadrature_data[1][1,i],
        quadrature_data[2][1,j])
end

function quadrule(op::Helmholtz3DOp, test_refspace, trial_refspace, i,
        test_element, j, trial_element, quadrature_data)

    test_quadpoints  = quadrature_data[1]
    trial_quadpoints = quadrature_data[2]

    return DoubleQuadStrategy(
        quadrature_data[1][1,i],
        quadrature_data[2][1,j])
end

function integrand(op::HH3DHyperSingularFDBIO,
        kernel, test_values, test_element, trial_values, trial_element)

    α = op.alpha
    β = op.beta

    G = kernel.green

    g, curlg = test_values
    f, curlf = trial_values

    nx = normal(test_element)
    ny = normal(trial_element)

    α*dot(nx,ny)*g*f*G + β*dot(curlg,curlf)*G
end

HH3DSingleLayerFDBIO(gamma) = HH3DSingleLayerFDBIO(one(gamma), gamma)

regularpart(op::HH3DSingleLayerFDBIO) = HH3DSingleLayerReg(op.alpha, op.gamma)
singularpart(op::HH3DSingleLayerFDBIO) = HH3DSingleLayerSng(op.alpha, op.gamma)

function integrand(op::Union{HH3DSingleLayerFDBIO,HH3DSingleLayerReg},
        kernel, test_values, test_element, trial_values, trial_element)

    α = op.alpha
    G = kernel.green
    g, curlg = test_values
    f, curlf = trial_values

    α*g*f*G
end

function innerintegrals!(op::HH3DSingleLayerSng,
        test_neighborhood, test_refspace, trial_refspace, test_elements,
        trial_element, zlocal, quadrature_rule::WiltonSEStrategy, dx)

    γ = op.gamma
    α = op.alpha

    s1, s2, s3 = trial_element.vertices

    x = cartesian(test_neighborhood)
    n = normalize((s1-s3)×(s2-s3))
    ρ = x - dot(x - s1, n) * n

    scal, vec = WiltonInts84.wiltonints(s1, s2, s3, x, Val{1})
    ∫G = (scal[2] - γ*scal[3] + 0.5*γ^2*scal[4]) / (4π)

    zlocal[1,1] += α * ∫G * dx
    return nothing
end

function integrand(biop::HH3DDoubleLayer,
        kernel, fp, mp, fq, mq)

    nq = normal(mq)
    fp[1] * dot(nq, -kernel.gradgreen) * fq[1]
end

function integrand(biop::HH3DDoubleLayerTransposed,
        kernel, fp, mp, fq, mq)

    np = normal(mp)
    fp[1] * dot(np, kernel.gradgreen) * fq[1]
end

module Helmholtz3D
    using ..BEAST

    singlelayer(;
        gamma=error("propagation constant is a required argument"),
        alpha=one(gamma)) = BEAST.HH3DSingleLayerFDBIO(alpha,gamma)

    hypersingular(;
        gamma=error("propagation constant is a required argument"),
        alpha=gamma^2,
        beta=one(gamma)) = BEAST.HH3DHyperSingularFDBIO(alpha, beta, gamma)

    planewave(;
        direction=error("direction is a required argument"),
        wavenumber=error("wavenumber is a required arguement"),
        amplitude=one(eltype(direction))) = BEAST.HH3DPlaneWave(direction, wavenumber, amplitude)

    doublelayer(;gamma=error("gamma missing"), alpha=one(gamma)) =
        BEAST.HH3DDoubleLayer(alpha, gamma)

    doublelayer_transposed(;gamma=error("gamma missing"), alpha=one(gamma)) =
        BEAST.HH3DDoubleLayerTransposed(alpha, gamma)
end

export Helmholtz3D


text:
# Programs testing branch behavior.

# This function tests branching forward when predicting that the branch will
# be taken. (Note - branch is always taken)
#
# Expect dmem addresses 0 and 5 to be written to, but not 1, 2, 3, and 4
@asap4asm function branch_forward_taken_1()
    MOVI(dmem[0], 1)
    BRL(exit)
    MOVI(dmem[1], 2)
    MOVI(dmem[2], 3)
    MOVI(dmem[3], 4)
    MOVI(dmem[4], 5)
    @label exit
    MOVI(dmem[5], 10)
end

# This function tests branching forward when predicting that the branch will
# be NOT taken. (Note - branch is always taken)
#
# This will force the pipeline to deal with the branch mispredict penalty.
#
# Expect dmem addresses 0 and 5 to be written to, but not 1, 2, 3, and 4
@asap4asm function branch_forward_untaken_1()
    MOVI(dmem[0], 1)
    BR(exit)
    MOVI(dmem[1], 2)
    MOVI(dmem[2], 3)
    MOVI(dmem[3], 4)
    MOVI(dmem[4], 5)
    @label exit
    MOVI(dmem[5], 10)
end

# A couple degenerate cases with no trailing instructions
@asap4asm function branch_forward_taken_2()
    MOVI(dmem[0], 1)
    BR(exit)
    @label exit
    MOVI(dmem[1], 1)
end
# A couple degenerate cases with no trailing instructions
@asap4asm function branch_forward_untaken_2()
    MOVI(dmem[0], 1)
    BRL(exit)
    @label exit
    MOVI(dmem[1], 1)
end

# -- Branch Chain -- #
# Try mixing forward and backward branches together.
# All branches are still taken, thouch with a mix of predicted taken and
# predicted not-taken.
@asap4asm function branch_stress_1()
    # Should only set dmem entries 0 to 4.
    MOVI(dmem[0], 1)
    @label label1
        BR(label4)
        MOVI(dmem[50], 1)
        MOVI(dmem[51], 1)
        MOVI(dmem[52], 1)

    @label label2
        MOVI(dmem[3], 1)
        BRL(exit)
        MOVI(dmem[59], 1)
        MOVI(dmem[60], 1)
        MOVI(dmem[61], 1)

    @label label3
        MOVI(dmem[2], 1)
        BR(label2)
        MOVI(dmem[56], 1)
        MOVI(dmem[57], 1)
        MOVI(dmem[58], 1)

    @label label4
        MOVI(dmem[1], 1)
        BRL(label3)
        MOVI(dmem[53], 1)
        MOVI(dmem[54], 1)
        MOVI(dmem[55], 1)

    @label exit
        MOVI(dmem[4], 1)

end

# --- Test using the Jump and Return feature --- #

# This should set the return_address register. This is just mainly a test to
# make sure that happens.
@asap4asm function test_jump_1()
    NOP()
    # The address of the branch instruction is address 2.
    #
    # Thus, the return address saved should be address 3.
    BR(exit, j)
    @label exit
    NOP()
    NOP()
    NOP()
    NOP()
end

# When branching from a RPT block, when the branch is at the end of the
# RPT, must make sure that the return address marks the BEGINNING of the RPT
# loop.
@asap4asm function test_jump_2()

    # Note, RPT is instruction 1, so the start of the RPT block is instruction
    # 2, which should be the return_address saved after the jump.
    RPT(10, nop3)
        NOP()
        NOP()
        BR(exit, j)
    END_RPT()
    # Some dummy move instructions that should be cleaned up by the
    # misprediction handling logic
    MOVI(dmem[1], 1)
    MOVI(dmem[2], 1)
    MOVI(dmem[3], 1)
    @label exit
    MOVI(dmem[0], 1)
end

# Test returning from a Jump
@asap4asm function test_return_1()
    @label main
        BRL(test_label, j)
        MOVI(dmem[1], 1, nop3)
        BRL(exit, nop2)

    @label test_label
        MOVI(dmem[0], 1)
        BRL(back)
        MOVI(dmem[5], 1)
        MOVI(dmem[6], 1)
        MOVI(dmem[7], 1)

    @label exit
        MOVI(dmem[2], 1)
end


text:
# Three DOF Vehicle Model
include("F_YR.jl")
include("F_YF.jl")
include("FZ_RL.jl")
include("FZ_RR.jl")
include("FZ_FL.jl")
include("FZ_FR.jl")
include("Ax_max.jl")
include("Ax_min.jl")

"""
ThreeDOFv2_expr(n) & sol, U = ThreeDOFv1(n,x,u)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 10/20/2017, Last Modified: 5/31/2017 \n
--------------------------------------------------------------------------------------\n
# this vehicle model is controlled using steering angle and speed
"""
function ThreeDOFv1{T<:Any}(n,x::Array{T,2},u::Array{T,2})
  if n.s.integrationMethod==:tm; L=size(x)[1]; else; L=size(x)[1]-1; end
  dx = Array{Any}(L,n.numStates)
  v = x[:,3]; r = x[:,4]; psi = x[:,5];

  # parameters
  ax = zeros(length(psi)); sa = u[:,1];
  pa=params[1]; UX=params[2];
  ux=getvalue(UX)*ones(L,1);           # assume UX is a constant
  @unpack_Vpara n.ocp.params[1]            # vehicle parameters

  # nonlinear tire model
  FYF=@NLexpression(n.ocp.mdl, [i = 1:L], (PD2*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*(FzF0 - (ax[i] - v[i]*r[i])*KZX))*sin(PC1*atan((((PK1*sin(2*atan(PK2*(FzF0 - (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX))^2 + EP^2)^(0.5))*0.001)+EP))*((atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i]) + PH2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PH1)) - ((PE2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PE1)*(1 - PE3)*(((atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i]) + PH2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PH1))/((((atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i]) + PH2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PH1)^2 + EP^2)^(0.5)))*((((PK1*sin(2*atan(PK2*(FzF0 - (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX))^2 + EP^2)^(0.5))*0.001)+EP))*((atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i]) + PH2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PH1)) - atan((((PK1*sin(2*atan(PK2*(FzF0 - (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzF0 - (ax[i] - v[i]*r[i])*KZX))^2 + EP^2)^(0.5))*0.001)+EP))*((atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i]) + PH2*(FzF0 - (ax[i] - v[i]*r[i])*KZX) + PH1)))))) + (PV2*(FzF0 - (ax[i] - v[i]*r[i])*KZX)^2 + PV1*(FzF0 - (ax[i] - v[i]*r[i])*KZX)));
  FYR=@NLexpression(n.ocp.mdl, [i = 1:L], (PD2*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*(FzR0 + (ax[i] - v[i]*r[i])*KZX))*sin(PC1*atan((((PK1*sin(2*atan(PK2*(FzR0 + (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX))^2+EP^2)^(0.5))*0.001)+EP))*((atan((v[i] - lb*r[i])/(ux[i]+EP))) + PH2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PH1)) - ((PE2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PE1)*(1 - PE3*(((atan((v[i] - lb*r[i])/(ux[i]+EP))) + PH2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PH1))/((((atan((v[i] - lb*r[i])/(ux[i]+EP))) + PH2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PH1)^2 + EP^2)^(0.5))))*((((PK1*sin(2*atan(PK2*(FzR0 + (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX))^2+EP^2)^(0.5))*0.001)+EP))*((atan((v[i] - lb*r[i])/(ux[i]+EP))) + PH2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PH1)) - atan((((PK1*sin(2*atan(PK2*(FzR0 + (ax[i] - v[i]*r[i])*KZX))))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)) + ((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)))/(((PD2*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PD1*PC1*(FzR0 + (ax[i] - v[i]*r[i])*KZX))^2+EP^2)^(0.5))*0.001)+EP))*((atan((v[i] - lb*r[i])/(ux[i]+EP))) + PH2*(FzR0 + (ax[i] - v[i]*r[i])*KZX) + PH1)))))) + (PV2*(FzR0 + (ax[i] - v[i]*r[i])*KZX)^2 + PV1*(FzR0 + (ax[i] - v[i]*r[i])*KZX)));

  # vertical tire load
  FZ_rl_con=@NLconstraint(n.ocp.mdl, [i=1:L], 0 <=  0.5*(FzR0 + KZX*(ax[i] - v[i]*r[i])) - KZYR*((FYF[i] + FYR[i])/m) - Fz_min)
  newConstraint!(n,FZ_rl_con,:FZ_rl_con);
  FZ_rr_con=@NLconstraint(n.ocp.mdl, [i=1:L], 0 <=  0.5*(FzR0 + KZX*(ax[i] - v[i]*r[i])) + KZYR*((FYF[i] + FYR[i])/m) - Fz_min)
  newConstraint!(n,FZ_rr_con,:FZ_rr_con);

  # linear tire and for now this also constrains the nonlinear tire model
  Fyf_con=@NLconstraint(n.ocp.mdl, [i=1:L], Fyf_min <=  (atan((v[i] + la*r[i])/(ux[i]+EP)) - sa[i])*Caf <= Fyf_max);
  newConstraint!(n,Fyf_con,:Fyf_con);
  Fyr_con=@NLconstraint(n.ocp.mdl, [i=1:L], Fyf_min <=   atan((v[i] - lb*r[i])/(ux[i]+EP))*Car <= Fyf_max);
  newConstraint!(n,Fyr_con,:Fyr_con);

  dx[:,1] = @NLexpression(n.ocp.mdl, [i=1:L], ux[i]*cos(psi[i]) - (v[i] + la*r[i])*sin(psi[i]));    # X position
  dx[:,2] = @NLexpression(n.ocp.mdl, [i=1:L], ux[i]*sin(psi[i]) + (v[i] + la*r[i])*cos(psi[i]));    # Y position
  dx[:,3] = @NLexpression(n.ocp.mdl, [i=1:L], (FYF[i] + FYR[i])/m - r[i]*ux[i]);                    # Lateral Speed
  dx[:,4] = @NLexpression(n.ocp.mdl, [i=1:L], (la*FYF[i]-lb*FYR[i])/Izz);                           # Yaw Rate
  dx[:,5] = @NLexpression(n.ocp.mdl, [i=1:L], r[i]);                                                # Yaw Angle
  return dx
end
"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 10/01/2016, Last Modified: 6/22/2017 \n
--------------------------------------------------------------------------------------\n
"""
function ThreeDOFv1(n,
                    x0::Vector,
                     t::Vector,
                     U::Matrix,
                    t0::Float64,
                    tf::Float64)

    @unpack_Vpara n.ocp.params[1]
    # create splines
    sp_SA = linearSpline(t,U[:,1])
    sp_U = linearSpline(t,U[:,2])

    f = (dx,x,p,t) -> begin
    # states
    V   = x[3]  # 3. Lateral Speed
    R   = x[4]  # 4. Yaw Rate
    PSI = x[5]  # 5. Yaw angle

    # controls
    SA  = sp_SA[t] # Steering Angle
    U  = sp_U[t]   # Longitudinal Speed

    # set variables for tire equations
    Ax = 0

    # diff eqs.
    dx[1]   = U*cos(PSI) - (V + la*R)*sin(PSI)    # X position
    dx[2] 	= U*sin(PSI) + (V + la*R)*cos(PSI)    # Y position
    dx[3]   = (@F_YF() + @F_YR())/m - R*U         # Lateral Speed
    dx[4]  	= (la*@F_YF()-lb*@F_YR())/Izz         # Yaw Rate
    dx[5]  	= R                                   # Yaw Angle
  end
  tspan = (t0,tf)
  prob = ODEProblem(f,x0,tspan)

  sol = DiffEqBase.solve(prob,Tsit5())
  U = [sp_SA,sp_U]
  return sol, U
end

"""
ThreeDOFv2_expr(n) & sol, U = ThreeDOFv2(n,x,u)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 10/20/2017, Last Modified:  5/31/2017 \n
--------------------------------------------------------------------------------------\n
# this vehicle model is controlled using steering rate and longitudinal jerk

"""
function ThreeDOFv2_expr(n)

  @unpack_Vpara n.ocp.params[1]

  # lateral tire load
  FYF=:(($PD2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))*sin($PC1*atan(((($PK1*sin(2*atan($PK2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j]) + $PH2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PH1)) - (($PE2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PE1)*(1 - $PE3)*(((atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j]) + $PH2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PH1))/((((atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j]) + $PH2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PH1)^2 + $EP^2)^(0.5)))*(((($PK1*sin(2*atan($PK2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j]) + $PH2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PH1)) - atan(((($PK1*sin(2*atan($PK2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j]) + $PH2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX) + $PH1)))))) + ($PV2*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)^2 + $PV1*($FzF0 - (ax[j] - v[j]*r[j])*$KZX)));
  FYR=:(($PD2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))*sin($PC1*atan(((($PK1*sin(2*atan($PK2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/(ux[j]+$EP))) + $PH2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PH1)) - (($PE2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PE1)*(1 - $PE3*(((atan((v[j] - $lb*r[j])/(ux[j]+$EP))) + $PH2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PH1))/((((atan((v[j] - $lb*r[j])/(ux[j]+$EP))) + $PH2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PH1)^2 + $EP^2)^(0.5))))*(((($PK1*sin(2*atan($PK2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/(ux[j]+$EP))) + $PH2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PH1)) - atan(((($PK1*sin(2*atan($PK2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/(ux[j]+$EP))) + $PH2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX) + $PH1)))))) + ($PV2*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)^2 + $PV1*($FzR0 + (ax[j] - v[j]*r[j])*$KZX)));

  dx = Array{Expr}(8)
  dx[1] = :(ux[j]*cos(psi[j]) - (v[j] + $la*r[j])*sin(psi[j]))   # X position
  dx[2] = :(ux[j]*sin(psi[j]) + (v[j] + $la*r[j])*cos(psi[j]))   # Y position
  dx[3] = :(($FYF + $FYR)/$m - r[j]*ux[j])                       # Lateral Speed
  dx[4] = :(($la*$FYF-$lb*$FYR)/$Izz)                            # Yaw Rate
  dx[5] = :(r[j])                                                # Yaw Angle
  dx[6] = :(sr[j])                                               # Steering Angle
  dx[7] = :(ax[j])                                               # Longitudinal Speed
  dx[8] = :(jx[j])                                               # Longitudinal Acceleration

  # rear vertical tire load
  FZ_RL=:(0.5*($FzR0 + $KZX*(ax[j] - v[j]*r[j])) - $KZYR*(($FYF + $FYR)/$m))
  FZ_RR=:(0.5*($FzR0 + $KZX*(ax[j] - v[j]*r[j])) + $KZYR*(($FYF + $FYR)/$m))
  FZ_RL_con=:(0 <= $FZ_RL - $Fz_min)
  FZ_RR_con=:(0 <= $FZ_RR - $Fz_min)

  # front vertical tire load
  FZ_FL=:(0.5*($FzR0 - $KZX*(ax[j] - v[j]*r[j])) - $KZYF*(($FYF + $FYR)/$m))
  FZ_FR=:(0.5*($FzR0 - $KZX*(ax[j] - v[j]*r[j])) + $KZYF*(($FYF + $FYR)/$m))
  FZ_FL_con=:(0 <= $FZ_FL - $Fz_min)
  FZ_FR_con=:(0 <= $FZ_FR - $Fz_min)

  # linear tire forces
  Fyf_linear=:($Fy_min <= (atan((v[j] + $la*r[j])/(ux[j]+$EP)) - sa[j])*$Caf <= $Fy_max);
  Fyr_linear=:($Fy_min <= atan((v[j] - $lb*r[j])/(ux[j]+$EP))*$Car <= $Fy_max);

  # nonlinear accleleration bounds
  min_ax=:(0 <= ax[j] - ($AXC[5]*ux[j]^3 + $AXC[6]*ux[j]^2 + $AXC[7]*ux[j] + $AXC[8]))
  max_ax=:(ax[j] - ($AXC[1]*ux[j]^3 + $AXC[2]*ux[j]^2 + $AXC[3]*ux[j] + $AXC[4]) <= 0)

  con=[FZ_RL_con,FZ_RR_con,FZ_FL_con,FZ_FR_con,Fyf_linear,Fyr_linear,min_ax,max_ax]

  # expression for cost function
  tire_expr = :(2 + tanh(-($FZ_RL - $a_t)/$b_t) + tanh(-($FZ_RR - $a_t)/$b_t))

  return dx,con,tire_expr
end

"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 10/01/2016, Last Modified: 6/22/2017 \n
--------------------------------------------------------------------------------------\n
"""
function ThreeDOFv2(n,
                   x0::Vector,
                   t::Vector,
                   U::Matrix,
                   t0::Float64,
                   tf::Float64)
    if length(t)!=size(U)[1]
      error(" \n The length of the time vector must match the length of the control input \n")
    end

    @unpack_Vpara n.ocp.params[1]

    # create splines
    sp_SR=linearSpline(t,U[:,1])
    sp_Jx=linearSpline(t,U[:,2])

    f = (dx,x,p,t) -> begin

    # states
    V   = x[3]  # 3. Lateral Speed
    R   = x[4]  # 4. Yaw Rate
    PSI = x[5]  # 5. Yaw angle
    SA  = x[6]  # 6. Steering Angle
    U   = x[7]  # 7. Longitudinal Speed
    Ax  = x[8]  # 8. Longitudinal Acceleration

    # controls
    SR  = sp_SR[t]
    Jx  = sp_Jx[t]

    # diff eqs.
    dx[1]   = U*cos(PSI) - (V + la*R)*sin(PSI)    # X position
    dx[2] 	= U*sin(PSI) + (V + la*R)*cos(PSI)    # Y position
    dx[3]   = (@F_YF() + @F_YR())/m - R*U         # Lateral Speed
    dx[4]  	= (la*@F_YF()-lb*@F_YR())/Izz         # Yaw Rate
    dx[5]  	= R                                   # Yaw Angle
    dx[6]   = SR                                  # Steering Angle
    dx[7]  	= Ax                                  # Longitudinal Speed
    dx[8]  	= Jx                                  # Longitudinal Acceleration
  end
  tspan = (t0,tf)
  prob = ODEProblem(f,x0,tspan)
  sol = DiffEqBase.solve(prob,Tsit5())
  U = [sp_SR,sp_Jx]
  return sol, U
end

"""
ThreeDOFv3_expr(n) & sol, U = ThreeDOFv3(n,x,u)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 11/3/2018, Last Modified:  11/3/2018 \n
--------------------------------------------------------------------------------------\n
# this vehicle model is controlled using steering rate and it has a constant speed
"""
function ThreeDOFv3_expr(n)

    @unpack_Vpara n.ocp.params[1]
    ux = n.ocp.params[5]["X0"]["ux"];
    # lateral tire load
    FYF=:(($PD2*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*($FzF0 - ( - v[j]*r[j])*$KZX))*sin($PC1*atan(((($PK1*sin(2*atan($PK2*($FzF0 - ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j]) + $PH2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PH1)) - (($PE2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PE1)*(1 - $PE3)*(((atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j]) + $PH2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PH1))/((((atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j]) + $PH2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PH1)^2 + $EP^2)^(0.5)))*(((($PK1*sin(2*atan($PK2*($FzF0 - ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j]) + $PH2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PH1)) - atan(((($PK1*sin(2*atan($PK2*($FzF0 - ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzF0 - ( - v[j]*r[j])*$KZX))^2 + $EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j]) + $PH2*($FzF0 - ( - v[j]*r[j])*$KZX) + $PH1)))))) + ($PV2*($FzF0 - ( - v[j]*r[j])*$KZX)^2 + $PV1*($FzF0 - ( - v[j]*r[j])*$KZX)));
    FYR=:(($PD2*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*($FzR0 + ( - v[j]*r[j])*$KZX))*sin($PC1*atan(((($PK1*sin(2*atan($PK2*($FzR0 + ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/($ux+$EP))) + $PH2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PH1)) - (($PE2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PE1)*(1 - $PE3*(((atan((v[j] - $lb*r[j])/($ux+$EP))) + $PH2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PH1))/((((atan((v[j] - $lb*r[j])/($ux+$EP))) + $PH2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PH1)^2 + $EP^2)^(0.5))))*(((($PK1*sin(2*atan($PK2*($FzR0 + ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/($ux+$EP))) + $PH2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PH1)) - atan(((($PK1*sin(2*atan($PK2*($FzR0 + ( - v[j]*r[j])*$KZX))))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)) + (($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)))/((($PD2*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PD1*$PC1*($FzR0 + ( - v[j]*r[j])*$KZX))^2+$EP^2)^(0.5))*0.001)+$EP))*((atan((v[j] - $lb*r[j])/($ux+$EP))) + $PH2*($FzR0 + ( - v[j]*r[j])*$KZX) + $PH1)))))) + ($PV2*($FzR0 + ( - v[j]*r[j])*$KZX)^2 + $PV1*($FzR0 + ( - v[j]*r[j])*$KZX)));

    dx = Array{Expr}(6)
    dx[1] = :($ux*cos(psi[j]) - (v[j] + $la*r[j])*sin(psi[j]))   # X position
    dx[2] = :($ux*sin(psi[j]) + (v[j] + $la*r[j])*cos(psi[j]))   # Y position
    dx[3] = :(($FYF + $FYR)/$m - r[j]*$ux)                       # Lateral Speed
    dx[4] = :(($la*$FYF-$lb*$FYR)/$Izz)                            # Yaw Rate
    dx[5] = :(r[j])                                                # Yaw Angle
    dx[6] = :(sr[j])                                               # Steering Angle

    ax = 0.0;
    # rear vertical tire load
    FZ_RL=:(0.5*($FzR0 + $KZX*( - v[j]*r[j])) - $KZYR*(($FYF + $FYR)/$m))
    FZ_RR=:(0.5*($FzR0 + $KZX*( - v[j]*r[j])) + $KZYR*(($FYF + $FYR)/$m))
    FZ_RL_con=:(0 <= $FZ_RL - $Fz_min)
    FZ_RR_con=:(0 <= $FZ_RR - $Fz_min)

    # front vertical tire load
    FZ_FL=:(0.5*($FzR0 - $KZX*( - v[j]*r[j])) - $KZYF*(($FYF + $FYR)/$m))
    FZ_FR=:(0.5*($FzR0 - $KZX*( - v[j]*r[j])) + $KZYF*(($FYF + $FYR)/$m))
    FZ_FL_con=:(0 <= $FZ_FL - $Fz_min)
    FZ_FR_con=:(0 <= $FZ_FR - $Fz_min)

    # linear tire forces
    Fyf_linear=:($Fy_min <= (atan((v[j] + $la*r[j])/($ux+$EP)) - sa[j])*$Caf <= $Fy_max);
    Fyr_linear=:($Fy_min <= atan((v[j] - $lb*r[j])/($ux+$EP))*$Car <= $Fy_max);

    con=[FZ_RL_con,FZ_RR_con,FZ_FL_con,FZ_FR_con,Fyf_linear,Fyr_linear]

    # expression for cost function
    tire_expr = :(2 + tanh(-($FZ_RL - $a_t)/$b_t) + tanh(-($FZ_RR - $a_t)/$b_t))

  return dx,con,tire_expr
end

"""
# NOTE could simplify even more and modify tire force equation to remove Ax
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 11/3/2018, Last Modified:  11/3/2018 \n
--------------------------------------------------------------------------------------\n
"""
function ThreeDOFv3(n,
                   x0::Vector,
                   t::Vector,
                   U::Matrix,
                   t0::Float64,
                   tf::Float64)
    if length(t)!=size(U)[1]
      error(" \n The length of the time vector must match the length of the control input \n")
    end

    @unpack_Vpara n.ocp.params[1]

    # create spline
    sp_SR=linearSpline(t,U[:,1])

    f = (dx,x,p,t) -> begin

    # states
    V   = x[3]  # 3. Lateral Speed
    R   = x[4]  # 4. Yaw Rate
    PSI = x[5]  # 5. Yaw angle
    SA  = x[6]  # 6. Steering Angle

    # params
    U   = u0_  # Longitudinal Speed
    Ax  = 0    #  Longitudinal Acceleration

    # controls
    SR  = sp_SR[t]

    # diff eqs.
    dx[1]   = U*cos(PSI) - (V + la*R)*sin(PSI)    # X position
    dx[2] 	= U*sin(PSI) + (V + la*R)*cos(PSI)    # Y position
    dx[3]   = (@F_YF() + @F_YR())/m - R*U         # Lateral Speed
    dx[4]  	= (la*@F_YF()-lb*@F_YR())/Izz         # Yaw Rate
    dx[5]  	= R                                   # Yaw Angle
    dx[6]   = SR                                  # Steering Angle
  end
  tspan = (t0,tf)
  prob = ODEProblem(f,x0,tspan)
  sol = DiffEqBase.solve(prob,Tsit5())
  U = [sp_SR]  # NOTE poor choice of variable names to use U twice.
  return sol, U
end


text:
# Forcing classes (areas, lines, points)

#=
Notes:
- The principle here is that the problem specification will supply the shape
 of the forcing region (area, line) and a model function to apply this forcing.
- The forcing model function should be able to make use of other system data, like physical
  parameters
=#

abstract type AbstractForcingModel end
abstract type AbstractRegionCache end

for ftype in [:Area,:Line,:Point]
    fmname = Symbol(string(ftype)*"ForcingModel")
    typename = lowercase(string(ftype))
    fmcache = string(ftype)*"RegionCache"

    eval(quote
      export $fmname

      @doc """
            $($fmname)(shape::Union{Body,BodyList,VectorData},model_function!)

        Bundles a `shape` (i.e., a `Body`, `BodyList`, or `VectorData`) and a `model_function!` (a function
        that returns the strength of the forcing) for $($typename)-type forcing.
        `model_function!` be in-place with a signature of the form

            model_function!(str,state,t,fcache,phys_params)

        where `str` is the strength to be returned, `state` the state vector,
        `t` is time, `fcache` is a corresponding `$($fmcache)`
        and `phys_params` are user-supplied physical parameters. Any of these can
        be utilized to compute the strength.
        """
        struct $fmname{RT,KT,MDT<:Function} <: AbstractForcingModel
          shape :: RT
          kwargs :: KT
          fcn :: MDT
          $fmname(shape,fcn;kwargs...) = new{typeof(shape),typeof(kwargs),typeof(fcn)}(shape,kwargs,fcn)
        end
      end)
end

"""
    PointForcingModel(point_function::Function,model_function!::Function)

Bundles a `point_function` (a function that returns the positions of forcing points)
and a `model_function` (a function that returns the strength of the forcing) for point-type forcing.

  `point_function` must have an out-of-place signature of the form

      point_function(state,t,fcache,phys_params)

  where `state` the state vector, `t` is time, `fcache` is a corresponding `PointRegionCache`
  and `phys_params` are user-supplied physical parameters. Any of these can
  be utilized to compute the positions. It must return either a vector for
  each coordinate or `VectorData`.

  `model_function!` must have an in-place signature of the form

      model_function!(str,state,t,fcache,phys_params)

  where `str` is the strength to be returned.
""" PointForcingModel(::Function,::Function)

#=
Region caches
=#

struct AreaRegionCache{MT,ST,CT<:BasicILMCache} <: AbstractRegionCache
    mask :: MT
    str :: ST
    cache :: CT
end

struct LineRegionCache{ACT,ST,CT<:BasicILMCache} <: AbstractRegionCache
    s :: ACT
    str :: ST
    cache :: CT
end

struct PointRegionCache{ST,CT<:PointCollectionCache} <: AbstractRegionCache
    str :: ST
    cache :: CT
end

"""
    AreaRegionCache(shape::Body/BodyList,cache::BasicILMCache)

Create an area region (basically, a mask) of the shape(s) `shape`, using
the data in `cache` to provide the details of the regularization.
""" AreaRegionCache(::Union{Body,BodyList},::BasicILMCache)

"""
    LineRegionCache(shape::Body/BodyList,cache::BasicILMCache)

Create a line region of the shape(s) `shape`, using
the data in `cache` to provide the details of the regularization.
""" LineRegionCache(::Union{Body,BodyList},::BasicILMCache)

"""
    PointRegionCache(pts::VectorData,cache::BasicILMCache[,kwargs])

Create a regularized point collection based on points `pts`, using the data
in `cache` to provide the details of the regularization. The `kwargs` can be used to override
the regularization choices, such as ddf.
""" PointRegionCache(::VectorData,::BasicILMCache)

# Some constructors of the region caches that distinguish scalar and vector data
for f in [:Scalar,:Vector]
    cname = Symbol("Surface"*string(f)*"Cache")
    pcname = Symbol(string(f)*"PointCollectionCache")
    gdtype = Symbol(string(f)*"GridData")
    @eval function AreaRegionCache(g::PhysicalGrid,shape::BodyList,data_prototype::$gdtype;kwargs...)
        cache = $cname(shape,g;kwargs...)
        m = mask(cache)
        str = similar_grid(cache)
        return AreaRegionCache{typeof(m),typeof(str),typeof(cache)}(m,str,cache)
    end

    @eval function LineRegionCache(g::PhysicalGrid,shape::BodyList,data_prototype::$gdtype;kwargs...)
        cache = $cname(shape,g;kwargs...)
        pts = points(shape)
        s = arcs(shape)
        str = similar_surface(cache)
        return LineRegionCache{typeof(s),typeof(str),typeof(cache)}(s,str,cache)
    end

    @eval function PointRegionCache(g::PhysicalGrid,pts::VectorData,data_prototype::$gdtype;kwargs...)
        cache = $pcname(pts,g;kwargs...)

        str = similar(cache.sdata_cache)
        return PointRegionCache{typeof(str),typeof(cache)}(str,cache)
    end

    @eval function PointRegionCache(g::PhysicalGrid,pts::Function,data_prototype::$gdtype;kwargs...)
        cache = $pcname(VectorData(0),g;kwargs...)

        str = similar(cache.sdata_cache)
        return PointRegionCache{typeof(str),typeof(cache)}(str,cache)
    end

end

for f in [:AreaRegionCache,:LineRegionCache]
    @eval $f(shape::Union{Body,BodyList},cache::BasicILMCache{N,SCA};scaling=SCA,kwargs...) where {N,SCA} = $f(cache.g,shape,similar_grid(cache);scaling=scaling,kwargs...)
    @eval $f(g,shape::Body,a...;kwargs...) = $f(g,BodyList([shape]),a...;kwargs...)
end

PointRegionCache(pts::Union{VectorData,Function},cache::AbstractBasicCache;kwargs...) =
      PointRegionCache(cache.g,pts,similar_grid(cache);kwargs...)

"""
    mask(ar::AreaRegionCache)

Return the mask for the given area region `ar`.
"""
mask(ar::AreaRegionCache) = ar.mask

"""
    arcs(lr::LineRegionCache)

Return the vector of arc length coordinates for the given line region `lr`.
"""
arcs(lr::LineRegionCache) = lr.s

"""
    points(pr::PointRegionCache)

Return the vector of coordinates of points associated with `pr`.
"""
points(pr::PointRegionCache) = points(pr.cache)

#=
Assembly of the forcing model and forcing region
=#

"""
    ForcingModelAndRegion

A type that holds the forcing model function, region, and cache
# Constructors

`ForcingModelAndRegion(model::AbstractForcingModel,cache::BasicILMCache)`

`ForcingModelAndRegion(modellist::Vector{AbstractForcingModel},cache::BasicILMCache)`

These forms gets called generally when building the extra cache. They
also gracefully generate an empty list of models if one passes along `nothing`
in the first argument.
""" ForcingModelAndRegion(::AbstractForcingModel,::BasicILMCache)

struct ForcingModelAndRegion{RT<:AbstractRegionCache,ST,MT,KT}
    region_cache :: RT
    shape :: ST
    fcn :: MT
    kwargs :: KT
end

for f in [:Area,:Line,:Point]
  modtype = Symbol(string(f)*"ForcingModel")
  regcache = Symbol(string(f)*"RegionCache")
  @eval function ForcingModelAndRegion(model::$modtype,cache::BasicILMCache)
      region_cache = $regcache(model.shape,cache;model.kwargs...)
      ForcingModelAndRegion(region_cache,model.shape,model.fcn,model.kwargs)
  end
end

function ForcingModelAndRegion(flist::Vector{<:AbstractForcingModel},cache::BasicILMCache)
   fmlist = ForcingModelAndRegion[]
   for f in flist
     push!(fmlist,ForcingModelAndRegion(f,cache))
   end
   return fmlist
end

ForcingModelAndRegion(::Any,cache::BasicILMCache) = ForcingModelAndRegion(AbstractForcingModel[],cache)

#=
Application of forcing
=#

"""
    apply_forcing!(out,x,t,fv::Vector{ForcingModelAndRegion},phys_params)

Return the total contribution of forcing in the vector `fv` to `out`,
based on the current state `x`, time `t`, and physical parameters in `phys_params`.
Note that `out` is zeroed before the contributions are added.
"""
function apply_forcing!(out,x,t,fr::Vector{<:ForcingModelAndRegion},phys_params)
    fill!(out,0.0)
    for f in fr
        _apply_forcing!(out,x,t,f,phys_params)
    end
    return out
end

"""
    apply_forcing!(dx,x,t,f::ForcingModelAndRegion,phys_params)

Return the contribution of forcing in `f` to the right-hand side `dx`
based on the current state `x`, time `t`, and physical parameters in `phys_params`.
"""
apply_forcing!(dx,x,t,fr::ForcingModelAndRegion,phys_params) = apply_forcing!(dx,x,t,[fr],phys_params)

#=
The following define how forcing of each type get applied. Each one
calls the forcing model to determine the strength.
=#

function _apply_forcing!(dx,x,t,fcache::ForcingModelAndRegion{<:AreaRegionCache},phys_params)
    @unpack region_cache, fcn = fcache
    @unpack str = region_cache

    fill!(str,0.0)
    fcn(str,x,t,region_cache,phys_params)

    dx .+= str*mask(region_cache)
    return dx
end

function _apply_forcing!(dx,x,t,fcache::ForcingModelAndRegion{<:LineRegionCache},phys_params)
    @unpack region_cache, fcn = fcache
    @unpack str, cache = region_cache

    fill!(str,0.0)
    fcn(str,x,t,region_cache,phys_params)

    fill!(cache.gdata_cache,0.0)
    regularize!(cache.gdata_cache,str,cache)
    dx .+= cache.gdata_cache
    return dx
end

function _apply_forcing!(dx,x,t,fcache::ForcingModelAndRegion{<:PointRegionCache,T},phys_params) where T
    @unpack region_cache, fcn = fcache
    @unpack str, cache = region_cache
    @unpack regop = cache

    fill!(str,0.0)
    fcn(str,x,t,region_cache,phys_params)

    fill!(cache.gdata_cache,0.0)
    regop(cache.gdata_cache,str)
    dx .+= cache.gdata_cache
    return dx
end

function _apply_forcing!(dx,x,t,fcache::ForcingModelAndRegion{<:PointRegionCache,<:Function},phys_params)
    @unpack region_cache, shape, fcn, kwargs = fcache
    @unpack cache = region_cache

    # `shape` is a function in this case, used to obtain the point coordinates
    # Use is to to generate an instantaneous PointRegionCache
    _pts = shape(x,t,region_cache,phys_params)
    typeof(_pts) <: Tuple ? pts = VectorData(_pts...) : pts = _pts
    new_region_cache = PointRegionCache(pts,cache;kwargs...)
    @unpack str, cache = new_region_cache
    @unpack regop = cache

    fill!(str,0.0)
    fcn(str,x,t,region_cache,phys_params)

    fill!(cache.gdata_cache,0.0)
    regop(cache.gdata_cache,str)
    dx .+= cache.gdata_cache
    return dx
end