text:
#=
--------------------------------------------------------------------------------------------------------------------------------
Subsampling: Jackknife
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    block_jackknife(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64)

Generate block jackknife (Kunsch, 1989) samples. This implementation is described in Pellegrino (2020).

This technique subsamples a time series dataset by removing, in turn, all the blocks of consecutive observations with a given size.

# Arguments
- `Y`: Observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: Block size as a percentage of number of observed periods. It is bounded between 0 and 1.

# References
Kunsch (1989) and Pellegrino (2020).
"""
function block_jackknife(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64)

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Dimensions
    n, T = size(Y);

    # Block size
    block_size = Int64(ceil(subsample*T));
    if block_size == 0
        error("subsample is too small!");
    end

    # Number of block jackknifes samples - as in Kunsch (1989)
    samples = T-block_size+1;

    # Initialise jackknife_data
    jackknife_data = JArray{Float64,3}(undef, n, T, samples);

    # Loop over j=1, ..., samples
    for j=1:samples

        # Index of missings
        ind_j = collect(j:j+block_size-1);

        # Block jackknife data
        jackknife_data[:, :, j] = Y;
        jackknife_data[:, ind_j, j] .= missing;
    end

    # Return jackknife_data
    return jackknife_data;
end

"""
    optimal_d(n::Int64, T::Int64)

Select the optimal value for d. See artificial_jackknife (...) for more details on d.

# Arguments
- `n`: Number of series
- `T`: Number of observations
"""
function optimal_d(n::Int64, T::Int64)
    objfun_array = zeros(n*T);

    for d=1:fld(n*T,2)
        objfun_array[d] = objfun_optimal_d(n, T, d);
    end

    return argmax(objfun_array);
end

"""
    objfun_optimal_d(n::Int64, T::Int64, d::Int64)

Objective function to select the optimal value for d.

# Arguments
- `n`: Number of series
- `T`: Number of observations
- `d`: Candidate d
"""
function objfun_optimal_d(n::Int64, T::Int64, d::Int64)
    fun = no_combinations(n*T, d) - (d>=n).*no_combinations(n*T-n, d-n)*T;

    for i=2:fld(d, n)
        fun -= (-1^(i-1))*no_combinations(T, i)*no_combinations(n*T-i*n, d-i*n);
    end

    return fun;
end

"""
    artificial_jackknife(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, max_samples::Int64, seed::Int64=1)

Generate artificial jackknife samples as in Pellegrino (2020).

The artificial delete-d jackknife is an extension of the delete-d jackknife for dependent data problems.
- This technique replaces the actual data removal step with a fictitious deletion, which consists of imposing `d`-dimensional (artificial) patterns of missing observations to the data.
- This approach does not alter the data order nor destroy the correlation structure.

# Arguments
- `Y`: Observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: `d` as a percentage of the original sample size. It is bounded between 0 and 1.
- `max_samples`: If `C(n*T,d)` is large, artificial_jackknife would generate `max_samples` jackknife samples.
- `seed`: Random seed (default: 1).

# References
Pellegrino (2020).
"""
function artificial_jackknife(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, max_samples::Int64, seed::Int64=1)

    # Dimensions
    n, T = size(Y);
    nT = n*T;

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Set d using subsample
    d = Int64(ceil(subsample*nT));

    # Error management
    if d == 0
        error("subsample is too small!");
    end

    # Warning
    if subsample > 0.5
        @warn "this algorithm might be unstable for `subsample` larger than 0.5!";
    end

    # Get vec(Y)
    vec_Y = Y[:] |> JVector{Float64};

    # Initialise loop (controls)
    C_nT_d = no_combinations(nT, d);
    samples = convert(Int64, min(C_nT_d, max_samples));
    zeros_vec = zeros(d);

    # Initialise loop (output)
    ind_missings = Array{Int64}(zeros(d, samples));
    jackknife_data = JArray{Float64,3}(undef, n, T, samples);

    # Set `rng`
    rng = StableRNG(seed);

    # Loop over j=1, ..., samples
    for j=1:samples

        # First draw
        if j == 1
            if samples == C_nT_d
                ind_missings[:,j] = rand_without_replacement(rng, nT, d);
            else
                ind_missings[:,j] = rand_without_replacement(rng, n, T, d);
            end

        # Iterates until ind_missings[:,j] is neither a vector of zeros, nor already included in ind_missings
        else
            while ind_missings[:,j] == zeros_vec || is_vector_in_matrix(ind_missings[:,j], ind_missings[:, 1:j-1])
                if samples == C_nT_d
                    ind_missings[:,j] = rand_without_replacement(rng, nT, d);
                else
                    ind_missings[:,j] = rand_without_replacement(rng, n, T, d);
                end
            end
        end

        # Add (artificial) missing observations
        vec_Y_j = copy(vec_Y);
        vec_Y_j[ind_missings[:,j]] .= missing;

        # Store data
        jackknife_data[:, :, j] = reshape(vec_Y_j, n, T);
    end

    # Return jackknife_data
    return jackknife_data;
end

#=
--------------------------------------------------------------------------------------------------------------------------------
Subsampling: Bootstrap
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    moving_block_bootstrap(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, samples::Int64, seed::Int64=1)

Generate moving block bootstrap samples.

The moving block bootstrap randomly subsamples a time series into ordered and overlapped blocks of consecutive observations.

# Arguments
- `Y`: Observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: Block size as a percentage of number of observed periods. It is bounded between 0 and 1.
- `samples`: Number of bootstrap samples.
- `seed`: Random seed (default: 1).

# References
Kunsch (1989) and Liu and Singh (1992).
"""
function moving_block_bootstrap(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, samples::Int64, seed::Int64=1)

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Dimensions
    n, T = size(Y);

    # Block size
    block_size = Int64(ceil(subsample*T));
    if block_size == 0
        error("subsample is too small!");
    end

    # Initialise bootstrap_data
    bootstrap_data = JArray{Float64,3}(undef, n, block_size, samples);

    # Set `rng`
    rng = StableRNG(seed);

    # Loop over j=1, ..., samples
    for j=1:samples

        # Starting point for the moving block
        ind_j = rand(rng, 1:T-block_size+1);

        # Bootstrap data
        bootstrap_data[:, :, j] .= Y[:, ind_j:ind_j+block_size-1];
    end

    # Return bootstrap_data
    return bootstrap_data;
end

"""
    stationary_block_bootstrap(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, samples::Int64, seed::Int64=1)

Generate stationary block bootstrap samples.

The stationary bootstrap is similar to the block bootstrap proposed in independently in Kunsch (1989) and Liu and Singh (1992).

There are two main differences:
- The blocks have random length
- In order to achieve stationarity, the stationary (block) bootstrap "wraps" the data around in a "circle" so that the first observation follows the last.

Note: Block size is exponentially distributed with mean `Int64(ceil(subsample*T))`.

# Arguments
- `Y`: Observed measurements (`nxT`), where `n` and `T` are the number of series and observations.
- `subsample`: Block size as a percentage of number of observed periods. It is bounded between 0 and 1.
- `samples`: Number of bootstrap samples.
- `seed`: Random seed (default: 1).

# References
Politis and Romano (1994).
"""
function stationary_block_bootstrap(Y::Union{FloatMatrix, JMatrix{Float64}}, subsample::Float64, samples::Int64, seed::Int64=1)

    # Check inputs
    check_bounds(subsample, 0, 1);

    # Dimensions
    n, T = size(Y);

    # Block length is exponentially distributed with mean
    avg_block_size = Int64(ceil(subsample*T));
    if avg_block_size == 0
        error("subsample is too small!");
    end

    # Initialise bootstrap_data
    bootstrap_data = JArray{Float64,3}(undef, n, T, samples);

    # Set `rng`
    rng = StableRNG(seed);

    # Loop over j=1, ..., samples
    for j=1:samples

        # Merge multiple blocks of random size
        ind_j    = zeros(T) |> Array{Int64,1};
        ind_j[1] = rand(rng, 1:T);

        # Loop over t=2,...,T
        for t=2:T

            # Let ind_j[t] be picked at random
            if rand(rng) < 1/avg_block_size;
                ind_j[t] = rand(rng, 1:T);

            # Let ind_j[t] be ind_j[t-1] + 1
            else
                if ind_j[t-1] == T
                    ind_j[t] = 1;
                else
                    ind_j[t] = ind_j[t-1] + 1;
                end
            end
        end

        # Generate j-th bootstrap sample
        bootstrap_data[:, :, j] .= Y[:, ind_j];
    end

    # Return bootstrap_data
    return bootstrap_data;
end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    MetropolisHastings

Metropolis-Hastings MCMC sampling algorithm.
!!!!!
"""
struct MetropolisHastings{
    Q<:ProposalDistSpec,
    W<:Real,
    WS<:AbstractWeightingScheme{W}
} <: MCMCAlgorithm
    proposalspec::Q
    weighting::WS

    MetropolisHastings(proposalspec::Q,weighting::WS) where {
        Q<:ProposalDistSpec, W<:Real, WS<:AbstractWeightingScheme{W}} =
        new{Q,W,WS}(proposalspec, weighting)
end

export MetropolisHastings

MetropolisHastings(proposalspec::ProposalDistSpec = MvTDistProposal()) =
    MetropolisHastings(proposalspec, RepetitionWeighting())

MetropolisHastings(weighting::AbstractWeightingScheme) =
    MetropolisHastings(MvTDistProposal(), weighting)

mcmc_compatible(::MetropolisHastings, ::AbstractProposalDist, ::NoParamBounds) = true

mcmc_compatible(::MetropolisHastings, proposaldist::AbstractProposalDist, bounds::HyperRectBounds) =
    issymmetric(proposaldist) || all(x -> x == hard_bounds, bounds.bt)

_sample_weight_type(::Type{MetropolisHastings{Q,W,WS}}) where {Q,W,WS} = W

mutable struct MHIterator{
    SP<:MCMCSpec,
    R<:AbstractRNG,
    Q<:AbstractProposalDist,
    SV<:DensitySampleVector
} <: MCMCIterator
    spec::SP
    rng::R
    info::MCMCIteratorInfo
    proposaldist::Q
    samples::SV
    nsamples::Int64
    stepno::Int64
end

function MHIterator(
    spec::MCMCSpec,
    info::MCMCIteratorInfo,
    initial_params::AbstractVector{P},
) where {P<:Real}
    stepno::Int64 = 0
    rng = spec.rngseed()
    reset_rng_counters!(rng, info.id, info.cycle, stepno)

    postr = spec.posterior
    npar = nparams(postr)
    alg = spec.algorithm

    params_vec = Vector{P}(undef, npar)
    if isempty(initial_params)
        initial_params!(params_vec, rng, postr, alg)
    else
        params_vec .= initial_params
    end
    !(params_vec in param_bounds(postr)) && throw(ArgumentError("Parameter(s) out of bounds"))

    proposaldist = alg.proposalspec(P, npar)

    # ToDo: Make numeric type configurable:

    log_posterior_value = apply_bounds_and_eval_posterior_logval_strict!(postr, params_vec)

    T = typeof(log_posterior_value)
    W = _sample_weight_type(typeof(alg))

    sample_info = MCMCSampleID(info.id, info.cycle, 1, CURRENT_SAMPLE)
    current_sample = DensitySample(params_vec, log_posterior_value, one(W), sample_info, nothing)

    samples = DensitySampleVector{P,T,W,MCMCSampleID,Nothing}(undef, 0, npar)
    push!(samples, current_sample)

    nsamples::Int64 = 0

    chain = MHIterator(
        spec,
        rng,
        info,
        proposaldist,
        samples,
        nsamples,
        stepno
    )

    chain
end

function reset_rng_counters!(chain::MHIterator)
    reset_rng_counters!(chain.rng, chain.info.id, chain.info.cycle, chain.stepno)
end

function (spec::MCMCSpec{<:MetropolisHastings})(
    chainid::Integer,
)
    P = float(eltype(param_bounds(spec.posterior)))

    cycle = 0
    tuned = false
    converged = false
    info = MCMCIteratorInfo(chainid, cycle, tuned, converged)

    MHIterator(spec, info, Vector{P}())
end

@inline _current_sample_idx(chain::MHIterator) = firstindex(chain.samples)
@inline _proposed_sample_idx(chain::MHIterator) = lastindex(chain.samples)

function _available_samples_idxs(chain::MHIterator)
    sampletype = chain.samples.info.sampletype
    @uviews sampletype begin
        from = firstindex(chain.samples)

        to = if samples_available(chain)
            lastidx = lastindex(chain.samples)
            @assert sampletype[from] == ACCEPTED_SAMPLE
            @assert sampletype[lastidx] == CURRENT_SAMPLE
            lastidx - 1
        else
            from - 1
        end

        r = from:to
        @assert all(x -> x > INVALID_SAMPLE, view(sampletype, r))
        r
    end
end

mcmc_spec(chain::MHIterator) = chain.spec

getrng(chain::MHIterator) = chain.rng

mcmc_info(chain::MHIterator) = chain.info

nsteps(chain::MHIterator) = chain.stepno

nsamples(chain::MHIterator) = chain.nsamples

current_sample(chain::MHIterator) = chain.samples[_current_sample_idx(chain)]

sample_type(chain::MHIterator) = eltype(chain.samples)

function samples_available(chain::MHIterator)
    i = _current_sample_idx(chain::MHIterator)
    chain.samples.info.sampletype[i] == ACCEPTED_SAMPLE
end

function get_samples!(appendable, chain::MHIterator, nonzero_weights::Bool)::typeof(appendable)
    if samples_available(chain)
        idxs = _available_samples_idxs(chain)
        samples = chain.samples

        @uviews samples begin
            # if nonzero_weights
                for i in idxs
                    if !nonzero_weights || samples.weight[i] > 0
                        push!(appendable, samples[i])
                    end
                end
            # else
            #     append!(appendable, view(samples, idxs))
            # end
        end
    end
    appendable
end

function next_cycle!(chain::MHIterator)
    chain.info = MCMCIteratorInfo(chain.info, cycle = chain.info.cycle + 1)
    chain.nsamples = 0
    chain.stepno = 0

    reset_rng_counters!(chain)

    resize!(chain.samples, 1)

    i = _current_sample_idx(chain)
    @assert chain.samples.info[i].sampletype == CURRENT_SAMPLE

    chain.samples.weight[i] = 1
    chain.samples.info[i] = MCMCSampleID(chain.info.id, chain.info.cycle, chain.stepno, CURRENT_SAMPLE)

    chain
end

function mcmc_step!(
    callback::AbstractMCMCCallback,
    chain::MHIterator
)
    alg = algorithm(chain)

    if !mcmc_compatible(alg, chain.proposaldist, param_bounds(getposterior(chain)))
        error("Implementation of algorithm $alg does not support current parameter bounds with current proposal distribution")
    end

    chain.stepno += 1
    reset_rng_counters!(chain)

    rng = getrng(chain)
    pstr = getposterior(chain)

    proposaldist = chain.proposaldist
    samples = chain.samples

    # Grow samples vector by one:
    resize!(samples, size(samples, 1) + 1)
    samples.info[lastindex(samples)] = MCMCSampleID(chain.info.id, chain.info.cycle, chain.stepno, PROPOSED_SAMPLE)

    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    @assert current != proposed

    accepted = @uviews samples begin
        current_params = samples.params[current]
        proposed_params = samples.params[proposed]

        # Propose new parameters:
        samples.weight[proposed] = 0
        proposal_rand!(rng, proposaldist, proposed_params, current_params)

        current_log_posterior = samples.logdensity[current]
        T = typeof(current_log_posterior)

        # Evaluate prior and likelihood with proposed parameters:
        proposed_log_posterior = apply_bounds_and_eval_posterior_logval!(T, pstr, proposed_params)

        samples.logdensity[proposed] = proposed_log_posterior

        p_accept = if proposed_log_posterior > -Inf
            # log of ratio of forward/reverse transition probability
            log_tpr = if issymmetric(proposaldist)
                T(0)
            else
                log_tp_fwd = distribution_logpdf(proposaldist, proposed_params, current_params)
                log_tp_rev = distribution_logpdf(proposaldist, current_params, proposed_params)
                T(log_tp_fwd - log_tp_rev)
            end

            p_accept_unclamped = exp(proposed_log_posterior - current_log_posterior - log_tpr)
            T(clamp(p_accept_unclamped, 0, 1))
        else
            zero(T)
        end

        @assert p_accept >= 0
        accepted = rand(chain.rng, float(typeof(p_accept))) < p_accept

        if accepted
            samples.info.sampletype[current] = ACCEPTED_SAMPLE
            samples.info.sampletype[proposed] = CURRENT_SAMPLE
            chain.nsamples += 1
        else
            samples.info.sampletype[proposed] = REJECTED_SAMPLE
        end

        delta_w_current, w_proposed = _mh_weights(alg, p_accept, accepted)
        samples.weight[current] += delta_w_current
        samples.weight[proposed] = w_proposed

        callback(1, chain)

        if accepted
            current_params .= proposed_params
            samples.logdensity[current] = samples.logdensity[proposed]
            samples.weight[current] = samples.weight[proposed]
            samples.info[current] = samples.info[proposed]
        end

        accepted
    end # @uviews

    if accepted
        resize!(samples, 1)
    end

    chain
end

function _mh_weights(
    algorithm::MetropolisHastings{Q,W,<:RepetitionWeighting},
    p_accept::Real,
    accepted::Bool
) where {Q,W}
    if accepted
        (0, 1)
    else
        (1, 0)
    end
end

function _mh_weights(
    algorithm::MetropolisHastings{Q,W,<:ARPWeighting},
    p_accept::Real,
    accepted::Bool
) where {Q,W}
    T = typeof(p_accept)
    if p_accept ≈ 1
        (zero(T), one(T))
    elseif p_accept ≈ 0
        (one(T), zero(T))
    else
        (T(1 - p_accept), p_accept)
    end
end

eff_acceptance_ratio(chain::MHIterator) = nsamples(chain) / nsteps(chain)

ext: jl


text:
abstract type RedisException<:Exception end

# Thrown if the client is unable to establish a connection to the server
struct ConnectionException <: RedisException
    message::AbstractString
end

# Thrown if the response from the server doesn't conform to RESP
struct ProtocolException <: RedisException
    message::AbstractString
end

# Thrown if the server returns an error response (RESP error)
# merl-dev: removed error_prefix as a separate member, the entire message
# is returned as a single string.  We could parse the message on the first word,
# however, as per http://redis.io/topics/protocol:
#   "This is called an Error Prefix and is a way to allow the client to understand
#   the kind of error returned by the server without to rely on the exact message given,
#   that **may change over the time**.
#   A client implementation may return different kind of exceptions for different
#   errors, or may provide a generic way to trap errors by directly providing the
#   error name to the caller as a string.
#   However, **such a feature should not be considered vital as it is rarely useful**,
#   and a limited client implementation may simply return a generic error condition,
#   such as false.
struct ServerException <: RedisException
    # error_prefix::AbstractString
    message::AbstractString
end

# Thrown if an error originates from the client
struct ClientException <: RedisException
    message::AbstractString
end


text:
module Dictionaries

using Random
using Indexing
using Base: @propagate_inbounds, Callable

export getindices, setindices!

export AbstractDictionary, AbstractIndices, IndexError, ArrayIndices, Indices, Dictionary, ArrayDictionary, MappedDictionary, ReverseIndices, ReverseDictionary, DictionaryView, FilteredDictionary, FilteredIndices, BroadcastedDictionary, FillDictionary

export dictionary, index, distinct, disjoint, isdictequal, filterview, sortkeys, sortpairs, sortkeys!, sortpairs!
export issettable, isinsertable, set!, unset!
export istokenizable, tokentype, tokens, tokenized, gettoken, gettokenvalue, istokenassigned, settokenvalue!, gettoken!, deletetoken!, sharetokens

include("AbstractDictionary.jl")
include("AbstractIndices.jl")

include("tokens.jl")
include("indexing.jl")
include("insertion.jl")
include("reverse.jl")
include("filter.jl")
include("foreach.jl")
include("map.jl")
include("broadcast.jl")
include("find.jl")
include("show.jl")
include("pairs.jl")

include("ArrayIndices.jl")
include("ArrayDictionary.jl")
include("Indices.jl")
include("Dictionary.jl")
include("FillDictionary.jl")

Base.@deprecate_binding HashIndices Indices true
Base.@deprecate_binding HashDictionary Dictionary true

end # module

# # TODO
#
# * Improved printing - don't calculate length (beyond some cutoff) if it is `SizeUnknown` and limit=true, fix indentiation problems for wider values
# * TODO: have `delete!` return next key, `deletetoken!` return next token.
#   For these kinds of algorithms, probably need: firstindex, firsttoken, nextind, prevind,
#   nexttoken, prevtoken, lastindex, lasttoken.
# * A surface interface for updates like https://github.com/JuliaLang/julia/pull/31367
# * More operations for "ordered" indices/sets (sort-based dictionaries and
#   B-trees). We can probably formalize an interface around a trait here. Certain operations
#   like slicing out an interval or performing a sort-merge co-iteration for `merge` become
#   feasible.


text:
__precompile__()

module Falcon

using OhMyJulia
using Insane
using StatsBase
using HypothesisTests
using BufferedStreams
using Libz
using JsonBuilder
using BioDataStructures
import Base: start, next, done, iteratorsize, eltype, view,
             getindex, setindex!, show, ==, hash, write

include("core/mut.jl")
include("core/read.jl")
include("core/bam.jl")
include("core/sam.jl")

include("index/pair.jl")
include("index/pileup.jl")
include("index/index.jl")

include("cover/bed.jl")
include("cover/cover.jl")

include("ui/stat.jl")

include("caller/rule.jl")
include("caller/driver.jl")
include("caller/vcf.jl")

include("bamqc/cover.jl")
end # module


text:
using vOptSpecific

using Test

println("Testing 2LAP...")
C1 = [3 9 0 0 6 10 7 5 16 11 ;
16 0 6 12 19 8 17 10 9 18 ;
2 7 11 15 8 10 10 12 8 6 ;
4 11 7 16 3 13 18 11 11 7 ;
14 19 7 0 4 19 8 13 9 10 ;
11 4 17 14 19 5 16 2 17 4 ;
8 14 7 15 2 11 0 1 14 11 ;
8 8 3 15 8 7 14 9 0 16 ;
11 3 12 8 9 11 6 5 5 15 ;
2 5 1 9 0 4 12 13 5 6 ]

C2 = [16 5 6 19 12 7 18 19 16 10 ;
15 7 13 7 7 2 10 5 0 8 ;
1 2 13 2 3 6 6 16 19 3 ;
14 7 8 1 7 13 8 17 12 16 ;
8 19 15 1 18 14 16 8 0 8 ;
8 1 10 14 15 5 0 14 1 11 ;
9 16 10 10 9 17 3 17 15 15 ;
5 15 14 0 16 12 14 4 12 6 ;
12 1 19 14 15 15 0 7 1 13 ;
10 10 1 0 0 10 10 3 19 17 ]

id = set2LAP(10, C1, C2)
z1,z2,solutions = vSolve(id)

@test z1 == [18,19,28,35,43,54,66,22,26,34,51]
@test solutions[1,:] == [3, 2, 1, 5, 4, 10, 7, 9, 8, 6]

id = load2LAP("../examples/2AP10-1A100.dat")
z1,z2,solutions = vSolve(id)

@test z1 == [146, 155, 185, 271, 363, 478, 572, 149, 220, 267, 307, 402, 433]
@test solutions[1,:] == [6, 3, 7, 9, 2, 10, 1, 5, 4, 8]

println("\nTesting 2UKP...")
id = set2UKP([1,2,3], [4,5,6], [7,8,9], 16)
z1, z2, w, solutions = vSolve(id)

@test z1 == [4]
@test w == [16]
@test solutions == [[1, 0, 1]]

id = load2UKP("../examples/2KP500-1A.DAT")
z1, z2, w, solutions = vSolve(id)
@assert length(z1) == 1682
@assert extrema(z1) == (16028, 20360)
@assert allunique(solutions)

println("\nTesting 2UMFLP...")
id = load2UMFLP("../examples/F50-51.txt")
z1, z2, facility_res, X, isEdge = vSolve(id)
@test length(z1) == 198
@test count(!iszero, X[21]) == 91
@test count(isEdge) == 188

println("\nTesting 2OSP...")
id = set2OSP(4, [2, 4, 3, 1], [1, 2, 4, 6])
z1, z2 , S = vSolve(id)

for i = 1:length(S)
    println("(", z1[i], ", ", z2[i], ") : ", S[i])
end

@test z1 == [20, 21, 27]
@test z2 == [8, 6, 5]
@test S == [[4, 1, 3, 2], [4, 1, 2, 3], [1, 2, 3, 4]]

max_stars_repo_name: vOptSolver/vOptSpecific

max_issues_repo_name: vOptSolver/vOptSpecific

max_forks_repo_name: vOptSolver/vOptSpecific


text:
import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Normal, truncated, Poisson, InverseWishart, MvNormal, Dirichlet, Categorical
using Logging
using StatsBase: mean, std
using Random: seed!
using PrettyPrint: pprintln

seed!(2021)

function eye(d :: Int)
    mat = zeros(Float64, d, d)
    for i in 1:d
        mat[i, i] = 1.0
    end
    mat
end

function clustering(t :: Trace, data :: Array{Float64, 2}, cluster_guess :: Int)
    # data has shape (D, N), where D is dimensionalty and N is num observations
    D = size(data, 1)
    N = size(data, 2)
    # first, choosing the number of clusters. Prior penalizes high number of clusters
    cluster_guess = max(cluster_guess, 1)
    n_clusters = sample(t, :n_clusters, truncated(Poisson(cluster_guess), 1, Inf))
    # second, given the number of clusters, define loc and cov for each
    covs = Array{Array{Float64, 2}, 1}()
    locs = Array{Array{Float64, 1}, 1}()
    id_mat = eye(D)
    for c in 1:n_clusters
        cov = sample(t, (:cov, c), InverseWishart(D + 1, id_mat))
        push!(covs, id_mat)
        loc = sample(t, (:loc, c), MvNormal(zeros(D), cov))
        push!(locs, loc)
    end
    # third, for each datapoint, assign to a cluster
    # note that each datapoint has its own local rv -- we could introduce an amortized model
    # that learns a mapping from datapoint features to cluster
    for n in 1:N
        cluster_prob = sample(t, (:cluster_prob, n), Dirichlet(n_clusters, 1.0))
        this_cluster = sample(t, (:this_cluster, n), Categorical(cluster_prob))
        observe(t, (:data, n), MvNormal(locs[this_cluster], covs[this_cluster]), data[:, n])
    end
end

function main()
    @info "Open-universe clustering inference"
    D = 4
    n1 = 5
    n2 = 12
    data1 = rand(D, n1)
    data2 = rand(D, n2) .+ 6
    data = hcat(data1, data2)
    @info "True number of clusters is 2"
    # pretend we don't know dgp, bias it to be higher than ground truth
    n_clusters_guess = 5
    @time results = mh(clustering; params = (data, n_clusters_guess), burn=2000, thin=100, num_iterations=100000)
    mean_n_clusters = mean(results, :n_clusters)
    std_n_clusters = std(results, :n_clusters)
    @info "posterior n_clusters ≈ $mean_n_clusters ± $std_n_clusters"
end

main()


text:
module StochasticIntegrators

    using HDF5
    using LinearAlgebra: dot, mul!
    using OffsetArrays
    using Reexport
    using SimpleSolvers

    @reexport using GeometricIntegrators.Common
    @reexport using GeometricIntegrators.Config
    @reexport using GeometricIntegrators.Utils

    import GeometricIntegrators.Equations: AbstractEquationSDE, AbstractEquationPSDE,
                                           SDE, PSDE, SPSDE, get_function_tuple

    import GeometricIntegrators.Integrators
    import GeometricIntegrators.Integrators: Integrator, Parameters

    import GeometricIntegrators.Integrators: integrate!, integrate_common!, get_internal_variables

    import GeometricIntegrators.Integrators: IntegratorCache, CacheDict, CacheType

    import GeometricIntegrators.Integrators: AbstractTableauERK, AbstractTableauIRK, CoefficientsRK

    import GeometricIntegrators.Integrators: create_internal_stage_vector, create_internal_stage_matrix,
                                             create_internal_stage_vector_with_zero, create_nonlinear_solver

    import GeometricIntegrators.Solutions

    import GeometricIntegrators.Solutions: AtomicSolution, Solution, ParallelSolution

    import GeometricIntegrators.Solutions: TimeSeries, DataSeries, SDataSeries, PDataSeries

    import GeometricIntegrators.Solutions: hdf5, timesteps, nsave, counter, offset, lastentry,
                                           DEFAULT_NSAVE, DEFAULT_NWRITE

    import GeometricIntegrators.Solutions: get_initial_conditions, get_initial_conditions!, set_initial_conditions!,
                                           get_solution, get_solution!, set_solution!,
                                           get_data!, set_data!, update!, compute_timeseries!,
                                           createHDF5, create_hdf5, create_hdf5!, save_attributes,
                                           init_timeteps_in_hdf5, init_solution_in_hdf5, copy_solution_to_hdf5

    import GeometricIntegrators.Tableaus: CoefficientsGLRK, CoefficientsLobattoIIIA, CoefficientsLobattoIIIB, CoefficientsLobattoIIID

    const DEFAULT_SCONV = :strong

    export SemiMartingale, WienerProcess, generate_wienerprocess!, conv

    include("solutions/wienerprocess.jl")

    export AtomicSolution, AtomicSolutionSDE, AtomicSolutionPSDE
    export update!, cut_periodic_solution!,
           get_increments, get_increments!, set_increments!

    include("solutions/atomic_solution_sde.jl")
    include("solutions/atomic_solution_psde.jl")

    export Solution, StochasticSolution
    export SolutionSDE, SSolutionSDE, PSolutionSDE, SolutionPSDE, SSolutionPSDE, PSolutionPSDE
    export hdf5, timesteps, nsave, counter, offset, lastentry
    export get_initial_conditions, get_initial_conditions!, set_initial_conditions!,
           get_solution, get_solution!, set_solution!,
           create_hdf5, create_hdf5!

    include("solutions/solution.jl")
    include("solutions/solution_sde.jl")
    include("solutions/solution_psde.jl")
    include("solutions/solution_hdf5.jl")

    include("solutions/atomic_solution_constructors.jl")

    export Parameters

    export Integrator, StochasticIntegrator
    export SDEIntegrator, PSDEIntegrator,
           StochasticIntegratorRK, StochasticIntegratorPRK

    export nstages, noisedims, integrate!

    export IntegratorSERK, TableauSERK
    export IntegratorSIRK, TableauSIRK
    export IntegratorSIPRK, TableauSIPRK
    export IntegratorSISPRK, TableauSISPRK
    export IntegratorWERK, TableauWERK
    export IntegratorWIRK, TableauWIRK

    include("integrators/abstract_integrators.jl")
    include("integrators/abstract_runge_kutta.jl")
    include("integrators/integrator_cache.jl")

    include("integrators/integrators_serk.jl")
    include("integrators/integrators_sirk.jl")
    include("integrators/integrators_siprk.jl")
    include("integrators/integrators_sisprk.jl")
    include("integrators/integrators_werk.jl")
    include("integrators/integrators_wirk.jl")
    include("integrators/common.jl")
    include("integrators/integrators.jl")

    include("tableaus/tableaus_sirk.jl")

    export TableauStochasticGLRK, TableauStochasticDIRK

    include("tableaus/tableaus_siprk.jl")

    export TableauStochasticStoermerVerlet, TableauStochasticSymplecticEuler

    include("tableaus/tableaus_sisprk.jl")

    export TableauStochasticLobattoIIIABD2, TableauModifiedStochasticStoermerVerlet

    include("tableaus/tableaus_serk.jl")

    export TableauPlaten, TableauBurrageR2, TableauBurrageCL
    export TableauBurrageE1, TableauBurrageG5, TableauStochasticHeun
    export TableauStochasticEuler

    include("tableaus/tableaus_werk.jl")

    export TableauRoesslerRS1, TableauRoesslerRS2

    include("tableaus/tableaus_wirk.jl")

    export TableauSRKw1, TableauSRKw2

end


text:
using Zygote
using Random
using DelimitedFiles
using LinearAlgebra
using GalacticOptim,Optim,NLopt#,BlackBoxOptim
using DifferentialEquations#,DiffEqFlux
#using ForwardDiff
using Plots
using LaTeXStrings
#using GLOQ
#using ReverseDiff
include("../src/GLOQ.jl")
pyplot()
# BlackBoxOptim somehow downgrade some packages and as a result breaks the auto-differentiation with Zygote
# we should avoid it.
# NLopt seems to be really fast. Their LBFGS with box constrains is faster than Optim's FMinbox(LBFGS)

N_states = 4;
freqs = [4.0108; 3.8830; 3.6287]
omegas = 2.0*pi.*freqs
gamma1   = [2.222222e-05; 4.761904e-05; 4.545455e-5]
gamma2   = [4.166667e-05; 6.8965517e-5; 2.3255814e-4]

# driving frequency
omr_t1_01 = 2.0*pi*4.0108
omr_t1_12 = 2.0*pi*3.8830
omr_t1_23 = 2.0*pi*3.6287

omr_ramsey_01 = 2.0*pi*(4.0108-5e-4)
omr_ramsey_12 = 2.0*pi*(3.8830-2.0e-3)
omr_ramsey_23 = 2.0*pi*(3.6287-2.0e-3)

omr_echo_01 = 2.0*pi*4.0108
omr_echo_12 = 2.0*pi*3.8830
omr_echo_23 = 2.0*pi*3.6287

# charge noise
char_noise = 2.0*pi.*[2e-6,6e-5,1.3e-4]
# width of the control signal
width = 21.25
TC_01 = 2.5*width

width = 17.00
TC_12 = 2.5*width

width = 17.00
TC_23 = 2.5*width

# in nano second
T_T1_01 = 125.0*1000.0
T_T1_12 = 125.0*1000.0
T_T1_23 = 125.0*1000.0
T_Ramsey_01 = 20.0*1000.0
T_Ramsey_12 = 20.0*1000.0
T_Ramsey_23 = 10.0*1000.0
T_Echo_01 = 40.0*1000.0
T_Echo_12 = 20.0*1000.0
T_Echo_23 = 10.0*1000.0

# data size used for characterization
n1use = 201
n2use = 201
n3use = 201
n4yse = 201
n4use = 201
n5use = 201
n6use = 201
n7use = 201
n8use = 201
n9use = 201

# time points evaluated
dt_t1_01 = T_T1_01/200
dt_t1_12 = T_T1_12/200
dt_t1_23 = T_T1_23/200
dt_ramsey_01 = T_Ramsey_01/200
dt_ramsey_12 = T_Ramsey_12/200
dt_ramsey_23= T_Ramsey_23/200
dt_echo_01 = T_Echo_01/200
dt_echo_12 = T_Echo_12/200
dt_echo_23 = T_Echo_23/200

t_t1_01 = zeros(Float64,n1use)
for i = 1:n1use
	t_t1_01[i] = (i-1)*dt_t1_01
end
T_evaluated_t1_01 = t_t1_01[end]

t_t1_12 = zeros(Float64,n2use)
for i = 1:n2use
	t_t1_12[i] = (i-1)*dt_t1_12
end
T_evaluated_t1_12 = t_t1_12[end]

t_t1_23 = zeros(Float64,n3use)
for i = 1:n3use
	t_t1_23[i] = (i-1)*dt_t1_23
end
T_evaluated_t1_23 = t_t1_23[end]

t_ramsey_01 = zeros(Float64,n4use)
for i = 1:n4use
	t_ramsey_01[i] = (i-1)*dt_ramsey_01
end
T_evaluated_ramsey_01 = t_ramsey_01[end]

t_ramsey_12 = zeros(Float64,n5use)
for i = 1:n5use
	t_ramsey_12[i] = (i-1)*dt_ramsey_12
end
T_evaluated_ramsey_12 = t_ramsey_12[end]

t_ramsey_23 = zeros(Float64,n6use)
for i = 1:n6use
	t_ramsey_23[i] = (i-1)*dt_ramsey_23
end
T_evaluated_ramsey_23 = t_ramsey_23[end]

t_echo_01 = zeros(Float64,n7use)
for i = 1:n7use
	t_echo_01[i] = (i-1)*dt_echo_01
end
T_evaluated_echo_01 = t_echo_01[end]

t_echo_12 = zeros(Float64,n8use)
for i = 1:n8use
	t_echo_12[i] = (i-1)*dt_echo_12
end
T_evaluated_echo_12 = t_echo_12[end]

t_echo_23 = zeros(Float64,n9use)
for i = 1:n9use
	t_echo_23[i] = (i-1)*dt_echo_23
end
T_evaluated_echo_23 = t_echo_23[end]

# initial conditions
u0_t1_01 = [1.0;0.0;0.0;0.0]
u0_t1_12 = [0.0;1.0;0.0;0.0]
u0_t1_23 = [0.0;0.0;1.0;0.0]
u0_ramsey_01 = [1.0;0.0;0.0;0.0]
u0_ramsey_12 = [0.0;1.0;0.0;0.0]
u0_ramsey_23 = [0.0;0.0;1.0;0.0]
u0_echo_01 = [1.0;0.0;0.0;0.0]
u0_echo_12 = [0.0;1.0;0.0;0.0]
u0_echo_23 = [0.0;0.0;1.0;0.0]

v0_t1_01 = [0.0;0.0;0.0;0.0]
v0_t1_12 = [0.0;0.0;0.0;0.0]
v0_t1_23 = [0.0;0.0;0.0;0.0]
v0_ramsey_01 = [0.0;0.0;0.0;0.0]
v0_ramsey_12 = [0.0;0.0;0.0;0.0]
v0_ramsey_23 = [0.0;0.0;0.0;0.0]
v0_echo_01 = [0.0;0.0;0.0;0.0]
v0_echo_12 = [0.0;0.0;0.0;0.0]
v0_echo_23 = [0.0;0.0;0.0;0.0]

########################################################
# ForwardSolves
########################################################
# Echo 0-1
rho_Echo01_u,rho_Echo01_v = GLOQ.EchoParityForwardSolve(u0_echo_01,v0_echo_01,
				 omegas,omr_echo_01,
				 char_noise,
				 gamma1,gamma2,
				 0,
				 TC_01,t_echo_01,N_states)
population_echo_01 = GLOQ.get_population(rho_Echo01_u)

# Echo 1-2
rho_Echo12_u,rho_Echo12_v = GLOQ.EchoParityForwardSolve(u0_echo_12,v0_echo_12,
				 omegas,omr_echo_12,
				 char_noise,
				 gamma1,gamma2,
				 1,
				 TC_12,t_echo_12,N_states)
population_echo_12 = GLOQ.get_population(rho_Echo12_u)

# Echo 2-3
rho_Echo23_u,rho_Echo23_v = GLOQ.EchoParityForwardSolve(u0_echo_23,v0_echo_23,
				 omegas,omr_echo_23,
				 char_noise,
				 gamma1,gamma2,
				 2,
				 TC_23,t_echo_23,N_states)
population_echo_23 = GLOQ.get_population(rho_Echo23_u)

# Ramsey 0-1
rho_Ramsey01_u,rho_Ramsey01_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_01,v0_ramsey_01,
				 omegas,omr_ramsey_01,
				 char_noise,
				 gamma1,gamma2,
				 0,
				 TC_01,t_ramsey_01,N_states)
population_ramsey_01 = GLOQ.get_population(rho_Ramsey01_u)

# Ramsey 1-2
rho_Ramsey12_u,rho_Ramsey12_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_12,v0_ramsey_12,
				 omegas,omr_ramsey_12,
				 char_noise,
				 gamma1,gamma2,
				 1,
				 TC_12,t_ramsey_12,N_states)
population_ramsey_12 = GLOQ.get_population(rho_Ramsey12_u)

# Ramsey 2-3
@time rho_Ramsey23_u,rho_Ramsey23_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_23,v0_ramsey_23,
				 omegas,omr_ramsey_23,
				 char_noise,
				 gamma1,gamma2,
				 2,
				 TC_23,t_ramsey_23,N_states)
population_ramsey_23 = GLOQ.get_population(rho_Ramsey23_u)

# T1 0-1
rho_T101_u,rho_T101_v = GLOQ.T1ParityForwardSolve(u0_t1_01,v0_t1_01,
				 omegas,omr_t1_01,
				 char_noise,
				 gamma1,gamma2,
				 0,
				 TC_01,t_t1_01,N_states)
population_t1_01 = GLOQ.get_population(rho_T101_u)

# T1 1-2
rho_T112_u,rho_T112_v = GLOQ.T1ParityForwardSolve(u0_t1_12,v0_t1_12,
				 omegas,omr_t1_12,
				 char_noise,
				 gamma1,gamma2,
				 1,
				 TC_12,t_t1_12,N_states)
population_t1_12 = GLOQ.get_population(rho_T112_u)

# T1 2-3
rho_T123_u,rho_T123_v = GLOQ.T1ParityForwardSolve(u0_t1_23,v0_t1_23,
				 omegas,omr_t1_23,
				 char_noise,
				 gamma1,gamma2,
				 2,
				 TC_23,t_t1_23,N_states)
population_t1_23 = GLOQ.get_population(rho_T123_u)

# Plotting
fnt = Plots.font("Helvetica",16)
lfnt = Plots.font("Helvetica",12)
Plots.default(titlefont=fnt,
			  guidefont=fnt,
			  tickfont=fnt,
			  legendfont=lfnt,
			  linewidth=2);

# Echo 0-1
fig_e01 = plot(t_echo_01./1000.0,population_echo_01,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Echo 0-1",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_e01)

# Echo 1-2
fig_e12 = plot(t_echo_12./1000.0,population_echo_12,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Echo 1-2",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_e12)

# Echo 2-3
fig_e23 = plot(t_echo_23./1000.0,population_echo_23,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Echo 2-3",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_e23)

# Ramsey 0-1
fig_r01 = plot(t_ramsey_01./1000.0,population_ramsey_01,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Ramsey 0-1",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_r01)

# Ramsey 1-2
fig_r12 = plot(t_ramsey_12./1000.0,population_ramsey_12,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Ramsey 1-2",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_r12)

# Ramsey 2-3
fig_r23 = plot(t_ramsey_23./1000.0,population_ramsey_23,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="Ramsey 2-3",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_r23)

# T1 0-1
fig_t01 = plot(t_t1_01./1000.0,population_t1_01,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="T1 0-1",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_t01)

# T1 1-2
fig_t12 = plot(t_t1_12./1000.0,population_t1_12,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="T1 1-2",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_t12)

# T1 1-2
fig_t23 = plot(t_t1_23./1000.0,population_t1_23,line=(:dash),
		   size=(1200,600),legend=:outerright,
		   title="T1 2-3",
		   label=[L"$\rho_{00}$, data" L"$\rho_{11}$, data" L"$\rho_{22}$, data" L"$\rho_{33}$, data"])
display(fig_t23)

# Save data
writedlm( string("synthetic_data/echo_01.txt"),  population_echo_01, ',')
writedlm( string("synthetic_data/echo_12.txt"),  population_echo_12, ',')
writedlm( string("synthetic_data/echo_23.txt"),  population_echo_23, ',')

#writedlm( string("synthetic_data/ramsey_01.txt"),  population_ramsey_01, ',')
#writedlm( string("synthetic_data/ramsey_12.txt"),  population_ramsey_12, ',')
#writedlm( string("synthetic_data/ramsey_23.txt"),  population_ramsey_23, ',')

writedlm( string("synthetic_data/t1_01.txt"),  population_t1_01, ',')
writedlm( string("synthetic_data/t1_12.txt"),  population_t1_12, ',')
writedlm( string("synthetic_data/t1_23.txt"),  population_t1_23, ',')


text:
function normalize!{T}(s::Array{T})
# Normalize s so that the squared L2 vector norm of each channel is at most 1

  (m, n, d) = size(s);
  for k=1:d;
    norm_s = vecnorm(s[:, :, k]);
    if norm_s*norm_s>1;
      s[:, :, k] /= norm_s;
    end
  end
end

max_stars_repo_name: bkong/FCSC


text:
"""
    get_neighbors(uc, N[, L = N])

This function searches for neighbors in the UnitCell uc up to order N.
It returns a nested vector of index offsets, which can be added to the indices
of a site to get one of its k-th neighbors.

The vector is indexed as such

    offset = neighbors[sublattice_index][k][running_index]

Optionally L can be set to change the size of the temporary lattice this
function creates. L has to be large enough to include all neighbors up to order
N.
"""
function get_neighbors(
        uc::UnitCell{D, T},
        N::Int64,
        L::Int64 = N
    ) where {D, T}

    # Get all positions
    # This is a bit awkward due to arbitrary dimensionality
    pos2d = Array{typeof(uc.basis[1])}(length(uc.basis), (2L+1)^D)
    for i in eachindex(uc.basis)
        pos2d[i, :] = get_array_from_UnitCell(uc, i, L)[:]
    end
    positions = reshape(pos2d, (length(uc.basis), [2L+1 for _ in 1:D]...))

    # neighbor[basis_index][neighbor_lvl] = Vector of offsets
    neighbors = [Vector{NTuple{D+1, Int64}}[] for _ in eachindex(uc.basis)]

    # Search for neighbors
    for i in eachindex(uc.basis)
        distances = [norm²(uc.basis[i] - p) for p in positions[:]]
        by_dist = sortperm(distances)

        # j = 1 is at 0 distance, skip this
        j = 2
        for lvl in 1:N
            level_distance = distances[by_dist[j]]
            nth_neighbors = NTuple{D+1, Int64}[]

            while distances[by_dist[j]] ≈ level_distance
                # linear index -> d-dimensional index
                index_pos = ind2sub(positions, by_dist[j])
                # we want relative index offsets.
                # the current pos is at index (i, L+1, ..., L+1)
                # because the lattice spans from (-L, L)
                push!(
                    nth_neighbors,
                    (
                        index_pos[1] - i,
                        map(x -> x - L - 1, index_pos[2:end])...
                    )
                )
                j += 1

                # If j > length(by_dist) we have checked all distances, but
                # haven't found enough neighbors. This (always?) means that the
                # temporary lattice is too small.
                j > length(by_dist) && throw(ErrorException(
                    "L is (perhaps) too small for the requested number of \
                    neighbor orders N."
                ))
            end
            push!(neighbors[i], nth_neighbors)
        end
    end

    neighbors
end

# less fault tolerant to use norm²
norm²(v::Point) = dot(v, v)

# shift every element of t by 'by'
shift(t::NTuple, by::Int64) = ([x + by for x in t]...)

function get_array_from_UnitCell(
        uc::UnitCell{D, T},
        alpha::Int64,
        L::Int64
    ) where {D,T}
    range = [2L+1 for _ in 1:D]
    reshape(
        [get_pos(
            uc,
            alpha,
            shift(ind2sub((range...), i), -L-1)
            # index -> D-dimensional index, shift range from (1, 2L+1) -> (-L, L)
        ) for i in 1:(2L+1)^D],
        range...
    )
end


text:
using Test, YaoBlocks, YaoArrayRegister, Flux.Optimise

function gaussian_pdf(x, μ::Real, σ::Real)
    pl = @. 1 / sqrt(2pi * σ^2) * exp(-(x - μ)^2 / (2 * σ^2))
    pl / sum(pl)
end

layer(nbit::Int, x::Symbol) = layer(nbit, Val(x))
layer(nbit::Int, ::Val{:first}) = chain(nbit, put(i => chain(Rx(0), Rz(0))) for i in 1:nbit)

layer(nbit::Int, ::Val{:last}) = chain(nbit, put(i => chain(Rz(0), Rx(0))) for i in 1:nbit)
layer(nbit::Int, ::Val{:mid}) = chain(nbit, put(i => chain(Rz(0), Rx(0), Rz(0))) for i in 1:nbit)

entangler(pairs) = chain(control(ctrl, target => X) for (ctrl, target) in pairs)

function build_circuit(n, nlayers, pairs)
    circuit = chain(n)
    push!(circuit, layer(n, :first))
    for i in 2:nlayers
        push!(circuit, cache(entangler(pairs)))
        push!(circuit, layer(n, :mid))
    end
    push!(circuit, cache(entangler(pairs)))
    push!(circuit, layer(n, :last))
    return circuit
end

struct RBFKernel
    σ::Float64
    m::Matrix{Float64}
end

function RBFKernel(σ::Float64, space)
    dx2 = (space .- space') .^ 2
    return RBFKernel(σ, exp.(-1 / 2σ * dx2))
end

kexpect(κ::RBFKernel, x, y) = x' * κ.m * y

get_prob(qcbm) = probs(zero_state(nqubits(qcbm)) |> qcbm)

function loss(κ, c, target)
    p = get_prob(c) - target
    return kexpect(κ, p, p)
end

function gradient(qcbm, κ, ptrain)
    n = nqubits(qcbm)
    prob = get_prob(qcbm)
    grad = zeros(Float64, nparameters(qcbm))

    count = 1
    for k in 1:2:length(qcbm), each_line in qcbm[k], gate in content(each_line)
        dispatch!(+, gate, π / 2)
        prob_pos = probs(zero_state(n) |> qcbm)

        dispatch!(-, gate, π)
        prob_neg = probs(zero_state(n) |> qcbm)

        dispatch!(+, gate, π / 2) # set back

        grad_pos = kexpect(κ, prob, prob_pos) - kexpect(κ, prob, prob_neg)
        grad_neg = kexpect(κ, ptrain, prob_pos) - kexpect(κ, ptrain, prob_neg)
        grad[count] = grad_pos - grad_neg
        count += 1
    end
    return grad
end

qcbm = build_circuit(6, 10, [1 => 2, 3 => 4, 5 => 6, 2 => 3, 4 => 5, 6 => 1])
dispatch!(qcbm, :random) # initialize the parameters

κ = RBFKernel(0.25, 0:2^6-1)
pg = gaussian_pdf(1:1<<6, 1 << 5 - 0.5, 1 << 4);
opt = ADAM()

function train(qcbm, κ, opt, target)
    history = Float64[]
    for _ in 1:100
        push!(history, loss(κ, qcbm, target))
        ps = parameters(qcbm)
        Optimise.update!(opt, ps, gradient(qcbm, κ, target))
        popdispatch!(qcbm, ps)
    end
    return history
end

history = train(qcbm, κ, opt, pg)
trained_pg = probs(zero_state(nqubits(qcbm)) |> qcbm)

using LinearAlgebra
@test norm(trained_pg - pg) < 0.1


text:
using NNlib, Test, LinearAlgebra, Logging
using NNlib: storage_type, storage_typejoin, is_strided,
    batched_mul_generic!, _unbatch, _copy_if_faster,
    BatchedAdjoint, BatchedTranspose

function bmm_test(a,b; transA = false, transB = false)
    bs = size(a,3)
    transA && (a = permutedims(a, [2,1,3]))
    transB && (b = permutedims(b, [2,1,3]))
    c = []
    for i = 1:bs
        push!(c, a[:,:,i]*b[:,:,i])
    end

    cat(c...; dims = 3)
end

function bmm_adjtest(a,b; adjA = false, adjB = false)
    bs = size(a,3)
    c = []
    for i = 1:bs
        ai = adjA ? adjoint(a[:,:,i]) : a[:,:,i]
        bi = adjB ? adjoint(b[:,:,i]) : b[:,:,i]
        push!(c, ai*bi)
    end

    cat(c...; dims = 3)
end

function half_batched_mul(x,y)
    @assert size(y,3) == 1
    d = size(x,2)
    x_mat = reshape(permutedims(x, (1,3,2)),:,d)
    y_mat = reshape(y,d,:)
    z_mat = x_mat * y_mat
    permutedims(reshape(z_mat, size(x,1), size(x,3), :), (1,3,2))
end

@testset "batched_mul: Float64 * $TB" for TB in [Float64, Float32]

    # Real
    A = randn(7,5,3)
    B = randn(TB, 5,7,3)
    C = randn(7,6,3)

    @test batched_mul(A, B) ≈ bmm_test(A, B)
    @test batched_mul(batched_transpose(A), batched_transpose(B)) ≈ bmm_test(A, B; transA = true, transB = true)
    @test batched_mul(batched_transpose(A), C) ≈ bmm_test(A, C; transA = true)
    @test batched_mul(A, batched_transpose(A)) ≈ bmm_test(A, A; transB = true)

    # Complex
    cA = randn(Complex{Float64}, 7,5,3)
    cB = randn(Complex{TB}, 5,7,3)
    cC = randn(Complex{Float64}, 7,6,3)

    @test batched_mul(cA, cB) ≈ bmm_adjtest(cA, cB)
    @test batched_mul(batched_adjoint(cA), batched_adjoint(cB)) ≈ bmm_adjtest(cA, cB; adjA = true, adjB = true)
    @test batched_mul(batched_adjoint(cA), cC) ≈ bmm_adjtest(cA, cC; adjA = true)
    @test batched_mul(cA, batched_adjoint(cA)) ≈ bmm_adjtest(cA, cA; adjB = true)

    # Wrappers which cancel
    @test batched_transpose(batched_transpose(A)) === A
    @test batched_transpose(PermutedDimsArray(A, (2,1,3))) === A
    @test batched_adjoint(batched_adjoint(cA)) === cA
    @test batched_transpose(batched_adjoint(cA)) isa NNlib.BatchedAdjoint

    # Integers
    TBi = TB==Float64 ? Int64 : Int32
    iA = rand(1:99, 7,5,3)
    iB = TB.(rand(1:99, 5,7,3))
    iC = zeros(Int, 7,6,3)
    @test batched_mul(iA, iB) == bmm_adjtest(iA, iB)
    @test batched_mul(cA, iB) ≈ bmm_adjtest(cA, iB)

    # Errors
    @test_throws DimensionMismatch batched_mul(rand(2,2,2), rand(TB, 2,2,10))
    @test_throws DimensionMismatch batched_mul(rand(2,2,2), rand(TB, 10,2,2))
    @test_throws Exception batched_mul!(zeros(2,2,10), rand(2,2,2), rand(TB, 2,2,2))

    # PermutedDimsArrays
    for perm in [(1,3,2), (2,1,3), (3,2,1)], fun in [identity, batched_adjoint], ty in [identity, complex]
        A = randn(ty(Float64), 4,4,4)
        B = randn(ty(TB), 4,4,4)
        @test batched_mul(fun(A), PermutedDimsArray(B, perm)) ≈ batched_mul(fun(A), permutedims(B, perm))
        @test batched_mul(fun(PermutedDimsArray(A, perm)), B) ≈ batched_mul(fun(permutedims(A, perm)), B)
        # when TB=Float64, only the case  perm=(2,1,3); fun=batched_adjoint; ty=complex;  goes to fallback
        # but all the perm=(3,2,1); cases copy their inputs.
    end

    # PermutedDimsArray output
    A′ = randn(4,3,2)
    B′ = batched_adjoint(randn(TB, 5,3,2))
    C1 = batched_mul(A′, B′) # size 4,5,2
    C2 = PermutedDimsArray(zeros(5,2,4), (3,1,2)) # size 4,5,2
    @test C1 ≈ batched_mul!(C2, A′, B′) # Float64: "Debug: transposing C = A * B into Cᵀ = Bᵀ * Aᵀ"
    @test C1 ≈ C2

    # 5-arg mul!
    @test 10 .* C1 ≈ batched_mul!(C2, A′, B′, 10)
    C2 .= 10
    @test C1 .+ 100 ≈ batched_mul!(C2, A′, B′, 1, 10)

    # Trivial batches for B
    D′ = randn(TB, 3,5,1)
    @test size(batched_mul(A′,D′)) == (4,5,2)
    @test batched_mul(A′,D′) ≈ half_batched_mul(A′, D′)

    # Large output, multi-threaded path
    if TB == Float64
        N = 50
        A = rand(N,N,N)
        B = rand(N,N,N)
        C = reshape(reduce(hcat, [vec(A[:,:,k] * B[:,:,k]) for k in 1:N]), N,N,N)
        @test C ≈ A ⊠ B

        D = rand(N,N,1)
        E = reshape(reduce(hcat, [vec(A[:,:,k] * D[:,:,1]) for k in 1:N]), N,N,N)
        @test E ≈ A ⊠ D
    end
end

perm_12(A) = PermutedDimsArray(A, (2,1,3))
perm_23(A) = PermutedDimsArray(A, (1,3,2))

@testset "batched_mul: trivial dimensions & unit strides, $T" for T in [Float64, ComplexF64]
    @testset "$tA(rand$((sA...,3))) ⊠ $tB(rand$((sB...,3)))" for
    tA in [identity, batched_adjoint, batched_transpose, perm_12, perm_23], sA in [(1,1), (1,3), (3,1), (3,3)],
    tB in [identity, batched_adjoint, batched_transpose, perm_12, perm_23], sB in [(1,1), (1,3), (3,1), (3,3)]

        A = tA(rand(T, sA..., 3))
        B = tB(rand(T, sB..., 3))
        size(A,2) == size(B,1) && size(A,3) == size(B,3) == 3 || continue

        C = cat(A[:,:,1] * B[:,:,1], A[:,:,2] * B[:,:,2], A[:,:,3] * B[:,:,3]; dims=3)
        @test A ⊠ B ≈ C
        @test_logs min_level=Logging.Debug A ⊠ B

        # In-place batched_mul!
        α, β = rand(T), rand(T)
        D = rand(T, size(C))
        @test batched_mul!(copy(D), A, B, α, β) ≈ α .* C .+ β .* D
        @test batched_mul_generic!(copy(D), A, B, α, β) ≈ α .* C .+ β .* D

        # ... and with weird LHS -- all to batched_mul_generic! right now
        C2 = batched_transpose(permutedims(C, (2,1,3)))
        C3 = batched_adjoint(permutedims(conj(C), (2,1,3)))
        @test C2 == C3 == C
        C2 .= D
        C3 .= D
        @test batched_mul!(C2, A, B, α, β) ≈ α .* C .+ β .* D
        @test C2 ≈ α .* C .+ β .* D
        @test batched_mul!(C3, A, B, α, β) ≈ α .* C .+ β .* D
        @test C3 ≈ α .* C .+ β .* D
    end
end

@testset "BatchedAdjOrTrans interface * $TB" for TB in [Float64, Float32]
    A = randn(7,5,3)
    B = randn(TB, 5,7,3)
    C = randn(7,6,3)

    function interface_tests(X, _X)
        @test length(_X) == length(X)
        @test size(_X) == (size(X, 2), size(X, 1), size(X, 3))
        @test axes(_X) == (axes(X, 2), axes(X, 1), axes(X, 3))
        #
        @test getindex(_X, 2, 3, 3) == getindex(X, 3, 2, 3)
        @test getindex(_X, 5, 4, 1) == getindex(X, 4, 5, 1)
        #
        setindex!(_X, 2.0, 2, 4, 1)
        @test getindex(_X, 2, 4, 1) == 2.0
        setindex!(_X, 3.0, 1, 2, 2)
        @test getindex(_X, 1, 2, 2) == 3.0

        _sim = similar(_X, TB, (2, 3))
        @test size(_sim) == (2, 3)
        @test typeof(_sim) == Array{TB, 2}

        _sim = similar(_X, TB)
        @test length(_sim) == length(_X)
        @test typeof(_sim) == Array{TB, 3}

        _sim = similar(_X, (2, 3))
        @test size(_sim) == (2, 3)
        @test typeof(_sim) == Array{Float64, 2}

        _sim = similar(_X)
        @test length(_sim) == length(_X)
        @test typeof(_sim) == Array{Float64, 3}

        @test parent(_X) == _X.parent
    end

    for (X, _X) in zip([A, B, C], map(batched_adjoint, [A, B, C]))
        interface_tests(X, _X)

        @test -_X == NNlib.BatchedAdjoint(-_X.parent)

        _copyX = copy(_X)
        @test _X == _copyX

        setindex!(_copyX, 2.0, 1, 2, 1)
        @test _X != _copyX
    end

    for (X, _X) in zip([A, B, C], map(batched_transpose, [A, B, C]))
        interface_tests(X, _X)

        @test -_X == NNlib.BatchedTranspose(-_X.parent)

        _copyX = copy(_X)
        @test _X == _copyX

        setindex!(_copyX, 2.0, 1, 2, 1)
        @test _X != _copyX
    end
end

@testset "batched_mul(ndims < 3), $TM" for TM in [ComplexF64, Int8]
    A = randn(ComplexF64, 3,3,3)
    M = rand(TM, 3,3) .+ im
    V = rand(TM, 3)

    # These are all reshaped and sent to batched_mul(3-array, 3-array)
    @test batched_mul(A, M) ≈ cat([A[:,:,k] * M for k in 1:3]...; dims=3)
    @test batched_mul(A, M') ≈ cat([A[:,:,k] * M' for k in 1:3]...; dims=3)
    @test A ⊠ transpose(M) ≈ cat([A[:,:,k] * transpose(M) for k in 1:3]...; dims=3)

    @test batched_mul(M, A) ≈ cat([M * A[:,:,k] for k in 1:3]...; dims=3)
    @test batched_mul(M', A) ≈ cat([M' * A[:,:,k] for k in 1:3]...; dims=3)
    @test transpose(M) ⊠ A ≈ cat([transpose(M) * A[:,:,k] for k in 1:3]...; dims=3)

    # batched_vec
    @test batched_vec(A, M) ≈ hcat([A[:,:,k] * M[:,k] for k in 1:3]...)
    @test batched_vec(A, M') ≈ hcat([A[:,:,k] * (M')[:,k] for k in 1:3]...)
    @test batched_vec(A, V) ≈ hcat([A[:,:,k] * V for k in 1:3]...)
end

@testset "storage_type" begin

    @test storage_type(transpose(reshape(view(rand(10), 2:9),4,:))) == Vector{Float64}
    @test storage_type(transpose(reshape(view(1:10,     2:9),4,:))) == UnitRange{Int}

    @test storage_typejoin(rand(2), rand(Float32, 2)) == Vector{<:Any}
    @test storage_typejoin(rand(2), rand(2,3)', rand(2,3,4)) == Array{Float64}
    @test storage_typejoin([1,2,3], 4:5) == AbstractVector{Int}

end

@testset "is_strided" begin

    M = ones(10,10)

    @test is_strided(M)
    @test is_strided(view(M, 1:2:5,:))
    @test is_strided(PermutedDimsArray(M, (2,1)))

    @test !is_strided(reshape(view(M, 1:2:10,:), 10,:))
    @test !is_strided((M.+im)')
    @test !is_strided(Diagonal(ones(3)))

    A = ones(2,2,2)

    @test is_strided(batched_adjoint(A))
    @test is_strided(batched_transpose(A))
    @test !is_strided(batched_adjoint(A .+ im))
    @test is_strided(batched_transpose(A .+ im))

end

FiniteDifferences.to_vec(x::BatchedAdjoint) = FiniteDifferences.to_vec(collect(x))
FiniteDifferences.to_vec(x::BatchedTranspose) = FiniteDifferences.to_vec(collect(x))

@testset "AutoDiff" begin
    M, P, Q = 13, 7, 11
    B = 3
    # Two 3-arrays
    gradtest(batched_mul, randn(rng, M, P, B), randn(rng, P, Q, B))
    gradtest(batched_mul, batched_adjoint(randn(rng, P, M, B)), randn(rng, P, Q, B))
    gradtest(batched_mul, randn(rng, M, P, B), batched_transpose(randn(rng, Q, P, B)))

    # One a matrix...
    gradtest(batched_mul, randn(rng, M, P), randn(rng, P, Q, B))
    gradtest(batched_mul, adjoint(randn(rng, P, M)), randn(rng, P, Q, B))
    gradtest(batched_mul, randn(rng, M, P), batched_adjoint(randn(rng, Q, P, B)))

    gradtest(batched_mul, randn(rng, M, P, B), randn(rng, P, Q))
    gradtest(batched_mul, batched_transpose(randn(rng, P, M, B)), randn(rng, P, Q))
    gradtest(batched_mul, randn(rng, M, P, B), transpose(randn(rng, Q, P)))

    # ... or equivalent to a matrix
    gradtest(batched_mul, randn(rng, M, P, 1), randn(rng, P, Q, B))
    gradtest(batched_mul, batched_transpose(randn(rng, P, M, 1)), randn(rng, P, Q, B))
    gradtest(batched_mul, randn(rng, M, P, 1), batched_transpose(randn(rng, Q, P, B)))

    gradtest(batched_mul, randn(rng, M, P, B), randn(rng, P, Q, 1))
    gradtest(batched_mul, batched_adjoint(randn(rng, P, M, B)), randn(rng, P, Q, 1))
    gradtest(batched_mul, randn(rng, M, P, B), batched_adjoint(randn(rng, Q, P, 1)))

    # batched_vec
    gradtest(batched_vec, randn(rng, M, P, B), randn(rng, P, B))
    gradtest(batched_vec, randn(rng, M, P, B), transpose(randn(rng, B, P)))

    gradtest(batched_vec, randn(rng, M, P, B), randn(rng, P))
end


text:
export Atan

"""
`Atan([domainType=Float64::Type,] dim_in::Tuple)`

Creates an inverse tangent non-linear operator with input dimensions `dim_in`:
```math
\\text{atan} ( \\mathbf{x} ).
```

"""
struct Atan{T,N} <: NonLinearOperator
	dim::NTuple{N,Int}
end

function Atan(DomainType::Type, DomainDim::NTuple{N,Int}) where {N}
	Atan{DomainType,N}(DomainDim)
end

Atan(DomainDim::NTuple{N,Int}) where {N} = Atan{Float64,N}(DomainDim)
Atan(DomainDim::Vararg{Int}) = Atan{Float64,length(DomainDim)}(DomainDim)

function mul!(y::AbstractArray{T,N}, L::Atan{T,N}, x::AbstractArray{T,N}) where {T,N}
	y .= atan.(x)
end

function mul!(y::AbstractArray,
              J::AdjointOperator{Jacobian{A,TT}},
              b::AbstractArray) where {T, N, A <: Atan{T,N}, TT <: AbstractArray{T,N}}
    L = J.A
    y .= conj.(1.0./(1.0.+ L.x.^2)).*b
end

fun_name(L::Atan) = "atan"

size(L::Atan) = (L.dim, L.dim)

domainType(L::Atan{T,N}) where {T,N} = T
codomainType(L::Atan{T,N}) where {T,N} = T


text:
# MTs with eltypes which have units from the Unitful package

using Test
using MomentTensors
using Unitful

@testset "Units" begin
    m = Unitful.m
    N = Unitful.N
    Nm = N*m

    @testset "Construction" begin
        @test MT((1, 2, 3, 4, 5, 6).*u"N*m") == MT(1Nm, 2Nm, 3Nm, 4Nm, 5Nm, 6Nm) == MT((1:6)*Nm)
        @test MT(1:6) != MT((1:6)*Nm)
        let mt = MT((1:6)*Nm)
            @test eltype(mt) == typeof(float(1Nm))
        end
    end

    @testset "getindex/setindex!" begin
        let mt = MT((1:6)*Nm)
            @test mt[1,1] == 1Nm
        end
    end

    @testset "Arithmetic" begin
        @test MT(ones(6).*u"N*m/s") + MT(2*ones(6).*u"N*m/s") == MT(3*ones(6).*u"N*m/s")
        let m1 = MT((11, 22, 33, 12, 13, 23).*Nm),
            m2 = MT((-11, -22, -33, -12, -13, -23).*Nm)
            @test m1 + m2 == MT((0, 0, 0, 0, 0, 0).*Nm)
            @test m1 == -m2
            @test -m1 == m2
            @test m1 + 1Nm == MT((12, 23, 34, 13, 14, 24).*Nm)
            @test m1 - 1Nm == MT((10, 21, 32, 11, 12, 22).*Nm)
            @test m1 + 3Nm == 3Nm + m1
            @test m1 - 4Nm == -(4Nm - m1)
            @test 2m1 == MT((22, 44, 66, 24, 26, 46).*Nm)
            @test 2m1 == m1*2
            @test 1/m1 == MT(1 ./ ((11, 22, 33, 12, 13, 23).*Nm)...)
            @test m1/2 == MT((5.5, 11, 16.5, 6, 6.5, 11.5).*Nm)
        end
    end

    @testset "Conversion" begin
        @test Array(MT((1:6)*Nm)) == Float64[1 4 5; 4 2 6; 5 6 3].*Nm
    end

    @testset "Decomposition" begin
        let m = MT([ 3 -5 10
                    -5  1  4
                    10  4  1].*Nm)
            eltp = typeof(float(1Nm))
            d = decompose(m)

            @testset "Unitful MT $f" for f in (:iso, :dev, :dc, :clvd)
                @test eltype(getfield(d, f)) == eltp
            end
            @testset "Unitful scalar $f" for f in (:iso_m0, :dev_m0, :m0)
                @test typeof(getfield(d, f)) == eltp
            end
            @testset "Dimensionless scalar $f" for f in (:prop_iso, :prop_dev, :prop_dc, :prop_clvd)
                @test typeof(getfield(d, f)) == typeof(ustrip(m[1,1]))
            end
        end
    end

    @testset "amplitude_v_azimuth" begin
        eltp = typeof(Float32(1N))
        p, sv, sh, j = amplitude_v_azimuth(MT(rand(Float32, 6).*N), 0, 0)
        for v in (p, sv, sh)
            @test eltype(v) == eltp
        end
        @test eltype(j) == Float32
    end

    @testset "radiation_pattern" begin
        eltp = typeof(Float32(1Nm))
        p, sv, sh, j = radiation_pattern(MT(rand(Float32, 6).*Nm), 0, 0)
        for v in (p, sv, sh)
            @test v isa eltp
        end
        @test j isa Float32
    end

    @testset "rotate" begin
        @test rotate(MT(rand(Float32, 6).*Nm), 1, 2, 3) isa MT{typeof(1f0*Nm)}
    end

    @testset "eps_non_dc" begin
        @test eps_non_dc(MT(rand(Float32, 6).*Nm)) isa typeof(ustrip(1f0*Nm))
    end
end


text:
using Luxor
using JSON
using LatticeTools

include("Kagome.jl")

lattice_data = JSON.parsefile("lattice-3.json")
sites = [(sitetype, [R...]) for (sitetype, R) in lattice_data["sites"]]
kagome = Kagome.make_kagome_obc(sites)
site_coordinates = []
site_types = String[]
for ((sitetype, sitesubscript), sitefc) in kagome.unitcell.sites
    sitecc = fract2carte(kagome.unitcell, sitefc)
    push!(site_coordinates, sitecc)
    push!(site_types, sitetype)
end
site_coordinates = hcat(site_coordinates...)
site_coordinates[2,:] = -site_coordinates[2,:]

Drawing(800, 300, "logo.svg")

origin()
translate(-450, 130)

let ratio = 100, radius = 13, radius_large = 18, offset = (0,0),
    sitecolor = Dict("A" => "#389826", "B" => "#9558B2", "C" => "#CB3C33")

    sethue("gray")
    setline(2)
    for (i, j) in kagome.nearest_neighbor_bonds
        ri = site_coordinates[:,i]
        rj = site_coordinates[:,j]
        line(
            Point(ri[1]*ratio + offset[1], ri[2]*ratio + offset[2]),
            Point(rj[1]*ratio + offset[1], rj[2]*ratio + offset[2]),
            :stroke
        )
    end
    for (sitetype, r) in zip(site_types, eachcol(site_coordinates))
        sethue("white")
        circle(Point(r[1]*ratio + offset[1], r[2]*ratio + offset[2]), radius_large, :fill)
        sethue(sitecolor[sitetype])
        circle(Point(r[1]*ratio + offset[1], r[2]*ratio + offset[2]), radius, :fill)
    end
end

translate(360, -160)
sethue("black")
fontsize(144)
fontface("Tamil MN")
texttrack("Lattice", Point(0, 0), -50, 144)
translate(0, 140)
texttrack("Tools.jl", Point(0, 0), -50, 144)

setcolor(sethue("#4063D8")..., 1.0)
#ellipse(286.3, -239.6, 25, 22.6, :fill)
ellipse(286, -238, 32, 32, :fill)

finish()
#preview()


text:
using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using SparseArrays
Random.seed!(233)

function sparse_solver(ii,jj,vv,f,method="SparseLU")
    sparse_solver_ = load_op_and_grad("./build/libSparseSolver","sparse_solver")
    ii,jj,vv,f = convert_to_tensor([ii,jj,vv,f], [Int64,Int64,Float64,Float64])
    sparse_solver_(ii,jj,vv,f,method)
end

# TODO: specify your input parameters
sess = Session(); init(sess)
for method in ["SparseLU", "SparseQR", "SimplicialLDLT", "SimplicialLLT"]
A = sprand(10,10,0.6)
global A = A'*A
global f = rand(10)
A_ = SparseTensor(A)
global ii, jj, vv = find(A_)
u = sparse_solver(ii,jj,vv,f,method)

@show norm(run(sess, u)-A\f)
end

# uncomment it for testing gradients
# error()
# py"""
# import traceback
# try:
#     $run($sess, $dy_)
# except Exception:
#     print(traceback.format_exc())
# """
# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m)
    return sum(sparse_solver(ii,jj,vv,m)^2)
end

# TODO: change `m_` and `v_` to appropriate values
# m_ = constant(rand(length(vv)))
# v_ = rand(length(vv))

m_ = constant(f)
v_ = rand(length(f))
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


text:
#####
##### Tester bots
#####

using TexasHoldem
import TexasHoldem
TH = TexasHoldem
const AGS = AbstractGameState

##### BotCheckFold
struct BotCheckFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCheckFold}, ::AGS, ::CheckRaiseFold) = check!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckFold}, ::AGS, ::CallRaiseFold) = fold!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckFold}, ::AGS, ::CallAllInFold) = fold!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckFold}, ::AGS, ::CallFold) = fold!(game, player)

##### BotCheckCall
struct BotCheckCall <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCheckCall}, ::AGS, ::CheckRaiseFold) = check!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckCall}, ::AGS, ::CallRaiseFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckCall}, ::AGS, ::CallAllInFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotCheckCall}, ::AGS, ::CallFold) = call!(game, player)

##### BotFlopRaise
struct BotFlopRaise <: AbstractAI end
# Call blind / check pre-flop, raise on flop

TH.player_option!(game::Game, player::Player{BotFlopRaise}, ::AGS, ::CheckRaiseFold) = check!(game, player)
TH.player_option!(game::Game, player::Player{BotFlopRaise}, ::AGS, ::CallRaiseFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotFlopRaise}, ::AGS, ::CallAllInFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotFlopRaise}, ::AGS, ::CallFold) = call!(game, player)

TH.player_option!(game::Game, player::Player{BotFlopRaise}, ::Flop, ::CheckRaiseFold) = raise_to!(game, player, bank_roll(player)/2)

##### BotRaiseAllIn
struct BotRaiseAllIn <: AbstractAI end
# Call blind / check pre-flop, raise on flop

TH.player_option!(game::Game, player::Player{BotRaiseAllIn}, ::AGS, ::CheckRaiseFold) = raise_all_in!(game, player)
TH.player_option!(game::Game, player::Player{BotRaiseAllIn}, ::AGS, ::CallRaiseFold) = raise_all_in!(game, player)
TH.player_option!(game::Game, player::Player{BotRaiseAllIn}, ::AGS, ::CallAllInFold) = raise_all_in!(game, player)

##### BotRaiseAlmostAllIn
struct BotRaiseAlmostAllIn <: AbstractAI end
# Call blind / check pre-flop, raise on flop

TH.player_option!(game::Game, player::Player{BotRaiseAlmostAllIn}, ::AGS, ::CheckRaiseFold) = raise_to!(game, player, 0.9*round_bank_roll(player))
TH.player_option!(game::Game, player::Player{BotRaiseAlmostAllIn}, ::AGS, ::CallRaiseFold) = raise_to!(game, player, 0.9*round_bank_roll(player))
TH.player_option!(game::Game, player::Player{BotRaiseAlmostAllIn}, ::AGS, ::CallAllInFold) = raise_to!(game, player, 0.9*round_bank_roll(player))

##### BotBetSB
struct BotBetSB <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotBetSB}, ::AGS, ::CheckRaiseFold) = raise_to!(game, player, TH.blinds(game).small)
TH.player_option!(game::Game, player::Player{BotBetSB}, ::PreFlop, ::CheckRaiseFold) = raise_to!(game, player, TH.blinds(game).small)
TH.player_option!(game::Game, player::Player{BotBetSB}, ::AGS, ::CallRaiseFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotBetSB}, ::AGS, ::CallFold) = call!(game, player)

##### BotBetBB
struct BotBetBB <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotBetBB}, ::PreFlop, ::CheckRaiseFold) = raise_to!(game, player, 2*TH.blinds(game).big)
TH.player_option!(game::Game, player::Player{BotBetBB}, ::AGS, ::CheckRaiseFold) = raise_to!(game, player, TH.blinds(game).big)
TH.player_option!(game::Game, player::Player{BotBetBB}, ::AGS, ::CallRaiseFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotBetBB}, ::AGS, ::CallFold) = call!(game, player)

##### BotCheckOnCallRaiseFold
struct BotCheckOnCallRaiseFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCheckOnCallRaiseFold}, ::AGS, ::CallRaiseFold) = check!(game, player)

##### BotCheckOnCallAllInFold
struct BotCheckOnCallAllInFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCheckOnCallAllInFold}, ::AGS, ::CallAllInFold) = check!(game, player)

##### BotCheckOnCallFold
struct BotCheckOnCallFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCheckOnCallFold}, ::AGS, ::CallFold) = check!(game, player)

##### BotCallOnCheckRaiseFold
struct BotCallOnCheckRaiseFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotCallOnCheckRaiseFold}, ::PreFlop, ::CallRaiseFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{BotCallOnCheckRaiseFold}, ::AGS, ::CheckRaiseFold) = call!(game, player)

##### BotRaiseOnCallFold
struct BotRaiseOnCallFold <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotRaiseOnCallFold}, ::AGS, ::CallFold) = raise_to!(game, player, bank_roll(player))

##### BotLimpAllIn
struct BotLimpAllIn <: AbstractAI end

TH.player_option!(game::Game, player::Player{BotLimpAllIn}, ::AGS, ::PlayerOptions) = call!(game, player)

##### BotNActions
struct BotNActions <: AbstractAI end

function TH.player_option!(game::Game, player::Player{BotNActions}, ::AGS, ::CheckRaiseFold)
    check!(game, player)
    n_check_actions[1]+=1
end

function TH.player_option!(game::Game, player::Player{BotNActions}, ::AGS, ::CallRaiseFold)
    call!(game, player)
    n_call_actions[1]+=1
end

##### BotPreFlopRaise
struct BotPreFlopRaise{FT} <: AbstractAI
    amt::FT
end
# Call blind / check pre-flop, raise on flop

TH.player_option!(game::Game, player::Player{<:BotPreFlopRaise}, ::AGS, ::CheckRaiseFold) = check!(game, player)
TH.player_option!(game::Game, player::Player{<:BotPreFlopRaise}, ::AGS, ::CallRaiseFold) = raise_to!(game, player, TH.life_form(player).amt)
TH.player_option!(game::Game, player::Player{<:BotPreFlopRaise}, ::AGS, ::CallAllInFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{<:BotPreFlopRaise}, ::AGS, ::CallFold) = call!(game, player)
TH.player_option!(game::Game, player::Player{<:BotPreFlopRaise}, ::PreFlop, ::CheckRaiseFold) = raise_to!(game, player, TH.life_form(player).amt)


text:
# This file is a part of Julia. License is MIT: http://julialang.org/license

include("strings/errors.jl")
include("strings/types.jl")
include("strings/basic.jl")
include("strings/search.jl")
include("strings/util.jl")
include("strings/io.jl")
include("strings/utf8proc.jl")
importall .UTF8proc


text:
# ------------------------------------------------------------------
# Copyright (c) 2017, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    CubeNeighborhood(domain, radius)

A cube (∞-norm ball) neighborhood with a given `radius` on a spatial `domain`.
"""
struct CubeNeighborhood{D<:AbstractDomain} <: AbstractNeighborhood{D}
  domain::D
  radius # we cannot use coordtype(D) here yet in Julia v0.6

  function CubeNeighborhood{D}(domain, radius) where {D<:AbstractDomain}
    @assert radius > 0 "neighborhood radius must be positive"
    @assert typeof(radius) == coordtype(domain) "radius and domain coordinate type must match"

    new(domain, radius)
  end
end

CubeNeighborhood(domain::D, radius) where {D<:AbstractDomain} = CubeNeighborhood{D}(domain, radius)

function (neigh::CubeNeighborhood{<:RegularGrid})(location::Int)
  # grid specs
  sz = size(neigh.domain)
  sp = spacing(neigh.domain)
  nd = ndims(neigh.domain)

  # cube center in Cartesian index format
  center = CartesianIndices(sz)[location]

  # number of units to reach the sides of the cube
  units  = @. floor(Int, neigh.radius / sp)

  # cube spans from top left to bottom right
  start  = @. max(center.I - units, 1)
  finish = @. min(center.I + units, sz)

  # number of points in the cube
  npoints = prod(@. finish - start + 1)

  # pre-allocate memory
  neighbors = Vector{Int}(undef, npoints)

  if nd == 1
    n = 1
    for i=start[1]:finish[1]
      @inbounds neighbors[n] = i
      n += 1
    end
  elseif nd == 2
    n = 1
    nx = sz[1]
    for j=start[2]:finish[2], i=start[1]:finish[1]
      @inbounds neighbors[n] = i + nx*(j-1)
      n += 1
    end
  elseif nd == 3
    n = 1
    nx = sz[1]
    nxny = sz[1]*sz[2]
    for k=start[3]:finish[3], j=start[2]:finish[2], i=start[1]:finish[1]
      @inbounds neighbors[n] = i + nx*(j-1) + nxny*(k-1)
      n += 1
    end
  else # higher dimensions
    irange = CartesianIndices(ntuple(i -> start[i]:finish[i], nd))

    for (n, ind) in enumerate(irange)
      @inbounds neighbors[n] = LinearIndices(sz)[ind]
    end
  end

  neighbors
end

function (neigh::CubeNeighborhood{<:PointSet})(location::Int)
  # retrieve domain
  ndomain = neigh.domain

  # center in real coordinates
  xₒ = coordinates(ndomain, location)

  # pre-allocate memory for neighbors coordinates
  x = MVector{ndims(ndomain),coordtype(ndomain)}(undef)

  neighbors = Vector{Int}()
  for loc in 1:npoints(ndomain)
    coordinates!(x, ndomain, loc)
    norm(x .- xₒ, Inf) ≤ neigh.radius && push!(neighbors, loc)
  end

  neighbors
end


text:
include("functions/GenerateCareerPaths.jl")
using Compose, GraphPlot

FirstLevel = Vector{AbstractLevel}()
g = DiGraph(4)

# 1 => world
add_edge!(g, 1, 2)
#add_edge!(g, 1, 4)

# 2 => 1B-BDL
add_edge!(g, 2, 3)
add_edge!(g, 2, 4)
#add_edge!(g, 2, 5)
#add_edge!(g, 2, 6)
push!(FirstLevel, AcademicLevel("1B-BDL", next = ["SLT", "1LT", "CPN"], min = 6, max = 6))

# 3 => 1B-B
add_edge!(g, 3, 4)
#add_edge!(g, 3, 5)
#add_edge!(g, 3, 6)
push!(FirstLevel, AcademicLevel("1B-B", next = ["1LT", "CPN", "CDT"], max = 14))

# 4 => 1A-BDL
#add_edge!(g, 4, 5)
#add_edge!(g, 4, 6)
#push!(FirstLevel, AcademicLevel("1A-BDL", next = ["SLT", "1LT"]))

# 5 => 1A-B
#add_edge!(g, 5, 6)
#push!(FirstLevel, AcademicLevel("1A-B", next = ["CPN", "CDT"]))

# 6 => Leaving org

levelnames = [FirstLevel[i].Name for i in 1:length(FirstLevel)]
draw(PNG("InitGraph.png", 25cm, 25cm), gplot(g, nodelabel=["In_world"; levelnames ;"Out_world"]))
draw(SVG("InitGraph.svg", 25cm, 25cm), gplot(g, nodelabel=["In_world"; levelnames ;"Out_world"]))

SecondLevel = Vector{AbstractLevel}()
push!(SecondLevel, Rank("SLT", 13, min = 2, max = 2))
push!(SecondLevel, Rank("1LT", 14, min = 3, max = 3))
push!(SecondLevel, Rank("CPN", 15, min = 4, max = 6))
push!(SecondLevel, Rank("CDT", 16, min = 9))

Levels = Array{Vector{AbstractLevel}}(2)
Levels[1] = FirstLevel
Levels[2] = SecondLevel

Paths = ConstructCareerPaths(g, Levels)

nbVertices = 0
for i in 1:length(Paths)
    nbVertices += length(Paths[i].Path)
end
PlotG = DiGraph(length(Paths) + nbVertices + length(Paths))

NodeLabels = Array{String}(length(Paths) + nbVertices + length(Paths))
i = 1
for path in Paths
    NodeLabels[i] = "World"
    i += 1
    for status in path.Path
        NodeLabels[i] = status.AcademicLevel.Name * " " * status.Rank.Name
        add_edge!(PlotG, i - 1, i)
        i += 1
    end
    add_edge!(PlotG, i - 1, i)
    NodeLabels[i] = "World"
    i += 1
end
draw(SVG("CareerPath.svg", 150cm, 150cm), gplot(PlotG, nodelabel=NodeLabels))
draw(PNG("CareerPath.png", 150cm, 150cm), gplot(PlotG, nodelabel=NodeLabels))

max_stars_repo_name: Omazaria/CareerPathBeDef

max_issues_repo_name: Omazaria/CareerPathBeDef

max_forks_repo_name: Omazaria/CareerPathBeDef


text:
using JuliaInterpreter, Test, Random

module JIVisible
module JIInvisible
end
end

@testset "Basics" begin
    @test JuliaInterpreter.is_doc_expr(:(@doc "string" sum))
    @test JuliaInterpreter.is_doc_expr(:(Core.@doc "string" sum))
    ex = quote
        """
        a docstring
        """
        sum
    end
    @test JuliaInterpreter.is_doc_expr(ex.args[2])
    @test !JuliaInterpreter.is_doc_expr(:(1+1))
    ex = :(@doc doc"""
       bla
       """)
    modexs, docexs = JuliaInterpreter.split_expressions(Main, ex; extract_docexprs=true)
    @test isempty(docexs)
    # https://github.com/JunoLab/Juno.jl/issues/271
    ex = quote
        """
        Special Docstring
        """
        module DocStringTest
        function foo()
            x = 4 + 5
        end
        end
    end
    modexs, docexs = JuliaInterpreter.split_expressions(Main, ex; extract_docexprs=true)
    m, ex = first(modexs)
    @test m == Main.DocStringTest
    dex = first(docexs[Main])
    @test dex.args[4] == :DocStringTest
    eval(dex)
    io = IOBuffer()
    show(io, @doc(Main.DocStringTest))
    @test occursin("Special", String(take!(io)))

    @test JuliaInterpreter.prepare_thunk(Main, :(export foo)) === nothing
    @test JuliaInterpreter.prepare_thunk(Base.Threads, :(global Condition)) === nothing

    @test !isdefined(Main, :JIInvisible)
    JuliaInterpreter.split_expressions(JIVisible, :(module JIInvisible f() = 1 end))
    @test !isdefined(Main, :JIInvisible)
    mktempdir() do path
        push!(LOAD_PATH, path)
        open(joinpath(path, "TmpPkg1.jl"), "w") do io
            println(io, """
                    module TmpPkg1
                    using TmpPkg2
                    end
                    """)
        end
        open(joinpath(path, "TmpPkg2.jl"), "w") do io
            println(io, """
                    module TmpPkg2
                    f() = 1
                    end
                    """)
        end
        @eval using TmpPkg1
        # Every package is technically parented in Main but the name may not be visible in Main
        @test isdefined(@__MODULE__, :TmpPkg1)
        @test !isdefined(@__MODULE__, :TmpPkg2)
        JuliaInterpreter.split_expressions(Main, quote
                                                     module TmpPkg2
                                                     f() = 2
                                                     end
                                                 end)
        @test isdefined(@__MODULE__, :TmpPkg1)
        @test !isdefined(@__MODULE__, :TmpPkg2)
    end
end

module Toplevel end

@testset "toplevel" begin
    modexs, _ = JuliaInterpreter.split_expressions(Toplevel, read_and_parse(joinpath(@__DIR__, "toplevel_script.jl")))
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end

    @test isconst(Toplevel, :StructParent)
    @test isconst(Toplevel, :Struct)
    @test isconst(Toplevel, :MyInt8)

    s = Toplevel.Struct([2.0])

    @test Toplevel.f1(0) == 1
    @test Toplevel.f1(0.0) == 2
    @test Toplevel.f1(0.0f0) == 3
    @test Toplevel.f2("hi") == -1
    @test Toplevel.f2(UInt16(1)) == UInt16
    @test Toplevel.f2(3.2) == 0
    @test Toplevel.f2(view([1,2], 1:1)) == 2
    @test Toplevel.f2([1,2]) == 3
    @test Toplevel.f2(reshape(view([1,2], 1:2), 2, 1)) == 4
    if VERSION >= v"1.1.0"
        @test Toplevel.f3(1, 1) == 1
    end
    @test Toplevel.f3(1, :hi) == 2
    @test Toplevel.f3(UInt16(1), :hi) == Symbol
    @test Toplevel.f3(rand(2, 2), :hi, :there) == 2
    @test_throws MethodError Toplevel.f3([1.0], :hi, :there)
    @test Toplevel.f4(1, 1.0) == 1
    @test Toplevel.f4(1, 1) == Toplevel.f4(1) == 2
    @test Toplevel.f4(UInt(1), "hey", 2) == 3
    @test Toplevel.f4(rand(2,2)) == 2
    @test Toplevel.f5(Int8(1); y=22) == 22
    @test Toplevel.f5(Int16(1)) == 2
    @test Toplevel.f5(Int32(1)) == 3
    @test Toplevel.f5(Int64(1)) == 4
    @test Toplevel.f5(rand(2,2); y=7) == 2
    @test Toplevel.f6(1, "hi"; z=8) == 1
    @test Toplevel.f7(1, (1, :hi)) == 1
    @test Toplevel.f8(0) == 1
    @test Toplevel.f9(3) == 9
    @test Toplevel.f9(3.0) == 3.0
    @test s("hello") == [2.0]
    @test Toplevel.Struct{Float32}(Dict(1=>"two")) == 4
    @test Toplevel.first_two_funcs == (Toplevel.f1, Toplevel.f2)
    if VERSION >= v"1.2.0-DEV.239"
        @test isconst(Toplevel, :first_two_funcs)
    else
        @test_broken isconst(Toplevel, :first_two_funcs)
    end
    @test Toplevel.myint isa Toplevel.MyInt8
    @test_throws UndefVarError Toplevel.ffalse(1)
    @test Toplevel.ftrue(1) == 3
    @test Toplevel.fctrue(0) == 1
    @test_throws UndefVarError Toplevel.fcfalse(0)
    @test !Toplevel.Consts.b2
    @test Toplevel.fb1true(0) == 1
    @test_throws UndefVarError Toplevel.fb1false(0)
    @test Toplevel.fb2false(0) == 1
    @test_throws UndefVarError Toplevel.fb2true(0)
    @test Toplevel.fstrue(0) == 1
    @test Toplevel.fouter(1) === 2
    @test Toplevel.feval1(1.0) === 1
    @test Toplevel.feval1(1.0f0) === 1
    @test_throws MethodError Toplevel.feval1(1)
    @test Toplevel.feval2(1.0, Int8(1)) == 2
    @test length(s) === nothing
    @test size(s) === nothing
    @test Toplevel.nbytes(Float32) == 4
    @test Toplevel.typestring(1.0) == "Float64"
    @test Toplevel._feval3(0) == 3
    @test Toplevel.feval_add!(0) == 1
    @test Toplevel.feval_min!(0) == 1
    @test Toplevel.paramtype(Vector{Int8}) == Int8
    @test Toplevel.paramtype(Vector) == Toplevel.NoParam
    @test Toplevel.Inner.g() == 5
    @test Toplevel.Inner.InnerInner.g() == 6
    @test isdefined(Toplevel, :Beat)
    @test Toplevel.Beat <: Toplevel.DatesMod.Period

    @test @interpret(Toplevel.f1(0)) == 1
    @test @interpret(Toplevel.f1(0.0)) == 2
    @test @interpret(Toplevel.f1(0.0f0)) == 3
    @test @interpret(Toplevel.f2("hi")) == -1
    @test @interpret(Toplevel.f2(UInt16(1))) == UInt16
    @test @interpret(Toplevel.f2(3.2)) == 0
    @test @interpret(Toplevel.f2(view([1,2], 1:1))) == 2
    @test @interpret(Toplevel.f2([1,2])) == 3
    @test @interpret(Toplevel.f2(reshape(view([1,2], 1:2), 2, 1))) == 4
    if VERSION >= v"1.1.0"
        @test @interpret(Toplevel.f3(1, 1)) == 1
    end
    @test @interpret(Toplevel.f3(1, :hi)) == 2
    @test @interpret(Toplevel.f3(UInt16(1), :hi)) == Symbol
    @test @interpret(Toplevel.f3(rand(2, 2), :hi, :there)) == 2
    @test_throws MethodError @interpret(Toplevel.f3([1.0], :hi, :there))
    @test @interpret(Toplevel.f4(1, 1.0)) == 1
    @test @interpret(Toplevel.f4(1, 1)) == @interpret(Toplevel.f4(1)) == 2
    @test @interpret(Toplevel.f4(UInt(1), "hey", 2)) == 3
    @test @interpret(Toplevel.f4(rand(2,2))) == 2
    @test @interpret(Toplevel.f5(Int8(1); y=22)) == 22
    @test @interpret(Toplevel.f5(Int16(1))) == 2
    @test @interpret(Toplevel.f5(Int32(1))) == 3
    @test @interpret(Toplevel.f5(Int64(1))) == 4
    @test @interpret(Toplevel.f5(rand(2,2); y=7)) == 2
    @test @interpret(Toplevel.f6(1, "hi"; z=8)) == 1
    @test @interpret(Toplevel.f7(1, (1, :hi))) == 1
    @test @interpret(Toplevel.f8(0)) == 1
    @test @interpret(Toplevel.f9(3)) == 9
    @test @interpret(Toplevel.f9(3.0)) == 3.0
    @test @interpret(s("hello")) == [2.0]
    @test @interpret(Toplevel.Struct{Float32}(Dict(1=>"two"))) == 4
    @test_throws UndefVarError @interpret(Toplevel.ffalse(1))
    @test @interpret(Toplevel.ftrue(1)) == 3
    @test @interpret(Toplevel.fctrue(0)) == 1
    @test_throws UndefVarError @interpret(Toplevel.fcfalse(0))
    @test @interpret(Toplevel.fb1true(0)) == 1
    @test_throws UndefVarError @interpret(Toplevel.fb1false(0))
    @test @interpret(Toplevel.fb2false(0)) == 1
    @test_throws UndefVarError @interpret(Toplevel.fb2true(0))
    @test @interpret(Toplevel.fstrue(0)) == 1
    @test @interpret(Toplevel.fouter(1)) === 2
    @test @interpret(Toplevel.feval1(1.0)) === 1
    @test @interpret(Toplevel.feval1(1.0f0)) === 1
    @test_throws MethodError @interpret(Toplevel.feval1(1))
    @test @interpret(Toplevel.feval2(1.0, Int8(1))) == 2
    @test @interpret(length(s)) === nothing
    @test @interpret(size(s)) === nothing
    @test @interpret(Toplevel.nbytes(Float32)) == 4
    @test @interpret(Toplevel.typestring(1.0)) == "Float64"
    @test @interpret(Toplevel._feval3(0)) == 3
    @test @interpret(Toplevel.feval_add!(0)) == 1
    @test @interpret(Toplevel.feval_min!(0)) == 1
    @test @interpret(Toplevel.paramtype(Vector{Int8})) == Int8
    @test @interpret(Toplevel.paramtype(Vector)) == Toplevel.NoParam
    @test @interpret(Toplevel.Inner.g()) == 5
    @test @interpret(Toplevel.Inner.InnerInner.g()) == 6
    # FIXME: even though they pass, these tests break Test!
    # @test @interpret(isdefined(Toplevel, :Beat))
    # @test @interpret(Toplevel.Beat <: Toplevel.DatesMod.Period)

    # Check that nested expressions are handled appropriately (module-in-block, internal `using`)
    ex = quote
       module Testing
       if true
           using JuliaInterpreter
       end
       end
   end
   modexs, _ = JuliaInterpreter.split_expressions(Toplevel, ex)
   for modex in modexs
       frame = JuliaInterpreter.prepare_thunk(modex)
       while true
           JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
       end
   end
   @test Toplevel.Testing.Frame === Frame
end

# incremental interpretation solves world-age problems
# Taken straight from Julia's test/tuple.jl
module IncTest
using Test

struct A_15703{N}
    keys::NTuple{N, Int}
end

struct B_15703
    x::A_15703
end
end

ex = quote
    @testset "issue #15703" begin
        function bug_15703(xs...)
            [x for x in xs]
        end

        function test_15703()
            s = (1,)
            a = A_15703(s)
            ss = B_15703(a).x.keys
            @test ss === s
            bug_15703(ss...)
        end

        test_15703()
    end
end
modexs, _ = JuliaInterpreter.split_expressions(IncTest, ex)
for (i, modex) in enumerate(modexs)
    frame = JuliaInterpreter.prepare_thunk(modex)
    while true
        JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
    end
    if i == length(modexs)
        @test isa(JuliaInterpreter.get_return(frame), Test.DefaultTestSet)
    end
end

@testset "Enum" begin
    ex = Expr(:toplevel,
        :(@enum EnumParent begin
              EnumChild0
              EnumChild1
          end))
    modexs, _ = JuliaInterpreter.split_expressions(Toplevel, ex)
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test isa(Toplevel.EnumChild1, Toplevel.EnumParent)
end

module LowerAnon
ret = Ref{Any}(nothing)
end

@testset "Anonymous functions" begin
    ex1 = quote
        f = x -> parse(Int16, x)
        ret[] = map(f, AbstractString[])
    end
    ex2 = quote
        ret[] = map(x->parse(Int16, x), AbstractString[])
    end
    modexs, _ = JuliaInterpreter.split_expressions(LowerAnon, ex1)
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test isa(LowerAnon.ret[], Vector{Int16})
    LowerAnon.ret[] = nothing
    modexs, _ = JuliaInterpreter.split_expressions(LowerAnon, ex2)
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test isa(LowerAnon.ret[], Vector{Int16})
    LowerAnon.ret[] = nothing

    ex3 = quote
        const BitIntegerType = Union{map(T->Type{T}, Base.BitInteger_types)...}
    end
    modexs, _ = JuliaInterpreter.split_expressions(LowerAnon, ex3)
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test isa(LowerAnon.BitIntegerType, Union)

    ex4 = quote
        y = 3
        z = map(x->x^2+y, [1,2,3])
        y = 4
    end
    modexs, _ = JuliaInterpreter.split_expressions(LowerAnon, ex4)
    for modex in modexs
        frame = JuliaInterpreter.prepare_thunk(modex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test LowerAnon.z == [4,7,12]
end

@testset "Docstrings" begin
    ex = quote
        """
        A docstring
        """
        f(x) = 1

        g(T::Type) = 1
        g(x) = 2

        """
        Docstring 2
        """
        g(T::Type)

        module Sub
        """
        Docstring 3
        """
        f(x) = 2
        end
    end
    Core.eval(Toplevel, Expr(:toplevel, ex.args...))
    modexs, docexprs = JuliaInterpreter.split_expressions(Toplevel, ex; extract_docexprs=true)
    for (mod, ex) in modexs
        frame = JuliaInterpreter.prepare_thunk(mod, ex)
        while true
            JuliaInterpreter.through_methoddef_or_done!(frame) === nothing && break
        end
    end
    @test length(docexprs[Toplevel]) == 2
    @test length(docexprs[Toplevel.Sub]) == 1
end

@testset "Self referential" begin
    # Revise issue #304
    ex = :(mutable struct Node t :: Node end)
    frame = JuliaInterpreter.prepare_thunk(Toplevel, ex)
    JuliaInterpreter.finish!(frame, true)
    @test Toplevel.Node isa Type
end

@testset "Non-frames" begin
    ex = Base.parse_input_line("""
            \"\"\"
            An expr that produces an `export nffoo` that doesn't produce a Frame
            \"\"\"
            module NonFrame
            nfbar(x) = 1
            @deprecate nffoo nfbar
            global CoolStuff
            end
            """)
    modexs, _ = JuliaInterpreter.split_expressions(Toplevel, ex)
    for (mod, ex) in modexs
        frame = JuliaInterpreter.prepare_thunk(mod, ex)
        frame === nothing && continue
        JuliaInterpreter.finish!(frame, true)
    end
    Core.eval(Toplevel, :(using .NonFrame))
    @test isdefined(Toplevel, :nffoo)
end

@testset "LOAD_PATH and modules" begin
    tmpdir = joinpath(tempdir(), randstring())
    mkpath(tmpdir)
    push!(LOAD_PATH, tmpdir)
    filename = joinpath(tmpdir, "NewModule.jl")
    open(filename, "w") do io
        print(io, """
        module NewModule
        f() = 1
        end""")
    end
    str = read(filename, String)
    ex = Base.parse_input_line(str)
    modexs, _ = JuliaInterpreter.split_expressions(Main, ex)
    @test !isempty(modexs)
    pop!(LOAD_PATH)
    rm(tmpdir, recursive=true)
end

@testset "`used` for abstract types" begin
    ex = :(abstract type AbstractType <: AbstractArray{Union{Int,Missing},2} end)
    frame = JuliaInterpreter.prepare_thunk(Toplevel, ex)
    JuliaInterpreter.finish!(frame, true)
    @test isabstracttype(Toplevel.AbstractType)
end

@testset "Recursive type definitions" begin
    # See https://github.com/timholy/Revise.jl/issues/417
    ex = :(struct RecursiveType x::Vector{RecursiveType} end)
    frame = JuliaInterpreter.prepare_thunk(Toplevel, ex)
    JuliaInterpreter.finish!(frame, true)
    @test Toplevel.RecursiveType(Vector{Toplevel.RecursiveType}()) isa Toplevel.RecursiveType
end

# https://github.com/timholy/Revise.jl/issues/420
module ToplevelParameters
Base.@kwdef struct MyStruct
   x::Array{<:Real, 1} = [.05]
end
end
@testset "Nested references in type definitions" begin
    ex = quote
        Base.@kwdef struct MyStruct
           x::Array{<:Real, 1} = [.05]
        end
    end
    frame = JuliaInterpreter.prepare_thunk(ToplevelParameters, ex)
    @test JuliaInterpreter.finish!(frame, true) === nothing
end


text:
using Mamba
include("BayesTest.jl")

type TTest <: BayesTest

    typ::Symbol
    prior::Distribution
    data::Dict{AbstractString, Any}
    BF::Float64 # BF_{01}
    posterior::Any # either distribution, or NaN
    params::Array{AbstractString, 2}

    function TTest(x1, x2; typ = :paired, scale::Real = √2/2)
        self = new()
        self.typ = typ;
        self.params = ["μ", "σ", "δ"]'
        self.data = Dict("x1" => x1, "x2" => x2)
        self.prior = Distributions.Cauchy(0, scale)
        typ == :paired ? PairedSamplesTTest(self) : TwoSampleTTest(self)
    end

    function PairedSamplesTTest(self::TTest)
        x1 = self.data["x1"]
        x2 = self.data["x2"]
        @assert length(x1) == length(x2) "Number of observations in each group must be equal"
        n = length(x1)
        df = n - 1

        t = mean(x1 - x2) / (√var(x1 - x2) / √n)
        self.BF = computeBF(t, n, df, self.prior.σ)
        return self
    end

    function TwoSampleTTest(self::TTest)
        x1 = self.data["x1"]
        x2 = self.data["x2"]
        s1 = var(x1); s2 = var(x2)
        n1 = length(x1); n2 = length(x2)
        n = Int((n1 + n2) / 2) # not sure, because BF doesn't do equal sample size

        pooled = √(s1/n1 + s2/n2)
        t = (mean(x1) - mean(x2)) / pooled
        df = (s1/n1 + s2/n2)^2 / ( (s1/n1)^2 / (n1 - 1) + (s1/n2)^2 / (n2 - 1) )
        self.BF = computeBF(t, n, df, self.prior.σ)
        return self
    end

    function computeBF(t::Real, n::Real, df::Real, γ::Real)
        f(g) = ((1 + n*g)^(-1/2) * (1 + t^2 / df*(1 + n*g))^(n-2) *
                (2π)^(-1/2) * g^(-3/2) * exp(-γ^2 / (2*g)))

        num = γ * hquadrature(f, 0.00001, 10)[1]
        denom = (1 + t^2/df)^(-n/2)
        return num / denom
    end
end

function update(Model::TTest, x1, x2)
    TTest(vcat(Model.data["x1"], x1),
          vcat(Model.data["x2"], x2),
          scale = Model.prior.σ)
end

function posterior(Model::TTest)
    model = Mamba.Model(
        x1 = Stochastic(1, (mu1, s1) ->  Normal(mu1, s1), false),
        x2 = Stochastic(1, (mu2, s2) ->  Normal(mu2, s2), false),
        mu1 = Stochastic(() -> Normal(0, 100)),
        mu2 = Stochastic(() -> Normal(0, 100)),
        s1 = Stochastic(() -> Uniform(0, 100)),
        s2 = Stochastic(() -> Uniform(0, 100)),
        δ = Logical((mu1, mu2, s1, s2) -> (mu2 - mu1) / ((s1 + s2) / 2))
    )

    scheme = [Mamba.NUTS([:mu1, :mu2, :s1, :s2])]
    line = Dict{Symbol, Any}(:x1 => Model.data["x1"], :x2 => Model.data["x2"])
    inits = [Dict{Symbol,Any}(
               :x1 => line[:x1],
               :x2 => line[:x2],
               :s1 => rand(Uniform(0, 100)),
               :s2 => rand(Uniform(0, 100)),
               :mu1 => rand(Normal(0, 100)),
               :mu2 => rand(Normal(0, 100))
             ) for i in 1:3]

    Mamba.setsamplers!(model, scheme)
    sim = Mamba.mcmc(model, line, inits, 5000, burnin = 500, chains = 2)
end


text:
"""
To use this abstract type and avaliable functions,
create a subtype and implement the necessary functions.

Example:
    mutable struct Regression <: LD_Score_Regression
        ...

        function Regression(...)
            new(...)
        end
    end

    function update_weights(reg::Regression, ...)
    end

See `Hsq.jl` for a case study example.
"""

abstract type LD_Score_Regression end

# TODO need to find jackknife library or implement it
#function _delete_vals_tot(reg::LD_Score_Regression, jknife, Nbar, M) end
#function _delete_vals_part(reg::LD_Score_Regression, jknife, Nbar, M) end
#function _coef(reg::LD_Score_Regression, jknife, Nbar) end
#function _cat(reg::LD_Score_Regression, jknife, M, Nbar, coef, coef_cov) end
#function _tot(reg::LD_Score_Regression, cat, cat_cov) end
#function _prop(reg::LD_Score_Regression, jknife, M, Nbar, cat, tot) end
#function _enrichment(reg::LD_Score_Regression, M, M_tot, cat, tot) end
#function _intercept(reg::LD_Score_Regression, jknife) end
#function _combine_twostep_jknives(
#    reg::LD_Score_Regression, step1_jknife, step2_jknife, M_tot, c; Nbar=1,
#) end
#function _delete_vals_tot(reg::LD_Score_Regression) end

function aggregate(reg::LD_Score_Regression, y, x, N, M; intercept=nothing)
    if intercept == nothing
        intercept = reg.__null_intercept__
    end

    num = M * (mean(y) - intercept)
    denom = mean(x .* N)
    return num / denom
end

function ld_score_regression(
    reg::LD_Score_Regression,
    y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights,
)
    n_snp, reg.n_annot = size(x) # TODO fix this

    M_tot = sum(M)
    # shape should be [n_snp, 1]
    x_tot = sum(x, dims = 2)
    reg.constrain_intercept = intercept != nothing
    reg.intercept = intercept
    reg.n_blocks = n_blocks

    tot_agg = aggregate(reg, y, x_tot, N, M_tot; intercept=intercept)
    initial_w = update_weights(
        reg, x_tot, w, N, M_tot, tot_agg, intercept,
    )
    Nbar = mean(N)
    x = N .* x / Nbar

    if !reg.constrain_intercept
        x, x_tot = append_intercept(x), append_intercept(x_tot)
        yp = y
    else
        yp = y - intercept
        intercept_se = "NA"
    end

    reg.twostep_filtered = nothing
    if step1_ii != nothing && reg.constrain_intercept
        throw(ErrorException(
            "twostep is not compatible with constrain_intercept."
        ))
    elseif !(step1_ii == nothing) && reg.n_annot > 1
        throw(ErrorException(
            "twostep not compatible with partitioned LD Score yet."
        ))
    elseif !(step1_ii == nothing)
        n1 = sum(step1_ii)
        reg.twostep_filtered = n_snp - n1
        s = dropdims(step1_ii; dims=2)
        x1 = zeros(0)
        for (i, j) in enumerate(s)
            if s[i] == 1
                append!(x1, x[i])
            end
        end
        yp1, w1, N1, initial_w1 = map(
            (x) -> reshape(x[step1_ii], (n1, 1)), (yp, w, N, initial_w)
        )
    end

    # TODO finish implementing this function

    return reg
end