text:
# Support output of TM-align: https://zhanggroup.org/TM-align/
# Zhang, Yang, and Jeffrey Skolnick. "TM-align: a protein structure alignment algorithm based on the TM-score."
# Nucleic acids research 33.7 (2005): 2302-2309.

struct MapAlign
    s2a::Vector{Int}    # sequence index => alignment index
    a2s::Vector{Int}    # alignment index => sequence index
end

struct StructAlign
    m1::MapAlign
    m2::MapAlign

    function StructAlign(m1::MapAlign, m2::MapAlign)
        axes(m1.a2s) == axes(m2.a2s) || error("the two maps disagree")
        return new(m1, m2)
    end
end

function parse_tma(filename)
    lines = readlines(filename)
    filter!(!isempty, lines)
    length(lines) == 3 || error("expecting 3 lines, got $(length(lines))")
    seq1 = [Residue(c) for c in lines[1]]
    seq2 = [Residue(c) for c in lines[3]]
    quality = [c for c in lines[2]]
    append!(quality, fill('-', length(seq1) - length(quality)))
    return seq1, seq2, quality
end

const alignsentinel = -1
function MapAlign(s::AbstractVector{PDBResidue}, a::AbstractVector{Residue}, quality)
    s2a, a2s = fill(alignsentinel, length(s)), Int[]  # we only retain the high-quality matches
    j = 0
    for (r, q) in zip(a, quality)
        r == Residue('-') && continue
        j += 1
        r == three2residue(s[j].id.name) || error("at position $j, residue was $(s[j].id.name) but alignment was $r")
        q == ':' || continue
        push!(a2s, j)
        s2a[j] = length(a2s)
    end
    return MapAlign(s2a, a2s)
end

function StructAlign(struct1::AbstractVector{PDBResidue}, struct2::AbstractVector{PDBResidue},
                     align1::AbstractVector{Residue},     align2::AbstractVector{Residue},
                     quality)
    StructAlign(MapAlign(struct1, align1, quality), MapAlign(struct2, align2, quality))
end

"""
    StructAlign(struct1::AbstractVector{PDBResidue}, struct2::AbstractVector{PDBResidue}, filename::AbstractString)

Create a structure-based alignment between `struct1` and `struct2`. `filename` is the name of the TM-align "results" file
(e.g., https://zhanggroup.org//TM-align/example/873772.html).

See also [`residueindex`](@ref).
"""
StructAlign(struct1::AbstractVector{PDBResidue}, struct2::AbstractVector{PDBResidue}, filename::AbstractString) =
    StructAlign(struct1, struct2, parse_tma(filename)...)

"""
    residueindex(sa::StructAlign, nothing, idx2)
    residueindex(sa::StructAlign, nothing, idx2, ±1)

Calculate the residue index in structure 1 corresponding to residue `idx2` in structure 2.
The second form allows you to find the nearest corresponding residue in the forward (+1) or reverse (-1)
directions, if `idx2` is not among the mapped residues.
"""
residueindex(sa::StructAlign, ::Nothing, idx2::Integer) = sa.m1.a2s[sa.m2.s2a[idx2]]
function residueindex(sa::StructAlign, ::Nothing, idx2::Integer, dir)
    dir ∈ (-1, 1) || error("dir must be ±1")
    i2 = sa.m2.s2a[idx2]
    while i2 == alignsentinel
        idx2 += dir
        i2 = sa.m2.s2a[idx2]
    end
     return sa.m1.a2s[i2]
end
ismapped(sa::StructAlign, ::Nothing, idx2::Integer) = sa.m2.s2a[idx2] != alignsentinel

"""
    residueindex(sa::StructAlign, idx1, nothing)
    residueindex(sa::StructAlign, idx1, nothing, ±1)

Calculate the residue index in structure 2 corresponding to residue `idx1` in structure 1.
The second form allows you to find the nearest corresponding residue in the forward (+1) or reverse (-1)
directions, if `idx1` is not among the mapped residues.
"""
residueindex(sa::StructAlign, idx1::Integer, ::Nothing) = sa.m2.a2s[sa.m1.s2a[idx1]]
function residueindex(sa::StructAlign, idx1::Integer, ::Nothing, dir)
    dir ∈ (-1, 1) || error("dir must be ±1")
    i1 = sa.m1.s2a[idx1]
    while i1 == alignsentinel
        idx1 += dir
        i1 = sa.m1.s2a[idx1]
    end
     return sa.m2.a2s[i1]
end
ismapped(sa::StructAlign, idx1::Integer, ::Nothing) = sa.m1.s2a[idx1] != alignsentinel


text:
@coeff function K_omega(reactor::Reactor)
  34.38 * tau_P / ( epsilon^2 * kappa * G )
end

ext: jl


text:
@testset "lmm" begin
    rng = MersenneTwister(123)
    FDM = FiniteDifferences.central_fdm(5, 1)
    N = 10
    in_dim = 3
    out_dim = 6
    x1 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
    x2 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
    H = rand(4, 6)

    k = LinearMixingModelKernel(
        [Matern32Kernel(), SqExponentialKernel(), FBMKernel(), Matern32Kernel()], H
    )
    @test k isa LinearMixingModelKernel
    @test k isa MOKernel
    @test k isa Kernel
    @test k(x1[1], x2[1]) isa Real

    @test string(k) == "Linear Mixing Model Multi-Output Kernel"
    @test repr("text/plain", k) == (
        "Linear Mixing Model Multi-Output Kernel. Kernels:\n" *
        "\tMatern 3/2 Kernel (metric = Euclidean(0.0))\n" *
        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
        "\tFractional Brownian Motion Kernel (h = 0.5)\n" *
        "\tMatern 3/2 Kernel (metric = Euclidean(0.0))"
    )

    k = LinearMixingModelKernel(SEKernel(), H)

    @test k isa LinearMixingModelKernel
    @test k isa MOKernel
    @test k isa Kernel
    @test length(k.K) == 4
    for kernel in k.K
        @test isa(kernel, SEKernel)
    end

    @test string(k) == "Linear Mixing Model Multi-Output Kernel"
    @test repr("text/plain", k) == (
        "Linear Mixing Model Multi-Output Kernel. Kernels:\n" *
        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
        "\tSquared Exponential Kernel (metric = Euclidean(0.0))"
    )

    test_ADs(k)
end


text:
"""
`mush(A,B)` takes two matrices of the same size and creates a new
matrix whose `i,j`-entry is `(A[i,j],B[i,j])`.
"""
function mush(A::Array{T,2}, B::Array{S,2}) where {S,T}
    TS = Tuple{S,T}
    r,c = size(A)
    if size(A) != size(B)
        error("Matrices must have same shape")
    end

    C = Array{TS,2}(undef,r,c)
    for i=1:r
        for j=1:c
            C[i,j] = (A[i,j],B[i,j])
        end
    end
    return C
end


text:
function likelihood(vm::VectorModel, doc::AbstractArray{Tw},
		window_length::Int, min_prob::Float64=1e-5)

	N = length(doc)
	if N == 1 return (0., 0) end

	z = zeros(T(vm))

	m = MeanCounter(Float64)

	for i in 1:N
		x = doc[i]

		window = window_length
		z[:] .= 0.

		expected_pi!(z, vm, x)

		for j in max(1, i - window):min(N, i + window)
			if i == j continue end

			y = doc[j]

			local_ll = Kahan(Float64)
			for s in 1:T(vm)
				if z[s] < min_prob continue end
				In = view(vm, x, s)

				add!(local_ll, z[s] * exp(Float64(log_skip_gram(vm, x, s, y))))
			end
			add!(m, log(sum(local_ll)))
		end
	end
	return mean(m), m.n
end

function likelihood(vm::VectorModel, dict::Dictionary, f::IO,
		window_length::Int; batch::Int=16777216)
	buffer = zeros(Int32, batch)
	j = 0
	ll = 0.
	while !eof(f)
		doc = read_words(f, dict, buffer, length(buffer), -1)
		if length(doc) == 0 break end
		#println(j)
		local_ll, n = likelihood(vm, doc, window_length)
		ll += local_ll
		j += n
	end
	return ll / j
end

function parallel_likelihood(vm::VectorModel, dict::Dictionary,
	path::AbstractString, window_length::Int, min_prob::Float64=1e-5;
	batch::Int=16777216, log::Union{AbstractString, Nothing}=nothing)
	nbytes = filesize(path)

	words_read = shared_zeros(Int64, (1,))
	stats = Array{Tuple{Float64, Int64}, 1}()

	function do_work(id::Int)
		file = open(path)

		bytes_per_worker = convert(Int, floor(nbytes / nworkers()))

		start_pos = bytes_per_worker * (id-1)
		end_pos = start_pos+bytes_per_worker

		seek(file, start_pos)
		align(file)
		buffer = zeros(Int32, batch)
		local_stats = Array{Tuple{Float64, Int64}, 1}()
		while true
			doc = read_words(file, dict, buffer, batch, end_pos)

			println("$(length(doc)) words read, $(position(file))/$end_pos")
			if length(doc) == 0 break end

			push!(local_stats, likelihood(vm, doc, window_length, min_prob))
		end

		close(file)
		return local_stats
	end

	refs = Array{Future, 1}(undef, nworkers())
	for i in 1:nworkers()
		refs[i] = remotecall(do_work, i+1, i)
	end

	for i in 1:nworkers()
		append!(stats, fetch(refs[i]))
	end

	log_file = if log == nothing stdout else open(log, "w") end
	total_n = 0
	total_mean = Kahan(Float64)
	for (ll, n) in stats
		total_n += n
		println(log_file, ll)
		d = ll - sum(total_mean)
		add!(total_mean, n * d / total_n)
	end
	if log_file != stdout close(log_file) end

	return sum(total_mean)
end


text:
using MartaCT

const _pre_types = (:Float32, :Float64)
const width = 5
const nϕ = 5

for T ∈ _pre_types, Img ∈ TestImages._gs_images, G ∈ Geometry._geometry_names
    @info "Executing... " T Img G
    @eval begin
        gs = $Img($T; width)
        FBPScanner(gs)
        g = $G(gs; nϕ)
        ct = FBPScanner(g, gs) |> project_image |> reconstruct_image
        ctc = calibrate_tomogram(ct, gs)
    end
end


text:
module GageStructs

mutable struct BoardInfo
	Size::UInt32   					  #Total size, in bytes, of the structure.
	BoardIndex::UInt32 				  # Board Index in the system
	BoardType::UInt32 				  # Numeric constant to represent Gage boards.
	strSerialNumber::Cstring;		  # String representing the board serial number.
	BaseBoardVersion::UInt32 		  # Version number of the base board.
	BaseBoardFirmwareVersion::UInt32   # Version number of the firmware on the base board.
	AddonBoardVersion::UInt32 		  # Version number of the Addon board.
	AddonBoardFirmwareVersion::UInt32  # Version number of the firmware of the Addon board.
	AddonFwOptions::UInt32 			  # Options of the current firmware image of Addon board
	BaseBoardFwOptions::UInt32 		  # Options of the current firmware image of Baseboard
	AddonHwOptions::UInt32 			  # Options of the Addon board pcb
	BaseBoardHwOptions::UInt32 		  # Options of the Baseboard pcb
	BoardInfo() = new(0,0,0,Ref{Cstring}(),0,0,0,0,0,0,0,0)
end

end


text:
function ensemble_history(e::IPM_Ensemble, bins=25)
    !e.sample_posterior && throw(ArgumentError("This ensemble has no posterior samples to show a history for!"))
    livec=vcat([model.log_Li for model in e.models],[model.log_Li for model in e.posterior_samples])
    show(histogram(livec, nbins=bins))
end

function e_backup(e::IPM_Ensemble, tuner::Permute_Tuner)
    cp(string(e.path,'/',"ens"),string(e.path,'/',"ens.bak"), force=true)
    serialize(string(e.path,'/',"ens"), e)
    serialize(string(e.path,'/',"tuner"), tuner)
end

function clean_ensemble_dir(e::IPM_Ensemble, model_pad::Integer; ignore_warn=false)
    !ignore_warn && e.sample_posterior && throw(ArgumentError("Ensemble is set to retain posterior samples and its directory should not be cleaned!"))
    for file in readdir(e.path)
        !(file in vcat([basename(model.path) for model in e.models],"ens",[string(number) for number in e.model_counter-length(e.models)-model_pad:e.model_counter-1])) && rm(e.path*'/'*file)
    end
end

function complete_evidence(e::IPM_Ensemble)
    log_Z=e.log_Zi[end]
    H=e.Hi[end]
    live_weight=lps(e.log_Xi[length(e.log_Li)], -log(length(e.models)))
    lis=sort([model.log_Li for model in e.models])
    liwis=sort(lps.(lis, live_weight))

    for (li,liwi) in zip(lis,liwis)
        last_Z=log_Z
        log_Z=logaddexp(log_Z,liwi)
        H=lps(
            (exp(lps(liwi,-log_Z)) * li),
            (exp(lps(last_Z,-log_Z)) * lps(H,last_Z)),
            -log_Z)
    end

    return log_Z, H
end

function measure_evidence(e::IPM_Ensemble)
    log_Z, H = complete_evidence(e)
    return measurement(log_Z,sqrt(abs(H)/length(e.models)))
end

function reset_ensemble!(e::IPM_Ensemble)
    new_e=deepcopy(e)
    for i in 1:length(e.models)
        if string(i) in [basename(record.path) for record in e.models]
            new_e.models[i]=e.models[findfirst(isequal(string(i)), [basename(record.path) for record in e.models])]
        else
            new_e.models[i]=e.posterior_samples[findfirst(isequal(string(i)), [basename(record.path) for record in e.posterior_samples])]
        end
    end

    new_e.contour=minimum([record.log_Li for record in new_e.models])

    new_e.log_Li=[new_e.log_Li[1]]
    new_e.log_Xi=[new_e.log_Xi[1]]
    new_e.log_wi=[new_e.log_wi[1]]
    new_e.log_Liwi=[new_e.log_Liwi[1]]
    new_e.log_Zi=[new_e.log_Zi[1]]
    new_e.Hi=[new_e.Hi[1]]

    new_e.posterior_samples=Vector{Model_Record}()

    new_e.model_counter=length(new_e.models)+1

    clean_ensemble_dir(new_e, 0; ignore_warn=true)
    isfile(e.path*"/tuner") && rm(e.path*"/tuner")
    serialize(e.path*"/ens", new_e)

    return new_e
end

function move_ensemble!(e::IPM_Ensemble,path::String)
    !isdir(path) && mkdir(path)
    for file in readdir(e.path)
        mv(e.path*'/'*file,path*'/'*file)
    end

    for (n,model) in enumerate(e.models)
        e.models[n]=Model_Record(path*'/'*basename(model.path), model.log_Li)
    end
    if e.sample_posterior
        for (n,model) in enumerate(e.posterior_samples)
            e.posterior_samples[n]=Model_Record(path*'/'*basename(model.path), model.log_Li)
        end
    end

    rm(e.path)
    e.path=path
    serialize(e.path*"/ens",e)
    return e
end

function copy_ensemble!(e::IPM_Ensemble,path::String)
    new_e=deepcopy(e)
    !isdir(path) && mkdir(path)
    for file in readdir(e.path)
        cp(e.path*'/'*file,path*'/'*file, force=true)
    end

    for (n,model) in enumerate(e.models)
        new_e.models[n]=Model_Record(path*'/'*basename(model.path), model.log_Li)
    end
    if e.sample_posterior
        for (n,model) in enumerate(e.posterior_samples)
            new_e.posterior_samples[n]=Model_Record(path*'/'*basename(model.path), model.log_Li)
        end
    end

    new_e.path=path
    serialize(new_e.path*"/ens",new_e)
    return new_e
end

function rewind_ensemble(e::IPM_Ensemble,rewind_idx)
    !e.sample_posterior && throw(ArgumentError("An ensemble not retaining posterior samples cannot be rewound!"))
    rewind_idx >= length(e.log_Li) && throw(ArgumentError("rewind_idx must be less than the current iterate!"))

    n=length(e.models)
    max_model_no=length(e.log_Li)+length(e.models)-1
    rewind_model_no=rewind_idx+length(e.models)-1
    new_e = deepcopy(e)

    rm_models=[string(name) for name in rewind_model_no+1:max_model_no]

    filter!(model->!(basename(model.path) in rm_models),new_e.models)
    filter!(model->!(basename(model.path) in rm_models),new_e.posterior_samples)

    while length(new_e.models) < n
        push!(new_e.models,pop!(new_e.posterior_samples))
    end
    new_e.contour=new_e.log_Li[rewind_idx]
    new_e.log_Li=new_e.log_Li[1:rewind_idx]
    new_e.log_Xi=new_e.log_Xi[1:rewind_idx]
    new_e.log_wi=new_e.log_wi[1:rewind_idx]
    new_e.log_Liwi=new_e.log_Liwi[1:rewind_idx]
    new_e.log_Zi=new_e.log_Zi[1:rewind_idx]
    new_e.Hi=new_e.Hi[1:rewind_idx]

    new_e.model_counter=length(new_e.models)+rewind_idx

    return new_e
end

function show_models(e::IPM_Ensemble,idxs)
    liperm=sortperm([model.log_Li for model in e.models],rev=true)
    for idx in idxs
        m=deserialize(e.models[liperm[idx]].path)
        show(m)
    end
end

function get_model(e::IPM_Ensemble,no)
    return deserialize(e.path*'/'*string(no))
end

function reestimate_ensemble!(e::IPM_Ensemble, wi_mode="trapezoidal")
    Ni=Int64.(round.(collect(-1:-1:-length(e.log_Li)+1)./e.log_Xi[2:end]))
    insert!(Ni,1,Ni[1])

    #fix any borked starting values
    e.log_Li[1]=-Inf #L0 = 0
	e.log_Xi[1]=0. #X0 = 1
	e.log_wi[1]=-Inf #w0 = 0
	e.log_Liwi[1]=-Inf #Liwi0 = 0
	e.log_Zi[1]=-Inf #Z0 = 0
	e.Hi[1]=0. #H0 = 0,

    for i in 1:length(e.log_Li)-1
        j=i+1

        e.log_Li[j]=e.posterior_samples[i].log_Li

        e.log_Xi[j]=-i/Ni[i]

        if wi_mode=="trapezoidal"
            e.log_wi[j]= logsubexp(e.log_Xi[i], -j/Ni[j]) - log(2) #log width of prior mass spanned by the last step-trapezoidal approx
        elseif wi_mode=="simple"
            e.log_wi[j]= logsubexp(e.log_Xi[i], e.log_Xi[j]) #log width of prior mass spanned by the last step-simple approx
        else
            throw(ArgumentError("Unsupported wi_mode!"))
        end
        e.log_Liwi[j]=lps(e.log_Li[j],e.log_wi[j]) #log likelihood + log width = increment of evidence spanned by iterate
        e.log_Zi[j]=logaddexp(e.log_Zi[i],e.log_Liwi[j])    #log evidence
        #information- dimensionless quantity
        Hj=lps(
            (exp(lps(e.log_Liwi[j],-e.log_Zi[j])) * e.log_Li[j]),
            (exp(lps(e.log_Zi[i],-e.log_Zi[j])) * lps(e.Hi[i],e.log_Zi[i])),
            -e.log_Zi[j])
        Hj === -Inf ? (e.Hi[j]=0.) : (e.Hi[j]=Hj)
    end
end


text:
module gmm

using Distributions
import Distributions: partype, logpdf
using Formatting: format
using ..Misc
using ..llh

using Random # randperm, MersenneTwister
using LinearAlgebra
using NNlib: softmax

export GMM, update, rmcomponents, add_noise_comp, addcomponents, responsibilities, importance_sample, is_eff_ss_per_comp, gmm_fit

@inline unsqueeze(xs, dim) = reshape(xs, (size(xs)[1:dim-1]..., 1, size(xs)[dim:end]...))
# @benchmark softmax(reduce(vcat, map(j -> AxUtil.gmm.log_gauss_llh(S, dGMM.mus[j,:], dGMM.sigmas[:,:,j], bypass=inactive_ixs[j]), 1:15)'))
# BenchmarkTools.Trial:
#   memory estimate:  25.60 MiB
#   allocs estimate:  712
#   --------------
#   minimum time:     14.678 ms (0.00% GC)
#   median time:      18.178 ms (0.00% GC)
#   mean time:        24.905 ms (19.52% GC)
#   maximum time:     473.150 ms (90.19% GC)
#   --------------
#   samples:          201
#   evals/sample:     1

# benchmark responsibilities(S, dGMM.mus, dGMM.sigmas, inactive_ixs)
# BenchmarkTools.Trial:
#   memory estimate:  29.47 MiB
#   allocs estimate:  394889
#   --------------
#   minimum time:     21.312 ms (0.00% GC)
#   median time:      23.391 ms (0.00% GC)
#   mean time:        33.607 ms (24.94% GC)
#   maximum time:     701.121 ms (95.72% GC)
#   --------------
#   samples:          149
#   evals/sample:     1

# function responsibilities(X::Matrix{T}, mus::Matrix{T}, sigmas::Array{T,3}, inactive_ixs) where T <: AbstractFloat
#     d, n = size(X)
#     k = size(mus, 1)
#     out = Matrix{T}(undef, k, n)
#     active_ixs = findall(.!inactive_ixs)
#
#     invLTs = [Matrix(inv(cholesky(sigmas[:,:,j]).L)) for j in 1:k]  # easier to just do inactive_ixs anyway...
#     normconsts = [-2*sum(log.(diag(invLT))) for invLT in invLTs]
#     @threads for i in 1:n
#         @inbounds begin
#             _max = -Inf   # not type stable
#             for j in 1:k
#                 if inactive_ixs[j]
#                     out[j, i] = 0. # not type stable; will ignore later
#                 else
#                     Δ = X[:, i] - mus[j, :]
#                     out[j, i] = -0.5 * (sum(x->x^2, invLTs[j] * Δ) + normconsts[j])
#                     _max = out[j, i] > _max ? out[j, i] : _max
#                 end
#             end
#             _sum = 0. # not type stable
#             for j in active_ixs
#                 out[j, i] = exp(out[j, i] - _max)
#                 _sum += out[j, i]
#             end
#             for j in active_ixs
#                 out[j, i] /= _sum
#             end
#         end # inbounds
#     end # loop over datapoints
#     return out
# end

# Definition / Constructor
struct GMM{T}
    mus::AbstractArray{T,2}      # k * d
    sigmas::AbstractArray{T,3}   # d * d * k
    pis::AbstractArray{T, 1}     # k
    GMM{T}(mus, sigmas, pis) where T <: Number = begin; x = new(mus, sigmas, pis); isvalidGMM(x); x; end
end

GMM(mus::AbstractArray{T,2}, sigmas::AbstractArray{T,3}, pis::AbstractArray{T, 1}) where T <: Number = GMM{T}(mus, sigmas, pis)

function isvalidGMM(x::GMM)
    @assert (x isa GMM{T} where T <: AbstractFloat) "type must be <: AbstractFloat"
    @assert (size(x.mus, 1) == size(x.sigmas,3) == length(x.pis)) "inconsistent ncomponents"
    @assert (size(x.mus, 2) == size(x.sigmas,1) == size(x.sigmas, 2)) "inconsistent dimension"
end

# Modify GMM 'objects'
update(d::GMM; mus=nothing, sigmas=nothing, pis=nothing) = GMM(something(mus, d.mus), something(sigmas, d.sigmas), something(pis, d.pis))

Base.convert(::Type{GMM}, d::MvNormal) = GMM{partype(d)}(unsqueeze(d.μ, 1), unsqueeze(Matrix(d.Σ), 3), [1.])
Base.convert(::Type{MixtureModel}, d::GMM) = MixtureModel([MvNormal(d.mus[j,:], d.sigmas[:,:,j]) for j in 1:ncomponents(d)], d.pis)
function Base.convert(::Type{GMM}, d::MixtureModel{Multivariate,Continuous,T}) where T <: MvNormal
    return GMM{partype(d)}(reduce(vcat, [d.components[j].μ' for j in 1:ncomponents(d)]),
               cat([Matrix(d.components[j].Σ) for j in 1:ncomponents(d)]..., dims=3), d.prior.p)
end
Base.copy(d::GMM) = GMM(copy(d.mus), copy(d.sigmas), copy(d.pis))

# remove components
function rmcomponents(d::GMM, ixs::Vector{T}) where T <: Signed
    bad_ixs = (sum(Flux.onehotbatch(Tuple(ixs), 1:ncomponents(d)), dims=2) .> 0)  |> dropdim2
    rmcomponents(d, bad_ixs)
end
rmcomponents(d::GMM, ixs::Vector{T}) where T <: Bool = rmcomponents(d, convert(BitArray, ixs))
rmcomponents(d::GMM, ixs::BitArray{1}) = GMM{partype(d)}(d.mus[.!ixs, :], d.sigmas[:, :, .!ixs], d.pis[.!ixs]/sum(d.pis[.!ixs]))

# add (>=0) noise component(s) into GMM, like Minka's classic clutter problem
function add_noise_comp(d::GMM; n::Signed=1, std_mult::AbstractFloat=2.0, pi_eat::AbstractFloat=0.02)
    @assert n >= 0
    (n == 0) && return d
    mus, sigmas, pis = d.mus, d.sigmas, d.pis
    p = size(d)
    center, cstdev = mean(mus, dims=1), std(mus, dims=1)
    new_mu = randn(n, p).*cstdev*0.2 .+ center
    mus = vcat(mus, new_mu)
    sigmas = cat(sigmas, repeat(mean(sigmas, dims=3)*std_mult, 1, 1, n), dims=3)
    pis = vcat(pis .* (1-pi_eat), ones(n) * pi_eat/n)
    GMM(mus, sigmas, pis)
end

# initialise new component(s) into GMM, by choosing existing point
function addcomponents(d::GMM; mus::AbstractArray{T,2}, sigmas::AbstractArray{T,3}, pis::AbstractArray{T, 1}) where T <: Number
    ∑π = sum(pis)
    @assert (0 < ∑π < 1) "sum(π) ∉ (0,1): π here is the target proportions: they must be non-zero and leave space for existing components."
    cpis = d.pis * (1-∑π)
    GMM(vcat(d.mus, mus), cat(d.sigmas, sigmas, dims=3), vcat(cpis, pis))
end

# GMM properties
partype(x::GMM{T}) where T <: AbstractFloat = T

Distributions.ncomponents(x::GMM) = length(x.pis)

Base.size(x::GMM) = size(x.mus, 2)

# GMM basic methods
Base.show(io::IO, z::GMM) = begin; print(io, format("{:d} component GMM in {:d} dimensions. Fields: " *
                            "`.mus`, `.sigmas`, and `.pis`.\n", ncomponents(z), size(z)));
                            show(convert(MixtureModel, z)); end

Base.rand(d::GMM, n::Int; shuffle=true) = sample_from_gmm(n, d.pis, d.mus, d.sigmas, shuffle=shuffle)

function sample_from_gmm(n, pis, mus, covs; shuffle=true)
    k, p = size(mus)
    Ns = rand(Multinomial(n, pis[:]))
    active_ixs = findall(Ns[:] .>= 1)

    ixs = hcat(vcat(1, 1 .+ cumsum(Ns[1:end-1], dims=1)), cumsum(Ns, dims=1))
    out = zeros(p, n)
    for j=active_ixs
        out[:, ixs[j,1]:ixs[j,2]] = rand(MvNormal(mus[j,:], covs[:,:,j]), Ns[j])
    end
    if shuffle
        out = out[:, randperm(n)]
    end
    return out
end

function logpdf(d::GMM, X::Matrix{T}; thrsh_comp=0.005) where T <: AbstractFloat
    return llh.gmm_llh(X, d.pis, d.mus, d.sigmas; thrsh_comp=thrsh_comp)
end

function logpdf(d::GMM, X::Vector{T}; thrsh_comp=0.005) where T <: AbstractFloat
    return llh.gmm_llh(reshape(X,:,1), d.pis, d.mus, d.sigmas; thrsh_comp=thrsh_comp)
end

responsibilities(X, d::GMM) = llh.responsibilities(X, d.mus, d.sigmas, d.pis)

# Importance sampling and ESS
function importance_sample(d::GMM, n::Int, log_f::Function; shuffle=false)
    S = rand(d, n; shuffle=shuffle)
    logW = log_f(S) - logpdf(d, S);
    return S, logW;
end

function is_eff_ss_per_comp(d::GMM, S::Matrix{T}, W::Vector{T}) where T <: AbstractFloat
    k = ncomponents(d)
    out = zeros(T, k)
    rs = softmax(reduce(vcat, map(j -> llh.log_gauss_llh(S, d.mus[j,:], d.sigmas[:,:,j], 1:k)')))
    out = map(j->is_eff_ss(W[rand(Categorical(rs[j,:]), length(W))]), 1:k)
    return out
end

is_eff_ss(W) = 1/sum((W./sum(W)).^2)

#====================================================================================
         Fitting to Empirical Data => Expectation Maximisation
====================================================================================#

function gmm_fit(X::Matrix{T}, d::GMM; max_iter=100, tol=1e-3, verbose=true, rm_inactive=false,
                 thrsh_comp=0.005, prior_strength=1.0) where T <: AbstractFloat
    gmm_fit(X, d.pis, d.mus, d.sigmas; max_iter=max_iter, tol=tol, verbose=verbose,
           rm_inactive=rm_inactive, thrsh_comp=thrsh_comp, prior_strength=prior_strength)
end

function gmm_fit(X::Matrix{T}, weights::Vector, d::GMM; max_iter=100, tol=1e-3, verbose=true,
                rm_inactive=false, thrsh_comp=0.005, prior_strength=1.0) where T <: AbstractFloat
    gmm_fit(X, weights, d.pis, d.mus, d.sigmas; max_iter=max_iter, tol=tol, verbose=verbose,
            rm_inactive=rm_inactive, thrsh_comp=thrsh_comp, prior_strength=prior_strength)
end

function gmm_fit(X, pi_prior, mu_prior, cov_prior; max_iter=100, tol=1e-3, verbose=true,
                 rm_inactive=false, thrsh_comp=0.005, prior_strength=1.0)
    gmm_fit(X, trues(size(X,2)), pi_prior, mu_prior, cov_prior; max_iter=max_iter,
            tol=tol, verbose=verbose, rm_inactive=rm_inactive, thrsh_comp=thrsh_comp, prior_strength=prior_strength)
end

function gmm_fit(X, weights, pi_prior, mu_prior, cov_prior; max_iter=100, tol=1e-3, verbose=true,
                 rm_inactive=false, thrsh_comp=0.005, prior_strength=1.0)
    p, n = size(X)
    k = length(pi_prior)
    @assert size(weights) == (n,)
    @assert size(mu_prior) == (k, p)
    @assert size(cov_prior) == (p, p, k)
    pis = pi_prior/sum(pi_prior)
    mus = copy(mu_prior)
    sigmas = copy(cov_prior)

    weights = weights / mean(weights)

    inactive_ixs = pi_prior[:] .< thrsh_comp
    pi_prior = copy(pi_prior) * prior_strength

    for i in range(1, stop=max_iter)
        # E-step
        rs = reduce(vcat, map(j -> llh.log_gauss_llh(X, mus[j,:], sigmas[:,:,j], bypass=inactive_ixs[j]), 1:k)')
        try
            rs .+= log.(pis)[:] # => DO NOT add log prior here
            catch e
            @warn "(gmm) rs and (log) pis are not conformable. The respective values are:"
            display(rs)
            display(log.(pis))
            rethrow(e)
        end
        @debug format("(gmm) ({:3d}/{:3d})", i, max_iter) llh=round(sum(log.(sum(pis .* (exp.(rs) .* weights'), dims=1))), digits=2)
        rs = softmax(rs)

        @debug format("(gmm) ({:3d}/{:3d}) original responsibilities", i, max_iter) rs=vec(sum(rs, dims=2))
        # reweight according to importance weights (see Adaptive IS in General Mix. Cappé et al. 2008)
        rs .*= weights'
        @debug format("(gmm) ({:3d}/{:3d}), wgtd responsibilities", i, max_iter) rs=vec(sum(rs, dims=2))

        # M-step
        Ns = vec(sum(rs, dims=2))
        inactive_ixs = Ns .< thrsh_comp
        active_ixs = .! inactive_ixs
        if any(inactive_ixs)
            pis[inactive_ixs] .= 0.0
            pi_prior[inactive_ixs] .= 0.0
        end
        pis = Ns + pi_prior[:]
        pis /= sum(pis)
        @debug format("(gmm) ({:3d}/{:3d})", i, max_iter) thrsh_comp=thrsh_comp n_inactive=sum(inactive_ixs) pis=pis
        # ==========>  .... SORT OUT X IS NOW p * n
        _mus = reduce(hcat, map(j -> sum(X .* rs[j:j,:], dims=2) .+ pi_prior[j]*mu_prior[j,:], findall(active_ixs)))'
        _mus ./= vec(Ns[active_ixs] + pi_prior[active_ixs])
        mus[active_ixs,:] = _mus

        @views for j in findall(active_ixs)
            Δx = X .- mus[j, :]
            Δμ = (mus[j,:] - mu_prior[j,:])'
            sigmas[:,:,j] = (Δx.*rs[j:j,:])*Δx' + pi_prior[j]*(Δμ'Δμ + cov_prior[:,:,j])
            sigmas[:,:,j] ./= (Ns[j] + pi_prior[j] + p + 2)     # normalizing terms from Wishart prior
            sigmas[:,:,j] = (sigmas[:,:,j] + sigmas[:,:,j]')/2 + eye(p)*1e-6   # hack: prevent collapse, sim. to prior.
        end
        # bson(format("dbg/gmm.bson"), pis=pis, sigmas=sigmas, mus=mus, X=X, weights=weights, active_ixs=active_ixs)

    end

    out = GMM{typeof(X[1])}(mus, sigmas, pis)
    if rm_inactive
        out = rmcomponents(out, inactive_ixs)
    end

    return out
end

end # module


text:
export applyPermutions, applyPermutionsRev
export permuteCombinations, permuteCombinationsName, flippings, flippingsName, convertMode2PermFlip

function convertMode2PermFlip(mode)
  if mode == 0
    perm = [2,1,3]
    flip = [1,2,3]
  elseif mode == 1
    perm = [2,1,3]
    flip = [2]
  elseif mode == 2
    perm = [2,3,1]
    flip = []
  else
    perm = [3,1,2]
    flip = [1,2]
  end
  return perm,flip
end

function applyPermutions(im, mode)
  perm, flip = convertMode2PermFlip(mode)

  im_ = permutedims(im,perm)
  for i in flip
    im_ = copyproperties(im, reverse(data(im_),dims=i))
  end
  return im_
end

function applyPermutionsRev(im, mode)
  perm, flip = convertMode2PermFlip(mode)
  im_ = im
  for i in flip
    im_ = copyproperties(im,reverse(data(im_),dims=i))
  end
  im_ = permutedims(im_,perm)
  return im_
end

function applyPermutions(im, perm, flip)
  im_ = permutedims(im,perm)
  for i in flip
    im_ = copyproperties(im, reverse(data(im_),dims=i))
  end
  return im_
end

function applyPermutionsRev(im, perm, flip)
  im_ = im
  for i in flip
    im_ = copyproperties(im,reverse(data(im_),dims=i))
  end
  im_ = permutedims(im_,perm)
  return im_
end

function permuteCombinations()
  perms = [
  [1,2,3],
  [1,3,2],
  [2,1,3],
  [2,3,1],
  [3,1,2],
  [3,2,1]]
 return perms
end

function permuteCombinationsName()
  permsName = [
  "[1,2,3]",
  "[1,3,2]",
  "[2,1,3] Isotropic 3D/Coronal MRI",
  "[2,3,1] Sagittal MRI",
  "[3,1,2] Transversal MRI",
  "[3,2,1]"]
 return permsName
end

function flippings()
  f = [
  Int64[],
  [1],
  [2],
  [3],
  [1,2],
  [1,3],
  [2,3],
  [1,2,3]]
 return f
end

function flippingsName()
  f =[
  "[] Transversal MRI",
  "[1]",
  "[2] Coronal MRI",
  "[3]",
  "[1,2] Sagittal MRI",
  "[1,3]",
  "[2,3]",
  "[1,2,3] Isotropic 3D"]
 return f
end


text:
################################################################################
## Multinomial Model
##   y ~ Multinomial(n, rho)
##   rho ~ Dirichlet(1, ..., 1)
################################################################################

using Mamba

## Data
n, k = 100, 5
rho0 = rand(Dirichlet(ones(k)))
y = rand(Multinomial(n, rho0))

## Log-transformed Posterior(rho) + Constant
logf = function(rho::DenseVector)
  logpdf(Multinomial(n, rho), y)
end

## MCMC Simulation with Slice Simplex Sampling
t = 10000
sim = Chains(t, k, names = map(i -> "rho[$i]", 1:k))
rho = SliceSimplexVariate(fill(1 / k, k), logf)
for i in 1:t
  sample!(rho)
  sim[i, :, 1] = rho
end
describe(sim)

p = plot(sim)
draw(p, filename = "slicesimplexplot")


text:
using LidJul
using Test

include("mit_implicit.jl")

# include("poisson2D.jl")


text:
"""
this script fetches sequential pairs of poses, fetches the big data (radar pings) tied to those poses, and then determines the pairwise factors that should be added between these sequential pairs
"""

using GraphPlot
using DistributedFactorGraphs
using IncrementalInference, RoME
using JSON2

# Where to fetch data
dfgDataFolder = joinpath(ENV["HOME"],"data","seagrant","rex")
datastore = FileDataStore("$dfgDataFolder/bigdata")
# Load the graph
fg = initfg()

# Reformat the data.
loadDFG!(fg, "$dfgDataFolder/fullsweep_updated.tar.gz");

# fetch variables containing a full sweep
allSweepVariables = filter(v -> :RADARSWEEP in getBigDataKeys(v), getVariables(fg)) |> sortDFG
fsvars = allSweepVariables .|> getLabel

# helper function to retrieve the radar sweep for a given variable
function fetchSweep(var::DFGVariable, store::FileDataStore)
    entry = getBigDataEntry(var, :RADARSWEEP)
    rawData = getBigData(datastore, entry)
    # raw data is json-encoded; this decoding should happen inside getBigData?
    rawdata = Vector{Float64}(JSON2.read(IOBuffer(rawData)))
    n = Int(sqrt(length(rawdata)))
    sweep = reshape(rawdata,(n,n))
    return sweep # That's pretty sweep if i say so myself...
end

# fetch all radar pings
sweeps = map(v -> fetchSweep(getVariable(fg, v), datastore), fsvars)
using Images, ImageView
# Filter the images
kg = Kernel.gaussian(7)
sweeps = map(s -> imfilter(s, kg), sweeps)
# Normalize
sweeps = map(s -> s/maximum(s), sweeps)

# Clamp out NaN's:
#map(clamp01nan, img)
#
# using ImageMagick
# for i in 5:length(sweeps)
#   s = sweeps[i]./maximum(sweeps[i])
#   save("/home/gearsad/SlamInDb/seagrant/$i.jpg", Gray.(s))
# end

# At this point we can load the sweeps; let's work on registration
# First step is to have a function that evaluates the cost of a given transform
# between two subsequent images.

using CoordinateTransformations
using ImageTransformations

# this function uses the sum of squared differences between the two images.
# To use low-passed versions of the images, simpy set the kernel argument to
# Kernel.gaussian(10) (or an appropriate size)
function getMismatch(a,b)
    sqrt(sum((a.-b).^2))
end
# sqrt(sum( imfilter( a.-b, Kernel.gaussian(5)).^2 ))
# sqrt(sum((imfilter(a, kernel).-imfilter(b, kernel)).^2))

# also do MMD version

# Next step is to define a function that applies a transform to the image. This
# transform consists of a translation and a rotation
function transformImage(img::Array{Float64,2}, dx::Real, dy::Real, dh::Real)
    tf = LinearMap(RotMatrix(dh))∘Translation(dx,dy)
    tf_img = warp(img, tf)

    # replace NaN w/ 0
    mask = findall(x->isnan(x),tf_img)
    tf_img[mask] .= 0.0
    return tf_img
end

# now we can combine the two into an evaluation function
function evaluateTransform(a::Array{Float64,2},b::Array{Float64,2}, dx::Float64, dy::Float64, dh::Float64)
    # transform image
    bp = transformImage(b,dx,dy,dh)
    # get matching padded views
    ap, bpp = paddedviews(0.0, a, bp)
    return getMismatch(ap,bpp)
end

## Building the graph
using DocStringExtensions
include("RadarFactor.jl")

using LinearAlgebra
using Optim

startsweep = 5
endsweep = 10
graphinit = false

# newfg = initfg()
newfg = generateCanonicalFG_ZeroPose2()
for i in 1:(endsweep-startsweep)
    addVariable!(newfg, Symbol("x$i"), Pose2, solvable=1)
end
for i in 1:(endsweep-startsweep)
    factor = AlignRadarPose2(sweeps[i+startsweep-1], sweeps[i+startsweep], MvNormal(zeros(3), diagm([5;5;pi/4])))
    addFactor!(newfg, Symbol.(["x$(i-1)", "x$i"]), factor, graphinit=graphinit, solvable=1)
end

# Run the initialization (very slow right now)
# ensureAllInitialized!(newfg)

# Factor debugging
# fs = getFactorFunction.(getFactor.(newfg, lsf(newfg)))
# fs = filter(f -> f isa AlignRadarPose2, fs)
# pf = convert.(PackedAlignRadarPose3, fs)
# convert.(AlignRadarPose2, pf)

# Save the graph
saveDFG(newfg, "$dfgDataFolder/segment_test.tar.gz");

lsf(newfg)
# this should run the radar alignment
pts = approxConv(newfg, :x0x1f1, :x1)

# solving will internally call ensureAllInitialized!(newfg)
tree, smt, hist = solveTree!(newfg)

## Looking at the results
using Plots

ppes = map(v -> getSuggestedPPE(getPPE(getVariable(newfg, v))), ls(newfg))
x = map(ppe -> ppe[1], ppes); y = map(ppe -> ppe[2], ppes); h = map(ppe -> ppe[3], ppes)
Plots.plot(x, y, title="Path Plot", lw=3)

## Stuff
using Optim

cost(x, im1, im2) = evaluateTransform(im1,im2,x[1],x[2],x[3])

# Plotting
xrange = -100.0:1.0:100.0
hrange = -pi:0.1:pi
val = reshape(
    [sweepx(sweeps[10],sweeps[11],xrange); sweepx(sweep_original[10],sweep_original[11],xrange)],
    length(xrange), 2)
Plots.plot(xrange,val)
# Heading
val = reshape(
    [sweeph(sweeps[10],sweeps[11],hrange); sweeph(sweep_original[10],sweep_original[11],hrange)],
    length(hrange), 2)
Plots.plot(hrange,val)

corr_func = (a,b)->sqrt(sum((a .- 0.5).*(b .- 0.5)))
val = reshape(
    [sweepx(sweeps[10],sweeps[11],xrange,diff_func=corr_func);
    sweepx(sweep_original[10],sweep_original[11],xrange,diff_func=corr_func)],
    length(xrange), 2)
Plots.plot(xrange,val)

## Sweep plotting
# sanity check: identity transform should yield zero cost
# @assert evaluateTransform(sweeps[11],sweeps[11],0.,0.,0.) == 0 "There's error with no transform!"

# let's try small displacements:
# sweepx(im1, im2, xrange) = (x->@show evaluateTransform(im1,im2,x,0.,0.)).(xrange)
# sweepy(im1, im2, yrange) = (y->@show evaluateTransform(im1,im2,0.,y,0.)).(yrange)
# sweeph(im1, im2, hrange) = (h->@show evaluateTransform(im1,im2,0.,0.,h)).(hrange)

# using Plots
# xrange = -10:0.1:10
# hrange = -pi:0.1:pi
# Plots.plot(xrange,sweepx(sweeps[10],sweeps[11],xrange))
# Plots.plot(xrange,sweepy(sweeps[10],sweeps[11],xrange))
# Plots.plot(hrange,sweeph(sweeps[10],sweeps[11],hrange))

# fs10 = imfilter(sweeps[10],Kernel.gaussian(3))
# fs11 = imfilter(sweeps[11],Kernel.gaussian(3))
# ffs10 = imfilter(fs10,Kernel.gaussian(3))
# ffs11 = imfilter(fs11,Kernel.gaussian(3))
#
# Plots.plot(xrange,sweepx(ffs10,ffs11,xrange))
# Plots.plot(xrange,sweepy(fs10,fs11,xrange))


text:
using Documenter
using LayeredArrays

# Setup for doctests in docstrings
DocMeta.setdocmeta!(LayeredArrays, :DocTestSetup, recursive = true,
    quote
        using LayeredArrays
    end
)

makedocs(;
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    modules = [LayeredArrays],
    sitename = "LayeredArrays.jl",
    pages=[
        "Home" => "index.md",
    ],
    doctest = true, # :fix
)

deploydocs(
    repo = "github.com/KeitaNakamura/LayeredArrays.jl.git",
    devbranch = "main",
)


text:
struct AutoType end

"""
    auto_parse(::Union{Type,AutoType}, val)

Automatically parse the value into a Julia value.
"""
auto_parse(ValType::Type, val) = parse(ValType, val)

auto_parse(::Type{AutoType}, val) = eval(Meta.parse(val))

auto_parse(::Type{String}, val) = String(strip(val))

"""
    parse_type(valtype, default_type::Type = AutoType)

Parse the type of `valtype`. If `valtype` has a type declaration,
like `parse_type("2::ComplexF64")`, it gets parsed as that type
declared, and returns `(ComplexF64, 2)`.

If `val` doesn't have a type declaration, it gets parsed into
`default_type`, which defaults to `AutoType`, so `parse_type("2")`
returns `(AutoType, 2)`
"""
function parse_type(valtype, default_type::Type = AutoType)
    # Check for a type decleration
    valtype_vec = split(valtype, "::"; limit = 2)
    ValType = default_type
    if length(valtype_vec) > 1
      # Type declaration
      ValType = eval(Meta.parse(valtype_vec[2]))
    end
    val = valtype_vec[1]
    return ValType, val
end

"""
    argsdict([args_list::Vector];
             first_arg::Int = 1,
             delim = '=',
             as_symbols::Bool = false,
             default_named_type::Type = ITensors.AutoType,
             save_positional::Bool = true,
             default_positional_type::Type = String,
             prefix::String = "_arg")

Parse the command line arguments such as `julia N=2 X=1e-12`
and put them in a dictionary, where the keys are before
the delimiter and the values are after, so
`Dict("N" => "2", "X" => "1e-12")`.
"""
function argsdict(args_list::Vector;
                  first_arg::Int = 1,
                  delim = '=',
                  as_symbols::Bool = false,
                  default_named_type::Type = AutoType,
                  save_positional::Bool = true,
                  default_positional_type::Type = String,
                  prefix::String = "_arg")
  KeyType = as_symbols ? Symbol : String
  parsed = Dict{KeyType, Any}()
  narg = 1
  for n in first_arg:length(args_list)
    a = args_list[n]

    # Check if it is a command line flag
    if startswith(a, "--")
      flag = a[3:end]
      if flag == "autotype" || flag == "a"
        default_positional_type = AutoType
        default_named_type = AutoType
      elseif flag == "stringtype" || flag == "s"
        default_positional_type = String
        default_named_type = String
      end
      continue
    end

    optval = split(a, delim)
    if length(optval) == 1
      if save_positional
        val = only(optval)
        parsed[KeyType("$prefix$narg")] =
          auto_parse(parse_type(val, default_positional_type)...)
        narg += 1
      else
        @warn "Ignoring argument $a since it does not have the delimiter \"$delim\"."
      end
      continue
    elseif length(optval) == 2
      opt, val = optval
    else
      error("Argument $a has more than one delimiter \"$delim\", which is not well defined.")
    end
    ValType, key = parse_type(opt, default_named_type)
    key = strip(key)
    ' ' in key && error("Option \"$key\" contains spaces, which is not well defined")
    typedkey = KeyType(key)
    typedval = auto_parse(ValType, val)
    parsed[typedkey] = typedval
  end
  return parsed
end

argsdict(; kwargs...) = argsdict(ARGS; kwargs...)


text:
# Module to perform simulations of the CIR process
# CIR = Cox-Ingersoll-Ross (1985)

module CIRProc

import Distributions: Gamma, Poisson, Chisq, rand, pdf, quantile
import Base: rand, mean, var

export CIRProcess, ssdist, scale, shape, rand, invmean, empmean,
    ExactCIRSim, EulerCIRSim, XCIRState,
    AbstractXCIRBound, XCIRTimeBound, XCIRParticleBound, XCIRXTauBound,
    AbstractXCIRSim, XCIRSim, XCIRDMSim, getdt, getmu, getcirp,
    FullXCIRTraj, StatsXCIRTraj, reset!, push!, first, last, lastt, mean, var,
    randtraj!, simpct, simproc

# ----------------------------------------------------------------------------
# CIR Process & Simulation
# ----------------------------------------------------------------------------

# -- Process --

immutable CIRProcess
    mu::Float64
    sig::Float64
    theta::Float64

    function CIRProcess(mu::Real, sig::Real, theta::Real)
        mu, sig, theta = float(mu), float(sig), float(theta)
        mu > 0 || error("mu needs to be positive")
        sig > 0 || error("sig needs to be positive")
        theta >= 0 || error("theta needs to be non-negative")
        new(mu, sig, theta)
    end

    CIRProcess() = CIRProcess(0.4, 0.2, 2)
end

# process parameters
shape(p::CIRProcess) = p.mu^2 / p.sig^2
scale(p::CIRProcess) = p.sig^2 / p.mu
pdf(p::CIRProcess, x::Real) = pdf(ssdist(p), x)
quantile(p::CIRProcess, x::Real) = quantile(ssdist(p), x)
# returns the steady state distribution
ssdist(p::CIRProcess) = Gamma(shape(p), scale(p))
# sample from steady state distribution
rand(p::CIRProcess) = rand(ssdist(p))
# mean of inverse steady state distribution
invmean(p::CIRProcess) = p.mu > p.sig ? p.mu / (p.mu * p.mu - p.sig * p.sig) : Inf

# for dgrat() and drcomp() in empmean
include("specialfuns.jl")

# empirical expectation of piece-wise linear function
#
# d is the process, x is the points at which the function is known, and fx is
# the function value at these points. x needs to be strictly increasing and
# cover all the parts of the steady-state distribution of p with mass.
#
# The function returns <f(x)>, x ~ ssdist(p)
function empmean{T1 <: Real, T2 <: Real}(p::CIRProcess, x::Vector{T1}, fx::Vector{T2})
    # computation based on assuming piece-wise linear fx and solving
    # expectation analytically over [x[j], x[j+1]] blocks
    res = 0.0
    const k, theta = shape(p), scale(p)
    dgratj, drcompj = dgrat(k, x[1] / theta)[1], drcomp(k, x[1] / theta)
    for j = 1:(length(x)-1)
        const dgratj1 = dgrat(k, x[j+1] / theta)[1]
        const drcompj1 = drcomp(k, x[j+1] / theta)
        res += (theta * (fx[j] - fx[j+1]) * (k * (dgratj - dgratj1) - drcompj + drcompj1) +
                (fx[j+1] * x[j] - fx[j] * x[j+1]) * (dgratj - dgratj1)) / (x[j+1] - x[j])
        dgratj, drcompj = dgratj1, drcompj1
    end
    res
end

# -- Simulations --
# simulations follow
# Anderson, Jaeckel, Kahl (2010): Simulation of Square-Root Processes
abstract AbstractCIRSim

# exact simulation
immutable ExactCIRSim <: AbstractCIRSim
    cirp::CIRProcess
    nt_dt::Float64
    d::Float64
    exp_theta_dt_nt_dt::Float64

    function ExactCIRSim(cirp::CIRProcess, dt::Float64)
        const exp_tmp = exp(-cirp.theta*dt)
        const nt_tmp = 2cirp.mu*exp_tmp/(cirp.sig^2*(1-exp_tmp))
        new(cirp, nt_tmp, 2cirp.mu^2/cirp.sig^2, exp_tmp/nt_tmp)
    end
end
rand(p::ExactCIRSim, curtau::Real) = p.exp_theta_dt_nt_dt *
    rand(Chisq(p.d + 2rand(Poisson(0.5curtau * p.nt_dt))))

# approximate simulation by truncated Euler method
immutable EulerCIRSim <: AbstractCIRSim
    cirp::CIRProcess
    theta_tau_dt::Float64
    eps_dt::Float64

    EulerCIRSim(cirp::CIRProcess, dt::Float64) =
        new(cirp, cirp.theta*dt, cirp.sig*sqrt(2cirp.theta*dt/cirp.mu))
end
function rand(p::EulerCIRSim, curtau::Float64)
    const taup = max(0.0, curtau)
    curtau + p.theta_tau_dt * (p.cirp.mu - taup) +
        p.eps_dt * sqrt(taup) * randn()
end

# generates a sample path of the CIR process
#
# p is a tau process
# n is the number of steps of simulate
function simcir(p::AbstractCIRSim, n::Int)
    taus = fill(NaN, n)
    taus[1] = rand(p.cirp)
    for i = 2:n
        taus[i] = rand(p, taus[i-1])
    end
    taus
end

# ----------------------------------------------------------------------------
# XCIR Process Simulation
# ----------------------------------------------------------------------------

# -- State --

# state of XCIR process
immutable XCIRState
    x::Float64
    tau::Float64
    XCIRState(x::Real, tau::Real) = new(float64(x), float64(tau))
    XCIRState(p::CIRProcess) = XCIRState(0.0, rand(p))
    XCIRState() = XCIRState(0.0, 1.0)
end

# -- Bounds --

# bound on XCIR process state & time
abstract AbstractXCIRBound

# bound on time only
immutable XCIRTimeBound <: AbstractXCIRBound
    t::Float64

    function XCIRTimeBound(t::Real)
        t > zero(t) || error("t needs to be positive")
        new(float(t))
    end
end
crossedbound(b::XCIRTimeBound, xtau::XCIRState, t::Real) = float(t) >= b.t

# bound on belief/particle only
immutable XCIRParticleBound <: AbstractXCIRBound
    x::Float64

    function XCIRParticleBound(x::Real)
        x > zero(x) || error("x needs to be positive")
        new(float(x))
    end
end
crossedbound(b::XCIRParticleBound, xtau::XCIRState, t::Real) = abs(xtau.x) >= b.x

# bound on particle, depending on tau
#
# taux a vector, specifying for each tau the bound in x (and -x). The
# corresponding taus are tau0, tau0+dtau, tau0+2dtau, ...
immutable XCIRXTauBound <: AbstractXCIRBound
    taux::Vector{Float64}
    dtau::Float64
    mintau::Float64
    imax::Float64

    function XCIRXTauBound(taux::Vector{Float64}, dtau::Real, mintau::Real)
        dtau, mintau = float(dtau), float(mintau)
        dtau > 0 || error("dtau needs to be postitive")
        mintau >= 0 || error("mintau needs to be non-negative")
        new(taux, dtau, mintau, float(length(taux) - 1))
    end
end
function crossedbound(b::XCIRXTauBound, xtau::XCIRState, t::Real)
    const x, tau = xtau.x, xtau.tau
    const i = (tau - b.mintau) / b.dtau
    const btau = i <= 0.0     ? b.taux[1] :
                 i >= b.imax  ? b.taux[end] :
                 isinteger(i) ? b.taux[iround(i) + 1] : begin
            const ix = ifloor(i)
            i -= float(ix)
            (1.0 - i) * b.taux[ix+1] + i * b.taux[ix+2]
        end
    abs(x) >= btau
end

# -- Simulation --

# simulating XCIRP process
abstract AbstractXCIRSim
getdt(s::AbstractXCIRSim) = s.dt
getmu(s::AbstractXCIRSim) = s.mu
getcirp(s::AbstractXCIRSim) = s.cirp

# simulating XCIRProcess with optimal (weighted) evidence accumulation
immutable XCIRSim <: AbstractXCIRSim
    cirp::CIRProcess
    mu::Float64
    dt::Float64
    cirs::AbstractCIRSim
    dmu::Float64

    function XCIRSim(cirp::CIRProcess, mu::Real, dt::Real, cirs::AbstractCIRSim)
        dt > zero(dt) || error("dt needs to be positive")
        new(cirp, float(mu), float(dt), cirs, float(mu * dt))
    end
    XCIRSim(cirp, mu, dt) = XCIRSim(cirp, mu, dt, EulerCIRSim(cirp, dt))
    XCIRSim(cirp, dt) = XCIRSim(cirp, 1.0, dt, EulerCIRSim(cirp, dt))
end
# draws a sample given current (x, tau) state
function rand(p::XCIRSim, curxtau::XCIRState)
    const tau = rand(p.cirs, curxtau.tau)
    const taupos = max(0.0, tau)
    XCIRState(curxtau.x + taupos * p.dmu + sqrt(p.dt * taupos) * randn(), tau)
end

# simulating unweighted XCIR process, diffusion model-style
immutable XCIRDMSim <: AbstractXCIRSim
    cirp::CIRProcess
    mu::Float64
    dt::Float64
    cirs::AbstractCIRSim
    dmu::Float64

    function XCIRDMSim(cirp::CIRProcess, mu::Real, dt::Real, cirs::AbstractCIRSim)
        dt > zero(dt) || error("dt needs to be positive")
        new(cirp, float(mu), float(dt), cirs, float(mu * dt))
    end
    XCIRDMSim(cirp, mu, dt) = XCIRDMSim(cirp, mu, dt, EulerCIRSim(cirp, dt))
    XCIRDMSim(cirp, dt) = XCIRDMSim(cirp, 1.0, dt, EulerCIRSim(cirp, dt))
end
# draws a sample given current (x, tau) state
function rand(p::XCIRDMSim, curxtau::XCIRState)
    const tau = rand(p.cirs, curxtau.tau)
    XCIRState(curxtau.x + p.dmu + sqrt(p.dt / max(eps(), tau)) * randn(), tau)
end

# -- State Trajectory --

# sample trajectory of XCIR process - maintains statistics about trajectory
abstract AbstractXCIRTraj

# stores full trajectory of XCIR process
type FullXCIRTraj <: AbstractXCIRTraj
    dt::Float64
    n::Int
    nincr::Int
    xtaus::Matrix{Float64}

    function FullXCIRTraj(xtau0::XCIRState, dt::Real)
        dt = float(dt)
        dt > 0.0 || error("dt needs to be positive")
        nincr = max(1, int(div(1, dt)))
        xtaus = fill(NaN, nincr, 2)
        xtaus[1,1] = xtau0.x
        xtaus[1,2] = xtau0.tau
        new(dt, 1, nincr, xtaus)
    end
    FullXCIRTraj(dt) = FullXCIRTraj(XCIRState(), dt)
end
function reset!(t::FullXCIRTraj, xtau0::XCIRState)
    t.xtaus = fill(NaN, t.nincr, 2)
    t.xtaus[1,1] = xtau0.x
    t.xtaus[1,2] = xtau0.tau
    t.n = 1;
end
reset!(t::FullXCIRTraj, p::CIRProcess) = reset!(t, rand(p))
function push!(t::FullXCIRTraj, xtau::XCIRState)
    if t.n == size(t.xtaus, 1)
        t.xtaus = [t.xtaus, fill(NaN, t.nincr, 2)]
    end
    t.n += 1
    t.xtaus[t.n,1] = xtau.x
    t.xtaus[t.n,2] = xtau.tau;
end
first(t::FullXCIRTraj) = XCIRState(t.xtaus[1,1], t.xtaus[1,2])
last(t::FullXCIRTraj) = XCIRState(t.xtaus[t.n,1], t.xtaus[t.n,2])
lastt(t::FullXCIRTraj) = (t.n - 1) * t.dt
mean(t::FullXCIRTraj) = XCIRState(mean(t.xtaus[1:t.n,1]), mean(t.xtaus[1:t.n,2]))
var(t::FullXCIRTraj) = XCIRState(var(t.xtaus[1:t.n,1]), var(t.xtaus[1:t.n,2]))

# only maintains trajectory statistics of XCIR process
type StatsXCIRTraj <: AbstractXCIRTraj
    dt::Float64
    n::Int
    x0::Float64
    tau0::Float64
    xn::Float64
    taun::Float64
    xsum::Float64
    tausum::Float64
    x2sum::Float64
    tau2sum::Float64

    function StatsXCIRTraj(xtau0::XCIRState, dt::Real)
        dt > zero(dt) || error("dt needs to be positive")
        new(float(dt), 1, xtau0.x, xtau0.tau, xtau0.x, xtau0.tau,
            xtau0.x, xtau0.tau, xtau0.x * xtau0.x, xtau0.tau * xtau0.tau)
    end
    StatsXCIRTraj(dt) = StatsXCIRTraj(XCIRState(), dt)
end
function reset!(t::StatsXCIRTraj, xtau0::XCIRState)
    t.x0, t.tau0, t.xn, t.taun = xtau0.x, xtau0.tau, xtau0.x, xtau0.tau
    t.xsum, t.x2sum = xtau0.x, xtau0.x * xtau0.x
    t.tausum, t.tau2sum = xtau0.tau, xtau0.tau * xtau0.tau
    t.n = 1;
end
reset!(t::StatsXCIRTraj, p::CIRProcess) = reset!(t, rand(p))
function push!(t::StatsXCIRTraj, xtau::XCIRState)
    t.xn, t.taun = xtau.x, xtau.tau
    t.xsum += xtau.x
    t.tausum += xtau.tau
    t.x2sum += xtau.x * xtau.x
    t.tau2sum += xtau.tau * xtau.tau
    t.n += 1;
end
first(t::StatsXCIRTraj) = XCIRState(t.x0, t.tau0)
last(t::StatsXCIRTraj) = XCIRState(t.xn, t.taun)
lastt(t::StatsXCIRTraj) = (t.n - 1) * t.dt
mean(t::StatsXCIRTraj) = XCIRState(t.xsum / t.n, t.tausum / t.n)
var(t::StatsXCIRTraj) =
    XCIRState(t.x2sum / t.n - (t.xsum / t.n)^2, t.tau2sum / t.n - (t.tausum / t.n)^2)

# -- Bounded Simulations --

# simulates a trajectory until the bound is reached
#
# The stimulation starts at last(t) and stops when the bound is crossed.
function randtraj!(t::AbstractXCIRTraj, s::AbstractXCIRSim, b::AbstractXCIRBound)
    xtau = last(t)
    while !crossedbound(b, xtau, lastt(t))
        xtau = rand(s, xtau)
        push!(t, xtau)
    end
end

# simulates process until bound is reached, returning some process state
#
# s is the simulator, b is the bound.
#
# The function returns (xtau, t) where xtau is the XCIRState and t is the time
# after bound crossing.
function rand(s::AbstractXCIRSim, b::AbstractXCIRBound)
    const dt = getdt(s)
    xtau = XCIRState(0.0, rand(getcirp(s)))
    t = 0.0
    while !crossedbound(b, xtau, t)
        xtau = rand(s, xtau)
        t += dt
    end
    xtau, t
end

# returns p(correct) and bound-crossing time for single trial
#
# s is the simulator, b is the bound
#
# The function returns (pc, t) where pc is the probability correct, and t is
# the bound-crossing time. pc is 1.0 if the sign of s.mu is the same as the
# sign of the final x.
function simpct(s::AbstractXCIRSim, b::AbstractXCIRBound)
    const xtau, t = rand(s, b)
    const pc = getmu(s) >= 0.0 ? (xtau.x >= 0.0 ? 1.0 : 0.0) :
                                 (xtau.x < 0.0 ? 1.0 : 0.0)
    pc, t
end
# Analytical solutions for optimal (weighted) simulation and bound on particle
function simpct(s::XCIRSim, b::XCIRParticleBound)
    const thetamin = 1e-10
    const pc = 1.0 / (1.0 + exp(-2.0b.x))
    # for small theta we can use the analytical solution
    const cirp = getcirp(s)
    if cirp.theta < thetamin
        return pc, invmean(cirp) * b.x * tanh(b.x)
    else
        const xtau, t = rand(s, b)
        return pc, t
    end
end

# simulates a set of bounded trials and returns statistics
#
# s is the XCIR simulator, b is the bound, and n is the number of trials.
#
# The function returns a dictionary with the following elements:
# - rt: vector of decision times per trial
# - choice: vector of choices, 1.0 for upper boundary, 0.0 for lower
# - tau0: vector of initial tau's
# - taun: vector of tau's at decision time
# - avgtau: vector of average tau's
function simproc(s::AbstractXCIRSim, b::AbstractXCIRBound, n::Int)
    const cirp = getcirp(s)
    t = StatsXCIRTraj(getdt(s))
    # simulate trajectories
    rt, choice = Array(Float64, n), Array(Float64, n)
    tau0, taun, avgtau = Array(Float64, n), Array(Float64, n), Array(Float64, n)
    for i = 1:n
        reset!(t, XCIRState(0.0, rand(cirp)))
        randtraj!(t, s, b)
        rt[i] = lastt(t)
        choice[i] = last(t).x >= 0.0  ? 1.0 : 0.0
        tau0[i], taun[i], avgtau[i] = first(t).tau, last(t).tau, mean(t).tau
    end
    ["rt" => rt, "choice" => choice,
     "tau0" => tau0, "taun" => taun, "avgtau" => avgtau]
end

end # module


text:
using NetworkDynamicalSystems
using Test

@testset "NetworkDynamicalSystems.jl" begin
    # Write your tests here.
end


text:
using DataFrames, CSV
data = CSV.read("data/purchaseData.csv", DataFrame)

newCol = DataFrame(Validated=ones(Int, size(data,1)))
newRow = DataFrame([["JOHN", "JACK"] [123456, 909595]], [:Name, :PhoneNo])
newData = DataFrame(Name=["JOHN", "ASHELY","MARYANNA"],
                    Job=["Lawyer", "Doctor","Lawyer"])

data = hcat(data, newCol)
println(first(data, 3), "\n")

data = vcat(data, newRow, cols=:union)
println(last(data, 3), "\n")

data = innerjoin(data, newData, on=:Name, matchmissing=:equal)
println(data, "\n")

select!(data,[:Name,:Job])
println(data, "\n")

unique!(data,:Job)
println(data)

max_stars_repo_name: newptcai/StatsWithJuliaBook

max_issues_repo_name: newptcai/StatsWithJuliaBook

max_forks_repo_name: newptcai/StatsWithJuliaBook


text:
const WCS_HEADERS_TEMPLATES = [
    "DATE",
    "MJD",

    "WCSAXESa",
    "WCAXna",
    "WCSTna",
    "WCSXna",
    "CRPIXja",
    "jCRPna",
    "jCRPXn",
    "TCRPna",
    "TCRPXn",
    "PCi_ja",
    "ijPCna",
    "TPn_ka",
    "TPCn_ka",
    "CDi_ja",
    "ijCDna",
    "TCn_ka",
    "TCDn_ka",
    "CDELTia",
    "iCDEna",
    "iCDLTn",
    "TCDEna",
    "TCDLTn",
    "CROTAi",
    "iCROTn",
    "TCROTn",
    "CUNITia",
    "iCUNna",
    "iCUNIn",
    "TCUNna",
    "TCUNIn",
    "CTYPEia",
    "iCTYna",
    "iCTYPn",
    "TCTYna",
    "TCTYPn",
    "CRVALia",
    "iCRVna",
    "iCRVLn",
    "TCRVna",
    "TCRVLn",
    "LONPOLEa",
    "LONPna",
    "LATPOLEa",
    "LATPna",
    "RESTFREQ",
    "RESTFRQa",
    "RFRQna",
    "RESTWAVa",
    "RWAVna",
    "PVi_ma",
    "iVn_ma",
    "iPVn_ma",
    "TVn_ma",
    "TPVn_ma",
    "PROJPm",
    "PSi_ma",
    "iSn_ma",
    "iPSn_ma",
    "TSn_ma",
    "TPSn_ma",
    "VELREF",
    "CNAMEia",
    "iCNAna",
    "iCNAMn",
    "TCNAna",
    "TCNAMn",
    "CRDERia",
    "iCRDna",
    "iCRDEn",
    "TCRDna",
    "TCRDEn",
    "CSYERia",
    "iCSYna",
    "iCSYEn",
    "TCSYna",
    "TCSYEn",
    "CZPHSia",
    "iCZPna",
    "iCZPHn",
    "TCZPna",
    "TCZPHn",
    "CPERIia",
    "iCPRna",
    "iCPERn",
    "TCPRna",
    "TCPERn",
    "WCSNAMEa",
    "WCSNna",
    "TWCSna",
    "TIMESYS",
    "TREFPOS",
    "TRPOSn",
    "TREFDIR",
    "TRDIRn",
    "PLEPHEM",
    "TIMEUNIT",
    "DATEREF",
    "MJDREF",
    "MJDREFI",
    "MJDREFF",
    "JDREF",
    "JDREFI",
    "JDREFF",
    "TIMEOFFS",
    "DATE-OBS",
    "DOBSn",
    "DATE-BEG",
    "DATE-AVG",
    "DAVGn",
    "DATE-END",
    "MJD-OBS",
    "MJDOBn",
    "MJD-BEG",
    "MJD-AVG",
    "MJDAn",
    "MJD-END",
    "JEPOCH",
    "BEPOCH",
    "TSTART",
    "TSTOP",
    "XPOSURE",
    "TELAPSE",
    "TIMSYER",
    "TIMRDER",
    "TIMEDEL",
    "TIMEPIXR",
    "OBSGEO-X",
    "OBSGXn",
    "OBSGEO-Y",
    "OBSGYn",
    "OBSGEO-Z",
    "OBSGZn",
    "OBSGEO-L",
    "OBSGLn",
    "OBSGEO-B",
    "OBSGBn",
    "OBSGEO-H",
    "OBSGHn",
    "OBSORBIT",
    "RADESYSa",
    "RADEna",
    "RADECSYS",
    "EPOCH",
    "EQUINOXa",
    "EQUIna",
    "SPECSYSa",
    "SPECna",
    "SSYSOBSa",
    "SOBSna",
    "VELOSYSa",
    "VSYSna",
    "VSOURCEa",
    "VSOUna",
    "ZSOURCEa",
    "ZSOUna",
    "SSYSSRCa",
    "SSRCna",
    "VELANGLa",
    "VANGna",
    "RSUN_REF",
    "DSUN_OBS",
    "CRLN_OBS",
    "HGLN_OBS",
    "HGLT_OBS",
    "NAXISn",
    "CROTAn",
    "PROJPn",
    "CPDISja",
    "CQDISia",
    "DPja",
    "DQia",
    "CPERRja",
    "CQERRia",
    "DVERRa",
    "A_ORDER",
    "B_ORDER",
    "AP_ORDER",
    "BP_ORDER",
    "A_DMAX",
    "B_DMAX",
    "A_p_q",
    "B_p_q",
    "AP_p_q",
    "BP_p_q",
    "CNPIX1",
    "PPO3",
    "PPO6",
    "XPIXELSZ",
    "YPIXELSZ",
    "PLTRAH",
    "PLTRAM",
    "PLTRAS",
    "PLTDECSN",
    "PLTDECD",
    "PLTDECM",
    "PLTDECS",
    "PLATEID",
    "AMDXm",
    "AMDYm",
    "WATi_m"
]

# Expand the headers containing lower case specifers into N copies
# Find all lower case templates
const WCS_HEADERS = Set(mapreduce(vcat, WCS_HEADERS_TEMPLATES) do template
    if any(islowercase, template)
        template_chars = Vector{Char}(template)
        chars = template_chars[islowercase.(template_chars)]
        out = String[template]
        for replace_target in chars
            newout = String[]
            for template in out
                for i in [""; string.(1:4); string.('a':'d')]
                    push!(newout, replace(template, replace_target=>i))
                end
            end
            append!(out, newout)
        end
        out
    else
        template
    end
end)

"""
    emptywcs()

Given an AbstractArray, return a blank WCSTransform of the appropriate
dimensionality.
"""
emptywcs(data::AbstractArray) = WCSTransform(ndims(data))
emptywcs(img::AstroImage) = WCSTransform(length(dims(img))+length(refdims(img)))

# """
#     filterwcsheader(hdrs::FITSHeader)

# Return a new FITSHeader containing WCS header from `hdrs`.
# This is useful for creating a new image with the same coordinates
# as another.
# """
# function filterwcsheader(hdrs::FITSHeader)
#     include_keys = intersect(keys(hdrs), WCS_HEADERS)
#     return FITSHeader(
#         include_keys,
#         map(key -> hdrs[key], include_keys),
#         map(key -> get_comment(hdrs, key), include_keys),
#     )
# end

"""
    wcsfromheader(img::AstroImage; relax=WCS.HDR_ALL, ignore_rejected=true)

Helper function to create a WCSTransform from an array and
FITSHeaders.
"""
function wcsfromheader(img::AstroImage; relax=WCS.HDR_ALL)
    # We only need to stringify WCS header. This might just be 4-10 header keywords
    # out of thousands.
    local wcsout

    io = IOBuffer()
    serializeheader(io, header(img))
    hdrstr = String(take!(io))

    # Load the headers without ignoring rejected to get error messages
    try
        wcsout = WCS.from_header(
            hdrstr;
            ignore_rejected=false,
            relax
        )
    catch err
        # Load them again ignoring error messages
        wcsout = WCS.from_header(
            hdrstr;
            ignore_rejected=true,
            relax
        )
        # If that still fails, the use gets the stack trace here
        # If not, print a warning about rejected header
        @warn "WCSTransform was generated by ignoring rejected header. It may not be valid." exception=err
    end

    if length(wcsout) == 0
        return [emptywcs(img)]
    else
        return wcsout
    end
end

# Map FITS stokes numbers to a symbol
function _stokes_symbol(i)
    return if i == 1
        :I
    elseif i == 2
        :Q
    elseif i == 3
        :U
    elseif i == 4
        :V
    elseif i == -1
        :RR
    elseif i == -2
        :LL
    elseif i == -3
        :RL
    elseif i == -4
        :LR
    elseif i == -5
        :XX
    elseif i == -6
        :YY
    elseif i == -7
        :XY
    elseif i == -8
        :YX
    else
        @warn "unknown FITS stokes number $i. See \"Representations of world coordinates in FITS\", Table 7."
        nothing
    end
end
function _stokes_name(symb)
    return if symb == :I
        "Stokes Unpolarized"
    elseif symb == :Q
        "Stokes Linear Q"
    elseif symb == :U
        "Stokes Linear U"
    elseif symb == :V
        "Stokes Circular"
    elseif symb == :RR
        "Right-right cicular"
    elseif symb == :LL
        "Left-left cicular"
    elseif symb == :RL
        "Right-left cross-cicular"
    elseif symb == :LR
        "Left-right cross-cicular"
    elseif symb == :XX
        "X parallel linear"
    elseif symb == :YY
        "Y parallel linear"
    elseif symb == :XY
        "XY cross linear"
    elseif symb == :YX
        "YX cross linear"
    else
        @warn "unknown FITS stokes key $symb. See \"Representations of world coordinates in FITS\", Table 7."
        ""
    end
end

# Smart versions of pix_to_world and world_to_pix
"""
    pix_to_world(img::AstroImage, pixcoords; all=false)

Given an astro image, look up the world coordinates of the pixels given
by `pixcoords`. World coordinates are resolved using WCS.jl and a
WCSTransform calculated from any FITS header present in `img`. If
no WCS information is in the header, or the axes are all linear, this will
just return pixel coordinates.

`pixcoords` should be the coordinates in your current selection
of the image. For example, if you select a slice like this:
```julia
julia> cube = load("some-3d-cube.fits")
julia> slice = cube[10:20, 30:40, 5]
```

Then to look up the coordinates of the pixel in the bottom left corner of
`slice`, run:
```julia
julia> world_coords = pix_to_world(img, [1, 1])
[10, 30]
```

If WCS information was present in the header of `cube`, then those coordinates
would be resolved using axis 1, 2, and 3 respectively.

To include world coordinates in all axes, pass `all=true`
```julia
julia> world_coords = pix_to_world(img, [1, 1], all=true)
[10, 30, 5]
```

!! Coordinates must be provided in the order of `dims(img)`. If you transpose
an image, the order you pass the coordinates should not change.
"""
function WCS.pix_to_world(img::AstroImage, pixcoords; wcsn=1, all=false, parent=false)
    if pixcoords isa Array{Float64}
        pixcoords_prepared = pixcoords
    else
        pixcoords_prepared = [Float64(c) for c in pixcoords]
    end
    D_out = wcs(img,wcsn).naxis
    if ndims(pixcoords_prepared) > 1
        worldcoords_out = similar(pixcoords_prepared, Float64, D_out, size(pixcoords_prepared,2))
    else
        worldcoords_out = similar(pixcoords_prepared, Float64, D_out)
    end

    # Find the coordinates in the parent array.
    # Dimensional data
    # pixcoords_floored = floor.(Int, pixcoords)
    # pixcoords_frac = (pixcoords .- pixcoords_floored) .* step.(dims(img))
    # parentcoords = getindex.(dims(img), pixcoords_floored) .+ pixcoords_frac
    if parent
        parentcoords = pixcoords
    else
        parentcoords = pixcoords .* step.(dims(img)) .+ first.(dims(img))
    end
    # WCS.jl is very restrictive. We need to supply a Vector{Float64}
    # as input, not any other kind of collection.
    # TODO: avoid allocation in case where refdims=() and pixcoords isa Array{Float64}
    if ndims(pixcoords_prepared) > 1
        parentcoords_prepared = zeros(wcs(img,wcsn).naxis, size(pixcoords_prepared,2))
    else
        parentcoords_prepared = zeros(wcs(img,wcsn).naxis)
    end
    # out = zeros(Float64, wcs(img,wcsn).naxis, size(pixcoords,2))
    for (i, dim) in enumerate(dims(img))
        j = wcsax(img, dim)
        parentcoords_prepared[j,:] .= parentcoords[i,:] .- 1
    end
    for dim in refdims(img)
        j = wcsax(img, dim)
        # Non numeric reference dims can be used, e.g. a polarization axis of symbols I, Q, U, etc.
        if eltype(dim) <: Number
            z = dim[1] - 1
        else
            # Find the index of the symbol into the parent cube
            parentrefdim = img.wcsdims[findfirst(d->name(d)==name(dim), img.wcsdims)]
            z = findfirst(==(first(dim)), collect(parentrefdim)) - 1
        end
        parentcoords_prepared[j,:] .= z
    end

    # Get world coordinates along all slices
    WCS.pix_to_world!(wcs(img, wcsn), parentcoords_prepared, worldcoords_out)

    # If user requested world coordinates in all dims, not just selected
    # dims of img
    if all
        return worldcoords_out
    end

    # Otherwise filter to only return coordinates along selected dims.
    if ndims(pixcoords_prepared) > 1
        world_coords_of_these_axes = zeros(length(dims(img)), size(pixcoords_prepared,2))
    else
        world_coords_of_these_axes = zeros(length(dims(img)))
    end
    for (i, dim) in enumerate(dims(img))
        j = wcsax(img, dim)
        world_coords_of_these_axes[i,:] .= worldcoords_out[j,:]
    end

    return world_coords_of_these_axes
end

##
function WCS.world_to_pix(img::AstroImage, worldcoords; parent=false, wcsn=1)
    if worldcoords isa Array{Float64}
        worldcoords_prepared = worldcoords
    else
        worldcoords_prepared = [Float64(c) for c in worldcoords]
    end
    D_out = wcs(img,wcsn).naxis
    if ndims(worldcoords_prepared) > 1
        out = similar(worldcoords_prepared, Float64, D_out, size(worldcoords_prepared,2))
    else
        out = similar(worldcoords_prepared, Float64, D_out)
    end
    return WCS.world_to_pix!(out, img, worldcoords_prepared; parent)
end
function WCS.world_to_pix!(pixcoords_out, img::AstroImage, worldcoords; wcsn=1, parent=false)
    # # Find the coordinates in the parent array.
    # # Dimensional data
    # worldcoords_floored = floor.(Int, worldcoords)
    # worldcoords_frac = (worldcoords .- worldcoords_floored) .* step.(dims(img))
    # parentcoords = getindex.(dims(img), worldcoords_floored) .+ worldcoords_frac
    # WCS.jl is very restrictive. We need to supply a Vector{Float64}
    # as input, not any other kind of collection.
    # TODO: avoid allocation in case where refdims=() and worldcoords isa Array{Float64}
    if ndims(worldcoords) > 1
        worldcoords_prepared = zeros(wcs(img,wcsn).naxis,size(worldcoords,2))
    else
        worldcoords_prepared = zeros(wcs(img,wcsn).naxis)
    end
    # TODO: we need to pass in ref dims locations as well, and then filter the
    # output to only include the dims of the current slice?
    # out = zeros(Float64, wcs(img,wcsn).naxis, size(worldcoords,2))
    for (i, dim) in enumerate(dims(img))
        j = wcsax(img, dim)
        worldcoords_prepared[j,:] = worldcoords[i,:]
    end
    for dim in refdims(img)
        j = wcsax(img, dim)
        # Non numeric reference dims can be used, e.g. a polarization axis of symbols I, Q, U, etc.
        if eltype(dim) <: Number
            z = dim[1]
        else
            # Find the index of the symbol into the parent cube
            parentrefdim = img.wcsdims[findfirst(d->name(d)==name(dim), img.wcsdims)]
            z = findfirst(==(first(dim)),collect(parentrefdim)) -1
        end
        worldcoords_prepared[j,:] .= z
    end

    # This returns the parent pixel coordinates.
    # TODO: switch to non-allocating version.
    pixcoords_out .= WCS.world_to_pix(wcs(img, wcsn), worldcoords_prepared)

    if !parent
        coordoffsets = zeros(wcs(img,wcsn).naxis)
        coordsteps = zeros(wcs(img,wcsn).naxis)
        for (i, dim) in enumerate(dims(img))
            j = wcsax(img, dim)
            coordoffsets[j] = first(dims(img)[i])
            coordsteps[j] = step(dims(img)[i])
        end
        for dim in refdims(img)
            j = wcsax(img, dim)
            # Non numeric reference dims can be used, e.g. a polarization axis of symbols I, Q, U, etc.
            if eltype(dim) <: Number
                coordoffsets[j] = first(dim)
                coordsteps[j] = step(dim)
            else
                # Find the index of the symbol into the parent cube
                parentrefdim = img.wcsdims[findfirst(d->name(d)==name(dim), img.wcsdims)]
                z = findfirst(==(first(dim)),collect(parentrefdim))
                coordoffsets[j] = z - 1
                coordsteps[j] = 1
            end
        end

        pixcoords_out .-= coordoffsets
        pixcoords_out .= (pixcoords_out .+ 1) ./ coordsteps
    end
    return pixcoords_out
end

## For now, we use a copied version of FITSIO's show method for FITSHeader.
# We have to be careful to format things in a way WCSLib will like.
# In particular, we can't put newlines after each 80 characters.
# FITSIO has to do this so users can see the header.

# functions for displaying header values in show(io, header)
hdrval_repr(v::Bool) = v ? "T" : "F"
hdrval_repr(v::String) = @sprintf "'%-8s'" v
hdrval_repr(v::Union{AbstractFloat, Integer}) = string(v)

function serializeheader(io, hdr::FITSHeader)
    n = length(hdr)
    for i=1:n
        if hdr.keys[i] == "COMMENT" || hdr.keys[i] == "HISTORY"
                lastc = min(72, lastindex(hdr.comments[i]))
                @printf io "%s %s" hdr.keys[i] hdr.comments[i][1:lastc]
                print(io, " "^(72-lastc))
        else
            @printf io "%-8s" hdr.keys[i]
            if hdr.values[i] === nothing
                print(io, "                      ")
                rc = 50  # remaining characters on line
            elseif hdr.values[i] isa String
                val = hdrval_repr(hdr.values[i])
                @printf io "= %-20s" val
                rc = length(val) <= 20 ? 50 : 70 - length(val)
            else
                val = hdrval_repr(hdr.values[i])
                @printf io "= %20s" val
                rc = length(val) <= 20 ? 50 : 70 - length(val)
            end
            if length(hdr.comments[i]) > 0
                lastc = min(rc-3, lastindex(hdr.comments[i]))
                @printf io " / %s" hdr.comments[i][1:lastc]
                rc -= lastc + 3
            end
            print(io, " "^rc)
        end
        if i == n
            print(io, "\nEND"*(" "^77))
        else
            print(io)
        end
    end
end


text:
module VOOTreeSearch

# using Revise
using POMDPs
using BasicPOMCP
using POMCPOW
using ParticleFilters
using Parameters
using MCTS
using CPUTime
using D3Trees
using Colors
using Random
using Printf
using POMDPPolicies
using Distributions
using Distances
using VDPTag2
using LinearAlgebra
using POMDPModelTools
using StaticArrays
using MCTS

using BasicPOMCP: convert_estimator

import Base: insert!
import POMDPs: action, solve, mean, rand, updater, currentobs, history
import POMDPModelTools: action_info, GenerativeBeliefMDP
import Distances: euclidean

import MCTS: next_action

export
    VOOActionGenerator,
    voo_sample,
    AbstractRolloutFirst,
    RolloutFirst,
    VOORolloutFirst

include("VOO.jl")

export
    VOSSOptions,
    VOOSparseSamplingSolver,
    VOWSSOptions,
    VSSPlanner,
    valuepairs

include("sparse_sampling.jl")
include("sparse_sampling_solver.jl")

export
	BoxND

include("new_actions.jl")

export
	LQGPOMDP,
    solve_LQG_one_step,
    LQGPolicy,
    RiccatiLQRPolicy,
    ExactLQRPolicy,
    RandomLQRPolicy,
    obs_weight,
    max_possible_weight,
	new_particle

include("lqg.jl")

end # module


text:
using Test
using MPI

# TODO: enable CUDA tests once OpenMPI has full support
ArrayType = Array

MPI.Init()

comm = MPI.COMM_WORLD

const rank = MPI.Comm_rank(comm)
const N = MPI.Comm_size(comm)

buf = ArrayType{Int}(fill(Int(rank),N))
received = ArrayType{Int}(fill(-1,N))

# Test Get using fence
win = MPI.Win_create(buf, comm)
MPI.Win_fence(win)
MPI.Get!(received, win; rank=(rank+1)%N)
MPI.Win_fence(win)

@test received == ArrayType{Int}(fill(Int(rank+1)%N, size(received)))

# Test locked window
if rank != 0
  MPI.Win_lock(win; rank=0, type=:exclusive)
  received[1] = rank
  MPI.Put!(view(received,1:1), win; rank=0, disp=rank)
  MPI.Win_unlock(win; rank=0)
else
  buf[1] = 0
end

MPI.Win_fence(win)

if rank == 0
  @test buf == collect(0:N-1)
end

MPI.Barrier(comm)

if rank == 1
  MPI.Win_lock(win; rank=1, type=:exclusive)
  fill!(buf,3)
  MPI.Win_unlock(win; rank=1)
end

MPI.Barrier(comm)

if rank == 0
  MPI.Win_lock(win; rank=0, type=:exclusive)
  fill!(buf, 2)
  MPI.Win_unlock(win; rank=0)
  MPI.Win_lock(win; rank=1, type=:exclusive)
  result = similar(buf)
  MPI.Get_accumulate!(buf, result, MPI.SUM, win; rank=1)
  MPI.Win_unlock(win; rank=1)
  @test all(result .== 3)
end

MPI.Barrier(comm)

if rank == 1
  MPI.Win_lock(win; rank=1, type=:exclusive)
  @test all(buf .== 5)
  fill!(buf,-2)
  MPI.Win_unlock(win; rank=1)
  MPI.Win_lock(win; rank=0, type=:exclusive)
  MPI.Accumulate!(buf, MPI.SUM, win; rank=0)
  MPI.Win_unlock(win; rank=0)
  MPI.Win_lock(win; rank=1, type=:exclusive)
  fill!(buf,1)
  MPI.Win_unlock(win; rank=1)
end

MPI.Barrier(comm)

if rank == 0
  MPI.Win_lock(win; rank=0, type=:exclusive)
  @test all(buf .== 0)
  MPI.Win_unlock(win; rank=0)
end

MPI.Barrier(comm)
MPI.free(win)
MPI.Barrier(comm)

# Try a dynamic window
win = MPI.Win_create_dynamic(comm)
MPI.Win_attach!(win, buf)

address_buf = Vector{Cptrdiff_t}(undef, 1)
address_win = MPI.Win_create(address_buf, comm)

# Address for dynamic window is global instead of relative to the start of the buffer, so we need to share it
MPI.Win_lock(address_win; rank=rank, type=:exclusive)
address_buf[1] = MPI.Get_address(buf)
MPI.Win_unlock(address_win; rank=rank)

MPI.Barrier(comm)

# Test atomic operation
if rank == 0
  received = Ref(0)
  to_send = Ref(0)
  for r in 0:N-1
    address = Ref{Cptrdiff_t}(0)
    MPI.Win_lock(address_win; rank=r, type=:exclusive)
    MPI.Get!(address, address_win; rank=r)
    MPI.Win_flush(address_win; rank=r)
    to_send[] = r+5
    MPI.Win_lock(win; rank=r, type=:exclusive)
    MPI.Fetch_and_op!(to_send, received, MPI.REPLACE, win; rank=r, disp=address[]+r*sizeof(Int))
    MPI.Win_flush(win; rank=r)
    @test received[] == r
    MPI.Win_unlock(win; rank=r)
    MPI.Win_unlock(address_win; rank=r)
  end
end

MPI.Barrier(comm)

@test buf[rank+1] == rank+5

MPI.Barrier(comm)
MPI.Win_detach!(win, buf)
MPI.free(win)
MPI.free(address_win)

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


text:
import CSV
using CSV

page_start =
"@def title = \"Papers\"
@def tags = [\"pdf\", \"publications\"]

# Archive of Papers\n"

page_content = ""
try
    reverse_file_list = sort(readdir("./_assets/papers"), rev=true)

    for csv_file in filter(f->contains(f, ".csv"), reverse_file_list)
        header = string("## ", replace(csv_file, ".csv"=>""), "\n@@doi-table")

        global page_content = string(page_content, header, "\n|Title|Authors|DOI/PDF|\n|-----|-------|------|\n")

        row_num = 0
        csv_reader = CSV.File(string("./_assets/papers/", csv_file), delim=";")
        for row in csv_reader
            global page_content = string(page_content, "|$(row.Title)|$(row.Authors)|[DOI](https://doi.org/$(row.DOI))/[PDF]($(row.pdf))|\n")
            row_num += 1
        end

        if row_num == 0
           global page_content = string(page_content, "| | | | |\n")
        end

       global page_content = string(page_content, "@@\n\n")
    end

    global page_content = string(page_start, page_content)

catch e
    @info e
    global page_content = string(page_start, "An error occured while reading the papers, please be patient, while we try to fix this issue.")
end

open("papers.md", "w") do io
    write(io, page_content)
end


text:
# First-order vector Stochastic Differential Equations

# NOTE: for scalar SDEs, checkstable can also be used
"""isstableode(A) -> Bool

Determine whether the matrix A corresponds to a stable ordinary differential equation
"""
isstableode(A) = all(real.(eigen(A).values).<0)

"""isdiagonalizable(A) -> Bool

Determine whether the matrix A is diagonalizable
"""
isdiagonalizable(A) = rank(eigen(A).vectors)==size(A,1)

abstract type TimeSeriesModel{T<:Number} end

abstract type AbstractSde{T<:Number} end

"""
Vector valued, first order stochastic differential equation with measurement model
$FIELDS
"""
struct VectorSde{T<:Number,U<:AbstractMatrix{T}} <: AbstractSde{T}
  A::U
  C::Matrix{T}
  Q::Matrix{T}
  R::Matrix{T}
  function VectorSde(A::U,C,Q,R) where {T,U<:AbstractMatrix{T}}
    p = LinearAlgebra.checksquare(A)
    q = size(C,1)
    size(C,2)==p || error("C must be q×",p)
    LinearAlgebra.checksquare(Q)==p || error("Q must be ",p,"×",p)
    LinearAlgebra.checksquare(R)==q || error("R must be ",q,"×",q)
    new{T,U}(A,C,Q,R)
  end
end

function Base.convert(::Type{AbstractSde{T}},x::VectorSde) where {T}
  return VectorSde(convert.(AbstractMatrix{T},[x.A,x.C,x.Q,x.R])...)
end

"State dimension"
statedim(m::VectorSde) = size(m.A,1)

"measurement dimension"
measdim(m::VectorSde) = size(m.R,1)

"Model order"
order(eq::VectorSde) = 1

isstable(eq::VectorSde) = isstableode(eq.A)

"""
    statcovz(m::VectorSde) -> R0

Stationary covariance matrix of the state variable `z`
"""
statcovz(m::VectorSde) = project_symmetric(lyap(m.A,m.Q))

"""
$SIGNATURES

Stationary covariance matrix of the observations `y`
"""
statcov(m::VectorSde) = m.C*statcovz(m)*m.C'+m.R

"""
Stationary state distribution
$SIGNATURES
"""
statstate(eq::AbstractSde) = State(statcovz(eq))

"""
    diag_vectorsde(m::VectorSde) -> VectorSde

Representation of the  vector SDE `m` with diagonalized `A` matrix.
Only if `scale` is `true, scale eigenvectors that the first component
is one.

NOTE: In general `diag_vectorsde` returns complex-valued models. Not all
functions work on complex-valued models, e.g. `rand`.
"""
function diag_vectorsde(m::VectorSde,scale=true)
  E = eigen(m.A)
  U = E.vectors
  if scale
    for j = 1:size(U,2)
      U[:,j] ./= U[1,j]
    end
  end
  A = Diagonal(E.values)
  C = m.C*U
  Q = U\m.Q/U'
  R = complex(m.R)
  VectorSde(A,C,Q,R)
end

function LinearAlgebra.lyap(A::Diagonal,C)
  a = A.diag
  p = length(a)
  ac = repeat(a,1,p)
  m = ac.+ac'
  -C./m
end

"Observation time and value"
struct Obs
  t::Float64
  y::Float64
end

## SDE process ————————————————————
"""
SDE process: Combines Stochastic Differential Equation with state
$FIELDS
"""
struct SdeProcess{T,U} <: TimeSeriesModel{T}
  "Stochastic Differential Equation"
  eq::U
  "State"
  state::State{T}
  function SdeProcess(eq::U,state::State{T}) where {T,U}
    if statedim(eq)≠dim(state) error("Model order mismatch between SDE and state") end
    new{T,U}(eq,state)
  end
end

SdeProcess(eq) = SdeProcess(eq,statstate(eq))

function Base.show(io::IO,proc::SdeProcess)
  print(io,proc.eq,"; state : ")
  if has_statstate(proc)
    print(io,"stationary")
  else
    print(io,proc.state)
  end
end

VectorSdeProcess = SdeProcess{T,U} where {T,U<:VectorSde}

function describe_model(proc::SdeProcess)
  describe_model(proc.eq)
  print("State: ")
  if has_statstate(proc)
    println("Stationary")
  else
    println(proc.state)
  end
end

order(proc::SdeProcess) = order(proc.eq)

Statistics.std(proc::SdeProcess) = std(proc.eq)

"Has stationary state?"
has_statstate(proc::SdeProcess;rtol=√eps()) = isstable(proc.eq) && isapprox(proc.state,statstate(proc.eq);rtol=rtol)

function Statistics.mean(procs::Vector{<:SdeProcess})
  eq = mean(getproperty.(procs,:eq))
  state = mean(getproperty.(procs,:state))
  return SdeProcess(eq,state)
end


text:
"""
    SIR

A type representing a Susceptible-Infected-Recovered model with infection rate `β` and recovery rate `δ`.
"""
struct SIR <: AbstractCompartimentalModel
    β::Real
    δ::Real
end

SI(sir::SIR) = SI(sir.β)

num_states(::SIR) = 3

# Metapopulation

function rate(i, state, mp::Metapopulation{SIR})
    β = mp.dynamics.β
    δ = mp.dynamics.δ
    D = mp.D
    h = mp.h
    β*state[i,1]*state[i,2] + δ*state[i,2] + sum(D .* state[i,:]) * (outdegree(h,i)>0)
end

function update_state_and_rates!(state, a, k, mp::Metapopulation{SIR})
    M = nv(mp.h)
    D = mp.D
    β = mp.dynamics.β
    δ = mp.dynamics.δ
    h = mp.h
    p = vcat(β*state[k,1]*state[k,2], δ*state[k,2], D.*state[k,:]*(outdegree(h,k)>0))
    j = sample(ProbabilityWeights(p, a[k]))
    if j == 1 # infection
        state[k,1] -= 1
        state[k,2] += 1
        a[k] = rate(k, state, mp)
    elseif j == 2 # recovery
        state[k,2] -= 1
        state[k,3] += 1
        a[k] = rate(k, state, mp)
    else # migration
        i = rand(outneighbors(h, k))
        state[k,j-2] -= 1
        state[i,j-2] += 1
        a[k] = rate(k, state, mp)
        a[i] = rate(i, state, mp)
    end
end

function meanfield_fun(mp::Metapopulation{SIR})
    L = normalized_laplacian(mp.h)
    β = mp.dynamics.β
    δ = mp.dynamics.δ
    D = mp.D
    f! = function(dx, x, p, t)
        infection = β*x[:,1].*x[:,2]
        recovery = δ*x[:,2]
        dx[:,1] .= .-infection .- D[1].*(L'*x[:,1])
        dx[:,2] .= infection .- recovery .- D[2].*(L'*x[:,2])
        dx[:,3] .= recovery .- D[3].*(L'x[:,3])
    end
    return f!
end

# Contact Process

function rate(k, state, cp::ContactProcess{SIR})
    if state[k] == 1
        l = length(filter(i->state[i]==2, inneighbors(cp.g, k)))
        return cp.dynamics.β * l
    elseif state[k] == 2
        return cp.dynamics.δ
    else
        return 0.0
    end
end

function update_state_and_rates!(state, a, k, cp::ContactProcess{SIR})
    g = cp.g
    β = cp.dynamics.β
    if state[k] == 1
        state[k] = 2
        for i in Iterators.filter(j->state[j]==1, outneighbors(g, k))
            a[i] += β
        end
    elseif state[k] == 2
        state[k] = 3
        for i in Iterators.filter(j->state[j]==1, outneighbors(g, k))
            a[i] -= β
        end
    end
    a[k] = rate(k, state, cp)
end

function update_output!(output, state, n, k, cp::ContactProcess{SIR})
    if state[k] == 2
        output[n] = output[n-1] + [-1, 1, 0]
    else
        output[n] = output[n-1] + [0, -1, 1]
    end
end

function meanfield_fun(cp::ContactProcess{SIR})
    A = adjacency_matrix(cp.g)
    β = cp.dynamics.β
    δ = cp.dynamics.δ
    f! = function(dx, x, p, t)
        infection = β.*x[:,1].*(A*x[:,2])
        dx[:,1] .= .- infection
        dx[:,2] .= infection .- δ.*x[:,2]
        dx[:,3] .= δ.*x[:,2]
    end
    return f!
end

# Metaplex

function rate(k, state, mpx::Metaplex{SIR})
    β = mpx.dynamics.β
    δ = mpx.dynamics.δ
    D = mpx.D
    x_i = state[:state_i]
    x_μ = state[:state_μ]
    od = outdegree(mpx.h, x_μ[k]) > 0
    if x_i[k] == 1
        l = length(filter(i->x_i[i]==2 && x_μ[i]==x_μ[k], inneighbors(mpx.g, k)))
        return β*l + D[1]*od
    elseif x_i[k] == 2
        return δ + D[2]*od
    else
        return D[3]*od
    end
end

function update_state_and_rates!(state, a, k, mpx::Metaplex{SIR})
    N = nv(mpx.g)
    M = nv(mpx.h)
    g = mpx.g
    h = mpx.h
    x_i = state[:state_i]
    x_μ = state[:state_μ]
    popcounts = state[:popcounts]
    β = mpx.dynamics.β
    δ = mpx.dynamics.δ
    D = mpx.D
    μ = x_μ[k]
    if x_i[k] == 1
        if rand()*a[k] < D[1] # susceptible node migrates
            ν = rand(outneighbors(h, μ))
            x_μ[k] = ν
            popcounts[μ,1] -= 1
            popcounts[ν,1] += 1
            a[k] = rate(k, state, mpx)
        else # infection
            x_i[k] = 2
            popcounts[μ,1] -= 1
            popcounts[μ,2] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1 && x_μ[j]==μ, outneighbors(g,k))
                a[i] += β
            end
        end
    elseif x_i[k] == 2
        if rand()*a[k] < D[2] # infected node migrates
            ν = rand(outneighbors(h,μ))
            x_μ[k] = ν
            popcounts[μ,2] -= 1
            popcounts[ν,2] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1, outneighbors(g,k))
                if x_μ[i] == μ
                    a[i] -= β
                elseif x_μ[i] == ν
                    a[i] += β
                end
            end
        else # recovery
            x_i[k] = 3
            popcounts[μ,2] -= 1
            popcounts[μ,3] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1 && x_μ[j]==μ, outneighbors(g,k))
                a[i] -= β
            end
        end
    else # recovered node migrates
        ν = rand(outneighbors(h,μ))
        x_μ[k] = ν
        popcounts[μ,3] -= 1
        popcounts[ν,3] += 1
        a[k] = rate(k, state, mpx)
    end
end

function meanfield_fun(mpx::Metaplex{SIR})
    N = nv(mpx.g)
    M = nv(mpx.h)
    A = adjacency_matrix(mpx.g)
    L = normalized_laplacian(mpx.h)
    β = mpx.dynamics.β
    δ = mpx.dynamics.δ
    D = mpx.D
    f! = function(dx, x, p, t)
        for i in 1:N
            dx[1,i,:] .= .-D[1]*L'*x[1,i,:]
            dx[2,i,:] .= .-D[2]*L'*x[2,i,:]
            dx[3,i,:] .= .-D[3]*L'*x[3,i,:]
        end
        for μ in 1:M
            infection = β.*x[1,:,μ].*(A*x[2,:,μ])
            recovery = δ.*x[2,:,μ]
            dx[1,:,μ] .-= infection
            dx[2,:,μ] .+= infection .- recovery
            dx[3,:,μ] .+= recovery
        end
    end
    return f!
end

# Heterogeneous Metaplex

function rate(k, state, mpx::HeterogeneousMetaplex{SIR})
    β = mpx.dynamics.β
    δ = mpx.dynamics.δ
    D = mpx.D
    x_i = state[:state_i]
    x_μ = state[:state_μ]
    μ = x_μ[k]
    od = outdegree(mpx.h, x_μ[k]) > 0
    if x_i[k] == 1
        l = length(filter(i->x_i[i]==2 && x_μ[i]==μ, inneighbors(mpx.g[μ], k)))
        return β*l + D[1]*od
    elseif x_i[k] == 2
        return δ + D[2]*od
    else
        return D[3]*od
    end
end

function update_state_and_rates!(state, a, k, mpx::HeterogeneousMetaplex{SIR})
    g = mpx.g
    h = mpx.h
    x_i = state[:state_i]
    x_μ = state[:state_μ]
    popcounts = state[:popcounts]
    β = mpx.dynamics.β
    D = mpx.D
    μ = x_μ[k]
    if x_i[k] == 1
        if rand()*a[k] < D[1] # susceptible node migrates
            ν = rand(outneighbors(h, μ))
            x_μ[k] = ν
            popcounts[μ,1] -= 1
            popcounts[ν,1] += 1
            a[k] = rate(k, state, mpx)
        else # infection
            x_i[k] = 2
            popcounts[μ,1] -= 1
            popcounts[μ,2] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1 && x_μ[j]==μ, outneighbors(g[μ],k))
                a[i] += β
            end
        end
    elseif x_i[k] == 2
        if rand()*a[k] < D[2] # infected node migrates
            ν = rand(outneighbors(h,μ))
            x_μ[k] = ν
            popcounts[μ,2] -= 1
            popcounts[ν,2] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1, outneighbors(g[μ],k))
                if x_μ[i] == μ
                    a[i] -= β
                elseif x_μ[i] == ν
                    a[i] += β
                end
            end
        else # recovery
            x_i[k] = 3
            popcounts[μ,2] -= 1
            popcounts[μ,3] += 1
            a[k] = rate(k, state, mpx)
            for i in Iterators.filter(j->x_i[j]==1 && x_μ[j]==μ, outneighbors(g[μ],k))
                a[i] -= β
            end
        end
    else # recovered node migrates
        ν = rand(outneighbors(h,μ))
        x_μ[k] = ν
        popcounts[μ,3] -= 1
        popcounts[ν,3] += 1
        a[k] = rate(k, state, mpx)
    end
end

function meanfield_fun(mpx::HeterogeneousMetaplex{SIR})
    N = nv(mpx.g[1])
    M = nv(mpx.h)
    A = adjacency_matrix.(mpx.g)
    L = normalized_laplacian(mpx.h)
    β = mpx.dynamics.β
    δ = mpx.dynamics.δ
    D = mpx.D
    f! = function(dx, x, p, t)
        for i in 1:N
            dx[1,i,:] .= .-D[1]*L'*x[1,i,:]
            dx[2,i,:] .= .-D[2]*L'*x[2,i,:]
            dx[3,i,:] .= .-D[3]*L'*x[3,i,:]
        end
        for μ in 1:M
            infection = β.*x[1,:,μ].*(A[μ]*x[2,:,μ])
            recovery = δ.*x[2,:,μ]
            dx[1,:,μ] .-= infection
            dx[2,:,μ] .+= infection .- recovery
            dx[3,:,μ] .+= recovery
        end
    end
    return f!
end


text:
"""
    FullSpaceModel{T, S} <: AbstractNLPModel{T, S}

A data type to store the full space `BlockNLPModel` as an `AbstractNLPModel`.
"""
mutable struct FullSpaceModel{T,S} <: AbstractNLPModel{T,S}
    meta::NLPModelMeta{T,S}
    counters::Counters
    blocknlp::AbstractBlockNLPModel
end

"""
    FullSpaceModel(
        m::AbstractBlockNLPModel
        x0::Union{Vector{Float64}, Nothing} = nothing,
        y0::Union{Vector{Float64}, Nothing} = nothing
    ) where {T}

Converts a `BlockNLPModel` to a `AbstractNLPModel` that can be solved with any standard NLP solver.

# Arguments

- `m::AbstractBlockNLPModel`: name of the BlockNLPModel
- `x0::Union{Vector{Float64}, Nothing}`: warm start primal solution
- `y0::Union{Vector{Float64}, Nothing}`: warm start dual solution
"""
function FullSpaceModel(
    m::AbstractBlockNLPModel;
    x0::Union{Vector{Float64},Nothing} = nothing,
    y0::Union{Vector{Float64},Nothing} = nothing,
)
    nb = m.problem_size.block_counter
    n_var = sum(m.blocks[i].meta.nvar for i = 1:nb)

    l_con::Vector{Float64} = []
    u_con::Vector{Float64} = []
    for i = 1:nb
        if m.blocks[i].meta.ncon > 0
            l_con = vcat(l_con, m.blocks[i].meta.lcon)
            u_con = vcat(u_con, m.blocks[i].meta.ucon)
        end
    end
    l_con = vcat(l_con, get_rhs_vector(m))
    u_con = vcat(u_con, get_rhs_vector(m))

    l_var::Vector{Float64} = []
    u_var::Vector{Float64} = []
    for i = 1:nb
        l_var = vcat(l_var, m.blocks[i].meta.lvar)
        u_var = vcat(u_var, m.blocks[i].meta.uvar)
    end

    isnothing(x0) && (x0 = NLPModels.get_x0(m))
    isnothing(y0) && (y0 = NLPModels.get_y0(m))

    meta = NLPModelMeta(
        n_var,
        ncon = n_constraints(m),
        nnzh = sum(m.blocks[i].meta.nnzh for i = 1:nb),
        nnzj = sum(m.blocks[i].meta.nnzj for i = 1:nb) + nnz(get_linking_matrix(m)),
        x0 = x0,
        y0 = y0,
        lvar = l_var,
        uvar = u_var,
        lcon = l_con,
        ucon = u_con,
        minimize = true,
        name = "full_space",
    )
    return FullSpaceModel(meta, Counters(), m)
end

function NLPModels.get_x0(m::AbstractBlockNLPModel)
    nb = m.problem_size.block_counter
    x0 = [NLPModels.get_x0(m.blocks[i]) for i = 1:nb]
    return x0 = reduce(vcat, x0)
end

function NLPModels.get_y0(m::AbstractBlockNLPModel)
    nb = m.problem_size.block_counter
    y0 = [NLPModels.get_y0(m.blocks[i]) for i = 1:nb]
    return y0 = reduce(vcat, push!(y0, zeros(Float64, m.problem_size.link_counter)))
end

function NLPModels.obj(nlp::FullSpaceModel, x::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    increment!(nlp, :neval_obj)
    return sum(
        obj(nlp.blocknlp.blocks[i].problem_block, x[nlp.blocknlp.blocks[i].var_idx]) for
        i = 1:nb
    )
end

function NLPModels.grad!(nlp::FullSpaceModel, x::AbstractVector, gx::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x gx
    increment!(nlp, :neval_grad)
    for i = 1:nb
        v = @view gx[nlp.blocknlp.blocks[i].var_idx]
        grad!(nlp.blocknlp.blocks[i].problem_block, x[nlp.blocknlp.blocks[i].var_idx], v)
    end
    return gx
end

function NLPModels.hess_structure!(
    nlp::FullSpaceModel,
    rows::AbstractVector{T},
    cols::AbstractVector{T},
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nnzh rows cols
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        temp_idx = idx + 1
        idx += block.meta.nnzh
        v1 = @view rows[temp_idx:idx]
        v2 = @view cols[temp_idx:idx]
        hess_structure!(block.problem_block, v1, v2)
        v1 .+= (block.var_idx[1] - 1)
        v2 .+= (block.var_idx[1] - 1)
    end
    return rows, cols
end

function NLPModels.hess_coord!(
    nlp::FullSpaceModel,
    x::AbstractVector{T},
    vals::AbstractVector{T};
    obj_weight = one(T),
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.nnzh vals
    increment!(nlp, :neval_hess)
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        temp_idx = idx + 1
        idx += block.meta.nnzh
        v = @view vals[temp_idx:idx]
        hess_coord!(block.problem_block, x[block.var_idx], v, obj_weight = obj_weight)
    end
    return vals
end

function NLPModels.hess_coord!(
    nlp::FullSpaceModel,
    x::AbstractVector{T},
    y::AbstractVector{T},
    vals::AbstractVector{T};
    obj_weight = one(T),
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.nnzh vals
    @lencheck nlp.meta.ncon y
    increment!(nlp, :neval_hess)
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon == 0
            temp_idx = idx + 1
            idx += block.meta.nnzh
            v = @view vals[temp_idx:idx]
            hess_coord!(block.problem_block, x[block.var_idx], v, obj_weight = obj_weight)
        else
            temp_idx = idx + 1
            idx += block.meta.nnzh
            v = @view vals[temp_idx:idx]
            hess_coord!(
                block.problem_block,
                x[block.var_idx],
                y[block.con_idx],
                v,
                obj_weight = obj_weight,
            )
        end
    end
    return vals
end

function NLPModels.cons!(nlp::FullSpaceModel, x::AbstractVector, cx::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.ncon cx
    increment!(nlp, :neval_cons)
    idx = 0
    for i = 1:nb
        if nlp.blocknlp.blocks[i].meta.ncon > 0
            temp_idx = idx + 1
            idx += nlp.blocknlp.blocks[i].meta.ncon
            v = @view cx[temp_idx:idx]
            cons!(
                nlp.blocknlp.blocks[i].problem_block,
                x[nlp.blocknlp.blocks[i].var_idx],
                v,
            )
        end
    end
    idx += 1
    A = get_linking_matrix_blocks(nlp.blocknlp)
    cx[idx:end] .= sum(A[j] * x[nlp.blocknlp.blocks[j].var_idx] for j = 1:nb)
    return cx
end

function NLPModels.jac_structure!(
    nlp::FullSpaceModel,
    rows::AbstractVector{T},
    cols::AbstractVector{T},
) where {T}
    @lencheck nlp.meta.nnzj rows cols
    nb = nlp.blocknlp.problem_size.block_counter
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon > 0
            temp_idx = idx + 1
            idx += block.meta.nnzj
            v1 = @view rows[temp_idx:idx]
            v2 = @view cols[temp_idx:idx]
            jac_structure!(block.problem_block, v1, v2)
            v1 .+= (block.con_idx[1] - 1)
            v2 .+= (block.var_idx[1] - 1)
        end
    end
    r, c, v = findnz(get_linking_matrix(nlp.blocknlp))
    rows[idx+1:end] .= r .+ nlp.blocknlp.problem_size.con_counter
    cols[idx+1:end] .= c
    return rows, cols
end

function NLPModels.jac_coord!(nlp::FullSpaceModel, x::AbstractVector, vals::AbstractVector)
    @lencheck nlp.meta.nvar x
    # @lencheck nlp.meta.nnzj vals # for madnlp
    increment!(nlp, :neval_jac)
    nb = nlp.blocknlp.problem_size.block_counter
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon > 0
            temp_idx = idx + 1
            idx += block.meta.nnzj
            v1 = @view vals[temp_idx:idx]
            jac_coord!(block.problem_block, x[block.var_idx], v1)
        end
    end
    r, c, v = findnz(get_linking_matrix(nlp.blocknlp))
    # vals[idx+1:end] .= v
    vals[idx+1:nlp.meta.nnzj] .= v # for madnlp
    return vals
end


text:
include("inverseCantorPairing.jl");
function inverseCantorPairingRecursive(z, q)
	pair = inverseCantorPairing(z);
	if q == 2
		return pair;
	else
		recursiveResult = inverseCantorPairingRecursive(pair[2], q-1);
		result = Array{Int64, 1}(q)
		result[1] = pair[1]
		result[2:q] = recursiveResult
		return result;
	end
end

max_stars_repo_name: driesverhees/dobble

max_issues_repo_name: driesverhees/dobble

max_forks_repo_name: driesverhees/dobble