text:
# Trajectory utilities
# struct Vectorized{T}
#     trajectories::T
# end

# (vectorized::Vectorized)(x, args...) = map(traj -> traj(x, args...), vectorized.trajectories)
# map_elements(f, vectorized::Vectorized) = Vectorized(map(f, vectorized.trajectories))

function map_subfunctions(f, traj::Piecewise)
    Piecewise(map(f, traj.subfunctions), traj.breaks; clamp=traj.clamp)
end

# Polynomial utilities
function scale_argument(x::Polynomial{N}, s) where {N}
    Polynomial(ntuple(i -> s^(i - 1) * x.coeffs[i], Val(N)))
end

# TODO: type piracy; should generalize to any N
function (p::Polynomial)(x, ::Val{2})
    pd = SUP.derivative(p)
    pdd = SUP.derivative(pd)
    p(x), pd(x), pdd(x)
end

# JuMP utilities
function simplify(model::JuMP.Model, x::JuMP.QuadExpr)
    canonical = JuMP.MOIU.canonical # TODO
    QuadExpr(model, canonical(JuMP.moi_function.(x)))
end

val(x) = JuMP.value(x)
val(x::Number) = x

function bezier(model::JuMP.Model, namefunc, degree::Union{Integer, Val})
    BezierCurve(ntuple(j -> @variable(model, base_name=namefunc(j)), degree + 1))
end

function bezier_svec(model::JuMP.Model, namefunc, degree::Union{Integer, Val}, length::Union{Integer, Val})
    SVector(ntuple(i -> bezier(model, j -> namefunc(i, j), degree), length))
end

function constrain_poly_equal(model, x::Polynomial, y::Polynomial)
    @constraint model SVector(x.coeffs) .== SVector(y.coeffs)
end

function constrain_poly_equal(model, x::BezierCurve, y::BezierCurve)
    @constraint model SVector(x.coeffs) .== SVector(y.coeffs)
end

function constrain_poly_equal(model, x::Polynomial{N}, y::Number) where N
    constrain_poly_equal(model, x, Polynomial{N}(Polynomial(y)))
end

function axis_array_vars(model::Model, namefunc, axes...; lower_bound=nothing, upper_bound=nothing)
    vars = map(Iterators.product(map(axis -> axis.val, axes)...)) do I
        @variable model base_name=namefunc(I...)
    end
    for var in vars
        if lower_bound !== nothing
            set_lower_bound(var, lower_bound)
        end
        if upper_bound !== nothing
            set_upper_bound(var, upper_bound)
        end
    end
    AxisArray(vars, axes...)
end

function constrain_l1_norm(model::JuMP.Model, x::AbstractVector{<:JuMP.AbstractJuMPScalar}, value; add_bounds)
    @assert add_bounds
    s = @variable model [1 : length(x)] lower_bound=0 upper_bound=value
    @constraint model s .>=  x
    @constraint model s .>= -x
    @constraint model sum(s) <= value
    model
end

# Geometry utilities
function extrude(polyhedron::Polyhedra.HRepresentation, thickness::Number; zmax=0)
    A = polyhedron.A
    b = polyhedron.b
    n = size(A, 1)
    Ā = [A zeros(n, 1); 0 0 1; 0 0 -1]
    b̄ = [b; zmax; -zmax + thickness]
    hrep(Ā, b̄, polyhedron.linset)
end

function polyhedron_extrema(polyhedron::Polyhedra.Polyhedron)
    Polyhedra.hasallrays(polyhedron) && error("Polyhedron is not compact")
    N = Polyhedra.fulldim(polyhedron)
    T = Polyhedra.coefficient_type(polyhedron)
    bounding_box = HyperRectangle{N, T}(collect(Point{N, T}, Polyhedra.points(polyhedron)))
    minimum(bounding_box), maximum(bounding_box)
end

# scrap
# function poly(model::JuMP.Model, namefunc, degree::Union{Integer, Val})
#     Polynomial(ntuple(j -> @variable(model, base_name=namefunc(j)), degree + 1))
# end

# function poly_svec(model::JuMP.Model, namefunc, degree::Union{Integer, Val}, length::Union{Integer, Val})
#     SVector(ntuple(i -> poly(model, j -> namefunc(i, j), degree), length))
# end

# JuMP.value(x::Polynomial) = Polynomial(JuMP.value.(x.coeffs))
# JuMP.value(x::BezierCurve) = BezierCurve(JuMP.value.(x.coeffs))
# JuMP.value(x::Piecewise) = Piecewise(JuMP.value.(x.subfunctions), JuMP.value.(x.breaks), x.clamp)

function align_z_axis(rot::Rotation{3}, z_axis::StaticVector{3})
    rot_z_axis = rot * SVector(0, 0, 1)
    rotation_between(rot_z_axis, z_axis) * rot
end

function closest_pose(previous::AffineMap, origin::StaticVector{3}, z_axis::StaticVector{3})
    rot = align_z_axis(previous.linear, z_axis)
    AffineMap(rot, origin)
end


text:
"""
    loss functions
"""
function discriminator_loss(d_real::AbstractArray, d_g_fake_adv::AbstractArray)
    real_loss = mse(1f0, mean(d_real; dims = (1, 2)))
    fake_loss = mse(0f0, mean(d_g_fake_adv; dims = (1, 2)))
    return real_loss + fake_loss
end

generator_adv_loss(d_g_fake_adv::AbstractArray) = mse(1f0, mean(d_g_fake_adv; dims = (1, 2)))

generator_rec_loss(real_img::AbstractArray, g_fake_rec::AbstractArray) = mse(real_img, g_fake_rec)

"""
    update discriminator
"""
function update_discriminator!(opt, dscr, real_img, g_fake_adv)
    @eval Flux.istraining() = true
    ps = params(dscr)
    grad = gradient(ps) do
        discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

"""
    update generator
"""
# function update_generator_rec!(opt, gen, real_img, prev_rec, noise_rec, alpha)
#     @eval Flux.istraining() = true
#     ps = params(gen)
#     grad = gradient(ps) do
#         g_fake_rec = gen(prev_rec, noise_rec)
#         alpha * generator_rec_loss(real_img, g_fake_rec)
#     end
#     update!(opt, ps, grad)
#     @eval Flux.istraining() = false
# end

# function update_generator_adv!(opt, dscr, gen, prev_adv, noise_adv)
#     @eval Flux.istraining() = true
#     ps = params(gen)
#     grad = gradient(ps) do
#         d_g_fake_adv = dscr(gen(prev_adv, noise_adv))
#         generator_adv_loss(d_g_fake_adv)
#     end
#     update!(opt, ps, grad)
#     @eval Flux.istraining() = false
# end

function update_generator!(opt, dscr, gen, real_img, prev_rec, prev_adv, noise_rec, noise_adv, alpha)
    @eval Flux.istraining() = true
    ps = params(gen)
    grad = gradient(ps) do
        g_fake_rec = gen(prev_rec, noise_rec)
        d_g_fake_adv = dscr(gen(prev_adv, noise_adv))
        generator_adv_loss(d_g_fake_adv) + alpha * generator_rec_loss(real_img, g_fake_rec)
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

"""
    train
"""
function train_epoch!(opt_dscr, opt_gen, st, loop_dscr, loop_gen,
        dscr, genp, prev_rec, noise_rec, real_img, amplifiers, alpha)
    # training
    # prev_rec is only used to infer array type

    # discriminator
    foreach(1:loop_dscr) do _
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        g_fake_adv = genp(noise_adv, st, false)
        # add noise to real
        # real_noise = 0.5f0 * amplifiers[st] * randn!(similar(real_img))
        # update_discriminator!(opt_dscr, dscr, real_img + real_noise, g_fake_adv)

        update_discriminator!(opt_dscr, dscr, real_img, g_fake_adv)
    end

    # generator
    foreach(1:loop_gen) do _
        # adv
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        prev_adv = genp(noise_adv, st - 1, true)
        # update_generator_adv!(opt_gen, dscr, genp.chains[st], prev_adv, last(noise_adv))
        # rec
        # update_generator_rec!(opt_gen, genp.chains[st], real_img, prev_rec, noise_rec, alpha)
        update_generator!(opt_gen, dscr, genp.chains[st], real_img, prev_rec, prev_adv, noise_rec, last(noise_adv), alpha)
    end

    noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
    g_fake_adv = genp(noise_adv, st, false)
    loss_dscr = discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    d_g_fake_adv = dscr(g_fake_adv)
    loss_gen_adv = generator_adv_loss(d_g_fake_adv)
    g_fake_rec = genp.chains[st](prev_rec, noise_rec)
    loss_gen_rec = generator_rec_loss(real_img, g_fake_rec)

    return loss_dscr, loss_gen_adv, loss_gen_rec
end

function estimate_noise_amplifier(prev_rec::AbstractArray{Float32,4}, real_img::AbstractArray{Float32,4},
        pad::Integer, amplifier_init::Float32)
    prev_rec_crop = @view prev_rec[1 + pad:end - pad, 1 + pad:end - pad, :, :]
    rmse = sqrt(mse(real_img, prev_rec_crop))
    return rmse * amplifier_init
end

function train!(dscrp::DiscriminatorPyramid, genp::GeneratorPyramid,
        real_img_p::Vector{T}, hp::HyperParams) where {T <: AbstractArray{Float32,4}}
    stages = Base.length(genp.image_shapes)
    generate_dirs(stages)
    save_scaled_reals(real_img_p)

    amplifiers = Float32[]

    # fixed noise for rec
    fixed_noise_rec = build_rec_pyramid(first(real_img_p), genp.noise_shapes, 1f0)
    fixed_noise_adv = similar(fixed_noise_rec)

    for st in 1:stages
        @info "Step $(st)"
        # reset optimizer
        opt_dscr = ADAM(hp.lr_dscr, (0.5, 0.999))
        opt_gen = ADAM(hp.lr_gen, (0.5, 0.999))

        # calculate noise amplifier
        prev_rec = genp(fixed_noise_rec, st - 1, true) # padded
        amp = st == 1 ? 1f0 : estimate_noise_amplifier(prev_rec, real_img_p[st], genp.pad, hp.amplifier_init)
        push!(amplifiers, amp)
        # add noise for adv
        fixed_noise_adv[st] = amp * randn_like(prev_rec, expand_dim(genp.noise_shapes[st]...))

        save_noise_amplifiers(st, amp)
        @info "Noise amplifier = $(amp)"

        @time for ep in 1:hp.max_epoch
            # reduce learnint rate
            if ep == hp.reduce_lr_epoch
                @info "Reduce learning rate"
                opt_dscr.eta /= 10
                opt_gen.eta /= 10
            end

            loss_dscr, loss_gen_adv, loss_gen_rec =
                train_epoch!(opt_dscr, opt_gen, st, hp.loop_dscr, hp.loop_gen,
                    dscrp.chains[st], genp, prev_rec, fixed_noise_rec[st], real_img_p[st], amplifiers, hp.alpha)

            # save image/loss
            if ep == 1 || ep % hp.save_image_every_epoch == 0 || ep == hp.max_epoch
                save_generated_images(genp, fixed_noise_rec, fixed_noise_adv, st, ep)
            end

            if ep == 1 || ep % hp.save_loss_every_epoch == 0 || ep == hp.max_epoch
                @info "Epoch $(ep)" loss_dscr loss_gen_adv loss_gen_rec
                save_training_loss(st, ep, loss_dscr, loss_gen_adv, loss_gen_rec)
            end
        end

        save_model_params(dscrp, genp, st)
    end

    return amplifiers, first(fixed_noise_rec)
end

ext: jl

max_stars_repo_name: matsueushi/SinGAN

max_issues_repo_name: matsueushi/SinGAN

max_forks_repo_name: matsueushi/SinGAN


text:
bstr = [path * "/SampleFiles/Bottle/B024*",
        path * "/SampleFiles/Bottle/B2021901700*",
        path * "/SampleFiles/Bottle/B203*",
        path * "/SampleFiles/Bottle/B20319115Rainfallmm"]

printstyled("  Bottle (UNAVCO strain data)\n", color=:light_green)
for i = 1:4
  S = read_data("bottle", bstr[i], nx_new=14400, nx_add=14400)
  fill_pbo!(S)
  if i == 4
    @test S.name[1] == "quarry203bwa2007"
    @test isapprox(S.fs[1], 1.0/(30*60))
    @test isapprox(S.loc[1].el, 814.4)
    @test S.units[1] == "mm"
  end
end

redirect_stdout(out) do
  S = SeisIO.read_bottle(bstr[4], 0, 14400, 14400)
end


text:
"""
    ProductOperation <: AbstractGroupOperation

Direct product group operation.
"""
struct ProductOperation <: AbstractGroupOperation end

const ProductGroup{T} = GroupManifold{ProductManifold{T},ProductOperation}

"""
    ProductGroup{T} <: GroupManifold{ProductManifold{T},ProductOperation}

Decorate a product manifold with a [`ProductOperation`](@ref).

Each submanifold must also be an [`AbstractGroupManifold`](@ref) or a decorated instance of
one. This type is mostly useful for equipping the direct product of group manifolds with an
[`Identity`](@ref) element.

# Constructor
    ProductGroup(manifold::ProductManifold)
"""
function ProductGroup(manifold::ProductManifold)
    if !all(M -> (is_decorator_group(M) === Val(true)), manifold.manifolds)
        error("All submanifolds of product manifold must be or decorate groups.")
    end
    op = ProductOperation()
    return GroupManifold(manifold, op)
end

show(io::IO, G::ProductGroup) = print(io, "ProductGroup$((G.manifold.manifolds...,))")

function submanifold_component(
    e::Identity{GT},
    ::Val{I},
) where {I,MT<:ProductManifold,GT<:GroupManifold{MT}}
    M = base_manifold(e.group)
    return Identity(submanifold(M, I))
end

function submanifold_components(
    e::Identity{GT},
) where {MT<:ProductManifold,GT<:GroupManifold{MT}}
    M = base_manifold(e.group)
    return Identity.(M.manifolds)
end

inv(G::ProductGroup, x) = inv(G.manifold, x)
inv(::GT, e::Identity{GT}) where {GT<:ProductGroup} = e
function inv(M::ProductManifold, x::ProductRepr)
    return ProductRepr(map(inv, M.manifolds, submanifold_components(M, x))...)
end
function inv(M::ProductManifold, x)
    y = allocate_result(M, inv, x)
    return inv!(M, y, x)
end

inv!(G::ProductGroup, y, x) = inv!(G.manifold, y, x)
function inv!(M::ProductManifold, y, x)
    map(inv!, M.manifolds, submanifold_components(M, y), submanifold_components(M, x))
    return y
end

identity(G::ProductGroup, x) = identity(G.manifold, x)
identity(::GT, e::Identity{GT}) where {GT<:ProductGroup} = e
function identity(M::ProductManifold, x::ProductRepr)
    return ProductRepr(map(identity, M.manifolds, submanifold_components(M, x))...)
end
function identity(M::ProductManifold, x)
    y = allocate_result(M, identity, x)
    return identity!(M, y, x)
end

identity!(G::ProductGroup, y, x) = identity!(G.manifold, y, x)
function identity!(M::ProductManifold, y, x)
    map(identity!, M.manifolds, submanifold_components(M, y), submanifold_components(M, x))
    return y
end

compose(G::ProductGroup, x, y) = compose(G.manifold, x, y)
compose(G::GT, ::Identity{GT}, x) where {GT<:ProductGroup} = x
compose(G::GT, x, ::Identity{GT}) where {GT<:ProductGroup} = x
compose(G::GT, e::E, ::E) where {GT<:ProductGroup,E<:Identity{GT}} = e
function compose(M::ProductManifold, x::ProductRepr, y::ProductRepr)
    return ProductRepr(map(
        compose,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
    )...)
end
function compose(M::ProductManifold, x, y)
    z = allocate_result(M, compose, x, y)
    return compose!(M, z, x, y)
end

compose!(G::ProductGroup, z, x, y) = compose!(G.manifold, z, x, y)
function compose!(M::ProductManifold, z, x, y)
    map(
        compose!,
        M.manifolds,
        submanifold_components(M, z),
        submanifold_components(M, x),
        submanifold_components(M, y),
    )
    return z
end

translate(G::ProductGroup, x, y, conv::ActionDirection) = translate(G.manifold, x, y, conv)
function translate(
    M::ProductManifold,
    x::ProductRepr,
    y::ProductRepr,
    conv::ActionDirection,
)
    return ProductRepr(map(
        translate,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
        repeated(conv),
    )...)
end
function translate(M::ProductManifold, x, y, conv::ActionDirection)
    z = allocate_result(M, translate, x, y)
    return translate!(M, z, x, y, conv)
end

function translate!(G::ProductGroup, z, x, y, conv::ActionDirection)
    return translate!(G.manifold, z, x, y, conv)
end
function translate!(M::ProductManifold, z, x, y, conv::ActionDirection)
    map(
        translate!,
        M.manifolds,
        submanifold_components(M, z),
        submanifold_components(M, x),
        submanifold_components(M, y),
        repeated(conv),
    )
    return z
end

function inverse_translate(G::ProductGroup, x, y, conv::ActionDirection)
    return inverse_translate(G.manifold, x, y, conv)
end
function inverse_translate(
    M::ProductManifold,
    x::ProductRepr,
    y::ProductRepr,
    conv::ActionDirection,
)
    return ProductRepr(map(
        inverse_translate,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
        repeated(conv),
    )...)
end
function inverse_translate(M::ProductManifold, x, y, conv::ActionDirection)
    z = allocate_result(M, inverse_translate, x, y)
    return inverse_translate!(M, z, x, y, conv)
end

function inverse_translate!(G::ProductGroup, z, x, y, conv::ActionDirection)
    return inverse_translate!(G.manifold, z, x, y, conv)
end
function inverse_translate!(M::ProductManifold, z, x, y, conv::ActionDirection)
    map(
        inverse_translate!,
        M.manifolds,
        submanifold_components(M, z),
        submanifold_components(M, x),
        submanifold_components(M, y),
        repeated(conv),
    )
    return z
end

function translate_diff(G::ProductGroup, x, y, v, conv::ActionDirection)
    return translate_diff(G.manifold, x, y, v, conv)
end
function translate_diff(
    M::ProductManifold,
    x::ProductRepr,
    y::ProductRepr,
    v::ProductRepr,
    conv::ActionDirection,
)
    return ProductRepr(map(
        translate_diff,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
        submanifold_components(M, v),
        repeated(conv),
    )...)
end
function translate_diff(M::ProductManifold, x, y, v, conv::ActionDirection)
    vout = allocate_result(M, translate_diff, v, x, y)
    return translate_diff!(M, vout, x, y, v, conv)
end

function translate_diff!(G::ProductGroup, vout, x, y, v, conv::ActionDirection)
    return translate_diff!(G.manifold, vout, x, y, v, conv)
end
function translate_diff!(M::ProductManifold, vout, x, y, v, conv::ActionDirection)
    map(
        translate_diff!,
        M.manifolds,
        submanifold_components(vout),
        submanifold_components(M, x),
        submanifold_components(M, y),
        submanifold_components(M, v),
        repeated(conv),
    )
    return vout
end

function inverse_translate_diff(G::ProductGroup, x, y, v, conv::ActionDirection)
    return inverse_translate_diff(G.manifold, x, y, v, conv)
end
function inverse_translate_diff(
    M::ProductManifold,
    x::ProductRepr,
    y::ProductRepr,
    v::ProductRepr,
    conv::ActionDirection,
)
    return ProductRepr(map(
        inverse_translate_diff,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
        submanifold_components(M, v),
        repeated(conv),
    )...)
end
function inverse_translate_diff(M::ProductManifold, x, y, v, conv::ActionDirection)
    vout = allocate_result(M, inverse_translate_diff, v, x, y)
    return inverse_translate_diff!(M, vout, x, y, v, conv)
end

function inverse_translate_diff!(G::ProductGroup, vout, x, y, v, conv::ActionDirection)
    return inverse_translate_diff!(G.manifold, vout, x, y, v, conv)
end
function inverse_translate_diff!(M::ProductManifold, vout, x, y, v, conv::ActionDirection)
    map(
        inverse_translate_diff!,
        M.manifolds,
        submanifold_components(vout),
        submanifold_components(M, x),
        submanifold_components(M, y),
        submanifold_components(M, v),
        repeated(conv),
    )
    return vout
end


text:
using PSSFSS
using PSSFSS.RWG
using StaticArrays: SArray
using Test

@testset "fufptrue" begin
    sheet = rectstrip(Lx=1, Ly=1, Px=1, Py=1, Nx=2, Ny=1, units=mm, fufp=true)
    rwgdata = setup_rwg(sheet)
    @test rwgdata.bfe == [8  6  9  1  2  5
                          8  6  9  3  4  7]
    @test rwgdata.bff == [1  2  3  2  4  1
                          2  3  4  1  3  4]
    @test rwgdata.ebf == [4, 5, 4, 5, 6, 2, 6, 1, 3]
    @test rwgdata.eci == [3, 3, 4, 4, 1, 0, 2, 0, 0]
    @test rwgdata.ufpm == [1  5   9  11
                           2  6  10  12
                           3  7   1   5
                           4  8   2   6]
    @test rwgdata.ufp2fp == [[1, 11], [2, 12], [3], [4], [5, 15], [6, 16], [7],
                             [8], [9], [10], [13], [14]]
    @test rwgdata.nufp == 12
end

@testset "fufpfalse" begin
    sheet = rectstrip(Lx=1, Ly=1, Px=1, Py=1, Nx=2, Ny=1, units=mm, fufp=false)
    rwgdata = setup_rwg(sheet)
    @test rwgdata.bfe == [8  6  9  1  2  5
                          8  6  9  3  4  7]
    @test rwgdata.bff == [1  2  3  2  4  1
                          2  3  4  1  3  4]
    @test rwgdata.ebf == [4, 5, 4, 5, 6, 2, 6, 1, 3]
    @test rwgdata.eci == [3, 3, 4, 4, 1, 0, 2, 0, 0]
    @test rwgdata.ufpm == reshape(1:16, (4,4))
    @test rwgdata.ufp2fp == [[i] for i in 1:16]
    @test rwgdata.nufp == 16
end

@testset "rwgbfft!" begin
    sheet = rectstrip(Lx=1, Ly=1, Px=1, Py=1, Nx=2, Ny=1, units=mm, fufp=false)
    rwgdat = setup_rwg(sheet)
    nbf = size(rwgdat.bff,2)
    ft = zeros(SArray{Tuple{2},ComplexF64,1,2}, (nbf))
    k = [0.,0.]; ψ₁ = 0.0; ψ₂ = 0.0
    rwgbfft!(ft, rwgdat, sheet, k, ψ₁, ψ₂)
    @test ft[1] ≈ [0.0001666666666666667 + 0.0im, -0.0003333333333333334 + 0.0im]
    @test ft[2] ≈ [0.0003333333333333334 + 0.0im, 0.0003333333333333334 + 0.0im]
    @test ft[3] ≈ [0.00016666666666666663 + 0.0im, -0.0003333333333333334 + 0.0im]

    sheet = rectstrip(Lx=1, Ly=1, Px=1, Py=1, Nx=2, Ny=1, units=inch, fufp=false)
    rwgdat = setup_rwg(sheet)
    nbf = size(rwgdat.bff,2)
    ft = zeros(SArray{Tuple{2},ComplexF64,1,2}, (nbf))
    k = [20.0,-30.3]; ψ₁ = 0.8; ψ₂ = -0.7
    rwgbfft!(ft, rwgdat, sheet, k, ψ₁, ψ₂)
    @test ft[1] ≈ [0.003986416685586114 - 0.0010511300702857503im, -0.008079847138694855 + 0.0021304773084817956im]
    @test ft[2] ≈ [0.008263557134402535 - 0.0010871639073400372im, 0.008284527320353664 - 0.0010899227712197804im]
    @test ft[3] ≈ [0.004122638196063312 - 1.5707327530113332e-5im, -0.008355946971820058 + 3.1836312009155855e-5im]
end


text:
using MAT
using DataStructures
using StatsBase
f = readlines("Facebook_Sets.txt")
graphs = split(f[1])
include("../src/pdensity_algorithms.jl")

# insert appropriate path to datasets
pathtoFB100 = homedir()*"/data/Facebook100/"

Ps = [collect(0.25:0.25:2); 1.05]
for i = 1:100
    graph = graphs[i]
    F = matread(pathtoFB100*"$graph.mat")
    A = F["A"]
    n = size(A,1)
    # Ssets = zeros(n,length(Ps))
    Times = zeros(length(Ps))
    Objectives = zeros(length(Ps))
    Sizes = zeros(length(Ps))
    EdgeDens = zeros(length(Ps))
    MeanDegree = zeros(length(Ps))
    Mean2Degree = zeros(length(Ps))
    MaxDegree = zeros(length(Ps))
    for i = 1:length(Ps)
        p = Ps[i]
        tic = time()
        S, objS, Objs, ranking = GenPeel(A,p)
        timers = time()-tic
        Times[i] = timers
        # Ssets[S,i] .= 1
        Objectives[i] = maximum(Objs)
        Sizes[i] = length(S)

        As = A[S,S]
        ds = vec(sum(As,dims = 2))
        edges = sum(As)/2
        posedges = binomial(length(S),2)
        EdgeDens[i] = edges/posedges
        MeanDegree[i] = StatsBase.mean(ds)
        Mean2Degree[i] = StatsBase.mean(ds.^2)
        mdeg = StatsBase.mean(ds)
        MaxDegree[i] = maximum(ds)

        println("$graph \t $p \t $(length(S)) \t $objS \t $timers \t $(mdeg)")
    end

    matwrite("fb_output/$(graph)_output.mat",Dict( #"Ssets"=>Ssets,
    "Ps"=>Ps,"Times"=>Times,"Objectives"=>Objectives,
    "Sizes"=>Sizes,"MeanDegree"=>MeanDegree,"Mean2Degree"=>Mean2Degree,
    "MaxDegree"=>MaxDegree,"EdgeDens"=>EdgeDens))

end

max_stars_repo_name: nveldt/GenMeanDSG

max_issues_repo_name: nveldt/GenMeanDSG

max_forks_repo_name: nveldt/GenMeanDSG


text:
# This implements the Adaptive Iterative Tresholding algorithm for
# compressed sensing / regularized regression. This implementation
# is based on the paper:
#  Yu Wang, Jinshan Zeng, Zhimin Peng, Xiangyu Chang, Zongben Xu,
#  "On Linear Convergence of Adaptively Iterative Thresholding
#  Algorithms for Compressed Sensing", arxiv, 2014, http://arxiv.org/abs/1408.6890
#
function adaptive_iterative_tresholding(A, b; kws...)
  M, N, k, xprev, x, An, acolnorms, sAnt, f, maxiterations, tolerance = unpack_ait_parameters(A, b; kws...)

  # Since we need one iteration before we can judge convergence
  # we check for convergence last in the loop.
  t = 0
  while true
    t += 1
    z = xprev .+ sAnt * (b .- An * xprev)
    zp = sortperm(abs(z), rev = true)
    tau = abs(z[zp[k+1]])

    # Treshold x
    for i in 1:N
      ip = zp[i]
      if i <= k
        x[ip] = f(tau, z[ip])
      else
        x[ip] = 0.0
      end
    end

    # Check convergence
    diff = norm(x .- xprev)
    if t >= maxiterations || diff < tolerance
      return (x ./ acolnorms, t)
    end

    # Switch their meaning => we need not allocate new memory in each iteration
    xprev, x = x, xprev
  end
end

function unpack_ait_parameters(A, b;
  k = -1,             # Sparsity level, this should be set as an upper bound on the number of non-zero coefficients of x. If set below the actual sparseness level the algorithm might not find right solution.
  s = 1.0,            # Step length, should be ~ in range [0.4, 1.3] but 1.0 seem to be generally good
  x0 = nothing,       # A starting point, will be randomized if not given
  f::Function = scad, # Defining function used in tresholding, by default SCAD
  tolerance = 1e-6,   # Convergence level, i.e. terminate algorithm if norm change is less than this between iterations
  maxIterations = 1e3 # Maximum number of iterations allowed. Often it converges in ~20-30 iterations if ever.
  )

  # Check and then initialize all parameters. Note that we use N for the num of covariates and
  # M for the number of cases. This is in line with the paper above but it is more common to use
  # N = M and P = N.
  M, N = size(A)
  @assert length(b) == M
  k = (k == -1) ? ifloor(sqrt(N)) : k
  @assert 1 <= k < N
  xprev = (x0 == nothing) ? zeros(N) : copy(x0)
  x = zeros(Float64, N)

  An, acolnorms = normalize_columns(A)
  sAnt = s * An'

  return (M, N, k, xprev, x, An, acolnorms, sAnt, f, maxIterations, tolerance)
end

function normalize_columns(A)
  acolnorms = map((j) -> norm(A[:,j]), 1:size(A, 2))
  An = broadcast(/, A, acolnorms')
  return (An, acolnorms)
end

# Inner AIT iteration loop. Assumes An has been normalized.
function ait_iteration_loop!(An, b, k, N, M, sAnt, tolerance, maxiterations, xprev, trainindices, Ntrain, f)
  t = 0
  x = zeros(Float64, N)

  # For now, later we will loop over the trainindices
  btrain = b[trainindices]
  Antrain = An[trainindices, :]
  sAnttrain = sAnt[:, trainindices]

  while true
    t += 1
    z = xprev .+ sAnttrain * (btrain .- Antrain * xprev)
    zp = sortperm(abs(z), rev = true)
    tau = abs(z[zp[k+1]])

    # Treshold x
    for i in 1:N
      ip = zp[i]
      if i <= k
        x[ip]::Float64 = f(tau::Float64, z[ip]::Float64)
      else
        x[ip]::Float64 = 0.0
      end
    end

    diff = norm(x .- xprev)
    if t >= maxiterations || diff < tolerance
      return (x, t)
    end

    xprev, x = x, xprev
  end
end

function ait2(A, b; kws...)
  M, N, k, xprev, x, An, acolnorms, sAnt, f, maxiterations, tolerance = unpack_ait_parameters(A, b; kws...)
  x, its = ait_iteration_loop!(An, b, k, N, M, sAnt, tolerance, maxiterations, xprev, collect(1:M), M, f)
  return (x ./ acolnorms, its)
end

function scad_penalty_tresholding(a, tau, u)
  absu = abs(u)
  if absu > a * tau
    u
  elseif absu > 2 * tau
    ((a-1)*u - sign(u)*a*tau)/(a-2)
  elseif absu > tau
    absu - sign(u)*tau
  else
    0.0
  end
end

scad(tau, u) = scad_penalty_tresholding(3.7, tau, u)
hard(tau, u) = abs(u) > tau ? u : 0.0

splitindices(n, ratio = 0.10) = begin
  randindices = shuffle(collect(1:n))
  numtest = ifloor(ratio*n)
  test = randindices[1:numtest]
  train = randindices[(numtest+1):end]
  return (train, test)
end

# Run AIt for multiple k-values from 2*sqrt(P) down to 1
# and then select the one with lowest MSE on a TestRatio
# holdout test set.
function ait_best_k(A, b, TestRatio = 0.20; kws...)
  N, P = size(A)

  # Split in train and test sets
  trainidxs, testidxs = splitindices(N, TestRatio)
  Atrain = A[trainidxs,:]
  btrain = b[trainidxs]
  Atest = A[testidxs,:]
  btest = b[testidxs]

  maxP = ifloor(2*sqrt(P))
  if maxP >= P
    maxP = P - 1
  end

  local bestcoefs = bestits = bestmse = Inf

  mses = Dict{Int64,Float64}()

  for k in maxP:-1:1
    coefs, its = adaptive_iterative_tresholding(Atrain, btrain; k = k, kws...)
    mses[k] = mse(btest, Atest * coefs)
    if mses[k] < bestmse
      bestcoefs = coefs
      bestits = its
      bestmse = mses[k]
    end
  end

  return bestcoefs, bestits, mses
end

# using PyPlot
# x = linspace(-5.0, 5.0, 200)
# y = map((x) -> scad_penalty_tresholding(3.7, 1.0, x), x)
# plot(x, y, color="red", linewidth=2.0, linestyle="--")

mse(y,yhat) = mean( (y .- yhat).^2 )

# N-fold cross-validation to select the sparseness level for AIT.
function nfoldCV_ait(A, b)
  M, N = size(A)

  num_ks = 5
  ks = ifloor(exp(log(N-1) / num_ks * reverse(1:num_ks)))
  xs = zeros(num_ks, N)
  x0 = zeros(N)
  its = zeros(Int64, num_ks)

  for i in 1:num_ks
    x0, its[i] = adaptive_iterative_tresholding(A, b; k=k, x0 = x0)
    xs[i,:] = x
  end
end

function ait_mse(A, b, k)
  betas, its = adaptive_iterative_tresholding(A, b; k = k)
  return (betas, mse(b, A * betas))
end

function adaptive_k_ait(A, b; kws...)
  N, P = size(A)

  mses = Dict{Int64, Float64}()
  betas = Dict{Int64, Vector{Float64}}()
  mse_for_k(k) = begin
    if !haskey(mses, k)
      betas[k], mses[k] = ait_mse(A, b, k)
    end
    mses[k]
  end

  binsearch(lowk, highk) = begin
    mselow = mse_for_k(lowk)
    msehigh = mse_for_k(highk)
    if mselow > 1e3 * msehigh
      if highk == lowk + 1
        return (highk, msehigh)
      end
      midk = lowk + ifloor((highk-lowk)/2)
      k, mseb = binsearch(lowk, midk)
      if k == -1
        return binsearch(midk, highk)
      else
        return (k, mseb)
      end
    else
      return (-1, nothing)
    end
  end

  k, mse = binsearch(1, ifloor(min(P/2, 2*sqrt(P))))

  return (k, mse, betas[k], mses)
end

function test_ait_on_random_problem(N, P, kstar)
  A = randn(N, P)
  nzw = map((i) -> rand(10.0:100.0) * (rand() < 0.5 ? 1.0 : (-1.0)), 1:kstar)
  betas = shuffle([nzw, zeros(P-kstar)])
  b = A * betas

  k, mse, betasp, mses = adaptive_k_ait(A, b)

  show(mses)
  println()

  return (N, P, kstar, k, mse, coeffquality(betas, betasp))
end

#for kstar in [3, 4, 5, 6]
#  for P in [10, 20, 50, 100]
#    println(test_ait_on_random_problem(100, P, kstar), "\n")
#  end
#end

# N, P = (rand(20:100), 20)
# kstar = ifloor(sqrt(P)) - 1
# A = randn(N, P)
# nzw = rand(1.0:10.0) * rand(-1.0:0.01:1.0, kstar)
# betas = shuffle([nzw, zeros(P-kstar)])
# b = A * betas
# k = -1.0
# s = 1.0
# x0 = nothing
# f = hard
# tolerance = 1e-5
# maxIterations = 1e4
#
# # Ensure compiled
# adaptive_iterative_tresholding(A[1:10,:], b[1:10])
# ait2(A[1:10,:], b[1:10])
#
# tic()
# betasp, its = adaptive_iterative_tresholding(A, b)
# t1 = toq()
# tic()
# betasp2, its2 = ait2(A, b)
# t2 = toq()
# nzwp = betasp[abs(betasp) .> 0.0]
# bhat = X' * betasp
# bhat2 = X' * betasp2
#
#
# show( [betas betasp betasp2] )
# println("\nN = ", N, ", P = ", P)
# println("MSE1: ", mse(b, bhat))
# println("MSE2: ", mse(b, bhat2))
# println("b1: ", coeffquality(betas, betasp))
# println("b2: ", coeffquality(betas, betasp2))
# println("Timediff: ", t1-t2)


text:
# using LinearAlgebraExtensions: vecofvec
vecofvec(A::AbstractMatrix) = [a for a in eachcol(A)] # one view allocates 64 bytes

# IDEA: use LazyArrays to write efficient gramian for VerticalRescaling and
# InputScaling kernels
############################ Lazy Kernel Matrix ###############################
# note: gramian specializations for special matrix structure has to be after definition of all kernels
# move to MyLazyArrays?
# K can be any kernel, so not necessarily MercerKernel
struct Gramian{T, K, U<:AbstractVector, V<:AbstractVector} <: AbstractMatrix{T}
    k::K # has to allow k(x[i], y[j]) evaluation ∀i,j
    x::U
    y::V
end

const BlockGramian{T, M, K, X, Y} = BlockFactorization{<:T, <:Gramian{M, K, X, Y}}

function Gramian(k, x::AbstractVector, y::AbstractVector = x)
    T = gramian_eltype(k, x[1], y[1])
    Gramian{T, typeof(k), typeof(x), typeof(y)}(k, x, y)
end
# defaults to euclidean dot product
Gramian(x::AbstractVector, y::AbstractVector) = Gramian(Dot(), x, y)

Base.size(K::Gramian) = (length(K.x), length(K.y))
Base.eltype(G::Gramian{T}) where {T} = T
# size of an element of a matrix of matrices
elsize(G::Gramian) = size(G[1, 1])
elsize(G::Gramian{<:Number}) = ()
function gramian_eltype(k::AbstractKernel, x, y)
    promote_type(eltype(k), eltype(eltype(x)), eltype(eltype(y)))
end
gramian_eltype(k, x, y) = typeof(k(x, y)) # default to evaluation

# indexing
# NOTE: @inline helps increase mvm performance by 50%
@inline function Base.getindex(G::Gramian, i::Integer, j::Integer)
    @boundscheck checkbounds(G, i, j) # add bounds check to G
    @inbounds G.k(G.x[i], G.y[j]) # remove boundscheck of x of x and y
end
function Base.getindex(G::Gramian, i, j::Integer)
    @boundscheck checkbounds(G, i, j) # add bounds check to G
    @inbounds @views G.k.(G.x[i], (G.y[j],))
end
function Base.getindex(G::Gramian, i::Integer, j)
    @boundscheck checkbounds(G, i, j) # add bounds check to G
    @inbounds @views G.k.((G.x[i],), G.y[j])
end
function Base.getindex(G::Gramian, i, j)
    @boundscheck checkbounds(G, i, j) # add bounds check to G
    @inbounds @views Matrix(gramian(G.k, G.x[i], G.y[j]))
end

# maintain laziness by default when adding diagonal
function Base.:+(D::Union{Diagonal, BlockDiagonalFactorization}, G::Union{Gramian, BlockGramian})
    LazyMatrixSum(D, G)
end
function Base.:+(G::Union{Gramian, BlockGramian}, D::Union{Diagonal, BlockDiagonalFactorization})
    LazyMatrixSum(G, D)
end

# IDEA: GPU
using LinearAlgebra: checksquare
using Base.Threads
# make generic multiply multi-threaded and SIMD-enabled
function Base.:*(G::Gramian, a::AbstractVector)
    T = promote_type(eltype(G), eltype(a))
    b = zeros(T, size(G, 1))
    mul!(b, G, a)
end
function Base.:*(G::Gramian, A::AbstractMatrix)
    T = promote_type(eltype(G), eltype(A))
    B = zeros(T, size(G, 1), size(A, 2))
    mul!(B, G, A)
end
# NOTE: these specialized multiplication functions scale better for large n
# but tend to be slightly slower than Julia's generic mul! for small n, likely because of parallelization
function LinearAlgebra.mul!(y::AbstractVector, G::Gramian, x::AbstractVector, α::Real = 1, β::Real = 0)
    n, m = size(G)
    @. y = iszero(β) ? 0 : β * y
    @threads for i in 1:n
        @simd for j in 1:m
            @inbounds y[i] += α * G[i, j] * x[j]
        end
    end
    return y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, G::Gramian, X::AbstractMatrix, α::Real = 1, β::Real = 0)
    @. Y = iszero(β) ? 0 : β * Y # makes sure returns correctly if β = 0 and y contains NaNs, e.g. through `similar` initialization
    @threads for j in 1:size(Y, 2)
        for i in 1:size(Y, 1)
            @simd for k in 1:size(X, 1) # indexing order so that dense arrays Y, X are indexed in a Cache-friendly way
                @inbounds Y[i, j] += α * G[i, k] * X[k, j]
            end
        end
    end
    return Y
end

# parallel matrix instantiation
function Base.Matrix(G::Gramian)
    n, m = size(G)
    M = Matrix{eltype(G)}(undef, n, m)
    Matrix!(M, G)
end
function Matrix!(M::AbstractMatrix, G)
    @threads for j in 1:size(G, 2)
        @simd for i in 1:size(G, 1)
            @inbounds M[i, j] = G[i, j]
        end
    end
    return M
end
Base.AbstractMatrix(G::Gramian) = Matrix(G)
Base.adjoint(G::Gramian) = Gramian(G.k, G.y, G.x)
Base.transpose(G::Gramian) = Gramian(G.k, G.y, G.x)

# by default, Gramians of matrix-valued kernels are BlockFactorizations, O(1) memory complexity
function gramian(k::MultiKernel, x::AbstractVector, y::AbstractVector, lazy::Val{true} = Val(true))
    G = Gramian(k, x, y)
    BlockFactorization(G, isstrided = true) # strided because every block has same size
end

# instantiates the blocks but respects structure O(n^2d) memory complexity for gradient kernel
function gramian(k::MultiKernel, x::AbstractVector, y::AbstractVector, lazy::Val{false}) # = Val(false))
    G = Gramian(k, x, y)
    G = Matrix(G)
    BlockFactorization(G, isstrided = true)
end

LinearAlgebra.issymmetric(G::Gramian) = (G.x ≡ G.y) || (G.x == G.y) # pointer check is constant time
LinearAlgebra.issymmetric(G::Gramian{<:Real, <:Constant}) = true
LinearAlgebra.ishermitian(G::Gramian) = issymmetric(G)
function LinearAlgebra.isposdef(G::Gramian)
    return typeof(G.k) <: Union{MercerKernel, MultiKernel} && issymmetric(G)
end

######################### smart pseudo-constructor #############################
# standard approach is a lazily represented kernel matrix
# by default, Gramians of matrix-valued kernels are BlockFactorizations
# TODO: don't return BlockFactorization if one dimension is one!
# elsize(G)[1]  == 1 && size(G)[1] == 1
function gramian(k, x::AbstractVector, y::AbstractVector)
    G = Gramian(k, x, y)
    eltype(G) <: AbstractMatOrFac ? BlockFactorization(G) : G
end
gramian(k, x::AbstractVector) = gramian(k, x, x)

gramian(x::AbstractVector, y::AbstractVector) = Gramian(x, y)
gramian(x::AbstractVector) = gramian(x, x)

# if matrix whose columns are datapoints is passed, convert to vector of vectors
gramian(k, x::AbstractMatrix) = gramian(k, vecofvec(x))
gramian(k, x::AbstractMatrix, y::AbstractMatrix) = gramian(k, vecofvec(x), vecofvec(y))

# if input trait is not used, ignore it
gramian(k, x, y, ::InputTrait) = gramian(k, x, y)
gramian(k, x, ::InputTrait) = gramian(k, x)

# IDEA: gramian(k, x, y, input = input_trait(k))
# makes it possible to take advantage of specialized implementations
# and custom kernels k, by defining input_trait(k)

# 1D stationary kernel on equi-spaced grid yields Toeplitz structure
# works if input_trait(k) <: Union{IsotropicInput, StationaryInput}
gramian(k, x::StepRangeLen, y::StepRangeLen) = gramian(k, x, y, input_trait(k))
gramian(k, x::StepRangeLen, y::StepRangeLen, ::GenericInput) = Gramian(k, x, y)

# IDEA: should this be in factorization? since dft still costs linear amount of information
# while gramian is usually lazy and O(1) in structure construction
function gramian(k, x::StepRangeLen, y::StepRangeLen, ::Union{IsotropicInput, StationaryInput})
    if x === y
        k1 = k.(x[1], x)
        SymmetricToeplitz(k1)
    elseif x.step == y.step
        k1 = k.(x, y[1])
        k2 = k.(x[1], y)
        Toeplitz(k1, k2)
    else
        Gramian(k, x, y)
    end
end

# 1D stationary kernel on equi-spaced grid with periodic boundary conditions
function gramian(k::StationaryKernel, x::StepRangeLen, ::PeriodicInput)
    k1 = k.(x[1], x)
    Circulant(k1)
end

###################### factorization of Gramian ################################
# IDEA: have to use special cholesky implementation to avoid instantiating G in low rank case
function LinearAlgebra.cholesky(G::Gramian, pivoting::Val{true}; check::Bool = false, tol::Real = 0.0)
    cholesky!(Symmetric(Matrix(G)), Val(true), check = check, tol = tol)
end

function LinearAlgebra.cholesky(G::Gramian, pivoting::Val{false} = Val(false); check::Bool = true)
    cholesky!(Symmetric(Matrix(G)), Val(false), check = check)
end

const DEFAULT_MAX_CHOLESKY_SIZE = 2^14
const DEFAULT_TOL = 1e-6

# TODO: add sparsification technique for exponentially decaying kernels in high d
function LinearAlgebra.factorize(G::Gramian; max_cholesky_size::Int = DEFAULT_MAX_CHOLESKY_SIZE,
                                             tol::Real = DEFAULT_TOL)
    n = checksquare(G)
    if n ≤ max_cholesky_size # defaults to pivoted cholesky to detect low rank structure
        cholesky(G, Val(true), check = false, tol = tol)
    else # IDEA: calculate preconditioner here?
        G # if instantiating G is not possible in memory, keep lazy and proceed via cg,
    end
end

# abstract type InferenceType end
# struct CholeskyInference <: InferenceType end
# struct ConjugateGradientInference <: InferenceType end
# struct SparseInference <: InferenceTypeend
#
# function LinearAlgebra.factorize(G::Gramian, ::CholeskyInference)
#     cholesky(G)
# end
#
# function LinearAlgebra.factorize(G::Gramian, ::SparseInference)
#     return -1
# end

###################### Specializations for BlockGramians #######################
function LinearAlgebra.:\(B::BlockGramian, b::AbstractVector)
    T = promote_type(eltype(B), eltype(b))
    x = zeros(T, size(B, 1))
    ldiv!(x, B, b)
end
# solve general BlockGramian via minimum residual solver
function LinearAlgebra.ldiv!(x::AbstractVector, B::BlockGramian, b::AbstractVector; kwargs...)
    minres!(x, B, b; kwargs...)
end

# carries out multiplication for general BlockFactorization
function BlockFactorizations.blockmul!(y::AbstractVecOfVecOrMat, G::Gramian, x::AbstractVecOfVecOrMat, α::Real = 1, β::Real = 0)
    Gijs = [G[1, 1] for _ in 1:Base.Threads.nthreads()] # pre-allocate storage for elements
    IT = input_trait(G.k)
    @threads for i in eachindex(y) # NOTE: without threading, this does not allocate anything
        @. y[i] = iszero(β) ? 0 : β * y[i] # makes sure returns correctly if β = 0 and y contains NaNs, e.g. through `similar` initialization
        Gij = Gijs[Base.Threads.threadid()]
        for j in eachindex(x)
            Gij = evaluate_block!(Gij, G, i, j, IT) # this is change to original blockmul!, allowing evaluation of Gramian's elements without additional allocations
            mul!(y[i], Gij, x[j], α, 1)
        end
    end
    return y
end

function LinearAlgebra.mul!(y::AbstractVecOfVecOrMat, G::Gramian, x::AbstractVecOfVecOrMat, α::Real = 1, β::Real = 0)
    BlockFactorizations.blockmul!(y, G, x, α, β)
end

# IDEA: @inline?
function evaluate_block!(Gij, G::Gramian, i::Int, j::Int, T = input_trait(G.k))
    evaluate_block!(Gij, G.k, G.x[i], G.y[j], T)
end

# if target is number, just evaluate the gramian
function evaluate_block!(Gij::Number, G::Gramian, i::Int, j::Int, T = input_trait(G.k))
    G[i, j]
end

# fallback
function evaluate_block!(Gij, k, x, y, T = input_trait(k))
    k(x, y)
end


text:
#
# Implementation of PRS similar to lassosum
#
# Mak et al.(2018), Polygenic scores via penalized regression on
# summary statistics
#

export z2cor,
    elnetg!,
    elnetg_path

"""
    z2cor(z, n)

Convert GWAS z-value to phenotype-genotype correlation coefficient.
"""
z2cor(z, n) = z ./ sqrt.(n .- 1 .+ z.^2)

"""
    elnetg!(X, r, λ, s, β = zeros(Float64, length(r));
            maxiter = 10_000, thresh = 1e-4)

Solve elastic net via coordinate descent.

# Arguments

- `X::Matrix` nsubj x nmarkers normalized genotype matrix
              need to pass in ac*sqrt((1-s)/nsubj),
              where ac is column normalized (μ=0, σ=1) allele count matrix
- `r::Vector` nmarkers x 1 vector of correlation between phenotype and genotypes
- `λ::Real`   shrinkage parameter for 1-norm
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnetg!(X, r, λ, s, β;
                 maxiter = 10_000, thresh = 1e-4,
                 verbose = false)

    p = length(r) # number of markers
    @assert p == size(X, 2) "Size mismatch: length(r) = $p, size(X, 2) = $(size(X, 2))"

    sd = std(X; dims = 1)[:]
    diag = ones(Float64, p) * (1.0 - s)
    diag[sd .== 0.0] .= 0.0

    yhat = X * β

    converged = false

    for k = 1:maxiter
        Δ = 0.0
        for j = 1:p
            β_j = β[j]
            β[j] = 0.0
            u_j = diag[j]*β_j + r[j]  - yhat'*X[:,j]

            abs_u_j = abs(u_j)
            #println("j=$j, abs_u_j = $abs_u_j")
            if abs_u_j > λ
                β[j] = sign(u_j) * (abs_u_j - λ) / (diag[j] + s)
            end
            Δ_j = β[j] - β_j
            Δ = max(Δ, abs(Δ_j))
            yhat += Δ_j * X[:,j]
        end

        if Δ < thresh
            converged = true
            break
        end

        if verbose
            print("\rλ = $λ, s = $s: Δ = $Δ  ($k/$maxiter)")
        end

    end
    obj = yhat'*yhat - 2*β'*r + 2*λ*sum(β) + s*β'*β

    if verbose
        println("\rλ = $λ, s = $s: obj = $obj, converged = $converged")
    end

    (converged = converged, obj = obj)
end

"""
    elnetg_path(X, r, λs, s;
                maxiter = 10_000, thresh = 1e-4)

Solve elastic net for path along `λs` using warm starts.

# Arguments

- `X::Matrix` nsubj x nmarkers normalized genotype matrix
              column normalized (μ = 0, σ = 1)
- `r::Vector` nmarkers x 1 vector of correlation coefficients
              between phenotype and genotypes
- `λs::Vector` shrinkage parameter path for 1-norm in decreasing order
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnetg_path(X, r, λs, s;
                     maxiter = 1000, thresh = 1e-4, verbose = false)

    nl = length(λs)
    p = length(r)

    β = zeros(Float64, p, nl)
    converged = zeros(Bool, nl)

    nsubj = size(X, 1)
    Xs = X .* sqrt((1.0 - s)/nsubj)
    ret = elnetg!(Xs, r, λs[1], s, view(β, :, 1);
                  maxiter = maxiter, thresh = thresh, verbose = verbose)
    converged[1] = ret.converged

    for l = 2:nl
        @assert λs[l] < λs[l-1] "λs not in decreasing order"
        β[:, l] .= β[:, l-1] # use warm start
        ret = elnetg!(Xs, r, λs[l], s, view(β, :, l);
                     maxiter = maxiter, thresh = thresh, verbose = verbose)
        converged[l] = ret.converged

    end

    (β = β,
     λs = λs,
     converged = converged,
     s = s)

end

#
# Implementations using sparse LD matrix instead of genotype matrix.
# These have convergence problems, though.
#

"""
    elnet(R, r, λ, s, β = zeros(Float64, length(r));
          maxiter = 10_000, thresh = 1e-4)

Solve elastic net via coordinate descent.

# Arguments

- `R::Matrix` p x p LD matrix
- `r::Vector` p x 1 vector of correlation between phenotype and genotypes
- `λ::Real`   shrinkage parameter for 1-norm
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnet!(R, r, λ, s, β;
                maxiter = 10_000, thresh = 1e-4,
                verbose = false)

    p = length(r)
    @assert p == size(β, 1) "Size mismatch: length(r) = $p, size(β, 1) = $(size(β, 1))"

    converged = false

    for k = 1:maxiter
        Δ = 0.0   # maximum change in β
        for j = 1:p
            β_j = β[j]
            β[j] = 0.0
            u_j = s ≈ 1 ? r[j] : (r[j] - (1.0-s) * β' * R[:,j])
            abs_u_j = abs(u_j)
            if abs_u_j > λ
                β[j] = sign(u_j) * (abs_u_j - λ) / (s + (1-s)*R[j,j])
            end
            Δ_j = β[j] - β_j
            Δ = max(Δ, abs(Δ_j))
        end

        if Δ < thresh
            converged = true
            break
        end

        if verbose
            print("\rλ = $λ, s = $s: Δ = $Δ  ($k/$maxiter)")
        end

    end

    converged
end

"""
    elnet_path(R, r, λs, s;
               maxiter = 10_000, thresh = 1e-4)

Solve elastic net for path along λs using warm starts.

HAS CONVERGENCE PROBLEMS. USE elnetg_path INSTEAD

# Arguments

- `R::Matrix` p x p LD matrix
- `r::Vector` p x 1 vector of correlation coefficients
              between phenotype and genotypes
- `λs::Vector` shrinkage parameter path for 1-norm in decreasing order
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnet_path(R, r, λs, s;
                    maxiter = 10_000, thresh = 1e-4)

    nl = length(λs)
    p = length(r)

    β = zeros(Float64, p, nl)
    converged = zeros(Bool, nl)

    converged[1] = elnet!(R, r, λs[1], s, view(β, :, 1);
                         maxiter = maxiter, thresh = thresh)

    for l = 2:nl
        @assert λs[l] < λs[l-1] "λs not in decreasing order"
        β[:, l] .= β[:, l-1] # use warm start
        converged[l] = elnet!(R, r, λs[l], s, view(β, :, l);
                              maxiter = maxiter, thresh = thresh)
    end

    (β = β,
     λs = λs,
     converged = converged,
     s = s)

end


text:
function group_files(freedim, dir;
        filter = (filename) -> true
    )
    freedim = string(freedim)
    files = Dict()
    for filename in readdir(dir) |> sort
        !filter(filename) && continue

        ret = tryparse_dfname(filename)
        isnothing(ret) && continue
        head, params, ext = ret

        !haskey(params, freedim) && continue
        freedim_val = params[freedim]
        delete!(params, freedim)

        f = get!(files, (freedim, head, params, ext), Dict())
        f[freedim_val] = joinpath(dir, filename)
    end
    files
end


text:
## Julia program for Slice sampler
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-22

function rnorm_slice(T)
    x = ones(T+1)
    w = ones(T+1)
    for t = 1:T
        w[t+1] = rand() * exp(-1.0 * x[t]^2/2)
        x[t+1] = rand() * 2 * sqrt(-2*log(w[t+1])) - sqrt(-2*log(w[t+1]))
    end
    return x[2:end]
end

## example
rnorm_slice(100)

max_stars_repo_name: szcf-weiya/MCMC

max_issues_repo_name: szcf-weiya/MCMC

max_forks_repo_name: szcf-weiya/MCMC


text:
export SoftmaxLoss, getSoftMaxLoss, getLabels, getMisfit

"""
sofmax loss function
"""
mutable struct SoftmaxLoss{T}
    shift :: T
    theta :: T
    addBias::Bool
end

function getSoftMaxLoss(TYPE::Type,shift=zero(TYPE), theta=convert(TYPE,1e-3), addBias=true)
	return SoftmaxLoss{TYPE}(shift,theta,addBias);
end

import Base.display
display(this::SoftmaxLoss) = println("SoftmaxLoss(shift=$(this.shift),theta=$(this.theta),addBias=$(this.addBias))")

function getMisfit(this::SoftmaxLoss{T},W::Array{T},Y::Array{T},C::Array{T},doDY=true,doDW=true) where {T}

    dWF = zeros(T,0); d2WF = zeros(T,0); dYF =zeros(T,0); d2YF = zeros(T,0)

    szY  = size(Y)
    nex  = szY[2]
    if this.addBias
        Y   = [Y; ones(T,1,nex)]
    else
        Y   = copy(Y)  # make sure we don't overwrite features
    end

    szW  = [size(C,1),size(Y,1)]

    W   = reshape(W,tuple(szW...))
    Y   -= this.shift;

    S   = exp.(W*Y)

    Cp   = getLabels(this,S)
    err  = countnz(C-Cp)/2
    F    = -sum(sum(C.*(W*Y))) + sum(log.(sum(S,1)))
    para = [F,nex,err]
    F   /= nex

    if doDW || doDY
        dF   = -C + S./sum(S,1)
        d2F  = U -> this.theta *U + (U.*S)./sum(S,1) -
                    S.*(repmat(sum(S.*U,1)./sum(S,1).^2,size(S,1),1))
    end
    if doDW
        dWF  = vec(dF*(Y'/nex))
        matW      = W -> reshape(W,tuple(szW...))
        d2WFmv    = U -> vec((d2F(matW(U/nex)*Y)*Y'))
        d2WF      = LinearOperator(prod(szW),prod(szW),true,true,d2WFmv)
    end
    if doDY
        if this.addBias
            W = W[:,1:end-1]
        end
        dYF  =   vec(W'*dF)/nex

        WI     = T -> W*T
        WIT    = T -> W'*T
        matY   = Y -> reshape(Y,szY)
        d2YFmv = T -> vec(WIT(((d2F(WI(matY(T/nex)))))))

        d2YF = LinearOperator(prod(szY),prod(szY),true,true,d2YFmv)
    end
    return F, para, dWF, d2WF, dYF, d2YF
end

function getLabels(this::SoftmaxLoss,W::Array{T},Y=nothing) where {T}
    if Y==nothing
        S = W
        nex = size(S,2)
    else
        ny = size(Y); nf = ny[1]; nex = ny[2]
        if this.addBias
            W      = reshape(W,:,nf+1)
            Y     = [Y; ones(T,1,nex)]
        end
        Y     = Y - this.shift
        S      = exp.(W*Y)
    end
    P      = S./sum(S,1)
    Cp     = zeros(T,size(P))
    for i=1:size(P,2)
        Cp[indmax(P[:,i]),i] = 1
    end
    return Cp
end


text:
"""
This uses a combination of Schumaker Splines to cover a 2 dimensional space.
We have a grid of splines. We first evaluate in one dimension (which leads us to a point between two adjacent splines).
Then we evaluate each of the two splines and interpolate.
### Members
* `IntStarts_` - A vector with the coordinates of each schumaker spline.
* `schumakers` - A vector of schumaker splines.
"""
struct Schumaker2d{T<:AbstractFloat}
    IntStarts_::Array{T,1}
    schumakers::Array{Schumaker{T},1}
    function Schumaker2d(x_row::Vector{R}, x_col::Vector{Y}, ygrid::Array{T,2}; bycol=true,
             extrapolation::Tuple{Schumaker_ExtrapolationSchemes,Schumaker_ExtrapolationSchemes} = (Curve,Curve)) where T<:Real where R<:Real where Y<:Real
        pro = promote_type(promote_type(T,Y),R)
        shape = size(ygrid)
        dimension = 1 + Integer(bycol)
        schums = Array{Schumaker{pro},1}(undef, shape[dimension])
        for i in 1:shape[dimension]
           ys = bycol ?  ygrid[:,i] : ygrid[i,:]
           xs = bycol ? x_row : x_col
           schums[i] = Schumaker(xs, ys; extrapolation = extrapolation)
        end
        return new{pro}(bycol ? x_col : x_row,schums)
    end
    function Schumaker2d(x_row::Vector{R}, x_col::Vector{Y}, ygrid::Array{T,2}; bycol=true,
             extrapolation::Tuple{Schumaker_ExtrapolationSchemes,Schumaker_ExtrapolationSchemes} = (Curve,Curve)) where T<:Integer where R<:Integer where Y<:Integer
        fl = typeof(AbstractFloat(1.0))
        return Schumaker2d( fl.(x_row), fl.(x_col), fl.(ygrid); bycol = bycol, extrapolation = extrapolation  )
    end
end

"""
    evaluate(spline::Schumaker2d, p1::Real, p2::Real)
### Inputs
* `spline` - A Schumaker2d
* `p1` - The coordinate in the first dimension.
* `p2` - The coordinate in the second dimension.
### Returns
* A scalar
"""
function evaluate(spline::Schumaker2d, p1::Real, p2::Real)
    distances   = abs.(spline.IntStarts_ .- p1)
    # Direct hits
    direct_hits = findall(distances .< 100*eps())
    if length(direct_hits) > 0
        return spline.schumakers[direct_hits[1]](p2)
    end
    closest = findall(distances .< sort(distances)[2] + 100*eps())[[1,2]]
    dists = distances[closest]
    ys = map(s -> s(p2), spline.schumakers[closest])
    y = sum(ys .* ( reverse(dists)  ) ./ sum(dists))
    return y
end
function (s::Schumaker2d)(p1::Real, p2::Real)
    return evaluate(s, p1, p2)
end


text:
using DataDeps

register(DataDep(
  "spambase",
  "http://archive.ics.uci.edu/ml/datasets/Spambase",
  "http://archive.ics.uci.edu/ml/machine-learning-databases/spambase/spambase.data",
  "b1ef93de71f97714d3d7d4f58fc9f718da7bbc8ac8a150eff2778616a8097b12",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=58,
        feature_indices=1:57,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
######################################################################
# vehicleset.jl
# v
######################################################################
type Vehicle
	x::Float64
	y::Float64
	max_v::Float64

	Vehicle(x::Float64, y::Float64, max_v::Float64) = new(x,y,max_v)
	Vehicle(x::Float64, y::Float64) = new(x,y,4.0)
end

typealias VehicleSet Vector{Vehicle}

type VehicleSet2
	vehicles::Array{Vehicle}
end


text:
using Oceananigans: tupleit

#####
##### Some validation tools
#####

# Tuple inflation for topologies with Flat dimensions
inflate_tuple(TX, TY, TZ, tup; default) = tup

inflate_tuple(::Type{Flat}, TY, TZ, tup; default) = tuple(default, tup[1], tup[2])
inflate_tuple(TY, ::Type{Flat}, TZ, tup; default) = tuple(tup[1], default, tup[2])
inflate_tuple(TY, TZ, ::Type{Flat}, tup; default) = tuple(tup[1], tup[2], default)

inflate_tuple(TX, ::Type{Flat}, ::Type{Flat}, tup; default) = (tup[1], default, default)
inflate_tuple(::Type{Flat}, TY, ::Type{Flat}, tup; default) = (default, tup[1], default)
inflate_tuple(::Type{Flat}, ::Type{Flat}, TZ, tup; default) = (default, default, tup[1])

inflate_tuple(::Type{Flat}, ::Type{Flat}, ::Type{Flat}, tup; default) = (default, default, default)

topological_tuple_length(TX, TY, TZ) = sum(T === Flat ? 0 : 1 for T in (TX, TY, TZ))

"""Validate that an argument tuple is the right length and has elements of type `argtype`."""
function validate_tupled_argument(arg, argtype, argname, len=3; greater_than=0)
    length(arg) == len      || throw(ArgumentError("length($argname) must be $len."))
    all(isa.(arg, argtype)) || throw(ArgumentError("$argname=$arg must contain $argtype s."))
    all(arg .> greater_than)  || throw(ArgumentError("Elements of $argname=$arg must be > $(greater_than)!"))
    return nothing
end

#####
##### Input validation
#####

function validate_topology(topology)
    for T in topology
        if !isa(T(), AbstractTopology)
            e = "$T is not a valid topology! " *
                "Valid topologies are: Periodic, Bounded, Flat."
            throw(ArgumentError(e))
        end
    end

    return topology
end

function validate_size(TX, TY, TZ, size)
    size = tupleit(size)
    validate_tupled_argument(size, Integer, "size", topological_tuple_length(TX, TY, TZ))
    return inflate_tuple(TX, TY, TZ, size, default=1)
end

# Note that the default halo size is specified to be 1 in the following function.
# This is easily changed but many of the tests will fail so this situation needs to be
# cleaned up.
function validate_halo(TX, TY, TZ, ::Nothing)
    halo = Tuple(1 for i = 1:topological_tuple_length(TX, TY, TZ))
    return validate_halo(TX, TY, TZ, halo)
end

function validate_halo(TX, TY, TZ, halo)
    halo = tupleit(halo)
    validate_tupled_argument(halo, Integer, "halo", topological_tuple_length(TX, TY, TZ))
    return inflate_tuple(TX, TY, TZ, halo, default=0)
end

coordinate_name(i) = i == 1 ? "x" : i == 2 ? "y" : "z"

function validate_dimension_specification(T, ξ, dir, FT)

    isnothing(ξ)         && throw(ArgumentError("Must supply extent or $dir keyword when $dir-direction is $T"))
    length(ξ) == 2       || throw(ArgumentError("$dir length($ξ) must be 2."))
    all(isa.(ξ, Number)) || throw(ArgumentError("$dir=$ξ should contain numbers."))
    ξ[2] >= ξ[1]         || throw(ArgumentError("$dir=$ξ should be an increasing interval."))

    return FT.(ξ)
end

function validate_rectilinear_domain(TX, TY, TZ, FT, extent, x, y, z)

    # Find domain endpoints or domain extent, depending on user input:
    if !isnothing(extent) # the user has specified an extent!

        (!isnothing(x) || !isnothing(y) || !isnothing(z)) &&
            throw(ArgumentError("Cannot specify both 'extent' and 'x, y, z' keyword arguments."))

        extent = tupleit(extent)

        validate_tupled_argument(extent, Number, "extent", topological_tuple_length(TX, TY, TZ))

        Lx, Ly, Lz = extent = inflate_tuple(TX, TY, TZ, extent, default=0)

        # An "oceanic" default domain:
        x = FT.((0, Lx))
        y = FT.((0, Ly))
        z = FT.((-Lz, 0))

    else # isnothing(extent) === true implies that user has not specified a length
        x = validate_dimension_specification(TX, x, :x, FT)
        y = validate_dimension_specification(TY, y, :y, FT)
        z = validate_dimension_specification(TZ, z, :z, FT)
    end

    return x, y, z
end

validate_dimension_specification(T, ξ::AbstractVector, dir, FT) = FT.(ξ)
validate_dimension_specification(T, ξ::Function,       dir, FT) = ξ
validate_dimension_specification(::Type{Flat}, ξ::AbstractVector, dir, FT) = FT.((ξ[1], ξ[1]))
validate_dimension_specification(::Type{Flat}, ξ::Function,       dir, FT) = FT.((ξ(1), ξ(1)))
validate_dimension_specification(::Type{Flat}, ξ::Tuple, dir, FT)  = FT.(ξ)
validate_dimension_specification(::Type{Flat}, ::Nothing, dir, FT) = FT.((0, 0))
validate_dimension_specification(::Type{Flat}, ξ::Number, dir, FT) = FT.((ξ, ξ))

default_horizontal_extent(T, extent) = (0, extent[i])
default_vertical_extent(T, extent) = (-extent[3], 0)

function validate_regular_grid_domain(TX, TY, TZ, FT, extent, x, y, z)

    # Find domain endpoints or domain extent, depending on user input:
    if !isnothing(extent) # the user has specified an extent!

        (!isnothing(x) || !isnothing(y) || !isnothing(z)) &&
            throw(ArgumentError("Cannot specify both 'extent' and 'x, y, z' keyword arguments."))

        extent = tupleit(extent)

        validate_tupled_argument(extent, Number, "extent", topological_tuple_length(TX, TY, TZ))

        Lx, Ly, Lz = extent = inflate_tuple(TX, TY, TZ, extent, default=0)

        # An "oceanic" default domain:
        x = (0, Lx)
        y = (0, Ly)
        z = (-Lz, 0)

    else # isnothing(extent) === true implies that user has not specified a length
        x = validate_dimension_specification(TX, x, :x, FT)
        y = validate_dimension_specification(TY, y, :y, FT)
        z = validate_dimension_specification(TZ, z, :z, FT)

        Lx = x[2] - x[1]
        Ly = y[2] - y[1]
        Lz = z[2] - z[1]
    end

    return FT(Lx), FT(Ly), FT(Lz), FT.(x), FT.(y), FT.(z)
end

function validate_vertically_stretched_grid_xy(TX, TY, FT, x, y)
    x = validate_dimension_specification(TX, x, :x, FT)
    y = validate_dimension_specification(TY, y, :y, FT)

    Lx = x[2] - x[1]
    Ly = y[2] - y[1]

    return FT(Lx), FT(Ly), FT.(x), FT.(y)
end

validate_unit_vector(ê::ZDirection) = ê

function validate_unit_vector(ê)
    length(ê) == 3 || throw(ArgumentError("unit vector must have length 3"))

    ex, ey, ez = ê

    ex^2 + ey^2 + ez^2 ≈ 1 ||
        throw(ArgumentError("unit vector `ê` must have ê[1]² + ê[2]² + ê[3]² ≈ 1"))

    return tuple(ê...)
end


text:
"""
The module `Blocks` contains common input-output components, referred to as blocks.

In general, input-output blocks follow the convention
```
     ┌───────────┐
 u   │  ẋ=f(x,u) │  y
────►│  y=g(x,u) ├────►
     │           │
     └───────────┘
```
where `u` are inputs, `x` are state variables and `y` are outputs. `x,u,y` are all implemented as `@variables` internally, `u` are marked as `[input=true]` and `y` are marked `[output=true]`.
"""
module Blocks
using ModelingToolkit, Symbolics, IfElse, OrdinaryDiffEq

@parameters t
Dₜ = Differential(t)

export Gain, Sum
include("math.jl")

export Saturation, DeadZone
include("nonlinear.jl")

export Constant, Integrator, Derivative, FirstOrder, SecondOrder, PID, StateSpace
include("continuous.jl")

end


text:
@doc raw"""
    LennardJones(; cutoff, nl_only, lorentz_mixing, weight_14, force_unit, energy_unit, skip_shortcut)

The Lennard-Jones 6-12 interaction. The potential is given by
```math
V_{ij}(r_{ij}) = 4\varepsilon_{ij} \left[\left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} - \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{6}\right]
```
and the force on each atom by
```math
\begin{aligned}
\vec{F}_i &= 24\varepsilon_{ij} \left(2\frac{\sigma_{ij}^{12}}{r_{ij}^{13}} - \frac{\sigma_{ij}^6}{r_{ij}^{7}}\right) \frac{\vec{r}_{ij}}{r_{ij}} \\
&= \frac{24\varepsilon_{ij}}{r_{ij}^2} \left[2\left(\frac{\sigma_{ij}^{6}}{r_{ij}^{6}}\right)^2 -\left(\frac{\sigma_{ij}}{r_{ij}}\right)^{6}\right] \vec{r}_{ij}
\end{aligned}
```
"""
struct LennardJones{S, C, W, F, E} <: GeneralInteraction
    cutoff::C
    nl_only::Bool
    lorentz_mixing::Bool
    weight_14::W
    force_unit::F
    energy_unit::E
end

function LennardJones(;
                        cutoff=NoCutoff(),
                        nl_only=false,
                        lorentz_mixing=true,
                        weight_14=1.0,
                        force_unit=u"kJ * mol^-1 * nm^-1",
                        energy_unit=u"kJ * mol^-1",
                        skip_shortcut=false)
    return LennardJones{skip_shortcut, typeof(cutoff), typeof(weight_14), typeof(force_unit), typeof(energy_unit)}(
        cutoff, nl_only, lorentz_mixing, weight_14, force_unit, energy_unit)
end

@inline @inbounds function force(inter::LennardJones{S, C},
                                    coord_i,
                                    coord_j,
                                    atom_i,
                                    atom_j,
                                    box_size,
                                    weight_14::Bool=false) where {S, C}
    dr = vector(coord_i, coord_j, box_size)
    r2 = sum(abs2, dr)

    if !S && iszero(atom_i.σ) || iszero(atom_j.σ)
        return ustrip.(zero(coord_i)) * inter.force_unit
    end

    # Lorentz-Berthelot mixing rules use the arithmetic average for σ
    # Otherwise use the geometric average
    σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ)
    ϵ = sqrt(atom_i.ϵ * atom_j.ϵ)

    cutoff = inter.cutoff
    σ2 = σ^2
    params = (σ2, ϵ)

    if cutoff_points(C) == 0
        f = force_divr_nocutoff(inter, r2, inv(r2), params)
    elseif cutoff_points(C) == 1
        r2 > cutoff.sqdist_cutoff && return ustrip.(zero(coord_i)) * inter.force_unit

        f = force_divr_cutoff(cutoff, r2, inter, params)
    elseif cutoff_points(C) == 2
        r2 > cutoff.sqdist_cutoff && return ustrip.(zero(coord_i)) * inter.force_unit

        if r2 < cutoff.activation_dist
            f = force_divr_nocutoff(inter, r2, inv(r2), params)
        else
            f = force_divr_cutoff(cutoff, r2, inter, params)
        end
    end

    if weight_14
        return f * dr * inter.weight_14
    else
        return f * dr
    end
end

@fastmath function force_divr_nocutoff(::LennardJones, r2, invr2, (σ2, ϵ))
    six_term = (σ2 * invr2) ^ 3

    return (24ϵ * invr2) * (2 * six_term ^ 2 - six_term)
end

@inline @inbounds function potential_energy(inter::LennardJones{S, C},
                                            s::System,
                                            i::Integer,
                                            j::Integer,
                                            weight_14::Bool=false) where {S, C}
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    r2 = sum(abs2, dr)

    if !S && iszero(s.atoms[i].σ) || iszero(s.atoms[j].σ)
        return ustrip(zero(s.box_size[1])) * inter.energy_unit
    end

    σ = inter.lorentz_mixing ? (s.atoms[i].σ + s.atoms[j].σ) / 2 : sqrt(s.atoms[i].σ * s.atoms[j].σ)
    ϵ = sqrt(s.atoms[i].ϵ * s.atoms[j].ϵ)

    cutoff = inter.cutoff
    σ2 = σ^2
    params = (σ2, ϵ)

    if cutoff_points(C) == 0
        pe = potential(inter, r2, inv(r2), params)
    elseif cutoff_points(C) == 1
        r2 > cutoff.sqdist_cutoff && return ustrip(zero(s.box_size[1])) * inter.energy_unit

        pe = potential_cutoff(cutoff, r2, inter, params)
    elseif cutoff_points(C) == 2
        r2 > cutoff.sqdist_cutoff && return ustrip(zero(s.box_size[1])) * inter.energy_unit

        if r2 < cutoff.activation_dist
            pe = potential(inter, r2, inv(r2), params)
        else
            pe = potential_cutoff(cutoff, r2, inter, params)
        end
    end

    if weight_14
        return pe * inter.weight_14
    else
        return pe
    end
end

@fastmath function potential(::LennardJones, r2, invr2, (σ2, ϵ))
    six_term = (σ2 * invr2) ^ 3

    return 4ϵ * (six_term ^ 2 - six_term)
end


text:
############################################################################################
#                               INFEASIBLE MODELS                                          #
############################################################################################
struct Infeasible{N,M,D<:AbstractModel} <: AbstractModel
    model::D
    _u::SVector{M,Int}  # inds to original controls
    _ui::SVector{N,Int} # inds to infeasible controls
end

struct InfeasibleLie{N,M,D<:AbstractModel} <: RobotDynamics.LieGroupModel
    model::D
    _u::SVector{M,Int}  # inds to original controls
    _ui::SVector{N,Int} # inds to infeasible controls
end

""" $(TYPEDEF)
An infeasible model is an augmented dynamics model that makes the system artifically fully
actuated by augmenting the control vector with `n` additional controls. The dynamics are
handled explicitly in discrete time:

``x_{k+1} = f(x_k,u_k,dt) + w_k``

where ``w_k`` are the additional `n`-dimensional controls. In practice, these are constrained
to be zero by the end of the solve.

# Constructors
```julia
InfeasibleModel(model::AbstractModel)
```
"""
const InfeasibleModel{N,M,D} = Union{Infeasible{N,M,D},InfeasibleLie{N,M,D}} where {N,M,D}

function InfeasibleModel(model::AbstractModel)
    n,m = size(model)
    _u  = SVector{m}(1:m)
    _ui = SVector{n}((1:n) .+ m)
    Infeasible(model, _u, _ui)
end

function InfeasibleModel(model::RobotDynamics.LieGroupModel)
    n,m = size(model)
    _u  = SVector{m}(1:m)
    _ui = SVector{n}((1:n) .+ m)
    InfeasibleLie(model, _u, _ui)
end

RobotDynamics.LieState(model::InfeasibleLie) = RobotDynamics.LieState(model.model)

# Generic Infeasible Methods
RobotDynamics.state_dim(model::InfeasibleModel{n}) where n = n
RobotDynamics.control_dim(model::InfeasibleModel{n,m}) where {n,m} = n+m

RobotDynamics.dynamics(::InfeasibleModel, x, u) =
    throw(ErrorException("Cannot evaluate continuous dynamics on an infeasible model"))

@generated function RobotDynamics.discrete_dynamics(::Type{Q}, model::InfeasibleModel{N,M},
        z::AbstractKnotPoint{T,N}) where {T,N,M,Q<:Explicit}
    _u = SVector{M}((1:M) .+ N)
    _ui = SVector{N}((1:N) .+ (N+M))
    quote
        x = state(z)
        dt = z.dt
        u0 = z.z[$_u]
        ui = z.z[$_ui]
        RobotDynamics.discrete_dynamics($Q, model.model, x, u0, z.t, dt) + ui
    end
end

@inline RobotDynamics.rotation_type(model::InfeasibleModel) where D = rotation_type(model.model)

@generated function RobotDynamics.discrete_jacobian!(::Type{Q}, ∇f, model::InfeasibleModel{N,M},
        z::AbstractKnotPoint{T,N}, cache=nothing) where {T,N,M,Q<:Explicit}

    ∇ui = [(@SMatrix zeros(N,N+M)) Diagonal(@SVector ones(N)) @SVector zeros(N)]
    _x = SVector{N}(1:N)
    _u = SVector{M}((1:M) .+ N)
    _z = SVector{N+M}(1:N+M)
    _ui = SVector{N}((1:N) .+ (N+M))
    zi = [:(z.z[$i]) for i = 1:N+M]
    NM1 = N+M+1
	NM = N+M
    ∇u0 = @SMatrix zeros(N,N)

    quote
        # Build KnotPoint for original model
        s0 = SVector{$NM1}($(zi...), z.dt)

        u0 = z.z[$_u]
        ui = z.z[$_ui]
		z_ = StaticKnotPoint(z.z[$_z], $_x, $_u, z.dt, z.t)
		∇f_ = uview(∇f, 1:N, 1:$NM)
        discrete_jacobian!($Q, ∇f_, model.model, z_)
		# ∇f[$_x, N+NM] .= ∇f_[$_x, N+M] # ∇dt
		∇f[$_x, $_ui] .= Diagonal(@SVector ones(N))
		return
		# ∇f[$_x,$_ui]
        # [∇f[$_x, $_z] $∇u0 ∇dt] + $∇ui
    end
end
function RD._discrete_jacobian!(::RD.ForwardAD, ::Type{Q}, ∇f, model::InfeasibleModel{N,M},
        z::AbstractKnotPoint{T,N}, cache=nothing) where {T,N,M,Q<:Explicit}
    RD.discrete_jacobian!(Q, ∇f, model, z, cache)
end

function RobotDynamics.state_diff(model::InfeasibleModel, x::SVector, x0::SVector)
	RobotDynamics.state_diff(model.model, x, x0)
end

function RobotDynamics.state_diff_jacobian!(G, model::InfeasibleModel, Z::Traj)
	RobotDynamics.state_diff_jacobian!(G, model.model, Z)
end

function RobotDynamics.∇²differential!(∇G, model::InfeasibleModel, x::SVector, dx::SVector)
	return ∇²differential!(∇G, model.model, x, dx)
end

RobotDynamics.state_diff_size(model::InfeasibleModel) = RobotDynamics.state_diff_size(model.model)

Base.position(model::InfeasibleModel, x::SVector) = position(model.model, x)

RobotDynamics.orientation(model::InfeasibleModel, x::SVector) = orientation(model.model, x)

"Calculate a dynamically feasible initial trajectory for an infeasible problem, given a
desired trajectory"
function infeasible_trajectory(model::InfeasibleModel{n,m}, Z0::Traj) where {T,n,m}
    x,u = zeros(model)
    ui = @SVector zeros(n)
    Z = [KnotPoint(state(z), [control(z); ui], z.dt, z.t) for z in Z0]
    N = length(Z0)
    for k = 1:N-1
        RobotDynamics.propagate_dynamics(RobotDynamics.RK4, model, Z[k+1], Z[k])
        x′ = state(Z[k+1])
        u_slack = state(Z0[k+1]) - x′
        u = [control(Z0[k]); u_slack]
        RobotDynamics.set_control!(Z[k], u)
        RobotDynamics.set_state!(Z[k+1], x′ + u_slack)
    end
    return Traj(Z)
end

############################################################################################
#  								INFEASIBLE CONSTRAINT 									   #
############################################################################################
""" $(TYPEDEF) Constraints additional ``infeasible'' controls to be zero.
Constructors: ```julia
InfeasibleConstraint(model::InfeasibleModel)
InfeasibleConstraint(n,m)
```
"""
struct InfeasibleConstraint{n} <: TO.ControlConstraint
	ui::SVector{n,Int}
	m::Int
	function InfeasibleConstraint(n::Int, m::Int)
		ui = SVector{n}((1:n) .+ m)
		new{n}(ui, m)
	end
end

InfeasibleConstraint(model::InfeasibleModel{n,m}) where {n,m} = InfeasibleConstraint(n,m)
RobotDynamics.control_dim(con::InfeasibleConstraint{n}) where n = n + con.m
@inline TO.sense(::InfeasibleConstraint) = TO.Equality()
@inline Base.length(::InfeasibleConstraint{n}) where n = n

function TO.evaluate(con::InfeasibleConstraint, u::SVector)
    ui = u[con.ui] # infeasible controls
end

function TO.jacobian!(∇c, con::InfeasibleConstraint{n}, u::SVector) where n
	for (i,j) in enumerate(con.ui)
		∇c[i,j] = 1
    end
    return true
end


text:
Pkg.add("Colors")
Pkg.add("ColorSchemes")
Pkg.add("NearestNeighbors")
Pkg.add("GeoStats")
Pkg.add("Images")
Pkg.add("Clustering")
Pkg.add("GeoJSON")
Pkg.add("LibGEOS")
Pkg.add("ArgParse")
Pkg.add("Lumberjack")
Pkg.add("SortingAlgorithms")
Pkg.add("ProgressMeter")
Pkg.add("DataArrays")
Pkg.add("Glob")
Pkg.add("LasIO")
Pkg.add("GDAL")


text:
using LinearAlgebra
using SparseArrays
using Plots
using ForwardDiff

"This routine runs a convergence test for the heat equation using either Backwards
Euler or Crank-Nicolson in time and finite differences in space. The exact solution
is generated using the method of manufactured solutions."

uexact(x,t) = log(2+sin(pi*x))*exp(-t)*t
dudx_exact(x,t) = ForwardDiff.derivative(x->uexact(x,t),x)
# du/dt - d^2u/dx^2 = f
f(x,t) = ForwardDiff.derivative(t->uexact(x,t),t) -
         ForwardDiff.derivative(x->dudx_exact(x,t),x)
u0(x) = uexact(x,0)
α(t) = uexact(-1,t)
β(t) = uexact(1,t)

# m = number of grid points
function build_FD(m,f,α,β)
    x = LinRange(-1,1,m+2)
    xint = x[2:end-1]
    h = x[2]-x[1]
    A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))

    function F(t)
        b = f.(xint,t)
        b[1] += α(t)/h^2
        b[m] += β(t)/h^2
        return b
    end
    return A,x,xint,h,F
end

function solve_backward_Euler(m, f, α, β, u0, T, dt)

    A,x,xint,h,F = build_FD(m,f,α,β)

    u = u0.(xint)
    Nsteps = ceil(Int,T/dt)
    dt = T/Nsteps
    for i = 1:Nsteps
        t = i*dt
        u .= (I + dt*A)\(u .+ dt*F(t))
    end
    return u,x,xint,h
end

function solve_Crank_Nicolson(m, f, α, β, u0, T, dt)

    A,x,xint,h,F = build_FD(m,f,α,β)

    u = u0.(xint)
    Nsteps = ceil(Int,T/dt)
    dt = T/Nsteps
    for i = 1:Nsteps
        tprev = (i-1)*dt
        t = i*dt
        u .= (I + .5*dt*A)\((I - .5*dt*A)*u + .5*dt*(F(t)+F(tprev)))
    end
    return u,x,xint,h
end

mvec = 2 .^ (2:8)
hvec = zeros(length(mvec))
hvec = Float64[]
err1 = Float64[]
err2 = Float64[]
for (i,m) in enumerate(mvec)

    T = 1.0
    dt = 5/m # proportional to h

    u,x,xint,h = solve_backward_Euler(m,f,α,β,u0,T,dt)
    push!(err1,maximum(@. abs(uexact(xint,T) - u)))

    u,x,xint,h = solve_Crank_Nicolson(m,f,α,β,u0,T,dt)
    push!(err2,maximum(@. abs(uexact(xint,T) - u)))

    push!(hvec,h)
end
plot(hvec,err1,marker=:dot,label="Max error (backwards Euler)")
plot!(hvec,err2,marker=:dot,label="Max error (Crank-Nicolson)")
plot!(hvec,hvec,linestyle=:dash,label="O(h)")
plot!(hvec,hvec.^2,linestyle=:dash,label="O(h^2)")
plot!(xaxis=:log,yaxis=:log,legend=:bottomright)


text:
using CSV, JuMP, Gurobi, LinearAlgebra

#Read the data
y15 = CSV.read("Y15_v4.csv");
T = y15[:,1]; #total periods
TPD  = y15[:,3]; #set of total demand
SS = [1,2,3];   #set of slices
WS = [0.25,0.5,0.25]; #set of weight of slices
S1 = y15[337:384,1]; #array of slice 1
S2= y15[4369:4416,1]; #array of slice 2
S3=y15[9457:9504,1]; #arracy of slice 3

TExt = y15[:,6]; #exterior temperature
pGW_15 = y15[:,4]; #2015 wind power
pGS_15 = y15[:,5]; #2015 solar power

#time period for operational model
#TO = y15[4086:4421,1]
#TTT = hcat(TO,TO,TO)

#time period for planning model
S1 = y15[337:384,1]; #array of slice 1
S2= y15[4369:4416,1]; #array of slice 2
S3=y15[9457:9504,1]; #arracy of slice 3
TT = hcat(S1,S2,S3); #total periods in all slices

gen_type = CSV.read("gen_type.csv");
G   = gen_type[1:4,1];  #set of conventional generator
W   = gen_type[1:5,2]; #set of wind gnereator
S   = gen_type[1:1,3]; #set of solar generator
ES  = gen_type[1:1,4]; #set of electricity
HS  = gen_type[1:1,5]; #set of heat store

gen_info = CSV.read("con_info.csv"); #information of conventional generator
EG = gen_info[:,1]; #efficiency of  gnereator
FG_CPX = gen_info[:,4]; #CAPEX of generator
FG_OM = gen_info[:,3]; #OPEX of generator
YG = gen_info[:,5]; #life of gnereator
EM = gen_info[:,7]; #co2 emission per tonee
CG = gen_info[:,2]; #cost of generator
CFuel = gen_info[:,8]; #cost of fuel
PGmax = gen_info[:,6]; #generator limits

wind_info = CSV.read("wind_info.csv");
WF = wind_info[:,9:14]; #wind pattern
EW = wind_info[1:5,2]; #efficiency of wind farm
FW_CPX = wind_info[1:5,5]; # CAPEX of wind
FW_OM = wind_info[1:5,4]; #OPEX of wind
YW = wind_info[1:5,6]; #life of wind
Wcap_15 = wind_info[1:5,7]; #wind capcity in 2015

solar_info = CSV.read("solar_info.csv");
ES = solar_info[:,1]; #efficiency of solarpv
FS_CPX = solar_info[:,4]; #CAPEX of solar
FS_OM = solar_info[:,3]; #OPEX of solar
YS = solar_info[:,5]; #life of solar
Scap_15 = solar_info[1,6]; #solar capacity in 2015

HS_info = CSV.read("Hstore_info.csv");
EHS_h = HS_info[:,1]; #efficiency of heating a store
EHS_c = HS_info[:,11];
CHeat = HS_info[:,2]; #cost of power go into the heat store
FHS_CPX = HS_info[:,4]; #fixed capex of heat store
YHS = HS_info[:,5]; #life of heat store
FHS_OM = HS_info[:,3]; #fixed om of heat store
Qmass = HS_info[1,6]; #Qmass
Ploss = HS_info[1,7]; #Ploss
THS_plus = HS_info[:,8]; #maximum allowed temperature of heat store
THS_minus = HS_info[:,9]; #minimum allowed temperature of heat store
PHS_plus_cap = HS_info[:,12]; #maximum power input

ES_info = CSV.read("Estore_info.csv");
CElect = ES_info[:,2]; #cost of power go into electricity store
EES = ES_info[:,1]; #efficiency of electricity store
FES_CPX = ES_info[:,4]; #capex of electricity store
YES = ES_info[:,5]; #life of electricity store
FES_OM = ES_info[:,3]; #fixed om of electricity store
ES_Ecap = ES_info[:,7]; #maximum energy level in the electricity store
FES_ECPX = ES_info[:,6]; #energy CAPX of electricity store
PES_plus_cap = ES_info[:,9]; #maximum power into electricity store es
PES_minus_cap = ES_info[:,10]; #maximum power out of electricity store es

other_info = CSV.read("other_info.csv");
CO2_tax = other_info[1,2]; #CO2 tax
CShed_load = other_info[2,2]; #cost of load shedding
CShed_gen = other_info[3,2]; #cost of generation shedding
Ht = other_info[4,2]; #numbder of hour in each period
CO2_max = other_info[5,2]; #maximum co2 emission allowed
H = other_info[6,2]; #number of hours in a day
Y = other_info[7,2]; #number of hours in a year

#CO2 cases
co2_info = CSV.read("co2_case.csv");
co2_case = co2_info[:,1]; #6 co2 cases
Z = co2_info[:,2]; #coefficient

#bus-line information
net_info = CSV.read("network_info.csv");
B = net_info[1:5,1]; #set of bus
BF = net_info[1:5,2]; #set of bus demand factors
L = net_info[1:7,3]; #set of line
FL_CPX = net_info[1:7,7]; #fixed capex of line
FL_OM = net_info[1:7,8]; #fixed opex of line
GLoc = net_info[1:4,10]; #generator location
WLoc = net_info[1:5,11]; #wind farm location
SLoc = net_info[1:1,12]; #solar loaction
HSLoc = net_info[1:1,21]; #heat store location
ESLoc = net_info[1:1,22]; #electricity store location
a = net_info[1:5,13:19]; #bus-line incidence matrix
LR = net_info[1:7,5]; #line capacity
Re_bus = net_info[1,20]; #reference bus
PLmax = net_info[1:7,4]; #line capacity
YL = net_info[1:7,6]; #life of line


text:
module Hyperopt

export Hyperoptimizer, @hyperopt, @phyperopt, @thyperopt, printmin, printmax, warn_on_boundary
export RandomSampler, LHSampler, CLHSampler, hyperband, Hyperband, hyperoptim, BOHB, Continuous, Categorical, UnorderedCategorical

using Base.Threads: threadid, nthreads
using LinearAlgebra, Statistics, Random
using ProgressMeter
using MacroTools
using MacroTools: postwalk, prewalk
using RecipesBase
using Distributed
using LatinHypercubeSampling
using ThreadPools
using Distributions: Normal, truncated
using MultiKDE
using Requires
using Printf

const DimensionType = LHCDimension

# # Types of dimensions
# const CategoricalDim = Categorical
# const ContinuousDim = Continuous
struct UnorderedCategorical <: DimensionType
    levels::Int64
end

abstract type Sampler end
Base.@kwdef mutable struct Hyperoptimizer{S<:Sampler, F}
    iterations::Int
    params
    candidates
    history = []
    results = []
    sampler::S = RandomSampler()
    objective::F = nothing
end

function Base.getproperty(ho::Hyperoptimizer, s::Symbol)
    s === :minimum && (return isempty(ho.results) ? NaN :  minimum(replace(ho.results, NaN => Inf)))
    s === :minimizer && (return isempty(ho.results) ? [] :  ho.history[argmin(replace(ho.results, NaN => Inf))])
    s === :maximum && (return isempty(ho.results) ? NaN :  maximum(replace(ho.results, NaN => Inf)))
    s === :maximizer && (return isempty(ho.results) ? [] :  ho.history[argmax(replace(ho.results, NaN => Inf))])
    return getfield(ho,s)
end

function Base.show(io::IO, ho::Hyperoptimizer)
    println(io, "Hyperoptimizer with")
    cands = ho.candidates
    candstring = map(keys(cands)) do k
        s = "  "*string(k)*" length: "
        if length(cands[k]) <= 3
            s*string(cands[k])
        else
            s*string(length(cands[k]))
        end
    end
    candstring = join(candstring, "\n")
    println(io, candstring)
    println(io, "  minimum / maximum: $((ho.minimum, ho.maximum))")
    println(io, "  minimizer:")
    for (i, v) in enumerate(ho.minimizer)
        @printf(io, "%9s ", string(ho.params[i]))
    end
    println(io)
    for (i, v) in enumerate(ho.minimizer)
        @printf(io, "%9.4g ", v)
    end
    println()
end

Base.propertynames(ho::Hyperoptimizer) = (:minimum, :minimizer, :maximum, :maximizer, fieldnames(Hyperoptimizer)...)

include("samplers.jl")

function Hyperoptimizer(iterations::Int, sampler::Sampler = RandomSampler(); kwargs...)
    params = keys(kwargs)
    candidates = []
    for kw in kwargs
        push!(candidates, kw[2])
    end
    Hyperoptimizer(iterations=iterations, params=params, candidates=candidates, sampler=sampler)
end

Base.length(ho::Hyperoptimizer) = ho.iterations

function Base.iterate(ho::Hyperoptimizer, state=1)
    state > ho.iterations && return nothing
    samples = ho.sampler(ho, state)
    push!(ho.history, samples)
    nt = (; Pair.((:i, ho.params...), (state, samples...))...)
    nt, state+1
end

"""
Hyperopt internal function

Returns
`params, candidates, sampler_, ho_, fun`
where
`params` is a vector of symbols
`candidates` is a vector of expresions creating candidate vectors
`sampler_` is an expression creating a sampler
`ho_` is either nothing or an expression for an existing hyperoptimizer
`fun` is a function from `(i, pars...) -> val` that executes the macrobody
"""
function preprocess_expression(ex)
    ex.head == :for || error("Wrong syntax, Use For-loop syntax, ex: @hyperopt for i=100, param=LinRange(1,10,100) ...")
    params     = []
    candidates = []
    sampler_ = :(RandomSampler())
    ho_ = nothing
    loop = ex.args[1].args # ex.args[1] = the arguments to the For loop
    i = 1
    while i <= length(loop)
        if @capture(loop[i], sampler = sam_) # A sampler was provided
            sampler_ = sam
            deleteat!(loop,i) # Remove the sampler from the args
            continue
        end
        if @capture(loop[i], ho = h_)
            ho_ = h
            deleteat!(loop,i)
            continue
        end
        @capture(loop[i], param_ = list_) || error("Wrong syntax, Use For-loop syntax, ex: @hyperopt for i=100, param=LinRange(1,10,100) ...")
        push!(params, param)
        push!(candidates, list)
        i += 1
    end
    params = (params...,)
    funname = gensym(:hyperopt_objective)
    state_ = :(state = nothing)
    fun = quote # produces a function(i, pars...)
        function $(funname)($(Expr(:tuple, esc.(params)...))...)
            $(esc(state_))
            $(esc(ex.args[2]))
        end
        function $(funname)($(esc(params[1])), $(esc(:state)))
            $(esc(ex.args[2]))
        end
    end

    params, candidates, sampler_, ho_, fun
end

function optimize(ho::Hyperoptimizer)
    try
        @showprogress "Hyperoptimizing" for nt = ho
            res = ho.objective(nt...)
            push!(ho.results, res)
        end
    catch e
        if e isa InterruptException
            @info "Aborting hyperoptimization"
        else
            rethrow(e)
        end
    end
    ho
end

function create_ho(params, candidates, sampler, ho_, objective_)
    esccands = esc.(candidates)
    allcands = Expr(:tuple, esccands...)
    quote
        objective, iters = $(objective_), $(esc(candidates[1]))
        ho = Hyperoptimizer(; iterations = iters, params = $(esc(params))[2:end], candidates = $(allcands)[2:end], sampler=$(esc(sampler)), objective)
        if $(esc(ho_)) isa Hyperoptimizer # get info from existing ho. the objective function might be changed, due to variables are captured into the closure, so the type of ho also changed.
            ho.sampler = $(esc(ho_)).sampler
            ho.history = $(esc(ho_)).history # it's important to use the same array, not copy.
            ho.results = $(esc(ho_)).results
        else
            s = ho.sampler
            if s isa Hyperband
                ho.iterations = length($(esc(candidates[end])))
                s = s.inner
            end
            s isa Union{LHSampler,CLHSampler} && init!(s, ho)
        end
        ho
    end
end

macro hyperopt(ex)
    pre = preprocess_expression(ex)
    ho_ = create_ho(pre...)
    :(optimize($ho_))
end

function pmacrobody(ex, params, ho_, pmap=pmap)
    quote
        function workaround_function()
            ho = $(ho_)
            # Getting the history right is tricky when using workers. The approach I've found to work is to
            # save the actual array (not copy) in hist, temporarily use a new array that will later be discarded
            # reassign the original array and then append the new history. If a new array is used, the change will not be visible in the original hyperoptimizer
            hist = ho.history
            ho.history = []
            res = $(pmap)(1:ho.iterations) do i
                $(Expr(:tuple, esc.(params)...)),_ = iterate(ho,i)
                res = $(esc(ex.args[2])) # ex.args[2] = Body of the For loop

                res, $(Expr(:tuple, esc.(params[2:end])...))
            end
            ho.history = hist
            append!(ho.results, getindex.(res,1))
            append!(ho.history, getindex.(res,2)) # history automatically appended by the iteration
            ho
        end
        workaround_function()
    end
end

"""
Same as `@hyperopt` but uses `Distributed.pmap` for parallel evaluation of the cost function.
"""
macro phyperopt(ex)
    pre = preprocess_expression(ex)
    ho_ = create_ho(pre...)
    pmacrobody(ex, pre[1], ho_)
end

"""
Same as `@hyperopt` but uses `ThreadPools.tmap` for multithreaded evaluation of the cost function.
"""
macro thyperopt(ex)
    pre = preprocess_expression(ex)
    ho_ = create_ho(pre...)
    pmacrobody(ex, pre[1], ho_, ThreadPools.tmap)
end

function Base.minimum(ho::Hyperoptimizer)
    m,i = findmin(replace(ho.results, NaN => Inf))
    m
end
function Base.maximum(ho::Hyperoptimizer)
    m,i = findmax(replace(ho.results, NaN => -Inf))
    m
end

function printmin(ho::Hyperoptimizer)
    for (param, value) in zip(ho.params, ho.minimizer)
        println(param, " = ", value)
    end
end

function printmax(ho::Hyperoptimizer)
    for (param, value) in zip(ho.params, ho.maximizer)
        println(param, " = ", value)
    end
end

"""
    warn_on_boundary(ho, sense = :min)

Prints a warning message for each parameter where the optimum was obtained on an extreme point of the sampled space.

Example: If parameter `a` can take values in 1:10 and the optimum was obtained at
`a = 1`, it's an indication that the parameter was constraind by the search space.
The warning is effective even if the lowest value of `a` that was sampled was higher than 1,
but the optimum occured on the lowest sampled value.
"""
function warn_on_boundary(ho, sense = :min)
    m = sense == :min ? ho.minimizer : ho.maximizer
    n_params = length(m)
    extremas = map(1:n_params) do i
        c = ho.candidates[i]
        if c isa AbstractArray{<:Real}
            extrema(getindex.(ho.history, i))
        else
            (m[i],)
        end
    end
    for i in eachindex(m)
        c = unique(ho.candidates[i])
        if m[i] ∈ extremas[i] && length(c) > 3
            println("Parameter $(ho.params[i]) obtained its optimum on an extremum of the sampled region: $(m[i])")
        end
    end
end

include("plotting.jl")

function __init__()
    Requires.@require Optim="429524aa-4258-5aef-a3af-852621145aeb" include("optim.jl")
end

end # module


text:
function CinI(prompt::String="")::Int
	print(prompt)
	while true
		try
			return parse(Int, readline())
		catch
			println("Input not an integer; Try again:")
		end
	end
end

function CinF(prompt::String="")::Float64
	print(prompt)
	while true
		try
			return parse(Float64, readline())
		catch
			println("Input not a float; Try again:")
		end
	end
end


text:
""" Convert a presentation math expression to a MML string or a list of MML string
"""
exlist_to_mml(nodes::Vector{CompoExpr}) = foldl(*, map(expr_to_mml, nodes), init="")

expr_to_mml(n::Sub) = mml_element(:msub, _2_exprs_to_mml(n.base, n.sub))
expr_to_mml(n::Sup) = mml_element(:msup, _2_exprs_to_mml(n.base, n.sup))
expr_to_mml(n::SubSup) = mml_element(:msubsup, _3_exprs_to_mml(n.base, n.sub, n.sup))
expr_to_mml(n::Under) = mml_element(:munder, _2_exprs_to_mml(n.base, n.under))
expr_to_mml(n::Over) = mml_element(:mover, _2_exprs_to_mml(n.base, n.over))
expr_to_mml(n::UnderOver) = mml_element(
        :munderover,
       _3_exprs_to_mml(n.base, n.under, n.over))
expr_to_mml(n::Sqrt) = mml_element(:msqrt, expr_to_mml(n.base))
expr_to_mml(n::Frac) = mml_element(:mfrac, _2_exprs_to_mml(n.num, n.den))
expr_to_mml(n::Group) = mml_element(:mrow, exlist_to_mml(n.val))
# we do not really use the `mfenced` element yet. Even firefox does not
# support it as of 2021 (but MathJax does). Per MathML 3.0 specification (2nd
# edition), the following construct is equivalent to `mfenced`.
expr_to_mml(n::Fenced) = mml_element(:mrow,
    mml_delim_pairs(n.ldelim, n.rdelim, exlist_to_mml(n.val)))
# MML Binom is a hack, there is no "built-in" MML element for a Binom
expr_to_mml(n::Binom) = mml_element(
    :mrow,
    mml_delim_pairs(
        '(', ')',
        mml_element_attr(
            :mfrac,
            _2_expr_to_mml(n.n, n.k),
            Dict("linethickness"=>"0")
        )))

function expr_to_mml(n::Environment)
    mat_delim = Dict("pmatrix"=>"()",
                    "bmatrix"=>"[]",
                    "vmatrix"=>"||",
                    "Vmatrix"=>"∥∥")
    if n.tag == "matrix"
        return make_mml_table(n.val)
    elseif n.tag in ("pmatrix", "bmatrix", "vmatrix")
        return mml_element(
            :mrow,
            mml_delim_pairs(mat_delim[n.tag][1], mat_delim[n.tag][2], make_mml_table(n.val)))
    elseif n.tag == "verbatim"
        # We quietly ignore the `verbatim` environment for now.
        # Perhaps we don't even need a `:mrow` tag, just return the content
        #return mml_element(:mrow, exlist_to_mml(n.val))
        return exlist_to_mml(n.val)
    else
        @error("Unknown environment: "*n.tag)
    end
end

function make_mml_table(m::Vector{CompoExpr})
    # t: table, r: row, c: cell
    r, c, t = "", "", ""
    for n in m
        if n isa Ampersand
            # create a mml cell whose content is `c`
            r *= mml_element(:mtd, c)
            c = ""
        elseif n isa Newline
            # create a mml row in `r`
            r *= mml_element(:mtd, c)
            t *= mml_element(:mtr, r)
            r, c = "", ""
        else
            # amalgamate elements in a cell
            c *= expr_to_mml(n)
        end
    end
    # finish the last row (assume there is no ampersand or newline
    # behind the last expression in the input vector)
    r *= mml_element(:mtd, c)
    t *= mml_element(:mtr, r)
    return mml_element(:mtable, t)
end

expr_to_mml(n::ID) = mml_element(:mi, n.val)
expr_to_mml(n::NUM) = mml_element(:mn, n.val)
expr_to_mml(n::OPID) = mml_element(:mo, n.val)
expr_to_mml(n::SPACE) = mml_element_attr(:mspace, "", Dict("width"=>n.val))
expr_to_mml(n::Ampersand) = mml_element(:mo, "&")
# The following should be the correct implementation, but as of 2021/7/1
# Firfox does not support the `linebreak="newline"` attribute.
#   "<mspace linebreak=\"newline\" width=\"2ex\" />"
#
# A dirty hack here. the <math> tag is wrapped by LaTeX2MathML.show().
# But we have no idea if the display style is `inline` or `block` here.
expr_to_mml(n::Newline) = "</math><br/><math>"
expr_to_mml(n) = mml_element(:unknown, repr(n))

mml_element(ele, content) = string("<", ele, ">", content, "</", ele, ">")

function mml_element_attr(ele, content, attr::Dict)
    s = string("<", ele, " ")
    for k in keys(attr)
        s *= k*"=\""*attr[k]*"\" "
    end
    s *= ">"*content*"</"*string(ele)*">"
    return s
end

function mml_element_attr(ele, content, attr::String)
    return string("<", ele, " ", attr, ">", content, "</",  ele, ">")
end

mml_delim_pairs(open, close, content) = (
    mml_element(:mo, open)*content*mml_element(:mo, close))

_2_exprs_to_mml(c1, c2) = expr_to_mml(c1)*expr_to_mml(c2)
_3_exprs_to_mml(c1, c2, c3) = _2_exprs_to_mml(c1, c2)*expr_to_mml(c3)

#############################################################################
## pretty_print_mml() is for debug
function pretty_print_mml(mml::String)
    # convert a utf-8 string to a vector of unicode chars so we can
    # index into it easily
    mml_chars = Char[]
    i = 1
    while i <= ncodeunits(mml)
        push!(mml_chars, mml[i])
        i = nextind(mml, i)
    end
    ele_stack = String[]
    t = 1
    while t <= length(mml_chars)
        #skip whitespaces
        while isspace(mml_chars[t]) t += 1 end
        mml_chars[t] != '<' && error("Unknown token: ", mml_chars[t])
        t = pretty_print_element(mml_chars, t, 0, ele_stack)
    end
end

const tbsz = 3
function pretty_print_element(mml::Vector{Char}, start::Int, indent::Int, ele_stack::Vector{String})
    # mml[start] must be '<'
    i, ele_name = start + 1, ""
    # get starting tag
    while mml[i] != '>'
        ele_name *= mml[i]
        i += 1
    end
    i += 1
    push!(ele_stack, ele_name)
    print(" "^indent*"<", ele_name, '>')
    if ele_name in ("mi", "mn", "mo", "mspace")
        # simple leaf elements
        while mml[i] != '<'
            print(mml[i])
            i += 1
        end
        # Now mml[i] == '<' in the ending tag of <mi>, <mn>, <mo>, <mspace>
        while mml[i] != '>' i += 1 end
        i += 1
        i <= length(mml) && while isspace(mml[i]) i += 1 end
        print("</", ele_name, ">")
        println()
        pop!(ele_stack)
        return i
    else
        # the starting tag name occupies a line
        println()
        # keep finding the next tag and print
        while true
            while isspace(mml[i]) i += 1 end
            if i > length(mml)
                break
            end
            if mml[i] == '<' && mml[i+1] != '/'
                i = pretty_print_element(mml, i, indent + tbsz, ele_stack)
            elseif mml[i+1] == '/'
                i += 2
                en2 = ""
                while mml[i] != '>'
                    en2 *= mml[i]
                    i += 1
                end
                en = pop!(ele_stack)
                # we could have sth like <mspace width=...> </mspace>
                !startswith(en, en2) && begin
                    println(ele_stack)
                    error("tag not match: ", en, " -- ", en2)
                end
                # mml[i] == '>'
                println(" "^indent*"</", en2, '>')
                i += 1
                i > length(mml) && return i
                while isspace(mml[i]) i += 1 end
                mml[i] != '<' && error("unknown tag: ", mml[i])
                return i
            else
                error("unknown tag: ", mml[i])
            end
        end
    end
end

###############################################################################
"""Convert a presentation math expression to a LaTeX string"""
exlist_to_tex(nodes::Vector{CompoExpr}) = foldl(*, map(expr_to_tex, nodes), init="")

# unicode identifiers are kept unchanged
expr_to_tex(ex::ID) = ex.val
expr_to_tex(ex::NUM) = ex.val
expr_to_tex(ex::OPID) = ex.val
expr_to_tex(ex::SPACE) = raw"\,"
expr_to_tex(::Ampersand) = "&"
expr_to_tex(::Newline) = "\\\\"

expr_to_tex(n::Sub) = "{$(expr_to_tex(n.base))}_{$(expr_to_tex(n.sub))}"
expr_to_tex(n::Sup) = "{$(expr_to_tex(n.base))}^{$(expr_to_tex(n.sup))}"
expr_to_tex(n::SubSup) =
    "{$(expr_to_tex(n.base))}_{$(expr_to_tex(n.sub))}^{$(expr_to_tex(n.sup))}"
expr_to_tex(n::Under) = "\\underset{$(n.under)}{$(n.base)}"
expr_to_tex(n::Over) = "\\underset{$(n.under)}{$(n.base)}"
expr_to_tex(n::UnderOver) =
    "\\underset{$(n.under)}{\\overset{$(n.over)}{$(n.base)}}"
expr_to_tex(n::Sqrt) = "\\sqrt{$(n.base)}"
expr_to_tex(n::Frac) = "\\frac{$(n.base)}"
expr_to_tex(n::Group) = "{$(exlist_to_tex(n.val))}"
expr_to_tex(n::Fenced) = "\\left$(n.ldelim)$(exlist_to_tex(n.val))\\right$(n.rdelim)"
expr_to_tex(n::Binom) = "\\binom{$(n.n)}{$(n.k)}"

function expr_to_tex(n::Environment)
    return "\\begin{$(n.tag)}"*exlist_to_tex(n.val)*"\\end{$(n.tag)}"
end


text:
using Distributions: Normal, MvNormal
using Models
using Models.TestUtils
using NamedDims: NamedDimsArray
using Test

@testset "Models.jl" begin
    include("traits.jl")
    include("test_utils.jl")
end


text:
###################
# plate generator #
###################

"""
Generator that makes many independent application of a kernel generator,
similar to 'map'.  The arguments are a tuple of vectors, each of of length N,
where N is the nubmer of applications of the kernel.
"""
struct Plate{T,U} <: Generator{PersistentVector{T},VectorTrace{T,U}}
    kernel::Generator{T,U}
end

accepts_output_grad(plate::Plate) = accepts_output_grad(plate.kernel)
has_argument_grads(plate::Plate) = has_argument_grads(plate.kernel)

function plate(kernel::Generator{T,U}) where {T,U}
    Plate{T,U}(kernel)
end

function get_static_argument_types(plate::Plate)
    [Vector{typ} for typ in get_static_argument_types(plate.kernel)]
end

function get_args_for_key(args::Tuple, key::Int)
    map((arg) -> arg[key], args)
end

function get_prev_and_new_lengths(args, trace)
    new_length = length(args[1])
    prev_args = get_call_record(trace).args
    prev_length = length(prev_args[1])
    (new_length, prev_length)
end

struct PlateChange{T}

    # the subset of kernel args which may have changed
    changed_args::Vector{Int}

    # the deltas to pass to each changed invocation of the kernel
    sub_changes::Vector{T}
end

include("simulate.jl")
include("assess.jl")
include("generate.jl")
include("update.jl")
include("extend.jl")
include("project.jl")
#include("backprop_params.jl")
include("backprop_trace.jl")

#include("ungenerate.jl")

export plate
export PlateChange

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen