text:
export sourceParticles

foreach([<=, >=, <, >]) do cmp
    MonteCarloMeasurements.register_primitive(cmp, eval)
end

export particles

particles(v::Vector{<:Real}) = Particles(v)
particles(d) = begin
    Particles(1000, d)
end

# particles(n :: NUTS_result) = particles(n.samples)
using IterTools

function particles(v::Vector{Vector{T} }) where {T}
    map(eachindex(v[1])) do j particles([x[j] for x in v]) end
end

function particles(s::Vector{NamedTuple{vs, T}}) where {vs, T}
    nt = NamedTuple()
    for k in keys(s[1])
        nt = merge(nt, [k => particles(getproperty.(s,k))])
    end
    nt
end

export parts

# Just a little helper function for particles
# https://github.com/baggepinnen/MonteCarloMeasurements.jl/issues/22
parts(d; N=1000) = particles(m)
parts(x::Normal{P} where {P <: AbstractParticles}; N=1000) = Particles(length(x.μ), Normal()) * x.σ + x.μ
parts(x::Sampleable{F,S}; N=1000) where {F,S} = Particles(N,x)
parts(x::Integer; N=1000) = parts(float(x))
parts(x::Real; N=1000) = parts(repeat([x],N))
parts(x::AbstractArray; N=1000) = Particles(x)
parts(p::Particles; N=1000) = p
parts(d::For; N=1000) = parts.(d.f.(d.θ...))

parts(d::iid; N=1000) = map(1:d.size) do j parts(d.dist) end
# size
# dist

# MonteCarloMeasurements.Particles(n::Int, d::For) =
# -(a::Particles{Float64,1000}, b::Array{Float64,1}) = [a-bj for bj in b]

# Base.promote(a::Particles{T,N}, b)  where {T,N} = (a,parts(b))
# Base.promote(a, b::Particles{T,N})  where {T,N} = (parts(a),b)

# promote_rule(::Type{A}, ::Type{B}) where {A <: Real, B <: AbstractParticles{T,N}} where {T} = AbsractParticles{promote_type(A,T),N} where {N}
# promote_rule(::Type{B}, ::Type{A}) where {A <: Real, B <: AbstractParticles{T,N}} where {T} = AbsractParticles{promote_type(A,T),N} where {N}

@inline function particles(m::JointDistribution)
    return _particles(m.model, m.args)
end

@gg function _particles(_m::Model, _args)
    type2model(_m) |> sourceParticles() |> loadvals(_args, NamedTuple())
end

@gg function _particles(_m::Model, _args::NamedTuple{()})
    type2model(_m) |> sourceParticles()
end

export sourceParticles
function sourceParticles()
    function(m::Model)

        _m = canonical(m)
        proc(_m, st::Assign)  = :($(st.x) = $(st.rhs))
        proc(_m, st::Sample)  = :($(st.x) = parts($(st.rhs)))
        proc(_m, st::Return)  = :(return $(st.rhs))
        proc(_m, st::LineNumber) = nothing

        vals = map(x -> Expr(:(=), x,x),variables(_m))

        wrap(kernel) = @q begin
            $kernel
            $(Expr(:tuple, vals...))
        end

        buildSource(_m, proc, wrap) |> flatten
    end
end


text:
## --- CRUST 1.0

    # Download CRUST 1.0 data and references from cloud
    function get_crust1()
        # Available variable names: "seafloorage", "seafloorage_sigma",
        # "seafloorrate", "information", and "reference".

        # Construct file paths
        filedir = joinpath(resourcepath,"crust1")
        referencepath = joinpath(filedir,"crust1.references.txt")
        vppath = joinpath(filedir,"crust1.vp")
        vspath = joinpath(filedir,"crust1.vs")
        rhopath = joinpath(filedir,"crust1.rho")
        bndpath = joinpath(filedir,"crust1.bnds")

        # Download HDF5 file from Google Cloud if necessary
        if ~isfile(referencepath)
            print("Downloading crust1 files from google cloud storage to $filedir\n")
            run(`mkdir -p $filedir`)
            download("https://storage.googleapis.com/statgeochem/crust1.references.txt", referencepath)
            download("https://storage.googleapis.com/statgeochem/crust1.vp", vppath)
            download("https://storage.googleapis.com/statgeochem/crust1.vs", vspath)
            download("https://storage.googleapis.com/statgeochem/crust1.rho", rhopath)
            download("https://storage.googleapis.com/statgeochem/crust1.bnds", bndpath)
        end

        return 0 # Success
    end
    export get_crust1

    # Get all point data (Vp, Vs, Rho, layer thickness) from Crust 1.0 layer
    function find_crust1_layer(lat,lon,layer)
        # Get Vp, Vs, Rho, and thickness for a given lat, lon, and crustal layer.

        if length(lat) != length(lon)
            error("lat and lon must be equal length\n")
        end

        if ~isa(layer,Integer) || layer < 1 || layer > 8
            error("""Error: layer must be an integer between 1 and 8.
            Available layers:
            1) water
            2) ice
            3) upper sediments   (VP, VS, rho not defined in all cells)
            4) middle sediments  "
            5) lower sediments   "
            6) upper crystalline crust
            7) middle crystalline crust
            8) lower crystalline crust
            Results are returned in form (Vp, Vs, Rho, thickness)
            """)
        end

        nlayers=9
        nlon=360
        nlat=180

        # Allocate data arrays
        vp = Array{Float64,3}(undef,nlayers,nlat,nlon)
        vs = Array{Float64,3}(undef,nlayers,nlat,nlon)
        rho = Array{Float64,3}(undef,nlayers,nlat,nlon)
        bnd = Array{Float64,3}(undef,nlayers,nlat,nlon)

        # Open data files
        vpfile = open(joinpath(resourcepath,"crust1","crust1.vp"), "r")
        vsfile = open(joinpath(resourcepath,"crust1","crust1.vs"), "r")
        rhofile = open(joinpath(resourcepath,"crust1","crust1.rho"), "r")
        bndfile = open(joinpath(resourcepath,"crust1","crust1.bnds"), "r")

        # Read data files into array
        for j=1:nlat
           for i=1:nlon
              vp[:,j,i] = delim_string_parse(readline(vpfile), ' ', Float64, merge=true)
              vs[:,j,i] = delim_string_parse(readline(vsfile), ' ', Float64, merge=true)
              rho[:,j,i] = delim_string_parse(readline(rhofile), ' ', Float64, merge=true)
              bnd[:,j,i] = delim_string_parse(readline(bndfile), ' ', Float64, merge=true)
          end
        end

        # Close data files
        close(vpfile)
        close(vsfile)
        close(rhofile)
        close(bndfile)

        # Avoid edge cases at lat = -90.0, lon = 180.0
        ilon = mod.(lon .+ 180, 360) .- 180
        ilat = max.(lat, -90+1e-9)

        # Convert lat and lon to index
        ilat = 91 .- ceil.(Int,ilat)
        ilon = 181 .+ floor.(Int,ilon)

        # Allocate output arrays
        vpout = Array{Float64}(undef,size(lat))
        vsout = Array{Float64}(undef,size(lat))
        rhoout = Array{Float64}(undef,size(lat))
        thkout = Array{Float64}(undef,size(lat))

        # Fill output arrays
        for j=1:length(lat)
            if isnan(lat[j]) || isnan(lon[j]) || lat[j] > 90 || lat[j] < -90 || lon[j] > 180 || lat[j] < -180
                vpout[j] = NaN
                vsout[j] = NaN
                rhoout[j] = NaN
                thkout[j] = NaN
            else
                vpout[j] = vp[layer,ilat[j],ilon[j]]
                vsout[j] = vs[layer,ilat[j],ilon[j]]
                rhoout[j] = rho[layer,ilat[j],ilon[j]]
                thkout[j] = bnd[layer,ilat[j],ilon[j]] - bnd[layer+1,ilat[j],ilon[j]]
            end
        end

        # The end
        return (vpout, vsout, rhoout, thkout)
    end
    export find_crust1_layer

    # Get seismic data (Vp, Vs, Rho) for crust 1.0 layer
    function find_crust1_seismic(lat,lon,layer)

        if length(lat) != length(lon)
            error("lat and lon must be equal length\n")
        end

        if ~isa(layer,Integer) || layer < 1 || layer > 9
            error("""Error: layer must be an integer between 1 and 9.
            Available layers:
            1) water
            2) ice
            3) upper sediments   (VP, VS, rho not defined in all cells)
            4) middle sediments  "
            5) lower sediments   "
            6) upper crystalline crust
            7) middle crystalline crust
            8) lower crystalline crust
            9) Top of mantle below crust
            Results are returned in form (Vp, Vs, Rho)
            """)
        end

        nlayers=9
        nlon=360
        nlat=180

        # Allocate data arrays
        vp = Array{Float64,3}(undef,nlayers,nlat,nlon)
        vs = Array{Float64,3}(undef,nlayers,nlat,nlon)
        rho = Array{Float64,3}(undef,nlayers,nlat,nlon)

        # Open data files
        vpfile = open(joinpath(resourcepath,"crust1","crust1.vp"), "r")
        vsfile = open(joinpath(resourcepath,"crust1","crust1.vs"), "r")
        rhofile = open(joinpath(resourcepath,"crust1","crust1.rho"), "r")

        # Read data files into array
        for j=1:nlat
           for i=1:nlon
              vp[:,j,i] = delim_string_parse(readline(vpfile), ' ', Float64, merge=true)
              vs[:,j,i] = delim_string_parse(readline(vsfile), ' ', Float64, merge=true)
              rho[:,j,i] = delim_string_parse(readline(rhofile), ' ', Float64, merge=true) * 1000 # convert to kg/m3
          end
        end

        # Close data files
        close(vpfile)
        close(vsfile)
        close(rhofile)

        # Avoid edge cases at lat = -90.0, lon = 180.0
        ilon = mod.(lon .+ 180, 360) .- 180
        ilat = max.(lat, -90+1e-9)

        # Convert lat and lon to index
        ilat = 91 .- ceil.(Int,ilat)
        ilon = 181 .+ floor.(Int,ilon)

        # Allocate output arrays
        vpout = Array{Float64}(undef,size(lat))
        vsout = Array{Float64}(undef,size(lat))
        rhoout = Array{Float64}(undef,size(lat))

        # Fill output arrays
        for j=1:length(lat)
            if isnan(lat[j]) || isnan(lon[j]) || lat[j] > 90 || lat[j] < -90 || lon[j] > 180 || lat[j] < -180
                vpout[j] = NaN
                vsout[j] = NaN
                rhoout[j] = NaN
            else
                vpout[j] = vp[layer,ilat[j],ilon[j]]
                vsout[j] = vs[layer,ilat[j],ilon[j]]
                rhoout[j] = rho[layer,ilat[j],ilon[j]]
            end
        end

        # The end
        return (vpout, vsout, rhoout)
    end
    export find_crust1_seismic

    # Get layer thickness for crust 1.0 layer
    function find_crust1_thickness(lat,lon,layer)
        # Layer thickness for a given lat, lon, and crustal layer.

        if length(lat) != length(lon)
            error("lat and lon must be equal length\n")
        end

        if ~isa(layer,Integer) || layer < 1 || layer > 8
            error("""Error: layer must be an integer between 1 and 8.
            Available layers:
            1) water
            2) ice
            3) upper sediments   (VP, VS, rho not defined in all cells)
            4) middle sediments  "
            5) lower sediments   "
            6) upper crystalline crust
            7) middle crystalline crust
            8) lower crystalline crust
            Result is thickness of the requested layer
            """)
        end

        nlayers=9
        nlon=360
        nlat=180

        # Allocate data arrays
        bnd = Array{Float64,3}(undef,nlayers,nlat,nlon)

        # Open data files
        bndfile = open(joinpath(resourcepath,"crust1","crust1.bnds"), "r")

        # Read data files into array
        for j=1:nlat
           for i=1:nlon
              bnd[:,j,i] = delim_string_parse(readline(bndfile), ' ', Float64, merge=true)
          end
        end

        # Close data files
        close(bndfile)

        # Avoid edge cases at lat = -90.0, lon = 180.0
        ilon = mod.(lon .+ 180, 360) .- 180
        ilat = max.(lat, -90+1e-9)

        # Convert lat and lon to index
        ilat = 91 .- ceil.(Int,ilat)
        ilon = 181 .+ floor.(Int,ilon)

        # Allocate output arrays
        thkout = Array{Float64}(undef,size(lat))

        # Fill output arrays
        for j=1:length(lat)
            if isnan(lat[j]) || isnan(lon[j]) || lat[j] > 90 || lat[j] < -90 || lon[j] > 180 || lat[j] < -180
                thkout[j] = NaN
            else
                thkout[j] = bnd[layer,ilat[j],ilon[j]]-bnd[layer+1,ilat[j],ilon[j]]
            end
        end

        # The end
        return thkout
    end
    export find_crust1_thickness

    # Get detph to layer base for crust 1.0 layer
    function find_crust1_base(lat,lon,layer)
        # Layer thickness for a given lat, lon, and crustal layer.

        if length(lat) != length(lon)
            error("lat and lon must be equal length\n")
        end

        if ~isa(layer,Integer) || layer < 1 || layer > 8
            error("""layer must be an integer between 1 and 8.
            Available layers:
            1) water
            2) ice
            3) upper sediments   (VP, VS, rho not defined in all cells)
            4) middle sediments  "
            5) lower sediments   "
            6) upper crystalline crust
            7) middle crystalline crust
            8) lower crystalline crust
            Result is depth from sea level to base of the requested layer
            """)
        end
        nlayers=9
        nlon=360
        nlat=180

        # Allocate data arrays
        bnd = Array{Float64,3}(undef,nlayers,nlat,nlon)

        # Open data files
        bndfile = open(joinpath(resourcepath,"crust1","crust1.bnds"), "r")

        # Read data files into array
        for j=1:nlat
           for i=1:nlon
              bnd[:,j,i] = delim_string_parse(readline(bndfile), ' ', Float64, merge=true)
          end
        end

        # Close data files
        close(bndfile)

        # Avoid edge cases at lat = -90.0, lon = 180.0
        ilon = mod.(lon .+ 180, 360) .- 180
        ilat = max.(lat, -90+1e-9)

        # Convert lat and lon to index
        ilat = 91 .- ceil.(Int,ilat)
        ilon = 181 .+ floor.(Int,ilon)

        # Allocate output arrays
        baseout = Array{Float64}(undef,size(lat))

        # Fill output arrays
        for j=1:length(lat)
            if isnan(lat[j]) || isnan(lon[j]) || lat[j] > 90 || lat[j] < -90 || lon[j] > 180 || lon[j] < -180
                baseout[j] = NaN
            else
                baseout[j] = bnd[layer+1,ilat[j],ilon[j]]
            end
        end

        # The end
        return baseout
    end
    export find_crust1_base

## --- End of File

ext: jl


text:
struct BSONWriter{D <: DenseVector{UInt8}}
    dst::D
    offset::Int

    function BSONWriter(dst::D) where D <: DenseVector{UInt8}
        offset = length(dst)
        resize!(dst, offset + 4)
        new{D}(dst, offset)
    end
end

function Base.close(writer::BSONWriter)
    dst = writer.dst
    push!(dst, 0x0)
    p = pointer(dst) + writer.offset
    GC.@preserve dst unsafe_store!(Ptr{Int32}(p), length(dst) - writer.offset)
    nothing
end

@inline wire_size_(::Type{T}) where T = missing
@inline wire_size_(::Type{T}) where T <: Union{
    Int32, Int64, Float64, Dec128, DateTime, BSONTimestamp, BSONObjectId
} = sizeof(T)
@inline wire_size_(::Type{T}) where T <: Union{Nothing, BSONMinKey, BSONMaxKey} = 0
@inline wire_size_(::Type{Bool}) = 1
@inline wire_size_(::Type{UUID}) = 21

@inline wire_size_(x) = sizeof(x)
@inline wire_size_(x::Nothing) = 0
@inline wire_size_(x::String) = sizeof(x) + 5
@inline wire_size_(x::BSONCode) = sizeof(x.code) + 5
@inline wire_size_(x::BSONSymbol) = sizeof(x.value) + 5
@inline wire_size_(x::UUID) = 21
@inline wire_size_(x::BSONUUIDOld) = 21
@inline wire_size_(x::Union{BSONBinary, UnsafeBSONBinary}) = 5 + length(x.data)
@inline wire_size_(x::BSONRegex) = 2 + sizeof(x.pattern) + sizeof(x.options)
@inline wire_size_(x::BSONMinKey) = 0
@inline wire_size_(x::BSONMaxKey) = 0
@inline wire_size_(x::BSONUndefined) = 0
@inline wire_size_(x::BSONDBPointer) = 17 + sizeof(x.collection)

@inline wire_store_(p::Ptr{UInt8}, x::T) where T = unsafe_store!(Ptr{T}(p), htol(x))
@inline wire_store_(::Ptr{UInt8}, ::Nothing) = nothing
@inline wire_store_(p::Ptr{UInt8}, x::Bool) = unsafe_store!(p, UInt8(x))
@inline wire_store_(p::Ptr{UInt8}, x::DateTime) = wire_store_(p, Dates.value(x) - Dates.UNIXEPOCH)
@inline wire_store_(p::Ptr{UInt8}, x::BSONTimestamp) = wire_store_(p, (x.counter % Int64) | ((x.time % Int64) << 32))
@inline wire_store_(p::Ptr{UInt8}, x::BSONObjectId) = unsafe_store!(Ptr{BSONObjectId}(p), x)
@inline wire_store_(::Ptr{UInt8}, ::BSONMinKey) = nothing
@inline wire_store_(::Ptr{UInt8}, ::BSONMaxKey) = nothing
@inline wire_store_(::Ptr{UInt8}, ::BSONUndefined) = nothing

@inline function wire_store_(p::Ptr{UInt8}, x::String)
    unsafe_store!(Ptr{Int32}(p), (sizeof(x) + 1) % Int32)
    GC.@preserve x unsafe_copyto!(p + 4, pointer(x), sizeof(x))
    unsafe_store!(p + 4 + sizeof(x), 0x0)
end

@inline wire_store_(p::Ptr{UInt8}, x::BSONCode) = wire_store_(p, x.code)

@inline wire_store_(p::Ptr{UInt8}, x::BSONSymbol) = wire_store_(p, x.value)

@inline function wire_store_(p::Ptr{UInt8}, x::UUID)
    unsafe_store!(Ptr{Int32}(p), Int32(16))
    unsafe_store!(p + 4, BSON_SUBTYPE_UUID)
    unsafe_store!(Ptr{UUID}(p + 5), x)
end

@inline function wire_store_(p::Ptr{UInt8}, x::BSONUUIDOld)
    unsafe_store!(Ptr{Int32}(p), Int32(16))
    unsafe_store!(p + 4, BSON_SUBTYPE_UUID_OLD)
    unsafe_store!(Ptr{UUID}(p + 5), x.value)
end

@inline function wire_store_(p::Ptr{UInt8}, x::Union{BSONBinary, UnsafeBSONBinary})
    unsafe_store!(Ptr{Int32}(p), length(x.data) % Int32)
    unsafe_store!(p + 4, x.subtype)
    GC.@preserve x unsafe_copyto!(p + 5, pointer(x.data), length(x.data))
end

@inline function wire_store_(p::Ptr{UInt8}, x::BSONRegex)
    GC.@preserve x unsafe_copyto!(p, pointer(x.pattern), sizeof(x.pattern))
    unsafe_store!(p + sizeof(x.pattern), 0x0)
    GC.@preserve x unsafe_copyto!(p + sizeof(x.pattern) + 1, pointer(x.options), sizeof(x.options))
    unsafe_store!(p + sizeof(x.pattern) + sizeof(x.options) + 1 , 0x0)
end

@inline function wire_store_(p::Ptr{UInt8}, x::BSONDBPointer)
    wire_store_(p, x.collection)
    unsafe_store!(Ptr{BSONObjectId}(p + 5 + sizeof(x.collection)), x.ref)
end

@inline len_(x::AbstractString) = sizeof(x)
@inline len_(x::Symbol) = ccall(:strlen, Csize_t, (Cstring,), Base.unsafe_convert(Ptr{UInt8}, x)) % Int

@inline function write_header_(dst::DenseVector{UInt8}, t::UInt8, name::Union{String, Symbol}, value_size::Integer)
    name_len = len_(name)
    offset = length(dst)
    resize!(dst, offset + name_len + value_size + 2)
    GC.@preserve name dst begin
        p = pointer(dst) + offset
        unsafe_store!(p, t)
        ccall(
            :memcpy,
            Cvoid,
            (Ptr{UInt8}, Ptr{UInt8}, Csize_t),
            p + 1, Base.unsafe_convert(Ptr{UInt8}, name), name_len % Csize_t
        )
        unsafe_store!(p + name_len + 1, 0x0)
    end
    offset + name_len + 2
end

function write_field_(writer::BSONWriter, value::T, name::Union{String, Symbol}) where T <: ValueField
    dst = writer.dst
    offset = write_header_(dst, bson_type_(T), name, wire_size_(value))
    p = pointer(dst) + offset
    GC.@preserve dst wire_store_(p, value)
    nothing
end

function write_field_(writer::BSONWriter, value::BSONCodeWithScope, name::Union{String, Symbol})
    dst = writer.dst
    offset = write_header_(dst, BSON_TYPE_CODE_WITH_SCOPE, name, wire_size_(value.code) + 4)
    GC.@preserve dst begin
        p = pointer(dst) + offset
        wire_store_(p + 4, value.code)
        mappings_writer = BSONWriter(dst)
        mappings_writer[] = value.mappings
        close(mappings_writer)
        p = pointer(dst) + offset
        unsafe_store!(Ptr{Int32}(p), length(dst) - offset)
    end
    nothing
end

function write_field_(writer::BSONWriter, generator::Function, name::Union{String, Symbol})
    dst = writer.dst
    write_header_(dst, BSON_TYPE_DOCUMENT, name, 0)
    element_writer = BSONWriter(dst)
    generator(element_writer)
    close(element_writer)
end

const SMALL_INDEX_STRINGS = [string(i) for i in 0:99]

function write_field_(writer::BSONWriter, values::Union{AbstractVector, Base.Generator}, name::Union{String, Symbol})
    dst = writer.dst
    write_header_(dst, BSON_TYPE_ARRAY, name, 0)
    element_writer = BSONWriter(dst)
    element_writer[] = values
    close(element_writer)
end

function Base.setindex!(writer::BSONWriter, values::Union{AbstractVector, Base.Generator})
    for (i, x) in enumerate(values)
        is = i <= length(SMALL_INDEX_STRINGS) ? SMALL_INDEX_STRINGS[i] : string(i - 1)
        writer[is] = x
    end
    nothing
end

@inline function write_field_(writer::BSONWriter, value::T, name::Union{String, Symbol}) where T
    if isstructtype(T)
        write_field_(writer, field_writer -> field_writer[] = value, name)
    else
        throw(ArgumentError("Unsupported type $T"))
    end
end

@inline function Base.setindex!(writer::BSONWriter, value::T, name::Union{String, Symbol}) where T
    RT = bson_representation_type(T)
    if RT != T
        write_field_(writer, bson_representation_convert(RT, value), name)
    else
        write_field_(writer, value, name)
    end
end

function Base.setindex!(writer::BSONWriter, fields::AbstractDict{<:Union{String, Symbol}})
    for (key, value) in fields
        writer[key] = value
    end
    nothing
end

@inline function Base.setindex!(writer::BSONWriter, field::Pair{String})
    writer[field.first] = field.second
    nothing
end

function Base.setindex!(writer::BSONWriter, fields::Tuple{Vararg{Pair{String, T} where T}})
    for (key, value) in fields
        writer[key] = value
    end
    nothing
end

@generated function Base.setindex!(writer::BSONWriter, fields::T) where T <: NamedTuple
    e = Expr(:block)
    for fn in fieldnames(T)
        fns = "$fn"
        push!(e.args, :(writer[$fns] = fields.$fn))
    end
    push!(e.args, nothing)
    e
end

@inline @generated function bson_write_simple(writer::BSONWriter, value::T) where T
    fieldcount(T) == 0 && return nothing
    totalsize = sum(wire_size_, fieldtypes(T)) + sum(sizeof, fieldnames(T)) + fieldcount(T) * 2
    if ismissing(totalsize)
        e = Expr(:block)
        for fn in fieldnames(T)
            fns = string(fn)
            push!(e.args, :(writer[$fns] = value.$fn))
        end
        e
    else
        e = Expr(:block)
        curoffset = 0
        for (ft, fn) in zip(fieldtypes(T), fieldnames(T))
            push!(e.args, :(unsafe_store!(p + $curoffset, $(bson_type_(ft)))))
            curoffset += 1
            fns = string(fn)
            fnl = sizeof(fns)
            push!(e.args, :(ccall(:memcpy, Cvoid, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), p + $curoffset, pointer($fns), $fnl)))
            curoffset += fnl
            push!(e.args, :(unsafe_store!(p + $curoffset, 0x0)))
            curoffset += 1
            push!(e.args, :(wire_store_(p + $curoffset, value.$fn)))
            curoffset += wire_size_(ft)
        end
        quote
            dst = writer.dst
            offset = length(dst)
            resize!(dst, offset + $totalsize)
            GC.@preserve dst begin
                p = pointer(dst) + offset
                $e
            end
        end
    end
end

@inline function bson_write_structtype(writer::BSONWriter, value)
    StructTypes.foreachfield(value) do i, name, FT, value
        writer[name] = value
    end
end

@inline function bson_write(writer::BSONWriter, value::T) where T
    v = bson_schema_version(T)
    if v !== nothing
        writer[bson_schema_version_field(T)] = v
    end
    if bson_simple(T)
        bson_write_simple(writer, value)
    else
        bson_write_structtype(writer, value)
    end
end

@inline function Base.setindex!(writer::BSONWriter, value::T) where T
    bson_write(writer::BSONWriter, value)
end


text:
export data

abstract type TensorStorage end

#
# Generic ITensor storage functions
#

storage_randn!(S::TensorStorage) = randn!(data(S))
storage_norm(S::TensorStorage) = norm(data(S))
storage_conj(S::T) where {T<:TensorStorage}= T(conj(data(S)))


text:
# Performing and plotting measurements during the simulation
using Statistics: mean
using RollingFunctions: rollmean, rolling
using Printf: @sprintf
using LaTeXStrings
using StatsPlots: plot, plot!, density, RGB, ylims, ylims!, savefig
# using StatsBase: fit, Histogram
using Dates: DateFormat, format, now
using DataFrames: DataFrame
using CSV

# Dummy function
null(args...) = nothing

# Compute translational kinetic energy
function avg_kinetic_energy(velocities, m)
    return 0.5 * m * mean( sum(velocities .^2, dims = 1) )
end

# Get a time-save string
filetime_fmt = DateFormat("yyyy-mm-dd-HH-MM-SS")
filetime() = format(now(), filetime_fmt)

# TODO: More versatile / customisable sensing

# Measure global quantities
struct GlobalSensor
    time :: Vector{Float64}
    ke :: Vector{Float64} # Per atom
    pe :: Vector{Float64} # Per atom
    F :: Vector{Float64}
    Nt :: Vector{Int64}
    cand :: Vector{Int64} # Total
    coll :: Vector{Int64} # Total
    Vsim :: Vector{Float64}
    n0 :: Vector{Float64} # Peak density

    function GlobalSensor(args...)
        # Check that all arrays have the same length
        if all(a -> length(a) == length(args[1]), args)
            return new(args...)
        else
            throw(ArgumentError("Sensor arrays must be equally sized."))
        end
    end
end

# Blank sensor
GlobalSensor() = GlobalSensor([zeros(0) for _ in 1:9]...)

# TODO: GridSensor

# Overload measurement
function measurer(sensor::GlobalSensor, p = 0, ωx = 0, ωy = 0, ωz = 0, centre = [0,0,0])
    ω_x, ω_y, ω_z = time_parametrize(ωx, ωy, ωz)
    centre = time_parametrize(centre)
    return (args...) -> measure(sensor, args..., p = p,
        ωx = ω_x, ωy = ω_y, ωz = ω_z, centre = centre)
end

# Save to CSV
function savecsv(sensor::GlobalSensor, filename)
    CSV.write(filename, DataFrame(
        time = sensor.time,
        ke = sensor.ke,
        pe = sensor.pe,
        F = sensor.F,
        Nt = sensor.Nt,
        cand = sensor.cand,
        coll = sensor.coll,
        Vsim = sensor.Vsim,
        n0 = sensor.n0
    ))
    return nothing
end

# Load from CSV
function loadsensor(filename)
    df = CSV.read(filename, DataFrame)
    return GlobalSensor(df.time, df.ke, df.pe, df.F, df.Nt,
                        df.cand, df.coll, df.Vsim, df.n0)
end

function trapped_cloud(cloud, conditions, T, ωx, ωy, ωz, p, centre)

    positions = view(positions, :, 1:N)
    velocities = view(velocities, :, 1:N)

    return positions, velocities
end

# Perform measurements
function measure(sensor::GlobalSensor, cloud, conditions, cand_count, coll_count, t, n0, V;
                 p = 0, ωx = (t) -> 0, ωy = (t) -> 0, ωz = (t) -> 0, centre = (t) -> [0,0,0])

    Nt = cloud.Nt
    positions = view(cloud.positions, :, 1:Nt)
    velocities = view(cloud.velocities, :, 1:Nt)

    species = conditions.species
    m = species.m

    # Update cloud potentials
    conditions.potential(positions, species, t, cloud.potens)

    # Current temperature for ellipsoid sizing
    T = (length(sensor.ke) > 0) ? 2 * last(sensor.ke) / (3 * kB) : Inf
    N = 0
    bound = - 2 * log(p)

    σx² = kB * T / (m * ωx(t)^2)
    σy² = kB * T / (m * ωy(t)^2)
    σz² = kB * T / (m * ωz(t)^2)

    cx, cy, cz = centre(t)

    pe_tot = 0.0
    ke_tot = 0.0

    for atom in 1:Nt
        x = cloud.positions[1, atom] - cx
        y = cloud.positions[2, atom] - cy
        z = cloud.positions[3, atom] - cz
        if (x^2 / σx² + y^2 / σy² + z^2 / σz²) < bound
            N += 1
            vx, vy, vz = view(velocities, :, atom)
            ke_tot += 0.5 * m * (vx^2 + vy^2 + vz^2)
            pe_tot += cloud.potens[atom]
        end
    end

    push!(sensor.time, t)
    push!(sensor.ke, ke_tot/N)
    push!(sensor.pe, pe_tot/N)
    push!(sensor.F, cloud.F)
    push!(sensor.Nt, Nt)
    push!(sensor.cand, cand_count)
    push!(sensor.coll, coll_count)
    push!(sensor.n0, n0)
    push!(sensor.Vsim, V)
end

# PLOTTING + ANALYSIS

# Convert a window time to window size
function window_time_size(tseries, window_time)
    return ceil(Int64, length(tseries) * window_time / last(tseries))
end

function temperature_data(sensor, window_size)
    instant_temp = 2 * sensor.ke / (3 * kB)
    rolling_temp = rollmean(instant_temp, window_size)

    final_temp = last(rolling_temp)

    return final_temp, instant_temp, rolling_temp
end

# Average temperature over time
function plot_temperature(sensor, window_time = 1e-2)
    window_size = window_time_size(sensor.time, window_time)
    T_final, instant_temp, rolling_temp = temperature_data(sensor, window_size)
    rolling_time = rollmean(sensor.time, window_size)

    linecolor = RGB(0.1, 0.1, 0.7)

    # Plot instantaneous
    plt = plot(sensor.time, instant_temp,
        title = "Average temperature (Final: $(@sprintf("%.3g", T_final)) K)",
        xlabel = "Time (s)",
        ylabel = "Temperature (K)",
        label = false,
        ylims = (0, Inf),
        linealpha = 0.5,
        linecolor = linecolor,
        dpi = 300)
    # Plot rolling
    plot!(plt, rolling_time, rolling_temp, linecolor = linecolor, label = false)

    return plt, T_final
end

# Total number of atoms over time
function plot_number(sensor)
    Nt = sensor.Nt
    Np = Nt .* sensor.F
    plt = plot(sensor.time, Np,
        title = "Total number of atoms",
        xlabel = "Time (s)",
        ylabel = "Number",
        label = false,
        ylims = (0, Inf),
        dpi = 300)
    return plt
end

# Average kinetic, potential and total energy per atom
function plot_energy(sensor, window_time = 1e-2)
    window_size = window_time_size(sensor.time, window_time)
    rolling_time = rollmean(sensor.time, window_size)

    # Final temp
    T, _, _ = temperature_data(sensor, window_size)

    instant_ke = sensor.ke / (kB * T) # Units kT (final)
    instant_pe = sensor.pe / (kB * T)

    # Rolling and total energy
    rolling_ke, rolling_pe = rollmean.([instant_ke, instant_pe], window_size)
    instant_e, rolling_e = instant_ke + instant_pe, rolling_ke + rolling_pe

    E = last(rolling_e) * kB * T # Final energy (J/atom)

    linecolors = hcat(RGB(0.1059, 0.6196, 0.4667),
                      RGB(0.851, 0.373, 0),
                      RGB(0.459, 0.439, 0.702))

    plt = plot(sensor.time, [instant_ke, instant_pe, instant_e],
        title="Average energy per atom (Final: $(@sprintf("%.3e", E)) J)",
        xlabel="Time (s)",
        ylabel=L"\textrm{Energy\:\:}(k_B T_f)",
        label=false,
        #ylims = (0, Inf),
        linecolor=linecolors,
        linealpha=0.5,
        dpi = 300)
    plot!(rolling_time, [rolling_ke, rolling_pe, rolling_e],
        label=["Kinetic" "Potential" "Total"],
        linecolor=linecolors)

    return plt
end

# Get speeds
function speeds(cloud)
    Nt = cloud.Nt
    velocities = view(cloud.velocities, :, 1:Nt)
    speeds = vec(sqrt.(sum(velocities .* velocities, dims=1)))
    return speeds
end

# Speed distribution
function plot_speed(cloud)
    plt = density(speeds(cloud),
        title="Speed distribution",
        label="Simulation",
        xlabel=L"\textrm{Speed\:\:}(ms^{-1})",
        ylabel="Probability density",
        dpi = 300)

    return plt
end

# Collision rate per atom
function plot_collrate(sensor, window_time = 1e-2)
    window_size = window_time_size(sensor.time, window_time)
    # Collision rates
    timesteps = vcat(first(sensor.time), diff(sensor.time))

    colls_per_atom = 2 * sensor.coll ./ sensor.Nt
    instant_collrate = colls_per_atom ./ timesteps

    rolling_time = rollmean(sensor.time, window_size)

    rolling_timesteps = (
        sensor.time[window_size:end] - sensor.time[1:end - window_size + 1]
    )
    rolling_collrate = (
        rolling(sum, colls_per_atom, window_size) ./ rolling_timesteps
    )

    final_rate = last(rolling_collrate)
    linecolor = RGB(0.851, 0.373, 0)
    # Note: per-atom collision rates of test and real particles are equal.

    plt = plot(rolling_time, rolling_collrate,
        title="""Average collision rate per atom
        (Final: $(@sprintf("%.3g", final_rate)) Hz)""",
        xlabel="Time (s)",
        ylabel="Rate (Hz)",
        label = false,
        #yaxis = :log10,
        #yminorticks = 10,
        linecolor=linecolor,
        dpi = 300)
    plot!(sensor.time, instant_collrate, label = false,
          linecolor = linecolor, linealpha = 0.5)

    return plt
end

# Density
function plot_density(sensor, window_time = 1e-2)
    window_size = window_time_size(sensor.time, window_time)
    rolling_time = rollmean(sensor.time, window_size)

    #T, _, _ = temperature_data(sensor, window_size)

    instant_V = sensor.Vsim
    rolling_V = rollmean(instant_V, window_size)

    #V = last(rolling_V)
    instant_n0 = sensor.n0 .* sensor.F
    rolling_n0 = rollmean(instant_n0, window_size)

    rolling_F = rollmean(sensor.F, window_size)
    rolling_Nt = rollmean(sensor.Nt, window_size)

    instant_n_avg = sensor.F .* sensor.Nt ./ instant_V
    rolling_n_avg = rolling_F .* rolling_Nt ./ rolling_V

    linecolors = hcat(RGB(0.1059, 0.6196, 0.4667),
                      RGB(0.851, 0.373, 0))

    plt = plot(sensor.time, [instant_n0, instant_n_avg],
        title="Number density",
        xlabel="Time (s)",
        ylabel=L"n (# $m^{-3}$)",
        label=false,
        linecolor=linecolors,
        linealpha=0.5,
        dpi = 300)
    plot!(rolling_time, [rolling_n0, rolling_n_avg],
        label=["Peak" "Average"],
        linecolor=linecolors)

    return plt
end

# Phase space density (WIP)
function plot_psd(sensor, m, window_time = 1e-2)
    window_size = window_time_size(sensor.time, window_time)
    rolling_time = rollmean(sensor.time, window_size)

    _, _, rolling_T = temperature_data(sensor, window_size)
    rolling_F = rollmean(sensor.F, window_size)
    rolling_n0 = rolling_F .* rollmean(sensor.n0, window_size)
    hbar = 6.626e-34 / (2 * π)
    rolling_psd = rolling_n0 .* (sqrt.(2*π*hbar^2 ./ (m * kB .* rolling_T) ) .^3)

    linecolor = RGB(0.1059, 0.6196, 0.4667)

    plt = plot(rolling_time, rolling_psd,
        title="Phase space density",
        xlabel="Time (s)",
        ylabel=L"PSD",
        label=false,
        linecolor=linecolor,
        dpi = 300)

    return plt
end


text:
using ContactStructures

plotlyjs()
# define a contact form for the standard contact structure
x = (x,y,z) -> x
std_form = dz + x * dy
plot_along_xy_grid(-2:0.5:2,-2:0.5:2,0,std_form;  markersize = 2.5,ε = 0.25)
plot!(xlims=(-3, 3), ylims=(-3,3), zlims=(-3,3),legend=false)


text:
################################################################################

"""

    writeHMM(modelHMM::Dict{S, Tuple{Array{T, 1}, Array{Array{U, 1}, 1}}}, shParams::Dict) where S <: String where T <: Int64 where U <: Float64

# Description
Write hidden markov model states and traceback wrapper.

See also: [`writePerformance`](@ref)
"""
function writeHMM(modelHMM::Dict{S, Tuple{Array{T, 1}, Array{Array{U, 1}, 1}}}, shParams::Dict) where S <: String where T <: Int64 where U <: Float64
  if haskey(shParams, "outdir") && haskey(shParams, "file")
    return writeHMM( string(shParams["outdir"], replace(shParams["file"], ".edf" => "_")), modelHMM )
  else
    @error "Variables are not defined in dictionary"
  end
end

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

"""

    writeHMM(filePrefix::S, modelHMM::Dict{S, Tuple{Array{T, 1}, Array{Array{U, 1}, 1}}}) where S <: String where T <: Int64 where U <: Float64

# Description
Write hidden markov model states and traceback wrapper.

See also: [`writePerformance`](@ref)
"""
function writeHMM(filePrefix::S, modelHMM::Dict{S, Tuple{Array{T, 1}, Array{Array{U, 1}, 1}}}) where S <: String where T <: Int64 where U <: Float64
  for (κ, υ) ∈ modelHMM
    filename = string( filePrefix, string(κ))
    writeHMM( string(filename, "_states", ".csv"), υ[1], κ)
    writeHMM( string(filename, "_traceb", ".csv"), υ[2])
  end
end

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

"""

    writeHMM(filename::S, statesHMM::Array{T, 1}, channel::S) where S <: String where T <: Number

# Description
Write hidden markov model states wrapper.

See also: [`writePerformance`](@ref)
"""
function writeHMM(filename::S, statesHMM::Array{T, 1}, channel::S) where S <: String where T <: Number
  CSV.write(filename, shiftHMM(statesHMM, channel))
end

"""

    writeHMM(filename::S, tracebHMM::Array{Array{T, 1}, 1}) where S <: String where T <: Number

# Description
Write hidden markov model traceback wrapper.

See also: [`writePerformance`](@ref)
"""
function writeHMM(filename::S, tracebHMM::Array{Array{T, 1}, 1}) where S <: String where T <: Number
  CSV.write(filename, shiftHMM(tracebHMM))
end

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

"""

    writePerformance(performanceDc::Dict{S, Array{T, 2}}) where S <: String where T <: Number

# Description
Transform model performance to table to write.

See also: [`writeHMM`](@ref)
"""
function writePerformance(performanceDc::Dict{S, Array{T, 2}}) where S <: String where T <: Number
  outAr = Array{Any, 2}(undef, length(performanceDc) + 1, 3)
  outAr[1, :] .= ["Electrode", "Sensitivity", "Specificity"]
  for (ζ, (κ, υ)) ∈ enumerate(performanceDc)
    outAr[ζ + 1, :] = [κ υ]
  end
  return outAr
end

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

"""

    writePerformance(filename::S, performanceDc::Dict{S, Array{T, 2}}, delim::S = ",") where S <: String where T <: Number

# Description
Write model performance to CSV file.

See also: [`writeHMM`](@ref)
"""
function writePerformance(filename::S, performanceDc::Dict{S, Array{T, 2}}, delim::S = ",") where S <: String where T <: Number
  writedlm(filename, writePerformance(performanceDc), delim)
end

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

"reorder hidden markov model states into table to write"
function shiftHMM(statesHMM::Array{T, 1}, channel::S) where S <: String where T <: Number
  return Tables.table(reshape(statesHMM, (length(statesHMM), 1)), header = [channel])
end

"reorder hidden markov model traceback vectors into table to write"
function shiftHMM(tracebHMM::Array{Array{T, 1}, 1}) where T <: Number
  outMt = Array{Float64, 2}(undef, length(tracebHMM[1]), length(tracebHMM))
  for ι ∈ 1:length(tracebHMM)
    outMt[:, ι] = tracebHMM[ι]
  end
  return Tables.table(outMt, header = [Symbol("S$i") for i = 1:size(outMt, 2)])
end

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

max_stars_repo_name: DanielRivasMD/MindReader

max_issues_repo_name: DanielRivasMD/MindReader

max_forks_repo_name: DanielRivasMD/MindReader


text:
# SPDX-License-Identifier: X11
# 2020-08-31
# Multiplication 3 (300pt)

((A, B) -> println(trunc(Int64, A * B)))(
  parse(Int64, readuntil(stdin, " ")),
  rationalize(parse(Float64, readline())))

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Functions related to the conversion from DCM to quaternion.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export dcm_to_quat

"""
    dcm_to_quat(dcm::DCM)

Convert the `dcm` to a quaternion.

The type of the quaternion will be automatically selected by the constructor
[`Quaternion`](@ref) to avoid `InexactError`.

# Remarks

By convention, the real part of the quaternion will always be positive.
Moreover, the function does not check if `dcm` is a valid direction cosine
matrix. This must be handle by the user.

This algorithm was obtained from [^1].

# Example

```jldoctest
julia> dcm = angle_to_dcm(pi / 2, 0.0, 0.0, :XYZ);

julia> q = dcm_to_quat(dcm)
Quaternion{Float64}:
  + 0.707107 + 0.707107⋅i + 0.0⋅j + 0.0⋅k
```

# References

[^1]: http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/
"""
function dcm_to_quat(dcm::DCM)
    if  tr(dcm) > 0
        # f = 4 * q0
        f = 2sqrt(tr(dcm) + 1)

        return Quaternion(
            f / 4,
            (dcm[2, 3] - dcm[3, 2]) / f,
            (dcm[3, 1] - dcm[1, 3]) / f,
            (dcm[1, 2] - dcm[2, 1]) / f
        )
    elseif (dcm[1, 1] > dcm[2, 2]) && (dcm[1, 1] > dcm[3, 3])
        # f = 4 * q1
        f = 2sqrt(1 + dcm[1, 1] - dcm[2, 2] - dcm[3, 3])

        # Real part.
        q0 = (dcm[2, 3] - dcm[3, 2]) / f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(
            s * q0,
            s * f / 4,
            s * (dcm[1, 2] + dcm[2, 1]) / f,
            s * (dcm[3, 1] + dcm[1, 3]) / f
        )
    elseif (dcm[2, 2] > dcm[3, 3])
        # f = 4 * q2
        f = 2sqrt(1 + dcm[2, 2] - dcm[1, 1] - dcm[3, 3])

        # Real part.
        q0 = (dcm[3, 1] - dcm[1, 3]) / f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(
            s * q0,
            s * (dcm[1, 2] + dcm[2, 1]) / f,
            s * f / 4,
            s * (dcm[3, 2] + dcm[2, 3]) / f
        )
    else
        # f = 4 * q3
        f = 2sqrt(1 + dcm[3, 3] - dcm[1, 1] - dcm[2, 2])

        # Real part.
        q0 = (dcm[1, 2] - dcm[2, 1]) / f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(
            s * q0,
            s * (dcm[1, 3] + dcm[3, 1]) / f,
            s * (dcm[2, 3] + dcm[3, 2]) / f,
            s * f / 4
        )
    end
end


text:
function extract_equation_data(obj::FieldEquationsMapper, arg0::jint, arg1::Vector{RealFieldElement})
    return jcall(obj, "extractEquationData", Vector{RealFieldElement}, (jint, Vector{RealFieldElement}), arg0, arg1)
end

function get_number_of_equations(obj::FieldEquationsMapper)
    return jcall(obj, "getNumberOfEquations", jint, ())
end

function get_total_dimension(obj::FieldEquationsMapper)
    return jcall(obj, "getTotalDimension", jint, ())
end

function insert_equation_data(obj::FieldEquationsMapper, arg0::jint, arg1::Vector{RealFieldElement}, arg2::Vector{RealFieldElement})
    return jcall(obj, "insertEquationData", void, (jint, Vector{RealFieldElement}, Vector{RealFieldElement}), arg0, arg1, arg2)
end

function map_state_and_derivative(obj::FieldEquationsMapper, arg0::RealFieldElement, arg1::Vector{RealFieldElement}, arg2::Vector{RealFieldElement})
    return jcall(obj, "mapStateAndDerivative", FieldODEStateAndDerivative, (RealFieldElement, Vector{RealFieldElement}, Vector{RealFieldElement}), arg0, arg1, arg2)
end


text:
using DataStructures
using NPZ
using Statistics
using VCFTools

struct Segment{T}
    snp_start::T
    snp_end::T
    sample_start::T
    sample_end::T
end

"""
    build_prefix_and_divergence_arrays(H::Array{Int8, 2})

Compute the prefix and divergence arrays (n_samples x n_snps+1) for the
Positional Burrows-Wheeler Transform on the haplotype matrix H (n_samples x n_snps).
"""
function build_prefix_and_divergence_arrays(H::Array{Int8, 2})
    n_samples, n_snps = size(H)
    u = v = p = q = index = match_start = one(Int32)
    ppa_buffer = Vector{Int32}(undef, n_samples)
    div_buffer = Vector{Int32}(undef, n_samples)
    ppa = Matrix{Int32}(undef, n_samples, n_snps + 1)
    div = Matrix{Int32}(undef, n_samples, n_snps + 1)
    ppa[:, 1] = 1:1:n_samples
    div[:, 1] = ones(Int32, n_samples)
    @inbounds for snp_index in 1:n_snps
        u = v = 1
        p = q = snp_index + 1
        for sample_index in 1:n_samples
            index = ppa[sample_index, snp_index]
            match_start = div[sample_index, snp_index]
            if match_start > p
                p = match_start
            end
            if match_start > q
                q = match_start
            end
            if H[index, snp_index] == 0
                ppa[u, snp_index + 1] = index
                div[u, snp_index + 1] = p
                u += 1
                p = 0
            else
                ppa_buffer[v] = index
                div_buffer[v] = q
                v += 1
                q = 0
            end
        end
        for sample_index in u:n_samples
            ppa[sample_index, snp_index + 1] = ppa_buffer[sample_index - u + 1]
            div[sample_index, snp_index + 1] = div_buffer[sample_index - u + 1]
        end
    end
    return ppa, div
end

function build_reverse_prefix_array(ppa::Array{Int32, 2})
    M, N = size(ppa)
    reverse_ppa = similar(ppa)
    Threads.@threads for j in 1:N
        @inbounds for i in 1:M
            reverse_ppa[ppa[i, j], j] = i
        end
    end
    return reverse_ppa
end

function build_segments(div::Array{Int32, 2}, ppa::Array{Int32, 2}, reverse_ppa::Array{Int32, 2})
    n_samples = size(div, 1)
    n_snps = size(div, 2) - 1
    segments_by_thread = [Vector{Segment{Int32}}() for t in 1:min(Threads.nthreads(), n_snps)]
    Threads.@threads for snp_index in 1:n_snps
        segments = segments_by_thread[Threads.threadid()]
        last_segment = Segment{Int32}(snp_index, snp_index, convert(Int32, n_samples + 1), convert(Int32, n_samples + 1))
        split_index = convert(Int32, n_samples + 1)
        @inbounds for sample_index in n_samples:-1:1
            match_start = div[sample_index, snp_index + 1]
            if match_start == snp_index
                continue
            elseif last_segment.snp_start == match_start && last_segment.sample_start < sample_index < last_segment.sample_end
                continue
            elseif match_start == snp_index + 1
                segment = Segment{Int32}(snp_index, snp_index, convert(Int32, sample_index), convert(Int32, split_index - 1))
                push!(segments, segment)
                last_segment = segment
                split_index = sample_index
            elseif snp_index == n_snps || match_start != div[reverse_ppa[ppa[sample_index, snp_index + 1], snp_index + 2], snp_index + 2]
                sample_start = sample_index
                sample_end = n_samples
                for neighbor_index in sample_index-1:-1:1
                    if match_start < div[neighbor_index, snp_index + 1]
                        sample_start = neighbor_index
                        break
                    end
                end
                for neighbor_index in sample_index+1:n_samples
                    if match_start < div[neighbor_index, snp_index + 1]
                        sample_end = neighbor_index - 1
                        break
                    end
                end
                segment = Segment{Int32}(match_start, snp_index, convert(Int32, sample_start), convert(Int32, sample_end))
                push!(segments, segment)
                last_segment = segment
            end
        end
    end
    return cat(segments_by_thread..., dims=1)
end

function build_segment_overlaps(segments::Vector{Segment{Int32}}, ppa::Array{Int32, 2})
    n_segments = length(segments)
    n_snps = size(ppa, 2) - 1
    segment_index_by_snp = [Vector{Int32}() for snp_index in 1:n_snps]
    segment_overlaps = [Vector{Int32}() for segment_index in 1:n_segments]
    for segment_index in 1:n_segments
        segment = segments[segment_index]
        for snp_index in segment.snp_start:segment.snp_end
            push!(segment_index_by_snp[snp_index], segment_index)
        end
    end
    Threads.@threads for segment_index in 1:n_segments
        segment = segments[segment_index]
        candidate_overlap_set = Set{Int32}()
        original_form_sample_indices = Vector{Int32}()
        @inbounds for sample_index in segment.sample_start:segment.sample_end
            push!(original_form_sample_indices, ppa[sample_index, segment.snp_end + 1])
        end
        @inbounds for snp_index in segment.snp_start:segment.snp_end
            for candidate_index in segment_index_by_snp[snp_index]
                # if segment_index == candidate_index || in(candidate_index, candidate_overlap_set)
                if in(candidate_index, candidate_overlap_set)
                    continue
                end
                push!(candidate_overlap_set, candidate_index)
                candidate_segment = segments[candidate_index]
                for sample_index in candidate_segment.sample_start:candidate_segment.sample_end
                    if in(ppa[sample_index, candidate_segment.snp_end + 1], original_form_sample_indices)
                        push!(segment_overlaps[segment_index], candidate_index)
                        break
                    end
                end
            end
        end
    end
    return segment_overlaps
end

function calculate_scores(coverage::Array{Int8, 2}, ppa::Array{Int32, 2}, segments::Vector{Segment{Int32}}, indices::Vector{Int32})
    scores = Vector{Int32}(undef, length(indices))
    Threads.@threads for (score_index, segment_index) in collect(enumerate(indices))
        segment = segments[segment_index]
        @inbounds original_form_sample_indices = [ppa[sample_index, segment.snp_end + 1] for sample_index in segment.sample_start:segment.sample_end]
        score = zero(Int32)
        for snp_index in segment.snp_start:segment.snp_end
            for original_form_sample_index in original_form_sample_indices
                @inbounds score += coverage[original_form_sample_index, snp_index]
            end
        end
        @inbounds scores[score_index] = score
    end
    return scores
end

function calculate_scores(coverage::Array{Int8, 2}, ppa::Array{Int32, 2}, segments::Vector{Segment{Int32}})
    indices = convert(Vector{Int32}, 1:length(segments) |> collect)
    return calculate_scores(coverage, ppa, segments, indices)
end

function update_coverage!(coverage::Array{Int8, 2}, ppa::Array{Int32, 2}, segments::Vector{Segment{Int32}}, segment_index::Int)
    segment = segments[segment_index]
    original_form_sample_indices = [ppa[sample_index, segment.snp_end + 1] for sample_index in segment.sample_start:segment.sample_end]
    for snp_index in segment.snp_start:segment.snp_end
        for original_form_sample_index in original_form_sample_indices
            @inbounds coverage[original_form_sample_index, snp_index] = 0
        end
    end
end

function update_scores!(scores::MutableBinaryMaxHeap{Int32}, coverage::Array{Int8, 2}, ppa::Array{Int32, 2}, segments::Vector{Segment{Int32}}, indices::Vector{Int32})
    updated_scores = calculate_scores(coverage, ppa, segments, indices)
    for (segment_index, score) in zip(indices, updated_scores)
        update!(scores, convert(Int, segment_index), score)
    end
end

function build_encoding_haplotypes(H::Array{Int8, 2}, ppa::Array{Int32, 2}, encoding_segments::Vector{Segment{Int32}})
    n_segments = length(encoding_segments)
    encoding_haplotypes = Vector{Vector{Int8}}(undef, n_segments)
    Threads.@threads for segment_index in 1:n_segments
        @inbounds segment = encoding_segments[segment_index]
        @inbounds encoding_haplotypes[segment_index] = H[ppa[segment.sample_start, segment.snp_end + 1], segment.snp_start:segment.snp_end]
    end
    return encoding_haplotypes
end

function encode_haplotype_matrix(H::Array{Int8, 2}, ppa::Array{Int32, 2}, encoding_segments::Vector{Segment{Int32}})
    n_samples = size(H, 1)
    n_segments = length(encoding_segments)
    H_encoded = zeros(Int8, n_samples, n_segments)
    Threads.@threads for segment_index in 1:n_segments
        @inbounds segment = encoding_segments[segment_index]
        for sample_index in segment.sample_start:segment.sample_end
            @inbounds H_encoded[ppa[sample_index, segment.snp_end + 1], segment_index] = 1
        end
    end
    return H_encoded
end

function encode_haplotype_matrix(H::Array{Int8, 2}, encoding_positions::Vector{Int32}, encoding_haplotypes::Vector{Vector{Int8}})
    n_samples = size(H, 1)
    n_haplotypes = length(encoding_haplotypes)
    H_encoded = ones(Int8, n_samples, n_haplotypes)
    Threads.@threads for haplotype_index in 1:n_haplotypes
        @inbounds snp_start = encoding_positions[haplotype_index]
        @inbounds haplotype = encoding_haplotypes[haplotype_index]
        for sample_index in 1:n_samples
            for snp_index in 1:length(haplotype)
                @inbounds if H[sample_index, snp_start + snp_index - 1] != haplotype[snp_index]
                    @inbounds H_encoded[sample_index, haplotype_index] = 0
                    break
                end
            end
        end
    end
    return H_encoded
end

function build_encoding(H::Array{Int8, 2}, max_size::Int)
    encoding_segments = Vector{Segment{Int32}}(undef, 0)
    ppa, div = build_prefix_and_divergence_arrays(H)
    reverse_ppa = build_reverse_prefix_array(ppa)
    segments = build_segments(div, ppa, reverse_ppa)
    segment_overlaps = build_segment_overlaps(segments, ppa)
    coverage = ones(Int8, size(H, 1), size(H, 2))
    scores = MutableBinaryMaxHeap{Int32}(calculate_scores(coverage, ppa, segments))
    score, segment_index = top_with_handle(scores)
    while score > 0 && length(encoding_segments) < max_size
        push!(encoding_segments, segments[segment_index])
        update_coverage!(coverage, ppa, segments, segment_index)
        update_indices = [index for index in segment_overlaps[segment_index] if scores[convert(Int, index)] > 0]
        update_scores!(scores, coverage, ppa, segments, update_indices)
        score, segment_index = top_with_handle(scores)
    end
    encoding_positions = [segment.snp_start for segment in encoding_segments]
    encoding_haplotypes = build_encoding_haplotypes(H, ppa, encoding_segments)
    H_encoded = encode_haplotype_matrix(H, ppa, encoding_segments)
    return encoding_positions, encoding_haplotypes, H_encoded
end

function build_encoding(H::Array{Int8, 2})
    return build_encoding(H, length(H))
end

H =  Array{Int8, 2}([0 1 0 1 0 1; 1 1 0 0 0 1; 1 1 1 1 1 1; 0 1 1 1 1 0; 0 0 0 0 0 0; 1 0 0 0 1 0; 1 1 0 0 0 1; 0 1 0 1 1 0])
# H =  Array{Int8, 2}([1 0 0 0; 0 0 1 0; 0 0 1 0; 1 0 1 0])
path = "/media/storage/1000_genomes/GRCh38/variants/chr20/yri.chr20.GRCh38.vcf"
path = "/media/storage/1000_genomes/GRCh38/variants/chr10/LATINO.chr10.phase3_shapeit2_mvncall_integrated_v5a.20130502.genotypes.biallelic.vcf.gz"
path = "/home/rlaboulaye/Work/projects/genome-retrieval/admixed_american_2.vcf"
path = "/home/rlaboulaye/Work/projects/genome-retrieval/ashkenazi.vcf"
H = convert_ht(Int8, path)
# filter out maf <= 1%
H = H[:, Statistics.mean(H, dims=1)[1, :] .> .01]
@time encoding_positions, encoding_haplotypes, H_encoded = build_encoding(H)
npzwrite("ashkenazi.chr10.haplosoup.npy", H_encoded)
npzwrite("ashkenazi.chr10.encoding_positions.npy", encoding_positions)

size(encoding_positions)

ppa, div = build_prefix_and_divergence_arrays(H)
reverse_ppa = build_reverse_prefix_array(ppa)
segments = build_segments(div, ppa, reverse_ppa)
segment_overlaps = build_segment_overlaps(segments, ppa)
coverage = ones(Int8, size(H, 1), size(H, 2))
scores = MutableBinaryMaxHeap{Int32}(calculate_scores(coverage, ppa, segments))

segment_index = 15
segment = segments[segment_index]
H[ppa[:, segment.snp_end + 1], :]
for segment_index in segment_overlaps[segment_index]
    segment = segments[segment_index]
    println(segment)
    println(H[ppa[:, segment.snp_end + 1], 1:segment.snp_end])
end


text:
module TestBasisPursuit
using Test
using LinearAlgebra
using SparseArrays
using CompressedSensing: bp, bp_candes, bp_ard, bpd, bpd_candes, bpd_ard, sparse_data, perturb

n, m = 32, 48
k = 3
A, x, b = sparse_data(n = n, m = m, k = k, rescaled = true)
δ = 1e-2
y = perturb(b, δ/2)

@testset "Basis Pursuit" begin
    # equality constrained l1 minimization
    xl = bp(A, b)
    @test xl.nzind == x.nzind

    xc = bp_candes(A, b)
    @test xc.nzind == x.nzind

    xard = bp_ard(A, b)
    @test xard.nzind == x.nzind
end

@testset "Basis Pursuit Denoising" begin
    xl = bpd(A, y, δ)
    droptol!(xl, 1e-2) # sometimes has spurious coefficients above perturbation level
    @test xl.nzind == x.nzind

    xc = bpd_candes(A, y, δ, maxiter = 3)
    droptol!(xc, 1e-6)
    @test xc.nzind == x.nzind

    xard = bpd_ard(A, y, δ, maxiter = 16)
    droptol!(xard, 1e-6)
    @test xard.nzind == x.nzind
end

using CompressedSensing: ista, fista
@testset "ISTA" begin
    λ = δ/10
    xista = ista(A, y, λ, maxiter = 1024, stepsize = 1e-1)
    # droptol!(xista, δ)
    # @test xista.nzind == x.nzind
    # @test xista.nzind ⊆ x.nzind
    @test norm(A*xista - y) < δ
    # TODO: FISTA
end

end