text:
using Test
using QuantumOptics
using Random, LinearAlgebra

@testset "mcwf" begin

# Define parameters for spin coupled to electric field mode.
ωc = 1.2
ωa = 0.9
g = 1.0
γ = 0.5
κ = 1.1

Ntrajectories = 500
T = Float64[0.:0.1:10.;]

# Define operators
fockbasis = FockBasis(8)
spinbasis = SpinBasis(1//2)
basis = tensor(spinbasis, fockbasis)

sx = sigmax(spinbasis)
sy = sigmay(spinbasis)
sz = sigmaz(spinbasis)
sp = sigmap(spinbasis)
sm = sigmam(spinbasis)

# Hamiltonian
Ha = embed(basis, 1, 0.5*ωa*sz)
Hc = embed(basis, 2, ωc*number(fockbasis))
Hint = sm ⊗ create(fockbasis) + sp ⊗ destroy(fockbasis)
H = Ha + Hc + Hint
Hdense = dense(H)
Hlazy = LazySum(Ha, Hc, Hint)

# Jump operators
Ja = embed(basis, 1, sqrt(γ)*sm)
Jc = embed(basis, 2, sqrt(κ)*destroy(fockbasis))
J = [Ja, Jc]
Jdense = map(dense, J)
Jlazy = [LazyTensor(basis, 1, sqrt(γ)*sm), LazyTensor(basis, 2, sqrt(κ)*destroy(fockbasis))]

# Initial conditions
Ψ₀ = spinup(spinbasis) ⊗ fockstate(fockbasis, 5)
ρ₀ = Ψ₀ ⊗ dagger(Ψ₀)

# Test mcwf
tout, Ψt = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7)
tout2, Ψt2 = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7)
@test Ψt == Ψt2
Ψ = Ψt[end]

tout, Ψt = timeevolution.mcwf(T, Ψ₀, H, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf(T, Ψ₀, H, J; seed=UInt(2), reltol=1e-6)
@test norm(Ψt[end]-Ψ) > 0.1

t_fout = Float64[]
Ψ_fout = []
function fout(t, x)
  push!(t_fout, t)
  push!(Ψ_fout, normalize(x))
end
timeevolution.mcwf(T, Ψ₀, H, J; seed=UInt(2), reltol=1e-6, fout=fout)
@test tout == t_fout && Ψt == Ψ_fout

# Test mcwf for irreducible input
tout, Ψt = timeevolution.mcwf(T, Ψ₀, Hlazy, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end] - Ψ) < 1e-5

tout, Ψt2 = timeevolution.mcwf(T, Ψ₀, H, Jlazy; seed=UInt(1), reltol=1e-6, display_beforeevent=true, display_afterevent=true)
@test norm(Ψt2[end] - Ψ) < 1e-5
@test length(Ψt2) > length(Ψt)

tout, Ψt = timeevolution.mcwf(T, Ψ₀, H, Jlazy./[sqrt(γ), sqrt(κ)]; seed=UInt(1), rates=[γ, κ], reltol=1e-6)
@test norm(Ψt[end] - Ψ) < 1e-5

# Test mcwf_h
tout, Ψt = timeevolution.mcwf_h(T, Ψ₀, H, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_h(T, Ψ₀, H, Jdense; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_h(T, Ψ₀, Hdense, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_h(T, Ψ₀, H, J; seed=UInt(2), reltol=1e-6)
@test norm(Ψt[end]-Ψ) > 0.1

# Test mcwf nh
Hnh = H - 0.5im*sum([dagger(J[i])*J[i] for i=1:length(J)])
Hnh_dense = dense(Hnh)

tout, Ψt = timeevolution.mcwf_nh(T, Ψ₀, Hnh, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_nh(T, Ψ₀, Hnh, Jdense; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_nh(T, Ψ₀, Hnh_dense, J; seed=UInt(1), reltol=1e-6)
@test norm(Ψt[end]-Ψ) < 1e-5

tout, Ψt = timeevolution.mcwf_nh(T, Ψ₀, Hnh, J; seed=UInt(2), reltol=1e-6)
@test norm(Ψt[end]-Ψ) > 0.1

# Test convergence to master solution
tout_master, ρt_master = timeevolution.master(T, ρ₀, H, J)

ρ_average = DenseOperator[0 * ρ₀ for i=1:length(T)]
for i=1:Ntrajectories
    tout, Ψt = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(i))
    for j=1:length(T)
        ρ_average[j] += (Ψt[j] ⊗ dagger(Ψt[j]))/Ntrajectories
    end
end
for i=1:length(T)
    err = tracedistance(ρt_master[i], ρ_average[i])
    @test err < 0.1
end

# Test single jump operator
J1 = [Ja]
J1_dense = map(dense, J1)
J2 = [Ja, 0 * Jc]
J2_dense = map(dense, J2)

tout_master, ρt_master = timeevolution.master(T, ρ₀, Hdense, J1_dense)

ρ_average_1 = DenseOperator[0 * ρ₀ for i=1:length(T)]
ρ_average_2 = DenseOperator[0 * ρ₀ for i=1:length(T)]
ρ_average_3 = DenseOperator[0 * ρ₀ for i=1:length(T)]
for i=1:Ntrajectories
    tout, Ψt_1 = timeevolution.mcwf(T, Ψ₀, Hdense, J1_dense; seed=UInt(i))
    tout, Ψt_2 = timeevolution.mcwf(T, Ψ₀, Hdense, J2_dense; seed=UInt(i))
    tout, Ψt_3 = timeevolution.mcwf(T, Ψ₀, Hdense, [J1_dense[1]/sqrt(γ)]; seed=UInt(i), rates=[γ])
    for j=1:length(T)
        ρ_average_1[j] += (Ψt_1[j] ⊗ dagger(Ψt_1[j]))/Ntrajectories
        ρ_average_2[j] += (Ψt_2[j] ⊗ dagger(Ψt_2[j]))/Ntrajectories
        ρ_average_3[j] += (Ψt_3[j] ⊗ dagger(Ψt_3[j]))/Ntrajectories
    end
end
for i=1:length(T)
    @test tracedistance(ρt_master[i], ρ_average_1[i]) < 0.1
    @test tracedistance(ρt_master[i], ρ_average_2[i]) < 0.1
    @test tracedistance(ρt_master[i], ρ_average_3[i]) < 0.1
end

# Test displaying before/after jump
tout, Ψt = timeevolution.mcwf([T[1],T[end]], Ψ₀, Hdense, J1; seed=2, display_beforeevent=true, display_afterevent=true)
for i=2:length(tout)-1
    if tout[i+1] == tout[i]
        @test Ψt[i+1].data ≈ normalize(J1[1]*Ψt[i]).data
    end
end

# Test equivalence to schroedinger time evolution for no decay
J = DenseOperator[]
tout_schroedinger, Ψt_schroedinger = timeevolution.schroedinger(T, Ψ₀, Hdense)
tout_mcwf, Ψt_mcwf = timeevolution.mcwf(T, Ψ₀, Hdense, J)
tout_mcwf_h, Ψt_mcwf_h = timeevolution.mcwf_h(T, Ψ₀, Hdense, J)
tout_mcwf_nh, Ψt_mcwf_nh = timeevolution.mcwf_nh(T, Ψ₀, Hdense, J)

for i=1:length(T)
    @test norm(Ψt_mcwf[i] - Ψt_schroedinger[i]) < 1e-4
    @test norm(Ψt_mcwf_h[i] - Ψt_schroedinger[i]) < 1e-4
    @test norm(Ψt_mcwf_nh[i] - Ψt_schroedinger[i]) < 1e-4
end

# Test diagonal jump operators
threespinbasis = spinbasis ⊗ spinbasis ⊗ spinbasis
Γ, γ1, γ2, = 1.0, 1/sqrt(2), 1/sqrt(3)
rates = [Γ γ1 γ2; γ1 Γ γ1; γ2 γ1 Γ]
J3 = [embed(threespinbasis, i, sm) for i=1:3]
H = sum(J3) + dagger(sum(J3))
d, diagJ = diagonaljumps(rates, J3)
ψ3 = spindown(spinbasis) ⊗ spindown(spinbasis) ⊗ spindown(spinbasis)
tout, ρ3_nondiag = timeevolution.master(T, ψ3, H, J3; rates=rates)
tout, ρ3_diag = timeevolution.master(T, ψ3, H, diagJ; rates=d)

ρ3_avg = DenseOperator[0*ρ3_diag[1] for i=1:length(T)]
for i=1:Ntrajectories
    tout, ψ3t = timeevolution.mcwf(T, ψ3, H, diagJ; rates=d)
    for j=1:length(T)
        ρ3_avg[j] += (ψ3t[j] ⊗ dagger(ψ3t[j]))/Ntrajectories
    end
end

dist = []
for i=1:length(tout)
  @test tracedistance(ρ3_nondiag[i], ρ3_diag[i]) < 1e-14
  @test tracedistance(ρ3_avg[i], ρ3_diag[i]) < 0.1
end

@test_throws ArgumentError timeevolution.mcwf(T, ψ3, H, J3; rates=rates)

J3_lazy = [LazyTensor(threespinbasis, i, sm) for i=1:3]
d, diagJ_lazy = diagonaljumps(rates, J3_lazy)
for i=1:3
    @test dense(diagJ_lazy[i]) == dense(diagJ[i])
end

# Test dynamic
H = sp + sm
J = [sm]
Jdagger = dagger.(J)
function Ht(t, psi)
    H*exp(-(5-t)^2), J, Jdagger
end
function Ht2(t, psi)
    H*exp(-(5-t)^2), J, Jdagger, [1.0]
end
ψ0 = spindown(spinbasis)
ρ0 = dm(ψ0)
tout_master, ρt_master = timeevolution.master_dynamic(T, ρ0, Ht)

ρ_average_1 = DenseOperator[0 * ρ0 for i=1:length(T)]
ρ_average_2 = DenseOperator[0 * ρ0 for i=1:length(T)]
ρ_average_3 = DenseOperator[0 * ρ0 for i=1:length(T)]
for i=1:Ntrajectories
    tout, Ψt_1 = timeevolution.mcwf_dynamic(T, ψ0, Ht; seed=UInt(i))
    tout, Ψt_2 = timeevolution.mcwf_dynamic(T, ψ0, Ht; seed=UInt(i))
    tout, Ψt_3 = timeevolution.mcwf_dynamic(T, ψ0, Ht2; seed=UInt(i))
    for j=1:length(T)
        ρ_average_1[j] += (Ψt_1[j] ⊗ dagger(Ψt_1[j]))/Ntrajectories
        ρ_average_2[j] += (Ψt_2[j] ⊗ dagger(Ψt_2[j]))/Ntrajectories
        ρ_average_3[j] += (Ψt_3[j] ⊗ dagger(Ψt_3[j]))/Ntrajectories
    end
end

# Test non-hermitian dynamic
H_nh = -0.5im*sum(Jdagger.*J)
function Ht_nh(t, psi)
    H*exp(-(5-t)^2) + H_nh, J, Jdagger
end

ρ_average_4 = DenseOperator[0 * ρ0 for i=1:length(T)]
for i=1:Ntrajectories
    tout, Ψt_4 = timeevolution.mcwf_nh_dynamic(T, ψ0, Ht_nh; seed=UInt(i))
    for j=1:length(T)
        ρ_average_4[j] += (Ψt_4[j] ⊗ dagger(Ψt_4[j]))/Ntrajectories
    end
end
for i=1:length(T)
    @test tracedistance(ρt_master[i], ρ_average_1[i]) < 0.1
    @test tracedistance(ρt_master[i], ρ_average_2[i]) < 0.1
    @test tracedistance(ρt_master[i], ρ_average_3[i]) < 0.1
    @test tracedistance(ρt_master[i], ρ_average_4[i]) < 0.1
end

# Test displaying of jumps
tout, Ψt, t_jump, j_index = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7, display_jumps=true)
tout, Ψt, t_jump2, j_index2 = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7, display_jumps=true)

@test length(j_index) == length(t_jump) == length(t_jump2) == length(j_index2)
@test j_index == j_index2
@test t_jump == t_jump2

ψ0 = spinup(spinbasis)⊗fockstate(fockbasis,0)
tout, ψt, t_jump, j_index = timeevolution.mcwf(T, ψ0, 0Hdense, Jdense; display_jumps=true, seed=2)

@test length(t_jump) == length(j_index) == 1
@test j_index[1] == 1

end # testset


text:
# Extend Ditributions type hiearchy to account for non-euclidean supports
abstract type NonEuclidean <: Dist.ValueSupport end

const UnivariateFiniteSuper = Dist.Distribution{Dist.Univariate,NonEuclidean}

# R - reference type <: Unsigned
# V - type of class labels (eg, Char in `categorical(['a', 'b'])`)
# P - raw probability type
# S - scitype of samples
# L - raw type of labels, eg `Symbol` or `String`

# Note that the keys of `prob_given_ref` need not exhaust all the
# refs of all classes but will be ordered (LittleDicts preserve order)
struct UnivariateFinite{S,V,R,P<:Real} <: UnivariateFiniteSuper
    scitype::Type{S}
    decoder::CategoricalDecoder{V,R}
    prob_given_ref::LittleDict{R,P,Vector{R}, Vector{P}}
end

struct UnivariateFiniteArray{S,V,R,P,N} <:
    AbstractArray{UnivariateFinite{S,V,R,P},N}
    scitype::Type{S}
    decoder::CategoricalDecoder{V,R}
    prob_given_ref::LittleDict{R,Array{P,N},Vector{R}, Vector{Array{P,N}}}
end

const UnivariateFiniteVector{S,V,R,P} = UnivariateFiniteArray{S,V,R,P,1}

## CONSTRUCTORS TO EXTEND MLJModelInterface

"""$(MMI.UNIVARIATE_FINITE_DOCSTRING)"""
UnivariateFinite(a...; kwargs...) = MMI.UnivariateFinite(a...; kwargs...)

## CHECKS AND ERROR MESSAGES

const Prob{P} = Union{P, AbstractArray{P}} where P <: Real

prob_error = ArgumentError("Probabilities must have `Real` type. ")

_err_01() = throw(DomainError("Probabilities must be in [0,1]."))
_err_sum_1() = throw(DomainError(
    "Probability arrays must sum to one along the last axis. Perhaps "*
"you meant to specify `augment=true`? "))
_err_dim(support, probs) = throw(DimensionMismatch(
"Probability array is incompatible "*
"with the number of classes, $(length(support)), which should "*
"be equal to `$(size(probs)[end])`, the last dimension "*
"of the array. Perhaps you meant to set `augment=true`? "))
_err_dim_augmented(support, probs) = throw(DimensionMismatch(
"Probability array to be augmented is incompatible "*
"with the number of classes, $(length(support)), which should "*
"be one more than `$(size(probs)[end])`, the last dimension "*
    "of the array. "))
_err_aug() = throw(ArgumentError(
    "Array cannot be augmented. There are "*
    "sums along the last axis exceeding one. "))

function _check_pool(pool)
    ismissing(pool) || pool == nothing ||
        @warn "Specified pool ignored, as class labels being "*
    "generated automatically. "
    return nothing
end
_check_probs_01(probs) =
    all(0 .<= probs .<= 1) || _err_01()
_check_probs_sum(probs::Vector{<:Prob{P}}) where P<:Real =
    all(x -> x≈one(P), sum(probs)) || _err_sum_1()
_check_probs(probs) = (_check_probs_01(probs); _check_probs_sum(probs))
_check_augmentable(support, probs) = _check_probs_01(probs) &&
    size(probs)[end] + 1 == length(support) ||
    _err_dim_augmented(support, probs)

## AUGMENTING ARRAYS TO MAKE THEM PROBABILITY ARRAYS

_unwrap(A::Array) = A
_unwrap(A::Vector) = first(A)

isbinary(support) = length(support) == 2

# augmentation inserts the sum-subarray *before* the array:
_augment_probs(support, probs) =
    _augment_probs(Val(isbinary(support)), support, probs,)
function _augment_probs(::Val{false},
                        support,
                        probs::AbstractArray{P,N}) where {P,N}
    _check_augmentable(support, probs)
    aug_size = size(probs) |> collect
    aug_size[end] += 1
    augmentation = _unwrap(one(P) .- sum(probs, dims=N))
    all(0 .<= augmentation .<= 1) || _err_aug()
    aug_probs = Array{P}(undef, aug_size...)
    aug_probs[fill(:, N - 1)..., 2:end] = probs
    aug_probs[fill(:, N - 1)..., 1] = augmentation
    return aug_probs
end
function _augment_probs(::Val{true},
                        support,
                        probs::AbstractArray{P,N}) where {P,N}
    _check_probs_01(probs)
    aug_size = [size(probs)..., 2]
    augmentation = one(P) .- probs
    all(0 .<= augmentation .<= 1) || _err_aug()
    aug_probs = Array{P}(undef, aug_size...)
    aug_probs[fill(:, N)..., 2] = probs
    aug_probs[fill(:, N)..., 1] = augmentation
    return aug_probs
end

## CONSTRUCTORS - FROM DICTIONARY

# The following constructor will get called by all the others. It
# returns a UnivariateFinite object *or* a UnivariateFiniteArray,
# depending on the values of the dictionary - scalar or array - which
# represent the probabilities, one for each class in the support.
function MMI.UnivariateFinite(
    ::FI,
    prob_given_class::AbstractDict{<:CategoricalValue, <:Prob};
    kwargs...)

    # this constructor ignores kwargs

    probs = values(prob_given_class) |> collect
    _check_probs_01.(probs)
    _check_probs_sum(probs)

    # retrieve decoder and classes from element
    class1         = first(keys(prob_given_class))
    parent_decoder = decoder(class1)
    parent_classes = classes(class1)

    # `LittleDict`s preserve order of keys, which we need for rand():

    support  = keys(prob_given_class) |> collect |> sort

    issubset(support, parent_classes) ||
        error("Categorical elements are not from the same pool. ")

    pairs = [int(c) => prob_given_class[c]
                for c in support]

    probs1 = first(values(prob_given_class))
    S = scitype(class1)
    if probs1 isa Real
        return UnivariateFinite(S, parent_decoder, LittleDict(pairs...))
    else
        return UnivariateFiniteArray(S, parent_decoder, LittleDict(pairs...))
    end
end

function MMI.UnivariateFinite(::FI, d::AbstractDict{V,<:Prob};
                              pool=nothing,
                              ordered=false) where V

    ismissing(pool) &&
        throw(ArgumentError(
            "You cannot specify `pool=missing` "*
            "if passing `UnivariateFinite` a dictionary"))

    pool === nothing && throw(ArgumentError(
        "You must specify `pool=c` "*
        "where `c` is a "*
        "`CategoricalArray`, `CategoricalArray` or "*
        "CategoricalPool`"))

    ordered && @warn "Ignoring `ordered` key-word argument as using "*
    "specified pool to order. "

    raw_support = keys(d) |> collect
    _classes = classes(pool)
    issubset(raw_support, _classes) ||
        error("Specified support, $raw_support, not contained in "*
              "specified pool, $(levels(classes)). ")
    support = filter(_classes) do c
        c in raw_support
    end

    prob_given_class = LittleDict([c=>d[get(c)] for c in support])

    return UnivariateFinite(FI(), prob_given_class)
end

## CONSTRUCTORS - FROM ARRAYS

# example: _get(A, 4) = A[:, :, 4] if A has 3 dims:
_get(probs::AbstractArray{<:Any,N}, i) where N = probs[fill(:,N-1)..., i]

# 1. Univariate Finite from a vector of classes or raw labels and
# array of probs; first, a dispatcher:
function MMI.UnivariateFinite(
    ::FI,
    support::AbstractVector,
    probs::Union{AbstractArray,Real};
    kwargs...)

    if support isa AbstractArray{<:CategoricalValue}
        if :pool in keys(kwargs)
            @warn "Ignoring value of `pool` as the specified "*
            "support defines one already. "
        end
        if :ordered in keys(kwargs)
            @warn "Ignoring value of `ordered` as the "*
            "specified support defines an order already. "
        end
    end

    return _UnivariateFinite(Val(isbinary(support)),
                             support,
                             probs;
                             kwargs...)
end

# The core method, ultimately called by 1.0, 1.1, 1.2, 1.3 below, or
# directly from the dispatcher 1. above
function _UnivariateFinite(support::AbstractVector{CategoricalValue{V,R}},
                           probs::AbstractArray{P},
                           N;
                           augment=false,
                           kwargs...) where {V,R,P<:Real}

    unique(support) == support ||
        error("Non-unique vector of classes specified")

    _probs = augment ? _augment_probs(support, probs) : probs

    # it's necessary to force the typing of the LittleDict otherwise it
    # flips to Any type (unlike regular Dict):

    if N == 0
        prob_given_class = LittleDict{CategoricalValue{V,R},P}()
    else
        prob_given_class =
            LittleDict{CategoricalValue{V,R}, AbstractArray{P,N}}()
    end
    for i in eachindex(support)
        prob_given_class[support[i]] = _get(_probs, i)
    end

    # calls dictionary constructor above:
    return MMI.UnivariateFinite(FI(), prob_given_class; kwargs...)
end

# 1.0 support does not consist of categorical elements:
function _UnivariateFinite(support::AbstractVector{L},
                           probs::AbstractArray{P},
                           N;
                           augment=false,
                           pool=nothing,
                           ordered=false) where {L,P<:Real}

    # If we got here, then L<:CategoricalValue is not true, ie L is a
    # raw label type

    if pool === nothing || ismissing(pool)
        if pool === nothing
            @warn "No `CategoricalValue` found from which to extract a "*
            "complete pool of classes. "*
            "Creating a new pool (ordered=$ordered). "*
            "You can:\n"*
            " (i) specify `pool=missing` to suppress this warning; or\n"*
            " (ii) use an existing pool by specifying `pool=c` "*
            "where `c` is a "*
            "`CategoricalArray`, `CategoricalValue` or "*
            "CategoricalPool`.\n"*
            "In case (i) "*
            "specify `ordered=true` if samples are to be `OrderedFactor`. "
        end
        v = categorical(support, ordered=ordered, compress=true)
        levels!(v, support)
        _support = classes(v)
    else
        _classes = classes(pool)
        issubset(support, _classes) ||
            error("Specified support, $support, not contained in "*
                  "specified pool, $(levels(classes)). ")
        _support = filter(_classes) do c
            c in support
        end
    end

    # calls core method:
    return _UnivariateFinite(_support, probs, N;
                             augment=augment, pool=pool, ordered=ordered)
end

# 1.1 generic (non-binary) case:
_UnivariateFinite(::Val{false},
                  support::AbstractVector,
                  probs::AbstractArray{<:Any,M};
                  augment=false,
                  kwargs...) where M =
                      _UnivariateFinite(support,
                                        probs,
                                        M - 1;
                                        augment=augment,
                                        kwargs...)

# 1.2 degenerate (binary) case:
_UnivariateFinite(::Val{true},
                  support::AbstractVector,
                  probs::AbstractArray{<:Any,M};
                  augment=false,
                  kwargs...) where M =
                      _UnivariateFinite(support,
                                        probs,
                                        augment ? M : M - 1;
                                        augment=augment,
                                        kwargs...)

# 1.3 corner case, probs a scalar:
_UnivariateFinite(::Val{true},
                  support::AbstractVector,
                  probs::Real;
                  kwargs...) =
                      UnivariateFinite(support, [probs,]; kwargs...)[1]

# 2. probablity only; unspecified support:
function MMI.UnivariateFinite(::FI,
                              probs::AbstractArray{<:Real,N};
                              pool=nothing,
                              augment=false,
                              kwargs...) where N
    _check_pool(pool)

    # try to infer number of classes:
    if N == 1
        if augment
            c = 2
        else
            c = length(probs)
        end
    elseif N == 2
        if augment
            c = size(probs, 2) + 1
        else
            c = size(probs, 2)
        end
    else
        throw(ArgumentError(
            "You need to explicitly specify a support for "*
            "probablility arrays of three "*
            "or more dimensions. "))
    end

    support = ["class_$i" for i in 1:c]

    return MMI.UnivariateFinite(FI(),
                                support,
                                probs;
                                pool=pool,
                                augment=augment,
                                kwargs...)
end

ext: jl


text:
abstract type Context end

struct HighwayDrivingContext{VM, VS} <: Context
    vehicle_model::VM
    vehicle_states::Vector{VS}
    roadway::Highway

    _fore_indices::Cached{Vector{Int}}
    _sorted_in_lanes::Cached{Vector{Vector{Int}}}
end

HighwayDrivingContext(vm, vss, rw) = HighwayDrivingContext(vm, vss, rw, Cached{Vector{Int}}())

# public interface
function headway(c::HighwayDrivingContext, idx)
    fidx = fore_index(c, idx)
    if fidx == 0
        return Inf
    else
        s = c.vehicle_states[idx]
        v = c.vehicle_model
        fs = c.vehicle_states[fidx]
        fv = c.vehicle_model
        return pos(fs)[1] - pos(s)[1] - length(v)/2 - length(fv)/2
    end
end

vehicle(c::HighwayDrivingContext, idx) = c.vehicle_model
state(c::HighwayDrivingContext, idx) = c.vehicle_states[idx]
n_vehicles(c::HighwayDrivingContext) = length(c.vehicle_states)

"Return a 2-tuple containing at most two lanes that the car is in."
function occupied_lanes(c::HighwayDrivingContext, idx)
    vs = c.vehicle_states[idx]
    vm = c.vehicle_model
    h = c.roadway
    y = pos(vs)[2]
    l = y + width(vm)/2
    r = y - width(vm)/2
    lo = clamp(ceil(Int, r/h.lanewidth), 1, h.n_lanes)
    hi = clamp(ceil(Int, l/h.lanewidth), 1, h.n_lanes)
    return (lo, hi)
end

"Iterator over each car index."
indexes(c::HighwayDrivingContext) = 1:n_vehicles(c)

# less-public interface
# function neighborhood(c::HighwayDrivingContext, idx)
#     ns = get(_calc_neighborhoods, c.neighborhoods, c)
#     return ns[idx]
# end

function fore_index(c::HighwayDrivingContext, idx)
    fis = get(c._fore_indices, c) do c
        # indices for cars in each lane sorted by x
        carsin = sorted_in_lanes(c)

        # sweep through the lanes and construct the fore indices
        fore_indices = zeros(Int, n_vehicles(c))
        for l in 1:c.roadway.n_lanes
            for j in 1:length(carsin[l])-1
                fore_indices[carsin[l][j]] = carsin[l][j+1]
            end
        end
        return fore_indices
    end
    return fis[idx]
end

function sorted_in_lanes(c::HighwayDrivingContext)
    return get(c._sorted_states, c) do c
        carsin = [Int[] for i in 1:c.roadway.n_lanes]

        # for each car, insert it into the lane list while maintaining order
        for i in indexes(c)
            l1, l2 = occupied_lanes(c, i)
            spot = searchsortedfirst(carsin[l1], i, by=ind->pos(c.vehicle_states[ind])[1])
            insert!(carsin[l1], spot, i)
            if l2 != l1
                spot = searchsortedfirst(carsin[l2], i, by=ind->pos(c.vehicle_states[ind])[1])
                insert!(carsin[l2], spot, i)
            end
        end
        return carsin
    end
end

"Index of the first car in the given lane"
function lane_head(c::HighwayDrivingContext, lane::Int)
    inlane = sorted_in_lanes(c)[lane]
    if isempty(inlane)
        return 0
    else
        return last(inlane)
    end
end

"Index of the last care in the given lane"
function lane_tail(c::HighwayDrivingContext, lane::Int)
    inlane = sorted_in_lanes(c)[lane]
    if isempty(inlane)
        return 0
    else
        return first(inlane)
    end
end

# internal
function _calc_neighborhoods

end


text:
"""
    Erlang(α,θ)

The *Erlang distribution* is a special case of a [`Gamma`](@ref) distribution with integer shape parameter.

```julia
Erlang()       # Erlang distribution with unit shape and unit scale, i.e. Erlang(1, 1)
Erlang(a)      # Erlang distribution with shape parameter a and unit scale, i.e. Erlang(a, 1)
Erlang(a, s)   # Erlang distribution with shape parameter a and scale s
```

External links

* [Erlang distribution on Wikipedia](http://en.wikipedia.org/wiki/Erlang_distribution)

"""
struct Erlang{T<:Real} <: ContinuousUnivariateDistribution
    α::Int
    θ::T
    Erlang{T}(α::Int, θ::T) where {T} = new{T}(α, θ)
end

function Erlang(α::Real, θ::Real; check_args::Bool=true)
    @check_args Erlang (α, isinteger(α)) (α, α >= zero(α))
    return Erlang{typeof(θ)}(α, θ)
end

function Erlang(α::Integer, θ::Real; check_args::Bool=true)
    @check_args Erlang (α, α >= zero(α))
    return Erlang{typeof(θ)}(α, θ)
end

function Erlang(α::Integer, θ::Integer; check_args::Bool=true)
    return Erlang(α, float(θ); check_args=check_args)
end

Erlang(α::Integer=1) = Erlang(α, 1.0; check_args=false)

@distr_support Erlang 0.0 Inf

#### Conversions
function convert(::Type{Erlang{T}}, α::Integer, θ::S) where {T <: Real, S <: Real}
    Erlang(α, T(θ), check_args=false)
end
function Base.convert(::Type{Erlang{T}}, d::Erlang) where {T<:Real}
    Erlang{T}(d.α, T(d.θ))
end
Base.convert(::Type{Erlang{T}}, d::Erlang{T}) where {T<:Real} = d

#### Parameters

shape(d::Erlang) = d.α
scale(d::Erlang) = d.θ
rate(d::Erlang) = inv(d.θ)
params(d::Erlang) = (d.α, d.θ)
@inline partype(d::Erlang{T}) where {T<:Real} = T

#### Statistics

mean(d::Erlang) = d.α * d.θ
var(d::Erlang) = d.α * d.θ^2
skewness(d::Erlang) = 2 / sqrt(d.α)
kurtosis(d::Erlang) = 6 / d.α

function mode(d::Erlang; check_args::Bool=true)
    α, θ = params(d)
    @check_args(
        Erlang,
        (α, α >= 1, "Erlang has no mode when α < 1"),
    )
    θ * (α - 1)
end

function entropy(d::Erlang)
    (α, θ) = params(d)
    α + loggamma(α) + (1 - α) * digamma(α) + log(θ)
end

mgf(d::Erlang, t::Real) = (1 - t * d.θ)^(-d.α)
cf(d::Erlang, t::Real)  = (1 - im * t * d.θ)^(-d.α)

#### Evaluation & Sampling

@_delegate_statsfuns Erlang gamma α θ

rand(rng, ::AbstractRNG, d::Erlang) = rand(rng, Gamma(Float64(d.α), d.θ))
sampler(d::Erlang) = Gamma(Float64(d.α), d.θ)


text:
module TDD

# Write your package code here.
    include("interval.jl")

end


text:
@testset "[nl_write] Temp file handling" begin
    # Turn on debug mode so files persist
    old_debug = AmplNLWriter.debug
    AmplNLWriter.setdebug(true)

    filename = "test"
    filepath = joinpath(AmplNLWriter.solverdata_dir, filename)
    probfile = "$filepath.nl"
    solfile = "$filepath.sol"
    AmplNLWriter.clean_solverdata()

    @testset "all temp files deleted successfully" begin
        @test length(readdir(AmplNLWriter.solverdata_dir)) == 1
    end

    solver = AmplNLSolver(Ipopt.amplexe, filename=filename)#BonminNLSolver(filename=filename)
    m = Model(solver=solver)
    @variable(m, x >= 0)
    @objective(m, Min, x)
    solve(m)

    @testset "temp files present after solve in debug mode" begin
        @test length(readdir(AmplNLWriter.solverdata_dir)) == 3
    end
    @testset "temp files used custom name" begin
        @test isfile(probfile) == true
        @test isfile(solfile) == true
    end

    # Reset debug mode and clean up
    AmplNLWriter.setdebug(old_debug)
    AmplNLWriter.clean_solverdata()
end


text:
using DataFrames
using RDatasets
using Distributions
using Gadfly

srand(619)

names(Normal)

params(dist1)

dist1.μ

dist1.σ

x = rand(dist1, 1000)

dist2 = Truncated(dist1, -4.0, 6.0)

plot(x = rand(dist1, 1000), Geom.histogram(bincount = 30))

max_stars_repo_path: Chapter 4/1.jl

max_stars_repo_name: PacktPublishing/Julia-for-data-science

max_issues_repo_path: Chapter 4/1.jl

max_issues_repo_name: PacktPublishing/Julia-for-data-science

max_forks_repo_path: Chapter 4/1.jl

max_forks_repo_name: PacktPublishing/Julia-for-data-science


text:
# ---
# title: JuliaRL\_CRR\_CartPole
# cover: assets/JuliaRL_CRR_CartPole_medium.png
# description: CRR applied to CartPole
# date: 2021-09-17
# author: "[Guoyu Yang](https://github.com/pilgrimygy)"
# ---

#+ tangle=true
using ReinforcementLearning
using StableRNGs
using Flux
using Flux.Losses

function RL.Experiment(
    ::Val{:JuliaRL},
    ::Val{:CRR},
    ::Val{:CartPole},
    type::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    rng = StableRNG(seed)
    UPDATE_FREQ = 10
    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(state(env)), length(action_space(env))

    trajectory_num = 10000
    dataset_size = 10000
    batch_size = 64

    init = glorot_uniform(rng)

    create_actor_critic() = ActorCritic(
        actor = Chain(
            Dense(ns, 128, relu; init),
            Dense(128, 128, relu; init),
            Dense(128, na; init),
        ),
        critic = Chain(
            Dense(ns, 128, relu; init),
            Dense(128, 128, relu; init),
            Dense(128, na; init),
        ),
        optimizer = ADAM(1e-3),
    )

    agent = Agent(
        policy = OfflinePolicy(
            learner = CRRLearner(
                approximator = create_actor_critic() |> cpu,
                target_approximator = create_actor_critic() |> cpu,
                γ = 0.99f0,
                batch_size = batch_size,
                continuous = false,
                policy_improvement_mode = :exp,
                ratio_upper_bound = 20.0f0,
                β = 1.0f0,
                advantage_estimator = :mean,
                m = 4,
                update_freq = UPDATE_FREQ,
            ),
            dataset = gen_JuliaRL_dataset(:BasicDQN, :CartPole, type; dataset_size = dataset_size),
            continuous = false,
            batch_size = batch_size,
        ),
        trajectory = CircularArraySARTTrajectory(
            capacity = 1000,
            state = Vector{Float32} => (ns,),
        ),
    )

    stop_condition = StopAfterStep(trajectory_num, is_show_progress=!haskey(ENV, "CI"))
    hook = TotalRewardPerEpisode()
    Experiment(agent, env, stop_condition, hook, "CRR <-> CartPole ($type dataset)")
end

#+ tangle=false
using Plots
pyplot() #hide
ex = E`JuliaRL_CRR_CartPole(medium)`
run(ex)
plot(ex.hook.rewards)
savefig("assets/JuliaRL_CRR_CartPole_medium.png") #hide

# ![](assets/JuliaRL_CRR_CartPole.png)


text:
using SpecialFunctions

"""
Base type of FractionalDiffEq algorithms
"""
abstract type FractionalDiffEqAlgorithm end

"""
@article{
title={A predictor-corrector approach for the numerical solution of fractional differential equations},
author={Diethelm, Kai and Ford, Neville J. and Freed, Alan D.}
doi={https://doi.org/10.1023/A:1016592219341}
}
"""

"""
Predict-Evaluate-Correct-Evaluate algorithm.

For more details, please refer to [Predictor-Corrector algorithms](https://en.wikipedia.org/wiki/Predictor%E2%80%93corrector_method)

This PECE algorithm is taken from Diethelm's paper.

```tex
@article{
title={A predictor-corrector approach for the numerical solution of fractional differential equations},
author={Diethelm, Kai and Ford, Neville J. and Freed, Alan D.}
doi={https://doi.org/10.1023/A:1016592219341}
}
```
"""
struct PECE <: FractionalDiffEqAlgorithm end
#TODO: Use Richardson extrapolation to refine the PECE algorithms

"""
    FDEProblem(f, α, u0, T, h)

Define a Fractional Differential in time interval [0, T] with initial value y(0)=u₀, α-order derivative and step size h.
"""
struct FDEProblem{T<:Float64}
    f
    α::T
    u0
    T
    h::T
end

"""
    solve(FDEProblem, PECE())

After define the FDEProblem, use **PECE(Predict-Evaluate-Correct-Evaluate) algorithm** to computing the Fractional Differential Equation
"""
function solve(FDE::FDEProblem, ::PECE)
    f, α, u0, T, h = FDE.f, FDE.α, FDE.u0, FDE.T, FDE.h
    N=Int64(T/h)
    y=zeros(N+1)
    leftsum = 0

    if floor(α)==0
        leftsum=u0
    elseif floor(α)==1
        leftsum=u0+T*u0
    end

    @fastmath @inbounds @simd for n in range(0, N, step=1)
        y[Int64(n+1)]=leftsum + h^α/gamma(α+2)*f((n+1)*h, predictor(f, y, α, n, h, u0, T))+h^α/gamma(α+2)*right(f, y, α, n, h)
    end

    return y
end

function right(f, y, α, n, h)
    temp = 0

    for j in range(0, n, step=1)
        temp+=A(j, n, α)*f(j*h, y[Int64(j+1)])
    end

    return temp
end

function predictor(f, y, α, n, h, u0, T)
     predict = 0
     leftsum = 0

     if floor(α)==0
        leftsum=u0
     elseif floor(α)==1
        leftsum=u0+T*u0
     end

     @fastmath @inbounds @simd for j in range(0, n, step=1)
        predict+=B(j, n, α, h)*f(j*h, y[Int64(j+1)])
     end

     #return u0+predict
     return leftsum+predict
end

function A(j, n, α)
    if j==0
        return n^(α+1)-(n-α)*(n+1)^α
    elseif 1 ≤ j ≤ n
        return (n-j+2)^(α+1)+(n-j)^(α+1)-2*(n-j+1)^(α+1)
    elseif j==n+1
        return 1
    end
end

function B(j, n, α, h)
    return h^α/α*((n+1-j)^α-(n-j)^α)
end


text:
using Documenter, ClimateModels, CairoMakie, Pkg
import PlutoSliderServer
Pkg.precompile()

lst=("defaults.jl","RandomWalker.jl","ShallowWaters.jl","Oceananigans.jl","Hector.jl","FaIR.jl","MITgcm.jl","Speedy.jl","CMIP6.jl","IPCC.jl")

makedocs(;
    modules=[ClimateModels],
    format=Documenter.HTML(),
    doctest=false,
    pages=[
        "Home" => "index.md",
        "Examples" => "examples.md",
        "Manual" => "functionalities.md",
    ],
    repo="https://github.com/gaelforget/ClimateModels.jl/blob/{commit}{path}#L{line}",
    sitename="ClimateModels.jl",
    authors="gaelforget <gforget@mit.edu>",
    assets=String[],
)

for i in lst
    fil_in=joinpath(@__DIR__,"..", "examples",i)
    fil_out=joinpath(@__DIR__,"build", "examples",i[1:end-2]*"html")
    PlutoSliderServer.export_notebook(fil_in)
    #somehow with pycall, this fails at first, but second call works
    i=="FaIR.jl" ? PlutoSliderServer.export_notebook(fil_in) : nothing
    mv(fil_in[1:end-2]*"html",fil_out)
    cp(fil_in,fil_out[1:end-4]*"jl")
end

fil_in=joinpath(@__DIR__,"build","ClimateModelsJuliaCon2021.jl")
PlutoSliderServer.export_notebook(fil_in)

deploydocs(;
    repo="github.com/gaelforget/ClimateModels.jl",
)


text:
struct Canvas{T} <: AbstractArray{Color{T},2}
	pixels::Matrix{Color{T}}
end
function Canvas{T}(width::Int,length::Int) where T
	Canvas(zeros(Color{T},width,length))
end

Base.size(canvas::Canvas) = size(canvas.pixels)
Base.getindex(canvas::Canvas,inds...) = getindex(canvas.pixels,inds...)
Base.setindex!(canvas::Canvas,inds...) = setindex!(canvas.pixels,inds...)


text:
function tensor_create_linear(T::Union{Type{Tensor{order, dim}}, Type{SymmetricTensor{order, dim}}}, f) where {order, dim}
    return Expr(:tuple, [f(i) for i=1:n_components(T)]...)
end

function tensor_create(::Type{Tensor{order, dim}}, f) where {order, dim}
    if order == 1
        ex = Expr(:tuple, [f(i) for i=1:dim]...)
    elseif order == 2
        ex = Expr(:tuple, [f(i,j) for i=1:dim, j=1:dim]...)
    elseif order == 4
        ex = Expr(:tuple, [f(i,j,k,l) for i=1:dim, j=1:dim, k = 1:dim, l = 1:dim]...)
    end
    return ex
end

function tensor_create(::Type{SymmetricTensor{order, dim}}, f) where {order, dim}
    ex = Expr(:tuple)
    if order == 2
        for j in 1:dim, i in j:dim
            push!(ex.args, f(i, j))
        end
    elseif order == 4
        for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
            push!(ex.args, f(i, j, k, l))
        end
    end
    return ex
end

# reduced two expressions by summing the products
# madd = true uses muladd instructions which is faster
# in some cases, like in double contraction
function reducer(ex1i, ex2i, madd=false)
    ex1, ex2 = remove_duplicates(ex1i, ex2i)
    N = length(ex1)
    expr = :($(ex1[1]) * $(ex2[1]))

    for i in 2:N
        expr = madd ? :(muladd($(ex1[i]), $(ex2[i]), $expr)) :
                      :($(expr) + $(ex1[i]) * $(ex2[i]))
    end
    return expr
end

function remove_duplicates(ex1in, ex2in)
    ex1out, ex2out = Expr[], Expr[]
    exout = Expr[]
    factors = ones(Int, length(ex1in))

    for (ex1ine, ex2ine) in zip(ex1in, ex2in)
        prod = :($ex1ine * $ex2ine)
        i = findfirst(isequal(prod), exout) # check if this product exist in the output
        if i == nothing # this product does not exist yet
            push!(ex1out, ex1ine)
            push!(ex2out, ex2ine)
            push!(exout, prod)
        else # found a duplicate
            factors[i] += 1
        end
    end
    for i in 1:length(ex1out)
        factors[i] != 1 && (ex1out[i] = :($(factors[i]) * $(ex1out[i])))
    end
    return ex1out, ex2out
end

# check symmetry and return
remove_duplicates(::Type{<:Tensor}, ex) = ex # do nothing if return type is a Tensor
function remove_duplicates(::Type{SymmetricTensor{order, dim}}, ex) where {order, dim}
    ex.args = ex.args[SYMMETRIC_INDICES[order][dim]]
    return ex
end

# return types
# double contraction
function dcontract end
@pure getreturntype(::typeof(dcontract), ::Type{<:FourthOrderTensor{dim}}, ::Type{<:FourthOrderTensor{dim}}) where {dim} = Tensor{4, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SymmetricTensor{4, dim}}, ::Type{<:SymmetricTensor{4, dim}}) where {dim} = SymmetricTensor{4, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:Tensor{4, dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SymmetricTensor{4, dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = SymmetricTensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:Tensor{4, dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:SymmetricTensor{4, dim}}) where {dim} = SymmetricTensor{2, dim}

# otimes
function otimes end
@pure getreturntype(::typeof(otimes), ::Type{<:Tensor{1, dim}}, ::Type{<:Tensor{1, dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(otimes), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = Tensor{4, dim}
@pure getreturntype(::typeof(otimes), ::Type{<:SymmetricTensor{2, dim}}, ::Type{<:SymmetricTensor{2, dim}}) where {dim} = SymmetricTensor{4, dim}


text:
using Plots

function velocityVerlet(ODE, lmbda0, q0, lmbdaf, dlmbda)
	N = floor(Int,(lmbdaf - lmbda0)/dlmbda )#Time steps
	lmbda = zeros(N)
	q = zeros(N, 6)
	lmbda[1] = lmbda0
	q[1,:] = q0
	for i in 1:N-1
		lmbda[i+1] = lmbda[i] + dlmbda
		v_half = q[i,4:6] + 0.5*ODE(q[i,:])[4:6]*dlmbda
		q[i+1,1:3] = q[i,1:3] + v_half*dlmbda
		q[i+1,4:6] = v_half + 0.5*ODE(q[i,:])[4:6]*dlmbda
	end
	return q
end

function geo(q0)
	G = 1.5607939e-13 # [lyr^3 sunMass^-1 yr^-1]
	M = G*1. # BH Mass in geometric units
	sfun = 2*M - q0[2]
	f = zeros(6)
	f[1:3] = q0[4:6]
	f[4] = 2*M*q0[1]*q0[2]/(sfun*q0[2])
	f[5] = M*sfun*q0[4]^2/(q0[2]^3) - M*q0[5]^2/(sfun*q0[2]) - sfun*q0[6]^2
	f[6] = -2*q0[5]*q0[6]/q0[2]
	return f
end

function flat(q0)
	f = zeros(6)
	f[1:3] = q0[4:6]
	f[4] = 0
	f[5] = 0
	f[6] = 0
	return f
end

t0 = 0.
r0 = 500.
ph0 = 0.

vr0 = -sqrt(2)/2
vph0 = sqrt(2)/2
vt0 = sqrt(vr0^2 + vph0^2)
q0 = [t0, r0, ph0, vt0, vr0, vph0]

Q = velocityVerlet(flat,0,q0,10,1)

x = Q[:,2].*cos.(Q[:,3])
y = Q[:,2].*sin.(Q[:,3])

println(Q[:,3])

display(plot(x,y))
readline()


text:
using JSON, PiecewiseDeterministicMarkovProcesses, LinearAlgebra, Random

const p0  = convert(Dict{AbstractString,Float64}, JSON.parsefile("../examples/ml.json")["type II"])
const p1  = ( JSON.parsefile("../examples/ml.json"))
include("morris_lecar_variables.jl")
const p_ml = ml(p0)

function F_ml!(xcdot, xc, xd, parms, t::Float64)
	# vector field used for the continuous variable
	#compute the current, v = xc[1]
	xcdot[1] = xd[2] / p_ml.N * (p_ml.g_Na * (p_ml.v_Na - xc[1])) + xd[4] / p_ml.M * (p_ml.g_K  * (p_ml.v_K  - xc[1]))  + (p_ml.g_L  * (p_ml.v_L  - xc[1])) + p_ml.I_app
	nothing
end

function R_ml!(rate, xc, xd, parms, t, issum::Bool)
	if issum == false
		rate[1] = p_ml.beta_na * exp(4.0 * p_ml.gamma_na * xc[1] + 4.0 * p_ml.k_na) * xd[1]
		rate[2] = p_ml.beta_na * xd[2]
		rate[3] = p_ml.beta_k * exp(p_ml.gamma_k * xc[1] + p_ml.k_k) * xd[3]
		rate[4] = p_ml.beta_k * exp(-p_ml.gamma_k * xc[1]  -p_ml.k_k) * xd[4]
		return 0.
	else
		return (p_ml.beta_na * exp(4.0 * p_ml.gamma_na * xc[1] + 4.0 * p_ml.k_na) * xd[1] +
		p_ml.beta_na * xd[2] +
		p_ml.beta_k * exp( p_ml.gamma_k * xc[1] + p_ml.k_k) * xd[3] +
		p_ml.beta_k * exp(-p_ml.gamma_k * xc[1] - p_ml.k_k) * xd[4])
	end
end

xc0 = vec([p1["v(0)"]])
xd0 = vec([Int(p0["N"]),	#Na closed
	0,			   			#Na opened
	Int(p0["M"]),			#K closed
	0])						#K opened

nu_ml = [[-1 1 0 0];[1 -1 0 1];[0 0 -1 1];[0 0 1 -1]]
parms = vec([0.])

tf = p1["t_end"]
tf = 350.

Random.seed!(123)
println("--> chv_optim - call")
pb = PDMP.PDMPProblem(F_ml!, R_ml!, nu_ml, xc0, xd0, parms, (0.0, tf))
# result =		PDMP.pdmp!(xc0,xd0, F_ml!, R_ml!, nu_ml, parms,0.0,tf,algo=:chv_optim,n_jumps = 6)
# result =  @time PDMP.pdmp!(xc0,xd0, F_ml!, R_ml!, nu_ml, parms,0.0,tf,algo=:chv_optim,n_jumps = 4500) #cpp = 100ms/2200 jumps
res = @time PDMP.solve(pb, CHV(Tsit5()), n_jumps = 2200, save_positions = (false, true))


text:
__precompile__()

module VortexModel
using Reexport

include("Vortex.jl")
include("TimeMarching.jl")

@reexport using .Vortex
@reexport using .TimeMarching

export Vortex

end


text:
# ---
# title: 1043. Partition Array for Maximum Sum
# id: problem1043
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/partition-array-for-maximum-sum/description/>
# hidden: true
# ---
#
# Given an integer array `arr`, you should partition the array into (contiguous)
# subarrays of length at most `k`. After partitioning, each subarray has their
# values changed to become the maximum value of that subarray.
#
# Return _the largest sum of the given array after partitioning._
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [1,15,7,9,2,5,10], k = 3
#     Output: 84
#     Explanation: arr becomes [15,15,15,9,10,10,10]
#
#
# **Example 2:**
#
#
#
#     Input: arr = [1,4,1,5,7,3,6,1,9,9,3], k = 4
#     Output: 83
#
#
# **Example 3:**
#
#
#
#     Input: arr = [1], k = 1
#     Output: 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 500`
#   * `0 <= arr[i] <= 109`
#   * `1 <= k <= arr.length`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
include("SortSparse.jl")
include("KoLesky.jl")
include("Cholesky.jl")
include("CovFuncs.jl")
include("Utils.jl")
using LinearAlgebra
using Distributions
using Distances
using Random
using Match
using JLD
Random.seed!(123)

N = 10000
d = 2
l = 0.5
σ = exp.(randn(N))
ρList = [1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0]
λ = 1.5

# Creating uniform random data
x = rand(d, N)

# Creating output arrays
NNZ = zeros(length(ρList))
KLTrueNaive = zeros(3, length(ρList))
KLNaiveTrue = zeros(3, length(ρList))
KLTrueApp = zeros(3, length(ρList))
KLAppTrue = zeros(3, length(ρList))
KLTrueLarge = zeros(3, length(ρList))
KLLargeTrue = zeros(3, length(ρList))
KLTrueExact = zeros(3, length(ρList))
KLExactTrue = zeros(3, length(ρList))

σ .= max.(σ, eps(Float64))

covfunc12!(r) = matern12!(r,l)
covfunc32!(r) = matern32!(r,l)
covfunc52!(r) = matern52!(r,l)

for (iOrder, maternOrder) in enumerate(["12", "32", "52"])
  # Make sure to avoid NANs
  @match maternOrder begin
    "12" => (covfunc! = covfunc12!)
    "32" => (covfunc! = covfunc32!)
    "52" => (covfunc! = covfunc52!)
  end

  for (iρ, ρ) in enumerate(ρList)

    # reordering points and forming skeletons
    colptr, rowval, P, revP, distances = sortSparseRev(x, ρ, 1)
    xOrd = x[:, P]
    skeletons = construct_skeletons(colptr, rowval, distances, λ)

    # Creating reference solution
    Θ = covfunc!(pairwise(Euclidean(), xOrd, dims=2)) + Diagonal(σ)

    # Creating naive vecchia approximation
    LNaive = assembleL(skeletons, xOrd, covfunc!, σ)
    ΘNaive = inv(Matrix(LNaive)') * inv(Matrix(LNaive))

    # Creating "smart" vecchia approximation
    L = assembleL(skeletons, xOrd, covfunc!)
    U = squareSparse(L) + spdiagm(0 => 1 ./ σ)
    icholU_high_level!(U)
    noisecov = NoiseCov(L, U, σ)
    ΘExact = exactMatrix(noisecov)
    ΘApp = appMatrix(noisecov)

    # Creating smart vecchia approximation with larger stencil for IC0
    ULarge = triu(L * L' + spdiagm(0 => 1 ./ σ))
    icholU_high_level!(ULarge)
    noisecovLarge = NoiseCov(L, ULarge, σ)
    ΘLarge = appMatrix(noisecovLarge)

    KLTrueNaive[iOrder, iρ] = KL(Θ, ΘNaive)
    KLNaiveTrue[iOrder, iρ] = KL(ΘNaive, Θ)

    KLTrueApp[iOrder, iρ] = KL(Θ, ΘApp)
    KLAppTrue[iOrder, iρ] = KL(ΘApp, Θ)

    KLTrueLarge[iOrder, iρ] = KL(Θ, ΘLarge)
    KLLargeTrue[iOrder, iρ] = KL(ΘLarge, Θ)

    KLTrueExact[iOrder, iρ] = KL(Θ, ΘExact)
    KLExactTrue[iOrder, iρ] = KL(ΘExact, Θ)
  end
end

save("./out/jld/KLVaryRho.jld",
     "ρList", ρList,
     "KLTrueNaive", KLTrueNaive,
     "KLNaiveTrue", KLNaiveTrue,
     "KLTrueApp", KLTrueApp,
     "KLAppTrue", KLAppTrue,
     "KLTrueLarge", KLTrueLarge,
     "KLLargeTrue", KLLargeTrue,
     "KLTrueExact", KLTrueExact,
     "KLExactTrue", KLExactTrue)


text:
# @testset "Threshold Poisson" begin
#     for rng in [nothing, Random.MersenneTwister()]
#         baserate = rand()
#         theta = rand()
#         deltav = rand()
#         v = rand()
#         t = 1
#         dt = 1
#         threshold = isnothing(rng) ? Threshold.Poisson(baserate, theta, deltav) : Threshold.Poisson(baserate, theta, deltav; rng = rng)
#         rho = baserate * exp((v - theta) / deltav)
#         if rand(threshold.rng) < rho * dt
#             @test evaluate!(threshold,t,v) == t
#         else
#             @test evaluate!(threshold,t,v) == zero(t)
#         end
#     end
# end

@testset "Threshold Ideal" begin
    v = rand()
    for vth in [v .+ (1-v)*(1 .- rand()), v * rand()] # Setting the (v,1] and [0,v) bound for vth
        t = rand()
        threshold = Threshold.Ideal(vth)
        if v < threshold.vth
            @test evaluate!(threshold, t, v) == zero(t)
        else
            @test evaluate!(threshold, t, v) == t
        end
    end
end


text:
# use functions from simulator_quantum_2d_square_renormalize.jl :
#	 renormalize

function renormalizeZ!{T}(simulator::Classical3dSquareSimulator{T}; whichIsTensorT::Int = 1)
	tensorT1, tensorT2 = getTensorT(simulator, whichIsTensorT)
	newTensorT, newTensorTtilde, normalizationFactor = renormalize(tensorT1, tensorT2, getDimM(simulator); whichIsTensorT = whichIsTensorT)

	setTensorT!(simulator.lattice, newTensorT)
	setTensorT!(simulator.lattice, newTensorTtilde; tilde = true)
	setNormalizationfactor!(simulator, normalizationFactor)
end

function renormalizeX!{T}(simulator::Classical3dSquareSimulator{T}; whichIsTensorT::Int = 1)
	tensorT1, tensorT2 = getTensorT(simulator, whichIsTensorT)
	newTensorT, newTensorTtilde, normalizationFactor = renormalize(permutedims(tensorT1,[3,4,5,6,1,2]), permutedims(tensorT2,[3,4,5,6,1,2]), getDimM(simulator); whichIsTensorT = whichIsTensorT)

	setTensorT!(simulator.lattice, permutedims(newTensorT,[5,6,1,2,3,4]))
	setTensorT!(simulator.lattice, permutedims(newTensorTtilde,[5,6,1,2,3,4]); tilde = true)
	setNormalizationfactor!(simulator, normalizationFactor)
end

function renormalizeY!{T}(simulator::Classical3dSquareSimulator{T}; whichIsTensorT::Int = 1)
	tensorT1, tensorT2 = getTensorT(simulator, whichIsTensorT)
	newTensorT, newTensorTtilde, normalizationFactor = renormalize(permutedims(tensorT1,[5,6,1,2,3,4]), permutedims(tensorT2,[5,6,1,2,3,4]), getDimM(simulator); whichIsTensorT = whichIsTensorT)

	setTensorT!(simulator.lattice, permutedims(newTensorT,[3,4,5,6,1,2]))
	setTensorT!(simulator.lattice, permutedims(newTensorTtilde,[3,4,5,6,1,2]); tilde = true)
	setNormalizationfactor!(simulator, normalizationFactor)
end


text:
@testset "FixedCapacityBinPacking2MILP: $(fct_type), $(n_bins) bin, 2 items, $(T)" for fct_type in ["vector of variables", "vector affine function"], n_bins in [1, 2], T in [Int, Float64]
    mock = MOIU.MockOptimizer(MILPModel{T}())
    model = COIB.FixedCapacityBinPacking2MILP{T}(mock)

    if T == Int
        @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Integer)
    end
    @test MOI.supports_constraint(
        model,
        MOI.ScalarAffineFunction{T},
        MOI.EqualTo{T},
    )
    @test MOIB.supports_bridging_constraint(
        model,
        MOI.VectorOfVariables,
        CP.BinPacking{CP.FIXED_CAPACITY_BINPACKING, T},
    )

    n_items = 2
    weights = T[3, 2]
    capas = T[5, 6][1:n_bins]

    if T == Int
        x_load_1, _ = MOI.add_constrained_variable(model, MOI.Integer())
        if n_bins == 1
            x_load_2 = nothing
        elseif n_bins == 2
            x_load_2, _ = MOI.add_constrained_variable(model, MOI.Integer())
        else
            @assert false
        end
    elseif T == Float64
        x_load_1 = MOI.add_variable(model)
        if n_bins == 1
            x_load_2 = nothing
        elseif n_bins == 2
            x_load_2 = MOI.add_variable(model)
        else
            @assert false
        end
    end
    x_bin_1, _ = MOI.add_constrained_variable(model, MOI.Integer())
    x_bin_2, _ = MOI.add_constrained_variable(model, MOI.Integer())

    fct = if fct_type == "vector of variables"
        if n_bins == 1
            MOI.VectorOfVariables([x_load_1, x_bin_1, x_bin_2])
        elseif n_bins == 2
            MOI.VectorOfVariables([x_load_1, x_load_2, x_bin_1, x_bin_2])
        else
            @assert false
        end
    elseif fct_type == "vector affine function"
        if n_bins == 1
            MOIU.vectorize([x_load_1, x_bin_1, x_bin_2])
        elseif n_bins == 2
            MOIU.vectorize([x_load_1, x_load_2, x_bin_1, x_bin_2])
        else
            @assert false
        end
    else
        @assert false
    end
    c = MOI.add_constraint(model, fct, CP.BinPacking{CP.FIXED_CAPACITY_BINPACKING}(n_bins, n_items, weights, capas))

    @test MOI.is_valid(model, x_load_1)
    if n_bins >= 2
        @test MOI.is_valid(model, x_load_2)
    end
    @test MOI.is_valid(model, x_bin_1)
    @test MOI.is_valid(model, x_bin_2)
    @test MOI.is_valid(model, c)

    bridge = MOIBC.bridges(model)[MOI.ConstraintIndex{MOI.VectorOfVariables, CP.BinPacking{CP.VARIABLE_CAPACITY_BINPACKING, T}}(-1)]

    @testset "Bridge properties" begin
        @test MOIBC.concrete_bridge_type(typeof(bridge), MOI.VectorOfVariables, CP.BinPacking{CP.VARIABLE_CAPACITY_BINPACKING, T}) == typeof(bridge)
        @test MOIB.added_constrained_variable_types(typeof(bridge)) == [(MOI.ZeroOne,)]
        @test MOIB.added_constraint_types(typeof(bridge)) == [
            (MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}),
            (MOI.ScalarAffineFunction{T}, MOI.LessThan{T}),
        ]

        @test MOI.get(bridge, MOI.NumberOfVariables()) == n_bins * n_items
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VariableIndex, MOI.ZeroOne}()) == n_bins * n_items
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}()) == n_bins + 2 * n_items
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, MOI.LessThan{T}}()) == n_bins

        @test Set(MOI.get(bridge, MOI.ListOfVariableIndices())) == Set(bridge.assign_var)
        @test Set(MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VariableIndex, MOI.ZeroOne}())) == Set(bridge.assign_con)
        @test Set(MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}())) == Set([bridge.assign_unique..., bridge.assign_number..., bridge.assign_load...])
        @test Set(MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, MOI.LessThan{T}}())) == Set(bridge.load_capacity)
    end

    @testset "Set of variables: one binary per item and per bin" begin
        @test length(bridge.assign_var) == n_items * n_bins
        for i in 1:(n_items * n_bins)
            @test MOI.is_valid(model, bridge.assign_var[i])
        end

        @test length(bridge.assign_con) == n_items * n_bins
        for i in 1:(n_items * n_bins)
            @test MOI.is_valid(model, bridge.assign_con[i])
            @test MOI.get(model, MOI.ConstraintFunction(), bridge.assign_con[i]) == bridge.assign_var[i]
            @test MOI.get(model, MOI.ConstraintSet(), bridge.assign_con[i]) == MOI.ZeroOne()
        end
    end

    @testset "One bin per item" begin
        @test length(bridge.assign_unique) == n_items
        for item in 1:n_items
            @test MOI.is_valid(model, bridge.assign_unique[item])
            f = MOI.get(model, MOI.ConstraintFunction(), bridge.assign_unique[item])
            @test length(f.terms) == n_bins
            @test MOI.get(model, MOI.ConstraintSet(), bridge.assign_unique[item]) == MOI.EqualTo(one(T))

            for bin in 1:n_bins
                t = f.terms[bin]
                @test t.coefficient === one(T)
                @test t.variable == bridge.assign_var[item, bin]
            end
        end
    end

    @testset "Relation between the integer and binary representation of bin assignment" begin
        @test length(bridge.assign_number) == n_items
        for item in 1:n_items
            @test MOI.is_valid(model, bridge.assign_number[item])
            f = MOI.get(model, MOI.ConstraintFunction(), bridge.assign_number[item])
            @test length(f.terms) == n_bins + 1
            @test MOI.get(model, MOI.ConstraintSet(), bridge.assign_number[item]) == MOI.EqualTo(zero(T))

            t = f.terms[1]
            @test t.coefficient === -one(T)
            @test t.variable == ((item == 1) ? x_bin_1 : x_bin_2)

            for bin in 1:n_bins
                t = f.terms[1 + bin]
                @test t.coefficient === T(bin)
                @test t.variable == bridge.assign_var[item, bin]
            end
        end
    end

    @testset "Load" begin
        @test length(bridge.assign_load) == n_bins
        for bin in 1:n_bins
            @test MOI.is_valid(model, bridge.assign_load[bin])
            f = MOI.get(model, MOI.ConstraintFunction(), bridge.assign_load[bin])
            @test length(f.terms) == n_items + 1
            @test MOI.get(model, MOI.ConstraintSet(), bridge.assign_load[bin]) == MOI.EqualTo(zero(T))

            t = f.terms[1]
            @test t.coefficient === -one(T)
            @test t.variable == ((bin == 1) ? x_load_1 : x_load_2)

            for item in 1:n_items
                t = f.terms[1 + item]
                @test t.coefficient === weights[item]
                @test t.variable == bridge.assign_var[item, bin]
            end
        end
    end

    @testset "Capacity" begin
        @test length(bridge.load_capacity) == n_bins
        for bin in 1:n_bins
            @test MOI.is_valid(model, bridge.load_capacity[bin])
            f = MOI.get(model, MOI.ConstraintFunction(), bridge.load_capacity[bin])
            @test length(f.terms) == 1
            @test MOI.get(model, MOI.ConstraintSet(), bridge.load_capacity[bin]) == MOI.LessThan(capas[bin])

            t = f.terms[1]
            @test t.coefficient === one(T)
            @test t.variable == ((bin == 1) ? x_load_1 : x_load_2)
        end
    end
end


text:
#===============================================================================
# Exercises and examples from the julia lang string documentation:
# http://docs.julialang.org/en/stable/manual/strings/
#
# Author: Carl Cortright
# Date: 12/23/2016
#
===============================================================================#
println("This is a test of Julia's string features... \n")
# Test a single character string
x = 'x'
println(typeof(x))

xToInt = Int(x)
println(typeof(xToInt))

# Try printing a for-all symbol
println('\u2200')

# String Indexing
println("Testing String Indexing:")
println(x[end])
helloWorld = "Hello world!"
println(helloWorld[1])
println(helloWorld[1:5])
println(helloWorld[7:end])

println("Interpolation: $helloWorld with \$100")
# Triple String Quotes + Common methods
println("Testing Triple Quote Strings:")
tripStr = """
This is a triplequoted string \r\n
"""
println(tripStr)
for c in tripStr
  println(c)
end

# Testing some common methods
println("Testing some common methods:")
println("l;skdjfklafjd" < "laksdjfklasdjf")
println("l;skdjfklafjd" == "laksdjfklasdjf")
println("l;skdjfklafjd" > "laksdjfklasdjf")
println(search(tripStr, "i"))
println(contains(tripStr, "blah"))
println(contains(tripStr, "triple"))
println(join(["hello", "world"], " ,", ", and "))
println(repeat("hi ", 100))

# Regex
println("Testing Regex expressions:")
myRegex = r".w"
println(typeof(myRegex))
println(ismatch(myRegex, "hello world"))
println(ismatch(myRegex, "no umbo"))

# Match anything that would be a Julia comment
commentr = r"^\s*(?:#|$)"
testline = "# hello"
m = match(commentr, testline)
if m === nothing
  println("\"$testline\" is not a comment")
else
  println("\"$testline\" is a comment")
  println("Offest is " * string(m.offset))
end

# Playing with byte array literals
barr = b"hello world"
println(barr[1])
println(barr)

# Identifying version numbers
if v"0.2" <= VERSION <= v"0.8"
  println("We have an acceptable Julia version!")
end

println("Our Julia version is " * string(VERSION))

max_issues_repo_name: carlcortright/Playground


text:
"""
    mutable struct Box{T} <: AbstractGeometry{T, 3}

Very simple rectengular box in cartesian coordinates.
"""
struct Box{T} <: AbstractVolumePrimitive{T, 3}
    x::Tuple{T, T}
    y::Tuple{T, T}
    z::Tuple{T, T}
    translate::Union{CartesianVector{T}, Missing}
end

function in(pt::CartesianPoint{T}, g::Box{T})::Bool where {T}
    ismissing(g.translate) ? nothing : pt -= g.translate
    return (g.x[1] <= pt[1] <= g.x[2]) && (g.y[1] <= pt[2] <= g.y[2]) && (g.z[1] <= pt[3] <= g.z[2])
end

@inline in(pt::CylindricalPoint, g::Box)::Bool = in(CartesianPoint(pt), g)

function Box{T}(dict::Union{Dict{Any, Any}, Dict{String, Any}}, inputunit_dict::Dict{String,Unitful.Units})::Box{T} where {T <: SSDFloat}

    if haskey(dict, "translate")
        translate = CartesianVector{T}(
            haskey(dict["translate"],"x") ? geom_round(ustrip(uconvert(u"m", T(dict["translate"]["x"]) * inputunit_dict["length"] ))) : 0.0,
            haskey(dict["translate"],"y") ? geom_round(ustrip(uconvert(u"m", T(dict["translate"]["y"]) * inputunit_dict["length"] ))) : 0.0,
            haskey(dict["translate"],"z") ? geom_round(ustrip(uconvert(u"m", T(dict["translate"]["z"]) * inputunit_dict["length"] ))) : 0.0)
    else
        translate = CartesianVector{T}(0.0, 0.0, 0.0)
    end

    if typeof(dict["x"]) <: SSDFloat
        xStart, xStop = geom_round(translate[1] + T(0.0)), geom_round(translate[1] + ustrip(uconvert(u"m", T(dict["x"]) * inputunit_dict["length"])))
    else
        xStart, xStop = geom_round(translate[1] + ustrip(uconvert(u"m", T(dict["x"]["from"]) * inputunit_dict["length"] ))), geom_round(translate[1] + ustrip(uconvert(u"m", T(dict["x"]["to"]) * inputunit_dict["length"])))
    end
    if typeof(dict["y"]) <: SSDFloat
        yStart, yStop = geom_round(translate[2] + T(0.0)), geom_round(translate[2] + ustrip(uconvert(u"m", T(dict["y"]) * inputunit_dict["length"])))
    else
        yStart, yStop = geom_round(translate[2] + ustrip(uconvert(u"m", T(dict["y"]["from"]) * inputunit_dict["length"] ))), geom_round(translate[2] + ustrip(uconvert(u"m", T(dict["y"]["to"]) * inputunit_dict["length"])))
    end
    if typeof(dict["z"]) <: SSDFloat
        zStart, zStop = geom_round(translate[3] + T(0.0)), geom_round(translate[3] + ustrip(uconvert(u"m", T(dict["z"]) * inputunit_dict["length"])))
    else
        zStart, zStop = geom_round(translate[3] + ustrip(uconvert(u"m", T(dict["z"]["from"]) * inputunit_dict["length"] ))), geom_round(translate[3] + ustrip(uconvert(u"m", T(dict["z"]["to"]) * inputunit_dict["length"])))
    end
    translate = missing
    return Box{T}(
        (xStart, xStop),
        (yStart, yStop),
        (zStart, zStop),
        translate )
end

function Geometry(T::DataType, t::Val{:box}, dict::Dict{Any, Any}, inputunit_dict::Dict{String,Unitful.Units})
    return Box{T}(dict, inputunit_dict)
end

function get_important_points(g::Box{T})::NTuple{3, Vector{T}} where {T <: SSDFloat}
    v1::Vector{T} = T[g.x[1], g.x[2]] #[g.x[1], g.x[2]]
    v2::Vector{T} = T[g.y[1], g.y[2]]
    v3::Vector{T} = T[g.z[1], g.z[2]]
    return v1, v2, v3
end

function get_important_points(g::Box{T}, ::Val{:r})::Vector{T} where {T <: SSDFloat}
    @warn "Not yet implemented"
    return T[]
end
function get_important_points(g::Box{T}, ::Val{:φ})::Vector{T} where {T <: SSDFloat}
    @warn "Not yet implemented"
    return T[]
end
function get_important_points(g::Box{T}, ::Val{:z})::Vector{T} where {T <: SSDFloat}
    return T[g.z[1], g.z[2]]
end
function get_important_points(g::Box{T}, ::Val{:x})::Vector{T} where {T <: SSDFloat}
    return T[g.x[1], g.x[2]]
end
function get_important_points(g::Box{T}, ::Val{:y})::Vector{T} where {T <: SSDFloat}
    return T[g.y[1], g.y[2]]
end

function vertices(cb::Box{T})::Vector{CartesianPoint{T}} where {T <: SSDFloat}
    v::Vector{CartesianPoint{T}} = CartesianPoint{T}[
        CartesianPoint{T}(cb.x[1], cb.y[1], cb.z[1]),
        CartesianPoint{T}(cb.x[2], cb.y[1], cb.z[1]),
        CartesianPoint{T}(cb.x[2], cb.y[2], cb.z[1]),
        CartesianPoint{T}(cb.x[1], cb.y[2], cb.z[1]),
        CartesianPoint{T}(cb.x[1], cb.y[1], cb.z[2]),
        CartesianPoint{T}(cb.x[2], cb.y[1], cb.z[2]),
        CartesianPoint{T}(cb.x[2], cb.y[2], cb.z[2]),
        CartesianPoint{T}(cb.x[1], cb.y[2], cb.z[2]),
    ]
    ismissing(cb.translate) ? nothing : v = map(x -> x + cb.translate, v)
    return v
end

function LineSegments(cb::Box{T})::Vector{LineSegment{T, 3, :cartesian}} where {T}
    v::Vector{CartesianPoint{T}} = vertices(cb)
    return LineSegment{T, 3, :cartesian}[
        LineSegment(v[1], v[2]),
        LineSegment(v[2], v[3]),
        LineSegment(v[3], v[4]),
        LineSegment(v[4], v[1]),
        LineSegment(v[1], v[5]),
        LineSegment(v[2], v[6]),
        LineSegment(v[3], v[7]),
        LineSegment(v[4], v[8]),
        LineSegment(v[5], v[6]),
        LineSegment(v[6], v[7]),
        LineSegment(v[7], v[8]),
        LineSegment(v[8], v[5])
    ]
end

@recipe function f(cb::Box{T}) where {T <: SSDFloat}
    label-->"Box"
    ls = LineSegments(cb)
    @series begin
        ls
    end
end

function sample(cb::Box{T}, stepsize::Vector{T})  where {T <: SSDFloat}
    samples  = CartesianPoint{T}[]
    for x in cb.x[1] : stepsize[1] : cb.x[2]
        for y in cb.y[1] :stepsize[2] : cb.y[2]
            for z in cb.z[1] : stepsize[3] : cb.z[2]
                push!(samples,CartesianPoint{T}(x, y, z))
            end
        end
    end
    ismissing(cb.translate) ? nothing : samples = map(x -> x + cb.translate, samples)
    return samples
end

function (+)(b::Box{T}, translate::Union{CartesianVector{T},Missing})::Box{T} where {T <: SSDFloat}
    if ismissing(translate)
        return b
    elseif ismissing(b.translate)
        return Box(b.x, b.y, b.z, translate)
    else
        return Box(b.x, b.y, b.z, b.translate + translate)
    end
 end


text:
module IntensityMetrics

using Unitful

export Medium, Excitation
export intensity, intensity_sppa, intensity_spta, mechanical_index
export pn_pressure, pp_pressure, ptp_pressure

include("models.jl")
include("metrics.jl")

end


text:
# Thanks to Mason Protter for this benchmark

module SUFib
using SymbolicUtils
using SymbolicUtils.Rewriters

@syms fib(x::Int)::Int

const rset = [
    @rule fib(0) => 0
    @rule fib(1) => 1
    @rule fib(~n) => fib(~n - 1) + fib(~n - 2)
] |> Chain |> Postwalk |> Fixpoint

compute_fib(n) = rset(fib(n))

end

module MTFib

using Metatheory
using Metatheory.EGraphs
@metatheory_init

const fibo = @theory begin
    x::$Int + y::$Int |> x+y
    fib(n::$Int) |> (n < 2 ? n : :(fib($(n-1)) + fib($(n-2))))
end;

using Suppressor

function compute_fib(n)
    @suppress begin # don't print crap
        g = EGraph(:(fib($n)))
        saturate!(g, fibo; timeout=7000)
        extran = addanalysis!(g, ExtractionAnalysis, astsize)
        extract!(g, extran)
    end
end

end

using BenchmarkTools

ns = 1:2:22

SU_ts = map(ns) do n
    @assert SUFib.compute_fib(n) isa Number
    b = @benchmarkable SUFib.compute_fib($n) seconds=0.2
    mean(run(b)).time / 1e9
end

MT_ts = map(ns) do n
    @assert MTFib.compute_fib(n) isa Number
    b = @benchmarkable MTFib.compute_fib($n) seconds=0.2
    mean(run(b)).time / 1e9
end

using Plots
plot(ns, SU_ts, label="SymbolicUtils", title="fib(n)", ylabel="time (s)", xlabel="n")
plot!(ns, MT_ts, label="Metatheory")


text:
## Some defaults if not already defined

DRT = isdefined(Main, :DRT) ? DRT : true
PPE = isdefined(Main, :PPE) ? PPE : true
REF = isdefined(Main, :REF) ? REF : true

##

dcx = 1200 # 1920
dcy = 771  # 1080

global scene
scene, layout = layoutscene(resolution = (dcx, dcy)) #10,

figaxes = [LAxis(scene) for i in 1:1, j in 1:1]
tightlimits!.(figaxes)
layout[1:1, 1:1] = figaxes

img = rotr90(load(ENV["HOME"]*"/Pictures/CharlesDockScale.png"))
# img = rotr90(load(ENV["HOME"]*"/Pictures/charlesLow.png"))
# img = rotr90(MakieGallery.loadasset("cow.png"))

for ax in figaxes
  image!(ax, img)
end

ax = figaxes[1,1]
# figaxes[1, 1].aspect = DataAspect()

# ??
xlims!(ax, 0, 260)
ylims!(ax, 0, 170)
# ax2.aspect = AxisAspect(1180/770)
ax.aspect = AxisAspect(dcx/dcy)
tightlimits!(ax)
# ??

ax.xgridvisible = false
ax.ygridvisible = false
ax.xticklabelsvisible = false
ax.yticklabelsvisible = false
ax.xticksvisible = false
ax.yticksvisible = false
# ax.xspinevisible = false
# ax.yspinevisible = false
# ax.xoppositespinevisible = false
# ax.yoppositespinevisible = false
# ax.xtrimspine = true
# ax.ytrimspine = true

# figaxes[1, 2].title = "AxisAspect(1)"
# figaxes[1, 2].aspect = AxisAspect(1)

scene

## new axis

ax2 = LAxis(scene)

arx = dcx # 1282
ary = dcy # 775
xlims!(ax2, 0, 300)
ylims!(ax2, 0, 200)
ax2.aspect = AxisAspect(arx/ary)
tightlimits!(ax2)

ax2.title = "Charles River Basin, MA"
ax2.xlabel = "East [meters]"
ax2.ylabel = "North [meters]"

scene

# ax2.ytickalign = 100
# ax2.yticklabelalign = :right # doesnt work

layout[1,1] = ax2

##

scl = 5.0
ofsx = 90*scl
ofsy = 145*scl

addLinesBelief!(fg, figaxes[1,1], TTm, scale=scl, origin=(ofsx,ofsy), drt=DRT, ppe=PPE, ref=REF, maskRef=(Second(420),Second(450)) )

scene
##

# resolution=(dcx,dcy)
# pl = plotVariableBeliefs(fg, r"x\d", sortVars=true, fade=2, fadeFloor=0.2, scene=figaxes[1,1], scale=scl, origin=(ofsx,750))
pl, Z = plotVariableBeliefs(fg, r"x\d", colormap=:blues, minColorBase=0.0,
                            sortVars=true, fade=1, extend=0.0,fadeFloor=0.2, scale=scl, origin=(ofsx,ofsy), N=100,
                            xmin=(0-ofsx)/scl, xmax=(dcx-ofsx)/scl,ymin=(0-ofsy)/scl,ymax=(dcy-ofsy)/scl, scene=figaxes[1,1])
#

scene

##


text:
using GitHub
using JLD2

function femtocleaner_prs(api, repo, auth)
    prs, _ = pull_requests(api, repo; params = Dict(
        :state => "all",
        :head => "$(GitHub.name(get(repo.owner))):fbot/deps"
    ), auth=auth)
    prs
end

function collect_statistics(app_id, key_path)
    println("# FemtoCleaner statistics at $(now())")
    jwt = GitHub.JWTAuth(app_id, key_path)
    app_name = get(GitHub.app(; auth=jwt).name)
    insts, _ = installations(jwt)
    api = GitHub.DEFAULT_API
    println("FemtoCleaner is installed in $(length(insts)) accounts:")
    for inst in insts
        println("\t", GitHub.name(get(inst.account)))
    end
    all_prs = Any[]
    all_repos = Any[]
    nrepos = 0
    ntotal_prs = 0
    ntotal_additions = 0
    ntotal_deletions = 0
    for inst in insts
        auth = create_access_token(api, inst, GitHub.JWTAuth(app_id, key_path))
        irepos, _ = repos(api, inst; auth=auth)
        foreach(irepos) do repo
            nrepos += 1
            println("Processing $(GitHub.name(repo))")
            reqs = try
                req_file = GitHub.file(repo, "REQUIRE"; auth=auth)
                Pkg.Reqs.parse(IOBuffer(base64decode(strip(get(req_file.content)))))
            catch e
                @show e
                nothing
            end
            push!(all_repos, (repo, reqs))
            prs = femtocleaner_prs(api, repo, auth)
            prs = map(pr->GitHub.pull_request(repo, pr; auth=auth), prs)
            append!(all_prs, Iterators.product((repo,),prs))
            ntotal_prs += length(prs)
            isempty(prs) && return
            ntotal_additions += sum(pr->get(pr.additions), prs)
            ntotal_deletions += sum(pr->get(pr.deletions), prs)
        end
    end
    println("Opened $ntotal_prs PRs in $nrepos repositories (+$ntotal_additions -$ntotal_deletions)")
    @save "data.jld2" all_repos all_prs
end


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

const HardData = Dict{CartesianIndex{3},Real}
coords(hd::HardData) = keys(hd)


text:
module NumericalContinuation

using DocStringExtensions

using TimerOutputs  # TODO: remove when finished
const to = TimerOutput()

include("signal.jl")
include("view_vector.jl")
include("problem_structure.jl")
include("monitor_function.jl")
include("options.jl")
include("closed_problem.jl")

include("utils.jl")

include("zero_problem.jl")

include("FourierCollocation.jl")

end # module


text:
#=
Common parameters for estimating the probability of for a diffusion with X(0) = x₀
satisfying X(T) ∈ B for the Muller potential.
=#

using LinearAlgebra
using Random
using BasicMD
using ForwardDiff
using Optim
using TestLandscapes: Muller

#  set up potential
V = x-> Muller(x)
# commented out due to threading issue with ForwardDiff
# cfg = ForwardDiff.GradientConfig(V, zeros(Float64,2));
∇V! = (gradV, X) -> ForwardDiff.gradient!(gradV, V, X);

# Identify minima for defining x₀ and target set B
mins=[optimize(V, x) for x in [[-0.5, 1.5],[0., 0.5],[0.5, 0.]]];
x₀ = copy(mins[1].minimizer);   # initial point
xt = copy(mins[3].minimizer);   # point defining target set B

T = 10;     # terminal time
Δt = 1e-4;  # time step
nΔt = Int(T/Δt);
β = 0.1;    # inverse temperature

# define observable as the indicator function on B = {x∣ |x-xt|≤r}
r = 0.5;
f = X-> Int(norm(X.-xt)≤r);

sampler = MALA(V,∇V!, β, Δt);

max_stars_repo_name: gideonsimpson/JuWeightedEnsemble

max_issues_repo_name: gideonsimpson/JuWeightedEnsemble

max_forks_repo_name: gideonsimpson/JuWeightedEnsemble


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""CustomResourceSubresourceStatus defines how to serve the status subresource for CustomResources. Status is represented by the `.status` JSON path inside of a CustomResource. When set, * exposes a /status subresource for the custom resource * PUT requests to the /status subresource take a custom resource object, and ignore changes to anything except the status stanza * PUT/POST/PATCH requests to the custom resource ignore changes to the status stanza

    IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus(;
    )

"""
mutable struct IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus <: SwaggerModel

    function IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus(;)
        o = new()
        o
    end
end # type IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus

const _property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus = Dict{Symbol,Symbol}()
const _property_types_IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus = Dict{Symbol,String}()
Base.propertynames(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus }) = collect(keys(_property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus))
Swagger.property_type(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus[name]))}
Swagger.field_name(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus }, property_name::Symbol) =  _property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus[property_name]

function check_required(o::IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus)
    true
end

function validate_property(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1beta1CustomResourceSubresourceStatus }, name::Symbol, val)
end


text:
using Printf
using LinearAlgebra
using Random

include("../pwgrid_04/PWGrid_v04.jl")
include("../pwgrid_04/read_kpts.jl")
include("../common/gen_lattice_pwscf.jl")
include("../common/ortho_gram_schmidt.jl")
include("../common/wrappers_fft.jl")

include("op_K.jl")
include("op_Vpot.jl")
include("op_H.jl")
include("diag_lobpcg.jl")
include("Kprec.jl")

include("dump_bandstructure.jl")

function test_main()
    LatConst = 5.0
    LatVecs = gen_lattice_fcc( LatConst )
    ecutwfc_Ry = 20.0
    kpts_red = read_kpts("../pwgrid_04/KPATH_FCC_60")

    pw = PWGrid( ecutwfc_Ry*0.5, LatVecs, kpts_red )

    Ns = pw.Ns
    Npoints = prod(Ns) # note that for the current implementation Ns = Ng
    Ngwx = pw.gkvec.Ngwx
    LatVecs = pw.LatVecs
    RecVecs = pw.RecVecs
    Ngw = pw.gkvec.Ngw
    Nkpts = size(kpts_red)[2]
    kpts = pw.gkvec.kpts

    @printf("Unit cell:\n")
    for i = 1:3
        for j = 1:3
            @printf("%10.5f ", LatVecs[i,j])
        end
        @printf("\n")
    end

    @printf("Reciprocal unit cell:\n")
    for i = 1:3
        for j = 1:3
            @printf("%10.5f ", RecVecs[i,j])
        end
        @printf("\n")
    end

    @printf("Number of G-vectors: %d\n", pw.gvec.Ng)
    @printf("Ngwx = %d\n", pw.gkvec.Ngwx)
    @printf("Ns   = (%d,%d,%d)\n", Ns[1], Ns[2], Ns[3])

    @printf("k-point list: (reduced)\n")
    for ik = 1:Nkpts
        @printf("%d %8.5f %8.5f %8.5f\n", ik, kpts_red[1,ik], kpts_red[2,ik], kpts_red[3,ik])
    end

    @printf("k-point list:\n")
    for ik = 1:Nkpts
        @printf("%d %8.5f %8.5f %8.5f\n", ik, kpts[1,ik], kpts[2,ik], kpts[3,ik])
    end

    Vpot = zeros(Npoints)  # no external potential

    Nstates = 4
    srand(1234)

    evals = zeros(Float64, Nstates, Nkpts)  # what is the optimal shape?

    for ik = 1:Nkpts

        Ngwk = Ngw[ik]
        psi  = rand(ComplexF64,Ngwk,Nstates)
        psi = ortho_gram_schmidt(psi)

        evals[:,ik], psi = diag_lobpcg( pw, Vpot, psi, ik, verbose=true, tol_avg=1e-7 )
    end

    dump_bandstructure( evals, kpts, filename="band_free_fcc.dat" )

end

test_main()


text:
export iauPn
"""
Convert a p-vector into modulus and unit vector.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  vector/matrix support function.

Given:
   p        double[3]      p-vector

Returned:
   r        double         modulus
   u        double[3]      unit vector

Notes:

   1. If p is null, the result is null.  Otherwise the result is a unit
   vector.

   2. It is permissible to re-use the same array for any of the
   arguments.

Called:
   iauPm        modulus of p-vector
   iauZp        zero p-vector
   iauSxp       multiply p-vector by scalar

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauPn(double p[3], double *r, double u[3])

function iauPn(p::AbstractVector{<:Real})
   # Allocate return values
   ref_r = Ref{Float64}(0.0)
   u = zeros(Float64, 3)

   ccall((:iauPn, libsofa_c), Cvoid,
        (Ptr{Cdouble}, Ref{Cdouble}, Ptr{Cdouble}),
         convert(Array{Float64, 1}, p), ref_r, u)

   return ref_r[], SVector{3}(u)
end


text:
const SimpleList = Vector{Any}
const SimpleDict = Dict{String, Any}

function getatkey(options::SimpleList, key)
    idx = findfirst(isequal(key)∘getkey, options)
    return isnothing(idx) ? nothing : getvalue(options[idx])
end

getkey(s::SimpleDict) = s["key"]
getvalue(s::SimpleDict) = s["value"]

# This simple table type is the preferred way to store tables
struct SimpleTable
    names::Vector{Symbol}
    columns::Vector{AbstractVector}
end

function SimpleTable(data)
    cols = Tables.columns(data)
    names = collect(Symbol, Tables.columnnames(cols))
    columns = AbstractVector[Tables.getcolumn(cols, name) for name in names]
    return SimpleTable(names, columns)
end

function SimpleTable(ps::Pair...)
    names = Symbol[first(p) for p in ps]
    columns = AbstractVector[last(p) for p in ps]
    return SimpleTable(names, columns)
end

Base.copy(s::SimpleTable) = SimpleTable(copy(s.names), copy(s.columns))

function mapcols!(f, s::SimpleTable)
    map!(f, s.columns, s.columns)
    return s
end

function mergedisjointcols!(s::SimpleTable, t::SimpleTable)
    sharedkeys = intersect(s.names, t.names)
    if isempty(sharedkeys)
        append!(s.names, t.names)
        append!(s.columns, t.columns)
        return s
    else
        msg = "Overwriting table is not allowed, " *
            "the following column names are repeated: " *
            join(sharedkeys, ", ")
        throw(ArgumentError(msg))
    end
end

Tables.istable(::SimpleTable) = true
Tables.schema(s::SimpleTable) = Tables.Schema(copy(s.names), eltype.(s.columns))

Tables.columnaccess(::SimpleTable) = true
Tables.columns(s::SimpleTable) = s

Tables.columnnames(s::SimpleTable) = copy(s.names)

function suggestnames(s::SimpleTable, name::Symbol)
    options = fuzzymatch(s.names, name)
    nearestnames = join(options, ", ")
    msg = "There isn't a variable called '$name' in your data; the nearest names appear to be: $nearestnames"
    throw(ArgumentError(msg))
end

# Port `DataFrame` behavior of suggesting alternative column names
function Tables.getcolumn(s::SimpleTable, name::Symbol)
    idx = findfirst(==(name), s.names)
    # Give descriptive error if the name is not found
    return isnothing(idx) ? suggestnames(s, name) : s.columns[idx]
end


text:
using ThreeBodyProblem
using DifferentialEquations
using Plots
using LinearAlgebra

# Our goal today is to compute a halo orbit in the Sun Earth system

## The first step is to define the system. ThreeBodyProblem.jl has built-in functions for common systems like Sun/Jupiter.
sys = sun_earth()
sys.name

## If you want to create a custom system, then you would use the following syntax
sys = System(ThreeBodyProblem.SUN, ThreeBodyProblem.EARTH)

## Now, we are going to compute our first guess at a Halo Orbit using Richardson's expansion
### The details are complicated, but if you're curious check out this paper
#### D. L. Richardson, “Analytic Construction Of Periodic Orbits About The Collinear Points,” Celest. Mech., vol. 22, no. 3, pp. 241–253, 1980, doi: 10.1007/BF01229511.
Az = 0.001
Lpt = 2 # Which libration point will I center about
NS = 1
npts = 100
t, rvs, T, Ax = rich3(sys, Az, Lpt, NS, npts)

## Let's plot the orbit to see what it looks like
plot(rvs[:,1],rvs[:,2],rvs[:,3], label="Richardson")

## Remember, this is a third order approximation
## What will happen if we actually integrate this trajectory?
tspan = (0., T)
prob = ODEProblem(CR3BPdynamics!,rvs[1,:],tspan,sys)
sol = solve(prob, reltol=1e-12)
plot!(sol,vars=(1,2,3),label="Actual",linecolor=:red)

## We can see that this orbit diverges before completing a period
## We need to use a differential corrector to hone in on the true periodic orbit
rv₀, ttf = differential_corrector(sys, rvs[1,:], myconst=3, tf=T)

Φ₀ = I(6)

# event function
condition(u, t, integrator) = u[2]
affect!(integrator) = terminate!(integrator)
cb = DifferentialEquations.ContinuousCallback(condition, affect!)
# sol = solve(prob, Tsit5(), calback=cb)

w₀ = vcat(rv₀, reshape(Φ₀,36,1))

prob = ODEProblem(CR3BPstm!,w₀,(0.0,20.0),sys)
tol = 1e-12
sol = solve(prob, reltol=tol, calback=cb)
plot(sol,vars=(1,2),xlims=[0.98,1.02],ylims=[-0.02,0.02])

w = sol[end]
rv = w[1:6]
Φ = reshape(w[7:42],6,6)
global T = 2*sol.t[end]

## Declare state vectors for L1 and L2 (with zero velocity)
L2 = computeL2(sys.μ,tol=1e-15)
rv2 = [L2; zeros(3)] # state at Lagrange point L2

## Declare state vectors for L1 and L2 (with zero velocity)
L1 = computeL1(sys.μ,tol=1e-15)
rv1 = [L1; zeros(3)] # state at Lagrange point L1

## The Lagrange points are equilibrium points in our dynamics. This means that an object placed there perfectly will stay there forever. But, Lagrange points L1, L2, and L3 are unstable, meaning if you perturb the object slightly, it may fall away. This is how we will generate the invariant manifolds.
## We first need to linearize the dynamics about these Lagrange points to determine which directions are stable and unstable

Φ₀ = I(6)
w₀ = [reshape(Φ₀,36,1);rv1]
tspan = (0.,1.)
prob = ODEProblem(CR3BPstm!,w₀,tspan,sys)
sol = solve(prob, reltol=1e-6)
Φₜ = Matrix(reshape(sol.u[end][1:36],6,6))
rvₜ = sol.u[end][37:42] # last time step should match very closely with rv₁
D,V = eigen(Φₜ,sortby=isreal)
Yw = real(V[:,findall(isreal, D)])
D = D[findall(isreal, D)]
Yws = Yw[:,findmin(real(D))[2]]
Ywu = Yw[:,findmax(real(D))[2]]

## We perturb our states in the stable and unstable directions
α = 1e-6

# L1 points
rv1up = rv1 + α*Ywu # unstable manifold + side
rv1un = rv1 - α*Ywu # unstable manifold - side
rv1sp = rv1 + α*Yws # stable manifold + side
rv1sn = rv1 - α*Yws # stable manifold - side

## Next, we integrate our initial conditions forward in time for the +x and -x directions and backward in time for the +y and -y directions
# Integrate trajectories for 2 Jupiter periods
tf = 2*2π
tspan_forward = (0.,tf) # make sure to add decimal points so tspan contains floating point values, not integers
tspan_back = (0.,-tf)
myreltol = 1e-12

# Set up the ODE Problems
prob1up = ODEProblem(CR3BPdynamics!,rv1up,tspan_forward,sys) # Unstable positive
prob1un = ODEProblem(CR3BPdynamics!,rv1un,tspan_forward,sys) # Unstable negative
prob1sp = ODEProblem(CR3BPdynamics!,rv1sp,tspan_back,sys) # Stable positive
prob1sn = ODEProblem(CR3BPdynamics!,rv1sn,tspan_back,sys) # Stable negative

# Solutions to the ODEs
sol1up = solve(prob1up, reltol=myreltol)
sol1un = solve(prob1un, reltol=myreltol)
sol1sp = solve(prob1sp, reltol=myreltol)
sol1sn = solve(prob1sn, reltol=myreltol)

plot(sys,scaled=true)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1,1],xlims=[-0.6,1.3],legend=:topleft,flip=false)

## We can see two strands moving in towards the sun, but the rest is a jumbled mess. Let's zoom in to see what's going on. Also, I need to make sure things are scaled correctly, or Jupiter will look way too big.
plot(sys)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.04,0.04],xlims=[0.93,1.06],legend=:outerright,flip=false)

# L2 points
rv2up = rv2 + α*Ywu # unstable manifold + side
rv2un = rv2 - α*Ywu # unstable manifold - side
rv2sp = rv2 + α*Yws # stable manifold + side
rv2sn = rv2 - α*Yws # stable manifold - side

## Now we look at the L2 point manifolds
prob2up = ODEProblem(CR3BPdynamics!,rv2up,tspan_forward,sys) # Unstable positive
prob2un = ODEProblem(CR3BPdynamics!,rv2un,tspan_forward,sys) # Unstable negative
prob2sp = ODEProblem(CR3BPdynamics!,rv2sp,tspan_back,sys) # Stable positive
prob2sn = ODEProblem(CR3BPdynamics!,rv2sn,tspan_back,sys) # Stable negative

sol2up = solve(prob2up, reltol=1e-6)
sol2un = solve(prob2un, reltol=1e-6)
sol2sp = solve(prob2sp, reltol=1e-6)
sol2sn = solve(prob2sn, reltol=1e-6)

plot(sys, scaled=true)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1.0,1.0],xlims=[-0.3,1.5],legend=:topright,flip=false)

## This time there are two strands moving out away from the Sun, along with another jumble around Jupiter. Let's zoom in again to take a look.
plot(sys)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.01,0.01],xlims=[0.975,1.025],legend=:outerright,flip=false)


text:
using Test,EasyHamiltonian,ITensors
using EasyHamiltonian:KitaevFermionsModel,tJModel

@testset "KitaevFermionsModel" begin
    model = KitaevFermionsModel(1.0,1.0,1.0)
    hamiltonian = Hamiltonian(model)
    @test isa(hamiltonian,Hamiltonian)
    @test isa(Matrix(hamiltonian),Matrix)
    @test isa(MPO(hamiltonian),MPO)
end

@testset "tJModel" begin
    model = tJModel(1.0,1.0)
    hamiltonian = Hamiltonian(model)
    @test isa(hamiltonian,Hamiltonian)
    @test isa(Matrix(hamiltonian),Matrix)
    @test isa(MPO(hamiltonian),MPO)
end


text:
# Run the package tests if we ask for it
if lowercase(get(ENV, "BINARYBUILDER_PACKAGE_TESTS", "false") ) == "true"
    cd("package_tests") do
        include(joinpath(pwd(), "runtests.jl"))
    end
end


text:
function makeMove(id::String, move::Move)
    @debug("stringifying move $move")
    moveString = tostring(move)
    @debug(moveString)
    moveUrl = baseUrl * "/bot/game/$id/move/$moveString"
    @debug(moveUrl)
    @debug("starting post request to move api")
    HTTP.request("POST", moveUrl, defaultHeader)
    @debug("made move")
end

function makeMove!(id::String, b::Board; variant = "standard", time=nothing)
    @debug "arrived in makeMove!"
    println(variant)
    time = isnothing(time) ? Second(1) : time
    move = Move(0)
    if variant=="threecheck"
        move = iterativeDeepening(Game(b), time, Ref(3); evalFunc=threeCheckPst, uci=true) #chooseMove(b; variant = variant)
    else
        move = iterativeDeepening(Game(b), time, Ref(3); evalFunc=pstAndMate, uci=true) #chooseMove(b; variant = variant)
    end
    #move = rand(moves(b))
    @debug("chose move $move")
    !isempty(id) && makeMove(id, move)
    domove!(b, move)
    @info("made! move")
    return tostring(move)
end

function streamGame(id::String)
    streamUrl = baseUrl * "/bot/game/stream/$id"
    variant = nothing
    myColor = nothing
    lastMove = nothing
    b = startboard()
    domove(b, "b2b3")
    @time JSON.parse("{}")
    @info("streaming Game with id $id")
    r = HTTP.open(gameCallback, "GET", streamUrl, defaultHeader) #=do http
        while !eof(http)
            s = String(readavailable(http))
            length(s) > 1 && @debug(length(s))
            if !isnothing(findfirst("{", s))
                @show b
                @debug("now parsing in streamGame")
                state = @time JSON.parse(s)
                if state["type"] == "gameFull"
                    @debug(s)
                    @debug("full game")
                    if state["white"]["id"] == myId
                        myColor = WHITE
                    elseif state["black"]["id"] == myId
                        myColor = BLACK
                    end
                    @show myColor
                    variant = state["variant"]["key"]
                    if state["initialFen"] != "startpos"
                        b = fromfen(state["initialFen"])
                        @debug("start is non default")
                    else
                        @debug("start is default")
                    end
                    if !isempty(state["state"]["moves"])
                        @debug("splitting and getting moves")
                        moves = string.(split(state["state"]["moves"]))
                        lastMove = moves[end]
                        @debug("applying moves to b")
                        @debug(b)
                        @debug(moves)
                        domoves!(b, moves...)
                        @debug("did all moves")
                    else
                        @debug("no moves to make")
                    end
                    if sidetomove(b) == myColor
                        @debug("starting to make move, still at streamGame")
                        lastMove = makeMove!(id, b, variant = variant)
                    end
                    @show b
                    @debug("reached end of gameFull evaluation")
                elseif state["type"] == "gameState"
                    @debug(s)
                    @debug("game status: $(state["status"])")
                    moves = string.(split(state["moves"]))
                    @debug("already $(length(moves)) moves done")
                    pos = findlast(x -> x == lastMove, moves)
                    @debug("checking if last move already done")
                    if moves[end] != lastMove
                        @debug("last move not yet done")
                        pos = isnothing(pos) ? 0 : pos
                        @show pos
                        @show lastMove
                        @debug(moves[max(pos, 1):end])
                        if pos + 1 == lastindex(moves)
                            domove!(b, moves[end])
                        elseif !isempty(moves)
                            @debug("moves not empty, applying...")
                            domoves!(b, moves[pos+1:end]...)
                        end
                        lastMove = moves[end]
                    end
                    if sidetomove(b) == myColor
                        @debug("starting to make move, still at streamGame")
                        lastMove = @time makeMove!(id, b, variant = variant)
                        @debug("to make move")
                    end
                    @show b
                    @debug("reached end of gameState evaluation")
                end
            end
        end

    end=#
    @debug("Game is over. exiting function streamGame of gameId $id")
end

function gameCallback(http::IO)
    id = ""
    variant = nothing
    myColor = BLACK
    lastMove = nothing
    b = startboard()
    while !eof(http)
        s = String(readavailable(http))
        length(s) > 1 && @debug(length(s))
        if !isnothing(findfirst("{", s))
            @debug b
            @debug("now parsing in streamGame")
            state = @time JSON.parse(s)
            println(state)
            if state["type"] == "gameFull"
                @debug(s)
                @debug("full game")
                id = state["id"]
                if state["white"]["id"] == myId
                    myColor = WHITE
                elseif state["black"]["id"] == myId
                    myColor = BLACK
                end
                @show myColor
                variant = state["variant"]["key"]
                if state["initialFen"] != "startpos"
                    b = fromfen(state["initialFen"])
                    @debug("start is non default")
                else
                    @debug("start is default")
                end
                if !isempty(state["state"]["moves"])
                    @debug("splitting and getting moves")
                    moves = string.(split(state["state"]["moves"]))
                    lastMove = moves[end]
                    @debug("applying moves to b")
                    @debug(b)
                    @debug(moves)
                    domoves!(b, moves...)
                    @debug(b)
                    @debug("did all moves")
                else
                    @debug("no moves to make")
                end
                if sidetomove(b) == myColor
                    @debug("starting to make move, still at streamGame")
                    @debug id
                    @debug b
                    @debug variant
                    lastMove = makeMove!(id, b, variant = variant)
                end
                @debug b
                @debug("reached end of gameFull evaluation")
            elseif state["type"] == "gameState"
                @debug(s)
                @debug("game status: $(state["status"])")
                moves = string.(split(state["moves"]))
                @debug("already $(length(moves)) moves done")
                pos = findlast(x -> x == lastMove, moves)
                @debug("checking if last move already done")
                if moves[end] != lastMove
                    @debug("last move not yet done")
                    pos = isnothing(pos) ? 0 : pos
                    @debug pos
                    @debug lastMove
                    @debug(moves[max(pos, 1):end])
                    if pos + 1 == lastindex(moves)
                        domove!(b, moves[end])
                    elseif !isempty(moves)
                        @debug("moves not empty, applying...")
                        domoves!(b, moves[pos+1:end]...)
                    end
                    lastMove = moves[end]
                end
                if sidetomove(b) == myColor
                    @debug("starting to make move, still at streamGame")
                    colorstring = myColor == WHITE ? "w" : "b"
                    leftTime = Millisecond(state[colorstring*"time"])
                    inc = Millisecond(state[colorstring*"inc"])
                    sumTime = leftTime ÷ 40 + inc
                    availableTime = sumTime > leftTime ? (leftTime * 5)÷10 : sumTime
                    println(availableTime)
                    lastMove = @time makeMove!(id, b; variant = variant, time=availableTime)
                    @debug("to make move")
                end
                @show b
                @debug("reached end of gameState evaluation")
            end
        end
    end

end

function printCB(http::IO)
    while !eof(http)
        s = String(readavailable(http))
        println(s)
    end
end

function runBot()
    eventUrl = baseUrl * "/stream/event"
    challengeUrl = baseUrl * "/challenge"
    retVal = 0
    #trigger precompilation
    @time JSON.parse("{}")
    @info("opening event stream")
    while true
        r = HTTP.open(eventsCallback, "GET", eventUrl, defaultHeader) #=do http
            while !eof(http)
                s = String(readavailable(http))
                length(s) == 1 && continue #print("\rbot still alive")
                if !isnothing(findfirst("{", s))
                    @debug("now parsing in runBot")
                    state = @time JSON.parse(s)
                    @debug(typeof(state))
                    @debug(state["type"])
                    if state["type"] == "challenge"
                        challenger = state["challenge"]["challenger"]["name"]
                        variant =
                            lowercase(state["challenge"]["variant"]["key"])
                        challengeId = state["challenge"]["id"]
                        @debug("""incoming $variant challenge
                                    from $challenger
                                    with id $challengeId""")
                        if lowercase(challenger) == lowercase(myId)
                            continue
                        end
                        if variant in [
                            "racingkings",
                            "horde",
                            "kingofthehill",
                            "antichess",
                            "crazyhouse",
                        ]
                            @debug("declaing challenge with variant $variant")
                            specificUrl =
                                challengeUrl * "/$challengeId/decline"
                            HTTP.request("POST", specificUrl, defaultHeader)
                        else
                            @debug("now accepting challenge")
                            specificUrl =
                                challengeUrl * "/$challengeId/accept"
                            HTTP.request("POST", specificUrl, defaultHeader)
                        end
                    elseif state["type"] == "gameStart"
                        @debug("now async starting stream")
                        @async streamGame(state["game"]["id"])
                    end
                    #write("state$(length(s)).bson", state)
                    #@debug("written to file")
                end
                #@debug(String(readavailable(http)))
            end
        end=#
        # sleep before trying to reconnect
        sleep(2)
    end
    return retVal
end

function eventsCallback(http::IO)
    eventUrl = baseUrl * "/stream/event"
    challengeUrl = baseUrl * "/challenge"
    @info "eventsCallback startetd"
    gameCounter = 0
    while !eof(http)
        s = String(readavailable(http))
        length(s) == 1 && continue #print("\rbot still alive")
        if !isnothing(findfirst("{", s))
            @debug("now parsing in runBot")
            state = @time JSON.parse(s)
            @debug(typeof(state))
            @debug(state["type"])
            if state["type"] == "challenge"
                challenger = state["challenge"]["challenger"]["name"]
                variant =
                    lowercase(state["challenge"]["variant"]["key"])
                challengeId = state["challenge"]["id"]
                @info("""incoming $variant challenge
                            from $challenger
                            with id $challengeId""")
                if lowercase(challenger) == lowercase(myId)
                    continue
                end
                if !(variant in ["standard", "threecheck"]) #=in [
                    "racingkings",
                    "horde",
                    "kingofthehill",
                    "antichess",
                    "crazyhouse",
                ]=#
                    @debug("declaing challenge with variant $variant")
                    specificUrl =
                        challengeUrl * "/$challengeId/decline"
                    HTTP.request("POST", specificUrl, defaultHeader)
                else
                    @debug("now accepting challenge")
                    specificUrl =
                        challengeUrl * "/$challengeId/accept"
                    HTTP.request("POST", specificUrl, defaultHeader)
                end
            elseif state["type"] == "gameStart"
                @debug("now async starting stream")
                @async streamGame(state["game"]["id"])
                gameCounter += 1
            end
            #write("state$(length(s)).bson", state)
            #@debug("written to file")
        end
        #@debug(String(readavailable(http)))
    end
    gameCounter
end

#=
r = HTTP.open("GET", "https://lichess.org/api/stream/event", ["Authorization"=>"Bearer $tok"]) do http
    while !eof(http)
        @debug(String(readavailable(http)))
    end
end
=#