text:
# Tree construction using pre-computed distances
#==function partition!{T<:Real}(v::Vector{T}, indices::Vector{Int}, i::Int, j::Int, p::T)
    while true
        while v[indices[i]] <= p
            i == j && return j
            i += 1
        end

        while v[indices[j]] > p
            i == j && return j - 1
            j -= 1
        end

        indices[i], indices[j] = indices[j], indices[i]
    end
end

function quickselect!{T<:Real}(v::Vector{T}, indices::Vector{Int}, i::Int, j::Int, k::Int)
    while true
        r = indices[rand(i:j)]
        p = v[r]
        n = partition!(v, indices, i, j, p)
        if (n > k) j = n elseif (n < k) i = n + 1 else return p end
    end
end

function nodedistances!{T<:Real}(_d::Vector{T}, _i::Vector{Int}, distance::Function, i::Int, j::Int, k::Int)
    for s in i:j
        _d[_i[s]] = distance(_i[k], _i[s])::T
    end
end

function Node{T<:Real}(_d::Vector{T}, _i::Vector{Int}, distance::Function, indices::Vector{Int}, i::Int, j::Int)
    j <  i && return Node{T}()
    j == i && return Node{T}(indices[_i[i]], i, j)

    r = rand(i:j)
    ri = _i[r]
    nodedistances!(_d, _i, distance, i, j, r)

    m = (i + j) >> 1
    d = quickselect!(_d, _i, i, j, m)

    left = Node(_d, _i, distance, indices, i, m)
    right = Node(_d, _i, distance, indices, m + 1, j)
    Node{T}(indices[ri], i, j, d, left, right)
end

function VPTree{T<:Real}(_d::Vector{T}, _i::Vector{Int}, distance::Function, indices::Vector{Int}, n::Int)
    init(_i, n)
    root = Node(_d, _i, distance, indices, 1, n)
    VPTree(n, root)
end

VPTree{T<:Real}(_d::Vector{T}, _i::Vector{Int}, distance::Function, indices::Vector{Int}) = VPTree(_d, _i, item_type, indices, length(indices))

function init(_i::Vector{Int}, n::Int)
    @inbounds @simd for i in 1:n
        _i[i] = i
    end
end

workmem{T<:Real}(::Type{VPTree}, n::Int, ::Type{T}) = return (Vector{T}(n), Vector{Int}(n)) ==#

# Tree construction using on the fly distance calculations

# Requirements for partitionu!:
#    p ∉ v
#    min(v) < p < max(v)
#==function partitionu!{T<:Real}(metric::Metric{T}, points, indices::Vector{Int}, vp::Int, d::T, i::Int, j::Int)
    while true
        @inbounds while evaluate(metric, points, vp, indices[i]) < d
            i += 1
        end

        @inbounds while evaluate(metric, points, vp, indices[j]) > d
            j -= 1
        end

        j + 1 == i && return j

        @inbounds indices[i], indices[j] = indices[j], indices[i]
    end
end

# Requirements for quickselectu!:
# j > i
function quickselectu!{T<:Real}(metric::Metric{T}, points, indices::Vector{Int}, vp::Int, i::Int, j::Int, k::Int)
    while true
        local d::T

        println([round(100*evaluate(metric, points, indices[vp], indices[s]), 2) for s in eachindex(indices)]')
        println(indices')
        println("vp: $vp\ti: $i\tj: $j\tk: $k\n")

        if i + 3 > j
            @inbounds d = (evaluate(metric, points, vp, indices[i]) + evaluate(metric, points, vp, indices[j])) / 2
        else
            r1, r2, r3 = rand(i:j-1), rand(i+1:j-1), rand(i+1:j)
            r1 == r3 && (r1 = i)
            @inbounds d = (evaluate(metric, points, vp, indices[r1])
                         + evaluate(metric, points, vp, indices[r2])
                         + evaluate(metric, points, vp, indices[r3])) / 3
        end

        n = partitionu!(metric, points, indices, vp, d, i, j)

        println([round(100*evaluate(metric, points, indices[vp], indices[s]), 2) for s in eachindex(indices)]')
        println(indices')
        println("vp: $vp\ti: $i\tj: $j\tk: $k\tn: $n\td: $(100d)\n\n")

        if n > k j = n elseif n < k i = n + 1 else return d end
    end
end==#

# Requirements for partitionu!:
#    p ∉ v
#    min(v) < p < max(v)
function partitionu!{T<:Real}(v::Vector{T}, indices::Vector{Int}, i::Int, j::Int, p::T)
    while true
        while v[indices[i]] < p
            i += 1
        end

        while v[indices[j]] > p
            j -= 1
        end

        j + 1 == i && return j

        indices[i], indices[j] = indices[j], indices[i]
    end
end

# Requirements for quickselectu!:
# j > i
function quickselectu!{T<:Real}(v::Vector{T}, indices::Vector{Int}, i::Int, j::Int, k::Int)
    while true
        local p::T

        if i + 3 > j
            p = (v[indices[i]] + v[indices[j]]) / 2
        else
            r1, r2, r3 = rand(i:j-1), rand(i+1:j-1), rand(i+1:j)
            r1 == r3 && (r1 = i)
            p = (v[indices[r1]] + v[indices[r2]] + v[indices[r3]]) / 3
        end

        n = partitionu!(v, indices, i, j, p)
        if (n > k) j = n elseif (n < k) i = n + 1 else return p end
    end
end

function knn{T<:Real}(tree::VPTree{T}, node::Node{T}, pq::PriorityQueue{Int, T}, τ::T, p::Int)
    d = distancesq(tree.bw, tree.points, node.index, p)

    if d < τ
        if !haskey(pq, node.index)
            dequeue!(pq)
            pq[node.index] = d
        end
        τ = peek(pq).second
    end

    node.isleaf && return τ

    μ = node.distance
    #if τ + d <= μ
    if !sumsqrtgt(τ, d, μ)
        τ = knn(tree, node.left, pq, τ, p)
    else
        τ = knn(tree, node.right, pq, τ, p)
        #if d <= τ + μ
        if sumsqrtgt(τ, μ, d)
            τ = knn(tree, node.left, pq, τ, p)
        end
    end

    return τ
end

knn{T<:Real}(tree::VPTree{T}, k::Int, pq::PriorityQueue{Int, T}, p::Int) = knn(tree, tree.root, pq, Inf, p)

function knn{T<:Real}(tree::VPTree{T}, k::Int, p::Int)
    pq = PriorityQueue(collect(-1:-1:-k), fill(Inf, k), Base.Order.Reverse)
    knn(tree, k, pq, p)
    return pq
end

knn{T<:Real}(tree::VPTree{T}, k::Int, pq::PriorityQueue{Int, T}, x::T, y::T, z::T) = knn(tree, tree.root, pq, Inf, x, y, z)

function dimmaxspread{T<:AbstractFloat}(bw::T, C::AbstractMatrix{T}, x::T, y::T, z::T)
    hbw = bw / 2
    nhbw = -hbw

    n = size(C, 2)

    xmin, ymin, zmin =  Inf,  Inf,  Inf
    xmax, ymax, zmax = -Inf, -Inf, -Inf

    for i in 1:n
        xd = x - C[1, i]
        yd = y - C[2, i]
        zd = z - C[3, i]
        if xd < nhbw xd += bw elseif xd > hbw xd -= bw end
        if yd < nhbw yd += bw elseif yd > hbw yd -= bw end
        if zd < nhbw zd += bw elseif zd > hbw zd -= bw end
        if xd > xmax xmax = xd elseif xd < xmin xmin = xd end
        if yd > ymax ymax = yd elseif yd < ymin ymin = yd end
        if zd > zmax zmax = zd elseif zd < zmin zmin = zd end
    end

    return (xmax - xmin, ymax - ymin, zmax - zmin)
    return indmax((xmax - xmin, ymax - ymin, zmax - zmin))
end

function furthestpoint{T<:AbstractFloat}(cs::PointSet{T}, bw::T, x::T, y::T, z::T)
    C = cs.coords

    max_i = zero(Int)
    max_r² = zero(T)

    for i in cs.first:cs.last
        r² = distanceSqPeriodic(bw, x, y, z, C[1, i], C[2, i], C[3, i])
        if r² > max_r²
            max_r² = r²
            max_i = i
        end
    end

    return max_i, max_r²
end

function twopoints{T<:AbstractFloat}(cs::PointSet{T}, bw::T, x::T, y::T, z::T)
    C = cs.coords

    i, r² = furthestpoint(cs, bw, x, y, z)
    j_l, r²_l = furthestpoint(PointSet(cs.coords, cs.indices, cs.first, i-1), bw, C[1, i], C[2, i], C[3, i])
    j_r, r²_r = furthestpoint(PointSet(cs.coords, cs.indices, i+1, cs.last), bw, C[1, i], C[2, i], C[3, i])
    j = ifelse(r²_l > r²_r, j_l, j_r)

    return i, j
end

# Rearranges the list of points in C and returns index f such that the first f points
# are closer to point i₁ and the remainder are closer to point i₂
function splitpoints!{T<:AbstractFloat}(cs::PointSet{T}, bw::T, i₁::Int, i₂::Int)
    C = cs.coords
    f = cs.first
    l = cs.last
    while f != l
        r₁ = distanceSqPeriodic(bw, C[1, i₁], C[2, i₁], C[3, i₁], C[1, f], C[2, f], C[3, f])
        r₂ = distanceSqPeriodic(bw, C[1, i₂], C[2, i₂], C[3, i₂], C[1, f], C[2, f], C[3, f])
        if r₂ > r₁
            cs.indices[l], cs.indices[f] = cs.indices[f], cs.indices[l]
            l -= 1
        else
            f += 1
        end
    end

    return f
end

function splitpoints!{T<:AbstractFloat}(cs::PointSet{T}, x::T, y::T, z::T, bw::T, r²_split::T)
    println("Split: ", r²_split)
    C = cs.coords
    f = cs.first
    l = cs.last
    while f != l
        r² = distanceSqPeriodic(bw, C[1, f], C[2, f], C[3, f], x, y, z)
        if r² > r²_split
            cs.indices[l], cs.indices[f] = cs.indices[f], cs.indices[l]
            l -= 1
        else
            f += 1
        end
    end

    return f
end

function mindistsqperiodic{T<:AbstractFloat}(cs::PointSet{T}, bw::T, x::T, y::T, z::T)
    r²max = zero(T)
    C = cs.coords
    for i in cs.first:cs.last
        r² = distanceSqPeriodic(bw, C[1, i], C[2, i], C[3, i], x, y, z)
        if r² > r²max
            r²max = r²
        end
    end

    return r²max
end

function nearestorigin{T<:AbstractFloat}(cs::PointSet{T}, bw::T)
    C = cs.coords
    i_min = zero(Int)
    r²min = Inf

    for i in cs.first:cs.last
        r² = distanceSqPeriodic(bw, C[1, i], C[2, i], C[3, i], zero(T), zero(T), zero(T))
        if r² < r²min
            r²min = r²
            i_min = i
        end
    end

    return i_min
end

volume(r) = (4/3) * π * r^3

function randpoint{T<:AbstractFloat}(cs::PointSet{T}, bw::T)
    C = cs.coords

    i1 = rand(cs.first:cs.last)
    i2 = i1
    while i2 == i1
        i2 = rand(cs.first:cs.last)
    end

    x = (C[1, i1] + C[1, i2])/2
    y = (C[2, i1] + C[2, i2])/2
    z = (C[3, i1] + C[3, i2])/2
    x, y, z = wrap(bw, x, y, z)

    return x, y, z
end

function findgoodpoint{T<:AbstractFloat}(cs::PointSet{T}, bw::T)
    C = cs.coords
    k = length(cs) >> 1 + 1
    i = rand(cs.first:cs.last)
    #rsq = @anon j -> distanceSqPeriodic(bw, C[1, i], C[2, i], C[3, i], C[1, j], C[2, j], C[3, j])
    @inbounds lessthan = @anon (j1, j2) -> j1 < j2
    #@time sort!(cs.indices, cs.first, cs.last, PartialQuickSort(k), Base.Order.Forward)
    #return cs.indices[k]
    @time @inbounds select!(cs.indices, k, lt = isless)
end

function findgoodpoint1{T<:AbstractFloat}(cs::PointSet{T}, bw::T, m::Int, n::Int)
    C = cs.coords
    x, y, z, rsq = randpoint(cs, bw)
    println("($x\t$y\t$z)")
    println("r: $(sqrt(rsq))")
    for p in 1:m
        s = 0
        for q in 1:n
            i = rand(cs.first:cs.last)
            if distanceSqPeriodic(bw, x, y, z, C[1, i], C[2, i], C[3, i]) > rsq
                s += 1
            else
                s -= 1
            end
        end
        x = s/n
        rsq *= 2^(2x/3)
        #rsq *= 1.0 + (2/3)*x - (1/9)*x^2 # + (4/81)*x^3 - (7/243)*x^4
        println("$s/$n\t\t$(sqrt(rsq))")
    end
end

function sortpoints!{T<:AbstractFloat}(cs::PointSet{T}, bw::T, i::T)
    C = cs.coords
    f = cs.first
    l = cs.last
    while f != l
        r² = distanceSqPeriodic(bw, C[1, f], C[2, f], C[3, f], C[1, i], C[2, i], C[3, i])
        if r² > r²_split
            cs.indices[l], cs.indices[f] = cs.indices[f], cs.indices[l]
            l -= 1
        else
            f += 1
        end
    end

    return f
end

function KDTree{T<:AbstractFloat}(C::AbstractMatrix{T}, indices::AbstractVector{Int}, depth::Int)
    k = size(C, 1)
    d = mod(depth - 1, k) + 1
    l = length(indices)
    i = l >>> 1 + 1

    i <= 1 && return
    select!(indices, i, by = j -> C[d, j])
    #println(round(Int, C[indices[1:i-1]]), '\t', round(Int, C[indices[i]]), '\t', round(Int, C[indices[i+1:l]]))
    KDTree(C, sub(indices, 1:i-1), depth + 1, node << 1 + 1)
    KDTree(C, sub(indices, i+1:l), depth + 1, node << 1)
end

function Octree{T<:AbstractFloat}(C::AbstractMatrix{T}, box_width::Float64, cube_width::Float64, NUM_MAX::Int64 = 8)
    bw = box_width
    hbw = bw / 2

    o = Octree(box_width, cube_width, NUM_MAX)

    n = size(C, 2)
    for i in 1:n
        x, y, z = wrap(bw, C[1, i], C[2, i], C[3, i])
        nx = floor(Int, o.cps * ((x + hbw) / bw)) + 1
        ny = floor(Int, o.cps * ((y + hbw) / bw)) + 1
        nz = floor(Int, o.cps * ((z + hbw) / bw)) + 1

        o.num[nx, ny, nz] += 1
        p = o.num[nx, ny, nz]
        p > o.NUM_MAX && error("NUM_MAX not high enough")

        o.indices[p, nx, ny, nz] = i

        o.coords[1, p, nx, ny, nz] = x
        o.coords[2, p, nx, ny, nz] = y
        o.coords[3, p, nx, ny, nz] = z
    end

    return o
end

function cube{T<:Integer}(o::Octree, nx::T, ny::T, nz::T)
    cps = o.cps

    if nx <= 0 nx += cps elseif nx > cps nx -= cps end
    if ny <= 0 ny += cps elseif ny > cps ny -= cps end
    if nz <= 0 nz += cps elseif nz > cps nz -= cps end

    return nx, ny, nz
end

function cube{T<:AbstractFloat}(o::Octree, x::T, y::T, z::T)
    bw = o.box_width
    hbw = bw / 2

    nx = floor(Int, o.cps * ((x + hbw) / bw)) + 1
    ny = floor(Int, o.cps * ((y + hbw) / bw)) + 1
    nz = floor(Int, o.cps * ((z + hbw) / bw)) + 1

    return cube(o, nx, ny, nz)
end

function cube{T<:AbstractFloat}(o::Octree, coord::AbstractVector{T})
    cube(o, coord[1], coord[2], coord[3])
end

#section(x::Float64, b::Float64, n::Int) = floor(Int, n * (x + b/2) / b)

function KDTree{T<:AbstractFloat}(C::AbstractMatrix{T})
    n = size(C, 2)
    indices = collect(1:n)
    KDTree(C, indices, 1, 1)
    levels = floor(Int, log2(n)) + 1
    nodes = Vector{Int}(2^levels)

    return indices, nodes
end

function median!{T<:Real}(v::Vector{T}, indices::Vector{Int}, i::Int, j::Int)
    m = (i + j) >> 1
    return quickselect!(v, indices, i, j, m)
end

function partitionpoints!{T<:AbstractFloat}(distances::Vector{T}, bw::T, C::Matrix{T}, indices::Vector{Int}, i::Int, j::Int, k::Int)
    distances!(distances, bw, C, indices, i, j, k)
    median!(distances, indices, i, j)
end

#println("$(indices')\t$(distances[indices]')\nk: $k\tki: $ki\td: $d\n")


text:
export i_sum, i_mul!, i_dot, i_axpy!, i_umm!, i_norm2

"""
    i_sum(out!, x)

get the sum of `x`.
"""
@i function i_sum(out!, x::AbstractArray)
	@invcheckoff for i=1:length(x)
		@inbounds out! += x[i]
	end
end

@i function i_sum(out!, f, x::AbstractArray)
	@invcheckoff for i=1:length(x)
		@inbounds out! += f(x[i])
	end
end

"""
    i_mul!(out!, x, y)

compute `x * y` (`x` and `y` are matrices, and store results in `out!`.
"""
@i function i_mul!(out!::AbstractMatrix{T}, x::AbstractMatrix{T}, y::AbstractMatrix{T}) where T
	@safe size(x, 2) == size(y, 1) || throw(DimensionMismatch())
	@invcheckoff @inbounds for k=1:size(y,2)
	    for j=1:size(x,2)
		    for i=1:size(x,1)
				out![i,k] += x[i,j] * y[j,k]
			end
		end
	end
end

@i function i_mul!(out!::AbstractVector{T}, x::AbstractMatrix, y::AbstractVector) where T
	@safe size(x, 2) == size(y, 1) || throw(DimensionMismatch())
	@invcheckoff @inbounds for j=1:size(x,2)
		yj ← zero(T)
		yj += y[j]
		for i=1:size(x,1)
			out![i] += x[i,j] * yj
		end
		yj -= y[j]
	end
end

@i function i_dot(out!, x, y)
    @safe @assert length(x) == length(y)
    @invcheckoff @inbounds for i=1:length(x)
        out! += x[i]' * y[i]
    end
end

"""
    i_norm2(out!, x)

get the squared norm of `x`.
"""
@i function i_norm2(out!, x)
    @invcheckoff @inbounds for i=1:length(x)
        out! += abs2(x[i])
    end
end

"""
    i_axpy!(a, x, y!)

compute `y! += a * x`, where `x` and `y` are vectors.
"""
@i function i_axpy!(a, X, Y)
    @safe @assert length(X) == length(Y)
    @invcheckoff @inbounds for i=1:length(Y)
        Y[i] += a * X[i]
    end
end

"""
    i_umm!(x!, θ)

Compute unitary matrix multiplication on `x`, where the unitary matrix is parameterized by (N+1)*N/2 `θ`s.
"""
@i function i_umm!(x!::AbstractArray, θ)
    M ← size(x!, 1)
    N ← size(x!, 2)
    k ← 0
    @safe @assert length(θ) == M*(M-1)/2
    for l = 1:N
        for j=1:M
            for i=M-1:-1:j
                INC(k)
                ROT(x![i,l], x![i+1,l], θ[k])
            end
        end
    end

    k → length(θ)
end

ext: jl


text:
module TestFilterCandidates

import Wordle
import Test: @test, @test_throws

constraints = Wordle.Constraints()
trial = Wordle.Trial(Wordle.Candidate("alone"), Wordle.Result("⬜⬜⬜⬜⬜"))
Wordle.update_constraints!(constraints, trial)

candidates = Wordle.candidate_pool(true)

@test Wordle.Candidate("alone") in candidates
Wordle.filter_candidates!(candidates, constraints)
@test !(Wordle.Candidate("alone") in candidates)

end # module


text:
"""model definition of point3d"""

struct StatePoint3D <: StatePoint
    x::Float64
    y::Float64
    z::Float64
end

function get_mid_status(p::StatePoint3D, q::StatePoint3D)::StatePoint3D
    return StatePoint3D((p.x + q.x) / 2, (p.y + q.y) / 2, (p.z + q.z) / 2)
end

function dist(q1::StatePoint3D, q2::StatePoint3D)::Float64
    return norm([q1.x - q2.x, q1.y - q2.y, q1.z - q2.z])
end

function dist(q_from::StatePoint3D, q_to::StatePoint3D, o::CircleObstacle3D)::Float64
    return dist([q_from.x, q_from.y, q_from.z], [q_to.x, q_to.y, q_to.z], [o.x, o.y, o.z])
end

function dist(q_from::StatePoint3D, q_to::StatePoint3D, q_static::StatePoint3D)::Float64
    return dist(
        [q_from.x, q_from.y, q_from.z],
        [q_to.x, q_to.y, q_to.z],
        [q_static.x, q_static.y, q_static.z],
    )
end

function dist(q::StatePoint3D, o::CircleObstacle3D)::Float64
    return norm([q.x - o.x, q.y - o.y, q.z - o.z])
end

function dist(
    a_from::StatePoint3D,
    a_to::StatePoint3D,
    b_from::StatePoint3D,
    b_to::StatePoint3D,
)::Float64
    return dist(
        [a_from.x, a_from.y, a_from.z],
        [a_to.x, a_to.y, a_to.z],
        [b_from.x, b_from.y, b_from.z],
        [b_to.x, b_to.y, b_to.z],
    )
end

function gen_connect(
    q::StatePoint3D,  # to identify type
    obstacles::Vector{CircleObstacle3D},
    rads::Vector{Float64},
)::Function

    # check: q \in C_free
    f(q::StatePoint3D, i::Int64)::Bool = begin
        if any(x -> (x < rads[i] || 1 - rads[i] < x), [q.x, q.y, q.z])
            return false
        end

        if any([dist(q, o) < o.r + rads[i] for o in obstacles])
            return false
        end

        return true
    end

    f(q_from::StatePoint3D, q_to::StatePoint3D, i::Int64) = begin
        if any(x -> (x < rads[i] || 1 - rads[i] < x), [q_to.x, q_to.y, q_to.z])
            return false
        end

        # check: collisions with static obstacles
        if any([dist(q_from, q_to, o) < o.r + rads[i] for o in obstacles])
            return false
        end

        return true
    end

    return f
end

function gen_uniform_sampling(q::StatePoint3D)::Function
    () -> begin
        return StatePoint3D(rand(3)...)
    end
end

function plot_motion!(q_from::StatePoint3D, q_to::StatePoint3D, rad::Float64, params)
    plot!([q_from.x, q_to.x], [q_from.y, q_to.y], [q_from.z, q_to.z]; params...)
end

function plot_start_goal!(q_init::StatePoint3D, q_goal::StatePoint3D, rad::Float64, params)
    plot!([q_init.x], [q_init.y], [q_init.z]; markershape = :hex, params...)
    plot!([q_goal.x], [q_goal.y], [q_goal.z]; markershape = :star, params...)
end

function plot_agent!(q::StatePoint3D, rad::Float64, color::String)
    plot_sphere!(q.x, q.y, q.z, rad, color)
end

function gen_random_instance_StatePoint3D(; params...)
    return gen_random_instance(StatePoint3D(0, 0, 0); params...)
end


text:
# # Hydropower Simulations with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl)

# **Originally Contributed by**: Clayton Barrows and Sourabh Dalvi

# ## Introduction

# PowerSimulations.jl supports simulations that consist of sequential optimization problems
# where results from previous problems inform subsequent problems in a variety of ways.
# This example demonstrates a few of the options for modeling hydropower generation.

# ## Dependencies
using SIIPExamples
pkgpath = dirname(dirname(pathof(SIIPExamples)))

# ### Modeling Packages
using PowerSystems
using PowerSimulations
const PSI = PowerSimulations
using D3TypeTrees

# ### Data management packages
using Dates
using DataFrames

# ### Optimization packages
using Cbc # solver
solver = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 1, "ratioGap" => 0.05)

# ### Data
# PowerSystems.jl links to some meaningless test data that is suitable for this example.
# The [make_hydro_data.jl](../../script/3_PowerSimulations_examples/make_hydro_data.jl) script
# loads three systems suitable for the examples here.

include(joinpath(pkgpath, "script", "3_PowerSimulations_examples", "make_hydro_data.jl"))

# This line just overloads JuMP printing to remove double underscores added by PowerSimulations.jl
PSI.JuMP._wrap_in_math_mode(str) = "\$\$ $(replace(str, "__"=>"")) \$\$"

# ## Two PowerSimulations features determine hydropower representation.
# There are two principal ways that we can customize hydropower representation in
# PowerSimulations. First, we can play with the formulation applied to hydropower generators
# using the `DeviceModel`. We can also adjust how simulations are configured to represent
# different decision making processes and the information flow between those processes.

# ### Hydropower `DeviceModel`s

# First, the assignment of device formulations to particular device types gives us control
# over the representation of devices. This is accomplished by defining `DeviceModel`
# instances. For hydro power representations, we have two available generator types in
# PowerSystems:

#nb TypeTree(PowerSystems.HydroGen)

# And in PowerSimulations, we have several available formulations that can be applied to
# the hydropower generation devices:

#nb TypeTree(PSI.AbstractHydroFormulation, scopesep = "\n", init_expand = 5)

# Let's see what some of the different combinations create. First, let's apply the
# `HydroDispatchRunOfRiver` formulation to the `HydroEnergyReservoir` generators, and the
# `FixedOutput` formulation to `HydroDispatch` generators.
#  - The `FixedOutput` formulation just acts
# like a load subtractor, forcing the system to accept it's generation.
#  - The `HydroDispatchRunOfRiver` formulation represents the the energy flowing out of
# a reservoir. The model can choose to produce power with that energy or just let it spill by.

devices = Dict{Symbol, DeviceModel}(
    :Hyd1 => DeviceModel(HydroEnergyReservoir, HydroDispatchRunOfRiver),
    :Hyd2 => DeviceModel(HydroDispatch, FixedOutput),
    :Load => DeviceModel(PowerLoad, StaticPowerLoad),
);

template = PSI.OperationsProblemTemplate(CopperPlatePowerModel, devices, Dict(), Dict());

op_problem = PSI.OperationsProblem(GenericOpProblem, template, c_sys5_hy_uc, horizon = 24)

# Now we can see the resulting JuMP model:

op_problem.psi_container.JuMPmodel

# The first two constraints are the power balance constraints that require the generation
# from the controllable `HydroEnergyReservoir` generators to be equal to the load (flat 10.0 for all time periods)
# minus the generation from the `HydroDispatch` generators [1.97, 1.983, ...]. The 3rd and 4th
# constraints limit the output of the `HydroEnergyReservoir` generator to the limit defined by the
# `max_activepwoer` time series. And the last 4 constraints are the lower and upper bounds of
# the `HydroEnergyReservoir` operating range.

#-

# Next, let's apply the `HydroDispatchReservoirBudget` formulation to the `HydroEnergyReservoir` generators.
devices = Dict{Symbol, DeviceModel}(
    :Hyd1 => DeviceModel(HydroEnergyReservoir, HydroDispatchReservoirBudget),
    :Load => DeviceModel(PowerLoad, StaticPowerLoad),
);

template = PSI.OperationsProblemTemplate(CopperPlatePowerModel, devices, Dict(), Dict());

op_problem = PSI.OperationsProblem(GenericOpProblem, template, c_sys5_hy_uc, horizon = 24)

# And, the resulting JuMP model:

op_problem.psi_container.JuMPmodel

# Finally, let's apply the `HydroDispatchReservoirStorage` formulation to the `HydroEnergyReservoir` generators.

devices = Dict{Symbol, DeviceModel}(
    :Hyd1 => DeviceModel(HydroEnergyReservoir, HydroDispatchReservoirStorage),
    :Load => DeviceModel(PowerLoad, StaticPowerLoad),
);

template = PSI.OperationsProblemTemplate(CopperPlatePowerModel, devices, Dict(), Dict());

op_problem = PSI.OperationsProblem(GenericOpProblem, template, c_sys5_hy_uc, horizon = 24)

#-

op_problem.psi_container.JuMPmodel

# ### Multi-Stage `SimulationSequence`
# The purpose of a multi-stage simulation is to represent scheduling decisions consistently
# with the time scales that govern different elements of power systems.

# #### Multi-Day to Daily Simulation:
# In the multi-day model, we'll use a really simple representation of all system devices
# so that we can maintain computational tractability while getting an estimate of system
# requirements/capabilities.

devices = Dict(
    :Generators => DeviceModel(ThermalStandard, ThermalDispatch),
    :Loads => DeviceModel(PowerLoad, StaticPowerLoad),
    :HydroEnergyReservoir =>
        DeviceModel(HydroEnergyReservoir, HydroDispatchReservoirStorage),
)
template_md = OperationsProblemTemplate(CopperPlatePowerModel, devices, Dict(), Dict());

# For the daily model, we can increase the modeling detail since we'll be solving shorter
# problems.

devices = Dict(
    :Generators => DeviceModel(ThermalStandard, ThermalDispatch),
    :Loads => DeviceModel(PowerLoad, StaticPowerLoad),
    :HydroEnergyReservoir =>
        DeviceModel(HydroEnergyReservoir, HydroDispatchReservoirStorage),
)
template_da = OperationsProblemTemplate(CopperPlatePowerModel, devices, Dict(), Dict());

#-

stages_definition = Dict(
    "MD" => Stage(
        GenericOpProblem,
        template_md,
        c_sys5_hy_wk,
        solver,
        system_to_file = false,
    ),
    "DA" => Stage(
        GenericOpProblem,
        template_da,
        c_sys5_hy_uc,
        solver,
        system_to_file = false,
    ),
)

# This builds the sequence and passes the the energy dispatch schedule for the `HydroEnergyReservoir`
# generator from the "MD" stage to the "DA" stage in the form of an energy limit over the
# synchronized periods.

sequence = SimulationSequence(
    step_resolution = Hour(48),
    order = Dict(1 => "MD", 2 => "DA"),
    feedforward_chronologies = Dict(("MD" => "DA") => Synchronize(periods = 2)),
    horizons = Dict("MD" => 2, "DA" => 24),
    intervals = Dict("MD" => (Hour(48), Consecutive()), "DA" => (Hour(24), Consecutive())),
    feedforward = Dict(
        ("DA", :devices, :HydroEnergyReservoir) => IntegralLimitFF(
            variable_source_stage = PSI.ACTIVE_POWER,
            affected_variables = [PSI.ACTIVE_POWER],
        ),
    ),
    cache = Dict(("MD", "DA") => StoredEnergy(HydroEnergyReservoir, PSI.ENERGY)),
    ini_cond_chronology = IntraStageChronology(),
);

#-

file_path = tempdir()

sim = Simulation(
    name = "hydro",
    steps = 1,
    stages = stages_definition,
    stages_sequence = sequence,
    simulation_folder = file_path,
)

build!(sim)

# We can look at the "MD" Model

sim.stages["MD"].internal.psi_container.JuMPmodel

# And we can look at the "DA" model

sim.stages["DA"].internal.psi_container.JuMPmodel

# And we can execute the simulation by running the following command
# ```julia
# sim_results = execute!(sim)
# ```
#-

# #### 3-Stage Simulation:

stages_definition = Dict(
    "MD" => Stage(
        GenericOpProblem,
        template_md,
        c_sys5_hy_wk,
        solver,
        system_to_file = false,
    ),
    "DA" => Stage(
        GenericOpProblem,
        template_da,
        c_sys5_hy_uc,
        solver,
        system_to_file = false,
    ),
    "ED" => Stage(
        GenericOpProblem,
        template_da,
        c_sys5_hy_ed,
        solver,
        system_to_file = false,
    ),
)

sequence = SimulationSequence(
    step_resolution = Hour(48),
    order = Dict(1 => "MD", 2 => "DA", 3 => "ED"),
    feedforward_chronologies = Dict(
        ("MD" => "DA") => Synchronize(periods = 2),
        ("DA" => "ED") => Synchronize(periods = 24),
    ),
    intervals = Dict(
        "MD" => (Hour(48), Consecutive()),
        "DA" => (Hour(24), Consecutive()),
        "ED" => (Hour(1), Consecutive()),
    ),
    horizons = Dict("MD" => 2, "DA" => 24, "ED" => 12),
    feedforward = Dict(
        ("DA", :devices, :HydroEnergyReservoir) => IntegralLimitFF(
            variable_source_stage = PSI.ACTIVE_POWER,
            affected_variables = [PSI.ACTIVE_POWER],
        ),
        ("ED", :devices, :HydroEnergyReservoir) => IntegralLimitFF(
            variable_source_stage = PSI.ACTIVE_POWER,
            affected_variables = [PSI.ACTIVE_POWER],
        ),
    ),
    cache = Dict(("MD", "DA") => StoredEnergy(HydroEnergyReservoir, PSI.ENERGY)),
    ini_cond_chronology = IntraStageChronology(),
);

#-

sim = Simulation(
    name = "hydro",
    steps = 1,
    stages = stages_definition,
    stages_sequence = sequence,
    simulation_folder = file_path,
)

#-

build!(sim)

# We can look at the "MD" Model

sim.stages["MD"].internal.psi_container.JuMPmodel

# And we can look at the "DA" model

sim.stages["DA"].internal.psi_container.JuMPmodel

# And we can look at the "ED" model

sim.stages["ED"].internal.psi_container.JuMPmodel

# And we can execute the simulation by running the following command
# ```julia
# sim_results = execute!(sim)
# ```
#-


text:
using Plots
theme(:default)
pyplot()

using DelimitedFiles

push!(LOAD_PATH,"src")
using DalitzPlotAnalysis
using amplitudes_compass
using PWAHelper
using SDMHelper
using FittingPWALikelihood

mass_bin_name = "1540_1560" # "2320_2340"
tslice = "t1"
# set names
for app in ["rd", "mc", "fu"]
    @eval $(Symbol("kinvar_"*app)) = joinpath("data","variables_$(mass_bin_name)_$(tslice)_"*$app*".txt")
    @eval $(Symbol("basisfunc_"*app)) = joinpath("data","functions_$(mass_bin_name)_$(tslice)_"*$app*".txt")
end
# get number of lines
nlines_rd = Meta.parse(split(read(`wc -l $(basisfunc_rd)`,String)," ")[1])
nlines_mc = Meta.parse(split(read(`wc -l $(basisfunc_mc)`,String)," ")[1])
nlines_fu = Meta.parse(split(read(`wc -l $(basisfunc_fu)`,String)," ")[1])
normfact = nlines_fu/nlines_mc*nlines_rd

##### Converting data data #####

isfile("data/$(mass_bin_name)_$(tslice)_rd.root") || error("Cannot find the file")
isfile("data/$(mass_bin_name)_$(tslice)_mc.root") || error("Cannot find the file")
@time let tmin = "0.1", tmax = "0.112853"
    run(pipeline(`src/Extract data/$(mass_bin_name)_$(tslice)_rd.root D $(tmin) $(tmax)`, stdout=kinvar_rd))
    run(pipeline(`src/Extract data/$(mass_bin_name)_$(tslice)_mc.root M $(tmin) $(tmax)`, stdout=kinvar_mc))
    run(pipeline(`src/Extract data/$(mass_bin_name)_$(tslice)_mc.root F $(tmin) $(tmax)`, stdout=kinvar_fu))
end

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

wavelist = get_wavelist(joinpath(pwd(),"src/wavelist_formated.txt");
    path_to_thresholds=joinpath(pwd(),"src/thresholds_formated.txt"),
    M3pi=Meta.parse(mass_bin_name[1:4])/1000)

Nwaves = size(wavelist,1)

wavenames = get_wavenames(wavelist)
wavebasis = get_wavebasis(wavelist)

@time precalculate_compass_basis(wavebasis, kinvar_rd, basisfunc_rd)
@time precalculate_compass_basis(wavebasis, kinvar_mc, basisfunc_mc)
@time precalculate_compass_basis(wavebasis, kinvar_fu, basisfunc_fu)

@time const BmatMC = let
    @time PsiMC = read_precalc_basis(basisfunc_mc)
    PsiTMC = transpose(PsiMC)
    @views v = sum(PsiTMC[:,e] * PsiTMC[:,e]' for e in 1:size(PsiMC,1)) ./ size(PsiMC,1);
    conj(v)
end
write_cmatrix(BmatMC, "data/integrmat_$(mass_bin_name)_$(tslice)_mc.txt")
const BmatMC = read_cmatrix("data/integrmat_$(mass_bin_name)_$(tslice)_mc.txt")

const noϵ =  [1] # flat wave
const posϵ = [i for (i,ϵ) in enumerate(wavelist[:,6]) if ϵ=="+"]
const negϵ = [i for (i,ϵ) in enumerate(wavelist[:,6]) if ϵ=="-"]

# Model description
const ModelBlocks = [noϵ, posϵ, negϵ, negϵ] #

# load precalculated data array
const PsiRD = read_precalc_basis(
    joinpath("data","functions_$(mass_bin_name)_$(tslice)_rd.txt"));

LLH, GRAD, LLH_and_GRAD!, HES = createLLHandGRAD(PsiRD, BmatMC, ModelBlocks);

const Npar = get_npars(ModelBlocks)

const pars0 = rand(Npar);
pars0 = rand(Npar);
pars0 .*= get_parameter_ranges(BmatMC, ModelBlocks)
normalize_pars!(pars0, BmatMC, ModelBlocks)

# minpars0 = vcat(readdlm("minpars_compass_$(mass_bin_name).txt")...)[1:Npar];
@time minpars = minimize(LLH, LLH_and_GRAD!;
    algorithm = :LD_LBFGS, verbose=1, starting_pars=pars0,
    llhtolerance=1e-4)

get_intesity(pars0,   BmatMC, ModelBlocks)
get_intesity(minpars, BmatMC, ModelBlocks)

minpars_new = minpars

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

BmatMC_n = [BmatMC[i,j]/sqrt(BmatMC[i,i]*BmatMC[j,j]) for i=1:Nwaves, j=1:Nwaves];
let BmatMC_n = [BmatMC[i,j]/sqrt(BmatMC[i,i]*BmatMC[j,j]) for i=1:Nwaves, j=1:Nwaves];
    heatmap(real(BmatMC_n))
end
PsiRD_n = copy(PsiRD)
for i in 1:size(PsiRD_n,2)
    PsiRD_n[:,i] .*= 1.0/sqrt(BmatMC[i,i])
end

LLH, GRAD, LLH_and_GRAD!, HES = createLLHandGRAD(PsiRD_n, BmatMC_n, ModelBlocks);

const pars0 = rand(Npar);
pars0 = rand(Npar);
pars0 .*= get_parameter_ranges(BmatMC_n, ModelBlocks)
normalize_pars!(pars0, BmatMC_n, ModelBlocks)

# minimization
@time minpars = minimize(LLH, LLH_and_GRAD!;
    algorithm = :LD_LBFGS, verbose=1, starting_pars=pars0,
    llhtolerance=1e-4)

# check of normalization
get_intesity(pars0, BmatMC_n, ModelBlocks)
get_intesity(minpars, BmatMC_n, ModelBlocks)

usual_minpars = minpars ./ abs.(extnd([sqrt(BmatMC[i,i]) for i in 1:size(BmatMC,1)],
    get_parameter_map(ModelBlocks, Nwaves)))

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

@time hes_anal = HES(minpars)
inv_hes = inv(hes_anal)
det(hes_anal[1:20,1:20])
diag(inv_hes)
diag_error = sqrt.(diag(inv_hes))

writedlm("invhes_compass_$(mass_bin_name)_$(tslice).txt", inv_hes)

plot(abs.(minpars[2:end-26]))

let tog = [[minpars[i],diag_error[i]] for i in 1:(length(minpars)-26)]
    stog = sort(tog, by=x->abs(x[1]), rev=true)
    hcat_stog = hcat(stog...)
    plot(abs.(hcat_stog[1,:]), yerr = hcat_stog[2,:], ylim=(0,.1))
end
savefig(joinpath("plots","official_comass_fit_parameters_nonegref.pdf"))

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

# start from random points on the contrained surface
for e in 13:100
    rand_pars =  get_rand_starting_values(1.5.*limits)
    normalize_pars!(rand_pars, BmatMC, ModelBlocks)
    @time pars_in_min = minimize(LLH, LLH_and_GRAD!;
        algorithm = :LD_LBFGS, verbose=1, starting_pars=rand_pars)
    writedlm("data/studies_of_minimas/m$(e).txt", pars_in_min)
end

minima = hcat([readdlm("data/studies_of_minimas/m$(e).txt")[:,1] for e in 1:95]...)
minimaLLH = [LLH(minima[:,i]) for i=1:size(minima,2)]

collect(1:95)[minimaLLH.-LLH(minpars0) .> 5]
histogram(minimaLLH.-LLH(minpars0), bins=linspace(-2,5,10))

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

@time weights = let mpars = minpars,
    Tmap = get_parameter_map(ModelBlocks, Nwaves),
    pblocks = make_pblock_inds(ModelBlocks)
        [cohsq(extnd(@view(PsiMC[i,:]),Tmap).*mpars,pblocks) for i in 1:size(PsiMC,1)];
end;

const MC1 = readdlm(kinvar_mc);
@time const MC3 = hcat([swap_kin_parameters(MC1[e,:]...) for e in 1:size(MC1,1)]...)';

histogram(sqrt.(vcat(MC1[:,2],MC3[:,2])), weights=vcat(weights,weights), bins=(linspace(0.3,2.2,100)))

const DT1 = readdlm(kinvar_rd);
@time const DT3 = hcat([swap_kin_parameters(DT1[e,:]...) for e in 1:size(DT1,1)]...)';

histogram2d(DT1[:,2], DT3[:,2], bins=linspace(0,6,100))

stephist(sqrt.(vcat(MC1[:,2],MC3[:,2])), weights=vcat(weights,weights)/size(MC1,1)*size(DT1,1),
    bins=linspace(0.3,2.2,100),lab="weighted MC", size=(800,500), xlab="M3pi (GeV)")
let v = stephist!(sqrt.(vcat(DT1[:,2],DT3[:,2])), bins=linspace(0.3,2.2,100),lab="data")
    yarr = v.subplots[1][2][:y]
    xarr = v.subplots[1][2][:x]
    @show length(xarr)
    x2 = [(xarr[2i]+xarr[2(i+1)])/2 for i in 1:(div(length(xarr),2)-1)];
    y2 = [yarr[2i] for i in 1:(div(length(yarr),2)-1)];
    plot!(x2, y2, err=sqrt.(y2), lab="", marker = (0, stroke(0.5, :gray)), l=nothing)
end
savefig(joinpath("plots","data_with_errors.png"))
savefig(joinpath("plots","data_with_errors.pdf"))

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

@time const BmatFU = let
    @time const PsiFU = read_precalc_basis(basisfunc_fu);
    PsiTFU = transpose(PsiFU)
    @views v = sum(PsiTFU[:,e] * PsiTFU[:,e]' for e in 1:size(PsiTFU,2)) ./ size(PsiFU,1);
    conj(v)
end
write_cmatrix(BmatFU, "data/integrmat_$(mass_bin_name)_$(tslice)_fu.txt")
BmatFU = read_cmatrix("data/integrmat_$(mass_bin_name)_$(tslice)_fu.txt")

BmatFU_n = [BmatFU[i,j]/sqrt(BmatFU[i,i]*BmatFU[j,j]) for i=1:Nwaves, j=1:Nwaves];
let BmatFU_n = [BmatFU[i,j]/sqrt(BmatFU[i,i]*BmatFU[j,j]) for i=1:Nwaves, j=1:Nwaves];
    heatmap(real(BmatFU_n))
end

# pars_to_SDM(minpars, BmatFU, ModelBlocks)

SDM = normfact*pars_to_SDM(minpars, BmatFU, ModelBlocks)

diag(SDM)

minpars./SDM_to_pars(SDM/size(PsiRD,1), BmatFU, ModelBlocks)

path_to_SDM = "/mnt/data/compass/2008/pwa_results/SDMs/0.100000-0.112853/";
SDM_RD = read_compass_SDM(joinpath(path_to_SDM,"sdm92."),
    path_to_wavelist   = joinpath("src", "wavelist_formated.txt");
    path_to_thresholds = joinpath("src","thresholds_formated.txt"),
    M3pi = Meta.parse(mass_bin_name[1:4])/1000)

SDM_RD_err = read_compass_SDM(joinpath(path_to_SDM,"sdm92-err."),
    path_to_wavelist   = joinpath("src", "wavelist_formated.txt");
    path_to_thresholds = joinpath("src","thresholds_formated.txt"),
    M3pi = Meta.parse(mass_bin_name[1:4])/1000)

minpars_rd = SDM_to_pars(SDM_RD/size(PsiRD,1), BmatFU, ModelBlocks)
@time minpars_from_rd = minimize(LLH, LLH_and_GRAD!;
    algorithm = :LD_LBFGS, verbose=1, starting_pars=minpars_rd)
minpars = minpars_from_rd

plot(hcat(real.(diag(SDM_RD)), real.(diag(SDM))), lab=["F.H.-D.R." "M.M."],
    xlab="# wave", title = "Diagonal of the SDM", size=(800,500))
savefig(joinpath("plots","sdm_results.png"))
savefig(joinpath("plots","sdm_results.pdf"))

let rd=real.(diag(SDM_RD)), wn = wavenames
    sort([(wn[i],rd[i]) for i in 1:length(rd)], by=x->x[2],rev=true)
end
let
    plot(real.(SDM_RD[2,:]), frame=:origin)
    plot!(real.(SDM[2,:]))
end
plot(hcat(real.(SDM_RD[2,:]), real.(SDM[2,:])), lab=["F.H.-D.R." "M.M."],
    xlab="# wave", title = "Real part of the second row of SDM", size=(800,500))
savefig(joinpath("plots","sdm2_results.png"))
savefig(joinpath("plots","sdm2_results.pdf"))

sum(diag(SDM))
sum(diag(SDM_RD))

plotPolarSDM(SDM, cutoff_scale=0.55)
savefig(joinpath("data","arrow_intensities.pdf"))

#####################################################################################
#####################################################################################
v = let Nb = 500
    # path to save
    path_to_tmp_res = joinpath("data","bootstrap_tmp")
    minpars0 = vcat(readdlm("minpars_compass_$(mass_bin_name).txt")...);
    # res = Matrix{Float64}(Nb,length(minpars0))
    # llh = Vector{Float64}(Nb)
    # b = 16
    @parallel for b in 1:Nb
        @show "Doing something"
        Nd = size(PsiRD,1)
        const _PsiRD = hcat([PsiRD[rand(1:Nd),:] for i in 1:Nd]...).'
        # @show size(_PsiRD)
        _LLH, _GRAD, _LLH_and_GRAD!, _HES = FittingPWALikelihood.createLLHandGRAD(_PsiRD, BmatMC, ModelBlocks);
        try
            @time _minpars = FittingPWALikelihood.minimize(_LLH, _LLH_and_GRAD!;
                verbose=0, starting_pars=minpars0) #,algorithm=:LD_MMA
            writedlm(joinpath(path_to_tmp_res,"BootstrapResults-$(b).txt"), _minpars)

            writedlm(joinpath(path_to_tmp_res,"llh-$(b).txt"), _LLH(_minpars))
        catch TheException
            @show TheException
        end
    end
end
# fetch(v[2])
# BootstrapResults = let Nb = 500, path_to_tmp_res = joinpath("data","bootstrap_tmp")
#     res = []
#     for b in 1:Nb
#         file = joinpath(path_to_tmp_res,"BootstrapResults-$(b).txt")
#         isfile(file) && push!(res, readdlm(file))
#     end
#     hcat(res...)'
# end
# writedlm(joinpath("data","bootstrap_tmp","bootstrap_main_$(mass_bin_name).txt"), BootstrapResults)
# llh = let Nb = 578, path_to_tmp_res = joinpath("data","bootstrap_tmp")
#     res = []
#     for b in 1:Nb
#         file = joinpath(path_to_tmp_res,"llh-$(b).txt")
#         isfile(file) && push!(res, readdlm(file))
#     end
#     hcat(res...)'[:,1]
# end
# writedlm(joinpath("data","bootstrap_tmp","bootstrap_llh_$(mass_bin_name).txt"), llh)
@time for b in 1:size(BootstrapResults,1)
    _br = BootstrapResults[b,:]
    _sdm = size(PsiRD,1)*pars_to_SDM(_br, BmatFU, ModelBlocks)
    path_to_tmp_res = joinpath("data","bootstrap_tmp","SDMs-1540_1560")
    writedlm(joinpath(path_to_tmp_res,"rb-$(b)-sdm.txt"), [real(_sdm) imag(_sdm)])
end
SDMs = [read_cmatrix(joinpath("data","bootstrap_tmp","SDMs_$(mass_bin_name)","rb-$(i)-sdm.txt"))
    for i in 1:497]

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

histogram(BootstrapResults[:,2], bins=100)
histogram(BootstrapResults[:,1], bins=100)

using PlotHelper
plotBSTsample(2, SDMs, SDM, SDM_RD, SDM_RD_err)
arg(z) = atan2(imag(z),real(z))
plotBSTsample(x->180/π*arg(x[2,15]*cis(π/2))-90, SDMs, SDM, SDM_RD, SDM_RD_err)
plotBSTsample(x->180/π*arg(x[2,15]), SDMs, SDM, SDM_RD, SDM_RD_err)

writedlm("/tmp/phase215.txt",broadcast(x->180/π*arg(x[2,15]*cis(π/2))-90, SDMs))

combres = vcat([constract_values(i, SDMs, SDM, SDM_RD, SDM_RD_err) for i in 1:Nwaves]...);
plotBSTsummary(combres) #  tosort=true, toannotate=true
savefig(joinpath("plots","bootstrap_combined_$(mass_bin_name)_t1.pdf"))

saveBSTSplot(joinpath("plots","bootstrap_summary_2300_t1.pdf"),
    SDMs, SDM, SDM_RD, SDM_RD_err)

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

plot(hcat(real.(diag(SDM_RD)), real.(diag(SDM1)), real.(diag(SDM9))),
    lab=["official" "min1" "min9"], lc=[:red :orange :green],
    xlab="# wave", ylab = "Diagonal of the SDM")
savefig("plots/two_minima_diagonal.pdf")

histogram([real(s[13,13]) for s in SDMs], bins=100, xlab=wavenames[13][3:end],
    c=:grey, lab="bootstrap")
vline!(broadcast(x->real(x[13,13]),[SDM_RD, SDM1, SDM9])',
    lab=["official" "min1" "min9"], lw = 2, lc=[:red :orange :green])
savefig("plots/two_minimas_sdm13.pdf")

histogram(broadcast(x->180/π*arg(x[2,15]*cis(π/2))-90+360,SDMs), bins=100,
    c=:grey, xlab="1++0+: f0(980)piP - rhopiS (deg.)", ylab = "Bootstrap Entries",
    lab="bootstrap")
vline!(broadcast(x->180/π*arg(x[2,15]*cis(π/2))-90+360,[SDM_RD, SDM1, SDM9])',
    lab=["official" "min1" "min9"], lw = 2, lc=[:red :orange :green])
savefig("plots/two_minimas_sdm215.pdf")


text:
using EffectiveWaves, Test

# This case is numerically challenging, because wavenumber() has many roots close together. Make sure spacing in ωs is small to help the optimisation method
@testset "strong scatterers and low frequency" begin
    medium = Medium(ρ=1.0, c=1.0)
    species = [
        Specie(ρ=5.,r=0.004, c=0.002, volfrac=0.2),
        Specie(ρ=0.3, r=0.002, c=0.01, volfrac=0.1)
    ]
    tol=1e-7
    # ωs = [0.001,0.003]
    ωs = [0.001]
    eff_medium = effective_medium(medium, species)
    k_eff_lows = ωs./eff_medium.c

    k_eff_φs = wavenumber_low_volfrac(ωs, medium, species)
    # num_wavenumbers =1 almost always finds the wavenubmer with the smallest attenuation

    k_effs_arr = [
        wavenumbers(ω, medium, species; tol=tol, num_wavenumbers=1)
    for ω in ωs]

    inds = [argmin(abs.(k_effs_arr[i] .- k_eff_φs[i])) for i in eachindex(ωs)]
    k_effs2 = [k_effs_arr[i][inds[i]] for i in eachindex(inds)]

    @test norm(k_effs2 - k_eff_lows)/norm(k_effs2) < 5e-7
    @test norm(k_effs2[1] - k_eff_lows[1])/norm(k_eff_lows[1]) < 5e-7
    @test norm(k_effs2 - k_eff_φs)/norm(k_effs2) < 0.01
end


text:
module Converting
  using ..Sources

  function convert(source::ConlluSource)
    parsed_sentences = []

    parse_sentence = (sentence_lines) -> begin
      words = filter(word -> (length(word) > 1) && !(occursin(r"^(# |(\d+\-\d+)|(\d+\.\d+))", word)), sentence_lines) .|>
        (word -> split(word, r"\t")) |>
        (words -> map(word -> [word[1], word[2], word[7], word[8]], words))

      words = map(words) do word
        dependent_word_number = word[1]
        dependent_word = word[2]
        dependent_word_label = string(dependent_word, "-", dependent_word_number)

        head_word_number = word[3]
        if head_word_number == "0"
          head_word = "ROOT"
        else
          head_word = words[findfirst(hword -> hword[1] == word[3], words)][2]
        end
        head_word_label = string(head_word, "-", head_word_number)

        dependency = word[4]

        return join([dependency, head_word_label, dependent_word_label], " ")
      end

      push!(parsed_sentences, join(words, "\n"))
    end

    open(source.filename) do file
      sentence_lines = []

      while !eof(file)
        line = readline(file)

        if !occursin(r"^\s*$", line)
          push!(sentence_lines, line)
        else
          parse_sentence(sentence_lines)
          sentence_lines = []
        end
      end
    end

    join(parsed_sentences, "\nSENTENCE_END\n")
  end

  function convert(source::SpacySource)
    sentences = open(f -> read(f, String), source.filename) |> (text -> split(text, r"SENTENCE END")) .|> strip
    deleteat!(sentences, length(sentences))

    sentences = map(sentences) do sentence
      map(enumerate(split(sentence, "\n"))) do (index, word)
        word = split(word)
        dependent_word = replace(word[1], r"TOKEN:" => "") |> strip
        dependency = replace(word[2], r"DEP:" => "") |> strip |> lowercase
        head_word = replace(word[3], r"HEAD:" => "") |> strip

        dependent_word = replace(dependent_word, r"-\d+" => "-$(index)")

        head_word_index = findfirst(word ->
        ((replace(split(word)[1], r"TOKEN:" => "")
            |> strip) == head_word), split(sentence, "\n"))
        if startswith(head_word, "ROOT")
          head_word_index = 0
        end

        head_word = replace(head_word, r"-\d+" => "-$(head_word_index)")

        return join([dependency, head_word, dependent_word], " ")
      end |> (sentence -> join(deleteat!(sentence, length(sentence)), "\n"))
    end |> (sentences -> join(sentences, "\nSENTENCE_END\n"))

    return sentences
  end

  function convert(source::StanzaSource)
    sentences = open(f -> read(f, String), source.filename) |> (text -> split(text, r"SENTENCE END")) .|> strip
    deleteat!(sentences, length(sentences))

    sentences = map(sentences) do sentence
      map(enumerate(split(sentence, "\n"))) do (index, word)
        word = split(word)
        dependent_word = replace(word[1], r"TOKEN:" => "") |> strip
        dependency = replace(word[2], r"DEP:" => "") |> strip |> lowercase
        head_word = replace(word[3], r"HEAD:" => "") |> strip

        dependent_word = replace(dependent_word, r"-\d+" => "-$(index)")

        head_word_index = findfirst(word ->
        ((replace(split(word)[1], r"TOKEN:" => "")
            |> strip) == head_word), split(sentence, "\n"))
        if startswith(head_word, "ROOT")
          head_word_index = 0
        end

        head_word = replace(head_word, r"-\d+" => "-$(head_word_index)")

        return join([dependency, head_word, dependent_word], " ")
      end |> (sentence -> join(sentence, "\n"))
    end |> (sentences -> join(sentences, "\nSENTENCE_END\n"))

    return sentences
  end

  function convert(source::CoreNLPSource)
    lines = readlines(source.filename) |> lines -> map(line -> match(r"^(.*)\((.*), (.*)\)", line), lines)

    new_lines = []

    foreach(lines) do line
      if line === nothing
        push!(new_lines, "SENTENCE_END")
      else
        push!(new_lines, join(line, " "))
      end
    end
    deleteat!(new_lines, length(new_lines))
    join(new_lines, "\n")
  end

  function extract_sentences_from_conllu(filename)
    sentence_lines = []

    open(filename) do file
      while !eof(file)
        line = readline(file)

        if occursin(r"# text = ", line)
          push!(sentence_lines, replace(line, r"# text = " => ""))
        end
      end
    end

    join(sentence_lines, "\n")
  end
end


text:
function parse_kw(ps::ParseState, ::Type{Val{Tokens.CONST}})
    startbyte = ps.t.startbyte

    # Parsing
    kw = INSTANCE(ps)
    format_kw(ps)
    @catcherror ps startbyte arg = @default ps parse_expression(ps)

    # Construction
    ret = EXPR{Const}(EXPR[kw, arg], ps.nt.startbyte - startbyte, Variable[], "")

    return ret
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.GLOBAL}})
    startbyte = ps.t.startbyte

    # Parsing
    kw = INSTANCE(ps)
    format_kw(ps)
    @catcherror ps startbyte arg = parse_expression(ps)

    # Construction
    # if arg isa EXPR{TupleH} && first(arg.punctuation) isa PUNCTUATION{Tokens.COMMA}
    #     ret = EXPR(Global, [kw, arg.args...], ps.nt.startbyte - startbyte, arg.punctuation)
    # else
    ret = EXPR{Global}(EXPR[kw, arg], ps.nt.startbyte - startbyte, Variable[], "")
    # end

    return ret
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.LOCAL}})
    startbyte = ps.t.startbyte

    # Parsing
    kw = INSTANCE(ps)
    format_kw(ps)
    @catcherror ps startbyte arg = @default ps parse_expression(ps)

    # Construction
    # if arg isa EXPR && arg.head == TUPLE && first(arg.punctuation) isa PUNCTUATION{Tokens.COMMA}
    #     ret = EXPR(Local, [kw, arg.args...], ps.nt.startbyte - startbyte, arg.punctuation)
    # else
    ret = EXPR{Local}(EXPR[kw, arg], ps.nt.startbyte - startbyte, Variable[], "")
    # end

    return ret
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.RETURN}})
    startbyte = ps.t.startbyte

    # Parsing
    kw = INSTANCE(ps)
    format_kw(ps)
    @catcherror ps startbyte args = @default ps closer(ps) ? NOTHING : parse_expression(ps)

    # Construction
    ret = EXPR{Return}(EXPR[kw, args], ps.nt.startbyte - startbyte, Variable[], "")

    return ret
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.END}})
    ret = EXPR{IDENTIFIER}(EXPR[], ps.nt.startbyte - ps.t.startbyte, Variable[], "end")
    if !ps.closer.square
        ps.errored = true
        return EXPR{ERROR}(EXPR[], 0, Variable[], "incorrect use of end")
    end
    return ret
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.ELSE}})
    ret = EXPR{IDENTIFIER}(EXPR[], ps.nt.startbyte - ps.t.startbyte, Variable[], "else")
    ps.errored = true
    return EXPR{ERROR}(EXPR[], 0, Variable[], "incorrect use of else")
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.ELSEIF}})
    ret = EXPR{IDENTIFIER}(EXPR[], ps.nt.startbyte - ps.t.startbyte, Variable[], "elseif")
    ps.errored = true
    return EXPR{ERROR}(EXPR[], 0, Variable[], "incorrect use of else")
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.CATCH}})
    ret = EXPR{IDENTIFIER}(EXPR[], ps.nt.startbyte - ps.t.startbyte, Variable[], "catch")
    ps.errored = true
    return EXPR{ERROR}(EXPR[], 0, Variable[], "incorrect use of catch")
end

function parse_kw(ps::ParseState, ::Type{Val{Tokens.FINALLY}})
    ret = IDENTIFIER(ps.nt.startbyte - ps.t.startbyte, :finally)
    ps.errored = true
    return EXPR{ERROR}(EXPR[], 0, Variable[], "incorrect use of finally")
end


text:
"""
Develop automatic differentiation (AD) support in Kinetic.jl
"""

using KitBase, ForwardDiff, ReverseDiff, BenchmarkTools, Plots

prim = [1.0, 0.0, 1.0]
vs = VSpace1D(-5, 5, 100)

# We build a unary Maxwellian function for testing.
Mu(u) = maxwellian(u, prim)

@btime fd = ForwardDiff.jacobian(Mu, vs.u)
@btime rd = ReverseDiff.jacobian(Mu, vs.u)
"""
On my workstation, the benchmark results are:
- Forward mode: 30.569 μs (15 allocations: 183.64 KiB)
- Reverse mode: 202.087 ms (6501013 allocations: 215.04 MiB)
So the foward mode seems to be much more efficient.
"""

@btime fd1 = ForwardDiff.derivative.(Mu, vs.u)
"""
15.390 μs (313 allocations: 9.15 KiB)
Using derivative is faster.
"""

# An advantage is that the reverse mode is not restricted to unary functions
@btime rd2 = ReverseDiff.jacobian(maxwellian, (vs.u, prim))
"""
405.675 ms (12996209 allocations: 429.89 MiB)
But it's even slower.
"""

# Check results
fd = ForwardDiff.jacobian(Mu, vs.u)
rd = ReverseDiff.jacobian(Mu, vs.u)
rd2 = ReverseDiff.jacobian(maxwellian, (vs.u, prim))
fd == rd == rd2[1]

dm = [rd2[1][i, i] for i in axes(fd, 1)]
plot(vs.u, dm)

#　These things can be built into a function

function ∂maxwellian(u::Real, ρ, U, λ)
    Mu = u -> maxwellian(u, ρ, U, λ)
    return ForwardDiff.derivative(Mu, u)
end

@btime ∂maxwellian.(vs.u, prim[1], prim[2], prim[3])

function ∂maxwellian(u::AbstractVector, prim)
    Mu = u -> maxwellian(u, prim)
    return ForwardDiff.derivative.(Mu, u)
end

@btime ∂maxwellian(vs.u, prim)


text:
using LinearAlgebra, NearestNeighbors
using ForwardDiff
using Printf
using SparseArrays
using JLD

export neighbour_distances

"""
Auxiliary function for neighbour_distances
"""
function lat_n(n1,n2)
       x1 = n1 + 0.5*n2
       x2 = n2*(0.5*sqrt(3.0))
       return [x1;x2]
       end
"""
Auxiliary function for neighbour_distances
"""
function n_r(col,r1,r2)
       r1 = r1 + 0.001
       r2 = r2+0.001
       II = findall(r1.< col .< r2)
       return col[II]
       end
"""
A function computing distances between neighbours in a triangular lattice with lattice constant equal to unity. The parameter ii specifies the size of the domain over which the neighbours are found.

One starts with a square lattice [-ii,ii]^2, perturbs it to obtain a triangular lattice, then computes sequences of neighbours.

The paramter ii should be chosen much larger than
the chosen cut-off value for interactions

neighbour_distances(ii) returns an array of first K neighbours. The ith entry is two-dimensional, first entry is the distance squared of ith neighbours, second entry is the number of ith neighbours
"""
function neighbour_distances(ii)
    col1 = Float64[]
    for i in -ii:ii
       for j in -ii:ii
           push!(col1,norm(lat_n(i,j)))
       end
    end
    p = sortperm(col1)
    col2 = col1[p];
    NN = Vector{Float64}[]
    for i in 1:(ii^2)
        xx = n_r(col2,sqrt(i-1),sqrt(i))
        yy = nothing
        if length(xx) > 0
            yy = [i;length(xx)]
            push!(NN,yy)
        end
    end
    return NN
end

tri = Val{:tri};
sqr = Val{:sqr};
hcmb = Val{:hmcb};
dc = Val{:dc}
lat_type = Union{tri,sqr,hcmb,dc}

"""
test1
"""
function domain(;r=5.0, lt::lat_type = tri())
    if typeof(lt) == dc
        return domain_silicon(;r=r)
    else

        dom = Vector{Float64}[]
        A = nothing
        if typeof(lt) == tri
            A = [1.0 0.5 0.0 ; 0.0 0.5*sqrt(3) 0.0; 0.0 0.0 0.0]
        elseif typeof(lt) == sqr
            A = [1.0 0.0 0.0 ; 0.0 1.0 0.0; 0.0 0.0 0.0]
        elseif typeof(lt) == hcmb
            A = [1.0 0.5 0.0 ; 0.0 0.5*sqrt(3) 0.0; 0.0 0.0 0.0]
        end
        rr = round(Int,3.0*r) + 3
        if typeof(lt) == hcmb
            shift = [0.5;sqrt(3)/6;0.0]
            for i in -rr:rr
                for j in -rr:rr
                    lp = A*[i;j;0.0]
                    push!(dom,lp)
                    push!(dom,lp.+shift)
                end
            end
        else
            for i in -rr:rr
                for j in -rr:rr
                    lp = A*[i;j;0.0]
                    push!(dom,lp)
               #     push!(dom,lp.+shift)
                end
            end
        end
        II = sortperm(norm.(dom))
        dom = dom[II]
        lc = norm(dom[1].-dom[2])
        dom = (1.0/lc)*dom
        II = findall(norm.(dom) .< r + 0.01)
        return dom[II]
    end
end

function domain_silicon(;r=5.0)
    #nn_d = 2.35167023739006;
    AA = [[0;0;0],[0;2;2], [2;0;2],[2;2;0],
          [3;3;3],[3;1;1],[1;3;1],[1;1;3]]
    dom = Vector{Float64}[]
    rr = round(Int,r/16) + 2
    for i=-rr:rr
        for j=-rr:rr
            for k =-rr:rr
                for x in AA
                    push!(dom,x.+[4*i;4*j;4*k])
                end
            end
        end
    end
     II = sortperm(norm.(dom))
     dom = dom[II]
     lc = norm(dom[1].-dom[2])
     dom = (1.0/lc)*dom
     II = findall(norm.(dom) .< r + 0.01)
     dom = dom[II]

     RM = load(@__DIR__()[1:end-3] * "/data/rot_mat_silicon.jld")["RM"]
     dom = [RM*x for x in dom]
    return dom
end

"""
test2
"""
function at_pairs(at;r=1.0)
    r += 0.01
    pairs = Vector{Int64}[]
    data = hcat(at...)
    balltree = BallTree(data)
    for i in 1:length(at)
        point = at[i]
        idxs = inrange(balltree, point, r, true)
        pp = setdiff(idxs,i)
        push!(pairs,pp)
    end
    return pairs
end

"""
Atoms environment:
    R       # domain parameter
    R_star  # interaction radius
    X       # lattice coordinates (ordered radially!)
    Ns      # neighbours in the interaction range
    Ifree   # last atom that is free to vary (before interface)
    Iclamp # first atom in the far field
    tri     # triangular or not (true or false)
"""
mutable struct AtmModel
    R       # domain parameter
    R_star  # interaction radius
    X       # lattice coordinates (ordered radially!)
    Ns      # neighbours in the interaction range
    Ifree   # last atom that is free to vary (before interface)
    Iclamp # first atom in the far field
    lt::lat_type     # lattice type (tri(), sqr() or hcmb()
    mode1    # mode 1 (true) or mode 3 (false)
    phi     # pair potential
    dphi
    ddphi
end

"""
A function to define an Atoms environment.

AtmModel(; R=5.0, R_star = 1.0, lt=tri(),mode1=true)

"""
function AtmModel(; R=5.0, R_star = 1.0, lt = tri(), mode1 = true)
    X = domain(r = R+(2*R_star), lt = lt)
    Ns = at_pairs(X, r = R_star)
    Iclamp = findfirst(length.(Ns) .< maximum(length.(Ns)))
    Ifree = findfirst([sum(Ns[i] .> Iclamp-0.5) for i in 1:length(X)] .> 0) - 1
    CC = [1.0; 2^(-1/6)]
    phi(r) = 4.0*CC[1]*((CC[2]/r)^(12.0) - (CC[2]/r)^(6.0))
    dphi(r) = ForwardDiff.derivative(phi,r)
    ddphi(r) = ForwardDiff.derivative(dphi,r);
    return AtmModel(R, R_star, X, Ns, Ifree, Iclamp, lt,mode1,phi,dphi,ddphi)
end

uhat_m1_iso(r,theta;kappa=2.0) = sqrt(r)*[(2*kappa-1)*cos(theta/2) - cos(3*theta/2); (2*kappa+1)*sin(theta/2) - sin(3*theta/2); 0]
uhat_m3(r,theta) = [0.0;0.0; sqrt(r)*sin(theta/2)]

function UHAT(at,lambda;kappa = 2.0)
    lc1 = norm(at.X[1]-at.X[2])
    XX = [x .- lc1.*[0.5+lambda;(sqrt(3)/4);0.0] for x in at.X]
    if at.mode1
        return [uhat_m1_iso(norm(x),angle(x[1]+im*x[2]),kappa=kappa) for x in XX]
    else
        return [uhat_m3(norm(x),angle(x[1]+im*x[2])) for x in XX]
    end
end

d1UHAT(at,x) = ForwardDiff.derivative(y->UHAT(at,y),x)
d1d1UHAT(at,x) = ForwardDiff.derivative(y->d1UHAT(at,y),x)

function site_energy_pp(at,i,U::AbstractVector{T}) where T
    E = 0.0
    for j in at.Ns[i]
        Y = at.X[i] - at.X[j] + U[i]-U[j]
        E+=at.phi(norm(Y))
    end
    return E
end

function energy_pp(at,U::AbstractVector{T}) where T
    site_E = []
    for i in 1:length(U)
        E = 0.0
        for j in at.Ns[i]
            Y = at.X[i] - at.X[j] + U[i]-U[j]
            E+=at.phi(norm(Y))
        end
        push!(site_E,E)
    end
    return sum(site_E)
end

function energy_pp2(at,U::AbstractVector{T}) where T
    site_E = []
    for i in 1:length(U)
        push!(site_E,site_energy_pp(at,i,U))
    end
    return sum(site_E)
end

function grad_pp(at,U::AbstractVector{T}) where T
    #dphi(r) = ForwardDiff.derivative(at.phi,r)
    grad = U .- U
    for i in 1:length(U)
        for j in at.Ns[i]
            DU = U[i]-U[j]
            DY = at.X[i] - at.X[j] + DU
            grad[i] += (at.dphi(norm(DY))/norm(DY))*DY
            grad[j] -= (at.dphi(norm(DY))/norm(DY))*DY
        end
    end
    return grad
end

function hess_pp(at,U::AbstractVector{T}) where T
#    dphi(r) = ForwardDiff.derivative(at.phi,r)
#    ddphi(r) = ForwardDiff.derivative(dphi,r)
    II = Int64[]
    J = Int64[]
    Z = Array{Float64,2}[]
    for i in 1:length(U)
        for j in at.Ns[i]
            DU = U[i]-U[j]
            DY = at.X[i] - at.X[j] + DU
            M = DY * DY'
            dy = norm(DY)
            hh0 = (at.ddphi(dy)/(dy)^2) - (at.dphi(dy)/(dy^3)) #+ dphi(dy)/dy
            hh = hh0 * M + (at.dphi(dy)/dy)*I
            append!(II, [i,   i,    j,    j])
            append!(J, [i,   j,    i,    j])
            push!(Z,hh)
            push!(Z,-hh)
            push!(Z,-hh)
            push!(Z,hh)
     #       append!(Z, [hh, -hh, -hh, hh])
        end
    end
    return sparse(II, J, Z, length(U), length(U))
end


text:
abstract type AbstractBinaryHeap{T,O<:Ordering} <: AbstractVector{T} end

typename(::Type{<:AbstractBinaryHeap}) = "binary heap"

"""
    BinaryHeap{T}(o = Forward)

yields an empty binary heap whose nodes are of type `T` and with ordering
specified by `o`.  A min-heap or a max-heap is built depending on whether `o`
is set to forward or reverse ordering.

A vector `vals` storing the nodes of the binary heap can be specified:

    BinaryHeap(vals, o = Forward)

A number of methods are available for a binary heap `h`:

    pop!(h)       # deletes root node of heap h and yields its value
    push!(h, x)   # pushes value x in heap h
    empty!(h)     # empties heap h
    isempty(h)    # yields whether heap h is empty
    delete!(h, i) # deletes i-th node from heap h
    peek(h)       # yields value of root node of heap h without deleting it
    first(h)      # idem

A binary heap `h` behaves like an abstract vector(with 1-based linear indices):

    length(h)   # the number of nodes in heap h
    h[i]        # the i-th node of heap h
    h[i] = x    # set the i-th node of heap h and heapify h

Note that `h[1]` is the root node of the heap `h` and that setting a node in
the heap may triggers reordering of the nodes to maintain the binary heap
structure.  In other words, after doing `h[i] = x`, do not assume that `h[i]`
yields `x`.

Operations that modify the heap, like deletion by `delete!(h,i)`, insertion by
`h[i] = x`, pushing by `push!(h,x)`, and extracting by `pop!(h)` are of
complexity `O(1)` in the best case, `O(log(n))` in the worst case, with `n =
length(h)` the number of nodes in the heap `h`.  Query the value of a given
node by `peek(h)`, `first(h)`, or `h[i]` is always of complexity `O(1)`.

Method `sizehint!(h,n)` may be called to anticipate that the heap may contains
`n` values.

"""
struct BinaryHeap{T,O} <: AbstractBinaryHeap{T,O}
    order::O         # ordering
    nodes::Vector{T} # storage for the nodes
    BinaryHeap{T}(o::O=default_ordering(BinaryHeap)) where {T,O<:Ordering} =
        new{T,O}(o, Vector{T}(undef, 0))
    BinaryHeap{T}(o::O, vals::AbstractVector) where {T,O<:Ordering} =
        heapify!(new{T,O}(o, vals))
end

"""
    h = FastBinaryHeap{T}(...)

yields a fast binary heap `h`.  Compared to `BinaryHeap{T}(...)`, the default
ordering is `FastForward` and the array backing the storage of the heap values
is never reduced to improve performances in some cases.  You may call
`resize!(h)` to explicitly reduce the storage to its minimum.

"""
mutable struct FastBinaryHeap{T,O} <: AbstractBinaryHeap{T,O}
    order::O         # ordering
    nodes::Vector{T} # storage for the nodes
    count::Int       # current number of nodes
    FastBinaryHeap{T}(o::O=default_ordering(FastBinaryHeap)) where {T,O<:Ordering} =
        new{T,O}(o, Vector{T}(undef, 0), 0)
    FastBinaryHeap{T}(o::O, vals::AbstractVector) where {T,O<:Ordering} =
        heapify!(new{T,O}(o, vals, length(vals)))
end

default_ordering(::Type{<:AbstractBinaryHeap}) = Forward
default_ordering(::Type{<:FastBinaryHeap}) = FastForward

# Outer constructors.
for type in (:BinaryHeap, :FastBinaryHeap)
    @eval begin
        $type{T}(vals::AbstractVector) where {T} =
            $type{T}(default_ordering($type), vals)
        $type(vals::AbstractVector{T}) where {T} = $type{T}(vals)
        $type(o::Ordering, vals::AbstractVector{T}) where {T} =
            $type{T}(o, vals)
    end
end

"""
    QuickHeaps.nodes(h)

yields the array backing the storage of the values in the binary heap `h`.

This method may be specialized for custom binary heap types.

"""
ordering(h::AbstractBinaryHeap) = getfield(h, :order)

"""
    QuickHeaps.ordering(h)

yields the ordering of the values in the binary heap `h`.

This method may be specialized for custom binary heap types.

"""
nodes(h::AbstractBinaryHeap) = getfield(h, :nodes)

length(h::FastBinaryHeap) = getfield(h, :count)
length(h::BinaryHeap) = length(nodes(h))
size(h::AbstractBinaryHeap) = (length(h),)
IndexStyle(::Type{<:AbstractBinaryHeap}) = IndexLinear()
sizehint!(h::AbstractBinaryHeap, n::Integer) = begin
    sizehint!(nodes(h), n)
    return h
end
isempty(h::AbstractBinaryHeap) = length(h) < 1

# Call `resize!(h)` with no other arguments to reduce the storage size.
resize!(h::AbstractBinaryHeap) = h # do nothing by default
resize!(h::FastBinaryHeap) = begin
    resize!(nodes(h), length(h))
    return h
end

# Heap indexing.  Note that linear 1-based indexing is assumed for the
# array storing the heap.
heap_left(i::Int) = 2*i
heap_right(i::Int) = 2*i + 1
heap_parent(i::Int) = div(i, 2)

@inline function getindex(h::AbstractBinaryHeap, i::Int)
    @boundscheck checkbounds(h, i)
    @inbounds r = getindex(nodes(h), i)
    return r
end

@inline @propagate_inbounds setindex!(h::AbstractBinaryHeap, x, i::Int) =
    setindex!(h, to_eltype(h, x), i)

@inline function setindex!(h::AbstractBinaryHeap{T},
                           x::T, i::Int) where {T}
    @boundscheck checkbounds(h, i)
    A = nodes(h)
    @inbounds y = A[i] # replaced node
    o = ordering(h)
    if lt(o, y, x)
        # Heap structure _above_ replaced node will remain valid, down-heapify
        # to fix the heap structure at and _below_ the node.
        unsafe_heapify_down!(o, A, i, x, length(h))
    else
        # Heap structure _below_ replaced node will remain valid, up-heapify to
        # fix the heap structure at and _above_ the node.
        unsafe_heapify_up!(o, A, i, x)
    end
    return h
end

first(h::AbstractBinaryHeap) = peek(h)

function peek(h::AbstractBinaryHeap)
    isempty(h) && throw_argument_error(typename(h), " is empty")
    @inbounds r = getindex(nodes(h), 1)
    return r
end

empty!(h::FastBinaryHeap) = (setfield!(h, :count, 0); h)
empty!(h::BinaryHeap) = (empty!(nodes(h)); h)

function pop!(h::AbstractBinaryHeap)
    n = length(h)
    n ≥ 1 || throw_argument_error(typename(h), " is empty")
    A = nodes(h)
    @inbounds x = A[1]
    if n > 1
        # Peek the last node and down-heapify starting at the root of the
        # binary heap to insert it.
        @inbounds y = A[n]
        unsafe_heapify_down!(ordering(h), A, 1, y, n - 1)
    end
    unsafe_shrink!(h, n - 1)
    return x
end

push!(h::AbstractBinaryHeap, ::Tuple{}) = h

function push!(h::AbstractBinaryHeap, args...)
    for x in args
        push!(h, x)
    end
    return h
end

push!(h::AbstractBinaryHeap, x) = push!(h, to_eltype(h, x))

function push!(h::AbstractBinaryHeap{T}, x::T) where {T}
    n = length(h) + 1
    unsafe_heapify_up!(ordering(h), unsafe_grow!(h, n), n, x)
    return h
end

delete!(h::AbstractBinaryHeap, i::Integer) = delete!(h, to_int(i))

function delete!(h::AbstractBinaryHeap, i::Int)
    n = length(h)
    in_range(i, n) || throw_argument_error("out of range index")
    if i < n
        # Replace the deleted node by the last node in the heap and
        # up-/down-heapify to restore the binary heap structure.
        A = nodes(h)
        o = ordering(h)
        @inbounds x = A[n] # node to replace deleted node
        @inbounds y = A[i] # deleted node
        if lt(o, y, x)
            # Heap structure _above_ deleted node will remain valid,
            # down-heapify to fix the heap structure at and _below_ the
            # node.
            unsafe_heapify_down!(o, A, i, x, n - 1)
        else
            # Heap structure _below_ deleted node will remain valid,
            # up-heapify to fix the heap structure at and _above_ the node.
            unsafe_heapify_up!(o, A, i, x)
        end
    end
    unsafe_shrink!(h, n - 1)
    return h
end

"""
    QuickHeaps.unsafe_grow!(h, n) -> A

grows the size of the binary heap `h` to be `n` and returns the array `A`
backing the storage of the nodes.  This method is *unsafe* because it does not
check its arguments and because it breaks the binary heap structure of the
array of nodes.

This method is called by `push!` to grow the size of the heap and shall be
specialized for any concrete sub-types of `QuickHeaps.AbstractBinaryHeap`.

"""
unsafe_grow!(h::BinaryHeap, n::Int) = resize!(nodes(h), n)
unsafe_grow!(h::FastBinaryHeap, n::Int) = begin
    A = nodes(h)
    length(A) < n && resize!(A, n)
    setfield!(h, :count, n)
    return A
end

"""
    QuickHeaps.unsafe_shrink!(h, n)

shrinks the size of the binary heap `h` to be `n`.  This method is *unsafe*
because it does not check its arguments.

This method is called by `delete!` to eventually reduce the size of the heap
and shall be specialized for any concrete sub-type of
[`QuickHeaps.AbstractBinaryHeap`](@ref).

"""
unsafe_shrink!(h::BinaryHeap, n::Int) = resize!(nodes(h), n)
unsafe_shrink!(h::FastBinaryHeap, n::Int) = setfield!(h, :count, n)

"""
    heapify!(h) -> h

reorders the nodes in the binary heap `h` in-place.  This method should be
called to initialize the heap or to re-order the heap if its contents have been
modified by other methods than `pop!` or `push!`.

The method can be called at a lower level to heapify (part of) an array storing
the heap values:

    heapify!([o=Base.Forward,] A, n=length(A)) -> A

reorders the `n` first elements of array `A` in-place to form a binary heap
according to the ordering specified by `o`.  The array `A` must have 1-based
linear indexing.  Arguments may be specified in any order.

"""
function heapify!(h::AbstractBinaryHeap)
    heapify!(ordering(h), nodes(h), length(h))
    return h
end

heapify!(o::Ordering, A::AbstractArray, n::Integer) = heapify!(o, A, to_int(n))

function heapify!(o::Ordering, A::AbstractArray, n::Int = length(A))
    # Heapify the n first elements of A.
    check_heap_storage(A, n)
    @inbounds for i in heap_parent(n):-1:1
        unsafe_heapify_down!(o, A, i, A[i], n)
    end
    return A
end

"""
    heapify([o=Base.Forward,] A, n=length(A))

yields an array with the `n` first values of array `A` stored in a binary heap
structure of ordering specified by `o`.  The storage of the returned heap is
a different array than `A`.  Arguments may be specified in any order.

"""
heapify(o::Ordering, A::AbstractArray, n::Integer) = heapify(o, A, to_int(n))

heapify(o::Ordering, A::AbstractArray{T}, n::Int = length(A)) where {T} =
    heapify!(o, copyto!(Vector{T}(undef, n), 1, A, 1, n))

"""
    isheap([o=Base.Forward,], A, n=length(A))

yields whether the `n` first elements of array `A` have a binary heap structure
ordered as specified by `o`.  Arguments may be specified in any order.

    isheap(obj; check=false)

yields whether object `obj` is a binary heap.  If keyword `check` is true, the
internal structure of `obj` is checked; otherwise, the type of `obj` is trusted
to determine whether it is a binary heap.

"""
isheap(o::Ordering, A::AbstractArray, n::Integer) = isheap(o, A, to_int(n))

function isheap(o::Ordering, A::AbstractArray, n::Int = length(A))
    check_heap_storage(A, n)
    @inbounds for i in 1:div(n, 2)
        l = heap_left(i)
        r = heap_right(i)
        if lt(o, A[l], A[i]) || (r ≤ n && lt(o, A[r], A[i]))
            return false
        end
    end
    return true
end

isheap(h::AbstractBinaryHeap; check::Bool = false) =
    if check
        isheap(ordering(h), nodes(h), length(h))
    else
        true
    end

# Cope with different ordering of arguments and using the same default ordering
# as in base Julia and DataStructures.
for func in (:heapify, :heapify!, :isheap)
    @eval begin
        function $func(A::AbstractArray,
                       o::Ordering = default_ordering(A),
                       n::Integer = length(A))
            return $func(o, A, n)
        end
        function $func(A::AbstractArray, n::Integer,
                       o::Ordering = default_ordering(A))
            return $func(o, A, n)
        end
    end
end

"""
    QuickHeaps.heapify_down!(o, A, i, x=A[i], n=lengh(A))

stores the value `x` in the `i`-th node of the binary heap built into the `n`
first elements of array `A` with ordering `o` and, if needed, moves down the
inserted value to maintain the binary heap structure.

This method is called to *heapify* an array in order to initialize or rebuild
the heap structure or to replace the value of the root node of the heap and
update the heap structure.

"""
function heapify_down!(o::Ordering, A::AbstractArray,
                       i::Integer, x = A[i], n::Integer = length(A))
    heapify_down!(o, A, to_int(i), to_eltype(A, x), to_int(n))
end

function heapify_down!(o::Ordering, A::AbstractArray{T},
                       i::Int, x::T, n::Int) where {T}
    check_heap_storage(A, n)
    in_range(i, n) || throw_argument_error("out of range index")
    unsafe_heapify_down!(o, A, i, x, n)
end

"""
    QuickHeaps.unsafe_heapify_down!(o, A, i, x=A[i], n=lengh(A))

This method is a fast but *unsafe* version of
[`QuickHeaps.heapify_down!`](@ref) which assumes that all arguments are
correct, that is `A` implements 1-based linear indexing, `0 ≤ n ≤ lengh(A)`,
and `1 ≤ i ≤ n`.

"""
@inline function unsafe_heapify_down!(o::Ordering,
                                      A::AbstractArray{T},
                                      i::Int,
                                      x::T = (@inbounds A[i]),
                                      n::Int = length(A)) where {T}
    @inbounds begin
        while (l = heap_left(i)) ≤ n
            j = (r = heap_right(i)) > n || lt(o, A[l], A[r]) ? l : r
            lt(o, A[j], x) || break
            A[i] = A[j]
            i = j
        end
        A[i] = x
    end
end

"""
   QuickHeaps.heapify_up!(o, A, i, x=A[i])

stores the value `x` in the `i`-th node of the binary heap built into the
`n ≥ i` first elements of array `A` with ordering `o` and, if needed, moves
up the value to maintain the heap structure.

"""
function heapify_up!(o::Ordering, A::AbstractArray,
                     i::Integer, x = A[i])
    heapify_up!(o, A, to_int(i), to_eltype(A, x))
end

function heapify_up!(o::Ordering, A::AbstractArray{T}, i::Int, x::T) where {T}
    check_heap_storage(A)
    in_range(i, length(A)) || error("out of range index")
    unsafe_heapify_up!(o, A, i, x)
end

"""
    QuickHeaps.unsafe_heapify_up!(o, A, i, x=A[i])

This methods is a fast but *unsafe* version of [`QuickHeaps.heapify_up!`](@ref)
which assumes that all arguments are correct, that is `A` implements 1-based
linear indexing and `1 ≤ i ≤ length(A)`.

"""
@inline function unsafe_heapify_up!(o::Ordering,
                                    A::AbstractArray{T},
                                    i::Int,
                                    x::T = (@inbounds A[i])) where {T}
    @inbounds begin
        while (j = heap_parent(i)) ≥ 1 && lt(o, x, A[j])
            A[i] = A[j]
            i = j
        end
        A[i] = x
    end
end

"""
    QuickHeaps.check_heap_storage(A)

throws an exception if array `A` is not suitable for storing a binary heap,
that is if `A` does not have 1-based linear indexing.

    QuickHeaps.check_heap_storage(A, n)

throws an exception if the first elements of array `A` are not suitable for
storing a binary heap of size `n`.

"""
check_heap_storage(A::AbstractArray) = begin
    has_standard_linear_indexing(A) || throw(ArgumentError(
        "array storing a binary heap must have 1-based linear indexing"))
    nothing
end

check_heap_storage(A::AbstractArray, n::Int) = begin
    # Check that array has linear indexing and that 0 ≤ n ≤ length(A).
    check_heap_storage(A)
    (n % UInt) ≤ (length(A) % UInt) || throw_argument_error(
        "out of range heap size")
    nothing
end


text:
"""
    simulate_random_snparray(s::String, n::Integer, p::Integer;
        [mafs::Vector{Float64}], [min_ma::Integer])

Creates a random SnpArray in the current directory without missing value,
where each SNP has ⫺5 (default) minor alleles.

Note: if supplied minor allele frequency is extremely small, it could take a
long time for the simulation to generate samples where at least `min_ma`
(defaults to 5) are present.

# Arguments:
- `s`: name of SnpArray that will be created in the current directory. To not
    create file, use `undef`.
- `n`: number of samples
- `p`: number of SNPs

# Optional Arguments:
- `mafs`: vector of desired minor allele freuqencies (uniform(0,0.5) by default)
- `min_ma`: the minimum number of minor alleles that must be present for each
    SNP (defaults to 5)
"""
function simulate_random_snparray(s::Union{String, UndefInitializer}, n::Int64,
    p::Int64; mafs::Vector{Float64}=zeros(Float64, p), min_ma::Int = 5)

    @assert all(0.0 .<= mafs .<= 0.5) "Minor allele frequencies not in (0, 0.5)"

    if mafs != zeros(Float64, p)
        return _random_snparray(s, n, p, mafs, min_ma=min_ma)
    end

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = 0
        while minor_alleles <= min_ma
            maf = 0.5rand()
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
        mafs[j] = maf
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(s, A1, A2)
end

"""
This function requires a vector of minor alleles frequencies to perform simulation.
It will create a random SnpArray in the current directory without missing value,
where each SNP has at least 5 minor alleles.
"""
function _random_snparray(s::Union{String, UndefInitializer}, n::Int64,
        p::Int64, mafs::Vector{Float64}; min_ma::Int = 5)
    all(0.0 .<= mafs .<= 0.5) || throw(ArgumentError("vector of minor allele frequencies must be in (0, 0.5)"))
    any(mafs .<= 0.0005) && @warn("Provided minor allele frequencies contain entries smaller than 0.0005, simulation may take long if sample size is small and min_ma = $min_ma is large")

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = mafs[j]
        while minor_alleles <= min_ma
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(s, A1, A2)
end

"""
Make a SnpArray from 2 BitArrays.
"""
function _make_snparray(s::Union{String, UndefInitializer}, A1::BitArray, A2::BitArray)
    n, p = size(A1)
    x = SnpArray(s, n, p)
    for i in 1:(n*p)
        c = A1[i] + A2[i]
        if c == 0
            x[i] = 0x00
        elseif c == 1
            x[i] = 0x02
        elseif c == 2
            x[i] = 0x03
        else
            throw(MissingException("matrix shouldn't have missing values!"))
        end
    end
    return x
end

"""
    simulate_correlated_snparray(s, n, p; block_length, hap, prob)

Simulates a SnpArray with correlation. SNPs are divided into blocks where each
adjacent SNP is the same with probability prob. There are no correlation between blocks.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional arguments:
- `block_length`: length of each LD block
- `hap`: number of haplotypes to simulate for each block
- `prob`: with probability `prob` an adjacent SNP would be the same.
"""
function simulate_correlated_snparray( s::Union{String, UndefInitializer},
    n::Int64, p::Int64; block_length::Int64=20, hap::Int=20, prob::Float64=0.75)

    @assert mod(p, block_length) == 0 "block_length ($block_length) is not divible by p ($p)"
    @assert 0 < prob < 1 "transition probably should be between 0 and 1, got $prob"

    x = SnpArray(s, n, p)
    haplotypes = zeros(hap, block_length)
    snps = zeros(block_length)
    blocks = Int(p / block_length)

    @inbounds for b in 1:blocks

        #create pool of haplotypes for each block
        _sample_haptotypes!(haplotypes, prob)

        for i in 1:n
            #sample 2 haplotypes with replacement from the pool of haplotypes
            row1 = rand(1:hap)
            row2 = rand(1:hap)
            for j in 1:block_length
                snps[j] = haplotypes[row1, j] + haplotypes[row2, j]
            end

            #copy haplotypes into x
            _copy_blocks!(x, i, snps, b, block_length)
        end
    end

    return x
end

function _sample_haptotypes!(haplotypes::Matrix, prob::Float64)
    n, p = size(haplotypes)
    fill!(haplotypes, 0)

    @inbounds for i in 1:n
        cur_row_sum = 0
        while cur_row_sum == 0
            curr = rand(0:1)
            haplotypes[i, 1] = curr
            cur_row_sum += curr
            for j in 2:p
                stay = rand(Bernoulli(prob)) #stay = 1 means retain the current value
                curr = (stay == 1 ? curr : 1 - curr)
                haplotypes[i, j] = curr
                cur_row_sum += curr
            end
        end
    end
end

function _copy_blocks!(x::SnpArray, row, snps, cur_block, block_length)
    #copy sampled snps into SnpArray
    @inbounds for k in 1:length(snps)
        c = snps[k]
        col = (cur_block - 1) * block_length + k
        if c == 0
            x[row, col] = 0x00
        elseif c == 1
            x[row, col] = 0x02
        elseif c == 2
            x[row, col] = 0x03
        else
            throw(ArgumentError("SNP values should be 0, 1, or 2 but was $c"))
        end
    end
end

"""
    simulate_random_response(x, k, d, l; kwargs...)

This function simulates a random response (trait) vector `y`. When the
distribution `d` is from Poisson, Gamma, or Negative Binomial, we simulate
`β ∼ N(0, 0.3)` to roughly ensure the mean of response `y` doesn't become too
large. For other distributions, we choose `β ∼ N(0, 1)`.

# Arguments
- `x`: Design matrix
- `k`: the true number of predictors.
- `d`: The distribution of the simulated trait (note `typeof(d) = UnionAll` but `typeof(d())` is an actual distribution: e.g. Normal)
- `l`: The link function. Input `canonicallink(d())` if you want to use the canonical link of `d`.

# Optional arguments
- `r`: The number of success until stopping in negative binomial regression, defaults to 10
- `α`: Shape parameter of the gamma distribution, defaults to 1
- `Zu`: Effect of non-genetic covariates. `Zu` should have dimension `n × 1`.
"""
function simulate_random_response(x::AbstractMatrix, k::Int,
    d::UnionAll, l::Link; r = 10, α = 1, Zu::AbstractVector = zeros(size(x, 1)))
    n, p = size(x)
    if (typeof(d) <: NegativeBinomial) || (typeof(d) <: Gamma)
        l == LogLink() || throw(ArgumentError("Distribution $d must use LogLink!"))
    end

    #simulate a random model β from a normal distribution
    true_b = zeros(p)
    if d == Poisson || d == Gamma || d == NegativeBinomial
        true_b[1:k] = rand(Normal(0, 0.3), k)
    else
        true_b[1:k] = randn(k)
    end
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, x * true_b + Zu)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        μ = linkinv.(l, x * true_b + Zu)
        clamp!(μ, -20, 20)
        prob = 1 ./ (1 .+ μ ./ r)
        y = [rand(d(r, i)) for i in prob] #number of failtures before r success occurs
    elseif d == Gamma
        μ = linkinv.(l, x * true_b + Zu)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)

    return y, true_b, correct_position
end
"""
    simulate_random_response(x, k, traits)

Simulates a response matrix `Y` where each row is an independent multivariate
Gaussian with length `trait`. There are `k` non-zero `β` over all traits. Each
trait shares `overlap` causal SNPs. The covariance matrix `Σ` is positive definite
and symmetric.

# Arguments
- `x`: Design matrix of dimension `n × p`. Each row is a sample.
- `k`: the total true number of causal SNPs (predictors)
- `traits`: Number of traits

# Optional arguments
- `Zu`: Effect of non-genetic covariates. `Zu` should have dimension `n × traits`.
- `overlap`: Number of causal SNPs shared by all traits. Shared SNPs does not have the same effect size.

# Outputs
- `Y`: Response matrix where each row is sampled from a multivariate normal with mean `μ[i] = X[i, :] * true_b` and variance `Σ`
- `Σ`: the symmetric, positive definite covariance matrix used
- `true_b`: A sparse matrix containing true beta values.
- `correct_position`: Non-zero indices of `true_b`
"""
function simulate_random_response(x::AbstractMatrix, k::Int, traits::Int;
    Zu::AbstractMatrix = zeros(size(x, 1), traits), overlap::Int = 0
    )
    n, p = size(x)
    traits * overlap ≤ k || throw(ArgumentError("traits * overlap cannot exceed k!"))

    #simulate a random model β. Each trait can have different number of causal SNPs
    true_b = zeros(p, traits)
    if overlap == 0
        causal_snps = sample(1:(traits * p), k, replace=false)
        true_b[causal_snps] = randn(k)
    else
        shared_snps = sample(1:p, overlap, replace=false)
        weight_vector = aweights(1 / (traits * (p - overlap)) * ones(traits * p))
        for i in 1:traits
            weight_vector[i*shared_snps] .= 0.0 # avoid sampling from shared snps
        end
        @assert sum(weight_vector) ≈ 1.0
        # simulate β for shared predictors
        for i in 1:traits
            true_b[shared_snps, i] = randn(overlap)
        end
        # simulate β for none shared predictors
        nonshared_snps = sample(1:(traits * p), weight_vector, k - traits * overlap, replace=false)
        true_b[nonshared_snps] = randn(k - traits * overlap)
    end
    correct_position = findall(x -> x != 0, true_b)

    # simulate random covariance matrix
    Σ = random_covariance_matrix(traits)

    # simulate multivariate normal phenotype for each sample
    μ = x * true_b + Zu

    # simulate response
    Y = zeros(n, traits)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return Y, Σ, true_b, correct_position
end

"""
    random_covariance_matrix(n::Int, [κ=10])

Generates a `n × n` positive definite, symmetric matrix with condition number less than `κ`
https://discourse.julialang.org/t/generate-a-positive-definite-matrix/48582

Both issymmetric(random_covariance_matrix(n)) and isposdef(random_covariance_matrix(n))
will always return true, and cond(random_covariance_matrix(n)) ≤ κ
"""
function random_covariance_matrix(n::Int, κ=10)
    d = Uniform(1, sqrt(κ))
    Q, _ = qr(randn(n, n))
    σ = rand(d, n)
    D = Diagonal(σ)
    A = Q * D * Q'
    return A' * A
end

"""
    adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})

Makes 1 SNP (a column of `x`) correlate with SNPs in `location` with correlation coefficient roughly `ρ`.

# Arguments:
- `x`: the snparray
- `ρ`: correlation coefficient
- `pos`: the position of the target SNP that everything would be correlated to
- `location`: All the SNPs that shall be correlated with the SNP at position `pos` with correlation `ρ`.
"""
function adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})
    @assert 0 <= ρ <= 1 "correlation coefficient must be in (0, 1) but was $ρ"

    for loc in location
        for i in 1:size(x, 1)
            rand(Bernoulli(ρ)) && (x[i, loc] = x[i, pos]) #make 2nd column the same as 1st ~ρ% of the time
        end
        corr = cor(x[:, loc], x[:, pos])
    end
end

"""
    make_bim_fam_files(x::SnpArray, y, name::String)

Creates .bim and .bed files from a SnpArray.

# Arguments:
- `x`: A SnpArray (i.e. `.bed` file on the disk) for which you wish to create corresponding `.bim` and `.fam` files.
- `name`: string that should match the `.bed` file (Do not include `.bim` or `.fam` extensions in `name`).
- `y`: Trait vector that will go in to the 6th column of `.fam` file.
"""
function make_bim_fam_files(x::SnpArray, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)
    @assert n == ly "dimension mismatch: phenotype data has length $ly but SnpArray has $n samples"

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1\tsnp$i\t0\t1\t1\t2\n")
        end
    end

    #create .fam file structure: https://www.cog-genomics.org/plink2/formats#fam
    traits = size(y, 2)
    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1")
            for j in 1:traits
                write(f, "\t$(y[i, j])")
            end
            write(f, "\n")
        end
    end
end


text:
using Test
using XLSXasJSON
using JSONPointer
using Tables

data_path = joinpath(@__DIR__, "data")

@testset "Tables.jl integration" begin
    f = joinpath(data_path, "example.xlsx")
    jws = JSONWorksheet(f, "Sheet1")

    @test Tables.istable(jws)
    @test Tables.isrowtable(jws)
    @test Tables.rowaccess(jws)
    @test Tables.columnaccess(jws) == false
    Tables.columnnames(jws) == jws.pointer

    @test Tables.rows(jws) == jws.data
    @test Tables.columns(jws) == map(i -> jws[:, i], 1:size(jws, 2))
    @test Tables.getcolumn(jws, 1) == jws[:, 1]
    @test Tables.getcolumn(jws, 2) == jws[:, 2]

    @test Tables.getcolumn(jws, j"/array_int") == Tables.getcolumn(jws, 3)
    @test Tables.getcolumn(jws, Symbol("/array_float")) == Tables.getcolumn(jws, 4)

    @test Tables.matrix(jws) == jws[:, :]
    mat_trans = Tables.matrix(jws, true)
    for i in 1:length(jws)
        @test mat_trans[1, i] == jws[i, 1]
        @test mat_trans[2, i] == jws[i, 2]
        @test mat_trans[3, i] == jws[i, 3]
        @test mat_trans[4, i] == jws[i, 4]
    end
end


text:
"""
Abstract supertype. `Objective`s map simulation outputs to predictions
that can be compared to target data using a loss function.

They must implement [`predictions`](@ref)and [`targets`](@ref) methods.
"""
abstract type Objective end

"""
    targets(obj::Objective)
Returns a targets array given an Objective. The targets must match the size and
dimensions of the prediction array returned by `predictions`.
"""
function targets end

"""
    predictions(obj::Objective, output::Output)

Methods that map an objective object and a simulation output to a
prediction array.
"""
function predictions end

predictions(obj::Objective, output) = output[end]

"""
    SimpleObjective(targets)

A basic objective that holds a target array uses the final frame of the
simulation as the prediction.
"""
struct SimpleObjective{T} <: Objective
    targets::T
end

targets(obj::SimpleObjective) = obj.targets

# """
    # PresenceAbsenceObjective(detectionthreshold, occurance, framesperstep)

# Implementation of a loss objective that converts cell data to regional
# presence/absence and compares to a target of regional occurance data.
# """
# struct PresenceAbsenceObjective{DT,RL,OC,FS,S} <: Objective
#     detectionthreshold::DT
#     occurance::OC
#     framesperstep::FS
# end

# targets(obj::PresenceAbsenceObjective) = obj.occurance
# predictions(obj::PresenceAbsenceObjective, output) = output[end]

"""
Implementation of a loss objective that converts cell data to regional
presence/absence and compares to a target of regional occurance data.

"""
struct RegionObjective{DT,RL,OC,FS,S} <: Objective
    detectionthreshold::DT
    regionlookup::RL
    occurance::OC
    framesperstep::FS
    # TODO fix this offset hack with real DateTime handling
    start::S
end

targets(obj::RegionObjective) = obj.occurance
predictions(obj::RegionObjective, output) = output[end]


text:
using LinearAlgebra
using Distributions
using Revise
using Random
using Printf
using StatsBase
using StaticArrays
using POMDPs
using BeliefUpdaters
using AutoViz
using AutomotiveDrivingModels
using AutomotivePOMDPs
using POMDPSimulators
using POMDPPolicies
using DeepQLearning
using Flux
using RLInterface
using POMDPModelTools
using TensorBoardLogger
using Distributions
# using Interact
# using Blink
using MCTS
using Reel

# LOGGER = Logger("log1", overwrite=true)
includet("environment.jl")
includet("generative_mdp.jl")
includet("masking.jl")
includet("cooperative_IDM.jl")
includet("belief_updater.jl")
includet("overlays.jl")
includet("rendering.jl")
includet("make_gif.jl")

rng = MersenneTwister(1)

# mdp = GenerativeMergingMDP(n_cars_main=5, observe_cooperation=false)

function test_state(s,v=0.0, acc=0.0,)
    ego = Vehicle(vehicle_state(35.0, merge_lane(mdp.env), 5.0, mdp.env.roadway), VehicleDef(), EGO_ID)
    veh1 = Vehicle(vehicle_state(s, main_lane(mdp.env), 4.9, mdp.env.roadway), VehicleDef(), EGO_ID + 1)
    veh2 = Vehicle(vehicle_state(s-15, main_lane(mdp.env), 4.9, mdp.env.roadway), VehicleDef(), EGO_ID + 2)
    scene = Scene()
    push!(scene, ego)
    push!(scene, veh1)
    push!(scene, veh2)
    return AugScene(scene, (acc=acc,))
end

rng = MersenneTwister(1)

mdp = GenerativeMergingMDP(random_n_cars=true, min_cars=10, max_cars=14, observe_cooperation=true)

mdp.driver_models[2] = CooperativeIDM(c=0.0)
set_desired_speed!(mdp.driver_models[2], 5.0)

s0 = test_state(85.0)

up = MergingUpdater(mdp)

s0 = initialstate(mdp, rng)
b0 = initialize_belief(up, s0)

super_render(mdp, s0, b0)
s = s0
b = b0
sp, r = generate_sr(mdp, s, 7, rng)
bp = update(up, b, 7, sp)
super_render(mdp, sp, bp)

s = sp
b = deepcopy(bp)
@show b.driver_types
a = 7
sp, r = generate_sr(mdp, s, a, rng)
bp = update(up, b, a, sp)
super_render(mdp, sp, bp)

pomdp = FullyObservablePOMDP(mdp)
up = MergingUpdater(mdp)
s0 = initialstate(mdp, rng)
b0 = initialize_belief(up, s0)
hr = HistoryRecorder(rng=rng, max_steps = 100)

hist= simulate(hr, pomdp, policy, up, b0, s0);

include("visualizer.jl")

includet("make_gif.jl")

make_gif(hist, mdp, hist.belief_hist)

function quick_evaluation(pomdp::POMDP, policy::Policy, rng::AbstractRNG, n_eval=1000)
    avg_r, avg_dr, c_rate, avg_steps, t_out = 0.0, 0.0, 0.0, 0.0, 0.0
    @showprogress for i=1:n_eval
        s0 = initialstate(pomdp, rng)
        up = MergingUpdater(pomdp.mdp)
        hr = HistoryRecorder(rng = rng, max_steps=100)
        hist = simulate(hr, pomdp, policy, up, b0, s0)
        avg_r += undiscounted_reward(hist)
        avg_dr += discounted_reward(hist)
        c_rate += undiscounted_reward(hist) <= pomdp.mdp.collision_cost
        t_out += n_steps(hist) >= hr.max_steps ? 1.0 : 0.0
        avg_steps += n_steps(hist)
    end
    avg_r /= n_eval
    avg_dr /= n_eval
    c_rate /= n_eval
    avg_steps /= n_eval
    t_out /= n_eval
    return avg_r, avg_dr, c_rate, avg_steps, t_out
end
policy = NNPolicy(mdp, policy.qnetwork, policy.action_map, policy.n_input_dims)
avg_r, avg_dr, c_rate, avg_steps, t_out = quick_evaluation(pomdp, policy, rng, 1000)

println("Collisions ", c_rate*100)
println("Avg steps ", avg_steps)
println("Time outs ", t_out*100)
println("avg disc reward ", avg_dr)
println("avg reward ", avg_r)

spbackup = deepcopy(sp)
sbackup = deepcopy(s)
bbackup = deepcopy(b)

neighbor = :fore
o = sp