text:
struct MALA{TV, TGV, TF<:AbstractFloat} <: FirstOrderMetropolisSampler
    V::TV
    ∇V!::TGV
    β::TF
    Δt::TF
    σ::TF
end

"""
    MALA(V, ∇V!, β, Δt)

Set up the MALA sampler for overdamped Langevin.

### Fields

* V     - Potential
* ∇V!   - In place gradient of the potential
* β     - Inverse temperature
* Δt    - Time step
"""
function MALA(V::TV,∇V!::TGV, β::TF, Δt::TF) where{TV, TGV, TF<:AbstractFloat}
    σ = sqrt(2 * Δt /β)
    return MALA(V, ∇V!, β, Δt, σ)
end

mutable struct MALAState{TF<:AbstractFloat, Tx} <:FirstOrderMetropolisSamplerState
    x::Tx
    x_proposal::Tx
    V::TF
    V_proposal::TF
    ∇V::Tx
    ∇V_proposal::Tx
    accept::Int
end

"""
    MALA_likelihood(X₀, X₁, gradV0, β, Δt)

Compute the likelihood of the MALA proposal X₀→X₁ for
potential with gradient gradV0 at inverse temperature β
and time step Δt
"""
function MALA_likelihood(X₀, X₁, gradV0, β, Δt)
    #norm of increment squared
    norminc2 = 0.0;
    for i = 1:length(X₀)
        norminc2 += (X₁[i]-X₀[i] + Δt * gradV0[i])^2;
    end
    return exp( -β * (norminc2) / (4*Δt) );
end

function InitState!(x₀, sampler::MALA)

    V = sampler.V(x₀);
    ∇V = similar(x₀);
    sampler.∇V!(∇V, x₀);
    return MALAState(x₀, copy(x₀),
        V, V, copy(∇V), copy(∇V), Int(0));
end

function InitState(x₀, sampler::MALA)

    V = sampler.V(x₀);
    ∇V = similar(x₀);
    sampler.∇V!(∇V, x₀);
    return MALAState(deepcopy(x₀), similar(x₀),
        V, V, copy(∇V), similar(∇V), Int(0));
end

function UpdateState!(state::MALAState, sampler::MALA)

    @. state.x_proposal = state.x - sampler.Δt * state.∇V + sampler.σ * randn();
    state.V_proposal = sampler.V(state.x_proposal);
    sampler.∇V!(state.∇V_proposal, state.x_proposal);
    g0 = MALA_likelihood(state.x, state.x_proposal, state.∇V, sampler.β, sampler.Δt);
    gp = MALA_likelihood(state.x_proposal, state.x, state.∇V_proposal, sampler.β, sampler.Δt);

    a = min(1, gp/g0 * exp(sampler.β*(state.V-state.V_proposal)));

    if rand()<a
        @. state.x = state.x_proposal;
        @. state.∇V = state.∇V_proposal;
        state.V = state.V_proposal;
        state.accept = 1;
    else
        state.accept = 0;
    end
    state
end

max_stars_repo_name: cortner/JuBasicMD


text:
struct Point
  x::Float64
  y::Float64
end

import Base: show
# vanilla version
show(io::IO, p::Point) = print(io, "($(p.x), $(p.y))")

# json vewersion
function show(io::IO, ::MIME"text/json", p::Point)
    print(io, "{x = $(p.x), y = $(p.y)}")
end

# plain text version
function show(io::IO, ::MIME"text/plain", p::Point)
    # extended version with compact option
    if get(io, :compact, true)
        print(io, "($(p.x), $(p.y))")
    else
        print(io, "Point($(p.x), $(p.y))")
    end
end

# XML version
function show(io::IO, ::MIME"text/xml", p::Point)
    println(io, "<point>")
    println(io, "\t<x>$(p.x)</x>")
    println(io, "\t<y>$(p.y)</y>")
    println(io, "</point>")
end

# Example
# julia> p = Point(10, 3)
# Point(10.0, 3.0)

# julia> p
# Point(10.0, 3.0)

# julia> display("text/json", p)
# {x = 10.0, y = 3.0}
# julia> display("text/xml", p)
# <point>
#         <x>10.0</x>
#         <y>3.0</y>
# </point>

# julia> display("text/plain", p)
# Point(10.0, 3.0)

# julia> repr("text/json", p)
# "{x = 10.0, y = 3.0}"

# Compact Example
# julia> io = IOContext(stdout, :compact => false)
# IOContext(Base.TTY(RawFD(15) open, 0 bytes waiting))

# julia> show(io, "text/plain", p)
# Point(10.0, 3.0)
# julia> p
# (10.0, 3.0)

# julia> io = IOContext(stdout, :compact => true)
# IOContext(Base.TTY(RawFD(15) open, 0 bytes waiting))

# julia> show(io, "text/plain", p)
# (10.0, 3.0)
# julia> p
# (10.0, 3.0)

ext: jl

max_stars_repo_name: FourMInfo/JuliaforBeginners

max_issues_repo_name: FourMInfo/JuliaforBeginners

max_forks_repo_name: FourMInfo/JuliaforBeginners


text:
using POMDPs
using ContinuousPOMDPTreeSearchExperiments
using POMDPToolbox
using BasicPOMCP
using POMCPOW
using ParticleFilters
using CPUTime
using VDPTag
using DataFrames
using ARDESPOT
using Distributions

function gen_sims(x::Vector{Float64}, n, k)
    c = max(0.0, x[1])
    k_act = max(1.0, x[2])
    inv_alpha_act = max(0.1, x[3])
    k_obs = max(1.0, x[4])
    inv_alpha_obs = max(0.1, x[5])

    rng = MersenneTwister(0)
    ro = ToNextMLSolver(rng)
    solver = POMCPOWSolver(tree_queries=10_000_000,
                           criterion=MaxUCB(c),
                           final_criterion=MaxTries(),
                           max_depth=10,
                           max_time=0.1,
                           k_action=k_act,
                           alpha_action=1/inv_alpha_act,
                           k_observation=k_obs,
                           alpha_observation=1/inv_alpha_obs,
                           estimate_value=FORollout(ro),
                           check_repeat_act=false,
                           check_repeat_obs=false,
                           next_action=RootToNextMLFirst(rng),
                           # default_action=ReportWhenUsed(TagAction(false, 0.0)),
                           default_action=TagAction(false, 0.0),
                           rng=rng
                          )

    pomdp = VDPTagPOMDP()
    planner = solve(solver, pomdp)

    sims = []
    for i in 1:n
        filter = ObsAdaptiveParticleFilter(deepcopy(pomdp),
                                           LowVarianceResampler(10_000),
                                           0.05,
                                           MersenneTwister(i+10000*k))

        srand(planner, i+40000*k)
        sim = Sim(deepcopy(pomdp),
                  deepcopy(planner),
                  filter,
                  rng=MersenneTwister(i+50_000*k),
                  max_steps=100,
                  metadata=Dict(:i=>i, :k=>k)
                 )

        push!(sims, sim)
    end

    return sims
end

#=
gen_syms(x::Vector{Float64}, n, k) = Any[k=>copy(x) for i in 1:n]

function run_parallel(sims::Vector{Any})
    ks = Int[]
    rewards = Float64[]
    for (k, v) in sims
        push!(ks, k)
        push!(rewards, -sum(x->x^2, v - [40.0, 30.0, 20.0, 3.1, 28.0]) + 10.0*rand())
    end
    return DataFrame(k=ks, reward=rewards)
end
=#

# start_mean = [40.0, 8.0, 20.0, 4.0, 20.0]
start_mean = [65.0, 13.0, 8.0, 2.0, 33.0]
start_cov = diagm([40.0^2, 6.0^2, 10.0^2, 6.0^2, 10.0^2])
d = MvNormal(start_mean, start_cov)
rng = MersenneTwister(15)
K = 200
m = 50

for i in 1:100
    sims = []
    params = Vector{Vector{Float64}}(K)
    for k in 1:K
        p = rand(d)
        params[k] = p
        k_sims = gen_sims(p, 40, k)
        append!(sims, k_sims)
    end
    results = run_parallel(sims)
    combined = by(results, :k) do df
        DataFrame(mean_reward=mean(df[:reward]))
    end
    @show mean(combined[:mean_reward])
    order = sortperm(combined[:mean_reward])
    elite = params[combined[:k][order[K-m:end]]]
    elite_matrix = Matrix{Float64}(5, m)
    for k in 1:m
        elite_matrix[:,k] = elite[k]
    end
    try
        d = fit(typeof(d), elite_matrix)
    catch ex
        if ex isa Base.LinAlg.PosDefException
            println("pos def exception")
            d = fit(typeof(d), elite_matrix += 0.01*randn(size(elite_matrix)))
        else
            rethrow(ex)
        end
    end
    println("iteration $i")
    @show mean(d)
    @show eigvals(cov(d))
    @show eigvecs(cov(d))
end


text:
export SecondJackknife, SecondJackknifeSet
export second_jackknife

mutable struct SecondJackknife <: MCObservable
    x_i  :: Vector{Float64}
    x_ij :: Matrix{Float64}
    b_i  :: Vector{Float64}
end

function SecondJackknife(x_ij :: Matrix{Float64})
    n = size(x_ij,1)
    x_i = zeros[n]
    b_i = zeros[n]
    for i in 1:n
        x_i[i] = x_ij[i,i]
        b_i[i] = x_ij[i,i]
        b_i[i] -= sum(x_ij[1:i-1,i])
        b_i[i] -= sum(x_ij[i,i+1:n])
    end
    return SecondJackknife(x_i, x_ij, b_i)
end

SecondJackknife(jk::SecondJackknife, f::Function) = SecondJackknife(map(f,jk.x_ij))

function SecondJackknife(b::BinningObservable)
    if count(b) > 0
        level = maxlevel(b)
        xs = b.bins[level]
        if b.nincomplete[level] > 0
            push!(xs, b.incompletes[level]/b.nincomplete[level])
        end
        s = sum(xs)
        n = length(xs)
        n1inv = 1.0/(n-1)
        n2inv = 1.0/(n-2)
        jk_ij = zeros(n,n)
        for i in 1:n
            xi = xs[i]
            jk_ij[i,i] = (s-xi)*n1inv
            for j in i+1:n
                jk_ij[i,j] = (s-xi-xs[j])*n2inv
            end
        end
        return SecondJackknife(jk_ij)
    else
        return Jackknife(zeros(0,0))
    end
end

count(jk::SecondJackknife) = length(jk.x_ij)
isempty(jk::SecondJackknife) = isempty(jk.x_ij)

function mean(jk::SecondJackknife)
    if isempty(jk)
        throw(DomainError())
    else
        return mean(jk.x_i + jk.b_i)
    end
end
function stderr(jk::SecondJackknife)
    n = count(jk)
    if n == 0
        throw(DomainError())
    elseif n == 1
        return Inf
    else
        sums = mapreduce(x->[x,x*x], +, jk.x_i + jk.b_i)
        sums /= n
        sigma2 = sums[2] - sums[1]*sums[1]
        sigma2 *= n-1
        return sqrt(sigma2)
    end
end

unary_functions = (
                   :-,
                   :sin, :cos, :tan,
                   :sind, :cosd, :tand,
                   :sinpi, :cospi,
                   :sinh, :cosh, :tanh,
                   :asin, :acos, :atan,
                   :asind, :acosd, :atand,
                   :sec, :csc, :cot,
                   :secd, :cscd, :cotd,
                   :asec, :acsc, :acot,
                   :asecd, :acscd, :acotd,
                   :sech, :csch, :coth,
                   :asinh, :acosh, :atanh,
                   :asech, :acsch, :acoth,
                   :sinc, :cosc,
                   :log, :log2, :log10, :log1p,
                   :exp, :exp2, :exp10, :expm1,
                   :abs, :abs2,
                   :sqrt, :cbrt,
                  )

for op in unary_functions
    eval( Expr(:import, :Base, op) )
    eval( Expr(:export, op) )
    @eval ($op)(jk::SecondJackknife) = SecondJackknife(jk, $op)
end

binary_functions = (
                    :+, :-, :*, :/, :\
                   )

for op in binary_functions
    eval( Expr(:import, :Base, op) )
    eval( Expr(:export, op) )
    @eval ($op)(jk::SecondJackknife, rhs::Real) = SecondJackknife(jk, lhs->($op)(lhs,rhs))
    @eval ($op)(lhs::Real, jk::SecondJackknife) = SecondJackknife(jk, rhs->($op)(lhs,rhs))
    op_bw = symbol("."*string(op))
    @eval ($op)(lhs::SecondJackknife, rhs::SecondJackknife) = SecondJackknife( ($op_bw)(lhs.xs, rhs.xs))
end

const SecondJackknifeSet = MCObservableSet{SecondJackknife}

function second_jackknife(obsset :: BinningObservableSet)
    JK = SecondJackknifeSet()
    for (k,v) in obsset
        JK[k] = SecondJackknife(v)
    end
    return JK
end


text:
#=
Utility functions used in the QuantEcon library

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-14

=#
meshgrid(x::AbstractVector, y::AbstractVector) = (repmat(x, 1, length(y))',
                                                  repmat(y, 1, length(x)))

fix(x::Real) = x >= 0 ? floor(Int, x) : ceil(Int, x)
fix!{T<:Real}(x::AbstractArray{T}, out::Array{Int}) = map!(fix, out, x)
fix{T<:Real}(x::AbstractArray{T}) = fix!(x, similar(x, Int))

ckron(A::Array, B::Array) = kron(A, B)
ckron(arrays::Array...) = reduce(kron, arrays)

function gridmake!(out, arrays::AbstractVector...)
    arr = arrays[1]
    typ = eltype(arr)
    l = 1; for a in arrays; l *= length(a); end
    n = length(arrays)
    @assert size(out) == (l, n)

    l_i = length(arr)
    m = Int(l / l_i)

    # fill this column
    row = 1
    @inbounds for el in arr, i = 1:m
        out[row, 1] = el
        row += 1
    end

    if n > 1
        # recursively call to fill upper right block for columns 2:end
        gridmake!(sub(out, 1:m, 2:n), arrays[2:end]...)

        # extract upper right block and fill in lower middle block for columns
        # 2:end
        @inbounds for j in 2:l_i
            out[(j-1)*m+1:(j)*m, 2:end] = sub(out, 1:m, 2:n)
        end
    end
    out
end

function gridmake(arrays::AbstractVector...)
    l = prod([length(a) for a in  arrays])
    T = reduce(promote_type, [eltype(a) for a in arrays])
    gridmake!(Array(T, l, length(arrays)), arrays...)
    out
end

# type stable version if all arrays have the same eltype
function gridmake{T}(arrays::AbstractVector{T}...)
    out = Array(T, prod([length(a) for a in  arrays]), length(arrays))
    gridmake!(out, arrays...)
end


text:
module Pangolin

export Handler3D, OpenGlRenderState, View,

       create_window, should_quit, activate, finish_frame, destroy_window,
       get_projection_view_mat, set_model_view_mat!, show

using PangolinCApi_jll

include("impl.jl")

end # module


text:
using DualNumbers

struct DualQuaternion{T<:Real} <: Number
  q0::Quaternion{T}
  qe::Quaternion{T}
  norm::Bool
end

DualQuaternion{T}(dq::DualQuaternion) where {T<:Real} = DualQuaternion{T}(dq.q0, dq.qe, dq.norm)
function DualQuaternion{T}(d1::Dual, d2::Dual, d3::Dual, d4::Dual, n::Bool=false) where {T<:Real}
  return DualQuaternion{T}(
    Quaternion(DualNumbers.value(d1), DualNumbers.value(d2), DualNumbers.value(d3), DualNumbers.value(d4), n),
    Quaternion(DualNumbers.epsilon(d1), DualNumbers.epsilon(d2), DualNumbers.epsilon(d3), DualNumbers.epsilon(d4)),
    n,
  )
end
function DualQuaternion{T}(q0::Quaternion) where {T<:Real}
  return DualQuaternion{T}(convert(Quaternion{T}, q0), zero(Quaternion{T}), q0.norm)
end
function DualQuaternion{T}(d::Dual) where {T<:Real}
  return DualQuaternion(
    Quaternion{T}(DualNumbers.value(d)),
    Quaternion{T}(DualNumbers.epsilon(d)),
    (DualNumbers.value(d)==one(DualNumbers.value(d))) & iszero(DualNumbers.epsilon(d)))
end
function DualQuaternion{T}(x::Real) where {T<:Real}
  return DualQuaternion(convert(Quaternion{T}, x), zero(Quaternion{T}), abs(x) == one(x))
end

DualQuaternion(q0::Quaternion, qe::Quaternion) = DualQuaternion(promote(q0, qe)..., false)
DualQuaternion(d1::Dual, d2::Dual, d3::Dual, d4::Dual, n::Bool=false) =
  DualQuaternion(Quaternion(DualNumbers.value(d1), DualNumbers.value(d2), DualNumbers.value(d3), DualNumbers.value(d4), n),
                  Quaternion(DualNumbers.epsilon(d1), DualNumbers.epsilon(d2), DualNumbers.epsilon(d3), DualNumbers.epsilon(d4)), n)
DualQuaternion(x::Real) = DualQuaternion(Quaternion(x), Quaternion(zero(x)), abs(x) == one(x))
DualQuaternion(d::Dual) = DualQuaternion(Quaternion(DualNumbers.value(d)), Quaternion(DualNumbers.epsilon(d)), (DualNumbers.value(d)==one(DualNumbers.value(d))) & iszero(DualNumbers.epsilon(d)))
DualQuaternion(q::Quaternion) = DualQuaternion(q, zero(q), q.norm)
DualQuaternion(a::Vector) = DualQuaternion(zero(Quaternion{typeof(a[1])}), Quaternion(a))

const DualQuaternionF16 = DualQuaternion{Float16}
const DualQuaternionF32 = DualQuaternion{Float32}
const DualQuaternionF64 = DualQuaternion{Float64}

promote_rule(::Type{DualQuaternion{T}}, ::Type{S}) where {T <: Real, S <: Real} = DualQuaternion{promote_type(T, S)}
promote_rule(::Type{DualQuaternion{T}}, ::Type{Dual{S}}) where {T <: Real, S <: Real} = DualQuaternion{promote_type(T, S)}
promote_rule(::Type{DualQuaternion{T}}, ::Type{Quaternion{S}}) where {T <: Real, S <: Real} = DualQuaternion{promote_type(T, S)}
promote_rule(::Type{DualQuaternion{T}}, ::Type{DualQuaternion{S}}) where {T <: Real, S <: Real} = DualQuaternion{promote_type(T, S)}

dualquat(q1, q2) = DualQuaternion(q1, q2)
dualquat(q1, q2, n) = DualQuaternion(q1, q2, n)
dualquat(d1, d2, d3, d4) = DualQuaternion(d1, d2, d3, d4)
dualquat(d1, d2, d3, d4, n) = DualQuaternion(d1, d2, d3, d4, n)
dualquat(x) = DualQuaternion(x)

function show(io::IO, dq::DualQuaternion)
  show(io, dq.q0)
  print(io, " + ( ")
  show(io, dq.qe)
  print(io, " )du")
end

Q0(dq::DualQuaternion) = dq.q0
Qe(dq::DualQuaternion) = dq.qe

(/)(dq::DualQuaternion, x::Real) = DualQuaternion(dq.q0 / x, dq.qe / x)

(/)(dq::DualQuaternion, d::Dual) =
  DualQuaternion(dual(dq.q0.s, dq.qe.s) / d,
                  dual(dq.q0.v1, dq.qe.v1) / d,
                  dual(dq.q0.v2, dq.qe.v2) / d,
                  dual(dq.q0.v3, dq.qe.v3) / d)

abs2(dq::DualQuaternion) = dq.norm ? dual(one(dq.q0.s)) :
  dual(abs2(dq.q0),
        2.0 * (dq.q0.s  * dq.qe.s  +
                dq.q0.v1 * dq.qe.v1 +
                dq.q0.v2 * dq.qe.v2 +
                dq.q0.v3 * dq.qe.v3))

abs(dq::DualQuaternion) = dq.norm ? dual(one(dq.q0.s)) : sqrt(abs2(dq))
float(dq::DualQuaternion{T}) where T = convert(DualQuaternion{float(T)}, dq)

conj(dq::DualQuaternion) = DualQuaternion(conj(dq.q0), conj(dq.qe), dq.norm)
dconj(dq::DualQuaternion) = DualQuaternion(dq.q0, -dq.qe, dq.norm)

inv(dq::DualQuaternion) = dq.norm ? conj(dq) : conj(dq) / abs2(dq)

function normalize(dq::DualQuaternion)
  if (dq.norm)
    return dq
  end
  a = abs(dq)
  if abs(a) > 0
    qa = dq / a
    dualquat(qa.q0, qa.qe, true)
  else
    dq
  end
end

function normalizea(dq::DualQuaternion{T}) where {T}
  if (dq.norm)
    return (dq, one(Dual{T}))
  end
  a = abs(dq)
  if abs(a) > 0
    qa = dq / a
    dualquat(qa.q0, qa.qe, true), a
  else
    dq, zero(Dual{T})
  end
end

(-)(dq::DualQuaternion) = DualQuaternion(-dq.q0, -dq.qe, dq.norm)

(+)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 + dw.q0, dq.qe + dw.qe)
(-)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 - dw.q0, dq.qe - dw.qe)
(*)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 * dw.q0,
                                                              dq.q0 * dw.qe + dq.qe * dw.q0,
                                                              dq.norm && dw.norm)
(*)(dq::DualQuaternion, d::Dual) = (*)(Base.promote(dq, d)...)
(*)(d::Dual, dq::DualQuaternion) = (*)(Base.promote(d, dq)...)
(/)(dq::DualQuaternion, dw::DualQuaternion) = dq * inv(dw)
(==)(q::DualQuaternion, w::DualQuaternion) = (q.q0 == w.q0) & (q.qe == w.qe) # ignore .norm field

function angleaxis(dq::DualQuaternion)
  tq = dq.qe * conj(dq.q0)
  t = [2.0 * tq.v1, 2.0 * tq.v2, 2.0 * tq.v3]
  q0s = dq.q0.s
  th0, s0 = angleaxis(dq.q0)
  sq0 = quat(0.0, s0)
  if abs(abs(q0s) - one(q0s)) == 0
    th = dual(th0, 0.5 * abs(quat(0, t)))
    th, dualquat(sq0)
  else
    th = dual(th0, 0.5 * dot(t, s0))
    s0c1 = cross(s0, t)
    tanth = tan(th0)
    s0c2 = (s0c1 / tanth + t) * 0.5
    sqev = cross(s0c2, s0)
    th, dualquat(sq0, quat(0.0, sqev))
  end
end

function angle(dq::DualQuaternion)
  th, ax = angleaxis(dq)
  th
end

function axis(dq::DualQuaternion)
  th, ax = angleaxis(dq)
  ax
end

function exp(dq::DualQuaternion)
  se = dual(dq.q0.s, dq.qe.s)
  se = exp(se)
  dq = dualquat(quat(0.0, imag_part(dq.q0)...), quat(0.0, imag_part(dq.qe)...))
  dq, th = normalizea(dq)
  if dq.norm
    dualquat(se) * (dualquat(cos(th)) + dq * dualquat(sin(th)))
  else
    dualquat(se)
  end
end

function log(dq::DualQuaternion)
  dq, a = normalizea(dq)
  sl = log(a)
  th, s = angleaxis(dq)
  s * dualquat(th) + dualquat(sl)
end

(^)(dq::DualQuaternion, dw::DualQuaternion) = exp(dw * log(dq))

function sqrt(dq::DualQuaternion)
  exp(0.5 * log(dq))
end

dualquatrand() = dualquat(quatrand(), quatrand())
ndualquatrand() = normalize(dualquatrand())

function rand(rng::AbstractRNG, ::Random.SamplerType{DualQuaternion{T}}) where {T<:Real}
    return DualQuaternion{T}(rand(rng, Quaternion{T}), rand(rng, Quaternion{T}), false)
end


text:
struct ProblemResultsExport
    name::Symbol
    duals::Set{ConstraintKey}
    expressions::Set{ExpressionKey}
    parameters::Set{ParameterKey}
    variables::Set{VariableKey}
    aux_variables::Set{AuxVarKey}
    optimizer_stats::Bool
    store_all_flags::Dict{Symbol, Bool}

    function ProblemResultsExport(
        name,
        duals,
        expressions,
        parameters,
        variables,
        aux_variables,
        optimizer_stats,
        store_all_flags,
    )
        duals = _check_fields(duals)
        expressions = _check_fields(expressions)
        parameters = _check_fields(parameters)
        variables = _check_fields(variables)
        aux_variables = _check_fields(aux_variables)
        new(
            name,
            duals,
            expressions,
            parameters,
            variables,
            aux_variables,
            optimizer_stats,
            store_all_flags,
        )
    end
end

function ProblemResultsExport(
    name;
    duals=Set{ConstraintKey}(),
    expressions=Set{ExpressionKey}(),
    parameters=Set{ParameterKey}(),
    variables=Set{VariableKey}(),
    aux_variables=Set{AuxVarKey}(),
    optimizer_stats=true,
    store_all_duals=false,
    store_all_expressions=false,
    store_all_parameters=false,
    store_all_variables=false,
    store_all_aux_variables=false,
)
    store_all_flags = Dict(
        :duals => store_all_duals,
        :expressions => store_all_expressions,
        :parameters => store_all_parameters,
        :variables => store_all_variables,
        :aux_variables => store_all_aux_variables,
    )
    return ProblemResultsExport(
        Symbol(name),
        duals,
        expressions,
        parameters,
        variables,
        aux_variables,
        optimizer_stats,
        store_all_flags,
    )
end

function _check_fields(fields)
    if !(typeof(fields) <: Set)
        fields = Set(fields)
    end

    return fields
end

should_export_dual(x::ProblemResultsExport, key) = _should_export(x, :duals, key)
should_export_expression(x::ProblemResultsExport, key) =
    _should_export(x, :expressions, key)
should_export_parameter(x::ProblemResultsExport, key) = _should_export(x, :parameters, key)
should_export_variable(x::ProblemResultsExport, key) = _should_export(x, :variables, key)
should_export_aux_variable(x::ProblemResultsExport, key) =
    _should_export(x, :aux_variables, key)

function _should_export(exports::ProblemResultsExport, field_name, key)
    exports.store_all_flags[field_name] && return true
    container = getproperty(exports, field_name)
    return key in container
end


text:
using Caesar

fg = initfg()

addVariable!(fg, :x0, Pose2)
addVariable!(fg, :x1, Pose2)

addFactor!(fg, [:x0], PriorPose2(MvNormal(zeros(3),0.001*Matrix(I,3,3))))
addFactor!(fg, [:x1], PriorPose2(MvNormal([0;1;0.0],0.001*Matrix(I,3,3))))

addVariable!(fg, :l1, Point2)

addFactor!(fg, [:x0;:l1], Pose2Point2BearingRange(Normal(pi/4,0.01),Normal(1.0, 1.0)) )
addFactor!(fg, [:x1;:l1], Pose2Point2BearingRange(Normal(-pi/4,0.01),Normal(1.0, 1.0)) )

drawGraph(fg)

tree = solveTree!(fg)

using RoMEPlotting
Gadfly.set_default_plot_size(35cm, 25cm)

drawPosesLandms(fg, spscale=0.3)

plotKDE(fg, :l1, levels=3)

plotLocalProduct(fg, :l1, levels=2)

#


text:
## Series
import BasisFunctions: plotgrid, postprocess
using GridArrays: MaskedGrid
# Postprocessing when the set is a frame: set values outside the domain to NaN
function postprocess(D::Domain, grid, vals, value=NaN)
    mgrid = MaskedGrid(grid, D)
    vals[.!(mask(mgrid))] .= value
    vals
end
postprocess(B::ExtensionFrame, args...) = postprocess(support(B), args...)
# Plotgrids are determined by the underlying set
plotgrid(B::ExtensionFrame, n) = plotgrid(basis(B),n)


text:
#=
tests for primitive cone barrier oracles
=#

using Test
import Random
import Random.randn
using LinearAlgebra
using SparseArrays
import ForwardDiff
import GenericLinearAlgebra.eigen # needed by ForwardDiff currently for test_barrier
import Hypatia
import Hypatia.PolyUtils
import Hypatia.Cones
import Hypatia.RealOrComplex

Random.randn(::Type{BigFloat}, dims::Integer...) = BigFloat.(randn(dims...))

Random.randn(::Type{Complex{BigFloat}}, dims::Integer...) =
    Complex{BigFloat}.(randn(ComplexF64, dims...))

# sanity check oracles
function test_oracles(
    cone::Cones.Cone{T};
    noise::T = T(1e-1),
    scale::T = T(1e-1),
    tol::Real = 1e3 * eps(T),
    init_only::Bool = false,
    init_tol::Real = tol,
    ) where {T <: Real}
    Random.seed!(1)
    dim = Cones.dimension(cone)
    Cones.setup_data!(cone)
    Cones.reset_data(cone)

    point = zeros(T, dim)
    Cones.set_initial_point!(point, cone)
    Cones.load_point(cone, point)
    @test Cones.is_feas(cone)
    @test cone.point == point

    dual_point = -Cones.grad(cone)
    Cones.load_dual_point(cone, dual_point)
    @test Cones.is_dual_feas(cone)
    @test cone.dual_point == dual_point
    @test Cones.get_proxsqr(cone, one(T), true) <= 1 # max proximity
    @test Cones.get_proxsqr(cone, one(T), false) <= dim # sum proximity

    prod_vec = zero(point)
    @test Cones.hess_prod!(prod_vec, point, cone) ≈ dual_point atol=tol rtol=tol

    # test centrality of initial point
    if isfinite(init_tol)
        @test point ≈ dual_point atol=init_tol rtol=init_tol
    end
    init_only && return

    hess = Cones.hess(cone)
    inv_hess = Cones.inv_hess(cone)
    @test hess * inv_hess ≈ I atol=tol rtol=tol

    # perturb and scale the initial point
    perturb_scale!(point, noise, scale)
    perturb_scale!(dual_point, noise, inv(scale))

    Cones.reset_data(cone)
    Cones.load_point(cone, point)
    @test Cones.is_feas(cone)
    Cones.load_dual_point(cone, dual_point)
    @test Cones.is_dual_feas(cone)

    # test gradient and Hessian oracles
    nu = Cones.get_nu(cone)
    grad = Cones.grad(cone)
    @test dot(point, grad) ≈ -nu atol=tol rtol=tol

    hess = Matrix(Cones.hess(cone))
    inv_hess = Matrix(Cones.inv_hess(cone))
    @test hess * inv_hess ≈ I atol=tol rtol=tol

    @test hess * point ≈ -grad atol=tol rtol=tol
    @test Cones.hess_prod!(prod_vec, point, cone) ≈ -grad atol=tol rtol=tol
    @test Cones.inv_hess_prod!(prod_vec, grad, cone) ≈ -point atol=tol rtol=tol

    prod_mat = zeros(T, dim, dim)
    @test Cones.hess_prod!(prod_mat, inv_hess, cone) ≈ I atol=tol rtol=tol
    @test Cones.inv_hess_prod!(prod_mat, hess, cone) ≈ I atol=tol rtol=tol

    psi = dual_point + grad
    proxsqr = dot(psi, Cones.inv_hess_prod!(prod_vec, psi, cone))
    @test Cones.get_proxsqr(cone, one(T), false) ≈ proxsqr atol=tol rtol=tol

    if hasproperty(cone, :use_hess_prod_slow)
        Cones.update_use_hess_prod_slow(cone)
        @test cone.use_hess_prod_slow_updated
        @test !cone.use_hess_prod_slow
        cone.use_hess_prod_slow = true
        @test Cones.hess_prod_slow!(prod_mat, inv_hess, cone) ≈ I atol=tol rtol=tol
    end

    if Cones.use_sqrt_hess_oracles(dim + 1, cone)
        prod_mat2 = Matrix(Cones.sqrt_hess_prod!(prod_mat, inv_hess, cone)')
        @test Cones.sqrt_hess_prod!(prod_mat, prod_mat2, cone) ≈ I atol=tol rtol=tol
        Cones.inv_sqrt_hess_prod!(prod_mat2, Matrix(one(T) * I, dim, dim), cone)
        @test prod_mat2' * prod_mat2 ≈ inv_hess atol=tol rtol=tol
    end

    # test third order deriv oracle
    if Cones.use_dder3(cone)
        @test -Cones.dder3(cone, point) ≈ grad atol=tol rtol=tol

        dir = perturb_scale!(zeros(T, dim), noise, one(T))
        dder3 = Cones.dder3(cone, dir)
        @test dot(dder3, point) ≈ dot(dir, hess * dir) atol=tol rtol=tol
    end

    return
end

# check some oracles agree with ForwardDiff
function test_barrier(
    cone::Cones.Cone{T},
    barrier::Function;
    noise::T = T(1e-1),
    scale::T = T(1e-1),
    tol::Real = 1e4 * eps(T),
    TFD::Type{<:Real} = T,
    ) where {T <: Real}
    Random.seed!(1)
    dim = Cones.dimension(cone)
    Cones.setup_data!(cone)

    point = zeros(T, dim)
    Cones.set_initial_point!(point, cone)
    perturb_scale!(point, noise, scale)

    Cones.reset_data(cone)
    Cones.load_point(cone, point)
    @test Cones.is_feas(cone)

    TFD_point = TFD.(point)

    fd_grad = ForwardDiff.gradient(barrier, TFD_point)
    @test Cones.grad(cone) ≈ fd_grad atol=tol rtol=tol

    dir = 10 * randn(T, dim)
    TFD_dir = TFD.(dir)

    barrier_dir(s, t) = barrier(s + t * TFD_dir)

    fd_hess_dir = ForwardDiff.gradient(s -> ForwardDiff.derivative(t ->
        barrier_dir(s, t), 0), TFD_point)
    @test Cones.hess(cone) * dir ≈ fd_hess_dir atol=tol rtol=tol
    prod_vec = zero(dir)
    @test Cones.hess_prod!(prod_vec, dir, cone) ≈ fd_hess_dir atol=tol rtol=tol

    if Cones.use_dder3(cone)
        fd_third_dir = ForwardDiff.gradient(s2 -> ForwardDiff.derivative(s ->
            ForwardDiff.derivative(t -> barrier_dir(s2, t), s), 0), TFD_point)
        @test -2 * Cones.dder3(cone, dir) ≈ fd_third_dir atol=tol rtol=tol
    end

    return
end

# show time and memory allocation for oracles
function show_time_alloc(
    cone::Cones.Cone{T};
    noise::T = T(1e-4),
    scale::T = T(1e-1),
    ) where {T <: Real}
    Random.seed!(1)
    dim = Cones.dimension(cone)
    println("dimension: ", dim)

    println("setup_data")
    @time Cones.setup_data!(cone)
    Cones.reset_data(cone)

    point = zeros(T, dim)
    Cones.set_initial_point!(point, cone)
    perturb_scale!(point, noise, scale)
    Cones.load_point(cone, point)
    @assert Cones.is_feas(cone)

    dual_point = -Cones.grad(cone)
    perturb_scale!(dual_point, noise, inv(scale))
    Cones.load_dual_point(cone, dual_point)
    @assert Cones.is_dual_feas(cone)

    Cones.reset_data(cone)

    Cones.load_point(cone, point)
    println("is_feas")
    @time Cones.is_feas(cone)

    Cones.load_dual_point(cone, dual_point)
    println("is_dual_feas")
    @time Cones.is_dual_feas(cone)

    println("grad")
    @time Cones.grad(cone)
    println("hess")
    @time Cones.hess(cone)
    println("inv_hess")
    @time Cones.inv_hess(cone)

    point1 = randn(T, dim)
    point2 = zero(point1)
    println("hess_prod")
    @time Cones.hess_prod!(point2, point1, cone)
    println("inv_hess_prod")
    @time Cones.inv_hess_prod!(point2, point1, cone)

    if hasproperty(cone, :use_hess_prod_slow)
        cone.use_hess_prod_slow_updated = true
        cone.use_hess_prod_slow = true
        println("hess_prod_slow")
        @time Cones.hess_prod_slow!(point2, point1, cone)
    end

    if Cones.use_sqrt_hess_oracles(dim + 1, cone)
        println("sqrt_hess_prod")
        @time Cones.sqrt_hess_prod!(point2, point1, cone)
        println("inv_sqrt_hess_prod")
        @time Cones.inv_sqrt_hess_prod!(point2, point1, cone)
    end

    if Cones.use_dder3(cone)
        println("dder3")
        @time Cones.dder3(cone, point1)
    end

    println("get_proxsqr")
    @time Cones.get_proxsqr(cone, one(T), true)

    return
end

function perturb_scale!(
    point::Vector{T},
    noise::T,
    scale::T,
    ) where {T <: Real}
    if !iszero(noise)
        @. point += 2 * noise * rand(T) - noise
    end
    if !isone(scale)
        point .*= scale
    end
    return point
end

# cone utilities

logdet_pd(W::Hermitian) = logdet(cholesky!(copy(W)))

new_vec(w::Vector, dw::Int, T::Type{<:Real}) = copy(w)

function new_vec(w::Vector, dw::Int, R::Type{Complex{T}}) where {T <: Real}
    wR = zeros(Complex{eltype(w)}, dw)
    Cones.vec_copyto!(wR, w)
    return wR
end

function new_herm(w::Vector, dW::Int, T::Type{<:Real})
    W = similar(w, dW, dW)
    Cones.svec_to_smat!(W, w, sqrt(T(2)))
    return Hermitian(W, :U)
end

function new_herm(w::Vector, dW::Int, R::Type{Complex{T}}) where {T <: Real}
    W = zeros(Complex{eltype(w)}, dW, dW)
    Cones.svec_to_smat!(W, w, sqrt(T(2)))
    return Hermitian(W, :U)
end

function rand_sppsd_pattern(dW::Int)
    sparsity = inv(sqrt(dW))
    (row_idxs, col_idxs, _) = findnz(tril!(sprand(Bool, dW, dW, sparsity)) + I)
    return (row_idxs, col_idxs)
end

function rand_herms(ds::Int, Rd::Vector, T::Type{<:Real})
    Ps = Vector{LinearAlgebra.HermOrSym{R, Matrix{R}} where {R <: RealOrComplex{T}}}(
        undef, length(Rd))
    A_1_half = randn(Rd[1], ds, ds)
    Ps[1] = Hermitian(A_1_half * A_1_half' + I, :U)
    for i in 2:length(Rd)
        Ps[i] = Hermitian(randn(Rd[i], ds, ds), :U)
    end
    return Ps
end

function rand_powers(T, d)
    Random.seed!(1)
    α = rand(T, d) .+ 1
    α ./= sum(α)
    return α
end

# real Ps for WSOS cones, use unit box domain
function rand_interp(num_vars::Int, halfdeg::Int, T::Type{<:Real})
    Random.seed!(1)
    domain = PolyUtils.BoxDomain{T}(-ones(T, num_vars), ones(T, num_vars))
    (d, _, Ps, _) = PolyUtils.interpolate(domain, halfdeg, sample = false)
    return (d, Ps)
end

# complex Ps for WSOS cones, use unit ball domain
function rand_interp(num_vars::Int, halfdeg::Int, R::Type{<:Complex{<:Real}})
    Random.seed!(1)
    gs = [z -> 1 - sum(abs2, z)]
    g_halfdegs = [1]
    (points, Ps) = PolyUtils.interpolate(
        R, halfdeg, num_vars, gs, g_halfdegs)
    d = length(points)
    return (d, Ps)
end

# cones

# Nonnegative
function test_oracles(C::Type{<:Cones.Nonnegative})
    for d in [1, 2, 6]
        test_oracles(C(d))
    end
end

function test_barrier(C::Type{<:Cones.Nonnegative})
    barrier = (s -> -sum(log, s))
    test_barrier(C(3), barrier)
end

show_time_alloc(C::Type{<:Cones.Nonnegative}) = show_time_alloc(C(9))

# PosSemidefTri
function test_oracles(C::Type{Cones.PosSemidefTri{T, R}}) where {T, R}
    for dW in [1, 2, 3, 5]
        test_oracles(C(Cones.svec_length(R, dW)))
    end
end

function test_barrier(C::Type{Cones.PosSemidefTri{T, R}}) where {T, R}
    dW = 3
    barrier(s) = -logdet_pd(Hermitian(new_herm(s, dW, R), :U))
    test_barrier(C(Cones.svec_length(R, dW)), barrier)
end

show_time_alloc(C::Type{Cones.PosSemidefTri{T, R}}) where {T, R} =
    show_time_alloc(C(Cones.svec_length(R, 4)))

# DoublyNonnegativeTri
function test_oracles(C::Type{Cones.DoublyNonnegativeTri{T}}) where T
    for dW in [1, 2, 5]
        test_oracles(C(Cones.svec_length(dW)), init_tol = sqrt(eps(T)))
    end
    for dW in [10, 20]
        test_oracles(C(Cones.svec_length(dW)), init_tol = sqrt(eps(T)),
        init_only = true)
    end
end

function test_barrier(C::Type{Cones.DoublyNonnegativeTri{T}}) where T
    dW = 3
    function barrier(s)
        W = new_herm(s, dW, T)
        offdiags = vcat([Cones.svec_length(i - 1) .+ (1:(i - 1)) for i in 2:dW]...)
        return -logdet_pd(Hermitian(W, :U)) - sum(log, s[offdiags])
    end
    test_barrier(C(Cones.svec_length(dW)), barrier)
end

show_time_alloc(C::Type{<:Cones.DoublyNonnegativeTri}) = show_time_alloc(C(10))

# PosSemidefTriSparse
function test_oracles(C::Type{<:Cones.PosSemidefTriSparse})
    for dW in [1, 2, 10, 25, 40]
        test_oracles(C(dW, rand_sppsd_pattern(dW)...))
    end
end

function test_barrier(C::Type{<:Cones.PosSemidefTriSparse{<:Cones.PSDSparseImpl, T, T}}) where T
    dW = 20
    (row_idxs, col_idxs) = rand_sppsd_pattern(dW)
    invrt2 = inv(sqrt(T(2)))
    function barrier(s)
        scal_s = copy(s)
        scal_s[row_idxs .!= col_idxs] .*= invrt2
        W = Matrix(sparse(row_idxs, col_idxs, scal_s, dW, dW))
        return -logdet_pd(Hermitian(W, :L))
    end
    test_barrier(C(dW, row_idxs, col_idxs), barrier)
end

function test_barrier(C::Type{<:Cones.PosSemidefTriSparse{<:Cones.PSDSparseImpl, T, Complex{T}}}) where T
    dW = 20
    (row_idxs, col_idxs) = rand_sppsd_pattern(dW)
    invrt2 = inv(sqrt(T(2)))
    function barrier(s)
        scal_s = zeros(Complex{eltype(s)}, length(row_idxs))
        idx = 1
        for i in eachindex(scal_s)
            if row_idxs[i] == col_idxs[i]
                scal_s[i] = s[idx]
                idx += 1
            else
                scal_s[i] = invrt2 * Complex(s[idx], s[idx + 1])
                idx += 2
            end
        end
        W = Matrix(sparse(row_idxs, col_idxs, scal_s, dW, dW))
        return -logdet_pd(Hermitian(W, :L))
    end
    test_barrier(C(dW, row_idxs, col_idxs), barrier)
end

show_time_alloc(C::Type{<:Cones.PosSemidefTriSparse}) =
    show_time_alloc(C(15, rand_sppsd_pattern(15)...))

# LinMatrixIneq
function test_oracles(C::Type{Cones.LinMatrixIneq{T}}) where T
    Random.seed!(1)
    Rd_list = [[T, T], [T, Complex{T}], [Complex{T}, T, T]]
    for ds in [2, 3, 4], Rd in Rd_list
        test_oracles(C(rand_herms(ds, Rd, T)), noise = T(1e-2), init_tol = Inf)
    end
end

function test_barrier(C::Type{Cones.LinMatrixIneq{T}}) where T
    Random.seed!(1)
    Ps = rand_herms(2, [T, Complex{T}], T)
    barrier(s) = -logdet_pd(Hermitian(sum(s[i] * Ps[i] for i in eachindex(s)), :U))
    test_barrier(C(Ps), barrier)
end

show_time_alloc(C::Type{Cones.LinMatrixIneq{T}}) where T =
    show_time_alloc(C(rand_herms(3, [T, Complex{T}, T, Complex{T}], T)))

# EpiNormInf
function test_oracles(C::Type{Cones.EpiNormInf{T, R}}) where {T, R}
    for dw in [1, 2, 5]
        test_oracles(C(1 + Cones.vec_length(R, dw)))
    end
end

function test_barrier(C::Type{Cones.EpiNormInf{T, R}}) where {T, R}
    dw = 2
    function barrier(s)
        u = s[1]
        w = new_vec(s[2:end], dw, R)
        return -sum(log(abs2(u) - abs2(wi)) for wi in w) + (dw - 1) * log(u)
    end
    test_barrier(C(1 + Cones.vec_length(R, dw)), barrier)
end

show_time_alloc(C::Type{<:Cones.EpiNormInf}) = show_time_alloc(C(9))

# EpiNormEucl
function test_oracles(C::Type{<:Cones.EpiNormEucl})
    for dw in [1, 2, 5]
        test_oracles(C(1 + dw))
    end
end

function test_barrier(C::Type{<:Cones.EpiNormEucl})
    function barrier(s)
        (u, w) = (s[1], s[2:end])
        return -log(abs2(u) - sum(abs2, w))
    end
    test_barrier(C(3), barrier)
end

show_time_alloc(C::Type{<:Cones.EpiNormEucl}) = show_time_alloc(C(9))

# EpiPerSquare
function test_oracles(C::Type{<:Cones.EpiPerSquare})
    for dw in [1, 2, 5]
        test_oracles(C(2 + dw))
    end
end

function test_barrier(C::Type{<:Cones.EpiPerSquare})
    function barrier(s)
        (u, v, w) = (s[1], s[2], s[3:end])
        return -log(2 * u * v - sum(abs2, w))
    end
    test_barrier(C(4), barrier)
end

show_time_alloc(C::Type{<:Cones.EpiPerSquare}) = show_time_alloc(C(9))

# EpiNormSpectral
function test_oracles(C::Type{<:Cones.EpiNormSpectral})
    for (dr, ds) in [(1, 1), (1, 2), (2, 2), (2, 4), (3, 4)]
        test_oracles(C(dr, ds))
    end
end

function test_barrier(C::Type{Cones.EpiNormSpectral{T, R}}) where {T, R}
    (dr, ds) = (2, 3)
    function barrier(s)
        u = s[1]
        W = reshape(new_vec(s[2:end], dr * ds, R), dr, ds)
        return -logdet_pd(Hermitian(abs2(u) * I - W * W')) + (dr - 1) * log(u)
    end
    test_barrier(C(dr, ds), barrier)
end

show_time_alloc(C::Type{<:Cones.EpiNormSpectral}) = show_time_alloc(C(2, 2))

# MatrixEpiPerSquare
function test_oracles(C::Type{<:Cones.MatrixEpiPerSquare})
    for (dr, ds) in [(1, 1), (1, 2), (2, 2), (2, 4), (3, 4)]
        test_oracles(C(dr, ds))
    end
end

function test_barrier(C::Type{Cones.MatrixEpiPerSquare{T, R}}) where {T, R}
    (dr, ds) = (2, 2)
    du = Cones.svec_length(R, dr)
    function barrier(s)
        U = new_herm(s[1:du], dr, R)
        v = s[du + 1]
        W = reshape(new_vec(s[(du + 2):end], dr * ds, R), dr, ds)
        return -logdet_pd(Hermitian(2 * v * U - W * W', :U)) + (dr - 1) * log(v)
    end
    test_barrier(C(dr, ds), barrier)
end

show_time_alloc(C::Type{<:Cones.MatrixEpiPerSquare}) = show_time_alloc(C(2, 2))

# GeneralizedPower
function test_oracles(C::Type{Cones.GeneralizedPower{T}}) where T
    for (du, dw) in [(2, 1), (3, 2), (4, 1), (2, 4)]
        test_oracles(C(rand_powers(T, du), dw))
    end
end

function test_barrier(C::Type{Cones.GeneralizedPower{T}}) where T
    (du, dw) = (2, 2)
    α = rand_powers(T, du)
    function barrier(s)
        (u, w) = (s[1:du], s[(du + 1):end])
        return -log(exp(2 * sum(α[i] * log(u[i]) for i in eachindex(u))) -
            sum(abs2, w)) - sum((1 - α[i]) * log(u[i]) for i in eachindex(u))
    end
    test_barrier(C(α, dw), barrier)
end

show_time_alloc(C::Type{Cones.GeneralizedPower{T}}) where T =
    show_time_alloc(C(rand_powers(T, 4), 5))

# HypoPowerMean
function test_oracles(C::Type{Cones.HypoPowerMean{T}}) where T
    for dw in [1, 2, 5]
        test_oracles(C(rand_powers(T, dw)), init_tol = 1e-2)
    end
    for dw in [15, 40, 100]
        test_oracles(C(rand_powers(T, dw)), init_tol = 1e-1, init_only = true)
    end
end

function test_barrier(C::Type{Cones.HypoPowerMean{T}}) where T
    α = rand_powers(T, 3)
    function barrier(s)
        (u, w) = (s[1], s[2:end])
        return -log(exp(sum(α[i] * log(w[i]) for i in eachindex(w))) - u) -
            sum(log, w)
    end
    test_barrier(C(α), barrier)
end

show_time_alloc(C::Type{Cones.HypoPowerMean{T}}) where T =
    show_time_alloc(C(rand_powers(T, 8)))

# HypoGeoMean
function test_oracles(C::Type{<:Cones.HypoGeoMean})
    for dw in [1, 2, 5]
        test_oracles(C(1 + dw))
    end
end

function test_barrier(C::Type{<:Cones.HypoGeoMean})
    function barrier(s)
        (u, w) = (s[1], s[2:end])
        sumlogw = sum(log, w)
        return -log(exp(sumlogw / length(w)) - u) - sumlogw
    end
    test_barrier(C(3), barrier)
end

show_time_alloc(C::Type{<:Cones.HypoGeoMean}) = show_time_alloc(C(9))

# HypoRootdetTri
function test_oracles(C::Type{Cones.HypoRootdetTri{T, R}}) where {T, R}
    for dW in [1, 2, 4]
        test_oracles(C(1 + Cones.svec_length(R, dW)))
    end
end

function test_barrier(C::Type{Cones.HypoRootdetTri{T, R}}) where {T, R}
    dW = 3
    function barrier(s)
        (u, w) = (s[1], s[2:end])
        logdet_W = logdet_pd(new_herm(w, dW, R))
        return -log(exp(logdet_W / dW) - u) - logdet_W
    end
    test_barrier(C(1 + Cones.svec_length(R, dW)), barrier)
end

show_time_alloc(C::Type{Cones.HypoRootdetTri{T, R}}) where {T, R} =
    show_time_alloc(C(1 + Cones.svec_length(R, 3)))

# HypoPerLog
function test_oracles(C::Type{<:Cones.HypoPerLog})
    for dw in [1, 2, 5]
        test_oracles(C(2 + dw), init_tol = 1e-5)
    end
    for dw in [15, 40, 100]
        test_oracles(C(2 + dw), init_tol = 1e-1, init_only = true)
    end
end

function test_barrier(C::Type{<:Cones.HypoPerLog})
    function barrier(s)
        (u, v, w) = (s[1], s[2], s[3:end])
        return -log(v * sum(log(wi / v) for wi in w) - u) - log(v) - sum(log, w)
    end
    test_barrier(C(4), barrier)
end

show_time_alloc(C::Type{<:Cones.HypoPerLog}) = show_time_alloc(C(10))

# HypoPerLogdetTri
function test_oracles(C::Type{Cones.HypoPerLogdetTri{T, R}}) where {T, R}
    for dW in [1, 2, 4]
        test_oracles(C(2 + Cones.svec_length(R, dW)), init_tol = 1e-4)
    end
    for dW in [8, 12]
        test_oracles(C(2 + Cones.svec_length(R, dW)), init_tol = 1e-1, init_only = true)
    end
end

function test_barrier(C::Type{Cones.HypoPerLogdetTri{T, R}}) where {T, R}
    dW = 3
    function barrier(s)
        (u, v, w) = (s[1], s[2], s[3:end])
        W = new_herm(w, dW, R)
        return -log(v * logdet_pd(W / v) - u) - log(v) - logdet_pd(W)
    end
    test_barrier(C(2 + Cones.svec_length(R, dW)), barrier)
end

show_time_alloc(C::Type{Cones.HypoPerLogdetTri{T, R}}) where {T, R} =
    show_time_alloc(C(2 + Cones.svec_length(R, 3)))

# EpiPerSepSpectral
function test_oracles(C::Type{<:Cones.EpiPerSepSpectral})
    for d in [1, 2, 3, 6], h_fun in sep_spectral_funs
        test_oracles(C(h_fun, d), init_tol = Inf)
    end
end

function test_barrier(C::Type{<:Cones.EpiPerSepSpectral{<:Cones.VectorCSqr}})
    for h_fun in sep_spectral_funs
        function barrier(s)
            (u, v, w) = (s[1], s[2], s[3:end])
            return -log(u - v * Cones.h_val(w ./ v, h_fun)) - log(v) - sum(log, w)
        end
        test_barrier(C(h_fun, 3), barrier)
    end
end

function test_barrier(C::Type{<:Cones.EpiPerSepSpectral{Cones.MatrixCSqr{T, R}}}) where {T, R}
    dW = 3
    for h_fun in sep_spectral_funs
        function barrier(s)
            (u, v, w) = (s[1], s[2], s[3:end])
            Wλ = eigen(new_herm(w, dW, R)).values
            return -log(u - v * Cones.h_val(Wλ ./ v, h_fun)) - log(v) - sum(log, Wλ)
        end
        test_barrier(C(h_fun, dW), barrier)
    end
end

show_time_alloc(C::Type{<:Cones.EpiPerSepSpectral{<:Cones.VectorCSqr}}) =
    show_time_alloc(C(first(sep_spectral_funs), 9))

show_time_alloc(C::Type{<:Cones.EpiPerSepSpectral{<:Cones.MatrixCSqr}}) =
    show_time_alloc(C(first(sep_spectral_funs), 4))

# EpiRelEntropy
function test_oracles(C::Type{<:Cones.EpiRelEntropy})
    for dw in [1, 2, 4]
        test_oracles(C(1 + 2 * dw), init_tol = 1e-5)
    end
    for dw in [15, 40, 100]
        test_oracles(C(1 + 2 * dw), init_tol = 1e-1, init_only = true)
    end
end

function test_barrier(C::Type{<:Cones.EpiRelEntropy})
    dw = 2
    function barrier(s)
        (u, v, w) = (s[1], s[2:(1 + dw)], s[(2 + dw):end])
        return -log(u - sum(wi * log(wi / vi) for (vi, wi) in zip(v, w))) -
            sum(log, v) - sum(log, w)
    end
    test_barrier(C(5), barrier)
end

show_time_alloc(C::Type{<:Cones.EpiRelEntropy}) = show_time_alloc(C(9))

# EpiTrRelEntropyTri
function test_oracles(C::Type{<:Cones.EpiTrRelEntropyTri})
    for dW in [1, 2, 4]
        test_oracles(C(1 + 2 * Cones.svec_length(dW)), init_tol = 1e-4)
    end
    for dW in [6, 10]
        test_oracles(C(1 + 2 * Cones.svec_length(dW)), init_tol = 1e-1,
            init_only = true)
    end
end

function test_barrier(C::Type{Cones.EpiTrRelEntropyTri{T}}) where T
    dW = 3
    dw = Cones.svec_length(dW)
    function barrier(s)
        (u, v, w) = (s[1], s[1 .+ (1:dw)], s[(2 + dw):end])
        V = new_herm(v, dW, T)
        W = new_herm(w, dW, T)
        return -log(u - dot(W, log(W) - log(V))) - logdet_pd(V) - logdet_pd(W)
    end
    test_barrier(C(1 + 2 * dw), barrier, TFD = BigFloat)
end

show_time_alloc(C::Type{<:Cones.EpiTrRelEntropyTri}) = show_time_alloc(C(13))

# WSOSInterpNonnegative
function test_oracles(C::Type{Cones.WSOSInterpNonnegative{T, R}}) where {T, R}
    for (num_vars, halfdeg) in [(1, 1), (1, 3), (2, 1), (2, 2), (3, 1)]
        (d, Ps) = rand_interp(num_vars, halfdeg, R)
        test_oracles(C(d, Ps), init_tol = Inf)
    end
end

function test_barrier(C::Type{Cones.WSOSInterpNonnegative{T, R}}) where {T, R}
    (d, Ps) = rand_interp(2, 1, R)
    barrier(s) = -sum(logdet_pd(Hermitian(P' * Diagonal(s) * P)) for P in Ps)
    test_barrier(C(d, Ps), barrier)
end

show_time_alloc(C::Type{Cones.WSOSInterpNonnegative{T, R}}) where {T, R} =
    show_time_alloc(C(rand_interp(3, 1, R)...))

# WSOSInterpPosSemidefTri
function test_oracles(C::Type{Cones.WSOSInterpPosSemidefTri{T}}) where T
    for (num_vars, halfdeg, ds) in [(1, 1, 1), (1, 1, 4), (2, 2, 1), (3, 1, 2)]
        (d, Ps) = rand_interp(num_vars, halfdeg, T)
        test_oracles(C(ds, d, Ps), init_tol = Inf)
    end
end

function test_barrier(C::Type{Cones.WSOSInterpPosSemidefTri{T}}) where T
    (d, Ps) = rand_interp(1, 1, T)
    ds = 3
    invrt2 = inv(sqrt(T(2)))
    function barrier(s)
        function ldΛP(P)
            dt = size(P, 2)
            Λ = similar(s, ds * dt, ds * dt)
            for i in 1:ds, j in 1:i
                sij = s[Cones.block_idxs(d, Cones.svec_idx(i, j))]
                Λij = P' * Diagonal(sij) * P
                if i != j
                    Λij .*= invrt2
                end
                Λ[Cones.block_idxs(dt, i), Cones.block_idxs(dt, j)] = Λij
            end
            return -logdet_pd(Hermitian(Λ, :L))
        end
        return sum(ldΛP, Ps)
    end
    test_barrier(C(ds, d, Ps), barrier)
end

show_time_alloc(C::Type{Cones.WSOSInterpPosSemidefTri{T}}) where T =
    show_time_alloc(C(2, rand_interp(2, 1, T)...))

# WSOSInterpEpiNormEucl
function test_oracles(C::Type{Cones.WSOSInterpEpiNormEucl{T}}) where T
    for (num_vars, halfdeg, ds) in [(1, 1, 1), (1, 2, 3), (2, 2, 2), (3, 1, 1)]
        (d, Ps) = rand_interp(num_vars, halfdeg, T)
        test_oracles(C(1 + ds, d, Ps), init_tol = Inf)
    end
end

function test_barrier(C::Type{Cones.WSOSInterpEpiNormEucl{T}}) where T
    (d, Ps) = rand_interp(1, 1, T)
    ds = 2
    invrt2 = inv(sqrt(T(2)))
    function barrier(s)
        function ldΛP(P)
            Λ = P' * Diagonal(s[1:d]) * P
            Λ1fact = cholesky(Hermitian(Λ, :L))
            PL1 = Λ1fact.L \ P'
            for i in 1:ds
                ΛLi = PL1 * Diagonal(s[Cones.block_idxs(d, 1 + i)]) * P
                Λ -= ΛLi' * ΛLi
            end
            return -logdet(Λ1fact) - logdet_pd(Hermitian(Λ))
        end
        return sum(ldΛP, Ps)
    end
    test_barrier(C(1 + ds, d, Ps), barrier)
end

show_time_alloc(C::Type{Cones.WSOSInterpEpiNormEucl{T}}) where T =
    show_time_alloc(C(3, rand_interp(2, 1, T)...))

# WSOSInterpEpiNormOne
function test_oracles(C::Type{Cones.WSOSInterpEpiNormOne{T}}) where T
    for (num_vars, halfdeg, ds) in [(1, 1, 1), (1, 2, 3), (2, 2, 2), (3, 1, 1)]
        (d, Ps) = rand_interp(num_vars, halfdeg, T)
        test_oracles(C(1 + ds, d, Ps), init_tol = Inf)
    end
end

function test_barrier(C::Type{Cones.WSOSInterpEpiNormOne{T}}) where T
    (d, Ps) = rand_interp(1, 1, T)
    ds = 2
    invrt2 = inv(sqrt(T(2)))
    function barrier(s)
        function ldΛP(P)
            Λ = P' * Diagonal(s[1:d]) * P
            Λ1fact = cholesky(Hermitian(Λ, :L))
            PL1 = Λ1fact.L \ P'
            ΛLs = [PL1 * Diagonal(s[Cones.block_idxs(d, 1 + i)]) * P for i in 1:ds]
            Λs = [Hermitian(Λ - ΛLi' * ΛLi) for ΛLi in ΛLs]
            return -logdet(Λ1fact) - sum(logdet_pd, Λs)
        end
        return sum(ldΛP, Ps)
    end
    test_barrier(C(1 + ds, d, Ps), barrier)
end

show_time_alloc(C::Type{Cones.WSOSInterpEpiNormOne{T}}) where T =
    show_time_alloc(C(3, rand_interp(2, 1, T)...))


text:
using NumericExtensions
using Base.Test

a1 = 2 * rand(8) - 1.0
a2 = 2 * rand(8, 7) - 1.0
a3 = 2 * rand(8, 7, 6) - 1.0
a4 = 2 * rand(8, 7, 6, 5) - 1.0

b1 = 2 * rand(8) - 1.0
b2 = 2 * rand(8, 7) - 1.0
b3 = 2 * rand(8, 7, 6) - 1.0
b4 = 2 * rand(8, 7, 6, 5) - 1.0

foldlsum(xs...) = foldl(Add(), 0., xs...)
foldlsumfdiff(f, x1, x2, dim) = foldl_fdiff(Add(), 0., f, x1, x2, dim)

@test_approx_eq foldlsum(a1, 1) sum(a1, 1)
@test_approx_eq foldlsum(a2, 1) sum(a2, 1)
@test_approx_eq foldlsum(a2, 2) sum(a2, 2)
@test_approx_eq foldlsum(a3, 1) sum(a3, 1)
@test_approx_eq foldlsum(a3, 2) sum(a3, 2)
@test_approx_eq foldlsum(a3, 3) sum(a3, 3)
@test_approx_eq foldlsum(a4, 1) sum(a4, 1)
@test_approx_eq foldlsum(a4, 2) sum(a4, 2)
@test_approx_eq foldlsum(a4, 3) sum(a4, 3)
@test_approx_eq foldlsum(a4, 4) sum(a4, 4)

do_foldlsum!(a, dim) = foldl!(zeros(Base.reduced_dims(size(a), dim)), Add(), 0., a, dim)

@test_approx_eq do_foldlsum!(a1, 1) sum(a1, 1)
@test_approx_eq do_foldlsum!(a2, 1) sum(a2, 1)
@test_approx_eq do_foldlsum!(a2, 2) sum(a2, 2)
@test_approx_eq do_foldlsum!(a3, 1) sum(a3, 1)
@test_approx_eq do_foldlsum!(a3, 2) sum(a3, 2)
@test_approx_eq do_foldlsum!(a3, 3) sum(a3, 3)
@test_approx_eq do_foldlsum!(a4, 1) sum(a4, 1)
@test_approx_eq do_foldlsum!(a4, 2) sum(a4, 2)
@test_approx_eq do_foldlsum!(a4, 3) sum(a4, 3)
@test_approx_eq do_foldlsum!(a4, 4) sum(a4, 4)

@test_approx_eq foldlsum(Abs2Fun(), a1, 1) sum(abs2(a1), 1)
@test_approx_eq foldlsum(Abs2Fun(), a2, 1) sum(abs2(a2), 1)
@test_approx_eq foldlsum(Abs2Fun(), a2, 2) sum(abs2(a2), 2)
@test_approx_eq foldlsum(Abs2Fun(), a3, 1) sum(abs2(a3), 1)
@test_approx_eq foldlsum(Abs2Fun(), a3, 2) sum(abs2(a3), 2)
@test_approx_eq foldlsum(Abs2Fun(), a3, 3) sum(abs2(a3), 3)
@test_approx_eq foldlsum(Abs2Fun(), a4, 1) sum(abs2(a4), 1)
@test_approx_eq foldlsum(Abs2Fun(), a4, 2) sum(abs2(a4), 2)
@test_approx_eq foldlsum(Abs2Fun(), a4, 3) sum(abs2(a4), 3)
@test_approx_eq foldlsum(Abs2Fun(), a4, 4) sum(abs2(a4), 4)

@test_approx_eq foldlsum(Multiply(), a1, b1, 1) sum(a1 .* b1, 1)
@test_approx_eq foldlsum(Multiply(), a2, b2, 1) sum(a2 .* b2, 1)
@test_approx_eq foldlsum(Multiply(), a2, b2, 2) sum(a2 .* b2, 2)
@test_approx_eq foldlsum(Multiply(), a3, b3, 1) sum(a3 .* b3, 1)
@test_approx_eq foldlsum(Multiply(), a3, b3, 2) sum(a3 .* b3, 2)
@test_approx_eq foldlsum(Multiply(), a3, b3, 3) sum(a3 .* b3, 3)
@test_approx_eq foldlsum(Multiply(), a4, b4, 1) sum(a4 .* b4, 1)
@test_approx_eq foldlsum(Multiply(), a4, b4, 2) sum(a4 .* b4, 2)
@test_approx_eq foldlsum(Multiply(), a4, b4, 3) sum(a4 .* b4, 3)
@test_approx_eq foldlsum(Multiply(), a4, b4, 4) sum(a4 .* b4, 4)

@test_approx_eq foldlsumfdiff(Abs2Fun(), a1, b1, 1) sum(abs2(a1 - b1), 1)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a2, b2, 1) sum(abs2(a2 - b2), 1)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a2, b2, 2) sum(abs2(a2 - b2), 2)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a3, b3, 1) sum(abs2(a3 - b3), 1)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a3, b3, 2) sum(abs2(a3 - b3), 2)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a3, b3, 3) sum(abs2(a3 - b3), 3)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a4, b4, 1) sum(abs2(a4 - b4), 1)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a4, b4, 2) sum(abs2(a4 - b4), 2)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a4, b4, 3) sum(abs2(a4 - b4), 3)
@test_approx_eq foldlsumfdiff(Abs2Fun(), a4, b4, 4) sum(abs2(a4 - b4), 4)


text:
is_ch_type(::Val{:FixedString})  = true
result_type(::Val{:FixedString}, len_str::String)  = Vector{String}

function read_col_data(sock::ClickHouseSock, num_rows::VarUInt,
                        ::Val{:FixedString}, len_str::String)
    len = parse(Int64, len_str)
    result = Vector{String}(undef, UInt64(num_rows))
    for i in 1:UInt64(num_rows)
        result[i] = String(Base.read(sock.io, len))
    end
    return result
end

function str_to_fix_len(str::String, len)
    length(str) == len && return str
    length(str) > len && error("Too large value \"$str\" for FixedString($len)")
    return str * repeat(" ", len - length(str))
end

function write_col_data(sock::ClickHouseSock, data::AbstractVector{String},
                         ::Val{:FixedString}, len_str::String)
    len = parse(Int64, len_str)
    for str in data
        Base.write(sock.io, str_to_fix_len(str, len))
    end
end


text:
### Solve linear system by applying the implicit ULV factorization
#
# as seen in
# Chandrasekaran, S., Gu, M., & Pals, T. (2006). A Fast $ULV$ Decomposition Solver for Hierarchically Semiseparable Representations.
# SIAM Journal on Matrix Analysis and Applications, 28(3), 603–622. https://doi.org/10.1137/S0895479803436652
#
# Written by Boris Bonev, Nov. 2020

## function for direct solution using the implicit ULV factorization
function ulvfactsolve(hssA::HssMatrix{T}, b::Matrix{T}) where T
  if isleaf(hssA)
    return hssA.D\b
  else
    z = zeros(size(hssA,2), size(b,2))
    _, _, _, _, _, _, _, QV = _ulvfactsolve!(hssA, b, z, 0; rootnode=true)
    z = _ulvfactsolve_topdown!(QV, z)
    return z
  end
end

# core routine to reduce the rows and triangularize the diagonal block via QR/LQ decompositions and immediately apply them to b
function _ulvreduce!(D::Matrix{T}, U::Matrix{T}, V::Matrix{T}, b::Matrix{T}) where T
  T <: Complex ? adj = 'C' : adj = 'T'
  m, n = size(D)
  m == size(U,1) || throw(DimensionMismatch("D and U have different first dimensions"))
  n == size(V,1) || throw(DimensionMismatch("D and V have different second dimensions"))
  k = min(size(U, 2), m)
  nk = min(m-k,n)
  ind = 1:m-k
  cind = m-k+1:m
  # can't be compressed, exit early
  if k >= m
    # @warn "Encountered a full-rank block with k=$(k). Clustering might not yield best performance!"
    u = zeros(size(V,2), size(b,2))
    zloc = Matrix{T}(undef, 0, size(b,2))
    lqf = (Matrix{T}(undef, 0, n), Vector{T}(undef, 0))
  else
    # form QL decomposition of the row generators and apply it
    qlf = geqlf!(U);
    U = tril(U[end-k+1:end,:]) # k x k block
    ormql!('L', adj, qlf..., D) # transform the diagonal block
    ormql!('L', adj, qlf..., b) # transform the right-hand side
    # Form the LQ decomposition of the first m-k rows of D
    lqf = gelqf!(D[1:end-k,:])
    L1 = tril(lqf[1])
    L1 = L1[:,1:nk]
    L2 = ormlq!('R', adj, lqf..., D[end-k+1:end,:]) # update the bottom block of the diagonal block
    zloc = trsm('L', 'L', 'N', 'N', 1., L1, b[ind,:])
    b = b[cind, :] .- L2[:,1:nk] * zloc # remove contribution in the uncompressed parts
    V = ormlq!('L', 'N', lqf..., V) # compute the updated off-diagonal generators on the right
    u = V[ind,:]' * zloc # compute update vector to be passed on
    # pass on uncompressed parts of the problem
    D = L2[:, nk+1:end]
    V = V[nk+1:end, :]
  end
  return D, U, V, b, zloc, u, m-k, nk, lqf
end

# tree traversal that handles the hierarchical ULV factorization
function _ulvfactsolve!(hssA::HssMatrix{T}, b::Matrix{T}, z::Matrix{T}, co::Int; rootnode=false) where T
  if isleaf(hssA)
    cols = collect(co .+ (1:size(hssA,2)))
    D = copy(hssA.D); U = copy(hssA.U); V = copy(hssA.V)
    D, U, V, b, zloc, u, mk, nk, lqf =_ulvreduce!(D, U, V, b)
    z[cols[1:mk], :] = zloc
    QV = BinaryNode((cols, lqf...))
  else
    m1, n1 = hssA.sz1; m2, n2 = hssA.sz2
    b1 = b[1:m1, :]; b2 = b[m1+1:end, :] # performance could be further improved by getting rid of those allocations
    b1, u1, D1, U1, V1, cols1, nk1, QV1 = _ulvfactsolve!(hssA.A11, b1, z, co)
    b2, u2, D2, U2, V2, cols2, nk2, QV2 = _ulvfactsolve!(hssA.A22, b2, z, co+n1)

    # merge nodes to form new diagonal block
    b = [b1; b2] .- [U1*hssA.B12*u2; U2*hssA.B21*u1]
    D = [D1 U1*hssA.B12*V2'; U2*hssA.B21*V1' D2]
    cols = [cols1[nk1+1:end]; cols2[nk2+1:end]] # to re-adjust local numbering

    U = [U1*hssA.R1; U2*hssA.R2]
    V = [V1*hssA.W1; V2*hssA.W2]

    # now if this is the rootnode we should directly solve?
    if rootnode
      z[cols, :] = D\b
      u = Matrix{T}(undef,0,size(b,2))
      mk, nk = size(D)
      QV = BinaryNode{Tuple{Vector{Int}, Matrix{T}, Vector{T}}}()
    else
      D, U, V, b, zloc, u, mk, nk, lqf = _ulvreduce!(D, U, V, b)
      u .= u .+ hssA.W1'*u1 .+ hssA.W2'*u2
      z[cols[1:mk], :] = zloc
      QV = BinaryNode((cols, lqf...))
    end
    QV.left = QV1; QV.right = QV2
  end
  return b, u, D, U, V, cols, nk, QV
end

function _ulvfactsolve_topdown!(QV::BinaryNode{Tuple{Vector{Int}, Matrix{T}, Vector{T}}}, z::Matrix{T}) where T
  T <: Complex ? adj = 'C' : adj = 'T'
  if !isnothing(QV.data)
    (cols, Q, tau) = QV.data
    z[cols, :] = ormlq!('L', adj, Q, tau, z[cols, :])
  end
  if !isnothing(QV.left) z = _ulvfactsolve_topdown!(QV.left, z) end
  if !isnothing(QV.right) z = _ulvfactsolve_topdown!(QV.right, z) end
  return z
end


text:
export TDLearner, DoubleLearner, DifferentialTDLearner, TDλReturnLearner

using .Utils: discount_rewards, discount_rewards_reduced
using Flux.Optimise: Descent
using LinearAlgebra: dot

"""
    TDLearner(approximator::Tapp, γ::Float64, optimizer::Float64; n::Int=0) where {Tapp<:AbstractVApproximator}
    TDLearner(approximator::Tapp, γ::Float64, optimizer::Float64; n::Int=0, method::Symbol=:SARSA) where {Tapp<:AbstractQApproximator}

The `TDLearner`(Temporal Difference Learner) use the latest `n` step experiences to update the `approximator`.
Note that `n` starts with `0`, which means looking forward for the next `n` steps.
`γ` is the discount rate of experience.
`optimizer` is the learning rate.

For [`AbstractVApproximator`](@ref), the only supported update method is `:SRS`, which means
only **S**tates, **R**ewards and next_**S**ates are used to update the `approximator`.

For [`AbstractQApproximator`](@ref), the following methods are supported:

- `:SARS` (aka Q-Learning)
- `:SARSA`
- `:ExpectedSARSA`
"""
mutable struct TDLearner{Tapp<:AbstractApproximator,method,O} <: AbstractLearner
    approximator::Tapp
    γ::Float64
    optimizer::O
    n::Int

    function TDLearner(
        ;
        approximator::Tapp,
        γ = 1.0,
        optimizer = Descent(1.0),
        n = 0,
        method::Symbol = :SARSA,
    ) where {Tapp<:AbstractApproximator}
        if approximator isa AbstractVApproximator
            if method != :SRS
                throw(ArgumentError("method [$method] is unsupported for a value approximator"))
            end
        elseif approximator isa AbstractQApproximator
            if !in(method, [:SARSA, :SARS, :ExpectedSARSA])
                throw(ArgumentError("Supported methods are $supported_methods , your input is $method"))
            end
        else
            throw(ArgumentError("unknown approximator"))
        end
        new{Tapp,method,typeof(optimizer)}(approximator, γ, optimizer, n)
    end
end

"""
    DoubleLearner(;L1, L2)

For now, this only supports [`TDLearner`](@ref) for `L1` and `L2`.
"""
Base.@kwdef struct DoubleLearner{T<:TDLearner} <: AbstractLearner
    L1::T
    L2::T
end

#####
# DifferentialTDLearner
#####

"""
    DifferentialTDLearner(;approximator::A, α::Float64, β::Float64, R̄::Float64 = 0.0, n::Int = 0)
"""
Base.@kwdef mutable struct DifferentialTDLearner{A<:AbstractApproximator} <: AbstractLearner
    approximator::A
    α::Float64
    β::Float64
    R̄::Float64 = 0.0
    n::Int = 0
end

function update!(learner::DifferentialTDLearner, transition)
    states, actions, rewards, terminals, next_states, next_actions = transition
    n, α, β, Q = learner.n, learner.α, learner.β, learner.approximator
    if length(states) ≥ n + 1
        s, a = states[1], actions[1]
        s′, a′ = next_states[end], next_actions[end]
        δ = sum(r -> r - learner.R̄, rewards) + Q(s′, a′) - Q(s, a)
        learner.R̄ += β * δ
        update!(Q, (s, a) => α * δ)
    end
end

extract_transitions(buffer, learner::DoubleLearner) =
    extract_transitions(buffer, learner.L1)

function update!(learner::DoubleLearner, args...)
    rand(Bool) ? update!(learner.L1, learner.L2, args...) :
    update!(learner.L2, learner.L1, args...)
end

(learner::DoubleLearner)(obs::Observation) = learner.L1(obs) .+ learner.L2(obs)

update!(learner::TDLearner{<:AbstractVApproximator,:SRS}, transitions) =
    update!(learner, transitions, nothing)

function update!(learner::TDLearner{<:AbstractVApproximator,:SRS}, transitions, weights)
    states, rewards, terminals, next_states = transitions.states,
        transitions.rewards,
        transitions.terminals,
        transitions.next_states
    n, γ, V, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        cum_weights = isnothing(weights) ? nothing : cumprod!(reverse(weights))
        for (i, G) in enumerate(gains)
            @views s = states[end-length(gains)+i]
            if isnothing(cum_weights)
                update!(V, s => apply!(optimizer, s, G - V(s)))
            else
                update!(V, s => apply!(optimizer, s, cum_weights[i] * (G - V(s))))
            end
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * V(next_states[end])
            @views s = states[end-n]
            w = isnothing(weights) ? 1.0 : reduce(*, weights)
            update!(V, s => apply!(optimizer, s, w * (G - V(s))))
        end
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDLearner{<:AbstractVApproximator,:SRS},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:SARSA}, transitions)
    states, actions, rewards, terminals, next_states, next_actions = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′, a′ = states[end-n],
                actions[end-n],
                next_states[end],
                next_actions[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * Q(s′, a′)
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::Union{TDLearner{<:AbstractQApproximator,:SARSA},DifferentialTDLearner},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            actions = action(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
            next_actions = action(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:ExpectedSARSA}, transitions)
    states, actions, rewards, terminals, next_states, probs_of_a′ = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * dot(Q(s′), probs_of_a′)
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:SARS}, transitions)
    states, actions, rewards, terminals, next_states = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * maximum(Q(s′))  # n starts with 0
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function update!(
    learner::T,
    target_learner::T,
    transitions,
) where {T<:TDLearner{<:AbstractQApproximator,:SARS}}
    states, actions, rewards, terminals, next_states = transitions
    n, γ, Q, Qₜ, optimizer = learner.n,
        learner.γ,
        learner.approximator,
        target_learner.approximator,
        learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * Qₜ(s′, argmax(Q(s′)))
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function priority(transition, learner::TDLearner{<:AbstractQApproximator,:SARS})
    s, a, r, d, s′ = transition
    γ, Q, opt = learner.γ, learner.approximator, learner.optimizer
    error = d ? apply!(opt, (s, a), r - Q(s, a)) :
            apply!(opt, (s, a), r + γ^(learner.n + 1) * maximum(Q(s′)) - Q(s, a))
    abs(error)
end

function update!(
    learner::TDLearner{<:AbstractQApproximator,:SARS},
    model::Union{TimeBasedSampleModel,ExperienceBasedSampleModel};
    plan_step = 1,
)
    for _ = 1:plan_step
        # @assert learner.n == 0 "n must be 0 here"
        transitions = extract_transitions(model, learner)
        if !isnothing(transitions)
            update!(learner, transitions)
        end
    end
end

function extract_transitions(
    model::Union{ExperienceBasedSampleModel,TimeBasedSampleModel},
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    if length(model.experiences) > 0
        map(Base.vect, sample(model))
    else
        nothing
    end
end

function update!(
    learner::TDLearner{<:AbstractQApproximator,:SARS},
    model::PrioritizedSweepingSampleModel;
    plan_step = 1,
)
    for _ = 1:plan_step
        # @assert learner.n == 0 "n must be 0 here"
        transitions = extract_transitions(model, learner)
        if !isnothing(transitions)
            update!(learner, transitions)
            s, _, _, _, _ = transitions
            s = s[]  # length(s) is assumed to be 1
            for (s̄, ā, r̄, d̄) in model.predecessors[s]
                P = priority((s̄, ā, r̄, d̄, s), learner)
                if P ≥ model.θ
                    model.PQueue[(s̄, ā)] = P
                end
            end
        end
    end
end

function extract_transitions(
    model::PrioritizedSweepingSampleModel,
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    if length(model.PQueue) > 0
        map(Base.vect, sample(model))
    else
        nothing
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            actions = action(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

#####
# TDλReturnLearner
#####
"""
    TDλReturnLearner(;approximator::Tapp, γ::Float64 = 1.0, α::Float64, λ::Float64)
"""
Base.@kwdef struct TDλReturnLearner{Tapp<:AbstractApproximator} <: AbstractLearner
    approximator::Tapp
    γ::Float64 = 1.0
    α::Float64
    λ::Float64
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDλReturnLearner{<:AbstractVApproximator},
)
    if isfull(buffer)
        @views (
            states = state(buffer)[1:end-1],
            rewards = reward(buffer)[2:end],
            terminals = terminal(buffer)[2:end],
            next_states = state(buffer)[2:end],
        )
    else
        nothing
    end
end

function update!(learner::TDλReturnLearner, transition)
    λ, γ, V, α = learner.λ, learner.γ, learner.approximator, learner.α
    states, rewards, terminals, next_states = transition
    T = length(states)
    for t = 1:T
        G = 0.0
        for n = 1:(T-t)
            G += λ^(n - 1) *
                 (discount_rewards_reduced(@view(rewards[t:t+n-1]), γ) +
                  γ^n * V(next_states[t+n-1]))
        end
        G *= 1 - λ
        G += λ^(T - t) *
             (discount_rewards_reduced(@view(rewards[t:T]), γ) +
              γ^(T - t + 1) * V(next_states[T]))
        sₜ = states[t]
        update!(V, sₜ => α * (G - V(sₜ)))
    end
end


text:
"""
	HalfIntSubArray{T,N,P,I,L}

`N`-dimensional view into a parent array (of type `P`) with an element type `T`,
restricted by a tuple of indices (of type `I`).
`L` is `true` for types that support fast linear indexing, and false otherwise.

Their behavior is analogous to `SubArray`s, except that they enable indexing with
half-integers.

Construct `HalfIntSubArray`s from [`AbstractHalfIntegerArray`](@ref)s
using the `view` function, or equivalently
the `@view` macro.

# Example
```jldoctest
julia> h = HalfIntArray(reshape(1:9, 3, 3), -1:1, -1:1)
3×3 HalfIntArray(reshape(::UnitRange{Int64}, 3, 3), -1:1, -1:1) with eltype Int64 with indices -1:1×-1:1:
 1  4  7
 2  5  8
 3  6  9

julia> hv = @view h[-1:1, 0]
3-element view(HalfIntArray(reshape(::UnitRange{Int64}, 3, 3), -1:1, -1:1), -1:1, 0) with eltype Int64:
 4
 5
 6

julia> hv[1]
4

julia> h = HalfIntArray(reshape(collect(1:4), 2, 2), -1//2:1//2, -1//2:1//2)
2×2 HalfIntArray(::Array{Int64,2}, -1/2:1/2, -1/2:1/2) with eltype Int64 with indices -1/2:1/2×-1/2:1/2:
 1  3
 2  4

julia> hv = @view h[:, 1//2]
2-element view(HalfIntArray(::Array{Int64,2}, -1/2:1/2, -1/2:1/2), :, 1/2) with eltype Int64 with indices -1/2:1/2:
 3
 4

julia> hv[1//2] = 10
10

julia> h
2×2 HalfIntArray(::Array{Int64,2}, -1/2:1/2, -1/2:1/2) with eltype Int64 with indices -1/2:1/2×-1/2:1/2:
 1   3
 2  10
```
"""
struct HalfIntSubArray{T,N,P,I,O,L} <: AbstractHalfIntegerArray{T,N}
	parent :: P
	indices :: I
	offset1:: O
	stride1:: Int

	function HalfIntSubArray{T,N,P,I,O,L}(parent, indices, offset1, stride1) where {T,N,P,I,O,L}
        Base.check_parent_index_match(parent, indices)
        new(parent, indices, offset1, stride1)
    end
end

# Compute the linear indexability of the indices, and combine it with the linear indexing of the parent
function HalfIntSubArray(parent::AbstractArray, indices::Tuple)
    HalfIntSubArray(IndexStyle(Base.viewindexing(indices), IndexStyle(parent)), parent,
    	Base.ensure_indexable(indices), Base.index_dimsum(indices...))
end
function HalfIntSubArray(::IndexCartesian, parent::P, indices::I, ::NTuple{N,Any}) where {P,I,N}
    HalfIntSubArray{eltype(P), N, P, I, Int, false}(parent, indices, 0, 0)
end
function HalfIntSubArray(::IndexLinear, parent::P, indices::I, ::NTuple{N,Any}) where {P,I,N}
    # Compute the stride and offset
    stride1 = Base.compute_stride1(parent, indices)
    offset1 = Base.compute_offset1(parent, stride1, indices)
    HalfIntSubArray{eltype(P), N, P, I, typeof(offset1), true}(parent, indices, offset1, stride1)
end

# Half-integer vectors use Cartesian indexing, the offset will be half-integer in general
function Base.compute_offset1(parent::AbstractHalfIntegerVector, stride1::Integer, I::Tuple{Real})
    HalfInt(I[1]) - stride1
end
Base.compute_stride1(s, inds, I::Tuple{AbstractUnitRange{HalfInt}, Vararg{Any}}) = s

Base.parent(h::HalfIntSubArray) = h.parent

Base.size(V::HalfIntSubArray) = map(n->Int(Base.unsafe_length(n)), axes(V))

Base.similar(V::HalfIntSubArray, T::Type, dims::Dims) = similar(V.parent, T, dims)

Base.sizeof(V::HalfIntSubArray) = length(V) * sizeof(eltype(V))

Base.copy(V::HalfIntSubArray) = V.parent[V.indices...]

Base.dataids(v::HalfIntSubArray) = Base.dataids(parent(v))

@static if VERSION >= v"1.2"
    reindex(V, I) = Base.reindex(V.indices, I)
else
    reindex(V, I) = Base.reindex(V, V.indices, I)
end

_maybeInt(::AbstractHalfIntegerArrayOrWrapper, i::Real) = i
_maybeInt(::AbstractArray, i::Real) = Int(i)

# In general, we simply re-index the parent indices by the provided ones
@inline function Base.getindex(V::HalfIntSubArray{T,N}, I::Vararg{Real,N}) where {T,N}
    @boundscheck checkbounds(V, I...)
    @inbounds r = V.parent[reindex(V, I)...]
    r
end
@propagate_inbounds function Base.getindex(V::HalfIntSubArray{T,N}, I::Vararg{Real}) where {T,N}
	@boundscheck checkbounds(V, I...)
	IN = trimtoN(I, Val(N))
	V[IN...]
end

# linear indexing for slow subarrays is through Cartesian indexing
@inline function Base.getindex(V::HalfIntSubArray{T,N}, i::Real) where {T,N}
	@boundscheck checkbounds(V, i)
    @boundscheck ensureInt(i)
	@inbounds V[CartesianIndicesHalfInt(V)[unsafeInt(i)]]
end

# But SubArrays with fast linear indexing pre-compute a stride and offset
FastSubArray{T,N,P,I,O} = HalfIntSubArray{T,N,P,I,O,true}
@inline function Base.getindex(V::FastSubArray, i::Real)
    @boundscheck checkbounds(V, i)
    ensureInt(i)
    @inbounds r = V.parent[_maybeInt(V.parent, V.offset1 + V.stride1*i)]
    r
end

# We can avoid a multiplication if the first parent index is a Colon or AbstractUnitRange,
# or if all the indices are scalars, i.e. the view is for a single value only
FastContiguousSubArray{T,N,P,I<:Union{Tuple{Union{Base.Slice, AbstractUnitRange}, Vararg{Any}},
                                      Tuple{Vararg{Real}}},O} = HalfIntSubArray{T,N,P,I,O,true}
@inline function Base.getindex(V::FastContiguousSubArray, i::Real)
    @boundscheck checkbounds(V, i)
    ensureInt(i)
    @inbounds r = V.parent[_maybeInt(V.parent, V.offset1 + i)]
    r
end
# For vector views with linear indexing, we disambiguate to favor the stride/offset
# computation as that'll generally be faster than (or just as fast as) re-indexing into a range.
@inline function Base.getindex(V::FastSubArray{<:Any, 1}, i::Real)
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[_maybeInt(V.parent, V.offset1 + V.stride1*i)]
    r
end
@inline function Base.getindex(V::FastContiguousSubArray{<:Any, 1}, i::Real)
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[_maybeInt(V.parent, V.offset1 + i)]
    r
end

# Indexed assignment follows the same pattern as `getindex` above
@inline function Base.setindex!(V::HalfIntSubArray{T,N}, x, I::Vararg{Real,N}) where {T,N}
    @boundscheck checkbounds(V, I...)
    @inbounds V.parent[reindex(V, I)...] = x
    V
end
@propagate_inbounds function Base.setindex!(V::HalfIntSubArray{T,N}, val, I::Vararg{Real}) where {T,N}
	@boundscheck checkbounds(V, I...)
	IN = trimtoN(I, Val(N))
	V[IN...] = val
	V
end
# linear indexing for slow subarrays is through Cartesian indexing
@inline function Base.setindex!(V::HalfIntSubArray{T,N}, val, i::Real) where {T,N}
	@boundscheck checkbounds(V, i)
	ensureInt(i)
	@inbounds V[CartesianIndicesHalfInt(V)[unsafeInt(i)]] = val
	V
end

@inline function Base.setindex!(V::FastSubArray, x, i::Real)
    @boundscheck checkbounds(V, i)
    ensureInt(i)
    @inbounds V.parent[_maybeInt(V.parent, V.offset1 + V.stride1*i)] = x
    V
end
@inline function Base.setindex!(V::FastContiguousSubArray, x, i::Real)
    @boundscheck checkbounds(V, i)
    ensureInt(i)
    @inbounds V.parent[_maybeInt(V.parent, V.offset1 + i)] = x
    V
end
@inline function Base.setindex!(V::FastSubArray{<:Any, 1}, x, i::Real)
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[_maybeInt(V.parent, V.offset1 + V.stride1*i)] = x
    V
end
@inline function Base.setindex!(V::FastContiguousSubArray{<:Any, 1}, x, i::Real)
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[_maybeInt(V.parent, V.offset1 + i)] = x
    V
end

Base.IndexStyle(::Type{<:FastSubArray}) = IndexLinear()
Base.IndexStyle(::Type{<:HalfIntSubArray}) = IndexCartesian()

# indices are taken from the range/vector
# Since bounds-checking is performance-critical and uses
# indices, it's worth optimizing these implementations thoroughly
Base.axes(S::HalfIntSubArray) = _indices_sub(S.indices...)
_indices_sub(::Real, I...) = _indices_sub(I...)
_indices_sub() = ()
function _indices_sub(i1::AbstractArray, I...)
    (Base.unsafe_indices(i1)..., _indices_sub(I...)...)
end

Base.unaliascopy(A::HalfIntSubArray) = typeof(A)(Base.unaliascopy(A.parent), map(Base.unaliascopy, A.indices), A.offset1, A.stride1)

# View

## SubArray creation
# We always assume that the dimensionality of the parent matches the number of
# indices that end up getting passed to it, so we store the parent as a
# ReshapedArray view if necessary. The trouble is that arrays of `CartesianIndex`
# can make the number of effective indices not equal to length(I).
_maybe_reshape_parent(A::AbstractArray, ::NTuple{1, Bool}) = reshape(A, Val(1))
_maybe_reshape_parent(A::AbstractArray{<:Any,1}, ::NTuple{1, Bool}) = reshape(A, Val(1))
_maybe_reshape_parent(A::AbstractArray{<:Any,N}, ::NTuple{N, Bool}) where {N} = A
_maybe_reshape_parent(A::AbstractArray, ::NTuple{N, Bool}) where {N} = reshape(A, Val(N))

function Base.view(A::AbstractHalfIntegerArrayOrWrapper, I::Vararg{Any,N}) where {N}
    J = map(i->Base.unalias(A,i), to_indices(A, I))
    @boundscheck checkbounds(A, J...)
    unsafe_view(_maybe_reshape_parent(A, Base.index_ndims(J...)), J...)
end

ViewIndex = Union{Real, AbstractArray}
function unsafe_view(A::AbstractArray, I::Vararg{ViewIndex})
    HalfIntSubArray(A, I)
end

# When we take the view of a view, it's often possible to "reindex" the parent
# view's indices such that we can "pop" the parent view and keep just one layer
# of indirection. But we can't always do this because arrays of `CartesianIndex`
# might span multiple parent indices, making the reindex calculation very hard.
# So we use _maybe_reindex to figure out if there are any arrays of
# `CartesianIndex`, and if so, we punt and keep two layers of indirection.
unsafe_view(V::HalfIntSubArray, I::Vararg{ViewIndex,N}) where {N} =
    _maybe_reindex(V, I)
_maybe_reindex(V, I) = _maybe_reindex(V, I, I)
_maybe_reindex(V, I, ::Tuple{AbstractArray{<:Base.AbstractCartesianIndex}, Vararg{Any}}) =
    HalfIntSubArray(V, I)
# But allow arrays of CartesianIndex{1}; they behave just like arrays of Ints
_maybe_reindex(V, I, A::Tuple{AbstractArray{<:Base.AbstractCartesianIndex{1}}, Vararg{Any}}) =
    _maybe_reindex(V, I, Base.tail(A))
_maybe_reindex(V, I, A::Tuple{Any, Vararg{Any}}) = _maybe_reindex(V, I, Base.tail(A))
function _maybe_reindex(V, I, ::Tuple{})
    @inbounds idxs = to_indices(V.parent, reindex(V, I))
    HalfIntSubArray(V.parent, idxs)
end

# hash

function Base.hash(v::HalfIntSubArray, h::UInt)
	h = hash(v.parent, h)
	h = hash(axes(v), h)
	h = hash(v.offset1, h)
	h = hash(v.stride1, h)
	h = hash(:HalfIntSubArray, h)
	h
end


text:
# module Jive

using Distributed # nprocs addprocs rmprocs
using Printf

function slash_to_path_separator(subpath::String)
    sep = Base.Filesystem.path_separator
    sep == "/" ? subpath : replace(subpath, "/" => sep)
end

function path_separator_to_slash(subpath::String)
    sep = Base.Filesystem.path_separator
    sep == "/" ? subpath : replace(subpath, sep => "/")
end

function get_all_files(dir, skip, targets)
    all_files = Vector{String}()
    filters = []
    start_idx = 1
    if !isempty(targets)
        for arg in targets
            if occursin('=', arg)
                name, val = split(arg, '=')
                if name == "start"
                    start_idx = parse(Int, val)
                end
            else
                subpath = path_separator_to_slash(arg)
                if subpath == "."
                elseif startswith(subpath, "./")
                    push!(filters, subpath[3:end])
                else
                    push!(filters, subpath)
                end
            end
        end
    end
    for (root, dirs, files) in walkdir(dir)
        for filename in files
            !endswith(filename, ".jl") && continue
            root == dir && "runtests.jl" == filename && continue
            subpath = path_separator_to_slash(relpath(normpath(root, filename), dir))
            any(x -> startswith(subpath, x), path_separator_to_slash.(skip)) && continue
            !isempty(filters) && !any(x -> startswith(subpath, x), filters) && continue
            push!(all_files, subpath)
        end
    end
    (all_files, start_idx)
end

"""
    runtests(dir::String; skip::Union{Vector{Any},Vector{String}}=[], node1::Union{Vector{Any},Vector{String}}=[], targets=ARGS, enable_distributed::Bool=true)

run the test files from the specific directory.

* `dir`: the root directory to traverse.
* `skip`: files or directories to skip.
* `node1`: run on node 1 during for the distributed tests.
* `targets`: filter targets and start. default is `ARGS`
* `enable_distributed`: option for distributed
"""
function runtests(dir::String; skip::Union{Vector{Any},Vector{String}}=[], node1::Union{Vector{Any},Vector{String}}=[], targets=ARGS, enable_distributed::Bool=true)
    (all_tests, start_idx) = get_all_files(dir, skip, targets)
    env_jive_procs = get(ENV, "JIVE_PROCS", "") # "" "auto" "0" "1" "2" "3" ...
    if ("0" == env_jive_procs) || !enable_distributed
        run(dir, all_tests, start_idx)
    else
        num_procs = nprocs()
        if isempty(env_jive_procs)
        elseif "auto" == env_jive_procs
            Sys.CPU_THREADS > num_procs && addprocs(Sys.CPU_THREADS - num_procs + 1)
        else
            jive_procs = parse(Int, env_jive_procs)
            jive_procs >= num_procs && addprocs(jive_procs - num_procs + 1)
        end
        if nprocs() > 1
            distributed_run(dir, all_tests, start_idx, path_separator_to_slash.(node1))
        else
            run(dir, all_tests, start_idx)
        end
    end
end

struct FinishedWithErrors <: Exception
end

function Base.showerror(io::IO, ex::FinishedWithErrors, bt; backtrace=true)
    printstyled(io, "Test run finished with errors", color=:red)
end

function report(io::IO, t0, anynonpass, n_passed)
    if anynonpass > 0
        printstyled(io, "❗  ", color=:red)
        print(io, "Test run finished with ")
        print(io, anynonpass, " ")
        print(io, anynonpass == 1 ? "error" : "errors", ".")
        Printf.@printf(io, "  (%.2f seconds)\n", (time_ns()-t0)/1e9)
        throw(FinishedWithErrors())
    elseif n_passed > 0
        printstyled(io, "✅  ", color=:green)
        print(io, "All ")
        printstyled(io, n_passed, color=:green)
        print(io, " ")
        print(io, n_passed == 1 ? "test has" : "tests have")
        print(io, " been completed.")
        Printf.@printf(io, "  (%.2f seconds)\n", (time_ns()-t0)/1e9)
    end
end

function jive_briefing(io::IO, numbering, subpath, msg, description)
    buf = IOBuffer()
    context = IOContext(buf, :color => true)
    printstyled(context, numbering, color=:underline)
    print(context, ' ', subpath)
    !isempty(msg) && print(context, ' ', msg)
    println(context)
    print(io, String(take!(buf)))
end

# code from https://github.com/JuliaLang/julia/blob/master/stdlib/Test/src/Test.jl
module CodeFromStdlibTest

using Test: TESTSET_PRINT_ENABLE, DefaultTestSet, Error, TestSetException, Random, get_testset_depth, get_testset, record, pop_testset, parse_testset_args, _check_testset, push_testset, get_test_counts, filter_errors
using ..Jive: jive_briefing
using Printf

if VERSION >= v"1.3.0-DEV.565"
    default_rng = Random.default_rng
else
    default_rng = () -> Random.GLOBAL_RNG
end

# print_counts
function jive_print_counts(io::IO, ts::DefaultTestSet, elapsedtime)
    passes, fails, errors, broken, c_passes, c_fails, c_errors, c_broken = get_test_counts(ts)

    nf = fails + c_fails
    if nf > 0
        print(io, repeat(' ', 4))
        printstyled(io, "Fail", " "; bold=true, color=Base.error_color())
        printstyled(io, nf, color=Base.error_color())
        println(io)
    end

    ne = errors + c_errors
    if ne > 0
        print(io, repeat(' ', 4))
        printstyled(io, "Error", " "; bold=true, color=Base.error_color())
        printstyled(io, ne, color=Base.error_color())
        println(io)
    end

    nb = broken + c_broken
    if nb > 0
        print(io, repeat(' ', 4))
        printstyled(io, "Broken", " "; bold=true, color=Base.warn_color())
        printstyled(io, nb, color=Base.warn_color())
        println(io)
    end

    np = passes + c_passes
    if np > 0
        print(io, repeat(' ', 4))
        printstyled(io, "Pass", " "; bold=true, color=:green)
        printstyled(io, np, color=:green)
        Printf.@printf(io, "  (%.2f seconds)\n", elapsedtime)
    end
end

# finish
function jive_finish(io::IO, ts::DefaultTestSet, elapsedtime)
    #Base.timev_print(t, memallocs)
    # If we are a nested test set, do not print a full summary
    # now - let the parent test set do the printing
    if get_testset_depth() != 0
        # Attach this test set to the parent test set
        parent_ts = get_testset()
        record(parent_ts, ts)
        return ts
    end
    passes, fails, errors, broken, c_passes, c_fails, c_errors, c_broken = get_test_counts(ts)
    total_pass   = passes + c_passes
    total_fail   = fails  + c_fails
    total_error  = errors + c_errors
    total_broken = broken + c_broken
    total = total_pass + total_fail + total_error + total_broken

    if TESTSET_PRINT_ENABLE[]
        jive_print_counts(io, ts, elapsedtime) # print_test_results(ts)
    end

    # Finally throw an error as we are the outermost test set
    if total != total_pass + total_broken
        # Get all the error/failures and bring them along for the ride
        efs = filter_errors(ts)
        throw(TestSetException(total_pass,total_fail,total_error, total_broken, efs))
    end

    # return the testset so it is returned from the @testset macro
    ts
end

# testset_beginend
function jive_testset_beginend(io, numbering, subpath, msg, args, tests, source)
    desc, testsettype, options = parse_testset_args(args[1:end-1])
    if desc === nothing
        desc = ""
    end
    # If we're at the top level we'll default to DefaultTestSet. Otherwise
    # default to the type of the parent testset
    if testsettype === nothing
        testsettype = :(get_testset_depth() == 0 ? DefaultTestSet : typeof(get_testset()))
    end

    # Generate a block of code that initializes a new testset, adds
    # it to the task local storage, evaluates the test(s), before
    # finally removing the testset and giving it a chance to take
    # action (such as reporting the results)
    ex = quote
        _check_testset($testsettype, $(QuoteNode(testsettype.args[1])))
        ts = $(testsettype)($desc; $options...)
        jive_briefing($(esc(io)), $(esc(numbering)), $(esc(subpath)), $(esc(msg)), ts.description)
        # this empty loop is here to force the block to be compiled,
        # which is needed for backtrace scrubbing to work correctly.
        while false; end
        push_testset(ts)
        # we reproduce the logic of guardseed, but this function
        # cannot be used as it changes slightly the semantic of @testset,
        # by wrapping the body in a function
        local RNG = default_rng()
        oldrng = copy(RNG)
        elapsedtime = @elapsed try
            # RNG is re-seeded with its own seed to ease reproduce a failed test
            Random.seed!(RNG.seed)
            $(esc(tests))
        catch err
            err isa InterruptException && rethrow()
            # something in the test block threw an error. Count that as an
            # error in this test set
            backtrace = VERSION >= v"1.2.0-DEV.459" ? Base.catch_stack() : stacktrace(catch_backtrace())
            record(ts, Error(:nontest_error, :(), err, backtrace, LineNumberNode(err.line, err.file)))
        finally
            copy!(RNG, oldrng)
        end
        pop_testset()
        jive_finish($(esc(io)), ts, elapsedtime) # finish(ts)
    end
    # preserve outer location if possible
    if tests isa Expr && tests.head === :block && !isempty(tests.args) && tests.args[1] isa LineNumberNode
        ex = Expr(:block, tests.args[1], ex)
    end
    return ex
end

# @testset
macro jive_testset(io, numbering, subpath, msg, args...)
    tests = args[end]
    return jive_testset_beginend(io, numbering, subpath, msg, args, tests, __source__)
end

end # module Jive.CodeFromStdlibTest

# code from https://github.com/JuliaLang/julia/blob/master/test/runtests.jl
module CodeFromJuliaTest

using ..CodeFromStdlibTest: @jive_testset, get_test_counts
using ..Jive: slash_to_path_separator, report, jive_briefing
using Test.Random # RandomDevice
using Distributed # @everywhere remotecall_fetch

function runner(worker, idx, num_tests, subpath, filepath)
    numbering = string(idx, /, num_tests)
    buf = IOBuffer()
    context = IOContext(buf, :color => true)
    ts = @jive_testset context numbering subpath " (worker: $worker)" "" begin
        Main.include(filepath)
    end
    (ts, buf)
end

function distributed_run(dir::String, tests::Vector{String}, start_idx::Int, node1::Vector{String})
    io = stdout
    printstyled(io, "Sys.CPU_THREADS", color=:cyan)
    printstyled(io, ": ", Sys.CPU_THREADS)
    printstyled(io, ", ")
    printstyled(io, "nworkers()", color=:cyan)
    printstyled(io, ": ", nworkers())
    println(io)

    idx = 0
    num_tests = length(tests)
    env = Dict{Int,Tuple{Int,String}}()

    n_passed = 0
    anynonpass = 0
    local t0 = time_ns()
    try
        node1_tests = []
        @everywhere @eval(Main, using Jive)
        @sync begin
            for worker in workers()
                @async begin
                    while length(tests) > 0
                        idx += 1
                        subpath = popfirst!(tests)
                        env[worker] = (idx, subpath)
                        if idx < start_idx
                            numbering = string(idx, /, num_tests)
                            jive_briefing(io, numbering, subpath, "--", "")
                            continue
                        end
                        if any(x -> startswith(subpath, x), node1)
                            push!(node1_tests, (idx, subpath))
                        else
                            filepath = normpath(dir, slash_to_path_separator(subpath))
                            f = remotecall(runner, worker, worker, idx, num_tests, subpath, filepath)
                            (ts, buf) = fetch(f)
                            print(io, String(take!(buf)))
                            passes, fails, errors, broken, c_passes, c_fails, c_errors, c_broken = get_test_counts(ts)
                            n_passed += passes + c_passes
                            anynonpass += ts.anynonpass
                        end
                    end # while length(tests) > 0
                    if worker != 1
                        # Free up memory =)
                        rmprocs(worker, waitfor=0)
                    end
                end # @async begin
            end # for worker in workers()
        end # @sync begin
        worker = myid()
        for (idx, subpath) in node1_tests
            filepath = normpath(dir, slash_to_path_separator(subpath))
            f = remotecall(runner, worker, worker, idx, num_tests, subpath, filepath)
            (ts, buf) = fetch(f)
            print(io, String(take!(buf)))
            passes, fails, errors, broken, c_passes, c_fails, c_errors, c_broken = get_test_counts(ts)
            n_passed += passes + c_passes
            anynonpass += ts.anynonpass
        end
    catch err
        anynonpass += 1
        if err isa CompositeException
            exception = first(err.exceptions)
            if exception isa CapturedException
                result = exception.ex
            elseif Symbol(typeof(exception)) === :TaskFailedException  # VERSION >= v"1.3.0-alpha.110"
                result = exception.task.result
            else
                result = nothing
            end
            if result isa RemoteException
                print(io, "⚠  ")
                showerror(io, result)
                println(io)
                println.(Ref(io), result.captured.ex.errors_and_fails)
                remote_worker = result.pid
                if haskey(env, remote_worker)
                    worker = myid()
                    (idx, subpath) = env[remote_worker]
                    printstyled(io, "Retrying"; bold=true, color=:cyan)
                    printstyled(io, " ")
                    numbering = string(idx, /, num_tests)
                    jive_briefing(io, numbering, subpath, string(" (worker: ", remote_worker => worker, ")"), "")
                    filepath = normpath(dir, slash_to_path_separator(subpath))
                    f = remotecall(runner, worker, worker, idx, num_tests, subpath, filepath)
                    res = fetch(f)
                    if res isa RemoteException
                    else
                        (ts, buf) = res
                        print(io, String(take!(buf)))
                    end
                end
            else
                print(io, "⚠  ")
                showerror(io, exception)
                println(io)
            end
        else
            print(io, "⚠  ")
            showerror(io, err)
            println(io)
        end
    finally
        GC.gc()
    end
    report(io, t0, anynonpass, n_passed)
end

end # module Jive.CodeFromJuliaTest

using .CodeFromJuliaTest: distributed_run
using .CodeFromStdlibTest: @jive_testset, get_test_counts

function run(dir::String, tests::Vector{String}, start_idx::Int)
    io = stdout
    n_passed = 0
    anynonpass = 0
    local t0 = time_ns()
    for (idx, subpath) in enumerate(tests)
        if idx < start_idx
            num_tests = length(tests)
            numbering = string(idx, /, num_tests)
            jive_briefing(io, numbering, subpath, "--", "")
            continue
        end
        filepath = normpath(dir, slash_to_path_separator(subpath))
        numbering = string(idx, /, length(tests))
        ts = @jive_testset io numbering subpath "" "" begin
            Main.include(filepath)
        end
        passes, fails, errors, broken, c_passes, c_fails, c_errors, c_broken = get_test_counts(ts)
        n_passed += passes + c_passes
        anynonpass += ts.anynonpass
    end
    report(io, t0, anynonpass, n_passed)
end

# module Jive