text:
# Bravais.jl -- Copyright 2015 James R. Garrison and contributors.
# Provided under the MIT License.

# NOTE: Most symbols used below are defined in the accompanying
# documentation.

# FIXME: need other constructors for e.g. OBC

# FIXME: have a way of getting the neighbors of a given site (instead
# of iterating through all of them).  for this, we may actually want
# to double count (or not).

# FIXME: in lattice w/ basis, be sure to error out anytime the "basis
# index" is out of range

module Bravais

using StaticArrays
using LinearAlgebra
using FastClosures

include("delegate.jl")

# Julia currently lacks row-major versions of sub2ind and ind2sub.  We
# want to use these so the ordering for returned tuples is correct;
# see discussion at <https://github.com/JuliaLang/julia/pull/10337#issuecomment-78206799>.
#
# Also, once julia easily supports it we can speed up ind2sub by
# precomputing some "magic numbers" for given lattice dimensions (see
# https://github.com/JuliaLang/julia/issues/8188#issuecomment-56763806).
rowmajor_ind2sub(dims, index) = reverse(Tuple(CartesianIndices(reverse(dims))[index]))

abstract type AbstractSiteNetwork <: AbstractVector{Vector{Int}} end
abstract type AbstractLattice{D} <: AbstractSiteNetwork end

abstract type AbstractBravaisLattice{D} <: AbstractLattice{D} end
abstract type AbstractLatticeWithBasis{D,Dp1} <: AbstractLattice{D} end

abstract type WrappedBravaisLattice{D} <: AbstractBravaisLattice{D} end
abstract type WrappedLatticeWithBasis{D,Dp1} <: AbstractLatticeWithBasis{D,Dp1} end

# NOTE: In the following, D is the number of dimensions, Dsq == D^2,
# and Dp1 == D+1.

struct BravaisLattice{D,Dsq} <: AbstractBravaisLattice{D}
    N_tot::Int  # total number of sites
    N::SVector{D,Int}  # lattice extent in each dimension
    M::SMatrix{D,D,Int,Dsq}  # "M" matrix (will be diagonal for non-helical boundary)
    η::SVector{D,Rational{Int}}  # twist in each direction (rational multiple of 2π)
    a::SMatrix{D,D,Float64,Dsq}  # primitive lattice vectors
    b::SMatrix{D,D,Float64,Dsq}  # reciprocal lattice vectors
    momenta::Vector{SVector{D,Rational{Int}}}
    strides::SVector{D,Int}

    function BravaisLattice{D,Dsq}(N::AbstractVector{Int},
                                   M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                                   η::AbstractVector{Rational{Int}}=zeros(SVector{D,Rational{Int}}),
                                   a::AbstractMatrix{Float64}=SMatrix{D,D}(1.0I)) where {D,Dsq}
        @assert Dsq == D * D

        # check N
        length(N) == D || throw(ArgumentError(""))
        all(d_i > 0 for d_i in N) || throw(ArgumentError(""))
        N_tot = prod(N)

        # check M
        size(M) == (D, D) || throw(ArgumentError(""))
        istril(M) || throw(ArgumentError(""))
        for i in 1:D
            M[i,i] == 0 || M[i,i] == N[i] || throw(ArgumentError(""))
            if M[i,i] != N[i]
                all(s == 0 for s in M[i,:]) || throw(ArgumentError(""))
                all(s == 0 for s in M[:,i]) || throw(ArgumentError(""))
            end
        end

        # check η
        length(η) == D || throw(ArgumentError(""))
        all(0 <= η_i < 1 for η_i in η) || throw(ArgumentError(""))
        all(M[i,i] != 0 || η[i] == 0 for i in 1:D) || throw(ArgumentError(""))

        # check a and generate b
        size(a) == (D, D) || throw(ArgumentError(""))
        b = 2pi * transpose(inv(a))

        # generate allowed momenta
        momenta_range = [s != 0 ? s : 1 for s in diag(M)]
        nmomenta = prod(momenta_range)
        momenta = sizehint!(SVector{D,Rational{Int}}[], nmomenta)
        momentum = zeros(MVector{D,Rational{Int}})
        for idx in 1:nmomenta
            # FIXME: getting idx should be easier, once
            # multidimensional iteration is supported
            ñ = [rowmajor_ind2sub(tuple(momenta_range...), idx)...] .- 1
            for i in 1:D
                if M[i,i] == 0
                    momentum[i] = 0
                else
                    tmp = sum(Rational{Int}[M[i,j] * momentum[j] for j in 1:i-1])
                    momentum[i] = (ñ[i] + η[i] - tmp) // M[i,i]
                end
            end
            push!(momenta, momentum)
        end

        # calculate strides
        strides = zeros(Int, D)
        s = 1
        for i in D:-1:1
            strides[i] = s
            s *= N[i]
        end

        new(N_tot, copy(N), copy(M), copy(η), copy(a), b, momenta, strides)
    end
end

Base.@pure square(x) = x * x

BravaisLattice{D}(args...) where {D} =
    BravaisLattice{D,square(D)}(args...)

BravaisLattice(N::StaticVector{D,Int}, args...) where {D} =
    BravaisLattice{D}(N, args...)

struct LatticeWithBasis{D,Dsq,Dp1} <: AbstractLatticeWithBasis{D,Dp1}
    N_tot::Int
    maxcoords::SVector{Dp1,Int}
    bravaislattice::BravaisLattice{D,Dsq}
    basis::Vector{SVector{D,Float64}}
    strides::SVector{Dp1,Int}

    function LatticeWithBasis{D,Dsq,Dp1}(N::AbstractVector{Int},
                                         M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                                         η::AbstractVector{Rational{Int}}=zeros(SVector{D,Rational{Int}}),
                                         a::AbstractMatrix{Float64}=SMatrix{D,D}(1.0I),
                                         basis::Vector{SVector{D,Float64}}=[zeros(SVector{D,Float64})]) where {D,Dsq,Dp1}
        @assert Dp1 == D + 1
        bravaislattice = BravaisLattice{D,Dsq}(N, M, η, a)

        # check basis
        nbasis = length(basis)
        nbasis > 0 || throw(ArgumentError(""))

        # determine N_tot and maxcoords, now that we know the basis size
        N_tot = prod(N) * nbasis
        maxcoords = vcat(N, [nbasis])

        # calculate strides
        strides = zeros(Int, length(maxcoords))
        strides[end] = 1
        s = nbasis
        for i in D:-1:1
            strides[i] = s
            s *= N[i]
        end

        new(N_tot, maxcoords, bravaislattice, copy(basis), strides)
    end
end

isbravais(::AbstractBravaisLattice) = true
isbravais(::AbstractLatticeWithBasis) = false

"Return the underlying Bravais lattice"
bravais(lattice::AbstractBravaisLattice) = lattice
bravais(lattice::LatticeWithBasis) = lattice.bravaislattice
bravais(lattice::WrappedLatticeWithBasis) = bravais(lattice.lattice)

const LatticeImplUnion{D,Dsq} = Union{BravaisLattice{D,Dsq}, LatticeWithBasis{D,Dsq}}
const WrappedLatticeUnion{D} = Union{WrappedBravaisLattice{D}, WrappedLatticeWithBasis{D}}

_strides(lattice::LatticeImplUnion) = lattice.strides
_strides(lattice::WrappedLatticeUnion) = _strides(lattice.lattice)

maxcoords(lattice::BravaisLattice) = lattice.N
maxcoords(lattice::LatticeWithBasis) = lattice.maxcoords
maxcoords(lattice::WrappedLatticeUnion) = maxcoords(lattice.lattice)

Base.size(lattice::LatticeImplUnion) = (lattice.N_tot,)
Base.size(lattice::WrappedLatticeUnion) = size(lattice.lattice)

function Base.getindex(lattice::Union{AbstractBravaisLattice{Dprime},AbstractLatticeWithBasis{D,Dprime} where D}, index::Integer) where {Dprime}
    @boundscheck checkbounds(lattice, index)
    # Alternatively: return [rowmajor_ind2sub(tuple(maxcoords(lattice)...), index)...] - 1
    strides = _strides(lattice)
    @assert length(strides) == Dprime
    rv = zeros(SVector{Dprime,Int})
    r = index - 1
    for i in 1:Dprime-1
        d, r = divrem(r, strides[i])
        rv = setindex(rv, d, i)
    end
    rv = setindex(rv, r, Dprime)
    return rv
end

function Base.in(site::AbstractVector{Int}, lattice::AbstractLattice)
    mc = maxcoords(lattice)
    length(mc) == length(site) || return false
    for i in 1:length(site)
        0 <= site[i] < mc[i] || return false
    end
    return true
end

EqualTo = Base.Fix2{typeof(isequal)}

function Base.findfirst(p::EqualTo{<:AbstractVector{Int}}, lattice::AbstractLattice)
    let site = p.x
        site ∉ lattice ? nothing : dot(site, _strides(lattice)) + 1
    end
end

import Base: findfirst
@deprecate findfirst(lattice::AbstractLattice, site::AbstractVector{Int}) findfirst(isequal(site), lattice)

function Base.iterate(lattice::Union{AbstractBravaisLattice{Dprime},AbstractLatticeWithBasis{D,Dprime} where D}, site::SVector{Dprime,Int}=zeros(SVector{Dprime,Int}))::Union{Nothing,NTuple{2,SVector{Dprime,Int}}} where {Dprime}
    newsite = site
    mc = maxcoords(lattice)
    # XXX FIXME @assert something # otherwise BoundsError!
    @assert length(mc) == Dprime >= 1
    if newsite[1] == mc[1]
        return nothing
    end
    for i in Dprime:-1:2
        newsite = setindex(newsite, newsite[i] + 1, i)
        if newsite[i] == mc[i]
            newsite = setindex(newsite, 0, i)
        else
            return site, newsite
        end
    end
    newsite = setindex(newsite, newsite[1] + 1, 1)
    return site, newsite
end

# FIXME: can we return readonly array views/proxies of some of the
# following things?  It seems to be difficult if not impossible at the
# moment -- even base library functions like eigfact return mutable
# arrays.

dimensions(lattice::BravaisLattice) = lattice.N  # FIXME: rename this `extent`?
dimensions(lattice::LatticeWithBasis) = dimensions(bravais(lattice))
dimensions(lattice::WrappedLatticeUnion) = dimensions(lattice.lattice)

ndimensions(::AbstractSiteNetwork) = 0
ndimensions(::AbstractLattice{D}) where {D} = D

twist(lattice::BravaisLattice) = lattice.η
twist(lattice::LatticeWithBasis) = bravais(lattice).η
twist(lattice::WrappedLatticeUnion) = twist(lattice.lattice)
repeater(lattice::BravaisLattice) = lattice.M
repeater(lattice::LatticeWithBasis) = bravais(lattice).M
repeater(lattice::WrappedLatticeUnion) = repeater(lattice.lattice)

ishelical(lattice::AbstractLattice) = !isdiag(repeater(lattice))

# We intentionally define many of the following functions for Bravais
# lattices only.  If one wants to query them for a lattice with a
# basis, call e.g. primvecs(bravais(lattice)).

"Return the primitive vectors of the direct lattice"
primvecs(lattice::BravaisLattice) = lattice.a

"Return the primitive vectors of the reciprocal lattice"
recivecs(lattice::BravaisLattice) = lattice.b

nmomenta(::AbstractSiteNetwork) = 1
nmomenta(lattice::BravaisLattice) = length(lattice.momenta)
nmomenta(lattice::WrappedBravaisLattice) = nmomenta(lattice.lattice)

eachmomentumindex(lattice::AbstractBravaisLattice) = 1:nmomenta(lattice)

# FIXME: need a way to iterate momenta

function momentum(::AbstractSiteNetwork, args...)
    # There is just one momentum
    @assert args[1] == 1
    SVector{0,Int}[]
end

momentum(lattice::BravaisLattice, idx::Integer) = lattice.momenta[idx]
function momentum(lattice::BravaisLattice{D}, idx::Integer, charge::Integer)::SVector{D,Rational{Int}} where {D}
    # "total momentum", really.  note that this may return things greater than one.
    x1 = momentum(lattice, idx)
    if charge == 1
        return x1
    end
    offsets = zeros(MVector{D,Rational{Int}})
    for i in 1:D
        if lattice.M[i,i] != 0
            offsets[i] = lattice.η[i] * (charge - 1)
            for j in 1:i-1
                offsets[i] -= lattice.M[i,j] * offsets[j]
            end
            offsets[i] = offsets[i] // lattice.M[i,i]
        end
    end
    return x1 + offsets
end

kdotr(::AbstractSiteNetwork, k, r) = 0 # always zero since there is just one momentum
kdotr(lattice::BravaisLattice, ksite::AbstractVector{Rational{Int}}, site::AbstractVector{Int}) = 2pi * dot(ksite, site)
kdotr(lattice::BravaisLattice, kidx::Integer, site::AbstractVector{Int}) = kdotr(lattice, momentum(lattice, kidx), site)
kdotr(lattice::BravaisLattice, k, ridx::Integer) = kdotr(lattice, k, lattice[ridx])

momentumspace(lattice::BravaisLattice, k::AbstractVector{Rational{Int}}) = lattice.b * k
momentumspace(lattice::BravaisLattice, kidx::Integer) = momentumspace(lattice, momentum(lattice, kidx))
# FIXME: make it possible to get momentum-space points translated to the first brillouin zone (momentumspace_bz)

realspace(lattice::BravaisLattice, site::AbstractVector{Int}) = lattice.a * site
# NOTE: the following function does not error out if wrap makes no sense due to OBC.
realspace(lattice::BravaisLattice, site::AbstractVector{Int}, wrap::AbstractVector{Int}) = lattice.a * (site + transpose(lattice.M) * wrap)

function realspace(lattice::LatticeWithBasis, site::AbstractVector{Int}, args...)
    length(site) == length(lattice.maxcoords) || throw(ArgumentError(""))
    return realspace(bravais(lattice), site[1:end-1], args...) + lattice.basis[site[end]+1]
end

realspace(lattice::LatticeImplUnion, ridx::Integer, args...) = realspace(lattice, lattice[ridx], args...)

function wraparound_site!(lattice::LatticeImplUnion{D}, site::AbstractVector{Int}) where {D}
    mc = maxcoords(lattice)
    length(site) == length(mc) || throw(ArgumentError(""))

    # For lattice w/ basis, make sure the last index is in range, as
    # we cannot wrap it around.
    if isa(lattice, LatticeWithBasis)
        if !(0 <= site[end] < mc[end])
            throw(ArgumentError("Site has invalid basis index."))
        end
    end

    wrap = zeros(Int, D)

    N = bravais(lattice).N
    M = bravais(lattice).M

    # FIXME: use a @generated function to unroll
    for i in D:-1:1
        if !(0 <= site[i] < N[i])
            if M[i,i] != 0
                # periodic/twisted BC's in this direction
                wrap[i] = fld(site[i], N[i])
                for j in 1:i
                    site[j] -= wrap[i] * M[i,j]
                end
            else
                # OBC in this direction
                #
                # XXX: This does not provide any exception guarantee.
                # If we wanted, we could easily undo what we had done
                # before throwing.  (Additionally, any mutating
                # functions that call this may also wish to undo their
                # side effects before propagating the exception.)
                throw(ArgumentError("This lattice has open boundary conditions in the direction we are trying to wraparound!"))
            end
        end
    end

    return site, wrap
end

function wraparound_site(lattice::LatticeImplUnion{D}, site::SVector{Dprime,Int} where Dprime) where {D}
    mc = maxcoords(lattice)
    length(site) == length(mc) || throw(ArgumentError(""))

    # For lattice w/ basis, make sure the last index is in range, as
    # we cannot wrap it around.
    if isa(lattice, LatticeWithBasis)
        if !(0 <= site[end] < mc[end])
            throw(ArgumentError("Site has invalid basis index."))
        end
    end

    wrap = zeros(SVector{D,Int})

    N = bravais(lattice).N
    M = bravais(lattice).M

    # FIXME: use a @generated function to unroll
    for i in D:-1:1
        if !(0 <= site[i] < N[i])
            if M[i,i] != 0
                # periodic/twisted BC's in this direction
                wrap = setindex(wrap, fld(site[i], N[i]), i)
                for j in 1:i
                    site = setindex(site, site[j] - wrap[i] * M[i,j], j)
                end
            else
                # OBC in this direction
                #
                # XXX: This does not provide any exception guarantee.
                # If we wanted, we could easily undo what we had done
                # before throwing.  (Additionally, any mutating
                # functions that call this may also wish to undo their
                # side effects before propagating the exception.)
                throw(ArgumentError("This lattice has open boundary conditions in the direction we are trying to wraparound!"))
            end
        end
    end

    return site, wrap
end

wraparound_site(lattice::WrappedLatticeUnion, site::SVector{Dprime,Int} where Dprime) = wraparound_site(lattice.lattice, site)

wraparound_site!(lattice::WrappedLatticeUnion, site::AbstractVector{Int}) = wraparound_site!(lattice.lattice, site)

function wraparound_site(lattice::Union{AbstractBravaisLattice{Dprime},AbstractLatticeWithBasis{D,Dprime} where D}, site::AbstractVector{Int}) where {Dprime}
    site_, wrap = wraparound_site!(lattice, MVector{Dprime,Int}(site))
    return SVector(site_), wrap
end

wraparound_site(lattice::AbstractLattice, index::Integer) = wraparound_site(lattice, lattice[index])

function wraparound(lattice::AbstractLattice, site_or_index::Union{AbstractVector{Int}, Integer})
    site, wrap = wraparound_site(lattice, site_or_index)
    idx = findfirst(isequal(site), lattice)::Int
    return idx, wrap
end

function wraparoundη(lattice::AbstractLattice, site_or_index::Union{AbstractVector{Int}, Integer})
    idx, wrap = wraparound(lattice, site_or_index)
    η = dot(wrap, twist(lattice))
    return idx, η
end

function translate_site!(lattice::AbstractLattice, site::AbstractVector{Int}, direction::Integer)
    if !isbravais(lattice) && direction > ndimensions(lattice)
        throw(ArgumentError("Cannot translate in the 'direction' of the basis index."))
    end
    site[direction] += 1
    return wraparound_site!(lattice, site)
end

function translate_site(lattice::Union{AbstractBravaisLattice{Dprime},AbstractLatticeWithBasis{D,Dprime} where D}, site::AbstractVector{Int}, direction::Integer) where {Dprime}
    site_, wrap = translate_site!(lattice, MVector{Dprime,Int}(site), direction)
    return SVector(site_), wrap
end

translate_site(lattice::AbstractLattice, index::Integer, direction::Integer) = translate_site(lattice, lattice[index], direction)

function translate(lattice::AbstractLattice, site_or_index::Union{AbstractVector{Int}, Integer}, direction::Integer)
    site, wrap = translate_site(lattice, site_or_index, direction)
    idx = findfirst(isequal(site), lattice)::Int
    return idx, wrap
end

function translateη(lattice::AbstractLattice, site_or_index::Union{AbstractVector{Int}, Integer}, direction::Integer)
    idx, wrap = translate(lattice, site_or_index, direction)
    η = dot(wrap, twist(lattice))
    return idx, η
end

function neighbors(f, lattice::AbstractLattice, neigh=Val{1})
    for ridx in 1:length(lattice)
        siteneighbors(f, lattice, ridx, neigh)
    end
end

function neighborsη(f, lattice::AbstractLattice, neigh=Val{1})
    @closure neighbors(lattice, neigh) do idx1, idx2, wrap
        η = dot(wrap, twist(lattice))
        f(idx1, idx2, η)
    end
end

#= Begin common code that will help with specific Bravais lattice implementations =#

sublattice_index(lattice::AbstractLattice, ridx::Int) = sublattice_index(lattice, lattice[ridx])

siteneighbors(f, lattice::AbstractLattice, site_or_index::Union{AbstractVector{Int}, Integer}) = siteneighbors(f, lattice, site_or_index, 1)

function _siteneighbors2d(f, lattice::AbstractLattice{2}, ridx::Integer, offsets)
    M = repeater(lattice)
    mc = maxcoords(lattice)

    site = lattice[ridx]

    for offset in offsets
        newsite = site + offset
        g(i) = M[i,i] == 0 && !(0 <= newsite[i] < mc[i])
        if g(1) || g(2)
            # In directions w/ OBC, do not provide the neighbors!
            #
            # NOTE: In e.g. a 3d lattice, similar logic will fail
            # if there's one OBC direction and helical BCs in the
            # two periodic directions.
            continue
        end
        newidx, wrap = wraparound(lattice, newsite)
        f(ridx, newidx, wrap)
    end
    nothing
end

#= Begin specific Bravais lattice implementations =#

function _hypercubic_sublattice_index(site::AbstractVector{Int})
    parity = 0
    for x in site
        parity = parity ⊻ x
    end
    return parity & 1
end

struct HypercubicLattice{D,Dsq} <: WrappedBravaisLattice{D}
    lattice::BravaisLattice{D,Dsq}
    bipartite::Bool

    function HypercubicLattice{D,Dsq}(N::AbstractVector{Int},
                                      M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                                      η::AbstractVector{Rational{Int}}=zeros(SVector{D,Rational{Int}})) where {D,Dsq}
        bravaislattice = BravaisLattice{D,Dsq}(N, M, η)
        bipartite = true
        for i in 1:D
            if M[i,i] != 0
                # Attempt to go across the boundary in the `i`
                # direction, and test if the site is on the same
                # sublattice after being wrapped around.
                site1 = setindex(zeros(SVector{D,Int}), N[i], i)
                site2, = wraparound_site(bravaislattice, site1)
                if _hypercubic_sublattice_index(site1) != _hypercubic_sublattice_index(site2)
                    bipartite = false
                    break
                end
            end
        end
        new(bravaislattice, bipartite)
    end
end

HypercubicLattice{D}(args...) where {D} =
    HypercubicLattice{D,square(D)}(args...)

HypercubicLattice(N::StaticVector{D,Int}, args...) where {D} =
    HypercubicLattice{D}(N, args...)

isbipartite(lattice::HypercubicLattice) = lattice.bipartite
istripartite(::HypercubicLattice) = false

function sublattice_index(lattice::HypercubicLattice, site::AbstractVector{Int})
    isbipartite(lattice) || throw(ArgumentError("Hypercubic lattice must be bipartite for it to have sublattice indices."))
    return _hypercubic_sublattice_index(site)
end

const ChainLattice = HypercubicLattice{1,1}
const SquareLattice = HypercubicLattice{2,4}
const CubicLattice = HypercubicLattice{3,9}

# FIXME: do the wrapping etc in a common function for all lattice
# types.  then the siteneighbors function (or `siteneighbordsimpl`)
# can be specified even more simply.

# BASICALLY: give two things: single count and double count offsets.
# also, for each basis index.  then again, this generic code is not
# going to be quite as optimized as e.g. the hypercubic lattice.  but
# that is okay, i can just have special code for that.

# XXX FIXME: we want to be able (in the end) to choose different
# primitive vectors so we can have a weird helical lattice.  how are
# we going to support this??

function siteneighbors(f, lattice::ChainLattice, ridx::Integer, ::Type{Val{1}})
    M = lattice.lattice.M
    mc = maxcoords(lattice)
    site = lattice[ridx]
    newsite = @SVector[site[1] + 1]
    if M[1,1] <= 1 && newsite[1] >= mc[1]
        # Do not provide neighbors across an open boundary.
        return
    end
    newidx, wrap = wraparound(lattice, newsite)
    f(ridx, newidx, wrap)
    nothing
end

function siteneighbors(f, lattice::ChainLattice, ridx::Integer, ::Type{Val{N}}) where {N}
    M = lattice.lattice.M
    mc = maxcoords(lattice)
    site = lattice[ridx]
    newsite = @SVector[site[1] + N]
    if M[1,1] <= 1 && newsite[1] >= mc[1]
        # Do not provide neighbors across an open boundary.
        return
    end
    newidx, wrap = wraparound(lattice, newsite)
    f(ridx, newidx, wrap)
    nothing
end

function siteneighbors(f, lattice::HypercubicLattice{D}, ridx::Integer, ::Type{Val{1}}) where {D}
    M = lattice.lattice.M
    mc = maxcoords(lattice)

    site = lattice[ridx]

    for i in 1:D
        newsite = setindex(site, site[i] + 1, i)
        if M[i,i] <= 1 && newsite[i] >= mc[i]
            # In directions w/ OBC or length one, do not provide the neighbors!
            continue
        end
        newidx, wrap = wraparound(lattice, newsite)
        f(ridx, newidx, wrap)
    end
end

function siteneighbors(f, lattice::SquareLattice, ridx::Integer, ::Type{Val{2}})
    offsets = @SVector[@SVector[1,1], @SVector[-1,1]]
    _siteneighbors2d(f, lattice, ridx, offsets)
end

function _triangular_sublattice_index(site::AbstractVector{Int})
    @assert length(site) == 2
    # FIXME: mod is slow
    return mod(site[2] - site[1], 3)
end

struct TriangularLattice <: WrappedBravaisLattice{2}
    lattice::BravaisLattice{2,4}
    tripartite::Bool

    function TriangularLattice(N::AbstractVector{Int},
                               M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                               η::AbstractVector{Rational{Int}}=zeros(SVector{2,Rational{Int}}))
        a = copy(transpose(@SMatrix[1.0 0; 0.5 sqrt(3)/2]))
        bravaislattice = BravaisLattice{2,4}(N, M, η, a)
        tripartite = true
        for i in 1:2
            if M[i,i] != 0
                # Attempt to go across the boundary in the `i`
                # direction, and test if the site is on the same
                # sublattice after being wrapped around.
                site1 = [0, 0]
                site1[i] = N[i]
                site2, = wraparound_site(bravaislattice, site1)
                if _triangular_sublattice_index(site1) != _triangular_sublattice_index(site2)
                    tripartite = false
                    break
                end
            end
        end
        new(bravaislattice, tripartite)
    end
end

isbipartite(::TriangularLattice) = false # although "technically" it might be if it's just a chain
istripartite(lattice::TriangularLattice) = lattice.tripartite

function sublattice_index(lattice::TriangularLattice, site::AbstractVector{Int})
    istripartite(lattice) || throw(ArgumentError("Triangular lattice must be tripartite for it to have sublattice indices."))
    return _triangular_sublattice_index(site)
end

# FIXME: many of these neighbor functions can use special handling when the lattice height or width is 1 in a direction.  or we could just forbid this.

function siteneighbors(f, lattice::TriangularLattice, ridx::Integer, ::Type{Val{1}})
    offsets = @SVector[@SVector[1,0], @SVector[0,1], @SVector[-1,1]]
    _siteneighbors2d(f, lattice, ridx, offsets)
end

function siteneighbors(f, lattice::TriangularLattice, ridx::Integer, ::Type{Val{2}})
    offsets = @SVector[@SVector[2,-1], @SVector[1,1], @SVector[-1,2]]
    _siteneighbors2d(f, lattice, ridx, offsets)
end

function siteneighbors(f, lattice::TriangularLattice, ridx::Integer, ::Type{Val{3}})
    offsets = @SVector[@SVector[2,0], @SVector[0,2], @SVector[-2,2]]
    _siteneighbors2d(f, lattice, ridx, offsets)
end

function rhombi(lattice::TriangularLattice)
    @SVector[
        @SVector[@SVector[0,0], @SVector[1,-1], @SVector[1,0], @SVector[0,1]],
        @SVector[@SVector[0,0], @SVector[0,1], @SVector[-1,1], @SVector[-1,0]],
        @SVector[@SVector[0,0], @SVector[-1,0], @SVector[0,-1], @SVector[1,-1]],
    ]
end

@delegate WrappedBravaisLattice.lattice [ Base.length, ndimensions, primvecs, recivecs, momentum, kdotr, momentumspace, realspace ]

#= End specific Bravais lattice implementations =#

#= Begin specific lattice w/ basis implementations =#

struct HoneycombLattice <: WrappedLatticeWithBasis{2,3}
    lattice::LatticeWithBasis{2,4,3}

    function HoneycombLattice(N::AbstractVector{Int},
                              M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                              η::AbstractVector{Rational{Int}}=zeros(SVector{2,Rational{Int}}))
        a = copy(transpose(@SMatrix[1.5 sqrt(3)/2; 0 sqrt(3)]))
        basis = [@SVector[0.0, 0], @SVector[1.0, 0]]
        return new(LatticeWithBasis{2,4,3}(N, M, η, a, basis))
    end
end

isbipartite(::HoneycombLattice) = true
istripartite(::HoneycombLattice) = false

function sublattice_index(::HoneycombLattice, site::AbstractVector{Int})
    retval = typeassert(site[end], Int)
    @assert 0 <= retval < 2 # for bipartite lattice
    return retval
end

function siteneighbors(f, lattice::HoneycombLattice, ridx::Integer, ::Type{Val{1}})
    site = lattice[ridx]
    if site[end] == 0
        offsets = @SVector[@SVector[0, 0, 1], @SVector[-1, 0, 1], @SVector[-1, 1, 1]]
        _siteneighbors2d(f, lattice, ridx, offsets)
    else
        @assert site[end] == 1
        # Commented out, otherwise we are double counting!
        #offsets = @SVector[@SVector[0, 0, -1], @SVector[1, 0, -1], @SVector[1, -1, -1]]
        #_siteneighbors2d(f, lattice, ridx, offsets)
    end
end

struct KagomeLattice <: WrappedLatticeWithBasis{2,3}
    lattice::LatticeWithBasis{2,4,3}

    function KagomeLattice(N::AbstractVector{Int},
                           M::AbstractMatrix{Int}=Diagonal(N), # assumes pbc
                           η::AbstractVector{Rational{Int}}=zeros(SVector{2,Rational{Int}}))
        a = copy(transpose(@SMatrix[2 0; 1 sqrt(3)]))
        basis = [@SVector[0.0, 0], @SVector[0.5, √3/2], @SVector[1.0, 0]]
        return new(LatticeWithBasis{2,4,3}(N, M, η, a, basis))
    end
end

isbipartite(::KagomeLattice) = false
istripartite(::KagomeLattice) = true

function sublattice_index(::KagomeLattice, site::AbstractVector{Int})
    retval = typeassert(site[end], Int)
    @assert 0 <= retval < 3 # for tripartite lattice
    return retval
end

function siteneighbors(f, lattice::KagomeLattice, ridx::Integer, ::Type{Val{1}})
    site = lattice[ridx]
    if site[end] == 0
        offsets = @SVector[@SVector[0, 0, 1], @SVector[0, -1, 1]]
    elseif site[end] == 1
        offsets = @SVector[@SVector[0, 0, 1], @SVector[-1, 1, 1]]
    else
        @assert site[end] == 2
        offsets = @SVector[@SVector[0, 0, -2], @SVector[1, 0, -2]]
    end
    _siteneighbors2d(f, lattice, ridx, offsets)
end

@delegate WrappedLatticeWithBasis.lattice [ Base.length, realspace ]

#= End specific lattice w/ basis implementations =#

#= Begin cache objects =#

"""
Precomputes the results of the translateη function on each site.

ltrc = LatticeTranslationCache(lattice, direction)
translateη(ltrc, site)
"""
struct LatticeTranslationCache{LatticeType<:AbstractLattice}
    lattice::LatticeType
    direction::Int
    cache::Vector{Tuple{Int,Rational{Int}}}

    function LatticeTranslationCache(lattice::LatticeType, direction::Integer) where {LatticeType<:AbstractLattice}
        cache = Tuple{Int,Rational{Int}}[]
        sizehint!(cache, length(lattice))
        for i in 1:length(lattice)
            push!(cache, translateη(lattice, i, direction))
        end
        new{LatticeType}(lattice, Int(direction), cache)
    end
end

struct SiteNetworkTranslationCache
    cache::Vector{Tuple{Int,Rational{Int}}}

    function SiteNetworkTranslationCache(permutation)
        isperm(permutation) || throw(ArgumentError("not a valid permutation"))
        cache = Tuple{Int,Rational{Int}}[]
        sizehint!(cache, length(permutation))
        for j in permutation
            # XXX: for simplicity we assume all phases are zero; but we can easily take it as argument
            push!(cache, (j, 0//1))
        end
        new(cache)
    end
end

translateη(ltrc::Union{LatticeTranslationCache,SiteNetworkTranslationCache}, j::Integer) = ltrc.cache[j]
translateη(ltrc::LatticeTranslationCache, site::AbstractVector{Int}) = translateη(ltrc, findfirst(isequal(site), ltrc.lattice)::Int)

#= End cache objects =#

export
    AbstractSiteNetwork,
    AbstractLattice,
    AbstractBravaisLattice,
    AbstractLatticeWithBasis,
    BravaisLattice,
    HypercubicLattice,
    ChainLattice,
    SquareLattice,
    CubicLattice,
    TriangularLattice,
    LatticeWithBasis,
    HoneycombLattice,
    KagomeLattice,
    bravais,
    isbravais,
    maxcoords,
    ndimensions,
    dimensions,
    twist,
    repeater,
    ishelical,
    primvecs,
    recivecs,
    nmomenta,
    eachmomentumindex,
    momentum,
    kdotr,
    realspace,
    wraparound_site!,
    wraparound_site,
    wraparound,
    wraparoundη,
    translate_site!,
    translate_site,
    translate,
    translateη,
    momentumspace,
    siteneighbors,
    neighbors,
    neighborsη,
    isbipartite,
    istripartite,
    sublattice_index,
    LatticeTranslationCache,
    SiteNetworkTranslationCache

end # module


text:
"""
    model_infill(search_range::Vector{Tuple{Float64,Float64}},plan::AbstractArray{T,2},
    samples::AbstractArray{T,2},sm_interpolant; options::Options=Options()) where T

Infill function that calculates the location of new samples based on the supplied options.
The returned options are updated to facilitate cycling through the infill objective
functions.

...
# Arguments
- `search_range::Vector{Tuple{Float64,Float64}`:
    a vector of tuples containing the lower and upper limits
    for each dimension. `length(search_range)` is equal to number
    of dimensions.
- `plan::AbstractArray{T,2}`:
    sample locations where each column corresponds to the location of one point.
    `size(plan) = (num_dimensions,num_samples)`.
- `samples::AbstractArray{T,2}`:
    function value at each sample location. each column contains one value from
    the corresponding plan location. `size(samples) = (1,num_samples)`.
- `options=Options()`:
    all options available to customize the surrogate infill.
...
"""
function model_infill(search_range::Vector{Tuple{Float64,Float64}},plan::AbstractArray{T,2},
        samples::AbstractArray{T,2},sm_interpolant; options::Options=Options()) where T

    @unpack rbf_opt_gens, num_infill_points, trace,
            infill_funcs, infill_iterations = options

    (trace == :verbose) && print_infill_head()

    (length(samples) != size(plan,2)) && error("plan and samples do not have the correct sizes")
    ((length(sm_interpolant(plan[:,1])) == 1) && (any(infill_funcs.==:std))) && error(":std infill criteria can not be used with one interpolant")

    # Extract the requested infill objective functions
    infill_obj_fun = infill_objective(sm_interpolant,plan,samples,infill_funcs[1:num_infill_points])

    # Find the infill optimization points
    @timeit_debug "infill opt" begin
    infill_plan,infill_type,infill_prediction,res_bboptim,options =
        infill_opt(search_range,infill_iterations,num_infill_points,infill_obj_fun,infill_funcs,plan,sm_interpolant,options)
    end

    # Add additional points to fill up the desired number of infill points if required
    infill_plan,infill_type,infill_prediction =
        infill_add(sm_interpolant,samples,plan,infill_prediction,search_range,infill_plan,
        infill_type,num_infill_points,infill_obj_fun,infill_funcs,res_bboptim)

    (trace == :verbose) && print_infill_tail(infill_type,num_infill_points,infill_plan,infill_prediction)

    return infill_plan, infill_type, infill_prediction, options
end

ext: jl


text:
module KubeMeta

include("client.jl")
include("informer.jl")

export TaskListWatcher, TaskInformer, run, K8sClient, request, client

include("controllers/Controllers.jl")
using .Controllers
export Controller, DifferentiableController, DifferentiableAgent, add_task, update_task, delete_task

end # module

max_stars_repo_name: pbarker/kubemeta

max_issues_repo_name: pbarker/kubemeta

max_forks_repo_name: pbarker/kubemeta


text:
# Autogenerated wrapper script for Pinchot_jll for armv7l-linux-gnueabihf-cxx03
export libpinchot

JLLWrappers.@generate_wrapper_header("Pinchot")
JLLWrappers.@declare_library_product(libpinchot, "libpinchot.so")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libpinchot,
        "lib/libpinchot.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# @testset "byte_stream_capacity.jl" begin
try
    ts = time_ns()
    let
        test = ByteStreamTestHarness("overwrite", 2)
        execute(test, with_bytes_written!(Write("cat"), 2))
        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(2))
        execute(test, RemainingCapacity(0))
        execute(test, BufferSize(2))
        execute(test, Peek("ca"))

        execute(test, with_bytes_written!(Write("t"), 0))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(2))
        execute(test, RemainingCapacity(0))
        execute(test, BufferSize(2))
        execute(test, Peek("ca"))
    end

    let
        test = ByteStreamTestHarness("overwrite-clear-overwrite", 2)

        execute(test, with_bytes_written!(Write("cat"), 2))
        execute(test, Pop(2))
        execute(test, with_bytes_written!(Write("tac"), 2))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(2))
        execute(test, BytesWritten(4))
        execute(test, RemainingCapacity(0))
        execute(test, BufferSize(2))
        execute(test, Peek("ta"))
    end
    let
        test = ByteStreamTestHarness("overwrite-pop-overwrite", 2)

        execute(test, with_bytes_written!(Write("cat"), 2))
        execute(test, Pop(1))
        execute(test, with_bytes_written!(Write("tac"), 1))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(1))
        execute(test, BytesWritten(3))
        execute(test, RemainingCapacity(0))
        execute(test, BufferSize(2))
        execute(test, Peek("at"))
    end

    let
        test = ByteStreamTestHarness("long-stream", 3)

        execute(test, with_bytes_written!(Write("abcdef"), 3))
        execute(test, Peek("abc"))
        execute(test, Pop(1))

        for _ in 0:997
            execute(test, RemainingCapacity(1))
            execute(test, BufferSize(2))
            execute(test, with_bytes_written!(Write("abc"), 1))
            execute(test, RemainingCapacity(0))
            execute(test, Peek("bca"))
            execute(test, Pop(1))

            execute(test, RemainingCapacity(1))
            execute(test, BufferSize(2))
            execute(test, with_bytes_written!(Write("bca"), 1))
            execute(test, RemainingCapacity(0))
            execute(test, Peek("cab"))
            execute(test, Pop(1))

            execute(test, RemainingCapacity(1))
            execute(test, BufferSize(2))
            execute(test, with_bytes_written!(Write("cab"), 1))
            execute(test, RemainingCapacity(0))
            execute(test, Peek("abc"))
            execute(test, Pop(1))
        end

        execute(test, EndInput())
        execute(test, Peek("bc"))
        execute(test, Pop(2))
        execute(test, Eof(true))
    end
    tc = (time_ns() - ts |> Int) / 1e9
catch e
    println(stderr, e)
end

max_stars_repo_name: AquaIndigo/JLSponge

max_issues_repo_name: AquaIndigo/JLSponge

max_forks_repo_name: AquaIndigo/JLSponge


text:
module WorkPool

using Printf
using Logging

export setLoggerFile, pool

mutable struct WorkerState
	name::String
	job_count::Int64
	errors::Int64
end

struct WorkPkg{T}
	name::String
	work::T
end

struct Ctrl
	q::Channel{Bool}
	t::Float64
end

show(io::IO, a::WorkPkg) = print(io, "WorkPkg $(a.name)")

function setLoggerFile(fn::String)
	io = open(fn, "w+")
	logger = SimpleLogger(io)
	global_logger(logger)
end

mutable struct Pool{T,E}
	num_workers::Int64
	channel_size::Int64
	workercb::Function ## function worker(work::T, result::E, done::Ctrl) where {T <: Any, E <: Any}

	source::Channel{WorkPkg{T}}
	sink::Channel{WorkPkg{E}}
	qc::Array{Ctrl}
	ta::Array{Task}
	started::Bool

	Pool{T,E}(num_workers, channel_size, workercb) where {T <: Any, E <: Any} = new(num_workers, channel_size, workercb, Channel{WorkPkg{T}}(channel_size), Channel{WorkPkg{E}}(channel_size), Array{Ctrl}(undef, num_workers), Array{Task}(undef, num_workers), false)

end

alreadyStartedException = ErrorException("Pool Already Started");
notStartedException = ErrorException("Pool NotStarted");

function startPool(p::Pool)

	if p.started
		throw(alreadyStartedException)
	else
		p.started = true
	end

	for i = 1:size(p.qc)[1]
		p.qc[i]=Ctrl(Channel{Bool}(1), convert(Int64, floor(rand()*5)))
	end

	@info(@sprintf("creating %d tasks\n", p.num_workers))
	for i in 1:p.num_workers
		ws = WorkerState(@sprintf("worker-%d",i), 0,0)
		p.ta[i] = @task(worker(ws, p.source, p.sink, p.workercb, p.qc[i]))

		if !istaskstarted(p.ta[i])
			schedule(p.ta[i])
		end
	end

end

# and a function `worker` which reads items from from source, processes the item read
# and writes a result to sink,
function worker(ws::WorkerState, source::Channel{WorkPkg{T}}, sink::Channel{WorkPkg{E}}, cb::Function, done::Ctrl) where {T <: Any, E <: Any}
    @info(@sprintf("worker[%s] startings\n", ws.name))
    while true
	if isready(done.q)
		_ = take!(done.q)
    		@info(@sprintf("worker [%s] ending \n", ws.name))
		return ws
	end

	if isready(source)
        	data = take!(source)
		@info(@sprintf("Worker[%s] received [%s] from source\n", ws.name, data.name))
		res = cb(data.work)
		put!(sink, WorkPkg{E}(@sprintf("%s-processed", data.name), res))
		@info(@sprintf("Worker[%s] sent [%s] to sink\n", ws.name, data.name))
		ws.job_count += 1
    	end

	sleep(done.t)
	yield()

    end
end

function push(p::Pool, w::WorkPkg{T}) where {T <: Any}

	if !p.started
		throw(notStartedException)
	end
	@info(@sprintf("sending job: %s\n", w.name))
	put!(p.source, w)
end

function pop(p::Pool)

	if !p.started
		throw(notStartedException)
	end

	if isready(p.sink)
		info = take!(p.sink)
		@info(@sprintf("received\n"))
		return info
	else
		@info(@sprintf("nothing to receive\n"))
		return nothing
	end
end

function shutdownPool(p::Pool)

	if !p.started
		throw(notStartedException)
	end

	for i=1:p.num_workers
		@info(@sprintf("signal worker %d to quit\n", i))
		put!(p.qc[i].q, true)
	end

	@info(@sprintf("waiting on done!\n"))
	while true
		all_done = true

		yield()
		for i = 1:(size(p.ta)[1])
			if !istaskdone(p.ta[i])
				#@info(@sprintf("task %d isn't done\n", i))
				all_done = false
			end
		end

		if (all_done)
			@info(@sprintf("done!\n"))
			break
		end
	end
	p.started = false
	return

end

function getFinalState(p::Pool)
	for i in 1:p.num_workers
		@info(@sprintf("worker %s quit, count, error results: %d, %d\n", p.ta[i].result.name, p.ta[i].result.job_count, p.ta[i].result.errors))
	end
end

end

max_stars_repo_name: rabarar/juliaSnippets

max_issues_repo_name: rabarar/juliaSnippets

max_forks_repo_name: rabarar/juliaSnippets


text:
export InlineString, InlineStringType

"""
    InlineString

A set of custom string types of various fixed sizes. Each inline string
is a custom primitive type and can benefit from being stack friendly by
avoiding allocations/heap tracking in the GC. When used in an array,
the elements are able to be stored inline since each one has a fixed
size. Currently support inline strings from 1 byte up to 255 bytes.
"""
abstract type InlineString <: AbstractString end

for sz in (1, 4, 8, 16, 32, 64, 128, 256)
    nm = Symbol(:InlineString, max(1, sz - 1))
    @eval begin
        """
            $($nm)

        Custom fixed-size string with a fixed size of $($sz) bytes.
        1 byte is used to store the length of the string. If an
        inline string is shorter than $($sz - 1) bytes, the entire
        string still occupies the full $($sz) bytes since they are,
        by definition, fixed size. Otherwise, they can be treated
        just like normal `String` values. Note that `sizeof(x)` will
        return the # of _codeunits_ in an $($nm) like `String`, not
        the total fixed size. For the fixed size, call `sizeof($($nm))`.
        $($nm) can be constructed from an existing `String` (`$($nm)(x::String)`),
        from a byte buffer with position and length (`$($nm)(buf, pos, len)`),
        or built iteratively by starting with `x = $($nm)()` and calling
        `WeakRefStrings.addcodeunit(x, b::UInt8)` which returns a new $($nm)
        with the new codeunit `b` appended.
        """
        primitive type $nm <: InlineString $(sz * 8) end
        export $nm
    end
end

_bswap(x::T) where {T <: InlineString} = T === InlineString1 ? x : Base.bswap_int(x)

const InlineStrings = Union{InlineString1,
                            InlineString3,
                            InlineString7,
                            InlineString15,
                            InlineString31,
                            InlineString63,
                            InlineString127,
                            InlineString255}

function Base.promote_rule(::Type{T}, ::Type{S}) where {T <: InlineString, S <: InlineString}
    T === InlineString1 && return S
    S === InlineString1 && return T
    T === InlineString3 && return S
    S === InlineString3 && return T
    T === InlineString7 && return S
    S === InlineString7 && return T
    T === InlineString15 && return S
    S === InlineString15 && return T
    T === InlineString31 && return S
    S === InlineString31 && return T
    T === InlineString63 && return S
    S === InlineString63 && return T
    T === InlineString127 && return S
    S === InlineString127 && return T
    return InlineString255
end

Base.promote_rule(::Type{T}, ::Type{String}) where {T <: InlineString} = String

Base.widen(::Type{InlineString1}) = InlineString3
Base.widen(::Type{InlineString3}) = InlineString7
Base.widen(::Type{InlineString7}) = InlineString15
Base.widen(::Type{InlineString15}) = InlineString31
Base.widen(::Type{InlineString31}) = InlineString63
Base.widen(::Type{InlineString63}) = InlineString127
Base.widen(::Type{InlineString127}) = InlineString255
Base.widen(::Type{InlineString255}) = String

Base.ncodeunits(::InlineString1) = 1
Base.ncodeunits(x::InlineString) = Int(Base.trunc_int(UInt8, x))
Base.codeunit(::InlineString) = UInt8

Base.@propagate_inbounds function Base.codeunit(x::T, i::Int) where {T <: InlineString}
    @boundscheck checkbounds(Bool, x, i) || throw(BoundsError(x, i))
    if T === InlineString1
        return Base.bitcast(UInt8, x)
    else
        return Base.trunc_int(UInt8, Base.lshr_int(x, 8 * (sizeof(T) - i)))
    end
end

function Base.String(x::T) where {T <: InlineString}
    len = ncodeunits(x)
    out = Base._string_n(len)
    if T === InlineString1
        GC.@preserve out unsafe_store!(pointer(out), codeunit(x, 1))
        return out
    end
    ref = Ref{T}(_bswap(x))
    GC.@preserve ref out begin
        ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
        unsafe_copyto!(pointer(out), ptr, len)
    end
    return out
end

function Base.Symbol(x::T) where {T <: InlineString}
    ref = Ref{T}(_bswap(x))
    return ccall(:jl_symbol_n, Ref{Symbol},
        (Ref{T}, Int), ref, sizeof(x))
end

# add a codeunit to end of string method
function addcodeunit(x::T, b::UInt8) where {T <: InlineString}
    if T === InlineString1
        return x, false
    end
    len = Base.trunc_int(UInt8, x)
    sz = Base.trunc_int(UInt8, sizeof(T))
    shf = Base.zext_int(Int16, max(0x01, sz - len - 0x01)) << 3
    x = Base.or_int(x, Base.shl_int(Base.zext_int(T, b), shf))
    return Base.add_int(x, Base.zext_int(T, 0x01)), (len + 0x01) >= sz
end

# from String
InlineString1(byte::UInt8=0x00) = Base.bitcast(InlineString1, byte)
(::Type{T})() where {T <: InlineString} = Base.zext_int(T, 0x00)

function (::Type{T})(x::AbstractString) where {T <: InlineString}
    if T === InlineString1
        sizeof(x) == 1 || stringtoolong(T, sizeof(x))
        return Base.bitcast(InlineString1, codeunit(x, 1))
    elseif typeof(x) === String && sizeof(T) <= sizeof(UInt)
        len = sizeof(x)
        len < sizeof(T) || stringtoolong(T, len)
        y = GC.@preserve x unsafe_load(convert(Ptr{T}, pointer(x)))
        sz = 8 * (sizeof(T) - len)
        return Base.or_int(Base.shl_int(Base.lshr_int(_bswap(y), sz), sz), Base.zext_int(T, UInt8(len)))
    else
        len = ncodeunits(x)
        len < sizeof(T) || stringtoolong(T, len)
        y = T()
        for i = 1:len
            @inbounds y, _ = addcodeunit(y, codeunit(x, i))
        end
        return y
    end
end

function (::Type{T})(buf::AbstractVector{UInt8}, pos, len) where {T <: InlineString}
    if T === InlineString1
        sizeof(x) == 1 || stringtoolong(T, sizeof(x))
        return Base.bitcast(InlineString1, buf[pos])
    else
        length(buf) < len && buftoosmall()
        len < sizeof(T) || stringtoolong(T, len)
        y = GC.@preserve buf unsafe_load(convert(Ptr{T}, pointer(buf, pos)))
        sz = 8 * (sizeof(T) - len)
        return Base.or_int(Base.shl_int(Base.lshr_int(_bswap(y), sz), sz), Base.zext_int(T, UInt8(len)))
    end
end

@noinline buftoosmall(n) = throw(ArgumentError("input buffer too short for requested length: $n"))
@noinline stringtoolong(T, n) = throw(ArgumentError("string too large ($n) to convert to $T"))

function InlineStringType(n::Integer)
    n > 255 && stringtoolong(InlineString, n)
    return n == 1  ? InlineString1   : n < 4  ? InlineString3  :
           n < 8   ? InlineString7   : n < 16 ? InlineString15 :
           n < 32  ? InlineString31  : n < 64 ? InlineString63 :
           n < 128 ? InlineString127 : InlineString255
end

InlineString(x::AbstractString)::InlineStrings = (InlineStringType(ncodeunits(x)))(x)

# between InlineStrings
function (::Type{T})(x::S) where {T <: InlineString, S <: InlineString}
    if T === S
        return x
    elseif T === InlineString1
        sizeof(x) == 1 || stringtoolong(T, sizeof(x))
        return Base.bitcast(InlineString1, codeunit(x, 1))
    elseif sizeof(T) < sizeof(S)
        # trying to compress
        len = sizeof(x)
        len > (sizeof(T) - 1) && stringtoolong(T, len)
        y = Base.trunc_int(T, Base.lshr_int(x, 8 * (sizeof(S) - sizeof(T))))
        return Base.add_int(y, Base.zext_int(T, UInt8(len)))
    else
        # promoting smaller InlineString to larger
        if S === InlineString1
            y = Base.shl_int(Base.zext_int(T, x), 8 * (sizeof(T) - sizeof(S)))
        else
            y = Base.shl_int(Base.zext_int(T, Base.lshr_int(x, 8)), 8 * (sizeof(T) - sizeof(S) + 1))
        end
        return Base.add_int(y, Base.zext_int(T, UInt8(sizeof(x))))
    end
end

(==)(x::T, y::T) where {T <: InlineString} = Base.eq_int(x, y)
function ==(x::String, y::T) where {T <: InlineString}
    sizeof(x) == sizeof(y) || return false
    ref = Ref{T}(_bswap(y))
    return ccall(:memcmp, Cint, (Ptr{UInt8}, Ref{T}, Csize_t),
            pointer(x), ref, sizeof(x)) == 0
end
==(y::InlineString, x::String) = x == y

function Base.hash(x::T, h::UInt) where {T <: InlineString}
    h += Base.memhash_seed
    ref = Ref{T}(_bswap(x))
    return ccall(Base.memhash, UInt,
        (Ref{T}, Csize_t, UInt32),
        ref, sizeof(x), h % UInt32) + h
end

function Base.write(io::IO, x::T) where {T <: InlineString}
    ref = Ref{T}(_bswap(x))
    return GC.@preserve ref begin
        ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
        Int(unsafe_write(io, ptr, reinterpret(UInt, sizeof(x))))::Int
    end
end

Base.print(io::IO, s::InlineString) = (write(io, s); nothing)

function Base.isascii(x::T) where {T <: InlineString}
    if T === InlineString1
        return codeunit(x, 1) < 0x80
    end
    len = ncodeunits(x)
    x = Base.lshr_int(x, 8 * (sizeof(T) - len))
    for _ = 1:len
        Base.trunc_int(UInt8, x) >= 0x80 && return false
        x = Base.lshr_int(x, 8)
    end
    return true
end

# copy/pasted from substring.jl
function Base.reverse(s::InlineString)
    # Read characters forwards from `s` and write backwards to `out`
    out = Base._string_n(sizeof(s))
    offs = sizeof(s) + 1
    for c in s
        offs -= ncodeunits(c)
        Base.__unsafe_string!(out, c, offs)
    end
    return out
end

@inline function Base.__unsafe_string!(out, x::T, offs::Integer) where {T <: InlineString}
    n = sizeof(x)
    ref = Ref{T}(_bswap(x))
    GC.@preserve ref out begin
        ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
        unsafe_copyto!(pointer(out, offs), ptr, n)
    end
    return n
end

Base.string(a::InlineString) = a
Base.string(a::InlineString...) = _string(a...)
Base.string(a::BaseStrs, b::InlineString) = _string(a, b)
Base.string(a::BaseStrs, b::BaseStrs, c::InlineString) = _string(a, b, c)
@inline function _string(a::Union{BaseStrs, InlineString}...)
    n = 0
    for v in a
        if v isa Char
            n += ncodeunits(v)
        else
            n += sizeof(v)
        end
    end
    out = Base._string_n(n)
    offs = 1
    for v in a
        offs += Base.__unsafe_string!(out, v, offs)
    end
    return out
end

function Base.repeat(x::T, r::Integer) where {T <: InlineString}
    r < 0 && throw(ArgumentError("can't repeat a string $r times"))
    r == 0 && return ""
    r == 1 && return s
    n = sizeof(x)
    out = Base._string_n(n * r)
    if n == 1 # common case: repeating a single-byte string
        @inbounds b = codeunit(x, 1)
        ccall(:memset, Ptr{Cvoid}, (Ptr{UInt8}, Cint, Csize_t), out, b, r)
    else
        for i = 0:r-1
            ref = Ref{T}(_bswap(x))
            GC.@preserve ref out begin
                ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
                unsafe_copyto!(pointer(out, i * n + 1), ptr, n)
            end
        end
    end
    return out
end

# copy/pasted from strings/util.jl
function Base.startswith(a::T, b::Union{String, SubString{String}, InlineString}) where {T <: InlineString}
    cub = ncodeunits(b)
    ncodeunits(a) < cub && return false
    ref = Ref{T}(_bswap(a))
    return GC.@preserve ref begin
        ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
        if Base._memcmp(ptr, b, sizeof(b)) == 0
            nextind(a, cub) == cub + 1
        else
            false
        end
    end
end

function Base.endswith(a::T, b::Union{String, SubString{String}, InlineString}) where {T <: InlineString}
    cub = ncodeunits(b)
    astart = ncodeunits(a) - ncodeunits(b) + 1
    astart < 1 && return false
    ref = Ref{T}(_bswap(a))
    return GC.@preserve ref begin
        ptr = convert(Ptr{UInt8}, Base.unsafe_convert(Ptr{T}, ref))
        if Base._memcmp(ptr + (astart - 1), b, sizeof(b)) == 0
            thisind(a, astart) == astart
        else
            false
        end
    end
end

Base.match(r::Regex, s::InlineString, i::Integer) = match(r, String(s), i)

# the rest of these methods are copy/pasted from Base strings/string.jl file
# for efficiency
Base.@propagate_inbounds function Base.isvalid(x::InlineString, i::Int)
    @boundscheck checkbounds(Bool, x, i) || throw(BoundsError(x, i))
    return @inbounds thisind(x, i) == i
end

Base.@propagate_inbounds function Base.thisind(s::InlineString, i::Int)
    i == 0 && return 0
    n = ncodeunits(s)
    i == n + 1 && return i
    @boundscheck Base.between(i, 1, n) || throw(BoundsError(s, i))
    @inbounds b = codeunit(s, i)
    (b & 0xc0 == 0x80) & (i-1 > 0) || return i
    @inbounds b = codeunit(s, i-1)
    Base.between(b, 0b11000000, 0b11110111) && return i-1
    (b & 0xc0 == 0x80) & (i-2 > 0) || return i
    @inbounds b = codeunit(s, i-2)
    Base.between(b, 0b11100000, 0b11110111) && return i-2
    (b & 0xc0 == 0x80) & (i-3 > 0) || return i
    @inbounds b = codeunit(s, i-3)
    Base.between(b, 0b11110000, 0b11110111) && return i-3
    return i
end

Base.@propagate_inbounds function Base.nextind(s::InlineString, i::Int)
    i == 0 && return 1
    n = ncodeunits(s)
    @boundscheck Base.between(i, 1, n) || throw(BoundsError(s, i))
    @inbounds l = codeunit(s, i)
    (l < 0x80) | (0xf8 ≤ l) && return i+1
    if l < 0xc0
        i′ = @inbounds thisind(s, i)
        return i′ < i ? @inbounds(nextind(s, i′)) : i+1
    end
    # first continuation byte
    (i += 1) > n && return i
    @inbounds b = codeunit(s, i)
    b & 0xc0 ≠ 0x80 && return i
    ((i += 1) > n) | (l < 0xe0) && return i
    # second continuation byte
    @inbounds b = codeunit(s, i)
    b & 0xc0 ≠ 0x80 && return i
    ((i += 1) > n) | (l < 0xf0) && return i
    # third continuation byte
    @inbounds b = codeunit(s, i)
    ifelse(b & 0xc0 ≠ 0x80, i, i+1)
end

Base.@propagate_inbounds function Base.iterate(s::InlineString, i::Int=firstindex(s))
    (i % UInt) - 1 < ncodeunits(s) || return nothing
    b = @inbounds codeunit(s, i)
    u = UInt32(b) << 24
    Base.between(b, 0x80, 0xf7) || return reinterpret(Char, u), i+1
    return iterate_continued(s, i, u)
end

function iterate_continued(s::InlineString, i::Int, u::UInt32)
    u < 0xc0000000 && (i += 1; @goto ret)
    n = ncodeunits(s)
    # first continuation byte
    (i += 1) > n && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 16
    # second continuation byte
    ((i += 1) > n) | (u < 0xe0000000) && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 8
    # third continuation byte
    ((i += 1) > n) | (u < 0xf0000000) && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b); i += 1
@label ret
    return reinterpret(Char, u), i
end

Base.@propagate_inbounds function Base.getindex(s::InlineString, i::Int)
    b = codeunit(s, i)
    u = UInt32(b) << 24
    Base.between(b, 0x80, 0xf7) || return reinterpret(Char, u)
    return getindex_continued(s, i, u)
end

function getindex_continued(s::InlineString, i::Int, u::UInt32)
    if u < 0xc0000000
        # called from `getindex` which checks bounds
        @inbounds isvalid(s, i) && @goto ret
        Base.string_index_err(s, i)
    end
    n = ncodeunits(s)

    (i += 1) > n && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 1
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 16

    ((i += 1) > n) | (u < 0xe0000000) && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 2
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 8

    ((i += 1) > n) | (u < 0xf0000000) && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 3
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b)
@label ret
    return reinterpret(Char, u)
end

Base.length(s::InlineString) = length_continued(s, 1, ncodeunits(s), ncodeunits(s))

Base.@propagate_inbounds function Base.length(s::InlineString, i::Int, j::Int)
    @boundscheck begin
        0 < i ≤ ncodeunits(s)+1 || throw(BoundsError(s, i))
        0 ≤ j < ncodeunits(s)+1 || throw(BoundsError(s, j))
    end
    j < i && return 0
    @inbounds i, k = thisind(s, i), i
    c = j - i + (i == k)
    length_continued(s, i, j, c)
end

@inline function length_continued(s::InlineString, i::Int, n::Int, c::Int)
    i < n || return c
    @inbounds b = codeunit(s, i)
    @inbounds while true
        while true
            (i += 1) ≤ n || return c
            0xc0 ≤ b ≤ 0xf7 && break
            b = codeunit(s, i)
        end
        l = b
        b = codeunit(s, i) # cont byte 1
        c -= (x = b & 0xc0 == 0x80)
        x & (l ≥ 0xe0) || continue

        (i += 1) ≤ n || return c
        b = codeunit(s, i) # cont byte 2
        c -= (x = b & 0xc0 == 0x80)
        x & (l ≥ 0xf0) || continue

        (i += 1) ≤ n || return c
        b = codeunit(s, i) # cont byte 3
        c -= (b & 0xc0 == 0x80)
    end
end

# Parsers.xparse
function Parsers.xparse(::Type{T}, source::Union{AbstractVector{UInt8}, IO}, pos, len, options::Parsers.Options, ::Type{S}=T)::Parsers.Result{S} where {T <: InlineString, S}
    res = Parsers.xparse(String, source, pos, len, options)
    code = res.code
    overflowed = false
    poslen = res.val
    if Parsers.invalid(code) || Parsers.sentinel(code)
        x = T()
    else
        poslen = res.val
        if T === InlineString1
            if poslen.len != 1
                overflowed = true
                x = T()
            else
                Parsers.fastseek!(source, poslen.pos)
                x = InlineString1(Parsers.peekbyte(source, poslen.pos))
                Parsers.fastseek!(source, pos + res.tlen - 1)
            end
        elseif Parsers.escapedstring(code) || !(source isa AbstractVector{UInt8})
            if poslen.len > (sizeof(T) - 1)
                overflowed = true
                x = T()
            else
                # manually build up InlineString
                i = poslen.pos
                maxi = i + poslen.len
                x = T()
                Parsers.fastseek!(source, i - 1)
                while i < maxi
                    b = Parsers.peekbyte(source, i)
                    if b == options.e
                        i += 1
                        Parsers.incr!(source)
                        b = Parsers.peekbyte(source, i)
                    end
                    x, overflowed = addcodeunit(x, b)
                    i += 1
                    Parsers.incr!(source)
                end
                Parsers.fastseek!(source, maxi)
            end
        else
            vlen = poslen.len
            if vlen > (sizeof(T) - 1)
                # @show T, vlen, sizeof(T)
                overflowed = true
                x = T()
            else
                # @show poslen.pos, vlen
                x = T(source, poslen.pos, vlen)
            end
        end
    end
    if overflowed
        code |= Parsers.OVERFLOW
    end
    return Parsers.Result{S}(code, res.tlen, x)
end


text:
using POMDPModels
using Test
using POMDPTesting

let
    problem = SimpleGridWorld()

    policy = RandomPolicy(problem)

    sim = HistoryRecorder(rng=MersenneTwister(1), max_steps=1000)

    hist = simulate(sim, problem, policy, GWPos(1,1))

    for (s, a) in zip(state_hist(hist), action_hist(hist))
        td = transition(problem, s, a)
        if td isa SparseCat
            @test sum(td.probs) ≈ 1.0 atol=0.01
            for p in td.probs
                @test p >= 0.0
            end
        end
    end

    sv = convert_s(Array{Float64}, GWPos(1,1), problem)
    @test sv == [1.0, 1.0]
    sv = convert_s(Array{Float64}, GWPos(5,3), problem)
    @test sv == [5.0, 3.0]
    s = convert_s(GWPos, sv, problem)
    @test s == GWPos(5, 3)

    av = convert_a(Array{Float64}, :up, problem)
    @test av == [1.0]
    a = convert_a(Symbol, av, problem)
    @test a == :up

    @test has_consistent_transition_distributions(problem)

    pol = FunctionPolicy(x->:up)
    stp = first(stepthrough(problem, pol, "s,a", max_steps=1))
    POMDPModelTools.render(problem, stp)
    POMDPModelTools.render(problem, NamedTuple())
    POMDPModelTools.render(problem, stp, color=s->reward(problem,s))
    POMDPModelTools.render(problem, stp, color=s->rand())
    POMDPModelTools.render(problem, stp, color=s->"yellow")
    POMDPModelTools.render(problem, stp, color=s->reward(problem,s), colormin=-1.0, colormax=1.0)

    ss = collect(states(problem))
    isd = initialstate(problem)
    for s in ss
        if !isterminal(problem, s)
            @test s in support(isd)
            @test pdf(isd, s) > 0.0
        end
    end
end

let
    @warn("NBInclude tests skipped")
    # @nbinclude(joinpath(dirname(@__FILE__), "..", "notebooks", "GridWorld Visualization.ipynb"))
end


text:
using Base: min, max

using .Geometry

export
    set!, expand!

mutable struct AABB{T <: AbstractFloat}
    # Top-left corner
    min::Geometry.Point{T}
    # Bottom-right corner
    max::Geometry.Point{T}

    function AABB{T}() where {T <: AbstractFloat}
        new(Point{T}(), Point{T}())
    end

    function AABB{T}(min::Geometry.Point{T}, max::Geometry.Point{T}) where {T <: AbstractFloat}
        new(copy(min), copy(max))
    end

end

# Set and expand if needed. p0,p1,p2 typically represent a triangle
function set!(aabb::AABB{T}, p0::Geometry.Point{T}, p1::Geometry.Point{T}, p2::Geometry.Point{T}) where {T <: AbstractFloat}
    aabb.min.x = min(p0.x, min(p1.x, p2.x));
    aabb.min.y = min(p0.y, min(p1.y, p2.y));

    aabb.max.x = max(p0.x, max(p1.x, p2.x));
    aabb.max.y = max(p0.y, max(p1.y, p2.y));
end

function set!(aabb::AABB{T}, vertices::AbstractArray{Geometry.Point{T}}) where {T <: AbstractFloat}
    minx = typemax(T);
    miny = typemax(T);
    maxx = typemin(T);
    maxy = typemin(T);

    for v in vertices
        minx = min(minx, v.x);
        maxx = max(maxx, v.x);
        miny = min(miny, v.y);
        maxy = max(maxy, v.y);
    end

    Geometry.set!(aabb.min, minx, miny);
    Geometry.set!(aabb.max, maxx, maxy);
end

function expand!(aabb::AABB{T}, vertices::AbstractArray{Geometry.Point{T}}) where {T <: AbstractFloat}
    minx = aabb.min.x;
    miny = aabb.min.y;
    maxx = aabb.max.x;
    maxy = aabb.max.y;

    for v in vertices
        minx = min(minx, v.x);
        maxx = max(maxx, v.x);
        miny = min(miny, v.y);
        maxy = max(maxy, v.y);
    end

    Geometry.set!(aabb.min, minx, miny);
    Geometry.set!(aabb.max, maxx, maxy);
end

max_stars_repo_name: wdevore/Ranger-Julia-SDL

max_issues_repo_name: wdevore/Ranger-Julia-SDL

max_forks_repo_name: wdevore/Ranger-Julia-SDL


text:
import Pkg
Pkg.test("BenchmarkTools")
Pkg.test("StaticArrays")
Pkg.test("UnicodePlots")

max_stars_repo_name: dev-juliacn/ubuntu

max_issues_repo_name: dev-juliacn/ubuntu

max_forks_repo_name: dev-juliacn/ubuntu


text:
macro d64_str(val::String)
    :(Double64(BigFloat($val)))
end

macro d32_str(val::String)
    :(Double32(BigFloat($val)))
end

macro d16_str(val::String)
    :(Double16(BigFloat($val)))
end

function splitnumstring(str::AbstractString, dlm)
    strs = String.(split(str, dlm))
    strs = map(x->(isempty(x) ? "0.0" : x), strs)
    n = length(strs)
    if n > 1
        strs
    elseif n == 1
        [strs[1], "0.0"]
    else
        ["0.0", "0.0"]
    end
end

# constructors
Double16(str::T) where {T<:AbstractString} = tryparse(Double16, str)
Double32(str::T) where {T<:AbstractString} = tryparse(Double32, str)
Double64(str::T) where {T<:AbstractString} = tryparse(Double64, str)

function Base.tryparse(::Type{Double64}, str::T) where {T<:AbstractString}
     str = strip(str)
     if startswith(str, "DoubleFloat{Float64}")
         str = string("Double64", str[length("DoubleFloat{Float64}")+1:end])
     end
     if !startswith(str, "Double64(")
         str = string("Double64(", str, ")")
     end
     return trytypedparse(Double64, str)
end

function Base.tryparse(::Type{Double32}, str::T) where {T<:AbstractString}
     str = strip(str)
     if startswith(str, "DoubleFloat{Float32}")
         str = string("Double32", str[length("DoubleFloat{Float32}")+1:end])
     end
     if !startswith(str, "Double32(")
         str = string("Double32(", str, ")")
     end
     return trytypedparse(Double32, str)
end

function Base.tryparse(::Type{Double16}, str::T) where {T<:AbstractString}
     str = strip(str)
     if startswith(str, "DoubleFloat{Float16}")
         str = string("Double16", str[length("DoubleFloat{Float16}")+1:end])
     end
     if !startswith(str, "Double16(")
         str = string("Double16(", str, ")")
     end
     return trytypedparse(Double16, str)
end

function trytypedparse(::Type{Double64}, str::T) where {T<:AbstractString}
     str = str[1+length("Double64("):length(str)-1]
     histr, lostr = splitnumstring(str, ", ")
     hi = Meta.parse(histr)
     lo = Meta.parse(lostr)
     result = Double64(hi, lo)
     return result
end

function trytypedparse(::Type{Double32}, str::T) where {T<:AbstractString}
     str = str[1+length("Double32("):end-1]
     histr, lostr = splitnumstring(str, ", ")
     hi = Meta.parse(histr)
     lo = Meta.parse(lostr)
     result = Double32(hi, lo)
     return result
end

function trytypedparse(::Type{Double16}, str::T) where {T<:AbstractString}
     str = str[1+length("Double16("):end-1]
     histr, lostr = splitnumstring(str, ", ")
     hi = Meta.parse(histr)
     lo = Meta.parse(lostr)
     result = Double16(hi, lo)
     return result
end


text:
module MedianFilter

export medfilt1, meanfilt

# ord = <  -> max-heap
# ord = >= -> min-heap

function heapify(n,i,f,swap,ord)
    l = 2*i
    r = 2*i+1
    if r <= n
        m = ord(f(r),f(l)) ? l : r
        if ord(f(i),f(m))
            swap(i,m)
            heapify(n,m,f,swap,ord)
        end
    elseif l == n && ord(f(i),f(l))
        swap(i,l)
    end
end

function bubble(n,i,f,swap,ord)
    if i <= 1
        return
    end
    m = div(i,2)
    if ord(f(m),f(i))
        swap(i,m)
        bubble(n,m,f,swap,ord)
    end
end

swp(a,i,j) = a[i],a[j]=a[j],a[i]

type MedHeap
    a
    ord
    f
    swap
    MedHeap(a,ord,values,lookup) = (m = new(a,ord);
        m.f = i -> values[m.a[i]];
        m.swap = (i,j) -> (swp(a,i,j); swp(lookup,m.a[i],m.a[j]));
        m
    )
end

loheap(w,values,lookup) = MedHeap(collect(1:div(w,2)),<,values,lookup)
hiheap(w,values,lookup) = MedHeap(collect(div(w,2)+1:w),>=,values,lookup)

heapify(h::MedHeap,i) = heapify(length(h.a),i,h.f,h.swap,h.ord)
bubble(h::MedHeap,i) = bubble(length(h.a),i,h.f,h.swap,h.ord)
# previous value x at i was replaced
function replaced(h,x,i)
    if h.ord(h.f(i),x)
        heapify(h,i)
    else
        bubble(h,i)
    end
end
top(h::MedHeap) = h.f(1)

type Med{T}
    # value buffer
    values::Vector{T}
    # lookup to heaps
    lookup::Vector{Int}
    # rolling index to values and lookup buffers
    index::Int
    # max heap
    lo::MedHeap
    # min heap
    hi::MedHeap
    Med(x::T,w) = (m = new(fill(x,w),collect(1:w),1);
        m.lo = loheap(w,m.values,m.lookup);
        m.hi = hiheap(w,m.values,m.lookup);
        m)
end

Med{T}(x::T,w) = Med{T}(x,w)

function push_med{T}(m::Med{T},x::T)
    # leaving value
    v = m.values[m.index]
    # current index in heap
    i = m.lookup[m.index]
    n = length(m.lo.a)
    # replace with incoming value
    m.values[m.index] = x
    # check which heap was modified
    a,b,j = i > n ? (m.hi,m.lo,i-n) : (m.lo,m.hi,i)
    # fix heap
    replaced(a,v,j)
    # check if top elements should be swapped
    if top(m.lo) > top(m.hi)
        m.lo.a[1],m.hi.a[1] = m.hi.a[1],m.lo.a[1]
        swp(m.lookup,m.lo.a[1],m.hi.a[1])
        heapify(b,1)
    end
    m.index = m.index == length(m.values) ? 1 : m.index+1
    # return current median
    bool(length(m.values) & 1) ? top(m.hi) : (top(m.lo) + top(m.hi))/2
end

function medfilt1(x,w)
    n = length(x)
    m = Med(x[1],w)
    w1 = div(w,2)
    for i = 1:w1-1
        push_med(m,x[i>n?n:i])
    end
    [push_med(m,x[i+w1>n?n:i+w1]) for i=1:n]
end

# central simple moving average along the first dimension, repeat edges
function meanfilt(x,w)
    # divide window
    w1 = div(w,2)+1
    w2 = w-w1
    n = size(x,1)
    s = size(x)[2:end]
    # value buffer
    b = [repmat(x[1,:]/w,w1),
         x[1:min(w2,n),:]./w,
         repmat(x[end,:]/w,w-min(w2,n)-w1)]
    # previous mean
    m = sum(b,1)
    r = zeros(size(x))
    j = 1
    for i = 1:n
        # next value
        p = reshape(x[min(i+w2,n),:],1,s...)./w
        r[i,:] = m - b[j,:] + p
        m = r[i,:]
        b[j,:] = p
        j = j == w ? 1 : j+1
    end
    r
end

end # module

max_stars_repo_name: jlep/MedianFilter


text:
using Base.Test.@test
using Modifyfield
using Modifyfield.@modify_field!
using Modifyfield.@modify_fields!
using Modifyfield.@modify_tuple_entry!
using Modifyfield.copy_and_modify_tup

immutable TestT
    intfld::Int
    boolfld::Bool
end

function testmodifyfield()
    n = 3
    a = Array(TestT,n)
    b = Array(TestT,n,n)
    for i = 1 : n
        a[i] = TestT(i,false)
        for j = 1 : n
            b[i,j] = TestT(j,true)
        end
    end
    for i = 1 : n
        @modify_field! a[i].intfld = i - 1
    end
    for i = 1 : n
        @test a[i].intfld == i - 1 && a[i].boolfld == false
    end
    for i = 1 : n
        for j = 1 : n
            @modify_field! b[i,j].boolfld = false
        end
    end
    for i = 1 : n
        for j = 1 : n
            @test b[i,j].intfld == j && b[i,j].boolfld == false
        end
    end
    nothing
end

function testmodifytuple()
    t = (5.5, 6.6, 7.7)
    @modify_tuple_entry! t[2] = true
    @test t == (5.5, true, 7.7)
    @modify_tuple_entry! t[3] = "a"
    @test t == (5.5, true, "a")
    @modify_tuple_entry! t[1] = Int
    @test t == (Int, true, "a")
    n = 10
    a = (Tuple{Int,Int})[]
    resize!(a,n)
    for j = 1 : n
        a[j] = (j,j)
    end
    for j = 1 : n
        @modify_tuple_entry! a[j][2] = j*j
    end
    for j = 1 : n
        @test a[j][1] == j && a[j][2] == j*j
    end
    nothing
end

function testmodifytuple2()
    t = (0,1,2)
    for i = 1 : 3
        # The macro call
        # @modify_tuple_entry! t[i] = i
        # fails (variable subscript).
        #
        # The function-call variant in
        # the next statement is dispatched at
        # run time instead of compile time because
        # the type of Val{i} is unknown at
        # compile time.  This leads to poorer
        # performance.
        t = copy_and_modify_tup(t, Val{i}, i)
    end
    @test t == (1,2,3)
    nothing
end

immutable TestT2
    a::Int
    b::Bool
    c::Float64
end

function testmodifyfields()
    v = TestT2[]
    push!(v, TestT2(4,false,65.2))
    push!(v, TestT2(3,true, 94.2))
    z = 64.5
    @modify_fields! v[1].(a=9,c=z+.5)
    @modify_fields! v[2].(c=z-0.5, b=false)
    @test v[1] == TestT2(9, false, 65.0)
    @test v[2] == TestT2(3, false, 64.0)
    @modify_fields! v[2].(a=-2,)
    @test v[2] == TestT2(-2, false, 64.0)
    nothing
end

println("starting tests...")
testmodifyfield()
testmodifytuple()
testmodifytuple2()
testmodifyfields()
println("tests finished")


text:
# this potential is an α-attractor model with a quadratic potential
# this run script is set to run on a cluster with 24 total physical CPU cores
# the simulations are scanned over the model, for various values of α, N0, and Nf
# each simulation, if it successfully computes perturbations is saved in a data folder

const NUMCPU = 24
Nfs = [2,5,10]
runs = 100
alphas = [1//6, 1//60]
N0s = 75:25:400

datafolder = "/data/ANflation/"
subfolder = "batch1"

if !isdir(datafolder)
    println("Data folder doesn't exist!")
elseif !isdir(joinpath(datafolder,subfolder))
    println("Creating $(joinpath(datafolder,subfolder)) ...")
    mkdir(joinpath(datafolder,subfolder))
end

using Distributed
if length(workers()) == 1
    addprocs(NUMCPU - 1)
end

@everywhere using Inflation
@everywhere using SymPy

@everywhere begin
    using Random
    using Roots
    using ForwardDiff
    m = 1e-5
    β = 1/2
    a = m^2 * (1-sqrt(β))^2
    b = m^2 * (1+sqrt(β))^2
    function MarchenkoPasturPDF(x)
        return sqrt( (b-x)*(x-a)) / (2*pi*x*β*m^2)
    end
    support = (a,b)
    function rejection_sample(pdf,num,support=(-Inf,Inf))
        # for now assume pdf has exactly one critical point
        # will probably fail without finite support
        a = support[1]
        b = support[2]
        dpdf(x) = ForwardDiff.derivative(pdf,x)
        xmax = find_zero(dpdf,(a,b))
        pdf_max = pdf(xmax)
        samples = Array{Float64}(undef,num)
        i = 1
        while i <= num
            x = a+(b-a)*rand()
            s = rand()*pdf_max
            if s < pdf(x)
                samples[i] = x
                i += 1
            end
        end
        return samples
    end
end

for d in Nfs
    function G(Phi,p)
        d = length(Phi)
        α = p[1]
        Phi2 = sum([f^2 for f in Phi])
        g = Diagonal(ones(Sym,d) .* 6*α / (1 - Phi2)^2)
        return g
    end

    function V(Phi,p)
        m = p[2:end]
        α = p[1]
        v = 0
        for i in 1:d
            v += α/2.0*(m[i]*Phi[i]^2)
        end
        return v
    end

    runs_to_run = []
    for N0 in N0s, alpha in alphas, run in 1:runs
        fname = "N$(N0)_a$(Int(1/alpha))_Nf$(d)_r$run.jld2"
        fullfname = joinpath(datafolder,subfolder,fname)
        if !isfile(fullfname)
            push!(runs_to_run,(N0,alpha,fullfname))
        end
    end

    if length(runs_to_run) == 0
        println("All runs done for d=$d")
        continue
    end
    println("Running $(length(runs_to_run)) runs at d=$d")
    # all params should be lowercase
    mparams = [symbols("m$i",real=true,positive=true) for i in 1:d]
    params = [symbols("α",real=true,positive=true),mparams...]

    options = SetupOptions(
                           attempt_HCA = false,
                           return_quotes = true
                          )
    quotes = inflation_setup(d,V,G,params,options=options)
    @everywhere funcs = Inflation.make_funcs($quotes)
    @everywhere runs_to_run = $runs_to_run
    @sync @distributed for (N0,alpha,fullfname) in runs_to_run
        println(fullfname)
        while !isfile(fullfname)
            mass_eigs = rejection_sample(MarchenkoPasturPDF,d,support)
            pvalues = [alpha,mass_eigs...]

            r0 =sqrt( (2*N0 + 3*sqrt(alpha) - 3*alpha) / (2*N0 + 3*sqrt(alpha)))
            ## get a vector on the sphere, use gaussian trick
            # N0 prior should be approximately true
            Phi0 = randn(d)
            Phi0 ./= norm(Phi0)
            Phi0 .*= r0

            Pi0 = zeros(d)
            rNe = 3*alpha/2*(1/(1-Phi0'*Phi0)-1/sqrt(alpha))
            println("estimated radial efolds: $rNe")
            println("total energy: $(V(Phi0,pvalues))")

            #analytics
            if d==2
                θ0 = Base.atan(Phi0[2],Phi0[1])
                R = pvalues[end]/pvalues[end-1]
                a = pvalues[1]
                println("R: $R, α: $a, θ0: $θ0")
                Nang = -1/(54*a)*( (18*a+4)*log(0.5*(-(R-1)*cos(2*θ0)+R+1)) - 2*(R+2) + 8*(-(R^2-1)*cos(2*θ0)+R^2+R+1) / (-(R-1)*cos(2*θ0)+R+1)^2 )
                println("Nang: $Nang")
                Neff = Nexit_to_end - Nang
                if Neff < 0
                    println("Exit horizon during angular motion, no analytic result.")
                else
                    anal_ns = 1 - 2/Neff
                    anal_r  = 12*a/Neff^2
                    println("Expected ns: $anal_ns")
                    println("Expected r: $anal_r")
                end
            else
                println("d>2, no analytic result")
            end

            sol = background_evolve(Phi0,Pi0,pvalues,funcs,verbose=true)
            Nend = sol.t[end]
            Nearly = Nend - Nexit_to_end - Nsubhorizon - Nkrange
            if Nearly > 0
                tsol = transport_perturbations(sol,pvalues,funcs,verbose=true)
                output_data(fullfname,funcs,pvalues,sol,tsol)
            else
                println("Not enough inflation! retrying...")
                continue
            end
        end
    end
end
rmprocs(NUMCPU-1)


text:
# Discrete OU process ALA Charlie Doering
function ou_transition_matrix(n)
    Mⱼₖ = zeros(n + 1, n + 1)
    δ(j, k) = (j == k) ? 1 : 0

    for j in 0:n, k in 0:n
        jj = j + 1
        kk = k + 1
        Mⱼₖ[jj, kk] =
            (-n * δ(j, k) + k * δ(j + 1, k) + (n - k) * δ(j - 1, k)) / 2
    end
    return Mⱼₖ
end

# Discrete Phase Transition ala Glenn Flierl
function uniform_phase(N)
    @assert N > 1
    T = zeros(N, N)
    for i in 1:N
        T[i, i] = -1.0
        T[i, i%N+1] = 1.0
    end
    return T
end

# Discrete Laplacian
function discrete_laplacian(n)
    Δ = zeros(n, n)
    for i in 2:n-1
        Δ[i, i] = -2
        Δ[i, i+1] = 1
        Δ[i, i-1] = 1
    end
    Δ[1, 1] = Δ[end, end] = -1
    Δ[1, 2] = 1
    Δ[end, end-1] = 1
    return Δ
end

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
# http://docs.julialang.org/en/stable/stdlib/numbers/#random-numbers

name_history = String[]

function new_name()
    generate_name() = join(map(x->Char(x), rand(65:90, 2))) * repr(rand(100:999))
    name = generate_name()
    while name in name_history
        name = generate_name()
    end
    push!(name_history, name)
    name
end

mutable struct Robot
    name::AbstractString

    Robot() = new(new_name())
end

function reset!(instance::Robot)
    instance.name = new_name()
    instance
end


text:
# ---
# title: 174. Dungeon Game
# id: problem174
# author: Indigo
# date: 2021-01-20
# difficulty: Hard
# categories: Binary Search, Dynamic Programming
# link: <https://leetcode.com/problems/dungeon-game/description/>
# hidden: true
# ---
#
# The demons had captured the princess ( **P** ) and imprisoned her in the
# bottom-right corner of a dungeon. The dungeon consists of M x N rooms laid out
# in a 2D grid. Our valiant knight ( **K** ) was initially positioned in the
# top-left room and must fight his way through the dungeon to rescue the
# princess.
#
# The knight has an initial health point represented by a positive integer. If
# at any point his health point drops to 0 or below, he dies immediately.
#
# Some of the rooms are guarded by demons, so the knight loses health (
# _negative_ integers) upon entering these rooms; other rooms are either empty (
# _0 's_) or contain magic orbs that increase the knight's health ( _positive_
# integers).
#
# In order to reach the princess as quickly as possible, the knight decides to
# move only rightward or downward in each step.
#
#
#
# **Write a function to determine the knight 's minimum initial health so that
# he is able to rescue the princess.**
#
# For example, given the dungeon below, the initial health of the knight must be
# at least **7** if he follows the optimal path `RIGHT-> RIGHT -> DOWN -> DOWN`.
#
# -2 (K) | -3 | 3
# ---|---|---
# -5 | -10 | 1
# 10 | 30 | -5 (P)
#
#
#
# **Note:**
#
#   * The knight's health has no upper bound.
#   * Any room can contain threats or power-ups, even the first room the knight enters and the bottom-right room where the princess is imprisoned.
#
#
## @lc code=start
using LeetCode

function calculate_minimum_hp(dungeon::Vector{Vector{Int}})
    m, n = length(dungeon), length(dungeon[1])
    dp = fill(typemax(Int) , m + 1, n + 1)
    dp[m+1, n] = dp[m, n+1] = 1
    for j in n:-1:1
        for i in m:-1:1
            dp[i, j] = max(min(dp[i+1, j], dp[i, j+1]) - dungeon[i][j], 1)
        end
    end
    dp[1, 1]
end
## @lc code=end


text:
#=
phase_vocoder.jl

Julia translation of Matlab code from:
https://sethares.engr.wisc.edu/vocoders/matlabphasevocoder.html
5/2005 Bill Sethares (original Matlab version)
2022-02-26 Jeff Fessler (this Julia version)
=#

export phase_vocoder

using FFTW: fft, ifft
#using Sound: hann, findpeaks

# Process multi-channel data (e.g., stereo) one channel at a time.
function phase_vocoder(
    y::AbstractMatrix{<:Real},
    sr::Real = framerate(y) ; # sampling rate (in Hz)
    T::DataType = Float32,
    kwargs...
)
    mapslices(x -> phase_vocoder(x, sr; T, kwargs...), y, dims=1)::Matrix{T}
end

"""
    phase_vocoder(x, sr = framerate(x); kwargs...)

Phase vocoder for time scaling of signal `x` having sampling rate `sr` in Hz.
The time-stretch is determined by the ratio of `hopin` and `hopout` variables.
For example, `hopin=242` and `hopout=161.3333` (integers are not required)
increases the tempo by `hopin/hopout = 1.5`.
To slow down a comparable amount, choose `hopin = 161.3333`, `hopout = 242`.

# Option
* `time::Real = length(x) * sr` : total time to process (in sec)
* `hopin::Real = 121` : hop length for input
* `hopout::Real = 2*hopin` : hop length for output
* `all2pi::Any = 2π*(0:100)` multiples of 2π (used in PV-style freq search)
* `max_peak::Int = 50` : parameters for peak finding: number of peaks
* `eps_peak::Real = 0.005` : minimum height of peaks
* `nfft::Int = 2^12` : fft length
* `win::AbstractVector{<:Real} = hann(nfft)` : window
* `T::DataType = Float32` : data type
"""
function phase_vocoder(
    x::AbstractVector{<:Real},
    sr::Real = framerate(x) ; # sampling rate (in Hz)
    time::Real = length(x) * sr, # total time to process (in sec)
    hopin::Real = 121, # hop length for input
    hopout::Real = 2*hopin, # hop length for output
    all2pi::Any = 2π*(0:100), # all multiples of 2π (used in PV-style freq search)
    max_peak::Int = 50, # parameters for peak finding: number of peaks
    eps_peak::Real = 0.005, # height of peaks
    nfft::Int = 2^12, # fft length
    win::AbstractVector{<:Real} = hann(nfft), # window
    chat::Int = 0,
    T::DataType = Float32,
)

    nfft2 = nfft ÷ 2
    N1 = min(length(x), Int(time*sr)) # length of processed audio in samples
    yy = zeros(T, ceil(Int, hopout/hopin) * N1) # output
    lenseg = floor(Int, (N1 - nfft) / hopin) # number of nfft segments to process

    ssf = sr * (0:nfft2) / nfft # frequency vector (k/N)*S
    phold = zeros(T, nfft2+1)
    phadvance = zeros(T, nfft2+1)
    dtin = hopin / sr # time advances dt per hop for input
    dtout = hopout / sr # time advances dt per hop for output

    for k in 0:lenseg-2 # main loop - process each beat separately
        (chat > 0) && mod1(k, chat) == 1 && println(k) # show where we are
        indin = round.(Int, k * hopin .+ (1:nfft))

        s = win .* x[indin] # get this frame and take FFT
        ffts = fft(s)
        mag = abs.(ffts[1:nfft2+1])
        ph = angle.(ffts[1:nfft2+1])

        # find peaks to define spectral mapping
        peaks = findpeaks(mag, max_peak, eps_peak)
        inds = sortperm(mag[peaks[:,2]])
        peaksort = peaks[inds,:]
        pc = peaksort[:,2]

        bestf = zeros(T, size(pc))
        for tk in 1:length(pc) # estimate frequency using PV strategy
            dtheta = (ph[pc[tk]] - phold[pc[tk]]) .+ all2pi
            fest = dtheta / (2π*dtin) # see pvanalysis.m for same idea
            indf = argmin(abs.(ssf[pc[tk]] .- fest))
            bestf[tk] = fest[indf] # find best freq estimate for each row
        end

        # generate output mag and phase
        magout = copy(mag)
        phout = copy(ph)
        for tk in 1:length(pc)
            fdes = bestf[tk] # reconstruct with original frequency
            freqind = peaksort[tk,1]:peaksort[tk,3] # indices of surrounding bins

            # specify magnitude and phase of each partial
            magout[freqind] .= mag[freqind]
            phadvance[peaksort[tk,2]] += 2π * fdes * dtout
            pizero = fill(Float64(π), length(freqind))
            pcent = peaksort[tk,2] - peaksort[tk,1] + 1
            indpc = (2 - mod(pcent,2)):2:length(freqind)
            pizero[indpc] .= 0
            phout[freqind] .= phadvance[peaksort[tk,2]] .+ pizero
        end

        # reconstruct time signal (stretched or compressed)
        compl = magout .* exp.(im * phout)
        compl[nfft2+1] = ffts[nfft2+1]
        compl = [compl; reverse(conj.(compl[2:nfft2]))]
        wave = real(ifft(compl))
        phold .= ph

        indout = round.(Int, (k * hopout) .+ (1:nfft))
        yy[indout] .+= wave
    end

    return yy
end


text:
# convergence statistics
close("all")
transD_GP.getchi2forall(opt)
ax = gcf().axes;
ax[2].set_ylim(10, 30)
χ² = aem.ndatalow + aem.ndatahigh
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
transD_GP.getchi2forall(optlog10λ)
ax = gcf().axes;
ax[2].set_ylim(10, 30)
ax[2].plot(xlim(), [χ²/2 , χ²/2], "--", color="gray")
# posterior
opt.xall[:] .= zall
transD_GP.plot_posterior(aem, opt, optlog10λ, burninfrac=0.5,
    figsize=(15,6), cmappdf="inferno", nbins=100, qp1=0.05, qp2=0.95, vmaxpc=0.5)
ax = gcf().axes
ax[1].invert_xaxis()
ax[1].step(log10.(ρ[2:end]), z[2:end], color="w")
ax[1].step(log10.(ρ[2:end]), z[2:end], color="k", linestyle="--")

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA


text:
using DataFrames
df = DataFrame()
df["pb"]=rand(1:50, 50);
df["t"]=rand(1:50, 50);
df["h"]=[0.0 for i in 1:50];

using Gadfly
plot(df, x = "pb", y = "t", color = "h", Geom.rectbin)


text:
# This code is derived from `build_sysimg.jl` (part of Julia) and should be kept aligned with it.

function default_sysimg_path(debug = false)
    ext = debug ? "sys-debug" : "sys"
    if Sys.isunix()
        dirname(Libdl.dlpath(ext))
    else
        normpath(Sys.BINDIR, "..", "lib", "julia")
    end
end

function compile_system_image(sysimg_path, cpu_target = nothing; debug = false)
    # Enter base and setup some useful paths
    base_dir = dirname(Base.find_source_file("sysimg.jl"))
    cd(base_dir) do
        # This can probably get merged with build_object.
        # At some point, I will need to understand build_object a bit better before doing that move, though!
        julia_cmd = Base.julia_cmd()
        julia = julia_cmd.exec[1]
        cpu_target = if cpu_target === nothing
            replace(julia_cmd.exec[2], "-C" => "")
        else
            cpu_target
        end
        cc = system_compiler
        # Ensure we have write-permissions to wherever we're trying to write to
        try
            touch("$sysimg_path.ji")
        catch
            error("Unable to modify $sysimg_path.ji, ensure that parent directory exists and is writable")
        end
        compiler_path = joinpath(dirname(sysimg_path), "basecompiler")
        compiler = "compiler/compiler.jl"

        # Start by building inference.{ji,o}
        inference_path = joinpath(dirname(sysimg_path), "inference")
        command = `$julia -C $cpu_target --output-ji $compiler_path.ji --output-o $compiler_path.o $compiler`
        @info "Building `inference.o`:\n$command"
        run(command)

        # Bootstrap off of that to create sys.{ji,o}
        command = `$julia -C $cpu_target --output-ji $sysimg_path.ji --output-o $sysimg_path.o -J $compiler_path.ji --startup-file=no sysimg.jl`
        @info "Building `sys.o`:\n$command"
        run(command)

        build_shared(
            "$sysimg_path.$(Libdl.dlext)", "$sysimg_path.o", false,
            ".", true, nothing, debug ? 2 : nothing, cc, nothing
        )
    end
end

"""
Returns the system image file stored in the backup folder.
If there is no backup, this function will automatically generate a system image
in the backup folder.
"""
function get_backup!(debug, cpu_target = nothing)
    target = julia_cpu_target(cpu_target)
    sysimg_backup = sysimgbackup_folder(target)
    isdir(sysimg_backup) || mkpath(sysimg_backup)
    if !all(x-> isfile(joinpath(sysimg_backup, x)), sysimage_binaries) # we have a backup
        compile_system_image(joinpath(sysimg_backup, "sys"), target; debug = debug)
    end
    return joinpath(sysimg_backup, "sys.$(Libdl.dlext)")
end


text:
using Oceananigans, Random, Printf

# Set `makeplot=true` to enable plotting.
makeplot = false

macro withplots(ex)
    makeplot ? :($(esc(ex))) : :(nothing)
end

@withplots using PyPlot

####
#### Model set-up
####

# Two cases from Van Roekel et al (JAMES, 2018)
parameters = Dict(:free_convection => Dict(:Fb=>3.39e-8, :Fu=>0.0,     :f=>1e-4, :N²=>1.96e-5),
                  :wind_stress     => Dict(:Fb=>0.0,     :Fu=>9.66e-5, :f=>0.0,  :N²=>9.81e-5))

# Simulation parameters
case = :free_convection
 N = 32                   # Resolution
 Δ = 0.5                  # Grid spacing
tf = hour/2               # Final simulation time

N², Fb, Fu, f = (parameters[case][p] for p in (:N², :Fb, :Fu, :f))
βT, g = 2e-4, 9.81
Fθ, dTdz = Fb / (g*βT), N² / (g*βT)

# Create boundary conditions. Note that temperature = buoyancy because βT=g=1.
ubcs = HorizontallyPeriodicBCs(top=BoundaryCondition(Flux, Fu))
Tbcs = HorizontallyPeriodicBCs(top=BoundaryCondition(Flux, Fθ), bottom=BoundaryCondition(Gradient, dTdz))

# Instantiate the model
model = Model(
           architecture = CPU(), # GPU() # this example will run on the GPU if cuda is available.
                   grid = RegularCartesianGrid(N = (N, N, N), L = (N*Δ, N*Δ, N*Δ)),
               coriolis = FPlane(f=f),
               buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=βT)),
                closure = AnisotropicMinimumDissipation(), # closure = ConstantSmagorinsky(),
    boundary_conditions = BoundaryConditions(u=ubcs, T=Tbcs)
)

# Set initial condition. Initial velocity and salinity fluctuations needed for AMD.
Ξ(z) = randn() * z / model.grid.Lz * (1 + z / model.grid.Lz) # noise
T0(x, y, z) = 20 + dTdz * z + dTdz * model.grid.Lz * 1e-6 * Ξ(z)
ϵ₀(x, y, z) = 1e-4 * Ξ(z)

set!(model, u=ϵ₀, w=ϵ₀, T=T0, S=ϵ₀)

####
#### Set up output
####

u(model) = Array{Float32}(model.velocities.u.data.parent)
v(model) = Array{Float32}(model.velocities.v.data.parent)
w(model) = Array{Float32}(model.velocities.w.data.parent)
T(model) = Array{Float32}(model.tracers.T.data.parent)

fields = Dict(:u=>u, :v=>v, :w=>w, :T=>T)
filename = @sprintf("%s_n%d", case, N)
field_writer = JLD2OutputWriter(model, fields; dir="data", prefix=filename, interval=hour/4, force=true)
push!(model.output_writers, field_writer)

####
#### Run the simulation
####

@withplots fig, axs = subplots()

function terse_message(model, walltime, Δt)
    wmax = maximum(abs, model.velocities.w.data.parent)
    cfl = Δt / Oceananigans.cell_advection_timescale(model)
    return @sprintf("i: %d, t: %.4f hours, Δt: %.1f s, wmax: %.6f ms⁻¹, cfl: %.3f, wall time: %s\n",
                    model.clock.iteration, model.clock.time/3600, Δt, wmax, cfl, prettytime(walltime))
end

# A wizard for managing the simulation time-step.
wizard = TimeStepWizard(cfl=0.2, Δt=1.0, max_change=1.1, max_Δt=90.0)

# Run the model
while model.clock.time < tf
    update_Δt!(wizard, model)
    walltime = @elapsed time_step!(model, 10, wizard.Δt)
    @printf "%s" terse_message(model, walltime, wizard.Δt)

    @withplots begin
        sca(axs); cla()
        imshow(rotr90(view(data(model.velocities.w), :, 2, :)))
        gcf(); pause(0.01)
    end
end


text:
# types.jl

import Base: ==, show

immutable App
    url::AbstractString
end

abstract QueryChain

immutable Setter <: QueryChain
    lhs::Vector
    rhs::Vector
end

immutable Getter <: QueryChain
    lhs::Vector
end

immutable ResultInfo
    typ::Symbol
    value::Any
    address::Union{Void,AbstractString}

    function ResultInfo(typ::Symbol, value::Any)
        ResultInfo(typ, value, nothing)
    end

    function ResultInfo(typ::Symbol, value::Any, address::Union{Void,AbstractString})
        if :symbol == typ && "nothing" == value
            new(typ, nothing, address)
        else
            new(typ, value, address)
        end
    end

    function ResultInfo(dict::Dict)
        if haskey(dict, "address")
           ResultInfo(Symbol(dict["typ"]), dict["value"], dict["address"])
        else
           ResultInfo(Symbol(dict["typ"]), dict["value"], nothing)
        end
    end
end

immutable QueryResult
    info::ResultInfo
    app::App
    verb::AbstractString
    param::Dict

    function QueryResult(info::ResultInfo, app::App, verb::AbstractString, param::Dict)
        new(info, app, verb, param)
    end

    function QueryResult(dict::Dict, app::App, verb::AbstractString, param::Dict)
        new(ResultInfo(dict), app, verb, param)
    end
end

type CallArg
    name::Symbol
    args::Vector
end

function ==(lhs::App, rhs::App)
    return lhs.url == rhs.url
end

function ==(lhs::Setter, rhs::Setter)
    return lhs.lhs == rhs.lhs && lhs.rhs == rhs.rhs
end

function ==(lhs::Getter, rhs::Getter)
    return lhs.lhs == rhs.lhs
end

function ==(lhs::ResultInfo, rhs::ResultInfo)
    return lhs.typ == rhs.typ && lhs.value == rhs.value && lhs.address == rhs.address
end

function ==(lhs::QueryResult, rhs::QueryResult)
    return lhs.info == rhs.info && lhs.app == rhs.app && lhs.verb == rhs.verb && lhs.param == rhs.param
end

function Base.show(io::IO, result::QueryResult)
    show(io, result.info)
end

function Base.show(io::IO, info::ResultInfo)
    if :symbol == info.typ
        if isa(info.value, Void)
            #print(io, repr(info.value))
        else
            print_with_color(:red, io, info.value)
        end
    elseif :string == info.typ
        show(io, info.value)
    elseif isa(info.value, AbstractArray)
        len = length(info.value)
        for (idx,el) in enumerate(info.value)
            print_with_color(:cyan, io, string(idx-1))
            write(io, Base.color_normal)
            write(io, " ")
            write(io, el)
            idx < len && write(io, "\n")
        end
    elseif isa(info.value, Dict)
        show(io, info.value)
    else
        write(io, info.value)
    end
end


text:
"""
    getBlockIds

Get the vector of block IDs that are assigned to the current MPI rank

# Arguments
- `nblocks`: number of blocks (by default, `DSPProblem.nblocks`)
- `master_has_subblocks`: indicate whether the master process solves (scenario) blocks or not (should always be `false`). This should not be modified.
"""
function getBlockIds(nblocks::Int = dspenv.nblocks; master_has_subblocks::Bool = true)::Vector{Int}

    if getVersionMajor(dspenv) >= 2
        master_has_subblocks = false
    end

    # processor info
    mysize = dspenv.comm_size
    myrank = dspenv.comm_rank
    # empty block ids
    proc_idx_set = Int[]
    # DSP is further parallelized with mysize > nblocks.
    modrank = myrank % nblocks
    # If we have more than one processor,
    # do not assign a sub-block to the master.
    if master_has_subblocks
        # assign sub-blocks in round-robin fashion
        for s = modrank:mysize:(nblocks-1)
            push!(proc_idx_set, s+1)
        end
    else
        if mysize > 1
            if myrank == 0
                return proc_idx_set
            end
            # exclude master
            mysize -= 1;
            modrank = (myrank-1) % nblocks
        end
        # assign sub-blocks in round-robin fashion
        for s = modrank:mysize:(nblocks-1)
            push!(proc_idx_set, s+1)
        end
    end

    # return assigned block ids
    return proc_idx_set
end
myblocks(nblocks::Int) = getBlockIds(nblocks)

"""
    getNumBlockCols

Get the number of columns for each block
"""
function getNumBlockCols()::Dict{Int,Int}
    # get number of block columns
    numBlockCols = Dict{Int,Int}()
    if dspenv.comm_size > 1
        num_proc_blocks = convert(Vector{Cint}, MPI.Allgather(length(dspenv.block_ids), dspenv.comm))
        #@show num_proc_blocks
        #@show collect(keys(blocks))
        block_ids = MPI.Allgatherv(convert(Vector{Cint}, dspenv.block_ids), num_proc_blocks, dspenv.comm)
        #@show block_ids
        ncols_to_send = Cint[dspenv.numCols[id] for id in dspenv.block_ids]
        #@show ncols_to_send
        ncols = MPI.Allgatherv(ncols_to_send, num_proc_blocks, dspenv.comm)
        #@show ncols
        @assert(length(block_ids) == dspenv.nblocks)
        for i in 1:dspenv.nblocks
            numBlockCols[block_ids[i]] = ncols[i]
        end
    else
        for id in dspenv.block_ids
            numBlockCols[id] = dspenv.numCols[id]
        end
    end
    return numBlockCols
end

"""
    setBlocks

The function sets the number of blocks (e.g., scenarios for stochastic program) and their Ids.
"""
function setBlocks()
    dspenv.block_ids = getBlockIds()
    @dsp_ccall("setIntPtrParam", Cvoid, (Ptr{Cvoid}, Ptr{UInt8}, Cint, Ptr{Cint}),
        dspenv.p, "ARR_PROC_IDX", convert(Cint, length(dspenv.block_ids)), convert(Vector{Cint}, dspenv.block_ids .- 1))
end

"""
    parallelize

Initialize DSP environment related to MPI, if MPI is initialized.
"""
function parallelize(comm)
    if @isdefined(MPI)
        if MPI.Initialized()
            dspenv.comm = comm
            dspenv.comm_size = MPI.Comm_size(dspenv.comm)
            dspenv.comm_rank = MPI.Comm_rank(dspenv.comm)
        else
            @error("MPI is not initialized.")
        end
    else
        @error("MPI.jl is not defined.")
    end
end

"""
    myrank

Returns the MPI rank (i.e., process id), starting from 0
"""
myrank() = dspenv.comm_rank

"""
    mysize

Returns the number of MPI processes
"""
mysize() = dspenv.comm_size


text:
using DataTrees
using Base.Test

@test TreeNode{String, String} <: TreeNode

tree = TreeNode{Int, String}("Tree", 0)
@test isa(tree, TreeNode)
@test isroot(tree)
@test isleaf(tree)
@test get(tree) == tree

for i in 1:3 addchild!(tree, "Child $i", i) end

print(tree)
@test !isnull(tree["Child 1"])
@test !isroot(tree["Child 2"])
@test isleaf(tree["Child 3"])
@test getdepth(tree) == 0

addchild!(tree["Child 1"], TreeNode{Int, String}("Grand child 1", 10))
t = tree["Child 1"]["Grand child 1"]
@test t.data == 10
@test isleaf(t)
@test !isleaf(t.parent)
@test !isroot(t.parent)
@test getdepth(t.parent) == 1
@test getdepth(t) == 2

@test print(tree.parent) == nothing

addsubtree!(tree, ["Child 2", "Foo", "bar"])
@test isleaf(tree[["Child 2", "Foo", "bar"]])

# tree from an array of edges
a = [1 0;
     2 1;
     3 1;
     4 0;
     5 4;
     6 4]
t_from_edges = buildtree(Int64, a)
@test isa(t_from_edges[4], TreeNode{Int,Int})

# tree from tuples with data
treedata = [(1, 0, "data1"), (2, 1, "data2")]
t_with_data = buildtree(treedata)
@test t_with_data[1].data == "data1"