text:
using Quantica.RegionPresets: Region

#######################################################################
# NeighborRange
#######################################################################
struct NeighborRange
    n::Int
end

"""
    nrange(n::Int)

Create a `NeighborRange` that represents a hopping range to distances corresponding to the
n-th nearest neighbors in a given lattice. Such distance is obtained by finding the n-th
closest pairs of sites in a lattice, irrespective of their sublattice.

    nrange(n::Int, lat::AbstractLattice)

Obtain the actual nth-nearest-neighbot distance between sites in lattice `lat`.

# See also
    `hopping`
"""
nrange(n::Int) = NeighborRange(n)

function nrange(n, lat::AbstractLattice{E,L}) where {E,L}
    sites = allsitepositions(lat)
    T = eltype(first(sites))
    dns = BoxIterator(zero(SVector{L,Int}))
    br = bravais(lat)
    # 640 is a heuristic cutoff for kdtree vs brute-force search
    if length(sites) <= 128
        dists = fill(T(Inf), n)
        for dn in dns
            iszero(dn) || ispositive(dn) || continue
            for (i, ri) in enumerate(sites), (j, rj) in enumerate(sites)
                j <= i && iszero(dn) && continue
                r = ri - rj + br * dn
                _update_dists!(dists, r'r)
            end
            isfinite(last(dists)) || acceptcell!(dns, dn)
        end
        dist = sqrt(last(dists))
    else
        tree = KDTree(sites)
        dist = T(Inf)
        for dn in dns
            iszero(dn) || ispositive(dn) || continue
            for r0 in sites
                r = r0 + br * dn
                dist = min(dist, _nrange(n, tree, r, nsites(lat)))
            end
            isfinite(dist) || acceptcell!(dns, dn)
        end
    end
    return dist
end

function _update_dists!(dists, dist::Real)
    len = length(dists)
    for (n, d) in enumerate(dists)
        isapprox(dist, d) && break
        if dist < d
            dists[n+1:len] .= dists[n:len-1]
            dists[n] = dist
            break
        end
    end
    return dists
end

function _nrange(n, tree, r::AbstractVector{T}, nmax) where {T}
    for m in n:nmax
        _, dists = knn(tree, r, 1 + m, true)
        popfirst!(dists)
        unique_sorted_approx!(dists)
        length(dists) == n && return maximum(dists)
    end
    return T(Inf)
end

#######################################################################
# Onsite/Hopping selectors
#######################################################################
abstract type Selector end

struct SiteSelector{S,I,M} <: Selector
    region::M
    sublats::S  # NTuple{N,NameType} (unresolved) or Vector{Int} (resolved on a lattice)
    indices::I  # Once resolved, this should be an Union{Integer,Not} container
end

struct HopSelector{S,I,D,T,M} <: Selector
    region::M
    sublats::S  # NTuple{N,Pair{NameType,NameType}} (unres) or Vector{Pair{Int,Int}} (res)
    dns::D
    range::T
    indices::I # Once resolved, this should be a Pair{Integer,Integer} container
end

struct ResolvedSelector{S<:Selector,L<:AbstractLattice} <: Selector
    selector::S
    lattice::L
end

struct Not{T} # Symbolizes an excluded elements
    i::T
end

"""
    not(i)

Wrapper indicating the negation or complement of `i`, typically used to encode excluded site
indices. See `siteselector` and `hopselector` for applications.

"""
not(i) = Not(i)
not(i...) = Not(i)

"""
    siteselector(; region = missing, sublats = missing, indices = missing)

Return a `SiteSelector` object that can be used to select sites in a lattice contained
within the specified region and sublattices. Only sites with index `i`, at position `r` and
belonging to a sublattice with name `s::NameType` will be selected if

    `region(r) && s in sublats && i in indices`

Any missing `region`, `sublat` or `indices` will not be used to constraint the selection.

The keyword `sublats` allows the following formats:

    sublats = :A                    # Sites on sublat :A only
    sublats = (:A,)                 # Same as above
    sublats = (:A, :B)              # Sites on sublat :A and :B

The keyword `indices` accepts a single integer, or a collection thereof. If several
collections are given, they are flattened into a single one. Possible combinations:

    indices = 1                     # Site 1 only
    indices = (1, )                 # Same as above
    indices = (1, 2, 3)             # Sites 1, 2 or 3
    indices = [1, 2, 3]             # Same as above
    indices = 1:3                   # Same as above
    indices = (1:3, 7, 8)           # Sites 1, 2, 3, 7 or 8

Additionally, indices or sublattices can be wrapped in `not` to exclude them (see `not`):

    sublats = not(:A)               # Any sublat different from :A
    sublats = not(:A, :B)           # Any sublat different from :A and :B
    indices = not(8)                # Any site index different from 8
    indices = not(1, 3:4)           # Any site index different from 1, 3 or 4
    indices = (not(3:4), 4:6)       # Any site different from 3 and 4, *or* equal to 4, 5 or 6
"""
siteselector(; region = missing, sublats = missing, indices = missing) =
    SiteSelector(region, sublats, indices)
siteselector(s::SiteSelector; region = s.region, sublats = s.sublats, indices = s.indices) =
    SiteSelector(region, sublats, indices)

"""
    hopselector(; range = missing, dn = missing, sublats = missing, indices = missing, region = missing)

Return a `HopSelector` object that can be used to select hops between two sites in a
lattice. Only hops between two sites, with indices `ipair = src => dst`, at positions `r₁ =
r - dr/2` and `r₂ = r + dr`, belonging to unit cells at integer distance `dn´` and to
sublattices `s₁` and `s₂` will be selected if:

    `region(r, dr) && s in sublats && dn´ in dn && norm(dr) <= range && ipair in indices`

If any of these is `missing` it will not be used to constraint the selection.

The keyword `range` admits the following possibilities

    max_range                   # i.e. `norm(dr) <= max_range`
    (min_range, max_range)      # i.e. `min_range <= norm(dr) <= max_range`

Both `max_range` and `min_range` can be a `Real` or a `NeighborRange` created with
`nrange(n)`. The latter represents the distance of `n`-th nearest neighbors.

The keyword `dn` can be a `Tuple`/`Vector`/`SVector` of `Int`s, or a tuple thereof.

The keyword `sublats` allows the following formats:

    sublats = :A => :B                  # Hopping from :A to :B sublattices, but not from :B to :A
    sublats = (:A => :B,)               # Same as above
    sublats = (:A => :B, :C => :D)      # Hopping from :A to :B or :C to :D
    sublats = (:A, :C) .=> (:B, :D)     # Broadcasted pairs, same as above
    sublats = (:A, :C) => (:B, :D)      # Direct product, (:A=>:B, :A=:D, :C=>:B, :C=>D)

The keyword `indices` accepts a single `src => dest` pair or a collection thereof. Any `src
== dest` will be neglected. Possible combinations:

    indices = 1 => 2                    # Hopping from site 1 to 2, but not from 2 to 1
    indices = (1 => 2, 2 => 1)          # Hoppings from 1 to 2 or from 2 to 1
    indices = [1 => 2, 2 => 1]          # Same as above
    indices = [(1, 2) .=> (2, 1)]       # Broadcasted pairs, same as above
    indices = [1:10 => 20:25, 3 => 30]  # Direct product, any hopping from sites 1:10 to sites 20:25, or from 3 to 30

Additionally, indices or sublattices can be wrapped in `not` to exclude them (see `not`):

    sublats = not(:A => :B, :B => :A)   # Any sublat pairs different from :A => :B or :B => :A
    sublats = not(:A) => :B             # Any sublat pair s1 => s2 with s1 different from :A and s2 equal to :B
    indices = not(8 => 9)               # Any site indices different from 8 => 9
    indices = 1 => not(3:4)             # Any site pair 1 => s with s different from 3, 4

"""
hopselector(; region = missing, sublats = missing, dn = missing, range = missing, indices = missing) =
    HopSelector(region, sublats, sanitize_dn(dn), sanitize_range(range), indices)
hopselector(s::HopSelector; region = s.region, sublats = s.sublats, dn = s.dns, range = s.range, indices = s.indices) =
    HopSelector(region, sublats, sanitize_dn(dn), sanitize_range(range), indices)

ensurenametype((s1, s2)::Pair) = nametype(s1) => nametype(s2)

sanitize_dn(dn::Missing) = missing
sanitize_dn(dn::Tuple{Vararg{Number,N}}) where {N} = (_sanitize_dn(dn),)
sanitize_dn(dn) = (_sanitize_dn(dn),)
sanitize_dn(dn::Tuple) = _sanitize_dn.(dn)
_sanitize_dn(dn::Tuple{Vararg{Number,N}}) where {N} = SVector{N,Int}(dn)
_sanitize_dn(dn::SVector{N}) where {N} = SVector{N,Int}(dn)
_sanitize_dn(dn::Vector) = SVector{length(dn),Int}(dn)

sanitize_range(::Missing) = missing
sanitize_range(r) = _shift_eps(r, 1)
sanitize_range(r::NTuple{2,Any}) = (_shift_eps(first(r), -1), _shift_eps(last(r), 1))

_shift_eps(r::Real, m) = ifelse(isfinite(r), float(r) + m * sqrt(eps(float(r))), float(r))
_shift_eps(r, m) = r

# API

function resolve(s::SiteSelector, lat::AbstractLattice)
    s = SiteSelector(s.region, resolve_sublats(s.sublats, lat), s.indices)
    return ResolvedSelector(s, lat)
end

function resolve(s::HopSelector, lat::AbstractLattice)
    s = HopSelector(s.region, resolve_sublat_pairs(s.sublats, lat), check_dn_dims(s.dns, lat), resolve_range(s.range, lat), s.indices)
    return ResolvedSelector(s, lat)
end

resolve_sublats(::Missing, lat) = missing # must be resolved to iterate over sublats
resolve_sublats(n::Not, lat) = Not(resolve_sublats(n.i, lat))
resolve_sublats(s, lat) = resolve_sublat_name.(s, Ref(lat))

resolve_range(r::Tuple, lat) = sanitize_range(_resolve_range.(r, Ref(lat)))
resolve_range(r, lat) = sanitize_range(_resolve_range(r, lat))
_resolve_range(r::NeighborRange, lat) = nrange(r.n, lat)
_resolve_range(r, lat) = r

function resolve_sublat_name(name::Union{NameType,Integer}, lat)
    i = findfirst(isequal(name), lat.unitcell.names)
    return i === nothing ? 0 : i
end

resolve_sublat_name(s, lat) =
    throw(ErrorException( "Unexpected format $s for `sublats`, see `onsite` for supported options"))

resolve_sublat_pairs(::Missing, lat) = missing
resolve_sublat_pairs(n::Not, lat) = Not(resolve_sublat_pairs(n.i, lat))
resolve_sublat_pairs(s::Tuple, lat) = resolve_sublat_pairs.(s, Ref(lat))
resolve_sublat_pairs(s::Vector, lat) = resolve_sublat_pairs.(s, Ref(lat))
resolve_sublat_pairs((src, dst)::Pair, lat) = _resolve_sublat_pairs(src, lat) => _resolve_sublat_pairs(dst, lat)
_resolve_sublat_pairs(n::Not, lat) = Not(_resolve_sublat_pairs(n.i, lat))
_resolve_sublat_pairs(p, lat) = resolve_sublat_name.(p, Ref(lat))

resolve_sublat_pairs(s, lat) =
    throw(ErrorException( "Unexpected format $s for `sublats`, see `hopping` for supported options"))

check_dn_dims(dns::Missing, lat::AbstractLattice{E,L}) where {E,L} = dns
check_dn_dims(dns::Tuple{Vararg{SVector{L,Int}}}, lat::AbstractLattice{E,L}) where {E,L} = dns
check_dn_dims(dns, lat::AbstractLattice{E,L}) where {E,L} =
    throw(DimensionMismatch("Specified cell distance `dn` does not match lattice dimension $L"))

# are sites at (i,j) and (dni, dnj) or (dn, 0) selected?
@inline function Base.in(i::Integer, rs::ResolvedSelector{<:SiteSelector, LA}) where {E,L,LA<:AbstractLattice{E,L}}
    dn0 = zero(SVector{L,Int})
    return ((i, i), (dn0, dn0)) in rs
end

@inline Base.in(((i, j), (dni, dnj))::Tuple{Tuple,Tuple}, rs::ResolvedSelector{<:SiteSelector}) =
    isonsite((i, j), (dni, dnj)) && (i, dni) in rs

Base.in((i, dni)::Tuple{Integer,SVector}, rs::ResolvedSelector{<:SiteSelector}) =
    isinindices(i, rs.selector.indices) &&
    isinregion(i, dni, rs.selector.region, rs.lattice) &&
    isinsublats(sublat_site(i, rs.lattice), rs.selector.sublats)

Base.in((j, i)::Pair{<:Integer,<:Integer}, rs::ResolvedSelector{<:HopSelector}) = (i, j) in rs

function Base.in(is::Tuple{Integer,Integer}, rs::ResolvedSelector{<:HopSelector, LA}) where {E,L,LA<:AbstractLattice{E,L}}
    dn0 = zero(SVector{L,Int})
    return (is, (dn0, dn0)) in rs
end

Base.in((inds, dns), rs::ResolvedSelector{<:HopSelector}) =
    !isonsite(inds, dns) && isinindices(indstopair(inds), rs.selector.indices) &&
    isinregion(inds, dns, rs.selector.region, rs.lattice) && isindns(dns, rs.selector.dns) &&
    isinrange(inds, dns, rs.selector.range, rs.lattice) &&
    isinsublats(indstopair(sublat_site.(inds, Ref(rs.lattice))), rs.selector.sublats)

isonsite((i, j), (dni, dnj)) = i == j && dni == dnj

isinregion(i::Int, ::Missing, lat) = true
isinregion(i::Int, dn, ::Missing, lat) = true

isinregion(i::Int, region::Union{Function,Region}, lat) =
    region(allsitepositions(lat)[i])
isinregion(i::Int, dn, region::Union{Function,Region}, lat) =
    region(allsitepositions(lat)[i] + bravais(lat) * dn)

isinregion(is::Tuple{Int,Int}, dns, ::Missing, lat) = true
function isinregion((row, col)::Tuple{Int,Int}, (dnrow, dncol), region::Union{Function,Region}, lat)
    br = bravais(lat)
    r, dr = _rdr(allsitepositions(lat)[col] + br * dncol, allsitepositions(lat)[row] + br * dnrow)
    return region(r, dr)
end

isindns((dnrow, dncol)::Tuple{SVector,SVector}, dns) = isindns(dnrow - dncol, dns)
isindns(dn::SVector{L,Int}, dns::Tuple{Vararg{SVector{L,Int}}}) where {L} = dn in dns
isindns(dn::SVector, dns::Missing) = true
isindns(dn, dns) =
    throw(ArgumentError("Cell distance dn in selector is incompatible with Lattice."))

isinrange(inds, dns, ::Missing, lat) = true
isinrange((row, col), (dnrow, dncol), range, lat) =
    _isinrange(allsitepositions(lat)[col] - allsitepositions(lat)[row] + bravais(lat) * (dncol - dnrow), range)
_isinrange(p, rmax::Real) = p'p <= rmax^2
_isinrange(p, (rmin, rmax)::Tuple{Real,Real}) =  rmin^2 <= p'p <= rmax^2

is_below_min_range(inds, dns, rsel::ResolvedSelector) =
    is_below_min_range(inds, dns, rsel.selector.range, rsel.lattice)
is_below_min_range((i, j), (dni, dnj), (rmin, rmax)::Tuple, lat) =
    norm(siteposition(i, dni, lat) - siteposition(j, dnj, lat)) < rmin
is_below_min_range(inds, dn, range, lat) = false

# Sublats are resolved, so they are equivalent to indices
isinsublats(i, j) = isinindices(i,j)

# Here we can have (1, 2:3)
isinindices(i::Integer, n::Not) = !isinindices(i, n.i)
isinindices(i::Integer, ::Missing) = true
isinindices(i::Integer, j::Integer) = i == j
isinindices(i::Integer, r::NTuple{N,Integer}) where {N} = i in r
isinindices(i::Integer, inds::Tuple) = any(is -> isinindices(i, is), inds)
isinindices(i::Integer, r) = i in r
# Here we cover ((1,2) .=> (3,4), 1=>2) and ((1,2) => (3,4), 1=>2)
isinindices(is::Pair, n::Not) = !isinindices(is, n.i)
isinindices(is::Pair, ::Missing) = true
# Here is => js could be (1,2) => (3,4) or 1:2 => 3:4, not simply 1 => 3
isinindices((j, i)::Pair, (js, is)::Pair) = isinindices(i, is) && isinindices(j, js)
# Here we support ((1,2) .=> (3,4), 3=>4) or ((1,2) .=> 3:4, 3=>4)
isinindices(pair::Pair, pairs) = any(p -> isinindices(pair, p), pairs)

# merge non-missing fields of s´ into s
merge_non_missing(s::SiteSelector, s´::SiteSelector) =
    SiteSelector(merge_non_missing.(
        (s.region,  s.sublats, s.indices),
        (s´.region, s´.sublats, s´.indices))...)
merge_non_missing(s::HopSelector, s´::HopSelector) =
    HopSelector(merge_non_missing.(
        (s.region,  s.sublats,  s.dns,  s.range, s.indices),
        (s´.region, s´.sublats, s´.dns, s´.range, s´.indices))...)
merge_non_missing(o, o´::Missing) = o
merge_non_missing(o, o´) = o´

Base.adjoint(s::SiteSelector) = s
function Base.adjoint(s::HopSelector)
    # is_unconstrained_selector(s) &&
    #     @warn("Taking the adjoint of an unconstrained hopping is likely unintended")
    region´ = _adjoint(s.region)
    sublats´ = _adjoint.(s.sublats)
    dns´ = _adjoint.(s.dns)
    range = s.range
    indices´ = _adjoint.(s.indices)
    return HopSelector(region´, sublats´, dns´, range, indices´)
end

_adjoint(::Missing) = missing
_adjoint(f::Function) = (r, dr) -> f(r, -dr)
_adjoint(t::Pair) = reverse(t)
_adjoint(t::Tuple) = _adjoint.(t)
_adjoint(t::AbstractVector) = _adjoint.(t)
_adjoint(t::SVector) = -t

# is_unconstrained_selector(s::HopSelector{Missing,Missing,Missing}) = true
# is_unconstrained_selector(s::HopSelector) = false

#######################################################################
# site, sublat, dn generators
#######################################################################

sitepositions(lat::AbstractLattice, s::SiteSelector) =
    sitepositions(resolve(s, lat))
sitepositions(rs::ResolvedSelector{<:SiteSelector}) =
    (s for (i, s) in enumerate(allsitepositions(rs.lattice)) if i in rs)

siteindices(lat::AbstractLattice, s::SiteSelector) =
    siteindices(resolve(s, lat))
siteindices(lat::AbstractLattice, ::Colon) = 1:nsites(lat)
siteindices(lat::AbstractLattice, is::AbstractRange) = intersect(1:nsites(lat), is)
siteindices(lat::AbstractLattice, i::Integer) = siteindices(lat, i:i)
siteindices(lat::AbstractLattice, is) = filter!(in(siteindices(lat, :)), sort(is))

siteindices(rs::ResolvedSelector{<:SiteSelector}) =
    (i for i in siteindex_candidates(rs) if i in rs)
siteindices(rs::ResolvedSelector{<:SiteSelector}, sublat::Int) =
    (i for i in siteindex_candidates(rs, sublat) if i in rs)
siteindices(rs::ResolvedSelector{<:SiteSelector}, sublat, dn) =
    (i for i in siteindex_candidates(rs, sublat) if (i, dn) in rs)
siteindices(rs::ResolvedSelector{<:SiteSelector}, dn::SVector) =
    (i for i in siteindex_candidates(rs) if (i, dn) in rs)

# Given a sublattice, which site indices should be checked by selector?
siteindex_candidates(rs) = eachindex(allsitepositions(rs.lattice))
siteindex_candidates(rs, sublat) =
    _siteindex_candidates(rs.selector.indices, siterange(rs.lattice, sublat))
# indices can be missing, 1, 2:3, (1,2,3) or (1, 2:3)
# we also support (1, (2,3)) and [1, 2, 3], useful for source_candidates below
_siteindex_candidates(::Missing, sr) = sr
# Better not exclude candidates with not, since that can lead to collecting a huge range
_siteindex_candidates(::Not, sr) = sr
_siteindex_candidates(i::Integer, sr) = ifelse(i in sr, (i,), ())
_siteindex_candidates(inds::AbstractUnitRange, sr) = intersect(inds, sr)
_siteindex_candidates(inds::NTuple{N,Integer}, sr) where {N} = filter(in(sr), inds)
_siteindex_candidates(inds::AbstractVector{<:Integer}, sr) = filter(in(sr), inds)
_siteindex_candidates(inds, sr) = Iterators.flatten(_siteindex_candidates.(inds, Ref(sr)))

source_candidates(rs::ResolvedSelector{<:HopSelector}, sublat) =
    _source_candidates(rs.selector.indices, siterange(rs.lattice, sublat))
_source_candidates(::Missing, sr) = sr
_source_candidates(::Not, sr) = sr
_source_candidates(inds, sr) = _siteindex_candidates(_recursivefirst(inds), sr)

_recursivefirst(p::Pair) = first(p)
_recursivefirst(p) = _recursivefirst.(p)

sublats(rs::ResolvedSelector{<:SiteSelector{Missing}}) = sublats(rs.lattice)

function sublats(rs::ResolvedSelector{<:SiteSelector})
    subs = sublats(rs.lattice)
    return (s for s in subs if isinsublats(s, rs.selector.sublats))
end

function sublats(rs::ResolvedSelector{<:HopSelector})
    subs = sublats(rs.lattice)
    return (s => d for s in subs for d in subs if isinsublats(s => d, rs.selector.sublats))
end

dniter(rs::ResolvedSelector{S,LA}) where {S,E,L,LA<:Lattice{E,L}} =
    dniter(rs.selector.dns, Val(L))
dniter(dns::Missing, ::Val{L}) where {L} = BoxIterator(zero(SVector{L,Int}))
dniter(dns, ::Val) = dns

#######################################################################
# Tightbinding types
#######################################################################
abstract type TightbindingModelTerm end
abstract type AbstractOnsiteTerm <: TightbindingModelTerm end
abstract type AbstractHoppingTerm <: TightbindingModelTerm end

struct TightbindingModel{N,T<:Tuple{Vararg{TightbindingModelTerm,N}}}
    terms::T
end

struct OnsiteTerm{F,S<:Union{SiteSelector, ResolvedSelector{<:SiteSelector}},C} <: AbstractOnsiteTerm
    o::F
    selector::S
    coefficient::C
end

struct HoppingTerm{F,S<:Union{HopSelector, ResolvedSelector{<:HopSelector}},C} <: AbstractHoppingTerm
    t::F
    selector::S
    coefficient::C
end

#######################################################################
# TightbindingModel API
#######################################################################
terms(t::TightbindingModel) = t.terms

TightbindingModel(ts::TightbindingModelTerm...) = TightbindingModel(ts)

# (m::TightbindingModel)(r, dr) = sum(t -> t(r, dr), m.terms)  # this does not filter by selector, so it's wrong

# External API #

Base.:*(x::Number, m::TightbindingModel) = TightbindingModel(x .* m.terms)
Base.:*(m::TightbindingModel, x::Number) = x * m
Base.:-(m::TightbindingModel) = TightbindingModel((-1) .* m.terms)

Base.:+(m::TightbindingModel, t::TightbindingModel) = TightbindingModel((m.terms..., t.terms...))
Base.:-(m::TightbindingModel, t::TightbindingModel) = m + (-t)

function Base.show(io::IO, m::TightbindingModel{N}) where {N}
    ioindent = IOContext(io, :indent => "  ")
    print(io, "TightbindingModel{$N}: model with $N terms", "\n")
    foreach(t -> print(ioindent, t, "\n"), m.terms)
end

#######################################################################
# TightbindingModelTerm API
#######################################################################
OnsiteTerm(t::OnsiteTerm, os::SiteSelector) =
    OnsiteTerm(t.o, os, t.coefficient)

(o::OnsiteTerm{<:Function})(r,dr) = o.coefficient * o.o(r)
(o::OnsiteTerm)(r,dr) = o.coefficient * o.o

HoppingTerm(t::HoppingTerm, os::HopSelector) =
    HoppingTerm(t.t, os, t.coefficient)

(h::HoppingTerm{<:Function})(r, dr) = h.coefficient * h.t(r, dr)
(h::HoppingTerm)(r, dr) = h.coefficient * h.t

sublats(t::TightbindingModelTerm, lat) = resolve_sublats(t.selector, lat)

sublats(m::TightbindingModel) = (t -> t.selector.sublats).(terms(m))

resolve(t::OnsiteTerm, lat::AbstractLattice) = OnsiteTerm(t.o, resolve(t.selector, lat), t.coefficient)
resolve(t::HoppingTerm, lat::AbstractLattice) = HoppingTerm(t.t, resolve(t.selector, lat), t.coefficient)
resolve(m::TightbindingModel, lat::AbstractLattice) = TightbindingModel(resolve.(m.terms, Ref(lat)))

displayparameter(::Type{<:Function}) = "Function"
displayparameter(::Type{T}) where {T} = "$T"

displayrange(r::Real) = round(r, digits = 6)
displayrange(::Missing) = "any"
displayrange(nr::NeighborRange) = "NeighborRange($(nr.n))"
displayrange(rs::Tuple) = "($(displayrange(first(rs))), $(displayrange(last(rs))))"

function Base.show(io::IO, o::OnsiteTerm{F,<:SiteSelector}) where {F}
    i = get(io, :indent, "")
    print(io,
"$(i)OnsiteTerm{$(displayparameter(F))}:
$(i)  Sublattices      : $(o.selector.sublats === missing ? "any" : o.selector.sublats)
$(i)  Coefficient      : $(o.coefficient)")
end

function Base.show(io::IO, h::HoppingTerm{F,<:HopSelector}) where {F}
    i = get(io, :indent, "")
    print(io,
"$(i)HoppingTerm{$(displayparameter(F))}:
$(i)  Sublattice pairs : $(h.selector.sublats === missing ? "any" : h.selector.sublats)
$(i)  dn cell distance : $(h.selector.dns === missing ? "any" : h.selector.dns)
$(i)  Hopping range    : $(displayrange(h.selector.range))
$(i)  Coefficient      : $(h.coefficient)")
end

# External API #
"""
    onsite(o; region = missing, sublats = missing)

Create an `TightbindingModel` with a single `OnsiteTerm` that applies an onsite energy `o`
to a `Lattice` when creating a `Hamiltonian` with `hamiltonian`.

The onsite energy `o` can be a number, an `SMatrix`, a `UniformScaling` (e.g. `3*I`) or a
function of the form `r -> ...` for a position-dependent onsite energy.

The dimension of `o::AbstractMatrix` must match the orbital dimension of applicable
sublattices (see also `orbitals` option for `hamiltonian`). If `o::UniformScaling` it will
be converted to an identity matrix of the appropriate size when applied to multiorbital
sublattices. Similarly, if `o::SMatrix` it will be truncated or padded to the appropriate
size.

    onsite(model::TightbindingModel; kw...)

Return a `TightbindingModel` with only the onsite terms of `model`. Any non-missing `kw` is
applied to all such terms.

# Keyword arguments

Keywords are the same as for `siteselector`. Only sites at position `r` in sublattice with
name `s::NameType` will be selected if `region(r) && s in sublats` is true. Any missing
`region` or `sublat` will not be used to constraint the selection.

The keyword `sublats` allows the following formats:

    sublats = :A           # Onsite on sublat :A only
    sublats = (:A,)        # Same as above
    sublats = (:A, :B)     # Onsite on sublat :A and :B

# Combining models

`OnsiteTerm`s and `HoppingTerm`s created with `onsite` or `hopping` can added or substracted
together or be multiplied by scalars to build more complicated `TightbindingModel`s, e.g.
`onsite(1) - 3 * hopping(2)`

# Examples

```jldoctest
julia> model = onsite(1, sublats = (:A,:B)) - 2 * hopping(2, sublats = :A=>:A)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : (:A, :B)
    Coefficient      : 1
  HoppingTerm{Int64}:
    Sublattice pairs : :A => :A
    dn cell distance : any
    Hopping range    : NeighborRange(1)
    Coefficient      : -2

julia> newmodel = onsite(model; sublats = :A) + hopping(model)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : A
    Coefficient      : 1
  HoppingTerm{Int64}:
    Sublattice pairs : :A => :A
    dn cell distance : any
    Hopping range    : NeighborRange(1)
    Coefficient      : -2

julia> LatticePresets.honeycomb() |> hamiltonian(onsite(r -> @SMatrix[1 2; 3 4]), orbitals = Val(2))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 1 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a, :a), (:a, :a))
  Element type     : 2 × 2 blocks (ComplexF64)
  Onsites          : 2
  Hoppings         : 0
  Coordination     : 0.0
```

# See also
    `hopping`
"""
onsite(o; kw...) = onsite(o, siteselector(; kw...))

onsite(o, selector::Selector) = TightbindingModel(OnsiteTerm(o, selector, 1))

onsite(m::TightbindingModel, selector::Selector) =
    TightbindingModel(_onlyonsites(selector, m.terms...))

_onlyonsites(s, t::OnsiteTerm, args...) =
    (OnsiteTerm(t, merge_non_missing(t.selector, s)), _onlyonsites(s, args...)...)
_onlyonsites(s, t::HoppingTerm, args...) = (_onlyonsites(s, args...)...,)
_onlyonsites(s) = ()

"""
    hopping(t; range = nrange(1), dn = missing, sublats = missing, indices = missing, region = missing, plusadjoint = false)

Create an `TightbindingModel` with a single `HoppingTerm` that applies a hopping `t` to a
`Lattice` when creating a `Hamiltonian` with `hamiltonian`.

The hopping amplitude `t` can be a number, an `SMatrix`, a `UniformScaling` (e.g. `3*I`) or
a function of the form `(r, dr) -> ...` for a position-dependent hopping (`r` is the bond
center, and `dr` the bond vector). If `sublats` is specified as a sublattice name pair, or
tuple thereof, `hopping` is only applied between sublattices with said names.

The dimension of `t::AbstractMatrix` must match the orbital dimension of applicable
sublattices (see also `orbitals` option for `hamiltonian`). If `t::UniformScaling` it will
be converted to a (possibly rectangular) identity matrix of the appropriate size when
applied to multiorbital sublattices. Similarly, if `t::SMatrix` it will be truncated or
padded to the appropriate size.

`OnsiteTerm`s and `HoppingTerm`s created with `onsite` or `hopping` can be added or
substracted together to build more complicated `TightbindingModel`s.

    hopping(model::TightbindingModel; kw...)

Return a `TightbindingModel` with only the hopping terms of `model`. Any non-missing `kw` is
applied to all such terms.

# Keyword arguments

Most keywords are the same as for `hopselector`. Only hoppings between two sites at
positions `r₁ = r - dr/2` and `r₂ = r + dr`, belonging to unit cells at integer distance
`dn´` and to sublattices `s₁` and `s₂` will be selected if: `region(r, dr) && s in sublats
&& dn´ in dn && norm(dr) <= range`. If any of these is `missing` it will not be used to
constraint the selection.

The keyword `range` admits the following possibilities

    max_range                   # i.e. `norm(dr) <= max_range`
    (min_range, max_range)      # i.e. `min_range <= norm(dr) <= max_range`

Both `max_range` and `min_range` can be a `Real` or a `NeighborRange` created with
`nrange(n)`. The latter represents the distance of `n`-th nearest neighbors. Note that the
`range` default for `hopping` (unlike for the more general `hopselector`) is `nrange(1)`,
i.e. first-nearest-neighbors.

The keyword `dn` can be a `Tuple`/`Vector`/`SVector` of `Int`s, or a tuple thereof. The
keyword `sublats` allows the following formats:

    sublats = :A => :B                 # Hopping from :A to :B sublattices
    sublats = (:A => :B,)              # Same as above
    sublats = (:A => :B, :C => :D)     # Hopping from :A to :B or :C to :D
    sublats = (:A, :C) .=> (:B, :D)    # Broadcasted pairs, same as above
    sublats = (:A, :C) => (:B, :D)     # Direct product, (:A=>:B, :A=:D, :C=>:B, :C=>D)

The keyword `plusadjoint` produces a model with the input hopping term plus its adjoint.
Note that this substitution is made before multiplying by any coefficient, so that
`im*hopping(..., plusadjoint = true) == im*(hopping(...) + hopping(...)')`.

# Combining models

`OnsiteTerm`s and `HoppingTerm`s created with `onsite` or `hopping` can added or substracted
together or be multiplied by scalars to build more complicated `TightbindingModel`s, e.g.
`onsite(1) - 3 * hopping(2)`

# Examples

```jldoctest
julia> model = 3 * onsite(1) - hopping(2, dn = ((1,2), (0,0)), sublats = :A=>:B)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : any
    Coefficient      : 3
  HoppingTerm{Int64}:
    Sublattice pairs : :A => :B
    dn cell distance : ([1, 2], [0, 0])
    Hopping range    : NeighborRange(1)
    Coefficient      : -1

julia> newmodel = onsite(model) + hopping(model, range = 2)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : any
    Coefficient      : 3
  HoppingTerm{Int64}:
    Sublattice pairs : :A => :B
    dn cell distance : ([1, 2], [0, 0])
    Hopping range    : 2.0
    Coefficient      : -1

julia> LatticePresets.honeycomb() |> hamiltonian(hopping((r,dr) -> cos(r[1]), sublats = (:A,:B) => (:A,:B)))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 5 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a,), (:a,))
  Element type     : scalar (ComplexF64)
  Onsites          : 0
  Hoppings         : 6
  Coordination     : 3.0
```

# See also
    `onsite`, `nrange`
"""
function hopping(t; plusadjoint = false, range = nrange(1), kw...)
    hop = hopping(t, hopselector(; range = range, kw...))
    return plusadjoint ? hop + hop' : hop
end
hopping(t, selector) = TightbindingModel(HoppingTerm(t, selector, 1))

hopping(m::TightbindingModel, selector::Selector) =
    TightbindingModel(_onlyhoppings(selector, m.terms...))

_onlyhoppings(s, t::OnsiteTerm, args...) = (_onlyhoppings(s, args...)...,)
_onlyhoppings(s, t::HoppingTerm, args...) =
    (HoppingTerm(t, merge_non_missing(t.selector, s)), _onlyhoppings(s, args...)...)
_onlyhoppings(s) = ()

Base.:*(x::Number, o::OnsiteTerm) =
    OnsiteTerm(o.o, o.selector, x * o.coefficient)
Base.:*(x::Number, t::HoppingTerm) = HoppingTerm(t.t, t.selector, x * t.coefficient)
Base.:*(t::TightbindingModelTerm, x::Number) = x * t
Base.:-(t::TightbindingModelTerm) = (-1) * t

Base.adjoint(t::TightbindingModel) = TightbindingModel(adjoint.(terms(t)))
Base.adjoint(t::OnsiteTerm{<:Function}) = OnsiteTerm(r -> t.o(r)', t.selector, t.coefficient')
Base.adjoint(t::OnsiteTerm) = OnsiteTerm(t.o', t.selector, t.coefficient')
Base.adjoint(t::HoppingTerm{<:Function}) = HoppingTerm((r, dr) -> t.t(r, -dr)', t.selector', t.coefficient')
Base.adjoint(t::HoppingTerm) = HoppingTerm(t.t', t.selector', t.coefficient')

#######################################################################
# offdiagonal
#######################################################################
"""
    offdiagonal(model, lat, nsublats::NTuple{N,Int})

Build a restricted version of `model` that applies only to off-diagonal blocks formed by
sublattice groups of size `nsublats`.
"""
offdiagonal(m::TightbindingModel, lat, nsublats) =
    TightbindingModel(offdiagonal.(m.terms, Ref(lat), Ref(nsublats)))

offdiagonal(o::OnsiteTerm, lat, nsublats) =
    throw(ArgumentError("No onsite terms allowed in off-diagonal coupling"))

function offdiagonal(t::HoppingTerm, lat, nsublats)
    rs = resolve(t.selector, lat)
    s = collect(sublats(rs))
    sr = sublatranges(nsublats...)
    filter!(spair -> findblock(first(spair), sr) != findblock(last(spair), sr), s)
    rs´ = ResolvedSelector(hopselector(rs.selector, sublats = s), lat)
    return HoppingTerm(t.t, rs´, t.coefficient)
end

sublatranges(i::Int, is::Int...) = _sublatranges((1:i,), is...)
_sublatranges(rs::Tuple, i::Int, is...) = _sublatranges((rs..., last(last(rs)) + 1: last(last(rs)) + i), is...)
_sublatranges(rs::Tuple) = rs

findblock(s, sr) = findfirst(r -> s in r, sr)

#######################################################################
# @onsite! and @hopping!
#######################################################################
abstract type AbstractModifier end
abstract type ElementModifier{N,S,F} <: AbstractModifier end

struct ParametricFunction{N,F,P<:Val}
    f::F
    params::P
end

ParametricFunction{N}(f::F, p::P) where {N,F,P} = ParametricFunction{N,F,P}(f, p)

(pf::ParametricFunction)(args...; kw...) = pf.f(args...; kw...)

struct OnsiteModifier{N,S<:Selector,F<:ParametricFunction{N}} <: ElementModifier{N,S,F}
    f::F
    selector::S
end

struct HoppingModifier{N,S<:Selector,F<:ParametricFunction{N}} <: ElementModifier{N,S,F}
    f::F
    selector::S
end

const UniformModifier = ElementModifier{1}
const UniformHoppingModifier = HoppingModifier{1}
const UniformOnsiteModifier = OnsiteModifier{1}

"""
    parameters(p::AbstractModifier...)

Return the parameter names for one or several `AbstractModifier` created with `@onsite!`,
`@hopping!` or `@block!`.
"""
parameters(ms::AbstractModifier...) = mergetuples(_parameters.(ms)...)
_parameters(m::AbstractModifier) = _parameters(m.f)
_parameters(pf::ParametricFunction) = pf.params

"""
    @onsite!(args -> body; kw...)

Create an `ElementModifier <: AbstractModifier`, to be used with `parametric`, that applies
`f = args -> body` to onsite energies specified by `kw` (see `onsite` for details on
possible `kw`s). The form of `args -> body` may be `(o; params...) -> ...` or `(o, r;
params...) -> ...` if the modification is position (`r`) dependent. Keyword arguments
`params` are optional, and include any parameters that `body` depends on that the user may
want to tune.

Note: unlike `onsite` and `hopping`, `ElementModifier`s cannot be combined (i.e. you cannot
do `@onsite!(...) + @hopping!(...)`). `ElementModifier`s are not model terms but
transformations of an existing Hamiltonian that are meant to be applied sequentially (the
order of application usually matters).

# See also
    `@hopping!`, `@block!`, `parametric`
"""
macro onsite!(kw, f)
    f, N, params = get_f_N_params(f, "Only @onsite!(args -> body; kw...) syntax supported. Mind the `;`.")
    return esc(:(Quantica.OnsiteModifier(Quantica.ParametricFunction{$N}($f, $(Val(params))), Quantica.siteselector($kw))))
end

macro onsite!(f)
    f, N, params = get_f_N_params(f, "Only @onsite!(args -> body; kw...) syntax supported.  Mind the `;`.")
    return esc(:(Quantica.OnsiteModifier(Quantica.ParametricFunction{$N}($f, $(Val(params))), Quantica.siteselector())))
end

"""
    @hopping!(args -> body; kw...)

Create an `ElementModifier <: AbstractModifier`, to be used with `parametric`, that applies `f = args -> body`
to hoppings energies specified by `kw` (see `hopping` for details on possible `kw`s). The
form of `args -> body` may be `(t; params...) -> ...` or `(t, r, dr; params...) -> ...` if
the modification is position (`r`, `dr`) dependent. Keyword arguments `params` are optional,
and include any parameters that `body` depends on that the user may want to tune.

Note: unlike `onsite` and `hopping`, `ElementModifier`s cannot be combined (i.e. you cannot
do `@onsite!(...) + @hopping!(...)`). `ElementModifier`s are not model terms but
transformations of an existing Hamiltonian that are meant to be applied sequentially (the
order of application usually matters).

# See also
    `@onsite!`, `@block!`, `parametric`
"""
macro hopping!(kw, f)
    f, N, params = get_f_N_params(f, "Only @hopping!(args -> body; kw...) syntax supported. Mind the `;`.")
    return esc(:(Quantica.HoppingModifier(Quantica.ParametricFunction{$N}($f, $(Val(params))), Quantica.hopselector($kw))))
end

macro hopping!(f)
    f, N, params = get_f_N_params(f, "Only @hopping!(args -> body; kw...) syntax supported. Mind the `;`.")
    return esc(:(Quantica.HoppingModifier(Quantica.ParametricFunction{$N}($f, $(Val(params))), Quantica.hopselector())))
end

# Extracts normalized f, number of arguments and kwarg names from an anonymous function f
function get_f_N_params(f, msg)
    (f isa Expr && f.head == :->) || throw(ArgumentError(msg))
    d = ExprTools.splitdef(f)
    kwargs = convert(Vector{Any}, get!(d, :kwargs, []))
    d[:kwargs] = kwargs  # in case it wasn't Vector{Any} originally
    if isempty(kwargs)
        params = ()
        push!(kwargs, :(_...))  # normalization : append _... to kwargs
    else
        params = get_kwname.(kwargs) |> Tuple
        if !isempty(params) && last(params) == :...
            params = Base.front(params)  # drop _... kwarg from params
        else
            push!(kwargs, :(_...))  # normalization : append _... to kwargs
        end
    end
    N = haskey(d, :args) ? length(d[:args]) : 0
    f´ = ExprTools.combinedef(d)
    return f´, N, params
end

get_kwname(x::Symbol) = x
get_kwname(x::Expr) = x.head === :kw ? x.args[1] : x.head  # x.head == :...

resolve(s::ResolvedSelector, lat) = s

resolve(o::OnsiteModifier, lat) = OnsiteModifier(o.f, resolve(o.selector, lat))

resolve(h::HoppingModifier, lat) = HoppingModifier(h.f, resolve(h.selector, lat))

resolve(t::Tuple, lat) = _resolve(lat, t...)
_resolve(lat, t, ts...) = (resolve(t, lat), _resolve(lat, ts...)...)
_resolve(lat) = ()

# Intended for resolved ElementModifier{N} only. The N is the number of arguments accepted.
@inline (o!::UniformOnsiteModifier)(o, r; kw...) = o!(o; kw...)
@inline (o!::UniformOnsiteModifier)(o, r, dr; kw...) = o!(o; kw...)
@inline (o!::UniformOnsiteModifier)(o; kw...) = o!.f(o; kw...)
@inline (o!::OnsiteModifier)(o, r, dr; kw...) = o!(o, r; kw...)
@inline (o!::OnsiteModifier)(o, r; kw...) = o!.f(o, r; kw...)

@inline (h!::UniformHoppingModifier)(t, r, dr; kw...) = h!(t; kw...)
@inline (h!::UniformHoppingModifier)(t; kw...) = h!.f(t; kw...)
@inline (h!::HoppingModifier)(t, r, dr; kw...) = h!.f(t, r, dr; kw...)

#######################################################################
# @block!
#######################################################################
struct BlockModifier{V,C<:Union{Missing,NTuple{<:Any,SVector}},F<:ParametricFunction{1}} <: AbstractModifier
    f::F
    dns::C
    rows::V
    cols::V
    function BlockModifier{V,C,F}(f, dns, rows, cols) where {V,C<:Union{Missing,NTuple{<:Any,SVector}},F<:ParametricFunction{1}}
        new(f, dns, rows, cols)
    end
end

BlockModifier(f, dns, sites) = BlockModifier(f, dns, sites, sites)

function BlockModifier(f::F, dns, rows, cols) where F<:ParametricFunction{1}
    dns´ = sanitize_dn(dns)
    rows´ = collect(rows)  # rows, cols from siteindices will be generators
    cols´ = collect(cols)  # we need collections of known length to use them in `view`
    V = typeof(rows´)
    C = typeof(dns´)
    return BlockModifier{V,C,F}(f, dns´, rows´, cols´)
end

"""
    @block!((block; params...) -> modified_block, sites; dn = missing)
    @block!((block; params...) -> modified_block, (rows, cols); dn = missing)

Create an `BlockModifier <: AbstractModifier`, to be used with `parametric`, that applies `f
= (block; ...) -> ...` to a block `h[dn][sites, sites]` or `h[dn][rows, cols]` of
hamiltonian `h`. Keyword arguments `params` are optional, and include any parameters that
`modified_block` depends on that the user may want to tune. If the keyword `dn = missing`,
the `dn` in `h[dn]` will be restricted to `dn = (0...)`. Otherwise the specified `dn`'s will
be modified (e.g. when `dn = (1,0)` or `dn = ((1,0), (-1,0))`).

Upon construction of a `ParametricHamiltonian` with a `@block!` modifier, a check is
performed that the whole block specified by `(rows, cols)` is stored in the sparse
Hamiltonian harmonics. If it is not, any non-zero element in `modified_block` will fail to
be applied to the harmonic in question, so a warning is issued. The warning can be ignored
if the user knows that all non-zero elements in `modified_block` are indeed stored in the
harmonic, either as finite matrix elements or structural zeros.

Special care should be taken when using `@block!` on Hamiltonians with different number of
orbitals in different sublattices. To avoid type-instabities in this case, the internal
representation of Hamiltonian harmonics uses a uniform `eltype` that is an `SMatrix{N,N}`
with `N` the maximum number of orbitals among the different sublattices (padded with zeros
in sublattices with less than `N` orbitals). The matrix `modified_block` should have this
same uniform `eltype`.

# See also
    `@onsite!`, `@hopping!`, `parametric`
"""
macro block!(kw, f, rows, cols...)
    f, N, params = get_f_N_params(f, "Only @block!(args -> body, inds...; dn = ...) syntax supported. Mind the `;`.")
    N == 1 || throw(ArgumentError("The function passed to `@block!` should be single-argument, with optional keywords."))
    return esc(:(Quantica.BlockModifier(Quantica.ParametricFunction{1}($f, $(Val(params))), Quantica.sanitize_dn(($kw,)[:dn]), $rows, $(cols...))))
end

macro block!(f, rows, cols...)
    f, N, params = get_f_N_params(f, "Only @block!(args -> body, inds...; dn = ...) syntax supported. Mind the `;`.")
    N == 1 || throw(ArgumentError("The function passed to `@block!` should be single-argument, with optional keywords."))
    return esc(:(Quantica.BlockModifier(Quantica.ParametricFunction{1}($f, $(Val(params))), missing, $rows, $(cols...))))
end

@inline (b!::BlockModifier)(h; kw...) = b!.f(h; kw...)

resolve(e::BlockModifier, lat) = e

Base.in(dn, t::BlockModifier) = t.dns === missing ? iszero(dn) : dn in t.dns

max_stars_repo_name: pablosanjose/Quantica

max_issues_repo_name: pablosanjose/Quantica

max_forks_repo_name: pablosanjose/Quantica


text:
using .Models

using ..ITensorAutoHOOT
using ..ITensorAutoHOOT: SubNetwork, neighboring_tensors

function insert_projectors(peps::PEPS, cutoff=1e-15, maxdim=100)
  psi_bra = addtags(linkinds, dag.(peps.data), "bra")
  psi_ket = addtags(linkinds, peps.data, "ket")
  tn = psi_bra .* psi_ket
  bmps = boundary_mps(tn; cutoff=cutoff, maxdim=maxdim)
  psi_bra_rot = addtags(linkinds, dag.(peps.data), "brarot")
  psi_ket_rot = addtags(linkinds, peps.data, "ketrot")
  tn_rot = psi_bra_rot .* psi_ket_rot
  bmps = boundary_mps(tn; cutoff=cutoff, maxdim=maxdim)
  bmps_rot = boundary_mps(tn_rot; cutoff=cutoff, maxdim=maxdim)
  # get the projector for each center
  Ny, Nx = size(peps.data)
  bonds_row = [(i, :) for i in 1:Ny]
  bonds_column = [(:, i) for i in 1:Nx]
  tn_split_row, tn_split_column = [], []
  projectors_row, projectors_column = Vector{Vector{ITensor}}(), Vector{Vector{ITensor}}()
  for bond in bonds_row
    tn_split, pl, pr = insert_projectors(tn, bmps; center=bond)
    push!(tn_split_row, tn_split)
    push!(projectors_row, vcat(reduce(vcat, pl), reduce(vcat, pr)))
  end
  for bond in bonds_column
    tn_split, pl, pr = insert_projectors(tn_rot, bmps_rot; center=bond)
    push!(tn_split_column, tn_split)
    push!(projectors_column, vcat(reduce(vcat, pl), reduce(vcat, pr)))
  end
  return tn_split_row, tn_split_column, projectors_row, projectors_column
end

function insert_projectors(peps::PEPS, center::Tuple, cutoff=1e-15, maxdim=100)
  # Square the tensor network
  psi_bra = addtags(linkinds, dag.(peps.data), "bra")
  psi_ket = addtags(linkinds, peps.data, "ket")
  tn = psi_bra .* psi_ket
  bmps = boundary_mps(tn; cutoff=cutoff, maxdim=maxdim)
  tn_split, pl, pr = insert_projectors(tn, bmps; center=center)
  return tn_split, vcat(reduce(vcat, pl), reduce(vcat, pr))
end

function inner_network(peps::PEPS, peps_prime::PEPS, projectors::Vector{ITensor})
  network = inner_network(peps::PEPS, peps_prime::PEPS)
  return vcat(network, projectors)
end

function inner_network(
  peps::PEPS, peps_prime::PEPS, projectors::Vector{<:ITensor}, ::typeof(tree_w_projectors)
)
  Ny, Nx = size(peps.data)
  function get_tree(i)
    line_tensors = vcat(peps.data[i, :], peps_prime.data[i, :])
    neighbor_projectors = neighboring_tensors(SubNetwork(line_tensors), projectors)
    return tree_w_projectors(peps.data[i, :], peps_prime.data[i, :], neighbor_projectors)
  end
  subnetworks = [get_tree(i) for i in 1:Ny]
  return SubNetwork(subnetworks)
end

function inner_network(
  peps::PEPS,
  peps_prime::PEPS,
  peps_prime_ham::PEPS,
  projectors::Vector{<:ITensor},
  mpo::MPO,
  coordinate::Tuple{<:Integer,Colon},
  ::typeof(tree_w_projectors),
)
  Ny, Nx = size(peps.data)
  function get_tree(i)
    if i == coordinate[1]
      line_tensors = vcat(peps.data[i, :], peps_prime_ham.data[i, :], mpo.data)
      neighbor_projectors = neighboring_tensors(SubNetwork(line_tensors), projectors)
      return tree_w_projectors(
        peps.data[i, :], peps_prime_ham.data[i, :], mpo.data, neighbor_projectors
      )
    else
      line_tensors = vcat(peps.data[i, :], peps_prime.data[i, :])
      neighbor_projectors = neighboring_tensors(SubNetwork(line_tensors), projectors)
      return tree_w_projectors(peps.data[i, :], peps_prime.data[i, :], neighbor_projectors)
    end
  end
  subnetworks = [get_tree(i) for i in 1:Ny]
  return SubNetwork(subnetworks)
end

function inner_network(
  peps::PEPS,
  peps_prime::PEPS,
  peps_prime_ham::PEPS,
  projectors::Vector{<:ITensor},
  mpo::MPO,
  coordinate::Tuple{Colon,<:Integer},
  ::typeof(tree_w_projectors),
)
  Ny, Nx = size(peps.data)
  function get_tree(i)
    if i == coordinate[2]
      line_tensors = vcat(peps.data[:, i], peps_prime_ham.data[:, i], mpo.data)
      neighbor_projectors = neighboring_tensors(SubNetwork(line_tensors), projectors)
      return tree_w_projectors(
        peps.data[:, i], peps_prime_ham.data[:, i], mpo.data, neighbor_projectors
      )
    else
      line_tensors = vcat(peps.data[:, i], peps_prime.data[:, i])
      neighbor_projectors = neighboring_tensors(SubNetwork(line_tensors), projectors)
      return tree_w_projectors(peps.data[:, i], peps_prime.data[:, i], neighbor_projectors)
    end
  end
  subnetworks = [get_tree(i) for i in 1:Nx]
  return SubNetwork(subnetworks)
end

function inner_networks(
  peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, projectors::Vector{ITensor}, Hs::Array
)
  network_list = inner_networks(peps, peps_prime, peps_prime_ham, Hs)
  return map(network -> vcat(network, projectors), network_list)
end

function inner_networks(
  peps::PEPS,
  peps_prime::PEPS,
  peps_prime_ham::PEPS,
  projectors::Vector{Vector{ITensor}},
  Hs::Vector{Models.LineMPO},
)
  @assert length(projectors) == length(Hs)
  function generate_each_network(projector, H)
    return inner_networks(peps, peps_prime, peps_prime_ham, projector, [H])[1]
  end
  return [generate_each_network(projector, H) for (projector, H) in zip(projectors, Hs)]
end

function inner_networks(
  peps::PEPS,
  peps_prime::PEPS,
  peps_prime_ham::PEPS,
  projectors::Vector{Vector{ITensor}},
  Hs::Vector{Models.LineMPO},
  ::typeof(tree_w_projectors),
)
  @assert length(projectors) == length(Hs)
  function generate_each_network(projector, H)
    return inner_network(
      peps, peps_prime, peps_prime_ham, projector, H.mpo, H.coord, tree_w_projectors
    )
  end
  return [generate_each_network(projector, H) for (projector, H) in zip(projectors, Hs)]
end

ext: jl


text:
Base.:+(x::ExactReal, y::ExactReal) = ExactReal(n -> (x(n+2) + y(n+2) + 2) >> 2)

Base.:-(x::ExactReal) = ExactReal(n->-x(n))
Base.:-(x::ExactReal, y::ExactReal) = x + (-y)
Base.:-(x::Union{Rational, Integer}, y::ExactReal) = ExactReal(x) - y

"Most significant digit"
function msd(x::ExactReal)
    x.msd > -1 && return x.msd

    n = 0

    while x(n) == 0
        n += 1
    end

    x.msd = n

    return n
end

function Base.:*(x::ExactReal, y::ExactReal)
    return ExactReal(n -> begin
        p = max(n - msd(y) + 4, n + 3 >> 1)
        q = max(n - msd(x) + 4, n + 3 >> 1)

        (1 + x(p) * y(q)) >> (p + q - n)
    end
    )
end

Base.:*(x::Rational, y::ExactReal) = ExactReal(x) * y

Base.sqrt(x::ExactReal) = ExactReal(n -> isqrt(x(2n)))

"Returns -1 if x < y and +1 if x > y"
function compare(x::ExactReal, y::ExactReal)
    n = 0

    while true
        xn, yn = x(n), y(n)
        xn < yn - 1 && return -1  # x < y
        xn > yn + 1 && return 1   # x > y

        n += 1
    end
end

Base.:<(x::ExactReal, y::ExactReal) = compare(x, y) == -1


text:
#!/usr/bin/env julia

using FactCheck

test_regex = r"^test_.*\.jl$"
test_dir = Pkg.dir("AudioIO", "test")

test_files = filter(n -> ismatch(test_regex, n), readdir(test_dir))
if length(test_files) == 0
    error("No test files found. Make sure you're running from the root directory")
end

for test_file in test_files
    include(test_file)
end

# return the overall exit status
exitstatus()


text:
using SafeTestsets
const LONGER_TESTS = false

const GROUP = get(ENV, "GROUP", "All")
const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR")
const is_TRAVIS = haskey(ENV,"TRAVIS")

@time begin
if GROUP == "All" || GROUP == "Core"
    @time @safetestset "DynamicHMC" begin include("dynamicHMC.jl") end
    @time @safetestset "Turing" begin include("turing.jl") end
    @time @safetestset "ABC" begin include("abc.jl") end
end
end

if GROUP == "Stan"
    if is_TRAVIS
      using Pkg
      Pkg.add("CmdStan")
    end
    @time @safetestset "Stan_String" begin include("stan_string.jl") end
    @time @safetestset "Stan" begin include("stan.jl") end
end


text:
"""
This module has different inversion schemes.
* Alternate Optimization
"""

"""
Parameters for alternating minimization of a single objective function, while updating different model parameters
"""
mutable struct ParamAM
	name::String			# prints name of the optimization
	max_roundtrips::Int64		# limit number of roundtrips
	max_reroundtrips::Int64		# there will be reroundtrips, if roundtrips don't converge
	noptim::Int64 			# number of optimizations in every roundtrip
	optim_func::Vector{Function}	# optimization functions, call them to perform optimizations
	reinit_func::Function		# re-initialize function that is executed if roundtrips fail to converge
	after_reroundtrip_func::Function # function executed after each reroundtrip (use it to store erros)
	after_roundtrip_func::Function  # function to be executed after certain number of roundtrips
	fvec::Matrix{Float64}		# functionals
	fvec_init::Vector{Float64}	#
	optim_tols::Vector{Float64}	# tolerance for each optimization
	roundtrip_tol::Float64		# tolerance for stopping roundtrips
	verbose::Bool
	log::DataFrames.DataFrame
end

function ParamAM(optim_func;
	       name="",
	       noptim=length(optim_func),
	       optim_tols=[1e-3 for iop in 1:noptim],
	       roundtrip_tol=1e-3,
	       min_roundtrips=10, # minimum roundtrips before checking rate of convergence
	       verbose=true,
	       reinit_func=x->randn(),
	       after_reroundtrip_func=x->randn(),
	       after_roundtrip_func=x->randn(),
	       max_reroundtrips=1, re_init_flag=true, max_roundtrips=1)

	fvec=zeros(noptim, 2)
	fvec_init=zeros(noptim)

	# print dataframe
	log=DataFrame([[] for i in 1:4*noptim+3],
		 vcat(
			[:trip],
			[Symbol(string("J",i)) for i in 1:noptim],
			[Symbol(string("δJ",i)) for i in 1:noptim],
			[Symbol(string("fcalls",i)) for i in 1:noptim],
			[Symbol(string("gcalls",i)) for i in 1:noptim],
			[Symbol(string("time",i)) for i in 1:noptim]
			))

	pa=ParamAM(name, max_roundtrips, max_reroundtrips, noptim, optim_func,
	  reinit_func, after_reroundtrip_func,
	  after_roundtrip_func,
	  fvec, fvec_init, optim_tols, roundtrip_tol, verbose, log)

	return pa

end

"""
Perform
alternating optimizations, updating different model parameters, computing a
same objective functional
"""
function go(pa::ParamAM, io=stdout)

	reroundtrip_converge=false
	itrr=0

	pa.verbose && write(io,string(pa.name, "\t alternate optimization\n"))
	rf=zeros(pa.noptim)

	while ((!reroundtrip_converge && itrr < pa.max_reroundtrips))
		itrr += 1
		pa.verbose && (itrr > 1) && write(io,string("failed to converge.. reintializing (",itrr,"/",pa.max_reroundtrips,")\n"))
		pa.verbose && write(io,"=========================================================================================\n")

		fill!(pa.fvec,0.0)
		fill!(pa.fvec_init,0.0)

		# execute re-initialization function
		(itrr > 1) && pa.reinit_func(nothing)

		itr=0
		roundtrip_converge=false

		# print
		if(pa.verbose)
			write(io,@sprintf( "trip\t|"))
			for iop in 1:pa.noptim
				write(io,@sprintf( "\t\top %d\t(%0.1e)\t|",iop, pa.roundtrip_tol))
			end
#			@printf("\tvar(op) (%0.1e)\t",pa.roundtrip_tol)
			write(io,@sprintf( "\n"))
		end
		if(io ≠ stdout)
			prog = ProgressThresh(pa.roundtrip_tol, "Minimizing:")
		end
		while !roundtrip_converge && itr < pa.max_roundtrips

			itr += 1

			# optimizations in each roundtrip
			for iop in 1:pa.noptim
				name=string("op ",iop," in each roundtrip")
				pa.fvec[iop,2]=pa.fvec[iop,1]
				pa.fvec[iop,1]=pa.optim_func[iop](nothing)
			end

			# store functionals at the first roundtrip
			if(iszero(pa.fvec_init))
				for iop in 1:pa.noptim
					pa.fvec_init[iop]=pa.fvec[iop,1]
				end
			end

			# normalize func for each optim, always with first
			for iop in 1:pa.noptim
				pa.fvec[iop,1] /= pa.fvec_init[1]
			end

			# compute the change in the functions
			for iop in 1:pa.noptim
				rf[iop]=abs(pa.fvec[iop,2]-pa.fvec[iop,1])
			end

			#=
			if((2<itr<5) ||(itr<40 && (mod(itr,5)==0)) || ((itr<500 && mod(itr,20)==0)) && (mod(itr,100)==0))
				plotam!(itr,pa.fvec[:,1],rf)
			elseif(itr==2)
				plotam(itr,pa.fvec[:,1],rf)
			end
			=#
			(io ≠ stdout) && ProgressMeter.update!(prog, maximum(rf))

			if(itr > 5)# do atleast 10 round trips before quitting
				roundtrip_converge=all(rf .< pa.roundtrip_tol)
			else
				roundtrip_converge=false
			end

			push!(pa.log[!,:trip], itr)
			for iop in 1:pa.noptim
				push!(pa.log[!,Symbol(string("J",iop))], pa.fvec[iop,1])
				if(itr>2)
					push!(pa.log[!,Symbol(string("δJ",iop))], rf[iop])
				end
			end

			# print info
			if(pa.verbose)
				if((itr<5) ||(itr<40 && (mod(itr,5)==0)) || (itr<500 && (mod(itr,20)==0)) || (itr>500 && (mod(itr,50)==0)) || roundtrip_converge)

					pa.after_roundtrip_func(nothing) # after each roundtrip, execute this
					#if(itr==1)
					#	show(pa.log)
					#else
					#	DataFrames.showrowindices(stdout, pa.log, [itr],
					#			  DataFrames.getmaxwidths(pa.log, 1:4, 1:4, :Row), 1, 4)
					#end
					write(io,@sprintf( "%d\t|",itr))
					for iop in 1:pa.noptim
						push!(pa.log[!,Symbol(string("J",iop))], pa.fvec[iop,1])
						if(itr>2)
							push!(pa.log[!,Symbol(string("δJ",iop))], rf[iop])
						end
						write(io,@sprintf( "\t%0.6e\t",pa.fvec[iop,1]))
						(itr==1) ? write(io,@sprintf( "\t\t|")) : write(io,@sprintf( "(%0.6e)\t|",rf[iop]))
					end
					write(io,@sprintf( "\n"))
				end
				flush(io)
			end
			#show(pa.log)

			# variance b/w objectives
#			rf=vecnorm(pa.fvec[:,1])/vecnorm(fill(pa.fvec[1,1],pa.noptim))
#			if(pa.verbose)
#				@printf("\t%0.6e\t|",rf)
#				@printf("\n")
#			end
		end
		pa.after_reroundtrip_func(nothing) # after each roundtrip, execute this
		reroundtrip_converge = all(pa.fvec[:,1] .< pa.optim_tols[:])
		pa.verbose && write(io,"=========================================================================================\n")
	end
	if(all(pa.fvec[:,1] .< pa.optim_tols[:]))
		message=string("CONVERGED in ",itrr," reroundtrips")
	elseif(itrr == pa.max_reroundtrips)
		message="NOT CONVERGED: reached maximum reroundtrips"
	end
	pa.verbose && write(io, string(pa.name, "\t", message, "\n"))
	return nothing
end


text:
# This file was generated, do not modify it. # hide
using PlotlyJS
z =  [10     10.625  12.5  15.625  20
     5.625  6.25    8.125 11.25   15.625
     2.5    3.125   5.    8.125   12.5
     0.625  1.25    3.125 6.25    10.625
     0      0.625   2.5   5.625   10]

data   = contour(; z=z)
layout = Layout(; title="Basic Contour Plot")
plt    = plot(data, layout)

fdplotly(json(plt)) # hide

max_stars_repo_name: mcognetta/franklindocs

max_issues_repo_name: mcognetta/franklindocs

max_forks_repo_name: mcognetta/franklindocs


text:
module operators

import Base: ==, +, -, *, /
import ..states

using Base.LinAlg.BLAS
using Base.Cartesian

using ..bases, ..states

export Operator, DenseOperator,
       tensor, dagger, projector,
       expect, dense_identity, embed

"""
Abstract base class for all operators.

All deriving operator classes have to define the fields
basis_l and basis_r defining the left and right side bases.

For fast time evolution also at least the function
gemv!(alpha, op::Operator, x::Ket, beta, result::Ket) should be
implemented. Many other generic multiplication functions can be defined in
terms of this function and are provided automatically.
"""
abstract Operator

"""
Dense array implementation of Operator.

The matrix consisting of complex floats is stored in the data field.
"""
type DenseOperator <: Operator
    basis_l::Basis
    basis_r::Basis
    data::Matrix{Complex{Float64}}
    DenseOperator(b1::Basis, b2::Basis, data) = length(b1) == size(data, 1) && length(b2) == size(data, 2) ? new(b1, b2, data) : throw(DimensionMismatch())
end

DenseOperator(b::Basis, data) = DenseOperator(b, b, data)
DenseOperator(b1::Basis, b2::Basis) = DenseOperator(b1, b2, zeros(Complex, length(b1), length(b2)))
DenseOperator(b::Basis) = DenseOperator(b, b)
DenseOperator(op::Operator) = full(op)

"""
Converting an arbitrary Operator into a DenseOperator.
"""
Base.copy(x::DenseOperator) = deepcopy(x)
Base.full(x::DenseOperator) = deepcopy(x)
Base.full(op::Operator) = op*dense_identity(op.basis_r)

Base.eltype(x::Operator) = Complex128
Base.zero{T<:Operator}(x::T) = T(x.basis_l, x.basis_r)

==(x::DenseOperator, y::DenseOperator) = (x.basis_l == y.basis_l) && (x.basis_r == y.basis_r) && (x.data == y.data)

# Arithmetic operations for dense DenseOperators
check_samebases(a::Operator, b::Operator) = ((a.basis_l!=b.basis_l) || (a.basis_r!=b.basis_r) ? throw(IncompatibleBases()) : nothing)

*(a::DenseOperator, b::Ket) = (check_multiplicable(a.basis_r, b.basis); Ket(a.basis_l, a.data*b.data))
*(a::Bra, b::DenseOperator) = (check_multiplicable(a.basis, b.basis_l); Bra(b.basis_r, b.data.'*a.data))
*(a::DenseOperator, b::DenseOperator) = (check_multiplicable(a.basis_r, b.basis_l); DenseOperator(a.basis_l, b.basis_r, a.data*b.data))
*(a::DenseOperator, b::Number) = DenseOperator(a.basis_l, a.basis_r, complex(b)*a.data)
*(a::Number, b::DenseOperator) = DenseOperator(b.basis_l, b.basis_r, complex(a)*b.data)

/(a::DenseOperator, b::Number) = DenseOperator(a.basis_l, a.basis_r, a.data/complex(b))

+(a::DenseOperator, b::DenseOperator) = (check_samebases(a,b); DenseOperator(a.basis_l, a.basis_r, a.data+b.data))

-(a::DenseOperator) = DenseOperator(a.basis_l, a.basis_r, -a.data)
-(a::DenseOperator, b::DenseOperator) = (check_samebases(a,b); DenseOperator(a.basis_l, a.basis_r, a.data-b.data))

# Fast in-place multiplication implementations
gemm!{T<:Complex}(alpha::T, a::Matrix{T}, b::Matrix{T}, beta::T, result::Matrix{T}) = BLAS.gemm!('N', 'N', alpha, a, b, beta, result)
gemv!{T<:Complex}(alpha::T, a::Matrix{T}, b::Vector{T}, beta::T, result::Vector{T}) = BLAS.gemv!('N', alpha, a, b, beta, result)
gemv!{T<:Complex}(alpha::T, a::Vector{T}, b::Matrix{T}, beta::T, result::Vector{T}) = BLAS.gemv!('T', alpha, b, a, beta, result)

gemm!(alpha, a::DenseOperator, b::DenseOperator, beta, result::DenseOperator) = gemm!(alpha, a.data, b.data, beta, result.data)
gemv!(alpha, a::DenseOperator, b::Ket, beta, result::Ket) = gemv!(alpha, a.data, b.data, beta, result.data)
gemv!(alpha, a::Bra, b::DenseOperator, beta, result::Bra) = gemv!(alpha, a.data, b.data, beta, result.data)

"""
Tensor product of operators.
"""
bases.tensor(a::DenseOperator, b::DenseOperator) = DenseOperator(tensor(a.basis_l, b.basis_l), tensor(a.basis_r, b.basis_r), kron(a.data, b.data))
bases.tensor(ops::DenseOperator...) = reduce(tensor, ops)

"""
Tensor product of a ket and a bra results in an operator.
"""
bases.tensor(a::Ket, b::Bra) = DenseOperator(a.basis, b.basis, reshape(kron(b.data, a.data), prod(a.basis.shape), prod(b.basis.shape)))

"""
Hermitian conjugate of the given operator.
"""
states.dagger(x::DenseOperator) = DenseOperator(x.basis_r, x.basis_l, x.data')

"""
Create a projection operator.
"""
projector(a::Ket, b::Bra) = tensor(a, b)
projector(a::Ket) = tensor(a, dagger(b))
projector(a::Bra) = tensor(dagger(a), b)

"""
p-norm of given operator.
"""
Base.norm(op::DenseOperator, p) = norm(op.data, p)

"""
Trace of given operator.
"""
Base.trace(op::DenseOperator) = trace(op.data)

"""
Normalized copy of given operator (trace is 1.).
"""
states.normalize(op::DenseOperator) = op/trace(op)

"""
Normalize the given operator.
"""
function states.normalize!(op::DenseOperator)
    u = 1./trace(op)
    for j=1:size(op.data,2), i=1:size(op.data,1)
        op.data[i,j] *= u
    end
    return op
end

"""
Operator exponential.
"""
function Base.expm(op::DenseOperator)
    if !multiplicable(op.basis_r, op.basis_l)
        throw(ArgumentError("Operator has to be multiplicable with itself."))
    end
    return DenseOperator(op.basis_l, op.basis_r, expm(op.data))
end

"""
Expectation value of the given operator for the specified state(s).
"""
expect(op::Operator, state::DenseOperator) = trace(op*state)
expect(op::Operator, states::Vector{DenseOperator}) = [expect(op, state) for state=states]
expect(op::Operator, state::Ket) = dagger(state)*(op*state)
expect(op::Operator, states::Vector{Ket}) = [expect(op, state) for state=states]

"""
Identity operator.
"""
dense_identity(b::Basis) = DenseOperator(b, b, eye(Complex, length(b)))
dense_identity(b1::Basis, b2::Basis) = DenseOperator(b1, b2, eye(Complex, length(b1), length(b2)))
dense_identity(op::Operator) = dense_identity(op.basis_l, op.basis_r)
Base.identity(op::DenseOperator) = dense_identity(op.basis_l, op.basis_r)

# Multiplication for Operators in terms of their gemv! implementation
function gemm!(alpha, M::Operator, b::DenseOperator, beta, result::DenseOperator)
    for i=1:size(b.data, 2)
        bket = Ket(b.basis_l, b.data[:,i])
        resultket = Ket(M.basis_l, result.data[:,i])
        gemv!(alpha, M, bket, beta, resultket)
        result.data[:,i] = resultket.data
    end
end

function gemm!(alpha, b::DenseOperator, M::Operator, beta, result::DenseOperator)
    for i=1:size(b.data, 1)
        bbra = Bra(b.basis_r, vec(b.data[i,:]))
        resultbra = Bra(M.basis_r, vec(result.data[i,:]))
        gemv!(alpha, bbra, M, beta, resultbra)
        result.data[i,:] = resultbra.data
    end
end

function *(op1::Operator, op2::DenseOperator)
    check_multiplicable(op1.basis_r, op2.basis_l)
    result = DenseOperator(op1.basis_l, op2.basis_r)
    gemm!(Complex(1.), op1, op2, Complex(0.), result)
    return result
end

function *(op1::DenseOperator, op2::Operator)
    check_multiplicable(op1.basis_r, op2.basis_l)
    result = DenseOperator(op1.basis_l, op2.basis_r)
    gemm!(Complex(1.), op1, op2, Complex(0.), result)
    return result
end

function *(op::Operator, psi::Ket)
    check_multiplicable(op.basis_r, psi.basis)
    result = Ket(op.basis_l)
    gemv!(Complex(1.), op, psi, Complex(0.), result)
    return result
end

function *(psi::Bra, op::Operator)
    check_multiplicable(psi.basis, op.basis_l)
    result = Bra(op.basis_r)
    gemv!(Complex(1.), psi, op, Complex(0.), result)
    return result
end

Base.prod{B<:Basis, T<:AbstractArray}(basis::B, operators::T) = (length(operators)==0 ? dense_identity(basis) : prod(operators))

"""
Tensor product of operators where all missing indices are identity operators.

Arguments
---------
basis
    CompositeBasis of the resulting operator.
indices
    Indices of the subsystems in which the given operators live.
operators
    Operators defined in the subsystems.
"""
embed(basis::CompositeBasis, indices::Vector{Int}, operators::Vector) = tensor([prod(basis.bases[i], operators[find(indices.==i)]) for i=1:length(basis.bases)]...)
embed{T<:Operator}(basis::CompositeBasis, index::Int, op::T) = embed(basis, Int[index], T[op])

# Partial trace for dense operators.
function _strides(shape::Vector{Int})
    N = length(shape)
    S = zeros(Int, N)
    S[N] = 1
    for m=N-1:-1:1
        S[m] = S[m+1]*shape[m+1]
    end
    return S
end

@generated function _ptrace{RANK}(rank::Array{Int,RANK}, a::Matrix{Complex128},
                                  shape_l::Vector{Int}, shape_r::Vector{Int},
                                  indices::Vector{Int})
    return quote
        a_strides_l = _strides(shape_l)
        result_shape_l = deepcopy(shape_l)
        result_shape_l[indices] = 1
        result_strides_l = _strides(result_shape_l)
        a_strides_r = _strides(shape_r)
        result_shape_r = deepcopy(shape_r)
        result_shape_r[indices] = 1
        result_strides_r = _strides(result_shape_r)
        N_result_l = prod(result_shape_l)
        N_result_r = prod(result_shape_r)
        result = zeros(Complex128, N_result_l, N_result_r)
        @nexprs 1 (d->(Jr_{$RANK}=1;Ir_{$RANK}=1))
        @nloops $RANK ir (d->1:shape_r[d]) (d->(Ir_{d-1}=Ir_d; Jr_{d-1}=Jr_d)) (d->(Ir_d+=a_strides_r[d]; if !(d in indices) Jr_d+=result_strides_r[d] end)) begin
            @nexprs 1 (d->(Jl_{$RANK}=1;Il_{$RANK}=1))
            @nloops $RANK il (k->1:shape_l[k]) (k->(Il_{k-1}=Il_k; Jl_{k-1}=Jl_k; if (k in indices && il_k!=ir_k) Il_k+=a_strides_l[k]; continue end)) (k->(Il_k+=a_strides_l[k]; if !(k in indices) Jl_k+=result_strides_l[k] end)) begin
                #println("Jl_0: ", Jl_0, "; Jr_0: ", Jr_0, "; Il_0: ", Il_0, "; Ir_0: ", Ir_0)
                result[Jl_0, Jr_0] += a[Il_0, Ir_0]
            end
        end
        return result
    end
end

"""
Partial trace of the given operator over the specified indices.
"""
function bases.ptrace(a::DenseOperator, indices::Vector{Int})
    rank = zeros(Int, [0 for i=1:length(a.basis_l.shape)]...)
    result = _ptrace(rank, a.data, a.basis_l.shape, a.basis_r.shape, indices)
    return DenseOperator(ptrace(a.basis_l, indices), ptrace(a.basis_r, indices), result)
end

bases.ptrace(a::DenseOperator, index::Int) = bases.ptrace(a, Int[index])

"""
Partial trace of the given state vector over the specified indices.
"""
bases.ptrace(a::Ket, indices) = bases.ptrace(tensor(a, dagger(a)), indices)
bases.ptrace(a::Bra, indices) = bases.ptrace(tensor(dagger(a), a), indices)

end # module


text:
@testset "BorgMOEA" begin
    @testset "Schaffer1" begin
        res = bboptimize(BlackBoxOptim.Schaffer1Family; Method=:borg_moea,
                         FitnessScheme=ParetoFitnessScheme{2}(is_minimizing=true),
                         SearchRange=(-10.0, 10.0), NumDimensions=2, ϵ=0.01,
                         MaxSteps=5000, TraceMode=:silent)
        @test BlackBoxOptim.IGD(BlackBoxOptim.Schaffer1Family.opt_value, pareto_frontier(res),
                                fitness_scheme(res), Val{length(best_candidate(res))}) < 0.05

        @testset "using MultithreadEvaluator" begin
            if Threads.nthreads() > 1
                opt = bbsetup(BlackBoxOptim.Schaffer1Family; Method=:borg_moea,
                              FitnessScheme=ParetoFitnessScheme{2}(is_minimizing=true),
                              SearchRange=(-10.0, 10.0), NumDimensions=2, ϵ=0.01,
                              MaxSteps=5000, TraceMode=:verbose, TraceInterval=1.0, NThreads=Threads.nthreads()-1)
                res = bboptimize(opt)
                @test isa(BlackBoxOptim.evaluator(lastrun(opt)), BlackBoxOptim.MultithreadEvaluator)
                @test BlackBoxOptim.IGD(BlackBoxOptim.Schaffer1Family.opt_value, pareto_frontier(res),
                           fitness_scheme(res), Val{length(best_candidate(res))}) < 0.05
            else
                @warn "BorgMOEA MultithreadEvaluator tests require >1 threads, $(Threads.nthreads()) found, use JULIA_NUM_THREADS"
                @test_skip Threads.nthreads() > 1
            end
        end
    end

    @testset "CEC09_UP8" begin
        res = bboptimize(BlackBoxOptim.CEC09_Unconstrained_Set[8]; Method=:borg_moea,
                         NumDimensions=4, ϵ=0.025,
                         MaxSteps=50000, TraceMode=:silent)
        @test BlackBoxOptim.IGD(BlackBoxOptim.CEC09_Unconstrained_Set[8].opt_value, pareto_frontier(res),
                                fitness_scheme(res), Val{length(best_candidate(res))}) < 0.13
    end
end


text:
//test 2 in julia
function a()
	x = 1
	while < x 4 do
		x + = x 1
	end
	print(x)
end

max_stars_repo_name: tylerholmes/Interpreter

max_issues_repo_name: tylerholmes/Interpreter

max_forks_repo_name: tylerholmes/Interpreter


text:
struct CubicDR{S<:Grid, T<:CartesianGrid, nx, d} <: AbstractDecisionRule{S, T, nx}
    grid_exo::S
    grid_endo::T
    itp::Vector{Array{Value{nx}, d}}
end

(dr::CubicDR)(args...) = evaluate(dr, args...)

#####
##### 1-argument decision rule
#####

function CubicDR(exo_grid::EmptyGrid, endo_grid::CartesianGrid{d}, i::Union{Val{nx}, Type{Val{nx}}}) where d where nx
    dims = endo_grid.n+2
    c = [zeros(Value{nx},dims...)]
    CubicDR{EmptyGrid, CartesianGrid{d}, nx, d}(exo_grid, endo_grid, c)
end

function set_values!(dr::CubicDR{EmptyGrid,CartesianGrid{d},n_x,d},  V::Vector{<:Array{Value{n_x}}}) where n_x where d
    n = dr.grid_endo.n
    C = dr.itp[1]
    ind = [2:(n[i]+1) for i=1:length(n)]
    # this segfaults, don't know why
    # C[ind...] = reshape(V[1], n)
    C[ind...] = V[1]
    prefilter!(C)
end

function evaluate(dr::CubicDR{EmptyGrid,CartesianGrid{d}}, points::AbstractVector{Point{d}}) where d
    a = SVector{d,Float64}(dr.grid_endo.min)
    b = SVector{d,Float64}(dr.grid_endo.max)
    n = SVector{d,Int64}(dr.grid_endo.n)
    C = dr.itp[1]
    return eval_UC_spline(a, b, n, C, points)
end

#
function evaluate(dr::CubicDR{EmptyGrid,CartesianGrid{d}}, y::Point{d}) where d
    return evaluate(dr, [y])[1]
end

####
#### 2 CartesianGrid continous arguments d.r.
####

# function CubicDR(exo_grid::EmptyGrid, endo_grid::CartesianGrid{d}, i::Union{Val{nx}, Type{Val{nx}}}) where d where nx
#     dims = endo_grid.n+2
#     c = [zeros(Value{nx},dims...)]
#     CubicDR{EmptyGrid, CartesianGrid, nx, d}(exo_grid, endo_grid, c)
# end

function CubicDR(exo_grid::CartesianGrid{d1}, endo_grid::CartesianGrid{d2}, i::Union{Val{nx}, Type{Val{nx}}}) where d1 where d2 where nx
    dims = cat(exo_grid.n+2, endo_grid.n+2; dims=1)
    c = [zeros(Value{nx},dims...)]
    CubicDR{CartesianGrid{d1}, CartesianGrid{d2}, nx, d1+d2}(exo_grid, endo_grid, c)
end
#
function set_values!(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2},n_x,d}, V::Vector{<:Array{Value{n_x}}}) where n_x where d where d1 where d2
    exog = dr.grid_exo
    endog = dr.grid_endo

    data = reshape( cat(V...; dims=1), endog.n..., exog.n...)
    d_m = ndims(exog)
    d_s = ndims(endog)
    perms = cat(((d_s+1):(d_s+d_m))..., (1:d_s)...; dims=1)
    RV = permutedims(data, perms)
    C = dr.itp[1]
    dims = size(C)
    inds = [2:(i-1) for i in dims]
    C[inds...] = RV
    prefilter!(C)
end

function evaluate(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2}, n_x, d}, z::AbstractVector{Point{d}}) where n_x where d where d1 where d2
    a = cat(dr.grid_exo.min, dr.grid_endo.min; dims=1)
    b = cat(dr.grid_exo.max, dr.grid_endo.max; dims=1)
    n = cat(dr.grid_exo.n, dr.grid_endo.n; dims=1)
    cc = dr.itp[1]
    res = eval_UC_spline(a, b, n, cc, z)
    return res
end

function evaluate(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2}}, x::AbstractVector{Point{d1}}, y::AbstractVector{Point{d2}}) where d1 where d2
    N = length(x)
    z = [ [x[i]; y[i]] for i=1:N]
    evaluate(dr, z)
end

function evaluate(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2}}, x::Point{d1}, y::AbstractVector{Point{d2}}) where d1 where d2
    N = length(y)
    z = [ [x; y[i]] for i=1:N]
    evaluate(dr, z)
end

function evaluate(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2}}, i::Int64, y::Union{Point{d2},AbstractVector{Point{d2}}}) where d1 where d2
    x = node(Point, dr.grid_exo, i)
    evaluate(dr, x, y)
end

# TODO replace by generic call?
function evaluate(dr::CubicDR{CartesianGrid{d1},CartesianGrid{d2}}, x::Point{d1}, y::Point{d2}) where d1 where d2
    dr([x],[y])[1]
end

####
#### UnstructuredGrid × CartesianGrid 2 continous arguments d.r.
####

function CubicDR(exo_grid::UnstructuredGrid{d1}, endo_grid::CartesianGrid{d2}, i::Union{Val{nx}, Type{Val{nx}}}) where d1 where d2 where nx
    c = [zeros(Value{nx},(endo_grid.n+2)...) for _=1:n_nodes(exo_grid)]
    CubicDR{UnstructuredGrid{d1}, CartesianGrid{d2}, nx, d2}(exo_grid, endo_grid, c)
end

function set_values!(dr::CubicDR{UnstructuredGrid{d1},CartesianGrid{d2}, n_x, d2}, values::Vector{<:Array{Value{n_x}}}) where d1 where d2 where n_x
    orders = dr.grid_endo.n
    inds = [2:(o+1) for o in orders]
    for (i,C) in enumerate(dr.itp)
        C[inds...] = reshape(values[i],orders...)
        prefilter!(C)
    end
end

function evaluate(dr::CubicDR{<:UnstructuredGrid,<:CartesianGrid}, i::Int, z::AbstractVector{Point{d}}) where d
    a = dr.grid_endo.min
    b = dr.grid_endo.max
    n = dr.grid_endo.n
    cc = dr.itp[i]
    splines.eval_UC_spline(a, b, n, cc, z)
end

# TODO replace by generic call?
function evaluate(dr::CubicDR{<:UnstructuredGrid,<:CartesianGrid}, i::Int, z::Point{d}) where d
    evaluate(dr, i, [z])[1]
end


text:
include("ast.jl")

mutable struct Parser
    lexer::Lexer
    current::Token
    next::Token
end

function Parser(lexer::Lexer)
    current = nexttoken(lexer)
    next = nexttoken(lexer)
    Parser(lexer, current, next)
end

function advance(p::Parser)
    p.current = p.next
    p.next = nexttoken(p.lexer)
end

function parse_identifier(p::Parser)
    if p.current.type != IDENT
        throw(DomainError(p.current, "identifier expected"))
    end
    ident = Identifier(p.current)
    advance(p)
    ident
end

function parse_expression(p)
    return Expression()
end

function parse_let(p::Parser)
    token = p.current
    advance(p)
    ident = parse_identifier(p)
    if p.current.type != ASSIGN
        throw(DomainError(p.current, "assignment expected"))
    end
    expr = parse_expression(p)
    while p.current.type != SEMICOLON
        advance(p)
    end
    advance(p)
    return LetStatement(token, ident, expr)
end

function parse_return(p::Parser)
    token = p.current
    advance(p)
    expr = parse_expression(p)
    while p.current.type != SEMICOLON
        advance(p)
    end
    advance(p)
    return ReturnStatement(token, expr)
end

function parse_if(p::Parser)

end

function parse_statement(p::Parser)
    if p.current.type == LET
        stmt = parse_let(p::Parser)
    elseif p.current.type == RETURN
        stmt = parse_return(p::Parser)
    elseif p.current.type == IF
        stmt = parse_if(p::Parser)
    else
        throw(DomainError(p.current, "statement expected"))
    end
    stmt
end

function parse(p::Parser)
    program = Program()
    local stmt::Statement

    while p.current.type != EOF
        push!(program.statements, parse_statement(p))
    end

    program
end

max_stars_repo_name: viert/monkeylang-julia

max_issues_repo_name: viert/monkeylang-julia

max_forks_repo_name: viert/monkeylang-julia


text:
module OpenIDConnect

using HTTP
using JSON
using MbedTLS
using Base64
using Random
using JWTs

const DEFAULT_SCOPES = ["openid", "profile", "email"]
const DEFAULT_STATE_TIMEOUT_SECS = 60
const DEFAULT_SKEW_SECS = 2*60
const STATE_PURGE_TRIGGER = 1024
const DEFAULT_KEY_REFRESH_SECS = 60*60

export OIDCCtx, flow_request_authorization_code, flow_get_authorization_code, flow_get_token, flow_validate_id_token, flow_refresh_token

"""
Holds an OpenID Connect context that can be used in subsequent OpenID request flows.
The context holds request states, and configuration options.
"""
struct OIDCCtx
    states::Dict{String,Float64}
    state_timeout_secs::Int
    allowed_skew_secs::Int
    openid_config::Dict{String,Any}
    http_tls_opts::Dict{Symbol,Any}
    validator::JWKSet
    key_refresh_secs::Int
    last_key_refresh::Float64
    client_id::String
    client_secret::String
    scopes::Vector{String}
    redirect_uri::String
    random_device::RandomDevice

    function OIDCCtx(issuer::String, redirect_uri::String, client_id::String, client_secret::String, scopes::Vector{String}=DEFAULT_SCOPES;
                        verify::Union{Nothing,Bool}=nothing, cacrt::Union{Nothing,String,MbedTLS.CRT}=nothing,
                        state_timeout_secs::Int=DEFAULT_STATE_TIMEOUT_SECS, allowed_skew_secs::Int=DEFAULT_SKEW_SECS, key_refresh_secs::Int=DEFAULT_KEY_REFRESH_SECS,
                        random_device::RandomDevice=RandomDevice())
        endswith(issuer, "/") || (issuer = issuer * "/")
        openid_config_url = issuer * ".well-known/openid-configuration"
        http_tls_opts = Dict{Symbol,Any}()

        if verify !== nothing
            http_tls_opts[:require_ssl_verification] = verify
        end

        if cacrt !== nothing
            isa(cacrt, String) && (cacrt = MbedTLS.crt_parse(base64decode(cacrt)))
            conf = MbedTLS.SSLConfig()
            MbedTLS.config_defaults!(conf)
            entropy = MbedTLS.Entropy()
            rng = MbedTLS.CtrDrbg()
            MbedTLS.seed!(rng, entropy)
            MbedTLS.rng!(conf, rng)
            MbedTLS.authmode!(conf, ((verify === nothing) || verify) ? MbedTLS.MBEDTLS_SSL_VERIFY_REQUIRED : MbedTLS.MBEDTLS_SSL_VERIFY_NONE)
            MbedTLS.ca_chain!(conf, cacrt)
            http_tls_opts[:tlsconfig] = conf
        end

        # fetch and store the openid config, along with the additional args for SSL
        openid_config = JSON.parse(String(HTTP.request("GET", openid_config_url; status_exception=true, http_tls_opts...).body))
        validator = JWKSet(openid_config["jwks_uri"])

        new(Dict{String,Float64}(), state_timeout_secs, allowed_skew_secs, openid_config, http_tls_opts, validator, key_refresh_secs, 0.0, client_id, client_secret, scopes, redirect_uri, random_device)
    end
end

authorization_endpoint(ctx::OIDCCtx) = ctx.openid_config["authorization_endpoint"]
token_endpoint(ctx::OIDCCtx) = ctx.openid_config["token_endpoint"]

function remember_state(ctx::OIDCCtx, state::String)
    ctx.states[state] = time()
    nothing
end

function validate_state(ctx::OIDCCtx, state::String)
    statestore = ctx.states
    if state in keys(statestore)
        t = statestore[state]
        delete!(statestore, state)
        if (time() - t) <= ctx.state_timeout_secs
            return true
        end
    end
    @info("encountered an unknown or expired state")
    if length(statestore) > STATE_PURGE_TRIGGER
        purge_states!(ctx)
    end
    false
end

function purge_states(ctx::OIDCCtx)
    tnow = time()
    tmout = ctx.state_timeout_secs
    filter!(nv->(tnow-nv[2])>tmout, ctx.states)
    nothing
end

"""
API calling error detected by this library
"""
struct APIError
    error::String
end

"""
Error returned from OpenID server
See section 3.1.2.6 of https://openid.net/specs/openid-connect-core-1_0.html
"""
struct AuthServerError
    error::String
    error_description::Union{Nothing,String}
    error_uri::Union{Nothing,String}
end

"""
Authentication request. Uses the authorization code flow.
Acceptable optional args as listed in section 3.1.2.1 of specifications (https://openid.net/specs/openid-connect-core-1_0.html)

Returns a String with the redirect URL.
Caller must perform the redirection.
"""
function flow_request_authorization_code(ctx::OIDCCtx; nonce=nothing, display=nothing, prompt=nothing, max_age=nothing, ui_locales=nothing, id_token_hint=nothing, login_hint=nothing, acr_values=nothing)
    @debug("oidc negotiation: initiating...")
    scopes = join(ctx.scopes, ' ')
    state = randstring(ctx.random_device, 10)
    remember_state(ctx, state)

    query = Dict("response_type"=>"code", "client_id"=>ctx.client_id, "redirect_uri"=>ctx.redirect_uri, "scope"=>scopes, "state"=>state)
    (nonce          === nothing) || (query["nonce"]         = String(nonce))
    (display        === nothing) || (query["display"]       = String(display))
    (prompt         === nothing) || (query["prompt"]        = String(prompt))
    (max_age        === nothing) || (query["max_age"]       = String(max_age))
    (ui_locales     === nothing) || (query["ui_locales"]    = String(ui_locales))
    (id_token_hint  === nothing) || (query["id_token_hint"] = String(id_token_hint))
    (login_hint     === nothing) || (query["login_hint"]    = String(login_hint))
    (acr_values     === nothing) || (query["acr_values"]    = String(acr_values))

    uri = HTTP.merge(HTTP.URIs.URI(authorization_endpoint(ctx)); query=query)
    return string(uri)
end

"""
Given the params from the redirected response from the authentication request, extract the authorization code.
See sections 3.1.2.5 and 3.1.2.6 of https://openid.net/specs/openid-connect-core-1_0.html.

Returns the authorization code on success.
Returns one of APIError or AuthServerError on failure.
"""
function flow_get_authorization_code(ctx::OIDCCtx, query)
    state = get(query, "state", nothing)
    if state === nothing
        return APIError("invalid request, no state found")
    end
    if validate_state(ctx, String(state)) === nothing
        return APIError("invalid or expired state")
    end

    code = get(query, "code", nothing)
    if code !== nothing
        return String(code)
    end

    errcode = get(query, "error", nothing)
    if errcode !== nothing
        return AuthServerError(errcode, get(query, "error_description", nothing), get(query, "error_uri", nothing))
    end
    return APIError("invalid request, no code or error found")
end

function parse_token_response(tok_res)
    @info("oidc: success response from token endpoint")
    resp_str = String(tok_res.body)

    if tok_res.status == 200
        return JSON.parse(resp_str)
    end

    try
        err_resp = JSON.parse(resp_str)
        errcode = get(err_resp, "error", nothing)
        if errcode !== nothing
            return AuthServerError(errcode, get(err_resp, "error_description", nothing), get(err_resp, "error_uri", nothing))
        end
    catch
        return APIError("unknown response from server: " * resp_str)
    end
end

"""
Token Request. Given the authorization code obtained, invoke the token end point and obtain an id_token, access_token, refresh_token.
See section 3.1.3.1 of https://openid.net/specs/openid-connect-core-1_0.html.

Returns a JSON object containing tokens on success.
Returns a AuthServerError or APIError object on failure.
"""
function flow_get_token(ctx::OIDCCtx, code)
    data = Dict("grant_type"=>"authorization_code",
                "code"=>String(code),
                "redirect_uri"=>ctx.redirect_uri,
                "client_id"=>ctx.client_id,
                "client_secret"=>ctx.client_secret)
    headers = Dict("Content-Type"=>"application/x-www-form-urlencoded")
    tok_res = HTTP.request("POST", token_endpoint(ctx), headers, HTTP.URIs.escapeuri(data); status_exception=false, ctx.http_tls_opts...)
    return parse_token_response(tok_res)
end

"""
Token Refresh. Given the refresh code obtained, invoke the token end point and obtain new tokens.
See section 12 of https://openid.net/specs/openid-connect-core-1_0.html.

Returns a JSON object containing tokens on success.
Returns a AuthServerError or APIError object on failure.
"""
function flow_refresh_token(ctx::OIDCCtx, refresh_token)
    data = Dict("grant_type"=>"refresh_token",
                "refresh_token"=>String(refresh_token),
                "client_id"=>ctx.client_id,
                "client_secret"=>ctx.client_secret)
    headers = Dict("Content-Type"=>"application/x-www-form-urlencoded")
    tok_res = HTTP.request("POST", token_endpoint(ctx), headers, HTTP.URIs.escapeuri(data); status_exception=false, ctx.http_tls_opts...)
    return parse_token_response(tok_res)
end

"""
Validate an OIDC token.
Validates both the structure and signature.
See section 3.1.3.7 of https://openid.net/specs/openid-connect-core-1_0.html
"""
flow_validate_id_token(ctx::OIDCCtx, id_token) = flow_validate_id_token(ctx, JWT(;jwt=String(id_token)))
function flow_validate_id_token(ctx::OIDCCtx, jwt::JWT)
    isvalid = false

    if issigned(jwt)
        try
            tokclaims = claims(jwt)
            issue_time = tokclaims["iat"] - ctx.allowed_skew_secs
            expiry_time = tokclaims["exp"] + ctx.allowed_skew_secs
            isvalid = issue_time <= round(Int, time()) <= expiry_time
        catch ex
            @info("invalid token format ($ex)")
        end

        if isvalid
            validator = ctx.validator
            if (time() - ctx.last_key_refresh) >= ctx.key_refresh_secs
                refresh!(validator)
            end
            isvalid = validate!(jwt, validator)
        end
    end

    return isvalid
end

end # module


text:
# First order methods trace, used by AcceleratedGradientDescent,
# ConjugateGradient, GradientDescent, LBFGS and MomentumGradientDescent
function common_trace!(tr, d, state, iteration, method::FirstOrderOptimizer, options, curr_time=time())
    dt = Dict()
    dt["time"] = curr_time
    if options.extended_trace
        dt["x"] = copy(state.x)
        dt["g(x)"] = copy(gradient(d))
        dt["Current step size"] = state.alpha
    end
    g_norm = maximum(abs, gradient(d))
    update!(tr,
            iteration,
            value(d),
            g_norm,
            dt,
            options.store_trace,
            options.show_trace,
            options.show_every,
            options.callback)
end


text:
# illustrate how to make and display the Shepp-Logan ellipse digital image phantom
using MIRT: jim, ellipse_im
x = ellipse_im(100)
jim(x, title="Shepp-Logan Phantom")


text:
#!!!! Run this script twice! the first run includes gpu compilation

using Pkg, Random, CUDAKernels, Adapt, CUDA, Rasters, BenchmarkTools, Plots
using DynamicGrids: CuGPU, CPU, GPU, SimData
using Rasters.LookupArrays
using Plots: px
theme(:vibrant)
CUDA.allowscalar(false)
# The directory for this project
basedir = dirname(Pkg.project().path)
# Inlude models and data scripts
include(joinpath(basedir, "src", "rules.jl"))
include(joinpath(basedir, "src", "data.jl"))
mkpath(joinpath(basedir, "output"))

# Define rules and simulation data
growth, allee, localdisp, wind, parasitism, localdisp_p, allee_p, carrycap = define_rules()
init_h, init_p, rate_h, rate_p, mask = load_data(basedir, carrycap)

##### GPU wind rule #####

# We can't use `rand` directly _inside_ a GPU kernel, as of February 2021.
# To get around this shortcoming we add a grid to hold random numbers and
# call the random generator from outside a kernel using a `SetGrid` rule that
# calls a function on the entire grid. All the work is still done on the GPU.
gpu_wind = SetNeighbors{Tuple{:rand,:H}}() do data, hood, (rand, H), I
    # Ignore empty cells
    H > zero(H) || return nothing
    # Randomise a dispersal event using cell value of rand
    rand < 0.02f0 || return nothing
    # Get random values for jump from the surrounding grid
    randi, randj = (I[1], I[2] + 1), (I[1] + 1, I[2] + 1)
    dest = I .+ (unsafe_trunc(Int, data[:rand][randi...] * 40 - 20),
                 unsafe_trunc(Int, data[:rand][randj...] * 40 - 20))
    # Update spotted cell if it's on the grid
    if isinbounds(data[:H], dest)
        add!(data[:H], H, dest...)
        sub!(data[:H], H, I...)
    end
    return nothing
end

_rand!(A) = rand!(A)
_rand!(A::CuArray) = CUDA.rand!(A)

# A rule that just fills the random grid
randomiser = SetGrid{Tuple{},:rand}() do randgrid
    # `parent` is so we write rand to the whole grid, including any padding
    _rand!(parent(randgrid))
end

##### Benchmark setup methods #####

function setupsim(rules::NamedTuple;
    tspan=DateTime(2020, 1):Week(1):DateTime(2022, 1),
    opt=NoOpt(), proc=SingleCPU(), size_ag,
    growthmin=-5.0f0, growthmax=0.2f0,
    output_type=ResultOutput,
)
    # Take top right corner to make a square - but most of australia
    ax = 1:200, lastindex(init_h, 2)-199:lastindex(init_h, 2)
    sze = size_ag[1]
    ag = size_ag[2]
    hpop = ag(init_h[ax...])
    @assert size(hpop, 1) == size(hpop, 2) == sze
    init = (;
        H=parent(ag(init_h[ax...])),
        P=parent(ag(init_p[ax...])),
        rand=rand(Float32, sze, sze)
    )
    o = output_type(init;
        aux=(; rH=ag(rate_h[ax..., :]), rP=ag(rate_p[ax..., :])),
        mask=Array(ag(mask[ax...])),
        tspan=tspan,
        style=Braile(), color=:blue, fps=100,
    )
    if proc isa GPU
        opt isa SparseOpt && error("Sparse GPU")
        proc_rules = rules[:gpu]
        # Copy data to GPU so it's not part of the timing - otherwise
        # we are just benchmarking CPU/GPU data transfer.
        o = adapt(CuArray, o)
        rs = Ruleset(proc_rules; opt=opt, proc=proc)
        sd = adapt(CuArray, SimData(o, rs))
    else
        proc_rules = rules[:cpu]
        rs = Ruleset(proc_rules; opt=opt, proc=proc)
        sd = SimData(o, rs)
    end
    return (; o, rs, sd)
end
function suite(rules)
    suite = BenchmarkGroup()
    for opt in keys(opts)
        suite[opt] = BenchmarkGroup([:proc, :size])
        for proc in keys(procs), s in sizes
            # SparseOpt is not implemented for GPU yet
            proc == :gpu && opt == :sparseopt && continue
            suite[opt][proc, first(s)] = @benchmarkable sim!(o, rs) setup=((o, rs, sd)=setupsim($rules; size_ag=$s, proc=$(procs[proc]), opt=$(opts[opt])))
        end
    end
    suite
end

##### Scenarios and benchmark #####

procs = (single=SingleCPU(), threaded=ThreadedCPU(), gpu=CuGPU(),)
opts = (noopt=NoOpt(), sparseopt=SparseOpt())
# We use a an anonymous function
sizes = (
    100 => A -> aggregate(Center(), A, (Y(2), X(2), Ti(1))),
    200 => A -> A,
    400 => A -> disaggregate(Center(), A, (Y(2), X(2), Ti(1))),
    800 => A -> disaggregate(Center(), A, (Y(4), X(4), Ti(1))),
    1600 => A -> disaggregate(Center(), A, (Y(8), X(8), Ti(1))),
)

# CPU/GPU rules are separated as the best optimisations are not the same on CPU and GPU.
# Chain has better effects on GPU than on GPU, and GPU needs random numbers generated
# in a GridRule on a separate grid, due to the current lack of a `rand` function
# inside GPU kernels.
rulegroups = (
    Wind=(cpu=(wind, allee, growth), gpu=(randomiser, gpu_wind, allee, growth)),
    Local=(cpu=(localdisp, allee, growth), gpu=(localdisp, allee, growth)),
    Combined=(cpu=(wind, localdisp, allee, growth), gpu=(randomiser, gpu_wind, localdisp, allee, growth)),
    Parasitism=(cpu=(wind, localdisp, growth, allee, localdisp_p, allee_p, parasitism),
                gpu=(randomiser, gpu_wind, localdisp, localdisp_p, Chain(allee, growth, allee_p, parasitism))),
)

suites = map(suite, rulegroups)
# map(tune!, suites)
results = map(su -> run(su, verbose=true), suites)

##### Benchmark plot #####

function plotbench(b, key)
    @show b
    p = plot()
    i = 2
    for prockey in Tuple(keys(procs))
        for optkey in Tuple(keys(opts))
            if optkey == :sparseopt
                prockey == :gpu && continue
                lab = uppercasefirst(string("Sparse ", prockey))
            else
                lab = if prockey == :gpu
                    uppercase(string(prockey))
                else
                    uppercasefirst(string(prockey))
                end
            end
            times = map(sizes) do (s, ag)
                median(b.data[string(optkey)].data[(string(prockey), s)].times) / 1e9
            end
            plot!(p, collect(map(first, sizes)), collect(times);
                color=i, label=lab, pallete=:Dark2_5, opacity=0.6
            )
            i += 1
        end
    end
    plot(p;
        title=string(key),
        ylims=(1e-3, 10.0),
        yaxis=:log,
        titlefontsize=10,
        guidefontsize=9,
        tickfontsize=8,
        linewidth=2,
        legend=(key == :Parasitism ? :bottomright : :none),
        tickfontcolor=RGB(0.2),
        legendfontcolor=RGB(0.2),
        minorgrid=false,
        grid=true,
        xrotation=60,
        xticks=[100, 200, 400, 800, 1600],
        xlabel="Grid size",
        ylabel=(key in (:Wind, :Combined) ? "Time in seconds" : ""),
    )
end
plot(map(plotbench, results, keys(results))...; dpi=600, layout=(2, 2), size=(530, 600))

savefig("output/benchmarks.png")


text:
##################################################################
##################### Objective with candidate storage
##################################################################

function objective_min_cost_storage_cordoba(pm::_PM.AbstractPowerModel)
    add_co2_cost = haskey(pm.setting, "add_co2_cost") && pm.setting["add_co2_cost"]
    return JuMP.@objective(pm.model, Min,
        sum(
            _FP.calc_gen_cost(pm, n, add_co2_cost)
            + _FP.calc_convdc_ne_cost(pm, n, add_co2_cost)
            + _FP.calc_ne_branch_cost(pm, n, add_co2_cost)
            + _FP.calc_branchdc_ne_cost(pm, n, add_co2_cost)
            + calc_ne_storage_cost_cordoba(pm, n, add_co2_cost)
        for n in _PM.nw_ids(pm))
    )
end

function calc_ne_storage_cost_cordoba(pm::_PM.AbstractPowerModel, n::Int, add_co2_cost::Bool)
    #addition of buy and sell price storage["cost_abs"],storage["cost_inj"]
    function calc_single_ne_storage_cost(i, cost_a,cost_i)
        #if (_PM.var(pm,n,:ps,i)>0)
            cost = -1*cost_i * _PM.var(pm, n, :sd_ne, i)*_PM.var(pm, n, :z_strg_ne, i)# takes profit
        #else
            cost += cost_a * _PM.var(pm, n, :sc_ne, i)*_PM.var(pm, n, :z_strg_ne, i)# buys energy
        #end
        return cost
    end

    if haskey(_PM.ref(pm, n), :ne_storage)
        ne_storage = _PM.ref(pm, n, :ne_storage)

        if !isempty(ne_storage)
            #cost for buy and sell MWh
            cost = sum(calc_single_ne_storage_cost(i,storage["cost_abs"],storage["cost_inj"]) for (i,storage) in ne_storage)
            cost += sum((storage["eq_cost"] + storage["inst_cost"])*_PM.var(pm, n, :z_strg_ne, i) for (i,storage) in ne_storage)
            if add_co2_cost
                cost += sum(storage["co2_cost"]*_PM.var(pm, n, :z_strg_ne, i) for (i,storage) in ne_storage)
            end
        end
    end
    return cost
end


text:
function define_memory(machine)
    define(machine, "allot",    op_allot)
    define(machine, "here",     op_here)
    define(machine, "@",        op_at)
    define(machine, "!",        op_exclamation)
end

function op_allot(machine, parent, myidx)
    cellcount = pop!(machine.stack)
    if cellcount == 0
        return 1
    end
    here = length(machine.memory)
    if cellcount < 0
        print(machine.out, "Err: Not implemented: deallocation")
        return 1
    end
    resize!(machine.memory, here + cellcount)
    machine.here = here + 1
    return 1
end

function op_here(machine, parent, myidx)
    push!(machine.stack, machine.here)
    return 1
end

function op_at(machine, parent, myidx)
    push!(machine.stack, machine.memory[pop!(machine.stack)])
    return 1
end

function op_exclamation(machine, parent, myidx)
    addr =  pop!(machine.stack)
    machine.memory[addr] = pop!(machine.stack)
    return 1
end

const op_memory =
"""
: variable immediate ' create , 1 allot ' here , ' , , ;
"""


text:
using Test
using AugmentedGaussianProcesses
using LinearAlgebra
using Statistics
using Distributions
const AGP = AugmentedGaussianProcesses
include("testingtools.jl")

nData = 100; nDim = 2
m = 50; b= 10
k = AGP.RBFKernel(0.1)

X = rand(nData,nDim)
y = rand(MvNormal(zeros(nData),AGP.kernelmatrix(X,k)+1e-3I))

floattypes = [Float64]
@testset "GP" begin
    for floattype in floattypes
        @test typeof(GP(X,y,k)) <: GP{eval(Meta.parse("GaussianLikelihood{"*string(floattype)*"}")),eval(Meta.parse("Analytic{"*string(floattype)*"}")),floattype,Vector{floattype}}
        model = GP(X,y,k,Autotuning=true,verbose=2)
        @test train!(model,iterations=50)
        @test mean(abs2.(predict_y(model,X)-y)) < 0.2
    end
end