text:
#######################################################################
# Hamiltonian
#######################################################################
struct HamiltonianHarmonic{L,M,A<:Union{AbstractMatrix{M},SparseMatrixBuilder{M}}}
    dn::SVector{L,Int}
    h::A
end

HamiltonianHarmonic{L,M,A}(dn::SVector{L,Int}, n::Int, m::Int) where {L,M,A<:SparseMatrixCSC{M}} =
    HamiltonianHarmonic(dn, sparse(Int[], Int[], M[], n, m))

HamiltonianHarmonic{L,M,A}(dn::SVector{L,Int}, n::Int, m::Int) where {L,M,A<:Matrix{M}} =
    HamiltonianHarmonic(dn, zeros(M, n, m))

struct Hamiltonian{LA<:AbstractLattice,L,M,A<:AbstractMatrix,
                   H<:HamiltonianHarmonic{L,M,A},
                   O<:Tuple{Vararg{Tuple{Vararg{NameType}}}}} # <: AbstractMatrix{M}
    lattice::LA
    harmonics::Vector{H}
    orbitals::O
end

function Hamiltonian(lat, hs::Vector{H}, orbs, n::Int, m::Int) where {L,M,H<:HamiltonianHarmonic{L,M}}
    sort!(hs, by = h -> abs.(h.dn))
    if isempty(hs) || !iszero(first(hs).dn)
        pushfirst!(hs, H(zero(SVector{L,Int}), empty_sparse(M, n, m)))
    end
    return Hamiltonian(lat, hs, orbs)
end

Base.show(io::IO, ham::Hamiltonian) = show(io, MIME("text/plain"), ham)
function Base.show(io::IO, ::MIME"text/plain", ham::Hamiltonian)
    i = get(io, :indent, "")
    print(io, i, summary(ham), "\n",
"$i  Bloch harmonics  : $(length(ham.harmonics)) ($(displaymatrixtype(ham)))
$i  Harmonic size    : $((n -> "$n × $n")(nsites(ham)))
$i  Orbitals         : $(displayorbitals(ham))
$i  Element type     : $(displayelements(ham))
$i  Onsites          : $(nonsites(ham))
$i  Hoppings         : $(nhoppings(ham))
$i  Coordination     : $(nhoppings(ham) / nsites(ham))")
    ioindent = IOContext(io, :indent => string("  "))
    issuperlattice(ham.lattice) && print(ioindent, "\n", ham.lattice.supercell)
end

Base.summary(h::Hamiltonian{LA}) where {E,L,LA<:Lattice{E,L}} =
    "Hamiltonian{<:Lattice} : Hamiltonian on a $(L)D Lattice in $(E)D space"

Base.summary(::Hamiltonian{LA}) where {E,L,T,L´,LA<:Superlattice{E,L,T,L´}} =
    "Hamiltonian{<:Superlattice} : $(L)D Hamiltonian on a $(L´)D Superlattice in $(E)D space"

Base.eltype(::Hamiltonian{<:Any,<:Any,M}) where {M} = M

Base.isequal(h1::HamiltonianHarmonic, h2::HamiltonianHarmonic) =
    h1.dn == h2.dn && h1.h == h2.h

displaymatrixtype(h::Hamiltonian) = displaymatrixtype(matrixtype(h))
displaymatrixtype(::Type{<:SparseMatrixCSC}) = "SparseMatrixCSC, sparse"
displaymatrixtype(::Type{<:Array}) = "Matrix, dense"
displaymatrixtype(A::Type{<:AbstractArray}) = string(A)
displayelements(h::Hamiltonian) = displayelements(blocktype(h))
displayelements(::Type{S}) where {N,T,S<:SMatrix{N,N,T}} = "$N × $N blocks ($T)"
displayelements(::Type{T}) where {T} = "scalar ($T)"
displayorbitals(h::Hamiltonian) =
    replace(replace(string(h.orbitals), "Symbol(\"" => ":"), "\")" => "")

SparseArrays.issparse(h::Hamiltonian{LA,L,M,A}) where {LA,L,M,A<:AbstractSparseMatrix} = true
SparseArrays.issparse(h::Hamiltonian{LA,L,M,A}) where {LA,L,M,A} = false

Base.parent(h::Hamiltonian) = h

# Internal API #

latdim(h::Hamiltonian{LA}) where {E,L,LA<:AbstractLattice{E,L}} = L

matrixtype(::Hamiltonian{LA,L,M,A}) where {LA,L,M,A} = A
blockeltype(::Hamiltonian{<:Any,<:Any,M}) where {M} = eltype(M)

# find SVector type that can hold all orbital amplitudes in any lattice sites
orbitaltype(orbs, type::Type{Tv} = Complex{T}) where {T,Tv} =
    _orbitaltype(SVector{1,Tv}, orbs...)
_orbitaltype(::Type{S}, ::NTuple{D,NameType}, os...) where {N,Tv,D,S<:SVector{N,Tv}} =
    (M = max(N,D); _orbitaltype(SVector{M,Tv}, os...))
_orbitaltype(t::Type{SVector{N,Tv}}) where {N,Tv} = t
_orbitaltype(t::Type{SVector{1,Tv}}) where {Tv} = Tv

# find SMatrix type that can hold all matrix elements between lattice sites
blocktype(orbs, type::Type{Tv}) where {Tv} =
    _blocktype(orbitaltype(orbs, Tv))
_blocktype(::Type{S}) where {N,Tv,S<:SVector{N,Tv}} = SMatrix{N,N,Tv,N*N}
_blocktype(::Type{S}) where {S<:Number} = S

blocktype(h::Hamiltonian{LA,L,M}) where {LA,L,M} = M

promote_blocktype(hs::Hamiltonian...) = promote_blocktype(blocktype.(hs)...)
promote_blocktype(s1::Type, s2::Type, ss::Type...) =
    promote_blocktype(promote_blocktype(s1, s2), ss...)
promote_blocktype(::Type{SMatrix{N1,N1,T1,NN1}}, ::Type{SMatrix{N2,N2,T2,NN2}}) where {N1,NN1,T1,N2,NN2,T2} =
    SMatrix{max(N1, N2), max(N1, N2), promote_type(T1, T2), max(NN1,NN2)}
promote_blocktype(T1::Type{<:Number}, T2::Type{<:Number}) = promote_type(T1, T2)
promote_blocktype(T::Type) = T

blockdim(h::Hamiltonian) = blockdim(blocktype(h))
blockdim(::Type{S}) where {N,S<:SMatrix{N,N}} = N
blockdim(::Type{T}) where {T<:Number} = 1

function nhoppings(ham::Hamiltonian)
    count = 0
    for h in ham.harmonics
        count += iszero(h.dn) ? (_nnz(h.h) - _nnzdiag(h.h)) : _nnz(h.h)
    end
    return count
end

function nonsites(ham::Hamiltonian)
    count = 0
    for h in ham.harmonics
        iszero(h.dn) && (count += _nnzdiag(h.h))
    end
    return count
end

_nnz(h::AbstractSparseMatrix) = count(!iszero, nonzeros(h)) # Does not include stored zeros
_nnz(h::DenseMatrix) = count(!iszero, h)

function _nnzdiag(s::SparseMatrixCSC)
    count = 0
    rowptrs = rowvals(s)
    nz = nonzeros(s)
    for col in 1:size(s,2)
        for ptr in nzrange(s, col)
            rowptrs[ptr] == col && (count += !iszero(nz[ptr]); break)
        end
    end
    return count
end
_nnzdiag(s::Matrix) = count(!iszero, s[i,i] for i in 1:minimum(size(s)))

Base.isequal(h1::Hamiltonian, h2::Hamiltonian) =
    isequal(h1.lattice, h2.lattice) && isequal(h1.harmonics, h2.harmonics) &&
    isequal(h1.orbitals, h2.orbitals)

# Iterators #

function nonzero_indices(h::Hamiltonian, rowrange = 1:size(h, 1), colrange = 1:size(h, 2))
    rowrange´ = rclamp(rowrange, 1:size(h, 1))
    colrange´ = rclamp(colrange, 1:size(h, 2))
    gen = ((har.dn, rowvals(har.h)[ptr], col)
                for har in h.harmonics
                for col in colrange´
                for ptr in nzrange_inrows(har.h, col, rowrange´)
                if !iszero(nonzeros(har.h)[ptr]))
    return gen
end

function nonzero_indices(har::HamiltonianHarmonic, rowrange = 1:size(h, 1), colrange = 1:size(h, 2))
    rowrange´ = rclamp(rowrange, 1:size(har, 1))
    colrange´ = rclamp(colrange, 1:size(har, 2))
    gen = ((rowvals(har.h)[ptr], col)
                for col in colrange´
                for ptr in nzrange_inrows(har.h, col, rowrange´)
                if !iszero(nonzeros(har.h)[ptr]))
    return gen
end

function nzrange_inrows(h, col, rowrange)
    ptrs = nzrange(h, col)
    rows = rowvals(h)
    ptrmin = first(ptrs)
    ptrmax = last(ptrs)

    for p in ptrs
        rows[p] in rowrange && break
        ptrmin = p + 1
    end

    if ptrmin < ptrmax
        for p in ptrmax:-1:ptrmin
            ptrmax = p
            rows[p] in rowrange && break
        end
    end

    return ptrmin:ptrmax
end

# External API #
"""
    hamiltonian(lat, model; orbitals, type)

Create a `Hamiltonian` by applying `model::TighbindingModel` to the lattice `lat` (see
`hopping` and `onsite` for details on building tightbinding models).

The number of orbitals on each sublattice can be specified by the keyword `orbitals`
(otherwise all sublattices have one orbital by default). The following, and obvious
combinations, are possible formats for the `orbitals` keyword:

    orbitals = :a                # all sublattices have 1 orbital named :a
    orbitals = (:a,)             # same as above
    orbitals = (:a, :b, 3)       # all sublattices have 3 orbitals named :a and :b and :3
    orbitals = ((:a, :b), (:c,)) # first sublattice has 2 orbitals, second has one
    orbitals = ((:a, :b), :c)    # same as above
    orbitals = (Val(2), Val(1))  # same as above, with automatic names
    orbitals = (:A => (:a, :b), :D => :c) # sublattice :A has two orbitals, :D and rest have one
    orbitals = :D => Val(4)      # sublattice :D has four orbitals, rest have one

The matrix sizes of tightbinding `model` must match the orbitals specified. Internally, we
define a block size `N = max(num_orbitals)`. If `N = 1` (all sublattices with one orbital)
the the Hamiltonian element type is `type`. Otherwise it is `SMatrix{N,N,type}` blocks,
padded with the necessary zeros as required. Keyword `type` is `Complex{T}` by default,
where `T` is the number type of `lat`.

    lat |> hamiltonian(model; kw...)

Functional `hamiltonian` form equivalent to `hamiltonian(lat, model[, funcmodel]; kw...)`.

# Indexing

Indexing into a Hamiltonian `h` works as follows. Access the `HamiltonianHarmonic` matrix at
a given `dn::NTuple{L,Int}` with `h[dn]`. Assign `v` into element `(i,j)` of said matrix
with `h[dn][i,j] = v` or `h[dn, i, j] = v`. Broadcasting with vectors of indices `is` and
`js` is supported, `h[dn][is, js] = v_matrix`.

To add an empty harmonic with a given `dn::NTuple{L,Int}`, do `push!(h, dn)`. To delete it,
do `deleteat!(h, dn)`.

# Examples

```jldoctest
julia> h = hamiltonian(LatticePresets.honeycomb(), hopping(@SMatrix[1 2; 3 4], range = 1/√3), orbitals = Val(2))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 5 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a, :a), (:a, :a))
  Element type     : 2 × 2 blocks (Complex{Float64})
  Onsites          : 0
  Hoppings         : 6
  Coordination     : 3.0

julia> push!(h, (3,3)) # Adding a new Hamiltonian harmonic (if not already present)
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 6 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a, :a), (:a, :a))
  Element type     : 2 × 2 blocks (Complex{Float64})
  Onsites          : 0
  Hoppings         : 6
  Coordination     : 3.0

julia> h[(3,3)][1,1] = @SMatrix[1 2; 2 1]; h[(3,3)] # element assignment
2×2 SparseArrays.SparseMatrixCSC{StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4},Int64} with 1 stored entry:
  [1, 1]  =  [1.0+0.0im 2.0+0.0im; 2.0+0.0im 1.0+0.0im]

julia> h[(3,3)][[1,2],[1,2]] .= Ref(@SMatrix[1 2; 2 1])
2×2 view(::SparseArrays.SparseMatrixCSC{StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4},Int64}, [1, 2], [1, 2]) with eltype StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4}:
 [1.0+0.0im 2.0+0.0im; 2.0+0.0im 1.0+0.0im]  [1.0+0.0im 2.0+0.0im; 2.0+0.0im 1.0+0.0im]
 [1.0+0.0im 2.0+0.0im; 2.0+0.0im 1.0+0.0im]  [1.0+0.0im 2.0+0.0im; 2.0+0.0im 1.0+0.0im]```
```

# See also:
    `onsite`, `hopping`, `bloch`, `bloch!`
"""
hamiltonian(lat, ts...; orbitals = missing, kw...) =
    _hamiltonian(lat, sanitize_orbs(orbitals, lat.unitcell.names), ts...; kw...)
_hamiltonian(lat::AbstractLattice, orbs; kw...) = _hamiltonian(lat, orbs, TightbindingModel(); kw...)
_hamiltonian(lat::AbstractLattice, orbs, m::TightbindingModel; type::Type = Complex{numbertype(lat)}, kw...) =
    hamiltonian_sparse(blocktype(orbs, type), lat, orbs, m; kw...)

hamiltonian(t::TightbindingModel...; kw...) = lat -> hamiltonian(lat, t...; kw...)

sanitize_orbs(o::Union{Val,NameType,Integer}, names::NTuple{N}) where {N} =
    ntuple(n -> sanitize_orbs(o), Val(N))
sanitize_orbs(o::NTuple{M,Union{NameType,Integer}}, names::NTuple{N}) where {M,N} =
    (ont = nametype.(o); ntuple(n -> ont, Val(N)))
sanitize_orbs(o::Missing, names) = sanitize_orbs((:a,), names)
sanitize_orbs(o::Pair, names) = sanitize_orbs((o,), names)
sanitize_orbs(os::NTuple{M,Pair}, names::NTuple{N}) where {N,M} =
    ntuple(Val(N)) do n
        for m in 1:M
            first(os[m]) == names[n] && return sanitize_orbs(os[m])
        end
        return (:a,)
    end
sanitize_orbs(os::NTuple{M,Any}, names::NTuple{N}) where {N,M} =
    ntuple(n -> n > M ? (:a,) : sanitize_orbs(os[n]), Val(N))

sanitize_orbs(p::Pair) = sanitize_orbs(last(p))
sanitize_orbs(o::Integer) = (nametype(o),)
sanitize_orbs(o::NameType) = (o,)
sanitize_orbs(o::Val{N}) where {N} = ntuple(_ -> :a, Val(N))
sanitize_orbs(o::NTuple{N,Union{Integer,NameType}}) where {N} = nametype.(o)
sanitize_orbs(p) = throw(ArgumentError("Wrong format for orbitals, see `hamiltonian`"))

Base.Matrix(h::Hamiltonian) = Hamiltonian(h.lattice, Matrix.(h.harmonics), h.orbitals)
Base.Matrix(h::HamiltonianHarmonic) = HamiltonianHarmonic(h.dn, Matrix(h.h))

Base.copy(h::Hamiltonian) = Hamiltonian(copy(h.lattice), copy.(h.harmonics), h.orbitals)
Base.copy(h::HamiltonianHarmonic) = HamiltonianHarmonic(h.dn, copy(h.h))

Base.size(h::Hamiltonian, n) = size(first(h.harmonics).h, n)
Base.size(h::Hamiltonian) = size(first(h.harmonics).h)
Base.size(h::HamiltonianHarmonic, n) = size(h.h, n)
Base.size(h::HamiltonianHarmonic) = size(h.h)

flatsize(h::Hamiltonian, n) = first(flatsize(h)) # h is always square

function flatsize(h::Hamiltonian)
    n = sum(sublatsites(h.lattice) .* length.(h.orbitals))
    return (n, n)
end

function LinearAlgebra.ishermitian(h::Hamiltonian)
    for hh in h.harmonics
        isassigned(h, -hh.dn) || return false
        hh.h == h[-hh.dn]' || return false
    end
    return true
end

bravais(h::Hamiltonian) = bravais(h.lattice)

issemibounded(h::Hamiltonian) = issemibounded(h.lattice)

nsites(h::Hamiltonian) = isempty(h.harmonics) ? 0 : nsites(first(h.harmonics))
nsites(h::HamiltonianHarmonic) = size(h.h, 1)

nsublats(h::Hamiltonian) = nsublats(h.lattice)

norbitals(h::Hamiltonian) = length.(h.orbitals)

function meandist(h::Hamiltonian)
    distsum = 0.0
    num = 0
    ss = sites(h.lattice)
    br = h.lattice.bravais.matrix
    for (dn, row, col) in nonzero_indices(h)
        if row != col
            num += 1
            rsrc = ss[col]
            rdst = ss[row] + br * dn
            distsum += norm(rsrc - rdst)
        end
    end
    return iszero(num) ? 0.0 : distsum / num
end

# External API #

"""
    transform!(f::Function, h::Hamiltonian)

Transform the site positions of the Hamiltonian's lattice in place without modifying the
Hamiltonian harmonics.
"""
function transform!(f, h::Hamiltonian)
    transform!(f, h.lattice)
    return h
end

# Indexing #

Base.push!(h::Hamiltonian{<:Any,L}, dn::NTuple{L,Int}) where {L} = push!(h, SVector(dn...))
Base.push!(h::Hamiltonian{<:Any,L}, dn::Vararg{Int,L}) where {L} = push!(h, SVector(dn...))
function Base.push!(h::Hamiltonian{<:Any,L,M,A}, dn::SVector{L,Int}) where {L,M,A}
    for hh in h.harmonics
        hh.dn == dn && return hh
    end
    hh = HamiltonianHarmonic{L,M,A}(dn, size(h)...)
    push!(h.harmonics, hh)
    return h
end

Base.getindex(h::Hamiltonian, dn::NTuple) = getindex(h, SVector(dn))
@inline function Base.getindex(h::Hamiltonian{<:Any,L}, dn::SVector{L,Int}) where {L}
    nh = findfirst(hh -> hh.dn == dn, h.harmonics)
    nh === nothing && throw(BoundsError(h, dn))
    return h.harmonics[nh].h
end
Base.getindex(h::Hamiltonian, dn::NTuple, i::Vararg{Int}) = h[dn][i...]
Base.getindex(h::Hamiltonian{LA, L}, i::Vararg{Int}) where {LA,L} = h[zero(SVector{L,Int})][i...]

Base.deleteat!(h::Hamiltonian{<:Any,L}, dn::Vararg{Int,L}) where {L} =
    deleteat!(h, toSVector(dn))
Base.deleteat!(h::Hamiltonian{<:Any,L}, dn::NTuple{L,Int}) where {L} =
    deleteat!(h, toSVector(dn))
function Base.deleteat!(h::Hamiltonian{<:Any,L}, dn::SVector{L,Int}) where {L}
    nh = findfirst(hh -> hh.dn == SVector(dn...), h.harmonics)
    nh === nothing || deleteat!(h.harmonics, nh)
    return h
end

Base.isassigned(h::Hamiltonian, dn::Vararg{Int}) = isassigned(h, SVector(dn))
Base.isassigned(h::Hamiltonian, dn::NTuple) = isassigned(h, SVector(dn))
Base.isassigned(h::Hamiltonian{<:Any,L}, dn::SVector{L,Int}) where {L} =
    findfirst(hh -> hh.dn == dn, h.harmonics) != nothing

## Boolean masking
"""
    &(h1::Hamiltonian{<:Superlattice}, h2::Hamiltonian{<:Superlattice})

Construct a new `Hamiltonian{<:Superlattice}` using an `and` boolean mask, i.e. with a
supercell that contains cells that are both in the supercell of `h1` and `h2`

    &(s1::Superlattice, s2::Superlattice}

Equivalent of the above for `Superlattice`s

# See also:
    `|`, `xor`
"""
(Base.:&)(s1::Hamiltonian{<:Superlattice}, s2::Hamiltonian{<:Superlattice}) =
    boolean_mask_hamiltonian(Base.:&, s1, s2)

"""
    |(h1::Hamiltonian{<:Superlattice}, h2::Hamiltonian{<:Superlattice})

Construct a new `Hamiltonian{<:Superlattice}` using an `or` boolean mask, i.e. with a
supercell that contains cells that are either in the supercell of `h1` or `h2`

    |(s1::Superlattice, s2::Superlattice}

Equivalent of the above for `Superlattice`s

# See also:
    `&`, `xor`
"""
(Base.:|)(s1::Hamiltonian{<:Superlattice}, s2::Hamiltonian{<:Superlattice}) =
    boolean_mask_hamiltonian(Base.:|, s1, s2)

"""
    xor(h1::Hamiltonian{<:Superlattice}, h2::Hamiltonian{<:Superlattice})

Construct a new `Hamiltonian{<:Superlattice}` using a `xor` boolean mask, i.e. with a
supercell that contains cells that are either in the supercell of `h1` or `h2` but not in
both

    xor(s1::Superlattice, s2::Superlattice}

Equivalent of the above for `Superlattice`s

# See also:
    `&`, `|`
"""
(Base.xor)(s1::Hamiltonian{<:Superlattice}, s2::Hamiltonian{<:Superlattice}) =
    boolean_mask_hamiltonian(Base.xor, s1, s2)

function boolean_mask_hamiltonian(f, s1::Hamiltonian{<:Superlattice}, s2::Hamiltonian{<:Superlattice})
    check_compatible_hsuper(s1, s2)
    return Hamiltonian(f(s1.lattice, s2.lattice), s1.harmonics, s1.orbitals)
end

function check_compatible_hsuper(s1, s2)
    compatible = isequal(s1.harmonics, s2.harmonics) && isequal(s1.orbitals, s2.orbitals)
    compatible || throw(ArgumentError("Hamiltonians are incompatible for boolean masking"))
    return nothing
end

#######################################################################
# auxiliary types
#######################################################################
struct IJV{L,M}
    dn::SVector{L,Int}
    i::Vector{Int}
    j::Vector{Int}
    v::Vector{M}
end

struct IJVBuilder{L,M,E,T,O,LA<:AbstractLattice{E,L,T}}
    lat::LA
    orbs::O
    ijvs::Vector{IJV{L,M}}
    kdtrees::Vector{KDTree{SVector{E,T},Euclidean,T}}
end

IJV{L,M}(dn::SVector{L} = zero(SVector{L,Int})) where {L,M} =
    IJV(dn, Int[], Int[], M[])

function IJVBuilder(lat::AbstractLattice{E,L,T}, orbs, ijvs::Vector{IJV{L,M}}) where {E,L,T,M}
    kdtrees = Vector{KDTree{SVector{E,T},Euclidean,T}}(undef, nsublats(lat))
    return IJVBuilder(lat, orbs, ijvs, kdtrees)
end

IJVBuilder(lat::AbstractLattice{E,L}, orbs, ::Type{M}) where {E,L,M} =
    IJVBuilder(lat, orbs, IJV{L,M}[])

function IJVBuilder(lat::AbstractLattice{E,L}, orbs, hs::Hamiltonian...) where {E,L}
    M = promote_blocktype(hs...)
    ijvs = IJV{L,M}[]
    builder = IJVBuilder(lat, orbs, ijvs)
    offset = 0
    for h in hs
        for har in h.harmonics
            ijv = builder[har.dn]
            push_block!(ijv, har, offset)
        end
        offset += size(h, 1)
    end
    return builder
end

Base.eltype(b::IJVBuilder{L,M}) where {L,M} = M

function Base.getindex(b::IJVBuilder{L,M}, dn::SVector{L2,Int}) where {L,L2,M}
    L == L2 || throw(error("Tried to apply an $L2-dimensional model to an $L-dimensional lattice"))
    for e in b.ijvs
        e.dn == dn && return e
    end
    e = IJV{L,M}(dn)
    push!(b.ijvs, e)
    return e
end

Base.length(h::IJV) = length(h.i)
Base.isempty(h::IJV) = length(h) == 0
Base.copy(h::IJV) = IJV(h.dn, copy(h.i), copy(h.j), copy(h.v))

function Base.resize!(h::IJV, n)
    resize!(h.i, n)
    resize!(h.j, n)
    resize!(h.v, n)
    return h
end

Base.push!(ijv::IJV, (i, j, v)::Tuple) = (push!(ijv.i, i); push!(ijv.j, j); push!(ijv.v, v))

function push_block!(ijv::IJV{L,M}, h::HamiltonianHarmonic, offset) where {L,M}
    I, J, V = findnz(h.h)
    for (i, j, v) in zip(I, J, V)
        push!(ijv, (i + offset, j + offset, padtotype(v, M)))
    end
    return ijv
end

#######################################################################
# hamiltonian_sparse
#######################################################################
function hamiltonian_sparse(Mtype, lat, orbs, model)
    builder = IJVBuilder(lat, orbs, Mtype)
    return hamiltonian_sparse!(builder, lat, orbs, model)
end

function hamiltonian_sparse!(builder::IJVBuilder{L,M}, lat::AbstractLattice{E,L}, orbs, model) where {E,L,M}
    applyterms!(builder, terms(model)...)
    n = nsites(lat)
    HT = HamiltonianHarmonic{L,M,SparseMatrixCSC{M,Int}}
    harmonics = HT[HT(e.dn, sparse(e.i, e.j, e.v, n, n)) for e in builder.ijvs if !isempty(e)]
    return Hamiltonian(lat, harmonics, orbs, n, n)
end

applyterms!(builder, terms...) = foreach(term -> applyterm!(builder, term), terms)

applyterm!(builder::IJVBuilder, term::Union{OnsiteTerm, HoppingTerm}) =
    applyterm!(builder, term, sublats(term, builder.lat))

function applyterm!(builder::IJVBuilder{L}, term::OnsiteTerm, termsublats) where {L}
    selector = term.selector
    lat = builder.lat
    for s in termsublats
        is = siterange(lat, s)
        dn0 = zero(SVector{L,Int})
        ijv = builder[dn0]
        offset = lat.unitcell.offsets[s]
        for i in is
            isinregion(i, dn0, selector.region, lat) || continue
            r = lat.unitcell.sites[i]
            v = toeltype(term(r, r), eltype(builder), builder.orbs[s], builder.orbs[s])
            push!(ijv, (i, i, v))
        end
    end
    return nothing
end

function applyterm!(builder::IJVBuilder{L}, term::HoppingTerm, termsublats) where {L}
    selector = term.selector
    L > 0 && checkinfinite(selector)
    lat = builder.lat
    for (s2, s1) in termsublats  # Each is a Pair s2 => s1
        is, js = siterange(lat, s1), siterange(lat, s2)
        dns = dniter(selector.dns, Val(L))
        for dn in dns
            foundlink = false
            ijv = builder[dn]
            for j in js
                sitej = lat.unitcell.sites[j]
                rsource = sitej - lat.bravais.matrix * dn
                itargets = targets(builder, selector.range, rsource, s1)
                for i in itargets
                    isselfhopping((i, j), (s1, s2), dn) && continue
                    isinregion((i, j), (dn, zero(dn)), selector.region, lat) || continue
                    foundlink = true
                    rtarget = lat.unitcell.sites[i]
                    r, dr = _rdr(rsource, rtarget)
                    v = toeltype(term(r, dr), eltype(builder), builder.orbs[s1], builder.orbs[s2])
                    push!(ijv, (i, j, v))
                end
            end
            foundlink && acceptcell!(dns, dn)
        end
    end
    return nothing
end

toeltype(t::Number, ::Type{T}, t1::NTuple{1}, t2::NTuple{1}) where {T<:Number} = T(t)
toeltype(t::SMatrix{N1,N2}, ::Type{S}, t1::NTuple{N1}, t2::NTuple{N2}) where {N1,N2,S<:SMatrix} =
    padtotype(t, S)
toeltype(u::UniformScaling, ::Type{T}, t1::NTuple{1}, t2::NTuple{1}) where {T<:Number} = T(u.λ)
toeltype(u::UniformScaling, ::Type{S}, t1::NTuple{N1}, t2::NTuple{N2}) where {N1,N2,S<:SMatrix} =
    padtotype(SMatrix{N1,N2}(u), S)
toeltype(t::Array, s, t1, t2) = throw(ArgumentError("Array input in model, please use StaticArrays instead (e.g. SA[1 0; 0 1] instead of [1 0; 0 1])"))
toeltype(t, s, t1, t2) = throw(DimensionMismatch("Dimension mismatch between model and Hamiltonian. Did you correctly specify the `orbitals` in hamiltonian? Consider also using `I` to cover non-uniform orbital dimensions"))

dniter(dns::Missing, ::Val{L}) where {L} = BoxIterator(zero(SVector{L,Int}))
dniter(dns, ::Val) = dns

function targets(builder, range::Real, rsource, s1)
    !isfinite(range) && return targets(builder, missing, rsource, s1)
    if !isassigned(builder.kdtrees, s1)
        sites = view(builder.lat.unitcell.sites, siterange(builder.lat, s1))
        (builder.kdtrees[s1] = KDTree(sites))
    end
    targetlist = inrange(builder.kdtrees[s1], rsource, range)
    targetlist .+= builder.lat.unitcell.offsets[s1]
    return targetlist
end

targets(builder, range::Missing, rsource, s1) = siterange(builder.lat, s1)

checkinfinite(selector) =
    selector.dns === missing && (selector.range === missing || !isfinite(selector.range)) &&
    throw(ErrorException("Tried to implement an infinite-range hopping on an unbounded lattice"))

isselfhopping((i, j), (s1, s2), dn) = i == j && s1 == s2 && iszero(dn)

#######################################################################
# unitcell/supercell for Hamiltonians
#######################################################################
function supercell(ham::Hamiltonian, args...; kw...)
    slat = supercell(ham.lattice, args...; kw...)
    return Hamiltonian(slat, ham.harmonics, ham.orbitals)
end

function unitcell(ham::Hamiltonian{<:Lattice}, args...; modifiers = (), kw...)
    sham = supercell(ham, args...; kw...)
    return unitcell(sham; modifiers = modifiers)
end

function unitcell(ham::Hamiltonian{LA,L}; modifiers = ()) where {E,L,T,L´,LA<:Superlattice{E,L,T,L´}}
    lat = ham.lattice
    sc = lat.supercell
    modifiers´ = resolve.(ensuretuple(modifiers), Ref(lat))
    mapping = OffsetArray{Int}(undef, sc.sites, sc.cells.indices...) # store supersite indices newi
    mapping .= 0
    foreach_supersite((s, oldi, olddn, newi) -> mapping[oldi, Tuple(olddn)...] = newi, lat)
    dim = nsites(sc)
    B = blocktype(ham)
    S = typeof(SparseMatrixBuilder{B}(dim, dim))
    harmonic_builders = HamiltonianHarmonic{L´,B,S}[]
    pinvint = pinvmultiple(sc.matrix)
    foreach_supersite(lat) do s, source_i, source_dn, newcol
        for oldh in ham.harmonics
            rows = rowvals(oldh.h)
            vals = nonzeros(oldh.h)
            target_dn = source_dn + oldh.dn
            super_dn = new_dn(target_dn, pinvint)
            wrapped_dn = wrap_dn(target_dn, super_dn, sc.matrix)
            newh = get_or_push!(harmonic_builders, super_dn, dim, newcol)
            for p in nzrange(oldh.h, source_i)
                target_i = rows[p]
                # check: wrapped_dn could exit bounding box along non-periodic direction
                checkbounds(Bool, mapping, target_i, Tuple(wrapped_dn)...) || continue
                newrow = mapping[target_i, Tuple(wrapped_dn)...]
                val = applymodifiers(vals[p], lat, (source_i, target_i), (source_dn, target_dn), modifiers´...)
                iszero(newrow) || pushtocolumn!(newh.h, newrow, val)
            end
        end
        foreach(h -> finalizecolumn!(h.h), harmonic_builders)
    end
    harmonics = [HamiltonianHarmonic(h.dn, sparse(h.h)) for h in harmonic_builders]
    unitlat = unitcell(lat)
    orbs = ham.orbitals
    return Hamiltonian(unitlat, harmonics, orbs)
end

function get_or_push!(hs::Vector{<:HamiltonianHarmonic{L,B,<:SparseMatrixBuilder}}, dn, dim, currentcol) where {L,B}
    for h in hs
        h.dn == dn && return h
    end
    newh = HamiltonianHarmonic(dn, SparseMatrixBuilder{B}(dim, dim))
    currentcol > 1 && finalizecolumn!(newh.h, currentcol - 1) # for columns that have been already processed
    push!(hs, newh)
    return newh
end

wrap_dn(olddn::SVector, newdn::SVector, supercell::SMatrix) = olddn - supercell * newdn

applymodifiers(val, lat, inds, dns) = val

function applymodifiers(val, lat, inds, dns, m::UniformModifier, ms...)
    selected = m.selector(lat, inds, dns)
    val´ = selected ? m.f(val) : val
    return applymodifiers(val´, lat, inds, dns, ms...)
end

function applymodifiers(val, lat, (row, col), (dnrow, dncol), m::OnsiteModifier, ms...)
    selected = m.selector(lat, (row, col), (dnrow, dncol))
    if selected
        r = sites(lat)[col] + bravais(lat) * dncol
        val´ = selected ? m(val, r) : val
    else
        val´ = val
    end
    return applymodifiers(val´, lat, (row, col), (dnrow, dncol), ms...)
end

function applymodifiers(val, lat, (row, col), (dnrow, dncol), m::HoppingModifier, ms...)
    selected = m.selector(lat, (row, col), (dnrow, dncol))
    if selected
        br = bravais(lat)
        r, dr = _rdr(sites(lat)[col] + br * dncol, sites(lat)[row] + br * dnrow)
        val´ = selected ? m(val, r, dr) : val
    else
        val´ = val
    end
    return applymodifiers(val´, lat, (row, col), (dnrow, dncol), ms...)
end

#######################################################################
# wrap
#######################################################################
"""
    wrap(h::Hamiltonian, axis::Int; factor = 1)

Build a new Hamiltonian wherein the Bravais `axis` is wrapped into a loop. If a `factor` is
given, the wrapped hoppings will be multiplied by said factor. This is useful to represent a
flux Φ through the loop, if `factor = exp(im * 2π * Φ/Φ₀)`.

    h |> wrap(axis; kw...)

Functional form equivalent to `wrap(h, axis; kw...)`.

# Examples

```jldoctest
julia> LatticePresets.honeycomb() |> hamiltonian(hopping(1, range = 1/√3)) |>
       unitcell((1,-1), (10, 10)) |> wrap(2)
Hamiltonian{<:Lattice} : Hamiltonian on a 1D Lattice in 2D space
  Bloch harmonics  : 3 (SparseMatrixCSC, sparse)
  Harmonic size    : 40 × 40
  Orbitals         : ((:a,), (:a,))
  Element type     : scalar (Complex{Float64})
  Onsites          : 0
  Hoppings         : 120
  Coordination     : 3.0
```
"""
function wrap(h::Hamiltonian{<:Lattice,L}, axis; factor = 1) where {L}
    1 <= axis <= L || throw(ArgumentError("wrap axis should be between 1 and the lattice dimension $L"))
    lattice´ = _wrap(h.lattice, axis)
    harmonics´ = _wrap(h.harmonics, axis, factor)
    return Hamiltonian(lattice´, harmonics´, h.orbitals)
end

wrap(axis; kw...) = h -> wrap(h, axis; kw...)

_wrap(lat::Lattice, axis) = Lattice(_wrap(lat.bravais, axis), lat.unitcell)

function _wrap(br::Bravais{E,L}, axis) where {E,L}
    mask = deleteat(SVector{L}(1:L), axis)
    return Bravais(br.matrix[:, mask], br.semibounded[mask])
end

function _wrap(harmonics::Vector{HamiltonianHarmonic{L,M,A}}, axis, factor) where {L,M,A}
    harmonics´ = HamiltonianHarmonic{L-1,M,A}[]
    for har in harmonics
        dn = har.dn
        dn´ = deleteat(dn, axis)
        factor´ = iszero(dn[axis]) ? 1 : factor
        add_or_push!(harmonics´, dn´, har.h, factor´)
    end
    return harmonics´
end

function add_or_push!(hs::Vector{<:HamiltonianHarmonic}, dn, matrix::AbstractMatrix, factor)
    for h in hs
        if h.dn == dn
            h.h .+= factor .* matrix
            return h
        end
    end
    newh = HamiltonianHarmonic(dn, matrix)
    push!(hs, newh)
    return newh
end

#######################################################################
# combine
#######################################################################
"""
    combine(hams::Hamiltonian...; coupling = missing)

Build a new Hamiltonian `h` that combines all `hams` as diagonal blocks, and applies
`coupling::Model`, if provided, to build the off-diagonal couplings. Note that the diagonal
blocks are not modified by the coupling model.
"""
combine(hams::Hamiltonian...; coupling = missing) = _combine(coupling, hams...)

_combine(::Missing, hams...) = _combine(TightbindingModel(), hams...)

function _combine(model::TightbindingModel, hams::Hamiltonian...)
    lat = combine((h -> h.lattice).(hams)...)
    orbs = tuplejoin((h -> h.orbitals).(hams)...)
    builder = IJVBuilder(lat, orbs, hams...)
    model´ = offdiagonal(model, lat, nsublats.(hams))
    ham = hamiltonian_sparse!(builder, lat, orbs, model´)
    return ham
end

#######################################################################
# Bloch routines
#######################################################################
"""
    similarmatrix(h::Hamiltonian)

Create an uninitialized matrix of the same type and size    of the Hamiltonian's matrix,
calling `optimize!(h)` first to produce an optimal work matrix in the sparse case.

    similarmatrix(h::Hamiltonian, T::Type{<:AbstractMatrix})

Specifies the desired type `T` of the uninitialized matrix.

    similarmatrix(h::Hamiltonian, method::AbstractDiagonalizeMethod)

Adapts the type of the matrix (e.g. dense/sparse) to the specified `method`
"""
function similarmatrix(h, ::Type{A´} = matrixtype(h)) where {A´<:AbstractMatrix}
    optimize!(h)
    return _similarmatrix(parent(h), matrixtype(h), A´)
end

# We only provide the type combinastions that make sense
_similarmatrix(h, ::Type{A}, ::Type{A´}) where {A´,A<:A´} =
    similar(h.harmonics[1].h)

_similarmatrix(h, ::Type{A}, ::Type{AbstractMatrix{T´}}) where {T<:Number,A<:AbstractMatrix{T},T´<:Number} =
    similar(h.harmonics[1].h, T´)

_similarmatrix(h, ::Type{A}, ::Type{AbstractMatrix{T´}}) where {N,T<:SMatrix{N,N},A<:AbstractMatrix{T},T´<:SMatrix{N,N}} =
    similar(h.harmonics[1].h, T´)

_similarmatrix(h, ::Type{A}, ::Type{A´}) where {N,T<:SMatrix{N,N},A<:AbstractMatrix{T},T´<:SMatrix{N,N},A´<:DenseMatrix{T´}} =
    similar(A´, size(h)...)

_similarmatrix(h, ::Type{A}, ::Type{A´}) where {T<:Number,A<:AbstractMatrix{T},T´<:Number,A´<:DenseMatrix{T´}} =
    similar(A´, size(h)...)

_similarmatrix(h, ::Type{A}, ::Type{A´}) where {T<:SMatrix,A<:AbstractMatrix{T},T´<:Number,A´<:DenseMatrix{T´}} =
    similar(A´, flatsize(h)...)

_similarmatrix(h, ::Type{A}, ::Type{AbstractMatrix{T´}}) where {T<:SMatrix,A<:AbstractMatrix{T},T´<:Number} =
    _flatten(h.harmonics[1].h, length.(h.orbitals), h.lattice, T´)

_similarmatrix(h, ::Type{A}, ::Type{A´}) where {T<:SMatrix,A<:AbstractSparseMatrix{T},T´<:Number,A´<:AbstractSparseMatrix{T´}} =
    _flatten(h.harmonics[1].h, length.(h.orbitals), h.lattice, T´)

"""
    optimize!(h::Hamiltonian)

Prepare a sparse Hamiltonian `h` to increase the performance of subsequent calls to
`bloch(h, ϕs)` and `bloch!(matrix, h, ϕs)` by minimizing memory reshufflings. It also
adds missing structural zeros to the diagonal to enable shifts by `α*I` (for
shift-and-invert methods).

No optimization will be performed on non-sparse Hamiltonians, or those defined on
`Superlattice`s, for which Bloch Hamiltonians are lazily evaluated.

Note that when calling `similarmatrix(h)` on a sparse `h`, `optimize!` is called first.

# See also:
    `bloch`, `bloch!`
"""
function optimize!(ham::Hamiltonian{<:Lattice,L,M,A}) where {L,M,A<:SparseMatrixCSC}
    h0 = first(ham.harmonics)
    n, m = size(h0.h)
    iszero(h0.dn) || throw(ArgumentError("First Hamiltonian harmonic is not the fundamental"))
    nh = length(ham.harmonics)
    builder = SparseMatrixBuilder{M}(n, m)
    for col in 1:m
        for i in eachindex(ham.harmonics)
            h = ham.harmonics[i].h
            for p in nzrange(h, col)
                v = i == 1 ? nonzeros(h)[p] : zero(M)
                row = rowvals(h)[p]
                pushtocolumn!(builder, row, v, false) # skips repeated rows
            end
        end
        pushtocolumn!(builder, col, zero(M), false) # if not present already, add structural zeros to diagonal
        finalizecolumn!(builder)
    end
    ho = sparse(builder)
    copy!(h0.h, ho) # Inject new structural zeros into zero harmonics
    return ham
end
# IDEA: could sum and subtract all harmonics instead
# Tested, it is slower

function optimize!(ham::Hamiltonian{<:Lattice,L,M,A}) where {L,M,A<:AbstractMatrix}
    # @warn "Hamiltonian is not sparse. Nothing changed."
    return ham
end

function optimize!(ham::Hamiltonian{<:Superlattice})
    # @warn "Hamiltonian is defined on a Superlattice. Nothing changed."
    return ham
end

"""
    bloch(h::Hamiltonian{<:Lattice}, ϕs)

Build the Bloch Hamiltonian matrix of `h`, for Bloch phases `ϕs = (ϕ₁, ϕ₂,...)` (or an
`SVector(ϕs...)`). In terms of Bloch wavevector `k`, `ϕs = k * bravais(h)`, it is defined as
`H(ϕs) = ∑exp(-im * ϕs' * dn) h_dn` where `h_dn` are Bloch harmonics connecting unit cells
at a distance `dR = bravais(h) * dn`.

    bloch(h::Hamiltonian{<:Lattice})

Build the intra-cell Hamiltonian matrix of `h`, without adding any Bloch harmonics.

    bloch(h::Hamiltonian{<:Lattice}, ϕs, axis::Int)

A nonzero `axis` produces the derivative of the Bloch matrix respect to `ϕs[axis]` (i.e. the
velocity operator along this axis), `∂H(ϕs) = ∑ -im * dn[axis] * exp(-im * ϕs' * dn) h_dn`

    bloch(matrix, h::Hamiltonian{<:Lattice}, ϕs::NTuple{L,Real}, dnfunc::Function)

Generalization that applies a prefactor `dnfunc(dn) * exp(im * ϕs' * dn)` to the `dn`
harmonic.

    bloch(ph::ParametricHamiltonian, pϕs, [axis])

Same as above, but with `pϕs = (p₁,...,pᵢ, ϕ₁, ..., ϕⱼ)`, with `p` values for
`parameters(ph)` and `ϕ` Bloch phases.

    h |> bloch(ϕs, ...)

Functional forms of `bloch`, equivalent to `bloch(h, ϕs, ...)`

# Notes

`bloch` allocates a new matrix on each call. For a non-allocating version of `bloch`, see
`bloch!`.

# Examples

```jldoctest
julia> h = LatticePresets.honeycomb() |> hamiltonian(onsite(1) + hopping(2)) |> bloch((0, 0))
2×2 SparseMatrixCSC{Complex{Float64},Int64} with 4 stored entries:
  [1, 1]  =  13.0+0.0im
  [2, 1]  =  6.0+0.0im
  [1, 2]  =  6.0+0.0im
  [2, 2]  =  13.0+0.0im
```

# See also:
    `bloch!`, `optimize!`, `similarmatrix`
"""
bloch(ϕs, axis = 0) = h -> bloch(h, ϕs, axis)
bloch(h::Hamiltonian, args...) = bloch!(similarmatrix(h), h, args...)

"""
    bloch!(matrix, h::Hamiltonian, ϕs, [axis])

In-place version of `bloch`. Overwrite `matrix` with the Bloch Hamiltonian matrix of `h` for
the specified Bloch phases `ϕs = (ϕ₁,ϕ₂,...)` (see `bloch` for definition and API).  A
conventient way to obtain a `matrix` is to use `similarmatrix(h,...)`, which will return an
`AbstractMatrix` of the same type as the Hamiltonian's. Note, however, that matrix need not
be of the same type (e.g. it can be dense with `Number` eltype for a sparse `h` with
`SMatrix` block eltype).

    bloch!(matrix, ph::ParametricHamiltonian, pϕs, [axis])

Same as above, but with `pϕs = (p₁,...,pᵢ, ϕ₁, ..., ϕⱼ)`, with `p` values for
`parameters(ph)` and `ϕ` Bloch phases.

# Examples

```jldoctest
julia> h = LatticePresets.honeycomb() |> hamiltonian(hopping(2I), orbitals = (Val(2), Val(1)));

julia> bloch!(similarmatrix(h), h, (0, 0))
2×2 SparseMatrixCSC{StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4},Int64} with 4 stored entries:
  [1, 1]  =  [12.0+0.0im 0.0+0.0im; 0.0+0.0im 12.0+0.0im]
  [2, 1]  =  [6.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]
  [1, 2]  =  [6.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]
  [2, 2]  =  [12.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]

julia> bloch!(similarmatrix(h, AbstractMatrix{ComplexF64}), h, (0, 0))
3×3 SparseMatrixCSC{Complex{Float64},Int64} with 9 stored entries:
  [1, 1]  =  12.0+0.0im
  [2, 1]  =  0.0+0.0im
  [3, 1]  =  6.0+0.0im
  [1, 2]  =  0.0+0.0im
  [2, 2]  =  12.0+0.0im
  [3, 2]  =  0.0+0.0im
  [1, 3]  =  6.0+0.0im
  [2, 3]  =  0.0+0.0im
  [3, 3]  =  12.0+0.0im

julia> ph = parametric(h, @hopping!((t; α) -> α * t));

julia> bloch!(similarmatrix(ph, AbstractMatrix{ComplexF64}), ph, (2, 0, 0))
3×3 SparseMatrixCSC{Complex{Float64},Int64} with 9 stored entries:
  [1, 1]  =  24.0+0.0im
  [2, 1]  =  0.0+0.0im
  [3, 1]  =  12.0+0.0im
  [1, 2]  =  0.0+0.0im
  [2, 2]  =  24.0+0.0im
  [3, 2]  =  0.0+0.0im
  [1, 3]  =  12.0+0.0im
  [2, 3]  =  0.0+0.0im
  [3, 3]  =  24.0+0.0im
```

# See also:
    `bloch`, `optimize!`, `similarmatrix`
"""
bloch!(matrix, h::Hamiltonian, ϕs = (), axis = 0) = _bloch!(matrix, h, toSVector(ϕs), axis)

function _bloch!(matrix::AbstractMatrix, h::Hamiltonian{<:Lattice,L,M}, ϕs, axis::Number) where {L,M}
    rawmatrix = parent(matrix)
    if iszero(axis)
        _copy!(rawmatrix, first(h.harmonics).h, h) # faster copy!(dense, sparse) specialization
        add_harmonics!(rawmatrix, h, ϕs, dn -> 1)
    else
        fill!(rawmatrix, zero(M)) # There is no guarantee of same structure
        add_harmonics!(rawmatrix, h, ϕs, dn -> -im * dn[axis])
    end
    return matrix
end

function _bloch!(matrix::AbstractMatrix, h::Hamiltonian{<:Lattice,L,M}, ϕs, dnfunc::Function) where {L,M}
    prefactor0 = dnfunc(zero(ϕs))
    rawmatrix = parent(matrix)
    if iszero(prefactor0)
        fill!(rawmatrix, zero(eltype(rawmatrix)))
    else
        _copy!(rawmatrix, first(h.harmonics).h, h)
        rmul!(rawmatrix, prefactor0)
    end
    add_harmonics!(rawmatrix, h, ϕs, dnfunc)
    return matrix
end

add_harmonics!(zerobloch, h::Hamiltonian{<:Lattice}, ϕs::SVector{0}, _) = zerobloch

function add_harmonics!(zerobloch, h::Hamiltonian{<:Lattice,L}, ϕs::SVector{L}, dnfunc) where {L}
    ϕs´ = ϕs'
    for ns in 2:length(h.harmonics)
        hh = h.harmonics[ns]
        hhmatrix = hh.h
        prefactor = dnfunc(hh.dn)
        iszero(prefactor) && continue
        ephi = prefactor * cis(-ϕs´ * hh.dn)
        _add!(zerobloch, hhmatrix, h, ephi)
    end
    return zerobloch
end

############################################################################################
######## _copy! and _add! call specialized methods in tools.jl #############################
############################################################################################

_copy!(dest, src, h) = copy!(dest, src)
_copy!(dst::AbstractMatrix{<:Number}, src::SparseMatrixCSC{<:Number}, h) = _fast_sparse_copy!(dst, src)
_copy!(dst::DenseMatrix{<:Number}, src::SparseMatrixCSC{<:Number}, h) = _fast_sparse_copy!(dst, src)
_copy!(dst::DenseMatrix{<:SMatrix{N,N}}, src::SparseMatrixCSC{<:SMatrix{N,N}}, h) where {N} = _fast_sparse_copy!(dst, src)
_copy!(dst::AbstractMatrix{<:Number}, src::SparseMatrixCSC{<:SMatrix}, h) = flatten_sparse_copy!(dst, src, h)
_copy!(dst::DenseMatrix{<:Number}, src::DenseMatrix{<:SMatrix}, h) = flatten_dense_copy!(dst, src, h)

_add!(dest, src, h, α) = _plain_muladd!(dest, src, α)
_add!(dst::AbstractMatrix{<:Number}, src::SparseMatrixCSC{<:Number}, h, α = 1) = _fast_sparse_muladd!(dst, src, α)
_add!(dst::AbstractMatrix{<:SMatrix{N,N}}, src::SparseMatrixCSC{<:SMatrix{N,N}}, h, α = I) where {N} = _fast_sparse_muladd!(dst, src, α)
_add!(dst::AbstractMatrix{<:Number}, src::SparseMatrixCSC{<:SMatrix}, h, α = I) = flatten_sparse_muladd!(dst, src, h, α)
_add!(dst::DenseMatrix{<:Number}, src::DenseMatrix{<:SMatrix}, h, α = I) = flatten_dense_muladd!(dst, src, h, α)

function flatten_sparse_copy!(dst, src, h)
    fill!(dst, zero(eltype(dst)))
    norbs = length.(h.orbitals)
    offsets = h.lattice.unitcell.offsets
    offsets´ = flatoffsets(offsets, norbs)
    coloffset = 0
    for s´ in sublats(h.lattice)
        N´ = norbs[s´]
        for col in siterange(h.lattice, s´)
            for p in nzrange(src, col)
                val = nonzeros(src)[p]
                row = rowvals(src)[p]
                rowoffset, M´ = flatoffsetorbs(row, h.lattice, norbs, offsets´)
                for j in 1:N´, i in 1:M´
                    dst[i + rowoffset, j + coloffset] = val[i, j]
                end
            end
            coloffset += N´
        end
    end
    return dst
end

function flatten_sparse_muladd!(dst, src, h, α = I)
    norbs = length.(h.orbitals)
    offsets = h.lattice.unitcell.offsets
    offsets´ = flatoffsets(offsets, norbs)
    coloffset = 0
    for s´ in sublats(h.lattice)
        N´ = norbs[s´]
        for col in siterange(h.lattice, s´)
            for p in nzrange(src, col)
                val = α * nonzeros(src)[p]
                row = rowvals(src)[p]
                rowoffset, M´ = flatoffsetorbs(row, h.lattice, norbs, offsets´)
                for j in 1:N´, i in 1:M´
                    dst[i + rowoffset, j + coloffset] += val[i, j]
                end
            end
            coloffset += N´
        end
    end
    return dst
end

function flatten_dense_muladd!(dst, src, h, α = I)
    norbs = length.(h.orbitals)
    offsets = h.lattice.unitcell.offsets
    offsets´ = flatoffsets(offsets, norbs)
    coloffset = 0
    for s´ in sublats(h.lattice)
        N´ = norbs[s´]
        for col in siterange(h.lattice, s´)
            rowoffset = 0
            for s in sublats(h.lattice)
                M´ = norbs[s]
                for row in siterange(h.lattice, s)
                    val = α * src[row, col]
                    for j in 1:N´, i in 1:M´
                        dst[i + rowoffset, j + coloffset] += val[i, j]
                    end
                    rowoffset += M´
                end
            end
            coloffset += N´
        end
    end
    return dst
end

function flatten_dense_copy!(dst, src, h)
    fill!(dst, zero(eltype(dst)))
    return flatten_dense_muladd!(dst, src, h, I)
end

# sublat offsets after flattening (without padding zeros)
flatoffsets(offsets, norbs) = _flatoffsets((0,), offsets, norbs...)
_flatoffsets(offsets´::NTuple{N,Any}, offsets, n, ns...) where {N} =
    _flatoffsets((offsets´..., offsets´[end] + n * (offsets[N+1] - offsets[N])), offsets, ns...)
_flatoffsets(offsets´, offsets) = offsets´

# offset of site i after flattening
@inline flatoffset(args...) = first(flatoffsetorbs(args...))

function flatoffsetorbs(i, lat, norbs, offsets´)
    s = sublat(lat, i)
    N = norbs[s]
    offset = lat.unitcell.offsets[s]
    Δi = i - offset
    i´ = offsets´[s] + (Δi - 1) * N
    return i´, N
end

"""
    flatten(h::Hamiltonian)

Flatten a multiorbital Hamiltonian `h` into one with a single orbital per site. The
associated lattice is flattened also, so that there is one site per orbital for each initial
site (all at the same position). Note that in the case of sparse Hamiltonians, zeros in
hopping/onsite matrices are preserved as structural zeros upon flattening.

    h |> flatten()

Functional form equivalent to `flatten(h)` of `h |> flatten` (included for consistency with
the rest of the API).

# Examples

```jldoctest
julia> h = LatticePresets.honeycomb() |>
           hamiltonian(hopping(@SMatrix[1; 2], range = 1/√3, sublats = :A =>:B),
           orbitals = (Val(1), Val(2)))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 3 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a,), (:a, :a))
  Element type     : 2 × 2 blocks (Complex{Float64})
  Onsites          : 0
  Hoppings         : 3
  Coordination     : 1.5

julia> flatten(h)
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 3 (SparseMatrixCSC, sparse)
  Harmonic size    : 3 × 3
  Orbitals         : ((:flat,), (:flat,))
  Element type     : scalar (Complex{Float64})
  Onsites          : 0
  Hoppings         : 6
  Coordination     : 2.0
```
"""
flatten() = h -> flatten(h)

function flatten(h::Hamiltonian)
    all(isequal(1), norbitals(h)) && return copy(h)
    harmonics´ = [flatten(har, h.orbitals, h.lattice) for har in h.harmonics]
    lattice´ = flatten(h.lattice, h.orbitals)
    orbitals´ = (_ -> (:flat, )).(h.orbitals)
    return Hamiltonian(lattice´, harmonics´, orbitals´)
end

flatten(h::HamiltonianHarmonic, orbs, lat) =
    HamiltonianHarmonic(h.dn, _flatten(h.h, length.(orbs), lat))

function _flatten(src::SparseMatrixCSC{<:SMatrix{N,N,T}}, norbs::NTuple{S,Any}, lat, ::Type{T´} = T) where {N,T,S,T´}
    offsets´ = flatoffsets(lat.unitcell.offsets, norbs)
    dim´ = last(offsets´)

    builder = SparseMatrixBuilder{T´}(dim´, dim´, nnz(src) * N * N)

    for col in 1:size(src, 2)
        scol = sublat(lat, col)
        for j in 1:norbs[scol]
            for p in nzrange(src, col)
                row = rowvals(src)[p]
                srow = sublat(lat, row)
                rowoffset´ = flatoffset(row, lat, norbs, offsets´)
                val = nonzeros(src)[p]
                for i in 1:norbs[srow]
                    pushtocolumn!(builder, rowoffset´ + i, val[i, j])
                end
            end
            finalizecolumn!(builder, false)
        end
    end
    matrix = sparse(builder)
    return matrix
end

function _flatten(src::DenseMatrix{<:SMatrix{N,N,T}}, norbs::NTuple{S,Any}, lat, ::Type{T´} = T) where {N,T,S,T´}
    offsets´ = flatoffsets(lat.unitcell.offsets, norbs)
    dim´ = last(offsets´)
    matrix = similar(src, T´, dim´, dim´)

    for col in 1:size(src, 2), row in 1:size(src, 1)
        srow, scol = sublat(lat, row), sublat(lat, col)
        nrow, ncol = norbs[srow], norbs[scol]
        val = src[row, col]
        rowoffset´ = flatoffset(row, lat, norbs, offsets´)
        coloffset´ = flatoffset(col, lat, norbs, offsets´)
        for j in 1:ncol, i in 1:nrow
            matrix[rowoffset´ + i, coloffset´ + j] = val[i, j]
        end
    end
    return matrix
end

function flatten(lat::Lattice, orbs)
    length(orbs) == nsublats(lat) || throw(ArgumentError("Msmatch between sublattices and orbitals"))
    unitcell´ = flatten(lat.unitcell, length.(orbs))
    bravais´ = lat.bravais
    lat´ = Lattice(bravais´, unitcell´)
end

function flatten(unitcell::Unitcell, norbs::NTuple{S,Int}) where {S}
    offsets´ = [flatoffsets(unitcell.offsets, norbs)...]
    ns´ = last(offsets´)
    sites´ = similar(unitcell.sites, ns´)
    i = 1
    for sl in 1:S, site in sites(unitcell, sl), rep in 1:norbs[sl]
        sites´[i] = site
        i += 1
    end
    names´ = unitcell.names
    unitcell´ = Unitcell(sites´, names´, offsets´)
    return unitcell´
end


text:
ENV["GKSwstype"] = "nul"

using Printf
using Statistics
using Plots

using Oceananigans
using Oceananigans.Operators: ζ₃ᶠᶠᶜ
using Oceananigans.Advection
using Oceananigans.AbstractOperations
using Oceananigans.OutputWriters
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Utils: prettytime
using Oceananigans.Models.HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel, ExplicitFreeSurface
using Oceananigans.OutputReaders: FieldTimeSeries

using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Oceananigans.Advection: WENOVectorInvariantVel, WENOVectorInvariantVort, VectorInvariant, VelocityStencil, VorticityStencil

using CUDA: device!

device!(2)

#####
##### The Bickley jet
#####

Ψ(y) = - tanh(y)
U(y) = sech(y)^2

# A sinusoidal tracer
C(y, L) = sin(2π * y / L)

# Slightly off-center vortical perturbations
ψ̃(x, y, ℓ, k) = exp(-(y + ℓ/10)^2 / 2ℓ^2) * cos(k * x) * cos(k * y)

# Vortical velocity fields (ũ, ṽ) = (-∂_y, +∂_x) ψ̃
ũ(x, y, ℓ, k) = + ψ̃(x, y, ℓ, k) * (k * tan(k * y) + y / ℓ^2)
ṽ(x, y, ℓ, k) = - ψ̃(x, y, ℓ, k) * k * tan(k * x)

"""
    run_bickley_jet(output_time_interval = 2, stop_time = 200, arch = CPU(), Nh = 64, ν = 0,
                    momentum_advection = VectorInvariant())

Run the Bickley jet validation experiment until `stop_time` using `momentum_advection`
scheme or formulation, with horizontal resolution `Nh`, viscosity `ν`, on `arch`itecture.
"""
function run_bickley_jet(; output_time_interval = 2, stop_time = 200, arch = CPU(), Nh = 64, ν = 0,
                           momentum_advection = VectorInvariant())

    Nφ = Int(Nh/2)

    grid = LatitudeLongitudeGrid(arch, size=(Nh, Nφ, 1),
                                 radius = 1,
                                 longitude = (-180, 180), latitude=(-80, 80), z=(0, 1), halo = (4, 4, 4),
                                 precompute_metrics = true)

    @inline toplft(x, y) = (((x > 45) & (x < 135)) & (((y > 15) & (y < 30)) | ((y > 45) & (y < 60))))
    @inline botlft(x, y) = (((x > 45) & (x < 135)) & (((y < -15) & (y > -30)) | ((y < -45) & (y > -60))))

    @inline toprgt(x, y) = (((x < - 45) & (x > - 135)) & (((y > 15) & (y < 30)) | ((y > 45) & (y < 60))))
    @inline botrgt(x, y) = (((x < - 45) & (x > - 135)) & (((y < -15) & (y > -30)) | ((y < -45) & (y > -60))))

    @inline bottom(x, y) = Int(toplft(x, y) | botlft(x, y) | toprgt(x, y) | botrgt(x, y))

    grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -1))

    free_surface = ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver, gravitational_acceleration=1)

    model = HydrostaticFreeSurfaceModel(momentum_advection = momentum_advection,
                                          tracer_advection = WENO5(),
                                                      grid = grid,
                                                   tracers = :c,
                                                   closure = nothing,
                                              free_surface = free_surface,
                                                  coriolis = HydrostaticSphericalCoriolis(rotation_rate=1),
                                                  buoyancy = nothing)

    # ** Initial conditions **
    #
    # u, v: Large-scale jet + vortical perturbations
    #    c: Sinusoid

    # Parameters
    ϵ = 0.1 # perturbation magnitude
    ℓ = 0.5 # Gaussian width
    k = 0.5 # Sinusoidal wavenumber

    dr(x) = deg2rad(x)

    # Total initial conditions
    uᵢ(x, y, z) = U(dr(y)*8) + ϵ * ũ(dr(x)*2, dr(y)*8, ℓ, k)
    vᵢ(x, y, z) = ϵ * ṽ(dr(x)*2, dr(y)*4, ℓ, k)
    cᵢ(x, y, z) = C(dr(y)*8, grid.Ly)

    set!(model, u=uᵢ, v=vᵢ, c=cᵢ)

    wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=10.0)

    c = sqrt(model.free_surface.gravitational_acceleration)
    Δt = 0.1 * Array(model.grid.Δxᶜᶠᵃ.parent)[5] / c

    simulation = Simulation(model, Δt=Δt, stop_time=stop_time)

    progress(sim) = @printf("Iter: %d, time: %s, Δt: %s, max|u|: %.3f, max|η|: %.3f \n",
                            iteration(sim), prettytime(sim), prettytime(sim.Δt),
                            maximum(abs, model.velocities.u), maximum(abs, model.free_surface.η))

    simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

    wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=10.0)

    simulation.callbacks[:wizard]   = Callback(wizard, IterationInterval(10))

    # Output: primitive fields + computations
    u, v, w, c = merge(model.velocities, model.tracers)

    ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᶜ, grid; computed_dependencies=(u, v))

    ζ = Field(ζ_op)
    outputs = merge(model.velocities, model.tracers, (ζ=ζ, η=model.free_surface.η))

    name = typeof(model.advection.momentum).name.wrapper
    if model.advection.momentum isa WENOVectorInvariantVel
        name = "WENOVectorInvariantVel"
    elseif model.advection.momentum isa WENOVectorInvariantVort
        name = "WENOVectorInvariantVort"
    end

    @show experiment_name = "bickley_jet_Nh_$(Nh)_$(name)"

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, outputs,
                                schedule =TimeInterval(output_time_interval),
                                filename = experiment_name,
                                overwrite_existing = true)

    @info "Running a simulation of an unstable Bickley jet with $(Nh)² degrees of freedom..."

    run!(simulation)

    return experiment_name
end

"""
    visualize_bickley_jet(experiment_name)

Visualize the Bickley jet data associated with `experiment_name`.
"""
function visualize_bickley_jet(experiment_name)

    @info "Making a fun movie about an unstable Bickley jet..."

    filepath = experiment_name * ".jld2"

    ζ_timeseries = FieldTimeSeries(filepath, "ζ", boundary_conditions=nothing, location=(Face, Face, Center))
    c_timeseries = FieldTimeSeries(filepath, "c", boundary_conditions=nothing, location=(Face, Center, Center))

    grid = c_timeseries.grid

    xζ, yζ, zζ = nodes(ζ_timeseries)
    xc, yc, zc = nodes(c_timeseries)

    anim = @animate for (i, iteration) in enumerate(c_timeseries.times)

        @info "    Plotting frame $i from iteration $iteration..."

        ζ = ζ_timeseries[i]
        c = c_timeseries[i]
        t = ζ_timeseries.times[i]

        ζi = interior(ζ)[:, :, 1]
        ci = interior(c)[:, :, 1]

        kwargs = Dict(
                      :aspectratio => 1,
                      :linewidth => 0,
                      :colorbar => :none,
                      :ticks => nothing,
                      :clims => (-1, 1),
                      :xlims => (-grid.Lx/2, grid.Lx/2),
                      :ylims => (-grid.Ly/2, grid.Ly/2)
                     )

        ζ_plot = Plots.heatmap(xζ, yζ, clamp.(ζi, -1, 1)'; color = :balance, kwargs...)
        c_plot = Plots.heatmap(xc, yc, clamp.(ci, -1, 1)'; color = :thermal, kwargs...)

        ζ_title = @sprintf("ζ at t = %.1f", t)
        c_title = @sprintf("u at t = %.1f", t)

        Plots.plot(ζ_plot, c_plot, title = [ζ_title c_title], size = (4000, 2000))
    end

    mp4(anim, experiment_name * ".mp4", fps = 8)
end

# experiment_name = run_bickley_jet(momentum_advection=VectorInvariant(), Nh=128)
# experiment_name = run_bickley_jet(momentum_advection=WENO5(zweno=true, vector_invariant=true), Nh=128)

advection_schemes = [WENO5(vector_invariant=VelocityStencil()),
                     WENO5(vector_invariant=VorticityStencil()),
                     VectorInvariant()]

for Nx in [64, 128, 256, 512, 1024]
    for advection in advection_schemes
        experiment_name = run_bickley_jet(arch=GPU(), momentum_advection=advection, Nh=Nx)
    end
end

ext: jl


text:
showsafe(x) = (x === nothing) ? "nothing" : x

function gitcmd(path::String, gitconfig::Dict)
    cmd = ["git", "-C", path]
    for (n,v) in gitconfig
        push!(cmd, "-c")
        push!(cmd, "$n=$v")
    end
    Cmd(cmd)
end

"""
Write TOML data (with sorted keys).
"""
function write_toml(file::String, data::Dict)
    open(file, "w") do io
        TOML.print(io, data, sorted=true)
    end
end


text:
"""
    AbstractSymmetry

Supertype for symmetries in a spin system.

# External Interface
- [`apply(symmetry, index)`](@ref)

# Internal Interface
Subtypes must implement:
- [`_order(::MyType)`](@ref)
- `(::MyType)(x)` should apply the symmetry operation to a BINARY index `x` in range 0:2^N-1
"""
abstract type AbstractSymmetry end

"""
    apply(symmetry, index)

Compute the action of the symmetry operation on a basis state given by index.
"""
function apply end

# Default: convert between 1-based indexing and bit representation of spin bases
apply(symmetry::AbstractSymmetry, index) = symmetry(index-1)+1

"""
    _order(symmetry)

Return the size of the group generated by the symmetry; i.e. the smallest, positive number k
such that `symmetry` applied k times is the identity function.
"""
function _order end

# Act as a scalar for broadcasting.
# Since a symmetry behaves like a function object, this enables standard broadcasting.
Base.Broadcast.broadcastable(as::AbstractSymmetry) = Ref(as)

Base.:(==)(symm1::T, symm2::T) where T<:AbstractSymmetry = hash(symm1) == hash(symm2)

const hashs_seed_symmetry = UInt === UInt64 ? 0xcd254b02e510da2d : 0x2d61926b
## Note that this hash function is independent on the order of fields
## This means that Swap(1,2) and Swap(2,1) are considered identically - which is correct I think
function Base.hash(symm::AbstractSymmetry, h::UInt)
    hv = hashs_seed_symmetry
    fields = fieldnames(typeof(symm))
    for f in fields
        hv ⊻= hash(getfield(symm, f))
    end
    return hash(hv, h)
end


text:
using ComplexMixtures
using PDBTools
using Plots
using LaTeXStrings
using EasyFit

function fig() # to simplify globals

# Plot defaults
plot_font = "Computer Modern"
default(
    fontfamily=plot_font,
    linewidth=2.5,
    framestyle=:box,
    label=nothing,
    grid=false,
    palette=:tab10
)
scalefontsizes(); scalefontsizes(1.3)

# Load pdb file of the system
system = readPDB("./simulation/equilibrated.pdb")

# Select the atoms corresponding to glycerol and water
popc = select(system,"resname POPC")
water = select(system,"water")
ethanol = select(system,"resname ETOH")

# Set the complete membrane as a single solute
solute = Selection(popc,nmols=1)

# Load previously computed data
mddf_water_POPC = load("./results/mddf_water_POPC.json")
mddf_ethanol_POPC = load("./results/mddf_ethanol_POPC.json")

# Set a solvent structure, uses to retrieve the atomic Contributions
# of ethanol to the MDDF
solvent = Selection(ethanol,natomspermol=9)

# Water-POPC MDDF and KB integral. Here we use `movavg` from `EasyFit` to soften
# the noise
plot(layout=(2,1))
x = mddf_water_POPC.d # distances
y = movavg(mddf_water_POPC.mddf,n=10).x
plot!(x,y,label="Water",
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF",subplot=1
)
plot!(xlim=(0,10),subplot=1)

y = movavg(mddf_water_POPC.kb,n=10).x/1000
plot!(x,y,xlabel=L"\textrm{Distance / \AA}",ylabel=L"\textrm{KB~/~L~mol^{-1}}",subplot=2)
plot!(xlim=(0,10),subplot=2)

# Ethanol-POPC MDDF and KB integral
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot!(x,y,label="Ethanol",
    xlabel=L"\textrm{Distance / \AA}",ylabel="MDDF",subplot=1
)
plot!(xlim=(0,10),subplot=1)

y = movavg(mddf_ethanol_POPC.kb,n=10).x/1000
plot!(x,y,xlabel=L"\textrm{Distance / \AA}",ylabel=L"\textrm{KB~/~L~mol^{-1}}",subplot=2)
plot!(xlim=(0,10),subplot=2)

savefig("./results/mddf_kb.png")

# Contributions of the ethanol groups
groups = [
    (select(ethanol,"name O1 or name HO1"),"Hydroxyl"),
    (select(ethanol,"not name O1 and not name HO1"),"Aliphatic chain"),
]
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot(x,y,label="Total")
for group in groups
    group_contrib = contrib(solvent,mddf_ethanol_POPC.solvent_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2])
end
plot!(
    xlim=(1,8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF"
)
savefig("./results/mddf_ethanol_groups.png")

# Contributions of POPC groups to ethanol-POPC distribution. Bellow is the list
# of the atom types of each POPC group. We need to list those to select specifically
# the atoms of each chemical group of interest.
groups = [
  (["N","C12","H12A","H12B","C13","H13A","H13B","H13C","C14",
    "H14A","H14B","H14C","C15","H15A","H15B","H15C","C11","H11A","H11B"],"Choline"),
  ([ "P","O13","O14","O12" ],"Phosphate"),
  (["O11","C1","HA","HB","C2","HS","O21","C3","HX","HY","O31"],"Glycerol"),
  (["O22","C21","H2R","H2S","C22","C23","H3R","H3S","C24","H4R","H4S",
    "C25","H5R","H5S","C26","H6R","H6S","C27","H7R","H7S","C28","H8R","H8S",
    "C29","H91","C210","H101","C211","H11R","H11S","C212","H12R","H12S",
    "C213","H13R","H13S","C214","H14R","H14S","C215","H15R","H15S",
    "C216","H16R","H16S","C217","H17R","H17S","C218","H18R","H18S","H18T"],"Oleoyl"),
  (["C31","O32","C32","H2X","H2Y","C33","H3X","H3Y","C34","H4X","H4Y",
    "C35","H5X","H5Y","C36","H6X","H6Y","C37","H7X","H7Y","C38","H8X",
    "H8Y","C39","H9X","H9Y","C310","H10X","H10Y","C311","H11X","H11Y",
    "C312","H12X","H12Y","C313","H13X","H13Y","C314","H14X","H14Y","C315",
    "H15X","H15Y","C316","H16X","H16Y","H16Z"],"Palmitoyl")
]
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot(x,y,label="Total")
for group in groups
    # Retrieve the contributions of the atoms of this group
    group_contrib = contrib(solute,mddf_ethanol_POPC.solute_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2])
end
plot!(
    xlim=(1.3,5),
    ylim=(0,1.8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF"
)
savefig("./results/mddf_popc_ethanol_groups.png")

# the same for POPC-water
x = mddf_water_POPC.d
y = movavg(mddf_water_POPC.mddf,n=10).x
plot(x,y,label="Total")
for group in groups
    group_contrib = contrib(solute,mddf_water_POPC.solute_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2])
end
plot!(
    xlim=(1.3,5),
    ylim=(0,1.8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF"
)
savefig("./results/mddf_popc_water_groups.png")

# Map of interactions of ethanol with Oleoyl groups. Now we split the
# Oleoyl chain into its components, ordered along the chain.
oleoyl_groups = [
  (["O22","C21"],"CO"),
  (["H2R","H2S","C22"],L"\textrm{CH_2}"),
  (["C23","H3R","H3S"],L"\textrm{CH_2}"),
  (["C24","H4R","H4S"],L"\textrm{CH_2}"),
  (["C25","H5R","H5S"],L"\textrm{CH_2}"),
  (["C26","H6R","H6S"],L"\textrm{CH_2}"),
  (["C27","H7R","H7S"],L"\textrm{CH_2}"),
  (["C28","H8R","H8S"],L"\textrm{CH_2}"),
  (["C29","H91"],L"\textrm{CH}"),
  (["C210","H101"],L"\textrm{CH}"),
  (["C211","H11R","H11S"],L"\textrm{CH_2}"),
  (["C212","H12R","H12S"],L"\textrm{CH_2}"),
  (["C213","H13R","H13S"],L"\textrm{CH_2}"),
  (["C214","H14R","H14S"],L"\textrm{CH_2}"),
  (["C215","H15R","H15S"],L"\textrm{CH_2}"),
  (["C216","H16R","H16S"],L"\textrm{CH_2}"),
  (["C217","H17R","H17S"],L"\textrm{CH_2}"),
  (["C218","H18R","H18S","H18T"],L"\textrm{CH_3}"),
]

gcontrib = zeros(length(mddf_ethanol_POPC.d),length(oleoyl_groups))
for (ig, group) in pairs(oleoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_ethanol_POPC.solute_atom,group[1]),n=10).x
end
labels = [ oleoyl_groups[i][2] for i in 1:length(oleoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_ethanol_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_ethanol_POPC.d)

plot(layout=(2,1))
contourf!(
    1:length(oleoyl_groups),
    mddf_ethanol_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    ylabel=L"r/\AA",xrotation=60,
    xticks=(1:length(oleoyl_groups),labels),subplot=1,
)

# And with Palmitoyl groups
palmitoyl_groups = [
  (["C31","O32"],"CO"),
  (["C32","H2X","H2Y"],L"\textrm{CH_2}"),
  (["C33","H3X","H3Y"],L"\textrm{CH_2}"),
  (["C34","H4X","H4Y"],L"\textrm{CH_2}"),
  (["C35","H5X","H5Y"],L"\textrm{CH_2}"),
  (["C36","H6X","H6Y"],L"\textrm{CH_2}"),
  (["C37","H7X","H7Y"],L"\textrm{CH_2}"),
  (["C38","H8X","H8Y"],L"\textrm{CH_2}"),
  (["C39","H9X","H9Y"],L"\textrm{CH_2}"),
  (["C310","H10X","H10Y"],L"\textrm{CH_2}"),
  (["C311","H11X","H11Y"],L"\textrm{CH_2}"),
  (["C312","H12X","H12Y"],L"\textrm{CH_2}"),
  (["C313","H13X","H13Y"],L"\textrm{CH_2}"),
  (["C314","H14X","H14Y"],L"\textrm{CH_2}"),
  (["C315","H15X","H15Y"],L"\textrm{CH_2}"),
  (["C316","H16X","H16Y","H16Z"],L"\textrm{CH_3}")
]
gcontrib = zeros(length(mddf_ethanol_POPC.d),length(palmitoyl_groups))
for (ig, group) in pairs(palmitoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_ethanol_POPC.solute_atom,group[1]),n=10).x
end
labels = [ palmitoyl_groups[i][2] for i in 1:length(palmitoyl_groups) ]

# This creates a contour plot, with some options to make the plot look nicer.
contourf!(
    1:length(palmitoyl_groups),
    mddf_ethanol_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    xlabel="Group",ylabel=L"r/\AA",xrotation=60,
    xticks=(1:length(palmitoyl_groups),labels),subplot=2,
)

annotate!( 14, 2.7, text("Oleoyl", :left, 12, plot_font), subplot=1)
annotate!( 12, 2.7, text("Palmitoyl", :left, 12, plot_font), subplot=2)

savefig("./results/map2D_ethanol_aliphatic_chains.png")

# And now the same for water-POPC
gcontrib = zeros(length(mddf_water_POPC.d),length(oleoyl_groups))
for (ig, group) in pairs(oleoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_water_POPC.solute_atom,group[1]),n=10).x
end
labels = [ oleoyl_groups[i][2] for i in 1:length(oleoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_water_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_water_POPC.d)

plot(layout=(2,1))
contourf!(
    1:length(oleoyl_groups),
    mddf_water_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    ylabel=L"r/\AA",xrotation=60,
    xticks=(1:length(oleoyl_groups),labels),subplot=1
)

gcontrib = zeros(length(mddf_water_POPC.d),length(palmitoyl_groups))
for (ig, group) in pairs(palmitoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_water_POPC.solute_atom,group[1]),n=10).x
end
labels = [ palmitoyl_groups[i][2] for i in 1:length(palmitoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_water_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_water_POPC.d)

contourf!(
    1:length(palmitoyl_groups),
    mddf_water_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    xlabel="Group",ylabel=L"r/\AA",xrotation=60,
    xticks=(1:length(palmitoyl_groups),labels),subplot=2
)

annotate!( 14, 2.7, text("Oleoyl", :left, 12, plot_font), subplot=1)
annotate!( 12, 2.7, text("Palmitoyl", :left, 12, plot_font), subplot=2)

savefig("./results/map2D_water_aliphatic_chains.png")

end; fig()


text:
using ShirleyRayTracer
using Test

function single_scancol()
    scene = Scene(Camera(Point3(13.,2.,3.), zero(Point3), Vec3(0,1,0), 20, 16/9, 0.1, 10.0))
    add!(scene, Sphere(Point3(0, 1, 0), 1.0, Dielectric(1.5)))
	add!(scene, Sphere(Point3(-4, 1, 0), 1.0, Lambertian(0.4,0.2,0.1)))
	add!(scene, Sphere(Point3(4, 1, 0), 1.0, Metal(0.7,0.6,0.5, 0.0)))
    @time ShirleyRayTracer.trace_scancol(scene, 600, 1, 1200, 100, 5)
    true
end

@testset "ShirleyRayTracer.jl" begin
    @test single_scancol()
end


text:
FLUX = 1.0mol/hr

ShootParamsCNE(;kwargs...) = Params(;name=:shoot, turnover_pars=TurnoverCNE, assimilation_pars=KooijmanSLAPhotosynthesis(), kwargs...)
RootParamsCNE(;kwargs...) = Params(;name=:root, turnover_pars=TurnoverCNE, assimilation_pars=NAssim(), kwargs...)
PlantCNE(;kwargs...) = Plant(;params=(ShootParams(), RootParams()), kwargs...)

ShootParamsCN(;kwargs...) = Params(;name=:shoot, assimilation_pars=KooijmanSLAPhotosynthesis(), kwargs...)
RootParamsCN(;kwargs...) = Params(;name=:root, assimilation_pars=NAssim(), kwargs...)
PlantCN(;kwargs...) = Plant(;params=(ShootParamsCN(), RootParamsCN()), kwargs...)

FvCBShootParams(;kwargs...) = Params(;name=:shoot, assimilation_pars=FvCBPhotosynthesis(), kwargs...)
FvCBRootParams(;kwargs...) = Params(;name=:root, assimilation_pars=NAssim(), kwargs...)
FvCBLeafParams3(;kwargs...) = Params(;name=:leaf, assimilation_pars=FvCBPhotosynthesis(),
                                    trans_pars=LosslessMultipleTranslocation(destnames=:stem), kwargs...)
FvCBRootParams3(;kwargs...) = Params(;name=:root, trans_pars=LosslessMultipleTranslocation(destnames=:stem),
                                    assimilation_pars=NAssim(), kwargs...)
FvCBStemParams3(;kwargs...) = Params(;name=:stem, assimilation_pars=nothing, trans_pars=LosslessMultipleTranslocation((:leaf,:root), 0.5), kwargs...)
FvCBNoRoots(;kwargs...) = Plant(;params=(FvCBShoot()), kwargs...)
FvCBPlant(;kwargs...) = Plant(;params=(FvCBShootParams(), FvCBRootParams()), kwargs...)
FvCBPlant3(;kwargs...) = Plant(;params=(FvCBLeafParams(), FvCBStemParams(), FvCBRootParams()), kwargs...)

ConstantShootParams(;kwargs...) = Params(turnover_pars=TurnoverCNE, assimilation_pars=ConstantCAssim())
ConstantRootParams(;kwargs...) = Params(turnover_pars=TurnoverCNE, assimilation_pars=ConstantNAssim())
ConstantPlantCNE(;kwargs...) = Plant(; params=(ConstantShootParamsCNE(), ConstantRootParamsCNE()), kwargs...)

ConstantShootParamsCN(;kwargs...) = Params(assimilation_pars=ConstantCAssim())
ConstantRootParamsCN(;kwargs...) = Params(assimilation_pars=ConstantNAssim())
ConstantPlantCN(;kwargs...) = Plant(; params=(ConstantShootParamsCN(), ConstantRootParamsCN()), kwargs...)

LeafParams(;kwargs...) = Params(;name=:leaf, trans_pars=LosslessMultipleTranslocation(destnames=:stem),
                                assimilation_pars=KooijmanSLAPhotosynthesis(), kwargs...)
StemParams(;kwargs...) = Params(;name=:stem, trans_pars=LosslessMultipleTranslocation(destnames=(:leaf, :root), proportions=0.5),
                                assimilation_pars=nothing, kwargs...)
RootParams3(;kwargs...) = Params(;name=:root, trans_pars=LosslessMultipleTranslocation(destnames=:stem),
                                 assimilation_pars=NAssim(), kwargs...)
Plant3(;kwargs...) = Plant(;params=(LeafParams(), StemParams(), RootParams()), kwargs...)


text:
using Revise
using PyCall
using PyPlot
using Infiltrator
using Bem2d

function fig_halfspace()
    mu = 30e9
    nu = 0.25

    # Free surface
    els = Elements(Int(1e5))
    nfreesurf = 120
    x1, y1, x2, y2 = discretizedline(-50, 0, 50, 0, nfreesurf)
    for i in 1:length(x1)
        els.x1[els.endidx + i] = x1[i]
        els.y1[els.endidx + i] = y1[i]
        els.x2[els.endidx + i] = x2[i]
        els.y2[els.endidx + i] = y2[i]
        els.name[els.endidx + i] = "freesurf"
    end
    standardize_elements!(els)

    # 45 degree dipping fault
    nfault = 1
    x1, y1, x2, y2 = discretizedline(-1, -1, 0, 0, nfault)
    for i in 1:length(x1)
        els.x1[els.endidx + i] = x1[i]
        els.y1[els.endidx + i] = y1[i]
        els.x2[els.endidx + i] = x2[i]
        els.y2[els.endidx + i] = y2[i]
        els.name[els.endidx + i] = "fault"
    end
    standardize_elements!(els)

    # Convenience data structures
    idx = getidxdict(els)
    partialsconst = initpartials(els)
    partialsquad = initpartials(els)

    # Constant slip fault
    partialsconst["disp"]["fault"]["freesurf"], _, partialsconst["trac"]["fault"]["freesurf"] = partialsconstdispstress(slip2dispstress, els, idx["fault"], idx["freesurf"], mu, nu)
    partialsconst["disp"]["freesurf"]["freesurf"], _, partialsconst["trac"]["freesurf"]["freesurf"] = partialsconstdispstress(slip2dispstress, els, idx["freesurf"], idx["freesurf"], mu, nu)
    faultslipconst = sqrt(2) / 2 * [1 ; 1]
    dispfullspaceconst = partialsconst["disp"]["fault"]["freesurf"] * faultslipconst
    dispfreesurfaceconst = inv(partialsconst["trac"]["freesurf"]["freesurf"]) * (partialsconst["trac"]["fault"]["freesurf"] * faultslipconst)
    xplotconst = els.xcenter[idx["freesurf"]]

    # Quadratic slip fault
    partialsquad["disp"]["fault"]["freesurf"], _, partialsquad["trac"]["fault"]["freesurf"] = partialsquaddispstress(slip2dispstress, els, idx["fault"], idx["freesurf"], mu, nu)
    partialsquad["disp"]["freesurf"]["freesurf"], _, partialsquad["trac"]["freesurf"]["freesurf"] = partialsquaddispstress(slip2dispstress, els, idx["freesurf"], idx["freesurf"], mu, nu)
    xplotquad = sort(els.xnodes[idx["freesurf"], :][:])
    faultslipquad = sqrt(2) ./ 2 * ones(6)
    dispfullspacequad = partialsquad["disp"]["fault"]["freesurf"] * faultslipquad
    dispfreesurfacequad = inv(partialsquad["trac"]["freesurf"]["freesurf"]) * (partialsquad["trac"]["fault"]["freesurf"] * faultslipquad)

    # Okada solution
    okadawrapper = pyimport("okada_wrapper")# from okada_wrapper import dc3dwrapper
    xokada = collect(LinRange(-5, 5, 1000))
    dispxokada = zeros(length(xokada))
    dispyokada = zeros(length(xokada))
    for i in 1:length(xokada)
        _, u, _ = okadawrapper.dc3dwrapper(2.0 / 3.0, [0, xokada[i] + 0.5, 0],
            0.5, 45, [-10000, 10000], [-sqrt(2) / 2, sqrt(2) / 2], [0.0, 1.0, 0.0])
        dispxokada[i] = u[2]
        dispyokada[i] = u[3]
    end

    fontsize = 6
    markersize = 4
    linewidth = 0.5
    close("all")
    figure(figsize = (4, 4))
    ax = subplot(2, 1, 1)
    plot(xokada, dispxokada, "-r", linewidth=linewidth, label="Okada", zorder=1)
    plot(xplotconst, dispfreesurfaceconst[1:2:end], ".k", markeredgewidth=linewidth, markersize=markersize, label = "CS BEM", zorder=3)
    plot(xplotquad, dispfreesurfacequad[1:2:end], "ok", markerfacecolor = "lightgray", markeredgewidth = 0.25, markersize=markersize, label = "3NQ BEM", zorder=2)
    gca().set_xlim([-5, 5]); gca().set_ylim([-1.0, 1.0])
    gca().set_xticks([-5, 0, 5]); gca().set_yticks([-1.00, 0.00, 1.00])
    legend(fontsize=fontsize)
    ax.tick_params("both", labelsize = fontsize)
    ylabel(L"$u_x$ (m)", fontsize=fontsize)

    ax = subplot(2, 1, 2)
    plot(xokada, dispyokada, "-r", linewidth=linewidth, label="Okada", zorder = 1)
    plot(xplotconst, dispfreesurfaceconst[2:2:end], ".k", markeredgewidth=linewidth, markersize=markersize, label = "CS BEM", zorder = 3)
    plot(xplotquad, dispfreesurfacequad[2:2:end], "ok", markerfacecolor = "lightgray", markeredgewidth = 0.25, markersize=markersize, label = "3NQ BEM", zorder = 2)
    gca().set_xlim([-5, 5]); gca().set_ylim([-1.0, 1.0])
    gca().set_xticks([-5, 0, 5]); gca().set_yticks([-1.00, 0.00, 1.00])
    legend(fontsize=fontsize)
    ax.tick_params("both", labelsize = fontsize)
    xlabel(L"$x$ (m)", fontsize=fontsize)
    ylabel(L"$u_y$ (m)", fontsize=fontsize)
    tight_layout()
    show()
end
fig_halfspace()


text:
module LogBench

using MicroLogging

struct TrivialLogger <: AbstractLogger; end

MicroLogging.min_enabled_level(::TrivialLogger) = MicroLogging.Warn
MicroLogging.shouldlog(::TrivialLogger, level, a...) = level >= MicroLogging.Error
MicroLogging.handle_message(::TrivialLogger, a...; kws...) = println(a...)

end

using BenchmarkTools
using MicroLogging

# Here's various loops which would be tight if they didn't have logging
# statements inside.  The benchmarks measure the impact of filtering in various
# ways.

@noinline function simple_println(N, dolog)
    s = 0.0
    for i=1:N
        if dolog
            println("A simple operation i=$i")
        end
        s += 0.3
    end
    s
end

@noinline function simple_println_ref(N, dolog)
    s = 0.0
    for i=1:N
        if dolog[]
            println("A simple operation i=$i")
        end
        s += 0.3
    end
    s
end

@noinline function log_debug_msgs(N)
    s = 0.0
    for i=1:N
        @debug "A simple operation" i
        s += 0.3
    end
    s
end

@noinline function log_info_msgs(N)
    s = 0.0
    for i=1:N
        @info "A simple operation" i
        s += 0.3
    end
    s
end

@noinline function log_warn_msgs(N)
    s = 0.0
    for i=1:N
        @warn "A simple operation" i
        s += 0.3
    end
    s
end

N = 1000

#-------------------------------------------------------------------------------
# The benchmarks compare the efficiency of early log filtering in the logging
# macros to the best case hand coded scenario: a boolean value guarding a call
# to println().
#
# In simple_println() the compiler knows the test is a loop invariant and might
# hoist it out entirely.
bench_println = @benchmark simple_println($N, false)

# In simple_println_ref(), the compiler can't hoist the load out of the loop.
# If the loop is simple enough this can make a huge difference.
bench_println_ref = @benchmark simple_println_ref($N, Ref(false))

#-------------------------------------------------------------------------------
# Benchmarks for the three ways that logs can be disabled using MicroLogging:
# 1) The global level filter
disable_logging(:debug)
bench_global_filter =
with_logger(LogBench.TrivialLogger()) do
    @benchmark log_debug_msgs($N)
end

# 2) The LogState level filter
# This could be a lot faster, but going through the TLS hash map is somewhat
# costly.
bench_logstate_filter =
with_logger(LogBench.TrivialLogger()) do
    @benchmark log_info_msgs($N)
end

# 3) The generic shouldlog() filter
bench_shouldlog_filter =
with_logger(LogBench.TrivialLogger()) do
    @benchmark log_warn_msgs($N)
end

# current_task().result = mylog

@info "Baseline print with untaken branch"
display(bench_println)

@info "Baseline print with untaken branch, loading from a ref"
display(bench_println_ref)

@info "global filter"
display(bench_global_filter)

@info "logstate filter"
display(bench_logstate_filter)

@info "shouldlog filter"
display(bench_shouldlog_filter)


text:
using PastaQ
using ITensors
using Test
using LinearAlgebra
using Random

""" HELPER FUNCTIONS """
function numgradslogZ(L::LPDO; accuracy=1e-8)
  M = L.X
  N = length(M)
  grad_r = []
  grad_i = []
  for j in 1:N
    push!(grad_r, zeros(ComplexF64, size(M[j])))
    push!(grad_i, zeros(ComplexF64, size(M[j])))
  end

  epsilon = zeros(ComplexF64, size(M[1]))
  # Site 1
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = 2.0 * lognorm(M)
    M[1] -= eps
    loss_m = 2.0 * lognorm(M)
    grad_r[1][i] = (loss_p - loss_m) / (accuracy)

    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = 2.0 * lognorm(M)
    M[1] -= eps
    loss_m = 2.0 * lognorm(M)
    grad_i[1][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  for j in 2:(N - 1)
    epsilon = zeros(ComplexF64, size(M[j]))
    for i in 1:length(epsilon)
      epsilon[i] = accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = 2.0 * lognorm(M)
      M[j] -= eps
      loss_m = 2.0 * lognorm(M)
      grad_r[j][i] = (loss_p - loss_m) / (accuracy)

      epsilon[i] = im * accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = 2.0 * lognorm(M)
      M[j] -= eps
      loss_m = 2.0 * lognorm(M)
      grad_i[j][i] = (loss_p - loss_m) / (im * accuracy)

      epsilon[i] = 0.0
    end
  end
  # Site N
  epsilon = zeros(ComplexF64, size(M[N]))
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = 2.0 * lognorm(M)
    M[N] -= eps
    loss_m = 2.0 * lognorm(M)
    grad_r[N][i] = (loss_p - loss_m) / (accuracy)

    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = 2.0 * lognorm(M)
    M[N] -= eps
    loss_m = 2.0 * lognorm(M)
    grad_i[N][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  return grad_r - grad_i
end

#numgradslogZ(M::MPS; kwargs...) = numgradslogZ(LPDO(M); kwargs...)

function numgradsnll(L::LPDO, data::Matrix{Pair{String,Pair{String,Int}}}, accuracy=1e-8)
  M = L.X
  N = length(M)
  grad_r = []
  grad_i = []
  for j in 1:N
    push!(grad_r, zeros(ComplexF64, size(M[j])))
    push!(grad_i, zeros(ComplexF64, size(M[j])))
  end

  epsilon = zeros(ComplexF64, size(M[1]))
  # Site 1
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = PastaQ.nll(L, data)
    M[1] -= eps
    loss_m = PastaQ.nll(L, data)
    grad_r[1][i] = (loss_p - loss_m) / (accuracy)

    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = PastaQ.nll(L, data)
    M[1] -= eps
    loss_m = PastaQ.nll(L, data)
    grad_i[1][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  for j in 2:(N - 1)
    epsilon = zeros(ComplexF64, size(M[j]))
    for i in 1:length(epsilon)
      epsilon[i] = accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = PastaQ.nll(L, data)
      M[j] -= eps
      loss_m = PastaQ.nll(L, data)
      grad_r[j][i] = (loss_p - loss_m) / (accuracy)

      epsilon[i] = im * accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = PastaQ.nll(L, data)
      M[j] -= eps
      loss_m = PastaQ.nll(L, data)
      grad_i[j][i] = (loss_p - loss_m) / (im * accuracy)

      epsilon[i] = 0.0
    end
  end

  # Site N
  epsilon = zeros(ComplexF64, size(M[N]))
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = PastaQ.nll(L, data)
    M[N] -= eps
    loss_m = PastaQ.nll(L, data)
    grad_r[N][i] = (loss_p - loss_m) / (accuracy)

    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = PastaQ.nll(L, data)
    M[N] -= eps
    loss_m = PastaQ.nll(L, data)
    grad_i[N][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  return grad_r - grad_i
end

function numgradsTP(L::LPDO; accuracy=1e-8)
  M = L.X
  N = length(M)
  grad_r = []
  grad_i = []

  for j in 1:N
    push!(grad_r, zeros(ComplexF64, size(M[j])))
    push!(grad_i, zeros(ComplexF64, size(M[j])))
  end

  epsilon = zeros(ComplexF64, size(M[1]))
  # Site 1
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = PastaQ.TP(L)
    M[1] -= eps
    loss_m = PastaQ.TP(L)
    grad_r[1][i] = (loss_p - loss_m) / (accuracy)

    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[1]))
    M[1] += eps
    loss_p = PastaQ.TP(L)
    M[1] -= eps
    loss_m = PastaQ.TP(L)
    grad_i[1][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  for j in 2:(N - 1)
    epsilon = zeros(ComplexF64, size(M[j]))
    for i in 1:length(epsilon)
      epsilon[i] = accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = PastaQ.TP(L)
      M[j] -= eps
      loss_m = PastaQ.TP(L)
      grad_r[j][i] = (loss_p - loss_m) / (accuracy)

      epsilon[i] = im * accuracy
      eps = ITensor(epsilon, inds(M[j]))
      M[j] += eps
      loss_p = PastaQ.TP(L)
      M[j] -= eps
      loss_m = PastaQ.TP(L)
      grad_i[j][i] = (loss_p - loss_m) / (im * accuracy)

      epsilon[i] = 0.0
    end
  end

  # Site N
  epsilon = zeros(ComplexF64, size(M[N]))
  for i in 1:length(epsilon)
    epsilon[i] = accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = PastaQ.TP(L)
    M[N] -= eps
    loss_m = PastaQ.TP(L)
    grad_r[N][i] = (loss_p - loss_m) / (accuracy)
    epsilon[i] = im * accuracy
    eps = ITensor(epsilon, inds(M[N]))
    M[N] += eps
    loss_p = PastaQ.TP(L)
    M[N] -= eps
    loss_m = PastaQ.TP(L)
    grad_i[N][i] = (loss_p - loss_m) / (im * accuracy)

    epsilon[i] = 0.0
  end

  return grad_r - grad_i
end

@testset "mpo-qpt: normalization" begin
  N = 5
  χ = 4
  U = randomprocess(N; χ=χ)
  Λ = LPDO(PastaQ.unitary_mpo_to_choi_mps(U))
  @test length(Λ) == N
  logZ = 2 * lognorm(Λ.X)
  sqrt_localZ = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localZ, localnorm=2)
  @test logZ ≈ N * log(2) + 2.0 * sum(log.(sqrt_localZ))
  @test abs2(norm(Λ.X)) ≈ 2^N
end

@testset "mpo-qpt: grad logZ" begin
  N = 5
  χ = 4

  Random.seed!(1234)
  U = randomprocess(N; χ=χ)
  Λ = LPDO(PastaQ.unitary_mpo_to_choi_mps(U))
  num_grad = numgradslogZ(Λ)

  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)
  @test norm(Λ.X)^2 ≈ 2^N
  alg_grad, _ = PastaQ.gradlogZ(Λ; sqrt_localnorms=sqrt_localnorms)

  alg_gradient = permutedims(ITensors.array(alg_grad[1]), [1, 3, 2])
  @test alg_gradient ≈ num_grad[1] rtol = 1e-3
  for j in 2:(N - 1)
    alg_gradient = permutedims(ITensors.array(alg_grad[j]), [2, 1, 3, 4])
    @test alg_gradient ≈ num_grad[j] rtol = 1e-3
  end
  alg_gradient = permutedims(ITensors.array(alg_grad[N]), [2, 1, 3])
  @test alg_gradient ≈ num_grad[N] rtol = 1e-3
end

@testset "mpo-qpt: grad nll" begin
  N = 4
  χ = 2
  nsamples = 10
  Random.seed!(1234)
  data_in = randompreparations(N, nsamples)
  data_out = PastaQ.convertdatapoints(randompreparations(N, nsamples))
  data = data_in .=> data_out

  U = randomprocess(N; χ=χ)
  Λ = LPDO(PastaQ.unitary_mpo_to_choi_mps(U))
  num_grad = numgradsnll(Λ, data)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)

  alg_grad, _ = PastaQ.gradnll(Λ, data; sqrt_localnorms=sqrt_localnorms)
  for j in 1:N
    @test ITensors.array(alg_grad[j]) ≈ num_grad[j] rtol = 1e-3
  end
end

@testset "mpo-qpt: grad TP" begin
  N = 3
  χ = 3
  Random.seed!(1234)

  Random.seed!(1234)
  U = randomprocess(N; χ=χ)
  Λ = LPDO(PastaQ.unitary_mpo_to_choi_mps(U))

  num_grad = numgradsTP(Λ; accuracy=1e-8)
  Γ_test = PastaQ.TP(Λ)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)

  alg_grad_logZ, logZ = PastaQ.gradlogZ(Λ; sqrt_localnorms=sqrt_localnorms)

  alg_grad, Γ = PastaQ.gradTP(Λ, alg_grad_logZ, logZ; sqrt_localnorms=sqrt_localnorms)

  @test Γ ≈ Γ_test
  alg_gradient = permutedims(ITensors.array(alg_grad[1]), [1, 3, 2])
  @test alg_gradient ≈ num_grad[1] rtol = 1e-5
  for j in 2:(N - 1)
    alg_gradient = permutedims(ITensors.array(alg_grad[j]), [1, 3, 2, 4])
    @test alg_gradient ≈ num_grad[j] rtol = 1e-5
  end
  alg_gradient = permutedims(ITensors.array(alg_grad[N]), [1, 3, 2])
  @test alg_gradient ≈ num_grad[N] rtol = 1e-5
end

@testset "mpo-qpt: full gradients" begin
  N = 3
  χ = 4
  nsamples = 10
  trace_preserving_regularizer = 0.1
  Random.seed!(1234)
  data_in = randompreparations(N, nsamples)
  data_out = PastaQ.convertdatapoints(randompreparations(N, nsamples))
  data = data_in .=> data_out

  U = randomprocess(N; χ=χ)
  Λ = LPDO(PastaQ.unitary_mpo_to_choi_mps(U))
  TP_distance = PastaQ.TP(Λ)
  logZ = log(tr(Λ))
  NLL = PastaQ.nll(Λ, data)
  num_gradZ = numgradslogZ(Λ)
  num_gradNLL = numgradsnll(Λ, data)
  num_gradTP = numgradsTP(Λ; accuracy=1e-5)

  num_grads = num_gradZ + num_gradNLL + trace_preserving_regularizer * num_gradTP

  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)

  ex_loss = PastaQ.nll(Λ, data) + 2 * lognorm(Λ.X)
  alg_grads, loss = PastaQ.gradients(
    Λ,
    data;
    sqrt_localnorms=sqrt_localnorms,
    trace_preserving_regularizer=trace_preserving_regularizer,
  )
  @test ex_loss ≈ loss
  for j in 1:N
    @test ITensors.array(alg_grads[j]) ≈ num_grads[j] rtol = 1e-3
  end
end

""" CHOI TESTS """

@testset "lpdo-qpt: normalization" begin
  N = 10
  χ = 4
  ξ = 3
  Λ = randomprocess(N; mixed=true, ξ=ξ, χ=χ)

  @test length(Λ) == N
  logZ = logtr(Λ)
  localZ = []
  normalize!(Λ; (sqrt_localnorms!)=localZ, localnorm=2)
  @test logZ ≈ N * log(2) + 2.0 * sum(log.(localZ))
  @test tr(Λ) ≈ 2^N
end

@testset "lpdo-qst: grad logZ" begin
  N = 5
  χ = 4
  ξ = 3

  Λ = randomprocess(N; mixed=true, χ=χ, ξ=ξ)
  num_grad = numgradslogZ(Λ)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)
  @test tr(Λ) ≈ 2^N
  alg_grad, _ = PastaQ.gradlogZ(Λ; sqrt_localnorms=sqrt_localnorms)

  alg_gradient = permutedims(ITensors.array(alg_grad[1]), [1, 2, 4, 3])
  for j in 2:(N - 1)
    alg_gradient = permutedims(ITensors.array(alg_grad[j]), [2, 3, 1, 4, 5])
    @test alg_gradient ≈ num_grad[j] rtol = 1e-3
  end
  alg_gradient = permutedims(ITensors.array(alg_grad[N]), [2, 3, 1, 4])
  @test alg_gradient ≈ num_grad[N] rtol = 1e-3
end

@testset "lpdo-qst: grad nll" begin
  N = 5
  χ = 4
  ξ = 3

  nsamples = 10
  Random.seed!(1234)
  data_in = randompreparations(N, nsamples)
  data_out = PastaQ.convertdatapoints(randompreparations(N, nsamples))
  data = data_in .=> data_out

  Λ = randomprocess(N; mixed=true, χ=χ, ξ=ξ)
  num_grad = numgradsnll(Λ, data)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)
  alg_grad, loss = PastaQ.gradnll(Λ, data; sqrt_localnorms=sqrt_localnorms)
  @test loss ≈ PastaQ.nll(Λ, data)

  alg_gradient = permutedims(ITensors.array(alg_grad[1]), [3, 4, 1, 2])
  @test alg_gradient ≈ num_grad[1] rtol = 1e-3
  for j in 2:(N - 1)
    alg_gradient = permutedims(ITensors.array(alg_grad[j]), [4, 5, 2, 3, 1])
    @test alg_gradient ≈ num_grad[j] rtol = 1e-3
  end
  alg_gradient = permutedims(ITensors.array(alg_grad[N]), [3, 4, 1, 2])
  @test alg_gradient ≈ num_grad[N] rtol = 1e-3
end

@testset "lpdo-qpt: grad TP" begin
  N = 3
  χ = 4
  ξ = 3
  Λ = randomprocess(N; mixed=true, χ=χ, ξ=ξ)

  num_grad = numgradsTP(Λ; accuracy=1e-8)
  Γ_test = PastaQ.TP(Λ)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)
  Γ_test = PastaQ.TP(Λ)

  alg_grad_logZ, logZ = PastaQ.gradlogZ(Λ; sqrt_localnorms=sqrt_localnorms)

  alg_grad, Γ = PastaQ.gradTP(Λ, alg_grad_logZ, logZ; sqrt_localnorms=sqrt_localnorms)

  @test Γ ≈ Γ_test
  alg_gradient = permutedims(ITensors.array(alg_grad[1]), [3, 1, 4, 2])
  @test alg_gradient ≈ num_grad[1] rtol = 1e-3
  for j in 2:(N - 1)
    alg_gradient = permutedims(ITensors.array(alg_grad[j]), [3, 1, 4, 2, 5])
    @test alg_gradient ≈ num_grad[j] rtol = 1e-3
  end
  alg_gradient = permutedims(ITensors.array(alg_grad[N]), [3, 1, 4, 2])
  @test alg_gradient ≈ num_grad[N] rtol = 1e-3
end

@testset "lpdo-qpt: full gradients" begin
  N = 3
  χ = 3
  ξ = 2

  trace_preserving_regularizer = 0.1
  nsamples = 10
  Random.seed!(1234)
  data_in = randompreparations(N, nsamples)
  data_out = PastaQ.convertdatapoints(randompreparations(N, nsamples))
  data = data_in .=> data_out

  Λ = randomprocess(N; mixed=true, χ=χ, ξ=ξ)
  num_grad = numgradsnll(Λ, data)
  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)
  alg_grad, loss = PastaQ.gradnll(Λ, data; sqrt_localnorms=sqrt_localnorms)

  TP_distance = PastaQ.TP(Λ)
  logZ = log(tr(Λ))
  NLL = PastaQ.nll(Λ, data)
  num_gradZ = numgradslogZ(Λ)
  num_gradNLL = numgradsnll(Λ, data)
  num_gradTP = numgradsTP(Λ; accuracy=1e-5)

  num_grads = num_gradZ + num_gradNLL + trace_preserving_regularizer * num_gradTP

  sqrt_localnorms = []
  normalize!(Λ; (sqrt_localnorms!)=sqrt_localnorms, localnorm=2)

  ex_loss = PastaQ.nll(Λ, data) + 2 * lognorm(Λ.X)
  alg_grads, loss = PastaQ.gradients(
    Λ,
    data;
    sqrt_localnorms=sqrt_localnorms,
    trace_preserving_regularizer=trace_preserving_regularizer,
  )
  @test ex_loss ≈ loss
  for j in 1:N
    @test ITensors.array(alg_grads[j]) ≈ num_grads[j] rtol = 1e-3
  end
end


text:
module DataGenerators

#generator macro
export @generator

# exported functions that are used to generate objects from a generator
export choose, robustchoose, generate, meta, choicepointinfo, choicemodel, setchoicemodel!, GenerationTerminatedException

# exported functions that are used to register and match generators
export register, generatorfor

# exported choice models application function
export setsimplechoicemodel!, setsamplerchoicemodel!, setnmcschoicemodel!, mctschoicemodel!
export SimpleChoiceModel, SamplerChoiceModel, NMCSChoiceModel, MCTSChoiceModel, GEChoiceModel, SwitchChoiceModel

# exported Choice Model functions
export paramranges, setparams!, getparams, numparams, estimateparams!, estimateconditionalmodel!

using DataGeneratorTranslators
using Distributions
import Base.show

const THIS_MODULE = current_module() # used when creating calls to this module in the macro

# list of number types that are supported directly by generator macro
# (could do this automatically as leaf subtypes of Integer and AbstractFloat, but some of these we can't really handle directly yet
# - e.g. BigInt, BigFloat - and it is possible that custom subtypes could have been added)
const GENERATOR_SUPPORTED_CHOOSE_NUMBER_TYPES = Type[Bool, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, Float16, Float32, Float64,]

# list of number types that are supported directly by generator macro
# (could do this automatically as leaf subtypes of AbstractString, but some of these we can't really handle directly yet)
const GENERATOR_SUPPORTED_CHOOSE_STRING_TYPES = Type[String,]

# all choose types supported directly by the generator maco
const GENERATOR_SUPPORTED_CHOOSE_TYPES = Type[GENERATOR_SUPPORTED_CHOOSE_NUMBER_TYPES; GENERATOR_SUPPORTED_CHOOSE_STRING_TYPES]

# generation and associated functions
include("generation.jl")

# @generator macro
include("generator_macro.jl")

# updating and querying registry metadata
include("registry.jl")

# choice models
include(joinpath("choice_model","choice_model.jl"))

end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "eudev"
version = v"3.2.9"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://dev.gentoo.org/~blueness/eudev/eudev-3.2.9.tar.gz", "89618619084a19e1451d373c43f141b469c9fd09767973d73dd268b92074d4fc")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/eudev*
apk add gperf
./configure --prefix=${prefix} --build=${MACHTYPE} --host=${target}
make
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = filter!(Sys.islinux, supported_platforms())

# The products that we will ensure are always built
products = [
    LibraryProduct("libudev", :libudev),
    ExecutableProduct("udevd", :udevd, "sbin"),
    ExecutableProduct("udevadm", :udevadm)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="gperf_jll", uuid="1a1c6b14-54f6-533d-8383-74cd7377aa70"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
module Gage
using Libdl
@show include("gagestructs.jl")
using Main.GageStructs: BoardInfo
##############################################
####         GageAPI Methods              ####
##############################################
# Julia translations of the Gage Driver API.
# Low-level operations. Require GaGe Drivers Installed and Visible to Julia.
function geterror(errorcode::Int32)
    errstring = Vector{UInt8}(undef,255)
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsGetErrorString)
    err = ccall(sym,Int32,(Int32,Ref{UInt8}, Int32), errorcode, errstring, 255)
    errstring[end] = 0
    if err < 0
        return err
    else
        return unsafe_string(pointer(errstring))
    end
end

function initialize()
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsInitialize)
    err = ccall(sym,Int32,())
    if err < 0
        throw(SystemError)
    end
    err
end

function getsystem(boardtype=UInt32(0), nchannels=UInt32(0),
     bitresolution=UInt32(0), systemindex=Int16(0))
    @show boardtype
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsGetSystem)
    result = Ref{UInt32}(0)
    @show result
    argtype = (result, UInt32, UInt32, UInt32, Int16)
    args = (result, boardtype, nchannels, bitresolution, systemindex)
    err = ccall(sym,Int32,
    (Ref{Cuint},UInt32, UInt32, UInt32, Int16),
        result,boardtype,nchannels,bitresolution,systemindex)
    if err < 1
        return geterror(err)
    else
        @show err
        return result
    end
end

function getsysteminfo(handle::UInt32)
    result = Ref{BoardInfo}()
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsGetSystemInfo)
    err = ccall(sym,Int32,(UInt32,Ref{BoardInfo}),handle,result)
    if err < 0
        return geterror(err)
    end
    err
end

function getstatus(hSys::UInt32)
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsGetStatus)
    err = ccall(sym,Int32,(UInt32,),hSys)
    if err < 0
        throw(SystemError)
    end
    err
end

function freesystem(hSys::UInt32)
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsFreeSystem)
    err = ccall(sym,Int32,(UInt32,), hSys)
    if err < 0
        return geterror(err)
    else
        return true
    end
end

function gagedo(hSys::UInt32, operation::Integer)
    lib = dlopen("CsSsm")
    sym =  dlsym(lib,:CsDo)
    err = ccall(sym,Int32,(UInt32,Int16), hSys, operation)
    if err < 0
        return geterror(err)
    else
        return true
    end
end

function transfer(hSys::UInt)
end

end