text:
module JuliaVariables
import LegibleLambdas
using MLStyle
using NameResolution

export Scope, ScopedVar, ScopedFunc, ScopedGenerator
export solve_from_local, solve
const DEBUG = false
@static if DEBUG
macro logger(call)
    @when :($f($ana, $(args...))) = call begin
        :($f($ana, $(args...)) = begin
            println($(QuoteNode(f)),'(', $string($objectid($ana), base = 62), ",", $(args...), ')')
            $NameResolution.$f($ana, $(args...))
        end)
    end
end
@logger child_analyzer!(ana, is_physical)
@logger enter!(ana, sym)
@logger require!(ana, sym)
@logger is_local!(ana, sym)
@logger is_global!(ana, sym)
end
@generated function field_update(main :: T, field::Val{Field}, value) where {T, Field}
    fields = fieldnames(T)
    quote
        $T($([field !== Field ? :(main.$field) : :value for field in fields]...))
    end
end

macro with(main, field::Symbol, value)
    :($field_update($main, $(Val(field)), $value)) |> esc
end

struct ScopedVar
    scope :: Scope
    sym   :: Symbol
end

struct ScopedFunc
    scope :: Scope
    func  :: Expr
end

struct ScopedGenerator
    scope :: Scope
    gen  :: Expr
end

function var_repr(var :: LocalVar)
    r = ""
    if var.is_mutable.x
        r *= "mut "
    end
    if var.is_shared.x
        r *= "cell "
    end
    r * string(var.sym)
end
function var_repr(var :: GlobalVar)
    "global $var"
end

function Base.show(io::IO, scopedvar::ScopedVar)
    scope = scopedvar.scope
    sym = scopedvar.sym
    if isempty(scope.bounds) && isempty(scope.freevars)
        Base.print(io, sym)
        return
    end
    var = var_repr(scope[sym])
    Base.print(io, "@", var)
end

function Base.show(io::IO, o::ScopedFunc)
    scope = o.scope
    func = o.func
    Base.print(io, "[", join(map(var_repr, values(scope.freevars) |> collect), ",") , "]", func)
end

function Base.show(io::IO, o::ScopedGenerator)
    scope = o.scope
    gen = o.gen
    Base.print(io, "[", join(map(var_repr, values(scope.freevars) |> collect), ",") , "]", gen)
end

islinenumbernode(x) = x isa LineNumberNode
rmlines(ex::Expr) = begin
    hd = ex.head
    tl = map(rmlines, filter(!islinenumbernode, ex.args))
    Expr(hd, tl...)
end

rmlines(ex::ScopedFunc) = begin
    func = ex.func |> rmlines
    @with ex func func
end

rmlines(ex::ScopedGenerator) = begin
    gen = ex.gen |> rmlines
    @with ex gen gen
end

rmlines(a) = a

function quick_lambda(ex)
    @when Expr(:call, args...) && if any(==(:_), args) end = ex begin
        quick_arg = gensym("quick_arg")
        λ = gensym("λ")
        args = Any[arg === :_ ? quick_arg : quick_lambda(arg)
                   for arg in args]
        ex = Expr(:call, args...)
        :($LegibleLambdas.@λ($quick_arg -> $ex))
    @when Expr(hd, args...) = ex
        Expr(hd, map(quick_lambda, args)...)
    @otherwise
        ex
    end
end

macro quick_lambda(ex)
    esc(quick_lambda(ex))
end

function with_fname(cont, ex)
    @match ex begin
        Expr(:call, f, args...) => Expr(:call, cont(f), args...)
        :($a where {$(args...)}) => :($(with_fname(cont, a)) where {$(args...)})
        :($a :: $t)              => :($(with_fname(cont, a)) :: $t)
        a => a
    end
end

is_func_sym(sym) = sym === :function

is_broadcast_fusing(sym) = begin
    s = string(sym)
    endswith(s, "=") && startswith(s, ".")
end

is_broadcast_sym(sym) = begin
    s = string(sym)
    startswith(s, ".") && Base.isoperator(Symbol(s[1:end]))
end

is_rhs_sym(sym) = begin
    (sym in Symbol[:ref, :.])
end

@data ScopeDesc begin
    Local()
    Global()
    Arg()
    Lexical()
end

struct CtxFlag
    default_scope :: ScopeDesc
    is_lhs        :: Bool
end
MLStyle.Record.@as_record CtxFlag

Base.@pure Base.:+(flag :: CtxFlag, desc :: Symbol) =
    @when CtxFlag(default_scope=default_scope, is_lhs=is_lhs) = flag begin
        @match desc begin
            :lexical => CtxFlag(Lexical(), is_lhs)
            :global => CtxFlag(Global(), is_lhs)
            :local => CtxFlag(Local(), is_lhs)
            :arg => CtxFlag(Arg(), is_lhs)
            :lhs => CtxFlag(default_scope, true)
            :rhs => CtxFlag(default_scope, false)
        end
    @otherwise
        error("impossible")
    end

CtxFlag() = CtxFlag(Lexical(), false)

function solve(ana, sym :: Symbol, ctx_flag::CtxFlag)
    default_scope = ctx_flag.default_scope
    @when Global() = default_scope begin
        is_global!(ana, sym)
    @when Local() || Arg() = default_scope
        is_local!(ana, sym)
    end
    if ctx_flag.is_lhs
        enter!(ana, sym)
    else
        require!(ana, sym)
    end
    if Arg() == default_scope
        sym
    else
        ScopedVar(ana.solved, sym)
    end
end

function solve(ana, ex, ctx_flag::CtxFlag = CtxFlag())
    @match ex begin
# give up analysing macrocall expressions.
        Expr(:macrocall, _...) => ex
# scoped constructs
        Expr(hd && if hd in (:function, :(=), :->) end, a,  b) =>
            begin
                is_fn = is_func_sym(hd)
                a = with_fname(a) do name
                    is_fn = true
                    !(name isa Symbol) && return name
                    solve(ana, name, ctx_flag + :lhs)
                end
                is_fn = is_fn || hd !== :(=)
                ana = if is_fn
                    child_analyzer!(ana, true)
                else
                    ana
                end
                a = solve(ana, a, ctx_flag + :lhs + :arg)
                b = solve(ana, b, CtxFlag())
                func = Expr(hd, a, b)
                is_fn ? ScopedFunc(ana.solved, func) : func
            end
        Expr(hd && if hd in (:for, :while) end, a, b) =>
            begin ctx_flag = CtxFlag()
                ana = child_analyzer!(ana, false)
                a = solve(ana, a, ctx_flag + :lhs + :local)
                b = solve(ana, b, ctx_flag)
                Expr(hd, a, b)
            end
        Expr(:try, attempt, a, blocks...) =>
            @quick_lambda begin ctx_flag = CtxFlag()
                attempt = solve(child_analyzer!(ana, false), attempt, ctx_flag)
                a = solve(child_analyzer!(ana, false), a, ctx_flag + :lhs) # catch exc
                blocks = map(solve(child_analyzer!(ana, false), _, ctx_flag), blocks)
                Expr(:try, attempt, a, blocks...)
            end
        Expr(:let, Expr(:block, binds...) || bind && Do(binds = [bind]), a) =>
            begin ctx_flag = CtxFlag() + :lhs + :arg
                cur_scope = ana
                binds = map(binds) do each
                    if !isa(each, LineNumberNode)
                        cur_scope = child_analyzer!(cur_scope, false)
                        solve(cur_scope, each, ctx_flag)
                    else
                        each
                    end
                end
                a = solve(cur_scope, a, CtxFlag())
                Expr(:let, Expr(:block, binds...), a)
            end
        Expr(:do, call, lam) =>
            begin
                ctx_flag = CtxFlag()
                call = solve(ana, call, ctx_flag)
                lam = solve(ana, lam, ctx_flag)
                Expr(:do, call, lam)
            end
        Expr(:generator, expr, binds...) =>
            @quick_lambda begin ctx_flag = CtxFlag()
                ana = child_analyzer!(ana, true)
                binds = map(solve(ana, _, ctx_flag + :lhs + :local), binds)
                expr = solve(ana, expr, ctx_flag)
                ScopedGenerator(ana.solved, Expr(:generator, expr, binds...))
            end
        Expr(:filter, cond, binds...) =>
            @quick_lambda begin
                binds = map(solve(ana, _, ctx_flag), binds)
                cond = solve(ana, cond, ctx_flag)
                Expr(:filter, cond, binds...)
            end
        Expr(hd && if is_broadcast_fusing(hd) end, a, b) =>
            begin
                a = solve(ana, a, ctx_flag + :lhs)
                b = solve(ana, b, ctx_flag + :rhs + :lexical)
                Expr(hd, a, b)
            end
        Expr(:local, args...) =>
            @quick_lambda  begin
                args = map(solve(ana, _, ctx_flag + :arg), args)
                Expr(:local, args...)
            end
        Expr(:global, args...) =>
            @quick_lambda  begin
                args = map(solve(ana, _, ctx_flag + :global), args)
                Expr(:global, args...)
            end
# non-scoped constructs
        :($a where {$(tps...)}) =>
            @quick_lambda begin
                if !ctx_flag.is_lhs && !isempty(tps)
                    ana = child_analyzer!(ana, false)
                    ctx_flag = CtxFlag()
                    tps = map(solve(ana, _, ctx_flag + :lhs), tps)
                else
                    tps = map(solve(ana, _, ctx_flag + :arg), tps)
                end
                a = solve(ana, a, ctx_flag)
                :($a where {$(tps...)})
            end
        :($a :: $b) =>
            begin
                a = solve(ana, a, ctx_flag)
                b = solve(ana, b, CtxFlag())
                :($a :: $b)
            end
        Expr(hd && if hd in (:module, :baremodule) end, args...) =>
           @quick_lambda Expr(hd, map(solve(ana, _, ctx_flag), args)...)
# namedtuple
# https://github.com/thautwarm/GG.jl/issues/8
        Expr(:tuple, args...) =>
            let args = map(args) do arg
                    @when Expr(:(=), k::Symbol, v) = arg begin
                        Expr(:(=), k, solve(ana, v, ctx_flag))
                    @otherwise
                        solve(ana, arg, ctx_flag)
                    end
                end
                Expr(:tuple, args...)
            end
# keyword arguments for tuples or calls
        Expr(:kw, k::Symbol, v) =>
            begin
                if ctx_flag.default_scope == Arg()
                    solve(ana, k, ctx_flag)
                end
                Expr(:kw, k, solve(ana, v, ctx_flag + :rhs))
            end
# broadcasting symbols
        Expr(:call, f :: Symbol, args...) &&
            if length(args) in (1, 2) && is_broadcast_sym(f)
            end =>
            @quick_lambda let args = map(solve(ana, _, ctx_flag), args)
                Expr(:call, f, args...)
            end
# addition assignment, https://github.com/thautwarm/JuliaVariables.jl/issues/6
        Expr(hd && if hd in (:(+=), :(*=), ) end, var, rhs...) =>
            @quick_lambda begin
                var = solve(ana, var, ctx_flag + :lhs)
                rhs = map(solve(ana, _, ctx_flag), rhs)
                Expr(hd, var, rhs...)
            end
        Expr(hd, args...) =>
            @quick_lambda begin
                if is_rhs_sym(hd)
                    ctx_flag += :rhs
                end
                args = map(solve(ana, _, ctx_flag), args)
                Expr(hd, args...)
            end
        a => a
    end
end

function solve(ex)
    ana = top_analyzer()
    ex = solve(ana, ex)
    anas = [(@with child parent nothing) for child in ana.children]
    foreach(run_analyzer, anas)
    ex
end

function solve_from_local(ex)
    ana = top_analyzer()
    ex = solve(ana, ex)
    run_analyzer(ana)
    ex
end

end # module


text:
module CompressedCellArraysTests

using Test
using Gridap.Arrays
using Gridap.Geometry
using Gridap.Fields

u(x) = x[1]^2 + x[2]

domain= (0,1,0,1)
cells = (3,3)
model = CartesianDiscreteModel(domain,cells)

Γ = BoundaryTriangulation(model)

cell_mat = [ones(3,3) for cell in 1:num_cells(Γ)]
cell_vec = [ones(3) for cell in 1:num_cells(Γ)]
cell_matvec = pair_arrays(cell_mat,cell_vec)

cell_to_bgcell = get_cell_to_bgcell(Γ)
ccell_to_first_cell = compress_ids(cell_to_bgcell)
@test ccell_to_first_cell == [1, 3, 4, 6, 7, 8, 10, 11, 13]

ccell_to_mat= compress_contributions(cell_mat,cell_to_bgcell,ccell_to_first_cell)
@test length(ccell_to_mat) == 8
@test ccell_to_mat[1] == 2*ones(3,3)
@test ccell_to_mat[1] == 2*ones(3,3)
@test ccell_to_mat[2] == ones(3,3)
@test ccell_to_mat[3] == 2*ones(3,3)
@test ccell_to_mat[3] == 2*ones(3,3)
test_array(ccell_to_mat,collect(ccell_to_mat))

cache = array_cache(ccell_to_mat)
@test getindex!(cache,ccell_to_mat,1) == 2*ones(3,3)
@test getindex!(cache,ccell_to_mat,1) == 2*ones(3,3)
@test getindex!(cache,ccell_to_mat,2) == ones(3,3)

ccell_to_vec = compress_contributions(cell_vec,cell_to_bgcell,ccell_to_first_cell)
@test ccell_to_vec[1] == 2*ones(3)
@test ccell_to_vec[1] == 2*ones(3)
@test ccell_to_vec[2] == ones(3)
@test ccell_to_vec[3] == 2*ones(3)
@test ccell_to_vec[3] == 2*ones(3)
test_array(ccell_to_vec,collect(ccell_to_vec))

cache = array_cache(ccell_to_vec)
@test getindex!(cache,ccell_to_vec,1) == 2*ones(3)
@test getindex!(cache,ccell_to_vec,1) == 2*ones(3)
@test getindex!(cache,ccell_to_vec,2) == ones(3)

ccell_to_matvec = compress_contributions(cell_matvec,cell_to_bgcell,ccell_to_first_cell)
@test ccell_to_matvec[1] == (2*ones(3,3), 2*ones(3))
@test ccell_to_matvec[1] == (2*ones(3,3), 2*ones(3))
@test ccell_to_matvec[2] == (ones(3,3)  , ones(3)  )
@test ccell_to_matvec[3] == (2*ones(3,3), 2*ones(3))
@test ccell_to_matvec[3] == (2*ones(3,3), 2*ones(3))
test_array(ccell_to_matvec,collect(ccell_to_matvec))

cache = array_cache(ccell_to_matvec)
@test getindex!(cache,ccell_to_matvec,1) == (2*ones(3,3), 2*ones(3))
@test getindex!(cache,ccell_to_matvec,1) == (2*ones(3,3), 2*ones(3))
@test getindex!(cache,ccell_to_matvec,2) == (ones(3,3)  , ones(3)  )

# now for blocks

mat = ArrayBlock(Matrix{Matrix{Float64}}(undef,2,2),[true false; true true])
mat[1,1] = ones(3,3)
mat[2,1] = ones(4,3)
mat[2,2] = ones(4,4)

vec = ArrayBlock(Vector{Vector{Float64}}(undef,2),[false, true])
vec[2] = ones(4)

cell_mat = [copy(mat) for cell in 1:num_cells(Γ)]
cell_vec = [copy(vec) for cell in 1:num_cells(Γ)]
cell_matvec = pair_arrays(cell_mat,cell_vec)
ccell_to_mat= compress_contributions(cell_mat,cell_to_bgcell,ccell_to_first_cell)
ccell_to_vec= compress_contributions(cell_vec,cell_to_bgcell,ccell_to_first_cell)
ccell_to_matvec= compress_contributions(cell_matvec,cell_to_bgcell,ccell_to_first_cell)
test_array(ccell_to_mat,collect(ccell_to_mat))
test_array(ccell_to_vec,collect(ccell_to_vec))
test_array(ccell_to_matvec,collect(ccell_to_matvec))

# now for blocks of blocks

mat11 = copy(mat)
mat = ArrayBlock(Matrix{MatrixBlock{Matrix{Float64}}}(undef,2,2),[true false; false false])
mat[1,1] = mat11

vec2 = copy(vec)
vec = ArrayBlock(Vector{VectorBlock{Vector{Float64}}}(undef,2),[false, true])
vec[2] = vec2

cell_mat = [copy(mat) for cell in 1:num_cells(Γ)]
cell_vec = [copy(vec) for cell in 1:num_cells(Γ)]
cell_matvec = pair_arrays(cell_mat,cell_vec)
ccell_to_mat= compress_contributions(cell_mat,cell_to_bgcell,ccell_to_first_cell)
ccell_to_vec= compress_contributions(cell_vec,cell_to_bgcell,ccell_to_first_cell)
ccell_to_matvec= compress_contributions(cell_matvec,cell_to_bgcell,ccell_to_first_cell)
test_array(ccell_to_mat,collect(ccell_to_mat))
test_array(ccell_to_vec,collect(ccell_to_vec))
test_array(ccell_to_matvec,collect(ccell_to_matvec))

end # module

ext: jl


text:
# Autogenerated wrapper script for libLLVM_jll for i686-w64-mingw32-cxx03
export libllvm, llvm_config

using Zlib_jll
JLLWrappers.@generate_wrapper_header("libLLVM")
JLLWrappers.@declare_library_product(libllvm, "libLLVM-13jl.dll")
JLLWrappers.@declare_executable_product(llvm_config)
function __init__()
    JLLWrappers.@generate_init_header(Zlib_jll)
    JLLWrappers.@init_library_product(
        libllvm,
        "bin\\libLLVM-13jl.dll",
        nothing,
    )

    JLLWrappers.@init_executable_product(
        llvm_config,
        "tools\\llvm-config.exe",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
function Index(path::AbstractString)
    db = SQLite.DB(path)
    Index(path, db, [], [], [], [])
end

@inline function _check_idx(b::Bgen)
    @assert b.idx !== nothing "bgen index (.bgi) is needed"
end

function select_region(idx::Index, chrom::AbstractString;
        start=nothing, stop=nothing)
    if start === nothing && stop === nothing
        q = "SELECT file_start_position FROM Variant WHERE chromosome=?"
        params = (chrom,)
    elseif stop === nothing
        q = "SELECT file_start_position FROM Variant" *
                    " WHERE chromosome=? AND position>=?"
        params = (chrom, start)
    else
        q = "SELECT file_start_position FROM Variant" *
                    " WHERE chromosome=? AND position>=? AND position<=?"
        params = (chrom, start, stop)
    end
    r = (DBInterface.execute(idx.db, q, params) |> columntable)[1]
end

"""
    select_region(bgen, chrom; start=nothing, stop=nothing)
Select variants from a region. Returns variant start offsets on the file.
Returns a `VariantIteratorFromOffsets` object.
"""
function select_region(b::Bgen, chrom::AbstractString;
        start=nothing, stop=nothing)
    _check_idx(b)
    offsets = select_region(b.idx, chrom; start=start, stop=stop)
    VariantIteratorFromOffsets(b, offsets)
end

function variant_by_rsid(idx::Index, rsid::AbstractString)
    q = "SELECT file_start_position FROM Variant WHERE rsid= ?"
    params = (rsid,)
    r = (DBInterface.execute(idx.db, q, params) |> columntable)[1]
    if length(r) == 0
        @error "variant with rsid $rsid not found"
    elseif length(r) > 1
        @error "multiple variant matches with $rsid"
    end
    return r[1]
end

"""
    variant_by_rsid(bgen, rsid)
Find a variant by rsid
"""
function variant_by_rsid(b::Bgen, rsid::AbstractString)
    _check_idx(b)
    offset = variant_by_rsid(b.idx, rsid)
    return Variant(b, offset)
end

function variant_by_pos(idx::Index, pos::Integer)
    q = "SELECT file_start_position FROM Variant WHERE position= ?"
    params = (pos,)
    r = (DBInterface.execute(idx.db, q, params) |> columntable)[1]
    if length(r) == 0
        @error "variant match at $pos not found"
    elseif length(r) > 1
        @error "multiple variant matches at $pos"
    end
    return r[1]
end

"""
    variant_by_pos(bgen, pos)
Get the variant of bgen variant given `pos` in the index file
"""
function variant_by_pos(b::Bgen, pos::Integer)
    _check_idx(b)
    offset = variant_by_pos(b.idx, pos)
    return Variant(b, offset)
end

function variant_by_index(idx::Index, first::Integer, last::Union{Nothing, Integer}=nothing)
    if last === nothing
        q = "SELECT file_start_position FROM Variant LIMIT 1 OFFSET ?"
        params = (first - 1,)
    else
        q = "SELECT file_start_position FROM Variant LIMIT ? OFFSET ?"
        params = (last - first + 1, first - 1)
    end
    r = (DBInterface.execute(idx.db, q, params) |> columntable)[1]
    return r
end

"""
    variant_by_index(bgen, n)
get the `n`-th variant (1-based).
"""
function variant_by_index(b::Bgen, first::Integer, last::Union{Nothing, Integer}=nothing)
    _check_idx(b)
    if last === nothing
        offset = variant_by_index(b.idx, first)[1]
        return Variant(b, offset)
    else
        offsets = variant_by_index(b.idx, first, last)
        return VariantIteratorFromOffsets(b, offsets)
    end
end

function offsets(idx::Index)
    if length(idx.offsets) != 0
        return idx.offsets
    end
    q = "SELECT file_start_position FROM Variant"
    r = (DBInterface.execute(idx.db, q) |> columntable)[1]
    resize!(idx.offsets, length(r))
    idx.offsets .= r
    return r
end

function rsids(idx::Index)
    if length(idx.rsids) != 0
        return idx.rsids
    end
    q = "SELECT rsid FROM Variant"
    r = (DBInterface.execute(idx.db, q) |> columntable)[1]
    resize!(idx.rsids, length(r))
    idx.rsids .= r
    return r
end

function chroms(idx::Index)
    if length(idx.chroms) != 0
        return idx.chroms
    end
    q = "SELECT chromosome FROM Variant"
    r = (DBInterface.execute(idx.db, q) |> columntable)[1]
    resize!(idx.chroms, length(r))
    idx.chroms .= r
    return r
end

function positions(idx::Index)
    if length(idx.positions) != 0
        return idx.positions
    end
    q = "SELECT position FROM Variant"
    r = (DBInterface.execute(idx.db, q) |> columntable)[1]
    resize!(idx.positions, length(r))
    idx.positions .= r
    return r
end


text:
using StatsBase, Distributions, DependentBootstrap
using ForecastEval

#Function for generating data
function gen_data(testtype::Symbol, N::Int)
	y = randn(N);
	if testtype == :mean
		x = Vector{Float64}[ randn(N) + 0.1*k for k = -10:1:10 ]
		xbc = randn(N)
	elseif testtype == :var
		x = Vector{Float64}[ sqrt(k) * randn(N) for k = 1.5:-0.05:0.5 ]
		xbc = randn(N)
	elseif testtype == :varsmall
		x = Vector{Float64}[ sqrt(k) * randn(N) for k = 1.2:-0.02:0.8 ]
		xbc = randn(N)
	elseif testtype == :same
		x = Vector{Float64}[ randn(N) for k = 1:20 ]
		xbc = randn(N)
	else
		error("Invalid testtype")
	end
	return(y, x, xbc)
end
vvtomat{T}(x::Vector{Vector{T}})::Matrix{T} = T[ x[m][n] for n = 1:length(x[1]), m = 1:length(x) ]
mattovv{T}(x::Matrix{T})::Vector{Vector{T}} = Vector{T}[ x[:, m] for m = 1:size(x, 2) ]

#----------------------------------------
# DEBUG TOOLS FOR DIEBOLD-MARIANO
#----------------------------------------
(y, x, xbc) = gen_data(:var, 1000);
dmB = Array(DMTest, length(x));
dmH = Array(DMTest, length(x));
for k = 1:length(x)
	xL = (x[k] - y).^2;
	xbcL = (xbc - y).^2;
	lD = xL - xbcL;
	dmB[k] = dm(lD, dmmethod=:boot);
	dmH[k] = dm(lD, dmmethod=:hac);
end
pvalB = Float64[ dmB[k].pvalue for k = 1:length(x) ];
bestB = Int[ dmB[k].bestinput for k = 1:length(x) ];
tailB = Int[ dmB[k].tailregion for k = 1:length(x) ];
statB = Float64[ dmB[k].teststat for k = 1:length(x) ];
println("pvalB = $pvalB");
println("bestB = $bestB");
println("tailB = $tailB");
println("statB = $statB");
pvalH = Float64[ dmH[k].pvalue for k = 1:length(x) ];
bestH = Int[ dmH[k].bestinput for k = 1:length(x) ];
tailH = Int[ dmH[k].tailregion for k = 1:length(x) ];
statH = Float64[ dmH[k].teststat for k = 1:length(x) ];
println("pvalH = $pvalH");
println("bestH = $bestH");
println("tailH = $tailH");
println("statH = $statH");

#----------------------------------------
# DEBUG TOOLS FOR REALITY CHECK
#----------------------------------------
M = 100;
rcOut = Array(RCTest, M);
for m = 1:M
	(y, x, xbc) = gen_data(:varsmall, 1000);
	xbcloss = (xbc - y).^2;
	ld = Vector{Float64}[ (x[k] - y).^2 - xbcloss for k = 1:length(x) ];
	rcOut[m] = rc(vvtomat(ld), alpha=0.05)
end
pvalRC = Float64[ rcOut[m].pvalue for m = 1:M ]
rejH0RC = Int[ rcOut[m].rejH0 for m = 1:M ]
println("pval = $pvalRC")
println("rejH0 = $rejH0RC")
println("rej prop = $(sum(rejH0RC)/M)")

#----------------------------------------
# DEBUG TOOLS FOR SPA TEST
#----------------------------------------
M = 200;
spaOut = Array(SPATest, M);
for m = 1:M
	(y, x, xbc) = gen_data(:varsmall, 1000);
	xbcloss = (xbc - y).^2;
	ld = Vector{Float64}[ (x[k] - y).^2 - xbcloss for k = 1:length(x) ];
	spaOut[m] = spa(vvtomat(ld), alpha=0.05)
end
pvalSPAu = Float64[ spaOut[m].pvalue_u for m = 1:M ];
pvalSPAc = Float64[ spaOut[m].pvalue_c for m = 1:M ];
pvalSPAl = Float64[ spaOut[m].pvalue_l for m = 1:M ];
pvalSPAauto = Float64[ spaOut[m].pvalueauto for m = 1:M ];
rejH0SPA = Int[ spaOut[m].rejH0 for m = 1:M ];
println("pval_u = $pvalSPAu")
println("pval_u rej prop = $(sum(pvalSPAu .< 0.05) / M)")
println("pval_l = $pvalSPAl")
println("pval_l rej prop = $(sum(pvalSPAl .< 0.05) / M)")
println("pval_c = $pvalSPAc")
println("pval_c rej prop = $(sum(pvalSPAc .< 0.05) / M)")
println("pval_auto = $pvalSPAauto")
println("pval_auto rej prop = $(sum(pvalSPAauto .< 0.05) / M)")
println("rejH0 = $rejH0SPA")
println("rej prop = $(sum(rejH0SPA)/M)")

#----------------------------------------
# DEBUG TOOLS FOR MCS
#----------------------------------------
M = 10;
mcsOut = Array(MCSTest, M);
for m = 1:M
	(y, x, xbc) = gen_data(:varsmall, 1000);
	xloss = Vector{Float64}[ (x[k] - y).^2 for k = 1:length(x) ];
	mcsOut[m] = mcs(vvtomat(xloss), alpha=0.05);
end
inQFLine = Vector{Int}[ mcsOut[m].inQF for m = 1:M ];
outQFLine = Vector{Int}[ mcsOut[m].outQF for m = 1:M ];
pValQFLine = Vector{Float64}[ mcsOut[m].pvalueQF for m = 1:M ];
inMTLine = Vector{Int}[ mcsOut[m].inMT for m = 1:M ];
outMTLine = Vector{Int}[ mcsOut[m].outMT for m = 1:M ];
pValMTLine = Vector{Float64}[ mcsOut[m].pvalueMT for m = 1:M ];
println("")
for m = 1:M ; println("$(m)) inQF = $(inQFLine[m])") ; end ;
println("")
for m = 1:M ; println("$(m)) outQF = $(outQFLine[m])") ; end ;
println("")
for m = 1:M ; println("$(m)) pValQF = $(pValQFLine[m])") ; end ;
println("")
println("")
for m = 1:M ; println("$(m)) inMT = $(inMTLine[m])") ; end ;
println("")
for m = 1:M ; println("$(m)) outMT = $(outMTLine[m])") ; end ;
println("")
for m = 1:M ; println("$(m)) pValMT = $(pValMTLine[m])") ; end ;


text:
using ImageFiltering, ImageCore, OffsetArrays
using Statistics, Test
using ImageFiltering: IdentityUnitRange

@testset "specialty" begin
    @testset "Laplacian" begin
        L1 = OffsetArray([1,-2,1],-1:1)
        L2 = OffsetArray([0 1 0; 1 -4 1; 0 1 0], -1:1, -1:1)
        kern = Kernel.Laplacian((true,))
        @test !isempty(kern)
        @test convert(AbstractArray, kern) == L1
        kern = Kernel.Laplacian()
        @test convert(AbstractArray, kern) == L2
        kern = Kernel.Laplacian((true,false))
        @test convert(AbstractArray, kern) == reshape(L1, -1:1, 0:0)
        kern = Kernel.Laplacian((false,true))
        @test convert(AbstractArray, kern) == reshape(L1, 0:0, -1:1)
        kern = Kernel.Laplacian([1], 2)
        @test convert(AbstractArray, kern) == reshape(L1, -1:1, 0:0)
        kern = Kernel.Laplacian([2], 2)
        @test convert(AbstractArray, kern) == reshape(L1, 0:0, -1:1)
        function makeimpulse(T, sz, x)
            A = zeros(T, sz)
            A[x] = oneunit(T)
            A
        end
        # 1d
        kern = Kernel.Laplacian((true,))
        for a in (makeimpulse(Float64, (5,), 3),
                  makeimpulse(Int, (5,), 3),
                  makeimpulse(UInt32, (5,), 3),
                  makeimpulse(UInt16, (5,), 3),
                  makeimpulse(UInt8, (5,), 3),
                  makeimpulse(Bool, (5,), 3),
                  makeimpulse(Gray{N0f8}, (5,), 3),
                  makeimpulse(RGB{Float32}, (5,), 3),
                  makeimpulse(RGB{N0f8}, (5,), 3))
            af = imfilter(a, kern)
            if eltype(a) == unsigned(Int)
                continue  # the concatenation below fails
            end
            T = eltype(a)
            @test af == [zero(T),a[3],-2a[3],a[3],zero(T)]
            af = imfilter(a, (kern,))
            @test af == [zero(T),a[3],-2a[3],a[3],zero(T)]
        end
        for a in (makeimpulse(Float64, (5,), 1),
                  makeimpulse(Gray{N0f8}, (5,), 1),
                  makeimpulse(RGB{Float32}, (5,), 1))
            for (border, edgecoef) in (("replicate", -1),
                                       (Fill(zero(eltype(a))), -2))
                af = imfilter(a, kern, border)
                T = eltype(a)
                @test af == [edgecoef*a[1],a[1],zero(T),zero(T),zero(T)]
            end
        end
        # 2d
        kern = Kernel.Laplacian((true,true))
        @test convert(AbstractArray, kern) == OffsetArray([0 1 0; 1 -4 1; 0 1 0],-1:1,-1:1)
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((3,3))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((3,3))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((3,3))))
            af = imfilter(a, kern)
            T = eltype(a)
            z = zero(T)
            c = oneunit(T)
            @test af == [z z z z z;
                         z z c z z;
                         z c -4c c z;
                         z z c z z;
                         z z z z z]
        end
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((3,1))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((3,1))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((3,1))))
            for (border, edgecoef) in (("replicate", -3),
                                       (Fill(zero(eltype(a))), -4))
                af = imfilter(a, kern, border)
                T = eltype(a)
                z = zero(T)
                c = oneunit(T)
                @test af == [z z z z z;
                             c z z z z;
                             edgecoef*c c z z z;
                             c z z z z;
                             z z z z z]
            end
        end
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((5,5))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((5,5))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((5,5))))
            for (border, edgecoef) in (("replicate", -2),
                                       (Fill(zero(eltype(a))), -4))
                af = imfilter(a, kern, border)
                T = eltype(a)
                z = zero(T)
                c = oneunit(T)
                @test af == [z z z z z;
                             z z z z z;
                             z z z z z;
                             z z z z c;
                             z z z c edgecoef*c]
            end
        end
        # 2d, but only computing the laplacian along the vertical
        kern = Kernel.Laplacian((true,false))
        @test convert(AbstractArray, kern) == OffsetArray(reshape([1, -2, 1], (3,1)),-1:1,0:0)
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((3,3))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((3,3))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((3,3))))
            af = imfilter(a, kern)
            T = eltype(a)
            z = zero(T)
            c = oneunit(T)
            @test af == [z z z z z;
                         z z c z z;
                         z z -2c z z;
                         z z c z z;
                         z z z z z]
        end
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((3,1))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((3,1))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((3,1))))
            for (border, edgecoef) in (("replicate", -2),
                                       (Fill(zero(eltype(a))), -2))
                af = imfilter(a, kern, border)
                T = eltype(a)
                z = zero(T)
                c = oneunit(T)
                @test af == [z z z z z;
                             c z z z z;
                             edgecoef*c z z z z;
                             c z z z z;
                             z z z z z]
            end
        end
        for a in (makeimpulse(Float64, (5,5), CartesianIndex((5,5))),
                  makeimpulse(Gray{N0f8}, (5,5), CartesianIndex((5,5))),
                  makeimpulse(RGB{Float32}, (5,5), CartesianIndex((5,5))))
            for (border, edgecoef) in (("replicate", -1),
                                       (Fill(zero(eltype(a))), -2))
                af = imfilter(a, kern, border)
                T = eltype(a)
                z = zero(T)
                c = oneunit(T)
                @test af == [z z z z z;
                             z z z z z;
                             z z z z z;
                             z z z z c;
                             z z z z edgecoef*c]
            end
        end
    end

    @testset "gaussian" begin
        function gaussiancmp(σ, xr)
            cmp = [exp(-x^2/(2σ^2)) for x in xr]
            cmp ./ sum(cmp)
        end
        for kern in (Kernel.gaussian(()), Kernel.gaussian((),()))
            @test ndims(kern) == 0
            @test kern[] == 1
            @test ImageFiltering.iscopy(kern)
        end
        for kern in (Kernel.gaussian((1.3,)), Kernel.gaussian((1.3,),(7,)))
            @test kern ≈ gaussiancmp(1.3, axes(kern,1))
        end
        @test KernelFactors.gaussian(2, 9) ≈ gaussiancmp(2, IdentityUnitRange(-4:4))
        k = KernelFactors.gaussian((2,3), (9,7))
        @test vec(k[1]) ≈ gaussiancmp(2, IdentityUnitRange(-4:4))
        @test vec(k[2]) ≈ gaussiancmp(3, IdentityUnitRange(-3:3))
        @test sum(KernelFactors.gaussian(5)) ≈ 1
        for k = (KernelFactors.gaussian((2,3)), KernelFactors.gaussian([2,3]), KernelFactors.gaussian([2,3], [9,7]))
            @test sum(k[1]) ≈ 1
            @test sum(k[2]) ≈ 1
        end
        @test Kernel.gaussian((2,), (9,)) ≈ gaussiancmp(2, IdentityUnitRange(-4:4))
        @test Kernel.gaussian((2,3), (9,7)) ≈ gaussiancmp(2, IdentityUnitRange(-4:4)).*gaussiancmp(3, IdentityUnitRange(-3:3))'
        @test sum(Kernel.gaussian(5)) ≈ 1
        for k = (Kernel.gaussian((2,3)), Kernel.gaussian([2,3]), Kernel.gaussian([2,3], [9,7]))
            @test sum(k) ≈ 1
        end
        # Bug noticed in Images issue #674
        k = KernelFactors.gaussian((3, 3, 0))
        @test k[3].data == OffsetArray([1.0], 0:0)
    end

    @testset "DoG" begin
        function gaussiancmp(σ, xr)
            cmp = [exp(-x^2/(2σ^2)) for x in xr]
            OffsetArray(cmp/sum(cmp), xr)
        end
        @test Kernel.DoG((2,), (3,), (9,)) ≈ gaussiancmp(2, -4:4) - gaussiancmp(3, -4:4)
        k = Kernel.DoG((2,3), (4,3.5), (9,7))
        k1 = gaussiancmp(2, -4:4) .* gaussiancmp(3, -3:3)'
        k2 = gaussiancmp(4, -4:4) .* gaussiancmp(3.5, -3:3)'
        @test k ≈ k1-k2
        @test abs(sum(Kernel.DoG(5))) < 1e-8
        @test Kernel.DoG(5) == Kernel.DoG((5,5))
        @test abs(sum(Kernel.DoG((5,), (7,), (21,)))) < 1e-8
        @test axes(Kernel.DoG((5,), (7,), (21,))) == (-10:10,)
    end

    @testset "LoG" begin
        img = rand(20,21)
        σs = (2.5, 3.2)
        kernel1 = Kernel.LoG(σs)
        kernel2 = (KernelFactors.IIRGaussian(σs)..., Kernel.Laplacian())
        imgf1 = imfilter(img, kernel1)
        imgf2 = imfilter(img, kernel2)
        @test cor(vec(imgf1), vec(imgf2)) > 0.8
        # Ensure that edge-trimming under successive stages of filtering works correctly
        ImageFiltering.fillbuf_nan[] = true
        kernel3 = (Kernel.Laplacian(), KernelFactors.IIRGaussian(σs)...)
        @test !any(isnan, imfilter(img, kernel3))
        ImageFiltering.fillbuf_nan[] = false
        @test Kernel.LoG(2.5) == Kernel.LoG((2.5,2.5))
    end

    @testset "moffat" begin
        α = rand()
        β = 2 * rand() + 3  #generate random float between 3 and 5 where commonly used
        @test Kernel.moffat(α,β,2) == Kernel.moffat(α,β,(2,2))

        fwhm = ceil(Int, (α*2*sqrt(2^(1/β) - 1)) * 4)
        @test Kernel.moffat(α, β) == Kernel.moffat(α, β, (fwhm, fwhm))
    end
end

nothing


text:
import FastTransforms: FTPlan

abstract type LinearOperator{T} end

struct NonlinearOperator{F}
    N::F
end

struct SemiLinearPDE{D, T, LL, NN}
    L::LL
    N::NN
    U0::NTuple{D, Matrix{T}}
    function SemiLinearPDE{D, T, LL, NN}(L::LL, N::NN, U0::NTuple{D, Matrix{T}}) where {D, T, LL, NN}
        if D > 1
            @assert LL <: Tuple
            @assert NN <: Tuple
            @assert length(L) == D
            @assert length(N) == D
        end
        new{D, T, LL, NN}(L, N, U0)
    end
end

SemiLinearPDE(L, N, U0::NTuple{D, Matrix{T}}) where {D, T} = SemiLinearPDE{D, T, typeof(L), typeof(N)}(L, N, U0)
SemiLinearPDE(L, N, U0::Matrix{T}) where {T} = SemiLinearPDE(L, N, (U0,))

# ϵ^2 Δ
struct LaplaceBeltrami{T} <: LinearOperator{T}
    ϵ::T
end

# ϵ^2 ℒ_{δ}(α)
struct NonlocalLaplaceBeltrami{T} <: LinearOperator{T}
    ϵ::T
    α::T
    δ::T
end

function create_linear_operator(LB::LaplaceBeltrami{T}, m::Integer, n::Integer) where T
    L = zeros(T, m, 2n-1)
    ϵ = LB.ϵ
    @inbounds for i = 1:m
        L[i,1] = -ϵ^2*i*(i-1)
    end
    @inbounds for j = 1:n-1
        for i = 1:m-j
            L[i,2j] = L[i+j,1]
            L[i,2j+1] = L[i+j,1]
        end
    end
    L
end

function create_linear_operator(NLB::NonlocalLaplaceBeltrami{T}, m::Integer, n::Integer) where T
    L = zeros(T, m, 2n-1)
    λ = evaluate_lambda(m, NLB.α, NLB.δ)
    ϵ = NLB.ϵ
    @inbounds for i = 1:m
        L[i,1] = ϵ^2*λ[i]
    end
    @inbounds for j = 1:n-1
        for i = 1:m-j
            L[i,2j] = L[i+j,1]
            L[i,2j+1] = L[i+j,1]
        end
    end
    L
end

function evaluate_nonlinear_operator!(NL::NonlinearOperator, P::FTPlan{T, 2, FastTransforms.SPHERE}, PS::FTPlan{T, 2, FastTransforms.SPHERESYNTHESIS}, PA::FTPlan{T, 2, FastTransforms.SPHEREANALYSIS}, NU::Matrix{T}) where T
    # Compute spherical harmonic coefficients of a nonlinear operator N(u).

    # sph2fourier
    lmul!(P, NU)
    # Fourier to function values on sphere
    lmul!(PS, NU)
    NU .= NL.N.(NU)
    # Function values on sphere to Fourier
    lmul!(PA, NU)
    # fourier2sph
    ldiv!(P, NU)

    return NU
end

function evaluate_nonlinear_operator!(NL::Tuple{NonlinearOperator{F1}, NonlinearOperator{F2}}, P::FTPlan{T, 2, FastTransforms.SPHERE}, PS::FTPlan{T, 2, FastTransforms.SPHERESYNTHESIS}, PA::FTPlan{T, 2, FastTransforms.SPHEREANALYSIS}, NU::NTuple{2, Matrix{T}}, TEMP::Matrix{T}) where {T, F1, F2}
    # Compute spherical harmonic coefficients of a nonlinear operator N(u).

    # sph2fourier
    lmul!(P, NU[1])
    lmul!(P, NU[2])
    # Fourier to function values on sphere
    lmul!(PS, NU[1])
    lmul!(PS, NU[2])
    TEMP .= NU[1]
    NU[1] .= NL[1].N.(NU[1], NU[2])
    NU[2] .= NL[2].N.(TEMP, NU[2])
    # Function values on sphere to Fourier
    lmul!(PA, NU[1])
    lmul!(PA, NU[2])
    # fourier2sph
    ldiv!(P, NU[1])
    ldiv!(P, NU[2])

    return NU
end


text:
"""
	get_leave_order(tree, MCCs)

Return an ordering of leaves in MCCs given a tree.

Output `order::Vector`: `order[i][label] == n` means that `label` appears at position `n` in MCC `i` for the input tree.

"""
function get_leaves_order(tree, MCCs)
	order = [Dict() for i in 1:length(MCCs)]
	for (i,mcc) in enumerate(MCCs)
		mcc_order = Dict{String, Int}()
		_get_leaves_order!(mcc_order, tree.root, mcc)
		order[i] = mcc_order
	end

	return order
end

function _get_leaves_order!(order::Dict, n::TreeNode, mcc::Vector{<:AbstractString})
	if n.isleaf && in(n.label, mcc)
		order[n.label] = length(order) + 1
	else
		for c in n.child
			_get_leaves_order!(order, c, mcc)
		end
	end

	return nothing

end

"""
	mcc_map(tree::Tree, MCCS)

Compute a map between nodes of `tree` and `MCCs`.
Output `M`: `M[label] == i` means node `label` (internal or leaf) belongs to MCC number `i`
  if `typeof(i)::Int` or does not belong to an MCC if `isnothing(i)`.

## Important
`tree` has to be resolved using `MCCs` for this to work.
"""
function mcc_map(tree::Tree, MCCs)
	mcc_map = Dict()
	mcc_map!(mcc_map, tree.root, MCCs)
	return mcc_map
end
#=
Strategy: message passing going up the tree
Each node `n` returns to parent  `(i, cnt)` where:
- `i` is the MCC to which `n` belongs
- `cnt` is the number of nodes of `MCCs[i]` in the offsprings of `n`, including `n` itself
  *e.g.*: for a leaf, `cnt` is 1; for an internal, it can be higher.
This method allows us to know when we've found all nodes of a given MCC
=#
function mcc_map!(mcc_map::Dict, n::TreeNode, MCCs)
	if n.isleaf
		i = which_mcc(n.label, MCCs)
		@assert !isnothing(i) "Leaf $(n.label) does not belong to any MCC."
		mcc_map[n.label] = i
		if length(MCCs[i]) == 1
			return (nothing, 0) # we're done with MCC i
		else
			return (i, 1) # in MCC i, representing one node
		end
	else
		i = nothing # in which mcc is `n`
		M = 0 # how many nodes of this MCC does it has as offsprings
		for c in n.child
			ic, cnt = mcc_map!(mcc_map, c, MCCs)
			if !isnothing(ic)
				if !isnothing(i)
					# if !isnothing(ic), we're not done with MCC[ic]
					# Meaning n has to belong to it
					@assert i == ic "MCCs $i and $ic conflict (resp. node $(n.label) and $(c.label)).
					Are trees resolved using the MCCs?"
				end
				i = ic
				M += cnt
			end
		end
		mcc_map[n.label] = i

		if isnothing(i) || length(MCCs[i]) == M
			return (nothing, 0) # we're done with mcc i
		elseif length(MCCs[i]) > M
			return (i, M) # in MCC i, representing M nodes
		else
			error("Reached root $(n.label) of MCC $i and have not found all nodes ($(M) out of $(length(MCCs[i])))")
		end
	end
end

"""
	which_mcc(n, MCCs)

Return the index of the MCC to which `n` belongs.
If `i` is returned, `in(n, MCCs[i]`) is true.
"""
function which_mcc(n::AbstractString, MCCs)
	for (i, mcc) in enumerate(MCCs)
		if in(n, mcc)
			return i
		end
	end
	return nothing
end

function sort_polytomies!(t1::Tree, t2::Tree, MCCs)
	mcc_order = get_leaves_order(t1, MCCs)
	_mcc_map = mcc_map(t2, MCCs)
	sort_polytomies!(t2.root, MCCs, _mcc_map, mcc_order)
	node2tree!(t2, t2.root)
	return nothing
end
function sort_polytomies!(n::TreeNode, MCCs, mcc_map, mcc_order)
	if n.isleaf
		i = mcc_map[n.label] # in MCCs[i]
		return mcc_order[i][n.label], 1 # rank in the polytomy, number of nodes in clade
	else
		i = mcc_map[n.label] # Id of the current MCC (can be `nothing`)

		# Construct ranking or children
		# positive if they are in same MCC as current node, negative otherwise
		# r = 1
		rank = Array{Any}(undef, length(n.child))
		for (k, c) in enumerate(n.child)
			if mcc_map[c.label] == i
				# child is in mcc. we know its rank in the MCC
				rank[k] = sort_polytomies!(c, MCCs, mcc_map, mcc_order)
			else
				# child is not in mcc: only count its ladder rank
				_trash, rladder = sort_polytomies!(c, MCCs, mcc_map, mcc_order)
				rank[k] = (nothing, rladder)
				# r += 1
			end
		end

		# # Sort children
		_sort_children!(n, rank)

		# # Return minimum value of the rank of nodes in the MCC.
		# # Will be used by parent to rank the current node
		return findmin(r -> isnothing(r[1]) ? Inf : r[1], rank)[1], sum(x[2] for x in rank)
	end
end

function _sort_children!(n, rank)
	# Custom sorting function
	## if the nodes are in the same MCC, rank them according to that
	## otherwise, rank according to ladder rank (biggest clade at the beginning)
	function _isless(x,y)
		if !isnothing(x[1]) && !isnothing(y[1])
			return x[1] < y[1] # those two numbers can't be equal (rank in mcc)
		else
			return x[2] < y[2]
		end
	end

	children_order = sortperm(rank, lt = _isless)
	n.child = n.child[children_order]
end

## BEFORE REMOVING `write_mccs!`:
## Check whether it's used in annotating auspice json files
# """
#     write_mccs!(trees::Dict, MCCs::Dict, key=:mcc_id)

# Write MCCs id to field `data.dat[key]` of tree nodes. Expect `trees` indexed by single segments, and `MCCs` indexed by pairs of segments.
# """
# function write_mccs!(trees::Dict, MCCs::Dict, key=:mcc_id; overwrite=false)
#     for ((i,j), mccs) in MCCs
#         k = Symbol(key,"_$(i)_$(j)")
#         write_mccs!(trees[i], mccs, k, overwrite=overwrite)
#         write_mccs!(trees[j], mccs, k, overwrite=overwrite)
#     end
# end
# """
#     write_mccs!(t::Tree, MCCs, key=:mcc_id)

# Write MCCs id to field `data.dat[key]` of tree nodes.
# """
# function write_mccs!(t::Tree{TreeTools.MiscData}, MCCs, key=:mcc_id; overwrite=false)
#     for (i,mcc) in enumerate(MCCs)
#         for label in mcc
#             t.lleaves[label].data.dat[key] = i
#         end
#         for n in Iterators.filter(n->!n.isleaf, values(t.lnodes))
#             if is_branch_in_mcc(n, mcc)
#                 if !overwrite && haskey(n.data.dat, key)
#                     error("Node $(n.label) already has an MCC attributed")
#                 end
#                 n.data.dat[key] = i
#             end
#         end
#     end
#     nothing
# end

"""
    is_branch_in_mccs(n::TreeNode, mccs::Array)

Is the branch from `n` to `n.anc` in an element of `mccs`?
"""
function is_branch_in_mccs(n::TreeNode, mccs::Dict)
    for mcc in values(mccs)
        if is_branch_in_mcc(n, mcc)
            return true
        end
    end
    return false
end
function is_branch_in_mccs(n::TreeNode, mccs)
    for mcc in mccs
        if is_branch_in_mcc(n, mcc)
            return true
        end
    end
    return false
end

"""
    is_branch_in_mcc(n::TreeNode, mcc::Array{<:AbstractString})

Is the branch from `n` to `n.anc` in `mcc`?
The clade defined by `n` has to intersect with `mcc`, and this intersection should be strictly smaller `mcc`.
"""
function is_branch_in_mcc(n::TreeNode, mcc::Array{<:AbstractString,1})
    # Simple check
    if n.isleaf
        return length(mcc) > 1 && in(n.label, mcc)
    end

    i = count(c -> c.isleaf && in(c.label, mcc), n)

    return (i > 0 && i < length(mcc))
end
"""
    find_mcc_with_branch(n::TreeNode, mccs::Dict)

Find the mcc to which the branch from `n` to `n.anc` belongs. If `mccs` is an array, return the pair `(index, value)`. If it is a dictionary, return the pair `(key, value)`. If no such mcc exists, return `nothing`.
Based on the same idea that `is_branch_in_mcc`.
"""
function find_mcc_with_branch(n::TreeNode, mccs::Dict)
    cl = [x.label for x in POTleaves(n)]
    for (key,mcc) in mccs
        if !isempty(intersect(cl, mcc)) && !isempty(setdiff(mcc, intersect(cl, mcc)))
            return (key, mcc)
        end
    end
    return nothing
end
function find_mcc_with_branch(n::TreeNode, mccs::Array)
    cl = [x.label for x in POTleaves(n)]
    for (i,mcc) in enumerate(mccs)
        if !isempty(intersect(cl, mcc)) && !isempty(setdiff(mcc, intersect(cl, mcc)))
            return (i,mcc)
        end
    end
    return nothing
end

"""
    is_linked_pair(n1, n2, mccs)
    is_linked_pair(n1::T, n2::T, mccs::Dict{Any,Array{T,1}}) where T
    is_linked_pair(n1::T, n2::T, mccs::Array{Array{T,1},1}) where T

Can I join `n1` and `n2` through common branches only? Equivalent to: is there an `m` in `mccs` such that `in(n1,m) && in(n2,m)`?
"""
function is_linked_pair(n1::T, n2::T, mccs::Dict{Any,Array{T,1}}) where T
    for mcc in values(mccs)
        if in(n1, mcc)
            return in(n2, mcc)
        elseif in(n2, mcc)
            return false
        end
    end
    return false
end
function is_linked_pair(n1, n2, mccs)
    for mcc in values(mccs)
        if in(n1, mcc)
            return in(n2, mcc)
        elseif in(n2, mcc)
            return false
        end
    end
    return false
end
function is_linked_pair(n1::T, n2::T, mccs::Array{Array{T,1},1}) where T
    for mcc in mccs
        if in(n1, mcc)
            return in(n2, mcc)
        elseif in(n2, mcc)
            return false
        end
    end
    return false
end

"""
    find_mcc_with_node(n::String, mccs::Array{Array{<:AbstractString,1},1})

Find MCC to which `n` belongs.
"""
function find_mcc_with_node(n::String, mccs)
    for m in mccs
        if in(n, m)
            return m
        end
    end
    return nothing
end
find_mcc_with_node(n::TreeNode, mccs) = find_mcc_with_node(n.label, mccs)
find_mcc_with_node(n, mccs) = find_mcc_with_node(n.label, mccs)

max_stars_repo_name: PierreBarrat/RecombTools

max_issues_repo_name: PierreBarrat/TreeKnit

max_forks_repo_name: PierreBarrat/TreeKnit


text:
# Source code for Fig.1
# power-law distribution in the number of reported COVID-19 cases
# in countries worldwide and in counties in the US
# this corresponds to Fig.1 in
# Blasius B. (2020) Power-law distribution in the number of confirmed COVID-19 cases. Chaos
# https://doi.org/10.1063/5.0013031

module A

using CSV, DataFrames, Polynomials, PyPlot, StatsBase

include("read_data.jl")
include("tools.jl")

function distribution_countries()
   println("Read countries")
   t,C_ts,D_ts,countries = read_all_data()
   (days, ncountries) = size(C_ts)
   println("days ", days)
   println("total _countries ", ncountries)

   NNN=61 # this is March 22
   C = C_ts[NNN,:]
   D = D_ts[NNN,:]

   ind = findall(x->x>0, C)
   C = sort(C[ind], rev=true)
   nC = length(C)
   nmaxC = maximum(C)
   println("countries with case, n ",nC)
   println("max case: ",nmaxC)

   ind = findall(x->x>0, D)
   D = sort(D[ind], rev=true)
   nD = length(D)
   nmaxD = maximum(D)
   println("countries with death: ",nD)
   println("maxdeath: ", nmaxD)

   # Histogram with bins 2,4,8,16,..
   ex1,ey1 = log2histogram(C)
   x1 = log.(ex1)
   y1 = log.(ey1)
   p = Polynomials.fit(x1,y1,1)
   muC=-coeffs(p)[2]
   println("linear fit ",p)
   ey1fit = exp.(p.(x1))

   alphaC = 1 + length(C) / sum(log.(C)) # calculate exponent by log-likelihood
   println("alphaC ", alphaC)
   alphaC, stC = estimate(C)
   println("alphaC ", alphaC, " pm ", stC)

   # #########
   ex1D,ey1D = log2histogram(D)
   x1D = log.(ex1D)
   y1D = log.(ey1D)
   p = Polynomials.fit(x1D,y1D,1)
   muD=-coeffs(p)[2]
   println("linear fit ",p)
   ey1fitD = exp.(p.(x1D))

   alphaD = 1 + length(D) / sum(log.(D)) # calculate exponent by log-likelihood
   println("alphaD ", alphaD)
   alphaD, stD = estimate(D)
   println("alphaD ", alphaD, " pm ", stD)

   ax = PyPlot.axes([fig_x1, fig_y1, fig_width_x, fig_width_y])
       scatter(ex1,ey1,s=50,label="confirmed cases")
       plot(ex1,ey1,"-",linewidth=2)
       loglog(ex1,ey1fit,"k", label=L"regression line, $\mu=1.18$")
       ylabel(L"frequency, $P_C(n)$",fontsize=18,labelpad=14)
       title("countries", fontsize=16)
       xlim(0.8,2e5)
       ylim(1e-7,1)
       xticks(fontsize=16)
       yticks([1e-7,1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=16)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "a)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix1, fig_y1+fig_Iy, fig_Iwidth_x, fig_Iwidth_y]) # Inset: cumulative distribution
       cc = 10 .^ range(0,5.5, length=50)
       loglog(cc, cumulpdf.(cc,alphaC,nmaxC), color="black", linewidth=2)
       scatter(C,(1:nC)/nC,s=20)
       myfont = Dict("fontname"=>"Sans","fontsize"=>14)
       xticks(fontsize=14)
       yticks(fontsize=14)
       xlim(1,1e5)
       ylim(5e-3,1)
       xlabel("cases, n",fontsize=16)
       ylabel("C(n)",fontsize=16)

   ax = PyPlot.axes([fig_x1, fig_y2, fig_width_x, fig_width_y])
       scatter(ex1D,ey1D,s=50,label="confirmed deaths")
       loglog(ex1D,ey1D,"-",linewidth=2)
       loglog(ex1D,ey1fitD,"k", label=L"regression line, $\mu=1.35$")
       xticks(fontsize=16)
       yticks([1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=16)
       ylim(1e-6,1)
       xlim(0.8,1e4)
       xlabel("cases, n",fontsize=18)
       ylabel(L"frequency, $P_D(n)$",fontsize=18,labelpad=14)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "b)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix1, fig_y2+fig_Iy, fig_Iwidth_x, fig_Iwidth_y]) # Inset: cumulative distribution
       cc = 10 .^ range(0,5, length=50)
       loglog(cc, cumulpdf.(cc,muD,1.5e4), color="black", linewidth=2)
       scatter(D,(1:nD)/nD,s=20)
       myfont = Dict("fontname"=>"Sans","fontsize"=>14)
       xticks(fontsize=14)
       yticks(fontsize=14)
       xlim(1,1e4)
       ylim(1e-2,1)
       xlabel("cases, n",fontsize=16)
       ylabel("C(n)",fontsize=16)
end

function distribution_counties()
   println(); println()
   println("Read US counties")
   # make sure to have access to the data:
   # this means, clone into the Johns Hopkins Covid19 repository
   # and replace the folder variable below with the correct folder to the data
   folder = "../../COVID-19/csse_covid_19_data/csse_covid_19_daily_reports/"
   df = CSV.read(folder*"03-31-2020.csv")

   countries = df.Country_Region
   C = df.Confirmed
   D = df.Deaths

   ind = findall(x->x=="US", countries)
   C = C[ind]
   D = D[ind]

   ind = findall(x->x>0, C)
   C = sort(C[ind], rev=true)
   nC = length(C)
   println("no of entries: ", nC)
   println("max ", maximum(C), " min ", minimum(C))
   println(C[1:3])
   n_ones = length(findall(x->x==1, C))
   println("counties with one case: ", n_ones)

   # Histogram with bins 2,4,8,16,..
   ex1,ey1 = log2histogram(C)
   x1 = log.(ex1)
   y1 = log.(ey1)
   p = Polynomials.fit(x1,y1,1)
   muC=-coeffs(p)[2]
   println("linear fit log2 ",p)
   ey1fit_log2 = exp.(p.(x1))

   alphaC = 1 + length(C) / sum(log.(C)) # calculate exponent by log-likelihood
   println("alphaC ", alphaC)
   alphaC, stC = estimate(C)
   println("alphaC ", alphaC, " pm ", stC)

   ind = findall(x->x>0, D)
   D = sort(D[ind], rev=true)
   nD = length(D)
   println("no of entries: ", nD)
   println("max ", maximum(D), " min ", minimum(D))
   println(D[1:3])
   n_ones = length(findall(x->x==1, D))
   println("counties with one case: ", n_ones)

   # Histogram with bins 2,4,8,16,..
   ex1D,ey1D = log2histogram(D)
   x1D = log.(ex1D)
   y1D = log.(ey1D)
   p = Polynomials.fit(x1D,y1D,1)
   muD=-coeffs(p)[2]
   println("linear fit log2 deaths ",p)
   ey1fit_log2D = exp.(p.(x1D))

   alphaD = 1 + length(D) / sum(log.(D)) # calculate exponent by log-likelihood
   println("alphaD ", alphaD)
   alphaD, stD = estimate(D)
   println("alphaD ", alphaD, " pm ", stD)

   ax = PyPlot.axes([fig_x2, fig_y1, fig_width_x, fig_width_y])
       scatter(ex1,ey1,s=50,label="confirmed cases")
       plot(ex1,ey1,"-",linewidth=2)
       loglog(ex1,ey1fit_log2,"k", label=L"regression line, $\mu=1.58$")
       title("US counties", fontsize=16)
       xlim(0.8,1e5)
       ylim(1e-8,1)
       xticks(fontsize=16)
       yticks([1e-8,1e-7,1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=16)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "c)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix2, fig_y1+fig_Iy, fig_Iwidth_x, fig_Iwidth_y]) # Inset: cumulative distribution
       cc = 10 .^ range(0,5.0, length=50)
       loglog(cc, cumulpdf.(cc,muC,7e4), color="black", linewidth=2)
       scatter(C,(1:nC)/nC,s=20)
       loglog(C,(1:nC)/nC,"-")
       myfont = Dict("fontname"=>"Sans","fontsize"=>14)
       xticks(fontsize=14)
       yticks(fontsize=14)
       xlim(1,8e4)
       ylim(3e-4,1)
       xlabel("cases, n",fontsize=16)
       ylabel("C(n)",fontsize=16)

   ax = PyPlot.axes([fig_x2, fig_y2, fig_width_x, fig_width_y])
       scatter(ex1D,ey1D,s=50,label="confirmed deaths")
       loglog(ex1D,ey1D,"-",linewidth=2)
       loglog(ex1D,ey1fit_log2D,"k", label=L"regression line, $\mu=1.83$")
       xticks(fontsize=16)
       yticks([1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=16)
       xlim(0.8,1.5e3)
       ylim(1e-6,1)
       xlabel("cases, n",fontsize=18)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "d)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix2, fig_y2+fig_Iy, fig_Iwidth_x, fig_Iwidth_y]) # Inset: cumulative distribution
       cc = 10 .^ range(0,3, length=50)
       loglog(cc, cumulpdf.(cc,muD,3e3), color="black", linewidth=2)
       scatter(D,(1:nD)/nD,s=20)
       xticks([1,10,100,1000],fontsize=14)
       yticks(fontsize=14)
       xlim(1,1.5e3)
       ylim(1e-3,1)
       xlabel("cases, n",fontsize=16)
       ylabel("C(n)",fontsize=16)
end

fig_Iy = 0.25
fig_Ix1 = 0.37
fig_Ix2 = 0.84
fig_Iwidth_y = 0.12
fig_Iwidth_x = 0.12

fig_sep_y   = 0.06
fig_width_y = 0.4
fig_width_x = 0.4
fig_x1 = 0.1
fig_x2 = 0.57
fig_y2=0.08
fig_y1=fig_y2 + fig_width_y + fig_sep_y

figure("distribution", figsize=(12,8))
clf()
distribution_countries()
distribution_counties()
#savefig("../figures/fig1.pdf")

end # module


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

using Markdown
using InteractiveUtils

# ╔═╡ ceff6126-51fc-11eb-3cf5-dd30d637e267
begin
	using Dates
	using MortalityTables
	using DifferentialEquations
	using Plots
	using ActuaryUtilities
end

# ╔═╡ cd5e7080-5348-11eb-2c22-455a3556cab0
md"
# Universal Life Policy Account Mechanics as a Differential Equation
"

# ╔═╡ de0ace1a-5348-11eb-35c9-5f0ff4009d97
md"""
## Introduction

This demonstrates an example of defining an universal life policy value roll like a discrete differential equation.

It then uses the [SciML DifferentialEquations](https://sciml.ai/) package to "solve" the policy projection given a single point, but also to see how the policy projection behaves under different premium and interest rate conditions./
"""

# ╔═╡ ef26ffb6-5348-11eb-04a8-c5cbfdcab72a
md"""
Let's use the 2001 CSO table as the basis for cost of insurance charges:
"""

# ╔═╡ 45dab59c-528d-11eb-357d-af3664466910
begin
	tables = MortalityTables.tables()
	cso = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"]
end

# ╔═╡ ffe9b6e0-5348-11eb-2d38-a5613eb24383
md"""
Next, policy mechanics are coded. It's essentially a discrete differential equation, so it leverages [`DifferentialEquations.jl`](https://diffeq.sciml.ai/dev/)

The projection is coded in the [Discrete DifferentialEquation format](https://diffeq.sciml.ai/dev/types/discrete_types/):

``u_{n+1} = f(u,p,t_{n+1})``

In the code below, this translates to:

- `u` is the *state* of the system. We  will track three variables to represent the `state`:
   - `state[1]` is the account value
   - `state[2]` is the premium paid
   - `state[3]` is the policy duration
- `p` are the parameters of the system.
- `t` is the time, which will represent days since policy issuance
"""

# ╔═╡ da199644-51fc-11eb-0aab-dd68b1fcddc6
function policy_projection(state,p,t)
    # grab the state from the inputs
    av = state[1]

    # calculated variables
    cur_date = p.issue_date + Day(t)
    dur = duration(p.issue_date,cur_date)
    att_age = p.issue_age + dur - 1

	# lapse if AV <= 0
    lapsed = (av <= 0.0 ) & (t > 1)

    if !lapsed

        monthly_coi_rate = (1 - (1-p.mort_assump[att_age]) ^ (1/12))

		## Periodic Policy elements

        # annual events
        if Dates.monthday(cur_date) == Dates.monthday(p.issue_date) ||
			cur_date ==p.issue_date + Day(1) # OR first issue date

            premium = p.annual_prem
        else
            premium = 0.0
        end

        # monthly_events
        if Dates.day(cur_date) == Dates.day(p.issue_date)
            coi = max((p.face - av) * monthly_coi_rate,0.0)
        else
            coi = 0.0
        end

        # daily events
        int(av) = av * ((1 + p.int_rate) ^ (1 / 360) - 1.0)

        # av
        new_av = max(0.0,av - coi + premium + int(av-coi))

				# new state
		return [new_av, premium, dur] # AV, Prem, Dur

    else
		# new state
		return [0.0, 0.0, dur] # AV, Prem, Dur

    end

end

# ╔═╡ 8a227236-5349-11eb-353d-1d3f94a968d0
md"The following function will create a named tuple of parameters given a varying `prem` (premium) and `int` (credit rate)."

# ╔═╡ 956811e4-5349-11eb-2b34-b5441def32a1
params(prem,int) = (
    int_rate = int,
    issue_date = Date(2010,1,1),
    face = 1e6,
    issue_age = 25,
    mort_assump = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"].ultimate,
    projection_years = 75,
    annual_prem = prem,
)

# ╔═╡ 9ee35710-5349-11eb-0d45-23bc292d8a3d
md"""
## Runing the system

This results in the following plot. The tracked output variables u1 and u2 represent the two vars that we tracked above: account value and cumulative premium.
"""

# ╔═╡ b65ef28c-5349-11eb-1baa-6575ad7fb946
begin
	p = params(
	        8000.0, # 8,000 annual premium
	        0.08    # 8% interest
	        )

	# calculate the number of days to project
	projection_end_date = p.issue_date + Year(p.projection_years)
	days_to_project = Dates.value(projection_end_date - p.issue_date)

	# the [0.0,..] are the initial conditions for the tracked variables
	 prob = DiscreteProblem(policy_projection,[0.0,0.0,0],(0,days_to_project),p)
	 proj = solve(prob,FunctionMap())

	 plot(proj)
end

# ╔═╡ c4637664-5349-11eb-0e81-59e39974f02e
md"""
## Moving up the ladder of abstraction

An excellent way to understand the behavior of a model is to [move up the ladder of abstraction](http://worrydream.com/LadderOfAbstraction/). Below, we will see what happens to the projection at varying levels of credit rates and annual premiums.
"""

# ╔═╡ 02cc352a-5409-11eb-0fda-336d5fa42fa9
begin
	prem_range = 1000.0:100.0:9000.0
	int_range = 0.02:0.0025:0.08

	function ending_av(ann_prem,int,days_to_project)
	    p = params(ann_prem,int)
	    prob = DiscreteProblem(policy_projection,[0.0,0.0,0],(0,days_to_project),p)
	    proj = solve(prob,FunctionMap())
	    end_av = proj[end][1]
	    if end_av == 0.0
	        lapse_time = findfirst(isequal(0.0),proj[1,2:end])
	    else
	        lapse_time = length(proj)
	    end
	    duration = proj[3,lapse_time]
	    end_age = p.issue_age + duration - 1.0

	    return end_av,end_age

	end

	end_age = zeros(length(prem_range),length(int_range))
	end_av = zeros(length(prem_range),length(int_range))

	# loop through each projection we did and fill our ranges with the ending AV and ending age
	for (i,vp) in enumerate(prem_range)
	    for (j,vi) in enumerate(int_range)
	        (end_av[i,j],end_age[i,j]) = ending_av(vp,vi,days_to_project)
	    end
	end
end

# ╔═╡ 280d25a6-5409-11eb-17d9-674f218f4351
begin
	using ColorSchemes # for Turbo colors, which emphasize readability

	viz = plot(layout=2) # side by side plot

	#
	contour!(viz[1],int_range,
	    prem_range,
	    end_av ./ 1e6, # scale to millions for readability

	    contour_labels=true,
	    c=cgrad(ColorSchemes.turbo.colors),
	    fill=true,
	    title="AV at age 100 (\$M)",
	    ylabel="Annual Premium (\$)",

	)

	contour!(viz[2],int_range,
	    prem_range,
	    end_age,

	    contour_labels=true,
	    c=cgrad(ColorSchemes.turbo.colors),
	    fill=true,
	    yaxis=false,
	    title="Age at Lapse",
	)

	annotate!(viz[2],0.055,7000,Plots.text("Doesn't lapse \nbefore age 100", 8, :white, :center))

end

# ╔═╡ 21e53f60-5409-11eb-291d-d903da8e16b4
md"
Now let's plot the result. Not surprising, interest has a huge effect on the policy projection. Premium is also a major influence.

One thing that's remarkable is how going from 2000 premium to just ~2200 of premium results in about a $5m difference at 8% interest. The power of compound interest!
"

# ╔═╡ 684f7efc-5409-11eb-0757-19b40c5d418e
md"
## Conclusion

This shows how universal life mechanics are a dynamic system. the growth/decay is governed by two competing feedback loops:

- Growth: the force of interest lets the balance grow exponentially over long periods of time

- Decay: low balances increase the net amount at risk and the resulting COI charges.

"

# ╔═╡ 7e80faf0-5409-11eb-33f6-4f6be6fea835
md"
## Endnotes

This is not meant to represent any particular insurance product, nor fully replicate typical account mechanics.
"

# ╔═╡ Cell order:
# ╟─cd5e7080-5348-11eb-2c22-455a3556cab0
# ╟─de0ace1a-5348-11eb-35c9-5f0ff4009d97
# ╠═ceff6126-51fc-11eb-3cf5-dd30d637e267
# ╟─ef26ffb6-5348-11eb-04a8-c5cbfdcab72a
# ╠═45dab59c-528d-11eb-357d-af3664466910
# ╟─ffe9b6e0-5348-11eb-2d38-a5613eb24383
# ╠═da199644-51fc-11eb-0aab-dd68b1fcddc6
# ╟─8a227236-5349-11eb-353d-1d3f94a968d0
# ╠═956811e4-5349-11eb-2b34-b5441def32a1
# ╟─9ee35710-5349-11eb-0d45-23bc292d8a3d
# ╠═b65ef28c-5349-11eb-1baa-6575ad7fb946
# ╟─c4637664-5349-11eb-0e81-59e39974f02e
# ╠═02cc352a-5409-11eb-0fda-336d5fa42fa9
# ╟─21e53f60-5409-11eb-291d-d903da8e16b4
# ╠═280d25a6-5409-11eb-17d9-674f218f4351
# ╟─684f7efc-5409-11eb-0757-19b40c5d418e
# ╟─7e80faf0-5409-11eb-33f6-4f6be6fea835

max_stars_repo_name: EarthGoddessDude/Learn

max_issues_repo_name: EarthGoddessDude/Learn

max_forks_repo_name: EarthGoddessDude/Learn


text:
@reexport using Colors, ColorVectorSpace, ColorSchemes

_ColorType = RGB{Float64}
_bgcolor = _ColorType(1,1,1) # White
_fgcolor = _ColorType(0,0,0) # Black
#_color = _ColorType(0,0,1) # Blue
#_gridcolor = _ColorType(0.5,0.5,0.5) # Gray
#_featurecolor = _ColorType(0.5,0.5,0.5) # Gray
_colorscheme = colorschemes[:jet]
_colorarray = Colorant[]

mutable struct _ColorConfig

end

"Set the color type for pixel images."
function colortype(T::DataType)
    @assert T <: Colorant
    global _ColorType = T
end # function

"Get the color type for pixel images."
colortype() = _ColorType

function _calculatecolorarray(ncolors::Integer)
    if ncolors <= 1
        global _colorarray = [ _colorscheme[1] ]
        return _colorarray
    end
    N = ncolors-1
    global _colorarray = [ _colorscheme[n/N] for n in 0:N ]
end

_calculatecolorarray(100)

"Update the color array to new size."
function updatecolorarray(ncolors::Integer)
    if ncolors == length(_colorarray) return end
    _calculatecolorarray(ncolors)
end # function

"Get the current Color Scheme."
colorarray() = _colorarray

"""
    colorscheme(symbol,invert=false)

Set the **ColorScheme** (to define a color map), from identifier symbol.

#### Arguments
- `sym::Symbol`: Valid colorscheme symbol identifier from **ColorSchemes.jl**.
- `invert::Bool`: If true, create the color scheme with reverse colors order.
"""
function colorscheme(sym::Symbol, invert::Bool=false)
    global _colorscheme = invert ? ColorScheme(reverse(colorschemes[sym].colors)) : colorschemes[sym]
    _calculatecolorarray(length(_colorarray))
    _colorscheme
end # function

"Set the color map (setting a **ColorScheme**), from identifier symbol."
colormap(sym::Symbol, invert::Bool=false) = colorscheme(sym,invert)

"Set the **ColorScheme** to define a color map."
function colorscheme(cs::ColorScheme, invert::Bool=false)
    global _colorscheme = invert ? ColorScheme(reverse(cs.colors)) : cs
    _calculatecolorarray(length(_colorarray))
end # function

"Set the color map (setting a **ColorScheme**), from identifier symbol."
colormap(cs::ColorScheme, invert::Bool=false) = colorscheme(sym,invert)

"Get the current color map as a **ColorScheme**."
colorscheme() = _colorscheme

"Get the current color map as a **ColorScheme**."
colormap() = _colorscheme

"Set the current color from n-th color in the color array."
color(n::Integer) = _backend.color(convert(_ColorType, _colorarray[n]))

"Set the current color from t-interpolated color in the color scheme."
colorinterpolation(t::Real) = _backend.color(convert(_ColorType, _colorscheme[t]))

"Set the current color using the current coloring function."
color(x::Real, y::Real) = _backend.color(convert(_ColorType, _coloringfunction(x,y)))

"Set the current color using the current coloring function."
color(z::Number) = _backend.color(convert(_ColorType, _coloringfunction(z)))

"Set the current color."
color(c::Colorant) = _backend.color(convert(_ColorType, c))

"Set the current color using RGBA."
color(r::Real,g::Real,b::Real,a::Real=1) =
_backend.color(convert(_ColorType, RGBA{Float64}(r,g,b,a)))

"Get the current color"
color() = _backend.color()

"Set the current background color."
bgcolor(c::Colorant) = global _bgcolor = convert(_ColorType, c)

"Set the current background color using RGBA."
bgcolor(r::Real,g::Real,b::Real,a::Real=1) =
global _bgcolor = convert(_ColorType, RGBA{Float64}(r,g,b,a))

"Get the current background color"
bgcolor() = _bgcolor

"Set the current foreground color."
fgcolor(c::Colorant) = global _fgcolor = convert(_ColorType, c)

"Set the current foreground color using RGBA."
fgcolor(r::Real,g::Real,b::Real,a::Real=1) =
global _fgcolor = convert(_ColorType, RGBA{Float64}(r,g,b,a))

"Get the current foreground color"
fgcolor() = _fgcolor

"""
    colorinterpolationbg(t, n)

Set the current color as the  linear color interpolation using the n-th color
in the current colormap to the backgrund color.

#### Arguments
- `t::Real`: Number between \$0\$ and \$1\$.
- `n::Int`: Color index in the color array from colormap.
"""
colorinterpolationbg(t::Real, n::Int) =
_backend.color(convert(_ColorType, (1-t)*_colorarray[n]+t*_bgcolor))

"""
    colorinterpolationfg(t, n)

Set the current color as the  linear color interpolation using the n-th color
in the current colormap to the foreground color.

#### Arguments
- `t::Real`: Number between \$0\$ and \$1\$.
- `n::Int`: Color index in the color array from colormap.
"""
colorinterpolationfg(t::Real, n::Int) =
_backend.color(convert(_ColorType, (1-t)*_colorarray[n]+t*_fgcolor))


text:
const WFLOAT = Sys.WORD_SIZE == 64 ? Float64 : Float32

#=========
Abstraction
==========#
abstract type RobustVariance <: CovarianceEstimator end
abstract type HAC{G} <: RobustVariance end
abstract type HC <: RobustVariance end
abstract type CRHC{V,D} <: RobustVariance end

#=========
HAC Types
=========#

abstract type BandwidthType end
#abstract type OptimalBandwidth end

struct NeweyWest<:BandwidthType end
struct Andrews<:BandwidthType end
struct Fixed<:BandwidthType end

#struct Optimal{G<:OptimalBandwidth}<:BandwidthType{G where G<:OptimalBandwidth} end

# struct Prewhiten end
# struct Unwhiten end
"""
    `TruncatedKernel`

# Constructors
    TruncatedKernel{Fixed}(x::Int)
    TruncatedKernel{Andrews}()
    TruncatedKernel{NeweyWest}()
# Note
- `Fixed`: fixed bandwidth
- `Andrews`: bandwidth selection a la Andrews
- `NeweyWest`: bandwidth selection a la Andrews
"""
struct TruncatedKernel{G<:BandwidthType}<:HAC{G}
  bw::Vector{WFLOAT}
  weights::Vector{WFLOAT}
end

"""
    `BartlettKernel`

# Constructors
    BartlettKernel(x::Int)
    BartlettKernel(::Type{Andrews})
    BartlettKernel(::Type{NeweyWest})
# Note
- `Andrews`: bandwidth selection a la Andrews
- `NeweyWest`: bandwidth selection a la Andrews
"""
struct BartlettKernel{G<:BandwidthType}<:HAC{G}
    bw::Vector{WFLOAT}
    weights::Vector{WFLOAT}
end

"""
    `ParzenKernel`

# Constructors
    ParzenKernel(x::Int)
    ParzenKernel(::Type{Andrews})
    ParzenKernel(::Type{NeweyWest})
# Note
- `Andrews`: bandwidth selection a la Andrews
- `NeweyWest`: bandwidth selection a la Andrews
"""
struct ParzenKernel{G<:BandwidthType}<:HAC{G}
    bw::Vector{WFLOAT}
    weights::Vector{WFLOAT}
end

"""
    `TukeyHanningKernel`

# Constructors
    TukeyHanningKernel(x::Int)
    TukeyHanningKernel(::Type{Andrews})
    TukeyHanningKernel(::Type{NeweyWest})
# Note
- `Andrews`: bandwidth selection a la Andrews
- `NeweyWest`: bandwidth selection a la Andrews
"""
struct TukeyHanningKernel{G<:BandwidthType}<:HAC{G}
    bw::Vector{WFLOAT}
    weights::Vector{WFLOAT}
end

"""
    `QuadraticSpectralKernel`

# Constructors
    QuadraticSpectralKernel(x::Int)
    QuadraticSpectralKernel(::Type{Andrews})
    QuadraticSpectralKernel(::Type{NeweyWest})
# Note
- `Andrews`: bandwidth selection a la Andrews
- `NeweyWest`: bandwidth selection a la Andrews
"""
struct QuadraticSpectralKernel{G<:BandwidthType}<:HAC{G}
    bw::Vector{WFLOAT}
    weights::Vector{WFLOAT}
end

mutable struct VARHAC <: RobustVariance
    maxlag::Int
    lagstrategy::Int
    selectionstrategy::Symbol
end

## ------------------------------------------------------------------------------------------
# Define HAC constructor
# Three type of constructors
# - HAC(::Andrews) or HAC(::NeweyWest) Optimal bandwidth selection a la Andrews or a la NeweyWest
# - HAC(bw) Fixed bandwidth
# - HAC() -> Optimal bandwidth selection a la Andrews

# Constructors
const kernels = [
    :BartlettKernel,
    :ParzenKernel,
    :QuadraticSpectralKernel,
    :TruncatedKernel,
    :TukeyHanningKernel,
]

for kerneltype in kernels
    @eval ($kerneltype(x::Number)) = ($kerneltype){Fixed}(WFLOAT[x], WFLOAT[])
    @eval ($kerneltype{Fixed}(x::Number)) = ($kerneltype){Fixed}(WFLOAT[x], WFLOAT[])
end

for kerneltype in kernels
    for opt in [:Andrews, :NeweyWest]
        if !(opt == :NeweyWest && kerneltype in [:TukeyHanningKernel, :TruncatedKernel])
            @eval ($kerneltype){$opt}() = ($kerneltype){$opt}(WFLOAT[0], WFLOAT[])
        else
            msg = "$kerneltype does not support Newey-West optimal bandwidth"
            @eval ($kerneltype){$opt}() = throw(ArgumentError($msg))
        end
    end
end

function Base.show(io::IO, x::HAC{T}) where T<:Union{Andrews, NeweyWest}
    print(typeof(x).name, "{", T,"}")
end
function Base.show(io::IO, x::HAC{T}) where T<:Fixed
    print(typeof(x).name, "{", T,"}(", first(x.bw), ")")
end
## Makes the default bandwidth selection
#Optimal() = Optimal{Andrews}()

## Accessor
bandwidth(x::HAC) = x.bw
kernelweights(x::HAC) = x.weights

#=========
HC Types
=========#
struct HC0  <: HC end
struct HC1  <: HC end
struct HC2  <: HC end
struct HC3  <: HC end
struct HC4  <: HC end
struct HC4m <: HC end
struct HC5  <: HC end

#=========
HC Types
=========#

mutable struct CRHC0{V, D}  <: CRHC{V,D}
    cl::V
    df::D
end

mutable struct CRHC1{V, D}  <: CRHC{V,D}
    cl::V
    df::D
end

mutable struct CRHC2{V, D}  <: CRHC{V,D}
    cl::V
    df::D
end

mutable struct CRHC3{V, D}  <: CRHC{V,D}
    cl::V
    df::D
end

## CRHC accessor
clusterindicator(x::CRHC) = x.cl

## CRHC Constructors
for tp in [:CRHC0, :CRHC1, :CRHC2, :CRHC3]
    @eval $(tp)() = $(tp)(nothing, nothing)
end

for tp in [:CRHC0, :CRHC1, :CRHC2, :CRHC3]
    @eval $(tp)(v::AbstractVector) = $(tp)(v, nothing)
end

struct CovarianceMatrix{T2<:Factorization, T3<:CovarianceMatrices.RobustVariance, F1, T1<:AbstractMatrix{F1}} <: AbstractMatrix{F1}
    F::T2       ## Factorization
    K::T3       ## RobustVariance, e.g. HC0()
    V::T1       ## Covariance matrix
end


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

abstract type AbstractFeasibilityAlgorithm end
abstract type AbstractFeasibilityStrategy end

"""
    NoFeasibilityAlgorithm

Empty functor object for running an L-shaped algorithm without dealing with second-stage feasibility.
"""
struct NoFeasibilityAlgorithm <: AbstractFeasibilityAlgorithm end

handle_feasibility(::NoFeasibilityAlgorithm) = false
num_cuts(::NoFeasibilityAlgorithm) = 0
active!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing
deactive!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing
restore!(::MOI.ModelLike, ::NoFeasibilityAlgorithm) = nothing

"""
    FeasibilityCutsMaster

Master functor object for using feasibility cuts in an L-shaped algorithm. Create by supplying a [`FeasibilityCuts`](@ref) object through `feasibility_strategy` in `LShaped.Optimizer` or set the [`FeasibilityStrategy`](@ref) attribute.
"""
struct FeasibilityCutsMaster{T <: AbstractFloat} <: AbstractFeasibilityAlgorithm
    cuts::Vector{SparseFeasibilityCut{T}}

    function FeasibilityCutsMaster(::Type{T}) where T <: AbstractFloat
        return new{T}(Vector{SparseFeasibilityCut{T}}())
    end
end

handle_feasibility(::FeasibilityCutsMaster) = true
worker_type(::FeasibilityCutsMaster) = FeasibilityCutsWorker
num_cuts(feasibility::FeasibilityCutsMaster) = length(feasibility.cuts)
active!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing
deactive!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing
restore!(::MOI.ModelLike, ::FeasibilityCutsMaster) = nothing

"""
    FeasibilityCutsWorker

Worker functor object for using feasibility cuts in an L-shaped algorithm. Create by supplying a [`FeasibilityCuts`](@ref) object through `feasibility_strategy` in `LShaped.Optimizer` or set the [`FeasibilityStrategy`](@ref) attribute.
"""
mutable struct FeasibilityCutsWorker <: AbstractFeasibilityAlgorithm
    objective::MOI.AbstractScalarFunction
    linking_constraints::Vector{MOI.ConstraintIndex}
    feasibility_variables::Vector{MOI.VariableIndex}
    aux_constraint::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}

    function FeasibilityCutsWorker(objective::MOI.AbstractScalarFunction,
                                   linking_constraints::Vector{MOI.ConstraintIndex},
                                   feasibility_variables::Vector{MOI.VariableIndex})
        return new(objective, linking_constraints, feasibility_variables, CI{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}(0))
    end
end

handle_feasibility(::FeasibilityCutsWorker) = true
num_cuts(::FeasibilityCutsWorker) = 0

function prepare!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Set objective to zero
    G = MOI.ScalarAffineFunction{Float64}
    MOI.set(model, MOI.ObjectiveFunction{G}(), zero(MOI.ScalarAffineFunction{Float64}))
    i = 1
    # Create auxiliary feasibility variables
    for ci in worker.linking_constraints
        i = add_auxiliary_variables!(model, worker, ci, i)
    end
    return nothing
end
function prepared(worker::FeasibilityCutsWorker)
    return length(worker.feasibility_variables) > 0
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    # Nothing to do for most most constraints
    return idx
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: AffineDecisionFunction, S <: MOI.AbstractScalarSet}
    G = MOI.ScalarAffineFunction{Float64}
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    # Positive feasibility variable
    pos_aux_var = MOI.add_variable(model)
    name = add_subscript(:v⁺, idx)
    MOI.set(model, MOI.VariableName(), pos_aux_var, name)
    push!(worker.feasibility_variables, pos_aux_var)
    # Nonnegativity constraint
    MOI.add_constraint(model, MOI.SingleVariable(pos_aux_var),
                       MOI.GreaterThan{Float64}(0.0))
    # Add to objective
    MOI.modify(model, MOI.ObjectiveFunction{G}(),
               MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
    # Add to constraint
    MOI.modify(model, ci, MOI.ScalarCoefficientChange(pos_aux_var, 1.0))
    # Negative feasibility variable
    neg_aux_var = MOI.add_variable(model)
    name = add_subscript(:v⁻, idx)
    MOI.set(model, MOI.VariableName(), neg_aux_var, name)
    push!(worker.feasibility_variables, neg_aux_var)
    # Nonnegativity constraint
    MOI.add_constraint(model, MOI.SingleVariable(neg_aux_var),
                       MOI.GreaterThan{Float64}(0.0))
    # Add to objective
    MOI.modify(model, MOI.ObjectiveFunction{G}(),
               MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
    # Add to constraint
    MOI.modify(model, ci, MOI.ScalarCoefficientChange(neg_aux_var, -1.0))
    # Return updated identification index
    return idx + 1
end

function add_auxiliary_variables!(model::MOI.ModelLike,
                                  worker::FeasibilityCutsWorker,
                                  ci::CI{F,S},
                                  idx::Integer) where {F <: VectorAffineDecisionFunction, S <: MOI.AbstractVectorSet}
    G = MOI.ScalarAffineFunction{Float64}
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    n = MOI.dimension(MOI.get(model, MOI.ConstraintSet(), ci))
    for (i, id) in enumerate(idx:(idx + n - 1))
        # Positive feasibility variable
        pos_aux_var = MOI.add_variable(model)
        name = add_subscript(:v⁺, id)
        MOI.set(model, MOI.VariableName(), pos_aux_var, name)
        push!(worker.feasibility_variables, pos_aux_var)
        # Nonnegativity constraint
        MOI.add_constraint(model, MOI.SingleVariable(pos_aux_var),
                           MOI.GreaterThan{Float64}(0.0))
        # Add to objective
        MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        # Add to constraint
        MOI.modify(model, ci, MOI.MultirowChange(pos_aux_var, [(i, 1.0)]))
    end
    for (i, id) in enumerate(idx:(idx + n - 1))
        # Negative feasibility variable
        neg_aux_var = MOI.add_variable(model)
        name = add_subscript(:v⁻, id)
        MOI.set(model, MOI.VariableName(), neg_aux_var, name)
        push!(worker.feasibility_variables, neg_aux_var)
        # Nonnegativity constraint
        MOI.add_constraint(model, MOI.SingleVariable(neg_aux_var),
                           MOI.GreaterThan{Float64}(0.0))
        # Add to objective
        MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        # Add to constraint
        MOI.modify(model, ci, MOI.MultirowChange(neg_aux_var, [(i, -1.0)]))
    end
    # Return updated identification index
    return idx + n + 1
end

function activate!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Set objective to zero
    G = MOI.ScalarAffineFunction{Float64}
    MOI.set(model, MOI.ObjectiveFunction{G}(), zero(MOI.ScalarAffineFunction{Float64}))
    obj_sense = MOI.get(model, MOI.ObjectiveSense())
    # Add auxiliary variables to objective and linking constraints
    idx = 0
    for ci in worker.linking_constraints
        dim = MOI.dimension(MOI.get(model, MOI.ConstraintSet(), ci))
        for i in 1:dim
            pos_aux_var = worker.feasibility_variables[idx + 2*(i-1) + 1]
            MOI.modify(model, MOI.ObjectiveFunction{G}(),
                   MOI.ScalarCoefficientChange(pos_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
            neg_aux_var = worker.feasibility_variables[idx + 2*(i-1) + 2]
            MOI.modify(model, MOI.ObjectiveFunction{G}(),
                       MOI.ScalarCoefficientChange(neg_aux_var, obj_sense == MOI.MAX_SENSE ? -1.0 : 1.0))
        end
        idx += 2*dim
    end
    if MOI.is_valid(model, worker.aux_constraint)
        MOI.delete(model, worker.aux_constraint)
    end
    return nothing
end

function deactivate!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Force auxiliary variables to zero
    func_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj = MOI.get(model, MOI.ObjectiveFunction{func_type}())
    worker.aux_constraint = MOI.add_constraint(model, obj, MOI.EqualTo{Float64}(0.0))
    # Restore objective
    F = typeof(worker.objective)
    MOI.set(model, MOI.ObjectiveFunction{F}(), worker.objective)
    return nothing
end

function restore!(model::MOI.ModelLike, worker::FeasibilityCutsWorker)
    # Remove aux constraint
    if MOI.is_valid(model, worker.aux_constraint)
        MOI.delete(model, worker.aux_constraint)
    end
    worker.aux_constraint = CI{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}(0)
    # Delete any feasibility variables
    if prepared(worker)
        MOI.delete(model, worker.feasibility_variables)
    end
    empty!(worker.feasibility_variables)
    # Restore objective
    F = typeof(worker.objective)
    MOI.set(model, MOI.ObjectiveFunction{F}(), worker.objective)
    return nothing
end

# API
# ------------------------------------------------------------
"""
    IgnoreFeasibility

Factory object for [`NoFeasibilityAlgorithm`](@ref). Passed by default to `feasibility_strategy` in `LShaped.Optimizer`.

"""
struct IgnoreFeasibility <: AbstractFeasibilityStrategy end

function master(::IgnoreFeasibility, ::Type{T}) where T <: AbstractFloat
    return NoFeasibilityAlgorithm()
end

function worker(::IgnoreFeasibility, ::Vector{MOI.ConstraintIndex}, ::MOI.ModelLike)
    return NoFeasibilityAlgorithm()
end
function worker_type(::IgnoreFeasibility)
    return NoFeasibilityAlgorithm
end

"""
    IgnoreFeasibility

Factory object for using feasibility cuts in an L-shaped algorithm.

"""
struct FeasibilityCuts <: AbstractFeasibilityStrategy end

function master(::FeasibilityCuts, ::Type{T}) where T <: AbstractFloat
    return FeasibilityCutsMaster(T)
end

function worker(::FeasibilityCuts, linking_constraints::Vector{MOI.ConstraintIndex}, model::MOI.ModelLike)
    # Cache objective
    func_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj = MOI.get(model, MOI.ObjectiveFunction{func_type}())
    return FeasibilityCutsWorker(obj, linking_constraints, Vector{MOI.VariableIndex}())
end
function worker_type(::FeasibilityCuts)
    return FeasibilityCutsWorker
end


text:
struct CPUEvent <: Event
    task::Union{Nothing, Core.Task}
end

function Event(::CPU)
    return CPUEvent(nothing)
end

wait(ev::CPUEvent, progress=nothing) = wait(CPU(), ev, progress)
function wait(::CPU, ev::CPUEvent, progress=nothing)
    ev.task === nothing && return

    if progress === nothing
        wait(ev.task)
    else
        while !Base.istaskdone(ev.task)
            progress()
            yield() # yield to the scheduler
        end
    end
end

function (obj::Kernel{CPU})(args...; ndrange=nothing, workgroupsize=nothing, dependencies=nothing)
    if ndrange isa Integer
        ndrange = (ndrange,)
    end
    if workgroupsize isa Integer
        workgroupsize = (workgroupsize, )
    end
    if dependencies isa Event
        dependencies = (dependencies,)
    end

    if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
        workgroupsize = (1024,) # Vectorization, 4x unrolling, minimal grain size
    end
    iterspace, dynamic = partition(obj, ndrange, workgroupsize)
    # partition checked that the ndrange's agreed
    if KernelAbstractions.ndrange(obj) <: StaticSize
        ndrange = nothing
    end

    t = __run(obj, ndrange, iterspace, args, dependencies)
    return CPUEvent(t)
end

# Inference barrier
function __run(obj, ndrange, iterspace, args, dependencies)
    return Threads.@spawn begin
        if dependencies !== nothing
            cpu_tasks = Core.Task[]
            for event in dependencies
                if event isa CPUEvent
                    push!(cpu_tasks, event.task)
                end
            end
            !isempty(cpu_tasks) && Base.sync_end(cpu_tasks)
            for event in dependencies
                if !(event isa CPUEvent)
                    wait(CPU(), event, ()->yield())
                end
            end
        end
        @sync begin
            # TODO: how do we use the information that the iteration space maps perfectly to
            #       the ndrange without incurring a 2x compilation overhead
            # if dynamic
                for block in iterspace
                    let ctx = mkcontextdynamic(obj, block, ndrange, iterspace)
                        Threads.@spawn Cassette.overdub(ctx, obj.f, args...)
                    end
                end
            # else
            #     for block in iterspace
            #         let ctx = mkcontext(obj, blocks, ndrange, iterspace)
            #             Threads.@spawn Cassette.overdub(ctx, obj.f, args...)
            #         end
            #     end
            # end
        end
    end
end

Cassette.@context CPUCtx

function mkcontext(kernel::Kernel{CPU}, I, _ndrange, iterspace)
    metadata = CompilerMetadata{ndrange(kernel), false}(I, _ndrange, iterspace)
    Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
end

function mkcontextdynamic(kernel::Kernel{CPU}, I, _ndrange, iterspace)
    metadata = CompilerMetadata{ndrange(kernel), true}(I, _ndrange, iterspace)
    Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Local_Linear), idx::CartesianIndex)
    indices = workitems(__iterspace(ctx.metadata))
    return @inbounds LinearIndices(indices)[idx]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Linear), idx::CartesianIndex)
    indices = blocks(__iterspace(ctx.metadata))
    return @inbounds LinearIndices(indices)[__groupindex(ctx.metadata)]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Linear), idx::CartesianIndex)
    I = @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
    @inbounds LinearIndices(__ndrange(ctx.metadata))[I]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Local_Cartesian), idx::CartesianIndex)
    return idx
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Cartesian), idx::CartesianIndex)
    __groupindex(ctx.metadata)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Cartesian), idx::CartesianIndex)
    return @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__validindex), idx::CartesianIndex)
    # Turns this into a noop for code where we can turn of checkbounds of
    if __dynamic_checkbounds(ctx.metadata)
        I = @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
        return I in __ndrange(ctx.metadata)
    else
        return true
    end
end

generate_overdubs(CPUCtx)

###
# CPU implementation of shared memory
###
@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(SharedMemory), ::Type{T}, ::Val{Dims}, ::Val) where {T, Dims}
    MArray{__size(Dims), T}(undef)
end

###
# CPU implementation of scratch memory
# - private memory for each workitem
# - memory allocated as a MArray with size `Dims + WorkgroupSize`
###
struct ScratchArray{N, D}
    data::D
    ScratchArray{N}(data::D) where {N, D} = new{N, D}(data)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(Scratchpad), ::Type{T}, ::Val{Dims}) where {T, Dims}
    return ScratchArray{length(Dims)}(MArray{__size((Dims..., __groupsize(ctx.metadata)...)), T}(undef))
end

Base.@propagate_inbounds function Base.getindex(A::ScratchArray, I...)
    return A.data[I...]
end

Base.@propagate_inbounds function Base.setindex!(A::ScratchArray, val, I...)
    A.data[I...] = val
end


text:
using CacheVariables, BSON, Test

## Add data directory, define data file path
dirpath = joinpath(@__DIR__, "data")
isdir(dirpath) && error("Test directory already has a data subdirectory.")
path = joinpath(dirpath, "test.bson")

## Test save functionality
@testset "Save" begin
    # 1. verify log message
    @test_logs (:info, "Saving to $path\nx\ny\nz") (@cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end)

    rm(path)
    out = @cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end

    # 3. did variables enter workspace correctly?
    @test x == [1, 2, 3]
    @test y == 4
    @test z == "test"
    @test out == "final output"
end

## Test load functionality
@testset "Load" begin
    # 1. set all variables to nothing
    x = nothing
    y = nothing
    z = nothing
    out = nothing

    # 2. file exists: load variables from it
    # verify log message
    @test_logs (:info, "Loading from $path\nx\ny\nz") (@cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end)

    # load variables
    out = @cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end

    # 3. did variables enter workspace correctly?
    @test x == [1, 2, 3]
    @test y == 4
    @test z == "test"
    @test out == "final output"
end

## Test overwrite behavior
@testset "Overwrite" begin
    # 1. change file contents
    bson(path; x = nothing, y = nothing, z = nothing, ans = nothing)

    # 2. add `true` to @cache call to overwrite
    # validate log message
    @test_logs (:info, "Overwriting $path\nx\ny\nz") (@cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end true)

    # overwrite data file
    overwrite = true
    @cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end overwrite

    # 3. check that correct data was overwritten
    x = nothing
    y = nothing
    z = nothing
    out = nothing
    out = @cache path begin
        x = collect(1:3)
        y = 4
        z = "test"
        "final output"
    end

    @test x == [1, 2, 3]
    @test y == 4
    @test z == "test"
    @test out == "final output"
end

## Test no variable case (where we just save ans)
@testset "Only ans" begin
    # 0. Clean up
    rm(path)

    # 1. verify log message
    @test_logs (:info, "Saving to $path\n") (@cache path 2 + 3)

    # 2. cache ans
    rm(path)
    out = @cache path 2 + 3

    # 3. did variables enter workspace correctly?
    @test out == 5

    # 4. set variable to nothing
    out = nothing

    # 5. file exists: load variables from it
    # verify log message
    @test_logs (:info, "Loading from $path\n") (@cache path 2 + 3)

    # load variables
    out = @cache path 2 + 3

    # 6. did variables enter workspace correctly?
    @test out == 5
end

## Clean up
rm(dirpath; recursive = true)