###### 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


module ModelHousehold

#
# This module provides a slightly abstracted view of a single Scottish household. The fields
# are predominantly based on the FRS, but could in principle come from other sources
# such as SHS, LCF and so on, or perhaps from a web form.
#
using Base: String
using Dates
import Base.show

using ScottishTaxBenefitModel
using .Definitions
using .Utils:
    has_non_z,
    todays_date,
    to_md_table

using .EquivalenceScales:
    EQ_P_Type,
    EQ_Person,
    EQScales,
    get_equivalence_scales,
    eq_dependent_child,
    eq_head,
    eq_other_adult,
    eq_spouse_of_head

import .EquivalenceScales.get_age
import .EquivalenceScales.eq_rel_to_hoh

import Base.isequal
import Base.hash

using .Uprating: uprate, UPRATE_MAPPINGS

export
    Household,
    BenefitUnit,
    BenefitUnits,
    BUAllocation,
    Person,
    People_Dict,

    adult_pids,
    age_then,
    between_ages,
    child_pids,
    default_bu_allocation,
    empl_status_in,
    ge_age,
    get_benefit_units,
    get_head,
    get_spouse,
    has_carer_member,
    has_children,
    has_disabled_member,
    has_income,
    interview_date,
    is_child,
    is_head,
    is_lone_parent,
    is_severe_disability,
    is_single,
    is_spouse,
    isless,
    le_age,
    make_benefit_unit,
    make_eq_scales!,
    num_adults,
    num_adults,
    num_carers,
    num_children,
    num_people,
    oldest_person,
    on_mt_benefits,
    pers_is_carer,
    pers_is_disabled,
    printpids,
    search,
    set_wage!,
    to_string,
    uprate!

mutable struct Person{RT<:Real}
    hid::BigInt # == sernum
    pid::BigInt # == unique id (year * 100000)+
    pno:: Int # person number in household
    is_hrp :: Bool
    default_benefit_unit:: Int
    is_standard_child :: Bool
    age:: Int

    sex::Sex
    ethnic_group::Ethnic_Group
    marital_status::Marital_Status
    highest_qualification::Qualification_Type

    sic::SIC_2007
    occupational_classification::Standard_Occupational_Classification
    public_or_private :: Employment_Sector
    principal_employment_type :: Employment_Type

    socio_economic_grouping::Socio_Economic_Group
    age_completed_full_time_education:: Int
    years_in_full_time_work:: Int
    employment_status::ILO_Employment
    actual_hours_worked::RT
    usual_hours_worked::RT

    age_started_first_job :: Int

    income::Incomes_Dict{RT}
    benefit_ratios :: Incomes_Dict{RT}

    jsa_type :: JSAType # FIXME change this name
    esa_type :: JSAType
    dla_self_care_type :: LowMiddleHigh
    dla_mobility_type :: LowMiddleHigh
    attendance_allowance_type :: LowMiddleHigh
    pip_daily_living_type :: PIPType
    pip_mobility_type ::  PIPType

    bereavement_type :: BereavementType
    had_children_when_bereaved :: Bool

    assets::Asset_Dict{RT}
    over_20_k_saving :: Bool
    pay_includes ::Included_In_Pay_Dict{Bool}

    # contracted_out_of_serps::Bool

    registered_blind::Bool
    registered_partially_sighted::Bool
    registered_deaf::Bool

    disabilities::Disability_Dict{Bool}

    health_status::Health_Status

    has_long_standing_illness :: Bool
    adls_are_reduced :: ADLS_Inhibited
    how_long_adls_reduced :: Illness_Length

    relationships::Relationship_Dict
    relationship_to_hoh :: Relationship
    is_informal_carer::Bool
    receives_informal_care_from_non_householder::Bool
    hours_of_care_received::RT
    hours_of_care_given::RT
    #
    # Childcare fields; assigned to children
    #
    hours_of_childcare :: RT
    cost_of_childcare :: RT
    childcare_type :: Child_Care_Type
    employer_provides_child_care :: Bool

    company_car_fuel_type :: Fuel_Type
    company_car_value :: RT
    company_car_contribution :: RT
    fuel_supplied :: RT
    onerand :: String
end

People_Dict = Dict{BigInt,Person{T}} where T<:Real
Pid_Array = Vector{BigInt}

mutable struct Household{RT<:Real}
    sequence:: Int # position in current generated dataset
    hid::BigInt
    data_year :: Int
    interview_year:: Int
    interview_month:: Int
    quarter:: Int
    tenure::Tenure_Type
    region::Standard_Region
    ct_band::CT_Band
    council_tax::RT
    water_and_sewerage ::RT
    mortgage_payment::RT
    mortgage_interest::RT
    years_outstanding_on_mortgage:: Int
    mortgage_outstanding::RT
    year_house_bought:: Int
    gross_rent::RT # rentg Gross rent including Housing Benefit  or rent Net amount of last rent payment
    rent_includes_water_and_sewerage::Bool
    other_housing_charges::RT
    gross_housing_costs::RT
    total_income::RT
    total_wealth::RT
    house_value::RT
    weight::RT
    council :: Symbol
    nhs_board :: Symbol
    bedrooms :: Int
    head_of_household :: BigInt
    people::People_Dict{RT}
    onerand :: String
    equivalence_scales :: EQScales{RT}
end

function to_string( hh :: Household ) :: String
    s = to_md_table( hh, exclude=[:people,:onerand])
    for pid in sort(collect(keys(hh.people)))
        s *= to_string( hh.people[pid] )
    end
    return s
end

function to_string( pers :: Person ) :: String
    s = """
    ### Person $(pers.pno)
    """
    s *= to_md_table( pers, depth=1,exclude=[:onerand])
    return s
end

"""
for equivalence scale implicit interface
"""
function get_age( p::Person ) :: Int
    return p.age;
end

"""
for the eq scale implicit interface
"""
function eq_rel_to_hoh( p :: Person ) :: EQ_P_Type
    if p.relationship_to_hoh == This_Person
        return eq_head
    elseif p.is_standard_child
        return eq_dependent_child
    elseif p.relationship_to_hoh in [Spouse,Cohabitee]
        return eq_spouse_of_head
    else
        return eq_other_adult
    end
end

function make_eq_scales!( hh :: Household{T} ) where T
    pl = collect(values(hh.people))
    # sz = size( pl )[1]
    # print( "num people $sz" )
    hh.equivalence_scales = get_equivalence_scales( T, pl )
end

function interview_date( hh :: Household ) :: Date
    return Date( hh.interview_year, hh.interview_month, 15 )
end

function uprate!( pid :: BigInt, year::Integer, quarter::Integer, person :: Person )
    for (t,inc) in person.income
            person.income[t] = uprate( inc, year, quarter, UPRATE_MAPPINGS[t])
    end
    for (a,ass) in person.assets
            person.assets[a] = uprate( ass, year, quarter, UPRATE_MAPPINGS[a])
    end
    person.cost_of_childcare = uprate( person.cost_of_childcare, year, quarter, upr_earnings )
end

function uprate!( hh :: Household )

    hh.water_and_sewerage  = uprate( hh.water_and_sewerage , hh.interview_year, hh.quarter, upr_housing_rents )
    hh.mortgage_payment = uprate( hh.mortgage_payment, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_interest = uprate( hh.mortgage_interest, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_outstanding = uprate( hh.mortgage_outstanding, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.gross_rent = uprate( hh.gross_rent, hh.interview_year, hh.quarter, upr_housing_rents )
    hh.other_housing_charges = uprate( hh.other_housing_charges, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.gross_housing_costs = uprate( hh.gross_housing_costs, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_income = uprate( hh.total_income, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_wealth = uprate( hh.total_wealth, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.house_value = uprate( hh.house_value, hh.interview_year, hh.quarter, upr_housing_oo )
    for (pid,person) in hh.people
        uprate!( pid, hh.interview_year, hh.quarter, person )
    end

end

function oldest_person( people :: People_Dict ) :: NamedTuple
    oldest = ( age=-999, pid=BigInt(0))
    for person in people
        if person.age > oldest.age
            oldest.age = person.age
            oldest.pid = person.pid
        end
    end
    oldest
end

PeopleArray = Vector{Person}

struct BenefitUnit{T<:Real}
    people :: People_Dict{T}
    head :: BigInt
    spouse :: BigInt
    adults :: Pid_Array
    children :: Pid_Array
end

BenefitUnits = Vector{BenefitUnit}
BUAllocation = Vector{PeopleArray}

function num_children( bu :: BenefitUnit ) :: Integer
    size( bu.children )[1]
end

function num_children( hh :: Household ) :: Integer
    count( hh, is_child )
end

"""
For Budget constraints: set the wage and make all the needed
corrections to employment status, hours worked, pensions(? pensions not done yet)
"""
function set_wage!( pers :: Person, gross :: Real, wage :: Real )
    pers.income[wages] = gross
    h = gross/wage
    pers.usual_hours_worked = h
    pers.employment_status = if h < 5
        Unemployed
    elseif h < 30
        Part_time_Employee
    else
        Full_time_Employee
    end
    # println( "made usual_hours_worked = $h gross=$gross wage=$wage pers.employment_status=$(pers.employment_status)")
end

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function default_bu_allocation( hh :: Household ) :: BUAllocation
    bua = BUAllocation()
    nbus = 0
    # how many bus
    for (pid,person) in hh.people
        nbus = max( nbus, person.default_benefit_unit )
    end
    ## create benefit units
    for buno in 1:nbus
        push!( bua, PeopleArray())
    end
    sz = size( bua )[1]
    ## allocate peobuallocationple to them
    for (pid,person) in hh.people
        pp = bua[person.default_benefit_unit]
        push!( pp, person )
    end
    ## sort people in each by pid
    for buno in 1:nbus
        sort!( bua[buno], lt=(left,right)->isless(right.age,left.age))
    end
    bua
end

"""

Create a benefit unit from an array of people
really only needed for some unit tests.

FIXME: don't use outside of tests! likely mess up children and spouses.

"""
function make_benefit_unit(
    people :: PeopleArray,
    head :: BigInt,
    spouse:: BigInt = -1 ) :: BenefitUnit
    npeople = size( people )[1]
    T = typeof( people[1].hours_of_care_received )
    pd = People_Dict{T}()
    children = Pid_Array()
    adults = Pid_Array()
    push!( adults, head )
    if spouse > 0
        push!( adults, spouse )
    end
    palloc = spouse <= 0 ? 2 : 3
    for n in palloc:npeople
        push!(children, people[n].pid)
        @assert people[n].age <= 21 # vague idiot check
    end
    for pers in people
        pd[pers.pid] = pers
    end
    return BenefitUnit( pd, head, spouse, adults, children )
end

function child_pids( hh :: Household ) :: Pid_Array
    chlds = Pid_Array()
    for (pid,pers) in hh.people
        if is_child( pers )
            push!( chlds, pid )
        end
    end
    return chlds
end

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function allocate_to_bus( T::Type, bua :: BUAllocation ) :: BenefitUnits
    nbus = size(bua)[1]
    bus = BenefitUnits(undef, nbus)
    for i in 1:nbus
        people = People_Dict{T}()
        head_pid :: BigInt = -1
        spouse_pid :: BigInt = -1
        children = Pid_Array()
        adults = Pid_Array()
        npeople = size( bua[i])[1]
        for p in 1:npeople
            person = bua[i][p]
            people[person.pid] = person
            if p == 1
                head_pid = person.pid
                push!( adults, head_pid )
            else
                # println( "on bu $i person $p relationships $(person.relationships)")
                reltohead = person.relationships[head_pid]
                if reltohead in [Spouse,Cohabitee,Civil_Partner]
                    spouse_pid = person.pid
                    push!( adults, spouse_pid )
                    # FIXME we need to remove these checks if
                    # we're using a non-default allocation to bus
                    # @assert person.age >= 16
                else
                    # @assert person.age <= 19
                    push!( children, person.pid )
                end
            end
        end
        new_bu = BenefitUnit( people, head_pid, spouse_pid, adults, children )
        bus[i] = new_bu
    end
    bus
end

function get_benefit_units(
    hh :: Household{T},
    allocator :: Function=default_bu_allocation ) :: BenefitUnits where T
    allocate_to_bus( T, allocator(hh))
end

function num_people( bu :: BenefitUnit )::Integer
    length( bu.people )
end

function num_adults( bu :: BenefitUnit )::Integer
    size( bu.adults )[1]
end

function adult_pids( hh :: Household )::Pid_Array
    pids = Pid_Array()
    for (pid,pers) in hh.people
        if ! is_child( pers )
            push!( pids, pid )
        end
    end
    return pids
end

function num_adults( hh )
    return size( adult_pids(hh))[1]
end

function num_people( hh :: Household ) :: Integer
    length( hh.people )
end

function get_head( bu :: BenefitUnit )::Person
    bu.people[bu.head]
end

function get_head( hh :: Household ) :: Person
    return hh.people[hh.head_of_household]
end

function is_head( bu :: BenefitUnit, pers :: Person  ) :: Bool
    return bu.head == pers.pid
end

function is_spouse( unit, pers :: Person  ) :: Bool
    head = get_head(unit)
    p = partner_of( head )
    if p === nothing
        return false
    end
    return p == pers.pid
end

function partner_of( pers :: Person ) :: Union{Nothing,BigInt}
    for (pid,rel) in pers.relationships
        if rel in [Spouse,Cohabitee]
            return pid
        end
    end
    return nothing
end

function get_spouse( hh :: Household ) ::Union{Nothing,Person}
    head = get_head( hh )
    p = partner_of( head )
    if p !== nothing
        return hh.people[p]
    end
    return nothing
end

function get_spouse( bu :: BenefitUnit )::Union{Nothing,Person}
    if bu.spouse <= 0
        return nothing
    end
    bu.people[bu.spouse]
end

function has_children( bu :: BenefitUnit ) :: Bool
    num_children( bu ) > 0
end

"""
relies on `is_child` being sensible
"""
function has_children( hh :: Household )::Bool
    for (pid, pers ) in hh.people
        if is_child( pers )
            return true
        end
    end
    return false
end

function is_lone_parent( bu :: BenefitUnit ) :: Bool
    return bu.spouse < 0 && has_children( bu )
end

function is_child( pers :: Person )
    # FIXME crude
    pers.is_standard_child
    # below won't work for hh  since
    # wouldn't be child if in own BU & we can't check that with 1 person
    #if pers.age <= 15
    #    return true
    #end
    #(pers.age <= 19) && (pers.employment_status in [Student])
end

#
# fixme just count people???
#
function is_lone_parent( hh :: Household,
    allocator :: Function=default_bu_allocation ) :: Bool

    bus = get_benefit_units( hh, allocator )
    if size(bus)[1] > 1
        return false
    end
    return bus[1].spouse < 0 && size( bus[1].children )[1] > 0
end

function is_single( bu :: BenefitUnit ) :: Bool
    num_people( bu ) == 1
end

function is_single( hh :: Household ) :: Bool
    num_people( hh ) == 1
end

"""
FIXME we're going to do this solely on benefit receipt/is_informal_carer flag
for now until we get the regressions done
we use historic benefits here since this uses actual data
"""
function pers_is_carer( pers :: Person, params ... ) :: Bool
    if has_non_z( pers.income, carers_allowance )
        return true
    end
    # FIXME check/parameterise this
    if (pers.is_informal_carer) && (pers.hours_of_care_given > 0)
        return true
    end
    return false
end

const DISABLE_BENEFITS = [
    severe_disability_allowance,
    attendance_allowance,
    incapacity_benefit,
    dlaself_care,
    dlamobility,
    personal_independence_payment_daily_living,
    personal_independence_payment_mobility]

"""

"""
function is_severe_disability( pers :: Person )
    return pers_is_disabled( pers ) ||
           pers.dla_self_care_type in [mid, high] ||
           pers.dla_mobility_type == high ||
           pers.attendance_allowance_type == high ||
           pers.pip_daily_living_type == enhanced_pip ||
           pers.pip_mobility_type == enhanced_pip ||
           pers.adls_are_reduced == reduced_a_lot ||
           pers.health_status == Very_Bad

end

"""
FIXME we're going to do this solely on benefit receipt
for now until we get the regressions done
we use historic benefits here since this uses actual data
"""
function pers_is_disabled( pers :: Person, params ... ) :: Bool
    if pers.registered_blind
        return true
    end
    if pers.employment_status in [Permanently_sick_or_disabled]
        return true
    end
    if pers.adls_are_reduced == reduced_a_lot
        return true
    end
    for dis in [mobility,
        dexterity,
        learning,
        memory,
        mental_health]
        if haskey(pers.disabilities,dis) && pers.disabilities[dis]
            return true
        end
    end

    for k in DISABLE_BENEFITS
        if has_non_z( pers.income, k )
            return true
        end
    end # loop
    return false
end

function empl_status_in( pers :: Person, statuses ...)
    return pers.employment_status in statuses
end

le_age( pers :: Person, age ... ) = pers.age <= age[1]

ge_age( pers :: Person, age ... ) = pers.age >= age[1]

between_ages( pers :: Person, age ... ) = age[1] >= pers.age <= age[2]

has_income( pers::Person, which :: Incomes_Type ) = haskey( pers.income, which )

function has_income( pers::Person, which ... )::Bool
    for k in which
        if has_income( pers, k )
            return true
        end
    end
    return false
end

function search( people :: People_Dict, func :: Function, params... ) :: Bool
    for ( pid,pers ) in people
        if func( pers, params... )
            return true
        end
    end
    return false
end

function count( people :: People_Dict, func :: Function, params... ) :: Int
    n = 0
    for ( pid,pers ) in people
        if func( pers, params... )
            n += 1
        end
    end
    return n
end

function search( bu :: BenefitUnit, func :: Function, params... ) :: Bool
    return search( bu.people, func, params... )
end

function search( hh :: Household, func :: Function, params ... ) :: Bool
    return search( hh.people, func, params... )
end

function count( bu :: BenefitUnit, func :: Function, params... ) :: Integer
     return count( bu.people, func, params... )
end

function count( hh :: Household, func :: Function, params ... ) :: Integer
    return count( hh.people, func, params... )
end

function on_mt_benefits( pers :: Person ) :: Bool
    if has_income( pers, [income_support, working_tax_credit, child_tax_credit, housing_benefit, universal_credit]...)
        return true
    end
    if pers.jsa_type == income_related_jsa
        return true
    end
    if pers.esa_type == income_related_jsa
        return true
    end
    return false
end

function on_mt_benefits( bu :: BenefitUnit ) :: Bool
    return search( bu, on_mt_benefits )
end

function on_mt_benefits( hh :: Household ) :: Bool
    return search( hh, on_mt_benefits )
end

function has_disabled_member( bu :: BenefitUnit ) :: Bool
    search( bu.people, pers_is_disabled )
end

function has_disabled_member( hh :: Household ) :: Bool
    search( hh.people, pers_is_disabled )
end

function has_carer_member( bu :: BenefitUnit ) :: Bool
    search( bu.people, pers_is_carer )
end

function has_carer_member( hh :: Household ) :: Bool
    search( hh.people, pers_is_carer )
end

function num_carers( bu :: BenefitUnit ) :: Int
    count( bu.people, pers_is_carer )
end

function carers( hh :: Household ) :: Int
    count( hh.people, pers_is_carer )
end

# simple diagnostic prints for testing allocation

function printpids( pers :: Person )
    println( "pid: $(pers.pid) age $(pers.age) sex $(pers.sex) default_benefit_unit $(pers.default_benefit_unit) ")
end

function printpids( bu::BenefitUnit)
      head = get_head( bu )
      print( "HEAD: ")
      printpids( head )
      spouse = get_spouse( bu )
      if spouse !== nothing
            print( "SPOUSE: ")
            printpids( spouse )

      end
      for chno in bu.children
            print( "CHILD_$(chno):")
            child = bu.people[chno]
            printpids( child )
      end
end

function printpids( people :: People_Dict )
    ks = sort(collect(keys(people)))
    for k in ks
        print( "PERSON_$k")
        printpids( people[k] )
    end
end

function printpids( buas::BUAllocation )
    nbus = size( buas )[1]
    for i in 1:nbus
        npeople = size( buas[i] )[1]
        for p in 1:npeople
            print( "PERS_$(p): ")
            printpids( buas[i][p])
        end
    end
end

function age_then( pers :: Person, when :: Int ) :: Int
    Utils.age_then( pers.age, when );
end

function age_then( pers :: Person, when :: Date = todays_date() ) :: Int
    Utils.age_then( pers.age, Dates.year( when ))
end

end # module


# everything related to a DG semidiscretization in 2D using MPI,
# currently limited to Lobatto-Legendre nodes

# TODO: MPI dimension agnostic
# TODO: MPI, adapt to different real types (and AD!)
mutable struct MPICache
  mpi_neighbor_ranks::Vector{Int}
  mpi_neighbor_interfaces::Vector{Vector{Int}}
  mpi_send_buffers::Vector{Vector{Float64}}
  mpi_recv_buffers::Vector{Vector{Float64}}
  mpi_send_requests::Vector{MPI.Request}
  mpi_recv_requests::Vector{MPI.Request}
  n_elements_by_rank::OffsetArray{Int, 1, Array{Int, 1}}
  n_elements_global::Int
  first_element_global_id::Int
end

function MPICache()
  mpi_neighbor_ranks = Vector{Int}(undef, 0)
  mpi_neighbor_interfaces = Vector{Vector{Int}}(undef, 0)
  mpi_send_buffers = Vector{Vector{Float64}}(undef, 0)
  mpi_recv_buffers = Vector{Vector{Float64}}(undef, 0)
  mpi_send_requests = Vector{MPI.Request}(undef, 0)
  mpi_recv_requests = Vector{MPI.Request}(undef, 0)
  n_elements_by_rank = OffsetArray(Vector{Int}(undef, 0), 0:-1)
  n_elements_global = 0
  first_element_global_id = 0

  MPICache(mpi_neighbor_ranks, mpi_neighbor_interfaces,
           mpi_send_buffers, mpi_recv_buffers,
           mpi_send_requests, mpi_recv_requests,
           n_elements_by_rank, n_elements_global,
           first_element_global_id)
end

# TODO: MPI dimension agnostic
function start_mpi_receive!(mpi_cache::MPICache)

  for (index, d) in enumerate(mpi_cache.mpi_neighbor_ranks)
    mpi_cache.mpi_recv_requests[index] = MPI.Irecv!(
      mpi_cache.mpi_recv_buffers[index], d, d, mpi_comm())
  end

  return nothing
end

# TODO: MPI dimension agnostic
function start_mpi_send!(mpi_cache::MPICache, mesh, equations, dg, cache)
  data_size = nvariables(equations) * nnodes(dg)^(ndims(mesh) - 1)

  for d in 1:length(mpi_cache.mpi_neighbor_ranks)
    send_buffer = mpi_cache.mpi_send_buffers[d]

    for (index, interface) in enumerate(mpi_cache.mpi_neighbor_interfaces[d])
      first = (index - 1) * data_size + 1
      last =  (index - 1) * data_size + data_size

      if cache.mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        @views send_buffer[first:last] .= vec(cache.mpi_interfaces.u[2, :, :, interface])
      else # local element in negative direction
        @views send_buffer[first:last] .= vec(cache.mpi_interfaces.u[1, :, :, interface])
      end
    end
  end

  # Start sending
  for (index, d) in enumerate(mpi_cache.mpi_neighbor_ranks)
    mpi_cache.mpi_send_requests[index] = MPI.Isend(
      mpi_cache.mpi_send_buffers[index], d, mpi_rank(), mpi_comm())
  end

  return nothing
end

# TODO: MPI dimension agnostic
function finish_mpi_send!(mpi_cache::MPICache)
  MPI.Waitall!(mpi_cache.mpi_send_requests)
end

# TODO: MPI dimension agnostic
function finish_mpi_receive!(mpi_cache::MPICache, mesh, equations, dg, cache)
  data_size = nvariables(equations) * nnodes(dg)^(ndims(mesh) - 1)

  # Start receiving and unpack received data until all communication is finished
  d, _ = MPI.Waitany!(mpi_cache.mpi_recv_requests)
  while d != 0
    recv_buffer = mpi_cache.mpi_recv_buffers[d]

    for (index, interface) in enumerate(mpi_cache.mpi_neighbor_interfaces[d])
      first = (index - 1) * data_size + 1
      last =  (index - 1) * data_size + data_size

      if cache.mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        @views vec(cache.mpi_interfaces.u[1, :, :, interface]) .= recv_buffer[first:last]
      else # local element in negative direction
        @views vec(cache.mpi_interfaces.u[2, :, :, interface]) .= recv_buffer[first:last]
      end
    end

    d, _ = MPI.Waitany!(mpi_cache.mpi_recv_requests)
  end

  return nothing
end

# This method is called when a SemidiscretizationHyperbolic is constructed.
# It constructs the basic `cache` used throughout the simulation to compute
# the RHS etc.
function create_cache(mesh::ParallelTreeMesh{2}, equations,
                      dg::DG, RealT, uEltype)
  # Get cells for which an element needs to be created (i.e. all leaf cells)
  leaf_cell_ids = local_leaf_cells(mesh.tree)

  elements = init_elements(leaf_cell_ids, mesh, equations, dg.basis, RealT, uEltype)

  interfaces = init_interfaces(leaf_cell_ids, mesh, elements)

  mpi_interfaces = init_mpi_interfaces(leaf_cell_ids, mesh, elements)

  boundaries = init_boundaries(leaf_cell_ids, mesh, elements)

  mortars = init_mortars(leaf_cell_ids, mesh, elements, dg.mortar)

  mpi_cache = init_mpi_cache(mesh, elements, mpi_interfaces, nvariables(equations), nnodes(dg))

  cache = (; elements, interfaces, mpi_interfaces, boundaries, mortars,
             mpi_cache)

  # Add specialized parts of the cache required to compute the volume integral etc.
  cache = (;cache..., create_cache(mesh, equations, dg.volume_integral, dg, uEltype)...)
  cache = (;cache..., create_cache(mesh, equations, dg.mortar, uEltype)...)

  return cache
end

function init_mpi_cache(mesh, elements, mpi_interfaces, nvars, nnodes)
  mpi_cache = MPICache()

  init_mpi_cache!(mpi_cache, mesh, elements, mpi_interfaces, nvars, nnodes)
  return mpi_cache
end

function init_mpi_cache!(mpi_cache, mesh, elements, mpi_interfaces, nvars, nnodes)
  mpi_neighbor_ranks, mpi_neighbor_interfaces =
    init_mpi_neighbor_connectivity(elements, mpi_interfaces, mesh)

  mpi_send_buffers, mpi_recv_buffers, mpi_send_requests, mpi_recv_requests =
    init_mpi_data_structures(mpi_neighbor_interfaces, ndims(mesh), nvars, nnodes)

  # Determine local and total number of elements
  n_elements_by_rank = Vector{Int}(undef, mpi_nranks())
  n_elements_by_rank[mpi_rank() + 1] = nelements(elements)
  MPI.Allgather!(MPI.UBuffer(n_elements_by_rank, 1), mpi_comm())
  n_elements_by_rank = OffsetArray(n_elements_by_rank, 0:(mpi_nranks() - 1))
  n_elements_global = MPI.Allreduce(nelements(elements), +, mpi_comm())
  @assert n_elements_global == sum(n_elements_by_rank) "error in total number of elements"

  # Determine the global element id of the first element
  first_element_global_id = MPI.Exscan(nelements(elements), +, mpi_comm())
  if mpi_isroot()
    # With Exscan, the result on the first rank is undefined
    first_element_global_id = 1
  else
    # On all other ranks we need to add one, since Julia has one-based indices
    first_element_global_id += 1
  end
  # TODO reuse existing structures
  @pack! mpi_cache = mpi_neighbor_ranks, mpi_neighbor_interfaces,
                     mpi_send_buffers, mpi_recv_buffers,
                     mpi_send_requests, mpi_recv_requests,
                     n_elements_by_rank, n_elements_global,
                     first_element_global_id
end

# Initialize connectivity between MPI neighbor ranks
function init_mpi_neighbor_connectivity(elements, mpi_interfaces, mesh::TreeMesh2D)
  tree = mesh.tree

  # Determine neighbor ranks and sides for MPI interfaces
  neighbor_ranks = fill(-1, nmpiinterfaces(mpi_interfaces))
  # The global interface id is the smaller of the (globally unique) neighbor cell ids, multiplied by
  # number of directions (2 * ndims) plus direction minus one
  global_interface_ids = fill(-1, nmpiinterfaces(mpi_interfaces))
  for interface_id in 1:nmpiinterfaces(mpi_interfaces)
    orientation = mpi_interfaces.orientations[interface_id]
    remote_side = mpi_interfaces.remote_sides[interface_id]
    # Direction is from local cell to remote cell
    if orientation == 1 # MPI interface in x-direction
      if remote_side == 1 # remote cell on the "left" of MPI interface
        direction = 1
      else # remote cell on the "right" of MPI interface
        direction = 2
      end
    else # MPI interface in y-direction
      if remote_side == 1 # remote cell on the "left" of MPI interface
        direction = 3
      else # remote cell on the "right" of MPI interface
        direction = 4
      end
    end
    local_element_id = mpi_interfaces.local_element_ids[interface_id]
    local_cell_id = elements.cell_ids[local_element_id]
    remote_cell_id = tree.neighbor_ids[direction, local_cell_id]
    neighbor_ranks[interface_id] = tree.mpi_ranks[remote_cell_id]
    if local_cell_id < remote_cell_id
      global_interface_ids[interface_id] = 2 * ndims(tree) * local_cell_id + direction - 1
    else
      global_interface_ids[interface_id] = (2 * ndims(tree) * remote_cell_id +
                                            opposite_direction(direction) - 1)
    end
  end

  # Get sorted, unique neighbor ranks
  mpi_neighbor_ranks = unique(sort(neighbor_ranks))

  # Sort interfaces by global interface id
  p = sortperm(global_interface_ids)
  neighbor_ranks .= neighbor_ranks[p]
  interface_ids = collect(1:nmpiinterfaces(mpi_interfaces))[p]

  # For each neighbor rank, init connectivity data structures
  mpi_neighbor_interfaces = Vector{Vector{Int}}(undef, length(mpi_neighbor_ranks))
  for (index, d) in enumerate(mpi_neighbor_ranks)
    mpi_neighbor_interfaces[index] = interface_ids[findall(x->(x == d), neighbor_ranks)]
  end

  # Sanity check that we counted all interfaces exactly once
  @assert sum(length(v) for v in mpi_neighbor_interfaces) == nmpiinterfaces(mpi_interfaces)

  return mpi_neighbor_ranks, mpi_neighbor_interfaces
end

# TODO: MPI dimension agnostic
# Initialize MPI data structures
function init_mpi_data_structures(mpi_neighbor_interfaces, ndims, nvars, n_nodes)
  data_size = nvars * n_nodes^(ndims - 1)
  mpi_send_buffers = Vector{Vector{Float64}}(undef, length(mpi_neighbor_interfaces))
  mpi_recv_buffers = Vector{Vector{Float64}}(undef, length(mpi_neighbor_interfaces))
  for index in 1:length(mpi_neighbor_interfaces)
    mpi_send_buffers[index] = Vector{Float64}(undef, length(mpi_neighbor_interfaces[index]) * data_size)
    mpi_recv_buffers[index] = Vector{Float64}(undef, length(mpi_neighbor_interfaces[index]) * data_size)
  end

  mpi_send_requests = Vector{MPI.Request}(undef, length(mpi_neighbor_interfaces))
  mpi_recv_requests = Vector{MPI.Request}(undef, length(mpi_neighbor_interfaces))

  return mpi_send_buffers, mpi_recv_buffers, mpi_send_requests, mpi_recv_requests
end

function rhs!(du, u, t,
              mesh::ParallelTreeMesh{2}, equations,
              initial_condition, boundary_conditions, source_terms,
              dg::DG, cache)
  # Start to receive MPI data
  @timed timer() "start MPI receive" start_mpi_receive!(cache.mpi_cache)

  # Prolong solution to MPI interfaces
  @timed timer() "prolong2mpiinterfaces" prolong2mpiinterfaces!(
    cache, u, mesh, equations, dg)

  # Start to send MPI data
  @timed timer() "start MPI send" start_mpi_send!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Reset du
  @timed timer() "reset ∂u/∂t" du .= zero(eltype(du))

  # Calculate volume integral
  @timed timer() "volume integral" calc_volume_integral!(
    du, u, mesh,
    have_nonconservative_terms(equations), equations,
    dg.volume_integral, dg, cache)

  # Prolong solution to interfaces
  # TODO: Taal decide order of arguments, consistent vs. modified cache first?
  @timed timer() "prolong2interfaces" prolong2interfaces!(
    cache, u, mesh, equations, dg)

  # Calculate interface fluxes
  @timed timer() "interface flux" calc_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg, cache)

  # Prolong solution to boundaries
  @timed timer() "prolong2boundaries" prolong2boundaries!(
    cache, u, mesh, equations, dg)

  # Calculate boundary fluxes
  @timed timer() "boundary flux" calc_boundary_flux!(
    cache, t, boundary_conditions, mesh, equations, dg)

  # Prolong solution to mortars
  @timed timer() "prolong2mortars" prolong2mortars!(
    cache, u, mesh, equations, dg.mortar, dg)

  # Calculate mortar fluxes
  @timed timer() "mortar flux" calc_mortar_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg.mortar, dg, cache)

  # Finish to receive MPI data
  @timed timer() "finish MPI receive" finish_mpi_receive!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Calculate MPI interface fluxes
  @timed timer() "MPI interface flux" calc_mpi_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg, cache)

  # Calculate surface integrals
  @timed timer() "surface integral" calc_surface_integral!(
    du, mesh, equations, dg, cache)

  # Apply Jacobian from mapping to reference element
  @timed timer() "Jacobian" apply_jacobian!(
    du, mesh, equations, dg, cache)

  # Calculate source terms
  @timed timer() "source terms" calc_sources!(
    du, u, t, source_terms, equations, dg, cache)

  # Finish to send MPI data
  @timed timer() "finish MPI send" finish_mpi_send!(cache.mpi_cache)

  return nothing
end

function prolong2mpiinterfaces!(cache, u,
                                mesh::ParallelTreeMesh{2},
                                equations, dg::DG)
  @unpack mpi_interfaces = cache

  @threaded for interface in eachmpiinterface(dg, cache)
    local_element = mpi_interfaces.local_element_ids[interface]

    if mpi_interfaces.orientations[interface] == 1 # interface in x-direction
      if mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        for j in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[2, v, j, interface] = u[v,          1, j, local_element]
        end
      else # local element in negative direction
        for j in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[1, v, j, interface] = u[v, nnodes(dg), j, local_element]
        end
      end
    else # interface in y-direction
      if mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        for i in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[2, v, i, interface] = u[v, i,          1, local_element]
        end
      else # local element in negative direction
        for i in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[1, v, i, interface] = u[v, i, nnodes(dg), local_element]
        end
      end
    end
  end

  return nothing
end

function calc_mpi_interface_flux!(surface_flux_values,
                                  mesh::ParallelTreeMesh{2},
                                  nonconservative_terms::Val{false}, equations,
                                  dg::DG, cache)
  @unpack surface_flux = dg
  @unpack u, local_element_ids, orientations, remote_sides = cache.mpi_interfaces

  @threaded for interface in eachmpiinterface(dg, cache)
    # Get local neighboring element
    element = local_element_ids[interface]

    # Determine interface direction with respect to element:
    if orientations[interface] == 1 # interface in x-direction
      if remote_sides[interface] == 1 # local element in positive direction
        direction = 1
      else # local element in negative direction
        direction = 2
      end
    else # interface in y-direction
      if remote_sides[interface] == 1 # local element in positive direction
        direction = 3
      else # local element in negative direction
        direction = 4
      end
    end

    for i in eachnode(dg)
      # Call pointwise Riemann solver
      u_ll, u_rr = get_surface_node_vars(u, equations, dg, i, interface)
      flux = surface_flux(u_ll, u_rr, orientations[interface], equations)

      # Copy flux to local element storage
      for v in eachvariable(equations)
        surface_flux_values[v, i, direction, element] = flux[v]
      end
    end
  end

  return nothing
end


# Copyright (C) 2021 Ilya Baran.  This program is distributed under the terms of the MIT license.

using OrdinaryDiffEq
using MAT

abstract type AbstractBlock end

struct BlockOutputDefinition
    outputs::Variables
    computeOutputs::Function # Takes variables, parameters, time range, solution (in Variables form), and outputs, and returns outputs
    includeRawSolution::Bool # If true, computeOutputs takes a raw solution before solution-as-Variables
end

mutable struct BlockExtraData
    timeRange::Union{AbstractRange, Nothing}
    discontinuities::Union{Function, Nothing} # Takes parameters, timeRange, returns vector of discontinuities
    outputDefinition::Union{BlockOutputDefinition, Nothing}
    BlockExtraData() = new(nothing, nothing, nothing)
end

struct Block <: AbstractBlock
    variables::Variables
    parameters::Variables
    reactions::Array{AbstractReaction}
    extraData::BlockExtraData
    Block(v, p, r) = new(deepcopy(v), deepcopy(p), deepcopy(r), BlockExtraData());
end

function getExtraData(block::AbstractBlock)::BlockExtraData block.extraData end

function getTimeRange(block::AbstractBlock)::AbstractRange
    range = getExtraData(block).timeRange;
    if range === nothing
        error("Expected time range but it was missing");
    end
    range;
end

function setTimeRange!(block::AbstractBlock, timeRange::AbstractRange) getExtraData(block).timeRange = timeRange; end

function getOutputDefinition(block::AbstractBlock)::BlockOutputDefinition
    outputs = getExtraData(block).outputDefinition;
    if outputs === nothing
        error("Expected output definition but it was missing");
    end
    outputs;
end

function setOutputDefinition!(block::AbstractBlock, outputs::Variables, computeOutputs::Function)
    getExtraData(block).outputDefinition = BlockOutputDefinition(outputs, computeOutputs, false);
end

function setOutputDefinition!(block::AbstractBlock, outputs::Variables, computeOutputs::Function, includeRawSolution::Bool)
    getExtraData(block).outputDefinition = BlockOutputDefinition(outputs, computeOutputs, includeRawSolution);
end

function renameOutputs!(block::AbstractBlock, renameFunction::Function)
    oldDef = getOutputDefinition(block);
    newOutputs = renameVariables(oldDef.outputs, renameFunction);
    if oldDef.includeRawSolution
        newComputeFunction = (v, p, t, r, s, o) -> begin
            oldOutputs = oldDef.computeOutputs(v, p, t, r, s, deepcopy(oldDef.outputs));
            o.values = oldOutputs.values;
            return o;
        end;
    else
        newComputeFunction = (v, p, t, s, o) -> begin
            oldOutputs = oldDef.computeOutputs(v, p, t, s, deepcopy(oldDef.outputs));
            o.values = oldOutputs.values;
            return o;
        end;
    end
    setOutputDefinition!(block, newOutputs, newComputeFunction, oldDef.includeRawSolution);
end

function getDiscontinuities(block::AbstractBlock)::Vector
    d = getExtraData(block).discontinuities;
    if d === nothing
        return []
    end
    return d(getParameters(block), getTimeRange(block));
end

function setDiscontinuities!(block::AbstractBlock, discontinuities::T) where T <: Function
    getExtraData(block).discontinuities = discontinuities;
end

function runBlock(block::AbstractBlock)
    timeRange = getTimeRange(block);
    floatInterval::Tuple{Float64, Float64} = convert(Tuple{Float64, Float64}, (minimum(timeRange), maximum(timeRange)));
    problem = ODEProblem((x, p, t) -> computeDerivatives(block, t, x), Vector{Float64}(getVariables(block).values), floatInterval);
    solve(problem, Rodas4P(), tstops = getDiscontinuities(block));
end

function setParameters!(block::AbstractBlock, nameToValue::AbstractDict{String, T}) where T
    for (name, value) in nameToValue
        setParameter!(block, name, value);
    end
end

function solutionToVariables(solution, block::AbstractBlock, timeRange::AbstractRange)
    resultArray = Array{Array{Float64, 1}, 1}(undef, length(getVariables(block).variables));
    for index = 1:length(resultArray)
        resultArray[index] = solution(timeRange, idxs=index).u;
    end
    return Variables(getVariables(block), resultArray);
end

function solutionToMatrix(solution, timeRange::AbstractRange)::Matrix
    return Matrix(Array(solution(timeRange)));
end

function solutionToMatlab(solution, timeRange::AbstractRange, filename::String)
    t = Vector(timeRange);
    y = solutionToMatrix(solution, timeRange);
    status = string(solution.retcode);
    file = matopen(filename, "w");
    write(file, "t", t);
    write(file, "y", y);
    write(file, "status", status);
    close(file);
end

function computeDerivatives(block::Block, t::Number, x::Vector)::Vector
    derivatives = Variables(block.variables, zero(t * x));
    variables = Variables(block.variables, x);
    extraVariables = Dict{String, Any}();
    for reaction = block.reactions
        apply!(reaction, derivatives, t, variables, block.parameters, extraVariables);
    end
    #println("Time = $t");#, x = $x, result = $(derivatives.values)");
    vType = typeof(t * x);
    return vType(derivatives.values);
end

function getVariables(block::Block)::Variables block.variables; end
function getParameters(block::Block)::Variables block.parameters; end
function setParameter!(block::Block, name::String, value)
    block.parameters[name] = value;
end
function setParameters!(block::Block, values::Vector)
    block.parameters.values = values;
end

function computeOutputs(blockWithOutputs::AbstractBlock)
    timeRange = getTimeRange(blockWithOutputs);
    solution = runBlock(blockWithOutputs);
    return computeOutputsFromSolution(blockWithOutputs, timeRange, solution);
end

# ============================================ BlockWithBindings =======================================
struct BlockWithBindings <: AbstractBlock
    subblock::AbstractBlock
    parameterToBinding::Dict{String, Function} # Function takes t, variables, and existing parameters and returns a value
    parameters::Variables
    extraData::BlockExtraData
end

function BlockWithBindings(subblock::AbstractBlock, parameterToBinding::AbstractDict{String, T}) where T
    subblock = deepcopy(subblock);
    subblockParameters::Variables = getParameters(subblock);
    blockParameterToBinding = Dict{String, Function}();
    for (parameter, binding) in parameterToBinding
        if !haskey(subblockParameters.nameToIndex, parameter)
            error("Cannot bind nonexistent parameter $(parameter)");
        end
        if isa(binding, Function)
            blockParameterToBinding[parameter] = binding;
        else
            setParameter!(subblock, parameter, binding);
        end
    end
    newParameters = variablesSubtract(subblockParameters, Set(keys(parameterToBinding)));
    extraData = BlockExtraData();
    extraData.timeRange = getExtraData(subblock).timeRange;
    return BlockWithBindings(subblock, blockParameterToBinding, newParameters, extraData);
end

# Keep the parameters that are provided, and bind the rest to their current values
function BlockWithBindings(subblock::AbstractBlock, parametersToKeep::AbstractArray{String})
    subblock = deepcopy(subblock);
    subblockParameters::Variables = getParameters(subblock);
    newParameters::Array{Variable} = [];
    newValues::Array = [];
    for parameter in parametersToKeep
        if !haskey(subblockParameters.nameToIndex, parameter)
            error("Cannot keep nonexistent parameter $(parameter)");
        end
        index = subblockParameters.nameToIndex[parameter];
        variable = subblockParameters.variables[index];
        push!(newParameters, variable);
        push!(newValues, subblockParameters.values[index]);
    end
    newParametersV = Variables(newParameters);
    newParametersV.values = newValues;
    extraData = BlockExtraData();
    extraData.timeRange = getExtraData(subblock).timeRange;
    return BlockWithBindings(subblock, Dict{String, Function}(), newParametersV, extraData);
end

function getVariables(block::BlockWithBindings)::Variables getVariables(block.subblock); end
function getParameters(block::BlockWithBindings)::Variables block.parameters; end
function setParameter!(block::BlockWithBindings, name::String, value)
    block.parameters[name] = value;
    setParameter!(block.subblock, name, value);
end
function setParameters!(block::BlockWithBindings, values::Vector)
    block.parameters.values = values;
    for (name, index) in block.parameters.nameToIndex
        setParameter!(block.subblock, name, block.parameters.values[index]);
    end
end

function computeDerivatives(block::BlockWithBindings, t::Number, x::Vector)::Vector
    blockVariables = Variables(getVariables(block), x);
    for (parameter, binding) in block.parameterToBinding
        setParameter!(block.subblock, parameter, binding(t, blockVariables, block.parameters));
    end
    return computeDerivatives(block.subblock, t, x);
end

function computeOutputsFromSolution(block::AbstractBlock, timeRange, solution)
    outputDefinition = getOutputDefinition(block)
    if outputDefinition.includeRawSolution
        return outputDefinition.computeOutputs(getVariables(block), getParameters(block), timeRange,
                                            solution, solutionToVariables(solution, block, timeRange),
                                            deepcopy(outputDefinition.outputs));
    else
        return outputDefinition.computeOutputs(getVariables(block), getParameters(block), timeRange,
                                            solutionToVariables(solution, block, timeRange),
                                            deepcopy(outputDefinition.outputs));
    end
end

function recursiveComputeOutputs(block::AbstractBlock, timeRange, solution)
    computeOutputsFromSolution(block, timeRange, solution);
end

function recursiveComputeOutputs(block::BlockWithBindings, timeRange, solution)
    outputDefinition = getExtraData(block).outputDefinition;
    if (outputDefinition === nothing)
        return recursiveComputeOutputs(block.subblock, timeRange, solution);
    end
    computeOutputsFromSolution(block, timeRange, solution);
end

function computeOutputs(blockWithOutputs::BlockWithBindings)
    timeRange = getTimeRange(blockWithOutputs);
    solution = runBlock(blockWithOutputs);
    recursiveComputeOutputs(blockWithOutputs, timeRange, solution);
end

function getDiscontinuities(block::BlockWithBindings)::Vector
    d = getExtraData(block).discontinuities;
    if d === nothing
        return getDiscontinuities(block.subblock);
    end
    return d(getParameters(block), getTimeRange(block));
end

# ============================================ BlockCombo =======================================

const EXTRA_NAME = "_extra_";

struct BlockCombo <: AbstractBlock
    subblocks::Vector{Tuple{String, AbstractBlock}}
    parameters::Variables
    parameterNameToBlocksAndNames::Dict{String, Vector{Tuple{String, String}}} # For setting parameters
    glueFunctions::Vector{Tuple{String, String, Function}} # Block name, parameter name, function that computes it
    # Glue functions take time, variables, parameters (of the entire block)
    variables::Variables
    extraData::BlockExtraData
end

function BlockCombo(subblocks::Vector{Tuple{String, T}},
                    glueFunctions::Vector{Tuple{String, String, F}},
                    extraParameters::Variables) where {T <: AbstractBlock, F <: Function}
    if (length(subblocks) < 1)
        error("Must have at least one subblock");
    end
    extraData = deepcopy(getExtraData(subblocks[1][2]));
    subblocks = deepcopy(subblocks);
    gluedParameters = Set((f[1], f[2]) for f in glueFunctions);

    # Gather parameters
    seenNames = Set{String}();
    parameterArray = Vector{Variable}();
    parameterValueArray = [];
    parameterNameToBlocksAndNames = Dict{String, Vector{Tuple{String, String}}}();

    namesAndParameterSets = [(b[1], getParameters(b[2])) for b in subblocks];
    push!(namesAndParameterSets, (EXTRA_NAME, extraParameters));

    for namesAndParameters in namesAndParameterSets
        parameters = namesAndParameters[2];
        for parameter::Variable in parameters.variables
            if (in((namesAndParameters[1], parameter.name), gluedParameters))
                continue;
            end
            if !in(parameter.name, seenNames)
                push!(seenNames, parameter.name);
                push!(parameterArray, parameter);
                push!(parameterValueArray, parameters.values[parameters.nameToIndex[parameter.name]]);
            end
            if (namesAndParameters[1] != EXTRA_NAME)
                push!(get!(parameterNameToBlocksAndNames, parameter.name, []), (namesAndParameters[1], parameter.name));
            end
        end
    end

    parameters = Variables(parameterArray);

    variableArray = Vector{Variable}();
    variableValueArray = [];

    for nameAndSubblock in subblocks
        variables = getVariables(nameAndSubblock[2]);
        for variable in variables.variables
            push!(variableArray, variable);
            push!(variableValueArray, variables.values[variables.nameToIndex[variable.name]]);
        end
    end

    variables = Variables(variableArray);
    variables.values = variableValueArray;

    result = BlockCombo(subblocks, parameters, parameterNameToBlocksAndNames, glueFunctions, variables, extraData);
    setParameters!(result, parameterValueArray); # Propagate to subblocks
    result;
end

function getSubblock(block::BlockCombo, name::String)
    for (subblockName, subblock) in block.subblocks
        if name == subblockName
            return subblock;
        end
    end
    error("Subblock $name not found");
end

function getVariables(block::BlockCombo)::Variables block.variables; end
function getParameters(block::BlockCombo)::Variables block.parameters; end
function setParameter!(block::BlockCombo, name::String, value)
    block.parameters[name] = value;
    if haskey(block.parameterNameToBlocksAndNames, name)
        blocksAndNames = block.parameterNameToBlocksAndNames[name];
        for blockAndName in blocksAndNames
            setParameter!(getSubblock(block, blockAndName[1]), blockAndName[2], value);
        end
    end
end

function setParameters!(block::BlockCombo, values::Vector)
    block.parameters.values = values;

    for parameter in block.parameters.variables
        blocksAndNames = get(block.parameterNameToBlocksAndNames, parameter.name, nothing);
        if blocksAndNames === nothing
            continue;
        end
        for blockAndName in blocksAndNames
            setParameter!(getSubblock(block, blockAndName[1]), blockAndName[2], values[block.parameters.nameToIndex[parameter.name]]);
        end
    end
end

function computeDerivatives(block::BlockCombo, t::Number, x::Vector)::Vector
    blockVariables = Variables(block.variables, x);
    for nameParameterAndGlue in block.glueFunctions
        value = nameParameterAndGlue[3](t, blockVariables, block.parameters);
        setParameter!(getSubblock(block, nameParameterAndGlue[1]), nameParameterAndGlue[2], value);
    end

    derivatives = [];
    variableIndex = 1;
    for (_, subblock) in block.subblocks
        nVariables = length(getVariables(subblock).variables);
        push!(derivatives, computeDerivatives(subblock, t, x[variableIndex:(variableIndex + nVariables - 1)]))
        variableIndex += nVariables;
    end

    vcat(derivatives...);
end


using StatsBase

"""
    function mutation!(pop::populations)

Mutate a base pair in a random species, chosen depending on subpopulation size.
If the species is extinct after mutation, it is removed. If mutation creates a new species,
it is added to the population, otherwise the subpopulation size of species the individual mutates to
is simply increased by one.
"""
function mutation!(pop::populations)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample(collect(1:length(pop.seqs)), Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate
    push!(pop.seqs, copy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    pop.freqs[mutation_sequence] -= 1
    # Choose random base to mutate
    mutation_base = rand(1:pop.l)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[end][mutation_base], collect(1:pop.n))
    # Choose random new base
    pop.seqs[end][mutation_base] = rand(poss_bases)
    copied = findfirst(x->~any(.~(x.==pop.seqs[end])), pop.seqs[1:end-1])
    # Look if mutated sequence already exists
    if ~isnothing(copied)
        pop.freqs[copied] += 1
        deleteat!(pop.freqs, length(pop.seqs))
        deleteat!(pop.seqs, length(pop.seqs))
    end
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs)
end

"""
    function mutation!(pop::driver_trailer_l)

Mutate a base pair in a random species, chosen depending on subpopulation size.
If the species is extinct after mutation, it is removed. If mutation creates a new species,
it is added to the population, otherwise the subpopulation size of species the individual mutates to
is simply increased by one.
"""
function mutation!(pop::driver_trailer_l)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample(collect(1:length(pop.seqs)), Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate
    push!(pop.seqs, deepcopy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    push!(pop.l, pop.l[mutation_sequence])
    pop.freqs[mutation_sequence] -= 1
    # Choose random base to mutate
    mutation_base = rand(1:pop.L)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[end][mutation_base], collect(1:pop.n))
    # Choose random new base
    pop.seqs[end][mutation_base] = rand(poss_bases)
    copied = findfirst(x->~any(.~(x.==pop.seqs[end])), pop.seqs[1:end-1])
    # Look if mutated sequence already exists
    if ~isnothing(copied)
        pop.freqs[copied] += 1
        deleteat!(pop.freqs, length(pop.seqs))
        deleteat!(pop.seqs, length(pop.seqs))
        deleteat!(pop.l, length(pop.l))
    end
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs, pop.l)
end

"""
    function driver_mutation!(pop::DT_population)

Mutate a random base in the driver sequence.
"""
function driver_mutation!(pop::DT_population)
    # Choose random base to mutate
    mutation_base = rand(1:length(pop.driver))
    # Exclude self mutations
    poss_bases = filter(x->x != pop.driver[mutation_base], collect(1:pop.m))
    # Choose random new base
    pop.driver[mutation_base] = rand(poss_bases)
end

"""
    function driver_mutation!(pop::driver_trailer)

Mutate the length of species and create a new species
"""
function length_mutation!(pop::driver_trailer_l)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample( collect(1:length(pop.seqs)),
                                Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate

    if rand() < 0.5 && pop.l[mutation_sequence] < pop.L
        var = 1
    elseif pop.l[mutation_sequence] > 2
        var = -1
    else
        return
    end

    push!(pop.seqs, copy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    push!(pop.l, pop.l[mutation_sequence])

    pop.freqs[mutation_sequence] -= 1
    pop.l[end] += var
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs, pop.l)
end

"""
    bp_substitution!(pop::mono_pop, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

Attempt a base pair substitution. The population has to consist of one species
only. Mutation is accepted with probability given by the Kimura fixation
probability.
"""
function bp_substitution!(pop::mono_pop, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

    # Choose random base to mutate
    mutation_base = rand(eachindex(pop.seqs))
    # Exclude self mutations
    possible_bases = filter(x->x != pop.seqs[mutation_base], collect(1:pop.n))
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.seqs[mutation_base] = rand(possible_bases)
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E, length(pop.seqs), fitness_function)
    F_mutant = fitness(E_mutant, length(pop.seqs), fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
    end
end

"""
    bp_substitution!(pop::populations, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

Attempt a base pair substitution. The population has to consist of one species
only. Mutation is accepted with probability given by the Kimura fixation
probability.
"""
function bp_substitution!(pop::populations, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}
    if length(pop.seqs) > 1
        throw(ArgumentError("Input population has more than one species."))
    end
    # Choose random base to mutate
    mutation_base = rand(1:pop.L)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[1][mutation_base], collect(1:pop.n))
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.seqs[1][mutation_base] = rand(poss_bases)
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E[1], pop.l[1], fitness_function)
    F_mutant = fitness(E_mutant[1], pop.l[1], fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
    end
end

"""
    l_substitution!(
        pop::driver_trailer_l,
        emat::Array{T, 2},
        fitness_function::fitness_functions) where {T<:Real}

Attempt a substitution of a length mutation. Population has to be consisting
of one species only.
"""
function l_substitution!(
    pop::driver_trailer_l,
    emat::Array{T, 2},
    fitness_function::fitness_functions) where {T<:Real}

    if length(pop.seqs) > 1
        throw(ArgumentError("Input population has more than one species."))
    end

    if rand() < 0.5 && pop.l[1] < pop.L
        var = 1
    elseif pop.l[1] > 2
        var = -1
    else
        return
    end
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.l[1] += var
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E[1], pop.l[1], fitness_function)
    F_mutant = fitness(E_mutant[1], temp_pop.l[1], fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.l = temp_pop.l
    end
end

"""
    l_substitution!(
        pop::mono_pop,
        emat::Array{T, 2},
        fitness_function::fitness_functions) where {T<:Real}

Attempt a substitution of a length mutation. Population has to be consisting
of one species only.
"""
function l_substitution!(
    pop::mono_pop,
    emat::Array{T, 2},
    fitness_function::fitness_functions) where {T<:Real}

    temp_pop = deepcopy(pop)
    if rand() < 0.5
        push!(temp_pop.seqs, rand(1:pop.n))
        push!(temp_pop.driver, rand(1:pop.m))
    elseif length(temp_pop.seqs) > 1
        pop!(temp_pop.seqs)
        pop!(temp_pop.driver)
    end
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E, length(pop.seqs), fitness_function)
    F_mutant = fitness(E_mutant, length(temp_pop.seqs), fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
        pop.driver = temp_pop.driver
    end
end

"""
    kimura_prob(s, N)

Probability of fixation of a mutation with selection coefficient
`s` in a monomorphic population of size `N`.
"""
function kimura_prob(s::Real, N::Int)
    if abs(s) <= 10^-8
        return 1 / N
    elseif s < -10^3
        return 0
    else
        return (1 - exp(-2s)) / (1 - exp(-2N * s))
    end
end


# Bit-parallel element nucleotide complementation
@inline function _complement_bitpar(a::A, head::UInt64, tail...) where {A<:NucleicAcidAlphabet{2}}
    return (~head, _complement_bitpar(a, tail...)...)
end

@inline _complement_bitpar(a::A) where {A<:NucleicAcidAlphabet{2}} = ()

@inline function pushfirst(x::Kmer{A,K,N}, nt) where {A,K,N}
    ntbits = UInt64(BioSequences.encode(A(), nt)) << (62 - (64N - 2K))
    #ntbits = UInt64(@inbounds BioSequences.twobitnucs[reinterpret(UInt8, nt) + 0x01]) << (62 - (64N - 2K))
    return Kmer{A,K,N}(_rightshift_carry(2, ntbits, x.data...))
end

@inline function pushlast(x::Kmer{A,K,N}, nt) where {A,K,N}
    ntbits = UInt64(BioSequences.encode(A(), nt))
    #ntbits = UInt64(@inbounds BioSequences.twobitnucs[reinterpret(UInt8, nt) + 0x01])
    _, newbits = _leftshift_carry(2, ntbits, x.data...)
    return Kmer{A,K,N}(newbits)
end

###
### Transformation methods
###

"""
    complement(seq::T) where {T<:Kmer}
Return the complement of a short sequence type `x`.
"""
@inline function BioSequences.complement(seq::T) where {T<:Kmer}
    return T(_complement_bitpar(Alphabet(seq), seq.data...))
end

"""
    reverse(seq::Kmer{A,K,N}) where {A,K,N}
Return the reverse of short sequence type variable `seq`.
"""
@inline function Base.reverse(seq::Kmer{A,K,N}) where {A,K,N}
#    rdata = _reverse(identity, BioSequences.BitsPerSymbol(seq), seq.data...)
    rdata = _reverse(BioSequences.BitsPerSymbol(seq), seq.data...)
    return Kmer{A,K,N}(rightshift_carry(rdata, 64N - 2K))
end

"""
    reverse_complement(x::Kmer)

Return the reverse complement of `x`.
"""
@inline function BioSequences.reverse_complement(x::Kmer{A,K,N}) where {A,K,N}
    return complement(reverse(x))
end

#=
@inline function reverse_complement2(seq::Kmer{A,K,N}) where {A,K,N}
    f = x -> complement_bitpar(x, A())
    rdata = _reverse(f, BioSequences.BitsPerSymbol(seq), seq.data...)
    return Kmer{A,K,N}(rightshift_carry(rdata, 64N - 2K))
end
=#

"""
    canonical(x::Kmer)

Return the canonical sequence of `x`.

A canonical sequence is the numerical lesser of a k-mer and its reverse complement.
This is useful in hashing/counting sequences in data that is not strand specific,
and thus observing the short sequence is equivalent to observing its reverse complement.
"""
@inline canonical(x::Kmer) = min(x, reverse_complement(x))

###
### Old Mer specific specializations of src/biosequence/transformations.jl
### - not currently transferred to new type.

# TODO: Sort this and decide on transferring to new NTuple based kmers or no.

#=
function swap(x::T, i, j) where {T<:AbstractMer}
    i = 2 * length(x) - 2i
    j = 2 * length(x) - 2j
    b = encoded_data(x)
    x = ((b >> i) ⊻ (b >> j)) & encoded_data_type(x)(0x03)
    return T(b ⊻ ((x << i) | (x << j)))
end

function Random.shuffle(x::T) where {T<:AbstractMer}
    # Fisher-Yates shuffle for mers.
    j = lastindex(x)
    for i in firstindex(x):(j - 1)
        j′ = rand(i:j)
        x = swap(x, i, j′)
    end
    return x
end
=#