text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

const StdUvDist = Union{StandardUvUniform, StandardUvNormal}
const StdMvDist = Union{StandardMvUniform, StandardMvNormal}

_adignore(f) = f()

function ChainRulesCore.rrule(::typeof(_adignore), f)
    result = _adignore(f)
    _nogradient_pullback(ΔΩ) = (NoTangent(), ZeroTangent())
    return result, _nogradient_pullback
end

macro _adignore(expr)
    :(_adignore(() -> $(esc(expr))))
end

function _pushfront(v::AbstractVector, x)
    T = promote_type(eltype(v), typeof(x))
    r = similar(v, T, length(eachindex(v)) + 1)
    r[firstindex(r)] = x
    r[firstindex(r)+1:lastindex(r)] = v
    r
end

function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
    result = _pushfront(v, x)
    function _pushfront_pullback(thunked_ΔΩ)
        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
        (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)])
    end
    return result, _pushfront_pullback
end

function _pushback(v::AbstractVector, x)
    T = promote_type(eltype(v), typeof(x))
    r = similar(v, T, length(eachindex(v)) + 1)
    r[lastindex(r)] = x
    r[firstindex(r):lastindex(r)-1] = v
    r
end

function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x)
    result = _pushback(v, x)
    function _pushback_pullback(thunked_ΔΩ)
        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
        (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)])
    end
    return result, _pushback_pullback
end

_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)]

_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1]

_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs)))

function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector)
    result = _rev_cumsum(xs)
    function _rev_cumsum_pullback(ΔΩ)
        ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ))
        (NoTangent(), ∂xs)
    end
    return result, _rev_cumsum_pullback
end

# Equivalent to `cumprod(xs)``:
_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))

function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector)
    result = _exp_cumsum_log(xs)
    function _exp_cumsum_log_pullback(ΔΩ)
        ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ))
        (NoTangent(), ∂xs)
    end
    return result, _exp_cumsum_log_pullback
end

"""
    abstract type DistributionTransform{VT<:AbstractValueShape,VF<:AbstractValueShape} <: Function

*Experimental feature, not part of stable public API.*

Transform variate values between distributions

Constructors:

```julia
DistributionTransform(target_dist, source_dist)
DistributionTransform(Uniform, source_dist)
DistributionTransform(Normal, source_dist)
```
"""
struct DistributionTransform{
    DT <: ContinuousDistribution,
    DF <: ContinuousDistribution,
    VT <: AbstractValueShape,
    VF <: AbstractValueShape,
} <: Function
    target_dist::DT
    source_dist::DF
end

function _distrafo_ctor_impl(target_dist::DT, source_dist::DF) where {DT<:ContinuousDistribution,DF<:ContinuousDistribution}
    @argcheck eff_totalndof(target_dist) == eff_totalndof(source_dist)
    VT = typeof(varshape(target_dist))
    VF = typeof(varshape(source_dist))
    DistributionTransform{DT,DF,VT,VF}(target_dist, source_dist)
end

function _distrafo_ctor_impl(target_dist::Distribution, source_dist::Distribution)
    @argcheck eff_totalndof(target_dist) == eff_totalndof(source_dist)
    DistributionTransform(target_dist, source_dist)
end

DistributionTransform(target_dist::Distribution{VF,Continuous}, source_dist::Distribution{VF,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{Multivariate,Continuous}, source_dist::Distribution{VF,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{VF,Continuous}, source_dist::Distribution{Multivariate,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{Multivariate,Continuous}, source_dist::Distribution{Multivariate,Continuous}) =
    _distrafo_ctor_impl(target_dist, source_dist)

show_distribution(io::IO, d::Distribution) = show(io, d)
function show_distribution(io::IO, d::NamedTupleDist)
    print(io, Base.typename(typeof(d)).name, "{")
    show(io, propertynames(d))
    print(io, "}(…)")
end

function Base.show(io::IO, trafo::DistributionTransform)
    print(io, Base.typename(typeof(trafo)).name, "(")
    show_distribution(io, trafo.target_dist)
    print(io, ", ")
    show_distribution(io, trafo.source_dist)
    print(io, ")")
end

Base.show(io::IO, M::MIME"text/plain", trafo::DistributionTransform) = show(io, trafo)

# apply_dist_trafo(trg_d, src_d, src_v)
function apply_dist_trafo end

(trafo::DistributionTransform)(x) = apply_dist_trafo(trafo.target_dist, trafo.source_dist, x)

InverseFunctions.inverse(trafo::DistributionTransform) = DistributionTransform(trafo.source_dist, trafo.target_dist)

import Base.inv
Base.@deprecate inv(trafo::DistributionTransform) inverse(trafo)

function ChangesOfVariables.with_logabsdet_jacobian(trafo::DistributionTransform, x)
    y = trafo(x)
    logpdf_src = logpdf(trafo.source_dist, x)
    logpdf_trg = logpdf(trafo.target_dist, y)
    ladj = logpdf_src - logpdf_trg
    # If logpdf_src and logpdf_trg are -Inf setting lafj to zero is safe:
    fixed_ladj = logpdf_src == logpdf_trg == -Inf ? zero(ladj) : ladj
    y, fixed_ladj
end

Base.:(∘)(::typeof(identity), f::DistributionTransform) = f
Base.:(∘)(f::DistributionTransform, ::typeof(identity)) = f

function Base.Broadcast.broadcasted(
    trafo::DistributionTransform,
    v_src::Union{ArrayOfSimilarVectors{<:Real},ShapedAsNTArray}
)
    broadcast_trafo(trafo, v_src)
end

function (trafo::DistributionTransform)(s::DensitySample)
    v, ladj = with_logabsdet_jacobian(trafo, s.v)
    logd = s.logd - ladj
    DensitySample(v, logd, s.weight, s.info, s.aux)
end

function Base.Broadcast.broadcasted(
    trafo::DistributionTransform,
    s_src::DensitySampleVector
)
    broadcast_trafo(trafo, s_src)
end

# Use ForwardDiff for univariate distribution transformations:
@inline function ChainRulesCore.rrule(::typeof(apply_dist_trafo), trg_d::Distribution{Univariate}, src_d::Distribution{Univariate}, src_v::Any)
    ChainRulesCore.rrule(fwddiff(apply_dist_trafo), trg_d, src_d, src_v)
end

import Base.∘
function ∘(a::DistributionTransform, b::DistributionTransform)
    @argcheck a.source_dist == b.target_dist
    DistributionTransform(a.target_dist, b.source_dist)
end

ValueShapes.varshape(trafo::DistributionTransform) = varshape(trafo.source_dist)

function (trafo::DistributionTransform)(vs::AbstractValueShape)
    @argcheck vs <= varshape(trafo)
    varshape(trafo.target_dist)
end

ValueShapes.unshaped(trafo::DistributionTransform) =
    DistributionTransform(unshaped(trafo.target_dist), unshaped(trafo.source_dist))

const _StdDistType = Union{Uniform, Normal}

_trg_disttype(::Type{Uniform}, ::Type{Univariate}) = StandardUvUniform
_trg_disttype(::Type{Uniform}, ::Type{Multivariate}) = StandardMvUniform
_trg_disttype(::Type{Normal}, ::Type{Univariate}) = StandardUvNormal
_trg_disttype(::Type{Normal}, ::Type{Multivariate}) = StandardMvNormal

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::Distribution{Univariate,Continuous})
    trg_dt = _trg_disttype(disttype, Univariate)
    trg_dt()
end

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::Distribution{Multivariate,Continuous})
    trg_dt = _trg_disttype(disttype, Multivariate)
    trg_dt(eff_totalndof(source_dist))
end

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::ContinuousDistribution)
    trg_dt = _trg_disttype(disttype, Multivariate)
    trg_dt(eff_totalndof(source_dist))
end

function DistributionTransform(disttype::Type{<:_StdDistType}, source_dist::ContinuousDistribution)
    trg_d = _trg_dist(disttype, source_dist)
    DistributionTransform(trg_d, source_dist)
end

function std_dist_from(src_d::Distribution)
    throw(ArgumentError("No standard intermediate distribution defined to transform from $(typeof(src_d).name)"))
end

function std_dist_to(trg_d::Distribution)
    throw(ArgumentError("No standard intermediate distribution defined to transform into $(typeof(trg_d).name)"))
end

@inline function _intermediate_std_dist(trg_d::Distribution, src_d::Distribution)
    _select_intermediate_dist(std_dist_to(trg_d), std_dist_from(src_d))
end

@inline _intermediate_std_dist(::Union{StdUvDist,StdMvDist}, src_d::Distribution) = std_dist_from(src_d)

@inline _intermediate_std_dist(trg_d::Distribution, ::Union{StdUvDist,StdMvDist}) = std_dist_to(trg_d)

function _intermediate_std_dist(::Union{StdUvDist,StdMvDist}, ::Union{StdUvDist,StdMvDist})
    throw(ArgumentError("Direct conversions must be used between standard intermediate distributions"))
end

@inline _select_intermediate_dist(a::D, ::D) where D<:Union{StdUvDist,StdMvDist} = a
@inline _select_intermediate_dist(a::D, ::D) where D<:Union{StandardUvUniform,StandardMvUniform} = a
@inline _select_intermediate_dist(a::Union{StandardUvUniform,StandardMvUniform}, ::Union{StdUvDist,StdMvDist}) = a
@inline _select_intermediate_dist(::Union{StdUvDist,StdMvDist}, b::Union{StandardUvUniform,StandardMvUniform}) = b

_check_conv_eff_totalndof(trg_d::Uniform, src_d::Uniform) = nothing

function _check_conv_eff_totalndof(trg_d::Distribution, src_d::Distribution)
    trg_d_n = eff_totalndof(trg_d)
    src_d_n = eff_totalndof(src_d)
    if trg_d_n != src_d_n
        throw(ArgumentError("Can't convert to $(typeof(trg_d).name) with $(trg_d_n) eff. DOF from $(typeof(src_d).name) with $(src_d_n) eff. DOF"))
    end
    nothing
end

function apply_dist_trafo(trg_d::Distribution, src_d::Distribution, src_v::Any)
    _check_conv_eff_totalndof(trg_d, src_d)
    intermediate_d = _intermediate_std_dist(trg_d, src_d)
    intermediate_d === trg_d && throw(ArgumentError("No transformation path between distributions"))
    intermediate_v = apply_dist_trafo(intermediate_d, src_d, src_v)
    apply_dist_trafo(trg_d, intermediate_d, intermediate_v)
end

function apply_dist_trafo(trg_d::DT, src_d::DT, src_v) where {DT <: StdMvDist}
    @argcheck src_v isa AbstractVector{<:Real}
    @argcheck length(trg_d) == length(src_d) == length(eachindex(src_v))
    return src_v
end

_dist_params_numtype(d::Distribution) = promote_type(map(typeof, Distributions.params(d))...)

@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) = _trafo_cdf_impl(_dist_params_numtype(d), d, x)

@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) = cdf(d, x)

@inline function _trafo_cdf_impl(::Type{<:Union{Integer,AbstractFloat}}, d::Distribution{Univariate,Continuous}, x::ForwardDiff.Dual{TAG}) where {N,TAG}
    x_v = ForwardDiff.value(x)
    u = cdf(d, x_v)
    dudx = pdf(d, x_v)
    ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x))
end

@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) = _trafo_quantile_impl(_dist_params_numtype(d), d, u)

@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) = _trafo_quantile_impl_generic(d, u)

@inline function _trafo_quantile_impl(::Type{<:Union{Integer,AbstractFloat}}, d::Distribution{Univariate,Continuous}, u::ForwardDiff.Dual{TAG}) where {TAG}
    x = _trafo_quantile_impl_generic(d, ForwardDiff.value(u))
    dxdu = inv(pdf(d, x))
    ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u))
end

# Workaround for Beta dist, ForwardDiff doesn't work for parameters:
@inline _trafo_quantile_impl_generic(d::Beta{T}, u::Real) where {T<:ForwardDiff.Dual} = convert(float(typeof(u)), NaN)
# Workaround for Beta dist, current quantile implementation only supports Float64:
@inline _trafo_quantile_impl_generic(d::Beta{T}, u::Union{Integer,AbstractFloat}) where {T<:Union{Integer,AbstractFloat}} = _trafo_quantile_impl(T, d, convert(promote_type(Float64, typeof(u)), u))
# Workaround for StatsFuns issues #133, caused by SpecialFunctions, fixed in SpecialFunctions v2.1.4:
@inline _trafo_quantile_impl_generic(d::Beta{T}, u::Float64) where {T<:Union{Integer,AbstractFloat}} = (d.α ≈ 1 && d.β ≈ 1 && u < 1e-19) ? u : convert(Float64, quantile(d, u))

@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) = quantile(d, u)

# Workaround for rounding errors that can result in quantile values outside of support of Truncated:
@inline function _trafo_quantile_impl_generic(d::Truncated{<:Distribution{Univariate,Continuous}}, u::Real)
    x = quantile(d, u)
    T = typeof(x)
    min_x = T(minimum(d))
    max_x = T(maximum(d))
    if x < min_x && isapprox(x, min_x, atol = 4 * eps(T))
        min_x
    elseif x > max_x && isapprox(x, max_x, atol = 4 * eps(T))
        max_x
    else
        x
    end
end

@inline function _eval_dist_trafo_func(f::typeof(_trafo_cdf), d::Distribution{Univariate,Continuous}, src_v::Real)
    R_V = float(promote_type(typeof(src_v), eltype(params(d))))
    if insupport(d, src_v)
        trg_v = f(d, src_v)
        convert(R_V, trg_v)
    else
        convert(R_V, NaN)
    end
end

@inline function _eval_dist_trafo_func(f::typeof(_trafo_quantile), d::Distribution{Univariate,Continuous}, src_v::Real)
    R_V = float(promote_type(typeof(src_v), eltype(params(d))))
    if 0 <= src_v <= 1
        trg_v = f(d, src_v)
        convert(R_V, trg_v)
    else
        convert(R_V, NaN)
    end
end

std_dist_from(src_d::Distribution{Univariate,Continuous}) = StandardUvUniform()

function apply_dist_trafo(::StandardUvUniform, src_d::Distribution{Univariate,Continuous}, src_v::Real)
    _eval_dist_trafo_func(_trafo_cdf, src_d, src_v)
end

std_dist_to(trg_d::Distribution{Univariate,Continuous}) = StandardUvUniform()

function apply_dist_trafo(trg_d::Distribution{Univariate,Continuous}, ::StandardUvUniform, src_v::Real)
    TV = float(typeof(src_v))
    # Avoid src_v ≈ 0 and src_v ≈ 1 to avoid infinite variate values for target distributions with infinite support:
    mod_src_v = ifelse(src_v == 0, zero(TV) + eps(TV), ifelse(src_v == 1, one(TV) - eps(TV), convert(TV, src_v)))
    _eval_dist_trafo_func(_trafo_quantile, trg_d, mod_src_v)
end

function _dist_trafo_rescale_impl(trg_d, src_d, src_v::Real)
    R = float(typeof(src_v))
    trg_offs, trg_scale = location(trg_d), scale(trg_d)
    src_offs, src_scale = location(src_d), scale(src_d)
    rescale_factor = trg_scale / src_scale
    (src_v - src_offs) * rescale_factor + trg_offs
end

@inline apply_dist_trafo(trg_d::Uniform, src_d::Uniform, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)
@inline apply_dist_trafo(trg_d::StandardUvUniform, src_d::Uniform, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)
@inline apply_dist_trafo(trg_d::Uniform, src_d::StandardUvUniform, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)

# ToDo: Use StandardUvNormal as standard intermediate dist for Normal? Would
# be useful if StandardUvNormal would be a better standard intermediate than
# StandardUvUniform for some other uniform distributions as well.
#
#     std_dist_from(src_d::Normal) = StandardUvNormal()
#     std_dist_to(trg_d::Normal) = StandardUvNormal()

@inline apply_dist_trafo(trg_d::Normal, src_d::Normal, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)
@inline apply_dist_trafo(trg_d::StandardUvNormal, src_d::Normal, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)
@inline apply_dist_trafo(trg_d::Normal, src_d::StandardUvNormal, src_v::Real) = _dist_trafo_rescale_impl(trg_d, src_d, src_v)

# ToDo: Optimized implementation for Distributions.Truncated <-> StandardUvUniform

@inline apply_dist_trafo(trg_d::StandardUvUniform, src_d::StandardUvUniform, src_v::Real) = src_v

@inline apply_dist_trafo(trg_d::StandardUvNormal, src_d::StandardUvNormal, src_v::Real) = src_v

@inline function apply_dist_trafo(trg_d::StandardUvUniform, src_d::StandardUvNormal, src_v::Real)
    apply_dist_trafo(StandardUvUniform(), Normal(), src_v)
end

@inline function apply_dist_trafo(trg_d::StandardUvNormal, src_d::StandardUvUniform, src_v::Real)
    apply_dist_trafo(Normal(), StandardUvUniform(), src_v)
end

@inline function apply_dist_trafo(trg_d::StandardMvUniform, src_d::StandardMvNormal, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvUniform(), StandardUvNormal(), src_v)
end

@inline function apply_dist_trafo(trg_d::StandardMvNormal, src_d::StandardMvUniform, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvNormal(), StandardUvUniform(), src_v)
end

std_dist_from(src_d::MvNormal) = StandardMvNormal(length(src_d))

function apply_dist_trafo(trg_d::StandardMvNormal, src_d::MvNormal, src_v::AbstractVector{<:Real})
    @argcheck length(trg_d) == length(src_d)
    A = cholesky(src_d.Σ).U
    transpose(A) \ (src_v - src_d.μ)
end

std_dist_to(trg_d::MvNormal) = StandardMvNormal(length(trg_d))

function apply_dist_trafo(trg_d::MvNormal, src_d::StandardMvNormal, src_v::AbstractVector{<:Real})
    @argcheck length(trg_d) == length(src_d)
    A = cholesky(trg_d.Σ).U
    transpose(A) * src_v + trg_d.μ
end

eff_totalndof(d::Dirichlet) = length(d) - 1
eff_totalndof(d::DistributionsAD.TuringDirichlet) = length(d) - 1

std_dist_to(trg_d::Dirichlet) = StandardMvUniform(eff_totalndof(trg_d))
std_dist_to(trg_d::DistributionsAD.TuringDirichlet) = StandardMvUniform(eff_totalndof(trg_d))

function apply_dist_trafo(trg_d::Dirichlet, src_d::StandardMvUniform, src_v::AbstractVector{<:Real})
    apply_dist_trafo(DistributionsAD.TuringDirichlet(trg_d.alpha), src_d, src_v)
end

function _dirichlet_beta_trafo(α::Real, β::Real, src_v::Real)
    R = float(promote_type(typeof(α), typeof(β), typeof(src_v)))
    convert(R, apply_dist_trafo(Beta(α, β), StandardUvUniform(), src_v))::R
end

_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)

function apply_dist_trafo(trg_d::DistributionsAD.TuringDirichlet, src_d::StandardMvUniform, src_v::AbstractVector{<:Real})
    # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
    # https://arxiv.org/abs/1010.3436

    @_adignore @argcheck length(trg_d) == length(src_d) + 1
    αs = _dropfront(_rev_cumsum(trg_d.alpha))
    βs = _dropback(trg_d.alpha)
    beta_v = fwddiff(_dirichlet_beta_trafo).(αs, βs, src_v)
    beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1))
    beta_v_ext = _pushback(beta_v, 0)
    fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext)
end

function _product_dist_trafo_impl(trg_ds, src_ds, src_v::AbstractVector{<:Real})
    fwddiff(apply_dist_trafo).(trg_ds, src_ds, src_v)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::Distributions.Product, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, src_d.v, src_v)
end

function apply_dist_trafo(trg_d::StandardMvUniform, src_d::Distributions.Product, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvUniform(), src_d.v, src_v)
end

function apply_dist_trafo(trg_d::StandardMvNormal, src_d::Distributions.Product, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvNormal(), src_d.v, src_v)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::StandardMvUniform, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, StandardUvUniform(), src_v)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::StandardMvNormal, src_v::AbstractVector{<:Real})
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, StandardUvNormal(), src_v)
end

_flat_ntd_elshape(d::Distribution) = ArrayShape{Real}(eff_totalndof(d))

function _flat_ntd_accessors(d::NamedTupleDist{names,DT,AT,VT}) where {names,DT,AT,VT}
    shapes = map(_flat_ntd_elshape, values(d))
    vs = NamedTupleShape(VT, NamedTuple{names}(shapes))
    values(vs)
end

function _flat_ntdistelem_to_stdmv(trg_d::StdMvDist, sd::Distribution, src_v_unshaped::AbstractVector{<:Real}, trg_acc::ValueAccessor)
    td = view(trg_d, ValueShapes.view_idxs(Base.OneTo(length(trg_d)), trg_acc))
    sv = trg_acc(src_v_unshaped)
    apply_dist_trafo(td, unshaped(sd), sv)
end

function _flat_ntdistelem_to_stdmv(trg_d::StdMvDist, sd::ConstValueDist, src_v_unshaped::AbstractVector{<:Real}, trg_acc::ValueAccessor)
    Bool[]
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::ValueShapes.UnshapedNTD, src_v::AbstractVector{<:Real})
    @argcheck length(src_d) == length(eachindex(src_v))
    trg_accessors = _flat_ntd_accessors(src_d.shaped)
    rs = map((acc, sd) -> _flat_ntdistelem_to_stdmv(trg_d, sd, src_v, acc), trg_accessors, values(src_d.shaped))
    vcat(rs...)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::NamedTupleDist, src_v::Union{NamedTuple,ShapedAsNT})
    src_v_unshaped = unshaped(src_v, varshape(src_d))
    apply_dist_trafo(trg_d, unshaped(src_d), src_v_unshaped)
end

function _stdmv_to_flat_ntdistelem(td::Distribution, src_d::StdMvDist, src_v::AbstractVector{<:Real}, src_acc::ValueAccessor)
    sd = view(src_d, ValueShapes.view_idxs(Base.OneTo(length(src_d)), src_acc))
    sv = src_acc(src_v)
    apply_dist_trafo(unshaped(td), sd, sv)
end

function _stdmv_to_flat_ntdistelem(td::ConstValueDist, src_d::StdMvDist, src_v::AbstractVector{<:Real}, src_acc::ValueAccessor)
    Bool[]
end

function apply_dist_trafo(trg_d::ValueShapes.UnshapedNTD, src_d::StdMvDist, src_v::AbstractVector{<:Real})
    @argcheck length(src_d) == length(eachindex(src_v))
    src_accessors = _flat_ntd_accessors(trg_d.shaped)
    rs = map((acc, td) -> _stdmv_to_flat_ntdistelem(td, src_d, src_v, acc), src_accessors, values(trg_d.shaped))
    vcat(rs...)
end

function apply_dist_trafo(trg_d::NamedTupleDist, src_d::StdMvDist, src_v::AbstractVector{<:Real})
    unshaped_result = apply_dist_trafo(unshaped(trg_d), src_d, src_v)
    varshape(trg_d)(unshaped_result)
end

@static if isdefined(Distributions, :ReshapedDistribution)
    const AnyReshapedDist = Union{Distributions.ReshapedDistribution,ValueShapes.ReshapedDist}
else
    const AnyReshapedDist = Union{Distributions.MatrixReshaped,ValueShapes.ReshapedDist}
end

eff_totalndof(d::AnyReshapedDist) = eff_totalndof(unshaped(d))
std_dist_from(src_d::AnyReshapedDist) = std_dist_from(unshaped(src_d))
std_dist_to(trg_d::AnyReshapedDist) = std_dist_to(unshaped(trg_d))

function apply_dist_trafo(trg_d::Distribution{Multivariate}, src_d::AnyReshapedDist, src_v::Any)
    src_vs = varshape(src_d)
    @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    apply_dist_trafo(trg_d, unshaped(src_d), unshaped(src_v, src_vs))
end

function apply_dist_trafo(trg_d::AnyReshapedDist, src_d::Distribution{Multivariate}, src_v::AbstractVector{<:Real})
    trg_vs = varshape(trg_d)
    @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    r = apply_dist_trafo(unshaped(trg_d), src_d, src_v)
    trg_vs(r)
end

function apply_dist_trafo(trg_d::AnyReshapedDist, src_d::AnyReshapedDist, src_v::AbstractVector{<:Real})
    trg_vs = varshape(trg_d)
    src_vs = varshape(src_d)
    @argcheck totalndof(trg_vs) == totalndof(src_vs)
    r = apply_dist_trafo(unshaped(trg_d), unshaped(src_d), unshaped(src_v, src_vs))
    v = trg_vs(r)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::UnshapedHDist, src_v::AbstractVector{<:Real})
    src_v_primary, src_v_secondary = _hd_split(src_d, src_v)
    trg_d_primary = typeof(trg_d)(length(eachindex(src_v_primary)))
    trg_d_secondary = typeof(trg_d)(length(eachindex(src_v_secondary)))
    trg_v_primary = apply_dist_trafo(trg_d_primary, _hd_pridist(src_d), src_v_primary)
    trg_v_secondary = apply_dist_trafo(trg_d_secondary, _hd_secdist(src_d, src_v_primary), src_v_secondary)
    vcat(trg_v_primary, trg_v_secondary)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::HierarchicalDistribution, src_v::Any)
    src_v_unshaped = unshaped(src_v, varshape(src_d))
    apply_dist_trafo(trg_d, unshaped(src_d), src_v_unshaped)
end

function apply_dist_trafo(trg_d::UnshapedHDist, src_d::StdMvDist, src_v::AbstractVector{<:Real})
    src_v_primary, src_v_secondary = _hd_split(trg_d, src_v)
    src_d_primary = typeof(src_d)(length(eachindex(src_v_primary)))
    src_d_secondary = typeof(src_d)(length(eachindex(src_v_secondary)))
    trg_v_primary = apply_dist_trafo(_hd_pridist(trg_d), src_d_primary, src_v_primary)
    trg_v_secondary = apply_dist_trafo(_hd_secdist(trg_d, trg_v_primary), src_d_secondary, src_v_secondary)
    vcat(trg_v_primary, trg_v_secondary)
end

function apply_dist_trafo(trg_d::HierarchicalDistribution, src_d::StdMvDist, src_v::AbstractVector{<:Real})
    unshaped_result = apply_dist_trafo(unshaped(trg_d), src_d, src_v)
    varshape(trg_d)(unshaped_result)
end


text:
# -------------------- Binary Search --------------------- #

function binary_search(; f::Function, a::T, b::T,
        niter=50, eps_x::T=1e-6) where T<:AbstractFloat
    fa = f(a)
    fb = f(b)
    mid = (a + b) / 2

    for i = 1 : niter
        mid = (a + b) / 2

        if abs(b - a) < eps_x
            return mid
        end

        fmid = f(mid)
        fmid == 0 && return mid

        if sign(fmid) == sign(fa)
            a = mid
            fa = fmid
        else
            b = mid
            fb = fmid
        end
    end
    return error("binary search fails to reach accuracy of ",
                 eps_x, " in ", niter, " iterations")
end
# -------------------------------------------------------- #

# -------------------- Secant Method --------------------- #

function secant_root_finder(; f::Function, a::T, b::T,
        niter=100, eps_x::T=1e-6) where T<:AbstractFloat

    if f(a) * f(b) >= 0
        return error("secant method fails:
                     condition f(a)f(b) < 0 not met.")
    end

    for i = 1 : niter
        x₀ = a - f(a) * (b - a) / (f(b) - f(a))
        f₀ = f(x₀)
        @debug x₀ f₀

        if abs(b - a) < eps_x
            return x₀
        end

        if f(a) * f₀ < 0
            b = x₀
        elseif f(b) * f₀ < 0
            a = x₀
        elseif f₀ == 0
            return x₀
        else
            return error("secant method fails.")
        end
    end

    return error("secant method fails to reach accuracy of ",
                 eps_x, " in ", niter, " iterations")
    # return a - f(a) * (b - a) / (f(b) - f(a))
end
# -------------------------------------------------------- #

# -------------------- Newton Raphson -------------------- #

cache_x₀ = 1e5
is_cache = 1

function newton_raphson(; f::Function, fder::Function,
    niter=10000, eps_x::T=1e-6)::T where T<:AbstractFloat

    x₀::AbstractFloat = 0
    x₁::AbstractFloat = 0
    ## TODO: what if not cached?
    if is_cache == 1
        x₀ = cache_x₀
    end

    xs = fill(0.0, niter)
    for i = 1 : niter
        @debug x₀ f(x₀)
        # println(x₀, " ", f(x₀))
        x₁ = x₀ - f(x₀) / fder(x₀)

        ## To avoid domain error: negative values in log
        if x₁ < eps_x
            x₁ = eps_x
        end

        if abs(x₁ - x₀) < eps_x || f(x₁) == 0
            is_cache = 1
            cache_x₀ = x₁
            return x₁
        end

        xs[i] = x₀
        x₀ = x₁
    end

    rgx = range(1e4, 1e7, length=10000)
    output = hcat(rgx, f.(rgx))
    writedlm("../dump/f.txt", output, ' ')
    writedlm("../dump/xs.txt", xs, ' ')

    return error("newton raphson method fails to reach accuracy of ",
                 eps_x, " in ", niter, " iterations Delta P = ",
                 abs(f(x₀)), " x₀ = ", x₀,
                 " delta x = ", abs(x₀ - xs[niter - 1]))
end
# -------------------------------------------------------- #

ext: jl

max_stars_repo_name: sofiabelen/Ideal-Gas-Compressible-Fluid

max_issues_repo_name: sofiabelen/Ideal-Gas-Compressible-Fluid

max_forks_repo_name: sofiabelen/Two-Phase-Filtration


text:
""
function ref_add_goc!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    _PM.apply_pm!(_ref_add_goc!, ref, data)
end

function _ref_add_goc!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    ref[:branch_sm_active] = Dict(x for x in ref[:branch] if haskey(x.second, "rate_a"))

    ref[:shunt_const] = Dict(x for x in ref[:shunt] if (!haskey(x.second, "dispatchable") || !x.second["dispatchable"]))
    ref[:shunt_var] = Dict(x for x in ref[:shunt] if (haskey(x.second, "dispatchable") && x.second["dispatchable"]))

    bus_shunts_const = Dict((i, []) for (i,bus) in ref[:bus])
    for (i,shunt) in ref[:shunt_const]
        push!(bus_shunts_const[shunt["shunt_bus"]], i)
    end
    ref[:bus_shunts_const] = bus_shunts_const

    bus_shunts_var = Dict((i, []) for (i,bus) in ref[:bus])
    for (i,shunt) in ref[:shunt_var]
        push!(bus_shunts_var[shunt["shunt_bus"]], i)
    end
    ref[:bus_shunts_var] = bus_shunts_var
end


text:
addprocs(2)

@everywhere instr = readlines("input.txt")
@everywhere reg = filter(x -> x != "1", unique(map(x -> x[2], split.(instr))))
@everywhere cmd = Dict("set" => "=", "add" => "+=", "mul" => "*=", "mod" => "%=")
@everywhere get(x) = x in reg ? eval(Symbol(x)) : parse(Int, x)

@everywhere function run(ch_in, ch_out, proc)
    foreach(x -> eval(parse("$x = 0")), reg)
    eval(parse("p = $proc"))
    pointer = 1
    num_sent = 0
    while pointer <= length(instr)
        d = split(instr[pointer])
        if d[1] == "snd"
            num_sent += 1
            put!(ch_out, get(d[2]))
            if proc == 1
                println("[$proc] Sent $num_sent")
            end
        elseif d[1] == "jgz" && get(d[2]) > 0
            pointer += get(d[3])
            continue
        elseif d[1] == "rcv"
            eval(parse("$(d[2]) = $(take!(ch_in))"))
        elseif d[1] != "jgz" && d[2] != "rcv"
            eval(parse("$(d[2]) $(cmd[d[1]]) $(d[3])"))
        end
        pointer += 1
    end
end

@sync begin
    ch1 = RemoteChannel(()->Channel{Int}(1e3))
    ch2 = RemoteChannel(()->Channel{Int}(1e3))
    @spawn run(ch1, ch2, 0)
    @spawn run(ch2, ch1, 1)
end


text:
export DataFrameSystem, getSource, getTarget, getFirstChild, getLastChild, getTrait
import .CONCEPT: getChildren, getParent, setProperty
import .KERNEL: coords, isHetero, collectAtoms
import Base: length, getindex, ==, show

"""
    DataFrameSystem
`DataFrameSystem` is a collection of `DataFrame`s to store `KERNEL` types in.
"""
mutable struct DataFrameSystem <: KernelInterface
    systems ::DataFrame
    chains  ::DataFrame
    residues::DataFrame
    atoms   ::DataFrame
    bonds   ::DataFrame
    DataFrameSystem(
        systems::DataFrame,
        chains::DataFrame,
        residues::DataFrame,
        atoms::DataFrame,
        bonds::DataFrame
    ) = new(systems, chains, residues, atoms, bonds)
end

const bond_names = [    #6
   :name_,
   :source_,
   :target_,
   :bond_order_,
   :bond_type_,
   :properties_
]

const bond_types = Type[
    String,
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},
    Order,
    BondType,
    Vector{Tuple{String,UInt8}}
]

const atom_names = [ #18
    :name_,
    :parent_,
    :type_name_,
    :element_,
    :radius_,
    :type_,
    :number_of_bonds_,
    :formal_charge_,
    :position_,
    :charge_,
    :velocity_,
    :force_,
    :occupancy_,
    :serial_,
    :temp_factor_,
    :selected_,
    :bonds_,
    :properties_
]

const atom_types = [
    Union{String,Nothing},                  #name_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #parent_
    Union{String,Nothing},                  #type_name_
    Union{Element,Nothing},                 #element_
    Union{Float64,Nothing},                 #radius_
    Union{UInt8,Nothing},                   #type_
    Union{UInt8,Nothing},                   #number_of_bonds_
    Union{Int64,Nothing},                   #formal_charge_
    SVector{3,Float64},      #position_
    Union{Float64,Nothing},                 #charge_
    Union{SVector{3,Float64},Nothing},      #velocity_
    Union{SVector{3,Float64},Nothing},      #force_
    Union{Float64, Nothing},                #occupancy_
    Union{Int64,Nothing},                   #serial_
    Union{Float64,Nothing},                 #temp_factor_
    Bool,                                   #selected_
    Union{Base.RefValue{Vector{DataFrameRow}}, Nothing},    #bonds_
    Vector{Tuple{String,UInt8}}             #properties_
]

#----------
const residue_names = [#17
    :name_,
    :number_of_children_,
    :parent_,
    :first_child_,
    :last_child_,
    :properties_,
    :contains_selection_,
    :number_of_selected_children_,
    :number_of_children_containing_selection_,
    :selection_stamp_,
    :modification_stamp_,
    :trait_,
    :insertion_code_,
    :is_disordered_,
    :res_number_,
    :is_hetero_,
    :selected_
]

const residue_types = [
    String,                  #name_
    Union{Int64,Nothing},                   #number_of_children_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #parent_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #first_child_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #last_child_
    Vector{Tuple{String,UInt8}},            #properties_
    Union{Bool,Nothing},            #contains_selection_
    Union{Int64,Nothing},           #number_of_selected_children_
    Union{Int64,Nothing},           #number_of_children_containing_selection_
    Union{TimeStamp,Nothing},               #selection_stamp_
    Union{TimeStamp,Nothing},               #modification_stamp_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #trait_
    Union{Char,Nothing},                    #insertion_code_
    Union{Bool,Nothing},                    #is_disordered_
    Union{Int64,Nothing},                   #res_number_
    Union{Bool,Nothing},                    #is_hetero_
    Bool                                    #selected_
]

#-------------------------

const chain_names = [#13
    :id_,
    :number_of_children_,
    :parent_,
    :first_child_,
    :last_child_,
    :properties_,
    :contains_selection_,
    :number_of_selected_children_,
    :number_of_children_containing_selection_,
    :selection_stamp_,
    :modification_stamp_,
    :trait_,
    :selected_

]

const chain_types = [
    Union{Char, Nothing}                                                    ,    #id_
    Union{Int64,Nothing}                                                    ,    #number_of_children_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #parent_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #first_child_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},            #last_child_
    Vector{Tuple{String,UInt8}}                                             ,    #properties_
    Union{Bool,Nothing}                                                     ,    #contains_selection_
    Union{Int64,Nothing}                                                    ,    #number_of_selected_children_
    Union{Int64,Nothing}                                                    ,    #number_of_children_containing_selection_
    Union{TimeStamp,Nothing}                                                ,    #selection_stamp_
    Union{TimeStamp,Nothing}                                                ,    #modification_stamp_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},           #trait_
    Bool                                                                #selected_
]
#----------------------------------

const system_names = [#13
    :name_,
    :number_of_children_,
    :first_child_,
    :last_child_,
    :properties_,
    :contains_selection_,
    :number_of_selected_children_,
    :number_of_children_containing_selection_,
    :selection_stamp_,
    :modification_stamp_,
    :trait_,
    :selected_,
    :dataframesystem_
]

const system_types = [
    Union{String,Nothing},                  #name_
    Union{Int64,Nothing},                   #number_of_children_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},           #first_child_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},           #last_child_
    Vector{Tuple{String,UInt8}},            #properties_
    Union{Bool,Nothing},                    #contains_selection_
    Union{Int64,Nothing},                   #number_of_selected_children_
    Union{Int64,Nothing},                   #number_of_children_containing_selection_
    Union{TimeStamp,Nothing},               #selection_stamp_
    Union{TimeStamp,Nothing},               #modification_stamp_
    Union{Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}, Nothing},           #trait_
    Bool,                                   #selected_
    Base.RefValue{DataFrameSystem}          #dataframesystem_
]

Base.show(io::IO, dfs::DataFrameSystem) = begin
    println(io, "systems : " ,dfs.systems)
    println(io, "chains  : " ,dfs.chains)
    println(io, "residues: " ,dfs.residues)
    println(io, "atoms   : " ,dfs.atoms)
    println(io, "bonds   : " ,dfs.bonds)
end

initBondsDF(bonds_df::DataFrame, names::Vector{Symbol}, types::Vector{Type}) = begin
    for (name, Type) in zip(names,types)
        setproperty!(bonds_df,name, Type[])
    end
    return bonds_df
end

initAtomDF(atoms_df::DataFrame, names::Vector{Symbol}, types::Vector{Type}) = begin

    for (name, Type) in zip(names,types)
        setproperty!(atoms_df,name, Type[])
    end
    return atoms_df
end

initResidueDF(residue_df::DataFrame, names::Vector{Symbol}, types::Vector{Type}) = begin

    for (name, Type) in zip(names,types)
        setproperty!(residue_df, name, Type[])
    end
    return residue_df
end

initChainDF(chain_df::DataFrame, names::Vector{Symbol}, types::Vector{Type}) = begin

    for (name, Type) in zip(names,types)
        setproperty!(chain_df, name, Type[])
    end
    return chain_df

end

initSystemDF(system_df::DataFrame, names::Vector{Symbol}, types::Vector{Type}) = begin

    for (name, Type) in zip(names,types)
        setproperty!(system_df, name, Type[])
    end
    return system_df
end

insert_systems_into_dfs(input_sys::AbstractSystem, systems_df::DataFrame) = begin
    @inbounds for sys in collectSystems(input_sys)
        push!(systems_df,  ( sys[:name_], sys[:number_of_children_], nothing, nothing,
                                    sys[:properties_], sys[:contains_selection_],
                                    sys[:number_of_selected_children_],
                                    sys[:number_of_children_containing_selection_],
                                    sys[:selection_stamp_], sys[:modification_stamp_],
                                    sys[:trait_], sys[:selected_],
                                    Base.RefValue{DataFrameSystem}()
                            )
        )
    end
end

insert_chains_into_dfs(systems_df::DataFrame, chains_df::DataFrame, input_sys::AbstractSystem) = begin
    #insert chains into DFS and set the correct references to parents, children etc.
    @inbounds for (sys_df_row, comp_sys) in zip(eachrow(systems_df), collectSystems(input_sys))
        for chain in getChildren(comp_sys)::Vector{Chain}
            push!(chains_df,  ( chain[:id_], chain[:number_of_children_], Ref(sys_df_row),
                                        nothing, nothing, chain[:properties_],
                                        chain[:contains_selection_],
                                        chain[:number_of_selected_children_],
                                        chain[:number_of_children_containing_selection_],
                                        chain[:selection_stamp_], chain[:modification_stamp_],
                                        chain[:trait_], chain[:selected_]
                                    )
            )
            #if isnothing(chain.previous_), then chain is a systems first_child_
            if isnothing(chain.previous_)
                sys_df_row.first_child_ = Ref(last(chains_df))
            end

            #if isnothing(chain.next_), then chain is a systems last_child_
            if isnothing(chain.next_)
                sys_df_row.last_child_ = Ref(last(chains_df))
            end
        end
    end
end

insert_residues_into_dfs(chains_df::DataFrame, residues_df::DataFrame, input_sys::AbstractSystem) = begin
    @inbounds for (chain_df_row, comp_chain) in zip(eachrow(chains_df), collectChains(input_sys))
        for residue in getChildren(comp_chain)::Vector{Residue}
            push!(residues_df,  ( residue[:name_], residue[:number_of_children_], Ref(chain_df_row),
                                        nothing, nothing, residue[:properties_],
                                        residue[:contains_selection_],
                                        residue[:number_of_selected_children_],
                                        residue[:number_of_children_containing_selection_],
                                        residue[:selection_stamp_], residue[:modification_stamp_],
                                        residue[:trait_],           residue[:insertion_code_],
                                        residue[:is_disordered_],   residue[:res_number_],
                                        residue[:is_hetero_],       residue[:selected_]

                                    )
            )
            #if isnothing(chain.previous_), then chain is a systems first_child_
            if isnothing(residue.previous_)
                chain_df_row.first_child_ = Ref(last(residues_df))
            end

            #if isnothing(chain.next_), then chain is a systems last_child_
            if isnothing(residue.next_)
                chain_df_row.last_child_ = Ref(last(residues_df))
            end
        end
    end
end

insert_atoms_into_dfs(residues_df::DataFrame, atoms_df::DataFrame, input_sys::AbstractSystem) = begin
    @inbounds for (res_df_row, comp_res) in zip(eachrow(residues_df), collectResidues(input_sys))
        for atom in getChildren(comp_res)::Vector{Atom}
            push!(atoms_df,  ( atom[:name_], Ref(res_df_row), atom[:type_name_],
                                        atom[:element_],        atom[:radius_], atom[:type_],
                                        atom[:number_of_bonds_],atom[:formal_charge_],
                                        atom[:position_],       atom[:charge_], atom[:velocity_],
                                        atom[:force_],          atom[:occupancy_], atom[:serial_],
                                        atom[:temp_factor_],    atom[:selected_],
                                        Ref(DataFrameRow[]),         atom[:properties_]
                                    )
            )
            #if isnothing(chain.previous_), then chain is a systems first_child_
            if isnothing(atom.previous_)
                res_df_row.first_child_ = Ref(last(atoms_df))
            end

            #if isnothing(chain.next_), then chain is a systems last_child_
            if isnothing(atom.next_)
                res_df_row.last_child_ = Ref(last(atoms_df))
            end
        end
    end
end

insert_bonds_into_dfs(atoms_df::DataFrame, bonds_df::DataFrame, input_sys::AbstractSystem) = begin
    for comp_bond in viewBonds(input_sys)
        source_atom_serial::Int64 = comp_bond.source_.serial_
        target_atom_serial::Int64 = comp_bond.target_.serial_

        df_bond::DataFrameRow = last(
                                        push!(bonds_df,
                                                (comp_bond[:name_],
                                                    Ref(atoms_df[source_atom_serial,:]),
                                                    Ref(atoms_df[target_atom_serial,:]),
                                                    comp_bond[:bond_order_], comp_bond[:bond_type_],
                                                    comp_bond[:properties_]
                                                )
                                        )
                                    )
        push!(atoms_df.bonds_[source_atom_serial][], df_bond)
        push!(atoms_df.bonds_[target_atom_serial][], df_bond)
    end
end

"""
    DataFrameSystem(input_sys::System)

Constructs a [`DataFrameSystem`](@ref) out of a System. The `DataFrames` will have the same fields
and types as the objects in the tree of composites rooted in `input_sys`.
"""
DataFrameSystem(input_sys::T where T<:AbstractSystem,
    atom_names::Vector{Symbol},     atom_types::Vector{Type},
    residue_names::Vector{Symbol},  residue_types::Vector{Type},
    chain_names::Vector{Symbol},    chain_types::Vector{Type},
    system_names::Vector{Symbol},   system_types::Vector{Type},
    bond_names::Vector{Symbol},     bond_types::Vector{Type})  = begin

    systems_df =    initSystemDF(DataFrame(),   system_names,   system_types)
    chains_df =     initChainDF(DataFrame(),    chain_names,    chain_types)
    residues_df =   initResidueDF(DataFrame(),  residue_names,  residue_types)
    atoms_df =      initAtomDF(DataFrame(),     atom_names,     atom_types)
    bonds_df =      initBondsDF(DataFrame(),    bond_names,     bond_types)

    insert_systems_into_dfs(input_sys, systems_df)

    insert_chains_into_dfs(systems_df, chains_df, input_sys)

    insert_residues_into_dfs(chains_df, residues_df, input_sys)

    insert_atoms_into_dfs(residues_df, atoms_df, input_sys)

    insert_bonds_into_dfs(atoms_df, bonds_df, input_sys)

    return DataFrameSystem(systems_df, chains_df, residues_df, atoms_df, bonds_df)
end

DataFrameSystem(input_sys::T where T<:AbstractSystem) = begin
    #temp::Vector{Type} = bond_types
    dfs::DataFrameSystem = DataFrameSystem(input_sys,
                                            atom_names,     atom_types,
                                            residue_names,  residue_types,
                                            chain_names,    chain_types,
                                            system_names,   system_types,
                                            bond_names,     bond_types
    )
    for row in eachrow(dfs.systems)
        row.dataframesystem_ = Ref(dfs)
    end

    return dfs
end

coords(dfr::DataFrameRow) = dfr.position_

getChildren(dfr::DataFrameRow) = begin
    propertynames(dfr)[2] == :parent_ && return nothing
    target_df::DataFrame = getfield(getFirstChild(dfr), :df)
    first_child_idx::Int = getfield(getFirstChild(dfr), :dfrow)
    last_child_idx::Int = getfield(getLastChild(dfr), :dfrow)
    return eachrow(target_df[first_child_idx:last_child_idx, :])
end

length(df::DataFrame) = nrow(df)

getindex(df::AbstractDataFrame, idx::Int) = @view df[idx,:]

getNext(dfr::DataFrameRow) = begin
    target_df::DataFrame = getfield(dfr, :df)
    idx::Int64 = getfield(dfr, :dfrow) + 1
    return @view target_df[idx,:]
end

getSource(dfr::DataFrameRow) = dfr[:source_][]
getTarget(dfr::DataFrameRow) = dfr[:target_][]

getParent(dfr::DataFrameRow) = isnothing(dfr[:parent_]) ? nothing : dfr[:parent_][]
getBonds(dfr::DataFrameRow) = isnothing(dfr[:bonds_]) ? nothing : dfr[:bonds_][]

getFirstChild(dfr::DataFrameRow) = isnothing(dfr[:first_child_]) ? nothing : dfr[:first_child_][]
getLastChild(dfr::DataFrameRow) = isnothing(dfr[:last_child_]) ? nothing : dfr[:last_child_][]

getTrait(dfr::DataFrameRow) = isnothing(dfr[:trait_]) ? nothing : dfr[:trait_][]

isHetero(dfr::DataFrameRow) = dfr.is_hetero_

setProperty(dfr::DataFrameRow, property::Tuple{String,UInt8}) = begin
    push!(dfr[:properties_], property)
end

hasProperty(dfr::DataFrameRow, property::Tuple{String,UInt8}) = begin
    return property in dfr[:properties_]
end

hasProperty(dfr::DataFrameRow, property::String) = begin
    for prop in dfr[:properties_]
        if prop[1] == property
            return true
        end
    return false
    end
end

setProperty(dfr::DataFrameRow, property::Tuple{String,Bool}) = setProperty(dfr,(property[1],UInt8(property[2])))

#= thanks to refs i can ignore this
(==)(df1::DataFrame, df2::DataFrame)= begin

end

(==)(dfr1::DataFrameRow, dfr2::DataFrameRow) = begin
    if DataFrames._equal_names(dfr1, dfr2)
        return getfield(dfr1, :dfrow) == getfield(dfr2, :dfrow)
    end
    return false
end

(==)(sdf1::SubDataFrame, sdf2::SubDataFrame) = begin
    if DataFrames._equal_names(sdf1, sdf2)
        return getfield(sdf1, :rows) == getfield(sdf2, :rows)
    end
    return false
end
=#

Base.show(io::IO, refrow::Base.RefValue{DataFrameRow{DataFrame, DataFrames.Index}}) =
    show(io, getfield(refrow[], :dfrow))

Base.show(io::IO, refrows::Base.RefValue{Vector{DataFrameRow}}) = show(io,refrows[])


text:
"""
    PointwiseDiagnostic

A diagnostic with the same dimensions as the original grid (DG or
interpolated). Mostly used to directly copy out prognostic or
auxiliary state variables.
"""
abstract type PointwiseDiagnostic <: DiagnosticVar end
function dv_PointwiseDiagnostic end

function dv_dg_points_length(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
)
    :(npoints)
end
function dv_dg_points_index(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
)
    :(ijk)
end

function dv_dg_elems_length(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
)
    :(nrealelem)
end
function dv_dg_elems_index(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
)
    :(e)
end

function dv_dg_dimnames(::ClimateMachineConfigType, ::Type{PointwiseDiagnostic})
    ("nodes", "elements")
end
function dv_dg_dimranges(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
)
    (:(1:npoints), :(1:nrealelem))
end

function dv_i_dimnames(::ClimateMachineConfigType, ::Type{PointwiseDiagnostic})
    :(tuple(collect(keys(dims))...))
end

function dv_op(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
    lhs,
    rhs,
)
    :($lhs = $rhs)
end

# Reduction for point-wise diagnostics would be a gather, but that will probably
# blow up memory. TODO.
function dv_reduce(
    ::ClimateMachineConfigType,
    ::Type{PointwiseDiagnostic},
    array_name,
)
    quote end
end

macro pointwise_diagnostic(impl, config_type, name, project = false)
    iex = quote
        $(generate_dv_interface(:PointwiseDiagnostic, config_type, name))
        $(generate_dv_function(:PointwiseDiagnostic, config_type, name, impl))
        $(generate_dv_project(:PointwiseDiagnostic, config_type, name, project))
    end
    esc(MacroTools.prewalk(unblock, iex))
end

"""
    @pointwise_diagnostic(
        impl,
        config_type,
        name,
        units,
        long_name,
        standard_name,
        project = false,
    )

Define `name` a point-wise diagnostic variable for `config_type`,
with the specified attributes and the given implementation. If
`project` is `true`, the variable will be projected along unit
vectors (for cubed shell topologies) after interpolation.

# Example

```julia
@pointwise_diagnostic(
    AtmosGCMConfigType,
    thv,
    "K",
    "virtual potential temperature",
    "virtual_potential_temperature",
) do (
    moisture::Union{EquilMoist, NonEquilMoist},
    atmos::AtmosModel,
    states::States,
    curr_time,
    cache,
)
    ts = get!(cache, :ts) do
        recover_thermo_state(atmos, states.prognostic, states.auxiliary)
    end
    virtual_pottemp(ts)
end
```
"""
macro pointwise_diagnostic(
    impl,
    config_type,
    name,
    units,
    long_name,
    standard_name,
    project = false,
)
    iex = quote
        $(generate_dv_interface(
            :PointwiseDiagnostic,
            config_type,
            name,
            units,
            long_name,
            standard_name,
        ))
        $(generate_dv_function(:PointwiseDiagnostic, config_type, name, impl))
        $(generate_dv_project(:PointwiseDiagnostic, config_type, name, project))
    end
    esc(MacroTools.prewalk(unblock, iex))
end


text:
using GLVisualize, GeometryTypes, GLAbstraction, GLWindow
w, r = Screen()
@async r()
a = GLNormalUVMesh(Cube{Float32}(Vec3f0(0), Vec3f0(1)))
view(visualize(a))


text:
############################# kernel algebra ###################################
# NOTE: output type inference of product, sum, and power not supported for
# user-defined kernels unless Base.eltype is defined for them
################################ Product #######################################
struct Product{T, AT<:Union{Tuple, AbstractVector}, IT} <: AbstractKernel{T}
    args::AT
    input_trait::IT # could keep track of the overall input trait
end
@functor Product
function Product(k::Union{Tuple, AbstractVector})
    T = promote_type(eltype.(k)...)
    IT = sum_and_product_input_trait(k)
    Product{T, typeof(k), typeof(IT)}(k, IT)
end
Product(k...) = Product(k)
(P::Product)(τ) = prod(k->k(τ), P.args) # IDEA could check for isotropy here
(P::Product)(x, y) = prod(k->k(x, y), P.args)
# (P::Product)(x, y) = isisotropic(P) ? P(difference(x, y)) : prod(k->k(x, y), P.args)
Product(k::AbstractKernel...) = Product(k)
Product(k::AbstractVector{<:AbstractKernel}) = Product(k...)
Base.prod(k::AbstractVector{<:AbstractKernel}) = Product(k)

Base.:*(k::AbstractKernel...) = Product(k)
Base.:*(c::Number, k::AbstractKernel) = Constant(c) * k
Base.:*(k::AbstractKernel, c::Number) = Constant(c) * k

################################### Sum ########################################
struct Sum{T, AT<:Union{Tuple, AbstractVector}, IT} <: AbstractKernel{T}
    args::AT
    input_trait::IT # could keep track of the overall input trait
end
@functor Sum
function Sum(k::Union{Tuple, AbstractVector})
    T = promote_type(eltype.(k)...)
    IT = sum_and_product_input_trait(k)
    Sum{T, typeof(k), typeof(IT)}(k, IT)
end
Sum(k...) = Sum(k)
(S::Sum)(τ) = sum(k->k(τ), S.args) # should only be called if S is stationary
(S::Sum)(x, y) = sum(k->k(x, y), S.args)
# (S::Sum)(τ) = isstationary(S) ? sum(k->k(τ), S.args) : error("One argument evaluation not possible for non-stationary kernel")
# (S::Sum)(x, y) = isstationary(S) ? S(difference(x, y)) : sum(k->k(x, y), S.args)
Base.sum(k::AbstractVector{<:AbstractKernel}) = Sum(k)

Base.:+(k::AbstractKernel...) = Sum(k)
Base.:+(k::AbstractKernel, c::Number) = k + Constant(c)
Base.:+(c::Number, k::AbstractKernel) = k + Constant(c)

################################## Power #######################################
struct Power{T, K, IT} <: AbstractKernel{T}
    k::K
    p::Int
    input_trait::IT # could keep track of the overall input trait
end
@functor Power
function Power(k, p::Int)
    T = promote_type(eltype(k))
    IT = input_trait(k)
    Power{T, typeof(k), typeof(IT)}(k, p, IT)
end
(P::Power)(τ) = P.k(τ)^P.p
(P::Power)(x, y) = P.k(x, y)^P.p
Base.:^(k::AbstractKernel, p::Number) = Power(k, p)

############################ Separable Product #################################
# product kernel, but separately evaluates component kernels on different parts of the input
# NOTE: input_trait(::SeparableProduct) defaults to GenericInput
struct SeparableProduct{T, K} <: AbstractKernel{T}
    args::K # kernel for input covariances
end
@functor SeparableProduct
SeparableProduct(k...) = SeparableProduct(k)
function SeparableProduct(k::Union{Tuple, AbstractVector})
    T = promote_type(eltype.(k)...)
    SeparableProduct{T, typeof(k)}(k)
end
# both x and y have to be vectors of inputs to individual kernels
# could also consist of tuples ... so restricting to AbstractVector might not be good
function (k::SeparableProduct)(x::AbstractVector, y::AbstractVector)
    d = checklength(x, y)
    length(k.args) == d || throw(DimensionMismatch("SeparableProduct needs d = $d kernels but has r = $(length(k.args))"))
    val = one(gramian_eltype(k, x, y))
    @inbounds @simd for i in eachindex(k.args)
        ki = k.args[i]
        val *= ki(x[i], y[i])
    end
    return val
end

# if we had kernel input type, could compare with eltype(X)
function gramian(k::SeparableProduct, X::LazyGrid, Y::LazyGrid)
    length(X.args) == length(Y.args) || throw(DimensionMismatch("length(X.args) = $(length(X.args)) ≠ $(length(Y.args)) = length(Y.args)"))
    length(k.args) == length(X.args) || throw(DimensionMismatch("SeparableProduct needs d = $(length(X.args)) kernels but has r = $(length(k.args))"))
    kronecker((gramian(kxy...) for kxy in zip(k.args, X.args, Y.args))...)
end
# IDEA: if points are not on a grid, can still evaluate dimensions separately,
# and take elementwise product. can lead to efficiency gains if constituent
# matrices are of low rank (see SKIP paper)
############################### Separable Sum ##################################
# what about separable sums? do they give rise to kronecker sums? yes!
# useful for "Additive Gaussian Processes" - Duvenaud 2011
# https://papers.nips.cc/paper/2011/file/4c5bde74a8f110656874902f07378009-Paper.pdf
# IDEA: could introduce have AdditiveKernel with "order" argument, that adds higher order interactions
# NOTE: input_trait(::SeparableSum) defaults to GenericInput
struct SeparableSum{T, K} <: AbstractKernel{T}
    args::K # kernel for input covariances
end
@functor SeparableSum
SeparableSum(k...) = SeparableSum(k)
function SeparableSum(k::Union{Tuple, AbstractVector})
    T = promote_type(eltype.(k)...)
    SeparableSum{T, typeof(k)}(k)
end
function (k::SeparableSum)(x::AbstractVector, y::AbstractVector)
    d = checklength(x, y)
    length(k.args) == d || throw(DimensionMismatch("SeparableProduct needs d = $d kernels but has r = $(length(k.args))"))
    val = one(gramian_eltype(k, x, y))
    @inbounds @simd for i in eachindex(k.args)
        ki = k.args[i]
        val += ki(x[i], y[i])
    end
    return val
end

# IDEA: does gramian have special structure, like kronecker sum?
# function gramian(k::SeparableSum, X::LazyGrid, Y::LazyGrid)
#     # ⊕((gramian(kxy...) for kxy in zip(K.args, X.args, Y.args))...)
#     x_lengths = [length(x) for x in X.args]
#     y_lengths = [length(y) for y in Y.args]
#     G = zeros(length(X), length(Y))
#     for i in eachindex(k.args)
#         ki, xi, yi = k.args[i], X.args[i], Y.args[i]
#         G .+= kron(gramian(ki, xi, yi))
#     end
#     return G
# end

# convenient constructor
# e.g. separable(*, k1, k2)
separable(::typeof(*), k...) = SeparableProduct(k)
separable(::typeof(+), k...) = SeparableSum(k)
# d-separable product of k
separable(::typeof(^), k::AbstractKernel, d::Int) = SeparableProduct(Fill(k, d))


text:
using Rankers.Storages
using LOLTools: ChampionV3, http_action
using LOLToolsDataDragon.DataDragon.en_US

api_key = get(ENV, "RIOT_TOKEN", "")
region = "kr"

read_from_file = true
champs = ChampionV3.champion_rotations(api_key, region; action=update_with(region, http_action, read_from_file=read_from_file, debug=true))
for championId in champs.freeChampionIds
    print(en_US.Champions[championId].name)
    print("  --  ")
end
println()


text:
import Base.copy

"""
A structure designed to track line search information from one iteration to
another. If we have f : ℜⁿ → ℜ, then we define h(θ) = f(x + θ*d) where x and d
are vectors of same dimension and θ is a scalar, more specifically our step size.

Tracked data can include:
 - x : our current step size
 - ht : h(θ) at the current iteration
 - gt : h'(θ)
 - h₀ : h(0)
 - g₀ : h'(0)
 - start_time: the time at which the line search algorithm started.

Unless they are defined otherwise, the default value for all parameter is NaN
(except for x). They can be updated through the update! function.

Example:
```
ls_a_t = LSAtT(1.0)
update!(ls_a_t, x = 0.0, h₀ = obj(h, 0.0), g₀ = grad(h, 0.0))
```
"""
mutable struct 	LSAtT <: AbstractState

    x            :: Number
    ht           :: Number  # h(θ)
    gt           :: Number  # h'(θ)
    h₀           :: Number  # h(0)
    g₀           :: Number  # h'(0)

    start_time   :: Number

 function LSAtT(t          :: Number;
                ht         :: Number = NaN,
                gt         :: Number = NaN,
                h₀         :: Number = NaN,
                g₀         :: Number = NaN,
                start_time :: Number = NaN)

  return new(t, ht, gt, h₀, g₀, start_time)
 end
end

function update!(ls_at_t :: LSAtT;
                 x       :: FloatVoid = nothing,
                 ht      :: FloatVoid = nothing,
                 gt      :: FloatVoid = nothing,
                 h₀      :: FloatVoid = nothing,
                 g₀      :: FloatVoid = nothing,
                 tmps    :: FloatVoid = nothing,
                 kwargs...)

    if x != nothing
       ls_at_t.x  = x
    end

    if ht != nothing
       ls_at_t.ht = ht
    end

    ls_at_t.gt = gt == nothing ? ls_at_t.gt : gt
    ls_at_t.h₀ = h₀ == nothing ? ls_at_t.h₀ : h₀
    ls_at_t.g₀ = g₀ == nothing ? ls_at_t.g₀ : g₀

    ls_at_t.start_time = tmps == nothing ? ls_at_t.start_time : tmps

    return ls_at_t
end

function copy(ls_at_t :: LSAtT)
    return LSAtT(copy(ls_at_t.x),
                 ht = copy(ls_at_t.ht),
                 gt = copy(ls_at_t.gt),
                 h₀ = copy(ls_at_t.h₀),
                 g₀ = copy(ls_at_t.g₀),
                 start_time = copy(ls_at_t.start_time))
end

# function convert_ls(T, ls_at_t :: LSAtT)
#     ls_a_t_T = LSAtT(T.(copy(ls_at_t.x)))
#
#     ls_a_t_T.ht = typeof(ls_at_t.ht) != Nothing ? convert.(T, ls_at_t.ht) : ls_at_t.ht
#     ls_a_t_T.gt = typeof(ls_at_t.gt) != Nothing ? convert.(T, ls_at_t.gt) : ls_at_t.gt
#     ls_a_t_T.h₀ = typeof(ls_at_t.h₀) != Nothing ? convert.(T, ls_at_t.h₀) : ls_at_t.h₀
#     ls_a_t_T.g₀ = typeof(ls_at_t.g₀) != Nothing ? convert.(T, ls_at_t.g₀) : ls_at_t.g₀
#
#     return ls_a_t_T
# end


text:
export
    gen_data_sqth_th,
    gen_data_sqth

function gen_data_sqth_th(;
    n_data, n_points=4096,
    r_range=(0., 2.), θ_range=(0., 2π), n̄_range=(0., 0.5),
    point_dim=1000, label_dim=100,
)
    args = Matrix{Float64}(undef, 6, n_data)
    points = Array{Float64, 3}(undef, 2, n_points, n_data)
    𝛒s = Array{ComplexF64, 3}(undef, label_dim, label_dim, n_data)
    σs = Matrix{Float64}(undef, n_points, n_data)

    for i in 1:n_data
        args[:, i] .= r, θ, n̄, c1, c2, c3 = rand_arg_sqth_th(r_range, θ_range, n̄_range)

        # points
        point_dim = (r > 1) ? point_dim : label_dim
        state = construct_state_sqth_th(r, θ, n̄, c1, c2, c3, point_dim)
        _, _, σs[:, i] = gaussian_state_sampler!(view(points, :, :, i), state, 0.)

        # 𝛒s
        𝛒s[:, :, i] = (r > 1) ? state.𝛒[1:label_dim, 1:label_dim] : state.𝛒
    end

    return points, 𝛒s, args, σs
end

function gen_data_sqth(;
    n_data, n_points=4096,
    r_range=(0., 2.), θ_range=(0., 2π), n̄_range=(0., 0.5),
    point_dim=1000, label_dim=100,
)
    args = Matrix{Float64}(undef, 3, n_data)
    points = Array{Float64, 3}(undef, 2, n_points, n_data)
    𝛒s = Array{ComplexF64, 3}(undef, label_dim, label_dim, n_data)
    σs = Matrix{Float64}(undef, n_points, n_data)

    for i in 1:n_data
        args[:, i] .= r, θ, n̄ = rand_arg_sqth(r_range, θ_range, n̄_range)

        # points
        point_dim = (r > 1) ? point_dim : label_dim
        state = construct_state_sqth(r, θ, n̄, point_dim)
        _, _, σs[:, i] = gaussian_state_sampler!(view(points, :, :, i), state, 0.)

        # 𝛒s
        𝛒s[:, :, i] = (r > 1) ? state.𝛒[1:label_dim, 1:label_dim] : state.𝛒
    end

    return points, 𝛒s, args, σs
end


text:
abstract type SimulationSystem <: System end

@component @with_kw struct Spring
    center::Point3f0 = zero(Point3f0)
    k::Float32       = 0.01f0
    damping::Float32 = 0.0001f0
end

struct Oscillator <: System end

function Overseer.requested_components(::Oscillator)
    return (Spatial, Spring, UpdatedComponents, TimingData)
end

function Overseer.update(::Oscillator, m::AbstractLedger)
    td = m[TimingData][1]
    dt = td.dtime
    @inbounds for e in @entities_in(m, Spatial && Spring)
        new_v = e.velocity - (e.position - e.center) * e.k - e.velocity * e.damping
        e[Spatial] = Spatial(e.position, new_v)
    end
    return register_update(m, Spatial)
end

@component struct Orbit
    omega::Float32
    center::Point3f0
    axis::Vec3f0
end

struct Rotator <: System end
Overseer.requested_components(::Rotator) = (Spatial, Orbit, TimingData)

function Overseer.update(::Rotator, m::AbstractLedger)
    dt = m[TimingData][1].dtime
    @inbounds for e in @entities_in(m, Rotation && Spatial)
        n          = e.axis
        r          = -e.center + e.position
        theta      = e.omega * dt
        nnd        = n * dot(n, r)
        e[Spatial] = Spatial(Point3f0(e.center + nnd + (r - nnd) * cos(theta) + cross(r, n) * sin(theta)), e.velocity)
    end
    return register_update(m, Spatial)
end
struct Mover <: System end

Overseer.requested_components(::Mover) = (Spatial, TimingData)

function Overseer.update(::Mover, m::AbstractLedger)
    dt = m[TimingData][1].dtime
    for e in @entities_in(m, Spatial)
        e[Spatial] = Spatial(e.position + e.velocity * dt, e.velocity)
    end
    return register_update(m, Spatial)
end


text:
using Test,
      DataFrames,
      Jib

include("core.jl")
include("encoder.jl")
include("decoder.jl")
include("roundtrip.jl")


text:
export prior

"""
    prior(m, xs...)

Returns the minimal model required to sample random variables `xs...`. Useful for extracting a prior distribution from a joint model `m` by designating `xs...` and the variables they depend on as the prior and hyperpriors.

# Example

```jldoctest
m = @model n begin
    α ~ Gamma()
    β ~ Gamma()
    θ ~ Beta(α,β)
    x ~ Binomial(n, θ)
end;
prior(m, :θ)

# output
@model begin
        β ~ Gamma()
        α ~ Gamma()
        θ ~ Beta(α, β)
    end
```
"""
function prior(m::Model, xs...)
    po = poset(m) #Creates a new SimplePoset, so no need to copy before mutating

    newvars = collect(xs)

    for x in xs
        append!(newvars, below(po,x))
    end

    newargs = arguments(m) ∩ newvars
    setdiff!(newvars, newargs)

    theModule = getmodule(m)
    m_init = Model(theModule, newargs, NamedTuple(), NamedTuple(), nothing)
    m = foldl(newvars; init=m_init) do m0,v
        merge(m0, Model(theModule, findStatement(m, v)))
    end
end


text:
##################################################################################
# This file is part of StateSpaceEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

export sssolve!
"""
    sssolve!(model; <options>)

Solve the steady state problem for the given model.

### Options
Standard options (default values are taken from `model.options`)
  * `verbose`
  * `tol`, `maxiter` - control the stopping criteria of the solver

Specific options
  * `presolve::Bool` - whether or not to use a presolve pass. Default is `true`.
  * `method::Symbol` - choose the solution algorithm. Valid options are `:nr`
    for Newton-Raphson, `:lm` for Levenberg-Marquardt, and `:auto`. The `:auto`
    method starts with the LM algorithm and automatically switches to NR when it
    starts to converge. Default is `:nr`.
"""
function sssolve!(model::Model;
    verbose::Bool=model.options.verbose,
    tol::Float64=model.options.tol,
    maxiter::Int64=model.options.maxiter,
    presolve::Bool=true,
    nropts=Options(linesearch=false),
    lmopts=Options(),
    method::Symbol=:nr
)::Vector{Float64}

    if method ∉ (:nr, :lm, :auto)
        error("Method should be one of :nr, :lm, or :auto, not :$method")
    end

    sd = SolverData(model; presolve=presolve)
    if sd.nvars == 0
        # Nothing left to solve for
        return sd.point
    end

    ss = model.sstate
    # vars = ss.vars[sd.solve_var]
    ss_nvars = 2*length(ss.vars)
    ss_neqns = ModelBaseEcon.neqns(ss)

    lm = LMData(model, sd; lmopts...)
    nr = NRData(model, sd; nropts...)

    if verbose
        if sd.neqns < ss_neqns
            @info "Presolved $(ss_neqns - sd.neqns) equations for $(ss_nvars - sd.nvars) variables."
        end
        @info "Solving $(sd.neqns) equations for $(sd.nvars) variables."
    end

    ssvals = copy(ss.values)
    xx = ssvals[sd.solve_var]
    dx = similar(xx)

    undef_vars = Set{Int64}()

    if method ∈ (:lm, :auto)
        @timer r0, j0 = global_SS_RJ(xx, sd)
        first_step_lm!(xx, dx, r0, j0, lm; verbose=verbose)
        nf = nr0 = norm(r0, Inf)
        if verbose
            @info "0, || Fx || = $(nr0), || Δx || = $(norm(dx, Inf)), lambda = $(lm.params[1])"
        end
        xx .-= dx
    else
        nr0 = 1.0
    end

    run_nr = (method == :nr)
    for it = 1:maxiter

        if any(isnan.(xx))
            error("Nan detected.")
        end

        @timer r0, j0 = global_SS_RJ(xx, sd)
        nf = norm(r0, Inf)

        if (method == :nr) || ((method == :auto) && run_nr)
            step_nr!(xx, dx, r0, j0, nr; verbose=verbose)
            last_step_method = :nr
        else
            step_lm!(xx, dx, r0, j0, lm; verbose=verbose)
            run_nr = (nf < nr0 < 1e-2) && (lm.params[1] <= 1e-3)
            last_step_method = :lm
        end
        nr0 = nf
        xx .-= dx

        # try updating the auxiliary variables
        if any(sd.solve_var[length(model.variables) * 2 + 1:end])
            try
                global_SS_R!(r0, xx, sd)
                n1 = sum(abs2, r0)
                ssvals[sd.solve_var] .= xx
                ssvals .= update_auxvars_ss(ssvals, model)
                global_SS_R!(r0, ssvals[sd.solve_var], sd)
                n2 = sum(abs2, r0)
                if n2 < n1
                    xx .= ssvals[sd.solve_var]
                    if verbose
                        @info "    -- updated auxiliary variables"
                    end
                end
            catch
            end
        end

        nx = norm(dx, Inf)
        if verbose
            if last_step_method == :nr
                @info "$it, NR, || Fx || = $(nf), || Δx || = $(nx)"
            else
                @info "$it, LM, || Fx || = $(nf), || Δx || = $(nx), lambda = $(lm.params[1])"
            end
        end
        if nx < tol || nf < tol
            break
        end
    end
    if verbose && nf > tol
        bad_eqn, bad_var = diagnose_sstate(model)
        lbva = length(bad_var)
        if lbva > 0
            bad_var = join(bad_var, ",")
            vars_str = "Unable to solve for $(lbva) variables:\n   $bad_var"
            @warn vars_str
        end
        lbeq = length(bad_eqn)
        if lbeq > 0
            bad_eqn = (lbeq > 10) ? (join(bad_eqn[1:10], "\n   ") * "\n   . . .") : join(bad_eqn, "\n   ")
            eqns_str = "System may be inconsistent. Couldn't solve $(lbeq) equations to the required accuracy:\n   $bad_eqn"
            @warn eqns_str
        end
    end
    xx[ abs.(xx) .< tol ] .= 0.0
    ss.values[sd.solve_var] .= xx
    ss.mask[sd.solve_var] .= true
    if run_nr
        foo = (1:ss_nvars)[sd.solve_var] # global indexes of active unknowns
        for v in findall(.!nr.updated)
            # mark unknowns that were not updated as not solved
            ss.mask[foo[v]] = false
        end
    end
    return ss.values
end
@assert precompile(sssolve!, (Model,))


text:
@testset "Cake2" begin
    img, p, q, Z = generate_surface_data(Cake2(), 0.5, [0,0,1], radius = 65)
    @test maximum(img) ≈ 0.5
    @test size(img)[1] ≈ 151
    @test img[10,10] ≈ 0.5
    @test size(p)[1] ≈ 151
    @test size(q)[2] ≈ 151
    @test p[32,54] ≈ 0.0
    @test p[79,100] ≈ 2.264
    @test q[50,77] ≈ 2.2932
    @test q[42,42] ≈ 0.0
    @test Z[10,10] ≈ 0.0
    @test Z[75,75] ≈ 0.0
end


text:
maybe_parent(v::SubArray) = parent(v)
maybe_parent(v) = v

function indices(lengths)
    inds = UnitRange{Int64}[]
    firstidx=1
    for (i,len) in enumerate(lengths)
        lastidx = firstidx + len - 1
        push!(inds, firstidx:lastidx)
        firstidx = lastidx + 1
    end
    return Tuple(inds)
end

recursive_length(x) = length(x)
recursive_length(a::AbstractVector{N}) where N<:Number = length(a)
recursive_length(a::AbstractVector) = recursive_length.(a) |> sum
recursive_length(nt::NamedTuple) = values(nt) .|> recursive_length |> sum

max_stars_repo_name: jonniedie/NamedViewVectors


text:
module CAS_QMDP

push!(LOAD_PATH, "./")

using DoubleUAVs_module, GridInterpolations, JLD, PGFPlots, Interact, PyPlot

export qmdp, write_alphas, write_alphas_to_csv, load_alphas, get_qvals
export viz_policy, viz_policy_multi
export norm_angle, advisories, advisory_to_action_dict, disc
export sigma_str, sigma_v, sigma_alert_deg, sigma_coc_deg
export actual_weight, discount, terminal_state_var
export is_terminal, terminal_state

const max_iteration = 3000
const alphaTol = 1e-7

"""
	save alpha vectors to jld file
"""
function write_alphas(alphas::Matrix{Float64}, filename::AbstractString)
	path = "../policies/"
	if !ispath(path)
		mkdir(path)
	end
	filename = path * filename * ".jld"
	println("Writing alpha vectors to: " * filename)
	jldopen(filename, "w") do file
    	write(file, "alphas", alphas)
	end
end

"""
	save alphas vectors to csv file
"""
function write_alphas_to_csv(alphas::Matrix{Float64}, filename::AbstractString)
	path = "../policies/"
	if !ispath(path)
		mkdir(path)
	end
	filename = path * filename * ".csv"
	println("Writing alpha vectors to: " * filename)
	open(filename, "w") do io
        for row in 1 : size(alphas, 1)
        	for col in 1 : size(alphas, 2)
        		if col != size(alphas, 2)
            		@printf(io, "%f, ", alphas[row, col])
            	else
            		@printf(io, "%f\n", alphas[row, col])
            	end
            end
        end
    end
end

"""
	read alpha vectors from file
"""
function load_alphas(filename::AbstractString)
	alphas = jldopen(filename, "r") do file
        read(file, "alphas")
    end
	return alphas
end

function update_max_alphas!(max_alphas::Vector{Float64}, alphas::Matrix{Float64})
	for istate in 1:num_states
		max_alphas[istate] = alphas[istate, 1]
		for iaction in 2:num_actions
			if max_alphas[istate] < alphas[istate, iaction]
				max_alphas[istate] = alphas[istate, iaction]
			end
		end
	end
	# max_alphas = maximum(alphas, 2)
end

function qmdp(alphas_filename::AbstractString;
	discount::Float64=0.95,
	if_sample::Bool=false,
	verbose::Bool=true,
	penalty_action::Float64=0.02,
	penalty_conflict::Float64=1.0,
	penalty_closeness::Float64=10.0,
	penalty_nmac::Float64=1000.0)

	println("Scheme: ", alphas_filename)
	println("Running value iteration alpha vectors approximation...")
	alphas = zeros(num_states, num_actions)
	cputime = 0
	max_alphas = zeros(num_states)

	for iter in 1:max_iteration
		residual = 0
		tic()
		for istate in 1:num_states
			# println("s = ", istate)
			state = terminal_state
			if istate != num_states
				state = ind2x(grid_states, istate)
			end

			for iaction in 1:num_actions
				# print("a = ", iaction)
				advisory = advisories[iaction]
				vnext = 0

				if if_sample == false
					# state transition through sigma-sampling
					for isigma in 1 : num_sigmas
						next_state = get_next_state(state, advisory, isigma=isigma, if_sample=if_sample)
						if is_terminal(next_state)
							vnext += max_alphas[end] * weights[isigma]
						else
							vnext += interpolate(grid_states, max_alphas, next_state) * weights[isigma]
						end
					end # for isigma
				else # if_sample == true
					# random sampled noise method
					next_state = get_next_state(state, advisory, if_sample=if_sample)
					if is_terminal(next_state)
						vnext += max_alphas[end]
					else
						vnext += interpolate(grid_states, max_alphas, next_state)
					end
				end

				prev_alpha_sa = alphas[istate, iaction]
				alphas[istate, iaction] =
					get_reward(state, advisory,
						penalty_action=penalty_action,
						penalty_conflict=penalty_conflict,
						penalty_closeness=penalty_closeness,
						penalty_nmac=penalty_nmac) + discount * vnext
				residual += (alphas[istate, iaction] - prev_alpha_sa)^2
			end # for iaction
		end # for istate

		update_max_alphas!(max_alphas, alphas)

		iterTime = toq()
		cputime = cputime + iterTime
        if verbose
            @printf("Iteration %d: residual = %.2e, Qmin = %.2f, cputime = %.2e\n", iter, residual, minimum(alphas), iterTime)
        end

        if iter % 10 == 0
        	write_alphas(alphas, alphas_filename)
        end

        if residual < alphaTol
        	break
        end

        if iter == max_iteration
        	println("Warning: maximum number of iterations reached;", "solution may be inaccurate")
        end
	end # for iter

	@printf("Value iteration done!\ncputime = %.2e sec\n\n", cputime)
	# modify the values of alphas[end, :] (the terminal state values)
	# to make it numerically small.
	# alphas[] .= minimum(alphas) ?
	return alphas
end

"""
	get belief from state
"""
function get_belief(state::Vector{Float64})
	belief = spzeros(num_states, 1)
	if is_terminal(state)
		belief[end] = 1.0
	else
		indices, interp_weights = interpolants(grid_states, state) # ? state[1:3]
        for i = 1:length(indices)
            belief[indices[i]] = interp_weights[i]
        end
	end
	return belief
end

"""
	get Q values given a state (observation if uncertainty in state presents)
"""
function get_qvals(alphas::Matrix{Float64}, state::Vector{Float64})
	belief = get_belief(state)

	qvals = zeros(num_actions)
	for iaction in 1:num_actions
		for ibelief in 1:length(belief.rowval) # dot product
			qvals[iaction] += belief.nzval[ibelief] * alphas[belief.rowval[ibelief], iaction]
		end
	end
	return qvals
end

# Auxiliary #################################################################################
"""
	Visualize policy using PGFPlots
	@ outdated @
"""
function viz_policy(alphas::Matrix{Float64}, phi::Float64, v_own::Float64, v_int::Float64;
						resolution::Int64=250, mode="interact")
	# phi in [deg]
	XMIN = -2000
	XMAX = 2000
	YMIN = XMIN
	YMAX = XMAX
	function get_heat(x::Float64, y::Float64)
    	rho = sqrt(x^2 + y^2)
    	theta = norm_angle(atan2(y, x))
    	state = [rho, theta, deg2rad(phi), v_own, v_int]
    	qvals = get_qvals(alphas, state)
    	action = advisory_to_action_dict[advisories[indmax(qvals)]] # [deg]
    	if action == advisory_to_action_dict[:COC]
    	    return 2.0
    	end
     	return action
	end

	if mode == "pgf"
		th = linspace(0, 2*pi, 50)
		viz = Axis(
    		[
    	    	Plots.Image(get_heat, (Int(XMIN), Int(XMAX)), (Int(YMIN), Int(YMAX)),
    	    	    xbins = resolution, ybins = resolution,
    	    	    colormap = ColorMaps.Named("RdBu"), colorbar = true),
    	    	Plots.Node(L">", 0, 0, style="rotate=0,font=\\Huge,color=magenta"),
    	    	Plots.Node(L">", 1800, 1800, style=string("rotate=", phi, ",font=\\Huge,color=red")),
    	    	Plots.Linear(sensing_range * cos.(th), sensing_range * sin.(th), style="dashed,color=red", mark="none")
    		],
    		title="Ownship Action with Intruder Heading = " * string(phi) * " deg",
    		xlabel="x (m)", ylabel="y (m)",
    		width="10cm", height="10cm")
		return viz

	elseif mode == "csv"
		open("../viz/viz_data.csv", "w") do io
			for y in linspace(YMIN, YMAX, resolution)
				for x in linspace(XMIN, XMAX, resolution)
            		@printf(io, "%s", get_heat(x, y))
            		if x < XMAX
            			@printf(io, ", ")
            		else
            			@printf(io, "\n")
            		end
    			end
			end
		end

	elseif mode == "interact"
		th = linspace(0, 2*pi, 50)
		@manipulate for p = 0. : 30. : 360,
						v0 = v_min : 10. : v_max,
						v1 = v_min : 10. : v_max
			function get_heat_interact(x::Float64, y::Float64)
    			rho = sqrt(x^2 + y^2)
    			theta = norm_angle(atan2(y, x))
    			state = [rho, theta, deg2rad(p), v0, v1]
    			qvals = get_qvals(alphas, state)
    			action = advisory_to_action_dict[advisories[indmax(qvals)]] # [deg]
    			if action == advisory_to_action_dict[:COC]
    			    return 2.0
    			end
     			return action
			end

			viz = Axis(
    			[
    	    		Plots.Image(get_heat_interact, (Int(XMIN), Int(XMAX)), (Int(YMIN), Int(YMAX)),
    	    		    xbins = resolution, ybins = resolution,
    	    		    colormap = ColorMaps.Named("RdBu"), colorbar = true),
    	    		Plots.Node(L">", 0, 0, style="rotate=0,font=\\Huge,color=magenta"),
    	    		Plots.Node(L">", 1800, 1800, style=string("rotate=", p, ",font=\\Huge,color=red")),
    	    		Plots.Linear(sensing_range * cos.(th), sensing_range * sin.(th), style="dashed,color=red", mark="none"),
    	    		Plots.Linear(collision_range * cos.(th), collision_range * sin.(th), style="color=red", mark="none")
    			],
    			title="Ownship Action with Intruder Heading = " * string(p) * " deg",
    			xlabel="x (m)", ylabel="y (m)",
    			width="10cm", height="10cm")
			return viz
		end # @manipulate for
	end # if mode
end

"""
	Visualize pairwise policy slice using PyPlot
"""
function viz_policy(alphas::Matrix{Float64}; resolution::Int64=250)
	XMIN = -2000
	XMAX = 2000
	YMIN = XMIN
	YMAX = XMAX
	viz = figure()

	save_button = button("Save Policy Slice", value=0)
	display(save_button)
	fig_name_input = textbox("Figure Name", value="policy_slice_filename")
	display(fig_name_input)

	@manipulate for p = 0. : 30. : 360,
					v0 = v_min : 10. : v_max,
					v1 = v_min : 10. : v_max
		withfig(viz) do

			function get_heat(x::Float64, y::Float64)
    			rho = sqrt(x^2 + y^2)
    			theta = norm_angle(atan2(y, x))
    			state = [rho, theta, deg2rad(p), v0, v1]
    			qvals = get_qvals(alphas, state)
    			action = advisory_to_action_dict[advisories[indmax(qvals)]] # [deg]
    			if action == advisory_to_action_dict[:COC]
    			    return 2.0
    			end
     			return action
			end

			action_map = zeros(resolution, resolution)
			x_arr = linspace(XMIN, XMAX, resolution)
			y_arr = linspace(YMIN, YMAX, resolution)

			for j in 1:resolution
				for i in 1:resolution
					action_map[j, i] = get_heat(x_arr[i], y_arr[j])
				end
			end

			action_map[1, 1] = 10.0
			action_map[1, 2] = -10.0

			action_map = flipdim(action_map, 1)

			imshow(action_map, cmap="RdBu", extent=(XMIN, XMAX, YMIN, YMAX))
			cbar = colorbar(label="Advisory", ticks=[-10, -5, 0, 2, 5, 10])
			cbar[:set_ticklabels](["-10", "-5", "0", "COC", "5", "10"])

			# plot Ownship triangle
			scatter(0, 0 , marker=">", color="green", s=180)

			# plot intruder arrow
			airspace_dim = XMAX - XMIN
			arrow_len = airspace_dim / 8
			arrow(0.7 * XMAX, 0.7 * YMAX,
				arrow_len * cos(deg2rad(p)), # dx
				arrow_len * sin(deg2rad(p)), # dy
				head_width = airspace_dim / 16,
				width = airspace_dim / 50,
				head_length = airspace_dim / 12,
				head_starts_at_zero = "true",
				facecolor = "red",
				length_includes_head = "true")

			# plot NMAC circle
			th = linspace(0, 2*pi, 50)
			PyPlot.plot(collision_range * cos.(th), collision_range * sin.(th), linestyle="--", color="red")

			# plot sensing circle
			PyPlot.plot(sensing_range * cos.(th), sensing_range * sin.(th), linestyle="--", color="red")

			title("Ownship Action with Intruder Heading = " * string(p) * " deg")
    		xlabel("x (m)")
    		ylabel("y (m)")

    		o_button = observe(save_button)
    		o_textbox = observe(fig_name_input)
			function save_slice(signal::Int64, filename=o_textbox[])
				if signal > 0
					savefig(filename * ".pdf")
				end
			end
			on(save_slice, o_button)

    		return viz

		end # withfig
	end # @manipulate for
end

"""
	Visualize multi-intruder policy slice using PyPlot
"""
function viz_policy_multi(alphas::Matrix{Float64}; num_intruders::Int64=1, resolution::Int64=250, method::Symbol=:closest)
	XMIN = -2000
	XMAX = 2000
	YMIN = XMIN
	YMAX = XMAX
	v = 50.0
	viz = figure()
	save_button = button("Save Policy Slice", value=0)
	display(save_button)
	fig_name_input = textbox("Figure Name", value="policy_slice_filename")
	display(fig_name_input)

	state_matrix = ones(state_dim, 5) * Inf # at most 5 intruders

	@manipulate for int_1_x = XMIN : 250 : XMAX,
					int_1_y = YMIN : 250 : YMAX,
					int_1_p = 0. : 30. : 360,
					p       = 0. : 30. : 360
					# more intruders to be added ...

		withfig(viz) do

			state_matrix[:, 1] .= [sqrt(int_1_x^2 + int_1_y^2), norm_angle(atan2(int_1_y, int_1_x)), deg2rad(int_1_p), v, v]
			# more intruders to be added ...

			function get_heat(x::Float64, y::Float64)
    			rho = sqrt(x^2 + y^2)
    			theta = norm_angle(atan2(y, x))
    			state_matrix[:, num_intruders + 1] .= [rho, theta, deg2rad(p), v, v]
    			if method == :closest
    				closest_intruder_ind = indmin(state_matrix[rho_ind, :])
    				state = Vector(state_matrix[:, closest_intruder_ind])
     				qvals = get_qvals(alphas, state)
     				action = advisories[indmax(qvals)] # symbol
     		    else # if method != :closest
     				q_matrix = Matrix{Float64}(num_actions, 0)
     				sum_qvals = zeros(num_actions)
     				for i in 1 : num_intruders + 1
     					state_i = Vector(state_matrix[:, i])
     					qvals = get_qvals(alphas, state_i)
     					q_matrix = hcat(q_matrix, qvals)
     					sum_qvals += qvals
     				end
     				if method == :maxmin
     					min_qvals = minimum(q_matrix, 2)
						the_ind = indmax(min_qvals)
					elseif method == :maxsum
						the_ind = indmax(sum_qvals)
					else
						error("No such method!")
					end # if method
     				action = advisories[the_ind]
     			end # if method

     			if action == :COC
    				return 2.0
    			end
     			return advisory_to_action_dict[action] # get numerical value of action
			end # get_heat

			# plot
			action_map = zeros(resolution, resolution)
			x_arr = linspace(XMIN, XMAX, resolution)
			y_arr = linspace(YMIN, YMAX, resolution)

			for j in 1:resolution
				for i in 1:resolution
					action_map[j, i] = get_heat(x_arr[i], y_arr[j])
				end
			end

			action_map[1, 1] = 10.0
			action_map[1, 2] = -10.0

			action_map = flipdim(action_map, 1)

			imshow(action_map, cmap="RdBu", extent=(XMIN, XMAX, YMIN, YMAX))
			cbar = colorbar(label="Advisory", ticks=[-10, -5, 0, 2, 5, 10])
			cbar[:set_ticklabels](["-10", "-5", "0", "COC", "5", "10"])

			# plot Ownship triangle
			scatter(0, 0 , marker=">", color="green", s=180)

			# plot intruder arrow
			airspace_dim = XMAX - XMIN
			arrow_len = airspace_dim / 10
			arrow(0.7 * XMAX, 0.7 * YMAX,
				arrow_len * cos(deg2rad(p)), # dx
				arrow_len * sin(deg2rad(p)), # dy
				head_width = airspace_dim / 16,
				width = airspace_dim / 50,
				head_length = airspace_dim / 14,
				head_starts_at_zero = "true",
				facecolor = "red",
				length_includes_head = "true")

			# plot existing intruder 1 arrow
			airspace_dim = XMAX - XMIN
			arrow_len = airspace_dim / 12
			arrow(int_1_x, int_1_y,
				arrow_len * cos(deg2rad(int_1_p)), # dx
				arrow_len * sin(deg2rad(int_1_p)), # dy
				head_width = airspace_dim / 20,
				width = airspace_dim / 70,
				head_length = airspace_dim / 18,
				head_starts_at_zero = "true",
				facecolor = "blue",
				length_includes_head = "true")
			# more intruders to be added ...

			# plot NMAC circle
			th = linspace(0, 2*pi, 50)
			PyPlot.plot(collision_range * cos.(th), collision_range * sin.(th), linestyle="--", color="red")

			# plot sensing circle
			PyPlot.plot(sensing_range * cos.(th), sensing_range * sin.(th), linestyle="--", color="red")

			title("Ownship Action with Intruder Heading = " * string(p) * " deg, " * string(method))
    		xlabel("x (m)")
    		ylabel("y (m)")

    		o_button = observe(save_button)
    		o_textbox = observe(fig_name_input)
			function save_slice(signal::Int64, filename=o_textbox[])
				if signal > 0
					savefig(filename * ".pdf")
				end
			end
			on(save_slice, o_button)

    		return viz

		end # withfig
	end
end

end # module

max_stars_repo_name: parachutel/DeepCorrectionCAS


text:
using Documenter, RHEOS

##

# convert Jupyter notebooks to markdown and do some simple preprocessing before
# feeding into Documenter.jl

function convert_to_markdown(file)
    run(`jupyter nbconvert examples/$file --to markdown --template docs/documenter.tpl --output-dir docs/src-staging`)
    return "docs/src-staging/$(replace(file, "ipynb"=>"md"))"
end

function convert_equations(file)
    contents = read(file, String)
    contents = replace(contents, r"\$\$(.*?)\$\$"s => s"""```math
    \g<1>
    ```""")
    contents = replace(contents, r"\* \$(.*?)\$" => s"* ``\g<1>``") # starting a line with inline math screws up tex2jax for some reason
    write(file, contents)
    return file
end

rm("docs/src-staging", force=true, recursive=true)
mkdir("docs/src-staging")
for file in readdir("examples")
    if endswith(file, "ipynb")
        file |> convert_to_markdown |> convert_equations
    elseif !startswith(file, ".")
        cp("examples/$file", "docs/src-staging/$file")
    end
end
symlink("../src/index.md","docs/src-staging/index.md")
symlink("../src/API.md","docs/src-staging/API.md")
symlink("../src/architecture.md","docs/src-staging/architecture.md")

mkdir("docs/src-staging/assets")
cp("docs/src/assets/logo.png", "docs/src-staging/assets/logo.png")
cp("docs/src/assets/diagram.svg", "docs/src-staging/assets/diagram.svg")
##

# highlight output cells (i.e. anything withouout a language specified) white

@eval Documenter.Writers.HTMLWriter function mdconvert(c::Markdown.Code, parent::MDBlockContext; kwargs...)
    @tags pre code
    language = isempty(c.language) ? "none" : c.language
    pre[".language-$(language)"](code[".language-$(language)"](c.code))
end

##

makedocs(modules=[RHEOS],
         doctest=false, clean=true,
         format =Documenter.HTML(),
         sitename=" ",
         source = "src-staging",
         authors="Louis Kaplan",
         pages = [
             "Home" => "index.md",
	     "Architecture" => "architecture.md",
	     "File I/O" => "fileIO.md",
	     "Preprocessing" => "preprocessing.md",
	     "Generating Data" => "generatingdata.md",
	     "Fit and Predict" => ["Time data" => "fitpredictTime.md",
				   "Frequency data" => "fitpredictFreq.md"],
             "Models" => ["Basic Elements" => "elements.md",
			 "Maxwell" => "fractionalMaxwell.md",
			 "Kelvin-Voigt" => "fractionalKelvinVoigt.md",
			 "Zener" => "fractionalZener.md",
			 "Poynting-Thomson" => "fractionalPT.md",
		         "Burgers" => "burgers.md",
			 "Create Your Model" => "createmodel.md"],
	     "Examples" => "examples.md",
            #  "Fitting Data" => "fittingdata.md",
            #  "Predicting Responses" => "predictingresponse.md",
              "API" => "API.md"
         ]
         )

deploydocs(
    repo = "github.com/JuliaRheology/RHEOS.jl.git",
    deps = nothing,
    make = nothing,
    target = "build"
)


text:
function DiffEqBase.__solve(prob::DiffEqBase.AbstractSteadyStateProblem,
                            alg::SteadyStateDiffEqAlgorithm,args...;
                            abstol=1e-8,kwargs...)

  if prob.f.mass_matrix != I
    error("This solver is not able to use mass matrices.")
  end

  if typeof(prob.u0) <: Number
    u0 = [prob.u0]
  else
    u0 = vec(deepcopy(prob.u0))
  end

  sizeu = size(prob.u0)
  p = prob.p

  if !isinplace(prob) && (typeof(prob.u0)<:AbstractVector || typeof(prob.u0)<:Number)
    f! = (du,u) -> (du[:] = prob.f(u,p,Inf); nothing)
  elseif !isinplace(prob) && typeof(prob.u0)<:AbstractArray
    f! = (du,u) -> (du[:] = vec(prob.f(reshape(u, sizeu),p,Inf)); nothing)
  elseif typeof(prob.u0)<:AbstractVector
    f! = (du,u) -> (prob.f(du,u,p,Inf); nothing)
  else # Then it's an in-place function on an abstract array
    f! = (du,u) -> (prob.f(reshape(du, sizeu),
                    reshape(u, sizeu),p,Inf);
                    du=vec(du); nothing)
  end

  # du = similar(u)
  # f = (u) -> (f!(du,u); du) # out-of-place version

  if typeof(alg) <: SSRootfind
    u = alg.nlsolve(f!,u0,abstol)
    resid = similar(u)
    f!(resid,u)
    DiffEqBase.build_solution(prob,alg,u,resid;retcode = :Success)
  else
    error("Algorithm not recognized")
  end
end

function DiffEqBase.__solve(prob::DiffEqBase.AbstractSteadyStateProblem,
                            alg::DynamicSS,args...;kwargs...)

  _prob = ODEProblem(prob.f,prob.u0,(0.0,Inf),prob.p)
  solve(_prob,alg.alg,args...;kwargs...,
        callback=TerminateSteadyState(alg.abstol,alg.reltol))
end


text:
module AutoFormat

# package code goes here

# minimal function export list
export format_demo

# main format function
function format_demo(input::String, output::String, tab_width::Int)
  result = String[]
  layer = 1
  block_start = ["^if" "^while" "^for" " begin " "^try" "^function" "^type" "^let" "^macro" " do " "^immutable"] # mask quote
  block_other = ["^else\$", "^elseif "]
  open(input, "r") do file
    for line in eachline(file)
      line = rstrip(lstrip(line))
      line = replace(line, "print\(", "#print\(")
      if sum(check_block_keyword(line, block_start))>0
        line = process_line(line)
        line = " "^(tab_width*(layer-1))*line
        layer = layer+1
      elseif sum(check_block_keyword(line, block_start))>0 && ismatch(r"end$", line)
        line = " "^(tab_width*(layer-1))*line
      elseif sum(check_block_keyword(line, block_other))>0
        layer = layer-1
        line = process_line(line)
        line = " "^(tab_width*(layer-1))*line
        layer = layer+1
      elseif ismatch(r"^end$", line)
        layer = layer-1
        line = process_line(line)
        line = " "^(tab_width*(layer-1))*line
      elseif ismatch(r"#print\(", line) || ismatch(r"^#", line)
        line = line
        line = " "^(tab_width*(layer-1))*line
      else
        line = process_line(line)
        line = " "^(tab_width*(layer-1))*line
      end
      line = replace(line, "#print\(", "print\(")
      result = push!(result, line*"\n")
    end
    file_output = open(output, "w+")
    write(file_output, result)
    close(file_output)
    print(output*" writing successfully")
  end
end

#add space (, = -= +=)
function add_space(line::String)
  line = replace(line, r"\s+", " ")
  line = replace(line, "-=", "MaSkEdMiNuS")
  line = replace(line, "+=", "MaSkEdPlUs")
  line = replace(line, r"\s=\s|=\s|\s=|=", " = ")
  line = replace(line, r"\s,\s|,\s|\s,|,", ", ")
  line = replace(line, r"\sMaSkEdPlUs\s|MaSkEdPlUs\s|\sMaSkEdPlUs|MaSkEdPlUs", " += ")
  line = replace(line, r"\sMaSkEdMiNuS\s|MaSkEdMiNuS\s|\sMaSkEdMiNuS|MaSkEdMiNuS", " -= ")
  return line
end

# format file line by line
function process_line(line::String)
  block_count = [1:length(split(line, "\""))]
  temp = split(line, "\"")
  temp[block_count[map(isodd, block_count)]] = map(add_space, temp[block_count[map(isodd, block_count)]])
  templine = join(temp, "\"")
  if ismatch(r";", templine) && !ismatch(r"\";\"", templine)
    templine = join(map(lstrip, map(rstrip, split(templine, ";"))), "\n")*"\n"
  end
  return templine
end

# check whether contain a block keyword
function check_block_keyword(line::String, list::Array)
  block_count = [1:length(split(line, "\""))]
  temp = split(line, "\"")
  templine = join(temp[block_count[map(isodd, block_count)]], "\"")
  return map((x)->ismatch(Regex(list[x]), templine), 1:length(list))
end

end # module


text:
using Documenter, RTFOOL

makedocs(
    format = :html,
    sitename = "RTFOOL",
    pages = [
        "index.md",
    ]
)