text:
using UnitfulCosmo
using Test

@testset "planck" begin
    # from Mukhanov. temperature, length, mass, time. no test for current yet
    @test isapprox(planck(2.725u"K"), 1.93e-32, rtol=1e-2)
    @test isapprox(unplanck(u"K", 1.93e-32), 2.725u"K", rtol=1e-2)
    @test isapprox(planck(310u"cm^-3"), 1.31e-96, rtol=1e-2)
    @test isapprox(unplanck(u"cm^-3", 1.31e-96), 310u"cm^-3", rtol=1e-2)
    @test isapprox(planck(1u"s"), 1.86e43, rtol=1e-2)
    @test isapprox(unplanck(u"s", 1.86e43), 1u"s", rtol=1e-2)
    @test isapprox(planck(4.45e5u"g*cm^-3"), 8.63e-89, rtol=1e-2)
    @test isapprox(unplanck(u"g*cm^-3", 8.63e-89), 4.45e5u"g*cm^-3", rtol=1e-2)
    @test_throws Unitful.DimensionError unplanck(u"Mpc", 1u"Mpc^-1")
end

@testset "mpc" begin
    @test mpc(1u"Mpc") ≈ 1u"Mpc"
    @test unmpc(u"Mpc", 1u"Mpc") ≈ 1u"Mpc"
    @test mpc(1u"yr") ≈ 3.064e-7u"Mpc" rtol=1e-2
    @test unmpc(u"yr", 3.064e-7u"Mpc") ≈ 1u"yr" rtol=1e-2
    @test mpc(10u"kg") ≈ 8.772e65u"Mpc^-1" rtol=1e-2
    @test unmpc(u"kg", 8.772e65u"Mpc^-1") ≈ 10u"kg" rtol=1e-2
    @test mpc(273u"K") ≈ 3.679e27u"Mpc^-1" rtol=1e-2
    @test unmpc(u"K", 3.679e27u"Mpc^-1") ≈ 273u"K" rtol=1e-2
    @test mpc(10u"A") ≈ 1.945391641e33u"Mpc^-1" rtol=1e-2
    @test unmpc(u"A", 1.945391641e33u"Mpc^-1") ≈ 10u"A" rtol=1e-2
    @test mpc(10u"kg"; base=u"kpc") ≈ 8.772e62u"kpc^-1" rtol=1e-2

    @test_throws Unitful.DimensionError unmpc(u"Mpc", 1u"Mpc^-1")
    @test_throws ArgumentError("mpc(q; base)` where `base` has dimensions `𝐌` has not yet been implemented. Please use a base with dimensions of `length` and then use `unmpc` to convert to your desired units."
        )  mpc(10u"kg"; base=u"kg")
end


text:
module IPPDSP

using IPPCore

IPPTypeDict = [ IPP8u   => "8u",
                IPP8s   => "8s",
                IPP16u  => "16u",
                IPP16s  => "16s",
                IPP32u  => "32u",
                IPP32s  => "32s",
                IPP16f  => "16f",
                IPP32f  => "32f",
                IPP64f  => "64f",
                IPP16uc => "16uc",
                IPP16sc => "16sc",
                IPP32uc => "32uc",
                IPP32sc => "32sc",
                IPP16fc => "16fc",
                IPP32fc => "32fc",
                IPP64fc => "64fc" ]

function IPPTypeSignature( IPPType )
    IPPTypeDict[ eval( IPPType ) ]
end

function IPPSuffix( IPPType )
   string( "_", IPPTypeSignature( IPPType ))
end

function IPPSuffix( IPPTypes::Tuple )
    lenTypes = length( IPPTypes )
    lenTypes == 0 && return ""
    lenTypes == 1 && return IPPSuffix( IPPTypes[1] )

    suffix = IPPTypeSignature( IPPTypes[1] )
    for Tsymb in IPPTypes[2:end]
        suffix *= typeof(Tsymb)<:String ? string( "_", Tsymb ) : IPPSuffix( Tsymb )
    end

    suffix
end

# common

macro ippscall(ippf, argtypes, args...)
    quote
        ret = ccall(($ippf, "libipps"), IppStatus,
            $argtypes, $(args...))
        ret == 0 || error(ippstatus_string(ret))
    end
end

# source files

include("Windows.jl")
include("Filtering.jl")

end # module

ext: jl


text:
export Network

function Network(g::SimpleGraph,
                 vertexf::Union{VertexFunction, Vector{<:VertexFunction}},
                 edgef::Union{EdgeFunction, Vector{<:EdgeFunction}};
                 accumulator=+, accdim=:auto,
                 execution=:seq,
                 verbose=false)
    verbose && println("Create dynamic network with $(nv(g)) vertices and $(ne(g)) edges:")
    @check execution ∈ EXECUTION_STYLES "Exectuion type $execution not supportet (choose from $(EXECUTION_STYLES))"

    im = IndexManager()

    vtypes, idxs = batch_identical(vertexf, collect(1:nv(g)))
    vertexbatches = collect(VertexBatch(im, i, t; verbose) for (i, t) in zip(idxs, vtypes))

    nl = NetworkLayer(im, g, edgef, accumulator, accdim; verbose)

    @assert isdense(im) "The index structure is not dense"
    return Network{execution, typeof(nl), typeof(vertexbatches)}(vertexbatches, nl, im)
end

function VertexBatch(im::IndexManager,
                     vertex_idxs::Vector{Int},
                     vertexf::VertexFunction;
                     verbose)
    dim = vertexf.dim
    pdim = vertexf.pdim
    (firstidx, pfirstidx) = register_vertices!(im, vertex_idxs, dim, pdim)

    verbose && println(" - VertexBatch: dim=$dim, pdim=$pdim, length=$(length(vertex_idxs))")

    VertexBatch(vertex_idxs,
                vertexf,
                dim,
                firstidx,
                pdim,
                pfirstidx)
end

function NetworkLayer(im::IndexManager, g::SimpleGraph,
                      edgef::Union{EdgeFunction, Vector{<:EdgeFunction}},
                      accumulator, accdim;
                      verbose)
    @check hasmethod(accumulator, NTuple{2, AbstractFloat}) "Accumulator needs `acc(AbstractFloat, AbstractFloat) method`"

    accdim_max = maxaccdim(edgef)
    if accdim === :auto
        accdim = accdim_max
        verbose && println(" - auto accumulation dimension = $accdim")
    end
    @check accdim<=accdim_max "For this system acc dim is limited to $accdim_max by the edgefunctions"

    edgecolors = color_edges_greedy(g)
    verbose && println(" - found $(length(unique(edgecolors))) edgecolors (optimum would be $(maximum(degree(g))))")

    colors = unique(edgecolors)

    idx_per_color = [findall(isequal(c), edgecolors) for c in colors]

    edgef_per_color = batch_by_idxs(edgef, idx_per_color)

    colorbatches = collect(ColorBatch(im, g, idx, ef; verbose)
                         for (idx, ef) in zip(idx_per_color, edgef_per_color))

    NetworkLayer(g, colorbatches, accumulator, accdim, CachePool())
end

maxaccdim(e::EdgeFunction) = accdim(e)
maxaccdim(e::AbstractVector{<:EdgeFunction}) = minimum(accdim.(e))

function ColorBatch(im::IndexManager, g::SimpleGraph,
                    edge_idxs::Vector{Int},
                    edgef::Union{EdgeFunction, Vector{<:EdgeFunction}};
                    verbose)
    verbose && println(" - ColorBatch with $(length(edge_idxs)) edges")

    eftypes, idxs = batch_identical(edgef, edge_idxs)
    edgebatches = collect(EdgeBatch(im, g, i, ef; verbose) for (i, ef) in zip(idxs, eftypes))

    return ColorBatch(edge_idxs, edgebatches)
end

function EdgeBatch(im::IndexManager, g::SimpleGraph,
                   edge_idxs::Vector{Int},
                   edgef::EdgeFunction;
                   verbose)
    dim = edgef.dim
    pdim = edgef.pdim

    if edgef isa StaticEdge
        (firstidx, pfirstidx) = register_edges!(im, edge_idxs, 0, pdim)
    else
        (firstidx, pfirstidx) = register_edges!(im, edge_idxs, dim, pdim)
    end

    vidx_array = Array{Int}(undef, 6*length(edge_idxs))
    empty!(vidx_array)
    alledges = collect(edges(g))
    for (i, eidx) in enumerate(edge_idxs)
        edge = alledges[eidx]
        src, dst = edge.src, edge.dst
        src_r = vertex_data_range(im, src)
        dst_r = vertex_data_range(im, dst)
        src_idx, dst_idx = src_r[begin], dst_r[begin]
        src_dim, dst_dim = length(src_r), length(dst_r)

        append!(vidx_array, src, src_idx, src_dim, dst, dst_idx, dst_dim)
    end
    @assert length(vidx_array) == 6*length(edge_idxs)

    verbose && println("   - EdgeBatch: dim=$dim, pdim=$pdim, length=$(length(edge_idxs))")

    EdgeBatch(edge_idxs,
              edgef,
              dim,
              firstidx,
              pdim,
              pfirstidx,
              vidx_array)
end

batch_by_idxs(v, idxs::Vector{Vector{Int}}) = [v for batch in idxs]
function batch_by_idxs(v::AbstractVector, batches::Vector{Vector{Int}})
    @assert length(v) == sum(length.(batches))
    [v[batch] for batch in batches]
end

batch_identical(el, idxs::Vector{Int}) = [el], [idxs]
function batch_identical(v::Vector{T}, indices::Vector{Int}) where {T}
    @assert length(v) == length(indices)
    idxs_per_type = Vector{Int}[]
    types = T[]
    for i in eachindex(v)
        found = false
        for j in eachindex(types)
            if v[i] === types[j]
                found = true
                push!(idxs_per_type[j], indices[i])
                break
            end
        end
        if !found
            push!(types, v[i])
            push!(idxs_per_type, [indices[i]])
        end
    end
    @assert length(types) == length(idxs_per_type)
    return types, idxs_per_type
end


text:
using Documenter

push!(LOAD_PATH, "../")
using SpeechDatasets

makedocs(
    sitename = "SpeechDatasets",
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    pages = [
        "Home" => "index.md",
        "Corpora" => Any[
            "ALFFA" => "alffa.md",
            "Babel" => "babel.md",
            "MaSS" => "massdataset.md",
            "Mboshi" => "mboshi.md",
            "TIMIT" => "timit.md",
        ]
    ]
)

deploydocs(
    repo = "github.com/lucasondel/SpeechDatasets.jl.git",
    devbranch = "main",
)

max_issues_repo_name: lucasondel/SpeechDatasets

max_forks_repo_name: lucasondel/SpeechDatasets


text:
using AlgebraPDF
using DelimitedFiles

using Plots
using LaTeXStrings
theme(:wong2, grid=false, frame=:box,
    guidefontvalign=:top, guidefonthalign=:right,
    foreground_color_legend = nothing,
    xlim=(:auto,:auto), ylim=(:auto,:auto),
    xlab=L"m(\Xi_c^+K^-)\,\,[\mathrm{GeV}]", ylab=L"\mathrm{Candidates}")
#
#
const xth = 2.96
const fitrange = (2.960, 3.180)

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

data = readdlm(joinpath("plots", "data.txt"))[:,1]
Nev = size(data,1)

# amplitude for the signal
Φ2(x) = sqrt(x-xth)
Γ(x,m,Γ₀) = Γ₀*Φ2(x)/Φ2(m)
breitwigner(x,m,Γ₀) = 1/(m^2-x^2-1im*m*Γ(x,m,Γ₀)) # complex function

# create
bw = FunctionWithParameters((x;p)->breitwigner(x,p.m,p.Γ), (m=3.05,Γ=0.02))
# check
bw(3.02) # works on a number
bw([3.05,3.065]) # works on a vector
pars(bw) # gives namedtuple
bw(3.02; p=(m=3.05,Γ=0.02)) # works parameters
bw(3.02, [3.05, 0.02]) #

# operators like `abs2`, `log` can be applied
bwsq = abs2(bw) # real function
plot(bwsq, fitrange...) # can plot, range is needed

# making PDF from the function just by normalizing
bwsq_norm = Normalized(bwsq, fitrange) # needs limits (ranges, support)
plot(bwsq_norm) # can plot

# fit would work already
testfit_summary = fit(bwsq_norm, data)
# same as
# fit(NegativeLogLikelihood(bwsq_norm, data))

# define a type SimpleBW and method `func` for dispatch
struct SimpleBW{P} <: AbstractFunctionWithParameters
    p::P
end
import AlgebraPDF:func
function func(bw::SimpleBW, x::NumberOrTuple; p=pars(bw))
    m,Γ = (getproperty(p,s) for s in keys(bw.p))
    breitwigner(x, m, Γ)
end

# create four instances of the same PDF with different parameters
signalamps = []
push!(signalamps, SimpleBW(( m1=3.00, Γ1=6.5e-3 )))
push!(signalamps, SimpleBW(( m2=3.05, Γ2=2.3e-3 )))
push!(signalamps, SimpleBW(( m3=3.06, Γ3=4.0e-3 )))
push!(signalamps, SimpleBW(( m4=3.09, Γ4=9.9e-3 )))
@time plot(); plot!.(abs2.(signalamps), fitrange...); plot!()

# |A|^2 * phase_space
phasespace = FunctionWithParameters((x;p)->Φ2(x), ∅)
signalpdfs = [Normalized(abs2(A)*phasespace, fitrange) for A in signalamps]
@time plot(); plot!.(signalpdfs, fill=0, α=0.3); plot!()

#  - BW S-wave
@makefuntype SimpleBWg(x;p) =
    1/(p.m0^2 - x^2 - 1im*p.g^2*Φ2(x))
#
signalpdf0 = Normalized(
        abs2(SimpleBWg((m0=2.95, g=0.01)))*phasespace,
        fitrange)
plot!(signalpdf0, fill=0, α=0.3)
#
backgrpdf = Normalized(phasespace, fitrange)
plot(backgrpdf)

#########################################################################
# full model
model0 =
    signalpdf0 * (f0=0.05Nev,) +
    signalpdfs[1] * (f1=0.25Nev,) + signalpdfs[2] * (f2=0.15Nev,) +
    signalpdfs[3] * (f3=0.15Nev,) + signalpdfs[4] * (f4=0.55Nev,) +
    backgrpdf * Ext(fb=0.3Nev,) # Ext - to be able to fix

model1 = fixpar(model0, :fb, 17.2)

@time fit_summary = fit(
    Extended(NegativeLogLikelihood(model1, data)))
#
let bins = range(fitrange..., length=80)
    Ns = scaletobinneddata(bins)
    stephist(data; bins, lc=:black, lab="Data")
    for i in 1:5
        plot!(fit_summary.best_model[i], Ns, 300,
            fillto=0.0, α=0.4, lab="Signal$(i)")
    end
    plot!(fit_summary.best_model, Ns, 300, lc=:red, lab="Total fit")
    #
    m = fit_summary.measurements
    for (i,(k,v)) in enumerate(zip(keys(m), m))
        annotate!(xth+0.01,25-1.2i,text("$k=$(string(v))",:left, 6))
    end
    plot!()
end
savefig(joinpath("plots", "sixcompfit.png"))


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBase.jl/blob/master/LICENSE

abstract type AbstractProblem end
abstract type FieldProblem<:AbstractProblem end
abstract type BoundaryProblem<:AbstractProblem end
abstract type MixedProblem<:AbstractProblem end

"""
General linearized problem to solve
    (K₁+K₂)Δu  +   C1'*Δλ = f₁+f₂
         C2Δu  +     D*Δλ = g
"""
mutable struct Assembly

    M :: SparseMatrixCOO  # mass matrix

    # for field assembly
    K :: SparseMatrixCOO   # stiffness matrix
    Kg :: SparseMatrixCOO  # geometric stiffness matrix
    f :: SparseMatrixCOO   # force vector
    fg :: SparseMatrixCOO  #

    # for boundary assembly
    C1 :: SparseMatrixCOO
    C2 :: SparseMatrixCOO
    D :: SparseMatrixCOO
    g :: SparseMatrixCOO
    c :: SparseMatrixCOO

    u :: Vector{Float64}  # solution vector u
    u_prev :: Vector{Float64}  # previous solution vector u
    u_norm_change :: Real  # change of norm in u

    la :: Vector{Float64}  # solution vector la
    la_prev :: Vector{Float64}  # previous solution vector u
    la_norm_change :: Real # change of norm in la

    removed_dofs :: Vector{Int64} # manually remove dofs from assembly
end

function Assembly()
    return Assembly(
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        SparseMatrixCOO(),
        [], [], Inf,
        [], [], Inf,
        [])
end

function empty!(assembly::Assembly)
    empty!(assembly.K)
    empty!(assembly.Kg)
    empty!(assembly.f)
    empty!(assembly.fg)
    empty!(assembly.C1)
    empty!(assembly.C2)
    empty!(assembly.D)
    empty!(assembly.g)
    empty!(assembly.c)
end

function isempty(assembly::Assembly)
    T = isempty(assembly.K)
    T &= isempty(assembly.Kg)
    T &= isempty(assembly.f)
    T &= isempty(assembly.fg)
    T &= isempty(assembly.C1)
    T &= isempty(assembly.C2)
    T &= isempty(assembly.D)
    T &= isempty(assembly.g)
    T &= isempty(assembly.c)
    return T
end

"""
Defines types for Problem variables.

# Examples

The type of 'elements' is Vector{Element}

Add elements into the Problem element list.
```@example
a = [1, 2, 3]
Problem.elements = a
```

"""
mutable struct Problem{P<:AbstractProblem}
    name :: AbstractString           # descriptive name for the problem
    dimension :: Int                 # degrees of freedom per node
    parent_field_name :: AbstractString # (optional) name of the parent field e.g. "displacement"
    elements :: Vector{Element}
    dofmap :: Dict{Element, Vector{Int64}} # connects the element local dofs to the global dofs
    assembly :: Assembly
    fields :: Dict{String, AbstractField}
    postprocess_fields :: Vector{String}
    properties :: P
end

"""
    Problem(problem_type, problem_name::String, problem_dimension)

Construct a new field problem where `problem_type` is the type of the problem
(Elasticity, Dirichlet, etc.), `problem_name` is the name of the problem and
`problem_dimension` is the number of DOF:s in one node (2 in a 2D problem, 3
in an elastic 3D problem, 6 in a 3D beam problem, etc.).

# Examples

Create a vector-valued (dim=3) elasticity problem:

```@example
prob1 = Problem(Elasticity, "this is my problem", 3)
```

"""
function Problem(::Type{P}, name::AbstractString, dimension::Int64) where P<:FieldProblem
    return Problem{P}(name, dimension, "none", [], Dict(), Assembly(), Dict(), Vector(), P())
end

"""
Construct a new boundary problem.

Examples
--------
Create a Dirichlet boundary problem for a vector-valued (dim=3) elasticity problem.

julia> bc1 = Problem(Dirichlet, "support", 3, "displacement")
solver.
"""
function Problem(::Type{P}, name, dimension, parent_field_name) where P<:BoundaryProblem
    return Problem{P}(name, dimension, parent_field_name, [], Dict(), Assembly(), Dict(), Vector(), P())
end

function get_formulation_type(::Problem)
    return :incremental
end

"""
    get_unknown_field_dimension(problem)

Return the dimension of the unknown field of this problem.
"""
function get_unknown_field_dimension(problem::Problem)
    return problem.dimension
end

"""
    get_unknown_field_name(problem)

Default function if unknown field name is not defined for some problem.
"""
function get_unknown_field_name(::P) where P<:AbstractProblem
    warn("The name of unknown field (e.g. displacement, temperature, ...) of the ",
         "problem type must be given by defining function `get_unknown_field_name`")
    return "N/A"
end

""" Return the name of the unknown field of this problem. """
function get_unknown_field_name(problem::Problem{P}) where P
    return get_unknown_field_name(problem.properties)
end

""" Return the name of the parent field of this (boundary) problem. """
function get_parent_field_name(problem::Problem{P}) where P<:BoundaryProblem
    return problem.parent_field_name
end

function get_unknown_field_name(::P) where P<:BoundaryProblem
    return "lambda"
end

is_field_problem(::Problem) = false
is_field_problem(::Problem{P}) where {P<:FieldProblem} = true
is_boundary_problem(::Problem) = false
is_boundary_problem(::Problem{P}) where {P<:BoundaryProblem} = true

function get_elements(problem::Problem)
    return problem.elements
end

"""
    update!(problem.properties, attr...)

Update properties for a problem.

# Example

```julia
update!(body.properties, "finite_strain" => "false")
```
"""
function update!(problem::P, attr::Pair{String, String}...) where P<:AbstractProblem
    for (name, value) in attr
        setfield!(problem, parse(name), parse(value))
    end
end

"""
    function initialize!(problem_type, element_name, time)

Initialize the element ready for calculation, where `problem_type` is the type
of the problem (Elasticity, Dirichlet, etc.), `element_name` is the name of a
constructed element (see Element(element_type, connectivity_vector)) and `time`
is the starting time of the initializing process.
"""
function initialize!(problem::Problem, element::Element, time::Float64)
    field_name = get_unknown_field_name(problem)
    field_dim = get_unknown_field_dimension(problem)
    nnodes = length(element)
    if field_dim == 1  # scalar field
        empty_field = tuple(zeros(nnodes)...)
    else  # vector field
        # FIXME: the most effective way to do
        # ([0.0,0.0], [0.0,0.0], ..., [0.0,0.0]) ?
        empty_field = tuple(map((x)->zeros(field_dim)*x, 1:nnodes)...)
    end

    # initialize primary field
    if !haskey(element, field_name)
        update!(element, field_name, time => empty_field)
    end

    # if a boundary problem, initialize also a field for the main problem
    is_boundary_problem(problem) || return
    field_name = get_parent_field_name(problem)
    if !haskey(element, field_name)
        update!(element, field_name, time => empty_field)
    end
end

function initialize!(problem::Problem, time::Float64=0.0)
    for element in get_elements(problem)
        initialize!(problem, element, time)
    end
end

"""
    update!(problem, assembly, u, la)

Update the problem solution vector for assembly.
"""
function update!(problem::Problem, assembly::Assembly, u::Vector, la::Vector)

    # resize & fill with zeros vectors if length mismatch with current solution

    if length(u) != length(assembly.u)
        info("resizing solution vector u")
        resize!(assembly.u, length(u))
        fill!(assembly.u, 0.0)
    end

    if length(la) != length(assembly.la)
        info("resizing lagrange multiplier vector la")
        resize!(assembly.la, length(la))
        fill!(assembly.la, 0.0)
    end

    # copy current solutions to previous ones and add/replace new solution
    # TODO: here we have couple of options and they need to be clarified
    # for total formulation we are solving total quantity Ku = f while in
    # incremental formulation we solve KΔu = f and u = u + Δu
    assembly.u_prev = copy(assembly.u)
    assembly.la_prev = copy(assembly.la)

    if get_formulation_type(problem) == :total
        assembly.u = u
        assembly.la = la
    elseif get_formulation_type(problem) == :incremental
        assembly.u += u
        assembly.la = la
    elseif get_formulation_type(problem) == :forwarddiff
        assembly.u += u
        assembly.la += la
    else
        info("$(problem.name): unknown formulation type, don't know what to do with results")
        error("serious failure with problem formulation: $(get_formulation_type(problem))")
    end

    # calculate change of norm
    assembly.u_norm_change = norm(assembly.u - assembly.u_prev)
    assembly.la_norm_change = norm(assembly.la - assembly.la_prev)
    return assembly.u, assembly.la
end

"""
    get_global_solution(problem, assembly)

Return a global solution (u, la) for a problem.

Notes
-----
If the length of solution vector != number of nodes, i.e. the field dimension is
something else than 1, reshape vectors so that their length matches to the
number of nodes. This helps to get nodal results easily.
"""
function get_global_solution(problem::Problem, assembly::Assembly)
    u = assembly.u
    la = assembly.la
    field_dim = get_unknown_field_dimension(problem)
    if field_dim == 1
        return u, la
    else
        nnodes = round(Int, length(u)/field_dim)
        u = reshape(u, field_dim, nnodes)
        u = Vector{Float64}[u[:,i] for i in 1:nnodes]
        la = reshape(la, field_dim, nnodes)
        la = Vector{Float64}[la[:,i] for i in 1:nnodes]
        return u, la
    end
end

"""
    update!(problem, assembly, elements, time)

Update a solution from the assebly to elements.
"""
function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:FieldProblem
    u, la = get_global_solution(problem, assembly)
    field_name = get_unknown_field_name(problem)
    # update solution u for elements
    for element in elements
        connectivity = get_connectivity(element)
        update!(element, field_name, time => tuple(u[connectivity]...))
    end
end

function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:BoundaryProblem
    u, la = get_global_solution(problem, assembly)
    parent_field_name = get_parent_field_name(problem) # displacement
    field_name = get_unknown_field_name(problem) # lambda
    # update solution and lagrange multipliers for boundary elements
    for element in elements
        connectivity = get_connectivity(element)
        update!(element, parent_field_name, time => tuple(u[connectivity]...))
        update!(element, field_name, time => tuple(la[connectivity]...))
    end
end

"""
    add_elements!(problem::Problem, elements)

Add new elements into the problem.
"""
function add_elements!(problem::Problem, elements)
    for element in elements
        push!(problem.elements, element)
    end
end

function get_assembly(problem::Problem)
    return problem.assembly
end

function length(problem::Problem)
    return length(problem.elements)
end

function update!(problem::Problem, field_name::AbstractString, data)
    #if haskey(problem.fields, field_name)
    #    update!(problem.fields[field_name], field_name::AbstractString, data)
    #else
    #    problem.fields[field_name] = Field(data)
    #end
    update!(problem.elements, field_name::AbstractString, data)
end

function haskey(problem::Problem, field_name::AbstractString)
    return haskey(problem.fields, field_name)
end

function getindex(problem::Problem, field_name::String)
    return problem.fields[field_name]
end

#""" Return field calculated to nodal points for elements in problem p. """
function (problem::Problem)(field_name::String, time::Float64)
    #if haskey(problem, field_name)
    #    return problem[field_name](time)
    #end
    f = Dict{Int64, Any}()
    for element in get_elements(problem)
        haskey(element, field_name) || continue
        for (c, v) in zip(get_connectivity(element), element(field_name, time))
            if haskey(f, c)
                if !isapprox(f[c], v)
                    info("several values for single node when returning field $field_name")
                    info("already have: $(f[c]), and trying to set $v")
                end
            else
                f[c] = v
            end
        end
    end
    #f == nothing && return f
    #update!(problem, field_name, time => f)
    return f
end

function push!(problem::Problem, elements...)
    push!(problem.elements, elements...)
end

function push!(problem::Problem, elements_::Vector...)
    for elements in elements_
        push!(problem.elements, elements...)
    end
end

"""
    set_gdofs!(problem, element)

Set element global degrees of freedom.
"""
function set_gdofs!(problem, element, dofs)
    problem.dofmap[element] = dofs
end

"""
    get_gdofs(problem, element)

Return the global degrees of freedom for element.

First make lookup from problem dofmap. If not defined there, make implicit
assumption that dofs follow formula `gdofs = [dim*(nid-1)+j for j=1:dim]`,
where `nid` is node id and `dim` is the dimension of problem. This formula
arranges dofs so that first comes all dofs of node 1, then node 2 and so on:
(u11, u12, u13, u21, u22, u23, ..., un1, un2, un3) for 3 dofs/node setting.
"""
function get_gdofs(problem::Problem, element::Element)
    if haskey(problem.dofmap, element)
        return problem.dofmap[element]
    end
    conn = get_connectivity(element)
    if length(conn) == 0
        error("element connectivity not defined, cannot determine global ",
              "degrees of freedom for element #: $(element.id)")
    end
    dim = get_unknown_field_dimension(problem)
    gdofs = [dim*(i-1)+j for i in conn for j=1:dim]
    return gdofs
end


text:
using ExoplanetsSysSim

using CSV
using DataFrames
using Glob
using PyPlot

include("../../SysSimExClusters/src/AMD_stability/solar_system.jl")

##### To load a table of the physical catalog with observed planets only:

dir = "/Users/hematthi/Documents/GradSchool/Research/ACI/Simulated_Data/AMD_system/Split_stars/Singles_ecc/Params11_KS/Distribute_AMD_per_mass/durations_norm_circ_singles_multis_GF2020_KS/GP_med/"
file = joinpath(dir, "physical_catalog_obs_only.csv")
table_planets = CSV.read(file)

##### To evaluate AMD stability for each system:

tid_all = unique(table_planets[!,:target_id])
n_sys = length(tid_all)

n_pl_all = Vector{Float64}(undef, n_sys)
amd_tot_all = Vector{Float64}(undef, n_sys)
amd_crit_all = Vector{Float64}(undef, n_sys)

for (i,tid) in enumerate(tid_all)
    id_sys = (1:size(table_planets,1))[table_planets[!,:target_id] .== tid]
    planets = table_planets[id_sys,:]
    n_pl_all[i] = size(planets, 1)

    m = planets[!,:planet_mass]
    μ = m ./ planets[!,:star_mass]
    P = planets[!,:period]
    a = semimajor_axis.(P, m .+ planets[!,:star_mass])
    e = planets[!,:ecc]
    i_m = planets[!,:incl_mut]

    AMD_tot = total_AMD_system(μ, a, e, i_m)[1]
    AMD_crit = critical_AMD_system(μ, a)

    amd_tot_all[i] = AMD_tot
    amd_crit_all[i] = AMD_crit
end

amd_ratio_crit_all = amd_tot_all ./ amd_crit_all

##### To make some plots:

logrange(x1, x2, n) = collect(10^logx for logx in range(log10(x1), log10(x2), length=n))

# Histogram of planet masses:
hist(table_planets[!,:planet_mass] ./ ExoplanetsSysSim.earth_mass, bins=logrange(1e-1, 1e3, 50), histtype="step")
ax = gca()
ax.set_xscale("log")
xlabel(L"$M_p$ ($M_\oplus$)", fontsize=20)

# Histogram of periods:
hist(table_planets[!,:period], bins=logrange(3, 300, 50), histtype="step")
ax = gca()
ax.set_xscale("log")
xlabel(L"$P$ (days)", fontsize=20)

# Histogram of eccentricities:
hist(table_planets[!,:ecc], bins=logrange(1e-3, 1, 50), histtype="step")
ax = gca()
ax.set_xscale("log")
xlabel(L"$e$", fontsize=20)

# Histogram of mutual inclinations (ignoring true singles which have zero mutual inclination):
im_all = table_planets[!,:incl_mut]
hist(im_all .* 180 ./π, bins=logrange(1e-2, 1e2, 50), histtype="step")
ax = gca()
ax.set_xscale("log")
xlabel(L"$i_m$ ($^\circ$)", fontsize=20)

# Histogram of total AMD/AMD_crit:
plot = PyPlot.figure(figsize=(8,5))
hist(amd_ratio_crit_all, bins=logrange(1e-5, 1, 50), histtype="step")
ax = gca()
ax.tick_params(axis="both", labelsize=20)
ax.set_xscale("log")
xlabel(L"AMD$_{tot}$/AMD$_{crit}$", fontsize=20)
ylabel("Systems", fontsize=20)
plot.tight_layout()

# Histogram of total AMD/AMD_crit split by (observed) multiplicity:
plot = PyPlot.figure(figsize=(8,5))
labels_m = [L"$m = 1$", L"$m = 2$", L"$m = 3$", L"$m = 4+$"]
bools_m = [n_pl_all .== 1, n_pl_all .== 2, n_pl_all .== 3, n_pl_all .>= 4]
hist([amd_ratio_crit_all[bools] for bools in bools_m], bins=logrange(1e-5, 1, 50), histtype="step", label=labels_m)
ax = gca()
ax.tick_params(axis="both", labelsize=20)
ax.set_xscale("log")
xlabel(L"AMD$_{tot}$/AMD$_{crit}$", fontsize=20)
ylabel("Systems", fontsize=20)
legend(loc="upper left", bbox_to_anchor=(0,1), frameon=false, fontsize=16)
plot.tight_layout()


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

using Markdown
using InteractiveUtils

# ╔═╡ d613968d-7549-401a-9f25-8169660f121f
using SymPy, PlutoUI

# ╔═╡ a5797465-247a-4e33-b733-46ad9a58ad00
TableOfContents(title="📚 Sadržaj", aside=true)

# ╔═╡ 71bf3e90-8277-11eb-188f-0f6601464f27
md"
# Integrali iracionalnih funkcija
"

# ╔═╡ 90465d99-b8e8-4ed3-9943-64aec908d6f7
md"""
## Eulerove i trigonometrijske supstitucije
"""

# ╔═╡ 37981771-05ce-40b5-a2a9-8031d41cdb6b
x=symbols("x",real=true)

# ╔═╡ 1dee94b6-fa66-462e-bbb3-0c941a59dd8e
# Ne može izračunati
integrate(1/(x+√(x^2+6*x+10)),x)

# ╔═╡ cc61e680-8110-11eb-2211-1d212f28c36d
t=symbols("t",real=true)

# ╔═╡ dacfbf80-8110-11eb-2981-4d7c579c8e48
I₀=1/(x+√(x^2+6x+10))

# ╔═╡ 2835b6ea-9c18-4259-82b9-f3eefea58772
I₁=simplify(subs(I₀,x,t-3))

# ╔═╡ 4b0bc7e5-c399-483b-8a4a-c3b4b35fc543
z=symbols("z",real=true)

# ╔═╡ a31fc5bd-cf3b-4248-8994-ce8adb4dd02f
# Eulerova supstitucija
I₂=subs(I₁,√(t^2+1),t+z)

# ╔═╡ 6594a67e-8111-11eb-29a5-eb1fd173be05
I₃=subs(I₂,t,(1-z^2)/(2*z))

# ╔═╡ d7184d11-2e04-434a-bb68-3231aa4a59e2
# Množi s dx
I₄=simplify(I₃*(-1)*(1+z^2)/(2*z^2))

# ╔═╡ f9ec108c-b5dd-46ef-b4ac-b6d97bbc2162
# Integriramo
I₅=integrate(I₄,z)

# ╔═╡ dc3b1f9b-5dde-41f5-9dda-07eb70b460dd
# Vratimo supstitucije
I₆=subs(I₅,z,√(t^2+1)-t)

# ╔═╡ 2bc4c1bf-d574-43f1-8788-3dcecdfd8218
I₇=subs(I₆,t,x+3)

# ╔═╡ 777be0b5-8d94-4347-b61c-d0827a0bf774
# Primjer 1.12
integrate(x^2*sqrt(4*x^2+9))

# ╔═╡ 08e1571e-3f23-456b-9eea-6f228c44980c
md"""
## Binomni integral
"""

# ╔═╡ 098a9b62-9ea4-4547-b6fb-f4fb6d1ee6cf
# Primjer 1.13
I₈=((3*x-x*x*x)^(1//3))

# ╔═╡ 94c8f287-8031-4cd3-a67a-186a5eebf09e
integrate(I₈,x)

# ╔═╡ 0cbc9b1a-7282-4075-a0a2-2b92595bc7c3
md"""
## Integriranje razvoja u red potencija
"""

# ╔═╡ ecf4a097-2e3f-496c-8e65-004e7f0740d1
s₁=series(exp(x),x)

# ╔═╡ 4ddece7e-33ee-4586-b08d-33cd39bb1e30
s₂=series(exp(x),x,0,10)

# ╔═╡ 30f2cb1b-5bb3-4db8-9239-bcde4ec2962f
md"""
Razvoj u red funkcije $f(x)=e^{-x^2}$
"""

# ╔═╡ 4b02de9e-ae81-4324-92e3-20f4b882a8d4
s₃=subs(s₂,x,(-x^2))

# ╔═╡ da968a92-6a56-4dff-97f6-5e87616d4291
s₄=integrate(s₃,x)

# ╔═╡ e8bcf956-422d-4829-9ab6-fda047805a5b
# Uklonimo član O(x^21)
s₆=s₄.removeO()

# ╔═╡ 58468833-ef55-4983-b637-cadd7bcdf9f3
# Izračunajmo vrijednost u točki x=0.1
subs(s₆,x,0.1)

# ╔═╡ 2ac4782d-6c74-4402-b6eb-17554ef27d03
md"""
###  Primjer 1.14 b)
"""

# ╔═╡ 17c63bc5-611b-4cfc-a916-c0018dd71ecb
# Maclaurin-ov red za sin(x)
T₁=series(sin(x),x,0,10)

# ╔═╡ 4375ec68-a2d1-4cf8-a2ab-88bb69717406
# Razvoj u red za sin(x)/x
T₂=simplify(s₁/x)

# ╔═╡ b9b5d733-931d-4f69-8f44-6bde41f7458b
# Ukonimo O() clan
T₃=s₂.removeO()

# ╔═╡ 85a68267-73a0-4c16-ae9f-db1d3ceac7bb
# Integrirajmo
T₄=integrate(s₃,x)

# ╔═╡ 3e684e89-3216-4ed3-a60c-782ab7234fcf

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"

[compat]
PlutoUI = "~0.7.14"
SymPy = "~1.0.52"
"""

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

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "a325370b9dd0e6bf5656a6f1a7ae80755f8ccc46"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.7.2"

[[CommonEq]]
git-tree-sha1 = "d1beba82ceee6dc0fce8cb6b80bf600bbde66381"
uuid = "3709ef60-1bee-4518-9f2f-acd86f176c50"
version = "0.2.0"

[[CommonSolve]]
git-tree-sha1 = "68a0743f578349ada8bc911a5cbd5a2ef6ed6d1f"
uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
version = "0.2.0"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "31d0151f5716b655421d9d75b7fa74cc4e744df2"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.39.0"

[[CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[Conda]]
deps = ["JSON", "VersionParsing"]
git-tree-sha1 = "299304989a5e6473d985212c28928899c74e9421"
uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
version = "1.5.2"

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

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "a32185f5428d3986f47c2ab78b1f216d5e6cc96f"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.5"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[HypertextLiteral]]
git-tree-sha1 = "72053798e1be56026b81d4e2682dbe58922e5ec9"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.0"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IrrationalConstants]]
git-tree-sha1 = "f76424439413893a832026ca355fe273e93bce94"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.0"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

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

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[LogExpFunctions]]
deps = ["ChainRulesCore", "DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "34dc30f868e368f8a17b728a1238f3fcda43931a"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.3"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "5a5bc6bf062f0f95e62d0fe0a2d99699fed82dd9"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.8"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

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

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

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

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoUI]]
deps = ["Base64", "Dates", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "d1fb76655a95bf6ea4348d7197b22e889a4375f4"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.14"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

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

[[PyCall]]
deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
git-tree-sha1 = "169bb8ea6b1b143c5cf57df6d34d022a7b60c6db"
uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
version = "1.92.3"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[RecipesBase]]
git-tree-sha1 = "44a75aa7a527910ee3d1751d1f0e4148698add9e"
uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
version = "1.1.2"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "793793f1df98e3d7d554b65a107e9c9a6399a6ed"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.7.0"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[SymPy]]
deps = ["CommonEq", "CommonSolve", "LinearAlgebra", "Markdown", "PyCall", "RecipesBase", "SpecialFunctions"]
git-tree-sha1 = "1ef257ecbcab8058595a68ca36a6844b41babcbd"
uuid = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"
version = "1.0.52"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

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

[[VersionParsing]]
git-tree-sha1 = "80229be1f670524750d905f8fc8148e5a8c4537f"
uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
version = "1.2.0"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╠═d613968d-7549-401a-9f25-8169660f121f
# ╠═a5797465-247a-4e33-b733-46ad9a58ad00
# ╟─71bf3e90-8277-11eb-188f-0f6601464f27
# ╟─90465d99-b8e8-4ed3-9943-64aec908d6f7
# ╠═37981771-05ce-40b5-a2a9-8031d41cdb6b
# ╠═1dee94b6-fa66-462e-bbb3-0c941a59dd8e
# ╠═cc61e680-8110-11eb-2211-1d212f28c36d
# ╠═dacfbf80-8110-11eb-2981-4d7c579c8e48
# ╠═2835b6ea-9c18-4259-82b9-f3eefea58772
# ╠═4b0bc7e5-c399-483b-8a4a-c3b4b35fc543
# ╠═a31fc5bd-cf3b-4248-8994-ce8adb4dd02f
# ╠═6594a67e-8111-11eb-29a5-eb1fd173be05
# ╠═d7184d11-2e04-434a-bb68-3231aa4a59e2
# ╠═f9ec108c-b5dd-46ef-b4ac-b6d97bbc2162
# ╠═dc3b1f9b-5dde-41f5-9dda-07eb70b460dd
# ╠═2bc4c1bf-d574-43f1-8788-3dcecdfd8218
# ╠═777be0b5-8d94-4347-b61c-d0827a0bf774
# ╟─08e1571e-3f23-456b-9eea-6f228c44980c
# ╠═098a9b62-9ea4-4547-b6fb-f4fb6d1ee6cf
# ╠═94c8f287-8031-4cd3-a67a-186a5eebf09e
# ╟─0cbc9b1a-7282-4075-a0a2-2b92595bc7c3
# ╠═ecf4a097-2e3f-496c-8e65-004e7f0740d1
# ╠═4ddece7e-33ee-4586-b08d-33cd39bb1e30
# ╟─30f2cb1b-5bb3-4db8-9239-bcde4ec2962f
# ╠═4b02de9e-ae81-4324-92e3-20f4b882a8d4
# ╠═da968a92-6a56-4dff-97f6-5e87616d4291
# ╠═e8bcf956-422d-4829-9ab6-fda047805a5b
# ╠═58468833-ef55-4983-b637-cadd7bcdf9f3
# ╟─2ac4782d-6c74-4402-b6eb-17554ef27d03
# ╠═17c63bc5-611b-4cfc-a916-c0018dd71ecb
# ╠═4375ec68-a2d1-4cf8-a2ab-88bb69717406
# ╠═b9b5d733-931d-4f69-8f44-6bde41f7458b
# ╠═85a68267-73a0-4c16-ae9f-db1d3ceac7bb
# ╠═3e684e89-3216-4ed3-a60c-782ab7234fcf
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002

max_stars_repo_name: ivanslapnicar/Matematika

max_issues_repo_name: ivanslapnicar/Matematika

max_forks_repo_name: ivanslapnicar/Matematika


text:
# Invertible network layer from Putzky and Welling (2019): https://arxiv.org/abs/1911.10914
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export NetworkLoop, NetworkLoop3D

"""
    L = NetworkLoop(n_in, n_hidden, maxiter, Ψ; k1=4, k2=3, p1=0, p2=1, s1=4, s2=1, ndims=2) (2D)

    L = NetworkLoop3D(n_in, n_hidden, maxiter, Ψ; k1=4, k2=3, p1=0, p2=1, s1=4, s2=1) (3D)

 Create an invertibel recurrent inference machine (i-RIM) consisting of an unrooled loop
 for a given number of iterations.

 *Input*:

 - 'n_in': number of input channels

 - `n_hidden`: number of hidden units in residual blocks

 - `maxiter`: number unrolled loop iterations

 - `Ψ`: link function

 - `k1`, `k2`: stencil sizes for convolutions in the residual blocks. The first convolution
   uses a stencil of size and stride `k1`, thereby downsampling the input. The second
   convolutions uses a stencil of size `k2`. The last layer uses a stencil of size and stride `k1`,
   but performs the transpose operation of the first convolution, thus upsampling the output to
   the original input size.

 - `p1`, `p2`: padding for the first and third convolution (`p1`) and the second convolution (`p2`) in
   residual block

 - `s1`, `s2`: stride for the first and third convolution (`s1`) and the second convolution (`s2`) in
   residual block

 - `ndims` : number of dimensions

 *Output*:

 - `L`: invertible i-RIM network.

 *Usage:*

 - Forward mode: `η_out, s_out = L.forward(η_in, s_in, d, A)`

 - Inverse mode: `η_in, s_in = L.inverse(η_out, s_out, d, A)`

 - Backward mode: `Δη_in, Δs_in, η_in, s_in = L.backward(Δη_out, Δs_out, η_out, s_out, d, A)`

 *Trainable parameters:*

 - None in `L` itself

 - Trainable parameters in the invertible coupling layers `L.L[i]`, and actnorm layers
   `L.AN[i]`, where `i` ranges from `1` to the number of loop iterations.

 See also: [`CouplingLayerIRIM`](@ref), [`ResidualBlock`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
"""
struct NetworkLoop <: InvertibleNetwork
    L::Union{AbstractArray{CouplingLayerIRIM, 1}, AbstractArray{CouplingLayerHINT}}
    AN::AbstractArray{ActNorm, 1}
    Ψ::Function
end

@Flux.functor NetworkLoop

# 2D Constructor
function NetworkLoop(n_in, n_hidden, maxiter, Ψ; k1=4, k2=3, p1=0, p2=1, s1=4, s2=1, type="additive", ndims=2)

    if type == "additive"
        L = Array{CouplingLayerIRIM}(undef, maxiter)
    elseif type == "HINT"
        L = Array{CouplingLayerHINT}(undef, maxiter)
    end

    AN = Array{ActNorm}(undef, maxiter)
    for j=1:maxiter
        if type == "additive"
            L[j] = CouplingLayerIRIM(n_in, n_hidden; k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, ndims=ndims)
        elseif type == "HINT"
            L[j] = CouplingLayerHINT(n_in, n_hidden; logdet=false, permute="both", k1=k1, k2=k2, p1=p1, p2=p2,
                                     s1=s1, s2=s2, ndims=ndims)
        end
        AN[j] = ActNorm(1)
    end

    return NetworkLoop(L, AN, Ψ)
end

# 3D Constructor
NetworkLoop3D(args...; kw...) = NetworkLoop(args...; kw..., ndims=3)

# 2D Forward loop: Input (η, s), Output (η, s)
function forward(η::AbstractArray{T, N}, s::AbstractArray{T, N}, d::AbstractArray, J, UL::NetworkLoop) where {T, N}

    # Dimensions
    n_in = size(s, N-1) + 1
    batchsize = size(s)[end]
    nn = size(s)[1:N-2]
    maxiter = length(UL.L)
    N0 = cuzeros(η, nn..., n_in-2, batchsize)

    for j=1:maxiter
        g = J'*(J*reshape(UL.Ψ(η), :, batchsize) - reshape(d, :, batchsize))
        g = reshape(g, nn..., 1, batchsize)
        gn = UL.AN[j].forward(g)   # normalize
        s_ = s + tensor_cat(gn, N0)

        ηs = UL.L[j].forward(tensor_cat(η, s_))
        η, s = tensor_split(ηs; split_index=1)
    end
    return η, s
end

# 2D Inverse loop: Input (η, s), Output (η, s)
function inverse(η::AbstractArray{T, N}, s::AbstractArray{T, N}, d::AbstractArray, J, UL::NetworkLoop) where {T, N}

    # Dimensions
    n_in = size(s, N-1) + 1
    batchsize = size(s)[end]
    nn = size(s)[1:N-2]
    maxiter = length(UL.L)

    N0 = cuzeros(η, nn..., n_in-2, batchsize)

    for j=maxiter:-1:1
        ηs_ = UL.L[j].inverse(tensor_cat(η, s))
        η, s_ = tensor_split(ηs_; split_index=1)

        g = J'*(J*reshape(UL.Ψ(η), :, batchsize) - reshape(d, :, batchsize))
        g = reshape(g, nn..., 1, batchsize)
        gn = UL.AN[j].forward(g)   # normalize
        s = s_ - tensor_cat(gn, N0)
    end
    return η, s
end

# 2D Backward loop: Input (Δη, Δs, η, s), Output (Δη, Δs, η, s)
function backward(Δη::AbstractArray{T, N}, Δs::AbstractArray{T, N},
    η::AbstractArray{T, N}, s::AbstractArray{T, N}, d::AbstractArray, J, UL::NetworkLoop; set_grad::Bool=true) where {T, N}

    # Dimensions
    n_in = size(s, N-1) + 1
    batchsize = size(s)[end]
    nn = size(s)[1:N-2]
    maxiter = length(UL.L)

    N0 = cuzeros(Δη, nn..., n_in-2, batchsize)
    typeof(Δs) == T && (Δs = 0 .* s)  # make Δs zero tensor

    # Initialize net parameters
    set_grad && (Δθ = Array{Parameter, 1}(undef, 0))

    for j = maxiter:-1:1
        if set_grad
            Δηs_, ηs_ = UL.L[j].backward(tensor_cat(Δη, Δs), tensor_cat(η, s))
        else
            Δηs_, Δθ_L, ηs_ = UL.L[j].backward(tensor_cat(Δη, Δs), tensor_cat(η, s); set_grad=set_grad)
            push!(Δθ, Δθ_L)
        end

        # Inverse pass
        η, s_ = tensor_split(ηs_; split_index=1)
        g = J'*(J*reshape(UL.Ψ(η), :, batchsize) - reshape(d, :, batchsize))
        g = reshape(g, nn..., 1, batchsize)
        gn = UL.AN[j].forward(g)   # normalize
        s = s_ - tensor_cat(gn, N0)

        # Gradients
        Δs2, Δs = tensor_split(Δηs_; split_index=1)
        Δgn = tensor_split(Δs; split_index=1)[1]
        Δg = UL.AN[j].backward(Δgn, gn)[1]
        Δη = reshape(J'*J*reshape(Δg, :, batchsize), nn..., 1, batchsize) + Δs2
    end
    set_grad ? (return Δη, Δs, η, s) : (Δη, Δs, Δθ, η, s)
end

## Jacobian-related utils
jacobian(::AbstractArray{T, 5}, ::AbstractArray{T, 5}, d::AbstractArray, J, UL::NetworkLoop) where T = throw(ArgumentError("Jacobian for NetworkLoop not yet implemented"))

adjointJacobian(Δη::AbstractArray{T, N}, Δs::AbstractArray{T, N},
                η::AbstractArray{T, N}, s::AbstractArray{T, N}, d::AbstractArray, J, UL::NetworkLoop;
                set_grad::Bool=true) where {T, N} =
            backward(Δη, Δs, η, s, d, J, UL; set_grad=false)

## Other utils
# Clear gradients
function clear_grad!(UL::NetworkLoop)
    maxiter = length(UL.L)
    for j=1:maxiter
        clear_grad!(UL.L[j])
        clear_grad!(UL.AN[j])
        UL.AN[j].s.data = nothing
        UL.AN[j].b.data = nothing
    end
end

# Get parameters (do not update actnorm weights)
function get_params(UL::NetworkLoop)
    maxiter = length(UL.L)
    p = get_params(UL.L[1])
    if maxiter > 1
        for j=2:maxiter
            p = cat(p, get_params(UL.L[j]); dims=1)
        end
    end
    return p
end


text:
@inline _sortperm(X::AbstractVecOrMat) = size(X,2)==1 ? sortperm(ndims(X)==1 ? X : vec(X), alg=RadixSort) : sortperm(collect(eachrow(X)))  # sort a data matrix

function Init!(o::StrBootTest{T}) where T  # for efficiency when varying r repeatedly to make CI, do stuff once that doesn't depend on r
	o.kX = o.kX₁ + o.kX₂
	o.kX₂==0 && (o.X₂ = zeros(T,o.Nobs,0))
  iszero(o.kY₂) && (o.Y₂ = zeros(T,o.Nobs,0))
  o.kZ = o.kX₁ + o.kY₂
  if o.LIML && o.kX₂==o.kY₂  # exactly identified LIML = 2SLS
  	o.κ = one(T)
  	o.LIML = false
  end
  if !(o.REst = length(o.R₁)>0)  # base model contains no restrictions?
    o.R₁ = zeros(T,0,o.kZ)
    o.r₁ = zeros(T,0)
  end
  isnan(o.κ) && (o.κ = o.kX₂>0 ? one(T) : zero(T))  # if κ in κ-class estimation not specified, it's 0 or 1 for OLS or 2SLS
  o.WRE = !(iszero(o.κ) || o.scorebs) || o.ARubin

  iszero(o.B) && (o.scorebs = true)

  o.haswt = !iszero(nrows(o.wt))
  o.sumwt = o.haswt ? sum(o.wt) : one(T)
  o._Nobs = o.haswt && o.fweights ? o.sumwt : T(o.Nobs)

  if !(iszero(o.NClustVar)  || o.issorted)
		o.subcluster = o.NClustVar - o.nerrclustvar
		p = _sortperm(view(o.ID, :, [collect(o.subcluster+1:o.NClustVar); collect(1:o.subcluster)]))  # sort by err cluster vars, then remaining boot cluster vars
		o.ID = ndims(o.ID)==1 ? o.ID[p] : o.ID[p,:]
		o.X₁ = ndims(o.X₁)==1 ? o.X₁[p] : o.X₁[p,:]
		o.X₂ = ndims(o.X₂)==1 ? o.X₂[p] : o.X₂[p,:]
		o.y₁ = o.y₁[p]
		o.Y₂ = ndims(o.Y₂)==1 ? o.Y₂[p] : o.Y₂[p,:]
		o.haswt && (o.wt = o.wt[p])
		isdefined(o, :FEID) && nrows(o.FEID)>0 && (o.FEID = o.FEID[p])
  end

  if o.WREnonARubin
	  if !iszero(o.NClustVar)
	    o.info✻, o.ID✻ = panelsetupID(o.ID, 1:o.nbootclustvar)
	  else
	    o.info⋂ = o.info✻ = Vector{UnitRange{Int64}}(undef, o.Nobs, 0)  # no clustering, so no collapsing by cluster
	  end
  elseif !iszero(o.NClustVar)
	  o.info✻ = panelsetup(o.ID, 1:min(o.NClustVar,o.nbootclustvar))  # bootstrap cluster grouping defs rel to original data
  else
	  info⋂ = info✻⋂ = o.info✻ = Vector{UnitRange{Int64}}(undef, o.Nobs, 0)  # causes no collapsing of data in panelsum() calls, only multiplying by weights if any
  end
  o.N✻ = nrows(o.info✻)

  if o.bootstrapt
    if o.NClustVar>0
	    minN = T(Inf)

	    combs = [x & 2^y > 0 for x in 2^o.nerrclustvar-1:-1:1, y in o.nerrclustvar-1:-1:0]  # represent all error clustering combinations. First is intersection of all error clustering vars
	    o.clust = Vector{StrClust{T}}(undef, nrows(combs))  # leave out no-cluster combination
	    o.NErrClustCombs = length(o.clust)

	    if o.NClustVar > o.nbootclustvar  # info for grouping by intersections of all bootstrap && clustering vars wrt data; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusters
		    if o.WREnonARubin && !o.granular
		      o.info✻⋂, ID✻⋂ = panelsetupID(o.ID, 1:o.NClustVar)
		    else
		      o.info✻⋂           = panelsetup(o.ID, 1:o.NClustVar)
		    end
				if o.subcluster>0 && nrows(o.info✻) ≠ nrows(o.info✻⋂)
					throw(ErrorException("\nThis program can only perform the subcluster bootstrap when the bootstrap clusters are nested within the (intersections of the) error clusters.\n"))
				end
			else
		    o.info✻⋂ = o.info✻  # info for grouping by intersections of all bootstrap && clustering vars wrt data; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusters
		    o.WREnonARubin && !o.granular && (ID✻⋂ = o.ID✻)
	    end

			Nall = length(o.info✻⋂)

	    if o.NClustVar > o.nerrclustvar  # info for intersections of error clustering wrt data
		    if o.WREnonARubin && !o.granular
		      o.info⋂, ID⋂ = panelsetupID(o.ID, o.subcluster+1:o.NClustVar)
		    else
		      o.info⋂      = panelsetup(o.ID, o.subcluster+1:o.NClustVar)
		    end
		    ID⋂_✻⋂ = length(o.info⋂)==o.Nobs ? o.ID : o.ID[first.(o.info⋂),:]  # version of ID matrix with one row for each all-error-cluster-var intersection instead of 1 row for each obs; gets resorted
		    o.ID✻⋂ = Nall==o.Nobs ? o.ID : o.ID[first.(o.info✻⋂),:]  # version of ID matrix with one row for each all-bootstrap && error cluster-var intersection instead of 1 row for each obs
	    else
		    o.info⋂ = o.info✻⋂  # info for intersections of error clustering wrt data
		    o.WREnonARubin && !o.granular && (ID⋂ = ID✻⋂)
		    o.ID✻⋂ = ID⋂_✻⋂ = nrows(o.info⋂)==o.Nobs ? o.ID : o.ID[first.(o.info⋂),:]  # version of ID matrix with one row for each all-error-cluster-var intersection instead of 1 row for each obs; gets resorted
	    end

	    o.BootClust = 2^(o.NClustVar - o.nbootclustvar)  # location of bootstrap clustering within list of cluster combinations

	    @inbounds for c ∈ 1:o.NErrClustCombs  # for each error clustering combination
		    ClustCols = o.subcluster .+ findall(@view combs[c,:])
		    even = isodd(length(ClustCols))  # not a typo

		    if isone(c)
		      if iszero(o.subcluster)
		    	  order = Vector{Int64}(undef,0)
		    	  info  = Vector{UnitRange{Int64}}(undef, 0)  # causes no collapsing of data in panelsum() calls
						N = Nall
		      else
		    	  order = _sortperm(@view ID⋂_✻⋂[:,ClustCols])
		    	  info  = panelsetup(view(ID⋂_✻⋂,order,:), ClustCols)
						N = nrows(info)
		      end
		    else
		      if any(combs[c, minimum(findall(combs[c,:] .≠ combs[c-1,:])):end])  # if this sort ordering same as last to some point and missing thereafter, no need to re-sort
		    	  order = _sortperm(@view ID⋂_✻⋂[:,ClustCols])
		    	  info = panelsetup(view(ID⋂_✻⋂,order,:), ClustCols)
		      else
		    	  order = Vector{Int64}(undef,0)
						info = panelsetup(ID⋂_✻⋂, ClustCols)
		      end
					N = nrows(info)
		    end

		    if o.small
					multiplier = T(N / (N-1))
					minN = min(minN,N)
				else
					multiplier = one(T)
				end
				o.clust[c] = StrClust{T}(N, multiplier, even, order, info)
	    end

	    (o.scorebs || !o.WREnonARubin) &&
		  	(o.ClustShare = o.haswt ? @panelsum(o, o.wt, o.info⋂)/o.sumwt : T.(length.(o.info⋂)./o.Nobs)) # share of observations by group

    else  # if no clustering, cast "robust" as clustering by observation
      o.clust = StrClust{T}(Nobs, small ? o._Nobs / (o._Nobs - one(T)) : one(T), true, Vector{Int64}(undef,0), Vector{UnitRange{Int64}}(undef,0))
      o.NErrClustCombs = one(Int16)
      (o.scorebs || !o.WREnonARubin) &&
    		(o.ClustShare = o.haswt ? o.wt/o.sumwt : fill(one(T)/o._Nobs,o._Nobs))
    end

		o.purerobust = o.robust && !o.scorebs && iszero(o.subcluster) && o.N✻==o.Nobs  # do we ever error-cluster *and* bootstrap-cluster by individual?
		o.granular   = o.WREnonARubin ? 2*o.Nobs*o.B*(2*o.N✻+1) < o.N✻*(o.N✻*o.Nobs+o.clust[1].N*o.B*(o.N✻+1)) :
											o.NClustVar>0 && !o.scorebs && (o.purerobust || (o.clust[1].N+o.N✻)*o.kZ*o.B + (o.clust[1].N-o.N✻)*o.B + o.kZ*o.B < o.clust[1].N*o.kZ^2 + o.Nobs*o.kZ + o.clust[1].N * o.N✻ * o.kZ + o.clust[1].N * o.N✻)

		if o.robust && !o.purerobust
			(o.subcluster>0 || o.granular) && !o.WREnonARubin &&
				(o.info⋂_✻⋂ = panelsetup(o.ID✻⋂, o.subcluster+1:o.NClustVar))  # info for error clusters wrt data collapsed to intersections of all bootstrapping && error clusters; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusterings
			((o.scorebs && o.B>0) || (o.WREnonARubin && !o.granular && o.bootstrapt)) &&
				(o.JN⋂N✻ = zeros(T, o.clust[1].N, o.N✻))
		end

		if o.WREnonARubin && o.robust && o.bootstrapt && !o.granular
			if iszero(length(ID✻⋂))
				_, ID✻⋂ = panelsetupID(o.ID,              1:o.NClustVar)
				_, ID⋂  = panelsetupID(o.ID, o.subcluster+1:o.NClustVar)
			end
			o.IDCT⋂✻   = Vector{Vector{Int64}}(undef, o.N✻)
			o.infoCT⋂✻ = Vector{Vector{UnitRange{Int64}}}(undef, o.N✻)
			for i ∈ 1:o.N✻
				tmp = ID✻⋂[o.info✻[i]]                    # ID numbers w.r.t. intersection of all bootstrap/error clusterings contained in bootstrap cluster i
				o.infoCT⋂✻[i] = o.info✻⋂[tmp[1]:tmp[end]] # for each of those ID's, panel info for the all-bootstrap/error-clusterings data row groupings
				o.IDCT⋂✻[i] = ID⋂[first.(o.infoCT⋂✻[i])]  # ID numbers of those groupings w.r.t. the all-error-clusterings grouping
			end
		end
  else
	  minN = T(nrows(o.info✻))
	end
o._ID✻⋂=ID✻⋂
	InitFEs(o)

	if o.B>0 && o.robust && o.granular && !o.purerobust && o.bootstrapt && !o.WREnonARubin
		if o.NFE>0 && !o.FEboot
			_, o.ID✻ = panelsetupID(o.ID   , 1:o.nbootclustvar)
		else
			_, o.ID✻_✻⋂  = panelsetupID(o.ID✻⋂, 1:o.nbootclustvar)
		end
	end

	o.enumerate = o.B>0 && o.auxtwtype==rademacher && o.N✻*log(2) < log(o.B)+1e-6  # generate full Rademacher set?
	o.enumerate && (o.maxmatsize = 0)

	o.Nw = iszero(o.maxmatsize) ? one(Int64) : ceil(Int64, (o.B+1) * Float64(max(nrows(o.ID✻), length(o.ID✻_✻⋂), o.N✻) * sizeof(T)) / o.maxmatsize / 1073741824) # 1073741824 = giga(byte)
	if isone(o.Nw)
		MakeWildWeights!(o, o.B, first=true)  # make all wild weights, once
		o.enumerate && (o.B = ncols(o.v) - 1)  # replications reduced to 2^G
		o.WeightGrp = [1:ncols(o.v)]
	else
    o.seed = rand(o.rng, UInt64)
		_B = ceil(Int64, (o.B+1) / o.Nw)
		o.Nw = ceil(Int64, (o.B+1) / _B)
		o.WeightGrp = [(i-1)*_B+1:i*_B for i ∈ 1:o.Nw]
		o.WeightGrp[end] = first(o.WeightGrp[end]):o.B+1
	end

	if o.ML
		o.dof = nrows(o.R)
	else
		if o.ARubin
			o.R = hcat(zeros(o.kX₂,o.kX₁), Matrix(I(o.kX₂)))  # attack surface is all endog vars
			o.R₁ = o.kX₁>0 && nrows(o.R₁)>0 ? hcat(o.R₁[:,1:o.kX₁], zeros(nrows(o.R₁),o.kX₂)) : zeros(0, o.kX)  # and convert model constraints from referring to X₁, Y₂ to X₁, X₂
		end
		o.dof = nrows(o.R)

		if !o.WRE && iszero(o.κ)  # regular OLS
			o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			o.Repl = StrEstimator{T}(true, false, zero(T), zero(T))  # XXX isDGP=1 for Repl? doesn't matter?
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁)  # DGP constraints: model constraints + null if imposed
			setR!(o.Repl, o, o.R₁)  # model constraints only
			InitVarsOLS!(o.DGP, o, o.Repl.R₁perp)
			InitTestDenoms!(o.DGP, o)
			o.M = o.DGP  # StrEstimator object from which to get A, AR, XAR
		elseif o.ARubin
			if o.willplot  # for plotting/CI purposes get original point estimate since not normally generated
				o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
				setR!(o.DGP, o, o.R₁, zeros(T,0,o.kZ))  # no-null model
				InitVarsIV!(o.DGP, o)
				EstimateIV!(o.DGP, o, o.r₁)
				o.confpeak = o.DGP.β̈  # estimated coordinate of confidence peak
			end

			o.DGP = StrEstimator{T}(true, false, zero(T), zero(T))
			setR!(o.DGP, o, o.R₁)
			InitVarsARubin!(o.DGP, o)
			InitTestDenoms!(o.DGP, o)
			o.M = o.DGP  # StrEstimator object from which to get A, AR, XAR
			o.kZ = o.kX

		elseif o.WREnonARubin

			o.DGP = StrEstimator{T}(true, T(o.kX₂ ≠ o.kY₂), zero(T), one(T))
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁, zeros(T,0,o.kZ))  # DGP constraints: model constraints + null if imposed
			InitVarsIV!(o.DGP, o)
			if !o.null  # if not imposing null, then DGP constraints, κ, Hessian, etc. do not vary with r and can be set now
				EstimateIV!(o.DGP, o, o.r₁)
				MakeResidualsIV!(o.DGP, o)
			end
			o.Repl = StrEstimator{T}(false, o.LIML, o.Fuller, o.κ)
			setR!(o.Repl, o, o.R₁, o.R)
			InitVarsIV!(o.Repl, o)
			EstimateIV!(o.Repl, o, o.r₁)
			InitWRE!(o)

			if o.bootstrapt
				o.δdenom_b = zeros(o.Repl.kZ, o.Repl.kZ)
				o.S✻UMZperp = Vector{Matrix{T}}(undef, o.Repl.kZ+1)
				o.S✻UPX     = Vector{Matrix{T}}(undef, o.Repl.kZ+1)
				o._J⋂ = zeros(o.clust[1].N, o.Repl.kZ)
				!o.granular && (o.S✻⋂u₁XinvXX = Matrix{Matrix{T}}(undef, o.Repl.kZ+1, o.N✻))
				if o.LIML || !o.robust
					o.YY✻_b   = zeros(o.Repl.kZ+1, o.Repl.kZ+1)
					o.YPXY✻_b = zeros(o.Repl.kZ+1, o.Repl.kZ+1)
				end
				o.NFE>0 && (o.bootstrapt || !isone(o.κ) || o.LIML) && (o.CTFEU = Vector{Matrix{T}}(undef, o.Repl.kZ+1))
			end

		else  # the score bootstrap for IV/GMM uses a IV/GMM DGP but then masquerades as an OLS test because most factors are fixed during the bootstrap. To conform, need DGP and Repl objects with different R, R₁, one with FWL, one not

			o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁, zeros(T,0,o.kZ))  # DGP constraints: model constraints + null if imposed
			InitVarsIV!(o.DGP, o)
			o.Repl = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			setR!(o.Repl, o, o.R₁, I)  # process replication restraints = model constraints only
			InitVarsIV!(o.Repl, o, o.Repl.R₁perp)
			EstimateIV!(o.Repl, o, o.r₁)  # bit inefficient to estimate in both objects, but maintains the conformity
			InitTestDenoms!(o.Repl, o)
			o.M = o.Repl  # StrEstimator object from which to get A, AR, XAR; DGP follows WRE convention of using FWL, Repl follows OLS convention of not; scorebs for IV/GMM mixes the two
			if !o.null  # if not imposing null, then DGP constraints, κ, Hessian, etc. do not vary with r and can be set now
				EstimateIV!(o.DGP, o, o.r₁)
				MakeResidualsIV!(o.DGP, o)
			end
		end
  end

  if !o.WREnonARubin && o.bootstrapt
		o.denom = [Matrix{T}(undef,0,0) for _ in 1:o.dof, _ in 1:o.dof]
		if o.robust
			o.Kcd =                       Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof)
			o.Jcd = iszero(o.B) ? o.Kcd : Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof)  # if B = 0, Kcd will be multiplied by v, which is all 1's, and will constitute Jcd
		end

		if o.robust && o.granular<o.NErrClustCombs && o.B>0
			inds = o.subcluster>0 ?
				        [CartesianIndex(j, i) for (j,v) ∈ enumerate(o.info⋂_✻⋂) for i ∈ v] :  # crosstab c,c* is wide
								o.NClustVar == o.nbootclustvar ?
										[CartesianIndex(i, i) for i ∈ 1:Nall] :  # crosstab c,c* is square
										[CartesianIndex(i, j) for (j,v) ∈ enumerate(o.clust[o.BootClust].info) for i ∈ v]  # crosstab c,c* is tall
			o.crosstab⋂✻ind = LinearIndices(FakeArray(Tuple(max(inds...))...))[inds]
		end
	end

  o.small && (o.dof_r = o.NClustVar>0 ? minN - one(T) : o._Nobs - (o.kZ + o.NFE))  # floating-point dof_r allows for fractional fweights, FWIW...

  o.sqrt = isone(o.dof)  # work with t/z stats instead of F/chi2?

  if o.small
		o.multiplier = (o.smallsample = (o._Nobs - o.kZ - o.FEdfadj * o.NFE) / (o._Nobs - o.robust)) / o.dof  # divide by # of constraints because F stat is so defined
  else
		o.multiplier = o.smallsample = 1
  end

  !(o.robust || o.ML) && (o.multiplier *= o._Nobs)  # will turn sum of squared errors in denom of t/z into mean
  o.sqrt && (o.multiplier = √o.multiplier)

  o.dist = fill(T(NaN), 1, o.B+1)
  (o.Nw>1 || o.WREnonARubin || (!o.null && o.dof≤2)) && (o.numer = fill(T(NaN), o.dof, o.B+1))

  if o.WREnonARubin
		if o.Repl.kZ>1
			o.bootstrapt && o.robust &&
      	(o.Zyg = Vector{Matrix{T}}(undef,o.Repl.kZ))
			o.numer_b = Vector{T}(undef,nrows(o.Repl.RRpar))
		end
		o.bootstrapt && o.robust &&
    	(o.crosstabBootind = o.Nobs==o.N✻ ? Vector(diagind(FakeArray(o.N✻,o.N✻))) :
			                                    LinearIndices(FakeArray(o.Nobs,o.N✻))[CartesianIndex.(1:o.Nobs, o.ID✻)])
	else
		o.poles = o.anchor = zeros(T,0)
		o.interpolable = o.bootstrapt && o.B>0 && o.null && o.Nw==1 && (iszero(o.κ) || o.ARubin)
		if o.interpolable
			o.∂numer∂r = Vector{Matrix{T}}(undef, o.q)
			o.interpolate_u = !(o.robust || o.ML)
			o.interpolate_u && (o.∂u∂r = Vector{Matrix{T}}(undef, o.q))
			if o.robust
				o.∂denom∂r   = [Matrix{Matrix{T}}(undef, o.dof, o.dof) for _ in 1:o.q]
				o.∂²denom∂r² = [Matrix{Matrix{T}}(undef, o.dof, o.dof) for _ in 1:o.q, _ in 1:o.q]
				o.∂Jcd∂r     = [Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof) for _ in 1:o.q]
			end
		end
  end
	o.initialized = true
	nothing
end

function InitFEs(o::StrBootTest{T}) where T
	if isdefined(o, :FEID) && nrows(o.FEID)>0
		p = _sortperm(o.FEID)
		sortID = o.FEID[p]
		i_FE = 1; o.FEboot = o.B>0 && !o.WREnonARubin && o.NClustVar>0; j = o.Nobs; o._FEID = ones(Int64, o.Nobs)
		o.invFEwt = zeros(T, o.NFE>0 ? o.NFE : o.Nobs)
		o.FEs = Vector{StrFE{T}}(undef, o.NFE>0 ? o.NFE : o.Nobs)
		@inbounds for i ∈ o.Nobs-1:-1:1
			if sortID[i] ≠ sortID[i+1]
				is = @view p[i+1:j]
				if o.haswt
					tmp  = o.wt[is]
					wt = tmp / (sumFEwt = sum(tmp))
				else
					sumFEwt = T(j - i)
					wt = fill(1/sumFEwt, j-i)
				end
				o.FEs[i_FE] = StrFE{T}(is, wt)
				if (o.B>0 && o.robust && o.granular < o.nerrclustvar) || (o.WREnonARubin && o.robust && o.granular && o.bootstrapt)
					o.invFEwt[i_FE] = one(T) / sumFEwt
				end

				j = i

				if o.FEboot  # are all of this FE's obs in same bootstrapping cluster? (But no need to check if B=0 for then CT_WE in 2nd term of (62) orthogonal to v = col of 1's)
					tmp = o.ID[is, 1:o.nbootclustvar]
					o.FEboot = all(tmp .== view(tmp, 1,:)')
				end
				i_FE += 1
			end
			o._FEID[p[i]] = i_FE
		end
		is = @view p[1:j]
		if o.haswt
			tmp = o.wt[is]
			wt = tmp / (sumFEwt = sum(tmp))
		else
			sumFEwt = T(j)
			wt = fill(1/sumFEwt, j)
		end
		o.FEs[i_FE] = StrFE{T}(is, wt)
		o.robust && ((o.B>0 && o.granular < o.nerrclustvar) || (o.WREnonARubin && o.granular && o.bootstrapt)) &&
			(o.invFEwt[i_FE] = 1 / sumFEwt)
		o.NFE = i_FE
		resize!(o.invFEwt, o.NFE)
		resize!(o.FEs    , o.NFE)
		if o.FEboot  # are all of this FE's obs in same bootstrapping cluster?
			tmp = o.ID[is, 1:o.nbootclustvar]
			o.FEboot = all(tmp .== @view tmp[1,:])
		end

		o.X₁ = partialFE(o, o.X₁)  # don't overwrite caller's data
		o.X₂ = partialFE(o, o.X₂)
		o.y₁ = partialFE(o, o.y₁)
		o.Y₂ = partialFE(o, o.Y₂)
	end
	nothing
end

# draw wild weight matrix of width _B. If first=true, insert column of 1s at front. For non-Anderson-Rubin WRE, subtract 1 from all
const ϕ = (1 + √5)/2

function MakeWildWeights!(o::StrBootTest{T}, _B::Integer; first::Bool=true) where T
  if _B>0  # in scoretest or waldtest WRE, still make v a col of 1's
    if o.enumerate
			o.v = o.WREnonARubin ? [zeros(o.N✻) count_binary(o.N✻, -2, 0)] :  # complete Rademacher set
								[ones( o.N✻) count_binary(o.N✻, -1, 1)]
		elseif o.auxtwtype==normal
			o.v = randn(o.rng, T, o.N✻, _B+first)
			o.WREnonARubin && (o.v .-= one(T))
		elseif o.auxtwtype==gamma
			tmp = quantile.(Gamma(4,.5), rand(o.rng, o.N✻, _B+first))
			o.v = T==Float64 ? tmp : T.(tmp)
			o.WREnonARubin && (o.v .-= one(T))
		elseif o.auxtwtype==webb
			o.v = rand(o.rng, T.([-√1.5, -1, -√.5, √.5, 1, √1.5] .- o.WREnonARubin), o.N✻, _B+first)
		elseif o.auxtwtype == mammen
			o.v = getindex.(Ref(T.([1-ϕ; ϕ] .- o.WREnonARubin)), ceil.(Int16, rand(o.rng, o.N✻, _B+first) ./ (ϕ/√5)))
		elseif o.WREnonARubin  # Rademacher
			o.v = -2rand(o.rng, Bool, o.N✻, _B+first)
		else
			o.v = rand(o.rng, Bool, o.N✻, _B+first) .- T(.5)  # rand(o.rng, Bool, o.N✻, _B+first) .- T(.5)
			o.v_sd = .5
		end

		first && !(o.enumerate && isone(o.v_sd)) && (o.v[:,1] .= o.WREnonARubin ? zero(T) : o.v_sd)  # keep original residuals in first entry to compute base model stat
  else
		o.v = Matrix{T}(undef,0,1)  # in places, ncols(v) indicates B -- 1 for classical tests
  end
	nothing
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 77311686-7a2b-11eb-2bc7-95b2a2211969
begin
	using Pkg
	Pkg.activate("MLJ_env", shared=true)

	using MLJ
	using CSV
	using DataFrames
	using PlutoUI
	using Random
	using Test
	using Printf
	using Plots
end

# ╔═╡ 142f3778-7c71-11eb-2739-4d408a8309e2
using Distributions

# ╔═╡ e9f8e942-7b29-11eb-1d64-d7e7ff02159c
begin
	include("./utils.jl");
	include("./dt_utils.jl");
end

# ╔═╡ 5435ff66-7a2b-11eb-1474-b96dcad21315
md"""
## C03w05: Exploring Ensemble Methods

In this assignment, we will explore the use of boosting. We will use the pre-implemented gradient boosted trees in Turi Create. We will:

  - Use SFrames to do some feature engineering.
  - Train a boosted ensemble of decision-trees (gradient boosted trees) on the LendingClub dataset.
  - Predict whether a loan will default along with prediction probabilities (on a validation set).
  - Evaluate the trained model and compare it with a baseline.
  - Find the most positive and negative loans using the learned model.
  - Explore how the number of trees influences classification performance.
"""

# ╔═╡ 9e476ba0-7c5e-11eb-0082-6d4ffa97add4
const DF = Union{DataFrame, SubDataFrame}

# ╔═╡ 771130dc-7a2b-11eb-170f-934c11fceede
md"""
#### Load LendingClub dataset
"""

# ╔═╡ 76f62e40-7a2b-11eb-095e-f36e207f06a2
begin
	loans =  CSV.File(("../../ML_UW_Spec/C03/data/lending-club-data.csv");
	header=true) |> DataFrame;
	first(loans, 3)
end

# ╔═╡ 76de16d4-7a2b-11eb-09c6-e36b6c893c1f
md"""
Like the previous assignment, we reassign the labels to have +1 for a safe loan, and -1 for a risky (bad) loan.
"""

# ╔═╡ 76c17c4a-7a2b-11eb-2c47-97106bd58f99
insertcols!(loans, :safe_loans => ifelse.(loans.bad_loans .== 0, 1, -1),
 	makeunique=true);

# ╔═╡ 76a83684-7a2b-11eb-1961-cb133db8c164
md"""
#### Selecting Features

In this assignment, we will be using a subset of features (categorical and numeric). The features we will be using are **described in the code comments** below. If you are a finance geek, the [LendingClub](https://www.lendingclub.com/) website has a lot more details about these features.

The features we will be using are described below.
"""

# ╔═╡ 768d6b08-7a2b-11eb-388a-8763c8da0a51
begin
	const Features = [
		:grade,                     # grade of the loan (categorical)
		:sub_grade_num,             # sub-grade of the loan as a number from 0 to 1
		:short_emp,                 # one year or less of employment
		:emp_length_num,            # number of years of employment
		:home_ownership,            # home_ownership status: own, mortgage or rent
		:dti,                       # debt to income ratio
		:purpose,                   # the purpose of the loan
		:payment_inc_ratio,         # ratio of the monthly payment to income
		:delinq_2yrs,               # number of delinquincies
		:delinq_2yrs_zero,          # no delinquincies in last 2 years
		:inq_last_6mths,            # number of creditor inquiries in last 6 months
		:last_delinq_none,          # has borrower had a delinquincy
		:last_major_derog_none,     # has borrower had 90 day or worse rating
		:open_acc,                  # number of open credit accounts
		:pub_rec,                   # number of derogatory public records
		:pub_rec_zero,              # no derogatory public records
		:revol_util,                # percent of available credit being used
		:total_rec_late_fee,        # total late fees received to day
		:int_rate,                  # interest rate of the loan
		:total_rec_int,             # interest received to date
		:annual_inc,                # annual income of borrower
		:funded_amnt,               # amount committed to the loan
		:funded_amnt_inv,           # amount committed by investors for the loan
		:installment,               # monthly payment owed by the borrower
  	]

	const Features_ST = Dict{Symbol, DataType}(
    	:grade => OrderedFactor{7},          # because 7 ≠ values
    	:home_ownership => Multiclass{4},
    	:purpose => Multiclass{12},

		:last_delinq_none => OrderedFactor{2},
		:last_major_derog_none => OrderedFactor{2},
		:delinq_2yrs_zero => OrderedFactor{2},
		:pub_rec_zero => OrderedFactor{2},

		## using autotype for the rest
  	)

	const Target = :safe_loans # prediction (y) (+1 means safe, -1 is risky)

	const Target_ST = Dict{Symbol, DataType}(
    	Target => Multiclass{2},
  	)

	## Considered categorical features
	const Cat_Features = [
		keys(Features_ST)...,
		:sub_grade_num,
		:short_emp,
		:emp_length_num,
		:delinq_2yrs,
		:inq_last_6mths,
		:open_acc,
		:pub_rec,
		:annual_inc,
		:funded_amnt,
		:funded_amnt_inv,
	] |> unique

  	## Extract the feature columns and Target column
 	select!(loans, [Features..., Target]);
 	length(names(loans)), names(loans)
end

# ╔═╡ 47382678-7bec-11eb-34ff-9b3613ca036c
md"""
#### Skipping observations with missing values

Recall from the lectures that one common approach to coping with missing values is to skip observations that contain missing values.
"""

# ╔═╡ 56d0dfee-7bec-11eb-2890-7577dda5d4de
begin
	s₀ = size(loans)[1];
	dropmissing!(loans);
	s₁ = size(loans)[1];
	## how many missing row were dropped:
	s₀ - s₁
end

# ╔═╡ 37575388-7c9f-11eb-1263-8d4dc28d711b
length(Cat_Features), Cat_Features

# ╔═╡ e73da3ea-7c92-11eb-3169-f5b8a11e4154
# reduce(vcat, loans[!, :pub_rec_zero]) |> unique

# ╔═╡ 7673d8e4-7a2b-11eb-02b4-2f837181b4e9
md"""
#### Subsample dataset to make sure classes are balanced

Just as we did in the previous assignment, we will undersample the larger class (safe loans) in order to balance out our dataset. This means we are throwing away many data points.
"""

# ╔═╡ ef0b9e82-7c61-11eb-2cdf-2bc585f76e61
begin
	function loan_ratios(safe_, risky_; n=size(loans)[1])
  		(size(safe_)[1] / n, size(risky_)[1] / n)
	end

	function partition_sr(df::DF; target=Target)
		"""Partition between safe and risky loan"""
		(df[df[!, target] .== 1, :], df[df[!, target] .!= 1, :])  # was -1
	end

	function resample!(safe_loans::DF, risky_loans::DF)
		"""
		Since there are fewer risky loans than safe loans, find the ratio of
		the sizes and use that percentage to downsample the safe loans.
		"""
		## => deps on sampling(...)
		n = size(safe_loans)[1]
		perc = size(risky_loans)[1] / n

		## down sample safe loans
		safe_loans = safe_loans[sampling(n, perc; seed=42), :];
		[risky_loans; safe_loans]
	end

end

# ╔═╡ 542db682-7a2e-11eb-35ad-25c82d6919e0
begin
	safe_loans₀, risky_loans₀ = partition_sr(loans)
	p_safe_loans₀, p_risky_loans₀ = loan_ratios(safe_loans₀, risky_loans₀);

	with_terminal() do
		@printf("Number of safe loans  : %d\n", size(safe_loans₀)[1])
		@printf("Number of risky loans : %d\n", size(risky_loans₀)[1])
		@printf("Percentage of safe loans:  %1.2f%%\n", p_safe_loans₀ * 100.)
		@printf("Percentage of risky loans: %3.2f%%\n", p_risky_loans₀ * 100.)
	end
end

# ╔═╡ 4f4034e6-7c5f-11eb-0530-b3e9f0d7cc94
loans_data = resample!(safe_loans₀, risky_loans₀);

# ╔═╡ 762690b8-7a2b-11eb-0257-f5a54932adc5
begin
	safe_loans₁, risky_loans₁ = partition_sr(loans_data)
	p_safe_loans₁, p_risky_loans₁ = loan_ratios(safe_loans₁, risky_loans₁;
		n=size(loans_data)[1]);

	with_terminal() do
		@printf("Number of safe loans  : %d\n", size(safe_loans₁)[1])
		@printf("Number of risky loans : %d\n", size(risky_loans₁)[1])
		@printf("Percentage of safe loans:  %1.2f%%\n", p_safe_loans₁ * 100.)
		@printf("Percentage of risky loans: %3.2f%%\n", p_risky_loans₁ * 100.)
	end
end

# ╔═╡ 3956455a-7c62-11eb-0ed6-cb2cda3ca879
md"""
##### Deal with Categorical features

initial mapping:
"""

# ╔═╡ 9ca8a008-7cad-11eb-1fdc-c31031290a79
with_terminal() do
	for n ∈ names(loans_data)
		@printf("%-25s => %50s\n", n, scitype(loans_data[!, n]))
	end
end

# ╔═╡ e1245fd0-7bf2-11eb-071c-0fb7dd06e885
begin
	## Scientific Type coercion
	coerce_map = [Features_ST..., Target_ST...]
	coerce!(loans_data, coerce_map...);
	coerce_map
end

# ╔═╡ a67aaa04-7cad-11eb-2a40-8585bad43994
with_terminal() do
	for n ∈ names(loans_data)
		@printf("%-25s => %50s\n", n, scitype(loans_data[!, n]))
	end
end

# ╔═╡ e209e38c-7c94-11eb-09fb-279bcffec0f7
md"""

For a number of the remaining features which are treated as Count there are few unique values in which case it might make more sense to recode them as OrderedFactor, this can be done with autotype.
"""

# ╔═╡ 9441f1d0-7c94-11eb-139d-67562ee93876
coerce!(loans_data, autotype(loans_data, :few_to_finite));

# ╔═╡ 9e7aa038-7c9f-11eb-32e9-b3579071450e
with_terminal() do
	for n ∈ names(loans_data)
		@printf("%-25s => %50s\n", n, scitype(loans_data[!, n]))
	end
end

# ╔═╡ 643ee9de-7bf4-11eb-1d64-d7e7ff02159c
md"""
### MLJ AdaBoostStumpClassifier
"""

# ╔═╡ f3999fe6-7c69-11eb-14c2-755155d7fe83
function hot_encode_fn(df::DF, cat_features)
	"""
	hot-encoding of categorical features
	"""
	ohe_mdl = OneHotEncoder(;
		features=cat_features,
		ordered_factor=false,
		ignore=false
	)

	ohe_mach = machine(ohe_mdl, df)
	fit!(ohe_mach)
	MLJ.transform(ohe_mach, df)  ## return df hotencoded
end

# ╔═╡ 72b7ce48-7caf-11eb-30bb-6561f21e7dcf
with_terminal() do
	for n ∈ names(loans_data)
		@printf("%-20s => %-60s\n", n, typeof(loans_data[!, n]))
	end
end

# ╔═╡ 97dcb9c4-7bf4-11eb-071c-0fb7dd06e885
begin
	attrs_to_ohe = Cat_Features
	loans_df_ohe = hot_encode_fn(loans_data, attrs_to_ohe);
	length(names(loans_df_ohe)), names(loans_df_ohe)
end

# ╔═╡ 7d4a709a-7caf-11eb-0666-893547dad910
with_terminal() do
	for n ∈ names(loans_df_ohe)
		@printf("%-25s => %-60s\n", n, typeof(loans_df_ohe[!, n]))
	end
end

# ╔═╡ a56e44e0-7a32-11eb-2587-6f5f8b8c21f3
md"""
#### Train-Validation split

We split the data into a train/validation split with 80% of the data in the training set and 20% of the data in the validation set.
"""

# ╔═╡ 6fad8afe-7a33-11eb-3208-6b289c3ccd23
begin
	train_data, validation_data = train_test_split(loans_df_ohe; split=0.8, seed=42);
	size(train_data), size(validation_data)
end

# ╔═╡ a5ff6eea-7c62-11eb-2cdf-2bc585f76e61
begin
	attrs = [Symbol("safe_loans__-1"), :safe_loans__1]

	Xₜ = select(train_data, Not(Target)) # Not(attrs))
	yₜ = select(train_data, Target)
	yₜ = yₜ[:, Target];   ## cast to vector
end

# ╔═╡ f6c23cf6-7cae-11eb-042d-1153a4f16773
with_terminal() do
	for n ∈ names(Xₜ)
		@printf("%-25s => %50s\n", n, scitype(Xₜ[!, n]))
	end
end

# ╔═╡ 14d15ed4-7caf-11eb-2cfe-c91267cec94f
with_terminal() do
	# for n ∈ names(Xₜ)
		@printf("%-20s => %50s\n", Target, scitype(yₜ))
	# end
end

# ╔═╡ e018a54a-7c5d-11eb-34ff-9b3613ca036c
@load AdaBoostStumpClassifier

# ╔═╡ 79ac6bc6-7c84-11eb-3e1d-515509843831
function abs_clf(X, y; n_iter=5)
	ada_clf = MLJDecisionTreeInterface.AdaBoostStumpClassifier(n_iter=n_iter)
	ada_mach = machine(ada_clf, X, y)
	fit!(ada_mach, rows=1:size(X)[1])
	ada_mach
end

# ╔═╡ 5c4a9dee-7bff-11eb-34ff-9b3613ca036c
ada_mach = abs_clf(Xₜ, yₜ; n_iter=5)

# ╔═╡ 322503ba-7c5e-11eb-3374-fb3932c68191
fitted_params(ada_mach)

# ╔═╡ edbb6bb8-7c87-11eb-3124-8506ee0b6781
ya_accuracy(ŷ, y) = sum(ŷ .== y) / length(y)

# ╔═╡ 7ee1f0dc-7c86-11eb-204d-4726e4d62fd4
# begin
	# evaluate!(ada_mach,
	#           resampling=Holdout(fraction_train=0.7, shuffle=true, rng=1234),
	#           measure=[Accuracy()],
	# 		    check_measure=false)
	# evaluate!(ada_mach, resampling=Holdout(), measure=[ya_accuracy],
	# 	operation=predict, 	verbosity=1)

	# r = evaluate!(ada_mach, resampling=CV(nfolds=3), measure=[ya_accuracy],
	# 	operation=predict, 	verbosity=0)
	# (r.measure, r.measurement)
# end

# ╔═╡ 7e79f75c-7c63-11eb-0429-374262cb7f0d
md"""
#### Making predictions

Let us consider a few positive and negative examples from the validation set. We will do the following:

  - Predict whether or not a loan is likely to default.
  - Predict the probability with which the loan is likely to default.
"""

# ╔═╡ b3a8e8b6-7c63-11eb-2254-d1728a2f197d
begin
  	v_safe_loans, v_risky_loans = partition_sr(validation_data; target=Target)
	sample_valid_data = vcat(v_safe_loans[1:2, :], v_risky_loans[1:2, :])
end

# ╔═╡ ddf3fcec-7c6c-11eb-0530-b3e9f0d7cc94
function do_predict(mach, Xₙ, yₙ)
	"""
	Note:
	using sample_valid_data:
	mode is: [1, 1, -1, 1]
	proba:   [0.606493, 0.542905, 0.606493, 0.393507]

	raw_proba (per level -1 / 1)
	  UnivariateFinite{ScientificTypes.Multiclass{2}}(-1=>0.394, 1=>0.606)
      UnivariateFinite{ScientificTypes.Multiclass{2}}(-1=>0.457, 1=>0.543)
	  UnivariateFinite{ScientificTypes.Multiclass{2}}(-1=>0.606, 1=>0.394)
      UnivariateFinite{ScientificTypes.Multiclass{2}}(-1=>0.394, 1=>0.606)

	this means that the probas for predicting [1, 1, -1, 1] are:
		[0.606493, 0.542905, 0.606493, 0.606493]  # see how last oen differs from proabs given above...

	Theerefore here, I am going to use for the proba, the following:
	  eachrow(pdf(ŷ_probs, levels(yₙ))) |> r -> maximum.(r)

	"""
	ŷ_probs = predict(ada_mach, Xₙ)
	ŷ = broadcast(mode, ŷ_probs)

	ratio_corr_pred = sum(ŷ .== yₙ) / size(Xₙ)[1]
	proba = eachrow(pdf(ŷ_probs, levels(yₙ))) |> r -> maximum.(r)

	# used: proba=broadcast(pdf, ŷ_probs, yₙ) which is not what I am expecting...
	(mode=ŷ, raw_proba=ŷ_probs, proba=proba,
		ratio_corr=ratio_corr_pred, pdf_per_level=pdf(ŷ_probs, levels(yₙ)))
end

# ╔═╡ 73b7d41c-7c66-11eb-224b-1d7474295bd2
begin
	# using Distributions
	Xsv = select(sample_valid_data, Not(Target))
	ysv = sample_valid_data.safe_loans

	predsv = do_predict(ada_mach, Xsv, ysv)

	with_terminal() do
		println("Predictions   mode: $(predsv.mode)")
		println("Predictions probas: $(round.(predsv.proba, digits=4))")
		# println("Predictions probas/level: $(predsv.pdf_per_level)")
		println("Ratio of correct predictions: $(round(predsv.ratio_corr, digits=4))")
	end
end

# ╔═╡ dfe9ab1e-7c71-11eb-21e1-7da6f486bb75
predsv.mode

# ╔═╡ e805a514-7c71-11eb-13cd-a73039c33692
predsv.raw_proba

# ╔═╡ 0773fae6-7c71-11eb-2638-8fe05034ba2d
predsv.pdf_per_level

# ╔═╡ c61a11c8-7c72-11eb-246e-15a9c4f4cb68
eachrow(predsv.pdf_per_level) |> r -> maximum.(r)

# ╔═╡ f093cc04-7c6e-11eb-1d27-2173c9ebeab7
## least likely to be a safe loan
begin
	ix_least = argmin(predsv.proba)
	(ix_least, round(predsv.proba[ix_least], digits=4))
end

# ╔═╡ 3c4f3a96-7c6c-11eb-0082-6d4ffa97add4
md"""
**Quiz Question**: What percentage of the predictions on sample_validation_data did model_5 get correct?
  - 75% (cf. above cell)

**Quiz Question**: According to our model, which loan is the LEAST likely to be a safe loan?

  - the second one

**Checkpoint**: Can you verify that for all the predictions with probability >= 0.5, the model predicted the label +1?
  - Ok see below.
"""

# ╔═╡ c7a5de14-7c75-11eb-1c8e-d718788cb40c
all(p -> p ≥ 0.5, predsv.proba[predsv.mode .== 1])

# ╔═╡ bf1ff544-7c6e-11eb-37f2-09f2102bbc7b
md"""
#### Evaluating the model on the validation data
"""

# ╔═╡ e27ffc5a-7c96-11eb-23b0-470fe22ceaea
begin
	Xᵥ = select(validation_data, Not(Target))
	yᵥ = select(validation_data, Target)
	yᵥ = yᵥ[:, Target];   ## cast to vector
end

# ╔═╡ 6d9cc902-7c70-11eb-3b71-ff1408d9d504
begin
	## on whole validation data
	preds = do_predict(ada_mach, Xᵥ, yᵥ)
	lim = 8

	with_terminal() do
		println("Predictions   mode: $(preds.mode[1:lim])...")
		println("Predictions probas: $(round.(preds.proba[1:lim], digits=4))")
		println("Ratio of correct predictions: $(round(preds.ratio_corr, digits=4))")
	end
end

# ╔═╡ 41919468-7c76-11eb-3389-471f19f9e92f
md"""
Calculate the number of *false positives* and the the number of *false negatives* made by the model where:

  - *false positives* are loans that were predicted to be safe but were actually risky
  - *false negatives* are loans that were were predicted to be risky but were actually safe.

"""

# ╔═╡ f698c078-7c76-11eb-0841-37338a78bd2c

fp = (preds.mode .== 1) |>                          ## predicted safe == 1
  ba -> (validation_data.safe_loans[ba] .== -1) |>  ## but actually risky == -1
  sum

# ╔═╡ 4bd262b2-7c77-11eb-310a-e1f639f135d5
fn = (preds.mode .== -1) |>                          ## predicted riksy == -1
  ba -> (validation_data.safe_loans[ba] .== 1) |>    ## but actually safe == 1
  sum

# ╔═╡ 9f9470a6-7c78-11eb-006f-8b06746cbc01
md"""
Using the costs defined above and the number of false positives and false negatives for the decision tree, we can calculate the total cost of the mistakes made by the decision tree model as follows:

cost = $10,000 * 1936  + $20,000 * 1503 = $49,420,000

The total cost of the mistakes of the model is $49.42M. That is a lot of money!.

**Quiz Question**: Using the same costs of the false positives and false negatives, what is the cost of the mistakes made by the boosted tree model (model_5) as evaluated on the validation_set?
"""

# ╔═╡ 9f6eb654-7c78-11eb-1a81-9ba2e6a0cd6d
cost = fp * 20000 + fn * 10000

# ╔═╡ c6261684-7c78-11eb-02b4-bb1911d035a7
md"""
#### Most positive & negative loans.

In this section, we will find the loans that are most likely to be predicted safe. We can do this in a few steps:

  1. Use the ada_mach (the model with 5 trees) and make probability predictions for all the loans in the validation dataset (*already done*).
  1. Add theses probability predictions as a column called predictions into the validation_data dataframe.
  1. Sort the data (in descreasing order) by the probability predictions.

"""

# ╔═╡ f9c2b1f2-7c80-11eb-333e-1da96d113d79
begin
	valid_ext_data = hcat(validation_data, preds.proba, makeunique=true)
	rename!(valid_ext_data, :x1 => :predictions)

	first(valid_ext_data, 3)
end

# ╔═╡ 9f01203c-7c80-11eb-0666-893547dad910
## Top 7 predictions
with_terminal() do
	@show sort(valid_ext_data, :predictions, rev=true)[1:7, [:grade, Target, :predictions]]
end

# ╔═╡ 9ee273a8-7c80-11eb-30bb-6561f21e7dcf
## bottom 7 predictions
with_terminal() do
	@show sort(valid_ext_data, :predictions, rev=false)[1:7, [:grade, Target, :predictions]]
end

# ╔═╡ d5d5c1f0-7c83-11eb-0987-bfc19d4850e3
md"""

### Effect of adding more trees

In this assignment, we will train 5 different ensemble classifiers in the form of gradient boosted trees. We will train models with 10, 50, 100, 200, and 500 trees. We use the n_iter parameter in AdaBoostStumpClassifier.

"""

# ╔═╡ 005e87fe-7c84-11eb-01f1-8fddeb393779
const N_ITERS = [10, 50,  75, 100, 125, 150, 200, 250, 300]

# ╔═╡ 0ca7bc74-7c84-11eb-252c-3bf3df6b6635
function try_models(X, y)
	hsh = Dict{Symbol, Any}()

  	for (ix, n_iter) ∈ enumerate(N_ITERS)
		sym = Symbol("model_", ix)
		hsh[sym] = Dict{Symbol, Any}(:ada_model => nothing, :n_iter => n_iter)
    	hsh[sym][:ada_model] = abs_clf(X, y; n_iter)
	end

	hsh
end

# ╔═╡ f9ff756a-7c85-11eb-0b45-0f1b3735844f
function eval_acc_model(mach, X, y)
	ŷ = predict_mode(mach, X)
	sum(ŷ .== y) / length(y)
end

# ╔═╡ 7dccb71e-7c8a-11eb-175d-0fc5400512d8
hsh = try_models(Xₜ, yₜ); # using the training data...

# ╔═╡ b33f43b2-7c85-11eb-21e7-f30c1eeb4785
begin
	## On vaLidation data
	best_model, best_acc = nothing, nothing
	m_acc = []

	for m ∈ keys(hsh)
  		acc = eval_acc_model(hsh[m][:ada_model], Xᵥ, yᵥ)
		push!(m_acc, acc)

  		if isnothing(best_acc) || acc > best_acc
    		global best_model = m
    		global best_acc = acc
		end
	end

	with_terminal() do
		for (ix, acc) ∈ enumerate(m_acc)
			@printf("- model_%d - validation accuracy %1.4f\n",
				ix, round(acc, digits=4))
		end
		@printf("\n+ best: %7s - validation accuracy %1.4f / max iter: %3d\n",
			best_model, round(best_acc, digits=4), hsh[best_model][:n_iter])
	end
end

# ╔═╡ 662e0a14-7c8e-11eb-245a-07169e1831ff
md"""
**Quiz Question**: Which model has the best accuracy on the validation_data?
  - model 4 (200 iter)

**Quiz Question**: Is it always true that the model with the most trees will perform best on test data?
  - No

"""

# ╔═╡ b9c2ca34-7c8e-11eb-1fa5-2d491d1dd7c6
begin
	train_errors, valid_errors = Float64[], Float64[]

	for m ∈ keys(hsh)
  		train_err = 1. - eval_acc_model(hsh[m][:ada_model], Xₜ, yₜ)
 		valid_err = 1. - eval_acc_model(hsh[m][:ada_model], Xᵥ, yᵥ)
  		push!(train_errors, train_err)
  		push!(valid_errors, valid_err)
	end

	train_errors, valid_errors
end

# ╔═╡ 98ba3708-7c90-11eb-0ba5-01dad28c3126
begin
	plot(N_ITERS, train_errors, color=[:lightblue], linewidth=3,
		label="train_error", leg=:topright)
	plot!(N_ITERS, valid_errors, color=[:orange], linewidth=3,
		label="valid_error", leg=:topright)
	plot!(ylims=(0.25, 0.45))
end

# ╔═╡ Cell order:
# ╟─5435ff66-7a2b-11eb-1474-b96dcad21315
# ╠═77311686-7a2b-11eb-2bc7-95b2a2211969
# ╠═9e476ba0-7c5e-11eb-0082-6d4ffa97add4
# ╟─771130dc-7a2b-11eb-170f-934c11fceede
# ╠═76f62e40-7a2b-11eb-095e-f36e207f06a2
# ╟─76de16d4-7a2b-11eb-09c6-e36b6c893c1f
# ╠═76c17c4a-7a2b-11eb-2c47-97106bd58f99
# ╠═e9f8e942-7b29-11eb-1d64-d7e7ff02159c
# ╟─76a83684-7a2b-11eb-1961-cb133db8c164
# ╠═768d6b08-7a2b-11eb-388a-8763c8da0a51
# ╟─47382678-7bec-11eb-34ff-9b3613ca036c
# ╠═56d0dfee-7bec-11eb-2890-7577dda5d4de
# ╠═37575388-7c9f-11eb-1263-8d4dc28d711b
# ╠═e73da3ea-7c92-11eb-3169-f5b8a11e4154
# ╟─7673d8e4-7a2b-11eb-02b4-2f837181b4e9
# ╠═ef0b9e82-7c61-11eb-2cdf-2bc585f76e61
# ╠═542db682-7a2e-11eb-35ad-25c82d6919e0
# ╠═4f4034e6-7c5f-11eb-0530-b3e9f0d7cc94
# ╠═762690b8-7a2b-11eb-0257-f5a54932adc5
# ╟─3956455a-7c62-11eb-0ed6-cb2cda3ca879
# ╠═9ca8a008-7cad-11eb-1fdc-c31031290a79
# ╠═e1245fd0-7bf2-11eb-071c-0fb7dd06e885
# ╠═a67aaa04-7cad-11eb-2a40-8585bad43994
# ╟─e209e38c-7c94-11eb-09fb-279bcffec0f7
# ╠═9441f1d0-7c94-11eb-139d-67562ee93876
# ╠═9e7aa038-7c9f-11eb-32e9-b3579071450e
# ╟─643ee9de-7bf4-11eb-1d64-d7e7ff02159c
# ╠═f3999fe6-7c69-11eb-14c2-755155d7fe83
# ╠═72b7ce48-7caf-11eb-30bb-6561f21e7dcf
# ╠═97dcb9c4-7bf4-11eb-071c-0fb7dd06e885
# ╠═7d4a709a-7caf-11eb-0666-893547dad910
# ╟─a56e44e0-7a32-11eb-2587-6f5f8b8c21f3
# ╠═6fad8afe-7a33-11eb-3208-6b289c3ccd23
# ╠═a5ff6eea-7c62-11eb-2cdf-2bc585f76e61
# ╠═f6c23cf6-7cae-11eb-042d-1153a4f16773
# ╠═14d15ed4-7caf-11eb-2cfe-c91267cec94f
# ╠═e018a54a-7c5d-11eb-34ff-9b3613ca036c
# ╠═79ac6bc6-7c84-11eb-3e1d-515509843831
# ╠═5c4a9dee-7bff-11eb-34ff-9b3613ca036c
# ╠═322503ba-7c5e-11eb-3374-fb3932c68191
# ╠═edbb6bb8-7c87-11eb-3124-8506ee0b6781
# ╠═7ee1f0dc-7c86-11eb-204d-4726e4d62fd4
# ╟─7e79f75c-7c63-11eb-0429-374262cb7f0d
# ╠═b3a8e8b6-7c63-11eb-2254-d1728a2f197d
# ╠═142f3778-7c71-11eb-2739-4d408a8309e2
# ╠═ddf3fcec-7c6c-11eb-0530-b3e9f0d7cc94
# ╠═73b7d41c-7c66-11eb-224b-1d7474295bd2
# ╠═dfe9ab1e-7c71-11eb-21e1-7da6f486bb75
# ╠═e805a514-7c71-11eb-13cd-a73039c33692
# ╠═0773fae6-7c71-11eb-2638-8fe05034ba2d
# ╠═c61a11c8-7c72-11eb-246e-15a9c4f4cb68
# ╠═f093cc04-7c6e-11eb-1d27-2173c9ebeab7
# ╟─3c4f3a96-7c6c-11eb-0082-6d4ffa97add4
# ╠═c7a5de14-7c75-11eb-1c8e-d718788cb40c
# ╟─bf1ff544-7c6e-11eb-37f2-09f2102bbc7b
# ╠═e27ffc5a-7c96-11eb-23b0-470fe22ceaea
# ╠═6d9cc902-7c70-11eb-3b71-ff1408d9d504
# ╟─41919468-7c76-11eb-3389-471f19f9e92f
# ╠═f698c078-7c76-11eb-0841-37338a78bd2c
# ╠═4bd262b2-7c77-11eb-310a-e1f639f135d5
# ╟─9f9470a6-7c78-11eb-006f-8b06746cbc01
# ╠═9f6eb654-7c78-11eb-1a81-9ba2e6a0cd6d
# ╟─c6261684-7c78-11eb-02b4-bb1911d035a7
# ╠═f9c2b1f2-7c80-11eb-333e-1da96d113d79
# ╠═9f01203c-7c80-11eb-0666-893547dad910
# ╠═9ee273a8-7c80-11eb-30bb-6561f21e7dcf
# ╟─d5d5c1f0-7c83-11eb-0987-bfc19d4850e3
# ╠═005e87fe-7c84-11eb-01f1-8fddeb393779
# ╠═0ca7bc74-7c84-11eb-252c-3bf3df6b6635
# ╠═f9ff756a-7c85-11eb-0b45-0f1b3735844f
# ╠═7dccb71e-7c8a-11eb-175d-0fc5400512d8
# ╠═b33f43b2-7c85-11eb-21e7-f30c1eeb4785
# ╟─662e0a14-7c8e-11eb-245a-07169e1831ff
# ╠═b9c2ca34-7c8e-11eb-1fa5-2d491d1dd7c6
# ╠═98ba3708-7c90-11eb-0ba5-01dad28c3126


text:
# resfile = "/home/yakir/downloads/20141106_30_Dtowards.res"

ncol(io) = size(read(io, "xdata"), 2)

function newresfile()
    @label resfileq
    dialog = Dialog()
    resfile = request("What is the path to the `.res` file?", dialog)
    if !isfile(resfile)
        @warn "$resfile is not a valid file, try again…"
        @goto resfileq
    end
    matopen(resfile) do io
        n = ncol(io)
        ids = Vector{UUID}(undef, n)
        for i in 1:n
            xyt = getcoordinates(io, i)
            @assert !isempty(xyt) "no coordinates in column $i"
            id = uuid1()
            writedlm(joinpath(pixelfolder, "$id.csv"), xyt)
            ids[i] = id
        end
        ids
    end
end

function getcoordinates(io, i)
    _x = read(io, "xdata")[:, i]
    x = nonzeros(_x)
    y = nonzeros(read(io, "ydata")[:,i])
    fr = read(io, "status")["FrameRate"]
    if length(x) == 1
        t = Float64.(findfirst(!iszero, _x))
        t /= fr
        [x[1] y[1] t]
    else
        t = Float64.(findall(!iszero, _x))
        t ./= fr
        [x y t]
    end
end