text:
function preParse(data_dictionary)
    regulatoryRules = deepcopy(data_dictionary["regulatoryRules"])
    regulatoryFull = vcat(data_dictionary["regulatoryGenes"], data_dictionary["regulatoryInputs1"], data_dictionary["regulatoryInputs2"])
    rulesList = Array{Any, 2}(nothing, size(regulatoryRules))
    for k = 1:length(regulatoryRules)
        fields = split(regulatoryRules[k], r"([\s~|&()])", keepempty = false)
        newfields = copy(fields)
        for j = 1:length(fields)
            word = fields[j]
            if word == "true"
                newfields[j] = ""
            elseif word == "false"
                newfield[j] = "false"
            else
                if ~isempty(findall(regulatoryFull .== word))
                    index = findall(regulatoryFull .== word)
                    newfields[j] = "initial_regstate("*string(index[1][1])*")"
                else
                    newfields[j] = ""
                end
            end
        end
            newRule = regulatoryRules[k]
            for j = 1:length(fields)
                newRule = replace(newRule, fields[j] => newfields[j])
            end
            rulesList[k] = newRule
    end
    for i = 1:length(rulesList)
        rulesList[i] = replace(rulesList[i], "||" => "|")
        rulesList[i] = replace(rulesList[i], "&&" => "&")
    end

    return rulesList
end

function process(expression::Expr)

    # new expression -
    new_expression = copy(expression)

    # get the args for this expression -
    args_array = expression.args
    number_of_arguments = length(args_array)
    for arg_index = 1:number_of_arguments

        # get the arg -
        arg = args_array[arg_index]

        if typeof(arg) == Symbol
            if arg == :|
                new_expression.args[arg_index] = :max
            elseif arg == :&
                new_expression.args[arg_index] = :min
            elseif arg == :x
                new_expression.args[arg_index] = :fn
            elseif arg == :initial_regstate
                new_expression.args[arg_index] = :fn
            elseif arg == :~
                new_expression.args[arg_index] = :~
            end
        elseif typeof(arg) == Expr
            new_expression.args[arg_index] = process(arg)
        end
    end

    return new_expression
end

function generate_rules_function(cobra_dictionary, path_to_output_rules_function::String)

    # Initialize -
    expression_array = Array{Any,1}()

    # Get the rules -
    rules_array = cobra_dictionary["rules"]

    # how many rules do we have?
    number_of_rules = length(rules_array)

    # create an array of expressions -
    for rule_string in rules_array

        # build an expression -
        local_expr = Meta.parse(rule_string)

        # process this expression -
        if local_expr == nothing
            push!(expression_array, nothing)
        else

            # work on this tree -
            converted_expression = process(local_expr)

            # grab -
            push!(expression_array, converted_expression)
        end
    end

    # dump reaction array to disk -
    open(path_to_output_rules_function, "w") do f

        # buffer -
        buffer = String[]
        tmp_value = "function calculate_rules_vector(data_dictionary::Dict{String,Any}, value_array::Array{Float64,1})"
        push!(buffer, tmp_value)
        push!(buffer, "")
        push!(buffer, "\t# Initialize - ")
        push!(buffer, "\tv = zeros($(number_of_rules))")
        push!(buffer, "")

        # write the buffer -
        for line_item in buffer
            write(f,"$(line_item)\n")
        end

        # define some *internal* functions -
        default_fn_buffer = String[]
        push!(default_fn_buffer,"\tfunction estimate_default(data_dictionary)")
        push!(default_fn_buffer,"\t\treturn 1.0")
        push!(default_fn_buffer,"\tend")
        push!(default_fn_buffer, "\tdefault() = estimate_default(data_dictionary)")
        push!(default_fn_buffer,"")

        # write the default_fn_buffer -
        for line_item in default_fn_buffer
            write(f,"$(line_item)\n")
        end

        # fn buffer -
        fn_buffer = String[]
        push!(fn_buffer, "\tfunction calculate_fn_value(x, value_array)")
        push!(fn_buffer, "\t\treturn value_array[x]")
        push!(fn_buffer,"\tend")
        push!(fn_buffer, "\tfn(x) = calculate_fn_value(x, value_array)")
        push!(fn_buffer,"")

        # write the default_fn_buffer -
        for line_item in fn_buffer
            write(f,"$(line_item)\n")
        end

        # write the rules -
        counter = 1
        for rule in expression_array

            line_item = ""

            if rule == nothing
                line_item = "\tv[$(counter)] = default()"
            else

                # get the rule string -
                rule_string = repr(rule)
                line_item = "\tv[$(counter)] = $(rule_string[3:end-1])"
            end

            counter = counter + 1
            write(f,"$(line_item)\n")
        end

        # return -
        write(f,"\n")
        write(f,"\treturn v\n")

        # write the end -
        write(f, "end\n")
    end
end

function generate_regulatory_function(cobra_dictionary, path_to_output_rules_function::String)

    # Initialize -
    expression_array = Array{Any,1}()

    # Get the rules -
    full_rules_array = preParse(cobra_dictionary)

    # how many rules do we have?
    number_of_rules = length(full_rules_array)

    # create an array of expressions -
    for rule_string in full_rules_array

        # build an expression -
        local_expr = Meta.parse(rule_string)
        println(local_expr)
        # process this expression -
        if local_expr == nothing
            push!(expression_array, nothing)
        else

            # work on this tree -
            converted_expression = process(local_expr)

            # grab -
            push!(expression_array, converted_expression)
        end
    end

    # dump reaction array to disk -
    open(path_to_output_rules_function, "w") do f

        # buffer -
        buffer = String[]
        tmp_value = "function calculate_full_rules_vector(data_dictionary::Dict{String,Any}, value_array::BitArray{2})"
        push!(buffer, tmp_value)
        push!(buffer, "")
        push!(buffer, "\t# Initialize - ")
        push!(buffer, "\tv = zeros($(number_of_rules))")
        push!(buffer, "")

        # write the buffer -
        for line_item in buffer
            write(f,"$(line_item)\n")
        end

        # define some *internal* functions -
        default_fn_buffer = String[]
        push!(default_fn_buffer,"\tfunction estimate_default(data_dictionary)")
        push!(default_fn_buffer,"\t\treturn 1.0")
        push!(default_fn_buffer,"\tend")
        push!(default_fn_buffer, "\tdefault() = true")
        push!(default_fn_buffer,"")

        # write the default_fn_buffer -
        for line_item in default_fn_buffer
            write(f,"$(line_item)\n")
        end

        # fn buffer -
        fn_buffer = String[]
        push!(fn_buffer, "\tfunction calculate_fn_value(x, value_array)")
        push!(fn_buffer, "\t\treturn value_array[x]")
        push!(fn_buffer,"\tend")
        push!(fn_buffer, "\tfn(x) = calculate_fn_value(x, value_array)")
        push!(fn_buffer,"")

        # write the default_fn_buffer -
        for line_item in fn_buffer
            write(f,"$(line_item)\n")
        end

        # write the rules -
        counter = 1
        for rule in expression_array

            line_item = ""

            if rule == nothing
                line_item = "\tv[$(counter)] = default()"
            else

                # get the rule string -
                rule_string = repr(rule)
                line_item = "\tv[$(counter)] = $(rule_string[3:end-1])"
            end

            counter = counter + 1
            write(f,"$(line_item)\n")
        end

        # return -
        write(f,"\n")
        write(f,"\treturn v\n")

        # write the end -
        write(f, "end\n")
    end
end


text:
using Test
using IFGF
using StaticArrays
using IFGF: chebeval_novec, chebeval_vec, chebeval_vec1d, chebtransform_fftw!, chebtransform_native!, HyperRectangle

for T in (Float64, ComplexF64)
    for dim in 1:3

        @testset "Cheyshev transform: T=$T, dim=$dim" begin
            vals = rand(T,ntuple(i->5+i,dim))
            c1   = chebtransform_fftw!(vals)
            c2   = chebtransform_native!(vals)
            @test c1 ≈ c2
        end

        @testset "Clenshaw summation: T=$T, dim=$dim" begin
            sz   = ntuple(i->3+i,dim)
            c    = rand(T,prod(sz))
            x     = @SVector rand(dim)
            rec = HyperRectangle(ntuple(i->-1.0,dim),ntuple(i->1.0,dim))
            SZ = Val(sz)
            v1 = chebeval_novec(c,x,rec,SZ)
            v2 = chebeval_vec1d(c,x,rec,SZ)
            @test v1 ≈ v2
        end

    end
end

ext: jl

max_stars_repo_name: WaveProp/IFGF

max_issues_repo_name: WaveProp/IFGF

max_forks_repo_name: WaveProp/IFGF


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

const CachableModel = Union{MOI.ModelLike, AbstractStochasticStructure}

struct SolutionCache{M <: MOI.ModelLike} <: MOI.ModelLike
    modattr::Dict{MOI.AbstractModelAttribute, Any}
    varattr::Dict{MOI.AbstractVariableAttribute, Dict{VI, Any}}
    conattr::Dict{MOI.AbstractConstraintAttribute, Dict{CI, Any}}

    function SolutionCache(model::MOI.ModelLike)
        M = typeof(model)
        cache = new{M}(Dict{MOI.AbstractModelAttribute, Any}(),
                       Dict{MOI.AbstractVariableAttribute, Dict{VI, Any}}(),
                       Dict{MOI.AbstractConstraintAttribute, Dict{CI, Any}}())
        return cache
    end

    function SolutionCache(model::MOI.ModelLike, variables::Vector{<:VI}, constraints::Vector{<:CI})
        M = typeof(model)
        cache = new{M}(Dict{MOI.AbstractModelAttribute, Any}(),
                       Dict{MOI.AbstractVariableAttribute, Dict{VI, Any}}(),
                       Dict{MOI.AbstractConstraintAttribute, Dict{CI, Any}}())
        cache_model_attributes!(cache, model)
        cache_variable_attributes!(cache, model, variables)
        cache_constraint_attributes!(cache, model, constraints)
        return cache
    end
end

function cache_model_attributes!(cache::SolutionCache, src::CachableModel)
    attributes = [MOI.ObjectiveValue(),
                  MOI.DualObjectiveValue(),
                  MOI.ObjectiveBound(),
                  MOI.RelativeGap(),
                  MOI.SolveTime(),
                  MOI.SimplexIterations(),
                  MOI.BarrierIterations(),
                  MOI.NodeCount(),
                  MOI.TerminationStatus(),
                  MOI.RawStatusString(),
                  MOI.PrimalStatus(),
                  MOI.DualStatus()]
    for attr in attributes
        try
            value = MOI.get(src, attr)
            MOI.set(cache, attr, value)
        catch
        end
    end
    return nothing
end

function cache_model_attributes!(cache::SolutionCache, src::CachableModel, stage::Integer, scenario_index::Integer)
    attributes = [MOI.ObjectiveValue(),
                  MOI.DualObjectiveValue(),
                  MOI.ObjectiveBound(),
                  MOI.RelativeGap(),
                  MOI.SolveTime(),
                  MOI.SimplexIterations(),
                  MOI.BarrierIterations(),
                  MOI.NodeCount(),
                  MOI.TerminationStatus(),
                  MOI.RawStatusString(),
                  MOI.PrimalStatus(),
                  MOI.DualStatus()]
    for attr in attributes
        try
            value = MOI.get(src, ScenarioDependentModelAttribute(stage, scenario_index, attr))
            MOI.set(cache, attr, value)
        catch
        end
    end
    return nothing
end

function cache_variable_attributes!(cache::SolutionCache, src::CachableModel, variables::Vector{<:VI})
    attr = MOI.VariablePrimal()
    for vi in variables
        try
            value = MOI.get(src, attr, vi)
            MOI.set(cache, attr, vi, value)
        catch
        end
    end
    return nothing
end
function cache_variable_attributes!(cache::SolutionCache, src::CachableModel)
    variables = MOI.get(src, MOI.ListOfVariableIndices())
    cache_variable_attributes!(cache, src, variables)
    return nothing
end
function cache_variable_attributes!(cache::SolutionCache, src::CachableModel, variables::Vector{<:VI}, stage::Integer, scenario_index::Integer)
    attr = MOI.VariablePrimal()
    for vi in variables
        try
            value = MOI.get(src, ScenarioDependentVariableAttribute(stage, scenario_index, attr), vi)
            MOI.set(cache, attr, vi, value)
        catch
        end
    end
    return nothing
end

function cache_constraint_attributes!(cache::SolutionCache, src::CachableModel, constraints::Vector{<:CI})
    conattrs = [MOI.ConstraintPrimal(), MOI.ConstraintDual(), MOI.ConstraintBasisStatus()]
    for ci in constraints
        for attr in conattrs
            try
                value = MOI.get(src, attr, ci)
                MOI.set(cache, attr, ci, value)
            catch
            end
        end
    end
    return nothing
end
function cache_constraint_attributes!(cache::SolutionCache, src::CachableModel)
    ctypes = filter(t -> is_decision_type(t[1]), MOI.get(src, MOI.ListOfConstraints()))
    constraints = mapreduce(vcat, ctypes) do (F, S)
        return MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
    end
    cache_constraint_attributes!(cache, src, constraints)
    return nothing
end
function cache_constraint_attributes!(cache::SolutionCache, src::CachableModel, constraints::Vector{<:CI}, stage::Integer, scenario_index::Integer)
    conattrs = [MOI.ConstraintPrimal(), MOI.ConstraintDual(), MOI.ConstraintBasisStatus()]
    for ci in constraints
        for attr in conattrs
            try
                value = MOI.get(src, ScenarioDependentConstraintAttribute(stage, scenario_index, attr), ci)
                MOI.set(cache, attr, ci, value)
            catch
            end
        end
    end
    return nothing
end

function MOI.get(cache::SolutionCache{M},
                 attr::MOI.AbstractModelAttribute) where M
    if !haskey(cache.modattr, attr)
        throw(ArgumentError("Origin model type $M did not support accessing the attribute $attr"))
    end
    return cache.modattr[attr]
end

function MOI.get(cache::SolutionCache{M},
                 attr::MOI.AbstractVariableAttribute,
                 vi::VI) where M
    attribute_dict = get(cache.varattr, attr, nothing)
    if attribute_dict === nothing
        throw(ArgumentError("Origin model type $M did not support accessing the attribute $attr"))
    end
    if !haskey(attribute_dict, vi)
        throw(ArgumentError("Origin model type $M did not support accessing the attribute $attr"))
    end
    return get(attribute_dict, vi, nothing)
end

function MOI.get(cache::SolutionCache{M},
                 attr::MOI.AbstractConstraintAttribute,
                 ci::CI) where M
    attribute_dict = get(cache.conattr, attr, nothing)
    if attribute_dict === nothing
        throw(ArgumentError("Origin model type $M did not support accessing the attribute $attr"))
    end
    if !haskey(attribute_dict, ci)
        throw(ArgumentError("Origin model type $M did not support accessing the attribute $attr"))
    end
    return get(attribute_dict, ci, nothing)
end

function MOI.set(cache::SolutionCache,
                 attr::MOI.AbstractModelAttribute,
                 value)
    cache.modattr[attr] = value
    return nothing
end

function MOI.set(cache::SolutionCache,
                 attr::MOI.AbstractVariableAttribute,
                 vi::VI,
                 value)
    if !haskey(cache.varattr, attr)
        cache.varattr[attr] = Dict{VI, Any}()
    end
    cache.varattr[attr][vi] = value
    return nothing
end

function MOI.set(cache::SolutionCache,
                 attr::MOI.AbstractConstraintAttribute,
                 ci::CI,
                 value)
    if !haskey(cache.conattr, attr)
        cache.conattr[attr] = Dict{VI, Any}()
    end
    cache.conattr[attr][ci] = value
    return nothing
end

MOI.supports(::SolutionCache, ::MOI.AbstractModelAttribute) = true
MOI.supports(::SolutionCache, ::Union{MOI.AbstractVariableAttribute, MOI.AbstractConstraintAttribute}, ::Type{<:MOI.Index}) = true


text:
module TaskOverheads

using DataFrames
using DataFrames: AbstractDataFrame
using Plots
using Statistics: median
using ThreadsAPIBenchmarks: load_task_overheads

load_rawdf(args...; kwargs...) = DataFrame(collect(load_task_overheads(args...; kwargs...)))

function plot_hists(rawdf::AbstractDataFrame)
    # ad-hoc filtering of extremes
    th_spawn = median(rawdf.spawn_overhead) * 100
    th_sync = median(rawdf.sync_overhead) * 100
    ok = (rawdf.spawn_overhead .< th_spawn) .& (rawdf.sync_overhead .< th_sync)
    rawdf = rawdf[ok, :]

    groups = groupby(rawdf, :nthreads, sort = true)
    colors = getindex.(Ref(cgrad(:vik)), range(0, 1, length = length(groups)))
    plt_spawn = plot()
    plt_sync = plot()

    for (i, (k, g)) in enumerate(pairs(groups))
        stephist!(
            plt_spawn,
            g.spawn_overhead,
            label = "N=$(k.nthreads)",
            color = colors[i],
            linewidth = 4,
        )
        stephist!(
            plt_sync,
            g.sync_overhead,
            label = "N=$(k.nthreads)",
            color = colors[i],
            linewidth = 4,
        )
    end

    return plot(
        plot!(plt_spawn, ylabel = "spawn"),
        plot!(plt_sync, ylabel = "sync", xlabel = "overhead [ns]"),
        yticks = nothing,
        xlim = (0, 15_000),
        layout = (2, 1),
        left_margin = 5 * Plots.mm,
    )
end

function plot_seq(rawdf::AbstractDataFrame)
    groups = groupby(rawdf, :nthreads, sort = true)
    plts = map(pairs(groups)) do (k, g)
        p = plot()
        scatter!(
            p,
            g.spawn_overhead,
            markershape = :xcross,
            markersize = 3,
            label = "spawn",
            ylabel = "overhead [ns] (N=$(k.nthreads))",
            xlabel = "trial",
        )
        scatter!(
            p,
            g.sync_overhead,
            markershape = :cross,
            markersize = 3,
            label = "sync",
        )
        return p
    end

    ymax = nextpow(10, max(maximum(rawdf.spawn_overhead), maximum(rawdf.sync_overhead)))
    ymin = prevpow(10, min(minimum(rawdf.spawn_overhead), minimum(rawdf.sync_overhead)))
    return plot(plts...; ylim = (ymin, ymax), yscale = :log10)
end

function analyze(args...; kwargs...)
    rawdf = load_rawdf(args...; kwargs...)
    return (
        rawdf = rawdf,
        hists = plot_hists(rawdf),
        seq = plot_seq(rawdf),
        #
    )
end

end # module TaskOverheads


text:
abstract type AbstractMechanism{T,Nn,Nb,Ne,Ni} end

mutable struct Mechanism{T,Nn,Nb,Ne,Ni} <: AbstractMechanism{T,Nn,Nb,Ne,Ni}
    origin::Origin{T}
    bodies::UnitDict{Base.OneTo{Int64},Body{T}}
    eqconstraints::UnitDict{UnitRange{Int64},<:EqualityConstraint{T}}
    ineqconstraints::UnitDict{UnitRange{Int64},<:InequalityConstraint{T}}

    graph::Graph{Nn}
    ldu::SparseLDU{T}

    # TODO remove once EqualityConstraint is homogenous
    normf::T
    normΔs::T

    Δt::T
    g::T

    α::T
    μ::T

    function Mechanism(origin::Origin{T},bodies::Vector{Body{T}},
            eqcs::Vector{<:EqualityConstraint{T}}, ineqcs::Vector{<:InequalityConstraint{T}};
            Δt::Real = .01, g::Real = -9.81, shapes::Vector{<:Shape{T}} = Shape{T}[]
        ) where T

        resetGlobalID()
        order = getGlobalOrder()

        for body in bodies
            if norm(body.m) == 0 || norm(body.J) == 0
                @info "Potentially bad inertial properties detected"
            end
        end

        Nb = length(bodies)
        Ne = length(eqcs)
        Ni = length(ineqcs)
        Nn = Nb + Ne
        # steps = Int64(ceil(tend / Δt))

        currentid = 1

        bdict = Dict{Int64,Int64}()
        for (ind, body) in enumerate(bodies)
            initknotpoints!(body.state, order)

            for eqc in eqcs
                eqc.parentid == body.id && (eqc.parentid = currentid)
                for (ind, bodyid) in enumerate(eqc.childids)
                    if bodyid == body.id
                        eqc.childids = setindex(eqc.childids, currentid, ind)
                        eqc.constraints[ind].childid = currentid
                    end
                end
            end

            for ineqc in ineqcs
                ineqc.parentid == body.id && (ineqc.parentid = currentid)
            end

            for shape in shapes
                for (i, id) in enumerate(shape.bodyids)
                    id == body.id && (shape.bodyids[i] = currentid)
                end
            end

            body.id = currentid
            currentid += 1

            bdict[body.id] = ind
        end

        eqdict = Dict{Int64,Int64}()
        for (ind, eqc) in enumerate(eqcs)
            eqc.id = currentid
            currentid += 1

            eqdict[eqc.id] = ind
        end

        ineqdict = Dict{Int64,Int64}()
        for (ind, ineqc) in enumerate(ineqcs)
            ineqc.id = currentid
            currentid += 1

            ineqdict[ineqc.id] = ind
        end

        normf = 0
        normΔs = 0

        graph = Graph(origin, bodies, eqcs, ineqcs)
        ldu = SparseLDU(graph, bodies, eqcs, ineqcs, bdict, eqdict, ineqdict)

        # storage = Storage{T}(steps, Nb, Ne)

        bodies = UnitDict(bodies)
        eqcs = UnitDict((eqcs[1].id):(eqcs[Ne].id), eqcs)
        if Ni > 0
            ineqcs = UnitDict((ineqcs[1].id):(ineqcs[Ni].id), ineqcs)
        else
            ineqcs = UnitDict(0:0, ineqcs)
        end

        α = 1
        μ = 1

        new{T,Nn,Nb,Ne,Ni}(origin, bodies, eqcs, ineqcs, graph, ldu, normf, normΔs, Δt, g, α, μ)
    end

    function Mechanism(origin::Origin{T},bodies::Vector{Body{T}},eqcs::Vector{<:EqualityConstraint{T}};
            Δt::Real = .01, g::Real = -9.81, shapes::Vector{<:Shape{T}} = Shape{T}[]
        ) where T

        ineqcs = InequalityConstraint{T}[]
        return Mechanism(origin, bodies, eqcs, ineqcs, Δt = Δt, g = g, shapes = shapes)
    end

    function Mechanism(origin::Origin{T},bodies::Vector{Body{T}},ineqcs::Vector{<:InequalityConstraint{T}};
            Δt::Real = .01, g::Real = -9.81, shapes::Vector{<:Shape{T}} = Shape{T}[]
        ) where T

        eqc = EqualityConstraint{T}[]
        for body in bodies
            push!(eqc, EqualityConstraint(Floating(origin, body)))
        end
        return Mechanism(origin, bodies, eqc, ineqcs, Δt = Δt, g = g, shapes = shapes)
    end

    function Mechanism(origin::Origin{T},bodies::Vector{Body{T}};
            Δt::Real = .01, g::Real = -9.81, shapes::Vector{<:Shape{T}} = Shape{T}[]
        ) where T

        eqc = EqualityConstraint{T}[]
        for body in bodies
            push!(eqc, EqualityConstraint(Floating(origin, body)))
        end
        return Mechanism(origin, bodies, eqc, Δt = Δt, g = g, shapes = shapes)
    end

    function Mechanism(filename::AbstractString; floating::Bool=false, type::Type{T} = Float64, Δt::Real = .01, g::Real = -9.81) where T
        origin, links, joints, shapes = parse_urdf(filename, floating, T)

        mechanism = Mechanism(origin, links, joints, shapes = shapes, Δt = Δt, g = g)

        graph = mechanism.graph
        xjointlist = Dict{Int64,SVector{3,T}}() # stores id, x in world frame
        qjointlist = Dict{Int64,UnitQuaternion{T}}() # stores id, q in world frame

        for id in graph.rdfslist # from root to leaves
            component = getcomponent(mechanism, id)
            if component isa Body
                body = component
                xbodylocal = body.state.xc
                qbodylocal = body.state.qc
                shape = getshape(shapes, id)

                preds = predecessors(graph, id);                    @assert length(preds) == 1
                parentid = preds[1]
                constraint = geteqconstraint(mechanism, parentid);  @assert length(constraint.constraints) == 2
                grandpreds = predecessors(graph, parentid);         @assert length(grandpreds) ∈ [0;1]

                if length(grandpreds) == 1 # predecessor is link
                    grandparentid = grandpreds[1]

                    parentbody = getbody(mechanism, grandparentid)
                    grandgrandpreds = predecessors(graph, grandparentid);               @assert length(grandgrandpreds) == 1
                    grandgrandparentid = grandgrandpreds[1]
                    parentconstraint = geteqconstraint(mechanism, grandgrandparentid);  @assert length(parentconstraint.constraints) == 2

                    xparentbody = parentbody.state.xc # in world frame
                    qparentbody = parentbody.state.qc # in world frame

                    xparentjoint = xjointlist[parentconstraint.id] # in world frame
                    qparentjoint = qjointlist[parentconstraint.id] # in world frame
                else # predecessor is origin
                    parentbody = origin

                    xparentbody = SA{T}[0; 0; 0]
                    qparentbody = one(UnitQuaternion{T})

                    xparentjoint = SA{T}[0; 0; 0]
                    qparentjoint = one(UnitQuaternion{T})
                end

                # urdf joint's x and q in parent's (parentbody) frame
                xjointlocal = vrotate(xparentjoint + vrotate(constraint.constraints[1].vertices[1], qparentjoint) - xparentbody, inv(qparentbody))
                qjointlocal = qparentbody \ qparentjoint * constraint.constraints[2].qoffset

                # store joint's x and q in world frame
                xjoint = xparentbody + vrotate(xjointlocal, qparentbody)
                qjoint = qparentbody * qjointlocal
                xjointlist[constraint.id] = xjoint
                qjointlist[constraint.id] = qjoint

                # difference to parent body (parentbody)
                qoffset = qjointlocal * qbodylocal

                # actual joint properties
                p1 = xjointlocal # in parent's (parentbody) frame
                p2 = vrotate(-xbodylocal, inv(qbodylocal)) # in body frame (xbodylocal and qbodylocal are both relative to the same (joint) frame -> rotationg by inv(body.q) gives body frame)
                constraint.constraints[1].vertices = (p1, p2)

                V3 = vrotate(constraint.constraints[2].V3', qjointlocal) # in parent's (parentbody) frame
                V12 = (svd(skew(V3)).Vt)[1:2,:]
                constraint.constraints[2].V3 = V3'
                constraint.constraints[2].V12 = V12
                constraint.constraints[2].qoffset = qoffset # in parent's (parentbody) frame

                # actual body properties
                setPosition!(body) # set everything to zero
                setPosition!(parentbody, body, p1 = p1, p2 = p2, Δq = qoffset)
                xbody = body.state.xc
                qbody = body.state.qc

                # shape relative
                if shape !== nothing
                    shape.xoffset = vrotate(xjoint + vrotate(shape.xoffset, qjoint) - xbody, inv(qbody))
                    shape.qoffset = qoffset \ qjointlocal * shape.qoffset
                end
            end
        end

        return mechanism, shapes
    end
end

mutable struct LinearMechanism{T,Nn,Nb,Ne,Ni} <: AbstractMechanism{T,Nn,Nb,Ne,Ni}
    ## Mechanism attributes
    origin::Origin{T}
    bodies::UnitDict{Base.OneTo{Int64},Body{T}}
    eqconstraints::UnitDict{UnitRange{Int64},<:EqualityConstraint{T}}
    ineqconstraints::UnitDict{UnitRange{Int64},<:InequalityConstraint{T}}

    graph::Graph{Nn}
    ldu::SparseLDU{T}

    # TODO remove once EqualityConstraint is homogenous
    normf::T
    normΔs::T

    Δt::T
    g::T

    α::T
    μ::T

    ## LinearMechanism attributes
    A::AbstractMatrix{T}
    Bu::AbstractMatrix{T}
    Bλ::AbstractMatrix{T}
    G::AbstractMatrix{T}

    xd::Vector{SVector{3,T}}
    vd::Vector{SVector{3,T}}
    qd::Vector{UnitQuaternion{T}}
    ωd::Vector{SVector{3,T}}
    # Fτd::Vector{SVector{3,T}}

    z::Vector{T}
    zsol::Vector{Vector{T}}
    Δz::Vector{T}
    λ::Vector{T}
    λsol::Vector{Vector{T}}
    Δλ::Vector{T}
    u::Vector{T}

    function LinearMechanism(mechanism::Mechanism{T,Nn,Nb,Ne,Ni}, xd, vd, qd, ωd, Fτd, eqcids) where {T,Nn,Nb,Ne,Ni}

        A, Bu, Bλ, G = linearsystem(mechanism, xd, vd, qd, ωd, Fτd, getid.(mechanism.bodies), eqcids)

        z = zeros(T,Nb*12)
        zsol = [zeros(T,Nb*12) for i=1:2]
        Δz = zeros(T,Nb*12)

        nc = 0
        for eqc in mechanism.eqconstraints
            nc += length(eqc)
        end
        λ = zeros(T,nc)
        λsol = [zeros(T,nc) for i=1:2]
        Δλ = zeros(T,nc)

        u = zeros(T,size(Bu)[2])

        new{T,Nn,Nb,Ne,Ni}([getfield(mechanism,i) for i=1:getfieldnumber(mechanism)]..., A, Bu, Bλ, G, xd, vd, qd, ωd, z, zsol, Δz, λ, λsol, Δλ, u)
    end

    function LinearMechanism(mechanism::Mechanism{T,Nn,Nb,Ne,Ni};
        xd = [mechanism.bodies[i].state.xc for i=1:Nb],
        vd = [mechanism.bodies[i].state.vc for i=1:Nb],
        qd = [mechanism.bodies[i].state.qc for i=1:Nb],
        ωd = [mechanism.bodies[i].state.ωc for i=1:Nb],
        eqcids = [],
        Fτd = []
    ) where {T,Nn,Nb,Ne,Ni}

        return LinearMechanism(mechanism, xd, vd, qd, ωd, Fτd, eqcids)
    end
end

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, mechanism::AbstractMechanism{T,Nn,Nb,Ne,0}) where {T,Nn,Nb,Ne}
    summary(io, mechanism)
    println(io, " with ", Nb, " bodies and ", Ne, " constraints")
    println(io, " Δt: "*string(mechanism.Δt))
    println(io, " g:  "*string(mechanism.g))
end

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, mechanism::AbstractMechanism{T,Nn,Nb,Ne,Ni}) where {T,Nn,Nb,Ne,Ni}
    summary(io, mechanism)
    println(io, " with ", Nb, " bodies, ", Ne, " equality constraints, and ", Ni, " inequality constraints")
    println(io, " Δt: "*string(mechanism.Δt))
    println(io, " g:  "*string(mechanism.g))
end


text:
using Zlib

# 1-byte unsigned integer -> Uint64
function size1unpack(bytes::Array{Uint8,1})
    size1unpack(bytes[1])
end

function size1unpack(byte::Uint8)
    uint64(byte)
end

# 2-bytes unsigned integer in big endian format -> Uint64
function size2unpack(bytes::Array{Uint8,1})
    uint64(reinterpret(Uint16, reverse(bytes))[1])
end

# 4-bytes unsigned integer in big endian format -> Uint64
function size4unpack(bytes::Array{Uint8,1})
    uint64(reinterpret(Uint32, reverse(bytes))[1])
end

function int1unpack(bytes::Array{Uint8,1})
    int1unpack(bytes[1])
end

function int1unpack(byte::Uint8)
    int(byte)
end

function int4unpack(bytes)
    int(reinterpret(Int32, reverse(bytes))[1])
end

function floatunpack(bytes)
    reinterpret(Float64, reverse(bytes))[1]
end

function charintpack(value::Integer, size::Int)
    bytes = zeros(Uint8, size)
    for i=1:size
        bytes[i] = uint8(value)
        value = value >>> 8
        end
    reverse(bytes)
end

function charintpack(value::Integer)
    bytes = []
    while value != 0
        bytes = vcat(bytes, uint8(value & 0xff))
        value = value >>> 8
    end
    uint8(bytes)
end

function charint4pack(integer::Integer)
    charintpack(integer, 4)
end

function charint2pack(integer::Integer)
    charintpack(integer, 2)
end

function charsignedint4pack(integer::Integer)
    charintpack(integer, 4)
end

function lencheck(bytes::Array{Uint8,1}, limit::Int)
    len = uint64(length(bytes))
    lencheck(len, len < limit, bytes)
end

function lencheck(len::Uint64, limit::Uint64, bytes::Array{Uint8,1})
    lencheck(limit, len < limit, bytes)
end

function lencheck(len::Uint64, pred::Bool, bytes::Array{Uint8,1})
    if pred
        throw(IncompleteData(bytes))
    end
    len
end

function compressterm(encodedterm, compression::Bool)
    compressterm(encodedterm, 6)
end

function compressterm(encodedterm, compression::Int)
    if compression < 0 || compression > 9
        throw(InvalidCompressionLevel(compression))
    end
    comp = Zlib.compress(encodedterm, compression)
    len = length(encodedterm)
    # XXX add check here for too small of length
    if length(comp) + 5 <= len
        vcat(int4pack(len), comp)
    end
end

function decompressterm(bytes::Array{Uint8,1})
    if length(bytes) < 16
        throw(IncompleteData(bytes))
    end
    sentlen = size4unpack(bytes[3:6])
    term = Zlib.decompress(bytes[7:end])
    actuallen = length(term)
    if actuallen != sentlen
        msg = "Header declared $sentlen bytes but got $actuallen bytes."
        throw(InvalidCompressedTag(msg))
    end
    return term
end


text:
@testset "Dijkstra" begin
    g4 = PathDiGraph(5)
    d1 = float([ 0 1 2 3 4; 5 0 6 7 8; 9 10 0 11 12; 13 14 15 0 16; 17 18 19 20 0])
    d2 = sparse(float([ 0 1 2 3 4; 5 0 6 7 8; 9 10 0 11 12; 13 14 15 0 16; 17 18 19 20 0]))

    for g in testdigraphs(g4)
      y = @inferred(dijkstra_shortest_paths(g, 2, d1))
      z = @inferred(dijkstra_shortest_paths(g, 2, d2))

      @test y.parents == z.parents == [0, 0, 2, 3, 4]
      @test y.dists == z.dists == [Inf, 0, 6, 17, 33]

      y = @inferred(dijkstra_shortest_paths(g, 2, d1; allpaths=true))
      z = @inferred(dijkstra_shortest_paths(g, 2, d2; allpaths=true))
      @test z.predecessors[3] == y.predecessors[3] == [2]

      @test @inferred(enumerate_paths(z)) == enumerate_paths(y)
      @test @inferred(enumerate_paths(z))[4] ==
          enumerate_paths(z,4) ==
          enumerate_paths(y,4) == [2,3,4]
    end

    gx = PathGraph(5)
    add_edge!(gx,2,4)
    d = ones(Int, 5,5)
    d[2,3] = 100
    for g in testgraphs(gx)
      z = @inferred(dijkstra_shortest_paths(g,1,d))
      @test z.dists == [0, 1, 3, 2, 3]
      @test z.parents == [0, 1, 4, 2, 4]
    end

    # small function to reconstruct the shortest path; I copied it from somewhere, can't find the original source to give the credits
    # @Beatzekatze on github
    spath(target, dijkstraStruct, source) = target == source ? target : [spath(dijkstraStruct.parents[target], dijkstraStruct, source) target]
    function spaths(ds, targets, source)
        shortest_paths = []
        for i in targets
            push!(shortest_paths,spath(i,ds,source))
        end
        return shortest_paths
    end

    G = LightGraphs.Graph()
    add_vertices!(G,4)
    add_edge!(G,2,1)
    add_edge!(G,2,3)
    add_edge!(G,1,4)
    add_edge!(G,3,4)
    add_edge!(G,2,2)
    w = [0. 3. 0. 1.;
        3. 0. 2. 0.;
        0. 2. 0. 3.;
        1. 0. 3. 0.]

    for g in testgraphs(G)
      ds = @inferred(dijkstra_shortest_paths(g,2,w))
      # this loop reconstructs the shortest path for vertices 1, 3 and 4
      @test spaths(ds, [1,3,4], 2) == Array[[2 1],
                                            [2 3],
                                            [2 1 4]]

    # here a selflink at source is introduced; it should not change the shortest paths
      w[2,2] = 10.0
      ds = @inferred(dijkstra_shortest_paths(g,2,w))
      shortest_paths = []
      # this loop reconstructs the shortest path for vertices 1, 3 and 4
      @test spaths(ds, [1,3,4], 2) == Array[[2 1],
                                            [2 3],
                                            [2 1 4]]
    end
end

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
using Distributions, SearchDiscoveryReplication, NLopt, GalacticOptim
using BSON
using PDMats # required for BSON to load data again
##########################################################################
## Loop simulation over different parameter values
function loopSimulation(βpVec,βoVec,csVec,cdVec)
	# Start at last index from filenames
	fnames = readdir("./gen")
	# only keep relevant filenames
	fnames = fnames[(occursin.("results_spec",fnames)) .&
				(occursin.("~",fnames) .== false) ]
	i = maximum([convert(Int,parse(Float64,fnames[i][13:14])) for i in 1:length(fnames)])
	i += 1

	# Loop over inputs
	for βp0 in βpVec, βo0 in βoVec, cs0 in csVec, cd0 in cdVec
		β0 = [3.5,βp0,βo0]

		d, wSD, outFI, outRS, outWf, outW = runSimulation(β0, cs0, cd0, dChars, outDummy, dV,
															nCons, nProdTot, nProdCons,nA0;
															seed =26 +i,
															options = estimOptions(algo = Opt(:LN_BOBYQA,1)))

		BSON.@save "gen/results_spec$i.bson" wSD outFI outRS outWf outW β0 cs0 cd0 dChars outDummy dV nA0
		println(i)
		i += 1
	end
end

cdVec = 0.01:0.05:0.25
csVec = 0.01:0.05:0.15
βpVec = [-1.0]
βoVec = 1.0:5.0
dChars = MvNormal([2., 3.5], [3.0, 1.0])
outDummy = true
dV = Normal(0, 0)
nCons = 1000
nProdTot = 100000
nProdCons = 30
nA0 = 1

loopSimulation(βpVec,βoVec,csVec,cdVec)

##########################################################################
## Summarize results (get same numbers as in tables)

function evaluateResultsLoop()
	fnames = readdir("./gen")
	fnames = fnames[(occursin.("results_spec",fnames)) .& (occursin.("~",fnames) .== false) ] # exclude unnecessary filenames

	# pre-allocate
	resSD = (zeros(7,length(fnames)),zeros(6,length(fnames)))
	resW = deepcopy(resSD)
	resWf = deepcopy(resSD)
	resRS = deepcopy(resSD)
	resFI = deepcopy(resSD)

	# Loop through results files, and fill in values for each of the models
	for i in eachindex(fnames)
		println("Checking spec $(fnames[i])")
		estimOut, welfareChangesOut, demandOut, dStats = gatherResults("gen/$(fnames[i])")
		for (r,m) in zip([resSD,resWf,resW,resRS,resFI],1:5)
			fillResults!(r,i,m,estimOut,welfareChangesOut,demandOut,dStats)
		end
	end
	return resSD,resWf,resW,resRS,resFI
end

function fillResults!(r,i,m,estimOut,welfareChangesOut,demandOut,dStats)

	# Same numbers as TABLE 2
	r[1][1,i] = dStats[m,1] # n Searches
	r[1][2,i] = (1 - dStats[m,2]) * 100 # Share purchases
	r[1][3,i] = estimOut[2,m] # price coefficient
	r[1][4:end,i] = estimOut[[1,3,4,5],m] / abs(estimOut[2,m])

	# Same numbers as TABLE 3
	r[2][1:3,i] = hcat(welfareChangesOut[1][m,2],demandOut[1][m,1],demandOut[2][m,1])
	r[2][4:6,i] = hcat(welfareChangesOut[1][m,3],demandOut[1][m,2],demandOut[2][m,2])
end

# resSD,resWf,resW,resRS,resFI = evaluateResultsLoop()
res = evaluateResultsLoop()
##########################################################################
## Plot results (for easy comparison)
using Plots

ns = length(res[1][1][1,:])
nm = length(res)
l = ["SD" "DS1" "DS2" "RS" "FI"]

# Table 2 numbers
for k in 1:7
	local d = zeros(ns,nm)
	for (i,r) in enumerate(res)
		d[:,i] = r[1][k,:]
	end
	p = plot(d, label = l, title = "Table 2: Column $k")
	png(p,"gen/fig$(k)_for_t2.png")
end

# Table 3 numbers
for k in 1:6
	local d = zeros(ns,nm)
	for (i,r) in enumerate(res)
		d[:,i] = r[2][k,:]
	end
	p = plot(d, label = l, title = "Table 3: Column $k")
	png(p,"gen/fig$(k)_for_t3.png")
end


text:
"""
    AbstractIndices{I} <: AbstractDictionary{I, I}

Abstract type for the unique keys of an `AbstractDictionary`. It is itself an `AbstractDictionary` for
which `getindex` is idempotent, such that `indices[i] = i`. (This is a generalization of
`Base.Slice`).

At minimum, an `AbstractIndices{I}` must implement:

 * The `iterate` protocol, returning unique values of type `I`.
 * `in`, such that `in(i, indices)` implies there is an element of `indices` which `isequal` to `i`.
 * Either `length`, or override `IteratorSize` to `SizeUnknown`.

While an `AbstractIndices` object is a dictionary, the value corresponding to each index is
fixed, so `issettable(::AbstractIndices) = false` and `setindex!` is never defined.

If arbitrary indices can be added or removed from the set, implement:

* `isinsertable(::AbstractIndices)` (returning `true`)
* `insert!(indices::AbstractIndices{I}, ::I}` (returning `indices`)
* `delete!(indices::AbstractIndices{I}, ::I}` (returning `indices`)
"""
abstract type AbstractIndices{I} <: AbstractDictionary{I, I}; end

@inline function Base.getindex(indices::AbstractIndices{I}, i::I) where {I}
    @boundscheck checkindex(indices, i)
    return i
end

function Base.setindex!(i::AbstractIndices{I}, ::I, ::I) where {I}
    error("Indices are not settable: $(typeof(i))")
end

function Base.isassigned(indices::AbstractIndices{I}, i::I) where {I}
    return i in indices
end

Base.keys(i::AbstractIndices) = i

@inline function Base.iterate(inds::AbstractIndices, s...)
    if istokenizable(inds)
        tmp = iteratetoken(keys(inds), s...)
        tmp === nothing && return nothing
        (t, s2) = tmp
        return (gettokenvalue(inds, t), s2)
    else
        error("All AbstractIndices must define `iterate`: $(typeof(inds))")
    end
end

function Base.in(i::I, indices::AbstractIndices{I}) where I
    if !istokenizable(indices)
        # A fallback definition based on iteration would be rediculously slow. There
        # shouldn't be many `AbstractIndex` types that rely on iteration for this.
        error("All AbstractIndices must define `in`: $(typeof(indices))")
    end

    return gettoken(indices, i)[1]
end

function Base.in(i, indices::AbstractIndices{I}) where I
    return convert(I, i) in indices
end

# Match the default setting from Base - the majority of containers will know their size
Base.IteratorSize(indices::AbstractIndices) = Base.HasLength()

function Base.length(indices::AbstractIndices)
    if Base.IteratorSize(indices) isa Base.SizeUnknown
        out = 0
        for _ in tokens(indices)
            out += 1
        end
        return out
    end

    error("All AbstractIndices must define `length` or else have `IteratorSize` of `SizeUnknown`: $(typeof(indices))")
end

Base.unique(i::AbstractIndices) = i

struct IndexError <: Exception
	msg::String
end

function Base.checkindex(indices::AbstractIndices{I}, i::I) where {I}
	if i ∉ indices
		throw(IndexError("Index $i not found in indices $indices"))
	end
end
Base.checkindex(indices::AbstractIndices{I}, i) where {I} = checkindex(indices, convert(I, i))

function checkindices(indices::AbstractIndices, inds)
    if !(inds ⊆ indices)
        throw(IndexError("Indices $inds are not a subset of $indices"))
    end
end

# Indices are isequal if they iterate in the same order
function Base.isequal(i1::AbstractIndices, i2::AbstractIndices)
    if sharetokens(i1, i2)
        return true
    end

    if length(i1) != length(i2)
        return false
    end

    for (j1, j2) in zip(i1, i2)
        if !isequal(j1, j2)
            return false
        end
    end

    return true
end

# Use `issetequal` semantics
function Base.:(==)(i1::AbstractIndices, i2::AbstractIndices)
    if sharetokens(i1, i2)
        return true
    end

    if length(i1) != length(i2)
        return false
    end

    for i in i1
        if !(i in i2)
            return false
        end
    end

    return true
end

# Lexical ordering based on iteration
function Base.isless(inds1::AbstractIndices, inds2::AbstractIndices)
    if sharetokens(inds1, inds2)
        return false # they are isequal
    end

    # We want to iterate... until one is longer than the other (no `zip`)
    tmp1 = iterate(inds1)
    tmp2 = iterate(inds2)
    while tmp1 !== nothing
        if tmp2 === nothing
            return false # shorter collections are isless in lexical ordering
        end
        (i1, s1) = tmp1
        (i2, s2) = tmp2
        !isequal(i1, i2)
        c = cmp(i1, i2)

        if c == -1
            return true
        elseif c == 1
            return false
        end

        tmp1 = iterate(inds1, s1)
        tmp2 = iterate(inds2, s2)
    end
    return tmp2 !== nothing # shorter collections are isless in lexical ordering, otherwise isequal
end

function Base.cmp(inds1::AbstractIndices, inds2::AbstractIndices)
    if sharetokens(inds1, inds2)
        return 0 # they are isequal
    end

    # We want to iterate... until one is longer than the other (no `zip`)
    tmp1 = iterate(inds1)
    tmp2 = iterate(inds2)
    while tmp1 !== nothing
        if tmp2 === nothing
            return 1 # shorter collections are isless in lexical ordering
        end
        (i1, s1) = tmp1
        (i2, s2) = tmp2
        !isequal(i1, i2)
        c = cmp(i1, i2)

        if c == -1
            return -1
        elseif c == 1
            return 1
        end

        tmp1 = iterate(inds1, s1)
        tmp2 = iterate(inds2, s2)
    end
    if tmp2 === nothing
        return 0
    else
        return -1 # shorter collections are isless in lexical ordering
    end
end

## Hashing - matches the dictionary case (matching keys and matching values)
function Base.hash(inds::AbstractIndices, h::UInt)
    h1 = h
    for i in inds
        h1 = hash(i, h1)
    end

    return hash(hash(UInt === UInt64 ? 0x8955a87bc313a509 : 0xa9cff5d1, h1), h1)
end


text:
using Turing, Random, MCMCChains
using CSV, DataFrames, StatsBase
using ParetoSmooth
using StatisticalRethinking

#Random.seed!(129111)

df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame)
df.D = zscore(df.Divorce)
df.M = zscore(df.Marriage)
df.A = zscore(df.MedianAgeMarriage)
data = (D=df.D, A=df.A)

function lin(a, b, c, x...)
    result = @. a + b * c
    for i in 1:2:length(x)
        @. result += x[i] * x[i+1]
    end
    return result
end

@model function m5_1t(A, D)
    a ~ Normal(0, 0.2)
    bA ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = lin(a, A[i], bA)
        D[i] ~ Normal(μ, σ)
    end
end

chn5_1t = sample(m5_1t(df.A, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

@model function m5_2t(M, D)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = lin(a, M[i], bM)
        D[i] ~ Normal(μ, σ)
    end
end

chn5_2t = sample(m5_2t(df.M, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

@model function m5_3t(A, M, D)
    a ~ Normal(0, 0.2)
    bA ~ Normal(0, 0.5)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = a + M[i] * bM + A[i] * bA
        D[i] ~ Normal(μ, σ)
    end
end

chn5_3t = sample(m5_3t(df.A, df.M, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

pw_lls5_1t = pointwise_log_likelihoods(m5_1t(df.A, df.D), chn5_1t)
pw_lls5_2t = pointwise_log_likelihoods(m5_2t(df.M, df.D), chn5_2t)
pw_lls5_3t = pointwise_log_likelihoods(m5_3t(df.A, df.M, df.D), chn5_3t)

loo_comparison = ParetoSmooth.loo_compare([pw_lls5_1t, pw_lls5_2t, pw_lls5_3t];
    model_names=[:m5_1t, :m5_2t, :m5_3t])

for (i, psis) in enumerate(loo_comparison.psis)
    psis |> display
    pk_plot(psis.pointwise(:pareto_k))
    savefig(joinpath(@__DIR__, "m5.$(i)t.png"))
end
println()
Particles(chn5_1t[:,[:a, :bA, :σ],:]) |> display
Particles(chn5_2t[:,[:a, :bM, :σ],:]) |> display
Particles(chn5_3t[:,[:a, :bA, :bM, :σ],:]) |> display
loo_comparison |> display


text:
function build_group(neurons::Vector{N}; name=:group) where {N}
    eqs = Array{Equation,1}(undef, 0)
    ODESystem(eqs, t, [], []; name=name, systems=neurons)
end

function add_connection(group, pre::ODESystem, post::ODESystem, syn::EmptyConnection; kwargs...)
    return group
end

function add_connection(group, pre::ODESystem, post::ODESystem, syn::Synapse; name=ModelingToolkit.getname(group), i=1)
    prename, postname = ModelingToolkit.getname.([pre, post])
    synapse_sys = syn(; name=Symbol(prename, :to, postname, get_name(syn)))

    oldeqs = ModelingToolkit.get_eqs(group)
    neweqs = [
        get_pre(syn, pre) ~ my_pre(syn, synapse_sys),
        get_post(syn, post) ~ my_post(syn, synapse_sys),
        syn_current(syn, synapse_sys) ~ getproperty(post, Symbol(post_connector(syn), i))
    ]

    eqs = cat(oldeqs, neweqs, dims=1)
    systems = cat(ModelingToolkit.get_systems(group), synapse_sys, dims=1)

    return connected = ODESystem(eqs, t, [], [];
        name=name,
        systems=systems)
end

function in_degree(edge_fn, num_nodes, j)
    sum(1:num_nodes) do i
        length(edge_fn(i, j)) - (typeof(edge_fn(i, j)) == Tuple{EmptyConnection})
    end
end

function build_network(nodefn, edgefn, which_nodes::AbstractRange; name=:network)
    eqs = Array{Equation,1}(undef, 0)
    num_nodes = length(which_nodes)
    neurons = [nodefn(j)(in_degree(edgefn, num_nodes, j)) for j in 1:num_nodes]
    network = ODESystem(eqs, t, [], []; name=name, systems=neurons)
    fill_counter = ones(Int64, num_nodes)
    for i in which_nodes
        for j in which_nodes
            for syn in edgefn(i, j)
                network = add_connection(network, neurons[i], neurons[j], syn; i=fill_counter[j])
                !(typeof(syn) == EmptyConnection) && (fill_counter[j] += 1)
            end
        end
    end
    return network |> structural_simplify
end

function add_all_connections(neurons, connections)
    grp = build_group(neurons; name=:STG)
    num_neurons = length(neurons)
    num_incoming_connections = [sum(x -> x.post_n == el, connections) for el in neurons]
    fill_counter = ones(Int64, num_neurons)
    for (el, j) in zip(neurons, 1:num_neurons)
        for dir_syn in filter(x -> x.post_n == el, connections)
            grp = add_connection(grp, dir_syn.pre_n, dir_syn.post_n, dir_syn.syn; i=fill_counter[j])
            fill_counter[j] += 1
        end
    end
    return structural_simplify(grp)
end


text:
function parse_pseudopotential_line(::Type{T}, pp_line::AbstractString) where T
    data = split(pp_line, ";")
    num_terms = parse(Int, data[1])
    length(data) == num_terms+2 || throw(ArgumentError("Unrecognizable potential string $(pp_line)"))
    n = Vector{Int}(undef, num_terms)
    β = Vector{T}(undef, num_terms)
    B = Vector{T}(undef, num_terms)
    map(1:num_terms) do i
        term = split(strip(data[i+1]),",")
        n[i] = parse(Int, term[1])
        β[i] = parse(T, term[2])
        B[i] = parse(T, term[3])
    end
    GaussianExpansion(n, β, B)
end

function parse_header(pp_string::String)
    lines = split(pp_string, "\n")
    gst_config = parse(Configuration{Orbital},
                       lstrip(lines[1], ['!', ' ']))
    i = findfirst(l -> first(l) != '!', lines)
    i === nothing && error("Could not find start of data")
    header = split(split(lines[i], ";")[1], ",")
    kind = header[1]

    element = String(header[2])
    Z = element_number(Symbol(element))
    nelec = num_electrons(gst_config)
    Z == nelec ||
        throw(ArgumentError("Atom number Z = $Z does not match number of electrons of ground state configuration $(gst_config) => $(nelec) electrons"))

    Q̃ = parse(Int, header[3])
    ncore = num_electrons(core(gst_config))
    Q̃ == ncore ||
        throw(ArgumentError("Charge modelled by pseudopotential, Q̃ = $(Q̃), does not match number of core electrons in ground state configuration: $(core(gst_config)) => $(ncore) electrons"))

    Q = Z-Q̃ # Long-range charge of potential

    element, gst_config, Q, header, lines[i+1:end]
end

parse_reference(lines) =
    map(lines) do line
        lstrip(line, ['!', '[', ']', ' ', '0':'9'...])
    end |> l -> join(l, "\n")

function parse_pseudopotential(::Type{T}, pp_string::String, relativistic::Bool) where T
    element, gst_config, Q, header, lines = parse_header(pp_string)

    ℓmax = parse(Int, header[4])
    ℓmax′ = parse(Int, header[5])
    ii = 0
    Vℓ = map(vcat(ℓmax,0:ℓmax-1)) do ℓ
        ℓ => parse_pseudopotential_line(T, lines[ii+=1])
    end |> v -> last.(sort(v, by=first))
    Vℓ′ = map(1:ℓmax′) do ℓ′
        parse_pseudopotential_line(T, lines[ii+=1])
    end |> Vector{GaussianExpansion{T}}
    reference = parse_reference(lines[ii+2:end])
    PseudoPotential{T,relativistic}(element, gst_config, Q, Vℓ, Vℓ′, reference)
end

function parse_relativistic_pseudopotential(::Type{T}, pp_string::String) where T
    element, gst_config, Q, header, lines = parse_header(pp_string)

    nκ₋ = parse(Int, header[4])
    nκ₊ = parse(Int, header[5])
    ii = 0
    V₋ = map(1:nκ₋) do ℓ
        parse_pseudopotential_line(T, lines[ii+=1])
    end
    V₊ = map(1:nκ₊) do ℓ′
        parse_pseudopotential_line(T, lines[ii+=1])
    end # |> Vector{GaussianExpansion{T}}
    reference = parse_reference(lines[ii+2:end])
    RelativisticPseudoPotential{T}(element, gst_config, Q, V₋, V₊, reference)
end

macro PP_str(pp_string, suffix="")
    if suffix == "" || suffix == "sso"
        parse_pseudopotential(Float64, pp_string, suffix=="sso")
    elseif suffix == "r"
        parse_relativistic_pseudopotential(Float64, pp_string)
    else
        throw(ArgumentError("Unkown pseudopotential suffix $(suffix)"))
    end
end


text:
using OrdinaryDiffEq, Test, ForwardDiff

# Test problem
# u' = p(t)u; p(t) = p if t <= tstop, 0 otherwise
# Target: du(1.0)/dtstop = exp(tstop)u0

u0 = 1.0
tspan = (0.0, 1.0)
p = 1.0

function stopping_cb(tstop)
    condition = (u,t,integrator) -> t == tstop
    affect! = integrator -> (println("Stopped!"); integrator.p = zero(integrator.p))
    DiscreteCallback(condition, affect!)
end

function test_fun(tstop)
    DualT = typeof(tstop)
    prob = ODEProblem((u,p,t) -> p*u, DualT(u0), DualT.(tspan), DualT(p))
    sol = solve(prob, Tsit5(), callback=stopping_cb(tstop), tstops=[tstop])
    sol(1.0)
end

@test ForwardDiff.derivative(test_fun, 0.5) ≈ exp(0.5)*u0 # Analytical solution: exp(tstop)*u0

function test_fun(tstop)
    DualT = typeof(tstop)
    prob = ODEProblem((u,p,t) -> p*u, DualT(u0), DualT.(tspan), DualT(p))
    sol = solve(prob, Tsit5(), callback=stopping_cb(tstop), tstops=[tstop], adaptive=false, dt=0.01)
    sol(1.0)
end

@test ForwardDiff.derivative(test_fun, 0.5) ≈ exp(0.5)*u0 # Analytical solution: exp(tstop)*u0


text:
@system NodalUnit(Organ) begin
    rank ~ ::int(override) # preserve
    leaf(context, phenology, rank) ~ ::Leaf
    sheath(context, phenology, rank) ~ ::Sheath

    mass(l=leaf.mass, s=sheath.mass) => (l + s) ~ track(u"g")
end


text:
#using Revise
using Distributed
@everywhere using Random
@everywhere using Statistics
using Plots
using StatsPlots
using SharedArrays

################################################################################
############################# PREPARATION ######################################

#%% Input Data #################################################################
# Individual observations in sequence:
@everywhere y = [607, 583, 521, 494, 369, 782, 570, 678, 467, 620, 425, 395, 346, 361, 310,
300, 382, 294, 315, 323, 421, 339, 398, 328, 335, 291, 329, 310, 294, 321, 286,
349, 279, 268, 293, 310, 259, 241, 243, 272, 247, 275, 220, 245, 268, 357, 273,
301, 322, 276, 401, 368, 149, 507, 411, 362, 358, 355, 362, 324, 332, 268, 259,
274, 248, 254, 242, 286, 276, 237, 259, 251, 239, 247, 260, 237, 206, 242, 361,
267, 245, 331, 357, 284, 263, 244, 317, 225, 254, 253, 251, 314, 239, 248, 250,
200, 256, 233, 427, 391, 331, 395, 337, 392, 352, 381, 330, 368, 381, 316, 335,
316, 302, 375, 361, 330, 351, 186, 221, 278, 244, 218, 126, 269, 238, 194, 384,
154, 555, 387, 317, 365, 357, 390, 320, 316, 297, 354, 266, 279, 327, 285, 258,
267, 226, 237, 264, 510, 490, 458, 425, 522, 927, 555, 550, 516, 548, 560, 545,
633, 496, 498, 223, 222, 309, 244, 207, 258, 255, 281, 258, 226, 257, 263, 266,
238, 249, 340, 247, 216, 241, 239, 226, 273, 235, 251, 290, 473, 416, 451, 475,
406, 349, 401, 334, 446, 401, 252, 266, 210, 228, 250, 265, 236, 289, 244, 327,
274, 223, 327, 307, 338, 345, 381, 369, 445, 296, 303, 326, 321, 309, 307, 319,
288, 299, 284, 278, 310, 282, 275, 372, 295, 306, 303, 285, 316, 294, 284, 324,
264, 278, 369, 254, 306, 237, 439, 287, 285, 261, 299, 311, 265, 292, 282, 271,
268, 270, 259, 269, 249, 261, 425, 291, 291, 441, 222, 347, 244, 232, 272, 264,
190, 219, 317, 232, 256, 185, 210, 213, 202, 226, 250, 238, 252, 233, 221, 220,
287, 267, 264, 273, 304, 294, 236, 200, 219, 276, 287, 365, 438, 420, 396, 359,
405, 397, 383, 360, 387, 429, 358, 459, 371, 368, 452, 358, 371];

# Indicator of the individual behind each observation
@everywhere ind = [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 8,
8, 8, 8, 8, 9, 9, 9, 9, 9, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
11, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12,
12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 14, 14, 14, 14,
14, 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16,
16, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19, 20, 20, 20,
20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21,
21, 21, 21, 21, 21, 21, 21, 21, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22,
22, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 24, 24, 24,
24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 25, 25, 25, 25, 25, 25, 25,
25, 25, 25, 25, 25, 25, 26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 28, 28, 28,
28, 28, 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30,
30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 33, 34, 34, 34, 34, 34,
34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34];

# Indicator for each individual j whether he/she is a child or not:
@everywhere child_j = [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1,
0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0];

# Indicator for each observations i whether it comes from a child or not:
@everywhere child_i = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,
1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];

#####
#%% Functions ##################################################################
# @everywhere to make them known to all workers
@everywhere function mySliceSampler(pdf_log, x_0, w, m, N = 2000, burnIn = 1000)
    # REQUIRES MODULES:
    # Random
    #
    # INPUT:
    # pdf_log --> log-pdf of the target distribution (a function)
    # x_init --> inital vector of x (dimension 1xD), where D dimension of pdf
    # w --> typical slice size, for a window (dimension 1xD)
    # m --> integer limiting the slice size to m*w
    # N --> number of sample points
    # burnIn --> number of burn-in samples (optional, results will be discarded)
    #
    # OUTPUT:
    # x_s --> sampled x values (NxD), only keep values after full permutations along dimensions
    # pdflog_x_s --> log-pdf of sampled values

    D = length(x_0)  # the dimension of the distribution (nr of coordinates)
    x_s = Array{Float64}(undef,N,D)   # Samples (NxD), preallocate array, with undefined values
    pdflog_x_s = Array{Float64}(undef,N)   # (unnormalised) log-prob of sampled points (Nx1)
                                           # for reducing evaluations
    pdflog_x_1 = 0.0;

    # to be able to input scalars:
    #if typeof(x_0) == Float64
    #    x_0 = [x_0]
    #end

    for ii in 1:(N+burnIn)
        # Update the new x progressively for each dimension:
        L   = 1*x_0  # R and L need to be of same dimension, because we apply the pdf upon those!
        R   = 1*x_0
        x_1 = 1*x_0  # "1*" needed, otherwise x_0 suddendly takes mysterious values
        for d in randperm(D)
            ### 1) Make the Slice
            # random "vertical" position
            if (ii == 1)
                #vertical = rand() * pdf_log(x_0) # in normal space
                z = pdf_log(x_0) + log(rand()) # in log-space
            else
                # saving evaluations, pdflog_x_1 was updated for current permutation in the end
                z = pdflog_x_1 + log(rand())
            end

            ### 2) Make the window, "Stepping-out" alogrithm
            # Randomly place a window, containing x_0:
            L[d] = x_0[d] - rand() * w[d]
            R[d] = L[d] + w[d]

            # Randomly share the max window size among left/right:
            J = floor(m*rand());
            K = (m-1) - J;

            # Extend window to the left, until outside slice or allowance seizes:
            while ((J > 0) && (z < pdf_log(L)))
                #println("Lefting")
                L[d] -= w[d]
                J -= 1
            end

            # Extend window to the right, until outside slice or allowance seizes:
            while ((K > 0) && (z < pdf_log(R)))
                #println("Righting")
                R[d] += w[d]
                K -= 1
            end

            ### 3) Sample from window
            # finding an allowable point:
            while true
                x_1[d] = L[d] + rand() * (R[d] - L[d])

                # this + breaking out of loop reduces the amount of log-pdf evaluations:
                pdflog_x_1 = pdf_log(x_1)

                if (pdflog_x_1 >= z)     # new value found
                    x_0[d] = x_1[d]     # update value for this dimension
                    break
                # Value was not within slice, shrink the interval:
                elseif (x_1[d] < x_0[d])
                    L[d] = x_1[d]
                elseif (x_1[d] > x_0[d])
                    R[d] = x_1[d]
                else
                    throw(ErrorException("Error during shrinking the interval"))
                end
            end

        end # end of permutating through dimensions

        ### 4) Update the chain:
        # Just overwrite the burnIn by using modular arithmetic:
        i = 1 + (ii-1)%N    # gives 1,2,3,...N,1,2,3,...,(N+burnIn)
        x_s[i,:] = x_1    # sample point, after one round of permutations
        pdflog_x_s[i] = pdflog_x_1 # now contains log_pdf of all new coordinates of new point, since has been updated at each permutation
        # note that x_0 has been updated during the permutations

    end

    return x_s, pdflog_x_s
end

## Jespers HDI:
@everywhere function hdi(theta_samp,alpha=0.05)
    cred_mass = 1.0-alpha
    ci = zeros(2)
    if length(size(theta_samp))>1   # multidimensional theta
        K,N = size(theta_samp)
        cis = zeros(2,K)
        for k in 1:K    # for each dimension

            ts = theta_samp[k,:]
            sind = sortperm(ts)     # get series of indices that sorts theta values
            sts = ts[sind]  # sorted theta values

            # Shifting a 95% index block from far left to far right and use the
            # samllest interval that can be found during the shift
            N = length(sind)    # number of sample points
            length_ci = Inf     # start with inf HDI interval and shrink
            for i in 1:Int(floor(N*alpha))  # iterate over the 5% lowest values (low to high)
                i2 = Int(floor(N*cred_mass)+i)  # the 5% highest values (low to high), one 95% block away on the other side
                prop_ci = [sts[i],sts[i2]]      # proposed interval
                length_prop_ci = prop_ci[2]-prop_ci[1]
                if length_prop_ci < length_ci
                    ci = prop_ci
                    length_ci = ci[2]-ci[1]
                end
            end
            cis[:,k] = ci

        end
        return cis
    else
        N = length(theta_samp)

        ts = theta_samp
        sind = sortperm(ts)
        sts = ts[sind]

        N = length(sind)
        length_ci = Inf
        for i in 1:Int(floor(N*alpha))
            i2 = Int(floor(N*cred_mass)+i)
            prop_ci = [sts[i],sts[i2]]
            length_prop_ci = prop_ci[2]-prop_ci[1]
            if length_prop_ci < length_ci
                ci = prop_ci
                length_ci = ci[2]-ci[1]
            end
        end
        return ci
    end
end

# A recursive application of the HDI algorithm to find the mode:
@everywhere function find_mode(sample_pts;runs=2)
    # re-apply the hdi algorithm, until there is only two points or so left
    # so the ci interval becomes smaller and smaller
    # in the end, return an average over the interval
    ci = zeros(2)
    i_left = 0
    i_right = 0

    # Make 50% credibility intervals
    alpha = 0.5;
    cred_mass = 1.0 - alpha;

    # sort the pts:
    sind = sortperm(sample_pts)
    sts = sample_pts[sind]
    N = length(sind)

    if N > runs
        # HDI algorithm, but return indices too:
        length_ci = Inf
        for i in 1:Int(floor(N*alpha))
            i2 = Int(floor(N*cred_mass)+i)
            prop_ci = [sts[i], sts[i2]]
            length_prop_ci = prop_ci[2] - prop_ci[1]
            if length_prop_ci < length_ci
                i_left  = i
                i_right = i2
                ci = prop_ci
                length_ci = ci[2] - ci[1]
            end
        end

        # Recursion:
        ci = find_mode(sts[i_left:i_right],runs=runs)

    else
        ci = [sts[1],sts[end]]
    end

    return ci
end

# A function to make the plots with HDIs
function makeDistributionPlot(X, color="blue"; ann=true, offset=0.0, scale=1.0)
    #%% Mode, Mean, HDIs
    ω = mean(find_mode(X))
    μ_bar = mean(X);
    left,right = hdi(X);

    h = fit(Histogram, X, nbins=100)
    w = h.weights      # weights of each bar
    e = h.edges[1]     # edges of the bars (must be +1 more than bars)

    riemannSum = sum(w.*diff(e))
    w = w ./ riemannSum

    e2 = similar([e[:]; e[:]])  # doubling of the edge vector
    e2[2:2:end] = e             # even indices
    e2[1:2:end] = e             # odd indices
    e2 = e2[2:(end-1)]          # pop the ends

    w2 = similar([w[:]; w[:]])  # doubling of the weight vector
    w2[2:2:end] = w             # even indices
    w2[1:2:end] = w             # odd indices
    w2 = w2 .* scale .+ offset

    w_left = w[sum(e .< left)] * scale + offset  # comparison leaves at max left edge of bar that contains the limit
    w_right = w[sum(e .< right)] * scale + offset
    w_ω = w[sum(e .< ω)] * scale + offset
    w_μ = w[sum(e .< μ_bar)] * scale + offset

    #%% Plotting
    plot(e2, w2,
         linealpha=0.1, linecolor=color,
         fillrange=offset, fillalpha=0.4, fillcolor=color,
         legend=false)

    # or do it by histogram, then we also have vertical lines
    # specify statsplot?
    #histogram(X, bins=100, normalize=:pdf,
    #          color=color, alpha=0.4, linealpha=0.1, legend=false)   # Comes from StatsPlots now

    if ann
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (left, w_left, text("$(round(left,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (right, w_right, text("$(round(right,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (μ_bar, 0.5*w_μ, text("mean $(round(μ_bar,sigdigits=4))",color,rotation=90,:bottom,:right)))
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot,
        annotations = (ω, w_ω, text("mode $(round(ω,sigdigits=4))",color,rotation=90,:bottom,:right)))
    else
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot)
    end
end

function makeDistributionPlot!(X, color="blue"; ann=true, offset=0.0, scale=1.0)
    #%% Mode, Mean, HDIs
    ω = mean(find_mode(X))
    μ_bar = mean(X);
    left,right = hdi(X);

    h = fit(Histogram, X, nbins=100)
    w = h.weights      # weights of each bar
    e = h.edges[1]     # edges of the bars (must be +1 more than bars)

    riemannSum = sum(w.*diff(e))
    w = w ./ riemannSum

    e2 = similar([e[:]; e[:]])  # doubling of the edge vector
    e2[2:2:end] = e             # even indices
    e2[1:2:end] = e             # odd indices
    e2 = e2[2:(end-1)]          # pop the ends

    w2 = similar([w[:]; w[:]])  # doubling of the weight vector
    w2[2:2:end] = w             # even indices
    w2[1:2:end] = w             # odd indices
    w2 = w2 .* scale .+ offset

    w_left = w[sum(e .< left)] * scale + offset  # comparison leaves at max left edge of bar that contains the limit
    w_right = w[sum(e .< right)] * scale + offset
    w_ω = w[sum(e .< ω)] * scale + offset
    w_μ = w[sum(e .< μ_bar)] * scale + offset

    #%% Plotting
    plot!(e2, w2,
         linealpha=0.1, linecolor=color,
         fillrange=offset, fillalpha=0.4, fillcolor=color,
         legend=false)

    # or do it by histogram, then we also have vertical lines
    # specify statsplot?
    #histogram!(X, bins=100, normalize=:pdf,
    #          color=color, alpha=0.4, linealpha=0.1, legend=false)   # Comes from StatsPlots now

    if ann
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (left, w_left, text("$(round(left,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (right, w_right, text("$(round(right,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (μ_bar, 0.5*w_μ, text("mean $(round(μ_bar,sigdigits=4))",color,rotation=90,:bottom,:right)))
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot,
        annotations = (ω, w_ω, text("mode $(round(ω,sigdigits=4))",color,rotation=90,:bottom,:right)))
    else
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot)
    end
end

# Taken from Jesper:
# https://en.wikipedia.org/wiki/Correlation_coefficient
function ACF(x, k=1)
  x_meanshift = x .- mean(x)
  zxk = x_meanshift[k+1:end]
  zyk = x_meanshift[1:end-k]
  return sum(zxk.*zyk)/sqrt(sum(zxk.^2)*sum(zyk.^2))
end

# Taken from Jesper:
function acovlim(x;lim=0.05)
  k = 0
  rhos = []
  rho = 1
  while rho>lim
    rho = ACF(x,k)
    push!(rhos,rho)
    k += 1
  end
  return rhos
end

# Taken from Jesper:
# ess -- effective sample size (Kruschke 2014, page 184)
function ess(x)
    if typeof(x)==Vector{Float64}
        n = length(x)
        acf = acovlim(x)
        return n/(1+2*sum(acf[2:end]))
    else
        m,n = size(x)
        list = zeros(m)
        for i in 1:m
            acf = acovlim(x[i,:])
            list[i] = n/(1+2*sum(acf[2:end]))
        end
        return list
    end
end

################################################################################

#####
#%% Priors, Likelihoods, Posteriors ############################################
#####
@everywhere function log_prior_flat(X::Float64)
    return 0.0  # or any other constant value, for the posterior it doesn't matter
end

@everywhere function log_prior_flat_positiveOnly(X::Float64)
    if X > 0
        return 0.0  # or any other constant value, for the posterior it doesn't matter
    else
        return -Inf
    end
    # Alternative of Jesper:
    #return log(tau > 0)
end

# Here we have to
# accomodate, that the means no longer are a constant, but an array, i.e. a
# a mean for each theta
@everywhere function log_prior_theta(θ::Array{Float64},μ::Array{Float64},σ::Float64)
    # Log-normal distribution
    # Logarithmise the normal pdfs in series:
    # We can leave out the 1/sqrt(2π) since it does not change sampling from the posterior
    return sum(-log(σ) .- 0.5 .* ((θ .- μ) ./ σ).^2)
end

@everywhere function log_lklhd_fct(jointParams::Array{Float64},zlogy::Array{Float64,1},ind::Array{Int64,1})
    # return the joint lklhd of the given input data
    J = length(jointParams) - 4
    θ = jointParams[1:J]
    μ_0, σ, τ, ϕ  = jointParams[(J+1):end]

    jointLklhd = 0.0

    if (σ > 0)
        # Logarithmise the normal pdfs in series:
        # summing up:
        # note, that we go through each observation, vectorising is hard:
        for i in 1:length(zlogy)
            jointLklhd += -log(σ) - 0.5 * ((zlogy[i] - θ[ind[i]]) / σ)^2
        end
    else
        jointLklhd = -Inf
    end

    return jointLklhd
end

@everywhere function log_priors(jointParams::Array{Float64},K_j::Array{Int64,1})
    # The joint log-pdf of all the priors in the model
    J = length(jointParams) - 4
    θ = jointParams[1:J]
    μ_0, σ, τ, ϕ  = jointParams[(J+1):end]

    μ = μ_0 .+ ϕ .* K_j.*1.0    # current sampled μ for each individual

    if (σ > 0) && (τ > 0)
        return log_prior_theta(θ,μ,τ) +
               log_prior_flat(μ_0) +
               log_prior_flat(ϕ) +
               log_prior_flat_positiveOnly(σ) +
               log_prior_flat_positiveOnly(τ)
    else
        return -Inf
    end
end

#%% Logarithmise and Transform Data ############################################
J = maximum(ind)    # number of individuals
I = length(y)       # number of observations

# Go to log-space:
@everywhere logy = log.(y)
@everywhere logMean = mean(logy)
@everywhere logStd = std(logy)

# Mean-centre and scale each element:
@everywhere zlogy = (logy .- logMean) ./ logStd

#%% Concretise the functions by including the data #############################
@everywhere log_lklhd(jointParams::Array{Float64}) =
            log_lklhd_fct(jointParams::Array{Float64}, zlogy, ind)
@everywhere log_posterior(jointParams::Array{Float64}) =
            log_lklhd(jointParams::Array{Float64}) +
            log_priors(jointParams::Array{Float64},child_j)

################################################################################
############################# MCMC #############################################
#%% Chain Setup
noOfChains = 4
N = 5*10^6                # total number of samples
N_chain = convert(Int64, N / noOfChains)   # samples per chain
burnIn = 10^4
w = ones(Float64,J+4) * 0.1         # typical window size
m = 100                             # multiplier for maximum window size

#%% MCMC run, make several chains
chain = SharedArray{Float64}(N_chain,J+4,noOfChains);
chainPdf = SharedArray{Float64}(N_chain,noOfChains);

# (using threads for parallel computing (other stuff didn't work)):
# Threads.@threads led to problems too
# now using @distributed, prefix with @sync to wait for completion
@sync @distributed for k in 1:noOfChains
    x_start = rand(Float64,J+4) * 0.5   # initial vector
    chain[:,:,k], chainPdf[:,k] = mySliceSampler(log_posterior,x_start,w,m,N_chain,burnIn);
end
# My machine: 4 cores, i7-7500U 2,7 GHz x 2 each, on Linux
# N=10^5 takes around 25 seconds.
# N=10^6 takes around 120 seconds.
# N=5*10^6 takes around 590 seconds.

################################################################################
############################# RESULTS ##########################################

##############
# Acess results and Transform back (undo mean-centering and scaling and go to non-log space)
# See derivation in PDF file instead

# Access the parameters and fuse the chains:
θ = Array{Float64,2}(undef,N,J)
for j in 1:J
    θ[:,j] = chain[:,j,:][:]
end
μ_0 = chain[:,J+1,:][:]
σ   = chain[:,J+2,:][:]
τ   = chain[:,J+3,:][:]
ϕ   = chain[:,J+4,:][:]

#####
# Autocorrelation and ESS
ess_θ = Array{Float64,1}(undef,J)
for j in 1:J
    ess_θ[j] = ess(θ[:,j])
    println("ESS of θ_$j = ",ess_θ[j])
end
println("Minimal ESS is of θ_$(argmin(ess_θ)) = ",ess(θ[:,argmin(ess_θ)]))
println("Maximal ESS is of θ_$(argmax(ess_θ)) = ",ess(θ[:,argmax(ess_θ)]))

println("ESS of μ_0 = ",ess(μ_0))
println("ESS of σ = ",ess(σ))
println("ESS of τ = ",ess(τ))
println("ESS of ϕ = ",ess(ϕ))

#####
# Un-scale and un-mean-centre:
θ_unscaled = θ .* logStd .+ logMean
# adults:
μ_0_unscaled = (μ_0 .+ 0) .* logStd .+ logMean
# kids only:
μ_ϕ_unscaled = (μ_0 .+ ϕ) .* logStd .+ logMean
ϕ_unscaled = ϕ .* logStd

σ_unscaled = σ .* logStd
τ_unscaled = τ .* logStd

#####
# Get into non-log space:
θ_unscaled_unLog = exp.(θ_unscaled .+ 0.5 .* repeat(σ_unscaled,1,J).^2);
μ_0_trans_unLog = exp.(μ_0_unscaled .+ 0.5 .* σ_unscaled.^2 .+ 0.5 .* τ_unscaled.^2);    # adults
μ_ϕ_trans_unLog = exp.(μ_0_unscaled .+ ϕ_unscaled .+ 0.5 .* σ_unscaled.^2 .+ 0.5 .* τ_unscaled.^2);  # kids

################################################################################
############################# TASKS ############################################

######################## Task 1 ####################################
# Effect of being a kid:
makeDistributionPlot(ϕ, "orange")
Plots.savefig(projDir*"/figs/phi_Slice.pdf")
makeDistributionPlot(ϕ_unscaled, "orange")
Plots.savefig(projDir*"/figs/phi_unscaled_Slice.pdf")

######################## Task 2 #############
makeDistributionPlot(τ,"blue")
Plots.savefig(projDir*"/figs/tau_Slice.pdf")

makeDistributionPlot(τ_unscaled, "blue")
Plots.savefig(projDir*"/figs/tau_unscaled_Slice.pdf")

######################## Task 3, Priors of expected log reaction #############
# prior for theta was
# theta[j] ~ normal(mu + phi*ISKID[j],tau)

prior_adult = mean(μ_0_unscaled) .+ mean(τ_unscaled) .* randn(N)
prior_kid = mean(μ_0_unscaled) .+ mean(ϕ_unscaled) .+ mean(τ_unscaled) .* randn(N)
makeDistributionPlot(prior_adult,"black")
makeDistributionPlot!(prior_kid,"red")
Plots.savefig(projDir*"/figs/priors_Slice.pdf")

# Logarithmic sample means
θ_mean_unscaled = zeros(J)
for j in 1:J
    θ_mean_unscaled[j] = mean(logy[ind .== j])
end

#%% Plot into one figure
# Get ordered indices:
sIdx = sortperm(θ_mean_unscaled)

plt = makeDistributionPlot(prior_adult,"black",ann=false)
makeDistributionPlot!(prior_kid,"red",ann=false)
top = 1.6;
for i in 1:J
    j = sIdx[i]
    if child_j[j] == 1
        color = "red"
    else
        color = "black"
    end
    makeDistributionPlot!(θ_unscaled[:,j],color,ann=false,offset=i*top/J,scale=1/100)

end
plt

Plots.savefig(projDir*"/figs/priors_postOverlay_Slice.pdf")

######################## Task 4, posterior prediction #############
#### a) knowing that it is a child

# Posterior predictive sampling:
idx = Int.(ceil.(rand(N).*N))   # choose random indices which will pick from the simulated posterior
# Sample zlogy according to the model:
zlogy_sim_adult = μ_0[idx] .+ randn(N).*τ[idx] .+ randn(N).*σ[idx]
zlogy_sim_kid = μ_0[idx] .+ ϕ[idx] .+ randn(N).*τ[idx] .+ randn(N).*σ[idx]
# Transform to non-standardised and non-log:
y_sim_adult = exp.(zlogy_sim_adult .* logStd .+ logMean)
y_sim_kid = exp.(zlogy_sim_kid .* logStd .+ logMean)

makeDistributionPlot(y_sim_adult,"black",ann=true)
makeDistributionPlot!(y_sim_kid,"red",ann=true)
Plots.savefig(projDir*"/figs/PP_known_Slice.pdf")

##############
#### b) not knowing that it is a child
# assuming a prior of equally likely to be adult or child (beta(1,1))
# then add a likelihood corresponding to a bernoulli processs, with number of
# heads equvivalent to number of kids among all individuals
# the posterior probability of having the amount of childs among
# our total individuals is given by a bernoulli process with a beta
# with parameters a=b=1, N = # of individuals, z = # of kids

nr_kids = sum(child_j)  # number of heads
a = nr_kids + 1         # head (kid) count here
b = J - nr_kids + 1     # non-heads (adults)
postBeingKid = rand(Beta(a,b),N)
areKids = postBeingKid .>= rand(N)  # only those who make it over the threshold

# Posterior predictive sampling
idx = Int.(ceil.(rand(N).*N))   # choose a random index which will pick from the simulated posterior
zlogy_sim_unknown = μ_0[idx] .+ ϕ[idx] .* areKids .+ randn(N).*τ[idx] .+ randn(N).*σ[idx]
y_sim_unknown = exp.(zlogy_sim_unknown .* logStd .+ logMean)
# Compare in plots:
makeDistributionPlot(y_sim_unknown,"blue",ann=true)
histogram!(y_sim_adult, bins=100, normalize=:pdf, alpha=0.1, linealpha=0.1, color="black")
histogram!(y_sim_kid, bins=100, normalize=:pdf, alpha=0.1, linealpha=0.1, color="red")
Plots.savefig(projDir*"/figs/PP_unknown_Slice.pdf")

##############
#### b-Version 2) not knowing that it is a child, set a fixed fraction
# using fraction 0.5
weight = 0.5
postBeingKid = weight
areKids = postBeingKid .>= rand(N)
idx = Int.(ceil.(rand(N).*N))   # choose a random index which will pick from the simulated posterior
zlogy_sim_unknown = μ_0[idx] .+ ϕ[idx] .* areKids .+ randn(N).*τ[idx] .+ randn(N).*σ[idx]
y_sim_unknown = exp.(zlogy_sim_unknown .* logStd .+ logMean)
makeDistributionPlot(y_sim_unknown,"blue",ann=true)
histogram!(y_sim_adult, bins=100, normalize=:pdf, alpha=0.1, linealpha=0.1, color="black")
histogram!(y_sim_kid, bins=100, normalize=:pdf, alpha=0.1, linealpha=0.1, color="red")
plot!(ann=(1200,0.0025,"Kids: $(weight*100) %"),grid=false)
Plots.savefig(projDir*"/figs/PP_unknown_fixed05_Slice.pdf")

#cur_colors = get_color_palette(:auto, plot_color(:white), 11)
# Limits for the figure
myXlims = (50,1200)
#myYlims = (0,0.006)

plt = StatsPlots.plot(size = (800, 1600),xlims=myXlims)

# Try various fixed ratios:
for (i, weight) in enumerate(0:.1:1)
    postBeingKid = weight
    areKids = postBeingKid .>= rand(N)
    idx = Int.(ceil.(rand(N).*N))   # choose a random index which will pick from the simulated posterior
    zlogy_sim_unknown = μ_0[idx] .+ ϕ[idx] .* areKids .+ randn(N).*τ[idx] .+ randn(N).*σ[idx]
    y_sim_unknown = exp.(zlogy_sim_unknown .* logStd .+ logMean)

    makeDistributionPlot!(y_sim_unknown,"blue",ann=false,offset=(i-1)*0.006,scale=1.0)
    makeDistributionPlot!(y_sim_adult,"black",ann=false,offset=(i-1)*0.006,scale=1.0)
    makeDistributionPlot!(y_sim_kid,"red",ann=false,offset=(i-1)*0.006,scale=1.0)
    plot!(ann=(myXlims[2]-200,(i-1)*0.006+0.003,"Kids: $(weight*100) %"),grid=false,yticks=false)
end
# Show plot
plt

Plots.savefig(projDir*"/figs/CompMixture_Slice.pdf")


text:
"""
	pb = BorderedProblem(;F, dxF, dpF, g, ∇g, dpg)

This composite type encodes a bordered problem, one by which we add a scalar constraint `g(x, p) = 0` to an equation `F(x, p) = 0`. This composite type thus allows to define the functional `G((x, p)) = [F(x, p) g(x, p)]` and solve `G = 0`.

You can then evaluate the functional using or `pb(z)` where `z = BorderedArray(x, p)` or `z = vcat(x, p)`, i.e. the last component of the vector is the Lagrange Multiplier.

# Arguments
The arguments correspond to the functions `F, g` and their derivatives.

# Simplified constructor

You can create such functional as `pb = BorderedProblem(F, g)`.

!!! warning "Multidimensional constraint (Really Experimental)"
    It is in fact possible, using this composite type, to define a bordered problem with constraint of dimension `npar > 1`. One has to pass the dimension to `pb = BorderedProblem(F, g, npar)` and possibly the derivatives as well. The second argument of `F,g` is `npar` dimensional (for now an `AbstractVector`). Finally, the only possible linear (bordered) solver in this case is `::MatrixBLS`.

"""
@with_kw struct BorderedProblem{Tf, Tdf, TdpF, Tg, Tdg, Tdpg, Tlens}
	F::Tf		# functional
	dxF::Tdf = (x0, p0) -> finiteDifferences(x -> F(x, p0), x0)	# partial Derivative w.r.t. first variable
	dpF::TdpF = (x0, p0) -> (F(x0, p0 + 1e-8) .- F(x0, p0)) .* 1e8	# partial Derivative w.r.t. first variable
	g::Tg		# scalar constraint
	∇g::Tdg = (x0, p0) -> finiteDifferences(x -> g(x, p0), x0)		# partial Derivative w.r.t. first variable
	dpg::Tdpg = (x0, p0) -> (g(x0, p0 + 1e-8) - g(x0, p0)) .* 1e8	# partial Derivative w.r.t. first variable
	npar::Int = 1													# number of parameters, equals length(g(x,p))
	lens::Tlens = @lens _[1]
end

BorderedProblem(F, g, lens::Lens, npar = 1) = BorderedProblem(F = F, g = g, lens = lens, npar = 1)

getParameterDim(pb::BorderedProblem) = pb.npar

function extractParameter(pb::BorderedProblem, x::AbstractVector)
	if getParameterDim(pb) == 1
		return x[end]
	else
		return @view x[end-pb.npar+1:end]
	end
end

extractVector(pb::BorderedProblem, x::AbstractVector) = @view x[1:end-getParameterDim(pb)]
setVector!(pb::BorderedProblem, out::AbstractVector, x::AbstractVector) = out[1:end-getParameterDim(pb)] .= x

etParameter!(pb::BorderedProblem, out::BorderedArray, p::Number) = out[end] = p
setParameter!(pb::BorderedProblem, out::AbstractVector, p) = out[end-getParameterDim(pb)+1:end] .= p

extractParameter(pb::BorderedProblem, x::BorderedArray) = x.p
extractVector(pb::BorderedProblem, x::BorderedArray) = x.u
setVector!(pb::BorderedProblem, out::BorderedArray, x) = copyto!(out.u, x)
setParameter!(pb::BorderedProblem, out::BorderedArray, p::Number) = out.p = p
setParameter!(pb::BorderedProblem, out::BorderedArray, p) = copyto!(out.p, p)

function (pb::BorderedProblem)(xe, par)
	out = similar(xe)
	# extract variables
	p = extractParameter(pb, xe)
	x = extractVector(pb, xe)
	# compute the residual
	setVector!(pb, out, pb.F(x, set(par, pb.lens, p)))
	setParameter!(pb, out, pb.g(x, set(par, pb.lens, p)))
	out
end

function (pb::BorderedProblem)(xe, par, dxe)
	out = similar(dxe)
	# extract variables
	p = extractParameter(pb, xe)
	x = extractVector(pb, xe)

	# extract the jacobians
	J = pb.dxF(x, par)

	@error "WIP"
	out
end

# Structure to hold the jacobian of the bordered problem
# mutable struct JacobianBorderedProblem{Tpb, Tj, Tdpf, Tdg, Tdpg}
# 	pb::Tpb
# 	J::Tj
# 	dpF::Tdpf
# 	∇g::Tdg
# 	dpg::Tdpg
# end

# # simplified constructor
# JacobianBorderedProblem(pb, x, p) = JacobianBorderedProblem(pb, pb.dxF(x, p), pb.dpF(x, p), pb.∇g(x, p), pb.dpg(x, p))
#
# JacobianBorderedProblem(pb, x, p) = JacobianBorderedProblem(pb, extractVector(pb, x), extractParameter(pb, x))
#
#
# function (Jpb::JacobianBorderedProblem)(x, p)
# 	# computation of the jacobian of the Bordered problem
# 	pb = Jpb.pb
# 	if Jpb.J isa AbstractArray
# 		copyto!(Jpb.J, 	 pb.dxF(x, p))
# 	else
# 		Jpb.J = pb.dxF(x, p)
# 	end
#
# 	if Jpb.dpF isa AbstractArray
# 		copyto!(Jpb.dpF, pb.dpF(x, p))
# 	else
# 		Jpb.dpF = pb.dpF(x, p)
# 	end
#
# 	if Jpb.∇g isa AbstractArray
# 		copyto!(Jpb.∇g,  pb.∇g(x, p))
# 	else
# 		Jpb.∇g = pb.∇g(x, p)
# 	end
# 	if getParameterDim(pb) == 1
# 		Jpb.dpg = pb.dpg(x, p)
# 	else
# 		copyto!(Jpb.dpg, pb.dpg(x, p))
# 	end
# 	return Jpb
# end
#
# (Jbp::JacobianBorderedProblem)(x) = Jbp(extractVector(Jbp.pb, x), extractParameter(Jbp.pb, x))
#
# @with_kw struct LinearSolverBorderedProblem{T, Ts, L <: AbstractBorderedLinearSolver} <: AbstractLinearSolver
# 	ls::L
# 	# these are used to alter the second linear equation in the bordered system
# 	xiu::T = 1.0
# 	xip::T = 1.0
# 	shift::Ts = nothing
# end
#
# # simplified constructor
# LinearSolverBorderedProblem(ls) = LinearSolverBorderedProblem(ls = ls)
#
# function (lsbdp::LinearSolverBorderedProblem)(J::JacobianBorderedProblem, x)
# 	# call the bordered linear solver
# 	ou, op, flag, it = lsbdp.ls(J.J, J.dpF, J.∇g, J.dpg, extractVector(J.pb, x), extractParameter(J.pb, x), lsbdp.xiu, lsbdp.xip; shift = lsbdp.shift)
# 	if x isa BorderedArray
# 		return BorderedArray(ou, op), flag, it
# 	else
# 		return vcat(ou, op), flag, it
# 	end
# end
#
# ####################################################################################################
# # newton functions
# """
# 	newtonBordered(pb::BorderedProblem, z0, options::NewtonPar{T, L, S}; kwargs...)
#
# This function solves the equation associated with the functional `pb` with initial guess
# """
# function newtonBordered(pb::Tpb, z0, par, options::NewtonPar{T, L, S}; kwargs...) where {T, L <: AbstractBorderedLinearSolver, S, Tpb <: BorderedProblem}
# 	@show z0
# 	Jac   = JacobianBorderedProblem(pb, z0, par)
# 	lsbpb = LinearSolverBorderedProblem(options.linsolver)
# 	options2 = @set options.linsolver = lsbpb
# 	return newton(pb, Jac, z0, par, options2; kwargs...)
# end
#
# """
# This is the newton solver used to solve `F(x, p) = 0` together
# with the scalar condition `n(x, p) = (x - x0) * xp + (p - p0) * lp - n0`
# """
# function _newtonPALC(F, Jh,
# 					z0::BorderedArray{vectype, T},
# 					tau0::BorderedArray{vectype, T},
# 					z_pred::BorderedArray{vectype, T},
# 					options::ContinuationPar{T},
# 					dottheta::DotTheta;
# 					linearbdalgo = BorderingBLS(),
# 					normN = norm,
# 					callback = cbDefault, kwargs...) where {T, vectype}
# 	# Extract parameters
# 	@error "WIP - will replace newtonPALC once performances improve"
# 	newtonOpts = options.newtonOptions
# 	@unpack tol, maxIter, verbose, alpha, almin, linesearch = newtonOpts
# 	@unpack theta, ds, finDiffEps = options
#
# 	N = (x, p) -> arcLengthEq(dottheta, minus(x, z0.u), p - z0.p, tau0.u, tau0.p, theta, ds)
# 	normAC = (resf, resn) -> max(normN(resf), abs(resn))
#
# 	pb = BorderedProblem(F = F, dxF = Jh, g = N, ∇g = (x, p) -> tau0.u, dpg = (x, p) -> tau0.p)
# 	Jac   = JacobianBorderedProblem(pb, z0)
# 	lsbpb = LinearSolverBorderedProblem(ls = linearbdalgo, xiu = theta / length(z0.u), xip = one(T) - theta)
# 	options2 = @set newtonOpts.linsolver = lsbpb
# 	return newton(pb, Jac, z0, options2; kwargs...)
# end
# ####################################################################################################
# # continuation function
# function continuationBordered(pb, z0, p0::Real, contParams::ContinuationPar, linearAlgo::AbstractBorderedLinearSolver; kwargs...) where {T, L <: AbstractBorderedLinearSolver, S}
# 	Jac   = p -> JacobianBorderedProblem(pb(p), z0)
# 	lsbpb = LinearSolverBorderedProblem(contParams.newtonOptions.linsolver)
# 	contParams2 = @set contParams.newtonOptions.linsolver = lsbpb
# 	return continuation((z, p) -> pb(p)(z),
# 						(z, p) -> Jac(p)(z),
# 						z0, p0, contParams2; kwargs...)
# end
#
# """
# continuationBordered(prob, z0, p0::Real, contParams::ContinuationPar; kwargs...)
#
# This is the continuation routine for finding the curve of solutions of a family of Bordered problems `p->prob(p)`.
#
# # Arguments
# - `p -> prob(p)` is a family such that `prob(p)::BorderedProblem` encodes the functional G
# - `z0` a guess for the constrained problem.
# - `p0` initial parameter, must be a real number
# - `contParams` same as for the regular `continuation` method
# """
# function continuationBordered(prob, z0, p0::Real, contParams::ContinuationPar; linearAlgo = BorderingBLS(), kwargs...) where {T, L <: AbstractBorderedLinearSolver, S}
# 	linearAlgo = @set linearAlgo.solver = contParams.newtonOptions.linsolver
# 	return continuationBordered(prob, z0, p0, contParams, linearAlgo; kwargs...)
# end


text:
# additional method to render the crosswalk environment with obstacles
function AutoViz.render!(rendermodel::RenderModel, env::CrosswalkEnv)
    roadway = gen_straight_roadway(2, env.params.roadway_length)
    AutoViz.render!(rendermodel, roadway)

    curve = env.crosswalk.curve
    n = length(curve)
    pts = Array{Float64}(undef, 2, n)
    for (i,pt) in enumerate(curve)
        pts[1,i] = pt.pos.x
        pts[2,i] = pt.pos.y
    end

    add_instruction!(rendermodel, render_dashed_line, (pts, colorant"white", env.crosswalk.width, 1.0, 1.0, 0.0, Cairo.CAIRO_LINE_CAP_BUTT))
    for obs in env.obstacles
        pts = Array{Float64}(undef, 2, obs.npts)
        for (i, pt) in enumerate(obs.pts)
            pts[1,i] = pt.x
            pts[2,i] = pt.y
        end

        add_instruction!(rendermodel, render_fill_region, (pts, colorant"gray"))
    end

    return rendermodel
end

function AutoViz.render!(rendermodel::RenderModel, env::UrbanEnv)
    # regular roadway
    roadway = Roadway(env.roadway.segments[1:end-length(env.crosswalks)])
    AutoViz.render!(rendermodel, roadway)

    for (i, cw) in enumerate(env.crosswalks)
        # crosswalk
        curve = cw.curve
        n = length(curve)
        pts = Array{Float64}(undef, 2, n)
        for (i,pt) in enumerate(curve)
            pts[1,i] = pt.pos.x
            pts[2,i] = pt.pos.y
        end
        add_instruction!(rendermodel, render_dashed_line, (pts, colorant"white", cw.width, 0.5, 0.7, 0.0, Cairo.CAIRO_LINE_CAP_BUTT))
    end
    # obstacles
    for obs in env.obstacles
        pts = Array{Float64}(undef, 2, obs.npts)
        for (i, pt) in enumerate(obs.pts)
            pts[1,i] = pt.x
            pts[2,i] = pt.y
        end

        add_instruction!(rendermodel, render_fill_region, (pts, colorant"gray"))
    end

    # render stop line
    stop_line = posg(Frenet(env.roadway[LaneTag(6,1)], env.params.stop_line), env.roadway)
    x_pos, y_pos = stop_line.x, stop_line.y
    stop_pts = zeros(2,2)
    stop_pts[1,:] =  [(x_pos - env.params.lane_width/2) , (x_pos + env.params.lane_width/2)]
    stop_pts[2,:] =  [y_pos, y_pos]
    add_instruction!(rendermodel, render_line, (stop_pts, colorant"white", 1.0, Cairo.CAIRO_LINE_CAP_BUTT))

    return rendermodel
end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
include("vcf_parser.jl")

"""
Load a VCF file into a Vcf object
A Vcf object has .header and a .data field
.data field is stored as a DataFrame
"""
function vcf_read(filename::AbstractString)
    stream = vcf_open(filename)
    if stream == false
        return false
    end
    hd = vcf_read_header(stream)
    data = vcf_read_data(stream, vcf_has_format_column(hd), '\t')

    return Vcf(hd, data)
end

"""
Write a Vcf object into a file
Vcf data field is tab-separated
"""
function vcf_write(filename::AbstractString, obj::Vcf)
    stream = vcf_open(filename, "w")
    if stream == false
        return false
    end
    vcf_write_header(stream, obj.header)
    vcf_write_data(stream, obj.data)
    close(stream)
end

function vcf_read_data(stream, has_format, separator = '\t')
    data = Variant[]
    while(true)
        try
            if eof(stream)
                return data
            end
            line = rstrip(readline(stream),'\n')
            var = vcf_parse_data_line(line, has_format, separator)
            if var == false
                continue
            end
            push!(data, var)
        catch e
            println(e)
            return data
        end
    end
end

# open a vcf stream, mode should be either r or w
function vcf_open(filename::AbstractString, mode::AbstractString="r")
    return opengene_open(filename, mode)
end

vcf_close(stream) = close(stream)

function vcf_read_header(stream)
    header = VcfHeader()
    while true
        if eof(stream)
            return header
        end
        line = readline(stream)
        if startswith(line, "##")
            vcf_add_meta_line(header, line)
        elseif startswith(line, "#CHROM")
            vcf_add_column_line(header, line)
            return header
        else
            error("VCF column header is not found, column header line is started with #CHROM")
            return false
        end
    end
end

# write a vcf header to a file
function vcf_write_header(stream, header::VcfHeader)
    try
        vcf_write_metas(stream, header)
        vcf_write_columns(stream, header)
    catch e
        println(e)
        return false
    end
end

# write vcf data field into stream
function vcf_write_data(stream, data)
    has_format = false
    for d in data
        if d.format != ""
            has_format = true
            break
        end
    end
    for d in data
        write(stream, d.chrom)
        write(stream, '\t', string(d.pos))
        write(stream, '\t', d.id)
        write(stream, '\t', d.ref)
        write(stream, '\t', d.alt)
        write(stream, '\t', d.qual)
        write(stream, '\t', d.filter)
        write(stream, '\t', d.info)
        if has_format
            write(stream, '\t', d.format)
        end
        for s in d.samples
            write(stream, '\t', s)
        end
        write(stream, '\n')
    end
end


text:
using ComoniconGUI
using Test

@testset "ComoniconGUI.jl" begin
    # Write your tests here.
end


text:
SKD = ScikitLearn.Skcore.pyimport("sklearn.datasets")

rmse(y, ŷ) = sqrt(sum(abs2, ŷ .- y) / length(y))

function simple_regression(n=100, p=3; sigma=0.1, σ=sigma, seed=6168901234)
    r = StableRNG(seed)
    X = hcat(randn(r, n, p), ones(n))
    θ = randn(r, p + 1)
    y = X * θ + σ * randn(r, n)
    ŷ = X * (X \ y)
    return X, y, rmse(y, ŷ)
end

function test_regression(m, X, y, ls)
    f, _, _ = MB.fit(m, 1, X, y)
    ŷ = MB.predict(m, f, X)
    return rmse(y, ŷ) / ls, f
end

function simple_regression_2(n=100, p=3;
                             sigma=0.1, σ=sigma, seed=6168901234)
    r  = StableRNG(seed)
    X  = hcat(randn(r, n, p), ones(n))
    θ  = randn(r, p + 1)
    y1 = X * θ + σ * randn(r, n)
    y2 = X * θ + σ * randn(r, n)
    ŷ  = X * (vcat(X, X) \ vcat(y1, y2))
    return X, hcat(y1, y2), (rmse(y1, ŷ)+rmse(y2, ŷ))/2
end

function test_regression_2(m, X, y, ls)
    f, _, _ = MB.fit(m, 1, X, y)
    ŷ = MB.predict(m, f, X)
    return sqrt(sum(abs2, vec(MB.matrix(y) .- MB.matrix(ŷ)))) / ls, f
end

function simple_binaryclf(n=100, p=3; sigma=0.1, seed=616866614)
    X, y = SKD.make_blobs(n_samples=n, n_features=p, centers=2, random_state=seed)
    return MB.table(X), MB.categorical(y)
end

function test_clf(m, X, y)
    f, _, _ = MB.fit(m, 1, X, y)
    if typeof(m) <: MB.Deterministic
        ŷ = MB.predict(m, f, X)
    else
        ŷ = MB.predict_mode(m, f, X)
    end
    return MB.accuracy(ŷ, y), f
end