text:
#=
naive+elimination linear system solver
solves linear system in naive.jl via the following procedure

eliminate s
-G*x + h*tau - s = zrhs
so if using primal barrier
z_k + mu*H_k*s_k = srhs_k --> s_k = (mu*H_k)\(srhs_k - z_k)
-->
-G_k*x + (mu*H_k)\z_k + h_k*tau = zrhs_k + (mu*H_k)\srhs_k
(if use_inv_hess = true)
-->
-mu*H_k*G_k*x + z_k + mu*H_k*h_k*tau = mu*H_k*zrhs_k + srhs_k
(if use_inv_hess = false)
or if using dual barrier
mu*H_k*z_k + s_k = srhs_k --> s_k = srhs_k - mu*H_k*z_k
-->
-G_k*x + mu*H_k*z_k + h_k*tau = zrhs_k + srhs_k

eliminate kap
-c'x - b'y - h'z - kap = taurhs
so
mu/(taubar^2)*tau + kap = kaprhs --> kap = kaprhs - mu/(taubar^2)*tau
-->
-c'x - b'y - h'z + mu/(taubar^2)*tau = taurhs + kaprhs

4x4 nonsymmetric system in (x, y, z, tau):
A'*y + G'*z + c*tau = xrhs
-A*x + b*tau = yrhs
(pr bar) -mu*H_k*G_k*x + z_k + mu*H_k*h_k*tau = mu*H_k*zrhs_k + srhs_k
(du bar) -G_k*x + mu*H_k*z_k + h_k*tau = zrhs_k + srhs_k
-c'x - b'y - h'z + mu/(taubar^2)*tau = taurhs + kaprhs
=#

abstract type NaiveElimSystemSolver{T <: Real} <: SystemSolver{T} end

function solve_subsystem4(
    syssolver::NaiveElimSystemSolver{T},
    solver::Solver{T},
    sol::Point{T},
    rhs::Point{T},
    ) where {T <: Real}
    rhs_sub = syssolver.rhs_sub
    dim4 = length(rhs_sub.vec)
    rhs_sub.vec .= rhs.vec[1:dim4]

    @inbounds for (k, cone_k) in enumerate(solver.model.cones)
        rhs_z_k = rhs.z_views[k]
        rhs_s_k = rhs.s_views[k]
        rhs_sub_z_k = rhs_sub.z_views[k]
        if Cones.use_dual_barrier(cone_k)
            # -G_k*x + mu*H_k*z_k + h_k*tau = zrhs_k + srhs_k
            rhs_sub_z_k .+= rhs_s_k
        elseif syssolver.use_inv_hess
            # -G_k*x + (mu*H_k)\z_k + h_k*tau = zrhs_k + (mu*H_k)\srhs_k
            Cones.inv_hess_prod!(rhs_sub_z_k, rhs_s_k, cone_k)
            rhs_sub_z_k .+= rhs_z_k
        else
            # -mu*H_k*G_k*x + z_k + mu*H_k*h_k*tau = mu*H_k*zrhs_k + srhs_k
            Cones.hess_prod!(rhs_sub_z_k, rhs_z_k, cone_k)
            rhs_sub_z_k .+= rhs_s_k
        end
    end
    # -c'x - b'y - h'z + kapbar/taubar*tau = taurhs + kaprhs
    rhs_sub.vec[end] += rhs.kap[]

    sol_sub = syssolver.sol_sub
    solve_inner_system(syssolver, sol_sub, rhs_sub)
    sol.vec[1:dim4] .= sol_sub.vec

    return sol
end

function setup_point_sub(
    syssolver::NaiveElimSystemSolver{T},
    model::Models.Model{T},
    ) where {T <: Real}
    syssolver.sol_sub = Point{T}()
    syssolver.rhs_sub = Point{T}()
    z_start = model.n + model.p
    dim_sub = size(syssolver.lhs_sub, 1)
    for point_sub in (syssolver.sol_sub, syssolver.rhs_sub)
        point_sub.vec = zeros(T, dim_sub)
        point_sub.z_views = [view(point_sub.vec, z_start .+ idxs) for
            idxs in model.cone_idxs]
    end
    return
end

#=
direct sparse
=#

mutable struct NaiveElimSparseSystemSolver{T <: Real} <: NaiveElimSystemSolver{T}
    use_inv_hess::Bool
    lhs_sub::SparseMatrixCSC{T}
    fact_cache::SparseNonSymCache{T}
    rhs_sub::Point{T}
    sol_sub::Point{T}
    hess_idxs::Vector
    function NaiveElimSparseSystemSolver{T}(;
        use_inv_hess::Bool = true,
        fact_cache::SparseNonSymCache{T} = SparseNonSymCache{T}(),
        ) where {T <: Real}
        syssolver = new{T}()
        if !use_inv_hess
            @warn("SymIndefSparseSystemSolver is not implemented with " *
                "`use_inv_hess` set to `false`, using `true` instead.")
        end
        syssolver.use_inv_hess = use_inv_hess
        syssolver.fact_cache = fact_cache
        return syssolver
    end
end

function load(
    syssolver::NaiveElimSparseSystemSolver{T},
    solver::Solver{T},
    ) where {T <: Real}
    syssolver.fact_cache.analyzed = false
    model = solver.model
    (n, p, q) = (model.n, model.p, model.q)
    cones = model.cones
    cone_idxs = model.cone_idxs

    # form sparse LHS without Hessians and inverse Hessians in z/z block
    # TODO check for inefficiency, may need to implement more manually
    spz(a, b) = spzeros(T, a, b)
    lhs_sub = hvcat((4, 3, 3, 4),
        spz(n, n), model.A', model.G', model.c,
        -model.A, spz(p, p + q), model.b,
        -model.G, spz(q, p + q), model.h,
        -model.c', -model.b', -model.h', 1)
    @assert lhs_sub isa SparseMatrixCSC{T}
    dropzeros!(lhs_sub)
    (Is, Js, Vs) = findnz(lhs_sub)

    # add I, J, V for Hessians and inverse Hessians
    if isempty(cones)
        hess_nz_total = 0
    else
        hess_nz_total = sum(Cones.use_dual_barrier(cone_k) ?
            Cones.hess_nz_count(cone_k) : Cones.inv_hess_nz_count(cone_k)
            for cone_k in cones)
    end
    H_Is = Vector{Int}(undef, hess_nz_total)
    H_Js = Vector{Int}(undef, hess_nz_total)
    offset = 1
    for (cone_k, idxs_k) in zip(cones, cone_idxs)
        z_start_k = n + p + first(idxs_k) - 1
        for j in 1:Cones.dimension(cone_k)
            nz_rows_kj = z_start_k .+ (Cones.use_dual_barrier(cone_k) ?
                Cones.hess_nz_idxs_col(cone_k, j) :
                Cones.inv_hess_nz_idxs_col(cone_k, j))
            len_kj = length(nz_rows_kj)
            IJV_idxs = offset:(offset + len_kj - 1)
            offset += len_kj
            @. H_Is[IJV_idxs] = nz_rows_kj
            @. H_Js[IJV_idxs] = z_start_k + j
        end
    end
    append!(Is, H_Is)
    append!(Js, H_Js)
    append!(Vs, ones(T, hess_nz_total))

    # convert integer types here rather than inside external wrappers
    # integer type supported by the sparse system solver library to be used
    Ti = int_type(syssolver.fact_cache)
    Is = convert(Vector{Ti}, Is)
    Js = convert(Vector{Ti}, Js)
    dim = size(lhs_sub, 1)
    lhs_sub = syssolver.lhs_sub = sparse(Is, Js, Vs, dim, dim)

    # cache indices of nonzeros of (inv) Hessians in sparse LHS nonzeros vector
    syssolver.hess_idxs = [Vector{Union{UnitRange, Vector{Int}}}(
        undef, Cones.dimension(cone_k)) for cone_k in cones]
    for (k, cone_k) in enumerate(cones)
        z_start_k = n + p + first(cone_idxs[k]) - 1
        for j in 1:Cones.dimension(cone_k)
            col = z_start_k + j
            # get nonzero rows in the current column of the LHS
            col_idx_start = lhs_sub.colptr[col]
            nz_rows = lhs_sub.rowval[col_idx_start:(lhs_sub.colptr[col + 1] - 1)]
            # get nonzero rows in column j of the Hessian or inverse Hessian
            nz_hess_indices = (Cones.use_dual_barrier(cone_k) ?
                Cones.hess_nz_idxs_col(cone_k, j) :
                Cones.inv_hess_nz_idxs_col(cone_k, j))
            # get index corresponding to first nonzero element of the current col
            first_H = findfirst(isequal(z_start_k + first(nz_hess_indices)),
                nz_rows)
            # indices of nonzero values for cone k column j
            syssolver.hess_idxs[k][j] = (col_idx_start + first_H - 2) .+
                (1:length(nz_hess_indices))
        end
    end

    setup_point_sub(syssolver, model)

    return syssolver
end

function update_lhs(syssolver::NaiveElimSparseSystemSolver, solver::Solver)
    for (k, cone_k) in enumerate(solver.model.cones)
        H_k = (Cones.use_dual_barrier(cone_k) ? Cones.hess(cone_k) :
            Cones.inv_hess(cone_k))
        for j in 1:Cones.dimension(cone_k)
            nz_rows = (Cones.use_dual_barrier(cone_k) ?
                Cones.hess_nz_idxs_col(cone_k, j) :
                Cones.inv_hess_nz_idxs_col(cone_k, j))
            @views copyto!(syssolver.lhs_sub.nzval[
                syssolver.hess_idxs[k][j]], H_k[nz_rows, j])
        end
    end
    tau = solver.point.tau[]
    syssolver.lhs_sub.nzval[end] = solver.mu / tau / tau

    solver.time_upfact += @elapsed update_fact(syssolver.fact_cache,
        syssolver.lhs_sub)

    return syssolver
end

function solve_inner_system(
    syssolver::NaiveElimSparseSystemSolver,
    sol::Point,
    rhs::Point,
    )
    inv_prod(syssolver.fact_cache, sol.vec, syssolver.lhs_sub, rhs.vec)
    return sol
end

#=
direct dense
=#

mutable struct NaiveElimDenseSystemSolver{T <: Real} <: NaiveElimSystemSolver{T}
    use_inv_hess::Bool
    lhs_sub::Matrix{T}
    lhs_sub_fact::Matrix{T}
    fact::LU{T, Matrix{T}}
    rhs_sub::Point{T}
    sol_sub::Point{T}
    function NaiveElimDenseSystemSolver{T}(;
        use_inv_hess::Bool = true,
        ) where {T <: Real}
        syssolver = new{T}()
        syssolver.use_inv_hess = use_inv_hess
        return syssolver
    end
end

function load(
    syssolver::NaiveElimDenseSystemSolver{T},
    solver::Solver{T},
    ) where {T <: Real}
    model = solver.model
    (n, p, q) = (model.n, model.p, model.q)

    # TODO check for inefficiency, may need to implement more manually
    dz(a, b) = zeros(T, a, b)
    lhs_sub = hvcat((4, 3, 3, 4),
        dz(n, n), model.A', model.G', model.c,
        -model.A, dz(p, p + q), model.b,
        -model.G, dz(q, p + q), model.h,
        -model.c', -model.b', -model.h', 1)
    @assert lhs_sub isa Matrix{T}
    syssolver.lhs_sub = lhs_sub
    syssolver.lhs_sub_fact = zero(lhs_sub)

    setup_point_sub(syssolver, model)

    return syssolver
end

function update_lhs(
    syssolver::NaiveElimDenseSystemSolver{T},
    solver::Solver{T},
    ) where {T <: Real}
    model = solver.model
    (n, p) = (model.n, model.p)
    lhs_sub = syssolver.lhs_sub

    @inbounds for (cone_k, idxs_k) in zip(model.cones, model.cone_idxs)
        z_rows_k = (n + p) .+ idxs_k
        if Cones.use_dual_barrier(cone_k)
            # -G_k*x + mu*H_k*z_k + h_k*tau = zrhs_k + srhs_k
            @views copyto!(lhs_sub[z_rows_k, z_rows_k], Cones.hess(cone_k))
        elseif syssolver.use_inv_hess
            # -G_k*x + (mu*H_k)\z_k + h_k*tau = zrhs_k + (mu*H_k)\srhs_k
            Hi_k = @views copyto!(lhs_sub[z_rows_k, z_rows_k],
                Cones.inv_hess(cone_k))
        else
            # -mu*H_k*G_k*x + z_k + mu*H_k*h_k*tau = mu*H_k*zrhs_k + srhs_k
            @views Cones.hess_prod!(lhs_sub[z_rows_k, 1:n],
                model.G[idxs_k, :], cone_k)
            @. lhs_sub[z_rows_k, 1:n] *= -1
            @views Cones.hess_prod!(lhs_sub[z_rows_k, end],
                model.h[idxs_k], cone_k)
        end
    end
    tau = solver.point.tau[]
    lhs_sub[end, end] = solver.mu / tau / tau

    solver.time_upfact += @elapsed syssolver.fact =
        nonsymm_fact_copy!(syssolver.lhs_sub_fact, syssolver.lhs_sub)

    if !issuccess(syssolver.fact)
        println("nonsymmetric linear system factorization failed")
    end

    return syssolver
end

function solve_inner_system(
    syssolver::NaiveElimDenseSystemSolver,
    sol::Point,
    rhs::Point,
    )
    ldiv!(sol.vec, syssolver.fact, rhs.vec)
    return sol
end


text:
struct Location
  ii::Vector{Union{Int,Symbol}}
end

Location() = Location([])

parent(l::Location) = Location(l.ii[1:end-1])
Base.getindex(l::Location, i...) = Location([l.ii..., i...])

Base.getindex(ex::EXPR, i::Integer) = ex.args[i]
Base.getindex(ex::AbstractEXPR, i) = collect(ex)[i]

function Base.getindex(ex::AbstractEXPR, l::Location)
  for i in l.ii
    ex = i isa Integer ? ex[i] : getproperty(ex, i)
  end
  return ex
end

# HACK
index_in_collected(ex, k) = findfirst(x -> x === getproperty(ex, k), collect(ex))

function charrange(ex::AbstractEXPR, l::Location)
    o = 0
    for i in l.ii
      i isa Symbol && (i = index_in_collected(ex, i))
      for j = 1:i-1
        o += ex[j].fullspan
      end
      ex = ex[i]
    end
    full = o .+ (1:ex.fullspan)
    inner = o .+ (1:ex.span)
    return full, inner
end

ext: jl


text:
# © 2016 Massachusetts Institute of Technology.  See LICENSE file for details.
abstract Instance
Base.keys(i::Instance) = keys(i._storage)
Base.values(i::Instance) = values(i._storage)
Base.getindex(i::Instance, a::ANode) = i._storage[a]

@deprecate evaluator(x...) instance(x...)

instance(n::ANode...) = instance(collect(ANode, n))
instance(T::Type, n::ANode...) = instance(T, collect(ANode, n))
instance(pool::StoragePool, n::ANode...) = instance(pool, collect(ANode, n))

instance{A<:ANode}(n::Vector{A}) = instance(Array{Float64}, n)
instance{A<:ANode}(T::Type, n::Vector{A}) = instance(NullPool(T), n)

output_node(n::ANode) = mutates(n) ? output_node(n.input[1]) : n

type RefCounter
  count::Int
  RefCounter() = new(0)
end

incr(r::RefCounter) = r.count += 1
decr(r::RefCounter) = (@assert r.count > 0; r.count -= 1; r.count == 0)
iszero(r::RefCounter) = r.count == 0

incr{T}(d::Dict{T,RefCounter}, x::T) = incr(get!(RefCounter, d, x))
decr{T}(d::Dict{T,RefCounter}, x::T) = decr(d[x])
iszero{T}(d::Dict{T,RefCounter}, x::T) = iszero(d[x])

function generate_function(sorted::Vector{ANode}, dict)
  b = :(begin end)

  if length(sorted) >= 20
    n = length(sorted)
    for i in 1:8
      ix = div((i-1)*n,8)+1:div(i*n,8)
      f = generate_function(sorted[ix], dict)
      push!(b.args, :($f()::Void))
    end
  else
    for i in sorted
      if sym(i) == :input || sym(i) == :const
        continue
      elseif sym(i) == :load
        if dict[i] != arg(i)
          push!(b.args, :(copy!($(dict[i]), $(arg(i)))))
        end
      else
        in = map(x->dict[x], i.input)
        if mutates(i)
          push!(b.args, :(Seep.do_forward!($i, $(in...))))
        else
          out = dict[i]
          push!(b.args, :(Seep.do_forward!($i, $out, $(in...))))
        end
      end
    end
  end

  push!(b.args, :nothing)

  @gensym SeepFunction
  @eval @compat $SeepFunction() = $b

  return SeepFunction
end

function instance{A<:ANode}(pool::StoragePool, n::Vector{A})
  sorted = toposort(n)

  T = atype(pool)
  dict = Dict{ANode,T}()
  userdict = Dict{ANode,T}()

  refcnt = Dict{ANode,RefCounter}()
  for a in n
    if sym(a) == :noop continue end
    incr(refcnt, output_node(a))
  end

  for a in sorted
    for b in a.input
      incr(refcnt, output_node(b))
    end
  end

  constants = ANode[]
  c_block = :(begin end)
  for i in sorted
    if sym(i) == :input
      userdict[i] = dict[i] = allocate(pool, i.size...)
    elseif sym(i) == :const
      if isa(arg(i), T)
        dict[i] =  arg(i)
      else
        incr(refcnt, i)
        t = dict[i] = allocate(pool, i.size...)
        push!(c_block.args, :(copy!($t, map(eltype($T), $(arg(i))))))
        push!(constants, i)
      end
    end
  end

  @gensym initialize_constants
  @eval $(initialize_constants)() = $c_block

  for i in sorted
    @assert !iszero(refcnt, output_node(i))

    if mutates(i)
      dict[i] = dict[i.input[1]]
    elseif sym(i) == :load && isa(arg(i), T)
      dict[i] = arg(i)
    elseif sym(i) != :input && sym(i) != :const
      dict[i] = allocate(pool, i.size...)
    end

    for a in i.input
      aa = output_node(a)
      if decr(refcnt, aa) && arg(aa) != dict[aa]
        deallocate(pool, dict[aa])
      end
    end
  end

  for i in n
    if sym(i) == :noop continue end
    userdict[i] = dict[i]

    aa = output_node(i)
    if decr(refcnt, aa) && arg(aa) != dict[aa]
      deallocate(pool, dict[aa])
    end
  end

  for i in constants
    decr(refcnt, i)
    deallocate(pool, dict[i])
  end

  for i in keys(refcnt)
    @assert iszero(refcnt, i)
  end

  local b
  if isempty(constants)
    b = :(begin end)
  else
    b = :(begin
      if $(pool).token[1] != $(QuoteNode(initialize_constants))
        $(pool).token[1] = $(QuoteNode(initialize_constants))
        $(initialize_constants)()
      end
    end)
  end

  fun = generate_function(sorted, dict)
  push!(b.args, :($fun()::Void))

  @gensym SeepInstance
  typeblock = quote
    immutable $SeepInstance <: Instance
      _storage::Dict{ANode,$T}
      _pool::StoragePool
    end
    Base.show(io::IO, ::$SeepInstance) = print(io, "Seep Instance with $($(length(sorted))) nodes.")
    @compat (::$SeepInstance)() = $b
    $SeepInstance($userdict, $pool)
  end

  for (n,a) in userdict
    if haskey(userdict, n) && n.name != ""
      if length(filter(x->x.name == n.name, sorted)) > 1
        warn("There are multiple nodes named $(n.name).")
        continue
      end
      push!(typeblock.args[2].args[3].args, :($(symbol(n.name))::$(typeof(a))))
      push!(typeblock.args[end].args, a)
    end
  end

  eval(typeblock)
end

# make sure run,dict = instance(...) works for a while, even though run and dict are the same thing
#Base.start(::Instance) = (Base.depwarn("replace run,dict=instance(...) with inst=instance(...)", :SeepStart); 1)
Base.start(::Instance) = 1
Base.done(::Instance, i::Int) = i > 2
Base.next(a::Instance, i::Int) = (a, i+1)


text:
using Test
using MPI
using Random

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CuArrays
    ArrayType = CuArray
else
    ArrayType = Array
end

MPI.Init()

comm = MPI.COMM_WORLD

function bcast_array(A, root)

end

root = 0
Random.seed!(17)
matsize = (17,17)

for T in Base.uniontypes(MPI.MPIDatatype)
    A = ArrayType(rand(T, matsize))
    B = MPI.Comm_rank(comm) == root ? A : similar(A)
    MPI.Bcast!(B, root, comm)
    @test B == A
end

# Char
A = ['s', 't', 'a', 'r', ' ', 'w', 'a', 'r', 's']
B = MPI.Comm_rank(comm) == root ? A : similar(A)
MPI.Bcast!(B, root, comm)
@test B == A

g = x -> x^2 + 2x - 1
if MPI.Comm_rank(comm) == root
    f = g
else
    f = nothing
end
f = MPI.bcast(f, root, comm)
@test f(3) == g(3)
@test f(5) == g(5)
@test f(7) == g(7)

A = Dict("foo" => "bar")
if MPI.Comm_rank(comm) == root
    B = A
else
    B = nothing
end
B = MPI.bcast(B, root, comm)
@test B["foo"] == "bar"

MPI.Finalize()
@test MPI.Finalized()


text:
function frenet_serret(coordinate::Matrix{T_Number}) where {T_Number<:Number}
    r = copy(coordinate)
    n_points = size(r, 1)
    if size(r, 2) == 2
        r = cat(r, zeros(eltype(r), size(r, 1)), dims = 2)
    end
    t = 1.0:size(r,1)
    dt = (maximum(t) - minimum(t)) / length(r)
    dr = gradient(r)
    ddr = gradient(dr)
    dddr = gradient(ddr)

    T = similar(dr)
    N = similar(dr)
    B = similar(dr)

    for i = 1:n_points

        T[i, :] = LinearAlgebra.normalize(dr[i, :])

        B[i, :] = LinearAlgebra.normalize(cross(dr[i, :], ddr[i, :]))

        N[i, :] = LinearAlgebra.normalize(cross(B[i, :], T[i, :]))

    end

    ds = Array{eltype(dr),1}(undef, n_points)
    curvature = similar(ds)
    torsion = similar(ds)

    for i in eachindex(ds)
        selected_dr = dr[i, :]
        selected_ddr = ddr[i, :]
        selected_dddr = dddr[i, :]
        ds[i] = norm(selected_dr)

        dr_cross_ddr = cross(selected_dr, selected_ddr)
        curvature_num = norm(dr_cross_ddr) # || dr X ddr ||
        curvature_denom = norm(selected_dr)^3 # || dr ||^3
        curvature[i] = curvature_num / curvature_denom

        torsion_num = dot(dr_cross_ddr, selected_dddr) # (dr X ddr) . dddr
        torsion_denom = norm(dr_cross_ddr)^2 # || dr X ddr ||^2
        torsion[i] = torsion_num / torsion_denom
    end
    speed = ds ./ dt
    return T,N,B,curvature,torsion,speed
end


text:
# Copyright 2015 Conrad Scott

import Base.length

immutable Dictionary{S <: AbstractString, D <: AbstractVector}
  d :: D

  function Dictionary(d :: AbstractVector{S})
    new(d)
  end
end

function Dictionary(d :: AbstractVector)
  Dictionary{eltype(d), typeof(d)}(d)
end

function length(d :: Dictionary)
  length(d.d)
end

function words(d :: Dictionary)
  length(d)
end


text:
module Observables

export Observable, on, off, onany, connect!, obsid, async_latest, throttle

import Base.Iterators.filter

# @nospecialize "blocks" codegen but not necessarily inference. This forces inference
# to drop specific information about an argument.
if isdefined(Base, :inferencebarrier)
    const inferencebarrier = Base.inferencebarrier
else
    inferencebarrier(x) = Ref{Any}(x)[]
end

const addhandler_callbacks = []
const removehandler_callbacks = []

abstract type AbstractObservable{T} end

function observe(::S) where {S<:AbstractObservable}
    error("observe not defined for AbstractObservable $S")
end

"""
    obs = Observable(val)
    obs = Observable{T}(val)

Like a `Ref`, but updates can be watched by adding a handler using [`on`](@ref) or [`map`](@ref).
"""
mutable struct Observable{T} <: AbstractObservable{T}
    listeners::Vector{Any}
    val::T
    inputs::Vector{Any}  # for map!ed Observables

    Observable{T}() where {T} = new{T}([])
    Observable{T}(val) where {T} = new{T}([], val)
    # Construct an Observable{Any} without runtime dispatch
    Observable{Any}(@nospecialize(val)) = new{Any}([], val)
end

function Base.copy(observable::Observable{T}) where T
    result = Observable{T}(observable[])
    on(observable) do value
        result[] = value
    end
    return result
end

Observable(val::T) where {T} = Observable{T}(val)

Base.eltype(::AbstractObservable{T}) where {T} = T

observe(x::Observable) = x

function Base.convert(::Type{Observable{T}}, x::AbstractObservable) where {T}
    result = Observable{T}(convert(T, x[]))
    on(x) do value
        result[] = convert(T, value)
    end
    return result
end

Base.convert(::Type{T}, x::T) where {T<:Observable} = x  # resolves ambiguity with convert(::Type{T}, x::T) in base/essentials.jl
Base.convert(::Type{T}, x) where {T<:Observable} = T(x)

function Base.getproperty(obs::Observable, field::Symbol)
    if field === :val
        return getfield(obs, field)
    elseif field === :listeners
        return getfield(obs, field)
    elseif field === :inputs
        return getfield(obs, field)
    elseif field === :id
        return obsid(obs)
    else
        error("Field $(field) not found")
    end
end

"""
    notify(observable::AbstractObservable)

Update all listeners of `observable`.
"""
function Base.notify(observable::AbstractObservable)
    val = observable[]
    for f in listeners(observable)
        Base.invokelatest(f, val)
    end
    return
end

function Base.show(io::IO, x::Observable{T}) where T
    println(io, "Observable{$T} with $(length(x.listeners)) listeners. Value:")
    if isdefined(x, :val)
        show(io, x.val)
    else
        print(io, "not assigned yet!")
    end
end

Base.show(io::IO, ::MIME"application/prs.juno.inline", x::Observable) = x

"""
    mutable struct ObserverFunction <: Function

Fields:

    f::Function
    observable::AbstractObservable
    weak::Bool

`ObserverFunction` is intended as the return value for `on` because
we can remove the created closure from `obsfunc.observable`'s listener vectors when
ObserverFunction goes out of scope - as long as the `weak` flag is set.
If the `weak` flag is not set, nothing happens
when the ObserverFunction goes out of scope and it can be safely ignored.
It can still be useful because it is easier to call `off(obsfunc)` instead of `off(observable, f)`
to release the connection later.
"""
mutable struct ObserverFunction <: Function
    f
    observable::AbstractObservable
    weak::Bool

    function ObserverFunction(@nospecialize(f), observable::AbstractObservable, weak)
        obsfunc = new(f, observable, weak)

        # If the weak flag is set, deregister the function f from the observable
        # storing it in its listeners once the ObserverFunction is garbage collected.
        # This should free all resources associated with f unless there
        # is another reference to it somewhere else.
        if weak
            finalizer(off, obsfunc)
        end

        obsfunc
    end
end

Base.precompile(obsf::ObserverFunction) = precompile(obsf.f, (eltype(obsf.observable),))
function Base.precompile(observable::Observable)
    tf = true
    T = eltype(observable)
    for f in observable.listeners
        precompile(f, (T,))
    end
    if isdefined(observable, :inputs)
        for obsf in observable.inputs
            tf &= precompile(obsf)
        end
    end
    return tf
end

"""
    on(f, observable::AbstractObservable; weak = false)

Adds function `f` as listener to `observable`. Whenever `observable`'s value
is set via `observable[] = val`, `f` is called with `val`.

Returns an [`ObserverFunction`](@ref) that wraps `f` and `observable` and allows to
disconnect easily by calling `off(observerfunction)` instead of `off(f, observable)`.
If instead you want to compute a new `Observable` from an old one, use [`map(f, ::Observable)`](@ref).

If `weak = true` is set, the new connection will be removed as soon as the returned `ObserverFunction`
is not referenced anywhere and is garbage collected. This is useful if some parent object
makes connections to outside observables and stores the resulting `ObserverFunction` instances.
Then, once that parent object is garbage collected, the weak
observable connections are removed automatically.

# Example

```jldoctest; setup=:(using Observables)
julia> obs = Observable(0)
Observable{$Int} with 0 listeners. Value:
0

julia> on(obs) do val
           println("current value is ", val)
       end
(::Observables.ObserverFunction) (generic function with 0 methods)

julia> obs[] = 5;
current value is 5
```
"""
function on(@nospecialize(f), observable::AbstractObservable; weak::Bool = false)
    push!(listeners(observable), f)
    for g in addhandler_callbacks
        g(f, observable)
    end
    # Return a ObserverFunction so that the caller is responsible
    # to keep a reference to it around as long as they want the connection to
    # persist. If the ObserverFunction is garbage collected, f will be released from
    # observable's listeners as well.
    return ObserverFunction(f, observable, weak)
end

"""
    off(observable::AbstractObservable, f)

Removes `f` from listeners of `observable`.

Returns `true` if `f` could be removed, otherwise `false`.
"""
function off(observable::AbstractObservable, @nospecialize(f))
    callbacks = listeners(observable)
    for (i, f2) in enumerate(callbacks)
        if f === f2
            deleteat!(callbacks, i)
            for g in removehandler_callbacks
                g(observable, f)
            end
            return true
        end
    end
    return false
end

function off(observable::AbstractObservable, obsfunc::ObserverFunction)
    f = obsfunc.f
    # remove the function inside obsfunc as usual
    off(observable, f)
end

"""
    off(obsfunc::ObserverFunction)

Remove the listener function `obsfunc.f` from the listeners of `obsfunc.observable`.
Once `obsfunc` goes out of scope, this should allow `obsfunc.f` and all the values
it might have closed over to be garbage collected (unless there
are other references to it).
"""
function off(obsfunc::ObserverFunction)
    off(obsfunc.observable, obsfunc)
end

"""
    observable[] = val

Updates the value of an `Observable` to `val` and call its listeners.
"""
function Base.setindex!(observable::Observable, val)
    observable.val = val
    notify(observable)
    return val
end

# For external packages that don't want to access an internal field
setexcludinghandlers!(obs::AbstractObservable, val) = observe(obs).val = val

####################################################
# tasks & channel api

Observable(val::Channel{T}) where {T} = Observable{T}(val)
Observable(val::Task) = Observable{Any}(val)
function Observable{Any}(val::Union{Task, Channel})   # ambiguity resolution
    observable = Observable{Any}()
    observable[] = val
    return observable
end
function Observable{T}(val::Union{Task, Channel}) where {T}
    observable = Observable{T}()
    observable[] = val
    return observable
end

function Base.setindex!(observable::Observable, val_async::Task)
    return @async begin
        try
            val = fetch(val_async)
            setindex!(observable, val)
        catch e
            Base.showerror(stderr, e)
            Base.show_backtrace(stderr, catch_backtrace())
        end
    end
end

function Base.setindex!(observable::Observable, channel::Channel)
    return @async begin
        try
            for val in channel
                setindex!(observable, val)
                yield()
            end
        catch e
            Base.showerror(stderr, e)
            Base.show_backtrace(stderr, catch_backtrace())
        end
    end
end

function Base.setindex!(observable::AbstractObservable, val)
    Base.setindex!(observe(observable), val)
end

"""
    observable[]

Returns the current value of `observable`.
"""
Base.getindex(observable::Observable) = observable.val

Base.getindex(observable::AbstractObservable) = getindex(observe(observable))

### Utilities

"""
    to_value(x::Union{Any, AbstractObservable})
Extracts the value of an observable, and returns the object if it's not an observable!
"""
to_value(x) = isa(x, AbstractObservable) ? x[] : x  # noninferrable dispatch is faster if there is only one Method

"""
    obsid(observable::Observable)

Gets a unique id for an observable.
"""
obsid(observable::Observable) = string(objectid(observable))
obsid(observable::AbstractObservable) = obsid(observe(observable))

listeners(observable::Observable) = observable.listeners
listeners(observable::AbstractObservable) = listeners(observe(observable))

struct OnUpdate{F, Args} <: Function
    f::F
    args::Args
end

(ou::OnUpdate)(_) = ou.f(map(to_value, ou.args)...)

"""
    onany(f, args...)

Calls `f` on updates to any observable refs in `args`.
`args` may contain any number of `Observable` objects.
`f` will be passed the values contained in the refs as the respective argument.
All other objects in `args` are passed as-is.

See also: [`on`](@ref).
"""
function onany(f::F, args...; weak::Bool = false) where F
    callback = OnUpdate(f, args)
    _onany(inferencebarrier(callback), args, weak)
end

@noinline function _onany(@nospecialize(callback), args, weak::Bool)
    # store all returned ObserverFunctions
    obsfuncs = ObserverFunction[]
    for observable in args
        if observable isa AbstractObservable
            obsfunc = on(callback, observable, weak = weak)
            push!(obsfuncs, obsfunc)
        end
    end

    # same principle as with `on`, this collection needs to be
    # stored by the caller or the connections made will be cut
    obsfuncs
end

struct MapUpdater{F, T} <: Function
    f::F
    observable::Observable{T}
end

function (mu::MapUpdater)(args...)
    mu.observable[] = mu.f(args...)
end

"""
    map!(f, observable::AbstractObservable, args...; update::Bool=true)

Updates `observable` with the result of calling `f` with values extracted from args.
`args` may contain any number of `Observable` objects.
`f` will be passed the values contained in the refs as the respective argument.
All other objects in `args` are passed as-is.

By default `observable` gets updated immediately, but this can be suppressed by specifying `update=false`.

# Example

We'll create an observable that can hold an arbitrary number:

```jldoctest map!; setup=:(using Observables)
julia> obs = Observable{Number}(3)
Observable{Number} with 0 listeners. Value:
3
```

Now,

```jldoctest map!
julia> obsrt1 = map(sqrt, obs)
Observable{Float64} with 0 listeners. Value:
1.7320508075688772
```

creates an `Observable{Float64}`, which will fail to update if we set `obs[] = 3+4im`.
However,

```jldoctest map!
julia> obsrt2 = map!(sqrt, Observable{Number}(), obs)
Observable{Number} with 0 listeners. Value:
1.7320508075688772
```

can handle any number type for which `sqrt` is defined.
"""
@inline function Base.map!(f::F, observable::AbstractObservable, os...; update::Bool=true) where F
    # note: the @inline prevents de-specialization due to the splatting
    obsfuncs = onany(MapUpdater(f, observable), os...)
    appendinputs!(observable, obsfuncs)
    if update
        observable[] = f(map(to_value, os)...)
    end
    return observable
end

function appendinputs!(observable, obsfuncs)  # latency: separating this from map! allows dropping the specialization on `f`
    if !isdefined(observable, :inputs)
        observable.inputs = obsfuncs
    else
        append!(observable.inputs, obsfuncs)
    end
    return observable
end

"""
    connect!(o1::AbstractObservable, o2::AbstractObservable)

Forwards all updates from `o2` to `o1`.

See also [`Observables.ObservablePair`](@ref).
"""
connect!(o1::AbstractObservable, o2::AbstractObservable) = map!(identity, o1, o2)

"""
    obs = map(f, arg1::AbstractObservable, args...)

Creates a new observable ref `obs` which contains the result of `f` applied to values
extracted from `arg1` and `args` (i.e., `f(arg1[], ...)`.
`arg1` must be an observable ref for dispatch reasons. `args` may contain any number of `Observable` objects.
`f` will be passed the values contained in the refs as the respective argument.
All other objects in `args` are passed as-is.

If you don't need the value of `obs`, and just want to run `f` whenever the
arguments update, use [`on`](@ref) or [`onany`](@ref) instead.

# Example

```jldoctest; setup=:(using Observables)
julia> obs = Observable([1,2,3]);

julia> map(length, obs)
Observable{$Int} with 0 listeners. Value:
3
```
"""
@inline function Base.map(f::F, arg1::AbstractObservable, args...; kwargs...) where F
    # note: the @inline prevents de-specialization due to the splatting
    if haskey(kwargs, :init)
        Base.depwarn("""
            the `init` keyword is deprecated, and the new implementation does not force `Observable{Any}` (it will use the output of `f`).
            Instead of `map(f, args...; init=initval)`, use
                `map!(f, Observable{T}(initval), args...; update=false)`
            To control just the eltype, use
                `map!(f, Observable{T}(), args...)`.
            """, :map)
        dest = Observable{Any}(kwargs[:init])
        return map!(f, dest, arg1, args...; update=false)
    end
    map!(f, Observable(f(arg1[], map(to_value, args)...)), arg1, args...; update=false)
end

"""
    async_latest(observable::AbstractObservable, n=1)

Returns an `Observable` which drops all but
the last `n` updates to `observable` if processing the updates
takes longer than the interval between updates.

This is useful if you want to pass the updates from,
say, a slider to a plotting function that takes a while to
compute. The plot will directly compute the last frame
skipping the intermediate ones.

# Example:
```
observable = Observable(0)
function compute_something(x)
    for i=1:10^8 rand() end # simulate something expensive
    println("updated with \$x")
end
o_latest = async_latest(observable, 1)
on(compute_something, o_latest) # compute something on the latest update

for i=1:5
    observable[] = i
end
```
"""
function async_latest(input::AbstractObservable{T}, n=1) where T
    buffer = T[]
    cond = Condition()
    lck  = ReentrantLock() # advisory lock for access to buffer
    output = Observable{T}(input[]) # output

    @async while true
        while true # while !isempty(buffer) but with a lock
            # transact a pop
            lock(lck)
            if isempty(buffer)
                unlock(lck)
                break
            end
            upd = pop!(buffer)
            unlock(lck)

            output[] = upd
        end
        wait(cond)
    end

    on(input) do val
        lock(lck)
        if length(buffer) < n
            push!(buffer, val)
        else
            while length(buffer) >= n
                pop!(buffer)
            end
            pushfirst!(buffer, val)
        end
        unlock(lck)
        notify(cond)
    end

    output
end

# TODO: overload broadcast on v0.6
include("observablepair.jl")
include("flatten.jl")
include("time.jl")
include("macros.jl")

# Look up the source location of `do` block Observable MethodInstances
function methodlist(@nospecialize(ft::Type))
    if ft <: OnUpdate
        ft = Base.unwrap_unionall(Base.unwrap_unionall(ft).parameters[1])
        if ft <: MapUpdater
            ft = Base.unwrap_unionall(Base.unwrap_unionall(ft).parameters[1])
        end
    end
    return Base.MethodList(ft.name.mt)
end

methodlist(mi::Core.MethodInstance) = methodlist(Base.unwrap_unionall(mi.specTypes).parameters[1])
methodlist(obsf::ObserverFunction) = methodlist(obsf.f)
methodlist(@nospecialize(f::Function)) = methodlist(typeof(f))

@deprecate notify! notify

end # module


text:
function get_samples(io::IOStream, n_samples::Integer)
    sample_header_length = read(io, UInt32)
    sample_n_check = read(io, UInt32)
    if sample_n_check != 0
        @assert n_samples == sample_n_check "Inconsistent number of samples"
    else
        @warn "Sample names unavailable. Do you have a separate '.sample' file?"
    end
    samples = String[]
    for i in 1:n_samples
        id_length = read(io, UInt16)
        id = read(io, id_length)
        push!(samples, String(id))
    end
    samples
end

function get_samples(path::String, n_samples::Integer)
    @assert endswith(path, ".sample") "Extension of the file should be .sample"
    io = open(path)
    readline(io) # header
    readline(io) # types
    samples = readlines(io)
    samples = map(x -> split(x, " ")[1], samples)
    @assert length(samples) == n_samples "Inconsistent number of samples"
    samples
end

function get_samples(n_samples::Integer)
    samples = [string(i) for i in 1:n_samples]
    samples
end


text:
abstract type AbstractUncertainDataset end

function summarise(ud::AbstractUncertainDataset)
    _type = typeof(ud)
    summary = "$_type"
    return summary
end

Base.show(io::IO, ud::AbstractUncertainDataset) = println(io, summarise(ud))

export AbstractUncertainDataset


text: foo() {}

max_stars_repo_name: jwilk/jtc

max_issues_repo_name: jwilk/jtc

max_forks_repo_name: jwilk/jtc


text:
test_string = "BFFFBBFRRR
FFFBBBFRRR
BBFFBBFRLL"

test_list = split(test_string,'\n')

list = readlines("day5_input.txt")

function seat_to_id(seat)
    seat_binary = seat
    for (k,v) in zip(('F', 'B', 'L', 'R'),('0', '1', '0', '1'))
        seat_binary = replace(seat_binary, k => v)
    end
    id = parse(Int, seat_binary, base=2)
    return id
end

function id_to_seat(id)
    seat_binary = bitstring(id)[end-9:end]
    seat_row = replace(replace(seat_binary[1:end-3], '1' => 'B'), '0' => 'F')
    seat_column = replace(replace(seat_binary[end-2:end], '1' => 'R'), '0' => 'L')
    seat = "B"^(7-length(seat_row)) * seat_row * 'L'^(3-length(seat_column)) * seat_column
    return seat
end

highest_id(list) = maximum(seat_to_id.(list))

highest_seat(list) = id_to_seat(highest_id(list))

println("Highest id of the test list: $(highest_id(test_list))")

println("Highest id of the problem list: $(highest_id(list))")

function myseat(list)
    sorted_list = sort(seat_to_id.(list))
    min_id = sorted_list[1]
    diff_list = sorted_list[2:end] - sorted_list[1:end-1]
    max, id_skip = findmax(diff_list)
    return min_id + id_skip
end

println("Your seat is $(myseat(list))")

# Answers:
# P1: 926
# P2: 657

function seat_to_id2(seat)
    f(x) = (x =='F' || x == 'L') ? 0 : 1
    return parse(Int, join(map(f, collect(seat))), base=2)
end

highest_id2(list) = maximum(seat_to_id2.(list))

println("Highest id of the problem list v2: $(highest_id(list))")

# method 1 is still better
#=
```julia
julia> @btime highest_id(list)
  1.158 ms (13538 allocations: 747.02 KiB)
926

julia> @btime highest_id2(list)
  1.359 ms (22844 allocations: 1.21 MiB)
926
```
 =#

seatnumber(line) = mapfoldl(c->c in "BR", (x, y)->2x+y, line, init=0)

# wow, look at the differences
#=
```julia
julia> @btime seat_to_id(line)
  1.276 μs (16 allocations: 896 bytes)
181

julia> @btime seat_to_id2(line)
  1.579 μs (27 allocations: 1.45 KiB)
181

julia> @btime seatnumber(line)
  281.837 ns (0 allocations: 0 bytes)
181
```
 =#

 # Here is a one-liner, which is actually a two-liner due to `;`,
 # that was started by [braden](bradeneliason#2230), using
 # the one-liner `seatnumber` function, and improved upon
 # after discussions with AaronRMatthis and Tom:

function seats(list)
    S=sort(mapfoldl.(∈("BR"),(x,y)->2x+y,list))
    return S[end],S[argmax(diff(S))]+1
end

nothing

# Much faster than mine:
#=
```julia
julia> @btime seats(list_day5)
  309.319 μs (14 allocations: 20.52 KiB)
(926, 657)
```
 =#


text:
using Rimu
using Rimu.BitStringAddresses
using Rimu.BitStringAddresses: remove_ghost_bits, has_ghost_bits
using Rimu.BitStringAddresses: occupied_modes, update_component
using Random
using StaticArrays
using Test

include("excitation_tests.jl")

@testset "BitString" begin
    @testset "num_chunks for Val(B)" begin
        @test_throws ArgumentError num_chunks(Val(0))
        @test_throws ArgumentError num_chunks(Val(-15))
        @test num_chunks(Val(1)) == (1, UInt8)
        @test num_chunks(Val(8)) == (1, UInt8)
        @test num_chunks(Val(9)) == (1, UInt16)
        @test num_chunks(Val(16)) == (1, UInt16)
        @test num_chunks(Val(17)) == (1, UInt32)
        @test num_chunks(Val(32)) == (1, UInt32)
        @test num_chunks(Val(33)) == (1, UInt64)
        @test num_chunks(Val(64)) == (1, UInt64)
        @test num_chunks(Val(65)) == (2, UInt64)
        @test num_chunks(Val(128)) == (2, UInt64)
        @test num_chunks(Val(129)) == (3, UInt64)
        @test num_chunks(Val(200)) == (4, UInt64)
    end

    @testset "Constructors" begin
        @testset "Constructing from small ints" begin
            @test BitString{5}(1) ==
                BitString{5}(0x1) ==
                BitString{5}(UInt16(1)) ==
                BitString{5}(big"1") ==
                BitString{5}(Int128(1))

            @test BitString{129}(1) ==
                BitString{129}(0x1) ==
                BitString{129}(UInt16(1)) ==
                BitString{129}(big"1") ==
                BitString{129}(Int128(1))
        end
        @testset "Constructing from 128-bit int. Number of chunks depends on B." begin
            bs128a = @inferred BitString{129}(0x00000decafc0ffee0000000deadbeefe)
            @test bs128a.chunks[1] ≡ zero(UInt64)
            @test bs128a.chunks[2] ≡ 0xdecafc0ffee
            @test bs128a.chunks[3] ≡ 0xdeadbeefe
            @test bs128a == BitString{129}(big"0xdecafc0ffee0000000deadbeefe")

            bs128b = @inferred BitString{128}(0x00000decafc0ffee0000000deadbeefe)
            @test bs128b.chunks[1] ≡ 0x00000decafc0ffee
            @test bs128b.chunks[2] ≡ 0x0000000deadbeefe
            @test bs128b == BitString{128}(big"0xdecafc0ffee0000000deadbeefe")
        end
        @testset "Zero" begin
            a = BitString{100}(0)
            @test all(iszero, chunks(a))
            @test zero(a) == a

            b = BitString{250}(0x0)
            @test all(iszero, chunks(b))
            @test zero(b) == b

            c = BitString{16}(big"0")
            @test all(iszero, chunks(c))
            @test zero(typeof(c)) == c
        end
        @testset "no ghost bits" begin
            @test !has_ghost_bits(BitString{5}(0b11111))
            @test !has_ghost_bits(BitString{5}(0b111111))
            @test !has_ghost_bits(BitString{100}(rand(UInt128)))
            @test !has_ghost_bits(BitString{129}(big"0xffffffffffffffffffffffffffffffffffff"))
        end
    end
    @testset "Counting operations" begin
        # Create rand integer with n ones
        function rand_int_n_ones(T, N)
            result = zero(T)
            while count_ones(result) < N
                result |= 1 << rand(1:120)
            end
            return result
        end
        for N in (1, 2, 4, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 120)
            i = rand_int_n_ones(UInt128, N)
            j = i & (UInt128(1) << 120 - 1)
            s = BitString{120}(i)
            @test count_ones(s) == count_ones(j)
            @test count_zeros(s) == count_zeros(j) - 8
            @test leading_ones(s) == leading_ones(j << 8)
            @test trailing_ones(s) == trailing_ones(j)
            @test leading_zeros(s) == min(leading_zeros(j << 8), 120)
            @test trailing_zeros(s) == min(trailing_zeros(j), 120)
        end
        for N in (1, 2, 4, 10, 20, 30, 40, 50, 60)
            i = rand_int_n_ones(UInt64, N)
            j = i & (1 << 40 - 1)
            s = BitString{40}(i)
            @test count_ones(s) == count_ones(j)
            @test count_zeros(s) == count_zeros(j) - 24
            @test leading_ones(s) == leading_ones(j << 24)
            @test trailing_ones(s) == trailing_ones(j)
            @test leading_zeros(s) == min(leading_zeros(j << 24), 40)
            @test trailing_zeros(s) == min(trailing_zeros(j), 40)
        end
    end
    @testset "Bitwise operations" begin
        function rand_bitstring(B)
            N, T = num_chunks(Val(B))
            s = rand(SVector{N,T})
            return remove_ghost_bits(BitString{B,N,T}(s))
        end

        for B in (11, 32, 65, 120, 200, 256)
            a = rand_bitstring(B)
            b = rand_bitstring(B)
            one = ~zero(a)

            @test iszero(a & ~a)
            @test a | ~a == ~zero(a)
            @test a ⊻ a == zero(a)
            @test (a << 7) >> 7 == a & ~zero(BitString{B}) >> 7
            @test (a >> 9) << 9 == a & ~zero(BitString{B}) << 9
            @test (a >> 23) >> 41 == a >>> 64
            @test (a << 23) << 41 == a << 64
            @test a >> B == zero(a)
            @test a << B == zero(a)
            @test a << -13 == a >> 13
            @test a >> -17 == a << 17
            @test a >> 0 == a
            @test a << 0 == a

            @test (a >> 13) << 13 == a & (one << 13)
            @test (a << 13) >> 13 == a & (one >> 13)

            @test !has_ghost_bits(~zero(a))
            @test !has_ghost_bits(a << 2)
            @test !has_ghost_bits(a | b)
            @test !has_ghost_bits(a ⊻ b)
            @test !has_ghost_bits(a & b)
            @test !has_ghost_bits(a >> 2)

            # This test checks that the result of show can be pasted into the REPL
            @test eval(Meta.parse(repr(a))) == a
        end
    end
end

using Rimu.Hamiltonians: num_occupied_modes, bose_hubbard_interaction, hopnextneighbour

@testset "BoseFS" begin
    middle_full = BoseFS{67,100}(
        BitString{166,3,UInt64}(SVector(1, ~UInt64(0), UInt64(1) << 63 | UInt64(2)))
    )
    middle_empty = BoseFS{10,150}(
        BitString{159,3,UInt64}(SVector{3,UInt64}(255, 0, 3))
    )
    two_full = BoseFS{136,136}(
        BitString{271,5,UInt64}(
            SVector(UInt64(0), UInt64(0), UInt64(255), ~UInt64(0), ~UInt64(0))
        )
    )
    @testset "onr" begin
        middle_full_onr = onr(middle_full)
        @test length(middle_full_onr) == 100
        @test middle_full_onr[63] == 66
        @test middle_full_onr[2] == 1
        @test middle_full_onr[1] == 0
        @test all(iszero, middle_full_onr[[3:62; 64:end]])

        middle_empty_onr = onr(middle_empty)
        @test length(middle_empty_onr) == 150
        @test middle_empty_onr[1] == 2
        @test middle_empty_onr[127] == 8
        @test all(iszero, middle_empty_onr[[2:126; 128:end]])

        two_full_onr = onr(two_full)
        @test length(two_full_onr) == 136
        @test two_full_onr[1] == 136
        @test all(iszero, two_full_onr[2:end])
    end
    @testset "num_occupied_modes" begin
        @test num_occupied_modes(middle_full) == 2
        @test num_occupied_modes(middle_empty) == 2
        @test num_occupied_modes(two_full) == 1
        @test num_occupied_modes(BoseFS((0, 0, 0, 0))) == 0
    end
    @testset "bose_hubbard_interaction" begin
        @test bose_hubbard_interaction(middle_full) == 66 * 65
        @test bose_hubbard_interaction(middle_empty) == 8 * 7 + 2
        @test bose_hubbard_interaction(two_full) == 136 * 135
    end
    @testset "occupied_modes" begin
        (bosons, orbital, bit), st = iterate(occupied_modes(middle_full))
        @test bosons == 1
        @test orbital == 2
        @test bit == 1
        (bosons, orbital, bit), st = iterate(occupied_modes(middle_full), st)
        @test bosons == 66
        @test orbital == 63
        @test bit == 63
        @test isnothing(iterate(occupied_modes(middle_full), st))

        @test collect(occupied_modes(middle_empty)) == [[2, 1, 0], [8, 127, 128]]
    end
    @testset "Randomized tests" begin
        # Note: the random number for these tests will be the same everytime. This is still
        # an ok way to look for errors.
        function rand_onr_bose(N, M)
            result = zeros(MVector{M,Int})
            for _ in 1:N
                result[rand(1:M)] += 1
            end
            return SVector(result)
        end
        # Should be exactly the same as onr, but slower.
        function onr2(bose::BoseFS{N,M}) where {N,M}
            result = zeros(MVector{M,Int32})
            for (n, i, _) in occupied_modes(bose)
                @assert n ≠ 0
                result[i] = n
            end
            return SVector(result)
        end
        # Should be exactly the same as hopnextneighbour, but slower.
        function hopnextneighbour2(bose::BoseFS{N,M}, chosen) where {N,M}
            o = MVector{M,Int32}(onr(bose))
            site = (chosen + 1) ÷ 2
            curr = 0
            i = 1
            while i ≤ M
                curr += o[i] > 0
                curr == site && break
                i += 1
            end
            if isodd(chosen)
                j = mod1(i + 1, M)
            else
                j = mod1(i - 1, M)
            end
            o[i] -= 1
            o[j] += 1
            return BoseFS{N}(SVector(o)), (o[i] + 1) * (o[j])
        end
        for (N, M) in ((16, 16), (10, 32), (64, 32), (200, 200), (200, 20), (20, 200))
            @testset "$N, $M" begin
                for _ in 1:10
                    input = rand_onr_bose(N, M)
                    bose = BoseFS(input)
                    @test num_particles(bose) == N
                    @test num_modes(bose) == M
                    @test onr(bose) == input
                    @test num_occupied_modes(bose) == count(!iszero, input)
                    @test bose_hubbard_interaction(bose) == sum(input .* (input .- 1))

                    @test onr2(bose) == input

                    @test all(
                        hopnextneighbour2(bose, i) == hopnextneighbour(bose, i)
                        for i in 1:num_occupied_modes(bose) * 2
                    )

                    occupied = [
                        find_occupied_mode(bose, i).mode
                        for i in 1:num_occupied_modes(bose)
                    ]
                    @test occupied == findall(!iszero, onr(bose))

                    check_single_excitations(bose, 64)
                    check_double_excitations(bose, 8)
                    check_triple_excitations(bose, 4)

                    # This test checks that the result of show can be pasted into the REPL
                    @test eval(Meta.parse(repr(bose))) == bose
                end
            end
        end
    end
end

@testset "FermiFS" begin
    small = SVector(1, 0, 0, 0, 0, 1, 1, 1, 0, 0)
    big = [rand(0:1) for _ in 1:70]
    giant = [rand() < 0.1 for _ in 1:200]
    giant_sparse = [1; 1; zeros(Int, 100); 1]
    giant_dense = [1; 0; ones(Int, 50); 0; 1; 1]
    @testset "onr, constructors" begin
        for o in (small, big, giant, giant_sparse, giant_dense)
            f = FermiFS(o)
            @test onr(f) == o
            @test typeof(f)(onr(f)) == f
        end

        @test_throws ErrorException FermiFS((1, 2, 3))
        @test_throws ErrorException FermiFS{2,3}((1, 1, 1))
    end
    @testset "occupied_modes" begin
        for o in (small, big, giant)
            f = FermiFS(o)
            sites = collect(occupied_modes(f))
            @test getproperty.(sites, :mode) == findall(≠(0), onr(f))
        end
    end
    @testset "Randomized Tests" begin
        function rand_onr_fermi(N, M)
            return SVector{M}(shuffle([ones(Int,N); zeros(Int,M - N)]))
        end
        for (N, M) in ((15, 16), (10, 29), (32, 60), (180, 200), (10, 200), (1, 20))
            @testset "$N, $M" begin
                for _ in 1:10
                    input = rand_onr_fermi(N, M)
                    fermi = FermiFS(input)

                    @test num_particles(fermi) == N
                    @test num_modes(fermi) == M
                    @test onr(fermi) == input

                    check_single_excitations(fermi, 64)
                    check_double_excitations(fermi, 8)
                    check_triple_excitations(fermi, 4)

                    f = FermiFS(input)
                    sites = Int[]
                    foreach(occupied_modes(f)) do i
                        push!(sites, i.mode)
                    end
                    @test sites == findall(≠(0), onr(f))
                end
            end
        end
    end
end

@testset "Multicomponent" begin
    @testset "CompositeFS" begin
        fs1 = CompositeFS(
            FermiFS((1,1,0,0,0,0)), FermiFS((1,1,1,0,0,0)), BoseFS((5,0,0,0,0,0))
        )
        @test num_modes(fs1) == 6
        @test num_components(fs1) == 3
        @test num_particles(fs1) == 10
        @test fs1.components[1] == FermiFS((1,1,0,0,0,0))
        @test fs1.components[2] == FermiFS((1,1,1,0,0,0))
        @test fs1.components[3] == BoseFS((5,0,0,0,0,0))
        @test eval(Meta.parse(repr(fs1))) == fs1

        fs2 = CompositeFS(
            FermiFS((0,1,1,0,0,0)), FermiFS((1,0,1,0,1,0)), BoseFS((1,1,1,1,1,0))
        )
        @test fs1 < fs2

        @test_throws ErrorException CompositeFS(BoseFS((1,1)), BoseFS((1,1,1)))

        @inferred update_component(fs1, FermiFS((0,0,0,1,1,0)), Val(1))
        @inferred update_component(fs1, FermiFS((0,0,1,1,1,0)), Val(2))
        @inferred update_component(fs1, BoseFS((1,1,1,1,1,0)), Val(3))
        @test_throws MethodError update_component(fs1, BoseFS((1,1,1,1,1,0)), Val(1))
        @test_throws MethodError update_component(fs1, FermiFS((1,1,1,1,1,0)), Val(1))

        fs3 = Rimu.BitStringAddresses.update_component(fs1, FermiFS((0,0,1,1,1,0)), Val(2))
        @test fs3 == CompositeFS(
            FermiFS((1,1,0,0,0,0)), FermiFS((0,0,1,1,1,0)), BoseFS((5,0,0,0,0,0)),
        )
        @test onr(fs3) == ([1,1,0,0,0,0], [0,0,1,1,1,0], [5,0,0,0,0,0])
        @test onr(fs3, LadderBoundaries(2, 3)) == (
            [1 0 0; 1 0 0], [0 1 1; 0 1 0], [5 0 0; 0 0 0]
        )
    end

    @testset "BoseFS2C" begin
        fs1 = BoseFS2C((1,1,1,0,0,0,0), (1,1,1,0,5,0,0))
        @test num_modes(fs1) == 7
        @test num_components(fs1) == 2
        @test num_particles(fs1) == 11
        @test eval(Meta.parse(repr(fs1))) == fs1
        @test onr(fs1) == ([1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 0, 5, 0, 0])

        fs2 = BoseFS2C(BoseFS((0,0,0,0,0,0,3)), BoseFS((0,2,1,0,5,0,0)))
        @test fs1 < fs2

        @test_throws MethodError BoseFS2C(BoseFS((1,1)), BoseFS((1,1,1)))
    end
end


text:
const UnionBuf = Vector{Tuple{Int64, Int64}}

function apply_rule!(g::EGraph, rule::UnequalRule,
        match::Match, matches::MatchesBuf, unions::UnionBuf,
        rep::Report, mod::Module)
    (_, pat, sub, id) = match
    lc = id
    rinst = instantiate(g, pat, sub, rule)
    rc = addexpr!(g, rinst).id
    # delete!(matches, match)
    if find(g, lc) == find(g, rc)
        @log "Contradiction!" rule
        rep.reason = Contradiction()
        return (false, rep)
    end
    return (true, nothing)
end

function apply_rule!(g::EGraph, rule::SymbolicRule,
        match::Match, matches::MatchesBuf, unions::UnionBuf,
        rep::Report,  mod::Module)
    (_, pat, sub, id) = match
    # println(sub)
    rinst = instantiate(g, pat, sub, rule)
    rc = addexpr!(g, rinst)
    push!(unions, (id, rc.id))
    return (true, nothing)
end

function apply_rule!(g::EGraph, rule::DynamicRule,
        match::Match, matches::MatchesBuf, unions::UnionBuf,
        rep::Report,  mod::Module)
    (_, pat, sub, id) = match
    f = Rules.getrhsfun(rule, mod)
    actual_params = [instantiate(g, PatVar(v, i), sub, rule) for (i, v) in enumerate(rule.patvars)]
    r = f(geteclass(g, id), g, actual_params...)
    rc = addexpr!(g, r)

    push!(unions, (id, rc.id))
    return (true, nothing)
end

using .ReportReasons
function eqsat_apply!(g::EGraph, matches::MatchesBuf,
        rep::Report, mod::Module, params::SaturationParams)::Report
    i = 0

    unions = UnionBuf()
    # println.(matches)
    for match ∈ matches
        i += 1

        # if i % 300 == 0
        if params.eclasslimit > 0 && g.numclasses > params.eclasslimit
            @log "E-GRAPH SIZEOUT"
            rep.reason = EClassLimit(params.eclasslimit)
            return rep
        end

        if reached(g, params.goal)
            @log "Goal reached"
            rep.reason = GoalReached()
            return rep
        end

        rule=match[1]
        # println("applied $rule")

        (ok, nrep) = apply_rule!(g, rule, match, matches, unions, rep, mod)
        if !ok
            return nrep
        end
        # println(rule)
        # println(sub)
        # println(l); println(r)
        # display(egraph.classes); println()
    end

    for (l,r) ∈ unions
        merge!(g, l, r)
    end
    return rep
end


text:
# # Instructions for Assignment01
#
# ## Introduction
#
# In this assignment,
# you will learn how to use github classroom
# to download, complete, and test your assignments.
# Make sure you have followed the instructions at <<URL>>
# to set up your account and join the classroom.

# To view additional instructions for this assignment, #jl
# see https://wellesley-bisc195.github.io/BISC195.jl/stable/Assignments/assignment01.html #jl

# To view a rendered version of this document, #jl
# see https://wellesley-bisc195.github.io/BISC195.jl/stable/Assignments/assignment01_code.html #jl

# ## Question 1
#
# Did you complete the instructions for modifying
# `some_variables.jl` and `some_functions.jl`?
# You should have changed the variables `n` and `a`
# to `my_name` and `my_age` respectively,
# and you should be able to execute `julia src/run.jl`
# without errors.
# Once you can, type `yes` between the quotations below
#
# Note: this syntax (`question1 = "yes"`)
# is called "variable assignment,"
# and we'll learn more about this in the next lesson.

question1 = ""

# ## Question 2
#
# What is the shell command to list the contents
# of the working directory,
# **including hidden files**?
#
# Enter the command between the quotes below.

question2 = ""

# ## Question 3
#
# What is a shell command to create a new directory
# called "q3" in your current working directory?
#
# Enter the command between the quotes below.

question3 = ""


text:
function ldl_dense!(A::Symmetric{T, Matrix{T}}) where {T <: Real}
  # use symmetric lower only
  @assert A.uplo == 'L'
  ldl_lowtri!(A.data)
end

function ldl_lowtri!(A::Matrix)
  n = size(A, 1)
  for j = 1:n
    djj = A[j, j]
    for k = 1:(j - 1)
      djj -= A[j, k]^2 * A[k, k]
    end
    A[j, j] = djj
    for i = (j + 1):n
      lij = A[i, j] # we use the lower block
      for k = 1:(j - 1)
        lij -= A[i, k] * A[j, k] * A[k, k]
      end
      A[i, j] = lij / djj
    end
  end
end

function solve_lowertri!(L, b)
  n = length(b)
  @inbounds for j = 1:n
    @simd for i = (j + 1):n
      b[i] -= L[i, j] * b[j]
    end
  end
end

function solve_lowertri_transpose!(LT, b)
  n = length(b)
  @inbounds for i = n:-1:1
    @simd for j = (i - 1):-1:1
      b[j] -= LT[i, j] * b[i]
    end
  end
end

function solve_diag!(D, b)
  n = length(b)
  @inbounds for j = 1:n
    b[j] /= D[j, j]
  end
end

# function solve(L, D, b)
#   x = copy(b)
#   solve_lowertri!(L, x)
#   solve_diag!(D, x)
#   solve_lowertri_transpose!(L, x)
#   return x
# end

function ldiv_dense!(LD, b)
  solve_lowertri!(LD, b)
  @views solve_diag!(Diagonal(LD[diagind(LD)]), b)
  solve_lowertri_transpose!(LD, b)
  return b
end

function ldiv_dense(x, LD, b)
  x .= b
  ldiv_dense!(LD, x)
end


text:
using AbstractGPs

mutable struct AbstractGPSurrogate{X, Y, GP, GP_P, S} <: AbstractSurrogate
    x::X
    y::Y
    gp::GP
    gp_posterior::GP_P
    Σy::S
 end

# constructor
function AbstractGPSurrogate(x, y; gp = GP(Matern52Kernel()), Σy = 0.1)
    AbstractGPSurrogate(x, y, gp, posterior(gp(x, Σy),y), Σy)
 end

# predictor
function (g::AbstractGPSurrogate)(val)
    return only(mean(g.gp_posterior([val])))
end

# for add point
# copies of x and y need to be made because we get
#"Error: cannot resize array with shared data " if we push! directly to x and y
function add_point!(g::AbstractGPSurrogate, new_x, new_y)
    if new_x in g.x
        println("Adding a sample that already exists, cannot build AbstracgGPSurrogate.")
        return
    end
    x_copy = copy(g.x)
    push!(x_copy, new_x)
    y_copy = copy(g.y)
    push!(y_copy, new_y)
    updated_posterior = posterior(g.gp(x_copy, g.Σy), y_copy)
    g.x, g.y, g.gp_posterior = x_copy, y_copy, updated_posterior
    nothing
end

function std_error_at_point(g::AbstractGPSurrogate, val)
    return sqrt(only(var(g.gp_posterior([val]))))
end

# Log marginal posterior predictive probability.
function logpdf_surrogate(g::AbstractGPSurrogate)
    return logpdf(g.gp_posterior(g.x), g.y)
end


text:
# # Custom gradients
#
# Here we will consider a custom gradient applied to a spherical geometry.
#
# We start by loading SpinDoctor and a Makie plotting backend.

# LSP indexing solution                                                          #src
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983 #src
if isdefined(@__MODULE__, :LanguageServer)                                       #src
    include("../src/SpinDoctor.jl")                                              #src
    using .SpinDoctor                                                            #src
end                                                                              #src

using SpinDoctor
using LinearAlgebra

if haskey(ENV, "GITHUB_ACTIONS")
    using CairoMakie
else
    using GLMakie
end

setup = SphereSetup(;
    name = "gaze-into-the-orb",
    ncell = 1,
    rmin = 2.0,
    rmax = 6.0,
    dmin = 0.2,
    dmax = 0.3,
    include_in = false,
    in_ratio = 0.6,
    ecs_shape = :no_ecs,
    ecs_ratio = 0.5,
)
coeffs = coefficients(
    setup;
    D = (; in = 0.002 * I(3), out = 0.002 * I(3), ecs = 0.002 * I(3)),
    T₂ = (; in = Inf, out = Inf, ecs = Inf),
    ρ = (; in = 1.0, out = 1.0, ecs = 1.0),
    κ = (; in_out = 1e-4, out_ecs = 1e-4, in = 0.0, out = 0.0, ecs = 0.0),
    γ = 2.67513e-4,
)

# We then proceed to build the geometry and finite element mesh.

mesh, = create_geometry(setup; recreate = true)
plot_mesh(mesh)

# The mesh looks good, so we may then proceed to assemble the biological model and the
# associated finite element matrices.

model = Model(; mesh, coeffs...)
matrices = assemble_matrices(model)

# The Bloch-Torrey PDE takes a magnetic field gradient pulse sequence as an input. We may
# define our custom three dimensional gradient sequence (given in T/m) as a simple Julia
# function. The echo time `TE` (given in microseconds) is also needed.

## Magnetic field gradient
TE = 5000.0
φ = -π / 6
R = [cos(φ) sin(φ) 0; -sin(φ) cos(φ) 0; 0 0 1]
g⃗(t) = 1.0 * R * [sin(2π * t / TE), sin(20π * t / TE) / 5, cos(2π * t / TE)]
gradient = GeneralGradient(; g⃗, TE)

# In order to follow the evolution of the solution during time stepping, we add a
# [`Plotter`](@ref) to a list of callbacks. Other available callbacks are [`Printer`](@ref)
# for showing time stepping information, and [`VTKWriter`](@ref) for saving the solution time
# series for visualization in ParaView.

plotter = Plotter{Float64}()

# We may then define the problem and solve for our gradient (with the callback).

btpde = BTPDE(; model, matrices)
ξ = solve(btpde, gradient; callbacks = [plotter])
plotter.fig

# [`MatrixFormalism`](@ref) also accepts [`GeneralGradient`](@ref)s. However, the ADC methods
# currently only support [`ScalarGradient`](@ref)s.


text:
function midpoint_radius(x::ArbFloat{P}) where {P}
    return midpoint(x), radius(x)
end

function midpoint_radius(midpt::ArbFloat{P}, radius::ArbFloat{P}) where {P}
    mid = midpoint(midpt)
    rad = midpoint(radius)
    ccall(@libarb(arb_add_error), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}), mid, rad)
    return mid
end

function midpoint_radius(mid::Float64, rad::Float64)
    m = convert(ArbFloat, mid)
    r = convert(ArbFloat, rad)
    return midpoint_radius(m, r)
end
function midpoint_radius(mid::ArbFloat{P}, rad::Float64) where {P}
    r = convert(ArbFloat{P}, rad)
    return midpoint_radius(mid, r)
end
function midpoint_radius(mid::ArfFloat{P}, rad::Float64) where {P}
    m = convert(ArbFloat{P}, mid)
    r = convert(ArbFloat{P}, rad)
    return midpoint_radius(m, r)
end
function midpoint_radius(mid::String, rad::String)
    P = precision(ArbFloat)
    m = ArbFloat(mid)
    r = ArbFloat(rad)
    return midpoint_radius(m, r)
end
function midpoint_radius(mid, rad)
    P = precision(ArbFloat)
    m = convert(ArbFloat{P}, mid)
    r = convert(ArbFloat{P}, rad)
    return midpoint_radius(m, r)
end

function bounding_midpoint(a::T) where {T <: ArbFloat}
    halflo = lowerbound(a) * 0.5
    halfhi = upperbound(a) * 0.5
    return midpoint(halflo + halfhi)
end

function bounding_radius(a::T) where {T <: ArbFloat}
    halflo = lowerbound(a) * 0.5
    halfhi = upperbound(a) * 0.5
    return midpoint(halfhi - halflo)
end

function bounding_midpoint_radius(a::T) where {T <: ArbFloat}
    halflo = lowerbound(a) * 0.5
    halfhi = upperbound(a) * 0.5
    mid = midpoint(halflo + halfhi)
    rad = midpoint(halfhi - halflo)
    return midpoint_radius(mid, rad)
end

#=
void arb_union(arb_t z, const arb_t x, const arb_t y, slong prec)¶
Sets z to a ball containing both x and y.
=#
union(a::T) where {T <: ArbFloat} = a

function union(a::T, b::T) where {T <: ArbFloat}
    P = precision(T)
    z = ArbFloat{P}{}
    ccall(@libarb(arb_union), Cvoid, (Ref{ArbFloat}, Ref{ArbFloat}, Ref{ArbFloat}, Clong), z, a, b, P)
    return z
end
function union(a::T, b::T, c::T) where {T <: ArbFloat}
    z1 = union(a,b)
    z2 = union(z1,c)
    return z2
end
function union(a::T, b::T, c::T, d::T) where {T <: ArbFloat}
    z1 = union(a,b)
    z2 = union(z1,c)
    z3 = union(z2,d)
    return z3
end

narrow(a::T) where {T <: ArbFloat} = a

function narrow(a::T, b::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    lo = max(alo,blo)
    hi = min(ahi,bhi)
    return union(lo,hi)
end

function narrow(a::T, b::T, c::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    clo, chi = bounds(c)
    lo = max(max(alo,blo),clo)
    hi = min(min(ahi,bhi),chi)
    return union(lo,hi)
end

function narrow(a::T, b::T, c::T, d::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    clo, chi = bounds(c)
    dlo, dhi = bounds(d)
    lo = max(max(alo,blo),max(clo,dlo))
    hi = min(min(ahi,bhi),min(chi,clo))
    return union(lo,hi)
end

intersect(a::T) where {T <: ArbFloat} = a

function intersect(a::T, b::T) where {T <: ArbFloat}
    P = precision(T)
    z = initializer(ArbFloat{P})
    if donotoverlap(a,b)
        ccall(@libarb(arb_indeterminate), Cvoid, (Ref{ArbFloat},), z)
    else
        alo,ahi = bounds(a)
        blo,bhi = bounds(b)
        lo,hi = minmax(max(alo,blo), min(ahi,bhi))
        bounded(z, lo, hi)
    end
    return z
end

function intersect(a::T, b::T, c::T) where {T <: ArbFloat}
   i1 = intersect(a,b)
   i2 = intersect(i1,c)
   return i2
end

function intersect(a::T, b::T, c::T, d::T) where {T <: ArbFloat}
   i1 = intersect(a,b)
   i2 = intersect(c,d)
   return intersect(i1,i2)
end

function bounded(z::ArbFloat{P}, lo::ArbFloat{P}, hi::ArbFloat{P}) where {P}
    lo2 = convert(ArfFloat{P}, lo)
    hi2 = convert(ArfFloat{P}, hi)
    ccall(@libarb(arb_set_interval_arf), Cvoid, (Ref{ArbFloat{P}}, Ref{ArfFloat{P}}, Ref{ArfFloat{P}}, Int64), z, lo2, hi2, P%Int64)
    return z
end
function bounded(lo::T, hi::T) where {T <: ArbFloat}
    P = precision(T)
    z = initializer(ArbFloat{P})
    return bounded(z,lo,hi)
end
function boundedrange(mid::T, rad::T) where {T <: ArbFloat}
    P = precision(T)
    z = initializer(ArbFloat{P})
    lo = mid-rad
    hi = mid+rad
    return bounded(z,lo,hi)
end


text:
# Adapted from https://github.com/karpathy/pytorch-normalizing-flows
using Revise
using ADCME
using PyCall
using PyPlot
using Random

# `nmoons` is adapted from https://github.com/wildart/nmoons
function nmoons(::Type{T}, n::Int=100, c::Int=2;
    shuffle::Bool=false, ε::Real=0.1, d::Int = 2,
    translation::Vector{T}=zeros(T, d),
    rotations::Dict{Pair{Int,Int},T} = Dict{Pair{Int,Int},T}(),
    seed::Union{Int,Nothing}=nothing) where {T <: Real}
    rng = seed === nothing ? Random.GLOBAL_RNG : MersenneTwister(Int(seed))
    ssize = floor(Int, n/c)
    ssizes = fill(ssize, c)
    ssizes[end] += n - ssize*c
    @assert sum(ssizes) == n "Incorrect partitioning"
    pi = convert(T, π)
    R(θ) = [cos(θ) -sin(θ); sin(θ) cos(θ)]
    X = zeros(d,0)
    for (i, s) in enumerate(ssizes)
    circ_x = cos.(range(zero(T), pi, length=s)).-1.0
    circ_y = sin.(range(zero(T), pi, length=s))
    C = R(-(i-1)*(2*pi/c)) * hcat(circ_x, circ_y)'
    C = vcat(C, zeros(d-2, s))
    dir = zeros(d)-C[:,end] # translation direction
    X = hcat(X, C .+ dir.*translation)
    end
    y = vcat([fill(i,s) for (i,s) in enumerate(ssizes)]...)
    if shuffle
        idx = randperm(rng, n)
        X, y = X[:, idx], y[idx]
    end
    # Add noise to the dataset
    if ε > 0.0
        X += randn(rng, size(X)).*convert(T,ε/d)
    end
    # Rotate dataset
    for ((i,j),θ) in rotations
        X[[i,j],:] .= R(θ)*view(X,[i,j],:)
    end
    return X, y
end

function sample_moons(n)
    X, _ = nmoons(Float64, n, 2, ε=0.05, d=2, translation=[0.25, -0.25])
    return Array(X')
end

#------------------------------------------------------------------------------------------
# RealNVP
function mlp(x, k, id)
    x = constant(x)
    variable_scope("layer$k$id") do
        x = dense(x, 24, activation="leaky_relu")
        x = dense(x, 24, activation="leaky_relu")
        x = dense(x, 24, activation="leaky_relu")
        x = dense(x, 1)
    end
    return x
end
flows = [AffineHalfFlow(2, mod(i,2)==1, x->mlp(x, i, 0), x->mlp(x, i, 1)) for i = 0:8]

#------------------------------------------------------------------------------------------
# NICE
# function mlp(x, k, id)
#     x = constant(x)
#     variable_scope("layer$k$id") do
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 1)
#     end
#     return x
# end
# flow1 = [AffineHalfFlow(2, mod(i,2)==1, missing, x->mlp(x, i, 1)) for i = 0:4]
# flow2 = [AffineConstantFlow(2, shift=false)]
# flows = [flow1;flow2]

# SlowMAF
#------------------------------------------------------------------------------------------
# function mlp(x, k, id)
#     x = constant(x)
#     variable_scope("layer$k$id") do
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 2)
#     end
#     return x
# end
# flows = [SlowMAF(2, mod(i,2)==1, [x->mlp(x, i, 0)]) for i = 0:3]

# MAF
#------------------------------------------------------------------------------------------
# flows = [MAF(2, mod(i,2)==1, [24, 24, 24], name="layer$i") for i = 0:3]

# IAF
#------------------------------------------------------------------------------------------
# flows = [IAF(2, mod(i,2)==1, [24, 24, 24], name="layer$i") for i = 0:3]
# prior = ADCME.MultivariateNormalDiag(loc=zeros(2))
# model = NormalizingFlowModel(prior, flows)

# Insert ActNorm to any of the flows
#------------------------------------------------------------------------------------------
# flow2 = [ActNorm(2, "ActNorm$i") for i = 1:length(flows)]
# flows = permutedims(hcat(flow2, flows))[:]
# # error()
# # msample = rand(model,1)
# # zs, prior_logprob, log_det = model([0.0040 0.4426])
# # sess = Session(); init(sess)
# # run(sess, msample)
# # run(sess,zs)

# GLOW
#------------------------------------------------------------------------------------------
# function mlp(x, k, id)
#     x = constant(x)
#     variable_scope("layer$k$id") do
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 24, activation="leaky_relu")
#         x = dense(x, 1)
#     end
#     return x
# end
# flows = [Invertible1x1Conv(2, "conv$i") for i = 0:2]
# norms = [ActNorm(2, "ActNorm$i") for i = 0:2]
# couplings = [AffineHalfFlow(2, mod(i, 2)==1, x->mlp(x, i, 0), x->mlp(x, i, 1)) for i = 0:length(flows)-1]
# flows = permutedims(hcat(norms, flows, couplings))[:]

#------------------------------------------------------------------------------------------
# Neural Splines Coupling
# function mlp(x, k, id)
#     x = constant(x)
#     variable_scope("fc$k$id") do
#         x = dense(x, 16, activation="leaky_relu")
#         x = dense(x, 16, activation="leaky_relu")
#         x = dense(x, 16, activation="leaky_relu")
#         x = dense(x, 3K-1)
#     end
#     return x
# end
# K = 8
# flows = [NeuralCouplingFlow(2, x->mlp(x, i, 0), x->mlp(x, i, 1), K) for i = 0:2]
# convs = [Invertible1x1Conv(2, "conv$i") for i = 0:2]
# norms = [ActNorm(2, "ActNorm$i") for i = 0:2]
# flows = permutedims(hcat(norms, convs, flows))[:]

#------------------------------------------------------------------------------------------

prior = ADCME.MultivariateNormalDiag(loc=zeros(2))
model = NormalizingFlowModel(prior, flows)

x = placeholder(rand(128,2))
zs, prior_logpdf, logdet = model(x)
log_pdf = prior_logpdf + logdet
loss = -sum(log_pdf)

model_samples = rand(model, 128*8)
sess = Session(); init(sess)
opt = AdamOptimizer(1e-4).minimize(loss)
sess = Session(); init(sess)
for i = 1:10000
    _, l = run(sess, [opt, loss], x=>sample_moons(128))
    if mod(i,100)==0
        @info i, l
    end
end

z = run(sess, model_samples[end])
x = sample_moons(128*8)
scatter(x[:,1], x[:,2], c="b", s=5, label="data")
scatter(z[:,1], z[:,2], c="r", s=5, label="prior --> posterior")
axis("scaled"); xlabel("x"); ylabel("y")#


text:
using Combinatorics: combinations

test1 = """
    mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
    mem[8] = 11
    mem[7] = 101
    mem[8] = 0
    """

test2 = """
    mask = 000000000000000000000000000000X1001X
    mem[42] = 100
    mask = 00000000000000000000000000000000X0XX
    mem[26] = 1
    """

function read_input(input)
    return map(readlines(input)) do line
        (opc, args...) = split(line, r"\[|\]? = "; keepempty = false)
        return parse_command(Val(Symbol(opc)), args...)
    end
end

parse_command(::Val{:mask}, mask) = mask
parse_command(::Val{:mem}, args...) = parse.(Int, args)

function run(execute::Function, prog::AbstractVector)::Dict
    return reduce(execute, prog; init = ("", Dict())) |> last
end

q1(prog::AbstractVector) = run(q1, prog) |> values |> sum
q1((_, mem), m::AbstractString) = (m, mem)
function q1((m, mem), (a, v)::Tuple)
    (or, and) = parse.(Int, replace.(Ref(m), 'X' .=> ['0', '1']); base = 2)
    mem[a] = (v & and) | or
    return (m, mem)
end

q2(prog::AbstractVector) = run(q2, prog) |> values |> sum
q2((_, mem), m::AbstractString) = (m, mem)
function q2((m, mem), (a, v)::Tuple)
    or = parse(Int, replace(m, 'X' => '0'); base = 2)
    and = parse(Int, replace(replace(m, '0' => '1'), 'X' => '0'); base = 2)
    a = (a & and) | or
    floats = [2 << (length(m) - i - 1) for i in findall(==('X'), m)]
    as = [reduce(|, [a; fs]) for fs in [[], combinations(floats)...]]
    setindex!.(Ref(mem), v, as)
    return (m, mem)
end

@assert q1(read_input(IOBuffer(test1))) == 165
@assert q2(read_input(IOBuffer(test2))) == 208

prog = read_input("day14.in")
println("Q1: ", q1(prog))
println("Q2: ", q2(prog))

max_stars_repo_name: dhanak/competitive-coding

max_issues_repo_name: dhanak/competitive-coding

max_forks_repo_name: dhanak/competitive-coding


text:
export BasicDQNLearner

using Random
using Flux

"""
    BasicDQNLearner(;kwargs...)

See paper: [Playing Atari with Deep Reinforcement Learning](https://www.cs.toronto.edu/~vmnih/docs/dqn.pdf)

This is the very basic implementation of DQN. Compared to the traditional Q learning, the only difference is that,
in the updating step it uses a batch of transitions sampled from an experience buffer instead of current transition.
And the `approximator` is usually a [`NeuralNetworkApproximator`](@ref).
You can start from this implementation to understand how everything is organized and how to write your own customized algorithm.
# Keywords
- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `loss_func`: the loss function to use. TODO: provide a default [`huber_loss`](@ref)?
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `seed=nothing`.
"""
mutable struct BasicDQNLearner{Q,F,R} <: AbstractLearner
    approximator::Q
    loss_func::F
    γ::Float32
    batch_size::Int
    min_replay_history::Int
    rng::R
    loss::Float32
end

Flux.functor(x::BasicDQNLearner) = (Q = x.approximator,), y -> begin
    x = @set x.approximator = y.Q
    x
end

(learner::BasicDQNLearner)(obs) =
    obs |>
    get_state |>
    x ->
        send_to_device(device(learner.approximator), x) |>
        learner.approximator |>
        send_to_host

function BasicDQNLearner(;
    approximator::Q,
    loss_func::F,
    γ = 0.99f0,
    batch_size = 32,
    min_replay_history = 32,
    seed = nothing,
) where {Q,F}
    rng = MersenneTwister(seed)
    BasicDQNLearner{Q,F,typeof(rng)}(
        approximator,
        loss_func,
        γ,
        batch_size,
        min_replay_history,
        rng,
        0.0,
    )
end

function RLBase.update!(learner::BasicDQNLearner, t::AbstractTrajectory)
    length(t) < learner.min_replay_history && return

    inds = rand(learner.rng, 1:length(t), learner.batch_size)
    batch = map(get_trace(t, :state, :action, :reward, :terminal, :next_state)) do x
        consecutive_view(x, inds)
    end

    Q, γ, loss_func, batch_size =
        learner.approximator, learner.γ, learner.loss_func, learner.batch_size
    s, r, t, s′ = map(
        x -> send_to_device(device(Q), x),
        (batch.state, batch.reward, batch.terminal, batch.next_state),
    )
    a = CartesianIndex.(batch.action, 1:batch_size)

    gs = gradient(params(Q)) do
        q = Q(s)[a]
        q′ = vec(maximum(Q(s′); dims = 1))
        G = r .+ γ .* (1 .- t) .* q′
        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end


text:
using NNlib: conv, ∇conv_data, depthwiseconv, output_size, DenseConvDims
import Flux

# pad dims of x with dims of y until ndims(x) == ndims(y)
_paddims(x::Tuple, y::Tuple) = (x..., y[(end - (length(y) - length(x) - 1)):end]...)

_convtransoutdims(isize, ksize, ssize, dsize, pad) = (isize .- 1).*ssize .+ 1 .+ (ksize .- 1).*dsize .- (pad[1:2:end] .+ pad[2:2:end])

expand(N, i::Tuple) = i
expand(N, i::Integer) = ntuple(_ -> i, N)
"""
    Conv(size, in => out,
        stride = 1, pad = 0, dilation = 1,
        omega_0 = 30, is_first = false)

Convolutional layer with SIREN initialization and additional scaling factor omega_0.
`size` should be a tuple like `(2, 2)`.
`in` and `out` specify the number of input and output channels respectively.

Data should be stored in WHCN order (width, height, # channels, batch size).
In other words, a 100×100 RGB image would be a `100×100×3×1` array,
and a batch of 50 would be a `100×100×3×50` array.

# Examples

Apply a `SIREN.Conv` layer to a 1-channel input using a 2×2 window size, giving us a
16-channel output.
```julia
size = (2,2)
in = 1
out = 16
SIREN.Conv(size, in => out)
```
"""
struct Conv{N,M,F,A,V}
  omega_0::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
end

function Conv(w::AbstractArray{T,N}, b::AbstractVector{T};
              omega_0 = 30, stride = 1, pad = 0, dilation = 1) where {T,N}
  stride = expand(Val(N-2), stride)
  pad = expand(Val(2*(N-2)), pad)
  dilation = expand(Val(N-2), dilation)
  return Conv(omega_0, w, b, stride, pad, dilation)
end

Conv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer};
     omega_0 = 30, is_first = false,  stride = 1, pad = 0, dilation = 1) where N =
  Conv(SIREN_init(omega_0, is_first, k..., ch...), zeros(ch[2]),
       omega_0=omega_0, stride = stride, pad = pad, dilation = dilation)

Flux.@functor Conv

function (c::Conv)(x::AbstractArray)
  # TODO: breaks gpu broadcast :(
  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
  omega_0, b = c.omega_0, reshape(c.bias, map(_->1, c.stride)..., :, 1)
  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
  sin.(omega_0 * (conv(x, c.weight, cdims) .+ b))
end

function Base.show(io::IO, l::Conv)
  print(io, "SIREN.Conv(", size(l.weight)[1:ndims(l.weight)-2])
  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
  print(io, ", omega_0=", l.omega_0)
  print(io, ")")
end

(a::Conv{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  invoke(a, Tuple{AbstractArray}, x)

(a::Conv{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
  a(T.(x))

max_issues_repo_name: MacKenzieHnC/SIREN

max_forks_repo_name: MacKenzieHnC/SIREN


text:
module TestMetadata

using Test
using MLJModels
import MLJBase
import MLJBase: Table, Continuous, Count, Finite, OrderedFactor, Multiclass

@testset "(de)serialization for TOML" begin
    d = Dict()
    d[:test] = Tuple{Union{Continuous,Missing},Finite}
    d["junk"] = Dict{Any,Any}("H" => Missing, :cross => "lemon",
                              :t => :w, "r" => "r",
                              "tuple" =>(nothing, Float64),
                              "vector" =>[1, 2, Int])
    d["a"] = "b"
    d[:f] = true
    d["j"] = :post

    @test MLJModels.decode_dic(MLJModels.encode_dic(d)) == d
end

@testset "inverting set-valued dictionaries" begin
    d = Dict(
        :x => Set([1, 2]),
        :y => Set([2, 3, 5]),
        :z => Set([4, 7]),
        :a => Set([8, 1]),
        :b => Set([4,]),
        :w => Set([3, 1, 2]),
        :t => Set([0,]))

    dinv = Dict(
        0 => Set([:t,]),
        1 => Set([:x, :a, :w]),
        2 => Set([:x, :y, :w]),
        3 => Set([:y, :w]),
        4 => Set([:z, :b]),
        5 => Set([:y,]),
        7 => Set([:z,]),
        8 => Set([:a,]))
    @test MLJModels.inverse(d) == dinv
end

metadata_file = joinpath(@__DIR__, "..", "src",
                         "registry", "Metadata.toml")
pca = MLJModels.Handle("PCA", "MultivariateStats")
cnst = MLJModels.Handle("ConstantRegressor", "MLJModels")
i = MLJModels.info_given_handle(metadata_file)[pca]

@testset "Handle constructors" begin
    @test MLJModels.Handle("PCA") ==
        MLJModels.Handle("PCA", "MultivariateStats")
    # TODO: add tests here when duplicate model names enter registry
end

@testset "building INFO_GIVEN_HANDLE" begin
    @test MLJModels.localmodeltypes(MLJBase) ==
        MLJModels.localmodeltypes(MLJModels)
    @test issubset(Set([DeterministicConstantClassifier,
                        DeterministicConstantRegressor,
                        ConstantClassifier,
                        ConstantRegressor,
                        FeatureSelector,
                        OneHotEncoder,
                        Standardizer,
                        UnivariateBoxCoxTransformer,
                        UnivariateStandardizer]),
                   MLJModels.localmodeltypes(MLJModels))
    @test MLJModels.info_given_handle(metadata_file)[pca][:name] == "PCA"
    d1 = MLJModels.info_given_handle(metadata_file)[cnst]
    d2 = MLJBase.info_dict(ConstantRegressor)
    for (k, v) in d1
        if v isa Vector
            @test Set(v) == Set(d2[k])
        else
            @test v == d2[k]
        end
    end
end

h = Vector{Any}(undef, 7)
h[1] = MLJModels.Handle("1", "a")
h[3] = MLJModels.Handle("1", "b")
h[2] = MLJModels.Handle("2", "b")
h[4] = MLJModels.Handle("3", "b")
h[5] = MLJModels.Handle("4", "c")
h[6] = MLJModels.Handle("3", "d")
h[7] = MLJModels.Handle("5", "e")
info_given_handle = Dict([h[j]=>i for j in 1:7]...)

@testset "building AMBIGUOUS_NAMES" begin
    @test Set(MLJModels.ambiguous_names(info_given_handle)) == Set(["1", "3"])
end

@testset "building PKGS_GIVEN_NAME" begin
    d = MLJModels.pkgs_given_name(info_given_handle)
    @test Set(d["1"]) == Set(["a", "b"])
    @test d["2"]==["b",]
    @test Set(d["3"]) == Set(["b", "d"])
    @test d["4"] == ["c",]
    @test d["5"] == ["e",]
end

@testset "building NAMES" begin
    model_names = MLJModels.model_names(info_given_handle)
    @test Set(model_names) == Set(["1", "2", "3", "4", "5"])
end

end
true


text:
function Program{S<:String}(shaders::Dict{GLenum,S})
	prog = Program()
	for (typ, src) = shaders
		AttachShader(prog, Shader(typ, src))
	end
	LinkProgram(prog)
	return prog
end

function Program()
	ret = ccall( (:glCreateProgram, lib), Program, ())
	if ret == 0
		GetError()
		error("Program creation failed")
	end
	return ret
end

function AttachShader(p::Program, s::Shader)
	ccall( (:glAttachShader, lib), Void, (Program, Shader), p, s)
	GetError()
end

function LinkProgram(p::Program)
	ccall( (:glLinkProgram, lib), Void, (Program,), p)
	if !GetProgram(p, LINK_STATUS)
		error(GetProgramInfoLog(s))
	end
end

function UseProgram(p::Program)
	ccall( (:glUseProgram, lib), Void, (Program,), p)
	GetError() # TODO: Benchmark overhead
end

const LINK_STATUS = 0x8B82

function GetProgram(p::Program, param::Integer)
	ret = GLint[-1]
	ccall( (:glGetProgramiv, lib), Void,
		(Program, GLenum, Ptr{GLint}), p, param, ret)
	ret = ret[1]
	if ret == -1
		GetError()
		error("no output")
	elseif contains((LINK_STATUS, DELETE_STATUS), param)
		return ret == 1
	else
		return ret
	end
end

function GetProgramInfoLog(p::Program)
	size = GetProgram(s, INFO_LOG_LENGTH)
	buf = Array(GLchar, size)
	ccall( (:glGetProgramInfoLog, lib), Void,
		(Program, GLsizei, Ptr{GLsizei}, Ptr{GLchar}),
		p, size, &size, buf)
	return bytestring(buf)
end


text:
module Utils

import GraphIdx: Const, enumerate_edges, Graph, Weights

"""
    sum2(x)

Sum of squares.

```jldoctest
julia> TreeLas.Utils.sum2([1, 11])
122

```
"""
function sum2(x::Array{F,N})::F where{F<:Real, N}
    local s::F = zero(F)
    for xi in x
        s += xi^2
    end
    s
end

"""
    primal_objective(x, y, graph [, mu])

Compute the objective function, i.e.
```math
\\frac{1}{2} \\sum_{i=1}^n μ_i (x_i - y_i)^2 + \\sum_{(i,j) \\in E} λ_{ij} |x_i - x_j|
```
Hereby the edges (and edge weights ``λ_{ij}``) are obtained via `enumerate_edges(graph)`.
"""
@inline function primal_objective(
    x::Array{F,N},
    y::Array{F,N},
    graph::G,
    mu::F2 = Const(1.0),
)::F where {F, N, G<:Graph, F1<:Weights{F}, F2<:Weights{F}}
    @assert length(y) == length(x)
    local n = length(x)
    local obj::Ref{F} = Ref{F}(0.5 * sum(mu[i] * (x[i] - y[i])^2 for i = 1:n))
    enumerate_edges(graph) do ei::Int, u::Int, v::Int, lam::Float64
        obj[] += lam * abs(x[u] - x[v])
    end
    return obj[]
end

end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
# Functions for querying network properties.

######### Accessors: #########

"""
    species(network)

Given a [`ReactionSystem`](@ref), return a vector of species `Variable`s.

Notes:
- If `ModelingToolkit.get_systems(network)` is not empty, may allocate. Otherwise returns
  `ModelingToolkit.ModelingToolkit.get_states(network)`.
"""
function species(network)
    isempty(ModelingToolkit.get_systems(network)) ? ModelingToolkit.ModelingToolkit.get_states(network) : states(network)
end

"""
    params(network)

Given a [`ReactionSystem`](@ref), return a vector of parameter `Variable`s.

Notes:
- If `ModelingToolkit.get_systems(network)` is not empty, may allocate. Otherwise returns
  `ModelingToolkit.get_ps(network)`.
"""
function params(network)
    isempty(ModelingToolkit.get_systems(network)) ? ModelingToolkit.get_ps(network) : parameters(network)
end

"""
    reactions(network)

Given a [`ReactionSystem`](@ref), return a vector of all `Reactions` in the system.

Notes:
- If `ModelingToolkit.get_systems(network)` is not empty, may allocate. Otherwise returns
  `ModelingToolkit.get_eqs(network)`.
"""
function reactions(network)
    isempty(ModelingToolkit.get_systems(network)) ? ModelingToolkit.get_eqs(network) : equations(network)
end

"""
    speciesmap(network)

Given a [`ReactionSystem`](@ref), return a Dictionary mapping from species to
species indices. (Allocates)
"""
function speciesmap(network)
    Dict(S => i for (i,S) in enumerate(species(network)))
end

"""
    paramsmap(network)

Given a [`ReactionSystem`](@ref), return a Dictionary mapping from parameters to
parameter indices. (Allocates)
"""
function paramsmap(network)
    Dict(p => i for (i,p) in enumerate(params(network)))
end

"""
    numspecies(network)

Return the number of species within the given [`ReactionSystem`](@ref).
"""
function numspecies(network)
    ns = length(ModelingToolkit.ModelingToolkit.get_states(network))
    for sys in ModelingToolkit.get_systems(network)
        ns += numspecies(ns)
    end
    ns
end

"""
    numreactions(network)

Return the number of reactions within the given [`ReactionSystem`](@ref).
"""
function numreactions(network)
    nr = length(ModelingToolkit.get_eqs(network))
    for sys in ModelingToolkit.get_systems(network)
        nr += numreactions(sys)
    end
    nr
end

"""
    numparams(network)

Return the number of parameters within the given [`ReactionSystem`](@ref).
"""
function numparams(network)
    np = length(ModelingToolkit.get_ps(network))
    for sys in ModelingToolkit.get_systems(network)
        np += numparams(ns)
    end
    np
end

"""
    dependents(rx, network)

Given a [`Reaction`](@ref) and a [`ReactionSystem`](@ref), return a vector of
`ModelingToolkit.Num`s corresponding to species the *reaction rate
law* depends on. E.g., for

`k*W, 2X + 3Y --> 5Z + W`

the returned vector would be `[W(t),X(t),Y(t)]`.

Notes:
- Allocates
- Does not check for dependents within any subsystems.
"""
function dependents(rx, network)
    if rx.rate isa Number
        return rx.substrates
    else
        rvars = ModelingToolkit.get_variables(rx.rate, species(network))
        return union!(rvars, rx.substrates)
    end
end

"""
    dependents(rx, network)

See documentation for [`dependents`](@ref).
"""
function dependants(network, rxidx)
    dependents(network, rxidx)
end

"""
    substoichmat(rn; smap=speciesmap(rn))

Returns the substrate stoichiometry matrix
"""
function substoichmat(rn; smap=speciesmap(rn))
    smat = zeros(Int,(numreactions(rn),numspecies(rn)))
    for (k,rx) in enumerate(reactions(rn))
        stoich = rx.substoich
        for (i,sub) in enumerate(rx.substrates)
            smat[k,smap[sub]] = stoich[i]
        end
    end
    smat
end

"""
    prodstoichmat(rn; smap=speciesmap(rn))

Returns the product stoichiometry matrix
"""
function prodstoichmat(rn; smap=speciesmap(rn))
    pmat = zeros(Int,(numreactions(rn),numspecies(rn)))
    for (k,rx) in enumerate(reactions(rn))
        stoich = rx.prodstoich
        for (i,prod) in enumerate(rx.products)
            pmat[k,smap[prod]] = stoich[i]
        end
    end
    pmat
end

######################## reaction network operators #######################

"""
    ==(rn1::Reaction, rn2::Reaction)

Tests whether two [`Reaction`](@ref)s are identical.

Notes:
- Ignores the order in which stoichiometry components are listed.
- *Does not* currently simplify rates, so a rate of `A^2+2*A+1` would be
    considered different than `(A+1)^2`.
"""
function (==)(rn1::Reaction, rn2::Reaction)
    isequal(rn1.rate, rn2.rate) || return false
    issetequal(zip(rn1.substrates,rn1.substoich), zip(rn2.substrates,rn2.substoich)) || return false
    issetequal(zip(rn1.products,rn1.prodstoich), zip(rn2.products,rn2.prodstoich)) || return false
    issetequal(rn1.netstoich, rn2.netstoich)
end

"""
    ==(rn1::ReactionSystem, rn2::ReactionSystem)

Tests whether the underlying species, parameters and reactions are the same in
the two [`ReactionSystem`](@ref)s. Ignores order network components were
defined.

Notes:
- *Does not* currently simplify rates, so a rate of `A^2+2*A+1` would be
    considered different than `(A+1)^2`.
- Flattens subsystems, and hence may allocate, when checking equality.
"""
function (==)(rn1::ReactionSystem, rn2::ReactionSystem)
    issetequal(species(rn1), species(rn2)) || return false
    issetequal(params(rn1), params(rn2)) || return false
    isequal(ModelingToolkit.independent_variable(rn1), ModelingToolkit.independent_variable(rn2)) || return false
    (numreactions(rn1) == numreactions(rn2)) || return false

    # the following fails for some reason, so need to use issubset
    #issetequal(equations(rn1), equations(rn2)) || return false
    (issubset(reactions(rn1),reactions(rn2)) && issubset(reactions(rn2),reactions(rn1))) || return false

    # BELOW SHOULD NOT BE NEEDED as species, params and equations flatten
    #issetequal(rn1.systems, rn2.systems) || return false
    # sys1 = rn1.systems; sys2 = rn2.systems
    # (issubset(sys1,sys2) && issubset(sys2,sys1)) || return false
    true
end

######################## functions to extend a network ####################

"""
    make_empty_network(; iv=Sym{ModelingToolkit.Parameter{Real}}(:t))

Construct an empty [`ReactionSystem`](@ref). `iv` is the independent variable,
usually time.
"""
function make_empty_network(; iv=Sym{ModelingToolkit.Parameter{Real}}(:t))
    ReactionSystem(Reaction[], iv, [], [], Equation[], gensym(:ReactionSystem), ReactionSystem[])
end

"""
    addspecies!(network::ReactionSystem, s::Symbolic; disablechecks=false)

Given a [`ReactionSystem`](@ref), add the species corresponding to the variable
`s` to the network (if it is not already defined). Returns the integer id of the
species within the system.

Notes:
- `disablechecks` will disable checking for whether the passed in variable is
  already defined, which is useful when adding many new variables to the system.
  *Do not disable checks* unless you are sure the passed in variable is a new
  variable, as this will potentially leave the system in an undefined state.
"""
function addspecies!(network::ReactionSystem, s::Symbolic; disablechecks=false)

    # we don't check subsystems since we will add it to the top-level system...
    curidx = disablechecks ? nothing : findfirst(S -> isequal(S, s), ModelingToolkit.ModelingToolkit.get_states(network))
    if curidx === nothing
        push!(ModelingToolkit.ModelingToolkit.get_states(network), s)
        return length(ModelingToolkit.ModelingToolkit.get_states(network))
    else
        return curidx
    end
end

"""
    addspecies!(network::ReactionSystem, s::Num; disablechecks=false)

Given a [`ReactionSystem`](@ref), add the species corresponding to the
variable `s` to the network (if it is not already defined). Returns the
integer id of the species within the system.

- `disablechecks` will disable checking for whether the passed in variable is
  already defined, which is useful when adding many new variables to the system.
  *Do not disable checks* unless you are sure the passed in variable is a new
  variable, as this will potentially leave the system in an undefined state.
"""
function addspecies!(network::ReactionSystem, s::Num; disablechecks=false)
    addspecies!(network, value(s), disablechecks=disablechecks)
end

"""
    addparam!(network::ReactionSystem, p::Symbolic; disablechecks=false)

Given a [`ReactionSystem`](@ref), add the parameter corresponding to the
variable `p` to the network (if it is not already defined). Returns the integer
id of the parameter within the system.

- `disablechecks` will disable checking for whether the passed in variable is
  already defined, which is useful when adding many new variables to the system.
  *Do not disable checks* unless you are sure the passed in variable is a new
  variable, as this will potentially leave the system in an undefined state.
"""
function addparam!(network::ReactionSystem, p::Symbolic; disablechecks=false)
    # we don't check subsystems since we will add it to the top-level system...
    if istree(p) && !(operation(p) isa Sym)
        error("If the passed in parameter is an expression, it must correspond to an underlying Variable.")
    end
    curidx = disablechecks ? nothing : findfirst(S -> isequal(S, p), ModelingToolkit.get_ps(network))
    if curidx === nothing
        push!(ModelingToolkit.get_ps(network), p)
        return length(ModelingToolkit.get_ps(network))
    else
        return curidx
    end
end

"""
    addparam!(network::ReactionSystem, p::Num; disablechecks=false)

Given a [`ReactionSystem`](@ref), add the parameter corresponding to the
variable `p` to the network (if it is not already defined). Returns the
integer id of the parameter within the system.

- `disablechecks` will disable checking for whether the passed in variable is
  already defined, which is useful when adding many new variables to the system.
  *Do not disable checks* unless you are sure the passed in variable is a new
  variable, as this will potentially leave the system in an undefined state.
"""
function addparam!(network::ReactionSystem, p::Num; disablechecks=false)
    addparam!(network, value(p); disablechecks=disablechecks)
end

"""
    addreaction!(network::ReactionSystem, rx::Reaction)

Add the passed in reaction to the [`ReactionSystem`](@ref). Returns the
integer id of `rx` in the list of `Reaction`s within `network`.

Notes:
- Any new species or parameters used in `rx` should be separately added to
    `network` using [`addspecies!`](@ref) and [`addparam!`](@ref).
"""
function addreaction!(network::ReactionSystem, rx::Reaction)
    push!(ModelingToolkit.get_eqs(network), rx)
    length(ModelingToolkit.get_eqs(network))
end

"""
    merge!(network1::ReactionSystem, network2::ReactionSystem)

Merge `network2` into `network1`.

Notes:
- Duplicate reactions between the two networks are not filtered out.
- [`Reaction`](@ref)s are not deepcopied to minimize allocations, so both
  networks will share underlying data arrays.
- Subsystems are not deepcopied between the two networks and will hence be
  shared.
- Returns `network1`.
- Does not currently handle pins.
"""
function merge!(network1::ReactionSystem, network2::ReactionSystem)
    isequal(ModelingToolkit.independent_variable(network1),
            ModelingToolkit.independent_variable(network2)) || error("Reaction networks must have the same independent variable to be mergable.")
    union!(network1.states, network2.states)
    union!(network1.ps, network2.ps)
    append!(network1.eqs, network2.eqs)
    append!(network1.systems, network2.systems)
    network1
end

"""
    merge(network1::ReactionSystem, network2::ReactionSystem)

Create a new network merging `network1` and `network2`.

Notes:
- Duplicate reactions between the two networks are not filtered out.
- [`Reaction`](@ref)s are not deepcopied to minimize allocations, so the new
  network will share underlying data arrays.
- Subsystems are not deepcopied between the two networks and will hence be
  shared.
- Returns the merged network.
- Does not currently handle pins.
"""
function merge(network1::ReactionSystem, network2::ReactionSystem)
    network = make_empty_network()
    merge!(network, network1)
    merge!(network, network2)
    network
end