text:
function pow1_test()
    #=
        max z
    s.t.
        x^0.9 * y^(0.1) >= |z| (i.e (x, y, z) are in the 3d power cone with a=0.9)
        x == 2
        y == 1
    =#
    model = TestModel{Float64}()

    a = 0.9
    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    vov = MOI.VectorOfVariables(v)

    vc = MOI.add_constraint(model, vov, MOI.PowerCone(a));

    cx = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[1])], 0.), MOI.EqualTo(2.))
    cy = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[2])], 0.), MOI.EqualTo(1.))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[3])], 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    return model
end

function pow2_test()
    #=
        max z
    s.t.
        x^0.9 * y^(0.1) >= |z| (i.e (x, y, z) are in the 3d power cone with a=0.9)
        x == 2
        y == 1
    =#
    a = 0.9
    model = TestModel{Float64}()
    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    vov = MOI.VectorOfVariables(v)

    vc = MOI.add_constraint(model, MOI.VectorAffineFunction{Float64}(vov), MOI.PowerCone(a))

    cx = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[1])], 0.), MOI.EqualTo(2.))
    cy = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[2])], 0.), MOI.EqualTo(1.))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[3])], 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    return model
end


text:
# Variables in functions have local scope

function sumPlusTwo(x, y)
  a = 2
  x + a + y
end

# println(a) # UndefVarError - a not defined outside its local scope
println(sumPlusTwo(1, 1))

# Loops have local scoped Variables
for i = 1:10
  a = 2
  m = a * i
  println("$a * $i = $m")
end

# UndefVarError for both of these
# println(a)
# println(m)

for i = 1:1
  a = 2
  println("Outer loop: a = $a")
  for i = 1:1
    a = 3
    local b = 1
    println("Inner loop a = $a")
    println("Inner loop b = $b")
  end
  println("Outer loop: a = $a")
  # println("Outer loop: b = $b") - b not defined here
end

ext: jl

max_stars_repo_name: cadamsmith/julia-programming

max_issues_repo_name: cadamsmith/julia-programming

max_forks_repo_name: cadamsmith/julia-programming


text:
using Test

using HalfSpaceTrees

@testset "basic" begin
    x = [Dict("x" => e, "y" => e, "z" => e) for e in [0.5, 0.45, 0.43, 0.44, 0.445, 0.45, 0.0]]
    hst = HalfSpaceTree(ntrees=10, height=3, windowsize=3)
    for e in x[1:3]
        learn!(hst, e)
    end
    @test score(hst, x[end - 1]) < 0.5
    @test score(hst, x[end]) > 0.5
    @test score(hst, Dict("x" => 0.0, "y" => 0.0)) > 0.5
    @test score(hst, Dict("a" => 0.0, "b" => 0.0)) > 0.5
end

@testset "edge cases" begin
    @test_throws ArgumentError HalfSpaceTree(ntrees=0)
    @test_throws ArgumentError HalfSpaceTree(height=0)
    @test_throws ArgumentError HalfSpaceTree(windowsize=0)
end


text:
"""
```
eqcond(m::THANK)
```
 Expresses the equilibrium conditions in canonical form using Γ0, Γ1, C, Ψ, and Π matrices.
Using the mappings of states/equations to integers defined in jpt.jl, coefficients are
specified in their proper positions.

Γ0 (n_states x n_states) holds coefficients of current time states.
Γ1 (n_states x n_states) holds coefficients of lagged states.
C  (n_states x 1) is a vector of constants
Ψ  (n_states x n_shocks_exogenous) holds coefficients of iid shocks.
Π  (n_states x n_states_expectational) holds coefficients of expectational states.
"""
function eqcond(m::THANK)
    endo = m.endogenous_states
    exo  = m.exogenous_shocks
    ex   = m.expected_shocks
    eq   = m.equilibrium_conditions

    Γ0 = zeros(n_states(m), n_states(m))
    Γ1 = zeros(n_states(m), n_states(m))
    C  = zeros(n_states(m))
    Ψ  = zeros(n_states(m), n_shocks_exogenous(m))
    Π  = zeros(n_states(m), n_shocks_expectational(m))

    ### ENDOGENOUS STATES ###

    ### 1. Production function

    # Sticky prices
    Γ0[eq[:eq_y], endo[:y_t]] = 1.
    Γ0[eq[:eq_y], endo[:k_t]] = -((m[:y_ss] + m[:F_ss]) / m[:y_ss]) * m[:α]
    Γ0[eq[:eq_y], endo[:L_t]] = -((m[:y_ss] + m[:F_ss]) / m[:y_ss]) * (1. - m[:α])

    # Flexible prices
    Γ0[eq[:eq_y_f], endo[:y_f_t]] = 1.
    Γ0[eq[:eq_y_f], endo[:k_f_t]] = -((m[:y_ss] + m[:F_ss]) / m[:y_ss]) * m[:α]
    Γ0[eq[:eq_y_f], endo[:L_f_t]] = -((m[:y_ss] + m[:F_ss]) / m[:y_ss]) * (1. - m[:α])

    ### 2. Cost minimization

    # Sticky prices
    Γ0[eq[:eq_L], endo[:Rk_t]] = 1.
    Γ0[eq[:eq_L], endo[:k_t]]  = 1.
    Γ0[eq[:eq_L], endo[:w_t]]  = -1.
    Γ0[eq[:eq_L], endo[:L_t]]  = -1.

    # Flexible prices
    Γ0[eq[:eq_L_f], endo[:Rk_f_t]] = 1.
    Γ0[eq[:eq_L_f], endo[:k_f_t]]  = 1.
    Γ0[eq[:eq_L_f], endo[:w_f_t]]  = -1.
    Γ0[eq[:eq_L_f], endo[:L_f_t]]  = -1.

    ### 3. Marginal cost

    # Sticky prices
    Γ0[eq[:eq_mc], endo[:mc_t]]  = 1.
    Γ0[eq[:eq_mc], endo[:Rk_t]] = -m[:α]
    Γ0[eq[:eq_mc], endo[:w_t]]  = -(1. - m[:α])

    # Flexible prices
    Γ0[eq[:eq_mc_f], endo[:mc_f_t]]  = 1.
    Γ0[eq[:eq_mc_f], endo[:Rk_f_t]] = -m[:α]
    Γ0[eq[:eq_mc_f], endo[:w_f_t]]  = -(1. - m[:α])

    ### 4. Phillips Curve

    # Sticky prices
    Γ0[eq[:eq_π], endo[:π_t]]   = 1.
    Γ0[eq[:eq_π], endo[:Eπ_t]]  = -m[:β] / (1 + m[:ι_p] * m[:β])
    Γ1[eq[:eq_π], endo[:π_t]]   = m[:ι_p] / (1 + m[:ι_p] * m[:β])
    Γ0[eq[:eq_π], endo[:mc_t]]   = -(1. - m[:β] * m[:ξ_p]) * (1. - m[:ξ_p]) / ((1 + m[:ι_p] * m[:β]) * m[:ξ_p])
    Γ0[eq[:eq_π], endo[:λ_p_t]] = -1.

    # Flexible prices
    Γ0[eq[:eq_R_f], endo[:mc_f_t]] = 1.

    ### 5. Consumption FOC

    # Sticky prices
    Γ0[eq[:eq_c], endo[:c_t]]  = 1
    Γ0[eq[:eq_c], endo[:c_h_t]]= -m[:θ] * (m[:c_h_ss] / m[:c_ss])
    Γ0[eq[:eq_c], endo[:c_s_t]]= -(1-m[:θ]) * (m[:c_s_ss] / m[:c_ss])

    Γ0[eq[:eq_c_h], endo[:c_h_t]] = 1
    Γ0[eq[:eq_c_h], endo[:w_t]] = -m[:w_ss] * m[:expL_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h], endo[:L_t]] = -m[:w_ss] * m[:expL_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h], endo[:t_h_t]] = -m[:y_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h], endo[:g_t]]  = m[:gL_ss] / m[:c_h_L_ss]

    # Flexible prices
    Γ0[eq[:eq_c_f], endo[:c_f_t]]  = 1
    Γ0[eq[:eq_c_f], endo[:c_h_f_t]]= -m[:θ] * (m[:c_h_ss] / m[:c_ss])
    Γ0[eq[:eq_c_f], endo[:c_s_f_t]]= -(1-m[:θ]) * (m[:c_s_ss] / m[:c_ss])

    Γ0[eq[:eq_c_h_f], endo[:c_h_f_t]] = 1
    Γ0[eq[:eq_c_h_f], endo[:w_f_t]] = -m[:w_ss] * m[:expL_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h_f], endo[:L_f_t]] = -m[:w_ss] * m[:expL_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h_f], endo[:t_h_f_t]] = -m[:y_ss] / m[:c_h_ss]
    Γ0[eq[:eq_c_h_f], endo[:g_t]] = m[:gL_ss] / m[:c_h_L_ss]

    ### 6. Euler equation

    # Sticky prices
    expγ = exp(m[:γ])
    Γ0[eq[:eq_λ], endo[:λ_t]]     = 1
    Γ0[eq[:eq_λ], endo[:λ_h_t]]   = -m[:θ] *(m[:λ_h_ss] / m[:λ_ss])
    Γ0[eq[:eq_λ], endo[:λ_s_t]]   = -(1 - m[:θ]) * (m[:λ_s_ss] / m[:λ_ss])

    Γ0[eq[:eq_λ_s], endo[:λ_s_t]] = 1
    Γ0[eq[:eq_λ_s], endo[:b_t]]   = -1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_λ_s], endo[:z_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])
    Γ0[eq[:eq_λ_s], endo[:c_s_t]] = expγ * m[:c_s_ss] / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])
    Γ1[eq[:eq_λ_s], endo[:c_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])

    Γ0[eq[:eq_λ_h], endo[:λ_h_t]] = 1
    Γ0[eq[:eq_λ_h], endo[:b_t]]   =  -1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_λ_h], endo[:z_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])
    Γ0[eq[:eq_λ_h], endo[:c_h_t]] = expγ * m[:c_h_ss] / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])
    Γ1[eq[:eq_λ_h], endo[:c_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])

    Γ0[eq[:eq_c_s], endo[:λ_s_t]] = 1
    Γ0[eq[:eq_c_s], endo[:R_t]]   = -1
    Γ0[eq[:eq_c_s], endo[:z_t]]   = m[:ρ_z]
    Γ0[eq[:eq_c_s], endo[:Eπ_t]]  = 1
    Γ0[eq[:eq_c_s], endo[:Eλ_s_t]]= -(m[:σ] * m[:λ_s_ss])     / (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])
    Γ0[eq[:eq_c_s], endo[:Eλ_h_t]]= -((1-m[:σ]) * m[:λ_h_ss]) / (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])
    Γ0[eq[:eq_c_s], endo[:Ey_t]]  = -(m[:σ′]) * m[:σ] * (m[:λ_s_ss] - m[:λ_h_ss]) /
                                        (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])
    # Flexible prices
    Γ0[eq[:eq_λ_f], endo[:λ_f_t]]     = 1.
    Γ0[eq[:eq_λ_f], endo[:λ_h_f_t]]   = -m[:θ] *(m[:λ_h_ss] / m[:λ_ss])
    Γ0[eq[:eq_λ_f], endo[:λ_s_f_t]]   = -(1 - m[:θ]) * (m[:λ_s_ss] / m[:λ_ss])

    Γ0[eq[:eq_λ_s_f], endo[:λ_s_f_t]] = 1
    Γ0[eq[:eq_λ_s_f], endo[:b_t]]     = -1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_λ_s_f], endo[:z_t]]     = m[:h] * m[:c_ss]  / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])
    Γ0[eq[:eq_λ_s_f], endo[:c_s_f_t]] = expγ * m[:c_s_ss] / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])
    Γ1[eq[:eq_λ_s_f], endo[:c_f_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_s_ss] - m[:h] * m[:c_ss])

    Γ0[eq[:eq_λ_h_f], endo[:λ_h_f_t]] = 1
    Γ0[eq[:eq_λ_h_f], endo[:b_t]]     = -1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_λ_h_f], endo[:z_t]]     = m[:h] * m[:c_ss]  / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])
    Γ0[eq[:eq_λ_h_f], endo[:c_h_f_t]] = expγ * m[:c_h_ss] / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])
    Γ1[eq[:eq_λ_h_f], endo[:c_f_t]]   = m[:h] * m[:c_ss]  / (expγ * m[:c_h_ss] - m[:h] * m[:c_ss])

    Γ0[eq[:eq_c_s_f], endo[:λ_s_f_t]] = 1
    Γ0[eq[:eq_c_s_f], endo[:R_f_t]]   = -1
    Γ0[eq[:eq_c_s_f], endo[:z_t]]     = m[:ρ_z]
    # Γ0[eq[:eq_c_s_f], endo[:Eπ_t]]  = 1  # matlab code says "CHECK 'epstar' "
    Γ0[eq[:eq_c_s_f], endo[:Eλ_s_f_t]]= -(m[:σ] * m[:λ_s_ss])     / (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])
    Γ0[eq[:eq_c_s_f], endo[:Eλ_h_f_t]]= -((1-m[:σ]) * m[:λ_h_ss]) / (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])
    Γ0[eq[:eq_c_s_f], endo[:Ey_f_t]]  = -(m[:σ′]) * m[:σ] * (m[:λ_s_ss] - m[:λ_h_ss]) /
                                        (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])

    ### transfers

    # Sticky prices
    Γ0[eq[:eq_t_h], endo[:t_h_t]]  = m[:θ]
    Γ0[eq[:eq_t_h], endo[:y_t]]    = -m[:τ_d]
    Γ0[eq[:eq_t_h], endo[:mc_t]]   = m[:τ_d]
    Γ0[eq[:eq_t_h], endo[:L_t]]    = m[:τ_d] * (1 - m[:α])
    Γ0[eq[:eq_t_h], endo[:k_t]]    = m[:τ_d] * m[:α] - m[:τ_k] * (m[:Rk_ss] * m[:k_ss] / m[:y_ss])
    Γ0[eq[:eq_t_h], endo[:Rk_t]]   = -m[:τ_k] * (m[:Rk_ss] * m[:k_ss] / m[:y_ss])

    # flexible prices
    Γ0[eq[:eq_t_h_f], endo[:t_h_f_t]]  = m[:θ]
    Γ0[eq[:eq_t_h_f], endo[:y_f_t]]    = -m[:τ_d]
    Γ0[eq[:eq_t_h_f], endo[:mc_f_t]]   = m[:τ_d]
    Γ0[eq[:eq_t_h_f], endo[:L_f_t]]    = m[:τ_d] * (1 - m[:α])
    Γ0[eq[:eq_t_h_f], endo[:k_f_t]]    = m[:τ_d] * m[:α] - m[:τ_k] * (m[:Rk_ss] * m[:k_ss] / m[:y_ss])
    Γ0[eq[:eq_t_h_f], endo[:Rk_f_t]]   = -m[:τ_k] * (m[:Rk_ss] * m[:k_ss] / m[:y_ss])

    ### 7. Capital utilization FOC

    # Sticky prices
    Γ0[eq[:eq_Rk], endo[:Rk_t]] = 1.
    Γ0[eq[:eq_Rk], endo[:u_t]]  = -m[:χ]

    # Flexible prices
    Γ0[eq[:eq_Rk_f], endo[:Rk_f_t]] = 1.
    Γ0[eq[:eq_Rk_f], endo[:u_f_t]] = -m[:χ]

    ### 8. Capital FOC

    #Sticky prices
    Γ0[eq[:eq_ϕ], endo[:ϕ_t]] = 1.
    Γ0[eq[:eq_ϕ], endo[:Eϕ_t]] = -m[:β] * exp(-m[:γ]) * (1 - m[:δ])
    Γ0[eq[:eq_ϕ], endo[:z_t]] = m[:ρ_z]
    Γ0[eq[:eq_ϕ], endo[:Eλ_s_t]] = -(1 - m[:β] * exp(-m[:γ]) * (1 - m[:δ]))
    Γ0[eq[:eq_ϕ], endo[:ERk_t]] = -(1 - m[:β] * exp(-m[:γ]) * (1 - m[:δ]))

    #Flexible prices
    Γ0[eq[:eq_ϕ_f], endo[:ϕ_f_t]] = 1
    Γ0[eq[:eq_ϕ_f], endo[:Eϕ_f_t]] = -m[:β] * exp(-m[:γ]) * (1 - m[:δ])
    Γ0[eq[:eq_ϕ_f], endo[:z_t]] = m[:ρ_z]
    Γ0[eq[:eq_ϕ_f], endo[:Eλ_s_f_t]] = -(1 - m[:β] * exp(-m[:γ]) * (1 - m[:δ]))
    Γ0[eq[:eq_ϕ_f], endo[:ERk_f_t]] = -(1 - m[:β] * exp(-m[:γ]) *(1 - m[:δ]))

    ### 9. Investment FOC

    #Sticky prices
    Γ0[eq[:eq_i], endo[:λ_s_t]] = 1 /( m[:S′′] * expγ^2)
    Γ0[eq[:eq_i], endo[:ϕ_t]] = -1 / (m[:S′′] * expγ^2)
    Γ0[eq[:eq_i], endo[:μ_t]] = -1 / (m[:S′′] * expγ^2)
    Γ0[eq[:eq_i], endo[:i_t]] = (1 + m[:β])
    Γ0[eq[:eq_i], endo[:z_t]] = (1 - m[:β] * m[:ρ_z])
    Γ0[eq[:eq_i], endo[:Ei_t]] = -m[:β]
    Γ1[eq[:eq_i], endo[:i_t]] = 1

    #Flexible prices
    Γ0[eq[:eq_i_f], endo[:λ_s_f_t]] = 1 / (m[:S′′] * expγ^2)
    Γ0[eq[:eq_i_f], endo[:ϕ_f_t]] = -1 / (m[:S′′] * expγ^2)
    Γ0[eq[:eq_i_f], endo[:μ_t]] = -1 / (m[:S′′] * expγ^2)
    Γ0[eq[:eq_i_f], endo[:i_f_t]] = (1 + m[:β])
    Γ0[eq[:eq_i_f], endo[:z_t]] = (1 - m[:β]*m[:ρ_z])
    Γ0[eq[:eq_i_f], endo[:Ei_f_t]] = -m[:β]
    Γ1[eq[:eq_i_f], endo[:i_f_t]] = 1

    ### Capital input

    #Sticky prices
    Γ0[eq[:eq_k], endo[:k_t]] = 1
    Γ0[eq[:eq_k], endo[:u_t]] = -1
    Γ0[eq[:eq_k], endo[:z_t]] = 1
    Γ1[eq[:eq_k], endo[:kbar_t]] =1

    #Flexible prices
    Γ0[eq[:eq_k_f], endo[:k_f_t]] = 1
    Γ0[eq[:eq_k_f], endo[:u_f_t]] = -1
    Γ0[eq[:eq_k_f], endo[:z_t]] = 1
    Γ1[eq[:eq_k_f], endo[:kbar_f_t]] = 1

    ### Capital Accumulation

    #Sticky prices
    Γ0[eq[:eq_kbar], endo[:kbar_t]] = 1
    Γ0[eq[:eq_kbar], endo[:μ_t]] = -(1 - (1 - m[:δ]) * exp(-m[:γ]))
    Γ0[eq[:eq_kbar], endo[:i_t]] = -(1 - (1 - m[:δ]) * exp(-m[:γ]))
    Γ1[eq[:eq_kbar], endo[:kbar_t]] = (1 - m[:δ]) * exp(-m[:γ])
    Γ0[eq[:eq_kbar], endo[:z_t]] = (1 - m[:δ]) * exp(-m[:γ])

    #Flexible Prices
    Γ0[eq[:eq_kbar_f], endo[:kbar_f_t]] = 1
    Γ0[eq[:eq_kbar_f], endo[:μ_t]] = -(1 - (1 - m[:δ]) * exp(-m[:γ]))
    Γ0[eq[:eq_kbar_f], endo[:i_f_t]] = -(1 - (1 - m[:δ]) * exp(-m[:γ]))
    Γ1[eq[:eq_kbar_f], endo[:kbar_f_t]] = (1 - m[:δ]) * exp(-m[:γ])
    Γ0[eq[:eq_kbar_f], endo[:z_t]] = (1 - m[:δ]) * exp(-m[:γ])

    ### wage Phillips curve

    #Sticky prices

    κw = ((1. - m[:ξ_w] * m[:β]) * (1. - m[:ξ_w])) / (m[:ξ_w] * (1. + m[:β]) * (1. + m[:ν] * (1. + 1. / m[:λ_w_ss])))

    Γ0[eq[:eq_w], endo[:w_t]] = 1
    Γ0[eq[:eq_w], endo[:Ew_t]] = -m[:β] / (1 + m[:β])
    Γ0[eq[:eq_w], endo[:wgap_t]] = κw
    Γ0[eq[:eq_w], endo[:π_t]] = (1 + m[:β] * m[:ι_w]) / (1 + m[:β])
    Γ0[eq[:eq_w], endo[:Eπ_t]] = -m[:β] / (1 + m[:β])
    Γ0[eq[:eq_w], endo[:z_t]] = (1 + m[:β] * m[:ι_w] - m[:β] * m[:ρ_z]) / (1 + m[:β])
    Γ0[eq[:eq_w], endo[:λ_w_t]] = -1
    Γ1[eq[:eq_w], endo[:w_t]] = 1 / (1 + m[:β])
    Γ1[eq[:eq_w], endo[:π_t]] = m[:ι_w] / (1 + m[:β])
    Γ1[eq[:eq_w], endo[:z_t]] = m[:ι_w] / (1 + m[:β])
    #Flexible prices
    Γ0[eq[:eq_w_f], endo[:wgap_f_t]] = 1

    ### Wage gap

    #Flexible prices
    Γ0[eq[:eq_wgap], endo[:wgap_t]] = 1
    Γ0[eq[:eq_wgap], endo[:w_t]] = -1
    Γ0[eq[:eq_wgap], endo[:b_t]] = 1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_wgap], endo[:L_t]] = m[:ν]
    Γ0[eq[:eq_wgap], endo[:λ_t]] = -1

    #Sticky prices
    Γ0[eq[:eq_wgap_f], endo[:wgap_f_t]] = 1
    Γ0[eq[:eq_wgap_f], endo[:w_f_t]] = -1
    Γ0[eq[:eq_wgap_f], endo[:b_t]] = 1 / ((1 - m[:ρ_b]) * (m[:σ] * m[:λ_s_ss] + (1 - m[:σ]) * m[:λ_h_ss])*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]) * (expγ * m[:c_h_ss] -
                                         m[:h] * m[:c_ss]) / (m[:h] * m[:c_ss] * (m[:σ] * m[:λ_s_ss] *
                                         (expγ * m[:c_h_ss] - m[:h] * m[:c_ss]) + (1 - m[:σ]) * m[:λ_h_ss]*
                                         (expγ * m[:c_s_ss] - m[:h] * m[:c_ss]))))
    Γ0[eq[:eq_wgap_f], endo[:L_f_t]] = m[:ν]
    Γ0[eq[:eq_wgap_f], endo[:λ_f_t]] = -1

    ### Monetary policy rule

    #Flexible prices
    Γ0[eq[:eq_R], endo[:R_t]] = 1
    Γ1[eq[:eq_R], endo[:R_t]] = m[:ρ_R]
    Γ0[eq[:eq_R], endo[:π_t]] = -(1 - m[:ρ_R]) *m[:ψ1]
    Γ0[eq[:eq_R], endo[:gdp_t]] = -(1 - m[:ρ_R]) * m[:ψ2] - m[:ψ3]
    Γ0[eq[:eq_R], endo[:gdp_f_t]] = (1 - m[:ρ_R]) * m[:ψ2] + m[:ψ3]
    Γ1[eq[:eq_R], endo[:gdp_t]] = -m[:ψ3]
    Γ1[eq[:eq_R], endo[:gdp_f_t]] = m[:ψ3]
    Γ0[eq[:eq_R], endo[:mp_t]] = -1

    ### definition of gdp

    #Sticky prices
    Γ0[eq[:eq_gdp], endo[:gdp_t]] = -1
    Γ0[eq[:eq_gdp], endo[:y_t]] = 1
    Γ0[eq[:eq_gdp], endo[:u_t]] = -m[:k_ss] * m[:Rk_ss] / m[:y_ss]

    #Flexible prices
    Γ0[eq[:eq_gdp_f], endo[:gdp_f_t]] = -1
    Γ0[eq[:eq_gdp_f], endo[:y_f_t]] = 1
    Γ0[eq[:eq_gdp_f], endo[:u_f_t]] = -m[:k_ss] * m[:Rk_ss] / m[:y_ss]

    ### market clearing

    #Sticky prices
    Γ0[eq[:eq_u], endo[:c_t]] = m[:c_ss] / m[:y_ss]
    Γ0[eq[:eq_u], endo[:i_t]] = (1 - m[:θ]) * m[:i_s_ss] / m[:y_ss]
    #Γ0[eq[:eq_u], endo[:y_t]] = -1 / m[:g_ss] matlab code has these two lines commented out
    #Γ0[eq[:eq_u], endo[:g_t]] = 1 / m[:g_ss]
    Γ0[eq[:eq_u], endo[:u_t]] = m[:k_ss] * m[:Rk_ss] / m[:y_ss]
    Γ0[eq[:eq_u], endo[:y_t]] = -1
    Γ0[eq[:eq_u], endo[:g_t]] = m[:gy_ss]

    #Flexible Prices
    Γ0[eq[:eq_u_f], endo[:c_f_t]] = m[:c_ss] / m[:y_ss]
    Γ0[eq[:eq_u_f], endo[:i_f_t]] = (1 - m[:θ]) * m[:i_s_ss] / m[:y_ss]
    Γ0[eq[:eq_u_f], endo[:y_f_t]] = -1
    Γ0[eq[:eq_u_f], endo[:g_t]] = m[:gy_ss]
    Γ0[eq[:eq_u_f], endo[:u_f_t]] = m[:k_ss] * m[:Rk_ss] / m[:y_ss]

    ### Exogenous shocks

    # z
    Γ0[eq[:eq_z], endo[:z_t]] = 1
    Γ1[eq[:eq_z], endo[:z_t]] = m[:ρ_z]
    Ψ[eq[:eq_z], exo[:z_sh]] = 1

    # g
    Γ0[eq[:eq_g], endo[:g_t]] = 1
    Γ1[eq[:eq_g], endo[:g_t]] = m[:ρ_g]
    Ψ[eq[:eq_g], exo[:g_sh]]  = 1

    # μ
    Γ0[eq[:eq_μ], endo[:μ_t]] = 1
    Γ1[eq[:eq_μ], endo[:μ_t]] = m[:ρ_μ]
    Ψ[eq[:eq_μ], exo[:μ_sh]]  = 1

    # b
    Γ0[eq[:eq_b], endo[:b_t]] = 1
    Γ1[eq[:eq_b], endo[:b_t]] = m[:ρ_b]
    Ψ[eq[:eq_b], exo[:b_sh]]  = 1

    # mp
    Γ0[eq[:eq_mp], endo[:mp_t]] = 1
    Γ1[eq[:eq_mp], endo[:mp_t]] = m[:ρ_mp]
    Ψ[eq[:eq_mp], exo[:R_sh]]  = 1

    # λ_w
    Γ0[eq[:eq_λ_w], endo[:λ_w_t]]   = 1
    Γ1[eq[:eq_λ_w], endo[:λ_w_t]]   = m[:ρ_λ_w]
    Γ0[eq[:eq_λ_w], endo[:λ_w_t1]]  = -1
    Γ1[eq[:eq_λ_w], endo[:λ_w_t1]]  = -m[:η_λ_w]
    Γ0[eq[:eq_λ_w1], endo[:λ_w_t1]] = 1
    Ψ[eq[:eq_λ_w1], exo[:λ_w_sh]]   = 1

    # λ_p
    Γ0[eq[:eq_λ_p], endo[:λ_p_t]]   = 1
    Γ1[eq[:eq_λ_p], endo[:λ_p_t]]   = m[:ρ_λ_p]
    Γ0[eq[:eq_λ_p], endo[:λ_p_t1]]  = -1
    Γ1[eq[:eq_λ_p], endo[:λ_p_t1]]  = -m[:η_λ_p]
    Γ0[eq[:eq_λ_p1], endo[:λ_p_t1]] = 1
    Ψ[eq[:eq_λ_p1], exo[:λ_p_sh]]   = 1

    ### Expectational terms

    # π sticky
    Γ0[eq[:eq_Eπ], endo[:π_t]] = 1
    Γ1[eq[:eq_Eπ], endo[:Eπ_t]] = 1
    Π[eq[:eq_Eπ], ex[:Eπ_sh]] = 1

    # c sticky
    Γ0[eq[:eq_Ec], endo[:c_t]] = 1
    Γ1[eq[:eq_Ec], endo[:Ec_t]] = 1
    Π[eq[:eq_Ec], ex[:Ec_sh]] = 1

    # c flexible
    Γ0[eq[:eq_Ec_f], endo[:c_f_t]] = 1
    Γ1[eq[:eq_Ec_f], endo[:Ec_f_t]] = 1
    Π[eq[:eq_Ec_f], ex[:Ec_f_sh]] = 1

    # λ sticky
    Γ0[eq[:eq_Eλ], endo[:λ_t]] = 1
    Γ1[eq[:eq_Eλ], endo[:Eλ_t]] = 1
    Π[eq[:eq_Eλ], ex[:Eλ_sh]] = 1

    # λ flexible
    Γ0[eq[:eq_Eλ_f], endo[:λ_f_t]] = 1
    Γ1[eq[:eq_Eλ_f], endo[:Eλ_f_t]] = 1
    Π[eq[:eq_Eλ_f], ex[:Eλ_f_sh]] = 1

    # ϕ sticky
    Γ0[eq[:eq_Eϕ], endo[:ϕ_t]] = 1
    Γ1[eq[:eq_Eϕ], endo[:Eϕ_t]] = 1
    Π[eq[:eq_Eϕ], ex[:Eϕ_sh]] = 1

    # ϕ flexible
    Γ0[eq[:eq_Eϕ_f], endo[:ϕ_f_t]] = 1
    Γ1[eq[:eq_Eϕ_f], endo[:Eϕ_f_t]] = 1
    Π[eq[:eq_Eϕ_f], ex[:Eϕ_f_sh]] = 1

    # Rk sticky
    Γ0[eq[:eq_ERk], endo[:Rk_t]] = 1
    Γ1[eq[:eq_ERk], endo[:ERk_t]] = 1
    Π[eq[:eq_ERk], ex[:ERk_sh]] = 1

    # Rk flexible
    Γ0[eq[:eq_ERk_f], endo[:Rk_f_t]] = 1
    Γ1[eq[:eq_ERk_f], endo[:ERk_f_t]] = 1
    Π[eq[:eq_ERk_f], ex[:ERk_f_sh]] = 1

    # i sticky
    Γ0[eq[:eq_Ei], endo[:i_t]] = 1
    Γ1[eq[:eq_Ei], endo[:Ei_t]] = 1
    Π[eq[:eq_Ei], ex[:Ei_sh]] = 1

    #  i flexible
    Γ0[eq[:eq_Ei_f], endo[:i_f_t]]  = 1
    Γ1[eq[:eq_Ei_f], endo[:Ei_f_t]] = 1
    Π[eq[:eq_Ei_f], ex[:Ei_f_sh]]   = 1

    # w sticky
    Γ0[eq[:eq_Ew], endo[:w_t]]  = 1
    Γ1[eq[:eq_Ew], endo[:Ew_t]] = 1
    Π[eq[:eq_Ew], ex[:Ew_sh]]   = 1

    # y sticky
    Γ0[eq[:eq_Ey], endo[:y_t]]  = 1
    Γ1[eq[:eq_Ey], endo[:Ey_t]] = 1
    Π[eq[:eq_Ey], ex[:Ey_sh]]    = 1

    # y flexible
    Γ0[eq[:eq_Ey_f], endo[:y_f_t]]  = 1
    Γ1[eq[:eq_Ey_f], endo[:Ey_f_t]] = 1
    Π[eq[:eq_Ey_f], ex[:Ey_f_sh]]   = 1

    # λ_s sticky
    Γ0[eq[:eq_Eλ_s], endo[:λ_s_t]]  = 1
    Γ1[eq[:eq_Eλ_s], endo[:Eλ_s_t]] = 1
    Π[eq[:eq_Eλ_s], ex[:Eλ_s_sh]]   = 1

    # λ_s flexible
    Γ0[eq[:eq_Eλ_s_f], endo[:λ_s_f_t]]  = 1
    Γ1[eq[:eq_Eλ_s_f], endo[:Eλ_s_f_t]] = 1
    Π[eq[:eq_Eλ_s_f], ex[:Eλ_s_f_sh]]   = 1

    # λ_h sticky
    Γ0[eq[:eq_Eλ_h], endo[:λ_h_t]]  = 1
    Γ1[eq[:eq_Eλ_h], endo[:Eλ_h_t]] = 1
    Π[eq[:eq_Eλ_h], ex[:Eλ_h_sh]]   = 1

    # λ_h flexible
    Γ0[eq[:eq_Eλ_h_f], endo[:λ_h_f_t]]  = 1
    Γ1[eq[:eq_Eλ_h_f], endo[:Eλ_h_f_t]] = 1
    Π[eq[:eq_Eλ_h_f], ex[:Eλ_h_f_sh]]   = 1

    ### Lagged variables

    # gdp
    Γ0[eq[:eq_gdp1], endo[:gdp_t1]] = 1
    Γ1[eq[:eq_gdp1], endo[:gdp_t]] = 1

    # c
    Γ0[eq[:eq_c1], endo[:c_t1]] = 1
    Γ1[eq[:eq_c1], endo[:c_t]] = 1

    # i
    Γ0[eq[:eq_i1], endo[:i_t1]] = 1
    Γ1[eq[:eq_i1], endo[:i_t]] = 1

    # w
    Γ0[eq[:eq_w1], endo[:w_t1]] = 1
    Γ1[eq[:eq_w1], endo[:w_t]] = 1

    # matlab comment: "TODO: throw in new lags?"
    # figure out if that still needs to be done
    return Γ0, Γ1, C, Ψ, Π

end


text:
using Distributions
using Unitful
using Unitful: mV
using Test

MilliVoltage = Quantity{Float64, dimension(mV), typeof(mV)}

# Two testsets manually written out, to show API usage & behaviour.
# One with implicit units, and one with explicit units.

@testset "Exponential" begin
    d = Exponential(8.0 * mV)
    @test unit(d) == mV
    @test eltype(d) == MilliVoltage
    @test partype(d) == MilliVoltage
    @test extrema(d) == (0mV, Inf*mV)
    @test islowerbounded(d)
    @test !isupperbounded(d)
    @test insupport(d, 3mV)
    @test rate(d) == 1/8mV
    @test mean(d) == 8mV
    @test median(d) ≈ 5.54517744 * mV
    @test mode(d) == 0mV
    @test var(d) == 64mV^2
    @test std(d) == 8mV
    @test kurtosis(d) == 6mV
    @test entropy(d) ≈ 3.07944154
    @test kldivergence(d, Exponential(7mV)) ≈ 0.0085313926
    @test_throws MethodError kldivergence(d, Exponential(7mV^2))
    @test pdf(d, 1mV)  ≈ 0.110312112 / mV
    @test cdf(d, 1mV)  ≈ 0.117503097
    @test ccdf(d, 1mV) ≈ 0.8824969
    @test logpdf(d, 1mV) ≈ -2.20444154
    @test logcdf(d, 1mV) ≈ -2.14129058
    @test logccdf(d, 1mV) ≈ -0.125
    @test quantile(d, 0.6) ≈ 7.33032585 * mV
    @test cquantile(d, 0.4) ≈ 7.33032585 * mV
    @test invlogcdf(d, -1) ≈ 3.66940116 * mV
    @test invlogccdf(d, -1) ≈ 8 * mV
    @test gradlogpdf(d, 1mV) == -0.125/mV
    @test gradlogpdf(d, -1mV) == 0/mV
    @test mgf(d, 1/mV) ≈ -0.142857143
    @test cf(d, 1/mV) ≈ 0.015384615 + 0.123076923im
    @test typeof(rand(d)) == MilliVoltage
    @test typeof(rand(d, (3,3))) == Matrix{MilliVoltage}
    @test fit(Exponential, [3,2,1]mV) == Exponential(2mV)
    @test fit(Exponential, [3,2,1]mV, float([1,1,2])) == Exponential(1.75mV)
end # @testset "Exponential"

@testset "LogNormal" begin
    d = LogNormal(3.0, 2.0, mV)
    @test unit(d) == mV
    @test eltype(d) == MilliVoltage
    @test partype(d) == Float64
    @test extrema(d) == (0mV, Inf*mV)
    @test islowerbounded(d)
    @test !isupperbounded(d)
    @test insupport(d, 3mV)
    @test meanlogx(d) == 3
    @test varlogx(d) == 4
    @test stdlogx(d) == 2
    @test mean(d) ≈ 148.413159 * mV
    @test median(d) ≈ 20.0855369 * mV
    @test mode(d) ≈ 0.367879441 * mV
    @test var(d) ≈ 1.18057782e6 * mV * mV
    @test std(d) ≈ 1086.54398 * mV
    @test kurtosis(d) ≈ 9.22055698e6 * mV
    @test entropy(d) ≈ 5.11208571
    @test pdf(d, 1mV)  ≈ 0.0647587978 / mV
    @test pdf(d, -1mV) == 0 / mV
    @test cdf(d, 1mV)  ≈ 0.0668072013
    @test ccdf(d, 1mV) ≈ 0.933192799
    @test logpdf(d, 1mV) ≈ -2.73708571
    @test logcdf(d, 1mV) ≈ -2.7059444
    @test logccdf(d, 1mV) ≈ -0.0691434556
    @test quantile(d, 0.6) ≈ 33.3378773 * mV
    @test cquantile(d, 0.4) ≈ 33.3378773 * mV
    @test invlogcdf(d, -1) ≈ 10.2271922 * mV
    @test invlogccdf(d, -1) ≈ 39.4466816 * mV
    @test gradlogpdf(d, 1mV) == -0.25/mV
    @test gradlogpdf(d, -1mV) == 0/mV
    @test typeof(rand(d)) == MilliVoltage
    @test typeof(rand(d, (3,3))) == Matrix{MilliVoltage}
    @test fit(LogNormal, Float32.([3,2,1])mV) == LogNormal(0.59725314f0, 0.55554837f0, mV)
end # @testset "LogNormal"


text:
@testset "Kernel Dependencies" begin
    function kernel(sig, waitval, A, val)
        i = workitemIdx().x
        AMDGPU.device_signal_wait(sig, waitval)
        A[i] = val
        return nothing
    end

    @testset "Barrier AND" begin
        for i in (0, 1, 5, 7)
            @testset "$i inputs" begin
                RA = ROCArray(zeros(Float64, 1))
                sig = AMDGPU.HSASignal(0)

                ret1 = [@roc(kernel(sig, 3, RA, 1.0)) for _ in 1:i]

                retb = barrier_and!(ret1)
                ret2 = @roc kernel(sig, 0, RA, 2.0)

                if i > 0
                    sleep(0.5)
                    @test Array(RA)[1] == 0.0
                    HSA.signal_store_release(sig.signal[], 3)
                    wait.(ret1)
                    @test Array(RA)[1] == 1.0
                end
                HSA.signal_store_release(sig.signal[], 0)
                # FIXME: wait(retb)
                wait(ret2)
                @test Array(RA)[1] == 2.0
            end
        end
    end

    #= FIXME
    @testset "Barrier OR" begin
        for i in (0, 1, 5, 7)
            @testset "$i inputs" begin
                RA = ROCArray(zeros(Float64, 1))
                sig = AMDGPU.HSASignal(0)

                ret1 = [@roc(kernel(sig, 7, RA, 5.0)) for _ in 1:i]
                pushfirst!(ret1, @roc(kernel(sig, 3, RA, 1.0)))

                retb = barrier_or!(ret1)
                ret2 = @roc kernel(sig, 0, RA, 2.0)

                if i > 0
                    sleep(0.5)
                    @test Array(RA)[1] == 0.0
                    HSA.signal_store_release(sig.signal[], 3)

                    wait(ret1[1])
                    @test Array(RA)[1] == 1.0
                end
                HSA.signal_store_release(sig.signal[], 0)
                sleep(0.5)
                @test Array(RA)[1] == 2.0
                wait(ret2)
                # FIXME: wait(retb)
                # clear waiting kernels
                HSA.signal_store_release(sig.signal[], 7)
            end
        end
    end
    =#
end


text:
module _Rootfinding

using ResumableFunctions
using ArgCheck: @argcheck

@resumable function bracket_bisect(xa, xb, ya=missing, yb=missing; growth_factor=2)
    @argcheck xa != xb
    @argcheck ismissing(ya) || !iszero(ya)
    @argcheck ismissing(yb) || !iszero(yb)
    @argcheck growth_factor >= 1

    if ismissing(ya)
        ya = @yield xa
    end
    if ismissing(yb)
        yb = @yield xb
    end

    while sign(ya) == sign(yb)
        x = xb + growth_factor*(xb - xa)
        y = @yield x
        xa, xb = xb, x
        ya, yb = yb, y
    end

    while true
        x = (xa + xb)/2
        y = @yield x
        if sign(y) == sign(ya)
            xa = x
            ya = y
        else
            xb = x
            yb = y
        end
    end
end

export bracket_bisect

end


text:
# test_le.sac was generated in SAC 101.6a with "fg seismogram; write test_le.sac"
sac_file    = path*"/SampleFiles/SAC/test_le.sac"
sac_be_file = path*"/SampleFiles/SAC/test_be.sac"
sac_pz_file = path*"/SampleFiles/SAC/test_sac.pz"
sac_pz_wc   = path*"/SampleFiles/SAC/test_sac.*"
uw_file     = path*"/SampleFiles/UW/00012502123W"
sac_pat     = path*"/SampleFiles/SAC/*.sac"
sac_pz_out1 = path*"/local_sac_1.pz"
sac_pz_out2 = path*"/local_sac_2.pz"
sac_pz_out3 = path*"/local_sac_3.pz"
sac_pz_out4 = path*"/local_sac_4.pz"
sac_pz_out5 = path*"/local_sac_5.pz"
f_stub      = "1981.088.10.38.23.460"
f_out       = f_stub * "..CDV...R.SAC"
f_out_new   = f_stub * ".VU.CDV..NUL.R.SAC"
sacv7_out   = "v7_out.sac"

printstyled("  SAC\n", color=:light_green)
printstyled("    read\n", color=:light_green)
@test_throws ErrorException verified_read_data("sac", uw_file)

SAC1 = verified_read_data("sac", sac_file)[1]
@test ≈(SAC1.fs, 100.0)
@test ≈(length(SAC1.x), 1000)

# SAC with mmap
printstyled("      mmap\n", color=:light_green)
SACm = read_data("sac", sac_file, memmap=true)[1]
@test SAC1 == SACm

SAC2 = verified_read_data("sac", sac_file, full=true)[1]
@test ≈(1/SAC1.fs, SAC2.misc["delta"])
@test ≈(length(SAC1.x), SAC2.misc["npts"])

printstyled("      wildcard\n", color=:light_green)
SAC = verified_read_data("sac", sac_pat, full=true)

printstyled("      bigendian\n", color=:light_green)
SAC3 = verified_read_data("sac", sac_be_file, full=true)[1]
@test ≈(1/SAC3.fs, SAC3.misc["delta"])
@test ≈(length(SAC3.x), SAC3.misc["npts"])

redirect_stdout(out) do
  sachdr(sac_be_file)
end

printstyled("    write\n", color=:light_green)
S = SeisData(SAC2)
writesac(S) # change 2019-07-15 to cover writesac on GphysData
@test safe_isfile(f_out)
@test any([occursin(string("wrote to file ", f_out), S.notes[1][i]) for i in 1:length(S.notes[1])])

printstyled("      reproducibility\n", color=:light_green)
SAC4 = verified_read_data("sac", f_out, full=true)[1]
for f in SeisIO.datafields
  (f in [:src, :notes, :misc]) && continue
  @test isequal(getfield(SAC2, f), getfield(SAC4, f))
end
fn = f_out[3:end]
writesac(SAC4, fname=fn)
@test safe_isfile(fn)
SAC5 = verified_read_data("sac", fn, full=true)[1]
for f in SeisIO.datafields
  (f in [:src, :notes, :misc]) && continue
  @test isequal(getfield(SAC2, f), getfield(SAC4, f))
  @test isequal(getfield(SAC4, f), getfield(SAC5, f))
end
safe_rm(fn)

printstyled("      logging\n", color=:light_green)
@test any([occursin("write", n) for n in S.notes[1]])
redirect_stdout(out) do
  show_writes(S, 1)
  show_writes(S[1])
  show_writes(S)
end

SAC1.id = "VU.CDV..NUL"
SAC1.name = "VU.CDV..NUL"
writesac(SAC1)
@test safe_isfile(f_out_new)
writesac(SAC1, fname="POTATO.SAC")
@test safe_isfile("POTATO.SAC")

# testing custom naming formats
writesac(SAC1, fname="test_write_1.sac")
@test safe_isfile("test_write_1.sac")
safe_rm("test_write_1.sac")

# testing that appending ".sac" to a file string works
writesac(SAC1, fname="test_write_2", v=1)
@test safe_isfile("test_write_2.sac")
safe_rm("test_write_2.sac")

redirect_stdout(out) do
  writesac(SAC1, v=1)
end
@test safe_isfile(f_out_new)
safe_rm(f_out_new)

printstyled("      skip if fs==0.0\n", color=:light_green)
SAC1.id = "VU.FS0..NUL"
SAC1.fs = 0.0
writesac(SAC1)
@test safe_isfile(f_stub*"VU.FS0..NUL.R.SAC") == false

# SACPZ
printstyled("    SACPZ\n", color=:light_green)
printstyled("      read\n", color=:light_green)
S = read_meta("sacpz", sac_pz_wc)
S = read_meta("sacpz", sac_pz_file)
writesacpz(sac_pz_out1, S)
T = read_meta("sacpz", sac_pz_out1)
for f in (:n, :id, :name, :loc, :fs, :gain, :units)
  @test isequal(getfield(S, f), getfield(T, f))
end
for f in fieldnames(typeof(S.resp[1]))
  for i = 1:S.n
    @test isapprox(getfield(S.resp[i], f), getfield(T.resp[i], f))
  end
end
U = deepcopy(S)
for i = 1:U.n
  U.resp[i] = MultiStageResp(3)
  U.resp[i].stage[1] = S.resp[i]
end

printstyled("      write\n", color=:light_green)
writesacpz(sac_pz_out2, U)
T = read_meta("sacpz", sac_pz_out2)
for f in (:n, :id, :name, :loc, :fs, :gain, :units)
  @test isequal(getfield(S, f), getfield(T, f))
end
for f in fieldnames(typeof(S.resp[1]))
  for i = 1:S.n
    @test isapprox(getfield(S.resp[i], f), getfield(T.resp[i], f))
  end
end
U[1] = SeisChannel(id = "UW.HOOD..ENE")
writesacpz(sac_pz_out3, U)
read_meta!(S, "sacpz", sac_pz_out3)
for f in (:n, :id, :name, :loc, :fs, :gain, :units)
  @test isequal(getfield(S, f), getfield(T, f))
end
S = breaking_seis()[1:3]
S.resp[1].resp = rand(ComplexF64, 12, 2)
S.resp[3].stage[2] = nothing
writesacpz(sac_pz_out4, S)

printstyled("        GphysChannel\n", color=:light_green)
writesacpz(sac_pz_out5, S[1])

safe_rm(sac_pz_out1)
safe_rm(sac_pz_out2)
safe_rm(sac_pz_out3)
safe_rm(sac_pz_out4)
safe_rm(sac_pz_out5)

printstyled("    SAC file v7 (SAC v102.0)\n", color=:light_green)
test_fs = 62.5
test_lat = 48.7456
test_lon = -122.4126

printstyled("      writesac(..., nvhdr=N)\n", color=:light_green)
C = read_data("sac", sac_file)[1]

writesac(C, fname=sacv7_out, nvhdr=6)
sz6 = stat(sacv7_out).size

writesac(C, fname=sacv7_out, nvhdr=7)
sz7 = stat(sacv7_out).size

# The only difference should be the addition of a length-22 Vector{Float64
@test sz7-sz6 == 176

# In fact, we can open the file, skip to byte sz6, and read in the array
io = open(sacv7_out, "r")
seek(io, sz6)
sac_buf_tmp = read(io)
close(io)
dv = reinterpret(Float64, sac_buf_tmp)

# ...and the variables we write to the header should be identical
@test C.fs == 1.0/dv[1]
@test C.loc.lon == dv[19]
@test C.loc.lat == dv[20]

printstyled("      big endian\n", color=:light_green)
io = open(sac_be_file, "r")
sac_raw = read(io)
close(io)
sac_raw[308] = 0x07

reset_sacbuf()
dv     = BUF.sac_dv
dv[1]  = 1.0/test_fs
dv[19] = test_lon
dv[20] = test_lat
dv    .= bswap.(dv)
sac_dbl_buf = reinterpret(UInt8, dv)

io = open(sacv7_out, "w")
write(io, sac_raw)
write(io, sac_dbl_buf)
close(io)

C = read_data("sac", sacv7_out, full=true)[1]
@test C.fs == test_fs
@test C.loc.lat == test_lat
@test C.loc.lon == test_lon

printstyled("      little endian\n", color=:light_green)
io = open(sac_file, "r")
sac_raw = read(io)
close(io)
sac_raw[305] = 0x07

reset_sacbuf()
dv[1]  = 1.0/test_fs
dv[19] = test_lon
dv[20] = test_lat
sac_dbl_buf = reinterpret(UInt8, dv)

io = open(sacv7_out, "w")
write(io, sac_raw)
write(io, sac_dbl_buf)
close(io)

C = read_data("sac", sacv7_out, full=true)[1]
@test C.fs == test_fs
@test C.loc.lat == test_lat
@test C.loc.lon == test_lon


text:
if length(ARGS) == 2
    pkgdir = ARGS[1]
    makefile = ARGS[2]
else
    makefile = joinpath(pwd(), "make.jl")
    pkgdir = normpath(joinpath(pwd(), ".."))
    @info("No directory specified. Falling back to pkgdir=`$(pkgdir)` and makefile=`$(makefile)`.")
end
docsdir = dirname(makefile)

using Pkg

include(joinpath(@__DIR__, "utils", "rewrite.jl"))
try
    Pkg.instantiate()
catch err
    @error("Error while instantiating docs directory:", err)
end
try
    Pkg.develop(PackageSpec(path=pkgdir))
catch err
    @error("Error while developing parent directory $(pkgdir):", err)
end
Pkg.status()

documenter_version = v"0.24.10"
try
    manifest = joinpath(docsdir, "Manifest.toml")
    if isfile(manifest)
        pm = Pkg.TOML.parsefile(manifest)
        global documenter_version = VersionNumber(first(pm["Documenter"])["version"])
    else
        @warn("No Mainfest.toml found at `$(manifest)`. Defaulting to $(documenter_version).")
    end
catch err
    @error(exception = err)
end

@info("Detected Documenter version $(documenter_version).")

expr, bpath = fix_makefile(makefile, documenter_version)

@info("`cd`ing to `$(docsdir)` and setting `tls[:SOURCE_PATH]` to `$(makefile)`.")
task_local_storage()[:SOURCE_PATH] = makefile
cd(docsdir)

@info("Evaluating the following `make` expr:")
@info(expr)

eval(expr)


text:
function findbingos(boards :: Vector{Vector{Int}}, drawn_nums)
    bingoboards :: Vector{Vector{Int}} = []
    for board in boards
        bingo = false
        for row in [1:5, 6:10, 11:15, 16:20, 21:25]
            if all([number in drawn_nums for number in board[row]])
                bingo = true
            end
        end

        for column in [1:5:21, 2:5:22, 3:5:23, 4:5:24, 5:5:25]
            if all([number in drawn_nums for number in board[column]])
                bingo = true
            end
        end

        if bingo
            push!(bingoboards, board)
        end
    end

    return bingoboards
end

function part1()
    lines = readlines("inputD4.txt")
    draws = [parse(Int, i) for i in split(lines[1], ",")]
    boards :: Vector{Vector{Int}} = []

    linenum = 3
    while linenum < length(lines)
        push!(boards, Int[])
        for inc in 0:4
            append!(last(boards), [parse(Int, i) for i in split(lines[linenum + inc])])
        end

        linenum += 6
    end

    drawn_nums = Int[]
    for draw in draws
        push!(drawn_nums, draw)
        bingos = findbingos(boards, drawn_nums)
        if length(bingos) == 1
            unmarked_sum = sum([n for n in bingos[1] if !(n in drawn_nums)])
            return unmarked_sum * draw
        end
    end
end

function part2()
    lines = readlines("inputD4.txt")
    draws = [parse(Int, i) for i in split(lines[1], ",")]
    boards :: Vector{Vector{Int}} = []

    linenum = 3
    while linenum < length(lines)
        push!(boards, Int[])
        for inc in 0:4
            append!(last(boards), [parse(Int, i) for i in split(lines[linenum + inc])])
        end

        linenum += 6
    end

    drawn_nums = Int[]
    bingosinorder = []
    for draw in draws
        push!(drawn_nums, draw)
        bingos = findbingos(boards, drawn_nums)
        newbingos = [b for b in bingos if !(b in bingosinorder)]
        for newbingo in newbingos
            push!(bingosinorder, newbingo)
        end

        if length(bingos) == length(boards)
            last_board = last(bingosinorder)
            unmarked_sum = sum([n for n in last_board if !(n in drawn_nums)])
            return unmarked_sum * draw
        end
    end
end

println(part1())
println(part2())


text:
using HDF5, JLD

include("CenterOfMass.jl")
include("setup_parameters.jl")

plt[:close]("all")

feasibility_constraint = 0.9

folder = "chance_constrained"

actuator_limit = 6.0
max_robots = 4
csqmi_ratio = 5.0
n_trial = 40
n_measurement = 20

prior = initialize_prior(circle_ps, resolution, interior_q, masses)

belief = deepcopy(prior)

type TrialData
  belief::Histogram
  estimate
  error
  applied_f
  f_hat
  csqmi
  robots_measurement
  robots
  current_feasibility
  true_feasibility
  feasibility_no_limit
end

data_array = Array{TrialData}(n_trial, n_measurement)

rand_com() = r_attach * rand_in_circle()[1:2] # hack
rand_mass() = min_mass + rand() * (max_mass - min_mass)

rand_parameters() = [rand_com(); rand_mass()]

thetas = Array{Float64}(3,0)
n_bad = 0
while size(thetas, 2) < n_trial
  theta = rand_parameters()

  gravity_w = theta[3]*g*point_to_wrench(theta[1:2])
  feasible = check_feasible_configuration(Array{Array{Float64}}(0), attachment_ws,
    gravity_w, actuator_limit, max_robots)

  if feasible
    thetas = hcat(thetas, theta)
  else
    n_bad += 1
    feasible = check_feasible_configuration(Array{Array{Float64}}(0), attachment_ws,
      gravity_w, 1e6, 10)
    if feasible
      #println("infeasible because of actuator limit")
    else
      #println("infeasible despite actuator limit")
    end
  end
end
println("$(n_bad) infeasible thetas encountered during instantiation")

for ii = 1:n_trial
  theta = thetas[:,ii]
  gravity_w = theta[3]*g*point_to_wrench(theta[1:2])
  belief = deepcopy(prior)
  robots = Array{Int64,1}()

  no_attach_count = 1

  for jj = 1:n_measurement
    Convex.clearmemory()

    println("\n$ii: Measurement: $(jj), robots: $(robots)\n")

    remaining_robots = setdiff(collect(1:length(attachment_ps)), robots)
    current_feasibility = feasibility_probability(belief, attachment_ws[robots], attachment_ws[remaining_robots], actuator_limit, max_robots)
    println("Current feasibility: $(current_feasibility)")

    @time @show csqmi_current, robots_current = maximize_csqmi_available_robots(
      robots, belief, circle_ws, attachment_ws, sigma, interior_q, actuator_limit)

    csqmi = csqmi_current
    robots_measurement = robots_current

    if length(robots) < max_robots
      inner_feasibility_constraint = feasibility_constraint^no_attach_count * current_feasibility

      if length(robots) == 0
        inner_feasibility_constraint = 0.0
      end

      @time @show csqmi_additional, robots_additional, additional_robot =
        maximize_csqmi_additional_robot_feasibility(robots, belief,
          circle_ws, attachment_ws, sigma, interior_q, actuator_limit,
          inner_feasibility_constraint)

      if length(robots) == 0 || csqmi_additional / csqmi_current > csqmi_ratio
        println("\nAdding robot $(additional_robot), csqmi: $(csqmi_additional) vs: $(csqmi_current)\n")

        csqmi = csqmi_additional
        robots_measurement = robots_additional
        push!(robots, additional_robot)

        no_attach_count = 1
      else
        no_attach_count += 1
      end
    end

    # check feasibility of current figuration against true theta
    all_robots = collect(1:length(attachment_ws))
    remaining_robots = setdiff(all_robots, robots)

    true_feasibility = check_feasible_configuration(attachment_ws[robots],
      attachment_ws[remaining_robots], gravity_w, actuator_limit, max_robots)
    feasibility_no_limit = check_feasible_configuration(attachment_ws[robots],
      attachment_ws[remaining_robots], gravity_w, 1e6, max_robots)
    if true_feasibility
      println("The current configuration is still feasible")
    else
      if feasibility_no_limit
        println("The current configuration NO LONGER feasible but is feasible if
          limits are removed")
      else
        println("The current configuration NO LONGER feasible")
      end
    end

    applied_w = mean(attachment_ws[robots_measurement])
    total_actuator_limit = length(robots_measurement) * actuator_limit

    applied_f = critical_force_iterative(applied_w, gravity_w,
      circle_ws, total_actuator_limit)

    f_hat = applied_f + sigma * randn()

    println("\nRobots measurement: $(robots_measurement)")
    println("Measurement: $(applied_f)/$(f_hat), total limit: $(total_actuator_limit)\n")

    critical_forces = get_critical_values(circle_ws, applied_w, belief,
      interior_q, total_actuator_limit)

    update_prior!(belief, critical_forces, f_hat, sigma)

    estimate = weighted_average(belief)

    error = normalized_error(estimate - theta)
    println("Estimate: $(estimate), actual: $(theta), normalized error: $(error)")

    data = TrialData(deepcopy(belief), deepcopy(estimate), error, applied_f, f_hat, csqmi,
      deepcopy(robots_measurement), deepcopy(robots), current_feasibility,
      true_feasibility, feasibility_no_limit)
    data_array[ii,jj] = data
  end
end

if ~isdir(folder)
  mkdir(folder)
end
@save "$(folder)/data" data_array csqmi_ratio thetas actuator_limit feasibility_constraint
Void

max_stars_repo_name: mcorah/CenterOfMass

max_issues_repo_name: mcorah/CenterOfMass

max_forks_repo_name: mcorah/CenterOfMass


text:
module SimulateGalaxy

# package code goes here
import GSL.hypergeom
import QuadGK.quadgk
import Optim, StatsBase

export NFWParameters
export SFWParameters
export RandomGalaxy,
    SphericalGalaxy,
    MetallicSphericalGalaxy,
    PlaneGalaxy,
    MetallicPlaneGalaxy,
    EuclideanGalaxy,
    MetallicEuclideanGalaxy,
    PartialEuclideanGalaxy,
    MetallicPartialEuclideanGalaxy,
    sample_plane,
    sample_euclidean,
    sample_partial_euclidean
export hypergeom_2F1_expanded,
    density_profile,
    gravitational_potential,
    gravitational_potential_integral,
    gravitational_potential0,
    gravitational_potential0_robust
export escape_velocity
export profile_density
export simulate_galaxy

include("NFWParameters.jl")
include("SFWParameters.jl")
include("RandomGalaxy.jl")
include("gravitational_potential.jl")
include("escape_velocity.jl")
include("profile_density.jl")
include("simulate_galaxy.jl")

end # module


text:
##########################
# FunctionCell Interface #
##########################
"""
	tree(nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a function cell that when piped data, trains a decision tree classifier using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `prune_purity` fraction in [0.0,1.0] of combined purity for leaves to be merged (-1.0 indicates no pruning)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
tree(nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	FunctionCell(tree, (nsubfeatures, maxdepth, prune_purity), ModelProperties(),
	      kwtitle("Decision Tree (classifier)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	tree(x, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a trained function cell by training a decision tree classifier using the DecisionTree.jl interface.
"""
# Training
tree(x::T where T<:CellDataL, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	tree((getx!(x), gety(x)), nsubfeatures, maxdepth, prune_purity; kwargs...)

tree(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	tree((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), nsubfeatures, maxdepth, prune_purity; kwargs...)

tree(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[tree] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	treedata = DecisionTree.build_tree(yenc, getobs(x[1])', nsubfeatures, maxdepth; kwargs...)

	# Prune if necessary
	if prune_purity >=0
		treedata = DecisionTree.prune_tree(treedata, prune_purity)
	end

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(tree, Model(treedata, modelprops), kwtitle("Decision Tree (classifier)", kwargs))
end

# Execution
tree(x::T where T<:CellData, model::Model{<:DecisionTree.LeafOrNode}) =
	datacell(tree(getx!(x), model), gety(x))
tree(x::T where T<:AbstractVector, model::Model{<:DecisionTree.LeafOrNode}) =
	tree(mat(x, LearnBase.ObsDim.Constant{2}()), model)
tree(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.LeafOrNode}) =
	# Always return probabilities
	DecisionTree.apply_tree_proba(model.data, getobs(x)', collect(1:model.properties.odim))'::Matrix{Float64} # labels are integer-coded, just use [1,2, ..., C]

##########################
# FunctionCell Interface #
##########################
"""
	treer(maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a function cell that when piped data, trains a decision tree regressor using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `maxlabels` number of samples used in each leaf to construct estimates (default 5)
  * `nsubfeatures` number of features to consider for node thresholdingi (default 0 e.g. all)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `prune_purity` fraction in [0.0,1.0] of combined purity for leaves to be merged (-1.0 indicates no pruning)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

!!! note
	The effects of `prune_purity` in the case of regression are not quite clear. Should work well for very
	low values e.g. leaves containing many distinct values are combined.

Read the `DecisionTree.jl` documentation for more information.
"""
treer(maxlabels::Int = 5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	FunctionCell(treer, (maxlabels, nsubfeatures, maxdepth, prune_purity), ModelProperties(),
	      kwtitle("Decision Tree (regressor)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	treer(x, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a trained function cell by training a decision tree regressor using the DecisionTree.jl interface.
"""
# Training
treer(x::T where T<:CellDataL, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	treer((getx!(x), gety(x)), maxlabels, nsubfeatures, maxdepth, prune_purity; kwargs...)

treer(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	treer((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), maxlabels, nsubfeatures, maxdepth, prune_purity; kwargs...)

treer(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[treer] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Convert targets to vector of Float first
	yenc = Vector{typeof(1.0)}(vec(x[2]))

	# Train model
	treedata = DecisionTree.build_tree(yenc, getobs(x[1])', maxlabels, nsubfeatures, maxdepth; kwargs...)

	# Prune if necessary
	if prune_purity >=0
		treedata = DecisionTree.prune_tree(treedata, prune_purity)
	end

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]),1)

	FunctionCell(treer, Model(treedata, modelprops), kwtitle("Decision Tree (regressor)", kwargs))
end

# Execution
treer(x::T where T<:CellData, model::Model{<:DecisionTree.LeafOrNode}) =
	datacell(treer(getx!(x), model), gety(x))
treer(x::T where T<:AbstractVector, model::Model{<:DecisionTree.LeafOrNode}) =
	treer(mat(x, LearnBase.ObsDim.Constant{2}()), model)
treer(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.LeafOrNode})::Matrix{Float64} =
	Matrix(DecisionTree.apply_tree(model.data, getobs(x)')')

##########################
# FunctionCell Interface #
##########################
"""
	randomforest(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1 [;kwargs])

Generates a function cell that when piped data, trains a random forest classifier using the DecisionTree.jl interface.
The trees are not pruned.

# Arguments (from `DecisionTree`)
  * `ntrees` number of trees (default 50)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `partialsampling` is the random fraction in [0.0,1.0] of training samples to be used for each tree (default 0.7)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
randomforest(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	FunctionCell(randomforest, (ntrees, nsubfeatures, partialsampling, maxdepth), ModelProperties(),
	      	     kwtitle("Random Forest (classifier, $ntrees trees)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	randomforest(x, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1 [;kwargs])

Generates a trained function cell by training a random forest classifier using the DecisionTree.jl interface.
"""
# Training
randomforest(x::T where T<:CellDataL, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	randomforest((getx!(x), gety(x)), ntrees, nsubfeatures, partialsampling, maxdepth; kwargs...)

randomforest(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	randomforest((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), ntrees, nsubfeatures, partialsampling, maxdepth; kwargs...)

randomforest(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
begin
	@assert nobs(x[1]) == nobs(x[2]) "[randomforest] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	rforestdata = DecisionTree.build_forest(yenc, getobs(x[1])', nsubfeatures, ntrees, partialsampling, maxdepth; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(randomforest, Model(rforestdata, modelprops), kwtitle("Random Forest (classifier, $ntrees trees)", kwargs))
end

# Execution
randomforest(x::T where T<:CellData, model::Model{<:DecisionTree.Ensemble}) =
	datacell(randomforest(getx!(x), model), gety(x))
randomforest(x::T where T<:AbstractVector, model::Model{<:DecisionTree.Ensemble}) =
	randomforest(mat(x, LearnBase.ObsDim.Constant{2}()), model)
randomforest(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.Ensemble})::Matrix{Float64} =
	DecisionTree.apply_forest_proba(model.data, getobs(x)', collect(1:model.properties.odim))' # labels are integer-coded, just use [1,2, ..., C]

##########################
# FunctionCell Interface #
##########################
"""
	randomforestr(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5 [;kwargs])

Generates a function cell that when piped data, trains a random forest regressor using the DecisionTree.jl interface.
The trees are not pruned.

# Arguments (from `DecisionTree`)
  * `ntrees` number of trees (default 50)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `partialsampling` is the random fraction in [0.0,1.0] of training samples to be used for each tree (default 0.7)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `maxlabels` number of samples used in each leaf to construct estimates (default 5)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
randomforestr(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	FunctionCell(randomforestr, (ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels), ModelProperties(),
	             kwtitle("Random Forest (regressor, $ntrees trees)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	randomforestr(x, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5 [;kwargs])

Generates a trained function cell by training a random forest regressor using the DecisionTree.jl interface.
"""
# Training
randomforestr(x::T where T<:CellDataL, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	randomforestr((getx!(x), gety(x)), ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels; kwargs...)

randomforestr(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	randomforestr((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels; kwargs...)

randomforestr(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
begin
	@assert nobs(x[1]) == nobs(x[2]) "[randomforestr] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Convert targets to vector of Float first
	yenc = Vector{typeof(1.0)}(vec(x[2]))

	# Train model
	rforestdata = DecisionTree.build_forest(yenc, getobs(x[1])', nsubfeatures, ntrees, maxlabels, partialsampling, maxdepth; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]),1)

	FunctionCell(randomforestr, Model(rforestdata, modelprops), kwtitle("Random Forest (regressor, $ntrees trees)", kwargs))
end

# Execution
randomforestr(x::T where T<:CellData, model::Model{<:DecisionTree.Ensemble}) =
	datacell(randomforestr(getx!(x), model), gety(x))
randomforestr(x::T where T<:AbstractVector, model::Model{<:DecisionTree.Ensemble}) =
	randomforestr(mat(x, LearnBase.ObsDim.Constant{2}()), model)
randomforestr(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.Ensemble}) =
	Matrix(DecisionTree.apply_forest(model.data, getobs(x)')')

##########################
# FunctionCell Interface #
##########################
"""
	adaboostump(niterations::Int=10 [;kwargs])

Generates a function cell that when piped data, trains an Adaboosted-stump classifier using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `niterations` number of iterations e.g ensemble size (default 10)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
adaboostump(niterations::Int=10; kwargs...) =
	FunctionCell(adaboostump, (niterations,), ModelProperties(),
	      kwtitle("Adaboost Stumps (classifier, $niterations stumps)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	adaboostump(x, niterations::Int=10 [;kwargs])

Generates a trained function cell by training an Adaboosted-stump classifier using the DecisionTree.jl interface.
"""
# Training
adaboostump(x::T where T<:CellDataL, niterations::Int=10; kwargs...) = adaboostump((getx!(x), gety(x)), niterations; kwargs...)

adaboostump(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, niterations::Int=10; kwargs...) =
	adaboostump((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), niterations; kwargs...)

adaboostump(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, niterations::Int=10; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[adaboostump] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	stumpdata = DecisionTree.build_adaboost_stumps(yenc, getobs(x[1])', niterations; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(adaboostump, Model(stumpdata, modelprops), kwtitle("Adaboost Stumps (classifier, $niterations stumps)", kwargs));
end

# Execution
adaboostump(x::T where T<:CellData, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	datacell(adaboostump(getx!(x), model), gety(x))
adaboostump(x::T where T<:AbstractVector, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	adaboostump(mat(x, LearnBase.ObsDim.Constant{2}()), model)
adaboostump(x::T where T<:AbstractMatrix, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	DecisionTree.apply_adaboost_stumps_proba(model.data[1], model.data[2],
					  getobs(x)', collect(1:model.properties.odim))'::Matrix{Float64}


text:
########################################################################################
# Integrated Brownian Motion
########################################################################################
"""Generate the discrete dynamics for a q-IBM model. INCLUDES AUTOMATIC PRECONDITIONING!

Careful: Dimensions are ordered differently than in `probnum`!"""
function ibm(d::Integer, q::Integer)
    F̃ = diagm(1 => ones(q))
    I_d = diagm(0 => ones(d))
    F = kron(F̃, I_d)  # In probnum the order is inverted

    @fastmath function A!(A::AbstractMatrix, h::Real)
        # Assumes that A comes from a previous computation => zeros and one-diag
        val = one(h)
        for i in 1:q
            val = val / i
            for j in 1:d*(q+1-i)
                @inbounds A[j,j+(d*i)] = val
            end
        end
    end

    @fastmath function _transdiff_ibm_element(row::Int, col::Int, h::Real)
        idx = 2 * q + 1 - row - col
        fact_rw = factorial(q - row)
        fact_cl = factorial(q - col)
        return h / (idx * fact_rw * fact_cl)
    end
    @fastmath function Q!(Q::AbstractMatrix, h::Real, σ²::Real=1.0)
        val = one(h)
        @simd for col in 0:q
            @simd for row in col:q
                val = _transdiff_ibm_element(row, col, h) * σ²
                @simd for i in 0:d-1
                    @inbounds Q[1 + col*d + i,1 + row*d + i] = val
                    @inbounds Q[1 + row*d + i,1 + col*d + i] = val
                end
            end
        end
    end

    return A!, Q!
end

"""Same as above, but without the automatic preconditioning"""
function vanilla_ibm(d::Integer, q::Integer)
    F̃ = diagm(1 => ones(q))
    I_d = diagm(0 => ones(d))
    F = kron(F̃, I_d)  # In probnum the order is inverted

    @fastmath function A!(A::AbstractMatrix, h::Real)
        # Assumes that A comes from a previous computation => zeros and one-diag
        val = one(h)
        for i in 1:q
            val = val * h / i
            for j in 1:d*(q+1-i)
                @inbounds A[j,j+(d*i)] = val
            end
        end
    end

    @fastmath function _transdiff_ibm_element(row::Int, col::Int, h::Real)
        idx = 2 * q + 1 - row - col
        fact_rw = factorial(q - row)
        fact_cl = factorial(q - col)
        return h^idx / (idx * fact_rw * fact_cl)
    end
    @fastmath function Q!(Q::AbstractMatrix, h::Real, σ²::Real=1.0)
        val = one(h)
        @simd for col in 0:q
            @simd for row in col:q
                val = _transdiff_ibm_element(row, col, h) * σ²
                @simd for i in 0:d-1
                    @inbounds Q[1 + col*d + i,1 + row*d + i] = val
                    @inbounds Q[1 + row*d + i,1 + col*d + i] = val
                end
            end
        end
    end

    return A!, Q!
end

max_stars_repo_name: nathanaelbosch/capos

max_issues_repo_name: nathanaelbosch/capos

max_forks_repo_name: nathanaelbosch/capos


text:
module Advection

export
    div_uc, div_ũu, div_ũv, div_ũw,

    momentum_flux_uu,
    momentum_flux_uv,
    momentum_flux_uw,
    momentum_flux_vu,
    momentum_flux_vv,
    momentum_flux_vw,
    momentum_flux_wu,
    momentum_flux_wv,
    momentum_flux_ww,
    advective_tracer_flux_x,
    advective_tracer_flux_y,
    advective_tracer_flux_z,

    CenteredSecondOrder,
    CenteredFourthOrder

using Oceananigans.Operators

abstract type AbstractAdvectionScheme end

include("centered_second_order.jl")
include("centered_fourth_order.jl")
include("momentum_advection_operators.jl")
include("tracer_advection_operators.jl")

end # module


text:
mutable struct Moon
    pos
    vel
end

function Moon(pos)
    Moon(pos, zeros(Int64, 3))
end

mutable struct System
    moons
end

function apply_gravity_pair!(moon1, moon2)
end

function step!(moons)
    for i=1:length(moons)
        for j=i+1:length(moons)
            for d=1:3
                @inbounds if moons[j].pos[d] > moons[i].pos[d]
                    @inbounds moons[i].vel[d] += 1
                    @inbounds moons[j].vel[d] -= 1
                elseif moons[j].pos[d] < moons[i].pos[d]
                    @inbounds moons[i].vel[d] -= 1
                    @inbounds moons[j].vel[d] += 1
                end
            end
        end
    end
    for j=1:length(moons)
        for d=1:3
            @inbounds moons[j].pos[d] += moons[j].vel[d]
        end
    end
end

function run!(system, steps)
    for i=1:steps
        step!(system.moons)
    end
end

function total_energy(system::System)
    sum(total_energy(moon) for moon in system.moons)
end

function potential_energy(moon::Moon)
    sum(abs.(moon.pos))
end

function kinetic_energy(moon::Moon)
    sum(abs.(moon.vel))
end

function total_energy(moon::Moon)
    potential_energy(moon) * kinetic_energy(moon)
end

test_moons = [Moon([-1,0,2]), Moon([2,-10,-7]), Moon([4,-8,8]), Moon([3,5,1])]

# input for day 12:
input_system = System([
    Moon([15, -2, -6]),
    Moon([-5, -4, -11]),
    Moon([0, -6, 0]),
    Moon([5, 9, 6]),
])

max_stars_repo_name: st--/advent-of-code

max_issues_repo_name: st--/advent-of-code

max_forks_repo_name: st--/advent-of-code


text:
##################### negative logarithm of density ##############################
function nld(G::Univariate, x::Number)
    ((x-mean(G))^2/var(G) + log(var(G)) + log(2π))/2
end

function nld(G::Multivariate, x::AbstractVector)
    x = difference(x, mean(G))
	nld(cov(G), x)
end

function nld(Σ::AbstractMatOrFacOrUni, x::AbstractVector)
    Σ = factorize(Σ) # since factorize on a Inverse is a no-op, saves unnecessary factorization in laplace's method
    Λ = inverse(Σ)
    (dot(x, Λ, x) + logdet(Σ) + length(x) * log(2π))/2
end

# probability density function
pdf(G::Gaussian, x) = exp(-nld(G, x))
pdf(x::Real) = exp(-x^2/2) / sqrt(2π)
function pdf(G::Uni, x::Real)
	σ = sqrt(G.Σ)
	pdf((x - G.μ) / σ) / σ
end

using SpecialFunctions: erf
cdf(x::Real) = (1 + erf(x / sqrt(2))) / 2
cdf(G::Uni, x::Real) = cdf((x - G.μ) / G.Σ)

# lazy
nld(G::Gaussian) = x -> nld(G, x)
pdf(G::Gaussian) = x -> pdf(G, x)
cdf(G::Uni) = x -> cdf(G, x)

entropy(G::Univariate) = log(2π*ℯ*cov(G)) / 2
entropy(G::Multivariate) = (length(G) * log(2π*ℯ) + logdet(cov(G))) / 2

# for Gaussian process
function nld(G::Gaussian, x::AbstractVector, y::AbstractVector)
	nld(marginal(G, x), y)
end
# negative log marginal likelihood
# with noiseless observations y
nlml(K::AbstractMatOrFac, y::AbstractVector) = nld(K, y)
nlml(k, x::AbstractVector, y::AbstractVector) = nlml(gramian(k, x), y)
nlml(G::Gaussian, x::AbstractVector, y::AbstractVector) = nlml(cov(G), x, y)

# with noisy observations y
# if first argument is not Gaussian, assume it's a kernel
function nlml(k, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	G = gramian(k, x)
	F = factorize(G, tol = tol) # this should be a pivoted cholesky or low rank
	y, Σ = mean(Gy), cov(Gy)
	K = Woodbury(Σ, F)
	nlml(K, y) # reduces to "noiseless" call
end
function nlml(G::Gaussian, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	if mean(G) != zero_mean # subtract prior mean, if applicable
		mean(Gy) .-= mean(G, x)
	end
	nlml(cov(G), x, Gy)
end
# convenience for passing noise covariance separately
function nlml(k, x::AbstractVector, y::AbstractVector, Σ::Union{Real, AbstractMatOrFacOrUni})
	nlml(k, x, 𝒩(y, Σ))
end

# value and gradient of nlml with noiseless observations
function pushforward(::typeof(nlml), k, x::AbstractVector, y::AbstractVector)
	K = gramian(k, x)
	pushforward(nlml, K, y)
end
# value and gradient of nlml with noisy observations
function pushforward(::typeof(nlml), k, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	K = gramian(k, x)
	K = factorize(K, tol = tol) # ideally, this is a pivoted cholesky or low rank
	y, Σ = mean(Gy), cov(Gy)
	K = Woodbury(Σ, K)
	pushforward(nlml, K, y)
end
function pushforward(::typeof(nlml), K::AbstractMatOrFac, y::AbstractVector)
	K = factorize(K)
    α = K \ y # coefficients α are the same as in conditional mean, could pre-allocate
	val = (dot(y, α) + logdet(K) + length(y) * log(2π))/2
	K⁻¹ = Matrix(inv(K))
	@. K⁻¹ -= α * α'
	val, (dK::AbstractMatOrFac) -> dot(K⁻¹, dK)/2 # this is O(n^2) per evaluation
end

# could be useful if certain hyper-parameter combinations lead to rank-deficiency
# function safe_nlml(Σ, y)
#     try nlml(Σ, y) catch e;
#         println(e)
#         Inf
#     end
# end

# # TODO: stochastic gradient via stochastic trace estimator and logdet
# # derivative for Multivariate nld
# function pushforward(::typeof(nlml), Σ::AbstractMatOrFac, y::AbstractVector,
# 										dense_inverse::Val{false} = Val(false))
# 	Σ = factorize(Σ)
#     α = Σ \ y # could pre-allocate here for Σ\∇,
# 	val = (dot(y, α) + logdet(Σ) + length(y) * log(2π))/2
#     val, function (dΣ)
# 				(tr(inverse(Σ)*dΣ) - dot(α, dΣ, α))/2 # this is O(n^3) per evaluation
# 			end
# end


text:
struct SelectValues{T}
    name::Symbol
    values::Vector{T}
    split::Bool
    filter::Bool
end

SelectValues(name, values, split = false) = SelectValues(name, values, split, true)

struct SelectPredicate{T}
    name::Symbol
    f::T
    split::Bool
    filter::Bool
end

SelectPredicate(name, f, split = false) = SelectPredicate(name, f, split, true)

struct Data2Select{T<:AbstractIndexedTable, N1, N2}
    table::T
    discrete::NTuple{N1, SelectValues}
    continuous::NTuple{N2, SelectPredicate}
end

struct SelectedData{T<:AbstractIndexedTable, N}
    table::T
    splitby::NTuple{N, Symbol}
end

SelectedData(d2s::Data2Select) =
    SelectedData(selectdata(d2s.table, d2s.discrete, d2s.continuous), Tuple(i.name for i in d2s.discrete if i.split))

Base.:(==)(a::SelectedData, b::SelectedData) = (a.table == b.table) && (a.splitby == b.splitby)

function selectdata(df, discrete, continuous)
    f = function(i)
        all(getfield(i, s.name) in s.values for s in discrete if s.filter) &&
        all(s.f(getfield(i, s.name)) for s in continuous if s.filter)
    end
    filter(f, df)
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

opticinterface(::Type{S}, insidematerial, outsidematerial, reflectance = zero(S), interfacemode = ReflectOrTransmit) where {S<:Real} = FresnelInterface{S}(insidematerial, outsidematerial, reflectance = reflectance, transmission = one(S) - reflectance, interfacemode = interfacemode)

"""
    SphericalLens(insidematerial, frontvertex, frontradius, backradius, thickness, semidiameter;  lastmaterial = OpticSim.GlassCat.Air, nextmaterial = OpticSim.GlassCat.Air, frontsurfacereflectance = 0.0, backsurfacereflectance = 0.0, frontdecenter = (0, 0), backdecenter = (0, 0), interfacemode = ReflectOrTransmit)

Constructs a simple cylindrical lens with spherical front and back surfaces. The side walls of the lens are absorbing.
"""
function SphericalLens(insidematerial::T, frontvertex::S, frontradius::S, backradius::S, thickness::S, semidiameter::S; lastmaterial::Q = OpticSim.GlassCat.Air, nextmaterial::R = OpticSim.GlassCat.Air, frontsurfacereflectance::S = zero(S), backsurfacereflectance::S = zero(S), frontdecenter::Tuple{S,S} = (zero(S), zero(S)), backdecenter::Tuple{S,S} = (zero(S), zero(S)), interfacemode = ReflectOrTransmit) where {R<:OpticSim.GlassCat.AbstractGlass,Q<:OpticSim.GlassCat.AbstractGlass,T<:OpticSim.GlassCat.AbstractGlass,S<:Real}
    @assert !isnan(frontradius)
    @assert !isnan(backradius)

    return ConicLens(insidematerial, frontvertex, frontradius, zero(S), backradius, zero(S), thickness, semidiameter, lastmaterial = lastmaterial, nextmaterial = nextmaterial, frontsurfacereflectance = frontsurfacereflectance, backsurfacereflectance = backsurfacereflectance, frontdecenter = frontdecenter, backdecenter = backdecenter, interfacemode = interfacemode)
end
export SphericalLens

"""
    ConicLens(insidematerial, frontvertex, frontradius, frontconic, backradius, backconic, thickness, semidiameter;  lastmaterial = OpticSim.GlassCat.Air, nextmaterial = OpticSim.GlassCat.Air, frontsurfacereflectance = 0.0, backsurfacereflectance = 0.0, frontdecenter = (0, 0), backdecenter = (0, 0), interfacemode = ReflectOrTransmit)

Constructs a simple cylindrical lens with front and back surfaces with a radius and conic term. The side walls of the lens are absorbing.
"""
function ConicLens(insidematerial::T, frontvertex::S, frontradius::S, frontconic::S, backradius::S, backconic::S, thickness::S, semidiameter::S; lastmaterial::Q = OpticSim.GlassCat.Air, nextmaterial::R = OpticSim.GlassCat.Air, frontsurfacereflectance::S = zero(S), backsurfacereflectance::S = zero(S), nsamples::Int = 17, frontdecenter::Tuple{S,S} = (zero(S), zero(S)), backdecenter::Tuple{S,S} = (zero(S), zero(S)), interfacemode = ReflectOrTransmit) where {R<:OpticSim.GlassCat.AbstractGlass,Q<:OpticSim.GlassCat.AbstractGlass,T<:OpticSim.GlassCat.AbstractGlass,S<:Real}
    @assert !isnan(frontradius)
    @assert !isnan(frontconic)

    return AsphericLens(insidematerial, frontvertex, frontradius, frontconic, nothing, backradius, backconic, nothing, thickness, semidiameter, lastmaterial = lastmaterial, nextmaterial = nextmaterial, frontsurfacereflectance = frontsurfacereflectance, backsurfacereflectance = backsurfacereflectance, nsamples = nsamples, frontdecenter = frontdecenter, backdecenter = backdecenter, interfacemode = interfacemode)
end
export ConicLens

"""
    AsphericLens(insidematerial, frontvertex, frontradius, frontconic, frontaspherics, backradius, backconic, backaspherics, thickness, semidiameter;  lastmaterial = OpticSim.GlassCat.Air, nextmaterial = OpticSim.GlassCat.Air, frontsurfacereflectance = 0.0, backsurfacereflectance = 0.0, frontdecenter = (0, 0), backdecenter = (0, 0), interfacemode = ReflectOrTransmit)

Cosntructs a simple cylindrical lens with front and back surfaces with a radius, conic and apsheric terms. The side walls of the lens are absorbing.
"""
function AsphericLens(insidematerial::T, frontvertex::S, frontradius::S, frontconic::S, frontaspherics::Union{Nothing,Vector{Pair{Int,S}}}, backradius::S, backconic::S, backaspherics::Union{Nothing,Vector{Pair{Int,S}}}, thickness::S, semidiameter::S; lastmaterial::Q = OpticSim.GlassCat.Air, nextmaterial::R = OpticSim.GlassCat.Air, frontsurfacereflectance::S = zero(S), backsurfacereflectance::S = zero(S), nsamples::Int = 17, frontdecenter::Tuple{S,S} = (zero(S), zero(S)), backdecenter::Tuple{S,S} = (zero(S), zero(S)), interfacemode = ReflectOrTransmit) where {R<:OpticSim.GlassCat.AbstractGlass,Q<:OpticSim.GlassCat.AbstractGlass,T<:OpticSim.GlassCat.AbstractGlass,S<:Real}
    @assert semidiameter > zero(S)
    @assert !isnan(frontradius)

    frontdecenter_l = abs(frontdecenter[1]) > eps(S) && abs(frontdecenter[2]) > eps(S) ? sqrt(frontdecenter[1]^2 + frontdecenter[2]^2) : zero(S)
    backdecenter_l = abs(backdecenter[1]) > eps(S) && abs(backdecenter[2]) > eps(S) ? sqrt(backdecenter[1]^2 + backdecenter[2]^2) : zero(S)

    if isinf(frontradius) && (frontaspherics === nothing)
        #planar
        lens_front = leaf(Plane(SVector{3,S}(0, 0, 1), SVector{3,S}(0, 0, frontvertex), vishalfsizeu = semidiameter, vishalfsizev = semidiameter, interface = interface = opticinterface(S, insidematerial, lastmaterial, frontsurfacereflectance, interfacemode)))
    elseif frontconic == zero(S) && (frontaspherics === nothing)
        #spherical
        if frontradius < zero(S)
            ϕmax = NaNsafeasin(min(abs((min(semidiameter, abs(frontradius)) + frontdecenter_l) / frontradius), one(S))) + S(π / 50)
            lens_front = leaf(SphericalCap(abs(frontradius), ϕmax, SVector{3,S}(0, 0, 1), SVector{3,S}(frontdecenter[1], frontdecenter[2], frontvertex), interface = opticinterface(S, insidematerial, lastmaterial, frontsurfacereflectance, interfacemode)))
        else
            offset = translation(S, frontdecenter[1], frontdecenter[2], frontvertex - frontradius)
            surf = Sphere(abs(frontradius), interface = opticinterface(S, insidematerial, lastmaterial, frontsurfacereflectance, interfacemode))
            lens_front = leaf(surf, offset)
        end
        if abs(frontradius) - frontdecenter_l <= semidiameter
            # if optical surface is smaller than semidiameter then use a plane to fill the gap
            plane = leaf(Plane(SVector{3,S}(0, 0, 1), SVector{3,S}(0, 0, frontvertex - frontradius), vishalfsizeu = semidiameter, vishalfsizev = semidiameter, interface = interface = opticinterface(S, insidematerial, lastmaterial, frontsurfacereflectance, interfacemode)))
            if frontradius > zero(S)
                lens_front = csgunion(lens_front, plane)
            else
                lens_front = csgintersection(lens_front, plane)
            end
        end
    else
        #conic or aspheric
        if frontaspherics !== nothing
            frontaspherics = [(i, -k) for (i, k) in frontaspherics]
        end
        surf = AcceleratedParametricSurface(ZernikeSurface(semidiameter + frontdecenter_l + S(0.01), radius = -frontradius, conic = frontconic, aspherics = frontaspherics), nsamples, interface = opticinterface(S, insidematerial, lastmaterial, frontsurfacereflectance, interfacemode))
        lens_front = leaf(surf, translation(S, frontdecenter[1], frontdecenter[2], frontvertex))
    end
    if isinf(backradius) && (backaspherics === nothing)
        #planar
        lens_rear = leaf(Plane(SVector{3,S}(0, 0, -1), SVector{3,S}(0, 0, frontvertex - thickness), vishalfsizeu = semidiameter, vishalfsizev = semidiameter, interface = interface = opticinterface(S, insidematerial, nextmaterial, backsurfacereflectance, interfacemode)))
    elseif backconic == zero(S) && (backaspherics === nothing)
        #spherical
        if backradius < zero(S)
            offset = translation(S, backdecenter[1], backdecenter[2], (frontvertex - thickness) - backradius)
            surf = Sphere(abs(backradius), interface = opticinterface(S, insidematerial, nextmaterial, backsurfacereflectance, interfacemode))
            lens_rear = leaf(surf, offset)
            if backradius == semidiameter
                # in this case we need to clip the sphere
                plane = leaf(Plane(SVector{3,S}(0, 0, 1), SVector{3,S}(0, 0, frontvertex - thickness + backradius)))
                lens_rear = csgintersection(lens_rear, plane)
            end
        else
            ϕmax = NaNsafeasin(min(abs((min(semidiameter, abs(backradius)) + backdecenter_l) / backradius), one(S))) + S(π / 50)
            lens_rear = leaf(SphericalCap(abs(backradius), ϕmax, SVector{3,S}(0, 0, -1), SVector{3,S}(backdecenter[1], backdecenter[2], frontvertex - thickness), interface = opticinterface(S, insidematerial, nextmaterial, backsurfacereflectance, interfacemode)))
        end
        if abs(backradius) - backdecenter_l <= semidiameter
            plane = leaf(Plane(SVector{3,S}(0, 0, -1), SVector{3,S}(0, 0, frontvertex - thickness - backradius), vishalfsizeu = semidiameter, vishalfsizev = semidiameter, interface = interface = opticinterface(S, insidematerial, lastmaterial, backsurfacereflectance, interfacemode)))
            if backradius < zero(S)
                lens_rear = csgunion(lens_rear, plane)
            else
                lens_rear = csgintersection(lens_rear, plane)
            end
        end
    else
        #conic or aspheric
        if backaspherics !== nothing
            backaspherics = Tuple{Int,S}.(backaspherics)
        end
        surf = AcceleratedParametricSurface(ZernikeSurface(semidiameter + backdecenter_l + S(0.01), radius = backradius, conic = backconic, aspherics = backaspherics), nsamples, interface = opticinterface(S, insidematerial, nextmaterial, backsurfacereflectance, interfacemode))
        lens_rear = leaf(surf, Transform{S}(zero(S), S(π), zero(S), backdecenter[1], backdecenter[2], frontvertex - thickness))
    end
    extra_front = frontradius >= zero(S) || isinf(frontradius) ? zero(S) : abs(frontradius) - sqrt(frontradius^2 - semidiameter^2)
    extra_back = backradius >= zero(S) || isinf(backradius) ? zero(S) : abs(backradius) - sqrt(backradius^2 - semidiameter^2)
    barrel_center = ((frontvertex + extra_front) + (frontvertex - thickness - extra_back)) / 2
    lens_barrel = leaf(Cylinder(semidiameter, (thickness + extra_back + extra_front) * 2, interface = FresnelInterface{S}(insidematerial, OpticSim.GlassCat.Air, reflectance = zero(S), transmission = zero(S))), translation(S, zero(S), zero(S), barrel_center))
    lens_csg = csgintersection(csgintersection(lens_front, lens_rear), lens_barrel)
    return lens_csg
end
export AsphericLens

"""
    FresnelLens(insidematerial, frontvertex, radius, thickness, semidiameter, groovedepth; conic = 0.0, aspherics = nothing, outsidematerial = OpticSim.GlassCat.Air)

Create a Fresnel lens as a CSG object, can be concave or convex. Groove positions are found iteratively based on `groovedepth`. For negative radii the vertex on the central surface is at `frontvertex`, so the total thickness of the lens is `thickness` + `groovedepth`.
**Aspherics currently not supported**.
"""
function FresnelLens(insidematerial::G, frontvertex::T, radius::T, thickness::T, semidiameter::T, groovedepth::T; conic::T = 0.0, aspherics::Union{Nothing,Vector{Pair{Int,T}}} = nothing, outsidematerial::H = OpticSim.GlassCat.Air, reverse::Bool = false) where {T<:Real,G<:OpticSim.GlassCat.AbstractGlass,H<:OpticSim.GlassCat.AbstractGlass}
    @assert abs(radius) > semidiameter
    interface = FresnelInterface{T}(insidematerial, outsidematerial)

    if aspherics !== nothing
        aspherics = [(i, -k) for (i, k) in aspherics]
    end

    # make the fresnel
    if conic == zero(T)
        if radius == zero(T) || isinf(radius)
            @error "Invalid radius"
        end
        # spherical lens which makes things much easier, we can just caluclate groove positions directly
        n = 1
        if radius > zero(T)
            sphere = Sphere(radius, interface = interface)
            fresnel = leaf(sphere, translation(T, zero(T), zero(T), frontvertex - radius))
        else
            ϕmax = NaNsafeasin(abs(semidiameter / radius)) + T(π / 50)
            sphere = SphericalCap(abs(radius), ϕmax, interface = interface)
            fresnel = leaf(sphere, translation(T, zero(T), zero(T), frontvertex))
        end
        while true
            offset = n * groovedepth
            cylrad = sqrt(offset * (2 * abs(radius) - offset))
            if cylrad >= semidiameter
                break
            end
            if radius > zero(T)
                cylinder = leaf(Cylinder(cylrad, groovedepth * 1.5, interface = interface), translation(T, zero(T), zero(T), frontvertex - groovedepth / 2))
                newsurf = csgdifference(leaf(sphere, translation(T, zero(T), zero(T), frontvertex - radius + offset)), cylinder)
                fresnel = csgunion(newsurf, fresnel)
            else
                cylinder = leaf(Cylinder(cylrad, groovedepth * 1.5, interface = interface), translation(T, zero(T), zero(T), frontvertex + groovedepth / 2))
                newsurf = csgunion(leaf(sphere, translation(T, zero(T), zero(T), frontvertex - offset)), cylinder)
                fresnel = csgintersection(newsurf, fresnel)
            end
            n += 1
        end
    elseif (aspherics === nothing)
        if radius == zero(T) || isinf(radius)
            @error "Invalid radius"
        end
        n = 1
        surface = AcceleratedParametricSurface(ZernikeSurface(semidiameter + T(0.01), radius = -radius, conic = conic), interface = interface)
        fresnel = leaf(surface, translation(T, zero(T), zero(T), frontvertex))
        while true
            offset = n * groovedepth
            if radius < zero(T)
                offset *= -1
            end
            cylrad = sqrt(2 * radius * offset - (conic + one(T)) * offset^2)
            if cylrad >= semidiameter
                break
            end
            if radius > zero(T)
                cylinder = leaf(Cylinder(cylrad, groovedepth * 1.5, interface = interface), translation(T, zero(T), zero(T), frontvertex - groovedepth / 2))
                newsurf = csgdifference(leaf(surface, translation(T, zero(T), zero(T), frontvertex + offset)), cylinder)
                fresnel = csgunion(newsurf, fresnel)
            else
                cylinder = leaf(Cylinder(cylrad, groovedepth * 1.5, interface = interface), translation(T, zero(T), zero(T), frontvertex + groovedepth / 2))
                newsurf = csgunion(leaf(surface, translation(T, zero(T), zero(T), frontvertex + offset)), cylinder)
                fresnel = csgintersection(newsurf, fresnel)
            end
            n += 1
        end
    else
        # TODO - CSG gets more complicated as the surface can be both convex and concave and we need to find the groove positions iteratively
        @error "Ashperics not supported for Fresnel lenses"
    end

    outer_barrel = leaf(Cylinder(semidiameter, thickness * 2, interface = interface), translation(T, zero(T), zero(T), frontvertex - thickness / 2))
    fresnel = csgintersection(outer_barrel, fresnel)
    backplane = leaf(Plane(zero(T), zero(T), -one(T), zero(T), zero(T), frontvertex - thickness, vishalfsizeu = semidiameter, vishalfsizev = semidiameter))
    fresnel = csgintersection(backplane, fresnel)
    if !reverse
        fresnel = leaf(fresnel, rotationd(T, zero(T), T(180.0), zero(T)))
    end
    return fresnel
end
export FresnelLens


text:
function savefig_monitoring_results(r::MonitoringResults, symbols=keys(r.results); ext=:pdf)
    for s in symbols
        savefig_monitoring_results(r, s; ext)
    end
    return nothing
end

function savefig_monitoring_results(r::MonitoringResults, s::Symbol; ext=:pdf)
    times = r.times
    values = r.results[s]
    title, ylabel = _symbol2title_and_label(s)
    ylims = _defaultylims(values)
    device_labels = [str for (str, uuid) in r.devices]

    f = CairoMakie.Figure(; resolution=(1000, 500))
    ax =
        f[1, 1] = CairoMakie.Axis(
            f; xlabel="Time [s]", ylabel=ylabel, title=title, ylims=ylims
        )
    CairoMakie.scatterlines!(times, getindex.(values, 1); label=device_labels[1])
    for i in 2:length(first(values))
        CairoMakie.scatterlines!(times, getindex.(values, i); label=device_labels[i])
    end
    f[1, 2] = CairoMakie.Legend(f, ax, "Devices"; framevisible=false)
    filename =
        replace(replace(replace(lowercase(title), " " => "_"), "(" => ""), ")" => "") * "_plot.$(string(ext))"
    CairoMakie.save(filename, f)
    return nothing
end


text:
module TestFolds

using CUDA
using Folds
using FoldsCUDA
using Test
using UnionArrays

const ==′ = isequal

function test_unionvector()
    xs = UnionVector(undef, CuVector, Union{Float32, Missing}, 3)
    fill!(xs, 1)
    CUDA.@allowscalar xs[2] = missing
    @test collect(xs) ==′ [1.0f0, missing, 1.0f0]
    @test Folds.sum(x -> x isa Missing ? 0.0f0 : x, xs, CUDAEx()) == 2

    @testset "setindex!" begin
        ys = range(1.0f0, 3.0f0, length = 3)
        Folds.map!(identity, xs, ys)
        @test collect(xs) == 1:3

        Folds.map!(==(missing), xs, ys)
        @test collect(xs) ==′ [missing, missing, missing]
    end
end

function test_unionmatrix()
    A = UnionArray(undef, CuVector, Union{Float32, Missing}, (3, 2))
    fill!(A, 1)
    CUDA.@allowscalar A[2, 2] = missing
    @test (CUDA.@allowscalar A[1, 1]) === 1.0f0
    @test (CUDA.@allowscalar A[2, 2]) === missing
    @test (CUDA.@allowscalar A[3, 2]) === 1.0f0
    @test Folds.sum(x -> x isa Missing ? 0.0f0 : x, A, CUDAEx()) == 5

    @testset "setindex!" begin
        B = reshape(range(1.0f0, 6.0f0, length = 6), 3, 2)
        Folds.map!(identity, A, B)
        @test collect(A) == B

        Folds.map!(==(missing), A, B)
        @test collect(A) ==′ fill(missing, 3, 2)
    end
end

end  # module


text:
using Sobol, StatsFuns

function simulate(mcProcess::SubordinatedBrownianMotion, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: SobolMode, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    drift = mcProcess.drift
    sigma = mcProcess.sigma
    @assert T > 0
    type_sub = typeof(quantile(mcProcess.subordinator_, 0.5))
    isDualZero = drift * zero(type_sub) * 0
    X = Matrix{typeof(isDualZero)}(undef, Nsim, Nstep + 1)
    @views X[:, 1] .= isDualZero
    seq = SobolSeq(Nstep)
    skip(seq, Nstep * Nsim)
    @inbounds for i = 1:Nsim
        vec = norminvcdf.(next!(seq))
        @inbounds for j = 1:Nstep
            dt_s = rand(mcBaseData.rng, mcProcess.subordinator_)
            @views X[i, j+1] .= X[i, j] .+ drift .* dt_s .+ sigma .* sqrt.(dt_s) .* vec[j]
        end
    end

    return X
end

function simulate(mcProcess::SubordinatedBrownianMotionVec, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: SobolMode, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    drift = mcProcess.drift
    sigma = mcProcess.sigma
    @assert T > 0
    type_sub = typeof(quantile(mcProcess.subordinator_, 0.5))
    zero_drift = drift(zero(type_sub), zero(type_sub) + dt) * 0
    isDualZero = zero_drift * zero(type_sub) * 0
    X = Matrix{typeof(isDualZero)}(undef, Nsim, Nstep + 1)
    @views X[:, 1] .= isDualZero
    seq = SobolSeq(Nstep)
    skip(seq, Nstep * Nsim)
    @inbounds for i = 1:Nsim
        t_s = zero(type_sub)
        vec = norminvcdf.(next!(seq))
        @inbounds for j = 1:Nstep
            dt_s = rand(mcBaseData.rng, mcProcess.subordinator_)
            tmp_drift = drift(t_s, dt_s)
            t_s += dt_s
            @views X[i, j+1] = X[i, j] + tmp_drift + sigma * sqrt(dt_s) * vec[j]
        end
    end

    return X
end


text:
using JSON
using Test

include("../src/biclusterseval.jl")

using .biclusterseval: initialize_input_on_gpus, get_biclusters

const DATA_DIR = "data/unibic"

function main()
    @testset "Score population" begin
        for (root, _, files) in walkdir(DATA_DIR)
            isempty(files) && continue

            println("""
            ####################################
            Test Case: '$(basename(root))'
            ####################################""")

            @testset "$(basename(root))" begin

            input_paths = Vector()
            biclusters_paths = Vector()
            for file in files
                if occursin("hidden", file)
                    push!(biclusters_paths, joinpath(root, file))
                else
                    push!(input_paths, joinpath(root, file))
                end
            end

            for (input_path, bicluster_path) in zip(input_paths, biclusters_paths)
                println("Testing: '$(basename(input_path))'")
                @testset "$(basename(input_path))" begin

                ground_truth = JSON.parsefile(bicluster_path)
                answers = Dict(map(o -> o["cols"] => o["rows"], ground_truth))

                population::Vector{Vector{Int64}} = collect(keys(answers))

                for chromo in population
                    chromo .+= 1
                end

                d_data = initialize_input_on_gpus(input_path, 1)

                biclusters = get_biclusters(d_data, population, 1, true, 0.85)

                for bicluster in biclusters
                    bicluster["cols"] .-= 1
                    bicluster["rows"] .-= 1

                    @test length(bicluster["rows"]) ==
                          length(answers[bicluster["cols"]])
                end
                end
            end
            end
        end
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
"""
    PoissonExpert(λ)

PMF:

```math
P(X = k) = \\frac{\\lambda^k}{k!} e^{-\\lambda}, \\quad \\text{ for } k = 0,1,2,\\ldots.
```

See also: [Poisson Distribution](https://en.wikipedia.org/wiki/Poisson_distribution) (Wikipedia)

"""
struct PoissonExpert{T<:Real} <: NonZIDiscreteExpert
    λ::T
    PoissonExpert{T}(λ::T) where {T<:Real} = new{T}(λ)
end

function PoissonExpert(λ::T; check_args=true) where {T <: Real}
    check_args && @check_args(PoissonExpert, λ >= zero(λ))
    return PoissonExpert{T}(λ)
end

## Outer constructors
# PoissonExpert(λ::Real) = PoissonExpert(promote(λ)...)
PoissonExpert(λ::Integer) = PoissonExpert(float(λ))
PoissonExpert() = PoissonExpert(1.0)

## Conversion
function convert(::Type{PoissonExpert{T}}, λ::S) where {T <: Real, S <: Real}
    PoissonExpert(T(λ))
end
function convert(::Type{PoissonExpert{T}}, d::PoissonExpert{S}) where {T <: Real, S <: Real}
    PoissonExpert(T(d.λ), check_args=false)
end
copy(d::PoissonExpert) = PoissonExpert(d.λ, check_args=false)

## Loglikelihood of Expoert
logpdf(d::PoissonExpert, x...) = isinf(x...) ? 0.0 : Distributions.logpdf.(Distributions.Poisson(d.λ), x...)
pdf(d::PoissonExpert, x...) = isinf(x...) ? -Inf : Distributions.pdf.(Distributions.Poisson(d.λ), x...)
logcdf(d::PoissonExpert, x...) = isinf(x...) ? 0.0 : Distributions.logcdf.(Distributions.Poisson(d.λ), x...)
cdf(d::PoissonExpert, x...) = isinf(x...) ? 1.0 : Distributions.cdf.(Distributions.Poisson(d.λ), x...)

## expert_ll, etc
expert_ll_exact(d::PoissonExpert, x::Real) = LRMoE.logpdf(LRMoE.PoissonExpert(d.λ), x)
function expert_ll(d::PoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    # d_exp = LRMoE.PoissonExpert(d.λ*exposure)
    expert_ll = (yl == yu) ? logpdf.(d, yl) : logcdf.(d, yu) + log1mexp.(logcdf.(d, ceil.(yl) .- 1) - logcdf.(d, yu))
    return expert_ll
end
function expert_tn(d::PoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    # d_exp = LRMoE.PoissonExpert(d.λ*exposure)
    expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, ceil.(tl) .- 1) - logcdf.(d, tu))
    return expert_tn
end
function expert_tn_bar(d::PoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    # d_exp = LRMoE.PoissonExpert(d.λ*exposure)
    expert_tn_bar = (tl == tu) ? log1mexp.(logpdf.(d, tl)) : log1mexp.(logcdf.(d, tu) + log1mexp.(logcdf.(d, ceil.(tl) .- 1) - logcdf.(d, tu)))
    return expert_tn_bar
end

exposurize_expert(d::PoissonExpert; exposure = 1) = PoissonExpert(d.λ*exposure)

## Parameters
params(d::PoissonExpert) = (d.λ)
function params_init(y, d::PoissonExpert)
    λ_init = mean(y)
    try
        return PoissonExpert(λ_init)
    catch;
        PoissonExpert()
    end
end

## KS stats for parameter initialization
function ks_distance(y, d::PoissonExpert)
    p_zero = sum(y .== 0.0) / sum(y .>= 0.0)
    return max(abs(p_zero-0.0), (1-0.0)*HypothesisTests.ksstats(y[y .> 0.0], Distributions.Poisson(d.λ))[2])
end

## Simululation
sim_expert(d::PoissonExpert) = Distributions.rand(Distributions.Poisson(d.λ), 1)[1]

## penalty
penalty_init(d::PoissonExpert) = [2.0 1.0]
no_penalty_init(d::PoissonExpert) = [1.0 Inf]
penalize(d::PoissonExpert, p) = (p[1]-1)*log(d.λ) - d.λ/p[2]

## statistics
mean(d::PoissonExpert) = mean(Distributions.Poisson(d.λ))
var(d::PoissonExpert) = var(Distributions.Poisson(d.λ))
quantile(d::PoissonExpert, p) = quantile(Distributions.Poisson(d.λ), p)

## Misc functions for E-Step

function _sum_densy_series(d::PoissonExpert, yl, yu)
    if isinf(yu)
        series = 0:(max(yl-1, 0))
        return d.λ - sum(pdf.(d, series) .* series)[1]
    else
        series = yl:yu
        return sum(pdf.(d, series) .* series)[1]
    end
end

function _int_obs_Y_raw(d::PoissonExpert, yl, yu)
    # if yl == yu
    #     return yl
    # else
    #     return _sum_densy_series(d, yl, yu)
    # end
    return _sum_densy_series(d, yl, yu)
end

function _int_lat_Y_raw(d::PoissonExpert, tl, tu)
    return maximum([d.λ - _sum_densy_series(d, tl, tu), 0])
end

## EM: M-Step
function EM_M_expert(d::PoissonExpert,
                    tl, yl, yu, tu,
                    exposure,
                    z_e_obs, z_e_lat, k_e;
                    penalty = true, pen_pararms_jk = [2.0 1.0])

    # Old parameters
    λ_old = d.λ

    # Further E-Step
    # yl_yu_unique = unique_bounds(yl, yu)
    # int_obs_Y_tmp = _int_obs_Y_raw.(d, yl_yu_unique[:,1], yl_yu_unique[:,2])
    # Y_e_obs = exp.(-expert_ll_pos) .* int_obs_Y_tmp[match_unique_bounds(hcat(vec(yl), vec(yu)), yl_yu_unique)]
    # nan2num(Y_e_obs, 0.0) # get rid of NaN

    # tl_tu_unique = unique_bounds(tl, tu)
    # int_lat_Y_tmp = _int_lat_Y_raw.(d, tl_tu_unique[:,1], tl_tu_unique[:,2])
    # Y_e_lat = exp.(-expert_tn_bar_pos) .* int_lat_Y_tmp[match_unique_bounds(hcat(vec(tl), vec(tu)), tl_tu_unique)]
    # nan2num(Y_e_lat, 0.0) # get rid of NaN

    Y_e_obs = fill(0.0, length(yl))
    Y_e_lat = fill(0.0, length(yl))

    for i in 1:length(yl)
        d_expo = exposurize_expert(d, exposure = exposure[i])
        expert_ll_pos = expert_ll.(d_expo, tl[i], yl[i], yu[i], tu[i])
        expert_tn_bar_pos = expert_tn_bar.(d_expo, tl[i], yl[i], yu[i], tu[i])
        Y_e_obs[i] = exp(-expert_ll_pos) * _int_obs_Y_raw(d_expo, yl[i], yu[i])
        Y_e_lat[i] = exp(-expert_tn_bar_pos) * _int_lat_Y_raw(d_expo, tl[i], tu[i])
    end

    nan2num(Y_e_obs, 0.0)
    nan2num(Y_e_lat, 0.0)

    # Update parameters
    term_zkz = (z_e_obs .* exposure) .+ (z_e_lat .* k_e .* exposure)
    term_zkz_Y = (z_e_obs .* Y_e_obs) .+ (z_e_lat .* k_e .* Y_e_lat)

    λ_new = penalty ? ((sum(term_zkz_Y)[1] - (pen_pararms_jk[1]-1)) / (sum(term_zkz)[1] + 1/pen_pararms_jk[2])) : (sum(term_zkz_Y)[1] / sum(term_zkz)[1])

    # Need to deal with zero mass: λ is very small
    if λ_new < 0.00001
        λ_new = λ_old
    end

    return PoissonExpert(λ_new)
end

## EM: M-Step, exact observations
function EM_M_expert_exact(d::PoissonExpert,
                    ye, exposure,
                    # expert_ll_pos,
                    z_e_obs;
                    penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    # Old parameters
    λ_old = d.λ

    # Update parameters
    term_zkz = z_e_obs .* exposure
    term_zkz_Y = (z_e_obs .* ye)

    λ_new = penalty ? ((sum(term_zkz_Y)[1] - (pen_pararms_jk[1]-1)) / (sum(term_zkz)[1] + 1/pen_pararms_jk[2])) : (sum(term_zkz_Y)[1] / sum(term_zkz)[1])

    # Need to deal with zero mass: λ is very small
    if λ_new < 0.00001
        λ_new = λ_old
    end

    return PoissonExpert(λ_new)
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE