###### This file is part of the MendelImpute.jl package.
###### It contains relevant code to solve the least squares criteria
###### `x ≈ hᵢ + hⱼ`.

"""
    compute_optimal_haplotypes!(haplotype1, haplotype2, compressed_Hunique, ...)

Computes the best haplotype pair `(hᵢ, hⱼ)` for each genotype vector `x` and
stores result in `haplotype1` and `haplotype2`.

# Arguments
- `haplotype1`: Person `i` strand1 haplotype in window `w` is `haplotype1[i][w]`
- `haplotype2`: Person `i` strand2 haplotype in window `w` is `haplotype2[i][w]`
- `compressed_Hunique`: A `CompressedHaplotypes` object
- `X`: the full genotype matrix possibly with missings. Each column is an
    individual.
- `X_pos`: Position of each SNP in `X`
- `typed_snps_error`: `typed_snps_error[i]` is the % of SNPs correctly imputed.
    `1` is best, `0` is worse.
- `stepscreen`: Boolean indicating whether to use a stepwise heuristic to screen
    for top haplotypes instead of performing global search
- `tf`: This option solves the least squares objective on only
    `tf` unique haplotypes.
- `scale_allelefreq` Boolean indicating whether to give rare SNPs more weight
    scaled by `wᵢ = 1 / √2p(1-p)` where max weight is 2.
- `max_haplotypes` Maximum number of haplotypes for using to global search.
    This number should be specified along with `stepscreen` or `tf`.
- `rescreen` Boolean for a more rigorous global search.
- `timers`

# Timers:
- `t1` = screening for top haplotypes (0 unless `tf` option is on)
- `t2` = BLAS3 mul! to get M and N
- `t3` = haplopair search
- `t4` = rescreen time (0 unless `rescreen = true`)
- `t5` = initializing missing
- `t6` = allocating and viewing
- `t7` = index conversion
"""
function compute_optimal_haplotypes!(
    haplotype1::AbstractVector,
    haplotype2::AbstractVector,
    compressed_Hunique::CompressedHaplotypes,
    X::AbstractMatrix,
    X_pos::AbstractVector,
    stepscreen::Union{Nothing, Int},
    tf::Union{Nothing, Int}, # thinning factor
    scale_allelefreq::Bool,
    max_haplotypes::Int,
    rescreen::Bool
    )
    # constants
    people = size(X, 2)
    ref_snps = length(compressed_Hunique.pos)
    max_width, max_d = MendelImpute.max_dim(compressed_Hunique)
    windows = length(haplotype1[1])
    threads = Threads.nthreads()
    inv_sqrt_allele_var = nothing
    winranges = compressed_Hunique.X_window_range
    overlap = compressed_Hunique.overlap

    # allocate working arrays
    timers = [zeros(7*8) for _ in 1:threads] # 8 for spacing
    timers[1][48] += @elapsed begin # time for allocating
        happair1 = [ones(Int32, people)               for _ in 1:threads]
        happair2 = [ones(Int32, people)               for _ in 1:threads]
        hapscore = [zeros(Float32, people)            for _ in 1:threads]
        Xwork    = [zeros(Float32, max_width, people) for _ in 1:threads]
        Hwork    = [zeros(Float32, max_width, max_d)  for _ in 1:threads]
        if !isnothing(tf)
            maxindx = [zeros(Int32, tf)         for _ in 1:threads]
            maxgrad = [zeros(Float32, tf)       for _ in 1:threads]
            Hk = [zeros(Float32, max_width, tf) for _ in 1:threads]
            Xi = [zeros(Float32, max_width)     for _ in 1:threads]
            M  = [zeros(Float32, tf, tf)        for _ in 1:threads]
            N  = [zeros(Float32, tf)            for _ in 1:threads]
            R  = [zeros(Float32, max_d, people) for _ in 1:threads]
        elseif !isnothing(stepscreen)
            maxindx = [zeros(Int32,   stepscreen) for _ in 1:threads]
            maxgrad = [zeros(Float32, stepscreen) for _ in 1:threads]
            M  = [zeros(Float32, max_d, max_d)    for _ in 1:threads]
            Nt = [zeros(Float32, max_d, people)   for _ in 1:threads]
        else
            M = [zeros(Float32, max_d,  max_d) for _ in 1:threads]
            N = [zeros(Float32, people, max_d) for _ in 1:threads]
        end
    end
    pmeter = Progress(windows, 5, "Computing optimal haplotypes...")

    # for w in 1:windows
    Threads.@threads for w in 1:windows
        id = Threads.threadid()
        t6 = @elapsed begin
            Hw_aligned = compressed_Hunique.CW_typed[w].uniqueH
            Xrange = extend_to_overlap_range(compressed_Hunique, w, overlap)
            Xw_aligned = view(X, Xrange, :)
            d = size(Hw_aligned, 2)
        end

        # weight snp by inverse allele variance if requested
        t6 += @elapsed if scale_allelefreq
            Hw_range = compressed_Hunique.Hstart[w]:(w ==
                windows ? ref_snps : compressed_Hunique.Hstart[w + 1] - 1)
            Hw_snp_pos = indexin(X_pos[winranges[w]],
                compressed_Hunique.pos[Hw_range])
            inv_sqrt_allele_var = compressed_Hunique.altfreq[Hw_snp_pos]
            map!(x -> x < 0.15 ? 1.98 : 1 / sqrt(2*x*(1-x)),
                inv_sqrt_allele_var, inv_sqrt_allele_var) # set min pᵢ = 0.15
        end

        # compute top haplotype pairs for each sample in current window
        if !isnothing(stepscreen) && d > max_haplotypes
            # find hᵢ via stepwise regression, then find hⱼ via global search
            t1, t2, t3, t4, t5, t6 = haplopair_stepscreen!(Xw_aligned,
                Hw_aligned, r=stepscreen,
                inv_sqrt_allele_var=inv_sqrt_allele_var, happair1=happair1[id],
                happair2=happair2[id], hapscore=hapscore[id],
                maxindx=maxindx[id], maxgrad=maxgrad[id], Xwork=Xwork[id],
                Hwork=Hwork[id], Nt=Nt[id], M=M[id])
        elseif !isnothing(tf) && d > max_haplotypes
            # haplotype thinning: search all (hᵢ, hⱼ) pairs where hᵢ ≈ x ≈ hⱼ
            t1, t2, t3, t4, t5, t6 = haplopair_thin_BLAS2!(Xw_aligned,
                Hw_aligned, allele_freq=inv_sqrt_allele_var, keep=tf,
                happair1=happair1[id], happair2=happair2[id],
                hapscore=hapscore[id], maxindx=maxindx[id], maxgrad=maxgrad[id],
                Xi=Xi[id], N=N[id], Hk=Hk[id], M=M[id], Xwork=Xwork[id],
                Hwork=Hwork[id], R=R[id])
        elseif rescreen
            # global search + searching ||x - hᵢ - hⱼ|| on observed entries
            t1, t2, t3, t4, t5, t6 = haplopair_rescreen!(Xw_aligned,
                Hw_aligned, happair1=happair1[id], happair2=happair2[id],
                hapscore=hapscore[id], Xwork=Xwork[id])
        else
            # global search
            t1, t2, t3, t4, t5, t6 = haplopair!(Xw_aligned, Hw_aligned,
                inv_sqrt_allele_var=inv_sqrt_allele_var, happair1=happair1[id],
                happair2=happair2[id], hapscore=hapscore[id], Xwork=Xwork[id],
                Hwork=Hwork[id], M=M[id], N = N[id])
        end

        # save result
        t7 = @elapsed save_haplotypes!(haplotype1, haplotype2, happair1[id],
            happair2[id], compressed_Hunique, w)

        # record timings and haplotypes (× 8 to avoid false sharing)
        timers[id][8]  += t1
        timers[id][16] += t2
        timers[id][24] += t3
        timers[id][32] += t4
        timers[id][40] += t5
        timers[id][48] += t6
        timers[id][56] += t7

        # update progress
        next!(pmeter)
    end

    return sum(timers) ./ threads
end

"""
    save_haplotypes!(haplotype1, haplotype2, happair1, happair2, ...)

Helper function to convert `happair`s (which index off unique haplotypes) to
indices of full haplotype pool, and store them in `haplotype`s.

# Arguments
- `haplotype1`: Person `i` strand1 haplotype in window `w` is `haplotype1[i][w]`
- `haplotype2`: Person `i` strand2 haplotype in window `w` is `haplotype2[i][w]`
- `happair1`: Optimal haplotype pair in strand1 of current window, indexes off
    of unique haplotypes.
- `happair2`: Optimal haplotype pair in strand2 of current window, indexes off
    of unique haplotypes.
- `compressed_Hunique`: A `CompressedHaplotypes` object
- `window` current window.
"""
function save_haplotypes!(
    haplotype1::Vector{Vector{Int32}},
    haplotype2::Vector{Vector{Int32}},
    happair1::Vector{Int32},
    happair2::Vector{Int32},
    compressed_Hunique::CompressedHaplotypes,
    window::Int,
    )

    people = length(haplotype1)
    @inbounds for i in 1:people
        haplotype1[i][window] = unique_idx_to_complete_idx(
            happair1[i], window, compressed_Hunique)
        haplotype2[i][window] = unique_idx_to_complete_idx(
            happair2[i], window, compressed_Hunique)
    end
    return nothing
end

# """
# Records optimal-redundant haplotypes for each window.

# Warning: This function is called in a multithreaded loop. If you modify this
# function you must check whether imputation accuracy is affected (when run with
# >1 threads).

# # Arguments:
# - `window_idx`: window in current chunk
# - `window_overall`: window index in terms of every windows
# """
# function compute_redundant_haplotypes!(
#     redundant_haplotypes::Vector{OptimalHaplotypeSet},
#     Hunique::CompressedHaplotypes,
#     happair1::AbstractVector,
#     happair2::AbstractVector,
#     window_idx::Int,
#     window_overall::Int,
#     storage1 = falses(nhaplotypes(Hunique)),
#     storage2 = falses(nhaplotypes(Hunique))
#     )

#     people = length(redundant_haplotypes)

#     @inbounds for k in 1:people
#         # convert happairs from unique idx to complete idx
#         Hi_idx = unique_idx_to_complete_idx(happair1[k], window_overall,
#             Hunique)
#         Hj_idx = unique_idx_to_complete_idx(happair2[k], window_overall,
#             Hunique)

#         # strand1
#         storage1 .= false
#         if haskey(Hunique.CW_typed[window_overall].hapmap, Hi_idx)
#             h1_set = Hunique.CW_typed[window_overall].hapmap[Hi_idx]
#             for i in h1_set
#                 storage1[i] = true
#             end
#         else
#             storage1[Hi_idx] = true # Hi_idx is singleton (i.e. unique)
#         end

#         # strand2
#         storage2 .= false
#         if haskey(Hunique.CW_typed[window_overall].hapmap, Hj_idx)
#             h2_set = Hunique.CW_typed[window_overall].hapmap[Hj_idx]
#             for i in h2_set
#                 storage2[i] = true
#             end
#         else
#             storage2[Hj_idx] = true # Hj_idx is singleton (i.e. unique)
#         end

#         # redundant_haplotypes[k].strand1[window_idx] = copy(storage1)
#         # redundant_haplotypes[k].strand2[window_idx] = copy(storage2)
#         if isassigned(redundant_haplotypes[k].strand1, window_idx)
#             redundant_haplotypes[k].strand1[window_idx] .= storage1
#             redundant_haplotypes[k].strand2[window_idx] .= storage2
#         else
#             redundant_haplotypes[k].strand1[window_idx] = copy(storage1)
#             redundant_haplotypes[k].strand2[window_idx] = copy(storage2)
#         end
#     end

#     return nothing
# end

"""
    screen_flanking_windows!(haplotype1, haplotype2, compressed_Hunique, X)

For each window's haplotype pair, tests whether adjacent window's haplotype
pairs produce better error on the observed entries.

# Arguments
- `haplotype1`: Person `i` strand1 haplotype in window `w` is `haplotype1[i][w]`
- `haplotype2`: Person `i` strand2 haplotype in window `w` is `haplotype2[i][w]`
- `compressed_Hunique`: A `CompressedHaplotypes` object
- `X`: the full genotype matrix possibly with missings. Each column is an
    individual.
"""
function screen_flanking_windows!(
    haplotype1::AbstractVector,
    haplotype2::AbstractVector,
    compressed_Hunique::CompressedHaplotypes,
    X::AbstractMatrix
    )

    people = length(haplotype1)
    haplotypes = nhaplotypes(compressed_Hunique)
    width = compressed_Hunique.width
    windows = length(haplotype1[1])
    winranges = compressed_Hunique.X_window_range

    for w in 1:windows
        Hw_aligned = compressed_Hunique.CW_typed[w].uniqueH
        Xw_aligned = view(X, winranges[w], :)

        for i in 1:people
            # calculate observed error for current pair
            h1_curr_complete = haplotype1[i][w] # complete index
            h2_curr_complete = haplotype2[i][w] # complete index
            h1_curr = complete_idx_to_unique_typed_idx(h1_curr_complete, w,
                compressed_Hunique) # unique index
            h2_curr = complete_idx_to_unique_typed_idx(h2_curr_complete, w,
                compressed_Hunique) # unique index
            curr_err = observed_error(Xw_aligned, i, Hw_aligned, h1_curr,
                h2_curr)

            # consider previous pair
            if w != 1
                h1_prev_complete = haplotype1[i][w - 1]
                h2_prev_complete = haplotype2[i][w - 1]
                h1_prev = complete_idx_to_unique_typed_idx(h1_prev_complete,
                    w, compressed_Hunique)
                h2_prev = complete_idx_to_unique_typed_idx(h2_prev_complete,
                    w, compressed_Hunique)
                prev_err = observed_error(Xw_aligned, i, Hw_aligned, h1_prev,
                    h2_prev)
                if prev_err < curr_err
                    h1_curr, h2_curr, curr_err = h1_prev, h2_prev, prev_err
                end
            end

            # consider next pair
            if w != windows
                h1_next_complete = haplotype1[i][w + 1]
                h2_next_complete = haplotype2[i][w + 1]
                h1_next = complete_idx_to_unique_typed_idx(h1_next_complete,
                    w, compressed_Hunique) # unique index
                h2_next = complete_idx_to_unique_typed_idx(h2_next_complete,
                    w, compressed_Hunique) # unique index
                next_err = observed_error(Xw_aligned, i, Hw_aligned, h1_next,
                    h2_next)
                if next_err < curr_err
                    h1_curr, h2_curr, curr_err = h1_next, h2_next, next_err
                end
            end

            # convert from unique idx back to complete idx
            H1_idx = unique_idx_to_complete_idx(h1_curr, w, compressed_Hunique)
            H2_idx = unique_idx_to_complete_idx(h2_curr, w, compressed_Hunique)
            haplotype1[i][w] = H1_idx
            haplotype2[i][w] = H2_idx
        end
    end

    return nothing
end

# uses dynamic programming. Only the first 1000 haplotype pairs will be saved.
# function compute_redundant_haplotypes!(
#     redundant_haplotypes::Vector{Vector{Vector{T}}},
#     Hunique::CompressedHaplotypes,
#     happair1::AbstractVector,
#     happair2::AbstractVector,
#     window_idx::Int,
#     window_overall::Int,
#     storage1 = falses(nhaplotypes(Hunique)),
#     storage2 = falses(nhaplotypes(Hunique))
#     ) where T <: Tuple{Int32, Int32}

#     people = length(redundant_haplotypes)

#     @inbounds for k in 1:people
#         # convert happairs from unique idx to complete idx
#         Hi_idx = unique_idx_to_complete_idx(happair1[k], window_overall, Hunique)
#         Hj_idx = unique_idx_to_complete_idx(happair2[k], window_overall, Hunique)

#         # find haplotypes that match Hi_idx and Hj_idx on typed snps
#         h1_set = get(Hunique.CW_typed[window_overall].hapmap, Hi_idx, Hi_idx)
#         h2_set = get(Hunique.CW_typed[window_overall].hapmap, Hj_idx, Hj_idx)

#         # save first 1000 haplotype pairs
#         for h1 in h1_set, h2 in h2_set
#             if length(redundant_haplotypes[k][window_idx]) < 1000
#                 push!(redundant_haplotypes[k][window_idx], (h1, h2))
#             else
#                 break
#             end
#         end
#     end

#     return nothing
# end

"""
    haplopair(X, H)

Calculate the best pair of haplotypes in `H` for each individual in `X`.
Missing data is initialized as 2x alternate allele freq.

# Input
* `X`: `p × n` genotype matrix possibly with missings. Each column is an individual.
* `H`: `p × d` haplotype matrix. Each column is a haplotype.

# Optional arguments
* `happair1`: optimal haplotype pairs for strand 1.
* `happair2`: optimal haplotype pairs for strand 2.
* `hapscore`: haplotyping score. 0 means best. Larger means worse.
* `inv_sqrt_allele_var`: If not nothing, SNP `i` will be scaled according to
    `inv_sqrt_allele_var[i]`.
* `Xwork`: A `Float32` version of `X`, missing entries are filled with `2maf`
* `Hwork`: A `Float32` version of `H`.
* `M`: `d x d` matrix with entries `M[i, j] = 2dot(Hwork[:, i], Hwork[:, j]) +
    sumabs2(Hwork[:, i]) + sumabs2(Hwork[:, j])`. Only the upper triangular
    part of `M` is used.
* `N`: `n x d` matrix `2Xwork'Hwork`.
"""
function haplopair!(
    X::AbstractMatrix, # p × n
    H::AbstractMatrix; # p × d
    # preallocated vectors
    happair1::AbstractVector = ones(Int32, size(X, 2)), # length n
    happair2::AbstractVector = ones(Int32, size(X, 2)), # length n
    hapscore::AbstractVector = Vector{Float32}(undef, size(X, 2)), # length n
    inv_sqrt_allele_var::Union{Nothing, AbstractVector} = nothing, # length p
    # preallocated matrices
    Xwork::AbstractMatrix{Float32}=Matrix{Float32}(undef,size(X,1),size(X,2)), # p × n
    Hwork::AbstractMatrix{Float32}=convert(Matrix{Float32}, H),                # p × d
    M::AbstractMatrix{Float32}=Matrix{Float32}(undef, size(H, 2), size(H, 2)), # d × d
    N::AbstractMatrix{Float32}=Matrix{Float32}(undef, size(X, 2), size(H, 2)), # n × d
    )
    p,  n = size(X)
    pp, d = size(H)
    p == pp || error("size of H doesn't match X! Expected $p got $pp")

    # create views for matrices
    t6 = @elapsed begin
        Mwork_view = view(M, 1:d, 1:d)
        Nwork_view = view(N, :, 1:d)
        Hwork_view = view(Hwork, 1:p, 1:d)
        Xwork_view = view(Xwork, 1:p, :)
        copyto!(Hwork_view, H) # update Hwork_view, others are updated later
    end

    # initializes missing
    t5 = @elapsed initXfloat!(Xwork_view, X)

    t2, t3 = haplopair!(Xwork_view, Hwork_view, Mwork_view, Nwork_view,
        happair1, happair2, hapscore, inv_sqrt_allele_var)
    t1 = t4 = 0.0 # no time spent on haplotype thinning or rescreening

    return t1, t2, t3, t4, t5, t6
end

"""
    haplopair!(X, H, M, N, happair, hapscore)

Calculate the best pair of haplotypes in `H` for each individual in `X`.
Overwite `M` by `M[i, j] = 2dot(H[:, i], H[:, j]) + sumabs2(H[:, i]) +
sumabs2(H[:, j])`, `N` by `2X'H`, `happair` by optimal haplotype pair, and
`hapscore` by objective value from the optimal haplotype pair.

# Input
* `X`: `p x n` genotype matrix. Each column is an individual.
* `H`: `p x d` haplotype matrix. Each column is a haplotype.
* `M`: overwritten by `M[i, j] = 2dot(H[:, i], H[:, j]) + sumabs2(H[:, i]) +
    sumabs2(H[:, j])`.
* `N`: overwritten by `n x d` matrix `2X'H`.
* `happair`: optimal haplotype pair. `X[:, k] ≈ H[:, happair[k, 1]] +
    H[:, happair[k, 2]]`.
* `hapscore`: haplotyping score. 0 means best. Larger means worse.
"""
function haplopair!(
    X::AbstractMatrix{Float32},
    H::AbstractMatrix{Float32},
    M::AbstractMatrix{Float32},
    N::AbstractMatrix{Float32},
    happair1::AbstractVector{Int32},
    happair2::AbstractVector{Int32},
    hapscore::AbstractVector{Float32},
    inv_sqrt_allele_var::Union{Nothing, AbstractVector}
    )

    p, n, d = size(X, 1), size(X, 2), size(H, 2)

    # assemble M (upper triangular only)
    t2 = @elapsed begin
        if !isnothing(inv_sqrt_allele_var)
            H .*= inv_sqrt_allele_var # wᵢ = 1/√2p(1-p)
        end
        mul!(M, Transpose(H), H)
        for j in 1:d, i in 1:(j - 1) # off-diagonal
            @inbounds M[i, j] = 2M[i, j] + M[i, i] + M[j, j]
        end
        for j in 1:d # diagonal
            @inbounds M[j, j] *= 4
        end

        # assemble N
        if !isnothing(inv_sqrt_allele_var)
            H .*= inv_sqrt_allele_var # wᵢ = 1/2p(1-p)
        end
        mul!(N, Transpose(X), H)
        @simd for I in eachindex(N)
            N[I] *= 2
        end
    end

    # computational routine
    t3 = @elapsed haplopair!(happair1, happair2, hapscore, M, N)

    # supplement the constant terms in objective
    t3 += @elapsed begin
        @inbounds for j in 1:n
            @simd for i in 1:p
                hapscore[j] += abs2(X[i, j])
            end
        end
    end

    return t2, t3
end

"""
    haplopair!(happair, hapscore, M, N)

Calculate the best pair of haplotypes pairs in the filtered haplotype panel
for each individual in `X` using sufficient statistics `M` and `N`.

# Note
The best haplotype pairs are column indices of the filtered haplotype panels.

# Input
* `happair1`: optimal haplotype in strand 1 for each individual.
* `happair2`: optimal haplotype in strand 2 for each individual.
* `hapmin`: minimum offered by the optimal haplotype pair.
* `M`: `d x d` matrix with entries `M[i, j] = 2dot(H[:, i], H[:, j]) +
    sumabs2(H[:, i]) + sumabs2(H[:, j])`, where `H` is the haplotype matrix
    with haplotypes in columns. Only the upper triangular part of `M` is used.
* `N`: `n x d` matrix `2X'H`, where `X` is the genotype matrix with individuals
    in columns.
"""
function haplopair!(
    happair1::AbstractVector{Int32},
    happair2::AbstractVector{Int32},
    hapmin::AbstractVector{Float32},
    M::AbstractMatrix{Float32},
    N::AbstractMatrix{Float32},
    )

    n, d = size(N)
    chunks = div(n, 4)
    fill!(hapmin, Inf32)
    @inbounds for k in 1:d, j in 1:k
        Mjk = M[j, k]
        i = 1
        # loop over individuals
        for chunk in 1:chunks # unrolled loop, equivalent to for i in 1:n
            i2, i3, i4 = i + 1, i + 2, i + 3
            score = Mjk - N[i, j] - N[i, k]
            if score < hapmin[i]
                hapmin[i], happair1[i], happair2[i] = score, j, k
            end
            score = Mjk - N[i2, j] - N[i2, k]
            if score < hapmin[i2]
                hapmin[i2], happair1[i2], happair2[i2] = score, j, k
            end
            score = Mjk - N[i3, j] - N[i3, k]
            if score < hapmin[i3]
                hapmin[i3], happair1[i3], happair2[i3] = score, j, k
            end
            score = Mjk - N[i4, j] - N[i4, k]
            if score < hapmin[i4]
                hapmin[i4], happair1[i4], happair2[i4] = score, j, k
            end
            i += 4
        end
        # handle remaining terms
        while i ≤ n
            score = Mjk - N[i, j] - N[i, k]
            if score < hapmin[i]
                hapmin[i], happair1[i], happair2[i] = score, j, k
            end
            i += 1
        end
    end
end

"""
    fillmissing!(Xm, Xwork, H, haplopairs)

Fill in missing genotypes in `X` according to haplotypes. Non-missing genotypes
remain same.

# Input
* `Xm`: `p x n` genotype matrix with missing values. Each column is an
    individual.
* `Xwork`: `p x n` genotype matrix where missing values are filled with sum of 2
    haplotypes.
* `H`: `p x d` haplotype matrix. Each column is a haplotype.
* `happair`: pair of haplotypes. `X[:, k] = H[:, happair[1][k]] +
    H[:, happair[2][k]]`.
"""
function fillmissing!(
    Xm::AbstractMatrix{Union{U, Missing}},
    Xwork::AbstractMatrix{T},
    H::AbstractMatrix{T},
    happairs::Vector{Vector{Tuple{Int32, Int32}}},
    ) where {T, U}

    p, n = size(Xm)
    best_discrepancy = typemax(eltype(Xwork))

    for j in 1:n, happair in happairs[j]
        discrepancy = zero(T)
        for i in 1:p
            if ismissing(Xm[i, j])
                tmp = H[i, happair[1]] + H[i, happair[2]]
                discrepancy += abs2(Xwork[i, j] - tmp)
                Xwork[i, j] = tmp
            end
        end
        if discrepancy < best_discrepancy
            best_discrepancy = discrepancy
        end
    end
    return best_discrepancy
end

"""
    initXfloat!(Xfloat, X)

Initializes the matrix `Xfloat` where missing values of matrix `X` by `2 x`
allele frequency and nonmissing entries of `X` are converted to type `Float32`
for subsequent BLAS routines.

# Input
* `X` is a `p x n` genotype matrix. Each column is an individual.
* `Xfloat` is the `p x n` matrix of X where missing values are filled by 2x
    allele frequency.
"""
function initXfloat!(
    Xfloat::AbstractMatrix,
    X::AbstractMatrix
    )

    T = Float32
    p, n = size(X)

    @inbounds for i in 1:p
        # allele frequency
        cnnz = zero(T)
        csum = zero(T)
        for j in 1:n
            if !ismissing(X[i, j])
                cnnz += one(T)
                csum += convert(T, X[i, j])
            end
        end
        # set missing values to 2freq, unless cnnz is 0
        imp = (cnnz == 0 ? zero(T) : csum / cnnz)
        for j in 1:n
            if ismissing(X[i, j])
                Xfloat[i, j] = imp
            else
                Xfloat[i, j] = convert(T, X[i, j])
            end
        end
    end

    any(isnan, Xfloat) && error("Xfloat has NaN during initialization!")
    any(isinf, Xfloat) && error("Xfloat has Inf during initialization!")
    any(ismissing, Xfloat) && error("Xfloat has Missing during initialization!")

    return nothing
end


function binary_op(op, W, @nospecialize(T))
  ty = LLVM_TYPES[T]
  if isone(W)
    V = T
  else
    ty = "<$W x $ty>"
    V = NTuple{W,VecElement{T}}
  end
  instrs = "%res = $op $ty %0, %1\nret $ty %res"
  call = :($LLVMCALL($instrs, $V, Tuple{$V,$V}, data(v1), data(v2)))
  W > 1 && (call = Expr(:call, :Vec, call))
  Expr(:block, Expr(:meta, :inline), call)
end

# Integer
for (op,f) ∈ [("add",:+),("sub",:-),("mul",:*),("shl",:<<)]
  ff = Symbol('v', op)
  fnsw = Symbol(ff,"_nsw")
  fnuw = Symbol(ff,"_nuw")
  fnw = Symbol(ff,"_nw")
  ff_fast = if op == "shl"
    @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($op, W, T)
  else
    ff_fast = :(FastMath.$(Symbol(op * "_fast")))
    @eval begin
      @generated $ff_fast(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($op, W, T)
      @inline Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = $ff_fast(v1, v2)
    end
  end
  @eval begin
    @generated $fnsw(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($(op * " nsw"), W, T)
    @generated $fnuw(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($(op * " nuw"), W, T)
    @generated $fnw(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($(op * " nsw nuw"), W, T)
    @generated $fnsw(v1::T, v2::T) where {T<:IntegerTypesHW} = binary_op($(op * " nsw"), 1, T)
    @generated $fnuw(v1::T, v2::T) where {T<:IntegerTypesHW} = binary_op($(op * " nuw"), 1, T)
    @generated $fnw(v1::T, v2::T) where {T<:IntegerTypesHW} = binary_op($(op * " nsw nuw"), 1, T)
  end
end
for (op,f) ∈ [("div",:÷),("rem",:%)]
  @eval begin
    @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:Integer} = binary_op((T <: Signed ? 's' : 'u') * $op, W, T)
  end
end
@inline div_fast(x::T, y::T) where {T <: SignedHW} = Base.sdiv_int(x,y)
@inline div_fast(x::T, y::T) where {T <: UnsignedHW} = Base.udiv_int(x,y)
@inline div_fast(x::AbstractSIMD, y::AbstractSIMD) = x ÷ y
@inline vcld(x, y) = (div_fast((x - one(x)), y) + one(x))
@inline function divrem_fast(x, y)
  d = div_fast(x, y)
  r = x - (d * y)
  d, r
end
for (op,sub) ∈ [
  ("ashr",:SignedHW),
  ("lshr",:UnsignedHW),
  ("lshr",:IntegerTypesHW),
  ("and",:IntegerTypesHW),
  ("or",:IntegerTypesHW),
  ("xor",:IntegerTypesHW)
  ]
  ff = sub === :UnsignedHW ? :vashr : Symbol('v', op)
  @eval begin
    @generated $ff(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:$sub}  = binary_op($op, W, T)
    @generated $ff(v1::T, v2::T) where {T<:$sub}  = binary_op($op, 1, T)
  end
end

for (op,f,s) ∈ [("ashr",:>>,0x01),("lshr",:>>,0x02),("lshr",:>>>,0x03),("and",:&,0x03),("or",:|,0x03),("xor",:⊻,0x03)]
  fdef = Expr(:where, :(Base.$f(v1::Vec{W,T}, v2::Vec{W,T})), :W)
  if iszero(s & 0x01)
    @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:UnsignedHW} = binary_op($op, W1, T, W2)
  elseif iszero(s & 0x02)
    @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:SignedHW} = binary_op($op, W1, T, W2)
  else
    @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:IntegerTypesHW} = binary_op($op, W1, T, W2)
  end
  if op === "ashr" # replace first iteration with <<
    f = :(<<)
  end
  @eval begin
    @inline Base.$f(v::Vec, i::Integer) = ((x, y) = promote(v, i); $f(x, y))
    @inline Base.$f(i::Integer, v::Vec) = ((x, y) = promote(i, v); $f(x, y))
  end
end

for (op,f,ff) ∈ [("fadd",:(+),:add_fast),("fsub",:(-),:sub_fast),("fmul",:(*),:mul_fast),("fdiv",:(/),:div_fast),("frem",:(%),:rem_fast)]

  fop_fast = f === :(/) ? "fdiv fast" : op * ' ' * fast_flags(true)
  fop_contract = op * ' ' * fast_flags(false)
  @eval begin
    @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = binary_op($fop_contract, W, T)
    @generated FastMath.$ff(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = binary_op($fop_fast, W, T)
  end
end

@inline Base.div(v1::AbstractSIMD{W,T}, v2::AbstractSIMD{W,T}) where {W,T<:FloatingTypes} = floor(signed(Base.uinttype(T)), FastMath.div_fast(v1, v2))

@inline Base.:(/)(a::AbstractSIMDVector{W}, b::AbstractSIMDVector{W}) where {W} = float(a) / float(b)
@inline FastMath.div_fast(a::AbstractSIMDVector{W}, b::AbstractSIMDVector{W}) where {W} = FastMath.div_fast(float(a), float(b))


@testset "testing catobs & getindex operations missing values" begin
    a = BagNode(ArrayNode(rand(3,4)), [1:4], nothing)
    e = BagNode(missing, AlignedBags([0:-1]), nothing)

    x = reduce(catobs, [a, e])
    @test x.data.data == a.data.data
    @test x.bags.bags == [1:4, 0:-1]

    x = reduce(catobs, [e, a])
    @test x.data.data == a.data.data
    @test x.bags.bags == [0:-1, 1:4]

    x = reduce(catobs, [e, e])
    @test ismissing(x.data)
    @test x.bags.bags == [0:-1, 0:-1]

    x = reduce(catobs, [a, e])
    @test  isnothing(x[2].metadata)
    Mill.emptyismissing!(true)
    @test  ismissing(x[2].data)
    Mill.emptyismissing!(false)
    @test  nobs(x[2].data) == 0
    @test  x[2].bags.bags == [0:-1]

    @test  x[1].data.data == a.data.data
    @test  x[1].bags.bags == [1:4]

    @test catobs(a, e).data.data == a.data.data
    @test catobs(a, e).bags.bags == [1:4, 0:-1]

    a = BagNode(a, [1:1], missing)
    @test catobs(a, e).bags.bags == [1:1, 0:-1]
    @test catobs(e, a).bags.bags == [0:-1, 1:1]
    @test catobs(e, a).data.data.data == a.data.data.data

    @test catobs(a,e,a,e).data.data.data == catobs(a,a).data.data.data
    @test catobs(a,e,a,e).bags.bags == [1:1, 0:-1, 2:2, 0:-1]

    @test_throws ErrorException BagNode(missing, AlignedBags([1:3]), nothing)
    @test_throws ErrorException BagNode(missing, Mill.ScatteredBags([[1,2,3]]), nothing)
end

@testset "testing catobs & getindex operations missing values for weighted" begin
    a = WeightedBagNode(ArrayNode(rand(3,4)), [1:4], [1.0, 0.0, 1.0, 0.5], nothing)
    e = WeightedBagNode(missing, AlignedBags([0:-1]), [], nothing)

    x = reduce(catobs, [a, e])
    @test x.data.data == a.data.data
    @test x.bags.bags == [1:4, 0:-1]
    @test x.weights == [1.0, 0.0, 1.0, 0.5]

    x = reduce(catobs, [e, a])
    @test x.data.data == a.data.data
    @test x.bags.bags == [0:-1, 1:4]
    @test x.weights == [1.0, 0.0, 1.0, 0.5]

    x = reduce(catobs, [e, e])
    @test ismissing(x.data)
    @test x.bags.bags == [0:-1, 0:-1]
    @test x.weights == []

    x = reduce(catobs, [a, e])
    @test  isnothing(x[2].metadata)
    Mill.emptyismissing!(true)
    @test  ismissing(x[2].data)
    Mill.emptyismissing!(false)
    @test  nobs(x[2].data) == 0
    @test  x[2].bags.bags == [0:-1]
    @test x[1].weights == [1.0, 0.0, 1.0, 0.5]
    @test x[2].weights == []
    @test  x[1].data.data == a.data.data
    @test  x[1].bags.bags == [1:4]
end

@testset "testing model operations missing values" begin
    a = BagNode(ArrayNode(rand(3, 4)), [1:4], nothing)
    e = BagNode(missing, AlignedBags([0:-1]), nothing)
    m = BagModel(ArrayModel(Dense(3, 2)), SegmentedMean(2), ArrayModel(Dense(2, 2)))

    @testset "BagNode" begin
        x = reduce(catobs, [a, e])
        @test m(x).data[:, 1] ≈ m(a).data
        @test m(x).data[:, 2] ≈ m(e).data
        @test m(x).data ≈ hcat(m(a).data, m(e).data)
    end

    @testset "ProductNode" begin
        t1 = ProductNode((a, a))
        t2 = ProductNode((a, e))
        t3 = ProductNode((e, a))
        t4 = ProductNode((e, e))
        tt = [t1, t2, t3, t4]
        x  = reduce(catobs, tt)
        tm = ProductModel((m, m))
        o = tm(x).data
        for i in 1:length(tt)
            @test o[:,i] ≈ tm(tt[i]).data
        end
        for i in 1:length(tt)
            @test o[:,i] ≈ tm(x[i]).data
        end
    end
end

@testset "testing reduction with missing values" begin
    a = ProductNode((x=ArrayNode(rand(3,4)), y=ArrayNode(rand(3,4))))
    b = ProductNode((x=ArrayNode(rand(3,4)), y=ArrayNode(rand(3,4))))

    c = ProductNode((
        x=BagNode(missing, AlignedBags([0:-1]), nothing), y=ArrayNode(rand(3,1)),
    ))
    d = ProductNode((
        x=BagNode(a, [1:4], nothing), y=ArrayNode(rand(3,1)),
    ))

    ba = BagNode(a, [1:4], nothing)
    bb = BagNode(b, [1:4], nothing)
    bc = BagNode(missing, AlignedBags([0:-1]), nothing)

    bba = BagNode(c, [1:1], nothing)
    bbb = BagNode(d, [1:1], nothing)
    bbc = BagNode(missing, AlignedBags([0:-1]), nothing)

    x = reduce(catobs, [c, d])
    y = reduce(catobs, [ba, bb, bc])
    z = reduce(catobs, [bba, bbb, bbc])

    mx = reflectinmodel(x, d -> Dense(d, 2, relu))
    my = reflectinmodel(y, d -> Dense(d, 2, relu))
    mz = reflectinmodel(z, d -> Dense(d, 2, relu))

    @test mx(x).data[:, 1] ≈ mx(c).data
    @test mx(x).data[:, 2] ≈ mx(d).data

    @test my(y).data[:, 1] ≈ my(ba).data
    @test my(y).data[:, 2] ≈ my(bb).data
    @test my(y).data[:, 3] ≈ my(bc).data

    @test mz(z).data[:, 1] ≈ mz(bba).data
    @test mz(z).data[:, 2] ≈ mz(bbb).data
    @test mz(z).data[:, 3] ≈ mz(bbc).data
end


"""
A matrix of colors is interpreted as an image
"""
_default(::Node{Array{RGBA{N0f8}, 2}}, ::Style{:default}, ::Dict{Symbol,Any})

function _default(main::MatTypes{T}, ::Style, data::Dict) where T <: Colorant
    @gen_defaults! data begin
        spatialorder = "yx"
    end
    if !(spatialorder in ("xy", "yx"))
        error("Spatial order only accepts \"xy\" or \"yz\" as a value. Found: $spatialorder")
    end
    ranges = get(data, :ranges) do
        const_lift(main, spatialorder) do m, s
            (0:size(m, s == "xy" ? 1 : 2), 0:size(m, s == "xy" ? 2 : 1))
        end
    end
    delete!(data, :ranges)
    @gen_defaults! data begin
        image = main => (Texture, "image, can be a Texture or Array of colors")
        primitive::GLUVMesh2D = const_lift(ranges) do r
            x, y = minimum(r[1]), minimum(r[2])
            xmax, ymax = maximum(r[1]), maximum(r[2])
            SimpleRectangle{Float32}(x, y, xmax - x, ymax - y)
        end
        preferred_camera      = :orthographic_pixel
        fxaa                  = false
        shader                = GLVisualizeShader(
            "fragment_output.frag", "uv_vert.vert", "texture.frag",
            view = Dict("uv_swizzle" => "o_uv.$(spatialorder)")
        )
    end
end
function _default(main::VectorTypes{T}, ::Style, data::Dict) where T <: Colorant
    @gen_defaults! data begin
        image                 = main => (Texture, "image, can be a Texture or Array of colors")
        primitive::GLUVMesh2D = SimpleRectangle{Float32}(0f0, 0f0, length(to_value(main)), 50f0) => "the 2D mesh the image is mapped to. Can be a 2D Geometry or mesh"
        preferred_camera      = :orthographic_pixel
        fxaa                  = false
        shader                = GLVisualizeShader(
            "fragment_output.frag", "uv_vert.vert", "texture.frag",
            view = Dict("uv_swizzle" => "o_uv.xy")
        )
    end
end

"""
A matrix of Intensities will result in a contourf kind of plot
"""
function gl_heatmap(main::MatTypes{T}, data::Dict) where T <: AbstractFloat
    main_v = to_value(main)
    @gen_defaults! data begin
        ranges = (0:size(main_v, 1), 0:size(main_v, 2))
    end
    prim = const_lift(data[:ranges]) do ranges
        x, y, xw, yh = minimum(ranges[1]), minimum(ranges[2]), maximum(ranges[1]), maximum(ranges[2])
        SimpleRectangle{Float32}(x, y, xw-x, yh-y)
    end
    delete!(data, :ranges)
    @gen_defaults! data begin
        intensity             = main => Texture
        color_map             = default(Vector{RGBA{N0f8}},s) => Texture
        primitive::GLUVMesh2D = prim
        color_norm            = const_lift(extrema2f0, main)
        stroke_width::Float32 = 0.05f0
        levels::Float32       = 5f0
        stroke_color          = RGBA{Float32}(1,1,1,1)
        shader                = GLVisualizeShader("fragment_output.frag", "uv_vert.vert", "intensity.frag")
        fxaa                  = false
    end
end

"""
Float matrix with the style distancefield will be interpreted as a distancefield.
A distancefield is describing a shape, with positive values denoting the inside
of the shape, negative values the outside and 0 the border
"""
function _default(main::MatTypes{T}, s::style"distancefield", data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        distancefield = main => Texture
        shape         = DISTANCEFIELD
        fxaa          = false
    end
    rect = SimpleRectangle{Float32}(0f0,0f0, size(to_value(main))...)
    _default((rect, Point2f0[0]), s, data)
end

export play
"""
    play(img, timedim, t)

Slice a 3D array along axis `timedim` at time `t`.
This can be used to treat a 3D array like a video and create an image stream from it.
"""
function play(array::Array{T, 3}, timedim::Integer, t::Integer) where T
    index = ntuple(dim-> dim == timedim ? t : Colon(), Val(3))
    array[index...]
end

"""
    play(buffer, video_stream, t)

Plays a video stream from VideoIO.jl. You need to supply the image `buffer`,
which will be reused for better performance.
"""
function play(buffer::Array{T, 2}, video_stream, t) where T
    eof(video_stream) && seekstart(video_stream)
    w, h = size(buffer)
    buffer = reinterpret(UInt8, buffer, (3, w,h))
    read!(video_stream, buffer) # looses type and shape
    return reinterpret(T, buffer, (w,h))
end

unwrap(img::AxisArrays.AxisArray) = unwrap(img.data)
unwrap(img::AbstractArray) = img

GLAbstraction.gl_convert(::Type{T}, img::AbstractArray) where {T} = _gl_convert(T, unwrap(img))
_gl_convert(::Type{T}, img::Array) where {T} = gl_convert(T, img)

"""
    play(img)

Turns an Image into a video stream
"""
function play(img::HasAxesArray{T, 3}) where T
    ax = ImageAxes.timeaxis(img)
    if ImageAxes.timeaxis(img) != nothing
        return const_lift(play, unwrap(img), ImageAxes.timedim(img), loop(1:length(ax)))
    end
    error("Image has no time axis: axes(img) = $(axes(img))")
end

"""
Takes a 3D image and decides if it is a volume or an animated Image.
"""
function _default(img::HasAxesArray{T, 3}, s::Style, data::Dict) where T
    # We could do this as a @traitfn, except that those don't
    # currently mix well with non-trait specialization.
    if ImageAxes.timeaxis(img) != nothing
        data[:spatialorder] = "yx"
        td = ImageAxes.timedim(img)
        video_signal = const_lift(play, unwrap(img), td, loop(1:size(img, timedim)))
        return _default(video_signal, s, data)
    else
        ps = ImageAxes.pixelspacing(img)
        spacing = Vec3f0(map(x-> x / maximum(ps), ps))
        pdims   = Vec3f0(map(length, indices(img)))
        dims    = pdims .* spacing
        dims    = dims/maximum(dims)
        data[:dimensions] = dims
        _default(unwrap(img), s, data)
    end
end
function _default(img::TOrSignal{T}, s::Style, data::Dict) where T <: AxisMatrix
    @gen_defaults! data begin
        ranges = const_lift(img) do img
            ps = ImageAxes.pixelspacing(img)
            spacing = Vec2f0(map(x-> x / maximum(ps), ps))
            pdims   = Vec2f0(map(length, indices(img)))
            dims    = pdims .* spacing
            dims    = dims / maximum(dims)
            (0:dims[1], 0:dims[2])
        end
    end
    _default(const_lift(unwrap, img), s, data)
end

"""
Displays 3D array as movie with 3rd dimension as time dimension
"""
function _default(img::AbstractArray{T, 3}, s::Style, data::Dict) where T
    video_signal = const_lift(play, unwrap(img), 3, loop(1:size(img, 3)))
    return _default(video_signal, s, data)
end

"""
Takes a shader as a parametric function. The shader should contain a function stubb
like this:
```GLSL
uniform float arg1; // you can add arbitrary uniforms and supply them via the keyword args
float function(float x) {
 return arg1*sin(1/tan(x));
}
```
"""
_default(func::String, s::Style{:shader}, data::Dict) = @gen_defaults! data begin
    color                 = default(RGBA, s) => Texture
    dimensions            = (120f0, 120f0)
    primitive::GLUVMesh2D = SimpleRectangle{Float32}(0f0,0f0, dimensions...)
    preferred_camera      = :orthographic_pixel
    fxaa                  = false
    shader                = GLVisualizeShader(
        "fragment_output.frag", "parametric.vert", "parametric.frag",
        view = Dict("function" => func)
    )
end

#Volumes
const VolumeElTypes = Union{Gray, AbstractFloat, Intensity}

const default_style = Style{:default}()

function _default(a::VolumeTypes{T}, s::Style{:iso}, data::Dict) where T <: VolumeElTypes
    data = @gen_defaults! data begin
        isovalue  = 0.5f0
        algorithm = IsoValue
    end
     _default(a, default_style, data)
end

function _default(a::VolumeTypes{T}, s::Style{:absorption}, data::Dict) where T<:VolumeElTypes
    data = @gen_defaults! data begin
        absorption = 1f0
        algorithm  = Absorption
    end
    _default(a, default_style, data)
end

function _default(a::VolumeTypes{T}, s::Style{:absorption}, data::Dict) where T<:RGBA
    data = @gen_defaults! data begin
        algorithm  = AbsorptionRGBA
    end
    _default(a, default_style, data)
end

using .GLAbstraction: StandardPrerender

struct VolumePrerender
    sp::StandardPrerender
end
VolumePrerender(a, b) = VolumePrerender(StandardPrerender(a, b))

function (x::VolumePrerender)()
    x.sp()
    glEnable(GL_CULL_FACE)
    glCullFace(GL_FRONT)
end

function _default(main::VolumeTypes{T}, s::Style, data::Dict) where T <: VolumeElTypes
    @gen_defaults! data begin
        volumedata       = main => Texture
        hull::GLUVWMesh  = AABB{Float32}(Vec3f0(0), Vec3f0(1))
        light_position   = Vec3f0(0.25, 1.0, 3.0)
        light_intensity  = Vec3f0(15.0)
        model            = Mat4f0(I)
        modelinv         = const_lift(inv, model)

        color_map        = default(Vector{RGBA}, s) => Texture
        color_norm       = color_map == nothing ? nothing : const_lift(extrema2f0, main)
        color            = color_map == nothing ? default(RGBA, s) : nothing

        algorithm        = MaximumIntensityProjection
        absorption       = 1f0
        isovalue         = 0.5f0
        isorange         = 0.01f0
        shader           = GLVisualizeShader("fragment_output.frag", "util.vert", "volume.vert", "volume.frag")
        prerender        = VolumePrerender(data[:transparency], data[:overdraw])
        postrender       = () -> begin
            glDisable(GL_CULL_FACE)
        end
    end
end

function _default(main::VolumeTypes{T}, s::Style, data::Dict) where T <: RGBA
    @gen_defaults! data begin
        volumedata       = main => Texture
        hull::GLUVWMesh  = AABB{Float32}(Vec3f0(0), Vec3f0(1))
        model            = Mat4f0(I)
        modelinv         = const_lift(inv, model)

        # These don't do anything but are needed for type specification in the frag shader
        color_map        = nothing
        color_norm       = nothing
        color            = color_map == nothing ? default(RGBA, s) : nothing

        algorithm        = AbsorptionRGBA
        shader           = GLVisualizeShader("fragment_output.frag", "util.vert", "volume.vert", "volume.frag")
        prerender        = VolumePrerender()
        postrender       = () -> begin
            glDisable(GL_CULL_FACE)
        end
    end
end

function _default(main::IndirectArray{T}, s::Style, data::Dict) where T <: RGBA
    @gen_defaults! data begin
        volumedata       = main.index => Texture
        hull::GLUVWMesh  = AABB{Float32}(Vec3f0(0), Vec3f0(1))
        model            = Mat4f0(I)
        modelinv         = const_lift(inv, model)

        color_map        = main.values => TextureBuffer
        color_norm       = nothing
        color            = nothing

        algorithm        = IndexedAbsorptionRGBA
        shader           = GLVisualizeShader("fragment_output.frag", "util.vert", "volume.vert", "volume.frag")
        prerender        = VolumePrerender()
        postrender       = () -> begin
            glDisable(GL_CULL_FACE)
        end
    end
end


module RigidBodySimTest

using Compat
using Compat.Test
using Compat.Random
using Compat.LinearAlgebra

using RigidBodySim

using RigidBodyDynamics

using DiffEqCallbacks: DiscreteCallback
using DiffEqBase: add_tstop!
using OrdinaryDiffEq: Rodas4P

using MechanismGeometries
using MeshCatMechanisms
using InteractBase: observe
using Blink: Window

if VERSION < v"0.7-"
    const seed! = srand
else
    import Random: seed!
end

function dynamics_allocations(dynamics::Dynamics, state::MechanismState) # introduce function barrier
    x = Vector(state)
    ẋ = similar(x)
    p = nothing
    t = 3.
    dynamics(ẋ, x, p, t)
    allocs = @allocated dynamics(ẋ, x, p, t)
end

@testset "Dynamics" begin
    seed!(134)
    mechanism = rand_tree_mechanism(Float64, [Revolute{Float64} for i = 1 : 30]...)
    dynamics = Dynamics(mechanism)
    state = MechanismState(mechanism)
    rand!(state)
    @test dynamics_allocations(dynamics, state) <= 80
end

@testset "compare to simulate" begin
    seed!(1)
    urdf = joinpath(@__DIR__, "urdf", "Acrobot.urdf")
    mechanism = parse_urdf(Float64, urdf)
    state = MechanismState(mechanism)
    rand!(state)
    x0 = Vector(state)
    final_time = 5.
    problem = ODEProblem(Dynamics(mechanism), state, (0., final_time))
    sol = solve(problem, Vern7(), abs_tol = 1e-10, dt = 0.05)
    Compat.copyto!(state, x0)
    ts, qs, vs = RigidBodyDynamics.simulate(state, final_time)
    @test [qs[end]; vs[end]] ≈ sol[end] atol = 1e-2
end

@testset "renormalization callback" begin
    mechanism = rand_tree_mechanism(Float64, QuaternionFloating{Float64})
    floatingjoint = first(joints(mechanism))
    state = MechanismState(mechanism)

    rand!(configuration(state))
    @test !RigidBodyDynamics.is_configuration_normalized(floatingjoint, configuration(state, floatingjoint))

    problem = ODEProblem(Dynamics(mechanism), state, (0., 1e-3))
    sol = solve(problem, Vern7(), dt = 1e-4, callback = configuration_renormalizer(state))

    Compat.copyto!(state, sol[end])
    @test RigidBodyDynamics.is_configuration_normalized(floatingjoint, configuration(state, floatingjoint))
end

@testset "RealtimeRateLimiter" begin
    du = [0; 0]
    u0 = [0.; 0.]
    dynamics = (u, p, t) -> eltype(u).(du)
    tmin = 10.1
    tmax = 12.2
    for max_rate in [2.0, 0.5]
        prob = ODEProblem(dynamics, u0, (tmin, tmax))
        rate_limiter = RealtimeRateLimiter(max_rate = max_rate)
        sol = solve(prob, Tsit5(); callback = rate_limiter)
        soltime = @elapsed solve(prob, Tsit5(); callback = rate_limiter)
        expected = (tmax - tmin) / max_rate
        @show soltime
        @show expected
        @test soltime ≈ expected atol = 0.5
    end
end

@testset "PeriodicController" begin
    urdf = joinpath(@__DIR__, "urdf", "Acrobot.urdf")
    mechanism = parse_urdf(Float64, urdf)
    state = MechanismState(mechanism)
    controltimes = Float64[]
    initialize = (c, u, t, integrator) -> empty!(controltimes)
    τ = similar(velocity(state))
    Δt = 0.25

    make_controller = function ()
        PeriodicController(τ, Δt, function (τ, t, state)
            push!(controltimes, t)
            τ[1] = sin(t)
            τ[2] = cos(t)
        end; initialize = initialize)
    end

    controller = make_controller()
    final_time = 25.3
    problem = ODEProblem(Dynamics(mechanism, controller), state, (0., final_time))

    # ensure that controller gets called at appropriate times:
    sol = solve(problem, Tsit5(), abs_tol = 1e-10, dt = 0.05)
    @test controltimes == collect(0. : Δt : final_time - rem(final_time, Δt))

    # ensure that we can solve the same problem again without errors and with a different integrator
    empty!(controltimes)
    sol = solve(problem, Vern7(), abs_tol = 1e-10, dt = 0.05)
    @test controltimes == collect(0. : Δt : final_time - rem(final_time, Δt))

    # issue #60
    empty!(controltimes)
    problem60 = ODEProblem(Dynamics(mechanism, (τ, t, state) -> controller(τ, t, state)), state, (0., final_time))
    @test_throws RigidBodySim.Control.PeriodicControlFailure solve(problem60, Vern7(), abs_tol = 1e-10, dt = 0.05)
    controller = controller = make_controller()
    problem60 = ODEProblem(Dynamics(mechanism, (τ, t, state) -> controller(τ, t, state)), state, (0., final_time))
    @test_throws RigidBodySim.Control.PeriodicControlFailure solve(problem60, Vern7(), abs_tol = 1e-10, dt = 0.05)
    problem60_fixed = ODEProblem(Dynamics(mechanism, (τ, t, state) -> controller(τ, t, state)), state, (0., final_time),
        callback = PeriodicCallback(controller))
    sol60 = solve(problem60_fixed, Vern7(), abs_tol = 1e-10, dt = 0.05)
    @test controltimes == collect(0. : Δt : final_time - rem(final_time, Δt))
    @test sol60.t == sol.t
    @test sol60.u == sol.u
end

@testset "visualizer callbacks" begin
    mechanism = rand_tree_mechanism(Float64, [Revolute{Float64} for i = 1 : 30]...)
    state = MechanismState(mechanism)
    gui = GUI(mechanism, MeshCatMechanisms.Skeleton(randomize_colors = true, inertias = true))
    vis = gui.visualizer
    controls = gui.controls
    open(vis, Window())
    open(controls, Window())
    open(gui)
    dt = 1e-4

    set_configuration!(vis, configuration(state))
    @test Vector(vis.state) == Vector(state)

    tfinal = 0.5
    problem = ODEProblem(Dynamics(mechanism), state, (0., tfinal), callback=CallbackSet(gui))

    # Simulate without interaction
    sol = solve(problem, RK4(), adaptive = false, dt = dt)
    @test sol.t[end] == tfinal
    @test Vector(vis.state) == sol.u[end]

    # Test simulation controls
    tfinal = 100.0
    problem = ODEProblem(Dynamics(mechanism), state, (0., tfinal), callback=CallbackSet(gui))

    # Simulate for 3 seconds (wall time) and then send a termination command
    @async (sleep(3.); observe(controls.terminate)[] += 1)
    sol = solve(problem, RK4(), adaptive = false, dt = dt)
    @test sol.t[end] > 2 * dt
    @test sol.t[end] < tfinal
    println("last(sol.t) after early termination 1: $(last(sol.t))")

    # Rinse and repeat with the same ODEProblem (make sure that we don't terminate straight away)
    observe(controls.terminate)[] += 1
    sleep(0.1)
    @async (sleep(3.); observe(controls.terminate)[] += 1)
    sol = solve(problem, RK4(), adaptive = false, dt = dt)
    @test sol.t[end] > 2 * dt
    @test sol.t[end] < tfinal
    println("last(sol.t) after early termination 2: $(last(sol.t))")

    # Pause and unpause a short simulation, make sure that the simulation takes longer than without pausing
    problem = ODEProblem(Dynamics(mechanism), state, (0., 1.), callback=CallbackSet(gui))
    tpause = 0.5
    pausecondition = Condition()
    havepaused = Ref(false)
    condition = (u, t, integrator) -> !havepaused[] && t >= tpause
    action = function (integrator)
        observe(controls.pause)[] += 1
        notify(pausecondition)
        havepaused[] = true
    end
    initialize = (c, t, u, integrator) -> (havepaused[] = false; add_tstop!(integrator, tpause))
    pauser = DiscreteCallback(condition, action, save_positions=(false, false), initialize = initialize)
    integrator = init(problem, RK4(), adaptive = false, dt = dt, callback=pauser)
    @async begin
        wait(pausecondition)
        yield() # wait for message to reach command handler callback
        integrator_time = integrator.t
        @test integrator_time < tpause + 0.1 # it takes a while for the pause message to reach command handler callback
        sleep(2.)
        @test integrator.t == integrator_time # make sure simulation remains paused
        havepaused[] = true
        observe(controls.pause)[] += 1 # unpause
    end
    solve!(integrator)
    @test havepaused[]

    # Simulate for 3 seconds wall time, then pause, and then terminate a second later to make sure terminating works while paused
    tfinal = 100.0
    problem = ODEProblem(Dynamics(mechanism), state, (0., tfinal), callback=CallbackSet(gui))
    @async (sleep(3.); observe(controls.pause)[] += 1)
    @async (sleep(4.); observe(controls.terminate)[] += 1)
    sol = solve(problem, RK4(), adaptive = false, dt = dt)
    @test sol.t[end] > 2 * dt
    @test sol.t[end] < tfinal
    println("last(sol.t) after early termination 3: $(last(sol.t))")
end

@testset "ODESolution animation" begin
    mechanism = rand_tree_mechanism(Float64, [Revolute{Float64} for i = 1 : 30]...)
    state = MechanismState(mechanism)
    vis = MechanismVisualizer(mechanism, MeshCatMechanisms.Skeleton(randomize_colors = true, inertias = true))
    open(vis, Window())
    wait(vis)

    final_time = 5.
    problem = ODEProblem(Dynamics(mechanism), state, (0., final_time))
    sol = solve(problem, Vern7(), abs_tol = 1e-10, dt = 0.05)

    # regular playback
    realtime_rate = 2.
    setanimation!(vis, sol, realtime_rate = realtime_rate, max_fps = 60.)
    sleep(final_time / realtime_rate)

    # broken:
    #=
    setanimation!(vis, sol, realtime_rate = 1000.)
    realtime_rate = 2.
    elapsed = @elapsed setanimation!(vis, sol, realtime_rate = realtime_rate, max_fps = 60.)
    @test elapsed ≈ final_time / realtime_rate atol = 0.1

    # premature termination
    termination_time = 1.5
    @async (sleep(termination_time); send_terminate_message())
    elapsed = @elapsed animate(vis, state, sol, realtime_rate = realtime_rate)
    @test elapsed ≈ termination_time atol = 0.1

    # pause and unpause
    pause_time = 1.0
    unpause_time = 3.5
    @async (sleep(pause_time); send_pause_message())
    @async (sleep(unpause_time); send_pause_message())
    elapsed = @elapsed animate(vis, state, sol, realtime_rate = realtime_rate)
    @show elapsed
    @test elapsed ≈ final_time / realtime_rate + (unpause_time - pause_time) atol = 0.2 # higher atol because of pause poll int

    # pause and terminate
    @async (sleep(pause_time); send_pause_message())
    @async (sleep(termination_time); send_terminate_message())
    elapsed = @elapsed animate(vis, state, sol, realtime_rate = realtime_rate)
    @show elapsed
    @test elapsed ≈ termination_time atol = 0.2 # higher atol because of pause poll int
    =#
end

@testset "Stiff integrator" begin
    urdf = joinpath(@__DIR__, "urdf", "Acrobot.urdf")
    mechanism = parse_urdf(Float64, urdf)
    state = MechanismState(mechanism)
    seed!(1)
    rand!(state)
    x0 = Vector(state)
    final_time = 1.
    problem = ODEProblem(Dynamics(mechanism), state, (0., final_time))

    sol_nonstiff = solve(problem, Vern7(), abs_tol = 1e-10, dt = 0.05)
    sol_stiff = solve(problem, Rodas4P(), abs_tol = 1e-10, dt = 0.05)
    @test last(sol_nonstiff) ≈ last(sol_stiff) atol = 1e-2
end

# notebooks
include("test_notebooks.jl")

end


@blasfun clblasSrot(N::Csize_t,
                         X::cl.CL_mem, offx::Csize_t, incx::Cint,
                         Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                         C::cl.CL_float, S::cl.CL_float,
                         n_queues::cl.CL_uint,
                         queues::Ptr{cl.CL_command_queue},
                         n_events_in_wait_list::cl.CL_uint,
                         event_wait_list::Ptr{cl.CL_event},
                         events::Ptr{cl.CL_event})

@blasfun2 clblasSrot(N::Csize_t,
                          X::cl.CL_mem, offx::Csize_t, incx::Cint,
                          Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                          C::cl.CL_float, S::cl.CL_float)

@blasfun clblasDrot(N::Csize_t,
                         X::cl.CL_mem, offx::Csize_t, incx::Cint,
                         Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                         C::cl.CL_double, S::cl.CL_double,
                         n_queues::cl.CL_uint,
                         queues::Ptr{cl.CL_command_queue},
                         n_events_in_wait_list::cl.CL_uint,
                         event_wait_list::Ptr{cl.CL_event},
                         events::Ptr{cl.CL_event})

@blasfun2 clblasDrot(N::Csize_t,
                          X::cl.CL_mem, offx::Csize_t, incx::Cint,
                          Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                          C::cl.CL_double, S::cl.CL_double)

@blasfun clblasCsrot(N::Csize_t,
                          X::cl.CL_mem, offx::Csize_t, incx::Cint,
                          Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                          C::cl.CL_float, S::cl.CL_float,
                          n_queues::cl.CL_uint,
                          queues::Ptr{cl.CL_command_queue},
                          n_events_in_wait_list::cl.CL_uint,
                          event_wait_list::Ptr{cl.CL_event},
                          events::Ptr{cl.CL_event})

@blasfun2 clblasCrot(N::Csize_t,
                          X::cl.CL_mem, offx::Csize_t, incx::Cint,
                          Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                          C::cl.CL_float, S::cl.CL_float)

@blasfun clblasZdrot(N::Csize_t,
                          X::cl.CL_mem, offx::Csize_t, incx::Cint,
                          Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                          C::cl.CL_double, S::cl.CL_double,
                          n_queues::cl.CL_uint,
                          queues::Ptr{cl.CL_command_queue},
                          n_events_in_wait_list::cl.CL_uint,
                          event_wait_list::Ptr{cl.CL_event},
                          events::Ptr{cl.CL_event})

@blasfun2 clblasZdrot(N::Csize_t,
                           X::cl.CL_mem, offx::Csize_t, incx::Cint,
                           Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                           C::cl.CL_double, S::cl.CL_double)


module ReinforcementLearningEnvironmentGym
using ReinforcementLearningBase
import ReinforcementLearningBase:interact!, reset!, getstate, plotenv, actionspace
using PyCall
const gym = PyNULL()

function __init__()
    copy!(gym, pyimport_conda("gym", ""))
    pyimport("pybullet_envs")
end

function gymspace2jlspace(s::PyObject)
    spacetype = s[:__class__][:__name__]
    if     spacetype == "Box"           BoxSpace(s[:low], s[:high])
    elseif spacetype == "Discrete"      DiscreteSpace(s[:n], 1)
    elseif spacetype == "MultiBinary"   MultiBinarySpace(s[:n])
    elseif spacetype == "MultiDiscrete" MultiDiscreteSpace(s[:nvec], 1)
    elseif spacetype == "Tuple"         map(gymspace2jlspace, s[:spaces])
    elseif spacetype == "Dict"          Dict(map((k, v) -> (k, gymspace2jlspace(v)), s[:spaces]))
    else error("Don't know how to convert [$(spacetype)]")
    end
end

struct GymEnv{Ta<:AbstractSpace, To<:AbstractSpace} <: AbstractEnv
    pyobj::PyObject
    observationspace::To
    actionspace::Ta
    state::PyObject
end

function GymEnv(name::String)
    pyenv = gym[:make](name)
    obsspace = gymspace2jlspace(pyenv[:observation_space])
    actspace = gymspace2jlspace(pyenv[:action_space])
    state = PyNULL()
    env = GymEnv(pyenv, obsspace, actspace, state)
    reset!(env) # state needs to be set to call defaultbuffer in RL
    env
end

function interact!(env::GymEnv, action)
    pycall!(env.state, env.pyobj[:step], PyVector, action)
    (observation=env.state[1], reward=env.state[2], isdone=env.state[3])
end

function interact!(env::GymEnv{DiscreteSpace}, action::Int)
    pycall!(env.state, env.pyobj[:step], PyVector, action - 1)
    (observation=env.state[1], reward=env.state[2], isdone=env.state[3])
end

function interact!(env::GymEnv{MultiDiscreteSpace}, action::AbstractArray{Int})
    pycall!(env.state, env.pyobj[:step], PyVector, action .- 1)
    (observation=env.state[1], reward=env.state[2], isdone=env.state[3])
end

"Not very useful, kept for compat"
function getstate(env::GymEnv)
    if pyisinstance(env.state, PyCall.@pyglobalobj :PyTuple_Type)
        (observation=env.state[1], isdone=env.state[3])
    else
        # env has just been reseted
        (observation=Float64.(env.state), isdone=false)
    end
end

reset!(env::GymEnv) = (observation=Float64.(pycall!(env.state, env.pyobj[:reset], PyArray)),)
plotenv(env::GymEnv) = env.pyobj[:render]()
actionspace(env::GymEnv) = env.actionspace

"""
    listallenvs(pattern = r"")

List all registered gym environment names. The optional argument `pattern`
allows to list all environment  that contain the `pattern` in their name.
"""
function listallenvs(pattern = r"")
    envs = sort(py"[spec.id for spec in $gym.envs.registry.all()]")
    if pattern != ""
        envs[findall(x -> occursin(pattern, x), envs)]
    else
        envs
    end
end

export GymEnv, listallenvs
end # module


if !(:ModelClockSystem in names(Main))
    include(normpath(joinpath(dirname(@__FILE__), "..", "..", "share", "ModelClockSystem.jl")))
end

if !(:ConfigCheck in names(Main))
    include(normpath(joinpath(dirname(@__FILE__), "..", "..", "share", "ConfigCheck.jl")))
end

if !(:CyclicData in names(Main))
    include(normpath(joinpath(dirname(@__FILE__), "..", "..", "share", "CyclicData.jl")))
end

if !(:LogSystem in names(Main))
    include(normpath(joinpath(dirname(@__FILE__), "..", "..", "share", "LogSystem.jl")))
end

if ! ( :DataManager in names(Main) )
    include(joinpath(@__DIR__, "..", "..", "share", "DataManager.jl"))
end

macro hinclude(path)
    return :(include(normpath(joinpath(@__DIR__, $path))))
end

module EMOM

    using LinearAlgebra
    using MPI
    using Dates
    using Printf
    using Formatting
    using SharedArrays
    using Distributed
    using SparseArrays
    using NCDatasets
    using JLD2
    using DataStructures

    using ..ModelClockSystem
    using ..ConfigCheck
    using ..CyclicData
    using ..LogSystem
    using ..DataManager

    macro hinclude(path)
        return :(include(normpath(joinpath(@__DIR__, $path))))
    end

    @hinclude("../../share/constants.jl")
    @hinclude("../../share/ocean_state_function.jl")

    # classes
    @hinclude("../../share/GridFile.jl")
    @hinclude("../../share/PolelikeCoordinate.jl")
    @hinclude("../../share/BasicMatrixOperators.jl")
    @hinclude("../../share/AdvancedMatrixOperators.jl")

    @hinclude("entry_list.jl")

    @hinclude("VerticalDiffusion.jl")

    @hinclude("Topography.jl")
    @hinclude("Workspace.jl")
    @hinclude("Env.jl")
    @hinclude("TempField.jl")
    @hinclude("Field.jl")
    @hinclude("Core.jl")

    @hinclude("ModelBlock.jl")

    @hinclude("Leonard1979.jl")
    @hinclude("setupForcing.jl")
    @hinclude("stepAdvection.jl")
    @hinclude("stepColumn.jl")
    @hinclude("checkBudget.jl")

    @hinclude("var_list.jl")
    @hinclude("var_desc.jl")

    @hinclude("snapshot_funcs.jl")

    @hinclude("updateDatastream.jl")
    @hinclude("updateBuoyancy.jl")

end


@testset "Mutation vertices" begin
    using Functors: fmap

    @testset "OutputsVertex with $label" for (label, cfun) in (
        (deepcopy, deepcopy),
        ("fmap", v -> (vc = fmap(deepcopy, v); NaiveNASlib.init!(vc, base(vc)); return vc)),
    )
        iv = OutputsVertex(InputVertex(1))

        @test inputs(iv) == []
        @test outputs(iv) == []

        @test issame(iv, cfun(iv))

        cv = OutputsVertex(CompVertex(x -> 2x, iv))
        NaiveNASlib.init!(cv, base(cv))

        @test inputs(cv) == [iv]
        @test outputs(iv) == [base(cv)]

        cvc = cfun(cv)

        ivc = inputs(cvc)[]
        @test issame(iv, ivc)
    end

    @testset "InputSizeVertex with $label" for (label, cfun) in (
        (deepcopy, deepcopy),
        ("fmap", v -> (vc = fmap(deepcopy, v); if vc isa OutputsVertex NaiveNASlib.init!(vc, base(vc)) end; return vc)),
    )
        iv = InputSizeVertex(InputVertex(1), 3)

        @test nout(iv) == 3
        @test nin(iv) == []
        @test inputs(iv) == []
        @test outputs(iv) == []

        @test issame(iv, cfun(iv))

        cv =  OutputsVertex(CompVertex(x -> 2x, iv))
        NaiveNASlib.init!(cv, base(cv))

        @test inputs(cv) == [iv]
        @test outputs(iv) == [base(cv)]

        cvc = cfun(cv)

        ivc = inputs(cvc)[]
        @test issame(iv, ivc)
    end

    @testset "Fixed size trait" begin
        using NaiveNASlib: vertex
        struct FixedSizeDummy end
        NaiveNASlib.nout(::FixedSizeDummy) = 5
        NaiveNASlib.nin(f::FixedSizeDummy) = [nout(f)]

        @testset "Fixed Size $stype" for stype in (SizeStack, SizeInvariant)
            iv = inputvertex("in", nout(FixedSizeDummy()) - 1)
            v1 = vertex(FixedSizeTrait(stype()), FixedSizeDummy(), iv)

            @test nout(v1) == nout(FixedSizeDummy())
            @test nin(v1) == nin(FixedSizeDummy())
        end
    end

    @testset "Pretty printing" begin
        import NaiveNASlib: show_less

        @testset "OutputsVertex" begin
            iv = OutputsVertex(InputVertex(1))
            cv = OutputsVertex(CompVertex(identity, iv))
            NaiveNASlib.init!(cv, base(cv))
            sv = OutputsVertex(CompVertex(hcat, iv, cv))
            NaiveNASlib.init!(sv, base(sv))

            @test showstr(show, iv) == "InputVertex(1, outputs=[CompVertex(identity), CompVertex(hcat)])"
            @test showstr(show, cv) == "CompVertex(identity, inputs=[InputVertex(1)], outputs=[CompVertex(hcat)])"
            @test showstr(show, sv) == "CompVertex(hcat, inputs=[InputVertex(1), CompVertex(identity)], outputs=[])"
        end

        @testset "InputSizeVertex" begin
            iv = InputSizeVertex("iv", 3)
            @test showstr(show_less, iv) == "iv"
        end

        @testset "NamedTrait" begin
            using NaiveNASlib: MutationVertex
            v1 = MutationVertex(CompVertex(identity, InputSizeVertex("input1", 3)), NamedTrait("mv", SizeInvariant()))

            @test showstr(show_less, v1) == "mv"

            v2 = MutationVertex(CompVertex(+, v1, InputSizeVertex("input2", 3)), NamedTrait("sv",SizeInvariant()))

            @test showstr(show, v2) == "MutationVertex(CompVertex(+, inputs=[mv, input2], outputs=[]), NamedTrait(\"sv\", SizeInvariant()))"
        end
    end
end


# simple adjacency list, which represents both nodes and timestamps as integers.

"""
    IntAdjacencyList(nv, nt)

Construct a graph represented by an adjacency list with `nv` nodes and `nt` timestamps, where both nodes and timestamps are represented by integers.

# Example

```jldoctest
julia> using EvolvingGraphs

julia> g = IntAdjacencyList(4,3)
Directed IntAdjacencyList (4 nodes, 0 static edges, 3 timestamps)

julia> add_edge!(g, 1, 2, 1)
Directed IntAdjacencyList (4 nodes, 1 static edges, 3 timestamps)

julia> add_edge!(g, 2, 3, 2)
Directed IntAdjacencyList (4 nodes, 2 static edges, 3 timestamps)

julia> num_edges(g)
2
```
"""
mutable struct IntAdjacencyList <: AbstractEvolvingGraph{Int, Tuple{Int,Int}, Int}
    is_directed::Bool
    nodes::UnitRange{Int}
    timestamps::Vector{Int}
    nnodes::Int      # number of nodes
    nedges::Int      # number of static edges
    forward_adjlist::Vector{Vector{Int}}
    backward_adjlist::Vector{Vector{Int}}
end
function IntAdjacencyList(nv::Int, nt::Int; is_directed::Bool = true)
    ts = Array{Int}(nv*nt)
    f_adj = Vector{Int}[]
    b_adj = Vector{Int}[]
    for i = 1:nv*nt
        push!(f_adj, Int[])
        push!(b_adj, Int[])
    end
    for t = 1:nt
        for v = 1:nv
            ts[v + nv*(t-1)] = t
        end
    end
    IntAdjacencyList(is_directed, 1:nv*nt, ts, nv, 0, f_adj, b_adj)
end

is_directed(g::IntAdjacencyList) = g.is_directed

"""
    evolving_graph_to_adj(g)

Convert an evolving graph `g` to an adjacency list.

# Example

```jldoctest
julia> using EvolvingGraphs

julia> g = EvolvingGraph()
Directed EvolvingGraph 0 nodes, 0 static edges, 0 timestamps

julia> add_bunch_of_edges!(g, [(1,2,2001),(2,3,2002), (2,4,2002), (3,1,2003)])
Directed EvolvingGraph 4 nodes, 4 static edges, 3 timestamps

julia> evolving_graph_to_adj(g)
Directed IntAdjacencyList (4 nodes, 4 static edges, 3 timestamps)
```
"""
function evolving_graph_to_adj(g::AbstractEvolvingGraph)
    g1 = IntAdjacencyList(num_nodes(g), num_timestamps(g), is_directed = is_directed(g))

    tmap = Dict(t => i for (i, t) in enumerate(unique_timestamps(g)))
    for e in edges(g)
        v1 = node_index(source(e))
        v2 = node_index(target(e))
        add_edge!(g1, v1, v2, tmap[e.timestamp])
    end
    g1
end

nodes(g::IntAdjacencyList) = collect(1:g.nnodes)
num_nodes(g::IntAdjacencyList) = g.nnodes
num_edges(g::IntAdjacencyList) = g.nedges
timestamps(g::IntAdjacencyList) = g.timestamps
unique_timestamps(g::IntAdjacencyList) = unique(g.timestamps)
num_timestamps(g::IntAdjacencyList) = length(unique_timestamps(g))
node_index(g::IntAdjacencyList, node::Int) = node

deepcopy(g::IntAdjacencyList) = IntAdjacencyList(is_directed(g),
                                             deepcopy(g.nodes),
                                             deepcopy(g.timestamps),
                                             g.nnodes,
                                             g.nedges,
                                             deepcopy(g.forward_adjlist),
                                             deepcopy(g.backward_adjlist))

function forward_neighbors(g::IntAdjacencyList, v::Int, t::Int)
    ns = g.nnodes
    n = v + ns*(t-1)
    nn = Tuple{Int, Int}[]
    for node in g.forward_adjlist[n]
        t = g.timestamps[node]
        v1 = node - ns*(t-1)
        push!(nn, (v1, t))
    end
    nn
end
forward_neighbors(g::IntAdjacencyList, vt::Tuple{Int,Int}) = forward_neighbors(g, vt[1], vt[2])

function backward_neighbors(g::IntAdjacencyList, v::Int, t::Int)
    ns = g.nnodes
    n = v + ns*(t-1)
    nn = Tuple{Int, Int}[]
    for node in g.backward_adjlist[n]
        t = g.timestamps[node]
        v1 = node - ns*(t-1)
        push!(nn, (v1, t))
    end
    nn
end
backward_neighbors(g::IntAdjacencyList, vt::Tuple{Int,Int}) = backward_neighbors(g, vt[1], vt[2])

function add_edge!(g::IntAdjacencyList, v1::Int, v2::Int, t::Int)
    ns = g.nnodes
    n1 = v1 + ns*(t-1)
    n2 = v2 + ns*(t-1)
    if _insertnode!(g.forward_adjlist[n1], n2)
        g.nedges += 1

        for (v, n) in ((v1, n1), (v2, n2))
            for i in 1:t-1
                len = length(g.forward_adjlist[v + ns*(i-1)]) +
                length(g.backward_adjlist[v + ns*(i-1)])
                if len > 0
                    _insertnode!(g.forward_adjlist[v + ns*(i-1)], n)
                    _insertnode!(g.backward_adjlist[n], (v+ns*(i-1)))
                end
            end
            for i in t+1:num_timestamps(g)
                #   println("node: $(v + ns*(i-1))")
                len = length(g.forward_adjlist[v + ns*(i-1)]) +
                length(g.backward_adjlist[v + ns*(i-1)])
                if len > 0
                    _insertnode!(g.forward_adjlist[n], (v+ns*(i-1)))
                    _insertnode!(g.backward_adjlist[v+ns*(i-1)], n)
                end
            end
        end
    end
    _insertnode!(g.backward_adjlist[n2], n1)
    g
end

# The function _insertnode! is from LightGraphs.jl
# The LightGraphs.jl package is licensed under the Simplified "2-clause" BSD License:

#Copyright (c) 2015: Seth Bromberger and other contributors.

#Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

#Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
#Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# returns true if insert succeeded, false if it was a duplicate
_insertnode!(v::Vector{Int}, x::Int) = isempty(splice!(v, searchsorted(v,x), x))

# has_edge on active nodes
function _has_edge(g::IntAdjacencyList, n1::Int, n2::Int)
    nn = length(g.nodes)
    if n1 > nn || n2 > nn
        return false
    end
    return n2 in g.forward_adjlist[n1]
end

function has_edge(g::IntAdjacencyList, v1::Int, v2::Int, t::Int)
    ns = g.nnodes
    n1 = v1 + ns*(t-1)
    n2 = v2 + ns*(t-1)
    return _has_edge(g, n1, n2)
end

function adjacency_matrix(g::IntAdjacencyList, t::Int)
    ns = nodes(g)
    num_ns = num_nodes(g)
    A = zeros(Float64, num_ns, num_ns)
    for n in ns
        fns = forward_neighbors(g, n, t)
        for (fn, ft) in fns
            if ft == t
                 A[n, fn] = 1.0
             end
         end
     end
     return A
end

function sparse_adjacency_matrix(g::IntAdjacencyList, t::Int)
    is = Int[]
    js = Int[]
    ns = nodes(g)
    num_ns = num_nodes(g)
    for n in ns
        fns = forward_neighbors(g, n, t)
        for (fn, ft) in fns
            if ft == t
                push!(is, n)
                push!(js, fn)
            end
        end
    end
    return sparse(is, js, ones(is, Float64), num_ns, num_ns)
end

function block_adjacency_matrix(g::IntAdjacencyList)
    ns = nodes(g)
    num_ns = num_nodes(g)
    uts = unique_timestamps(g)
    num_ts = num_timestamps(g)
    is = Int[]
    js = Int[]
    vs = Float64[]
    for n in ns
        for t in uts
            fns = forward_neighbors(g, n, t)
            for (fn, ft) in fns
                push!(is, n + num_ns*(t-1))
                push!(js, fn + num_ns*(ft-1))
                abs(ft - t) == 0 ? push!(vs, 1.0) : push!(vs, 1./abs(ft-t))
            end
        end
    end
    return sparse(is, js, vs, num_ns*num_ts, num_ns*num_ts)
end


## Copyright (c) 2013 Miguel Bazdresch
##
## This file is distributed under the 2-clause BSD License.

# llplot() is our workhorse plotting function
function llplot()
    global gnuplot_state
    global gaston_config

    # select current figure
    c = findfigure(gnuplot_state.current)
    if c == 0
        println("No current figure")
        return
    end
    fig = gnuplot_state.figs[c]
    config = fig.conf

    # if figure has no data, stop here
    if isempty(fig.isempty)
        return
    end

    # Reset gnuplot settable options.
    gnuplot_send("\nreset\n")

    # Build terminal setup string and send it to gnuplot
    gnuplot_send(termstring())

    # Datafile filename. This is where we store the coordinates to plot.
    # This file is then read by gnuplot to do the actual plotting. One file
    # per figure handle is used; this avoids polutting /tmp with too many files.
    filename = joinpath(tempdir(),
                        "gaston-$(gaston_config.tmpprefix)-$(fig.handle)")
    f = open(filename,"w")

    # Send appropriate coordinates and data to gnuplot, depending on
    # whether we are doing 2-d, 3-d or image plots.

    # 2-d plot: Z is empty or plostyle is {,rgb}image
    if isempty(fig.curves[1].Z) ||
        fig.curves[1].conf.plotstyle == "image" ||
        fig.curves[1].conf.plotstyle == "rgbimage"
        # create data file
        for i in fig.curves
            ps = i.conf.plotstyle
            if ps == "errorbars" || ps == "errorlines"
                if isempty(i.E.yhigh)
                    # ydelta (single error coordinate)
                    writedlm(f,[i.x i.y i.E.ylow],' ')
                else
                    # ylow, yhigh (double error coordinate)
                    writedlm(f,[i.x i.y i.E.ylow i.E.yhigh],' ')
                end
            elseif ps == "financebars"
                # data is written to tmparr, which is then written to disk
                tmparr = zeros(length(i.x),5)
                # output matrix
                for col = 1:length(i.x)
                    tmparr[col,1] = i.x[col]
                    tmparr[col,2] = i.F.open[col]
                    tmparr[col,3] = i.F.low[col]
                    tmparr[col,4] = i.F.high[col]
                    tmparr[col,5] = i.F.close[col]
                end
                writedlm(f,tmparr,' ')
            elseif ps == "image"
                # data is written to tmparr, which is then written to disk
                tmparr = zeros(length(i.x)*length(i.y),3)
                tmparr_row_index = 1  # index into tmparr row
                # output matrix
                for row = 1:length(i.y)
                    x = length(i.x)
                    for col = 1:length(i.x)
                        tmparr[tmparr_row_index,1] = i.x[col]
                        tmparr[tmparr_row_index,2] = i.y[row]
                        tmparr[tmparr_row_index,3] = i.Z[row,col]
                        tmparr_row_index = tmparr_row_index+1
                        x = x-1
                    end
                end
                writedlm(f,tmparr,' ')
            elseif ps == "rgbimage"
                # data is written to tmparr, which is then written to disk
                tmparr = zeros(length(i.x)*length(i.y), 5)
                tmparr_row_index = 1
                # output matrix
                for col = 1:length(i.x)
                    y = length(i.y)
                    for row = 1:length(i.y)
                        tmparr[tmparr_row_index,1] = i.x[col]
                        tmparr[tmparr_row_index,2] = i.y[row]
                        tmparr[tmparr_row_index,3] = i.Z[y,col,1]
                        tmparr[tmparr_row_index,4] = i.Z[y,col,2]
                        tmparr[tmparr_row_index,5] = i.Z[y,col,3]
                        tmparr_row_index = tmparr_row_index+1
                        y = y-1
                    end
                end
                writedlm(f,tmparr,' ')
            else
                writedlm(f,[i.x i.y],' ')
            end
            write(f,"\n\n")
        end
        close(f)
        # send figure configuration to gnuplot
        gnuplot_send_fig_config(config)
        # Send user command to gnuplot
        !isempty(fig.gpcom) && gnuplot_send(fig.gpcom)
        # send plot command to gnuplot
        gnuplot_send(linestr(fig.curves, "plot", filename))

    # 3-d plot: Z is not empty and plotstyle is not {,rgb}image
    elseif !isempty(fig.curves[1].Z) &&
            fig.curves[1].conf.plotstyle != "image" &&
            fig.curves[1].conf.plotstyle != "rgbimage"
        # create data file
        for i in fig.curves
            # data is written to tmparr, which is then written to disk
            tmparr = zeros(1, 3)
            tmparr_row_index = 1
            for row in 1:length(i.x)
                for col in 1:length(i.y)
                    tmparr[1,1] = i.x[row]
                    tmparr[1,2] = i.y[col]
                    tmparr[1,3] = i.Z[row,col]
                    writedlm(f,tmparr,' ')
                end
                write(f,"\n")
            end
            write(f,"\n\n")
        end
        close(f)
        # send figure configuration to gnuplot
        gnuplot_send_fig_config(config)
        # Send user command to gnuplot
        !isempty(fig.gpcom) && gnuplot_send(fig.gpcom)
        # send command to gnuplot
        gnuplot_send(linestr(fig.curves, "splot",filename))
    end

    # Make sure gnuplot is done; if terminal is text, read data
    # reset error handling
    err = ""
    gnuplot_state.gp_lasterror = err
    gnuplot_state.gp_error = false

    attempt_stderr = 20
    attempt_stdout = 100
    sleep_interval = 0.05
    sleep_increment = 1.2

    gnuplot_send("printerr \"GastonDone\"\n")
    sleep(sleep_interval)

    # wait for stderr channel to be ready
    si = sleep_interval
    count = 0
    while !isready(ChanStdErr)
        sleep(si)
        si = sleep_increment * si
        count = count + 1
        count > attempt_stderr &&
        error("Gnuplot is taking too long to respond.")
    end

    # read all data in channel
    si = sleep_interval
    while isready(ChanStdErr)
        err = err * take!(ChanStdErr)
        sleep(si)
        si = sleep_increment * si
    end

    # check for errors while plotting
    if err != "GastonDone\n"
        gnuplot_state.gp_lasterror = err
        gnuplot_state.gp_error = true
        @warn("Gnuplot returned an error message:\n  $err)")
    end

    # if there was no error and text terminal, read all data from stdout
    if err == "GastonDone\n"
        if (gaston_config.terminal ∈ term_text)
            # wait for stdout to be ready
            si = sleep_interval
            count = 0
            while !isready(ChanStdOut)
                sleep(si)
                si = sleep_increment * si
                count = count + 1
                count > attempt_stdout &&
                error("Gnuplot is taking too long to respond.")
            end
            svgdata = ""
            si = sleep_interval
            while isready(ChanStdOut)
                svgdata = svgdata * take!(ChanStdOut)
                sleep(si)
                si = sleep_increment * si
            end
            fig.svg = svgdata
        end
    end

    return nothing

end