text:
# Author: Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
# Date: March, 2020
# Copyright: Georgia Institute of Technology, 2020
#
# Test seismic modeling with invertible neural network

using LinearAlgebra, Statistics, PyPlot, Test
using NNlib, Flux
import Flux.Optimise.update!
push!(LOAD_PATH, "./src/");      using WaveOpsAVP
push!(LOAD_PATH, "./src/inn1D"); using InvertibleNetworks1D
using Random; Random.seed!(1)

## Generating synthetics

# Model
nz = 256
dz = R(5)
m_geom = ModelGeom(nz, dz)
mtrue = R(1/2000^2)*ones(R, nz, 1)
mtrue[128-30:128+30, 1] .= R(1/2200^2)
Mtrue = Model(m_geom, mtrue)

# Data parameters
nt = 1001
dt = R(0.001)
θ = range(R(0), stop = R(sqrt(3)/2), length = 11)
p = sin.(θ)*sqrt(max(mtrue...))
dat_pars = DataParams(nt, dt, p)

# Source parameters
wav = cfreq_wavelet(50, nt, dt, (R(5), R(10), R(30), R(60)))
src_pars = SrcParams(wav, dt)

# ABC parameters
abc_size = (50, 50)
abc_fact = R(1e-4)
abc_geom = ABC_Geom(nz, abc_size, abc_fact)

# Modeling operator
acc_mod = 4
F = seismod_fun(dat_pars, src_pars; abc_geom = abc_geom, acc = acc_mod)

# Generate synthetics
dat = F.eval(Mtrue)

## Setting generator (multiscale invertible net architecture)

# Input size
nc = 1
batchsize = 32

# Hyperbolic net generator
depth = 2
depthH = 20
hidden_factor = 4
logdet = true
# G = HyperbolicNet1D(depth, depthH, nz, hidden_factor, batchsize; logdet = logdet);
G = HINT1D(depth, nz, hidden_factor; logdet = logdet);

# Initialization
Z0 = randn(R, nz, nc, batchsize)
G.forward(Z0)

# Tuning factor
logdet ? (Xtuning = G.forward(randn(R, nz, nc, 1000))[1]) : (Xtuning = G.forward(randn(R, nz, nc, 1000)))
normXinf = Array{R, 1}(undef, 1000)
for i = 1:1000
    normXinf[i] = norm(Xtuning[:, :, i], Inf)
end
fact_tune = R(0.1)/mean(normXinf)

## Loss function

# Log-likelihood
mloglike_fun(m::Model; compute_grad::Bool = true) = mLogLikelihood_gauss(dat_n, m, σ2_n, F; compute_grad = compute_grad)

# Log-prior
# D = derivative_linop(Mbg.geom)
D = LinFunctionMod2Mod(m->m, m->m)
# M0 = Model(Mbg.geom, zeros(R, nz, 1))
M0 = Mbg
Dx = D.eval(Mtrue-M0)
# (gauss)
σ2_pr = R(1e-8)^2
mlogpr_fun(m::Model; compute_grad::Bool = false) = mLogPrior_gauss(m, σ2_pr; m0 = M0, D = D, compute_grad = compute_grad)
# # (L1smooth)
# ϵ2 = R(0.001)*norm(Dx.m, Inf)^2
# σ2_pr = R(1e-3)*sum(sqrt.(Dx.m.^R(2).+ϵ2))[1]
# mlogpr_fun(m::Model; compute_grad::Bool = true) = mLogPrior_L1smooth(m, σ2_pr, ϵ2; m0 = M0, D = D, compute_grad = compute_grad)

# Log-posterior
mlogpost_fun(m::Model) = mLogPost(m, mloglike_fun, mlogpr_fun; compute_grad = true)

# Pre- and post-processing for log-posterior
m0 = R(1/2000^2)*ones(R, nz, 1)
prec(X::Array{R, 3}) = Model(m_geom, m0.*(R(1).+fact_tune*reshape(X, nz, batchsize))) # contrast preconditioning
post(g::Model) = fact_tune*reshape(m0.*g.m, (nz, 1, batchsize))

# Log-posterior
σ2_dat = mean(norm(dat).^R(2))
M_pr = Model(m_geom, zeros(R, size(m0)))
σ2_pr = R(1)
D = derivative_linop(m_geom)
mlogpost(m::Model) = mlogpost_gauss(dat, m, σ2_dat, σ2_pr, F; m0 = M_pr, D = D, compute_grad = true)

# Overloading log-posterior for tensor-type input
function mlogpost(X::Array{R, 3})
    f, g_ = mlogpost(prec(X))
    return f, post(g_)
end

# # Final objective
# function loss(Z::Array{R, 3})
#     nb = size(Z, 3)
#     logdet ? ((X, lgdet) = G.forward(Z)) : (X = G.forward(Z))
#     f, ΔX = mlogpost(X)
#     ΔZ, Z = G.backward(ΔX/nb, X)
#     logdet ? (return sum(f)/nb-lgdet, ΔZ) : (return sum(f)/nb, ΔZ)
# end

# Final objective (only mlogpost)
function loss(X::Array{R, 3})
    nb = size(X, 3)
    f, ΔX = mlogpost(X)
    return sum(f)/nb, ΔX/nb
end

## Gradient test (wrt input)

_, dZ0 = loss(Z0)

ΔZ = randn(R, nz, nc, batchsize)
h = R(1e-3)
a = (R(1/12)*loss(Z0-2*h*ΔZ)[1]-R(2/3)*loss(Z0-h*ΔZ)[1]+R(2/3)*loss(Z0+h*ΔZ)[1]-R(1/12)*loss(Z0+2*h*ΔZ)[1])/h
b = dot(dZ0, ΔZ)
@show err = norm(a-b)/norm(a)

# ## Parameter algebra
#
# import Base.+, Base.-, Base.*, LinearAlgebra.dot
# import InvertibleNetworks.Parameter
# function +(θ1::Array{Parameter, 1}, θ2::Array{Parameter, 1})
#     θ = deepcopy(θ1)
#     for i = 1:length(θ1)
#         θ[i].data = θ1[i].data+θ2[i].data
#         θ[i].grad = nothing
#     end
#     return θ
# end
# function -(θ1::Array{Parameter, 1}, θ2::Array{Parameter, 1})
#     θ = deepcopy(θ1)
#     for i = 1:length(θ1)
#         θ[i].data = θ1[i].data-θ2[i].data
#         θ[i].grad = nothing
#     end
#     return θ
# end
# function *(t::R, θ::Array{Parameter, 1})
#     θt = deepcopy(θ)
#     for i = 1:length(θ)
#         θt[i].data = t*θ[i].data
#         θt[i].grad = nothing
#     end
#     return θt
# end
# function dot(θ1::Array{Parameter, 1}, θ2::Array{Parameter, 1})
#     a = R(0)
#     for i = 1:length(θ1)
#         a += dot(θ1[i].data, θ2[i].data)
#     end
#     return a
# end
#
#
# ## Gradient test (wrt parameters)
#
# # Redefine loss as a function of θ
# clear_grad!(G)
# θ_ = get_params(G)
# function loss_(θ)
#     for i = 1:length(θ)
#         θ_[i].data = θ[i].data
#         θ_[i].grad = nothing
#     end
#     return loss(Z0)[1]
# end
#
# # Computing gradient
# θ0 = Array{Parameter, 1}(undef, length(θ_))
# for i = 1:length(θ0)
#     θ0[i] = Parameter(deepcopy(θ_[i].data), nothing)
# end
# f0 = loss_(θ0)
# dθ0 = Array{Parameter, 1}(undef, length(θ0))
# for i = 1:length(θ0)
#     dθ0[i] = Parameter(deepcopy(θ_[i].grad), nothing)
# end
#
# # Finite-difference approx
# Δθ = Array{Parameter, 1}(undef, length(θ0))
# for i = 1:length(θ0)
#     Δθ[i] = Parameter(nothing, nothing)
#     Δθ[i].data = randn(R, size(dθ0[i].data))
#     Δθ[i].data = Δθ[i].data*norm(θ0[i].data, Inf)/norm(Δθ[i].data, Inf)
# end
# h = R(1e-2)
# a = (-R(1/12)*loss_(θ0+2*h*Δθ)+R(2/3)*loss_(θ0+h*Δθ)-R(2/3)*loss_(θ0-h*Δθ)+R(1/12)*loss_(θ0-2*h*Δθ))/h
#
# # Analytic
# b = dot(dθ0, Δθ)
#
# # Relative error
# @show err = norm(a-b)/norm(a)


text:
"""
    PoincareInvariants

A Julia library for the computation of Poincaré integral invariants.
"""
module PoincareInvariants

using Reexport

export AbstractPoincareInvariant, compute!, getdim, getpoints, getpointnum

"""
    AbstractPoincareInvariant

supertype of `FirstPoincareInvariant` and `SecondPoincareInvariant`.
"""
abstract type AbstractPoincareInvariant end

"""
    compute!(pinv::AbstractPoincareInvariant, args...)

computes a Poincaré invariant.
"""
function compute! end

"""
    getpoints(pinv::AbstractPoincareInvariant)

returns points on which to evaluate the phase space line or surface parameterisation
so as to `compute!` `pinv`.
"""
function getpoints end

"""
    getpointnum(pinv::AbstractPoincareInvariant)

returns number of points to sample in phase space to `compute!` `pinv`.
"""
function getpointnum end

"""
    getdim(pinv::AbstractPoincareInvariant)

returns dimension of phase space to `compute!` `pinv` in.
"""
function getdim end

include("utils.jl")

include("CanonicalSymplecticStructures.jl")
include("FirstPoincareInvariants.jl")
include("SecondPoincareInvariants/SecondPoincareInvariants.jl")

end  # module

ext: jl


text:
function optimize_iteratively(obj::EulerObject, state::EulerState; xtol=1/200, ftol=1.0e-7, maxtime=30, init_step=.1,
	alg = :GN_CRS2_LM, pop = 0, verbose = false)

	order = sortperm(obj.sizes, rev=true)

	# foreach object, in order
	orig_lb = copy(obj.lb)
	orig_ub = copy(obj.ub)
	minf = 0.0; minx = []; ret = 0.0;
	for obj_idx in order
		if (verbose) println("optimizing #$(obj_idx)") end
		# lock every other objects' bounds to the current value
		obj′ = deepcopy(obj)
		obj′.lb = copy(state); obj′.ub = copy(state)
		thisrange = obj.specs[obj_idx].statepos
		if (verbose) println("preserving range: $thisrange") end
		obj′.lb[thisrange] = orig_lb[thisrange]
		obj′.ub[thisrange] = orig_ub[thisrange]
		if (verbose) println("lb = $(obj′.lb)\nub = $(obj′.ub)") end
		# optimize
		(minf,minx,ret) = optimize(obj′, state, xtol=xtol, ftol=ftol, maxtime=maxtime,
			init_step=init_step, alg=alg, pop=pop)
		println("got $minf at $minx (returned $ret)")
		state = minx
	end
	minf, state, ret
end

"""
    optimize(obj::EulerObject, state::EulerState;
             xtol=1/200, ftol=1.0e-7, maxtime=30, init_step=.1, alg=:GN_CRS2_LM, pop=0) ->
             loss, state, result

Iteratively adjust `obj`'s parameters, starting from `state`, to minimize the
difference in the areas of the shapes therein and the data they represent.
Returned are this difference, the final state, and reason the iterations were
stopped, respectively.
"""
function optimize(obj::EulerObject, state::EulerState; xtol=1/200, ftol=1.0e-7, maxtime=30, init_step=.1,
	alg = :GN_CRS2_LM, pop = 0)
	opt = Opt(alg, length(state))
	lower_bounds!(opt, obj.lb)
	upper_bounds!(opt, obj.ub)
	xtol_abs!(opt, xtol)
	ftol_abs!(opt, ftol)
	min_objective!(opt, obj.evalfn)
	maxtime!(opt, maxtime)
	initial_step!(opt, init_step)
	population!(opt, pop)
	NLopt.optimize(opt, state)
end

function eval_euler_state(obj::EulerObject, state::EulerState; verbose::Int64=0, px::Int64=200)
	# given this state vector and the object, do the following:
	# generate a 2-D bitmap from each set
	# foreach element of the DisjointSet, calculate the size of the overlap of the bitmaps
	# compare the overlaps with the target, returning the error metric

	# draws different shapes depending on spec
	bitmaps = [make_bitmap(state[obj.specs[i].statepos], obj.sizes[i], obj.specs[i], px)
				for i in 1:length(obj.labels)]
	not_bitmaps = map(bm->.!bm, bitmaps)

	# iterate through the powerset index
	# ignore 000
	overlaps = zeros(2^length(obj.labels)-1)
	for psi = 1:length(overlaps)
		bpsi = bitstring(psi)
		compare_str = bpsi[(lastindex(bpsi) - length(obj.labels) + 1):end]
		if (verbose > 0)
			@show compare_str
		end
		# if we have 101, then take the overlap of bitmaps[1], !bitmaps[2], and bitmaps[3]
		overlap = trues(px,px)
		for bm in 1:length(compare_str)
			# TODO: create and!(a,b) that's like bitarray:956, but in-place, for speed
			overlap = overlap .& (compare_str[bm] == '1' ? bitmaps[bm] : not_bitmaps[bm])
		end
		if (verbose > 0) showbitmap(overlap) end
		overlaps[psi] = sum(overlap)
	end
	if (verbose > 0)
		@show overlaps
	end
	# TODO: use better error metric
	# compare overlaps with obj.target
	overlaps_norm = overlaps ./ sum(overlaps)
	target_norm = obj.target.counts[2:end] ./ sum(obj.target.counts[2:end])
	sum((overlaps_norm .- target_norm).^2)
end

# on a field of [0,1) x [0,1)
function make_bitmap(params, size, spec, px)
	if spec.shape == :circle
		make_bitmap_circle(params[1], params[2], size, px)
	elseif spec.shape == :square
		make_bitmap_square(params[1], params[2], size, px)
	elseif spec.shape == :triangle
		make_bitmap_triangle(params[1], params[2], size, px)
	elseif spec.shape == :rectangle
		make_bitmap_rectangle(params[1], params[2], params[3], size, px)
	else
		error("unknown shape: ", spec.shape)
	end
end

function make_bitmap_triangle(x, y, r, px) # r = radius of circle around the triangle
	bm = falses(px,px)
	pixel = 1/px

	# this shape isn't symmetric -- need to be careful that 0,0 is UPPER-LEFT corner, like Cairo

	# TODO: might be able to share some of this with the Cairo rendering

	# walk columns from left to right, doing the trig to find the number of bits
	# to set to true
	h = cos(pi/6) * r # half the side of the triangle
	j = sin(pi/6) * r # distance from center to side
	for xoffset in -h:pixel:h
		# figure out the range for this column
		p = h - abs(xoffset) # dist from the corner to this column, mirrored
		yoffset = (j - (r+j) * p/h, j) # congruent triangles give ratio of vertical offset from j

		# covert into bitmap coords
		xoffset_bm =  round(Integer, (x + xoffset) * px + 1)
		yrange_bm = round(Integer, max(1, (y + yoffset[1]) * px + 1)) : round(Integer, min(px, (y + yoffset[2]) * px + 1))
		bm[yrange_bm, xoffset_bm] .= true
	end
	bm
end

function make_bitmap_rectangle(x, y, ecc, halfsz, px)
	bm = falses(px, px) # DRY...
	pixel = 1/px

	hw = 4 ^ (ecc - 0.5) # height/width scaling (sqrt of ratio)
	# compute column and row ranges, then fast assignment
	xrange_bm = round(Integer, max(1, (x - halfsz/hw) * px + 1)) : round(Integer, min(px, (x + halfsz/hw) * px + 1))
	yrange_bm = round(Integer, max(1, (y - halfsz*hw) * px + 1)) : round(Integer, min(px, (y + halfsz*hw) * px + 1))
	bm[yrange_bm, xrange_bm] .= true # matrices are indexed Y,X...
	bm
end

function make_bitmap_square(x, y, halfsz, px)
	make_bitmap_rectangle(x, y, .5, halfsz, px)
	# bm = falses(px, px) # DRY...
	# pixel = 1/px

	# # compute column and row ranges, then fast assignment
	# xrange_bm = round(Integer, max(1, (x - halfsz) * px + 1)) : round(Integer, min(px, (x + halfsz) * px + 1))
	# yrange_bm = round(Integer, max(1, (y - halfsz) * px + 1)) : round(Integer, min(px, (y + halfsz) * px + 1))
	# bm[yrange_bm, xrange_bm] .= true # matrices are indexed Y,X...
	# bm
end

function make_bitmap_circle(x, y, r, px)
	bm = falses(px,px)
	pixel = 1/px

	# walk rows from -r to +r, doing the trig to find the number of bits
	# left and right of x to set to true
	for yoffset in -r:pixel:r
		# figure out the range for this row
		alpha = r * sqrt(1 - (yoffset/r)^2) # a big of trig
		#@show alpha
		# convert into bitmap coordinates
		yoffset_bm = round(Integer, (y + yoffset) * px + 1)
		#@show yoffset_bm
		if 1 <= yoffset_bm <= px 	# if Y is inside the box
			# convert X into bitmap coordinates, bounding
			xrange_bm = round(Integer, max(1,(x - alpha) * px + 1)) : round(Integer, min(px,(x + alpha) * px + 1))
			#@show xrange_bm
			if (length(xrange_bm) > 0)
				@inbounds bm[yoffset_bm, xrange_bm] .= true
			end
		end
	end
	bm
end

function showbitmap(bm)
	for r in 1:size(bm,1)
		for c in 1:size(bm,2)
			print(bm[r,c] ? '￭' : '·')
		end
		println("")
	end
end


text:
@recipe function plot_qs(
        qs::VersionSpacePQs,
        history::Vector{Int},
        pools::Vector{Symbol}; grid_resolution = 100, axis_overhang = 0.2)
    data = qs.lb_occ.data
    grid_range, grid_data = OCALPlots.get_grid(extrema(data)..., grid_resolution, axis_overhang)
    grid_scores = reshape(qs_score(qs, grid_data, labelmap(fill(:U, size(grid_range,2)))), grid_resolution, grid_resolution)

    kernel_gamma = MLKernels.getvalue(qs.lb_occ.kernel_fct.alpha)
    title := "gamma=$(round(kernel_gamma, digits=2)), C_LB=$(round(qs.lb_occ.C, digits=2)), C_UB=$(round(qs.ub_occ.C, digits=2)) i=$(length(history))"

    cbar := false
    @series begin
        seriestype := :contourf
        seriescolor --> :PuBu
        levels := range(0.0, maximum(grid_scores), length=10)
        grid_range, grid_range, grid_scores
    end

    markeralpha --> 0.7
    markersize --> 5

    sub_idx_history = [i ∈ history for i in 1:size(data, 2)]

    colors = Dict(:inlier => (U = :lightblue, Lin = :orange, Lout = :orange),
                  :outlier => (U = :darkblue, Lin = :orange, Lout = :orange))
    shapes = (U = :circle, Lin = :square, Lout = :utriangle)

    for (k,v) in filter(d -> first(d)!=:Lout, labelmap(pools))
        @series begin
            markeralpha := 1.0
            label := "inlier-$k"
            seriestype := :scatter
            markercolor := colors[:inlier][k]
            markershape := shapes[k]
            OCALPlots.split_2d_array(data, [i ∈ v for i in 1:size(data,2)] .& (labels .== :inlier))
        end
    end

    for (k,v) in filter(d -> first(d)!=:Lin, labelmap(pools))
        @series begin
            markeralpha := 1.0
            label := "outlier-$k"
            seriestype := :scatter
            markercolor := colors[:outlier][k]
            markershape := shapes[k]
            OCALPlots.split_2d_array(data, [i ∈ v for i in 1:size(data,2)] .& (labels .== :outlier))
        end
    end

    if !isempty(history)
        query_id = history[end]
        @series begin
            label := "current-selection"
            seriestype := :scatter
            markercolor := :red
            markeralpha := 1.0
            markershape := :star5
            markersize := 12
            OCALPlots.split_2d_array(data, 1:size(data, 2) .== query_id)
        end
    end
end

max_stars_repo_name: holtri/lama

max_issues_repo_name: holtri/lama

max_forks_repo_name: holtri/lama


text:
import CSV

# This file implements using Cullen's Evaluated Atomic Data Library 1997 for emission probabilities.
# The file relax.csv contains five columns, the atomic number, 3 sub-shell indices and a probability.
# The first two columns indicate the atom and the sub-shell ionized.  The third and fourth columns
# contain the inner and outer sub-shells involved in the X-ray transition and the fifth contains
# the probability of an ionization in column 2 producing an X-ray in columns 3 and 4.  This is what
# Perkins would call the Total yield - An ionization in column 2 will initiate a cascade of transitions.
# When the atom finally returns to the ground state, it could have emitted zero or more x-rays and
# zero or more Augers.  A K shell vacancy could relax via L3 which could relax via M5 and so on.

struct CullenEADL <: NeXLAlgorithm end

let transitions_data, xrayweights_data

    function loadAltWeights()
        @info "Loading EADL transition rate data."
        nn = (
            1, # Shell index
            (2 for _ in 1:3)...,
            (3 for _ in 1:5)...,
            (4 for _ in 1:7)...,
            (5 for _ in 1:9)...,
            (6 for _ in 1:11)...,
            (7 for _ in 1:13)...
        ) # = collect(Iterators.flatten(collect(i for _ in 1:(2i-1)) for i in 1:7 ))
        trans = Dict{Tuple{Int,Int},Int}()
        xrw = [ Dict{NTuple{3, Int}, Float64}() for _ in 1:99 ]
        for row in CSV.File(joinpath(dirname(pathof(@__MODULE__)), "..", "data", "relax.csv"))
            z, ionized, inn, out, weight = row.ZZ, row.II, row.NN, row.OO, row.PP
            if (z <= 99) && hasedge(z, inn, FFASTDB) && hasedge(z, out, FFASTDB) && (nn[inn] != nn[out])
                # There seems to be a problem with the L2-M1 and L3-M1 weights which I resolve with this ad-hoc fix.
                if (out == 5) && ((inn == 4) || (inn == 3))
                    if z >= 29
                        weight *= max(0.1, 0.1 + ((0.9 * (z - 29.0)) / (79.0 - 29.0)))
                    else
                        weight *= max(0.1, 0.2 - ((0.1 * (z - 22.0)) / (29.0 - 22.0)))
                    end
                end
                xrw[z][(ionized, inn, out)] = weight
                trans[(inn, out)] = get(trans, (inn, out), 0) + 1
            end
        end
        # Add these which aren't in Cullen (The weights are from the Practical Handbook of Spectroscopy pg 834)
        for x in ((3, 1, 1, 2, 1.06e-4), (4, 1, 1, 2, 3.2e-4), (5, 1, 1, 3, 7.1e-4))
            xrw[x[1]] = Dict( (x[2], x[3], x[4]) => x[5] )
        end
        # Now compute some of the normalized weights
        xrw2 = [ Dict{NTuple{3, Int}, NTuple{4,Float64}}() for _ in 1:99 ]
        for (z, id) in enumerate(xrw)
            sum_s = Dict{Int, Float64}() # Normalize over all in shell
            sum_ss = Dict{Int, Float64}() # Normalize over all in sub-shell
            max_s = Dict{Int, Float64}() # Normalize most intense in sub-shell to 1.0
            for ((ionized, inn, _), w) in id
                # Assume large overvoltage so ignore differences in ionization rates
                if inn==ionized
                    icx = ionizationcrosssection(z, inn, 4.0*edgeenergy(z, inn, FFASTDB), Bote2009)
                    sum_ss[inn] = get(sum_ss, inn, 0.0) + w
                    sum_s[nn[inn]] = get(sum_s, nn[inn], 0.0) + icx*w
                    max_s[nn[inn]] = max(icx*w, get(max_s, nn[inn], 0.0))
                end
            end
            for ((ionized, inn, out), w) in id
                icx = ionizationcrosssection(z, inn, 4.0*edgeenergy(z, inn, FFASTDB), Bote2009)
                if inn == ionized
                    # Yield, norm[sub_shell], norm[shell], norm[max=1]
                    xrw2[z][(ionized, inn, out)] = (w, w/sum_ss[inn], (icx*w)/sum_s[nn[inn]], (icx*w)/max_s[nn[inn]])
                else
                    xrw2[z][(ionized, inn, out)] = (w, 0.0, 0.0, 0.0 )
                end
            end
        end
        return ( trans, xrw2 )
    end
    transitions_data, xrayweights_data = loadAltWeights()
    global transitions()::Dict{Tuple{Int,Int},Int} = transitions_data
    global xrayweights()::Vector{Dict{NTuple{3, Int}, NTuple{4, Float64}}} = xrayweights_data
end

fluorescenceyield(z::Int, ionized::Int, inner::Int, outer::Int, ::Type{CullenEADL})::Float64 =
    get(xrayweights()[z], (ionized, inner, outer), (0.0, 0.0, 0.0, 0.0))[1]
fluorescenceyield(z::Int, inner::Int, outer::Int, ::Type{CullenEADL})::Float64 =
    get(xrayweights()[z], (inner, inner, outer), (0.0, 0.0, 0.0, 0.0))[1]
function fluorescenceyield(z::Int, inner::Int, ::Type{CullenEADL})
    sum=0.0
    for ((ion,inn,_), ws) in xrayweights()[z]
        if (ion==inn) && (inner==inn)
            sum+=ws[1]
        end
    end
    return sum
end

isAvailable(z::Int, inner::Int, outer::Int, ::Type{CullenEADL}) =
    fluorescenceyield(z, inner, inner, outer, CullenEADL) > 0.0


text:
using BinaryProvider # requires BinaryProvider 0.3.0 or later

# Example taken from
# https://github.com/JuliaIO/ImageMagick.jl/blob/master/deps/build.jl
dependencies = [
    "build_Zlib.v1.2.11.jl",
    "build_GEOS.v3.7.2.jl",
    "build_SQLite.v3.28.0.jl",
    "build_PROJ.v6.1.0.jl",
    # "build_LibCURL.v7.64.1.jl"
]

for elem in dependencies
    # it's a bit faster to run the build in an anonymous module instead of
    # starting a new julia process
    m = Module(:__anon__)
    Core.include(m, (joinpath(@__DIR__, elem)))
end

# Parse some basic command-line arguments
const verbose = "--verbose" in ARGS
const prefix = Prefix(get([a for a in ARGS if a != "--verbose"], 1, joinpath(@__DIR__, "usr")))
products = [
    LibraryProduct(prefix, ["libgdal"], :libgdal),
    ExecutableProduct(prefix, "gdal_contour", :gdal_contour_path),
    ExecutableProduct(prefix, "gdal_grid", :gdal_grid_path),
    ExecutableProduct(prefix, "gdal_rasterize", :gdal_rasterize_path),
    ExecutableProduct(prefix, "gdal_translate", :gdal_translate_path),
    ExecutableProduct(prefix, "gdaladdo", :gdaladdo_path),
    ExecutableProduct(prefix, "gdalbuildvrt", :gdalbuildvrt_path),
    ExecutableProduct(prefix, "gdaldem", :gdaldem_path),
    ExecutableProduct(prefix, "gdalinfo", :gdalinfo_path),
    ExecutableProduct(prefix, "gdallocationinfo", :gdallocationinfo_path),
    ExecutableProduct(prefix, "gdalmanage", :gdalmanage_path),
    ExecutableProduct(prefix, "gdalsrsinfo", :gdalsrsinfo_path),
    ExecutableProduct(prefix, "gdaltindex", :gdaltindex_path),
    ExecutableProduct(prefix, "gdaltransform", :gdaltransform_path),
    ExecutableProduct(prefix, "gdalwarp", :gdalwarp_path),
    ExecutableProduct(prefix, "nearblack", :nearblack_path),
    ExecutableProduct(prefix, "ogr2ogr", :ogr2ogr_path),
    ExecutableProduct(prefix, "ogrinfo", :ogrinfo_path),
    ExecutableProduct(prefix, "ogrlineref", :ogrlineref_path),
    ExecutableProduct(prefix, "ogrtindex", :ogrtindex_path),
]

# Download binaries from hosted location
bin_prefix = "https://github.com/JuliaGeo/GDALBuilder/releases/download/v3.0.0-3"

# Listing of files generated by BinaryBuilder:
download_info = Dict(
    Linux(:aarch64, libc=:glibc) => ("$bin_prefix/GDAL.v3.0.0.aarch64-linux-gnu.tar.gz", "ebe236883d8480fe7482fd1d27c2535c41e808abe1e5cf8c4550be538c3b77e2"),
    Linux(:aarch64, libc=:musl) => ("$bin_prefix/GDAL.v3.0.0.aarch64-linux-musl.tar.gz", "0425ae5159cd02c41a526fcb8469b2676bc2f5b4120bea014c8383e291d0c4eb"),
    Linux(:armv7l, libc=:glibc, call_abi=:eabihf) => ("$bin_prefix/GDAL.v3.0.0.arm-linux-gnueabihf.tar.gz", "8e9af6c82436788805d59b7c62b26f319fe4460898186939bda3d2931dff28d3"),
    Linux(:armv7l, libc=:musl, call_abi=:eabihf) => ("$bin_prefix/GDAL.v3.0.0.arm-linux-musleabihf.tar.gz", "08c659da2d0bdd9b5c67458e9506ac2ef92ee5c6143f45f0650ade3d6e7f6722"),
    Linux(:i686, libc=:glibc) => ("$bin_prefix/GDAL.v3.0.0.i686-linux-gnu.tar.gz", "99bd95d5de4dcd7ac9248654c7e53d9b7cb65aaf47352d923ba5da2c41336017"),
    Linux(:i686, libc=:musl) => ("$bin_prefix/GDAL.v3.0.0.i686-linux-musl.tar.gz", "7cbd4de983f1b6c470cfcf208febaaa5467e91f46644fac034ad70e0603e42d1"),
    # removed compiler_abi as suggested in https://github.com/JuliaPackaging/BinaryBuilder.jl/issues/407#issuecomment-473688254
    # such that GCC4 platforms will also pick up this GCC7 build, ref https://github.com/JuliaPackaging/BinaryBuilder.jl/issues/407
    Windows(:i686) => ("$bin_prefix/GDAL.v3.0.0.i686-w64-mingw32-gcc7.tar.gz", "8d056a9a4781cc13e572b689c581ef9ab48d5f397b2ca3882466e168aabb5b45"),
    Linux(:powerpc64le, libc=:glibc) => ("$bin_prefix/GDAL.v3.0.0.powerpc64le-linux-gnu.tar.gz", "29772014a34d5300c950399f6d07dcc960feff6671a30fc21a834c7482d4e359"),
    MacOS(:x86_64) => ("$bin_prefix/GDAL.v3.0.0.x86_64-apple-darwin14.tar.gz", "df0dd3dfec28cf98b7a8d4fdf9f770599bd3bde2fe223f33e903aa052a4f0ada"),
    Linux(:x86_64, libc=:glibc) => ("$bin_prefix/GDAL.v3.0.0.x86_64-linux-gnu.tar.gz", "22e576ebb9c45a8571869f5ebaca0ab1c6f443aa281c0fcdf5176c75d9912742"),
    Linux(:x86_64, libc=:musl) => ("$bin_prefix/GDAL.v3.0.0.x86_64-linux-musl.tar.gz", "f2ab4b8ff55621f22963f6749c210876b5be14c3cee724ece7dfc73ea9b3a14e"),
    FreeBSD(:x86_64) => ("$bin_prefix/GDAL.v3.0.0.x86_64-unknown-freebsd11.1.tar.gz", "0b4b817a13fccb3e65b32df535679163aa0d061077542ceda00764d34c96159e"),
    # removed compiler_abi as suggested in https://github.com/JuliaPackaging/BinaryBuilder.jl/issues/407#issuecomment-473688254
    # such that GCC4 platforms will also pick up this GCC7 build, ref https://github.com/JuliaPackaging/BinaryBuilder.jl/issues/407
    Windows(:x86_64) => ("$bin_prefix/GDAL.v3.0.0.x86_64-w64-mingw32-gcc7.tar.gz", "aa1b84e77de61eda6f3e9f569ff8e7bab3b99d38fa2edddb9069ef82ba3d5256"),
)

# Install unsatisfied or updated dependencies:
unsatisfied = any(!satisfied(p; verbose=verbose) for p in products)
dl_info = choose_download(download_info, platform_key_abi())
if dl_info === nothing && unsatisfied
    # If we don't have a compatible .tar.gz to download, complain.
    # Alternatively, you could attempt to install from a separate provider,
    # build from source or something even more ambitious here.
    error("Your platform (\"$(Sys.MACHINE)\", parsed as \"$(triplet(platform_key_abi()))\") is not supported by this package!")
end

# If we have a download, and we are unsatisfied (or the version we're
# trying to install is not itself installed) then load it up!
if unsatisfied || !isinstalled(dl_info...; prefix=prefix)
    # Download and install binaries
    install(dl_info...; prefix=prefix, force=true, verbose=verbose)
end

# Write out a deps.jl file that will contain mappings for our products
write_deps_file(joinpath(@__DIR__, "deps_gdal.jl"), products, verbose=verbose)

function include_deps(name)
    """
    module $name
        import Libdl
        path = joinpath(@__DIR__, $(repr(string("deps_", name, ".jl"))))
        isfile(path) || error("$name wasn't build correctly. Please run Pkg.build(\\\"GDAL\\\")")
        include(path)
    end
    using .$name
    """
end

open("deps.jl", "w") do io
    for dep in (:zlib, :geos, :sqlite, :proj, #=:curl,=# :gdal)
        println(io, include_deps(dep))
    end
    println(io, """
    const libgdal = gdal.libgdal
    const libproj = proj.libproj
    function check_deps()
        zlib.check_deps()
        geos.check_deps()
        sqlite.check_deps()
        proj.check_deps()
        # curl.check_deps()
        gdal.check_deps()
    end
    """)
end


text:
module HDF5saveUtils
using HDF5

export saveManualModif,saveMaskbyName,saveMaskBeforeVisualization

"""
saving to the HDF5 data from manual modifications
requires the manual modification array to be called manualModif
fid- object referencing HDF5 database
"""
function saveManualModif(fid,patienGroupName , mainScrollDat)
    manualModif= filter((it)->it.name=="manualModif" ,mainScrollDat.dataToScroll)[1].dat
    group = fid[patienGroupName]
    if(!haskey(group, "manualModif"))
        write(group, "manualModif", manualModif)
        dset = group["manualModif"]
        write_attribute(dset, "dataType", "manualModif")
        write_attribute(dset, "min", minimum(manualModif))
        write_attribute(dset, "max", max(maximum(manualModif), 1))

    else
        delete_object(group, "manualModif")
        write(group, "manualModif", manualModif)
        dset = group["manualModif"]
        write_attribute(dset, "dataType", "manualModif")
        write_attribute(dset, "min", minimum(manualModif))
        write_attribute(dset, "max", max(maximum(manualModif), 1))

    end#if

end#saveManualModif

"""
sava data that is ready for visualization into HDF5

fid - the reference object to hdf5
patienGroupName - name of the group (by hdf5 terminology) in which we want to put our data
mainScrollDat - the object holding all required data for displaying including arrays data
name - name that will identify this array
dataType - wheather it is CT scan ... (described in visualizationFromHdf5 file)
"""
function saveMaskbyName(fid,patienGroupName , mainScrollDat, name, dataType)
    arr= filter((it)->it.name==name ,mainScrollDat.dataToScroll)[1].dat
    group = fid[patienGroupName]
    if(!haskey(group, name))
        write(group, name, arr)
        dset = group[name]
        write_attribute(dset, "dataType", dataType)
        write_attribute(dset, "min", minimum(arr))
        write_attribute(dset, "max", max(maximum(arr), 1))

    else
        delete_object(group, name)
        write(group, name, arr)
        dset = group[name]
        write_attribute(dset, "dataType", dataType)
        write_attribute(dset, "min", minimum(arr))
        write_attribute(dset, "max", max(maximum(arr), 1))

    end#if
end#saveMaskbyName

"""
sava data that is not yet ready for visualization into HDF5

fid - the reference object to hdf5
patienGroupName - name of the group (by hdf5 terminology) in which we want to put our data
arr - 3D array with data we want to save
name - name that will identify this array
dataType - wheather it is CT scan ... (described in visualizationFromHdf5 file)
"""
function saveMaskBeforeVisualization(fid,patienGroupName , arr, name, dataType)

    group = fid[patienGroupName]
    if(!haskey(group, name))
        write(group, name, arr)
        dset = group[name]
        write_attribute(dset, "dataType", dataType)
        write_attribute(dset, "min", minimum(arr))
        write_attribute(dset, "max", max(maximum(arr), 1))

    else
        delete_object(group, name)
        write(group, name, arr)
        dset = group[name]
        write_attribute(dset, "dataType", dataType)
        write_attribute(dset, "min", minimum(arr))
        write_attribute(dset, "max", max(maximum(arr), 1))

    end#if
end#saveMaskbyName

end #HDF5saveUtils


text:
# XXX create optimization model after anyModel has been initialized
"""
Create constraints of the model's underlying optimization problem.
"""
function createOptModel!(anyM::anyModel)

	# <editor-fold desc="create technology related variables and constraints"

	techIdx_arr = collect(keys(anyM.parts.tech))
    parDef_dic = defineParameter(anyM.options,anyM.report)

    # XXX get dimension of expansion and capacity variables and mapping of capacity constraints
    tsYear_dic = Dict(zip(anyM.supTs.step,collect(0:anyM.options.shortExp:(length(anyM.supTs.step)-1)*anyM.options.shortExp)))
    prepVar_dic = Dict{Int,Dict{Symbol,NamedTuple}}()
    prepareTechs!(techIdx_arr,prepVar_dic,tsYear_dic,anyM)
    if any(getindex.(anyM.report,1) .== 3) print(getElapsed(anyM.options.startTime)); errorTest(anyM.report,anyM.options) end

    # remove technologies without any potential capacity variables
    techIdx_arr = collect(keys(prepVar_dic))
    foreach(x -> delete!(anyM.parts.tech, x),setdiff(collect(keys(anyM.parts.tech)),techIdx_arr))

    # XXX create all technology related elements

    # creates dictionary that assigns combination of superordinate dispatch timestep and dispatch level to dispatch timesteps
    allLvlTsDis_arr = unique(getfield.(values(anyM.cInfo),:tsDis))
    ts_dic = Dict((x[1], x[2]) => anyM.sets[:Ts].nodes[x[1]].lvl == x[2] ? [x[1]] : getDescendants(x[1],anyM.sets[:Ts],false,x[2]) for x in Iterators.product(anyM.supTs.step,allLvlTsDis_arr))

    # creates dictionary that assigns combination of expansion region and dispatch level to dispatch region
    allLvlR_arr = union(getindex.(getfield.(getfield.(values(anyM.parts.tech),:balLvl),:exp),2),map(x -> x.rDis,values(anyM.cInfo)))

    allRExp_arr = union([getfield.(getNodesLvl(anyM.sets[:R],x),:idx) for x in allLvlR_arr]...)
    r_dic = Dict((x[1], x[2]) => (anyM.sets[:R].nodes[x[1]].lvl <= x[2] ? getDescendants(x[1], anyM.sets[:R],false,x[2]) : getAncestors(x[1],anyM.sets[:R],:int,x[2])[end]) |> (z -> typeof(z) <: Array ? z : [z]) for x in Iterators.product(allRExp_arr,allLvlR_arr))

    produceMessage(anyM.options,anyM.report, 3," - Determined dimension of expansion and capacity variables for technologies")

    # constraints for technologies are prepared in threaded loop and stored in an array of dictionaries
	techCnsDic_arr = Array{Dict{Symbol,cnsCont}}(undef,length(techIdx_arr))
	tech_itr = collect(enumerate(techIdx_arr))

	@threads for (idx,t) in tech_itr
		techCnsDic_arr[idx] = createTech!(t,anyM.parts.tech[t],prepVar_dic[t],copy(parDef_dic),ts_dic,r_dic,anyM)
	end

    # loops over array of dictionary with constraint container for each technology to create actual jump constraints
    for (idx,cnsDic) in enumerate(techCnsDic_arr), cnsSym in keys(cnsDic)
        anyM.parts.tech[techIdx_arr[idx]].cns[cnsSym] = createCns(cnsDic[cnsSym],anyM.optModel)
    end
    produceMessage(anyM.options,anyM.report, 1," - Created variables and constraints for all technologies")

	# </editor-fold>

	# <editor-fold desc="create exchange related variables and constraints"

	prepExc_dic = Dict{Symbol,NamedTuple}()
	partExc = anyM.parts.exc
	partLim = anyM.parts.lim

	# obtain dimensions of expansion variables for exchange
	potExc_df = prepareExcExpansion!(partExc,partLim,prepExc_dic,tsYear_dic,anyM)

	# obtain capacity dimensions solely based on expansion variables
	prepareCapacity!(partExc,prepExc_dic,prepExc_dic[:expExc].var,:capaExc,anyM)
	addResidualCapaExc!(partExc,prepExc_dic,potExc_df,anyM)

	if !all(map(x -> isempty(x),values(prepExc_dic[:capaExc])))
		# create expansion and capacity variables
		createExpCap!(partExc,prepExc_dic,anyM)
		# create capacity constraint
		createCapaExcCns!(partExc,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all variables and constraints related to expansion and capacity for exchange")
		# create dispatch related variables
		createExcVar!(partExc,ts_dic,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all dispatch variables for exchange")
		# create capacity restrictions
		createRestrExc!(ts_dic,partExc,anyM)
		produceMessage(anyM.options,anyM.report, 2," - Created all capacity restrictions for exchange")
		produceMessage(anyM.options,anyM.report, 1," - Created variables and constraints for exchange")
	end

	# </editor-fold>

	createTradeVarCns!(anyM.parts.trd,anyM)
	createEnergyBal!(techIdx_arr,anyM)
	createLimitCns!(techIdx_arr,anyM.parts.lim,anyM)

	produceMessage(anyM.options,anyM.report, 1," - Completed model creation")
end


text:
# This file is auto-generated by AWSMetadata.jl
using AWS
using AWS.AWSServices: personalize_runtime
using AWS.Compat
using AWS.UUIDs

"""
    get_personalized_ranking(campaign_arn, input_list, user_id)
    get_personalized_ranking(campaign_arn, input_list, user_id, params::Dict{String,<:Any})

Re-ranks a list of recommended items for the given user. The first item in the list is
deemed the most likely item to be of interest to the user.  The solution backing the
campaign must have been created using a recipe of type PERSONALIZED_RANKING.

# Arguments
- `campaign_arn`: The Amazon Resource Name (ARN) of the campaign to use for generating the
  personalized ranking.
- `input_list`: A list of items (by itemId) to rank. If an item was not included in the
  training dataset, the item is appended to the end of the reranked list. The maximum is 500.
- `user_id`: The user for which you want the campaign to provide a personalized ranking.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"context"`: The contextual metadata to use when getting recommendations. Contextual
  metadata includes any interaction information that might be relevant when getting a user's
  recommendations, such as the user's current location or device type.
- `"filterArn"`: The Amazon Resource Name (ARN) of a filter you created to include items or
  exclude items from recommendations for a given user. For more information, see Filtering
  Recommendations.
- `"filterValues"`: The values to use when filtering recommendations. For each placeholder
  parameter in your filter expression, provide the parameter name (in matching case) as a key
  and the filter value(s) as the corresponding value. Separate multiple values for one
  parameter with a comma.  For filter expressions that use an INCLUDE element to include
  items, you must provide values for all parameters that are defined in the expression. For
  filters with expressions that use an EXCLUDE element to exclude items, you can omit the
  filter-values.In this case, Amazon Personalize doesn't use that portion of the expression
  to filter recommendations. For more information, see Filtering Recommendations.
"""
get_personalized_ranking(campaignArn, inputList, userId; aws_config::AbstractAWSConfig=global_aws_config()) = personalize_runtime("POST", "/personalize-ranking", Dict{String, Any}("campaignArn"=>campaignArn, "inputList"=>inputList, "userId"=>userId); aws_config=aws_config)
get_personalized_ranking(campaignArn, inputList, userId, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = personalize_runtime("POST", "/personalize-ranking", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("campaignArn"=>campaignArn, "inputList"=>inputList, "userId"=>userId), params)); aws_config=aws_config)

"""
    get_recommendations(campaign_arn)
    get_recommendations(campaign_arn, params::Dict{String,<:Any})

Returns a list of recommended items. The required input depends on the recipe type used to
create the solution backing the campaign, as follows:   RELATED_ITEMS - itemId required,
userId not used   USER_PERSONALIZATION - itemId optional, userId required    Campaigns that
are backed by a solution created using a recipe of type PERSONALIZED_RANKING use the API.

# Arguments
- `campaign_arn`: The Amazon Resource Name (ARN) of the campaign to use for getting
  recommendations.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"context"`: The contextual metadata to use when getting recommendations. Contextual
  metadata includes any interaction information that might be relevant when getting a user's
  recommendations, such as the user's current location or device type.
- `"filterArn"`: The ARN of the filter to apply to the returned recommendations. For more
  information, see Filtering Recommendations. When using this parameter, be sure the filter
  resource is ACTIVE.
- `"filterValues"`: The values to use when filtering recommendations. For each placeholder
  parameter in your filter expression, provide the parameter name (in matching case) as a key
  and the filter value(s) as the corresponding value. Separate multiple values for one
  parameter with a comma.  For filter expressions that use an INCLUDE element to include
  items, you must provide values for all parameters that are defined in the expression. For
  filters with expressions that use an EXCLUDE element to exclude items, you can omit the
  filter-values.In this case, Amazon Personalize doesn't use that portion of the expression
  to filter recommendations. For more information, see Filtering Recommendations.
- `"itemId"`: The item ID to provide recommendations for. Required for RELATED_ITEMS recipe
  type.
- `"numResults"`: The number of results to return. The default is 25. The maximum is 500.
- `"userId"`: The user ID to provide recommendations for. Required for USER_PERSONALIZATION
  recipe type.
"""
get_recommendations(campaignArn; aws_config::AbstractAWSConfig=global_aws_config()) = personalize_runtime("POST", "/recommendations", Dict{String, Any}("campaignArn"=>campaignArn); aws_config=aws_config)
get_recommendations(campaignArn, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = personalize_runtime("POST", "/recommendations", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("campaignArn"=>campaignArn), params)); aws_config=aws_config)


text:
# A group of amphipods notice your fancy submarine and flag you down. "With
# such an impressive shell," one amphipod says, "surely you can help us with a
# question that has stumped our best scientists."
#
# They go on to explain that a group of timid, stubborn amphipods live in a
# nearby burrow. Four types of amphipods live there: Amber (A), Bronze (B),
# Copper (C), and Desert (D). They live in a burrow that consists of a hallway
# and four side rooms. The side rooms are initially full of amphipods, and the
# hallway is initially empty.
#
# They give you a diagram of the situation (your puzzle input), including
# locations of each amphipod (A, B, C, or D, each of which is occupying an
# otherwise open space), walls (#), and open space (.).
#
# For example:
#
# #############
# #...........#
# ###B#C#B#D###
#   #A#D#C#A#
#   #########
#
# The amphipods would like a method to organize every amphipod into side rooms
# so that each side room contains one type of amphipod and the types are sorted
# A-D going left to right, like this:
#
# #############
# #...........#
# ###A#B#C#D###
#   #A#B#C#D#
#   #########
#
# Amphipods can move up, down, left, or right so long as they are moving into
# an unoccupied open space. Each type of amphipod requires a different amount
# of energy to move one step: Amber amphipods require 1 energy per step, Bronze
# amphipods require 10 energy, Copper amphipods require 100, and Desert ones
# require 1000. The amphipods would like you to find a way to organize the
# amphipods that requires the least total energy.
#
# However, because they are timid and stubborn, the amphipods have some extra
# rules:
#
# - Amphipods will never stop on the space immediately outside any room. They
#   can move into that space so long as they immediately continue moving.
#   (Specifically, this refers to the four open spaces in the hallway that are
#   directly above an amphipod starting position.)
# - Amphipods will never move from the hallway into a room unless that room is
#   their destination room and that room contains no amphipods which do not
#   also have that room as their own destination. If an amphipod's starting
#   room is not its destination room, it can stay in that room until it leaves
#   the room.  (For example, an Amber amphipod will not move from the hallway
#   into the right three rooms, and will only move into the leftmost room if
#   that room is empty or if it only contains other Amber amphipods.)
# - Once an amphipod stops moving in the hallway, it will stay in that spot
#   until it can move into a room. (That is, once any amphipod starts moving,
#   any other amphipods currently in the hallway are locked in place and will
#   not move again until they can move fully into a room.)
#
# In the above example, the amphipods can be organized using a minimum of 12521
# energy. One way to do this is shown below.
#
# Starting configuration:
#
# #############
# #...........#
# ###B#C#B#D###
#   #A#D#C#A#
#   #########
#
# One Bronze amphipod moves into the hallway, taking 4 steps and using 40
# energy:
#
# #############
# #...B.......#
# ###B#C#.#D###
#   #A#D#C#A#
#   #########
#
# The only Copper amphipod not in its side room moves there, taking 4 steps and
# using 400 energy:
#
# #############
# #...B.......#
# ###B#.#C#D###
#   #A#D#C#A#
#   #########
#
# A Desert amphipod moves out of the way, taking 3 steps and using 3000 energy,
# and then the Bronze amphipod takes its place, taking 3 steps and using 30
# energy:
#
# #############
# #.....D.....#
# ###B#.#C#D###
#   #A#B#C#A#
#   #########
#
# The leftmost Bronze amphipod moves to its room using 40 energy:
#
# #############
# #.....D.....#
# ###.#B#C#D###
#   #A#B#C#A#
#   #########
#
# Both amphipods in the rightmost room move into the hallway, using 2003 energy
# in total:
#
# #############
# #.....D.D.A.#
# ###.#B#C#.###
#   #A#B#C#.#
#   #########
#
# Both Desert amphipods move into the rightmost room using 7000 energy:
#
# #############
# #.........A.#
# ###.#B#C#D###
#   #A#B#C#D#
#   #########
#
# Finally, the last Amber amphipod moves into its room, using 8 energy:
#
# #############
# #...........#
# ###A#B#C#D###
#   #A#B#C#D#
#   #########
#
# What is the least energy required to organize the amphipods?

# using Pkg; Pkg.add("DataStructures")
using DataStructures

# cost of moving each amphipod
const amphipodcost = [1, 10, 100, 1000]

# position of the rooms from the hallway
const roomposl = 3:2:9

# win condition
const win = [1 1; 2 2; 3 3; 4 4]

# read input
rooms = zeros(Int8, 4, 2)
roomposition = 1
foreach(eachline("input.txt")) do l
  m = match(r"([A-D])#([A-D])#([A-D])#([A-D])", l)
  isnothing(m) && return

  rooms[:, roomposition] = map(c -> c[1] - 'A' + 1, m)
  global roomposition += 1
end

# Attempt to move amphipods home - mutates rooms and hallway, returns cost
function movehome!(rooms, hallway)
  cost = 0

  # keep trying to move amphipods home until we can't
  while true
    progress = false

    # for each amphipod in the hallway, see if we can move them home
    for (i,c) in enumerate(hallway)
      # check for an amphipod, and see if their room is available
      c == 0 && continue
      rooms[c,:] != [0,0] && rooms[c,:] != [0,c] && continue

      # assume amphipod is coming from the left of the room
      roompos = roomposl[c]
      hallwaypos = (i+1):roompos
      if i > roompos
        # amphipod is coming from the right
        hallwaypos = roompos:(i-1)
      end

      # check if the hallway is clear to move through
      any(!=(0), hallway[hallwaypos]) && continue

      # cost is the length of the hallway, plus 1 to represent moving into the
      # room, times the amphipod's movement energy.
      cost += (length(hallwaypos) + 1) * amphipodcost[c]

      # update hallway and room
      hallway[i] = 0
      if rooms[c,2] == 0
        # room was empty so had to move one additional space
        rooms[c,2] = c
        cost += amphipodcost[c]
      else
        rooms[c,1] = c
      end

      # makin' progress
      progress = true
    end

    !progress && break
  end

  return cost
end

# adds a state to the queue, or, if such a state already exists, update the
# cost if the new cost is less
function enqueueorupdate!(queue, state, cost)
  if haskey(queue, state)
    if cost < queue[state]
      queue[state] = cost
    end
  else
    enqueue!(queue, state, cost)
  end
end

# Dijkstra's, sorta... if you consider each "neighbor" to be all possible moves
# from the current move.
State = Tuple{Matrix{Int8}, Vector{Int8}}
queue = PriorityQueue{State, Int}((rooms, zeros(Int8, 11)) => 0)
visited = Set{State}()
while length(queue) > 0
  total_cost = peek(queue).second
  state = dequeue!(queue)
  push!(visited, state)

  local (rooms,hallway) = copy.(state)

  # are we winning yet?
  if rooms == win
    println("Result: ", total_cost)
    break
  end

  # first try moving amphipods home
  total_cost += movehome!(rooms, hallway)

  # if we've reached a win condition, push it into the queue - there's nothing
  # more to do, but there may be other ways to win that are shorter
  if rooms == win
    enqueueorupdate!(queue, (rooms, hallway), total_cost)
    continue
  end

  # now find all possible moves from the current state
  for (i,room) in enumerate(eachrow(rooms))
    # if the room is empty or already has the right amphipods, skip it
    (room == [0,0] || room == [0,i] || room == [i,i]) && continue

    # find the amphipod that's moving
    newrooms = copy(rooms)
    movingamphipod = room[1]
    movement = 0
    if movingamphipod == 0
      movingamphipod = room[2]
      movement += 1
      newrooms[i,2] = 0
    else
      newrooms[i,1] = 0
    end

    # find where they can move to
    hallwayexit = roomposl[i]
    for pos in [1,2,4,6,8,10,11]
      # positions we move through in the hallway
      hallwaypos = pos <= hallwayexit ? (pos:hallwayexit) : (hallwayexit:pos)

      # check if hallway is empty
      any(!=(0), hallway[hallwaypos]) && continue

      # cost is number of moves to exit the room, plus number of hallway spaces
      cost = total_cost + (length(hallwaypos) + movement) * amphipodcost[movingamphipod]

      # update hallway
      newhallway = copy(hallway)
      newhallway[pos] = movingamphipod

      # update queue
      newstate = (newrooms, newhallway)
      !in(newstate, visited) && enqueueorupdate!(queue, newstate, cost)
    end
  end
end


text:
get_params(state::EulerState, spec::EulerSpec) = tuple(state[spec.statepos]...)

reds =   [1,  0, 0, .7, 0, .7, .2, .7]
greens = [0, .8, 0, 0, .7, .7, .2, .7]
blues =  [0,  0, 1, .7, .7, 0, .2, .7]

function setup_shape!(cr, spec, state, size, px)
	if spec.shape == :circle
		x,y = get_params(state, spec)
		move_to(cr, x*px+size*px, y*px)
		arc(cr, x*px, y*px, size*px, 0, 2*pi)
	elseif spec.shape == :rectangle
		x,y,ecc = get_params(state, spec)
		halfside = size
		hw = 4 ^ (ecc - 0.5)
		rectangle(cr, (x - halfside/hw)*px, (y - halfside*hw)*px, 2halfside*px/hw, 2halfside*px*hw)
	elseif spec.shape == :square
		x,y = get_params(state, spec)
		halfside = size
		rectangle(cr, (x - halfside)*px, (y - halfside)*px, 2halfside*px, 2halfside*px)
	elseif spec.shape == :triangle
		x,y = get_params(state, spec)
		# convert into our triangle notation
		r = size
		h = cos(pi/6) * r # half the side of the triangle
		j = sin(pi/6) * r # distance from center to side
		# start at lower-left, then top, then lower-right, then close
		move_to(cr, (x-h)*px, (y+j)*px)
		line_to(cr, x*px, (y-r)*px)
		line_to(cr, (x+h)*px, (y+j)*px)
		close_path(cr)
	end
	x,y
end

"""
    render(fn, obj::EulerObject, state::EulerState; verbose=0, px=500.0)

Draw `obj` to file `fn` using `state`.  `fn` should end in ".svg" or ".pdf".
"""
function render(fn, obj::EulerObject, state::EulerState; verbose=0, px=500.0)
	if isa(fn, AbstractString)
		bn,ext = splitext(fn)
		if ext == ".svg"
			c = CairoSVGSurface(fn, px, px);
		elseif ext ==".pdf"
			c = CairoPDFSurface(fn, px, px);
		else
			throw(ArgumentError("Unable to render to output with extension $(ext)"))
		end
	else
		c = CairoSVGSurface(fn, px, px);
	end
	cr = CairoContext(c);
	set_line_width(cr, px / 200.0);
	select_font_face(cr, "Sans", Cairo.FONT_SLANT_NORMAL,
	                     Cairo.FONT_WEIGHT_NORMAL);
	set_font_size(cr, px / 25.0);

	@assert length(obj.labels) <= length(reds)

	# foreach object, set the color, calc the center, then draw the appropriate object
	for i = 1:length(obj.labels)
		if verbose>0 @show obj.specs[i] end

		set_source_rgba(cr, reds[i], greens[i], blues[i], .3)
		setup_shape!(cr, obj.specs[i], state, obj.sizes[i], px)
		fill(cr)
	end
	for i = 1:length(obj.labels) # put edges and labels on top
		set_source_rgba(cr, reds[i], greens[i], blues[i], 1)

		# could use set_operator(CAIRO_OPERATOR_DEST_OVER) to draw underneath existing stuff
		x,y = setup_shape!(cr, obj.specs[i], state, obj.sizes[i], px)
		stroke(cr);

		extents = text_extents(cr, obj.labels[i]);
		move_to(cr, x*px-(extents[3]/2 + extents[1]), y*px-(extents[4]/2 + extents[2]));
		show_text(cr, obj.labels[i]);
	end
	finish(c)
end


text:
scene, layout  = layoutscene()

lscene = layout[1, 1] = Axis(scene, xlabel = "South to North [m]",
        xlabelcolor = :black, ylabel = "Depth [m]",
        ylabelcolor = :black, xlabelsize = 40, ylabelsize = 40,
        xticklabelsize = 25, yticklabelsize = 25,
        xtickcolor = :black, ytickcolor = :black,
        xticklabelcolor  = :black, yticklabelcolor = :black,
        titlesize = 50)
statestring = "ν ≈  f v'b' / ∂ᶻb / ∂ᶻu"
fieldindex = findall(x->x==statestring, statenames)
u = log10.(abs.(states[fieldindex[1]]))
b = states[4]
statenames[4]

xlims = (0, 1e6)
ylims = (-3000, 0)
xlims = Array(range(xlims[1], xlims[2], length = 4))
ylims = Array(range(ylims[1], ylims[2], length = 4))

state = u .* 1.0
cmap_rgb = to_colormap(:balance);
clims = quantile.(Ref(state[:]), (0.2,0.8))
heatmap1 = heatmap!(lscene, xlims, ylims, state, interpolate = true, colormap = cmap_rgb, colorrange = clims)
contour!(lscene,0..1e6,-3000..0, b, levels = 40, linewidth = 4, color = :black, alpha = 0.5)
display(scene)
##
ii = 11
state = states[ii]
b = states[4]
scene = Figure()
cmap_rgb = to_colormap(:viridis);
clims = (quantile(state[:], 0.05) , quantile(state[:], 0.95))
scene = heatmap(state, interpolate = true, colormap = cmap_rgb, colorrange = clims)
cp = contour!(scene, b, levels = 30, linewidth = 4, color = :black, alpha = 0.5)
axis = scene[Axis]
axis.names.axisnames = ["South to North [m]", "Depth [m]"]
newscene = title(scene, statenames[ii] * " with " * statenames[4] * " contours")

ls = colorlegend(scene[end-1], raw = true, camera = campixel!)
scene_final = vbox(newscene, ls)
display(scene_final)

max_stars_repo_name: sandreza/Mesoscale

max_issues_repo_name: sandreza/Mesoscale

max_forks_repo_name: sandreza/Mesoscale


text:
using Test
import LabViewXML

@testset "LabViewXML Tests" begin
    d = LabViewXML.readxml(joinpath(@__DIR__, "testdata/testdata.xml"))
    @test d["MyI8"] == Int8(4)
    @test d[""] == Int32(-1)
    @test d["MyCluster"]["MyU8Array"] == UInt8[1,2,3]
end


text:
using Test, OMEinsum
using ForwardDiff
using TensorNetworkEvolve.TensorAD

@testset "jacobians 1" begin
    n = 3
    for T in [Float64, ComplexF64]
        for (f, args, kwargs) in [
            (ein"ii->", (randn(T,4,4),), ()),
            (reshape, (randn(T,4,4), 2, 8), ()),
            (transpose, (randn(T,4,4),), ()),
            (TensorAD.accum, (randn(T,4,4), (:,:), randn(T,4,4)), ()),
            (cat, (randn(T,4,4), randn(T,4, 2)), (;dims=2)),
            (copy, (randn(T,4,4),), ()),
            (conj, (randn(T,4,4),), ()),
            (real, (randn(T,4,4),), ()),
            (imag, (randn(T,4,4),), ()),
            (+, (randn(T,4,4), randn(T,4, 4)), ()),
            (*, (randn(T,4,4), randn(T,4, 4)), ()),
            (*, (3.0, randn(T,4, 4)), ()),
            (-, (randn(T,4,4), randn(T,4, 4)), ()),
            (-, (randn(T,4,4),), ()),
            (getindex, (randn(T,4,4), 3:4, 2:3), ()),
            (Base.broadcast, (*, randn(T,4,4), randn(T,4, 4)), ()),
            (Base.broadcast, (/, randn(T,4,4), randn(T,4, 4)), ()),
            (Base.broadcast, (^, randn(T,4,4), 3), ()),
            (Base.broadcast, (^, randn(T,4,4), n), ()),
            (Base.broadcast, (sin, randn(T,4,4)), ()),
            (Base.broadcast, (cos, randn(T,4,4)), ()),
            (Base.broadcast, (sign, randn(T,4,4)), ()),
            (Base.broadcast, (abs, randn(T,4,4)), ()),
            (Base.broadcast, (sqrt, 10 .+ randn(T,4,4)), ()),
            (Base.broadcast, (inv, randn(T,4,4)), ()),
            (Base.broadcast, (Complex, randn(T,4,4)), ()),
        ]
            @info "Differentiating function: $f, arg types: $(typeof(args))"
            @test match_jacobian(f, args...; realpart=true, kwargs...)
            @test match_random(f, args...; realpart=true, kwargs...)
            if T <: Complex
                @info "Differentiating function: $f, arg types: $(typeof(args)), imaginary part"
                @test match_jacobian(f, args...; realpart=false, kwargs...)
                @test match_random(f, args...; realpart=false, kwargs...)
            end
        end
    end
end

@testset "jacobians 1" begin
    n = 3
    _mul(x::Array{T,0}, y::Array{T,0}) where T = fill(x[] * y[])
    _mul(x::Array{T,0}, y::Array) where T = x[] * y
    _mul(x::Array, y::Array{T,0}) where T = x * y[]
    _mul(x, y) = x * y
    _fill(x, args...) = fill(x, args...)
    _fill(x::Array, args...) = fill(x[], args...)
    for T in [Float64, ComplexF64]
        for (f, args, kwargs) in [
            (sum, (randn(T,4,4),), ()),
            (_fill, (fill(randn(T)), 2, 3), ()),
            (_mul, (fill(randn(T)), randn(T,4, 4)), ()),
            (_mul, (randn(T,4,4), fill(randn(T))), ()),
            (_mul, (fill(T(2.0)), fill(T(2.0))), ()),
        ]
            @info "Differentiating function: $f, arg types: $(typeof(args))"
            @test match_jacobian(f, args...; realpart=true, kwargs...)
            @test match_random(f, args...; realpart=true, kwargs...)
            if T <: Complex
                @info "Differentiating function: $f, arg types: $(typeof(args)), imaginary part"
                @test match_jacobian(f, args...; realpart=false, kwargs...)
                @test match_random(f, args...; realpart=false, kwargs...)
            end
        end
    end
end


text:
module Thermodynamics

using ThermodynamicsBase
using ForwardDiff: derivative
using Reexport
using NLSolve: nlsolve
using CSV

include("backends/idealgas.jl")
include("backends/multiparameterhelmholtz.jl")
include("backends/pengrobinson.jl")

export phasestate, spec_state!, mixture, thermocalc, thermocalc!, makethermodata, makethermoplot

end # module Thermodynamics


text:
"""
Momentum transfer coefficients.

* [`dyn_visc_dryair`](@ref)
* [`dyn_visc_vapor`](@ref)
* [`dyn_visc_moistair`](@ref)

# References

* [T08] Tsilingiris, P. T. (2008). Thermophysical and transport properties of
  humid air at temperature range between 0 and 100°C. *Energy Conversion and
  Management*, 49(5), 1098--1110.
  <https://doi.org/10.1016/j.enconman.2007.09.015>
"""
module Momentum

include("../docstring_style.jl")

using Canopy.Constants: M_w, M_d
using Canopy.Water: vapor_mole_frac

export dyn_visc_dryair,
    dyn_visc_vapor,
    dyn_visc_moistair

"""
Calculate the dynamic viscosity of dry air [kg m^-1 s^-1] from temperature [K].
Applicable between 250 and 600 K.

!!! note

    The influence of pressure is not considered because it is negligible when
    pressure < a few MPa.
"""
dyn_visc_dryair(temp) = @evalpoly(
    temp, -9.8601e-1, 9.080125e-2, -1.17635575e-4,
    1.2349703e-7, -5.7971299e-11) * 1e-6

"""
Calculate the dynamic viscosity of water vapor [kg m^-1 s^-1] from temperature
[K]. Applicable between 273 and 393 K.

!!! note

    The influence of pressure is not considered because it is negligible when
    pressure < a few MPa.
"""
dyn_visc_vapor(temp) = -2.869368957406498e-06 + 4.000549451e-8 * temp

"""
Calculate viscosity interaction terms from temperature [K]. This is a helper
function for calculating moist air properties.
"""
function viscosity_interaction_terms(temp)
    # `_d` - dry air; `_w` - water vapor
    M_dw = M_d / M_w
    M_wd = M_w / M_d
    mu_d = dyn_visc_dryair(temp)
    mu_w = dyn_visc_vapor(temp)
    # interaction coefficients
    phi_dw = sqrt(2.) * 0.25 * (1. + M_dw) ^ -0.5 *
        (1. + sqrt(mu_d / mu_w) * M_wd^0.25) ^ 2.
    phi_wd = sqrt(2.) * 0.25 * (1. + M_wd) ^ -0.5 *
        (1. + sqrt(mu_w / mu_d) * M_dw^0.25) ^ 2.
    return phi_dw, phi_wd, mu_d, mu_w
end

"""
Calculate the dynamic viscosity of moist air [kg m^-1 s^-1] from temperature
[K], pressure [Pa], and relative humidity [0--1].
"""
function dyn_visc_moistair(temp, pressure, RH)
    x_w = vapor_mole_frac(temp, pressure, RH)
    phi_dw, phi_wd, mu_d, mu_w = viscosity_interaction_terms(temp)
    return (1. - x_w) * mu_d / ((1. - x_w) + x_w * phi_dw) +
        x_w * mu_w / (x_w + (1. - x_w) * phi_wd)
end

end  # module


text:
#!/usr/bin/env julia
ENV["JULIA_NUM_THREADS"] = 16
ENV["JULIA_EDITOR"] = "code"
ENV["PYTHON"] = "python"

atreplinit() do repl
    if isfile("Project.toml")
        try
            @eval using Revise
            @async Revise.wait_steal_repl_backend()
        catch
            try
                @eval using Pkg
                haskey(Pkg.installed(), "Revise") || @eval Pkg.add("Revise")
            catch
            end
        end
    end
end

try
    import AbstractTrees
    AbstractTrees.children(x::Type) = subtypes(x)
    using AbstractTrees: print_tree
catch
    try
        using Pkg
        haskey(Pkg.installed(), "AbstractTrees") || @eval Pkg.add("AbstractTrees")
    catch
    end
end

if isfile("Project.toml")
    using Pkg
    Pkg.activate(".")
else
    using LinearAlgebra
    # From https://discourse.julialang.org/t/how-to-pass-multiple-arguments-to-a-function-using/29117/3
    →(args, f) = f(args...)
    Base.round(x::AbstractMatrix, digits::Int = 15) = round.(x, digits = digits)
end


text:
module test

struct ComponentInstanceParameters{T <: NamedTuple}
    nt::T
end

struct ComponentInstanceVariables{T <: NamedTuple}
    nt::T
end

struct ComponentInstance{V <: NamedTuple, P <: NamedTuple}
    vars::ComponentInstanceVariables{V}
    params::ComponentInstanceParameters{P}
end

function ComponentInstanceParameters(names, types, values)
    NT = NamedTuple{names, types}
    ComponentInstanceParameters{NT}(NT(values))
end

@inline function Base.getproperty(obj::ComponentInstanceParameters{T}, name::Symbol) where {T}
    nt = getfield(obj, :nt)
    return fieldtype(T, name) <: Ref ? getproperty(nt, name)[] : getproperty(nt, name)
end

using BenchmarkTools

ci = ComponentInstanceParameters((a=1., b=2.))

foo(ci) = ci.a + ci.b

println("@btime ci.a + ci.b")
@btime foo($ci)

function run_timestep(p, v, d, t)
end

end # module


text:
export batchnorm, bnmoments, bnparams
using AutoGrad: AutoGrad, @primitive1, value, recording
using Statistics: mean, var

# dy20200804: why not have bnparams and bnmoments in a single struct?
# dy20200804: what array types / sizes are they supposed to use? which ones are Param?
# moments and params need to have the same atype as x, params needs to be wrapped in Param.
#TODO: improve the documentation. Explain what batchnorm is.

mutable struct BNMoments
    momentum::AbstractFloat
    mean
    var
    meaninit
    varinit
end

"""
    batchnorm(x[, moments, params]; kwargs...)

perform batch normalization on `x` with optional mean and variance in `moments` and scaling
factor and bias in `params`. See https://arxiv.org/abs/1502.03167 for reference.

2d, 4d and 5d inputs are supported. Mean and variance are computed over dimensions (2,),
(1,2,4) and (1,2,3,5) for 2d, 4d and 5d arrays, respectively.

`moments` stores running mean and variance to be used at inference time.  It is optional in
training mode, but mandatory in test mode.  Training and test modes can be controlled by the
`training` keyword argument which defaults to `Knet.training()`.

`params` stores the optional affine parameters gamma and beta.  `bnparams` function can be
used to initialize `params`.

# Example

    # Inilization, C is an integer
    moments = bnmoments()
    params = bnparams(C)
    ...
    # size(x) -> (H, W, C, N)
    y = batchnorm(x, moments, params)
    # size(y) -> (H, W, C, N)

# Keywords

 `eps=1e-5`: The epsilon parameter added to the variance to avoid division by 0.

 `training=Knet.training()`: When `training` is true, the mean and variance of `x` are used
 and `moments` argument is modified if it is provided. When `training` is false, mean and
 variance stored in the `moments` argument are used.

"""
function batchnorm(x, moments::Union{BNMoments, Nothing}=nothing, params=nothing;
                   training=AutoGrad.recording(), o...)
    xnd = ndims(x)
    a = (x,)
    if params !== nothing
        g = reshape(_bnscale(params), _wsize(x))
        b = reshape(_bnbias(params), _wsize(x))
        a = (g, b, x)
    end
    if xnd == 2
        return batchnorm2(a...; o...,
                          moments=moments,
                          training=training,
                          cache=BNCache())
    elseif xnd in [4, 5]
        return batchnorm4(a...; o...,
                          moments=moments,
                          training=training,
                          cache=BNCache())
    else
        error("Unsupported input dimension ", xnd)
    end
end

"""
    bnmoments(;momentum=0.1, mean=nothing, var=nothing, meaninit=zeros, varinit=ones)

Return a `BNMoments` object, a data structure used to store running mean and running
variance of batch normalization with the following fields:

* `momentum=0.1`: A real number between 0 and 1 to be used as the scale of last
mean and variance. The existing running mean or variance is multiplied by (1-momentum).

* `mean=nothing`: The running mean.

* `var=nothing`: The running variance.

* `meaninit=zeros`: The function used for initialize the running mean. Should either be
`nothing` or of the form `([eltype], dims...)->data`. `zeros` is a good option.

* `varinit=ones`: The function used for initialize the running variance. Should either be
`nothing` or `([eltype], dims...)->data`. `ones` is a good option.

This constructor can be used directly load moments from data. `meaninit` and `varinit` are
called if `mean` and `var` are nothing. Type and size of the `mean` and `var` are determined
automatically from the inputs in the `batchnorm` calls.

"""
bnmoments(;momentum=0.1, meaninit=zeros, varinit=ones, mean=nothing, var=nothing) =
    BNMoments(momentum, mean, var, meaninit, varinit)

"""
    bnparams(etype, channels::Integer)

Return a single 1d array that contains both scale and bias of batchnorm, where the first
half is scale and the second half is bias.

`bnparams(channels)` calls `bnparams(Float64, channels)`, following Julia convention.

"""
function bnparams(etype, channels::Integer)
    buf = Array{etype}(undef,2channels)
    buf[1:channels] .= 1
    buf[channels+1:end] .= 0
    return buf
end

bnparams(channels::Integer) = bnparams(Float64, channels)

#=
LOW-LEVEL API that won't be exported by default
=#
@inline _bnscale(param) = param[1:div(length(param), 2)]
@inline _bnbias(param) = param[div(length(param), 2)+1:end]

# Dimension helpers
@inline _wsize(y) = ((1 for _=1:ndims(y)-2)..., size(y)[end-1], 1)
@inline _reddims(y) = ((i for i=1:ndims(y)-2)..., ndims(y))

# A black box data type for storing the bn state
mutable struct BNCache
    mean
    ivar
    dx
    dg
    db
end

BNCache() = BNCache(nothing, nothing, nothing, nothing, nothing)

# CPU Implementation
function batchnorm4(g, b, x; o...)
    return _batchnorm4_fused(g, b, x; o...)
end

function batchnorm4(x; o...)
    return _batchnorm4_fused(nothing,nothing,x; o...)
end

function _batchnorm4_fused(g, b, x; eps=1e-5, training=AutoGrad.recording(), cache=nothing, moments=nothing, o...)
    T = eltype(x)
    y = copy(x)
    eps = T(eps)
    dims = _reddims(y)
    if moments!==nothing
        _lazy_init!(moments, x)
    end
    affine = (g !== nothing)
    if training
        mu = mean(y, dims=dims)
        sigma2 = var(y, dims=dims, corrected=false, mean=mu)
        _update_moments!(moments, mu, sigma2)
        sigma2 .= 1 ./ sqrt.(sigma2 .+ eps)
        ivar = sigma2
        if affine
            y .= g .* (y .- mu) .* ivar .+ b
        else
            y .= (y .- mu) .* ivar
        end
        if cache !== nothing
            cache.ivar = ivar
            cache.mean = mu
        end
    else
        @assert (moments!==nothing)  "Moments must be provided in test mode"
        ivar = 1 ./ sqrt.(moments.var .+ eps)
        if affine
            y .= g .* (y .- moments.mean) .* ivar .+ b
        else
            y .= (y .- moments.mean) .* ivar
        end
    end
    return y
end

function _update_moments!(moments, mu, sigma2)
    moments === nothing && return
    moments.mean .= moments.momentum .* mu .+
        (1 .- moments.momentum) .* moments.mean
    moments.var .= moments.momentum .* sigma2 .+
        (1 .- moments.momentum) .* moments.var
end

# TODO: consider automatic type conversion
function _lazy_init!(m::BNMoments, x)
    x = value(x)
    buf_size = (ndims(x) > 2) ? _wsize(x) : (size(x,1), 1)
    tx = typeof(x)
    ex = eltype(x)
    m.momentum = ex(m.momentum)
    if m.mean == nothing
        m.mean = tx(m.meaninit(ex, buf_size...))
    end
    if m.var == nothing
        m.var = tx(m.varinit(ex, buf_size...))
    end
end

# Implement batchnorm2 using batchnorm4, with autograd

function batchnorm2(g, b, x; moments=nothing, training=AutoGrad.recording(), o...)
    # TODO: This support should be added when needed
    if training == false && (isa(g, AutoGrad.Value) || isa(x, AutoGrad.Value) || isa(b, AutoGrad.Value))
        error("Test mode backward is not supported with 2d inputs")
    end
    @inline _pad4(x) = reshape(x, (1,1,size(x,1),size(x,2)))
     # process moments
    if moments !== nothing
        _lazy_init!(moments, x)
        moments.mean = _pad4(moments.mean)
        moments.var = _pad4(moments.var)
    end
    x = _pad4(x)
    args = (x,)
    if g !== nothing
        g = _pad4(g)
        b = _pad4(b)
        args = (g, b, x)
    end
    y = mat(batchnorm4(args...;
                       moments=moments,
                       training=training,
                       o...))
    if moments !== nothing
        moments.mean = mat(moments.mean)
        moments.var = mat(moments.var)
    end
    return y
end

batchnorm2(x;o...) = batchnorm2(nothing, nothing, x; o...)

# CPU backward
function batchnorm4_back(g, x, dy; eps=1e-5, training=AutoGrad.recording(), cache=nothing, moments=nothing,  o...)
    T = eltype(x)
    eps = T(eps)
    dims = _reddims(x)
    if training
        mu, ivar = _get_cache_data(cache, x, eps)
        x_mu = x .- mu
        # equations from the original paper
        dyivar = dy .* ivar
        if g !== nothing
            dg = sum(x_mu .* dyivar, dims=dims)
            db = sum(dy, dims=dims)
            dyivar .*= g
        else
            dg, db = nothing, nothing
        end
        m = prod(d->size(x,d), dims) # size(x, dims...))
        dsigma2 = -sum(dyivar .* x_mu .* ivar.^2, dims=dims) ./ 2
        dmu = -sum(dyivar, dims=dims) .- 2dsigma2 .* sum(x_mu, dims=dims) ./ m
        dx = dyivar .+ dsigma2 .* 2x_mu ./ m .+ dmu ./ m
    else #same reasoning with the gpu version
        ivar = 1 ./ sqrt.(moments.var .+ eps)
        dx = (g !== nothing) ? (dy .* g .* ivar) : (dy .* ivar)
        if g !== nothing
            dg = sum(dy .* (x .- moments.mean) .* ivar, dims=dims)
            db = sum(dy, dims=dims)
        else
            dg, db = nothing, nothing
        end
    end
    return dg, db, dx
end

function _get_cache_data(cache, x, eps)
    if cache !== nothing
        mu = cache.mean
        ivar = cache.ivar
    else
        mu = mean(x, dims=_reddims(x))
        ivar = 1 ./ sqrt.(var(x, dims=_reddims(x); mean=mu, correct=false) .+ eps)
    end
    return mu, ivar
end

function batchnorm4g(g,x,dy; cache=nothing, o...)
    dg, db, dx = batchnorm4_back(g, x, dy; cache=cache, o...)
    if cache !== nothing
        cache.dx, cache.db = dx, db
    else
        @warn("Calling batchnorm with affine without cache is not recommended for performance")
    end
    return dg
end

function batchnorm4b(dy; cache=nothing, o...)
    if cache == nothing || cache.db == nothing
        return sum(dy, dims=_reddims(dy))
    else
        return cache.db
    end
end

function batchnorm4x(g, x, dy; cache=nothing, o...)
    if cache !== nothing && cache.dx !== nothing
        return cache.dx
    else
        return batchnorm4_back(g, x, dy; cache=cache, o...)[3]
    end
end

function batchnorm4x(x, dy; o...)
    return batchnorm4_back(nothing, x, dy; o...)[3]
end

@primitive1 batchnorm4(x;o...),dy batchnorm4x(x, dy; o...)

# 4d with affine
@primitive1 batchnorm4(g, b, x; o...),dy batchnorm4g(g, x, dy; o...) batchnorm4b(dy; o...) batchnorm4x(g, x, dy; o...)


text:
"""
The Result type represents result codes for all 5 characters in a candidate.
"""
struct Result
    codes::NTuple{5, ResultCode}
end

"""
Convert a Wordle emoji string to a Result object.
"""
function Result(s::String)
    if length(s) !== 5
        throw(ArgumentError("Input must be 5 characters long"))
    end
    chars = Char[]
    for char in s
        push!(chars, char)
    end
    Result(ntuple(i -> ResultCode(chars[i]), 5))
end


text:
using Colors
using Images
using Boltzmann
using MNIST
using ImageView
using Base.Test

function run_mnist()
    X, y = testdata()  # test data is smaller, no need to downsample
    normalize_samples!(X)
    binarize!(X;threhold=0.01)

    TrainSet = X[:,1:10000]
    ValidSet = []
    HiddenUnits = 256;
    Epochs = 15;
    Gibbs = 1;
    LearnRate = 0.05

    m = BernoulliRBM(28*28, HiddenUnits, (28,28); momentum=0.0)
    mwdQuad = BernoulliRBM(28*28, HiddenUnits, (28,28); momentum=0.0)
    mwdLin = BernoulliRBM(28*28, HiddenUnits, (28,28); momentum=0.0)

    # Attempt without weight decay
    info("Running Without Weight Decay")
    m, mon = fit(m, TrainSet;n_iter=Epochs,lr=LearnRate,validation=ValidSet,n_gibbs=Gibbs)
    SaveMonitor(m,mon,"test_nodecay.pdf")

    # Attempt with L2 weight decay
    info("Running With L2-Decay")
    mwQuad, mon = fit(mwdQuad, TrainSet;n_iter=Epochs,weight_decay="l2",decay_magnitude=0.05,lr=LearnRate,validation=ValidSet,n_gibbs=Gibbs)
    SaveMonitor(mwQuad,mon,"test_l2decay.pdf")

    # Attempt with L1 weight decay
    info("Running With L1-Decay")
    mwLin, mon = fit(mwdLin, TrainSet;n_iter=Epochs,weight_decay="l1",decay_magnitude=0.05,lr=LearnRate,validation=ValidSet,n_gibbs=Gibbs)
    SaveMonitor(mwLin,mon,"test_l1decay.pdf")

    return m
end

run_mnist()


text:
using JuMP, VariationalInequality

# write your own tests here

m = VIPModel()

A = [25; 25; 75; 25; 25]
B = [0.010; 0.010; 0.001; 0.010; 0.010]
T14 = 100
p = 3

@variable(m, h[i=1:p] >= 0)

@NLexpression(m, F1, A[1]+B[1]*h[1]^2 + A[4]+B[4]*(h[1]+h[2])^2 )
@NLexpression(m, F2, A[2]+B[2]*(h[2]+h[3])^2 + A[3]+B[3]*h[2]^2 + A[4]+B[4]*(h[1]+h[2])^2 )
@NLexpression(m, F3, A[2]+B[2]*(h[2]+h[3])^2 + A[5]+B[5]*(h[3])^2 )

@NLconstraint(m, sum{h[i], i=1:p} == T14)

F = [F1, F2, F3]
correspond(m, F, h)
sol, Fval, gap = solveVIP(m, algorithm=:extra_gradient, max_iter=1000, step_size=0.01)

@assert 0<= gap1 < 1e-6

@show sol
@show gap


text:
"""
    MH(n_iters::Int)

Metropolis-Hastings sampler.

Usage:

```julia
MH(100, (:m, (x) -> Normal(x, 0.1)))
```

Example:

```julia
# Define a simple Normal model with unknown mean and variance.
@model gdemo(x) = begin
  s ~ InverseGamma(2, 3)
  m ~ Normal(0, sqrt(s))
  x[1] ~ Normal(m, sqrt(s))
  x[2] ~ Normal(m, sqrt(s))
  return s, m
end

chn = sample(gdemo([1.5, 2]), MH(1000))
```
"""
mutable struct MH{space} <: InferenceAlgorithm
    n_iters   ::  Int       # number of iterations
    proposals ::  Dict{Symbol,Any}  # Proposals for paramters
end
function MH(n_iters::Int, proposals::Dict{Symbol, Any}, space::Tuple)
    return MH{space}(n_iters, proposals)
end

function MH(n_iters::Int, space...)
    new_space = ()
    proposals = Dict{Symbol,Any}()

    # parse random variables with their hypothetical proposal
    for element in space
        if isa(element, Symbol)
            new_space = (new_space..., element)
        else
            @assert isa(element[1], Symbol) "[MH] ($element[1]) should be a Symbol. For proposal, use the syntax MH(N, (:m, (x) -> Normal(x, 0.1)))"
            new_space = (new_space..., element[1])
            proposals[element[1]] = element[2]
        end
    end
    MH(n_iters, proposals, new_space)
end

function Sampler(alg::MH, model::Model, s::Selector)
    alg_str = "MH"

    # Sanity check for space
    if (s.tag == :default) && !isempty(getspace(alg))
        @assert issubset(get_pvars(model), getspace(alg)) "[$alg_str] symbols specified to samplers ($getspace(alg)) doesn't cover the model parameters ($(get_pvars(model)))"
        if !(issetequal(get_pvars(model), getspace(alg)))
            @warn("[$alg_str] extra parameters specified by samplers don't exist in model: $(setdiff(getspace(alg), get_pvars(model)))")
        end
    end

    info = Dict{Symbol, Any}()
    info[:proposal_ratio] = 0.0
    info[:prior_prob] = 0.0
    info[:violating_support] = false

    return Sampler(alg, info, s)
end

function propose(model, spl::Sampler{<:MH}, vi::VarInfo)
    spl.info[:proposal_ratio] = 0.0
    spl.info[:prior_prob] = 0.0
    spl.info[:violating_support] = false
    return runmodel!(model, vi ,spl)
end

function step(model, spl::Sampler{<:MH}, vi::VarInfo, is_first::Val{true})
    return vi, true
end

function step(model, spl::Sampler{<:MH}, vi::VarInfo, is_first::Val{false})
  if spl.selector.tag != :default # Recompute joint in logp
    runmodel!(model, vi)
  end
  old_θ = copy(vi[spl])
  old_logp = getlogp(vi)

  Turing.DEBUG && @debug "Propose new parameters from proposals..."
  propose(model, spl, vi)

  Turing.DEBUG && @debug "computing accept rate α..."
  is_accept, _ = mh_accept(-old_logp, -getlogp(vi), spl.info[:proposal_ratio])

  Turing.DEBUG && @debug "decide wether to accept..."
  if is_accept && !spl.info[:violating_support]  # accepted
    is_accept = true
  else                      # rejected
    is_accept = false
    vi[spl] = old_θ         # reset Θ
    setlogp!(vi, old_logp)  # reset logp
  end

  return vi, is_accept
end

function sample(model::Model, alg::MH;
                save_state=false,         # flag for state saving
                resume_from=nothing,      # chain to continue
                reuse_spl_n=0,            # flag for spl re-using
                )

  spl = reuse_spl_n > 0 ?
        resume_from.info[:spl] :
        Sampler(alg, model)
    if resume_from != nothing
        spl.selector = resume_from.info[:spl].selector
    end
  alg_str = "MH"

  # Initialization
  time_total = 0.0
  n = reuse_spl_n > 0 ?
      reuse_spl_n :
      alg.n_iters
  samples = Array{Sample}(undef, n)
  weight = 1 / n
  for i = 1:n
    samples[i] = Sample(weight, Dict{Symbol, Any}())
  end

    vi = if resume_from == nothing
        VarInfo(model)
    else
        resume_from.info[:vi]
    end

  if spl.selector.tag == :default
    runmodel!(model, vi, spl)
  end

  # MH steps
  accept_his = Bool[]
  PROGRESS[] && (spl.info[:progress] = ProgressMeter.Progress(n, 1, "[$alg_str] Sampling...", 0))
  for i = 1:n
    Turing.DEBUG && @debug "$alg_str stepping..."

    time_elapsed = @elapsed vi, is_accept = step(model, spl, vi, Val(i == 1))
    time_total += time_elapsed

    if is_accept # accepted => store the new predcits
        samples[i].value = Sample(vi, spl).value
    else         # rejected => store the previous predcits
        samples[i] = samples[i - 1]
    end

    samples[i].value[:elapsed] = time_elapsed
    push!(accept_his, is_accept)

    PROGRESS[] && (ProgressMeter.next!(spl.info[:progress]))
  end

  println("[$alg_str] Finished with")
  println("  Running time        = $time_total;")
  accept_rate = sum(accept_his) / n  # calculate the accept rate
  println("  Accept rate         = $accept_rate;")

  if resume_from != nothing   # concat samples
    pushfirst!(samples, resume_from.info[:samples]...)
  end
  c = Chain(0.0, samples)       # wrap the result by Chain
  if save_state               # save state
    c = save(c, spl, model, vi, samples)
  end

  c
end

function assume(spl::Sampler{<:MH}, dist::Distribution, vn::VarName, vi::VarInfo)
    if isempty(getspace(spl.alg)) || vn.sym in getspace(spl.alg)
        if ~haskey(vi, vn) error("[MH] does not handle stochastic existence yet") end
        old_val = vi[vn]

        if vn.sym in keys(spl.alg.proposals) # Custom proposal for this parameter
            proposal = spl.alg.proposals[vn.sym](old_val)
            if typeof(proposal) == Distributions.Normal{Float64} # If Gaussian proposal
                σ = std(proposal)
                lb = support(dist).lb
                ub = support(dist).ub
                stdG = Normal()
                r = rand(TruncatedNormal(proposal.μ, proposal.σ, lb, ub))
                # cf http://fsaad.scripts.mit.edu/randomseed/metropolis-hastings-sampling-with-gaussian-drift-proposal-on-bounded-support/
                spl.info[:proposal_ratio] += log(cdf(stdG, (ub-old_val)/σ) - cdf(stdG,(lb-old_val)/σ))
                spl.info[:proposal_ratio] -= log(cdf(stdG, (ub-r)/σ) - cdf(stdG,(lb-r)/σ))
            else # Other than Gaussian proposal
                r = rand(proposal)
                if (r < support(dist).lb) | (r > support(dist).ub) # check if value lies in support
                    spl.info[:violating_support] = true
                    r = old_val
                end
                spl.info[:proposal_ratio] -= logpdf(proposal, r) # accumulate pdf of proposal
                reverse_proposal = spl.alg.proposals[vn.sym](r)
                spl.info[:proposal_ratio] += logpdf(reverse_proposal, old_val)
            end

        else # Prior as proposal
            r = rand(dist)
            spl.info[:proposal_ratio] += (logpdf(dist, old_val) - logpdf(dist, r))
        end

        spl.info[:prior_prob] += logpdf(dist, r) # accumulate prior for PMMH
        vi[vn] = vectorize(dist, r)
        setgid!(vi, spl.selector, vn)
    else
        r = vi[vn]
    end

    # acclogp!(vi, logpdf(dist, r)) # accumulate pdf of prior
    r, logpdf(dist, r)
end

function assume(  spl::Sampler{<:MH},
                  dists::Vector{D},
                  vn::VarName,
                  var::Any,
                  vi::VarInfo
                ) where D<:Distribution
    error("[Turing] MH doesn't support vectorizing assume statement")
end

function observe(spl::Sampler{<:MH}, d::Distribution, value::Any, vi::VarInfo)
    return observe(nothing, d, value, vi)  # accumulate pdf of likelihood
end

function observe( spl::Sampler{<:MH},
                  ds::Vector{D},
                  value::Any,
                  vi::VarInfo
                )  where D<:Distribution
    return observe(nothing, ds, value, vi) # accumulate pdf of likelihood
end


text:
# ======================
# Domains
# ======================

# The following domains are used throughout the tests in Tables 3-5 in
# [1] Althoff, M., Grebenyuk, D., & Kochdumper, N. (2018). Implementation of Taylor models in CORA 2018.
#     In Proc. of the 5th International Workshop on Applied Verification for Continuous and Hybrid Systems.
a = Interval(-4.5, -0.3)
b = Interval(0.4, 0.9)
c = Interval(3.8, 7.8)
d = Interval(8.0, 10.0)
e = Interval(-10.0, 8.0)
f = Interval(1.0, 2.0)

# ======================
# Relative precision
# ======================

# Dictionary to hold the vector of relative precision intervals for each benchmark
RELPREC = Dict{String, Any}()

# This function measures the relative precision of the result in a more informative way than
# taking the scalar overestimation because it evaluates the precision of the lower and the
# upper range bounds separately, see Eq. (20) in [1].
function relative_precision(x, x_ref)
    x_low, x_high = inf(x), sup(x)
    x_ref_low, x_ref_high = inf(x_ref), sup(x_ref)
    rel_low = -(x_low - x_ref_low) / (x_ref_high - x_ref_low)
    rel_high = (x_high - x_ref_high) / (x_ref_high - x_ref_low)
    return 100 * Interval(rel_low, rel_high)
end

# ==========================
# Methods to compute bounds
# ==========================

# taylor model in one variable
function bounds_TM(func::Function, dom::Interval, ord::Int)
    x0 = Interval(mid(dom))
    x = TaylorModel1(ord, x0, dom)
    return evaluate(func(x), dom - x0)
end

# taylor model in N variables
function bounds_TM(func::Function, dom::IntervalBox{N}, ord) where {N}
    x0 = mid(dom)
    set_variables(Float64, "x", order=2ord, numvars=N)
    x = [TaylorModelN(i, ord, IntervalBox(x0), dom) for i=1:N]
    return evaluate(func(x...), dom - x0)
end

# normalized taylor model in one variable
function bounds_TM_NORM(func::Function, dom::Interval, ord::Int)
    x0 = Interval(mid(dom))
    x = TaylorModel1(ord, x0, dom)
    xnorm = normalize_taylor(x.pol, dom - x0, true)
    xnormTM = TaylorModel1(xnorm, 0..0, 0..0, -1..1)
    return evaluate(func(xnormTM), -1..1)
end

# normalized taylor model in N variables
function bounds_TM_NORM(func::Function, dom::IntervalBox{N}, ord::Int) where {N}
    x0 = mid(dom)
    set_variables(Float64, "x", order=2ord, numvars=N)

    zeroBox = IntervalBox(0..0, N)
    symBox = IntervalBox(-1..1, N)

    x = [TaylorModelN(i, ord, IntervalBox(x0), dom) for i=1:N]
    xnorm = [normalize_taylor(xi.pol, dom - x0, true) for xi in x]
    xnormTM = [TaylorModelN(xi_norm, 0..0, zeroBox, symBox) for xi_norm in xnorm]
    return evaluate(func(xnormTM...), symBox)
end

# interval arithmetic substitution
function bounds_IA(f, dom)
    return f(dom...)
end


text:
"""
    ksettings_input_test(ksettings::KalmanSettings, Y::JArray, B::FloatMatrix, R::Union{UniformScaling{Float64}, SymMatrix}, C::FloatMatrix, D::FloatMatrix, Q::SymMatrix, X0::FloatArray, P0::SymMatrix, DQD::SymMatrix, n::Int64, T::Int64, m::Int64; compute_loglik::Bool=true, store_history::Bool=true)

Return true if the entries of ksettings are correct (false otherwise).
"""
function ksettings_input_test(ksettings::KalmanSettings, Y::JArray, B::FloatMatrix, R::Union{UniformScaling{Float64}, SymMatrix}, C::FloatMatrix, D::FloatMatrix, Q::SymMatrix, X0::FloatArray, P0::SymMatrix, DQD::SymMatrix, n::Int64, T::Int64, m::Int64; compute_loglik::Bool=true, store_history::Bool=true)

    return ~(false in [ksettings.Y.data === Y;
                      ksettings.B == B;
                      ksettings.R == R;
                      ksettings.C == C;
                      ksettings.D == D;
                      ksettings.Q == Q;
                      ksettings.X0 == X0;
                      round.(ksettings.P0, digits=10) == P0;
                      ksettings.DQD == DQD;
                      ksettings.Y.n == n;
                      ksettings.Y.T == T;
                      ksettings.m == m;
                      ksettings.compute_loglik == compute_loglik;
                      ksettings.store_history == store_history]);
end

"""
    test_kalman_output(ksettings::KalmanSettings, kstatus::DynamicKalmanStatus, benchmark_data::Tuple)
    test_kalman_output(ksettings::KalmanSettings, kstatus::SizedKalmanStatus, benchmark_data::Tuple)

kalman_test internals.
"""
function test_kalman_output(ksettings::KalmanSettings, kstatus::DynamicKalmanStatus, benchmark_data::Tuple)

    benchmark_n, benchmark_T, benchmark_m, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior,
        benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik, benchmark_X0_sm, benchmark_P0_sm,
            benchmark_X_sm, benchmark_P_sm = benchmark_data;

    for t=1:ksettings.Y.T

        # Run filter
        kfilter!(ksettings, kstatus);

        # A-priori
        @test round.(kstatus.X_prior, digits=10) == benchmark_X_prior[t];
        @test round.(kstatus.P_prior, digits=10) == benchmark_P_prior[t];

        # A-posteriori
        @test round.(kstatus.X_post, digits=10) == benchmark_X_post[t];
        @test round.(kstatus.P_post, digits=10) == benchmark_P_post[t];

        # 12-step ahead forecast
        @test round.(kforecast(ksettings, kstatus.X_post, 12)[end], digits=10) == benchmark_X_fc[t];
        @test round.(kforecast(ksettings, kstatus.X_post, kstatus.P_post, 12)[1][end], digits=10) == benchmark_X_fc[t];
        @test round.(kforecast(ksettings, kstatus.X_post, kstatus.P_post, 12)[2][end], digits=10) == benchmark_P_fc[t];
    end

    # Final value of the loglikelihood
    @test round.(kstatus.loglik, digits=10) == benchmark_loglik;
end

function test_kalman_output(ksettings::KalmanSettings, kstatus::SizedKalmanStatus, benchmark_data::Tuple)

    benchmark_n, benchmark_T, benchmark_m, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior,
        benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik, benchmark_X0_sm, benchmark_P0_sm,
            benchmark_X_sm, benchmark_P_sm = benchmark_data;

    for t=1:ksettings.Y.T

        # Run filter
        kfilter!(ksettings, kstatus);

        # A-priori
        @test round.(kstatus.online_status.X_prior, digits=10) == benchmark_X_prior[t];
        @test round.(kstatus.online_status.P_prior, digits=10) == benchmark_P_prior[t];

        # A-posteriori
        @test round.(kstatus.online_status.X_post, digits=10) == benchmark_X_post[t];
        @test round.(kstatus.online_status.P_post, digits=10) == benchmark_P_post[t];

        # 12-step ahead forecast
        @test round.(kforecast(ksettings, kstatus.online_status.X_post, 12)[end], digits=10) == benchmark_X_fc[t];
        @test round.(kforecast(ksettings, kstatus.online_status.X_post, kstatus.online_status.P_post, 12)[1][end], digits=10) == benchmark_X_fc[t];
        @test round.(kforecast(ksettings, kstatus.online_status.X_post, kstatus.online_status.P_post, 12)[2][end], digits=10) == benchmark_P_fc[t];
    end

    # Final value of the loglikelihood
    @test round.(kstatus.online_status.loglik, digits=10) == benchmark_loglik;
end

"""
    kalman_test(Y::JArray, B::FloatMatrix, R::Union{UniformScaling{Float64}, SymMatrix}, C::FloatMatrix, D::FloatMatrix, Q::SymMatrix, benchmark_data::Tuple)

Run a series of tests to check whether the kalman.jl functions work.
"""
function kalman_test(Y::JArray, B::FloatMatrix, R::Union{UniformScaling{Float64}, SymMatrix}, C::FloatMatrix, D::FloatMatrix, Q::SymMatrix, benchmark_data::Tuple)

    # Benchmark data
    benchmark_n, benchmark_T, benchmark_m, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior,
        benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik, benchmark_X0_sm, benchmark_P0_sm,
            benchmark_X_sm, benchmark_P_sm = benchmark_data;

    for test_family_id=1:3

        # Tests on KalmanSettings (min. number of arguments)
        if test_family_id == 1
            D_I = Matrix(1.0I, benchmark_m, benchmark_m);
            test_family_input = (Y, B, R, C, benchmark_DQD);
            test_family_benchmark = (Y, B, R, C, D_I, benchmark_DQD, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_n, benchmark_T, benchmark_m);

        # Tests on KalmanSettings (std. number of arguments)
        elseif test_family_id == 2
            test_family_input = (Y, B, R, C, D, Q);
            test_family_benchmark = (Y, B, R, C, D, Q, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_n, benchmark_T, benchmark_m);

        # Tests on KalmanSettings (full number of arguments)
        elseif test_family_id == 3
            test_family_input = (Y, B, R, C, D, Q, benchmark_X0, benchmark_P0);
            test_family_benchmark = (Y, B, R, C, D, Q, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_n, benchmark_T, benchmark_m);
        end

        # Tests on KalmanSettings (full number of arguments)
        ksettings1 = KalmanSettings(test_family_input..., compute_loglik=true, store_history=true);
        @test ksettings_input_test(ksettings1, test_family_benchmark..., compute_loglik=true, store_history=true);

        ksettings2 = KalmanSettings(test_family_input..., compute_loglik=false, store_history=true);
        @test ksettings_input_test(ksettings2, test_family_benchmark..., compute_loglik=false, store_history=true);

        ksettings3 = KalmanSettings(test_family_input..., compute_loglik=true, store_history=false);
        @test ksettings_input_test(ksettings3, test_family_benchmark..., compute_loglik=true, store_history=false);

        ksettings4 = KalmanSettings(test_family_input..., compute_loglik=false, store_history=false);
        @test ksettings_input_test(ksettings4, test_family_benchmark..., compute_loglik=false, store_history=false);

        ksettings5 = KalmanSettings(test_family_input...);
        @test ksettings_input_test(ksettings5, test_family_benchmark...);

        # Select default ksettings
        ksettings = ksettings5;

        # Initialise kstatus
        for kstatus in [DynamicKalmanStatus(), SizedKalmanStatus(size(Y,2))]

            test_kalman_output(ksettings, kstatus, benchmark_data);

            # Kalman smoother
            X_sm, P_sm, X0_sm, P0_sm = ksmoother(ksettings, kstatus);

            for t=1:ksettings.Y.T
                @test round.(X_sm[t], digits=10) == benchmark_X_sm[t];
                @test round.(P_sm[t], digits=10) == benchmark_P_sm[t];
            end

            @test round.(X0_sm, digits=10) == benchmark_X0_sm;
            @test round.(P0_sm, digits=10) == benchmark_P0_sm;
        end
    end
end

@testset "univariate model" begin

    # Initialise data and state-space parameters
    Y = [0.35 0.62 missing missing 1.11 missing 2.76 2.73 3.45 3.66];
    B = ones(1,1);
    R = Symmetric(1e-4*ones(1,1));
    C = 0.9*ones(1,1);
    D = ones(1,1);
    Q = Symmetric(ones(1,1));

    # Correct estimates: DQD
    benchmark_DQD = copy(Q);

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/univariate/benchmark_X0");
    benchmark_P0 = Symmetric(read_test_input("./input/univariate/benchmark_P0"));

    # Correct estimates: a priori
    benchmark_X_prior = read_test_input("./input/univariate/benchmark_X_prior");
    benchmark_P_prior = read_test_input("./input/univariate/benchmark_P_prior");

    # Correct estimates: a posteriori
    benchmark_X_post = read_test_input("./input/univariate/benchmark_X_post");
    benchmark_P_post = read_test_input("./input/univariate/benchmark_P_post");

    # Correct estimates: 12-step ahead forecast
    benchmark_X_fc = read_test_input("./input/univariate/benchmark_X_fc");
    benchmark_P_fc = read_test_input("./input/univariate/benchmark_P_fc");

    # Correct estimates: loglikelihood
    benchmark_loglik = read_test_input("./input/univariate/benchmark_loglik")[1];

    # Correct estimates: kalman smoother (smoothed initial conditions)
    benchmark_X0_sm = read_test_input("./input/univariate/benchmark_X0_sm");
    benchmark_P0_sm = read_test_input("./input/univariate/benchmark_P0_sm");

    # Correct estimates: kalman smoother
    benchmark_X_sm = read_test_input("./input/univariate/benchmark_X_sm");
    benchmark_P_sm = read_test_input("./input/univariate/benchmark_P_sm");

    # Benchmark data
    benchmark_data = (1, size(Y,2), 1, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior, benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik,
                      benchmark_X0_sm, benchmark_P0_sm, benchmark_X_sm, benchmark_P_sm);

    # Run tests
    kalman_test(Y, B, R, C, D, Q, benchmark_data);
    kalman_test(Y, B, 1e-4I, C, D, Q, benchmark_data); # same test with UniformScaling R
end

@testset "multivariate model" begin

    # Initialise data and state-space parameters
    Y = [0.72 missing 1.86 missing missing 2.52 2.98 3.81 missing 4.36;
         0.95 0.70 missing missing missing missing 2.84 3.88 3.84 4.63];

    B = [1.0 0.0; 1.0 1.0];
    R = Symmetric(1e-4*Matrix(I,2,2));
    C = [0.9 0.0; 0.0 0.1];
    D = Matrix(I,2,2) |> FloatMatrix;
    Q = Symmetric(D);

    # Correct estimates: DQD
    benchmark_DQD = copy(Q);

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/multivariate/benchmark_X0");
    benchmark_P0 = Symmetric(read_test_input("./input/multivariate/benchmark_P0"));

    # Correct estimates: a priori
    benchmark_X_prior = read_test_input("./input/multivariate/benchmark_X_prior");
    benchmark_P_prior = read_test_input("./input/multivariate/benchmark_P_prior");

    # Correct estimates: a posteriori
    benchmark_X_post = read_test_input("./input/multivariate/benchmark_X_post");
    benchmark_P_post = read_test_input("./input/multivariate/benchmark_P_post");

    # Correct estimates: 12-step ahead forecast
    benchmark_X_fc = read_test_input("./input/multivariate/benchmark_X_fc");
    benchmark_P_fc = read_test_input("./input/multivariate/benchmark_P_fc");

    # Correct estimates: loglikelihood
    benchmark_loglik = read_test_input("./input/multivariate/benchmark_loglik")[1];

    # Correct estimates: kalman smoother (smoothed initial conditions)
    benchmark_X0_sm = read_test_input("./input/multivariate/benchmark_X0_sm");
    benchmark_P0_sm = read_test_input("./input/multivariate/benchmark_P0_sm");

    # Correct estimates: kalman smoother
    benchmark_X_sm = read_test_input("./input/multivariate/benchmark_X_sm");
    benchmark_P_sm = read_test_input("./input/multivariate/benchmark_P_sm");

    # Benchmark data
    benchmark_data = (2, size(Y,2), 2, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior, benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik,
                      benchmark_X0_sm, benchmark_P0_sm, benchmark_X_sm, benchmark_P_sm);

    # Run tests
    kalman_test(Y, B, R, C, D, Q, benchmark_data);
    kalman_test(Y, B, 1e-4I, C, D, Q, benchmark_data); # same test with UniformScaling R
end

@testset "multivariate model (non-diagonal)" begin

    # Initialise data and state-space parameters
    Y = [0.49 missing 1.01	-0.41 1.33 1.5 -0.36 -1.14 0.11 missing;
         0.02 missing 0.20 0.14 missing 0.85 missing -0.57 missing 0.11;
         missing missing 1.15 -0.65 1.32 1.28 -0.57 -1.03 0.14 1.53];

    B = [1 0; 0.5 -0.25; 0.9 0.2];
    R = Symmetric(1e-4*Matrix(I,3,3));
    C = [0.6 0.1; 0.25 -0.2];
    D = [1.0 0.0; 0.5 0.86602540];
    Q = Symmetric(1.0*Matrix(I,2,2));

    # Correct estimates: DQD
    benchmark_DQD = Symmetric(D*Q*D');

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/multivariate_non_diagonal/benchmark_X0");
    benchmark_P0 = Symmetric(read_test_input("./input/multivariate_non_diagonal/benchmark_P0"));

    # Correct estimates: a priori
    benchmark_X_prior = read_test_input("./input/multivariate_non_diagonal/benchmark_X_prior");
    benchmark_P_prior = read_test_input("./input/multivariate_non_diagonal/benchmark_P_prior");

    # Correct estimates: a posteriori
    benchmark_X_post = read_test_input("./input/multivariate_non_diagonal/benchmark_X_post");
    benchmark_P_post = read_test_input("./input/multivariate_non_diagonal/benchmark_P_post");

    # Correct estimates: 12-step ahead forecast
    benchmark_X_fc = read_test_input("./input/multivariate_non_diagonal/benchmark_X_fc");
    benchmark_P_fc = read_test_input("./input/multivariate_non_diagonal/benchmark_P_fc");

    # Correct estimates: loglikelihood
    benchmark_loglik = read_test_input("./input/multivariate_non_diagonal/benchmark_loglik")[1];

    # Correct estimates: kalman smoother (smoothed initial conditions)
    benchmark_X0_sm = read_test_input("./input/multivariate_non_diagonal/benchmark_X0_sm");
    benchmark_P0_sm = read_test_input("./input/multivariate_non_diagonal/benchmark_P0_sm");

    # Correct estimates: kalman smoother
    benchmark_X_sm = read_test_input("./input/multivariate_non_diagonal/benchmark_X_sm");
    benchmark_P_sm = read_test_input("./input/multivariate_non_diagonal/benchmark_P_sm");

    # Benchmark data
    benchmark_data = (3, size(Y,2), 2, benchmark_X0, benchmark_P0, benchmark_DQD, benchmark_X_prior, benchmark_P_prior, benchmark_X_post, benchmark_P_post, benchmark_X_fc, benchmark_P_fc, benchmark_loglik,
                      benchmark_X0_sm, benchmark_P0_sm, benchmark_X_sm, benchmark_P_sm);

    # Run tests
    kalman_test(Y, B, R, C, D, Q, benchmark_data);
    kalman_test(Y, B, 1e-4I, C, D, Q, benchmark_data); # same test with UniformScaling R
end