text:
### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ b5b8a606-2fc1-448e-b1e0-bdb5b0203268
using JSON

# ╔═╡ 3971acc6-2b7d-11ec-02fc-c3c4044112a8
include("src/show_volume.jl")

# ╔═╡ fd0aa0d8-9eeb-47fa-bde7-767d0d3b9941
vdata2 = [z/128 + .3 * (rand() - .5) for x in 1:128, y in 1:128, z in 1:128];

# ╔═╡ 99a6c6c4-b322-476a-aec9-fa828447c3b4
show_volume(vdata2; default_renderstyle="mip")

# ╔═╡ f1f49f78-3ceb-4964-85d1-24afe5000aed
io = open("vol.flat", "r");

# ╔═╡ e0279c1c-9cbf-4b7a-a1d8-62421e97cea0
s = read(io, String);

# ╔═╡ dcb8f2ef-288d-4388-92ff-8f4179f8e4c2
data = JSON.parse(s);

# ╔═╡ 6b79fe45-c141-47e5-b98a-c3ea1f95f97a
vdata = Array{Float64, 3}(reshape(data, 128, 128, 256));

# ╔═╡ be28a9f7-c96a-4d01-a593-6430da4e38c1
show_volume(vdata; cmap=[(i==1)*(j-1) + (i==3)*(256-j) for i in 1:3, j in 1:256], default_cmap="custom", clim1=0, clim2=1, default_renderstyle="iso", isothreshold=.08)

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"

[compat]
JSON = "~0.21.2"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "a8709b968a1ea6abc2dc1967cb1db6ac9a00dfb6"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.5"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
"""

# ╔═╡ Cell order:
# ╠═3971acc6-2b7d-11ec-02fc-c3c4044112a8
# ╠═fd0aa0d8-9eeb-47fa-bde7-767d0d3b9941
# ╠═99a6c6c4-b322-476a-aec9-fa828447c3b4
# ╠═b5b8a606-2fc1-448e-b1e0-bdb5b0203268
# ╠═f1f49f78-3ceb-4964-85d1-24afe5000aed
# ╠═e0279c1c-9cbf-4b7a-a1d8-62421e97cea0
# ╠═dcb8f2ef-288d-4388-92ff-8f4179f8e4c2
# ╠═6b79fe45-c141-47e5-b98a-c3ea1f95f97a
# ╠═be28a9f7-c96a-4d01-a593-6430da4e38c1
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
##############
# Owner Type #
##############

type Owner <: GitLabType
    name::Nullable{GitLabString}
    username::Nullable{GitLabString}
    id::Nullable{Int}
    state::Nullable{GitLabString}
    avatar_url::Nullable{HttpCommon.URI}
    web_url::Nullable{HttpCommon.URI}
    ownership_type::Nullable{GitLabString}

#=
    email::Nullable{GitLabString}
    bio::Nullable{GitLabString}
    company::Nullable{GitLabString}
    location::Nullable{GitLabString}
    gravatar_id::Nullable{GitLabString}
    public_repos::Nullable{Int}
    owned_private_repos::Nullable{Int}
    total_private_repos::Nullable{Int}
    public_gists::Nullable{Int}
    private_gists::Nullable{Int}
    followers::Nullable{Int}
    following::Nullable{Int}
    collaborators::Nullable{Int}
    html_url::Nullable{HttpCommon.URI}
    updated_at::Nullable{Dates.DateTime}
    created_at::Nullable{Dates.DateTime}
    date::Nullable{Dates.DateTime}
    hireable::Nullable{Bool}
    site_admin::Nullable{Bool}
=#
end

function Owner(data::Dict)
    o = json2gitlab(Owner, data)
    isnull(o.username) ? o.ownership_type = Nullable("Organization") : o.ownership_type = Nullable("User")
    o
end

Owner(username::AbstractString, isorg = false) = Owner(
    Dict("username" => isorg ? "" : username,
         "name" => isorg ? username : "",
         "ownership_type" => isorg ? "Organization" : "User"))
## Owner(username::AbstractString) = Owner(Dict("username" => username))

## namefield(owner::Owner) = owner.ownership_type == "Organization" ? owner.name : owner.username
namefield(owner::Owner) = isorg(owner) ? owner.name : owner.username

## typprefix(isorg) = isorg ? "orgs" : "users"
typprefix(isorg) = isorg ? "projects" : "users"

#############
# Owner API #
#############

isorg(owner::Owner) = isnull(owner.ownership_type) ? true : get(owner.ownership_type, "") == "Organization"

owner(owner_obj::Owner; options...) = owner(name(owner_obj), isorg(owner_obj); options...)

function owner(owner_obj, isorg = false; options...)
    ## TODO Need to look for a cleaner way of doing this ! Returns an array even while requesting a specific user
    if isorg
        result = gh_get_json("/api/v3/projects/search/$(owner_obj)"; options...)
        return Owner(result[1]["owner"])
    else
        result = gh_get_json("/api/v3/users?username=$(owner_obj)"; options...)
        return Owner(result[1])
    end
end

function users(; options...)
    results, page_data = gh_get_paged_json("/api/v3/users"; options...)
    return map(Owner, results), page_data
end

function orgs(owner; options...)
    ## results, page_data = gh_get_paged_json("/api/v3/users/$(name(owner))/projects"; options...)
    results, page_data = gh_get_paged_json("/api/v3/projects"; options...)
    return map(Owner, results), page_data
end

#= TODO: There seems to be no equivalent for these APIs
function followers(owner; options...)
    results, page_data = gh_get_paged_json("/api/v3/users/$(name(owner))/followers"; options...)
    return map(Owner, results), page_data
end

function following(owner; options...)
    results, page_data = gh_get_paged_json("/api/v3/users/$(name(owner))/following"; options...)
    return map(Owner, results), page_data
end
=#

repos(owner::Owner; options...) = repos(name(owner), isorg(owner); options...)

function repos(owner, isorg = false; options...)
    ## results, page_data = gh_get_paged_json("/api/v3/$(typprefix(isorg))/$(name(owner))/repos"; options...)
    results, page_data = gh_get_paged_json("/api/v3/projects/owned"; options...)
    return map(Repo, results), page_data
end

ext: jl


text:
module Simulations

Base.Experimental.@optlevel 3

export AbstractSimulation, CTSimulation
export simulate

using ..Monads
using ..CTImages
using ..Utils
using ProgressMeter, LinearAlgebra, IntervalSets
import Random, Distributions, StatsBase

abstract type AbstractSimulation end

function simulate end
function simulate_blocks end
function simulate_single end

include("simulations/sinogram_sampling.jl")

@inline function Random.rand(
    rng::Random.AbstractRNG,
    s::Random.SamplerTrivial{ORI{T}},
) where {T <: AbstractFloat}
    a, b = endpoints(s[])
    t = rand(rng, T)
    a * (one(T) - t) + b * t
end

end # module


text:
# Functions and operators on models
"""
    _model_dicts(model::Model)

    input: a Model object
outputs:
     - SDict: a dictionary with keys `"+","-","0","bullet"`
    and corresponding values `findall(model.C .> 0)`, `findall(model.C .< 0)`,
    `findall(model.C .== 0)`, `findall(model.C .!= 0)`, respectively.
     - TDict: a dictionary of submatrices of `T` with keys
    `"ℓm"` with ``ℓ,m∈{+,-,0,bullet}`` and corresponding values
    `model.T[S[ℓ],S[m]]`.
"""
function _model_dicts(model::Model)
    nPhases = n_phases(model)
    SDict = Dict{String,Array}("S" => 1:nPhases)
    SDict["+"] = findall(rates(model) .> 0)
    SDict["-"] = findall(rates(model) .< 0)
    SDict["0"] = findall(rates(model) .== 0)
    SDict["bullet"] = findall(rates(model) .!= 0)

    TDict = Dict{String,Array}("T" => model.T)
    for ℓ in ["+" "-" "0" "bullet"], m in ["+" "-" "0" "bullet"]
        TDict[ℓ*m] = model.T[SDict[ℓ], SDict[m]]
    end

    return SDict, TDict
end

"""
Construct and evaluate ``Ψ(s)`` for a triditional SFM.

Uses newtons method to solve the Ricatti equation
``D⁺⁻(s) + Ψ(s)D⁻⁺(s)Ψ(s) + Ψ(s)D⁻⁻(s) + D⁺⁺(s)Ψ(s) = 0.``

    psi_fun_x( model::Model; s = 0, MaxIters = 1000, err = 1e-8)

# Arguments
- `model`: a Model object
- `s::Real`: a value to evaluate the LST at
- `MaxIters::Int`: the maximum number of iterations of newtons method
- `err::Float64`: an error tolerance for terminating newtons method. Terminates
    when `max(Ψ_{n} - Ψ{n-1}) .< eps`.

# Output
- `Ψ(s)::Array{Float64,2}` the matrix ``Ψ``
"""
function psi_fun_x( model::Model; s = 0, MaxIters = 1000, err = 1e-8)
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"] - s * LinearAlgebra.I)
    # construct the generator Q(s)
    Q =
        (1 ./ abs.(rates(model)[SDict["bullet"]])) .* (
            TDict["bulletbullet"] - s * LinearAlgebra.I -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    model_without_zero_phases = BoundedFluidQueue(Q,model.S[SDict["bullet"]],model.b)

    ~, QDict = _model_dicts(model_without_zero_phases)

    Ψ = zeros(Float64, length(SDict["+"]), length(SDict["-"]))
    A = QDict["++"]
    B = QDict["--"]
    D = QDict["+-"]
    # use netwons method to solve the Ricatti equation
    for n in 1:MaxIters
        Ψ = LinearAlgebra.sylvester(A,B,D)
        if maximum(abs.(sum(Ψ,dims=2).-1)) < err
            break
        end
        A = QDict["++"] + Ψ * QDict["-+"]
        B = QDict["--"] + QDict["-+"] * Ψ
        D = QDict["+-"] - Ψ * QDict["-+"] * Ψ
    end

    return Ψ
end

"""
Construct the vector ``ξ`` containing the distribution of the phase at the time
when ``X(t)`` first hits `0`.

    xi_x( model::Model, Ψ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `psi_fun_x`

# Output
- the vector `ξ`
"""
function xi_x( model::Model, Ψ::Array)
    # the system to solve is [ξ 0](-[B₋₋ B₋₀; B₀₋ B₀₀])⁻¹[B₋₊; B₀₊]Ψ = ξ
    # writing this out and using block inversion (as described on wikipedia)
    # we can solve this in the following way
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    A =
        -(
            invT₋₋ * TDict["-+"] * Ψ + invT₋₀ * TDict["0+"] * Ψ +
            LinearAlgebra.I
        )
    b = zeros(1, size(TDict["--"], 1))
    A[:, 1] .= 1.0 # normalisation conditions
    b[1] = 1.0 # normalisation conditions

    ξ = b / A

    return ξ
end

"""
Construct the stationary distribution of the SFM

    stationary_distribution_x( model::Model, Ψ::Array, ξ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `psi_fun_x`
- `ξ`: an array as returned from `xi_x`

# Output
- `pₓ::Array{Float64,2}`: the point masses of the SFM
- `πₓ(x)` a function with two methods
    - `πₓ(x::Real)`: for scalar inputs, returns the stationary density evaluated
        at `x` in all phases.
    - `πₓ(x::Array)`: for array inputs, returns an array with the same shape
        as is output by Coeff2Dist.
- `K::Array{Float64,2}`: the matrix in the exponential of the density.
"""
function stationary_distribution_x( model::Model, Ψ::Array, ξ::Array)
    # using the same block inversion trick as in xi_x
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    Q =
        (1 ./ abs.(rates(model)[SDict["bullet"]])) .* (
            TDict["bulletbullet"] -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    model_without_zero_phases = BoundedFluidQueue(Q,model.S[SDict["bullet"]],model.b)

    ~, QDict = _model_dicts(model_without_zero_phases)

    K = QDict["++"] + Ψ * QDict["-+"]

    A = -[invT₋₋ invT₋₀]

    # unnormalised values
    αpₓ = ξ * A

    απₓ = αpₓ *
        [TDict["-+"]; TDict["0+"]] *
        -inv(K) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]]))

    απₓ0 = -απₓ * [TDict["+0"];TDict["-0"]] * T00inv

    # normalising constant
    α = sum(αpₓ) + sum(απₓ) + sum(απₓ0)

    # normalised values
    # point masses
    pₓ = αpₓ/α
    # density method for scalar x-values
    idx = [findall(rates(model).>0);
        findall(rates(model).<0);
        findall(rates(model).==0)]
    function πₓ(x::Real)
        out = zeros(n_phases(model))
        out[idx] = (pₓ *
        [TDict["-+"]; TDict["0+"]] *
        exp(K*x) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(rates(model) .!= 0)) [TDict["+0"];TDict["-0"]] * -T00inv])
        return out
    end
    # density method for arrays so that πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function πₓ(x::Array)
        temp = πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), n_phases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    # CDF method for scalar x-values
    function Πₓ(x::Real)
        out = zeros(n_phases(model))
        out[idx] = [zeros(1,sum(rates(model).>0)) pₓ] .+
        pₓ *
        [TDict["-+"]; TDict["0+"]] *
        (exp(K*x) - LinearAlgebra.I) / K *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(rates(model) .!= 0)) [TDict["+0"];TDict["-0"]] * -T00inv]
        return out
    end
    # CDF method for arrays so that Πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function Πₓ(x::Array)
        temp = Πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), n_phases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    return pₓ, πₓ, Πₓ, K
end

function stationary_distribution_x( model::Model)
    Ψ = psi_fun_x( model)
    ξ = xi_x( model, Ψ)
    pₓ, πₓ, Πₓ, K = stationary_distribution_x(model,Ψ,ξ)
    return pₓ, πₓ, Πₓ, K
end


text:
module OneDTMazeUtils

using SparseArrays
using Distributions

import ..TMazeCumulantSchedules
import ..OneDTmazeConst
import ..OneDTMaze
import ..Learner
import ..check_goal
import ..range_check
import ..get_action_probs
import ..GVFHordes
import ..update
import ..Curiosity
import ..GVFSRHordes
import ..SRCreationUtils
import ..FeatureCreator
using ...StatsBase

const TMCS = TMazeCumulantSchedules
const ODTMC = OneDTmazeConst
const SRCU = Curiosity.SRCreationUtils

#####
# GVF Parameter Functions
####

struct GoalTermination <: GVFHordes.GVFParamFuncs.AbstractDiscount
    γ::Float64
end

function Base.get(gt::GoalTermination; state_tp1, kwargs...)
    any([check_goal(OneDTMaze, i, state_tp1) for i in 1:4]) ? 0.0 : gt.γ
end

struct GoalPolicy <: GVFHordes.GVFParamFuncs.AbstractPolicy
    goal::Int
end
(π::GoalPolicy)(s) = sample(Weights([get(π, state_t=s, action_t=a) for a ∈ 1:4]))

function Base.get(π::GoalPolicy; state_t, action_t, kwargs...)
    cur_x = state_t[1]
    cur_y = state_t[2]
    if π.goal == 1
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 2
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 3
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.RIGHT == action_t
        end
    elseif π.goal == 4
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.RIGHT == action_t
        end
    end
end

####
# Demon Creation
####

function create_demons(parsed, demon_projected_fc = nothing)
    action_space = 4
    demons = if parsed["demon_learner"] != "SR"
        GVFHordes.Horde(
            [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                           GoalTermination(0.9),
                           GoalPolicy(i)) for i in 1:4])
    elseif parsed["demon_learner"] == "SR"
        @assert demon_projected_fc != nothing
        pred_horde =  GVFHordes.Horde(
                [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                     GVFHordes.GVFParamFuncs.ConstantDiscount(0.0),
                     GoalPolicy(i)) for i in 1:4])

        SF_policies = [GoalPolicy(i) for i in 1:4]
        SF_discounts = [GoalTermination(0.9) for i in 1:4]
        num_SFs = length(SF_policies)
        # SF_horde = if parsed["demon_reward_feature_type"] == "state"
        #     SRCU.create_state_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        # else
        # SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)

        # end
        GVFSRHordes.SRHorde(pred_horde, SF_horde, num_SFs, demon_projected_fc)
    else
        throw(ArgumentError("Cannot create demons"))
    end
    return demons
end

function make_behaviour_gvf(learner, γ, state_constructor, expl_strat) #discount, state_constructor_func, learner, exploration_strategy)
    function b_π(state_constructor, learner, exploration_strategy; kwargs...)
        s = state_constructor(kwargs[:state_t])
        preds = learner(s)
        return exploration_strategy(preds)[kwargs[:action_t]]
    end
    GVF_policy = GVFHordes.GVFParamFuncs.FunctionalPolicy((;kwargs...) -> b_π(state_constructor, learner, expl_strat; kwargs...))
    BehaviourGVF = GVFHordes.GVF(GVFHordes.GVFParamFuncs.RewardCumulant(),
                       GoalTermination(γ),
                       GVF_policy)
end

# ####
struct IdealStateActionDemonFeatures <: FeatureCreator
    num_actions::Int
end
function project_features(fc::IdealStateActionDemonFeatures, s_t, a_t, s_tp1)
    # new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, s_tp1) for i in 1:4]))
    goal_ind = findfirst([check_goal(OneDTMaze, i, s_tp1) for i in 1:4])
    reward_feature = zeros(Int(fc.num_actions*4))
    if !(goal_ind isa Nothing)
        reward_feature[Int((goal_ind-1)*fc.num_actions + a_t)] = 1
    end
    return sparsevec(reward_feature)
end
(FP::IdealStateActionDemonFeatures)(s_t,a_t,s_tp1) = project_features(FP, s_t,a_t,s_tp1)

Base.size(FP::IdealStateActionDemonFeatures) = 16
####
# Ideal Feature Creator
####
struct IdealDemonFeatures <: FeatureCreator
end

function project_features(fc::IdealDemonFeatures, state,action,state_tp1)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::IdealDemonFeatures)(state, action, next_state) = project_features(FP, state, action, next_state)
Base.size(FP::IdealDemonFeatures) = 4

struct MarthaIdealDemonFeatures <: FeatureCreator
    num_actions::Int
end

function project_features(fc::MarthaIdealDemonFeatures, state)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state, ODTMC.ACTION_STEP + ODTMC.EPSILON + 0.00001) for i in 1:4]))
    # alt_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::MarthaIdealDemonFeatures)(state) = project_features(FP, state)
Base.size(FP::MarthaIdealDemonFeatures) = 4

####
# GPI Feature Creation
####
struct TMazeEncoding <: FeatureCreator
    num_segments::Int
    num_partitions::Int
    num_actions::Int
    function TMazeEncoding()
        new(7,3,4)
    end
end
Base.size(FP::TMazeEncoding) = FP.num_segments * FP.num_partitions * FP.num_actions
function project_features(fc::FeatureCreator, obs, a_t, obs_tp1)
    segment = Inf
    partition = Inf
    x,y = obs

    #Segment 1: Top Right
    if (x == 0.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 1
        partition = (y - 0.8) / segment_length
    #Segment 2: Bottom Right
    elseif (x == 0.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 2
        partition = (y - 0.8) / segment_length
    #Segment 3: Middle Left Branch
    elseif (y == 0.8 && range_check(x,0.0,0.5))
        segment_length = (0.5 - 0.0)
        segment = 3
        partition = (x) / segment_length
    # Segment 4: Middle Right Branch
    elseif (y == 0.8 && range_check(x,0.5,1.0))
        segment_length = (1.0-0.5)
        segment = 4
        partition = (x - 0.5)/segment_length
    #Segment 5: Bottom Branch
    elseif (x == 0.5)
        segment_length = (0.8 - 0.0)
        segment = 5
        partition = (y - 0.0)/segment_length
    # #Segment 6: Top Left Branch
    elseif (x == 1.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 6
        partition = (y - 0.8) / segment_length
        #Segment 7: Bottom Left Branch
    elseif (x == 1.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 7
        partition = (y - 0.8) / segment_length
    else
        @warn "Not A Valid State (x,y)"
    end

    state_action = zeros(fc.num_segments * fc.num_partitions * fc.num_actions)
    # Need maximum since if agent is right on boundary, ceil (0.0) is 0.
    partition_offset = maximum([Int(ceil(partition * fc.num_partitions)),1])
    if partition_offset == 0
        @show x,y,segment
    end
    state_ind = (segment-1)*fc.num_partitions + partition_offset
    state_action_ind = state_ind * fc.num_actions + a_t

    state_action[state_action_ind] = 1
    return sparsevec(state_action)
end
(FP::TMazeEncoding)(state, action, state_tp1) = project_features(FP, state, action, state_tp1)

####
# Behaviour policies
####

Base.@kwdef struct RoundRobinPolicy <: Learner
    update = Nothing
end

Curiosity.update!(learner::RoundRobinPolicy, args...) = nothing

Base.get(π::RoundRobinPolicy; state_t, action_t, kwargs...) =
    get_action_probs(π, state_t, nothing)[action_t]

function Curiosity.get_action_probs(π::RoundRobinPolicy, features, state)
    cur_x = state[1]
    cur_y = state[2]
    ret = zeros(4)

    if cur_x == 0.5
        if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON) # Middle Junction
            ret[ODTMC.LEFT] = 0.5
            ret[ODTMC.RIGHT] = 0.5
        else # Middle Hallway
            ret[ODTMC.UP] = 1.0
        end
    elseif cur_y == 0.8 && range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON) # Left Junction
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_y == 0.8 && range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_x == 0.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x == 1.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x < 0.5
        ret[ODTMC.LEFT] = 1.0
    else
        ret[ODTMC.RIGHT] = 1.0
    end
    ret
end

####
# Cumulant Schedules
####
DrifterDistractor(parsed) = begin
    c_dist = Uniform(parsed["constant_target"][1],parsed["constant_target"][2])
    c1,c2 = rand(c_dist,2)
    if "drifter" ∈ keys(parsed)
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter"][1],
            parsed["drifter"][2],
            parsed["distractor"][1],
            parsed["distractor"][2])
    else
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter_init"],
            parsed["drifter_std"],
            parsed["distractor_mean"],
            parsed["distractor_std"])
    end
end

function get_cumulant_schedule(parsed)
    sched = parsed["cumulant_schedule"]
    if parsed["cumulant_schedule"] == "DrifterDistractor"
        DrifterDistractor(parsed)
    elseif parsed["cumulant_schedule"] == "Constant"
        if parsed["cumulant"] isa Number
            TMCS.Constant(parsed["cumulant"])
        else
            TMCS.Constant(parsed["cumulant"]...)
        end

    else
        throw("$(sched) Not Implemented")
    end
end

function get_true_values(env::Curiosity.OneDTMaze, eval_set)
    copy_eval_est = deepcopy(eval_set)
    num_gvfs = 4
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    for i in 1:num_gvfs
        copy_eval_est[i, :] .*= goal_cumulants[i]
    end
    return copy_eval_est
end
function get_true_values(env::Curiosity.OneDTMaze, eval_set, gvf_idx)
    copy_eval_est = deepcopy(eval_set)
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    copy_eval_est .*= goal_cumulants[gvf_idx]
    return copy_eval_est
end

end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
module ComputeGraphTests

using Test

using QXContexts

include("utils.jl")

@testset "Compute Graph Tests" begin

    @testset "Test commands" begin
        read_io = c -> begin
            io = IOBuffer()
            write(io, c)
            String(take!(io))
        end

        contract_example = "ncon t3 1,2,4 t1 1,2,3 t2 3,4"
        c = ContractCommand(contract_example)
        @test strip(read_io(c)) == contract_example # strip to remove \n
        @test inputs(c) == [:t1, :t2]
        @test output(c) == :t3
        @test length(params(c)) == 0

        load_example = "load t1 data_1 2,2"
        c = LoadCommand(load_example)
        @test strip(read_io(c)) == load_example # strip to remove \n
        @test inputs(c) == []
        @test output(c) == :t1
        @test length(params(c)) == 0

        save_example = "save result t3"
        c = SaveCommand(save_example)
        @test strip(read_io(c)) == save_example # strip to remove \n
        @test inputs(c) == [:t3]
        @test output(c) == :result
        @test length(params(c)) == 0

        reshape_example = "reshape t2 t1 1;2,3"
        c = ReshapeCommand(reshape_example)
        @test strip(read_io(c)) == reshape_example # strip to remove \n
        @test inputs(c) == [:t1]
        @test output(c) == :t2
        @test length(params(c)) == 0

        view_example = "view t1_s t1 v1 1 2"
        c = ViewCommand(view_example)
        @test strip(read_io(c)) == view_example # strip to remove \n
        @test inputs(c) == [:t1]
        @test output(c) == :t1_s
        @test length(params(c)) == 1
        @test params(c)[:v1] == 2

        output_example = "output t1 2 2"
        c = OutputCommand(output_example)
        @test strip(read_io(c)) == output_example # strip to remove \n
        @test inputs(c) == []
        @test output(c) == :t1
        @test length(params(c)) == 1
        @test params(c)[:o2] == 2
    end

    @testset "Test build tree" begin
        tree = build_tree(sample_cmds)
        @test length(tree) == length(sample_cmds)
        @test output(tree) == :result

        # test get commands function
        @test length(get_commands(tree, LoadCommand)) == 2
        @test length(get_commands(tree, Union{LoadCommand, ReshapeCommand})) == 3

        # test params
        @test Set(collect(keys(params(tree)))) == Set([:o1, :v1])

        # create compute graph
        cg = ComputeGraph(tree, deepcopy(sample_tensors))
        @test output(cg) == :result
    end

    @testset "Test dsl" begin
        mktempdir() do path
            tree = build_tree(sample_cmds)
            fn = joinpath(path, "foo.qx")
            open(fn, "w") do io
                write(io, tree)
            end

            tree2, metadata = parse_dsl(fn)
            @test length(tree) == length(tree2)

            fn2 = joinpath(path, "foo2.qx")
            open(fn2, "w") do io
                write(io, tree2)
            end

            dsl_1 = open(fn, "r") do io read(io, String) end
            dsl_2 = open(fn2, "r") do io read(io, String) end

            @test dsl_1 == dsl_2
        end
    end
end

end


text:
##################
# `FieldStorage` #
##################

abstract type FieldStorage{D} end

#=== `Mutable` ===#

mutable struct Mutable{D} <: FieldStorage{D}
    data::D
    Mutable{D}() where D = new{D}()
    Mutable{D}(data) where D = new{D}(data)
end

load(x::Mutable) = x.data
store!(x::Mutable, y) = (x.data = y)

#=== `Immutable` ===#

struct Immutable{D} <: FieldStorage{D}
    data::D
    Immutable{D}(data) where D = new{D}(data)
end

load(x::Immutable) = x.data
store!(x::Immutable, y) = error("cannot mutate immutable field")

##########
# `Meta` #
##########

struct NoMetaData end
struct NoMetaMeta end

_metadataconvert(T, x::NoMetaData) = x
_metadataconvert(T, x) = convert(T, x)

_metametaconvert(T, x::NoMetaMeta) = x
_metametaconvert(T, x) = convert(T, x)

struct Meta{D,M#=<:Union{Tuple,NamedTuple,Array,ModuleMeta}=#}
    data::Union{D,NoMetaData}
    meta::Union{M,NoMetaMeta}
    Meta(data::D, meta::M) where {D,M} = Meta{D,M}(data, meta)
    Meta{D,M}(data, meta) where {D,M} = new{D,M}(_metadataconvert(D, data), _metametaconvert(M, meta))
end

const NOMETA = Meta(NoMetaData(), NoMetaMeta())

# These defined to allow conversion of `Meta{NoMetaData,NoMetaMeta}`
# into whatever metatype is expected by a container.
Base.convert(::Type{M}, meta::M) where {M<:Meta} = meta

function Base.convert(::Type{Meta{D,M}}, meta::Meta) where {D,M}
    metadata = _metadataconvert(D, meta.data)
    metameta = _metametaconvert(M, meta.meta)
    return Meta{D,M}(metadata, metameta)
end

################
# `ModuleMeta` #
################
# note that `BindingMeta` was defined earlier in src/context.jl

struct ModuleMeta{D,M}
    name::Meta{D,M}
    bindings::BindingMetaDict
end

Base.convert(::Type{M}, meta::M) where {M<:ModuleMeta} = meta

function Base.convert(::Type{ModuleMeta{D,M}}, meta::ModuleMeta) where {D,M}
    return ModuleMeta(convert(Meta{D,M}, meta.name), meta.bindings)
end

# TODO: For fast methods (~ns), this fetch can cost drastically more than the primal method
# invocation. We easily have the module at compile time, but we don't have access to the
# actual context object (just the type). This `@pure` is vtjnash-approved. It should allow
# the compiler to optimize away the fetch once we have support for it, e.g. loop invariant
# code motion.
Base.@pure @noinline function fetch_tagged_module(context::Context, m::Module)
    return Tagged(context, m, Meta(NoMetaData(), fetch_modulemeta(context, m)))
end

Base.@pure @noinline function fetch_modulemeta(context::Context, m::Module)
    if haskey(context.bindingscache, m)
        bindings = context.bindingscache[m]
    else
        bindings = Cassette.BindingMetaDict()
        context.bindingscache[m] = bindings
    end
    return ModuleMeta(NOMETA, bindings::BindingMetaDict)
end

Base.@pure @noinline function _fetch_bindingmeta!(context::Context,
                                                  m::Module,
                                                  bindings::BindingMetaDict,
                                                  name::Symbol)
    return get!(bindings, name) do
        bindingmeta = BindingMeta()
        # If `!(isdefined(m, name))`, there must be a context-observable assigment to
        # the primal binding before we can access it again. This is okay because a) it's
        # obvious that the primal program can't access bindings without defining them,
        # and b) we already don't allow non-context-observable assignments that could
        # invalidate the context's metadata.
        if isdefined(m, name)
            bindingmeta.data = initmeta(context, getfield(m, name), NoMetaData())
        end
        return bindingmeta
    end
end

function fetch_bindingmeta!(context::Context,
                             m::Module,
                             bindings::BindingMetaDict,
                             name::Symbol,
                             primal)
    M = metatype(typeof(context), typeof(primal))
    return convert(M, _fetch_bindingmeta!(context, m, bindings, name).data)::M
end

############################
# `metatype` specification #
############################

#=
The optimal metatype specification for a given non-leaftype is the union of the metatypes of
its subtypes, and thus can be implemented as:

```
function metatype(::Type{C}, ::Type{T}) where {C,T}
    if isconcretetype(T)
        return Meta{metadatatype(C, T),metametatype(C, T)}
    else
        return Union{map(X -> metatype(C, X), subtypes(T))...}
    end
end
```

However, this is an expensive implementation of this operation. Thus, our actual `metatype`
implementation returns a still-correct, but extremely pessimistic metatype with the
benefit that metatype computation is very fast. If, in the future, `metatype` is
parameterized on world age, then we can call `subtypes` at compile time, and compute a more
optimally bounded metatype.
=#
function metatype(::Type{C}, ::Type{T}) where {C<:Context,T}
    if isconcretetype(T) || T <: Type
        return Meta{metadatatype(C, T),metametatype(C, T)}
    end
    return Meta
end

function _fieldtypes_for_metatype(T::Type)
    ftypes = Any[]
    for i in 1:fieldcount(T)
        # TODO improve recursion detection; this only spots 1-degree cycles right now
        ftype = fieldtype(T, i)
        ftype = ftype === T ? Any : ftype
        push!(ftypes, ftype)
    end
    return ftypes
end

doesnotneedmetatype(::Type{T}) where {T} = isbitstype(T)
doesnotneedmetatype(::Type{Symbol}) = true
doesnotneedmetatype(::Type{<:Type}) = true

doesnotneedmeta(x) = isbits(x)
doesnotneedmeta(::Symbol) = true
doesnotneedmeta(::Type) = true

#=== metadatatype ===#

"""
```
metadatatype(::Type{<:Context}, ::Type{T})
```

Overload this Cassette method w.r.t. a given context to define the type of metadata that
can be tagged to values of type `T` within that context.

By default, this method is set such that associating metadata with any tagged value is
disallowed.

Cassette uses `metadatatype` to statically compute a context-specific metadata type hiearchy
for all tagged values within overdubbed programs. To gain a mental model for this mechanism,
consider a simple struct definition as follows:

```
struct Foo
    x::Int
    y::Complex{Int}
end
```

Now, Cassette can use `metadatatype` to determine type constraints for metadata structures
associated with tagged values of type `Foo`. In psuedo-Julia-code, these metadata structures
might look something like the following for `Foo`:

```
struct IntMeta
    data::metadatatype(Ctx, Int)
    meta::Cassette.NoMetaMeta
end

struct ComplexIntMeta
    data::metadatatype(Ctx, Complex{Int})
    meta::NamedTuple{(:re,:im),Tuple{IntMeta,IntMeta}}
end

struct FooMeta
    data::metadatatype(Ctx, Foo)
    meta::NamedTuple{(:x,:y),Tuple{IntMeta,ComplexIntMeta}
end
```

# Examples

```
julia> Cassette.@context Ctx;

# any value of type `Number` can now be tagged with metadata of type `Number`
julia> Cassette.metadatatype(::Type{<:Ctx}, ::Type{<:Number}) = Number

# any value of type `T<:Number` can now be tagged with metadata of type `T`
julia> Cassette.metadatatype(::Type{<:Ctx}, ::Type{T}) where {T<:Number} = T

# any value of type `T<:Number` can now be tagged with metadata of type `promote_type(T, M)`
# where `M` is the type of the trace-local metadata associated with the context
julia> Cassette.metadatatype(::Type{<:Ctx{M}}, ::Type{T}) where {M<:Number,T<:Number} = promote_type(T, M)
```
"""
metadatatype(::Type{<:Context}, ::DataType) = NoMetaData

#=== metametatype ===#

@static if isdefined(Base, :ismutabletype)
    using Base: ismutabletype
else
    function ismutabletype(@nospecialize(t::Type))
        t = Base.unwrap_unionall(t)
        return isa(t, DataType) && t.mutable
    end
end

@generated function metametatype(::Type{C}, ::Type{T}) where {C<:Context,T}
    if T <: Type || fieldcount(T) == 0
        body = :(NoMetaMeta)
    elseif !(isconcretetype(T))
        body = :(error("cannot call metametatype on non-concrete type ", $T))
    else
        F = ismutabletype(T) ? :Mutable : :Immutable
        ftypes = [:($F{metatype(C, $S)}) for S in _fieldtypes_for_metatype(T)]
        tuplemetatype = :(Tuple{$(ftypes...)})
        if T <: Tuple
            body = tuplemetatype
        else
            fnames = Expr(:tuple, map(Base.Meta.quot, fieldnames(T))...)
            body = :(NamedTuple{$fnames,$tuplemetatype})
        end
    end
    return quote
        $body
    end
end

@generated function metametatype(::Type{C}, ::Type{T}) where {C<:Context,T<:Array}
    return :(Array{metatype(C, $(eltype(T))),$(ndims(T))})
end

function metametatype(::Type{C}, ::Type{Module}) where {C<:Context}
    return ModuleMeta{metadatatype(C, Symbol), metametatype(C, Symbol)}
end

#=== initmetameta ===#

function _metametaexpr(::Type{C}, ::Type{V}, metaexprs::Vector) where {C,V}
    if V <: Type || fieldcount(V) == 0 || (all(x == :NOMETA for x in metaexprs) && doesnotneedmetatype(V))
        return :(NoMetaMeta())
    else
        F = ismutabletype(V) ? :Mutable : :Immutable
        metatypes = [:(metatype(C, $S)) for S in _fieldtypes_for_metatype(V)]
        metaconverts = [:(convert($(metatypes[i]), $(metaexprs[i]))) for i in 1:fieldcount(V)]
        metametafields = [:($F{$(metatypes[i])}($(metaconverts[i]))) for i in 1:fieldcount(V)]
        if !(V <: Tuple)
            fnames = fieldnames(V)
            for i in 1:fieldcount(V)
                metametafields[i] = :($(fnames[i]) = $(metametafields[i]))
            end
        end
        return Expr(:tuple, metametafields...)
    end
end

initmetameta(context::Context, value::Module) = fetch_modulemeta(context, value)

function initmetameta(context::C, value::Array{V}) where {C<:Context,V}
    M = metatype(C, V)
    if M <: typeof(NOMETA)
        return NoMetaMeta()
    else
        return fill!(similar(value, M), NOMETA)
    end
end

@generated function initmetameta(context::C, value::V) where {C<:Context,V}
    return quote
        $(_metametaexpr(C, V, [:NOMETA for i in 1:fieldcount(V)]))
    end
end

#=== initmeta ===#

function initmeta(context::C, value::V, metadata::D) where {C<:Context,V,D}
    return Meta{metadatatype(C, V),metametatype(C, V)}(metadata, initmetameta(context, value))
end

############
# `Tagged` #
############

struct Tagged{T<:Tag,V,D,M,C<:ContextTagged}
    context::C
    value::V
    meta::Meta{D,M}
    function Tagged(context::C, value::V, meta::Meta) where {T<:Tag,V,C<:ContextTagged{T}}
        D = metadatatype(C, V)
        M = metametatype(C, V)
        return new{T,V,D,M,C}(context, value, convert(Meta{D,M}, meta))
    end
end

#=== `Tagged` API ===#

"""
```
tag(value, context::Context, metadata = Cassette.NoMetaData())
```

Return `value` tagged w.r.t. `context`, optionally associating `metadata` with the returned
`Tagged` instance.

Any provided `metadata` must obey the type constraints determined by Cassette's
[`metadatatype`](@ref) method.

Note that `hastagging(typeof(context))` must be `true` for a value to be tagged w.r.t. to
`context`.

See also: [`untag`](@ref), [`enabletagging`](@ref), [`hastagging`](@ref)
"""
function tag(value, context::Context, metadata = NoMetaData())
    return Tagged(context, value, initmeta(context, value, metadata))
end

function tag(value, context::ContextUntagged, metadata = NoMetaData())
    error("cannot `tag` a value w.r.t. a `context` if `!hastagging(typeof(context))`")
end

"""
```
untag(x, context::Context)
```

Return `x` untagged w.r.t. `context` if `istagged(x, context)`, otherwise return
`x` directly.

In other words, `untag(tag(x, context), context) === x` is always `true`.

If `!istagged(x, context)`, then `untag(x, context) === x` is `true`.

See also: [`tag`](@ref), [`istagged`](@ref)
"""
untag(x, context::Context) = untag(x, context.tag)
untag(x::Tagged{T}, tag::T) where {T<:Tag} = x.value
untag(x, ::Union{Tag,Nothing}) = x

"""
```
untagtype(::Type{T}, ::Type{C<:Context})
```

Return `typeof(untag(::T, ::C))`.

In other words, `untagtype(typeof(tag(x, context)), typeof(context)) === typeof(x)` is always
`true`.

If `!istaggedtype(T, C)`, then `untagtype(T, C) === T` is `true`.
"""
untagtype(X::Type, ::Type{C}) where {C<:Context} = untagtype(X, tagtype(C))
untagtype(::Type{<:Tagged{T,V}}, ::Type{T}) where {T<:Tag,V} = V
untagtype(X::Type, ::Type{<:Union{Tag,Nothing}}) = X

"""
```
metadata(x, context::Context)
```

Return the `metadata` attached to `x` if `hasmetadata(x, context)`, otherwise return
`Cassette.NoMetaData()`.

In other words, `metadata(tag(x, context, m), context) === m` is always `true`.

If `!hasmetadata(x, context)`, then `metadata(x, context) === Cassette.NoMetaData()` is
`true`.
"""
metadata(x, context::Context) = metadata(x, context.tag)
metadata(x::Tagged{T}, tag::T) where {T<:Tag} = x.meta.data
metadata(::Any, ::Union{Tag,Nothing}) = NoMetaData()

metameta(x, context::Context) = metameta(x, context.tag)
metameta(x::Tagged{T}, tag::T) where {T<:Tag} = x.meta.meta
metameta(::Any, ::Union{Tag,Nothing}) = NoMetaMeta()

"""
```
istagged(x, context::Context)
```

Return `true` if `x` is tagged w.r.t. `context`, return `false` otherwise.

In other words, `istagged(tag(x, context), context)` is always `true`.

See also: [`tag`](@ref), [`istaggedtype`](@ref)
"""
istagged(x, context::Context) = istagged(x, context.tag)
istagged(x::Tagged{T}, tag::T) where {T<:Tag} = true
istagged(::Any, ::Union{Tag,Nothing}) = false

"""
```
istaggedtype(::Type{T}, ::Type{C<:Context})
```

Return `typeof(istagged(::T, ::C))`.

In other words, `istaggedtype(typeof(tag(x, context)), typeof(context))` is always `true`.

See also: [`tag`](@ref), [`istagged`](@ref)
"""
istaggedtype(X::Type, ::Type{C}) where {C<:Context} = istaggedtype(X, tagtype(C))
istaggedtype(::Type{<:Tagged{T}}, ::Type{T}) where {T<:Tag} = true
istaggedtype(::DataType, ::Type{<:Union{Tag,Nothing}}) = false

"""
```
hasmetadata(x, context::Context)
```

Return `true` if `!isa(metadata(x, context), Cassette.NoMetaData)`, return `false` otherwise.

In other words, `hasmetadata(tag(x, context, m), context)` is always `true` and
`hasmetadata(tag(x, context), context)` is always `false`.

See also: [`metadata`](@ref)
"""
hasmetadata(x, context::Context) = hasmetadata(x, context.tag)
hasmetadata(x, tag::Union{Tag,Nothing}) = !isa(metadata(x, tag), NoMetaData)

hasmetameta(x, context::Context) = hasmetameta(x, context.tag)
hasmetameta(x, tag::Union{Tag,Nothing}) = !isa(metameta(x, tag), NoMetaMeta)

#########################
# Core._apply iteration #
#########################

destructstate(ctx::ContextTagged{T}, state::Tagged{T,<:Tuple}) where {T} = (tagged_getfield(ctx, state, 1), tagged_getfield(ctx, state, 2))
destructstate(ctx, state) = untag(state, ctx)

Base.iterate(t::Tagged) = destructstate(t.context, overdub(t.context, iterate, t))
Base.iterate(t::Tagged, state) = destructstate(t.context, overdub(t.context, iterate, t, state))

##################################
# `tagged_new`/`tagged_splatnew` #
##################################

@generated function tagged_new(context::C, ::Type{T}, args...) where {C<:Context,T}
    argmetaexprs = Any[]
    for i in 1:fieldcount(T)
        if i <= nfields(args) && istaggedtype(args[i], C)
            push!(argmetaexprs, :(args[$i].meta))
        else
            push!(argmetaexprs, :NOMETA)
        end
    end
    untagged_args = [:(untag(args[$i], context)) for i in 1:nfields(args)]
    newexpr = (T <: Tuple) ? Expr(:tuple, untagged_args...) : Expr(:new, T, untagged_args...)
    onlytypeargs = true
    for arg in args
        if !(arg <: Type)
            onlytypeargs = false
            break
        end
    end
    if (all(x == :NOMETA for x in argmetaexprs) && doesnotneedmetatype(T)) || onlytypeargs
        return newexpr
    else
        metametaexpr = _metametaexpr(C, T, argmetaexprs)
        return quote
            M = metatype(C, T)
            return Tagged(context, $newexpr, Meta(NoMetaData(), $metametaexpr))
        end
    end
end

@inline tagged_splatnew(context::Context, T::Type, args) = tagged_new(context, T, args...)

@generated function tagged_new_array(context::C, ::Type{T}, args...) where {C<:Context,T<:Array}
    untagged_args = [:(untag(args[$i], context)) for i in 1:nfields(args)]
    return quote
        return tag($(T)($(untagged_args...)), context)
    end
end

@generated function tagged_new_module(context::C, args...) where {C<:Context}
    if istaggedtype(args[1], C)
        return_expr = quote
            Tagged(context, tagged_module.value,
                   Meta(NoMetaData(),
                        ModuleMeta(args[1].meta, tagged_module.meta.meta.bindings)))
        end
    else
        return_expr = :(tagged_module)
    end
    return quote
        new_module = Module(args...)
        tagged_module = fetch_tagged_module(context, new_module)
        return $return_expr
    end
end

@generated function tagged_new_tuple(context::C, args...) where {C<:Context}
    T = Tuple{[untagtype(args[i], C) for i in 1:nfields(args)]...}
    return quote
        tagged_new(context, $T, args...)
    end
end

# like `tagged_new_tuple`, but will not necessarily tag non-`doesnotneedmeta` elements of `args`
@generated function _tagged_new_tuple_unsafe(context::C, args...) where {C<:Context}
    if all(!istaggedtype(arg, C) for arg in args)
        return quote
            Core.tuple(args...)
        end
    else
        return quote
            tagged_new_tuple(context, args...)
        end
    end
end

##################################
# `tagged_*` intrisic primitives #
##################################

_untag_all(context::Context, a) = untag(a, context)
_untag_all(context::Context, a, b) = (untag(a, context), untag(b, context))
_untag_all(context::Context, a, b, c) = (untag(a, context), untag(b, context), untag(c, context))
_untag_all(context::Context, a, b, c, rest...) = (untag(a, context), untag(b, context), untag(c, context), _untag_all(context, rest...)...)

#=== tagged_nameof ===#

tagged_nameof(context::Context, x) = nameof(untag(x, context))

function tagged_nameof(context::ContextTagged{T}, x::Tagged{T,Module}) where {T}
    name_value = nameof(x.value)
    name_meta = hasmetameta(x, context) ? x.meta.meta.name : NOMETA
    return Tagged(context, name_value, name_meta)
end

#=== tagged_globalref ===#

@inline function tagged_globalref(context::ContextTagged{T},
                                  m::Tagged{T},
                                  name,
                                  primal) where {T}
    if hasmetameta(m, context) && !istagged(primal, context)
        return _tagged_globalref(context, m, name, primal)
    else
        return primal
    end
end

# assume that `context` === `primal` TODO is this assumption valid?
@inline function tagged_globalref(context::ContextTagged{T},
                                  m::Tagged{T},
                                  name,
                                  primal::ContextTagged{T}) where {T}
    return primal
end

@inline function _tagged_globalref(context::ContextTagged{T},
                                   m::Tagged{T},
                                   name,
                                   primal) where {T}
    untagged_name = untag(name, context)
    if isconst(m.value, untagged_name) && doesnotneedmeta(primal)
        # It's very important that this fast path exists and is taken with no runtime
        # overhead; this is the path that will be taken by, for example, access of simple
        # named function bindings.
        return primal
    else
        meta = fetch_bindingmeta!(context, m.value, m.meta.meta.bindings, untagged_name, primal)
        return Tagged(context, primal, meta)
    end
end

#=== tagged_globalref_set_meta! ===#

@inline function tagged_globalref_set_meta!(context::ContextTagged{T}, m::Tagged{T}, name::Symbol, primal) where {T}
    bindingmeta = _fetch_bindingmeta!(context, m.value, m.meta.meta.bindings, name)
    bindingmeta.data = istagged(primal, context) ? primal.meta : NOMETA
    return nothing
end

#=== tagged_getfield ===#

tagged_getfield(context::Context, x, name) = tagged_getfield(context, x, name, false)

tagged_getfield(context::ContextTagged{T}, x, name, boundscheck) where {T} = getfield(x, untag(name, context), untag(boundscheck, context))

function tagged_getfield(context::ContextTagged{T}, x::Tagged{T}, name, boundscheck) where {T}
    untagged_boundscheck = untag(boundscheck, context)
    untagged_name = untag(name, context)
    x_value = untag(x, context)
    y_value = getfield(x_value, untagged_name, untagged_boundscheck)
    if isa(x_value, Module)
        return tagged_globalref(context, x, untagged_name, getfield(x_value, untagged_name))
    elseif hasmetameta(x, context)
        y_meta = load(getfield(x.meta.meta, untagged_name, untagged_boundscheck))
        doesnotneedmeta(y_value) && y_meta === NOMETA && return y_value
        return Tagged(context, y_value, y_meta)
    elseif doesnotneedmeta(y_value)
        return y_value
    else
        return Tagged(context, y_value, NOMETA)
    end
end

#=== tagged_setfield! ===#

tagged_setfield!(context::ContextTagged{T}, x, name, y) where {T} = setfield!(x, untag(name, context), y)

function tagged_setfield!(context::ContextTagged{T}, x::Tagged{T}, name, y) where {T}
    untagged_name = untag(name, context)
    y_value = untag(y, context)
    y_meta = istagged(y, context) ? y.meta : NOMETA
    setfield!(x.value, untagged_name, y_value)
    if hasmetameta(x, context)
        store!(getfield(x.meta.meta, untagged_name), y_meta)
    end
    return y
end

#=== tagged_arrayref ===#

function tagged_arrayref(context::ContextTagged{T}, boundscheck, x, inds...) where {T}
    return Core.arrayref(untag(boundscheck, context), x, _untag_all(context, inds...)...)
end

function tagged_arrayref(context::ContextTagged{T}, boundscheck, x::Tagged{T}, inds...) where {T}
    untagged_boundscheck = untag(boundscheck, context)
    untagged_inds = _untag_all(context, inds...)
    y_value = Core.arrayref(untagged_boundscheck, untag(x, context), untagged_inds...)
    if hasmetameta(x, context)
        y_meta = Core.arrayref(untagged_boundscheck, x.meta.meta, untagged_inds...)
    else
        y_meta = NOMETA
    end
    return Tagged(context, y_value, y_meta)
end

#=== tagged_arrayset ===#

function tagged_arrayset(context::ContextTagged{T}, boundscheck, x, y, inds...) where {T}
    return Core.arrayset(untag(boundscheck, context), x, y, _untag_all(context, inds...)...)
end

function tagged_arrayset(context::ContextTagged{T}, boundscheck, x::Tagged{T}, y, inds...) where {T}
    untagged_boundscheck = untag(boundscheck, context)
    untagged_inds = _untag_all(context, inds...)
    y_value = untag(y, context)
    y_meta = istagged(y, context) ? y.meta : NOMETA
    Core.arrayset(untagged_boundscheck, untag(x, context), y_value, untagged_inds...)
    if hasmetameta(x, context)
        Core.arrayset(untagged_boundscheck, x.meta.meta, convert(eltype(x.meta.meta), y_meta), untagged_inds...)
    end
    return x
end

#=== tagged_growbeg! ===#

tagged_growbeg!(context::ContextTagged{T}, x, delta) where {T} = Base._growbeg!(x, untag(delta, context))

function tagged_growbeg!(context::ContextTagged{T}, x::Tagged{T}, delta) where {T}
    delta_untagged = untag(delta, context)
    Base._growbeg!(x.value, delta_untagged)
    if hasmetameta(x, context)
        Base._growbeg!(x.meta.meta, delta_untagged)
        x.meta.meta[1:delta_untagged] .= Ref(NOMETA)
    end
    return nothing
end

#=== tagged_growend! ===#

tagged_growend!(context::ContextTagged{T}, x, delta) where {T} = Base._growend!(x, untag(delta, context))

function tagged_growend!(context::ContextTagged{T}, x::Tagged{T}, delta) where {T}
    delta_untagged = untag(delta, context)
    Base._growend!(x.value, delta_untagged)
    if hasmetameta(x, context)
        old_length = length(x.meta.meta)
        Base._growend!(x.meta.meta, delta_untagged)
        x.meta.meta[(old_length + 1):(old_length + delta_untagged)] .= Ref(NOMETA)
    end
    return nothing
end

#=== tagged_growat! ===#

function tagged_growat!(context::ContextTagged{T}, x, i, delta) where {T}
    return Base._growat!(x, untag(i, context), untag(delta, context))
end

function tagged_growat!(context::ContextTagged{T}, x::Tagged{T}, i, delta) where {T}
    i_untagged = untag(i, context)
    delta_untagged = untag(delta, context)
    Base._growat!(x.value, i_untagged, delta_untagged)
    if hasmetameta(x, context)
        Base._growat!(x.meta.meta, i_untagged, delta_untagged)
        x.meta.meta[i_untagged:(i_untagged + delta_untagged - 1)] .= Ref(NOMETA)
    end
    return nothing
end

#=== tagged_deletebeg! ===#

tagged_deletebeg!(context::ContextTagged{T}, x, delta) where {T} = Base._deletebeg!(x, untag(delta, context))

function tagged_deletebeg!(context::ContextTagged{T}, x::Tagged{T}, delta) where {T}
    delta_untagged = untag(delta, context)
    Base._deletebeg!(x.value, delta_untagged)
    hasmetameta(x, context) && Base._deletebeg!(x.meta.meta, delta_untagged)
    return nothing
end

#=== tagged_deleteend! ===#

tagged_deleteend!(context::ContextTagged{T}, x, delta) where {T} = Base._deleteend!(x, untag(delta, context))

function tagged_deleteend!(context::ContextTagged{T}, x::Tagged{T}, delta) where {T}
    delta_untagged = untag(delta, context)
    Base._deleteend!(x.value, delta_untagged)
    hasmetameta(x, context) && Base._deleteend!(x.meta.meta, delta_untagged)
    return nothing
end

#=== tagged_deleteat! ===#

function tagged_deleteat!(context::ContextTagged{T}, x, i, delta) where {T}
    return Base._deleteat!(x, untag(i, context), untag(delta, context))
end

function tagged_deleteat!(context::ContextTagged{T}, x::Tagged{T}, i, delta) where {T}
    i_untagged = untag(i, context)
    delta_untagged = untag(delta, context)
    Base._deleteat!(x.value, i_untagged, delta_untagged)
    hasmetameta(x, context) && Base._deleteat!(x.meta.meta, i_untagged, delta_untagged)
    return nothing
end

#=== tagged_typeassert ===#

function tagged_typeassert(context::ContextTagged{T}, x, typ) where {T}
    return Core.typeassert(x, untag(typ, context))
end

function tagged_typeassert(context::ContextTagged{T}, x::Tagged{T}, typ) where {T}
    untagged_result = Core.typeassert(untag(x, context), untag(typ, context))
    return Tagged(context, untagged_result, x.meta)
end

#=== tagged_sitofp ===#

function tagged_sitofp(context::ContextTagged{T}, F, x) where {T}
    return Base.sitofp(untag(F, context), x)
end

function tagged_sitofp(context::ContextTagged{T}, F, x::Tagged{T}) where {T}
    return Tagged(context, Base.sitofp(untag(F, context), x.value), x.meta)
end

#=== tagged_sle_int ===#

tagged_sle_int(context::Context, x, y) = Base.sle_int(untag(x, context), untag(y, context))

###################
# Pretty Printing #
###################

Base.show(io::IO, meta::Union{NoMetaMeta,NoMetaData}) = print(io, "_")

function Base.show(io::IO, meta::Meta)
    if isa(meta.data, NoMetaData) && isa(meta.meta, NoMetaMeta)
        print(io, "_")
    else
        if isa(meta.meta, NamedTuple)
            tmp = IOBuffer()
            write(tmp, "(")
            i = 1
            for (k, v) in pairs(meta.meta)
                print(tmp, k, " = ", load(v))
                if i == length(meta.meta)
                    print(tmp, ")")
                else
                    print(tmp, ", ")
                    i += 1
                end
            end
            metametastr = String(take!(tmp))
        elseif isa(meta.meta, Tuple)
            tmp = IOBuffer()
            write(tmp, "(")
            i = 1
            for v in meta.meta
                print(tmp, load(v))
                if i == length(meta.meta)
                    print(tmp, ")")
                else
                    print(tmp, ", ")
                    i += 1
                end
            end
            metametastr = String(take!(tmp))
        else
            metametastr = sprint(show, meta.meta)
        end
        print(io, "Meta(", meta.data, ", ", metametastr, ")")
    end
end

Base.show(io::IO, x::Tagged) = print(io, "Tagged(", x.context.tag, ", ", x.value, ", ", x.meta, ")")

Base.show(io::IO, ::Tag{N,X,E}) where {N,X,E} = print(io, "Tag{", N, ",", X, ",", E, "}()")


text:
function DiffEqBase.addsteps!(k,t,uprev,u,dt,f,p,cache::Union{Rosenbrock23ConstantCache,Rosenbrock32ConstantCache},always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if length(k)<2 || always_calc_begin
    @unpack tf,uf,d = cache
    γ = dt*d
    dT = ForwardDiff.derivative(tf, t)
    if typeof(uprev) <: AbstractArray
      J = ForwardDiff.jacobian(uf, uprev)
      W = I - γ*J
    else
      J = ForwardDiff.derivative(uf, uprev)
      W = 1 - γ*J
    end
    k₁ = W\(f(uprev,p,t) + dt*d*dT)
    f₁ = f(uprev+dt*k₁/2,p,t+dt/2)
    k₂ = W\(f₁-k₁) + k₁
    copyat_or_push!(k,1,k₁)
    copyat_or_push!(k,2,k₂)
  end
  nothing
end

function DiffEqBase.addsteps!(k,t,uprev,u,dt,f,p,cache::Union{Rosenbrock23Cache,Rosenbrock32Cache},always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if length(k)<2 || always_calc_begin
    @unpack k₁,k₂,k₃,du1,du2,f₁,vectmp,vectmp2,vectmp3,fsalfirst,fsallast,dT,J,W,tmp,uf,tf,linsolve_tmp,linsolve_tmp_vec = cache
    @unpack c₃₂,d = cache.tab
    uidx = eachindex(uprev)

    # Assignments
    sizeu  = size(u)
    #mass_matrix = integrator.f.mass_matrix
    γ = dt*d
    dto2 = dt/2
    dto6 = dt/6

    #@. linsolve_tmp = @muladd fsalfirst + γ*dT
    @tight_loop_macros for i in uidx
      @inbounds linsolve_tmp[i] = @muladd fsalfirst[i] + γ*dT[i]
    end

    if DiffEqBase.has_invW(f)
      f.invW(W,u,p,γ,t) # W == inverse W
      mul!(vectmp,W,linsolve_tmp_vec)
    else
      ### Jacobian does not need to be re-evaluated after an event
      ### Since it's unchanged
      for i in 1:length(u), j in 1:length(u)
        @inbounds W[i,j] = @muladd I[i,j]-γ*J[i,j]
      end
      cache.linsolve(vectmp,W,linsolve_tmp_vec,true)
    end

    recursivecopy!(k₁,reshape(vectmp,size(u)...))
    @. tmp = uprev + dto2*k₁
    f(f₁,tmp,p,t+dto2)

    #if mass_matrix == I
      tmp .= k₁
    #else
    #  mul!(tmp,mass_matrix,k₁)
    #end

    @. linsolve_tmp = f₁ - tmp
    if DiffEqBase.has_invW(f)
      mul!(vectmp2, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp2, W, linsolve_tmp_vec)
    end

    tmp2 = reshape(vectmp2, sizeu...)
    @. k₂ = tmp2 + k₁

    copyat_or_push!(k,1,k₁)
    copyat_or_push!(k,2,k₂)
  end
  nothing
end

function DiffEqBase.addsteps!(k,t,uprev,u,dt,f,p,cache::Rodas4ConstantCache,always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if length(k)<2 || always_calc_begin
    @unpack tf,uf = cache
    @unpack a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,C21,C31,C32,C41,C42,C43,C51,C52,C53,C54,C61,C62,C63,C64,C65,gamma,c2,c3,c4,d1,d2,d3,d4 = cache.tab

    # Precalculations
    dtC21 = C21/dt
    dtC31 = C31/dt
    dtC32 = C32/dt
    dtC41 = C41/dt
    dtC42 = C42/dt
    dtC43 = C43/dt
    dtC51 = C51/dt
    dtC52 = C52/dt
    dtC53 = C53/dt
    dtC54 = C54/dt
    dtC61 = C61/dt
    dtC62 = C62/dt
    dtC63 = C63/dt
    dtC64 = C64/dt
    dtC65 = C65/dt

    dtd1 = dt*d1
    dtd2 = dt*d2
    dtd3 = dt*d3
    dtd4 = dt*d4
    dtgamma = dt*gamma

    # Time derivative
    tf.u = uprev
    dT = ForwardDiff.derivative(tf, t)

    # Jacobian
    uf.t = t
    if typeof(uprev) <: AbstractArray
      J = ForwardDiff.jacobian(uf, uprev)
      W = I/dtgamma - J
    else
      J = ForwardDiff.derivative(uf, uprev)
      W = 1/dtgamma - J
    end

    du = f(uprev, p, t)

    linsolve_tmp =  du + dtd1*dT

    k1 = W\linsolve_tmp
    u = uprev  + a21*k1
    du = f(u, p, t+c2*dt)

    linsolve_tmp =  du + dtd2*dT + dtC21*k1

    k2 = W\linsolve_tmp
    u = uprev  + a31*k1 + a32*k2
    du = f(u, p, t+c3*dt)

    linsolve_tmp =  du + dtd3*dT + (dtC31*k1 + dtC32*k2)

    k3 = W\linsolve_tmp
    u = uprev  + a41*k1 + a42*k2 + a43*k3
    du = f(u, p, t+c4*dt)

    linsolve_tmp =  du + dtd4*dT + (dtC41*k1 + dtC42*k2 + dtC43*k3)

    k4 = W\linsolve_tmp
    u = uprev  + a51*k1 + a52*k2 + a53*k3 + a54*k4
    du = f(u, p, t+dt)

    linsolve_tmp =  du + (dtC52*k2 + dtC54*k4 + dtC51*k1 + dtC53*k3)

    k5 = W\linsolve_tmp

    @unpack h21,h22,h23,h24,h25,h31,h32,h33,h34,h35 = cache.tab
    k₁ =  h21*k1 + h22*k2 + h23*k3 + h24*k4 + h25*k5
    k₂ =  h31*k1 + h32*k2 + h33*k3 + h34*k4 + h35*k5
    copyat_or_push!(k,1,k₁)
    copyat_or_push!(k,2,k₂)
  end
  nothing
end

function DiffEqBase.addsteps!(k,t,uprev,u,dt,f,p,cache::Rodas4Cache,always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if length(k)<2 || always_calc_begin

    @unpack du,du1,du2,tmp,vectmp,vectmp2,vectmp3,vectmp4,vectmp5,vectmp6,dT,J,W,uf,tf,linsolve_tmp,linsolve_tmp_vec,jac_config,fsalfirst = cache
    @unpack a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,C21,C31,C32,C41,C42,C43,C51,C52,C53,C54,C61,C62,C63,C64,C65,gamma,c2,c3,c4,d1,d2,d3,d4 = cache.tab

    # Assignments
    sizeu  = size(u)
    uidx = eachindex(uprev)
    #mass_matrix = integrator.f.mass_matrix
    mass_matrix = I
    atmp = du # does not work with units - additional unitless array required!

    # Precalculations
    dtC21 = C21/dt
    dtC31 = C31/dt
    dtC32 = C32/dt
    dtC41 = C41/dt
    dtC42 = C42/dt
    dtC43 = C43/dt
    dtC51 = C51/dt
    dtC52 = C52/dt
    dtC53 = C53/dt
    dtC54 = C54/dt
    dtC61 = C61/dt
    dtC62 = C62/dt
    dtC63 = C63/dt
    dtC64 = C64/dt
    dtC65 = C65/dt

    dtd1 = dt*d1
    dtd2 = dt*d2
    dtd3 = dt*d3
    dtd4 = dt*d4
    dtgamma = dt*gamma

    @tight_loop_macros for i in uidx
      @inbounds linsolve_tmp[i] = @muladd fsalfirst[i] + dtgamma*dT[i]
    end

    if DiffEqBase.has_invW(f)
      f.invW_t(W,u,p,dtgamma,t) # W == inverse W
    else
      ### Jacobian does not need to be re-evaluated after an event
      ### Since it's unchanged
      for i in 1:length(u), j in 1:length(u)
        @inbounds W[i,j] = @muladd I[i,j]/dtgamma-J[i,j]
      end
    end

    if DiffEqBase.has_invW(f)
      mul!(vectmp, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp, W, linsolve_tmp_vec, true)
    end

    k1 = reshape(vectmp, sizeu...)
    @. tmp = uprev + a21*k1
    f( du,  tmp, p, t+c2*dt)

    if mass_matrix == I
      @. linsolve_tmp = du + dtd2*dT + dtC21*k1
    else
      @. du1 = dtC21*k1
      mul!(du2,mass_matrix,du1)
      @. linsolve_tmp = du + dtd2*dT + du2
    end

    if DiffEqBase.has_invW(f)
      mul!(vectmp2, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp2, W, linsolve_tmp_vec)
    end

    k2 = reshape(vectmp2, sizeu...)
    @. tmp = uprev + a31*k1 + a32*k2
    f( du,  tmp, p, t+c3*dt)

    if mass_matrix == I
      @. linsolve_tmp = du + dtd3*dT + (dtC31*k1 + dtC32*k2)
    else
      @. du1 = dtC31*k1 + dtC32*k2
      mul!(du2,mass_matrix,du1)
      @. linsolve_tmp = du + dtd3*dT + du2
    end

    if DiffEqBase.has_invW(f)
      mul!(vectmp3, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp3, W, linsolve_tmp_vec)
    end

    k3 = reshape(vectmp3, sizeu...)
    @. tmp = uprev + a41*k1 + a42*k2 + a43*k3
    f( du,  tmp, p, t+c4*dt)

    if mass_matrix == I
      @. linsolve_tmp = du + dtd4*dT + (dtC41*k1 + dtC42*k2 + dtC43*k3)
    else
      @. du1 = dtC41*k1 + dtC42*k2 + dtC43*k3
      mul!(du2,mass_matrix,du1)
      @. linsolve_tmp = du + dtd4*dT + du2
    end

    if DiffEqBase.has_invW(f)
      mul!(vectmp4, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp4, W, linsolve_tmp_vec)
    end

    k4 = reshape(vectmp4, sizeu...)
    @. tmp = uprev + a51*k1 + a52*k2 + a53*k3 + a54*k4
    f( du,  tmp, p, t+dt)

    if mass_matrix == I
      @. linsolve_tmp = du + (dtC52*k2 + dtC54*k4 + dtC51*k1 + dtC53*k3)
    else
      @. du1 = dtC52*k2 + dtC54*k4 + dtC51*k1 + dtC53*k3
      mul!(du2,mass_matrix,du1)
      @. linsolve_tmp = du + du2
    end

    if DiffEqBase.has_invW(f)
      mul!(vectmp5, W, linsolve_tmp_vec)
    else
      cache.linsolve(vectmp5, W, linsolve_tmp_vec)
    end

    k5 = reshape(vectmp5, sizeu...)

    k_tmp = reshape(vectmp6, sizeu...)

    @unpack h21,h22,h23,h24,h25,h31,h32,h33,h34,h35 = cache.tab
    @. k_tmp = h21*k1 + h22*k2 + h23*k3 + h24*k4 + h25*k5
    copyat_or_push!(k,1,copy(k_tmp))

    @. k_tmp = h31*k1 + h32*k2 + h33*k3 + h34*k4 + h35*k5
    copyat_or_push!(k,2,copy(k_tmp))

  end
  nothing
end


text:
################################################################
import Pkg
"""
Should be removed once Pkg allows adding test dependencies to the current environment

Used in Github Action workflow yaml file
"""
function addtestdep()
    if isfile("test/Test.toml")
        testToml = Pkg.Types.parse_toml("test/Test.toml")
    elseif isfile("test/Project.toml")
        testToml = Pkg.Types.parse_toml("test/Project.toml")
    else
        error("please add a Test.toml or Project.toml to the /test directory for test dependencies")
    end

    for (name, uuid) in testToml["deps"]
        Pkg.add(Pkg.PackageSpec(name = name, uuid = uuid))
    end
end


text:
######### sensis EUR du 11 oct 2017
using CSV, DataFrames
using FileIO
using Optim
using VegaLite
using Rotolo
using Base.Markdown

global_theme = """
	text-align: justify;
	border-style:solid;
	padding: 5px
  """

methods(Media.render, (Rotolo.Juno, Any))

methods(Media.render, (Atom.Inline, Void))

methods(Media.render, (Atom.Editor, Any))

methods(Media.render, (Atom.Inline, Any))

@session MontantsEquiv6 global_theme
# @endsession

@redirect Markdown.MD
sleep(3)

@redirect Katex

Rotolo.currentSession.redirected_types

@container header style=>"text-align: center;"

@redirect Base.Markdown.MD

md"# Montants équivalents"

typeof(md"# Montants équivalents")

methods(Juno.render, (Juno.Editor, Base.Markdown.MD))
methods(Atom.render, (Atom.Editor, Base.Markdown.MD))

@style md"*From Wikipedia, the free encyclopedia*" style=>"font-size:x-small"

Katex("f(x,y)=(x^2+y-11)^2+(x+y^2-7)^2")

5+6

@redirect Int

5+65

# delete!(Rotolo.currentSession.redirected_types, Int)

"abcd"

@redirect String

"abcd"

Rotolo.loadmsg(Katex)
Katex("f(x,y)=(x^2+y-11)^2+(x+y^2-7)^2")

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

fvar(s) = s' * covmat * s
dfvar(s) = covmat * s + (s' * covmat)'

function minfunc(initial::Vector{Float64}, vidx::Int64, h::Float64)
  nv = copy(initial)
  nv[vidx] += h
  fvar(nv)
end

function stephedge1(minfunc::Function, initial::Vector, mask::BitArray)
  ovar = round(Int64, minfunc(initial, 1, 0.))

  minzer = zeros(length(initial))
  minmum = [ 1e20 for i in 1:length(initial) ]
  for i in collect(1:length(initial))[mask]
    res = optimize(x -> minfunc(initial, i, x), -1e6, 1e6 )
    minzer[i] = Optim.minimizer(res)
    minmum[i] = Optim.minimum(res)
  end
  vidx = indmin(minmum)
  amnt = round(Int64, minzer[vidx])
  nvar = round(Int64, minmum[vidx])
  println("old var = $ovar, mat $(sensimat[vidx]) ($vidx), amount = $amnt, new var. = $nvar")

  newvec = copy(initial)
  newvec[vidx] += minzer[vidx]
  newmask = copy(mask)
  newmask[vidx] = false
  vidx, newvec, newmask
end

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

sensis = CSV.read(IOBuffer("""
  Mats	EUR.EIB.6M	USD.LIB.3M
  3M	-1202	4589
  6M	-8415	4684
  1Y	-1626	-13340
  18M	-15253	17372
  2Y	-9917	15681
  3Y	5901	-11285
  4Y	3657	-15135
  5Y	65133	24432
  6Y	-30523	-16872
  7Y	-41440	-12484
  8Y	49300	-1001
  9Y	-5135	-1847
  10Y	-21788	6180
  11Y	11166	-9559
  12Y	17216	-4027
  13Y	-2389	1057
  14Y	-11172	2661
  15Y	2677	1759
  16Y	-6815	2324
  17Y	-14232	5430
  18Y	4067	5638
  19Y	23327	645
  20Y	-18821	-118
  21Y	-21295	0
  22Y	-6	0
  23Y	-16188	0
  24Y	-2756	0
  25Y	50119	0
  26Y	4388	0
  27Y	-364	0
  28Y	-19	0
  """), delim='\t')
sensimat = Vector(String.(sensis[:,1]))

###########  EUR  #########################################

dat = load("c:/temp/histovarEUR.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]

midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat) # 1.86e-9 , 3.61e-7

sensivec = Vector{Float64}(sensis[Symbol("EUR.EIB.6M")])
initial, mask = copy(sensivec), trues(length(sensivec))
_, initial, mask = stephedge1(minfunc, initial, mask)
# old var = 50, mat 5Y (8), amount = -31042, new var. = 12
# old var = 12, mat 18M (4), amount = 20089, new var. = 6
# old var = 6, mat 6M (2), amount = 7675, new var. = 5
# old var = 5, mat 4Y (7), amount = -5172, new var. = 4
# old var = 4, mat 2Y (5), amount = 3887, new var. = 4
# old var = 4, mat 28Y (31), amount = -3262, new var. = 4
# old var = 4, mat 7Y (10), amount = 1876, new var. = 4
# old var = 4, mat 3Y (6), amount = -2555, new var. = 3
# old var = 3, mat 17Y (20), amount = 1735, new var. = 3
# old var = 3, mat 27Y (30), amount = -1825, new var. = 3
# old var = 3, mat 16Y (19), amount = 1494, new var. = 3
# old var = 3, mat 25Y (28), amount = -1546, new var. = 3
# old var = 3, mat 8Y (11), amount = 1103, new var. = 3

std₀ = sqrt(50) * 10000 # 71k€
sqrt(12) * 10000        # 35k€
sqrt(6) * 10000         # 25k€

-31042 * 10000 / 5 / 0.98 # 5Y  prêt de 63 M€
20089 * 10000 / 1.5       # 18M emprunt de 134 M€
7675 * 10000 / 0.5       # 6M emprunt de 154 M€

using VegaLite

pd = DataFrame(mats = sensimat, std= sqrt.(diag(covmat)) * 10000)
pd |> markbar() |>
  encoding(xordinal(field=:mats, vlscale(domain=sensimat)),
  yquantitative(field=:std))

###########  USD  #########################################

dat = load("c:/temp/histovarUSD.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]

cidx = isnull.(sensis[3]) .== false
sensimat = Vector{String}(sensis[cidx,1])

midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat) # 1.86e-9 , 3.61e-7

pd = DataFrame(mats = sensimat, std= sqrt.(diag(covmat)) * 10000)
pd |> markbar() |>
  encoding(xordinal(field=:mats, vlscale(domain=sensimat)),
  yquantitative(field=:std))

sensivec = Vector{Float64}(sensis[cidx, Symbol("USD.LIB.3M")])
initial, mask = copy(sensivec), trues(length(sensivec))
_, initial, mask = stephedge1(minfunc, initial, mask)

# old var = 32, mat 18M (4), amount = -13699, new var. = 25
# old var = 25, mat 7Y (10), amount = 11879, new var. = 16
# old var = 16, mat 2Y (5), amount = -6602, new var. = 14
# old var = 14, mat 1Y (3), amount = 6651, new var. = 12
# old var = 12, mat 3M (1), amount = -5328, new var. = 9
# old var = 9, mat 3Y (6), amount = 4410, new var. = 8
# old var = 8, mat 6M (2), amount = -3390, new var. = 7
# old var = 7, mat 20Y (23), amount = -3719, new var. = 6
# old var = 6, mat 8Y (11), amount = 3072, new var. = 5
# old var = 5, mat 21Y (24), amount = -2938, new var. = 5
# old var = 5, mat 4Y (7), amount = 2501, new var. = 5

std₀ = sqrt(32) * 10000 # 57k€
sqrt(25) * 10000        # 50k€
sqrt(16) * 10000         # 40k€
sqrt(14) * 10000         # 37k€

-13699 * 10000 / 1.5  # 18M  prêt de 91M€
11879 * 10000 / 7     # 7Y   emprunt de 17 M€
-6602 * 10000 / 2     # 2Y   prêt de 33 M€

#############  find simplified equivalment pose   ############################

# subset of sensis to minimize
sensimat = Vector{String}(sensis[1])
sensimat2 = String["3M", "6M", "1Y", "3Y", "7Y", "12Y"]
mmask = indexin(sensimat2, sensimat)
sensi₀ = Vector{Float64}(sensis[3])

dat = load("c:/temp/histovarUSD.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]
midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat)

sensi₁ = copy(sensi₀)
function minfunc(x::Vector{Float64})
  sensi₁[mmask] = x
  sensi₁' * covmat * sensi₁
end

result = optimize(minfunc, zeros(length(mmask)), BFGS())
Optim.minimizer(result)
Optim.minimum(result)

sensi₁ - sensi₀

sqrt(sensi₀' * covmat * sensi₀) * 10000 # 56 k€
sqrt(sensi₁' * covmat * sensi₁) * 10000 # 37 k€

pdtfr = vcat(DataFrame(mat=sensimat,  sensi=sensi₀, src = "orig"),
             DataFrame(mat=sensimat2, sensi=(sensi₀-sensi₁)[mmask], src="simplified"))

pdtfr |>
  markbar() |>
  encoding(columnordinal(field=:mat),
           xordinal(field=:src),
           colornominal(field=:src),
           yquantitative(field=:sensi))

max_stars_repo_path: test/tests nov 2017.jl

max_issues_repo_path: test/tests nov 2017.jl

max_forks_repo_path: test/tests nov 2017.jl


text:
# Autogenerated wrapper script for lib4ti2_jll for x86_64-linux-musl-cxx11
export circuits, exe4ti2gmp, exe4ti2int32, exe4ti2int64, genmodel, gensymm, graver, groebner, hilbert, lib4ti2, lib4ti2int32, lib4ti2int64, libzsolve, markov, minimize, normalform, output, ppi, qsolve, rays, walk, zbasis, zsolve

using GMP_jll
using GLPK_jll
JLLWrappers.@generate_wrapper_header("lib4ti2")
JLLWrappers.@declare_file_product(circuits)
JLLWrappers.@declare_executable_product(exe4ti2gmp)
JLLWrappers.@declare_executable_product(exe4ti2int32)
JLLWrappers.@declare_executable_product(exe4ti2int64)
JLLWrappers.@declare_executable_product(genmodel)
JLLWrappers.@declare_executable_product(gensymm)
JLLWrappers.@declare_file_product(graver)
JLLWrappers.@declare_file_product(groebner)
JLLWrappers.@declare_file_product(hilbert)
JLLWrappers.@declare_library_product(lib4ti2, "lib4ti2gmp.so.0")
JLLWrappers.@declare_library_product(lib4ti2int32, "lib4ti2int32.so.0")
JLLWrappers.@declare_library_product(lib4ti2int64, "lib4ti2int64.so.0")
JLLWrappers.@declare_library_product(libzsolve, "libzsolve.so.0")
JLLWrappers.@declare_file_product(markov)
JLLWrappers.@declare_file_product(minimize)
JLLWrappers.@declare_file_product(normalform)
JLLWrappers.@declare_executable_product(output)
JLLWrappers.@declare_executable_product(ppi)
JLLWrappers.@declare_file_product(qsolve)
JLLWrappers.@declare_file_product(rays)
JLLWrappers.@declare_file_product(walk)
JLLWrappers.@declare_file_product(zbasis)
JLLWrappers.@declare_executable_product(zsolve)
function __init__()
    JLLWrappers.@generate_init_header(GMP_jll, GLPK_jll)
    JLLWrappers.@init_file_product(
        circuits,
        "bin/circuits",
    )

    JLLWrappers.@init_executable_product(
        exe4ti2gmp,
        "bin/4ti2gmp",
    )

    JLLWrappers.@init_executable_product(
        exe4ti2int32,
        "bin/4ti2int32",
    )

    JLLWrappers.@init_executable_product(
        exe4ti2int64,
        "bin/4ti2int64",
    )

    JLLWrappers.@init_executable_product(
        genmodel,
        "bin/genmodel",
    )

    JLLWrappers.@init_executable_product(
        gensymm,
        "bin/gensymm",
    )

    JLLWrappers.@init_file_product(
        graver,
        "bin/graver",
    )

    JLLWrappers.@init_file_product(
        groebner,
        "bin/groebner",
    )

    JLLWrappers.@init_file_product(
        hilbert,
        "bin/hilbert",
    )

    JLLWrappers.@init_library_product(
        lib4ti2,
        "lib/lib4ti2gmp.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        lib4ti2int32,
        "lib/lib4ti2int32.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        lib4ti2int64,
        "lib/lib4ti2int64.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libzsolve,
        "lib/libzsolve.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_file_product(
        markov,
        "bin/markov",
    )

    JLLWrappers.@init_file_product(
        minimize,
        "bin/minimize",
    )

    JLLWrappers.@init_file_product(
        normalform,
        "bin/normalform",
    )

    JLLWrappers.@init_executable_product(
        output,
        "bin/output",
    )

    JLLWrappers.@init_executable_product(
        ppi,
        "bin/ppi",
    )

    JLLWrappers.@init_file_product(
        qsolve,
        "bin/qsolve",
    )

    JLLWrappers.@init_file_product(
        rays,
        "bin/rays",
    )

    JLLWrappers.@init_file_product(
        walk,
        "bin/walk",
    )

    JLLWrappers.@init_file_product(
        zbasis,
        "bin/zbasis",
    )

    JLLWrappers.@init_executable_product(
        zsolve,
        "bin/zsolve",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
"""
	mut_crossmap_to_tree!(t::Tree, json::AbstractString, segment;
		cmkey=:cmmuts,
		selfmutkey=:selfmuts,
		jsonkey=:aa_muts,
		muttype=Char,
	)

Add mutation crossmapping `json` file from augur to `t`. Data is added to `n.data.dat[cmkey][s]` for each tree node `n` and segment `s != segment` (read from the json file).
Mutations for key `segment` are added to `n.data.dat[selfmutkey]`.
"""
function mut_crossmap_to_tree!(t::Tree, json::AbstractString, segment;
		cmkey=:cmmuts,
		selfmutkey = :selfmuts,
		jsonkey=:aa_muts,
		muttype=Char,
	)
	cm = JSON3.read(read(json, String), Dict)["nodes"]
	# Look for `jsonkey` in cm and map it to tree
	TreeTools.map_dict_to_tree!(t, cm, key=jsonkey)
	for n in values(t.lnodes)
		n.data.dat[cmkey] = Dict(s => [TreeTools.parse_mutation(m, muttype) for m in M] for (s,M) in n.data.dat[jsonkey])
		n.data.dat[selfmutkey] = n.data.dat[cmkey][segment]
		delete!(n.data.dat[cmkey], segment)
		delete!(n.data.dat, jsonkey)
	end
end

"""
	fasta_crossmap_to_trees!(f::Function, trees::Dict, selfkey=:selfseq, cmkey=:cmseq)

Store sequences of several fasta files onto trees. For each pair `(s, t)` of `trees`:
- Sequences in fasta file `f(s)` are stored on leaves of `t` under `n.data.dat[selfkey]`
- For each other key `σ` of `trees`, sequences in fasta file `f(σ)` are stored on leaves of `t` under `n.data.dat[cmkey][σ]`.
"""
function fasta_crossmap_to_trees!(f::Function, trees::Dict, selfkey=:selfseq, cmkey=:cmseq)
	for (sref,t) in trees
		TreeTools.fasta2tree!(t, f(sref), selfkey)
		for s in Iterators.filter(!=(sref), keys(trees))
			TreeTools.fasta2tree!(t, f(s), (cmkey, s))
		end
	end
end

"""
	crossmap_sequences!(trees::Dict, selfkey = :selfseq, cmkey = :cmseq)

Map sequences of the leaves of each tree to the leaves of all others.
"""
function crossmap_sequences!(trees::Dict, selfkey = :selfseq, cmkey = :cmseq)
	for (sref,tref) in trees
		for n in leaves(tref)
			n.data.dat[cmkey] = Dict()
			for s in Iterators.filter(!=(sref), keys(trees))
				n.data.dat[cmkey][s] = trees[s].lleaves[n.label].data.dat[selfkey]
			end
		end
	end
end

"""
	ancestral_sequences!(trees::Dict;
		self = true,
		selfkey_seq = :selfseq,
		selfkey_muts = :selfmuts,
		crossmapped = true,
		cmkey_seq = :cmseq,
		cmkey_muts = :cmmuts,
	)

Compute ancestral *mutations* using Fitch algorithm. (*Note*: name should be changed).
"""
function ancestral_sequences!(trees::Dict;
		self = true,
		selfkey_seq = :selfseq,
		selfkey_muts = :selfmuts,
		crossmapped = true,
		cmkey_seq = :cmseq,
		cmkey_muts = :cmmuts,
	)
	for (sref,tref) in trees
		# self sequences and mutations
		self && self_ancestral_states!(tref, selfkey_seq, selfkey_muts)
		# crossmapped sequences and mutations
		crossmapped && crossmapped_ancestral_states!(trees, sref, cmkey_seq, cmkey_muts)
	end
end

function self_ancestral_states!(tree, seqkey=:selfseq, mutkey=:selfmuts)
	TreeAlgs.Fitch.fitch_mutations!(tree, mutkey, seqkey)
end

function crossmapped_ancestral_states!(trees, sref, seqkey=:cmseq, mutkey=:cmmuts)
	tref = trees[sref]
	for s in Iterators.filter(!=(sref), keys(trees))
		TreeAlgs.Fitch.fitch_mutations!(tref, (mutkey, s), (seqkey, s))
	end
end

"""
	get_mut_dict(f::Function, tree:Tree)

Construct dictionary mapping mutations of segment `s` in tree `t` to the number of times they appear. For `n` in `t.lnodes`, looks for mutations in `f(n)`.
"""
function get_mut_dict(f::Function, t::Tree)
	md = Dict{TreeTools.Mutation,Int64}()
	for n in values(t.lnodes)
		for m in f(n)
			md[TreeTools.parse_mutation(m)] = get(md, TreeTools.parse_mutation(m), 0) + 1
		end
	end
	return md
end

"""
    suspicious_mutations(md_ref, md_new ; n = 100)

Following mutations `m` are suspicious:
- `m` appears in `md_new` but not in `md_ref`
- `m` appears only once in `md_ref`, but multiple times in `md_new`
"""
function suspicious_mutations(md_ref, md_new)
    unique_mut = Set{Tuple{Int64, Int64, Int64}}()
    new_mut = Set{Tuple{Int64, Int64, Int64}}()
    for (k,v) in md_new
        if get(md_ref, k, -1) == -1
            push!(new_mut, k)
        elseif v > 1 && md_ref[k] == 1
            push!(unique_mut, k)
        end
    end
    return (new_mut, unique_mut)
end

"""
	suspicious_branches!(trees::Dict{T,<:Tree}, cmkey=cmmuts where T

Mark suspicious branches in `trees` using cross-mapped mutations. Store the number of suspicious mutations for key `s` on the branch above tree node `n` in `n.data.dat[:suspicious_muts][s]`.
"""
function suspicious_branches!(trees::Dict{T,<:Tree}, selfmutkey = :selfmuts, cmkey=:cmmuts) where T
	# Set of mutations on branches for segment s on its own tree
	mut_dict = Dict(s=>get_mut_dict(n -> n.data.dat[selfmutkey], t) for (s,t) in trees)
	#
	for (s,t) in trees
		for n in values(t.lnodes)
			n.data.dat[:suspicious_muts] = Dict{T, Int64}()
		end
	end
	#
	for (sref,tref) in trees, s in Iterators.filter(!=(sref), keys(trees))
		_suspicious_branches!(tref, s, mut_dict, cmkey)
	end
end
function _suspicious_branches!(t::Tree, s, md::Dict, cmkey)
	for n in values(t.lnodes)
		n.data.dat[:suspicious_muts][s] = 0

		for mut in n.data.dat[cmkey][s]
			if !haskey(md[s], mut)
				n.data.dat[:suspicious_muts][s] += 1
			end
		end
	end
	nothing
end

"""
"""
function prune_suspicious_mccs!(trees::Dict, suspmut_key=:suspicious_muts, suspmut_threshold=2)
	if length(trees) > 2
		error("Received $(length(trees)) trees. Only for pairs of trees.")
	end
	pMCCs = []
	ct = deepcopy(trees)
	pmccs = _prune_suspicious_mccs!(trees, suspmut_key, suspmut_threshold)
	while !mapreduce(t->length(nodes(t))==1, *, values(trees); init = true) && !isempty(pmccs)
		append!(pMCCs, pmccs)
		ct = deepcopy(trees)
		pmccs = _prune_suspicious_mccs!(trees, suspmut_key, suspmut_threshold)
	end
	return pMCCs
end

function _prune_suspicious_mccs!(trees, suspmut_key=:suspicious_muts, suspmut_threshold=2)
	S = collect(keys(trees))
	# Get mccs
	mccs = naive_mccs(collect(values(trees)))
	# For the root of each, see if it has more than 2 suspicious mutations
	idx = Int64[]
	for (i,m) in enumerate(mccs)
		for (sref,t) in trees
			a = lca(t, m)
			for s in Iterators.filter(!=(sref), S) # Just one value normally
				if haskey(a.data.dat,suspmut_key) && a.data.dat[suspmut_key][s] >= suspmut_threshold
					@debug "$(a.label) (mcc $i) with $(a.data.dat[suspmut_key][s]) suspicious mutations in $s for tree of $sref"
					@debug "Will prune $m"
					push!(idx, i)
				end
			end
		end
	end
	unique!(idx)
	# Prune those that do
	for i in idx
		for (s,t) in trees
			if !lca(t, mccs[i]).isroot
				prunesubtree!(t, mccs[i])
			end
		end
	end
	# Return pruned mccs
	return mccs[idx]
end

max_stars_repo_name: PierreBarrat/OldTreeKnit

max_issues_repo_name: PierreBarrat/OldTreeKnit

max_forks_repo_name: PierreBarrat/OldTreeKnit


text:
#hideall
using Plots, Joe, Random
using Plots.PlotMeasures
Plots.reset_defaults()
gr()
default(
    titlefont  = font("JuliaMono", default(:titlefontsize),  ),
    guidefont  = font("JuliaMono",  default(:guidefontsize),  ),
    tickfont   = font("JuliaMono", default(:tickfontsize),   ),
    legendfont = font("JuliaMono",  default(:legendfontsize), ),
    left_margin = 30px,
    right_margin = 30px,
    bottom_margin = 30px
)
# ### 例64
# Bostonデータセットを決定木で回帰分析します。
# $\alpha$の値は0で、最低要素数だけを指定します。
using Joe
using ScikitLearn
@sk_import datasets: load_boston
df = load_boston()
X =df["data"];y=df["target"];
node = Joe.decisiontree(X,y;n_min=50);
# グラフ化はのちほど検討。
# ### 例65
# 10-fold CVで$\alpha \in [0,1.5]$の範囲での最適化を検討します。
α_seq = 0:0.1:1.5
n = 100
X65 = X[begin:n,:]; y65 = y[begin:n];
s = Int(n/10)
out = zeros(length(α_seq))
for (i,α) in enumerate(α_seq)
    global n, X65, y65, out # hide
    local node # hide
    SS = 0
    for h in 0:9
        test = h*s+1:(h+1)*s
        train = setdiff(1:n,test)
        node = Joe.decisiontree(X65[train,:],y65[train];α=α)
        for t in test
            SS += (y[t] - Joe.dtvalue(X65[t,:],node))^2
        end
    end
    out[i] = SS/n
end
p65_1 = plot(α_seq,out,xlabel="α",ylabel="二乗誤差",
            title="CVで最適なα(N=100)",label=false);
# 同様にn_minも$n_{min} \in \{ 1,2,\dots, 15\}$の範囲で最適化を検討します。

n_min_seq = 1:15
out2 = zeros(length(n_min_seq))
for (i,n_min) in enumerate(n_min_seq)
    global n, X65, y65, out2 # hide
    local node # hide
    SS = 0
    for h in 0:9
        test = h*s+1:(h+1)*s
        train = setdiff(1:n,test)
        node = Joe.decisiontree(X65[train,:],y65[train];n_min=n_min)
        for t in test
            SS += (y[t] - Joe.dtvalue(X65[t,:],node))^2
        end
    end
    out2[i] = SS/n
end
p65_2 = plot(n_min_seq,out2,xlabel="n_min",ylabel="二乗誤差",
            title="CVで最適なn_min(N=100)",label=false);

p65 = plot(p65_1,p65_2,layout=(1,2))
savefig(p65,joinpath(@OUTPUT,"fig7-5.svg")) # hide
# \fig{fig7-5}
# ## 分類の決定木
# ### 例66 Fisherのあやめ（何回目？）
iris = dataset("datasets","iris")
x = iris[!,1:4] |> Matrix
targets=unique(iris.Species)
# あやめの種類をIntに変換します。
label = Dict(target=>i for (i,target) in enumerate(targets))
y = [label[i] for i in iris.Species];


text:
import RIOPA
import Test: @testset, @test, @test_throws

@testset "Ctrl-cfg" begin
    config = RIOPA.default_config()
    dscfg_raw = config[:datasets][1]
    streamcfg = RIOPA.Ctrl.configure_stream(dscfg_raw[:data_streams][1])

    @test streamcfg.name == "Level_0"
    @test streamcfg.payload_groups[1].range == (1000, 1200)
    @test streamcfg.payload_groups[1].ratio == 0.5
    @test streamcfg.payload_groups[2].range == (2000, 2400)
    @test streamcfg.payload_groups[2].ratio == 0.5

    ds = RIOPA.Ctrl.configure_dataset(dscfg_raw)
    @test ds.cfg.name == "data 1"
    @test ds.cfg.basename == "data_one"
    @test ds.cfg.datagen_backend_tag == RIOPA.DataGen.DefaultDataGenTag()
    @test ds.cfg.io_backend_tag == RIOPA.HDF5IOTag()
    @test ds.cfg.nsteps == 10
    @test ds.cfg.compute_seconds == 1.0
    @test length(ds.cfg.streams) == 2
    streamcfg = ds.cfg.streams[1]
    @test streamcfg.name == "Level_0"
    @test streamcfg.payload_groups[1].range == (1000, 1200)
    @test streamcfg.payload_groups[1].ratio == 0.5
    @test streamcfg.payload_groups[2].range == (2000, 2400)
    @test streamcfg.payload_groups[2].ratio == 0.5
    streamcfg = ds.cfg.streams[2]
    @test streamcfg.name == "Level_1"
    @test streamcfg.payload_groups[1].range == (2000, 2500)
    @test streamcfg.payload_groups[1].ratio == 1 // 4
    @test streamcfg.payload_groups[2].range == (4000, 4800)
    @test streamcfg.payload_groups[2].ratio == 3 // 4
end

mutable struct TestTag <: RIOPA.TagBase
    datacount::Int32
    times::Vector{Float64}
end

TestTag() = TestTag(0, [])

function RIOPA.DataGen.generate!(tag::TestTag, ds::RIOPA.Ctrl.DataSet)
    # ds.data = rand(Float64, 10^6)
    tag.datacount += 1
end

function RIOPA.IO.perform_step(tag::TestTag, ds::RIOPA.Ctrl.DataSet)
    # println(
    #     "time: ",
    #     time() - ds.timestamp,
    #     "s; ",
    #     ds.cfg.name,
    #     ": step ",
    #     ds.curr_step,
    # )
    push!(tag.times, time() - ds.timestamp)
end

@testset "Controller" begin
    # Run once to get rid of compiler time
    tag1 = TestTag()
    ds_configs =
        [RIOPA.Ctrl.DataSetConfig("test1", "test1", tag1, tag1, 1, 0.25, [])]
    RIOPA.Ctrl.Controller(map(RIOPA.Ctrl.DataSet, ds_configs))()

    tag1 = TestTag()
    tag3 = TestTag()
    ds_configs = [
        RIOPA.Ctrl.DataSetConfig("test1", "test1", tag1, tag1, 6, 1.0, []),
        RIOPA.Ctrl.DataSetConfig("test2", "test2", tag3, tag3, 2, 3.0, []),
    ]
    RIOPA.Ctrl.Controller(map(RIOPA.Ctrl.DataSet, ds_configs))()

    @test tag1.datacount == 6
    for t in tag1.times
        @test abs(t - 1.0) < 0.1
    end

    @test tag3.datacount == 2
    for t in tag3.times
        @test abs(t - 3.0) < 0.1
    end
end


text:
using Test
using MDCDL
using Random

@testset "Permutation" begin
    Random.seed!(39328582)
    @testset "inversion" begin
        for d = 1:5
            for n = 1:20
                df = (rand(1:4, d)...,)
                x = rand((4 .* df)...)

                pvx = MDCDL.mdarray2polyphase(x, df)

                nBlocks = pvx.nBlocks
                expctd = copy(pvx.data)
                actual = MDCDL.irotatedimspv(MDCDL.rotatedimspv(pvx.data, nBlocks[1]), nBlocks[1])
                @test actual == expctd
            end
        end
    end

    @testset "cyclic group" begin
        for d = 1:5
            for n = 1:20
                df = (rand(1:4, d)...,)
                x = rand((4 .* df)...)

                expctd = copy(x)

                pvx = MDCDL.mdarray2polyphase(x, df)
                pvdata = pvx.data
                nBlocks = pvx.nBlocks

                for blk in nBlocks
                    pvdata = MDCDL.rotatedimspv(pvdata, blk)
                end

                rpvx = MDCDL.PolyphaseVector(pvdata, nBlocks)
                actual = MDCDL.polyphase2mdarray(rpvx, df)
                @test actual == expctd
            end
        end
    end

    @testset "pv vs md" begin
        for d = 1:5
            for n = 1:20
                df = (fill(1,d)...,)
                szx = rand(1:16, d)
                x = rand(szx...)

                expctd = permutedims(x, [collect(2:d)..., 1])

                pvx = MDCDL.mdarray2polyphase(x, df)
                nBlocks = pvx.nBlocks

                permdata = MDCDL.rotatedimspv(pvx.data, nBlocks[1])
                permpvx = PolyphaseVector(permdata, ([nBlocks[2:d]..., nBlocks[1]]...,))

                actual = MDCDL.polyphase2mdarray(permpvx, df)
                @test actual == expctd
            end
        end
    end

    @testset "PolyphaseVector Type" begin
        for d = 1:5
            for n = 1:20
                df = (fill(1,d)...,)
                szx = rand(1:16, d)
                x = rand(szx...)
                nrot = rand(1:d)

                expctd = permutedims(x, circshift(collect(1:d), -nrot))

                pvx = MDCDL.mdarray2polyphase(x, df)
                permpvx = foldl(1:nrot, init = pvx) do prevpvx, k
                    MDCDL.rotatedimspv(prevpvx)
                end
                actual = MDCDL.polyphase2mdarray(permpvx, df)
                @test actual == expctd
            end
        end
    end
end

max_stars_repo_name: nixir/MDCDL

max_issues_repo_name: nixir/MDCDL

max_forks_repo_name: nixir/MDCDL


text:
const LEAF_FEATURE = -1::Int64

immutable FlatNode
    feature::Int64
    value::Float64
    leftChild::Int64
end

immutable FlatTree
    nodes::Vector{FlatNode}

    function FlatTree(t::Union(Inner, Leaf))
        nodes = [FlatNode(0, 0, 0)]
        function recur(l::Leaf, current_index::Int64)
            nodes[current_index] = FlatNode(LEAF_FEATURE, l.value, 0)
        end

        function recur(i::Inner, current_index::Int64)
            left_child = length(nodes) + 1
            push!(nodes, FlatNode(0, 0, 0), FlatNode(0, 0, 0))
            nodes[current_index] = FlatNode(i.feature, i.splitValue, left_child)
            recur(i.left, left_child)
            recur(i.right, left_child + 1)
        end
        recur(t, 1)
        return new(nodes)
    end
end

function evaluate(t::FlatTree, features::Vector{Float64})
    current = 1
    while true
        n = t.nodes[current]
        if n.feature == LEAF_FEATURE
            return n.value
        end
        current = features[n.feature] < n.value ? n.leftChild : n.leftChild + 1
    end
end

immutable FlatForest <: Evaluator
    trees::Vector{FlatTree}

    function FlatForest(f::Forest)
        return new([FlatTree(t) for t in f.trees])
    end
end

function evaluate(f::FlatForest, features::Vector{Float64})
    sum = 0.0
    for t in f.trees
        sum += evaluate(t, features)
    end
    return sum
end


text:
"""
  Abstract type for an extendabe preconditoioner working together
  with ExtandableSparseMatrix.
  Still beta, so not in the published documentation.

  Any such preconditioner should have the following fields
````
  extmatrix
  pattern_timestamp
````
and  methods
````
  update!(precon)
````
  The idea is that, depending if the matrix pattern has changed,
  different steps are needed to update the preconditioner.

  Moreover, they have the ExtendableSparseMatrix as a field, ensuring
  consistency after construction.
"""
abstract type AbstractExtendablePreconditioner{Tv,Ti} end

need_symbolic_update(precon::AbstractExtendablePreconditioner)=precon.extmatrix.pattern_timestamp>precon.pattern_timestamp
timestamp!(precon::AbstractExtendablePreconditioner)= precon.pattern_timestamp=time()

include("jacobi.jl")
include("ilu0.jl")


text:
#TODO merge into StatsBase.jl & potentially make types there abstract

using Random
using StatsBase: make_alias_table!

struct OneToInf <: AbstractVector{Int} end
Base.size(::OneToInf) = (typemax(Int),)
Base.getindex(::OneToInf, x::Integer) = x
Base.iterate(::OneToInf, state=1) = (state, state+1)
Base.eltype(::Type{OneToInf}) = Int
Base.IteratorSize(::Type{OneToInf}) = Base.IsInfinite()
Base.show(io::IO, X::OneToInf) = print(io, "OneToInf()")
Base.:(==)(::OneToInf, ::OneToInf) = true

struct AliasTable{T, S} <: Random.Sampler{T} where {S <: AbstractVector{T}}
    accept::Vector{Float64}
    alias::Vector{Int}
    support::S
end

function AliasTable(probs::AbstractVector{<:Real}, support::AbstractVector=OneToInf())
    AliasTable!(Float64.(probs), support)
end
function AliasTable!(probs::AbstractVector{Float64}, support::AbstractVector=OneToInf())
    Base.has_offset_axes(probs, support) && throw(ArgumentError("offset arrays are not supported but got an array with index other than 1"))
    n = length(probs)
    n > 0 || throw(ArgumentError("The input probability vector is empty."))
    checkbounds(Bool, support, axes(probs, 1)) || throw(BoundsError("probabilities extend past support"))
    alias = similar(probs, Int)
    sum_probs = sum(probs)
    0 < sum_probs < Inf || throw(ArgumentError("sum(probs) = $sum_probs"))
    make_alias_table!(probs, sum_probs, probs, alias)
    AliasTable{eltype(support), typeof(support)}(probs, alias, support)
end

# with an alias table that has a high proportion of guaranteed acceptance, this can be
# optimized to reduce the number of times u needs to be computed.
# StatsBase.make_alias_table! does not produce such a table.
function Random.rand(rng::AbstractRNG, s::AliasTable)
    i = rand(rng, eachindex(s.accept))
    u = rand(rng)
    @inbounds s.support[u < s.accept[i] ? i : s.alias[i]]
end

Base.:(==)(a::AliasTable, b::AliasTable) =
    a.accept == b.accept && a.support == b.support && a.alias[a.accept .!= 1] == b.alias[b.accept .!= 1]

# I think this is faster to make, faster to sample, more general, and better integrated with
# Random than Distributions.jl's version and much more usable & complete than StatsBase.jl's
# version.

#=
using BenchmarkTools, StatsBase

#FBD
source = rand(1000);
@btime AliasTable($source); # 7.646 μs (4 allocations: 31.75 KiB)
at = AliasTable(source)
@btime rand($at, 1000); # 12.331 μs (1 allocation: 7.94 KiB)

#Distributions
using Distributions
@btime sampler(Categorical($source ./ sum($source))); # 10.933 μs (11 allocations: 47.77 KiB)
at2 = sampler(Categorical(source ./ sum(source)))
@btime rand($at2, 1000); # 22.882 μs (1 allocation: 7.94 KiB)

#StatsBase
function f(weights)
    probs = similar(source)
    alias = similar(source, Int)
    make_alias_table!(weights, sum(source), probs, alias)
    probs, alias
end
@btime f($source);
probs, alias = f(source) # 7.525 μs (7 allocations: 31.81 KiB)
#no samplng method available.

=#

max_stars_repo_name: LilithHafner/FBD

max_issues_repo_name: LilithHafner/FBD

max_forks_repo_name: LilithHafner/FBD


text:
using PortfolioBedtest
using Documenter

DocMeta.setdocmeta!(PortfolioBedtest, :DocTestSetup, :(using PortfolioBedtest); recursive=true)

makedocs(;
    modules=[PortfolioBedtest],
    authors="Andrey Oskin",
    repo="https://github.com/Arkoniak/PortfolioBedtest.jl/blob/{commit}{path}#{line}",
    sitename="PortfolioBedtest.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://Arkoniak.github.io/PortfolioBedtest.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/Arkoniak/PortfolioBedtest.jl",
)


text:
N = 10
ph = discrete_phs_closed(N, 0, 1);
num_freq = frequencies(ph)

@test abs(num_freq[2]-pi) < 1e-13

max_stars_repo_name: flavioluiz/phjulia

max_issues_repo_name: flavioluiz/phjulia

max_forks_repo_name: flavioluiz/phjulia


text:
module StructuralAnalysis

include("2DTruss.jl")
include("LocalK.jl")

end


text:
using CSV
using Plots
using DataFrames
using AnnotationFusion

file = joinpath(dirname(dirname(@__FILE__)), "data/behavioral_codes.csv")
data = CSV.read(file, DataFrame)

raters = :annotatorID
items = :sessionID # Contains info on columns {stage, initiator, person_rated}
ratings = :dim1

annotations = dense(data, raters, items, ratings)

μ = fuse(annotations, :sessionID, Mean())

te = fuse(annotations, :sessionID, TE())

plot(μ); plot!(te)


text:
function is_leap_year(year::Int)
  if year % 400 == 0
	return true
  end

  if year % 100 == 0
	return false
  end

  if year % 4 == 0
	return true
  end

  return false

end

max_stars_repo_name: dceresoli/exercism

max_issues_repo_name: dceresoli/exercism

max_forks_repo_name: dceresoli/exercism