text:
export FieldmapNFFTOp, InhomogeneityData, createInhomogeneityData_

mutable struct InhomogeneityData{T}
  A_k::Matrix{Complex{T}}
  C_k::Matrix{Complex{T}}
  times::Vector{T}
  Cmap::Matrix{Complex{T}}
  t_hat::T
  z_hat::Complex{T}
  method::String
end

mutable struct FieldmapNFFTOp{T,F1,F2,D} <:AbstractLinearOperator{Complex{T}}
  nrow :: Int
  ncol :: Int
  symmetric :: Bool
  hermitian :: Bool
  prod! :: Function
  tprod! :: F1
  ctprod! :: F2
  nprod :: Int
  ntprod :: Int
  nctprod :: Int
  args5 :: Bool
  use_prod5! :: Bool
  allocated5 :: Bool
  Mv5 :: Vector{Complex{T}}
  Mtu5 :: Vector{Complex{T}}
  plans
  idx::Vector{Vector{Int64}}
  circTraj::Bool
  shape::NTuple{D,Int64}
  cparam::InhomogeneityData{T}
end

"""
    FieldmapNFFTOp(shape::NTuple{D,Int64}, tr::Trajectory,
                        correctionmap::Array{Complex{T},D};
                        method::String="nfft",
                        echoImage::Bool=true,
                        alpha::Float64=1.75,
                        m::Float64=3.0,
                        K=20,
                        kargs...) where D

generates a `FieldmapNFFTOp` which evaluates the MRI Fourier signal encoding operator,
including B0-inhomogeneities using time-segmented NFFTs.

# Arguments:
* `shape::NTuple{D,Int64}`             - size of image to encode/reconstruct
* `tr::Trajectory`                     - Trajectory with the kspace nodes to sample
* `correctionmap::Array{Complex{T},D}` - fieldmap for the correction of off-resonance effects
* (`method::String="nfft"`)            - method to use for time-segmentation when correctio field inhomogeneities
* (`echoImage::Bool=false`)            - if true sampling times will only be considered relative to the echo time
                                         this results in complex valued image even for real-valued input.
* (`alpha::Float64=1.75`)              - oversampling factor for interpolation
* (`m::Float64=3.0`)                   - truncation size of interpolation kernel
* (`K=20`)                             - number of translates for LeastSquares approaches
                                         (not NFFT-approach) to time-segmentation
* (`kargs`)                            - additional keyword arguments
"""
function FieldmapNFFTOp(shape::NTuple{D,Int64}, tr::Trajectory,
                        correctionmap::Array{Complex{T},D};
                        method::String="nfft",
                        echoImage::Bool=true,
                        alpha::Float64=1.25,
                        m::Float64=2.0,
                        K=20,
                        K_tol::Float64=1.e-3,
                        step::Int64=10,
                        kargs...) where {T,D}

  nodes,times = kspaceNodes(tr), readoutTimes(tr)
  if echoImage
    times = times .- echoTime(tr)
  end
  nrow = size(nodes,2)
  ncol = prod(shape)

 # create and truncate low-rank expansion
  cparam = createInhomogeneityData_(vec(times), correctionmap; K=K, alpha=alpha, m=m, method=method, K_tol=K_tol, numSamp=numSamplingPerProfile(tr),step=step)
  K = size(cparam.A_k,2)

  @debug "K = $K"

  plans = Vector{NFFT.NFFTPlan{T,D,1}}(undef,K)
  idx = Vector{Vector{Int64}}(undef,K)
  for κ=1:K
    idx[κ] = findall(x->x!=0.0, cparam.A_k[:,κ])
    plans[κ] = plan_nfft(nodes[:,idx[κ]], shape, m=3, σ=1.25, precompute = NFFT.FULL)
  end

  d = [zeros(Complex{T}, length(idx[κ])) for κ=1:K ]

  circTraj = isCircular(tr)

  mul!(res, x::Vector{T}) where T =
     (res .= produ(x,nrow,ncol,shape,plans,idx,cparam,circTraj,d))
  ctmul!(res, y::Vector{T}) where T =
     (res .= ctprodu(y,shape,plans,idx,cparam,circTraj,d))

  return FieldmapNFFTOp{T,Nothing,Function,D}(nrow, ncol, false, false
            , mul!
            , nothing
            , ctmul!, 0, 0, 0, false, false, false, Complex{T}[], Complex{T}[]
            , plans, idx, circTraj, shape, cparam)
end

function Base.copy(S::FieldmapNFFTOp{T,Nothing,Function,D}) where {T,D}
  K=length(S.plans)
  plans = [copy(S.plans[i]) for i=1:K]
  idx = deepcopy(S.idx)

  d = [zeros(Complex{T}, length(idx[κ])) for κ=1:K ]

  cparam = deepcopy(S.cparam)

  mul!(res, x::Vector{T}) where T =
     (res .= produ(x,S.nrow,S.ncol,S.shape,plans,idx,cparam,S.circTraj,d))
  ctmul!(res, y::Vector{T}) where T =
     (res .= ctprodu(y,S.shape,plans,idx,cparam,S.circTraj,d))

  D_ = length(S.shape)
  return FieldmapNFFTOp{T,Nothing,Function,D_}(S.nrow, S.ncol, false, false
            , mul!
            , nothing
            , ctmul!, 0, 0, 0, false, false, false, Complex{T}[], Complex{T}[]
            , plans, idx, S.circTraj, S.shape, cparam)
end

# function produ{T<:Float64}(x::Vector{T}, numOfNodes::Int, numOfPixel::Int, shape::Tuple, plan::Vector{NFFTPlan{Float64,2,1}}, cparam::InhomogeneityData, density::Vector{Float64}, symmetrize::Bool)
function produ(x::Vector{T}, numOfNodes::Int, numOfPixel::Int, shape::Tuple, plan,
               idx::Vector{Vector{Int64}}, cparam::InhomogeneityData,
               shutter::Bool, d) where T
  K = size(cparam.A_k,2)
  s = zeros(T,numOfNodes)

  if shutter
    circularShutter!(reshape(x, shape), 1.0)
  end

  # Preprocessing step when time and correctionMap are centered
  if cparam.method == "nfft"
    x_ = x .* exp.(-vec(cparam.Cmap) * cparam.t_hat )
  else
    x_ = copy(x)
  end

  sp = Threads.SpinLock()
  produ_inner(K,cparam.C_k, cparam.A_k, shape, d, s, sp, plan, idx, x_)

  # Postprocessing step when time and correctionMap are centered
  if cparam.method == "nfft"
      s .*= exp.(-cparam.z_hat*(cparam.times .- cparam.t_hat) )
  end
  return s
end

function produ_inner(K, C, A, shape, d, s, sp, plan, idx, x_)
  @sync for κ=1:K
    Threads.@spawn begin
      p_ = C[κ,:] .* x_
      mul!(d[κ], plan[κ], reshape(p_, shape))

      lock(sp)
      for k=1:length(idx[κ])
        l=idx[κ][k]
        s[l] += d[κ][k]*A[l,κ]
      end
      unlock(sp)
    end
  end

  return
end

# function ctprodu{T<:Complex{T}}(x::Vector{T}, shape::Tuple, plan::Vector{NFFTPlan{Float64,2,1}, cparam::InhomogeneityData, density::Vector{Float64}, symmetrize::Bool)
function ctprodu(x::Vector{Complex{T}}, shape::Tuple, plan, idx::Vector{Vector{Int64}},
                 cparam::InhomogeneityData{T}, shutter::Bool, d) where T

  y = zeros(Complex{T},prod(shape))
  K = size(cparam.A_k,2)

  x_ = copy(x)

  # Preprocessing step when time and correctionMap are centered
  if cparam.method == "nfft"
      x_ .*= conj.(exp.(-cparam.z_hat*(cparam.times .- cparam.t_hat)))
  end

  sp = Threads.SpinLock()
  ctprodu_inner(K,cparam.C_k, cparam.A_k, shape, d, y, sp, plan, idx, x_)

  if cparam.method == "nfft"
    y .*=  conj(exp.(-vec(cparam.Cmap) * cparam.t_hat))
  end

  if shutter
    circularShutter!(reshape(y, shape), 1.0)
  end

  return y
end

function ctprodu_inner(K, C::Matrix{Complex{T}}, A::Matrix{Complex{T}}, shape, d, y, sp, plan, idx, x_) where T

  @sync for κ=1:K
    Threads.@spawn begin
       for k=1:length(idx[κ])
         d[κ][k] = conj.(A[idx[κ][k],κ]) * x_[idx[κ][k]]
       end
       p_ = zeros(Complex{T}, shape)

       mul!(p_, adjoint(plan[κ]), d[κ])

       lock(sp)
       for k=1:length(p_)
         y[k] += conj(C[κ,k]) * p_[k]
       end
       unlock(sp)
    end
  end

  return
end

function createInhomogeneityData_(times::Vector{T},
                                  correctionmap::Array{Complex{T},D};
                                  K::Int64=20,
                                  K_tol::Float64=1.e-3,
                                  alpha::Float64=1.75,
                                  m::Float64 = 4.0,
                                  method="nfft",
                                  numSamp::Int64=length(times),
                                  step::Int64=10) where {T,D}

    if method == "nfft"
      A,C = get_AC_coefficients_nfft(K,vec(times),vec(correctionmap), alpha, m)
    elseif method == "leastsquare"
      A,C = get_AC_coefficients_lsqr(K,vec(times),vec(correctionmap))
    elseif method == "hist"
      A = get_A_coefficients_hist_lsqr(K,vec(times),vec(correctionmap))
      C = get_C_coefficients_hist_lsqr(K,vec(times),vec(correctionmap))
    elseif method == "leastsquareFS"
      A,C = get_AC_coefficients_lsqr_fs(K,vec(times),vec(correctionmap))
    elseif method == "histFS"
      A = get_A_coefficients_hist_lsqr_fs(K,vec(times),vec(correctionmap))
      # C = get_C_coefficients_hist_lsqr_fs(K,vec(times),vec(correctionmap))
      C = get_C_coefficients_lsqr(vec(times),vec(correctionmap),A,numSamp)
    elseif method == "svd"
      A = get_A_coefficients_svd(K,vec(times)[1:numSamp],vec(correctionmap); K_tol=K_tol)
      C = get_C_coefficients_lsqr(vec(times),vec(correctionmap),A,numSamp,step)
      rep = div(length(times),numSamp)
      A = repeat(A,outer=(rep,1))
    else
      error("approximation scheme $(interp) is not yet implemented")
    end
    if size(A,2) != size(C,1)
      @info "method: $(method), A: $(size(A)), C: $(size(C))"
      error("Consistency check failed! A and C are not compatible")
    end

    t_hat = (times[1] + times[end])/2
    z_hat = minimum(real(correctionmap)) + maximum(real(correctionmap))
    z_hat += 1im*(minimum(imag(correctionmap)) + maximum(imag(correctionmap)))
    z_hat *= T(0.5)

    return InhomogeneityData(A,C,vec(times),correctionmap,t_hat,z_hat,method)
end


text:
# TwoStageBenders.jl
#
# Author
#   Oscar Dowson
#
# Description
#   A simple example reading a StochOptFormat file and solving it via Benders
#   decomposition.
#
#   The code is intended for pedagogical use. It does not contain robust checks
#   or nice error messages, preferring to throw assertion errors.
#
# Usage
#   julia TwoStageBenders.jl [problem]
#   julia TwoStageBenders.jl ../problems/newsvendor.sof.json
#
# Notes
#   You need to install Julia, and have the following packages installed:
#       Clp, JSON, JSONSchema, JuMP, and SHA.
module TwoStageBenders

import Clp
import JSON
import JSONSchema
import JuMP
import Printf
import SHA

const SCHEMA_FILENAME =
    joinpath(dirname(dirname(@__DIR__)), "sof-latest.schema.json")

const RESULT_SCHEMA_FILENAME =
    joinpath(dirname(dirname(@__DIR__)), "sof_result.schema.json")

struct TwoStageProblem
    sha256::String
    first::JuMP.Model
    second::JuMP.Model
    state_variables::Set{String}
    validation_scenarios::Vector{Dict{String, Any}}
end

"""
    TwoStageProblem(filename::String; validate::Bool = true)

Create a two-stage stochastic program by reading `filename` in StochOptFormat.

If `validate`, validate `filename` against the StochOptFormat schema.
"""
function TwoStageProblem(filename::String; validate::Bool = true)
    sha_256 = open(filename) do io
        bytes2hex(SHA.sha2_256(io))
    end
    data = JSON.parsefile(filename)
    if validate
        _validate(data; schema_filename = SCHEMA_FILENAME)
    end
    @assert(data["version"]["major"] == 0)
    @assert(data["version"]["minor"] == 2)
    first, second = _get_stage_names(data)
    problem = TwoStageProblem(
        sha_256,
        _mathoptformat_to_jump(data, first),
        _mathoptformat_to_jump(data, second),
        Set{String}(),
        get(data, "validation_scenarios", Dict{String, Any}[]),
    )
    for k in keys(problem.first.ext[:state_variables])
        push!(problem.state_variables, k)
    end
    _initialize_first_stage(data, first, problem.first)
    return problem
end

function _validate(data::Dict; schema_filename::String)
    schema = JSONSchema.Schema(JSON.parsefile(schema_filename))
    return JSONSchema.validate(data, schema)
end

function _initialize_first_stage(data::Dict, first::String, sp::JuMP.Model)
    for (name, init) in data["root"]["state_variables"]
        x = JuMP.variable_by_name(sp, sp.ext[:state_variables][name]["in"])
        JuMP.fix(x, init["initial_value"]; force = true)
    end
    JuMP.@variable(sp, -1e6 <= theta <= 1e6)
    JuMP.set_objective_function(sp, JuMP.objective_function(sp) + theta)
    return
end

function _get_stage_names(data::Dict)
    @assert length(data["nodes"]) == 2
    @assert length(data["root"]["successors"]) == 1
    first_node, probability = first(data["root"]["successors"])
    @assert probability == 1.0
    successors = data["nodes"][first_node]["successors"]
    @assert length(successors) == 1
    second_node, probability = first(successors)
    @assert probability == 1.0
    @assert length(data["nodes"][second_node]["successors"]) == 0
    return first_node, second_node
end

function _mathoptformat_to_jump(data, name)
    node = data["nodes"][name]
    sp = data["subproblems"][node["subproblem"]]
    model = JuMP.MOI.FileFormats.Model(format = JuMP.MOI.FileFormats.FORMAT_MOF)
    io = IOBuffer()
    write(io, JSON.json(sp["subproblem"]))
    seekstart(io)
    JuMP.MOI.read!(io, model)
    subproblem = JuMP.Model(Clp.Optimizer)
    JuMP.MOI.copy_to(subproblem, model)
    JuMP.set_silent(subproblem)
    subproblem.ext[:state_variables] = sp["state_variables"]
    subproblem.ext[:realizations] = node["realizations"]
    _convert_realizations(subproblem.ext[:realizations])
    return subproblem
end

# These `_convert_realization` functions are Julia-specific helper functions to
# convert the support dictionaries from `Dict{Symbol, Any}` to
# `Dict{Symbol, Float64}`.
function _convert_realizations(realizations)
    for r in realizations
        r["support"] = convert(Dict{String, Float64}, r["support"])
    end
end

# These are helper functions to extract the incoming and outgoing state
# variables from a JuMP model.
function _incoming_state(sp::JuMP.Model, name::String)
    return JuMP.variable_by_name(sp, sp.ext[:state_variables][name]["in"])
end
function _outgoing_state(sp::JuMP.Model, name::String)
    return JuMP.variable_by_name(sp, sp.ext[:state_variables][name]["out"])
end

function _solve_first_stage(problem::TwoStageProblem)
    sp = problem.first
    JuMP.optimize!(sp)
    @assert JuMP.termination_status(sp) == JuMP.MOI.OPTIMAL
    return Dict(
        "objective" => JuMP.objective_value(sp) - JuMP.value(sp[:theta]),
        "primal" => Dict(
            JuMP.name(x) => JuMP.value(x) for x in JuMP.all_variables(sp)
            # Don't leak the cost-to-go variable back to the user.
            if x != sp[:theta]
        )
    )
end

function _solve_second_stage(
    problem::TwoStageProblem,
    state_variables::Dict{String, Float64},
    random_variables::Dict{String, Float64},
)
    sp = problem.second
    for name in problem.state_variables
        JuMP.fix(_incoming_state(sp, name), state_variables[name]; force = true)
    end
    for (name, w) in random_variables
        JuMP.fix(JuMP.variable_by_name(sp, name), w; force = true)
    end
    JuMP.optimize!(sp)
    return Dict(
        "objective" => JuMP.objective_value(sp),
        "primal" => Dict(
            JuMP.name(x) => JuMP.value(x) for x in JuMP.all_variables(sp)),
    )
end

function _add_benders_optimality_cut(
    problem::TwoStageProblem,
    state_variables::Dict{String, Float64},
    probabilities::Vector{Float64},
    objectives::Vector{Float64},
    dual_variables::Vector{Dict{String, Float64}},
)
    cut_term = JuMP.@expression(problem.first, sum(
            p * (
                y + sum(
                    π[name] * (_outgoing_state(problem.first, name) - x_val)
                    for (name, x_val) in state_variables
                )
            ) for (π, y, p) in zip(dual_variables, objectives, probabilities)
        )
    )
    if JuMP.objective_sense(problem.first) == JuMP.MOI.MAX_SENSE
        JuMP.@constraint(problem.first, problem.first[:theta] <= cut_term)
    else
        JuMP.@constraint(problem.first, problem.first[:theta] >= cut_term)
    end
    return
end

"""
    train(problem::TwoStageProblem; iteration_limit = 20)

Train a two-stage problem using Benders decomposition.
"""
function train(problem::TwoStageProblem; iteration_limit = 20)
    println("Iteration | Lower Bound | Upper Bound | Gap (abs)")
    for iter = 1:iteration_limit
        ret_first = _solve_first_stage(problem)
        probabilities = Float64[]
        objectives = Float64[]
        dual_variables = Dict{String, Float64}[]
        x = Dict(
            name => ret_first["primal"][s["out"]]
            for (name, s) in problem.second.ext[:state_variables]
        )
        for realization in problem.second.ext[:realizations]
            ret = _solve_second_stage(problem, x, realization["support"])
            push!(probabilities, realization["probability"])
            push!(objectives, ret["objective"])
            push!(dual_variables, Dict(
                name => JuMP.shadow_price(
                    JuMP.FixRef(_incoming_state(problem.second, name))
                ) for name in problem.state_variables
            ))
        end
        deterministic_bound = JuMP.objective_value(problem.first)
        stat_bound = ret_first["objective"] +
            sum(p * o for (p, o) in zip(probabilities, objectives))
        is_min = JuMP.objective_sense(problem.first) == JuMP.MOI.MIN_SENSE
        gap = abs(deterministic_bound - stat_bound)
        Printf.@printf(
            "%9d | % 5.4e | % 5.4e | %4.3e\n",
            iter,
            is_min ? deterministic_bound : stat_bound,
            is_min ? stat_bound : deterministic_bound,
            gap,
        )
        if gap < 1e-6
            println("Terminating training: convergence")
            return
        end
        _add_benders_optimality_cut(
            problem, x, probabilities, objectives, dual_variables
        )
    end
    println("Terminating training: iteration limit")
    return
end

"""
    evaluate(
        problem::TwoStageProblem;
        scenarios = problem.validation_scenarios,
    )

Evaluate the policy after training `problem` on `scenarios`. By default, these
are the `validation_scenarios` contained in the StochOptFormat file, but you can
pass a different set if necessary.
"""
function evaluate(
    problem::TwoStageProblem;
    scenarios = problem.validation_scenarios,
    filename::Union{Nothing, String} = nothing
)
    solutions = Vector{Dict{String, Any}}[]
    for s_dict in scenarios
        scenario = s_dict["scenario"]
        @assert length(scenario) == 2
        first_sol = _solve_first_stage(problem)
        incoming_state = Dict(
            name => first_sol["primal"][s["out"]]
            for (name, s) in problem.second.ext[:state_variables]
        )
        second_sol = _solve_second_stage(
            problem,
            incoming_state,
            convert(Dict{String, Float64}, scenario[2]["support"]),
        )
        push!(solutions, [first_sol, second_sol])
    end
    solution = Dict(
        "problem_sha256_checksum" => problem.sha256,
        "scenarios" => solutions,
    )
    _validate(solution; schema_filename = RESULT_SCHEMA_FILENAME)
    if filename !== nothing
        open(filename, "w") do io
            write(io, JSON.json(solution))
        end
    end
    return solution
end

export
    TwoStageProblem,
    train,
    evaluate

end # module TwoStageBenders

if endswith(@__FILE__, PROGRAM_FILE)
    TSSP = TwoStageBenders
    @assert length(ARGS) == 1
    filename = ARGS[1]
    problem = TSSP.TwoStageProblem(filename)
    ret = TSSP.train(problem; iteration_limit = 20)
    solutions = TSSP.evaluate(problem; filename = "sol_jl.json")
    if endswith(filename, "news_vendor.sof.json")
        # Check solutions
        @assert solutions["scenarios"][1][1]["objective"] ≈ -10
        @assert solutions["scenarios"][1][2]["objective"] ≈ 15
        @assert solutions["scenarios"][2][2]["objective"] ≈ 15
        @assert solutions["scenarios"][3][2]["objective"] ≈ 13.5
    end
end

ext: jl

max_stars_repo_name: odow/StochOptFormat

max_issues_repo_name: odow/StochOptFormat

max_forks_repo_name: odow/StochOptFormat


text:
# AUTO GENERATED FILE - DO NOT EDIT

export 'feffery'_antdcarousel

"""
    'feffery'_antdcarousel(;kwargs...)
    'feffery'_antdcarousel(children::Any;kwargs...)
    'feffery'_antdcarousel(children_maker::Function;kwargs...)

An AntdCarousel component.

Keyword arguments:
- `children` (a list of or a singular dash component, string or number; optional)
- `id` (String; optional)
- `autoplay` (Bool; optional)
- `className` (String; optional)
- `dotPosition` (a value equal to: 'top', 'bottom', 'left', 'right'; optional)
- `easing` (String; optional)
- `effect` (a value equal to: 'scrollx', 'fade'; optional)
- `loading_state` (optional): . loading_state has the following type: lists containing elements 'is_loading', 'prop_name', 'component_name'.
Those elements have the following types:
  - `is_loading` (Bool; optional): Determines if the component is loading or not
  - `prop_name` (String; optional): Holds which property is loading
  - `component_name` (String; optional): Holds the name of the component that is loading
- `style` (Dict; optional)
"""
function 'feffery'_antdcarousel(; kwargs...)
        available_props = Symbol[:children, :id, :autoplay, :className, :dotPosition, :easing, :effect, :loading_state, :style]
        wild_props = Symbol[]
        return Component("'feffery'_antdcarousel", "AntdCarousel", "feffery_antd_components", available_props, wild_props; kwargs...)
end

'feffery'_antdcarousel(children::Any; kwargs...) = 'feffery'_antdcarousel(;kwargs..., children = children)
'feffery'_antdcarousel(children_maker::Function; kwargs...) = 'feffery'_antdcarousel(children_maker(); kwargs...)

max_stars_repo_name: ReilYoung/feffery-antd-components

max_issues_repo_name: ReilYoung/feffery-antd-components

max_forks_repo_name: ReilYoung/feffery-antd-components


text:
using BenchmarkTools
using NiBundleAdjustment
using NiBundleAdjustment: vec2cam, vec2scam, compute_reproj_err
using StaticArrays

function load(b, n, m, p)
    dir_in = joinpath(dirname(dirname(@__FILE__)), "data", "ba")
    fn = "ba$(b)_n$(n)_m$(m)_p$(p)"
    fn_in = joinpath(dir_in, fn)
    NiBundleAdjustment.read_ba_instance(string(fn_in,".txt"))
end

cams, X, w, obs, feats = load(4, 372, 47423, 204472)
CAMS = [vec2cam(cams[:,i]) for i = 1:size(cams,2)]
XX = [P3(X[:,i]...) for i=1:size(X, 2)]
FEATS = [P2(feats[:,i]...) for i=1:size(feats, 2)]
println("Normal Objective")
SCAMS = [NiBundleAdjustment.vec2scam(cams[:,i]) for i = 1:size(cams,2)]
display(@benchmark NiBundleAdjustment.ba_objective($SCAMS, $X, $w, $obs, $feats))

println()
println("Reversible Objective")
reproj_err! = zeros(P2{Float64}, size(feats, 2))
w_err! = zero(w)
reproj_err_cache! = zeros(P2{Float64}, size(feats, 2))
display(@benchmark NiBundleAdjustment.ba_objective!($reproj_err!, $w_err!, $reproj_err_cache!, $CAMS, $XX, $w, $obs, $FEATS))

println("ForwardDiff Gradient")
display(@benchmark compute_ba_J(Val(:ForwardDiff), $cams, $X, $w, $obs, $feats))

println()
println("NiLang Gradient")
display(@benchmark compute_ba_J(Val(:NiLang), $CAMS, $XX, $w, $obs, $FEATS))

using KernelAbstractions
using CuArrays
println()
println("NiLang-CUDA Gradient")
display(@benchmark compute_ba_J_multithread(Val(:NiLang), $CAMS, $XX, $w, $obs, $FEATS))
display(@benchmark compute_ba_J_cuda(Val(:NiLang), $(CuArray(CAMS)), $(CuArray(XX)), $(CuArray(w)), $(CuArray(obs)), $(CuArray(FEATS))))
println()


text:
const HDF_FILENAME = "simulation_store.h5"
const HDF_SIMULATION_ROOT_PATH = "simulation"
const OPTIMIZER_STATS_PATH = "optimizer_stats"

mutable struct Dataset
    dataset::HDF5.Dataset
    column_dataset::HDF5.Dataset
    write_index::Int
end

Dataset(dataset, column_dataset) = Dataset(dataset, column_dataset, 1)

# This only applies if chunks are enabled, and that will only likely happen if we enable
# compression.
# The optimal number of chunks to store in memory will vary widely.
# The HDF docs recommend keeping chunk byte sizes between 10 KiB - 1 MiB.
# We want to make it big enough to compress duplicate values.
# The downside to making this larger is that any read causes the
# entire chunk to be read.
# If one variable has 10,000 components and each value is a Float64 then one row would
# consume 10,000 * 8 = 78 KiB
DEFAULT_MAX_CHUNK_BYTES = 128 * KiB

"""
Stores HDF5 datasets for one problem.
"""
mutable struct ProblemDatasets
    duals::Dict{Symbol, Dataset}
    parameters::Dict{Symbol, Dataset}
    variables::Dict{Symbol, Dataset}
end

function ProblemDatasets()
    return ProblemDatasets(
        Dict{Symbol, Dataset}(),
        Dict{Symbol, Dataset}(),
        Dict{Symbol, Dataset}(),
    )
end

"""
Stores simulation data in an HDF file.
"""
mutable struct HdfSimulationStore <: SimulationStore
    file::HDF5.File
    params::SimulationStoreParams
    # The key order is the problem execution order.
    datasets::OrderedDict{Symbol, ProblemDatasets}
    # The key is the problem name.
    optimizer_stats_datasets::Dict{Symbol, HDF5.Dataset}
    optimizer_stats_write_index::Dict{Symbol, Int}
    cache::ResultCache
end

function HdfSimulationStore(file_path::AbstractString, mode::AbstractString)
    if !(mode == "w" || mode == "r")
        throw(IS.ConflictingInputsError("mode can only be 'w' or 'r'"))
    end

    if isfile(file_path) && mode == "w"
        throw(IS.ConflictingInputsError("$file_path already exists"))
    end

    file = HDF5.h5open(file_path, mode)
    if mode == "w"
        HDF5.create_group(file, HDF_SIMULATION_ROOT_PATH)
        @debug "Created store" file_path
    end

    datasets = OrderedDict{Symbol, ProblemDatasets}()
    cache = ResultCache()
    store = HdfSimulationStore(
        file,
        SimulationStoreParams(),
        datasets,
        Dict{Symbol, HDF5.Dataset}(),
        Dict{Symbol, Int}(),
        cache,
    )
    mode == "r" && _deserialize_attributes!(store)

    finalizer(_check_state, store)
    return store
end

"""
Construct and open an HdfSimulationStore.

When reading or writing results in a program you should use the method that accepts a
function in order to guarantee that the file handle gets closed.

# Examples

```julia
# Assumes a simulation has been executed in the './rts' directory with these parameters.
path = "./rts"
problem = :ED
var_name = :P__ThermalStandard
timestamp = DateTime("2020-01-01T05:00:00")
store = open_store(HdfSimulationStore, path)
df = PowerSimulations.read_result(DataFrame, store, problem, :variables, var_name, timestamp)
```
"""
function open_store(
    ::Type{HdfSimulationStore},
    directory::AbstractString,
    mode = "r",
    filename = HDF_FILENAME,
)
    return HdfSimulationStore(joinpath(directory, filename), mode)
end

function open_store(
    func::Function,
    ::Type{HdfSimulationStore},
    directory::AbstractString,
    mode = "r",
    filename = HDF_FILENAME,
)
    store = nothing
    try
        store = HdfSimulationStore(joinpath(directory, filename), mode)
        return func(store)
    finally
        !isnothing(store) && close(store)
    end
end

function Base.close(store::HdfSimulationStore)
    flush(store)
    HDF5.close(store.file)
    @debug "Close store file handle" store.file
end

function Base.isopen(store::HdfSimulationStore)
    return isnothing(store.file) ? false : HDF5.isopen(store.file)
end

function Base.flush(store::HdfSimulationStore)
    for (key, param_cache) in store.cache.data
        _flush_data!(param_cache, store, key, false)
    end

    flush(store.file)
    @debug "Flush store"
end

get_params(store::HdfSimulationStore) = store.params

"""
Return the problem names in order of execution.
"""
list_problems(store::HdfSimulationStore) = keys(store.datasets)

"""
Return the fields stored for the `problem` and `container_type` (duals/parameters/variables).
"""
function list_fields(store::HdfSimulationStore, problem::Symbol, container_type::Symbol)
    container = getfield(store.datasets[problem], container_type)
    return keys(container)
end

function write_optimizer_stats!(
    store::HdfSimulationStore,
    problem,
    stats::OptimizerStats,
    timestamp,  # Unused here. Matches the interface for InMemorySimulationStore.
)
    problem_name = Symbol(problem)
    dataset = _get_dataset(OptimizerStats, store, problem_name)

    # Uncomment for performance measures of HDF Store
    #TimerOutputs.@timeit RUN_SIMULATION_TIMER "Write optimizer stats" begin
    dataset[:, store.optimizer_stats_write_index[problem_name]] = to_array(stats)
    #end

    store.optimizer_stats_write_index[problem_name] += 1
    return
end

"""
Read the optimizer stats for a problem execution.
"""
function read_problem_optimizer_stats(
    store::HdfSimulationStore,
    simulation_step,
    problem,
    execution_index,
)
    optimizer_stats_write_index =
        (simulation_step - 1) * store.params.problems[problem].num_executions + index
    dataset = _get_dataset(OptimizerStats, store, problem)
    return OptimizerStats(dataset[:, optimizer_stats_write_index])
end

"""
Return the optimizer stats for a problem as a DataFrame.
"""
function read_problem_optimizer_stats(store::HdfSimulationStore, problem)
    dataset = _get_dataset(OptimizerStats, store, problem)
    data = permutedims(dataset[:, :])
    stats = [to_namedtuple(OptimizerStats(data[i, :])) for i in axes(data)[1]]
    return DataFrames.DataFrame(stats)
end

function initialize_problem_storage!(
    store::HdfSimulationStore,
    params,
    problem_reqs,
    flush_rules,
)
    store.params = params
    root = store.file[HDF_SIMULATION_ROOT_PATH]
    problems_group = _get_group_or_create(root, "problems")
    set_max_size!(store.cache, flush_rules.max_size)
    set_min_flush_size!(store.cache, flush_rules.min_flush_size)
    @debug "initialize_problem_storage" store.cache

    for problem in keys(store.params.problems)
        store.datasets[problem] = ProblemDatasets()
        problem_group = _get_group_or_create(problems_group, string(problem))
        for type in STORE_CONTAINERS
            group = _get_group_or_create(problem_group, string(type))
            for (name, reqs) in getfield(problem_reqs[problem], type)
                dataset = _create_dataset(group, string(name), reqs)
                # Columns can't be stored in attributes because they might be larger than
                # the max size of 64 KiB.
                col = _make_column_name(name)
                HDF5.write_dataset(group, col, string.(reqs["columns"]))
                column_dataset = group[col]
                datasets = getfield(store.datasets[problem], type)
                datasets[name] = Dataset(dataset, column_dataset)
                key = make_cache_key(problem, type, name)
                add_param_cache!(store.cache, key, get_rule(flush_rules, key))
            end
        end

        num_stats = params.num_steps * params.problems[problem].num_executions
        columns = fieldnames(PSI.OptimizerStats)
        num_columns = length(columns)
        dataset = HDF5.create_dataset(
            problem_group,
            OPTIMIZER_STATS_PATH,
            HDF5.datatype(Float64),
            HDF5.dataspace((num_columns, num_stats)),
        )
        HDF5.attributes(dataset)["columns"] = [string(x) for x in columns]
        store.optimizer_stats_datasets[problem] = dataset
        store.optimizer_stats_write_index[problem] = 1
        @debug "Initialized optimizer_stats_datasets $problem ($num_columns, $num_stats)"
    end

    # This has to run after problem groups are created.
    _serialize_attributes(store, problems_group, problem_reqs)
end

log_cache_hit_percentages(x::HdfSimulationStore) = log_cache_hit_percentages(x.cache)

"""
Return DataFrame, DenseAxisArray, or Array for a model result at a timestamp.
"""
function read_result(
    ::Type{DataFrames.DataFrame},
    store::HdfSimulationStore,
    problem_name,
    type,
    name,
    timestamp::Dates.DateTime,
)
    key = make_cache_key(problem_name, type, name)
    data, columns = _read_data_columns(store, key, timestamp)
    return DataFrames.DataFrame(data, Symbol.(columns))
end

function read_result(
    ::Type{JuMP.Containers.DenseAxisArray},
    store::HdfSimulationStore,
    problem_name,
    type,
    name,
    timestamp::Dates.DateTime,
)
    key = make_cache_key(problem_name, type, name)
    data, columns = _read_data_columns(store, key, timestamp)
    return JuMP.Containers.DenseAxisArray(data, size(data[1]), columns)
end

function read_result(
    ::Type{<:Array},
    store::HdfSimulationStore,
    problem_name,
    type,
    name,
    timestamp::Dates.DateTime,
)
    key = make_cache_key(problem_name, type, name)
    if is_cached(store.cache, key, timestamp)
        data = read_result(store.cache, key, timestamp)
    else
        # PERF: If this will be commonly used then we need to remove reading of columns.
        data, _ = read_result(store, key, timestamp)
    end
end

function read_result(store::HdfSimulationStore, key, timestamp::Dates.DateTime)
    simulation_step, execution_index = _get_indices(store, key.problem, timestamp)
    return read_result(store, key, simulation_step, execution_index)
end

function read_result(
    store::HdfSimulationStore,
    key,
    simulation_step::Int,
    execution_index::Int,
)
    @assert key.type in STORE_CONTAINERS "$(key.type)"

    !isopen(store) && throw(ArgumentError("store must be opened prior to reading"))

    horizon = store.params.problems[key.problem].horizon
    num_executions = store.params.problems[key.problem].num_executions
    if execution_index > num_executions
        throw(
            ArgumentError(
                "execution_index = $execution_index cannot be larger than $num_executions",
            ),
        )
    end

    dataset = _get_dataset(store, key)
    dset = dataset.dataset
    row_index = (simulation_step - 1) * num_executions + execution_index
    columns = dataset.column_dataset[:]

    TimerOutputs.@timeit RUN_SIMULATION_TIMER "Read dataset" begin
        num_dims = ndims(dset)
        if num_dims == 3
            data = dset[:, :, row_index]
            #elseif num_dims == 4
            #    data = dset[:, :, :, row_index]
        else
            error("unsupported dims: $num_dims")
        end
    end

    return data, columns
end

"""
Write a model result for a timestamp to the store.
"""
function write_result!(
    store::HdfSimulationStore,
    problem_name,
    container_type,
    name,
    timestamp,
    data,
    columns = nothing,  # Unused here. Matches the interface for InMemorySimulationStore.
)
    key = make_cache_key(problem_name, container_type, name)
    param_cache = get_param_cache(store.cache, key)

    cur_size = get_size(store.cache)
    add_result!(param_cache, timestamp, to_array(data), is_full(store.cache, cur_size))

    if get_dirty_size(param_cache) >= get_min_flush_size(store.cache)
        discard = !should_keep_in_cache(param_cache)

        # PERF: A potentially significant performance improvement would be to queue several
        # flushes and submit them in parallel.
        size_flushed = _flush_data!(param_cache, store, key, discard)

        @debug "flushed data" key size_flushed discard
    end

    # Disabled because this is currently a noop.
    #if is_full(store.cache)
    #    _flush_data!(store.cache, store)
    #end

    @debug "write_result" get_size(store.cache)
    return
end

function _check_state(store::HdfSimulationStore)
    if has_dirty(store.cache)
        error("BUG!!! dirty cache is present at shutdown: $(store.file)")
    end
end

function _compute_chunk_count(dims, dtype; max_chunk_bytes = DEFAULT_MAX_CHUNK_BYTES)
    bytes_per_element = sizeof(dtype)

    if length(dims) == 2
        size_row = dims[1] * bytes_per_element
    elseif length(dims) == 3
        size_row = dims[1] * dims[2] * bytes_per_element
    elseif length(dims) == 4
        size_row = dims[1] * dims[2] * dims[3] * bytes_per_element
    else
        error("unsupported dims = $dims")
    end

    chunk_count = minimum((trunc(max_chunk_bytes / size_row), dims[end]))
    if chunk_count == 0
        error(
            "HDF Max Chunk Bytes is smaller than the size of a row. Please increase it. " *
            "max_chunk_bytes=$max_chunk_bytes dims=$dims " *
            "size_row=$size_row",
        )
    end

    chunk_dims = [x for x in dims]
    chunk_dims[end] = chunk_count
    return chunk_dims
end

function _create_dataset(group, name, reqs)
    dataset = HDF5.create_dataset(
        group,
        string(name),
        HDF5.datatype(Float64),
        HDF5.dataspace(reqs["dims"]),
        # We are choosing to optimize read performance in the first implementation.
        # Compression would slow that down.
        #chunk = _compute_chunk_count(reqs["dims"], Float64),
        #shuffle = (),
        #deflate = 3,
    )
    @debug "Created dataset for" group name size(dataset)
    return dataset
end

function _deserialize_attributes!(store::HdfSimulationStore)
    group = store.file["simulation"]
    initial_time = Dates.DateTime(HDF5.read(HDF5.attributes(group)["initial_time"]))
    step_resolution =
        Dates.Millisecond(HDF5.read(HDF5.attributes(group)["step_resolution_ms"]))
    num_steps = HDF5.read(HDF5.attributes(group)["num_steps"])
    store.params = SimulationStoreParams(initial_time, step_resolution, num_steps)
    empty!(store.datasets)
    for problem in HDF5.read(HDF5.attributes(group)["problem_order"])
        problem_group = store.file["simulation/problems/$problem"]
        problem_name = Symbol(problem)
        store.params.problems[problem_name] = SimulationStoreProblemParams(
            HDF5.read(HDF5.attributes(problem_group)["num_executions"]),
            HDF5.read(HDF5.attributes(problem_group)["horizon"]),
            Dates.Millisecond(HDF5.read(HDF5.attributes(problem_group)["interval_ms"])),
            Dates.Millisecond(HDF5.read(HDF5.attributes(problem_group)["resolution_ms"])),
            HDF5.read(HDF5.attributes(problem_group)["end_of_interval_step"]),
            HDF5.read(HDF5.attributes(problem_group)["base_power"]),
            Base.UUID(HDF5.read(HDF5.attributes(problem_group)["system_uuid"])),
        )
        store.datasets[problem_name] = ProblemDatasets()
        for type in STORE_CONTAINERS
            group = problem_group[string(type)]
            for name in keys(group)
                if !endswith(name, "columns")
                    dataset = group[name]
                    column_dataset = group[_make_column_name(name)]
                    item = Dataset(dataset, column_dataset)
                    getfield(store.datasets[problem_name], type)[Symbol(name)] = item
                end
                key = make_cache_key(problem_name, type, Symbol(name))
                add_param_cache!(store.cache, key, CacheFlushRule())
            end
        end

        store.optimizer_stats_datasets[problem_name] = problem_group[OPTIMIZER_STATS_PATH]
        store.optimizer_stats_write_index[problem_name] = 0
    end

    @debug "deserialized store params and datasets" store.params
end

function _serialize_attributes(store::HdfSimulationStore, problems_group, problem_reqs)
    params = store.params
    group = store.file["simulation"]
    HDF5.attributes(group)["problem_order"] = [string(k) for k in keys(params.problems)]
    HDF5.attributes(group)["initial_time"] = string(params.initial_time)
    HDF5.attributes(group)["step_resolution_ms"] =
        Dates.Millisecond(params.step_resolution).value
    HDF5.attributes(group)["num_steps"] = params.num_steps

    for problem in keys(params.problems)
        problem_group = store.file["simulation/problems/$problem"]
        HDF5.attributes(problem_group)["num_executions"] =
            params.problems[problem].num_executions
        HDF5.attributes(problem_group)["horizon"] = params.problems[problem].horizon
        HDF5.attributes(problem_group)["resolution_ms"] =
            Dates.Millisecond(params.problems[problem].resolution).value
        HDF5.attributes(problem_group)["end_of_interval_step"] =
            params.problems[problem].end_of_interval_step
        HDF5.attributes(problem_group)["interval_ms"] =
            Dates.Millisecond(params.problems[problem].interval).value
        HDF5.attributes(problem_group)["base_power"] = params.problems[problem].base_power
        HDF5.attributes(problem_group)["system_uuid"] =
            string(params.problems[problem].system_uuid)
    end
end

function _flush_data!(cache::ParamResultCache, store::HdfSimulationStore, key, discard)
    !has_dirty(cache) && return 0
    dataset = _get_dataset(store, key)
    timestamps, data = get_data_to_flush!(cache, get_min_flush_size(store.cache))
    num_results = length(timestamps)
    @assert_op num_results == size(data)[end]
    end_index = dataset.write_index + length(timestamps) - 1
    write_range = (dataset.write_index):end_index
    TimerOutputs.@timeit RUN_SIMULATION_TIMER "Write $(key.type) array to HDF" begin
        _write_dataset!(dataset.dataset, data, write_range)
    end

    discard && discard_results!(cache, timestamps)

    dataset.write_index += num_results
    size_flushed = cache.size_per_entry * num_results

    @debug "Flushed cache results to HDF5" key size_flushed
    return size_flushed
end

function _flush_data!(cache::ResultCache, store::HdfSimulationStore)
    # PERF: may need to optimize memory management
    # Do we need flush down to some ~70% watermark?
    # What are GC implications of doing replacing entries one at a time vs freeing large
    # chunks at once?
end

function _get_columns_dataset(store::HdfSimulationStore, key)
    return getfield(store.datasets[key.problem], key.type)[key.name].columns
end

function _get_dataset(::Type{OptimizerStats}, store::HdfSimulationStore, problem_name)
    return store.optimizer_stats_datasets[problem_name]
end

function _get_dataset(store::HdfSimulationStore, problem_name::Symbol)
    return store.datasets[problem_name]
end

function _get_dataset(store::HdfSimulationStore, key)
    return getfield(store.datasets[key.problem], key.type)[key.name]
end

function _get_group_or_create(parent, group_name)
    if haskey(parent, group_name)
        group = parent[group_name]
    else
        group = HDF5.create_group(parent, group_name)
        @debug "Created group" group
    end

    return group
end

_make_column_name(name) = string(name) * "__columns"

function _get_indices(store::HdfSimulationStore, problem, timestamp)
    time_diff = Dates.Millisecond(timestamp - store.params.initial_time)
    step = time_diff ÷ store.params.step_resolution + 1
    if step > store.params.num_steps
        throw(
            ArgumentError("timestamp = $timestamp is beyond the simulation: step = $step"),
        )
    end
    problem_params = store.params.problems[problem]
    initial_time = store.params.initial_time + (step - 1) * store.params.step_resolution
    time_diff = timestamp - initial_time
    if time_diff % problem_params.interval != Dates.Millisecond(0)
        throw(ArgumentError("timestamp = $timestamp is not a valid problem timestamp"))
    end
    execution_index = time_diff ÷ problem_params.interval + 1
    return step, execution_index
end

_get_root(store::HdfSimulationStore) = store.file[HDF_SIMULATION_ROOT_PATH]

function _read_column_names(::Type{OptimizerStats}, store::HdfSimulationStore)
    dataset = _get_dataset(OptimizerStats, store)
    return HDF5.read(HDF5.attributes(dataset), "columns")
end

function _read_data_columns(store, key, timestamp)
    if is_cached(store.cache, key, timestamp)
        data = read_result(store.cache, key, timestamp)
        column_dataset = _get_dataset(store, key).column_dataset
        columns = column_dataset[:]
    else
        data, columns = read_result(store, key, timestamp)
    end

    return data, columns
end

function _read_length(::Type{OptimizerStats}, store::HdfSimulationStore)
    dataset = _get_dataset(OptimizerStats, store)
    return HDF5.read(HDF5.attributes(dataset), "columns")
end

function _write_dataset!(dataset, array, row_range)
    if ndims(array) == 2
        dataset[:, 1, row_range] = array
    elseif ndims(array) == 3
        dataset[:, :, row_range] = array
        #elseif ndims(array) == 4
        #    dataset[:, :, :, row_range] = array
    else
        error("ndims not supported: $(ndims(array))")
    end

    @debug "wrote dataset" dataset row_range
end


text:
using Test
using BenchmarkTools

# Set up problem
model, obj0 = Dynamics.cartpole_analytical
n,m = model.n, model.m

obj = copy(obj0)
obj.x0 = [0;0;0;0.]
obj.xf = [0.5;pi;0;0]
obj.tf = 2.
u_bnd = 50
x_bnd = [0.6,Inf,Inf,Inf]
obj_con = ConstrainedObjective(obj,u_min=-u_bnd, u_max=u_bnd, x_min=-x_bnd, x_max=x_bnd)
obj_con = ConstrainedObjective(obj,u_min=-u_bnd, u_max=u_bnd)
dt = 0.1

# Initialize trajectory
solver = Solver(model,obj,dt=dt,integration=:rk3_foh)
n,m,N = get_sizes(solver)
U0 = ones(1,N)*1
X0 = line_trajectory(obj.x0, obj.xf, N)

# Get results structures
obj_con = to_static(obj_con)
obj = to_static(obj)

#######################
#    BACKWARDS PASS   #
#######################

solver = Solver(model,obj,dt=dt,integration=:rk3)
res = UnconstrainedResults(n,m,N)
reS = UnconstrainedStaticResults(n,m,N)
reV = UnconstrainedVectorResults(n,m,N)

solver = Solver(model,obj_con,dt=dt,integration=:rk3)
res = ConstrainedResults(n,m,obj_con.p,N)
reS = ConstrainedStaticResults(n,m,obj_con.p,N)
reV = ConstrainedVectorResults(n,m,obj_con.p,N)

# println("Default")
rollout!(res,solver)
J_prev1 = cost(solver, res)
TrajectoryOptimization.update_jacobians!(res, solver)
update_constraints!(res, solver)
# println("Static")
rollout!(reS,solver)
J_prev2 = cost(solver, reS)
TrajectoryOptimization.update_jacobians!(reS, solver)
update_constraints!(reS, solver)
# println("Vecs")
rollout!(reV,solver)
J_prev3 = _cost(solver, reV)
TrajectoryOptimization.update_jacobians!(reV, solver)
update_constraints!(reV, solver)
to_array(reV.X) == to_array(reS.X) == res.X
to_array(reV.fx) == to_array(reS.fx) == res.fx
if res isa ConstrainedResults
    to_array(reS.C) == res.C == to_array(reV.C)
end

# Zero order hold
println("Original")
@btime Δv1 = backwardpass!(res, solver)

println("Vector Results")
@btime Δv3 = backwardpass!(reV, solver)
@btime Δv5 = backwardpass_new!(reV, solver)
# @btime Δv5 = backwardpass_new_ts!(reV, solver)

println("Static Results")
@btime Δv2 = backwardpass!(reS, solver)
@btime Δv4 = backwardpass_new!(reS, solver)
# @btime Δv4 = backwardpass_foh_new_ts!(reS, solver)

# First order hold
solver = Solver(model,obj,dt=dt,integration=:rk3_foh)
res = UnconstrainedResults(n,m,N)
reS = UnconstrainedStaticResults(n,m,N)
reV = UnconstrainedVectorResults(n,m,N)

solver = Solver(model,obj_con,dt=dt,integration=:rk3_foh)
res = ConstrainedResults(n,m,obj_con.p,N)
reS = ConstrainedStaticResults(n,m,obj_con.p,N)
reV = ConstrainedVectorResults(n,m,obj_con.p,N)

# println("Default")
rollout!(res,solver)
J_prev1 = cost(solver, res)
TrajectoryOptimization.update_jacobians!(res, solver)
update_constraints!(res, solver)
# println("Static")
rollout!(reS,solver)
J_prev2 = cost(solver, reS)
TrajectoryOptimization.update_jacobians!(reS, solver)
update_constraints!(reS, solver)
# println("Vecs")
rollout!(reV,solver)
J_prev3 = _cost(solver, reV)
TrajectoryOptimization.update_jacobians!(reV, solver)
update_constraints!(reV, solver)
to_array(reV.X) == to_array(reS.X) == res.X
to_array(reV.fx) == to_array(reS.fx) == res.fx
if res isa ConstrainedResults
    to_array(reS.C) == res.C == to_array(reV.C)
end

println("Original")
@btime Δv1 = backwardpass_foh!(res, solver)

println("Vector Results")
@btime Δv3 = backwardpass_foh!(reV, solver)
@btime Δv5 = backwardpass_foh_new!(reV, solver)
@btime Δv5 = backwardpass_foh_new_ts!(reV, solver)

println("Static Results")
@btime Δv2 = backwardpass_foh!(reS, solver)
@btime Δv4 = backwardpass_foh_new!(reS, solver)
@btime Δv4 = backwardpass_foh_new_ts!(reS, solver)

# Type Stability checks

# Zero order hold
Δv1 = backwardpass!(res, solver)
Δv2 = backwardpass!(reV, solver)
Δv3 = backwardpass!(reS, solver)
Δv4 = backwardpass_new!(reV, solver)
Δv5 = backwardpass_new!(reS, solver)
Δv1 ≈ Δv2 ≈ Δv3 ≈ Δv4 ≈ Δv5

@code_warntype backwardpass!(reS, solver)
@code_warntype backwardpass!(reV, solver)

@code_warntype backwardpass_new!(reS, solver)
@code_warntype backwardpass_new!(reV, solver)

Δv1 = backwardpass_foh!(res, solver)
Δv2 = backwardpass_foh_new!(reV, solver)
Δv3 = backwardpass_foh_new_ts!(reV, solver)
Δv4 = backwardpass_foh_new!(reS, solver)
Δv5 = backwardpass_foh_new_ts!(reS, solver)
Δv1 ≈ Δv2 ≈ Δv3 ≈ Δv4 ≈ Δv5

@code_warntype backwardpass_foh!(reS, solver)
@code_warntype backwardpass_foh!(reV, solver)

@code_warntype backwardpass_foh_new!(reS, solver)
@code_warntype backwardpass_foh_new!(reV, solver)

@code_warntype backwardpass_foh_new_ts!(reS, solver)
@code_warntype backwardpass_foh_new_ts!(reV, solver)

function backwardpass_new!(res::SolverVectorResults,solver::Solver)
    N = solver.N; n = solver.model.n; m = solver.model.m;
    Q = solver.obj.Q; Qf = solver.obj.Qf; xf = solver.obj.xf;
    R = getR(solver)
    dt = solver.dt

    if solver.model.m != length(res.U[1])
        m += n
    end

    # pull out values from results
    X = res.X; U = res.U; K = res.K; d = res.d; S = res.S; s = res.s

    # Initialize expected change in cost-to-go
    S[N],s[N] = terminal_step(res,solver)
    Δv = [0.0 0.0]

    k = N-1

    # Backward pass
    while k >= 1
        fx, fu = res.fx[k], res.fu[k]
        if res isa ConstrainedIterResults
            Qfun = cost_to_go_expansion(S[k+1],s[k+1],X[k],U[k],fx,fu,Q,R,xf,dt,res.ρ[1],
                res.C[k],res.Cx[k],res.Cu[k],res.LAMBDA[k],res.Iμ[k])
        else
            Qfun = cost_to_go_expansion(S[k+1],s[k+1],X[k],U[k],fx,fu,Q,R,xf,dt,res.ρ[1])
        end

        K[k],d[k],S[k],s[k],Δv_,reg = calculate_optimal_gains(Qfun)
        Δv += Δv_

        if reg
            if solver.opts.verbose
                println("regularized (normal bp)")
            end
            regularization_update!(res,solver,true)
            k = N-1
            Δv = [0.0 0.0]
            S[N],s[N] = terminal_step(res,solver)
            continue
        end

        k = k - 1;
    end

    return Δv
end

function backwardpass_foh_new!(res::SolverVectorResults,solver::Solver)
    n,m,N = get_sizes(solver)
    dt = solver.dt

    if solver.model.m != length(res.U[1])
        m += n
    end

    Q = solver.obj.Q
    Qf = solver.obj.Qf
    xf = solver.obj.xf
    R = getR(solver)

    K = res.K
    b = res.b
    d = res.d

    X,U = res.X, res.U

    # Initialization
    S,s = terminal_step(res,solver)
    Δv = [0. 0.]

    k = N-1
    while k >= 1

        # Extract Variables from results
        Ac1, Bc1 = res.Ac[k], res.Bc[k]
        Ac2, Bc2 = res.Ac[k+1], res.Bc[k+1]
        Ad, Bd, Cd = res.fx[k], res.fu[k], res.fv[k]

        xdot1 = res.xdot[k]
        xdot2 = res.xdot[k+1]

        # Calculate the Midpoint
        x,u = X[k], U[k]
        xm = 0.5*X[k] + dt/8.0*xdot1 + 0.5*X[k+1] - dt/8.0*xdot2
        um = (U[k] + U[k+1])/2.0
        y,v = X[k+1],U[k+1]

        # Get second order Taylor series of the Lagrangian
        if res isa ConstrainedIterResults
            Qfun = cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,res.ρ[1],
                res.C[k+1],res.Cx[k+1],res.Cu[k+1],res.LAMBDA[k+1],res.Iμ[k+1])
        else
            Qfun = cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,res.ρ[1])
        end

        # Calculate the optimal control gains
        K[k+1],b[k+1],d[k+1],Δv_,reg = calculate_optimal_gains!(S,s,Qfun,res.ρ[1])
        Δv += Δv_

        # Check regularization flag
        if reg
            if solver.opts.verbose
                println("*NOT implemented* regularized (foh bp)")
            end
            regularization_update!(res,solver,true)
            k = N-1
            S,s = terminal_step(res,solver)
            Δv = [0. 0.]
            continue
        end

        # at last time step, optimize over final control
        if k == 1
            # Calculate gains at first time step
            if res isa ConstrainedIterResults
                K[k],b[k],d[k],Δv_,reg = calculate_optimal_gains!(S,s,res.C[k],res.Cx[k],res.Cu[k],res.LAMBDA[k],res.Iμ[k])
            else
                K[k],b[k],d[k],Δv_,reg = calculate_optimal_gains!(S,s)
            end

            # Check regularization flag
            if reg
                if solver.opts.verbose
                    println("regularized (foh bp)")
                end
                regularization_update!(res,solver::Solver,true)
                k = N-1
                S,s = terminal_step(res,solver)
                Δv = [0. 0.]
                continue
            end

            res.s[1] = s[1:n]
            Δv += Δv_
        end

        k = k - 1;
    end

    return Δv
end

function backwardpass_foh_new_ts!(res,solver)
    n,m,N = get_sizes(solver)
    dt = solver.dt

    if solver.model.m != length(res.U[1])
        m += n
    end

    Q = solver.obj.Q
    Qf = solver.obj.Qf
    xf = solver.obj.xf
    R = getR(solver)

    K = res.K
    b = res.b
    d = res.d

    X,U = res.X, res.U

    # Initialization
    S,s = terminal_step(res,solver)
    Δv = [0. 0.]

    k = N-1
    while k >= 1
        # Extract Variables from results
        Ac1, Bc1 = res.Ac[k], res.Bc[k]
        Ac2, Bc2 = res.Ac[k+1], res.Bc[k+1]
        Ad, Bd, Cd = res.fx[k], res.fu[k], res.fv[k]

        xdot1 = res.xdot[k]
        xdot2 = res.xdot[k+1]

        # Calculate the Midpoint
        x,u = X[k], U[k]
        xm = 0.5*X[k] + dt/8.0*xdot1 + 0.5*X[k+1] - dt/8.0*xdot2
        um = (U[k] + U[k+1])/2.0
        y,v = X[k+1],U[k+1]

        # Get second order Taylor series of the Lagrangian
        if res isa ConstrainedIterResults
            Qfun = cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,res.ρ[1],
            res.C[k+1],res.Cx[k+1],res.Cu[k+1],res.LAMBDA[k+1],res.Iμ[k+1])
        else
            Qfun = cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,res.ρ[1])
        end

        # Calculate the optimal control gains
        K[k+1],b[k+1],d[k+1],Δv_,reg = calculate_optimal_gains!(S,s,Qfun,res.ρ[1])
        Δv += Δv_

        # Check regularization flag
        if reg
            if solver.opts.verbose
                println("*NOT implemented* regularized (foh bp)")
            end
            regularization_update!(res,solver,true)
            k = N-1
            S,s = terminal_step(res,solver)
            Δv = [0. 0.]
            continue
        end

        # at last time step, optimize over final control
        if k == 1
            # Calculate gains at first time step
            if res isa ConstrainedIterResults
                Qbar = S_to_Qbar(S,s,res.C[k],res.Cx[k],res.Cu[k],res.LAMBDA[k],res.Iμ[k])
            else
                Qbar = S_to_Qbar(S,s,n,m)
            end
            K[k],b[k],d[k],Δv_,reg = calculate_optimal_gains!(S,s,Qbar)

            # Check regularization flag
            if reg
                if solver.opts.verbose
                    println("regularized (foh bp)")
                end
                regularization_update!(res,solver::Solver,true)
                k = N-1
                S,s = terminal_step(res,solver)
                Δv = [0. 0.]
                continue
            end

            res.s[1] = s[1:n]
            Δv += Δv_
        end

        k = k - 1;

    end

    return Δv
end

"""
$(SIGNATURES)
Calculate the cost to go at the terminal time step
"""
function terminal_step(res::SolverVectorResults,solver::Solver)::Tuple{Matrix{Float64},Vector{Float64}}
    n,m,N = get_sizes(solver)
    Qf = solver.obj.Qf
    xf = solver.obj.xf
    if solver.control_integration == :foh
        # Initialization
        S = zeros(n+m,n+m)
        s = zeros(n+m)

        # Boundary conditions
        S[1:n,1:n] = Qf
        s[1:n] = Qf*(res.X[N]-xf)

        # Terminal constraints
        if res isa ConstrainedIterResults
            C = res.C; Iμ = res.Iμ; LAMBDA = res.LAMBDA
            CxN = res.Cx_N
            S[1:n,1:n] += CxN'*res.IμN*CxN
            s[1:n] += CxN'*res.IμN*res.CN + CxN'*res.λN
        end
    elseif solver.control_integration == :zoh

        # Boundary Conditions
        S = Array(Qf)
        s = Array(Qf*(res.X[N] - xf))

        # Terminal constraints
        if res isa ConstrainedIterResults
            C = res.C; Iμ = res.Iμ; LAMBDA = res.LAMBDA
            CxN = res.Cx_N
            S += CxN'*res.IμN*CxN
            s += CxN'*res.IμN*res.CN + CxN'*res.λN
        end
    end
    return S,s
end

#########################################
#        COST TO GO EXPANSIONS          #
#########################################
"""
$(SIGNATURES)
Perform a 2nd order taylor series expansion of the the stage cost for first
    order hold
"""
function _lagrangian_2nd_order(S::Matrix{Float64},s::Vector{Float64},x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ)
    # Expansion of stage cost L(x,u,y,v) -> dL(dx,du,dy,dv)
    Lx = dt/6*Q*(x - xf) + 4*dt/6*(I/2 + dt/8*Ac1)'*Q*(xm - xf)
    Lu = dt/6*R*u + 4*dt/6*((dt/8*Bc1)'*Q*(xm - xf) + 0.5*R*um)
    Ly = dt/6*Q*(y - xf) + 4*dt/6*(I/2 - dt/8*Ac2)'*Q*(xm - xf)
    Lv = dt/6*R*v + 4*dt/6*((-dt/8*Bc2)'*Q*(xm - xf) + 0.5*R*um)

    Lxx = dt/6*Q + 4*dt/6*(I/2 + dt/8*Ac1)'*Q*(I/2 + dt/8*Ac1)
    Luu = dt/6*R + 4*dt/6*((dt/8*Bc1)'*Q*(dt/8*Bc1) + 0.5*R*0.5)
    Lyy = dt/6*Q + 4*dt/6*(I/2 - dt/8*Ac2)'*Q*(I/2 - dt/8*Ac2)
    Lvv = dt/6*R + 4*dt/6*((-dt/8*Bc2)'*Q*(-dt/8*Bc2) + 0.5*R*0.5)

    Lxu = 4*dt/6*(I/2 + dt/8*Ac1)'*Q*(dt/8*Bc1)
    Lxy = 4*dt/6*(I/2 + dt/8*Ac1)'*Q*(I/2 - dt/8*Ac2)
    Lxv = 4*dt/6*(I/2 + dt/8*Ac1)'*Q*(-dt/8*Bc2)
    Luy = 4*dt/6*(dt/8*Bc1)'*Q*(I/2 - dt/8*Ac2)
    Luv = 4*dt/6*((dt/8*Bc1)'*Q*(-dt/8*Bc2) + 0.5*R*0.5)
    Lyv = 4*dt/6*(I/2 - dt/8*Ac2)'*Q*(-dt/8*Bc2)

    # Unpack cost-to-go P, then add L + P
    Sy = s[1:n]
    Sv = s[n+1:n+m]
    Syy = S[1:n,1:n]
    Svv = S[n+1:n+m,n+1:n+m]
    Syv = S[1:n,n+1:n+m]

    Ly += Sy
    Lv += Sv
    Lyy += Syy
    Lvv += Svv
    Lyv += Syv

    TX,TU = typeof(Lx),typeof(Lu)
    TXX,TUU,TXU = typeof(Lxx),typeof(Luu),typeof(Lxu)

    return Lx,Lu,Ly::TX,Lv::TU, Lxx,Luu,Lyy::TXX,Lvv::TUU,Lxu,Lxy,Lxv,Luy,Luv,Lyv::TXU
end

"""
$(SIGNATURES)
Find the 2nd order taylor series expansion of the cost to go at the current
    time step
"""
# ZERO ORDER HOLD
function cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ)
    Qx,Qu,Qxx,Quu,Qux = _cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ)
    return (x=Qx,u=Qu, xx=Qxx,uu=Quu,ux=Qux)
end

# (with constraints)
function cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ,
        C,Cx,Cu,λ,Iμ)
    Qx,Qu,Qxx,Quu,Qux = _cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ)

    Qx +=  Cx'Iμ*C + Cx'λ
    Qu +=  Cu'Iμ*C + Cu'λ
    Qxx += Cx'Iμ*Cx
    Quu += Cu'Iμ*Cu
    Qux += Cu'Iμ*Cx

    return (x=Qx,u=Qu, xx=Qxx,uu=Quu,ux=Qux)
end

function _cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ)
    lx = dt*Q*(x - xf)
    lu = dt*R*u
    lxx = dt*Q
    luu = dt*R

    # Gradients and Hessians of Taylor Series Expansion of Q
    Qx = lx + fx's
    Qu = lu + fu's
    Qxx = lxx + fx'S*fx

    Quu = luu + fu'S*fu + ρ*I
    Qux = fu'S*fx

    return Qx,Qu,Qxx,Quu,Qux
end

# FIRST ORDER HOLD
function cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ)

    Lx,Lu,Ly,Lv, Lxx,Luu,Lyy,Lvv, Lxu,Lxy,Lxv,Luy,Luv,Lyv = _lagrangian_2nd_order(S,s,
        x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ)

    # Substitute in discrete dynamics dx = (Ad)dx + (Bd)du1 + (Cd)du2
    Qx = vec(Lx) + Ad'*vec(Ly)
    Qu = vec(Lu) + Bd'*vec(Ly)
    Qv = vec(Lv) + Cd'*vec(Ly)

    Qxx = Lxx + Lxy*Ad + Ad'Lxy' + Ad'Lyy*Ad
    Quu = Luu + Luy*Bd + Bd'Luy' + Bd'Lyy*Bd + ρ*I
    Qvv = Lvv + Lyv'Cd + Cd'Lyv  + Cd'Lyy*Cd + ρ*I
    Qxu = Lxu + Lxy*Bd + Ad'Luy' + Ad'Lyy*Bd
    Qxv = Lxv + Lxy*Cd + Ad'Lyv  + Ad'Lyy*Cd
    Quv = Luv + Luy*Cd + Bd'Lyv  + Bd'Lyy*Cd

    # return (x=Qx,u=Qu,v=Qv, xx=Qxx,uu=Quu,vv=Qvv, xu=Qxu,xv=Qxv,uv=Quv)
    return Qx,Qu,Qv, Qxx,Quu,Qvv, Qxu,Qxv,Quv
end

# (with constraints)
function cost_to_go_expansion(S,s,x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ,
        C,Cy,Cv,LAMBDA,Iμ)

    Lx,Lu,Ly,Lv, Lxx,Luu,Lyy,Lvv, Lxu,Lxy,Lxv,Luy,Luv,Lyv = _lagrangian_2nd_order(S,s,
        x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ)

    Ly += (Cy'*Iμ*C + Cy'*LAMBDA)
    Lv += (Cv'*Iμ*C + Cv'*LAMBDA)
    Lyy += Cy'*Iμ*Cy
    Lvv += Cv'*Iμ*Cv
    Lyv += Cy'*Iμ*Cv

    # Substitute in discrete dynamics dx = (Ad)dx + (Bd)du1 + (Cd)du2
    Qx = vec(Lx) + Ad'*vec(Ly)
    Qu = vec(Lu) + Bd'*vec(Ly)
    Qv = vec(Lv) + Cd'*vec(Ly)

    Qxx = Lxx + Lxy*Ad + Ad'Lxy' + Ad'Lyy*Ad
    Quu = Luu + Luy*Bd + Bd'Luy' + Bd'Lyy*Bd + ρ*I
    Qvv = Lvv + Lyv'Cd + Cd'Lyv  + Cd'Lyy*Cd + ρ*I
    Qxu = Lxu + Lxy*Bd + Ad'Luy' + Ad'Lyy*Bd
    Qxv = Lxv + Lxy*Cd + Ad'Lyv  + Ad'Lyy*Cd
    Quv = Luv + Luy*Cd + Bd'Lyv  + Bd'Lyy*Cd

    # return (x=Qx,u=Qu,v=Qv, xx=Qxx,uu=Quu,vv=Qvv, xu=Qxu,xv=Qxv,uv=Quv)
    return Qx,Qu,Qv, Qxx,Quu,Qvv, Qxu,Qxv,Quv
end

#############################################
#     CALCULATE OPTIMAL CONTROL GAINS       #
#############################################
"""
$(SIGNATURES)
Calculate the optimal control gains given the 2nd order expansion of the cost-to-go
"""
# FIRST ORDER HOLD
function calculate_optimal_gains(Q::NamedTuple)
    # regularization
    if !isposdef(Hermitian(Array(Q.uu)))  # need to wrap Array since isposdef doesn't work for static arrays
        K = zeros(Q.ux)
        d = zeros(Q.u)
        S = zeros(Q.xx)
        s = zeros(Q.x)

        reg = true
    else
        # Compute gains
        K = Q.uu\Q.ux
        d = Q.uu\Q.u
        S = Q.xx - Q.ux'K
        s = Q.x - Q.ux'd

        reg = false
    end

    Δv = [vec(Q.u)'*vec(d) 0.5*vec(d)'*Q.uu*vec(d)]
    return K,d,S,s,Δv,reg
end

# SECOND ORDER HOLD
function calculate_optimal_gains!(S,s,Q::Tuple,ρ::Float64)
    Qx,Qu,Qv, Qxx,Quu,Qvv, Qxu,Qxv,Quv = Q
    n = length(Qx)
    m = length(Qu)

    # Qvv = Hermitian(Qvv)
    # regularization
    if !isposdef(Array(Qvv))
        K = zeros(Qxv')
        b = zeros(Quv')
        d = zeros(Qv)
        Δv = [0. 0.]
        reg = true
        return K,b,d,Δv,reg
    else
        reg = false
    end

    K = -Qvv\Qxv'
    b = -Qvv\Quv'
    d = -Qvv\vec(Qv)

    Qx_ = vec(Qx) + K'*vec(Qv) + Qxv*vec(d) + K'Qvv*d
    Qu_ = vec(Qu) + b'*vec(Qv) + Quv*vec(d) + b'*Qvv*d
    Qxx_ = Qxx + Qxv*K + K'*Qxv' + K'*Qvv*K
    Quu_ = Quu + Quv*b + b'*Quv' + b'*Qvv*b + ρ*I
    Qxu_ = Qxu + K'*Quv' + Qxv*b + K'*Qvv*b

    # cache (approximate) cost-to-go at timestep k
    s[1:n] = Qx_
    s[n+1:n+m] = Qu_
    S[1:n,1:n] = Qxx_
    S[n+1:n+m,n+1:n+m] = Quu_
    S[1:n,n+1:n+m] = Qxu_
    S[n+1:n+m,1:n] = Qxu_'

    # line search terms
    Δv = [-vec(Qv)'vec(d) 0.5*vec(d)'Qvv*vec(d)]

    return K,b,d,Δv,reg
end

# (Final time step)
function calculate_optimal_gains!(S,s,Q::NTuple{4,Any})
    Qx_,Qu_,Quu_,Qxu_ = Q
    n = length(Qx_)
    m = length(Qu_)

    if !isposdef(Array(Quu_))
        K = zero(Qxu_')
        b = zero(Quu_)
        d = zero(Qu_)
        Δv = [0. 0.]
        reg = true
        return K,b,d,Δv,reg
    else
        reg = false
    end

    K = -Quu_\Qxu_'
    b = zero(Quu_)
    d = -Quu_\vec(Qu_)

    s[1:n] = Qx_ + Qxu_*vec(d)

    Δv = [-vec(Qu_)'*vec(d) 0.5*vec(d)'*Quu_*vec(d)]

    return K,b,d,Δv,reg
end

"""
$(SIGNATURES)
Calculate the 2nd order expansion of the final (k=1) cost-to-go for first order hold case
"""
function S_to_Qbar(S::Matrix{Float64},s::Vector{Float64},n::Int,m::Int)
    Qx_ = s[1:n]
    Qu_ = s[n+1:n+m]
    Quu_ = S[n+1:n+m,n+1:n+m]
    Qxu_ = S[n+1:n+m,1:n]'
    return Qx_,Qu_,Quu_,Qxu_
end

# (with constraints)
function S_to_Qbar(S::Matrix{Float64},s::Vector{Float64},C,Cx,Cu,LAMBDA,Iμ)
    n = size(Cx,2)
    m = size(Cu,2)

    Qx_ = s[1:n]
    Qu_ = s[n+1:n+m]
    Quu_ = S[n+1:n+m,n+1:n+m]
    Qxu_ = S[n+1:n+m,1:n]'

    Qx_ += (Cx'*Iμ*C + Cx'*LAMBDA)
    Qu_ += (Cu'*Iμ*C + Cu'*LAMBDA)
    Quu_ += Cu'*Iμ*Cu
    Qxu_ += Cx'*Iμ*Cu

    return Qx_,Qu_,Quu_,Qxu_
end

# function cost_to_go_expansion(S,s,x::MVector{N,Float64},u::MVector{M,Float64},xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ) where {N,M}
#     println("Static version")
#     Lx,Lu,Ly,Lv, Lxx,Luu,Lyy,Lvv, Lxu,Lxy,Lxv,Luy,Luv,Lyv = _lagrangian_2nd_order(S,s,
#         x,u,xm,um,y,v,Ac1,Bc1,Ac2,Bc2,Ad,Bd,Cd,Q,R,xf,dt,ρ)
#
#     # Substitute in discrete dynamics dx = (Ad)dx + (Bd)du1 + (Cd)du2
#     Qx = vec(Lx) + Ad'*vec(Ly)
#     Qu = vec(Lu) + Bd'*vec(Ly)
#     Qv = vec(Lv) + Cd'*vec(Ly)
#
#     Qxx = Lxx + Lxy*Ad + Ad'Lxy' + Ad'Lyy*Ad
#     Quu = Luu + Luy*Bd + Bd'Luy' + Bd'Lyy*Bd + ρ*I
#     Qvv = Lvv + Lyv'Cd + Cd'Lyv  + Cd'Lyy*Cd + ρ*I
#     Qxu = Lxu + Lxy*Bd + Ad'Luy' + Ad'Lyy*Bd
#     Qxv = Lxv + Lxy*Cd + Ad'Lyv  + Ad'Lyy*Cd
#     Quv = Luv + Luy*Cd + Bd'Lyv  + Bd'Lyy*Cd
#
#     # return (x=Qx,u=Qu,v=Qv, xx=Qxx,uu=Quu,vv=Qvv, xu=Qxu,xv=Qxv,uv=Quv)
#     return Qx::SVector{N,Float64},Qu::SVector{M,Float64},Qv::SVector{M,Float64},
#         Qxx::SMatrix{N,N,Float64},Quu::SMatrix{M,M,Float64},Qvv::SMatrix{M,M,Float64}, Qxu::SMatrix{N,M,Float64},Qxv::SMatrix{N,M,Float64},Quv::SMatrix{M,M,Float64}
# end

# function cost_to_go_expansion(S,s,x::MVector{N,Float64},u::MVector{M,Float64},fx,fu,Q,R,xf,dt,ρ) where {N,M}
#     Qx,Qu,Qxx,Quu,Qux = _cost_to_go_expansion(S,s,x,u,fx,fu,Q,R,xf,dt,ρ)
#     return (x=Qx,u=Qu, xx=Qxx,uu=Quu,ux=Qux)
# end

# function calculate_optimal_gains!(S,s,Q::NamedTuple,ρ::Float64)
#     n = length(Q.x)
#     m = length(Q.u)
#
#     # Qvv = Hermitian(Qvv)
#     # regularization
#     if !isposdef(Array(Q.vv))
#         K = zeros(Q.xv')
#         b = zeros(Q.uv')
#         d = zeros(Q.v)
#         Δv = [0. 0.]
#         reg = true
#         return K,b,d,Δv,reg
#     else
#         reg = false
#     end
#
#     K = -Q.vv\Q.xv'
#     b = -Q.vv\Q.uv'
#     d = -Q.vv\vec(Q.v)
#
#     Qx_ = vec(Q.x) + K'*vec(Q.v) + Q.xv*vec(d) + K'Q.vv*d
#     Qu_ = vec(Q.u) + b'*vec(Q.v) + Q.uv*vec(d) + b'*Q.vv*d
#     Qxx_ = Q.xx + Q.xv*K + K'*Q.xv' + K'*Q.vv*K
#     Quu_ = Q.uu + Q.uv*b + b'*Q.uv' + b'*Q.vv*b + ρ*I
#     Qxu_ = Q.xu + K'*Q.uv' + Q.xv*b + K'*Q.vv*b
#
#     # cache (approximate) cost-to-go at timestep k
#     s[1:n] = Qx_
#     s[n+1:n+m] = Qu_
#     S[1:n,1:n] = Qxx_
#     S[n+1:n+m,n+1:n+m] = Quu_
#     S[1:n,n+1:n+m] = Qxu_
#     S[n+1:n+m,1:n] = Qxu_'
#
#     # line search terms
#     Δv = [-vec(Q.v)'vec(d) 0.5*vec(d)'Q.vv*vec(d)]
#
#     return K,b,d,Δv,reg
# end


text:
struct QuietRealTimePresenter <: RealTimePresenter end
present(::QuietRealTimePresenter, ::Scenario) = nothing
present(::QuietRealTimePresenter, ::Gherkin.ScenarioStep) = nothing
present(::QuietRealTimePresenter, ::Gherkin.ScenarioStep, ::StepExecutionResult) = nothing
present(::QuietRealTimePresenter, ::Gherkin.Feature) = nothing

struct ColorConsolePresenter <: RealTimePresenter
    io::IO
    colors::Dict{Type, Symbol}

    function ColorConsolePresenter(io::IO = STDOUT)
        colors = Dict{Type, Symbol}(
            NoStepDefinitionFound => :yellow,
            NonUniqueMatch => :magenta,
            SuccessfulStepExecution => :green,
            StepFailed => :red,
            UnexpectedStepError => :light_magenta,
            SkippedStep => :light_black
        )

        new(io, colors)
    end
end

stepformat(step::Given) = "Given $(step.text)"
stepformat(step::When) = " When $(step.text)"
stepformat(step::Then) = " Then $(step.text)"

stepcolor(presenter::Presenter, step::StepExecutionResult) = presenter.colors[typeof(step)]

stepresultmessage(step::StepFailed) = ["FAILED: " * step.assertion]
stepresultmessage(nomatch::NoStepDefinitionFound) = ["No match for '$(stepformat(nomatch.step))'"]
stepresultmessage(nonunique::NonUniqueMatch) = vcat(["Multiple matches found:"], ["  In " * location.filename for location in nonunique.locations])
stepresultmessage(::SuccessfulStepExecution) = []
stepresultmessage(::SkippedStep) = []
stepresultmessage(unexpected::UnexpectedStepError) = vcat(["Exception: $(string(unexpected.ex))"], ["  " * string(x) for x in unexpected.stack])
stepresultmessage(::StepExecutionResult) = []

function present(presenter::ColorConsolePresenter, feature::Feature)
    println()
    print_with_color(:white, presenter.io, "Feature: $(feature.header.description)\n")
end

function present(presenter::ColorConsolePresenter, scenario::Scenario)
    println()
    print_with_color(:blue, presenter.io, "  Scenario: $(scenario.description)\n")
end

function present(presenter::ColorConsolePresenter, step::Gherkin.ScenarioStep)
    print_with_color(:light_cyan, presenter.io, "    $(stepformat(step))")
end

function present(presenter::ColorConsolePresenter, step::Gherkin.ScenarioStep, result::StepExecutionResult)
    color = stepcolor(presenter, result)
    print_with_color(color, presenter.io, "\r    $(stepformat(step))\n")

    resultmessage = stepresultmessage(result)
    if !isempty(resultmessage)
        # The result message may be one or more lines. Indent them all.
        s = [" " ^ 8 * x * "\n" for x in resultmessage]
        println()
        println(join(s))
    end
end


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities
const MOIB = MOI.Bridges

import CSDP
const optimizer = CSDP.Optimizer()
MOI.set(optimizer, MOI.Silent(), true)

@testset "SolverName" begin
    @test MOI.get(optimizer, MOI.SolverName()) == "CSDP"
end

@testset "supports_default_copy_to" begin
    @test MOIU.supports_allocate_load(optimizer, false)
    @test !MOIU.supports_allocate_load(optimizer, true)
end

# UniversalFallback is needed for starting values, even if they are ignored by CSDP
const cache = MOIU.UniversalFallback(MOIU.Model{Float64}())
const cached = MOIU.CachingOptimizer(cache, optimizer)
const bridged = MOIB.full_bridge_optimizer(cached, Float64)
const config = MOIT.TestConfig(atol=1e-4, rtol=1e-4)

@testset "Options" begin
    param = MOI.RawParameter(:bad_option)
    err = MOI.UnsupportedAttribute(param)
    @test_throws err CSDP.Optimizer(bad_option = 1)
end

@testset "Unit" begin
    MOIT.unittest(bridged, config, [
        # `NUMERICAL_ERROR` on Mac: https://travis-ci.org/JuliaOpt/CSDP.jl/jobs/601302777#L217-L219
        "solve_unbounded_model",
        # `NumberOfThreads` not supported.
        "number_threads",
        # `TimeLimitSec` not supported.
        "time_limit_sec",
        # SingleVariable objective of bridged variables, will be solved by objective bridges
        "solve_time", "raw_status_string", "solve_singlevariable_obj",
        # Quadratic functions are not supported
        "solve_qcp_edge_cases", "solve_qp_edge_cases",
        # Integer and ZeroOne sets are not supported
        "solve_integer_edge_cases", "solve_objbound_edge_cases",
        "solve_zero_one_with_bounds_1",
        "solve_zero_one_with_bounds_2",
        "solve_zero_one_with_bounds_3"])
end
@testset "Continuous Linear" begin
    # See explanation in `MOI/test/Bridges/lazy_bridge_optimizer.jl`.
    # This is to avoid `Variable.VectorizeBridge` which does not support
    # `ConstraintSet` modification.
    MOIB.remove_bridge(bridged, MOIB.Constraint.ScalarSlackBridge{Float64})
    MOIT.contlineartest(bridged, config, [
        # Finds `MOI.ALMOST_OPTIMAL` instead of `MOI.OPTIMAL`
        "linear10b",
        # Empty constraint
        "linear15"
    ])
end
@testset "Continuous Conic" begin
    MOIT.contconictest(bridged, config, [
        # Finds `MOI.OPTIMAL` instead of `MOI.INFEASIBLE`.
        "soc3",
        # Empty constraint `c4`
        "psdt2",
        # See https://github.com/coin-or/Csdp/issues/11
        "rotatedsoc1v",
        # Missing bridges
        "rootdets",
        # Does not support power and exponential cone
        "pow", "dualpow", "logdet", "exp", "dualexp"])
end


text:
import Test: @test, @testset, @test_broken
import Luna: Raman, Maths, Nonlinear, PhysData, Grid
import QuadGK: quadgk
import FFTW

import Luna: set_fftw_mode
set_fftw_mode(:estimate)

grid = Grid.RealGrid(1.0, 800e-9, (200e-9, 2000e-9), 40e-12)
ρ = PhysData.density(:N2, 5.0)
h = Raman.raman_response(grid.to, :N2)
T = collect(range(-5e-12, stop=20e-12, length=2^16))
dt = T[2] - T[1]

# test response function is zero for 0 or negative time
@test h(0.0, ρ) == 0.0
@test all(h.(T[T .<= 0.0], ρ) .== 0.0)

hsdo = Raman.RamanRespNormedSingleDampedOscillator(1.0, 2π*100e12, 100e-15)
# test integral of normalised response function is 1
@test isapprox(quadgk(x -> hsdo(x, 1.0), -1e-12, 2e-12)[1], 1.0, rtol=1e-8)

# test frequencies of max gain
hv = Raman.raman_response(grid.to, :N2, rotation=false)
RP! = Nonlinear.RamanPolarField(grid.to, hv)
E = Maths.gauss.(grid.to, fwhm=20e-15)
P = similar(E)
RP!(P, E, ρ) # get polarisation induced by E
# test maximum of frequency response occurs at correct frequency
# note hω grid is twice as long at grid.to, so we need a new freq. grid
ωR = 2π .* collect(0:(length(grid.to) - 1)) ./ ((grid.to[2] - grid.to[1]) * length(grid.to) * 2)
@test isapprox(ωR[argmin(imag.(RP!.hω))], PhysData.raman_parameters(:N2).Ωv, rtol=1e-4)

# same, but for envelopes
grid = Grid.EnvGrid(1.0, 800e-9, (200e-9, 2000e-9), 40e-12)
hv = Raman.raman_response(grid.to, :N2, rotation=false)
RP! = Nonlinear.RamanPolarEnv(grid.to, hv)
E = complex.(Maths.gauss.(grid.to, fwhm=20e-15))
P = similar(E)
RP!(P, E, ρ) # get polarisation induced by E
# test maximum of frequency response occurs at correct frequency
# note hω grid is twice as long at grid.to, so we need a new freq. grid (we only make first half of it)
ωR = 2π .* collect(0:(length(grid.to) - 1)) ./ ((grid.to[2] - grid.to[1]) * length(grid.to) * 2)
@test isapprox(ωR[argmin(imag.(RP!.hω))], PhysData.raman_parameters(:N2).Ωv, rtol=1e-4)

# same, but for hydrogen
grid = Grid.EnvGrid(1.0, 800e-9, (200e-9, 2000e-9), 40e-12)
hv = Raman.raman_response(grid.to, :H2, rotation=false)
RP! = Nonlinear.RamanPolarEnv(grid.to, hv)
E = complex.(Maths.gauss.(grid.to, fwhm=20e-15))
P = similar(E)
ρ = PhysData.density(:H2, 5.0)
RP!(P, E, ρ) # get polarisation induced by E
# test maximum of frequency response occurs at correct frequency
# note hω grid is twice as long at grid.to, so we need a new freq. grid (we only make first half of it)
ωR = 2π .* collect(0:(length(grid.to) - 1)) ./ ((grid.to[2] - grid.to[1]) * length(grid.to) * 2)
@test isapprox(ωR[argmin(imag.(RP!.hω))], PhysData.raman_parameters(:H2).Ωv, rtol=1e-4)

# same, but for hydrogen rotation
grid = Grid.EnvGrid(1.0, 800e-9, (200e-9, 2000e-9), 40e-12)
hv = Raman.raman_response(grid.to, :H2, vibration=false)
RP! = Nonlinear.RamanPolarEnv(grid.to, hv)
E = complex.(Maths.gauss.(grid.to, fwhm=20e-15))
P = similar(E)
ρ = PhysData.density(:H2, 5.0)
RP!(P, E, ρ) # get polarisation induced by E
# test maximum of frequency response occurs at correct frequency
# note hω grid is twice as long at grid.to, so we need a new freq. grid (we only make first half of it)
ωR = 2π .* collect(0:(length(grid.to) - 1)) ./ ((grid.to[2] - grid.to[1]) * length(grid.to) * 2)
@test isapprox(ωR[argmin(imag.(RP!.hω))], 2π*17.45e12, rtol=1e-4)


text:
using Chrono
using TimeZones
using Base.Test

let s = "2014-01-09T21:48:00.921000"
    @test parsedate(s) == DateTime(s)
end

let s = "2014-01-09T21:48:00.921000+05:30"
    tzd = parsedatetz(s)
    utct = string(ZonedDateTime(tzd, TimeZone("UTC")))
    utct = utct[1:(length(utct)-6)]
    @test tzd.utc_datetime == DateTime(utct)
end

function test_durations()
    h2 = Hours(2)
    @test (h2 + h2) == Hours(4)
    d3 = Days(3)
    @test (h2 + d3) == Hours(74)
    @test convert(Hours, h2) == h2
    @test convert(Seconds, h2) == Seconds(7200)
    @test div(d3, h2) == Hours(36)
    @test rem(d3, h2) == Hours(0)
    @test rem(d3, Hours(5)) == Hours(2)
    @test (d3 < h2) == false
    @test (d3 == h2) == false
    @test d3 != h2
    @test d3 !== h2
    @test d3 > h2
    @test d3 > Hours(48)
    @test (d3 > Hours(72)) == false
    @test d3 >= Hours(72)
    @test d3 == Hours(72)
    @test isequal(d3, Hours(72))
end

test_durations()

max_stars_repo_name: tanmaykm/Chrono

max_issues_repo_name: tanmaykm/Chrono

max_forks_repo_name: tanmaykm/Chrono


text:
include("Lathe.jl")
import .Lathe.stats: z

using Test

"""
Stats
"""
# z sanity check
@test size(Lathe.stats.z([1 2 3]), 2) == 3
@test Lathe.stats.z([1 2 3])[2] == 0


text:
using ComoniconCast
using Test

@testset "cast" begin
    include("cast.jl")
end

@testset "markdown" begin
    include("markdown.jl")
end


text:
__precompile__()
module SExpressions

using Reexport

include("Lists/Lists.jl")
include("Keywords/Keywords.jl")
include("Parser/Parser.jl")

import .Parser: parse, parseall, parsefile

macro sx_str(x::String)
    QuoteNode(parse(x))
end

export @sx_str

end  # module SExpressions


text:
using FreeTypeAbstraction: hadvance, leftinkbound, inkwidth, get_extent, ascender, descender

one_attribute_per_char(attribute, string) = (attribute for char in string)

function one_attribute_per_char(font::NativeFont, string)
    return (find_font_for_char(char, font) for char in string)
end

function attribute_per_char(string, attribute)
    n_words = 0
    if attribute isa AbstractVector
        if length(attribute) == length(string)
            return attribute
        else
            n_words = length(split(string, r"\s+"))
            if length(attribute) == n_words
                i = 1
                return map(collect(string)) do char
                    f = attribute[i]
                    char == "\n" && (i += 1)
                    return f
                end
            end
        end
    else
        return one_attribute_per_char(attribute, string)
    end
    error("A vector of attributes with $(length(attribute)) elements was given but this fits neither the length of '$string' ($(length(string))) nor the number of words ($(n_words))")
end

"""
    layout_text(
        string::AbstractString, textsize::Union{AbstractVector, Number},
        font, align, rotation, justification, lineheight
    )

Compute a GlyphCollection for a `string` given textsize, font, align, rotation, model, justification, and lineheight.
"""
function layout_text(
        string::AbstractString, textsize::Union{AbstractVector, Number},
        font, align, rotation, justification, lineheight, color, strokecolor, strokewidth
    )

    ft_font = to_font(font)
    rscale = to_textsize(textsize)
    rot = to_rotation(rotation)

    fontperchar = attribute_per_char(string, ft_font)
    textsizeperchar = attribute_per_char(string, rscale)

    glyphcollection = glyph_collection(string, fontperchar, textsizeperchar, align[1],
        align[2], lineheight, justification, rot, color, strokecolor, strokewidth)

    return glyphcollection
end

"""
    glyph_collection(str::AbstractString, font_per_char, fontscale_px, halign, valign, lineheight_factor, justification, rotation, color)

Calculate the positions for each glyph in a string given a certain font, font size, alignment, etc.
This layout in text coordinates, relative to the anchor point [0,0] can then be translated and
rotated to wherever it is needed in the plot.
"""
function glyph_collection(str::AbstractString, font_per_char, fontscale_px, halign, valign,
        lineheight_factor, justification, rotation, color, strokecolor, strokewidth)

    isempty(str) && return GlyphCollection(
        [], [], Point3f[],FreeTypeAbstraction.FontExtent{Float32}[],
        Vec2f[], Float32[], RGBAf[], RGBAf[], Float32[])

    # collect information about every character in the string
    charinfos = broadcast((c for c in str), font_per_char, fontscale_px) do char, font, scale
        unscaled_extent = get_extent(font, char)
        lineheight = Float32(font.height / font.units_per_EM * lineheight_factor * scale)
        unscaled_hi_bb = height_insensitive_boundingbox(unscaled_extent, font)
        hi_bb = Rect2f(
            Makie.origin(unscaled_hi_bb) * scale,
            widths(unscaled_hi_bb) * scale)
        (char = char, font = font, scale = scale, hadvance = hadvance(unscaled_extent) * scale,
            hi_bb = hi_bb, lineheight = lineheight, extent = unscaled_extent)
    end

    # split the character info vector into lines after every \n
    lineinfos = let
        last_line_start = 1
        lineinfos = typeof(view(charinfos, last_line_start:last_line_start))[]
        for (i, ci) in enumerate(charinfos)
            if ci.char == '\n' || i == length(charinfos)
                push!(lineinfos, view(charinfos, last_line_start:i))
                last_line_start = i+1
            end
        end
        lineinfos
    end

    # calculate the x positions of each character in each line
    xs = map(lineinfos) do line
        cumsum([
            0f0;
            [l.hadvance for l in line[1:end-1]]
        ])
    end

    # calculate linewidths as the last origin plus hadvance for each line
    linewidths = map(lineinfos, xs) do line, xx
        nchars = length(line)
        # if the last and not the only character is \n, take the previous one
        # to compute the width
        i = (nchars > 1 && line[end].char == '\n') ? nchars - 1 : nchars
        xx[i] + line[i].hadvance
    end

    # the maximum width is needed for justification
    maxwidth = maximum(linewidths)

    # how much each line differs from the maximum width for justification correction
    width_differences = maxwidth .- linewidths

    # shift all x values by the justification amount needed for each line
    # if justification is automatic it depends on alignment
    float_justification = if justification === automatic
        if halign == :left || halign == 0
            0.0f0
        elseif halign == :right || halign == 1
            1.0f0
        elseif halign == :center || halign == 0.5
            0.5f0
        else
            0.5f0
        end
    elseif justification == :left
        0.0f0
    elseif justification == :right
        1.0f0
    elseif justification == :center
        0.5f0
    else
        Float32(justification)
    end

    xs_justified = map(xs, width_differences) do xsgroup, wd
        xsgroup .+ wd * float_justification
    end

    # each character carries a "lineheight" metric given its font and scale and a lineheight scaling factor
    # make each line's height the maximum of these values in the line
    lineheights = map(lineinfos) do line
        maximum(l -> l.lineheight, line)
    end

    # compute y values by adding up lineheights in negative y direction
    ys = cumsum([0.0; -lineheights[2:end]])

    # compute x values after left/center/right alignment
    halign = if halign isa Number
        Float32(halign)
    elseif halign == :left
        0.0f0
    elseif halign == :center
        0.5f0
    elseif halign == :right
        1.0f0
    else
        error("Invalid halign $halign. Valid values are <:Number, :left, :center and :right.")
    end
    xs_aligned = [xsgroup .- halign * maxwidth for xsgroup in xs_justified]

    # for y alignment, we need the largest ascender of the first line
    # and the largest descender of the last line
    first_line_ascender = maximum(lineinfos[1]) do l
        ascender(l.font) * l.scale
    end

    last_line_descender = minimum(lineinfos[end]) do l
        descender(l.font) * l.scale
    end

    # compute the height of all lines together
    overall_height = first_line_ascender - ys[end] - last_line_descender

    # compute y values after top/center/bottom/baseline alignment
    ys_aligned = if valign == :baseline
        ys .- first_line_ascender .+ overall_height .+ last_line_descender
    else
        va = if valign isa Number
            Float32(valign)
        elseif valign == :top
            1f0
        elseif valign == :bottom
            0f0
        elseif valign == :center
            0.5f0
        else
            error("Invalid valign $valign. Valid values are <:Number, :bottom, :baseline, :top, and :center.")
        end

        ys .- first_line_ascender .+ (1 - va) .* overall_height
    end

    # compute the origins for each character by rotating each character around the common origin
    # which is the alignment anchor and now [0, 0]
    # use 3D coordinates already because later they will be required in that format anyway
    charorigins = [Ref(rotation) .* Point3f.(xsgroup, y, 0) for (xsgroup, y) in zip(xs_aligned, ys_aligned)]

    # return a GlyphCollection, which contains each character's origin, height-insensitive
    # boundingbox and horizontal advance value
    # these values should be enough to draw characters correctly,
    # compute boundingboxes without relayouting and maybe implement
    # interactive features that need to know where characters begin and end
    per_char(attr) = collect(attribute_per_char(str, attr)) # attribute_per_char returns generators
    return GlyphCollection(
        [x.char for x in charinfos],
        [x.font for x in charinfos],
        reduce(vcat, charorigins),
        [x.extent for x in charinfos],
        [Vec2f(x.scale) for x in charinfos],
        per_char(rotation), # rotations is used as one rotation per string above. TODO, allow one rotation per char
        per_char(color),
        per_char(strokecolor),
        per_char(strokewidth)
    )
end

# function to concatenate vectors with a value between every pair
function padded_vcat(arrs::AbstractVector{T}, fillvalue) where T <: AbstractVector{S} where S
    n = sum(length.(arrs))
    arr = fill(convert(S, fillvalue), n + length(arrs) - 1)

    counter = 1
    @inbounds for a in arrs
        for v in a
            arr[counter] = v
            counter += 1
        end
        counter += 1
    end
    arr
end

function alignment2num(x::Symbol)
    (x == :center) && return 0.5f0
    (x in (:left, :bottom)) && return 0.0f0
    (x in (:right, :top)) && return 1.0f0
    return 0.0f0 # 0 default, or better to error?
end

# Backend data

_offset_to_vec(o::VecTypes) = to_ndim(Vec3f, o, 0)
_offset_to_vec(o::Vector) = to_ndim.(Vec3f, o, 0)
_offset_at(o::Vec3f, i) = o
_offset_at(o::Vector, i) = o[i]
Base.getindex(x::ScalarOrVector, i) = x.sv isa Vector ? x.sv[i] : x.sv

function text_quads(position::VecTypes, gc::GlyphCollection, offset, transfunc)
    p = apply_transform(transfunc, position)
    pos = [to_ndim(Point3f, p, 0) for _ in gc.origins]

    atlas = get_texture_atlas()
    pad = GLYPH_PADDING[] / PIXELSIZE_IN_ATLAS[]
    off = _offset_to_vec(offset)

    char_offsets = Vector{Vec3f}(undef, length(pos)) # TODO can this be Vec2f?
    quad_offsets = Vector{Vec2f}(undef, length(pos))
    scales = Vector{Vec2f}(undef, length(pos))
    uvs = Vector{Vec4f}(undef, length(pos))

    for i in eachindex(pos)
        glyph_bb, extent = FreeTypeAbstraction.metrics_bb(
            gc.glyphs[i], gc.fonts[i], gc.scales[i]
        )
        uvs[i] = glyph_uv_width!(atlas, gc.glyphs[i], gc.fonts[i])
        scales[i] = widths(glyph_bb) .+ gc.scales[i] .* 2pad
        char_offsets[i] = gc.origins[i] .+ _offset_at(off, i)
        quad_offsets[i] = minimum(glyph_bb) .- gc.scales[i] .* pad
    end

    # pos is the (space) position given to text (with transfunc applied)
    # char_offsets are 3D offsets in marker space, including:
    #   - offsets passed to text
    #   - character origins (relative to the string origin pos)
    # quad_offsets are 2D offsets of the quad origin to the character origins
    #   (these transform differently than character offsets, also marker space)
    # uvs are the texture coordinates of each character
    # scales is the size of the quad including textsize, character size and texture atlas padding
    return pos, char_offsets, quad_offsets, uvs, scales
end

function text_quads(position::Vector, gcs::Vector{<: GlyphCollection}, offset, transfunc)
    ps = apply_transform(transfunc, position)
    pos = [to_ndim(Point3f, p, 0) for (p, gc) in zip(ps, gcs) for _ in gc.origins]

    atlas = get_texture_atlas()
    pad = GLYPH_PADDING[] / PIXELSIZE_IN_ATLAS[]
    off = _offset_to_vec(offset)

    char_offsets = Vector{Vec3f}(undef, length(pos)) # TODO can this be Vec2f?
    quad_offsets = Vector{Vec2f}(undef, length(pos))
    scales = Vector{Vec2f}(undef, length(pos))
    uvs = Vector{Vec4f}(undef, length(pos))

    k = 1
    for (j, gc) in enumerate(gcs), i in eachindex(gc.origins)
        glyph_bb, extent = FreeTypeAbstraction.metrics_bb(
            gc.glyphs[i], gc.fonts[i], gc.scales[i]
        )
        uvs[k] = glyph_uv_width!(atlas, gc.glyphs[i], gc.fonts[i])
        scales[k] = widths(glyph_bb) .+ gc.scales[i] * 2pad
        char_offsets[k] = gc.origins[i] .+ _offset_at(off, j)
        quad_offsets[k] = minimum(glyph_bb) .- gc.scales[i] .* pad
        k += 1
    end

    return pos, char_offsets, quad_offsets, uvs, scales
end


text:
export d4rl_dataset_params

function d4rl_pybullet_dataset_params()
    dataset = keys(D4RL_PYBULLET_URLS)
    repo = "d4rl-pybullet"
    @info dataset repo
end

const D4RL_PYBULLET_URLS = Dict(
    "hopper-bullet-mixed-v0" => "https://www.dropbox.com/s/xv3p0h7dzgxt8xb/hopper-bullet-mixed-v0.hdf5?dl=1",
    "walker2d-bullet-random-v0" => "https://www.dropbox.com/s/1gwcfl2nmx6878m/walker2d-bullet-random-v0.hdf5?dl=1",
    "hopper-bullet-medium-v0" => "https://www.dropbox.com/s/w22kgzldn6eng7j/hopper-bullet-medium-v0.hdf5?dl=1",
    "walker2d-bullet-mixed-v0" => "https://www.dropbox.com/s/i4u2ii0d85iblou/walker2d-bullet-mixed-v0.hdf5?dl=1",
    "halfcheetah-bullet-mixed-v0" => "https://www.dropbox.com/s/scj1rqun963aw90/halfcheetah-bullet-mixed-v0.hdf5?dl=1",
    "halfcheetah-bullet-random-v0" => "https://www.dropbox.com/s/jnvpb1hp60zt2ak/halfcheetah-bullet-random-v0.hdf5?dl=1",
    "walker2d-bullet-medium-v0" => "https://www.dropbox.com/s/v0f2kz48b1hw6or/walker2d-bullet-medium-v0.hdf5?dl=1",
    "hopper-bullet-random-v0" => "https://www.dropbox.com/s/bino8ojd7iq4p4d/hopper-bullet-random-v0.hdf5?dl=1",
    "ant-bullet-random-v0" => "https://www.dropbox.com/s/2xpmh4wk2m7i8xh/ant-bullet-random-v0.hdf5?dl=1",
    "halfcheetah-bullet-medium-v0" => "https://www.dropbox.com/s/v4xgssp1w968a9l/halfcheetah-bullet-medium-v0.hdf5?dl=1",
    "ant-bullet-medium-v0" => "https://www.dropbox.com/s/6n79kwd94xthr1t/ant-bullet-medium-v0.hdf5?dl=1",
    "ant-bullet-mixed-v0" => "https://www.dropbox.com/s/pmy3dzab35g4whk/ant-bullet-mixed-v0.hdf5?dl=1"
)

function d4rl_pybullet_init()
    repo = "d4rl-pybullet"
    for ds in keys(D4RL_PYBULLET_URLS)
        register(
            DataDep(
                repo* "-" * ds,
                """
                Credits: https://github.com/takuseno/d4rl-pybullet
                The following dataset is fetched from the d4rl-pybullet.
                """,
                D4RL_PYBULLET_URLS[ds],
            )
        )
    end
    nothing
end


text:
###################################################################################################
#
# ResNet50 model for Flux
#   Reference: [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
#
####################################################################################################
using Flux

include("config.jl")

"""
    The identity block is the block that has no conv layer at shortcut.

    # Arguments
        kernel_size: defualt 3, the kernel size of middle conv layer at main path
        filters: list of integers, the filterss of 3 conv layer at main path
Returns the block model.
"""
function identity_block(kernel_size, filters)
    filter0, filter1, filter2, filter3 = filters
    Chain(
        Conv((1, 1), filter0=>filter1),
        BatchNorm(filter1),
        x -> relu.(x),
        Conv(kernel_size, filter1=>filter2, pad=(1, 1)),
        BatchNorm(filter2),
        x -> relu.(x),
        Conv((1, 1), filter2=>filter3),
        BatchNorm(filter3),
    ) |> gpu
end

struct IDBlock
    kernel_size::Tuple
    filters::AbstractArray
    model::Chain
    IDBlock(kernel_size, filters) = new(kernel_size, filters, identity_block(kernel_size, filters))
end

function (ib::IDBlock)(x)
    relu.(ib.model(x) + x)
end

"""
    conv_block is the block that has a conv layer at shortcut

    # Arguments
        input_tensor: input for the model
        kernel_size: the kernel size of middle conv layer at main path
        filters: list of integers, the filterss of 3 conv layer at main path
    Note: the shortcut should have strides=(2, 2)
Returns block model.
"""
function plain_model(kernel_size, filters; strides=(2, 2))
    filter0, filter1, filter2, filter3 = filters
    Chain(
        Conv((1, 1), filter0=>filter1, stride=strides),
        BatchNorm(filter1),
        x -> relu.(x),
        Conv(kernel_size, filter1=>filter2, pad=(1, 1)),
        BatchNorm(filter2),
        x -> relu.(x),
        Conv((1, 1), filter2=>filter3),
        BatchNorm(filter3),
    ) |> gpu
end

struct ConvBlock
    kernel_size::Tuple
    filters::AbstractArray
    strides::Tuple
    plain::Chain
    shortcut::Chain
    ConvBlock(kernel_size, filters; strides=(2, 2)) = new(kernel_size, filters, strides,
        plain_model(kernel_size, filters, strides=strides),
        Chain(Conv((1, 1), filters[2]=>filters[4], stride=strides), BatchNorm(filters[4])) |> gpu)
end

function (cb::ConvBlock)(x)
    relu.(cb.plain(x) + cb.shortcut(x))
end

"""
    Zero-padding layer for input

    # Arguments
        inputs: input tensor
        padding: symmetric padding values for height and width [default=(1, 1)]
Returns add rows and columns of zeros at the top, bottom, left and right side of an input tensor.
"""
function zero_padding(inputs; padding=(1, 1))
    row, col, c, s = size(inputs)
    h, w = padding
    result = vcat(inputs, zeros(h, row, c, s))
    result = hcat(result, zeros(col + h, w, c, s))
    result = hcat(zeros(row + h, w, c, s), result)
    vcat(zeros(h, col + w * 2, c, s), result)
end

ResNetModel = Chain(
    x -> zero_padding(x, padding=(3, 3)),
    Conv((7, 7), 1=>64, stride=(2, 2)),
    BatchNorm(64),
    x -> relu.(x),
    x -> maxpool(x, (3, 3); stride=(2, 2)),
    ConvBlock((3, 3), [64, 64, 64, 256], strides=(1, 1)),
    IDBlock((3, 3), [256, 64, 64, 256]),
    IDBlock((3, 3), [256, 64, 64, 256]),
    ConvBlock((3, 3), [256, 128, 128, 512]),
    IDBlock((3, 3), [512, 128, 128, 512]),
    IDBlock((3, 3), [512, 128, 128, 512]),
    IDBlock((3, 3), [512, 128, 128, 512]),
    ConvBlock((3, 3), [512, 256, 256, 1024]),
    IDBlock((3, 3), [1024, 256, 256, 1024]),
    IDBlock((3, 3), [1024, 256, 256, 1024]),
    IDBlock((3, 3), [1024, 256, 256, 1024]),
    IDBlock((3, 3), [1024, 256, 256, 1024]),
    IDBlock((3, 3), [1024, 256, 256, 1024]),
    ConvBlock((3, 3), [1024, 512, 512, 2048]),
    IDBlock((3, 3), [2048, 512, 512, 2048]),
    IDBlock((3, 3), [2048, 512, 512, 2048]),
    x -> meanpool(x, (7, 7)),
    x -> reshape(x, :, size(x, 4)),
    Dense(2048, NUM_CLASSES),
    softmax
) |> gpu


text:
@testset "1-Particle Basis"  begin

##

using ACE
using Printf, Test, LinearAlgebra, StaticArrays
using ACE: evaluate, evaluate_d, Rn1pBasis, Ylm1pBasis,
      EuclideanVectorState, Product1pBasis
using Random: shuffle
using ACEbase.Testing: dirfdtest, fdtest, print_tf

##

@info "Build a 1p basis from scratch"

maxdeg = 5
r0 = 1.0
rcut = 3.0

trans = PolyTransform(1, r0)   # r -> x = 1/r^2
J = transformed_jacobi(maxdeg, trans, rcut; pcut = 2)   #  J_n(x) * (x - xcut)^pcut
Rn = Rn1pBasis(J)
Ylm = Ylm1pBasis(maxdeg)
B1p = Product1pBasis( (Rn, Ylm) )
ACE.init1pspec!(B1p)

nX = 10
Xs = rand(EuclideanVectorState, Rn, nX)
cfg = ACEConfig(Xs)

A = evaluate(B1p, cfg)

@info("test against manual summation")
A1 = sum( evaluate(B1p, X) for X in Xs )
println(@test A1 ≈ A)

@info("test permutation invariance")
println(@test A ≈ evaluate(B1p, ACEConfig(shuffle(Xs))))

##

@info("Ylm1pBasis gradients")
Y = ACE.alloc_B(Ylm)
dY = ACE.alloc_dB(Ylm)
tmpd = ACE.alloc_temp_d(Ylm)
ACE.evaluate_ed!(Y, dY, tmpd, Ylm, Xs[1])

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(Y))
   F = x -> sum(ACE.evaluate(Ylm, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_d(Ylm, EuclideanVectorState(SVector{3}(x))) .* c))
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()

##

@info("Rn1pBasis gradients")

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(Rn))
   F = x -> sum(ACE.evaluate(Rn, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_d(Rn, EuclideanVectorState(SVector{3}(x))) .* c))
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()

##

@info("Product basis evaluate_ed! tests")

tmp_d = ACE.alloc_temp_d(B1p)
A1 = ACE.alloc_B(B1p)
A2 = ACE.alloc_B(B1p)
ACE.evaluate!(A1, tmp_d, B1p, cfg)
dA = ACE.alloc_dB(B1p, length(cfg))
ACE.evaluate_ed!(A2, dA, tmp_d, B1p, cfg)
println(@test A1 ≈ A2)

##
@info("Product basis gradient test")

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(B1p))
   F = x -> sum(ACE.evaluate(B1p, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_ed(B1p, EuclideanVectorState(SVector{3}(x)))[2] .* c))
   F(x0)
   dF(x0)
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()
##

end


text:
#=

don't want to fall into the same trap as with JAMES!

APBE and Pygpe both are more suited for atom-structure-like sims,
don't seem to support simple dielectrics?
fenics requires it in a "weak form",

=#

using PhysicalConstants.CODATA2018

VacuumElectricPermittivity

Nx, Ny, Nz = 100

A = zeros(Nx,Ny,Nz)

dielectric = ones(Nx,Ny,Nz)
charge_distribution = zeros(Nx,Ny,Nz)

kappa_external =

function poisson_boltzmann()
    for x = 1:Nx-2
        for y = 1:Nx-2
            for z = 1:Nx-2
        (2.0 / )
        #wait, julia is 1 indexed?
        # wtf?
        a1 =
    end

end