text:
"""
An AbstractYield is an object which can be used as an argument to:

- zero-coupon spot rates via [`zero`](@ref)
- forward zero rates via [`forward`](@ref)
- discount factor via [`discount`](@ref)
- accumulation factor via [`accumulation`](@ref)

It can be be constructed via:

- zero rate curve with [`Zero`](@ref)
- forward rate curve with [`Forward`](@ref)
- par rate curve with [`Par`](@ref)
- typical OIS curve with [`OIS`](@ref)
- typical constant maturity treasury (CMT) curve with [`CMT`](@ref)
"""
abstract type AbstractYield end

abstract type AbstractYieldCurve <: AbstractYield end

"""
A `YieldCurveFitParameters` is a structure which contains associated parameters for a yield curve fitting procedure. The type of the object determines the method, and the values of the object determine the parameters.

If the fitting data and the rates are passed as `<:Real` numbers instead of a type of `Rate`s, the default interpretation may vary depending on the fitting type/parameter. See the individual docstrings of the types for more information.

Available types are:

- [`Bootstrap`](@ref)
- [`NelsonSiegel`](@ref)
- [`NelsonSiegelSvensson`](@ref)
"""
abstract type YieldCurveFitParameters end

Base.Broadcast.broadcastable(x::T) where {T<:YieldCurveFitParameters} = Ref(x)


text:
#========================================================================================#
"""
	Objectives

A collection of displayable objective functions from De Jong and Watson.

Author: Niall Palfreyman, 10/3/2022.
"""
module Objectives

export Objective, dim, domain, sample, evaluate, depict

using CairoMakie

#-----------------------------------------------------------------------------------------
# Module types:

"""
	Objective

An Objective encapsulates an objective function for GAs. An Objective instance is a
function f: C^n -> R together with a (by default infinite) domain and dimensionality.
"""
struct Objective
	functn::Function			# The objective function
	domain::Vector{Vector}		# Domain specification
end

"Construct a new Objective from the test suite with default dimension and donain."
function Objective( fun::Int = 1)
	if ~(fun in 1:length(testSuite))
		fun = 1
	end
	Objective(testSuite[fun]...)
end

"Construct a new Objective from the test suite with newly provided donain."
function Objective( fun::Int, dom, dim::Int=1)
	if ~(fun in 1:length(testSuite))
		fun = 1
	end
	Objective(testSuite[fun][1],dom,dim)
end

"""
	Objective(fun,dom,dim)

Create an Objective function with the given domain. If the domain is a vector of two
Numbers, it is a range that should be replicated over several dimensions; otherwise,
if it is a vector of such ranges, it is a complete domain and the dimensionality should
be ignored.
"""
function Objective( fun::Function, dom::Vector, dim::Int)
	if dom[1] isa Vector
		# This is a ready-made domain - ignore the given dimension:
		if any(length.(dom) .!= 2) || any(first.(a).>=last.(a))
			# The domain structure is invalid - abort:
			error( "Domain specification is invalid")
		end
		Objective( fun, dom)
	else
		# This a 1-dimensional range that should be replicated over dim dimensions:
		if length(dom) != 2 || dom[1]>=dom[2] || dim < 1
			# The domain structure is invalid - abort:
			error( "Single domain/dimension specification is invalid")
		end
		Objective( fun, [dom for _ in 1:dim])
	end
end

#-----------------------------------------------------------------------------------------
# Module methods:

#-----------------------------------------------------------------------------------------
"Return dimension of Objective function domain."
function dim( obj::Objective)
	length(obj.domain)
end

#-----------------------------------------------------------------------------------------
"Return domain of Objective function."
function domain( obj::Objective)
	obj.domain
end

#-----------------------------------------------------------------------------------------
"Return a sample vector of n random points in function domain."
function sample( obj::Objective, n::Int=1)
	if any( abs.(collect(Iterators.flatten(obj.domain))).==Inf)
		error("Cannot sample an infinite domain")
	end

	dim = length(obj.domain)
	lo = (x->x[1]).(obj.domain)
	dif = (x->x[2]).(obj.domain) - lo

	[lo + dif.*rand(dim) for _ in 1:n]
end

#-----------------------------------------------------------------------------------------
"Evaluate Objective function for the given collection of arguments."
function evaluate( obj::Objective, args)
	obj.functn.(args)
end

#-----------------------------------------------------------------------------------------
"""
	depict( obj, dims, centre, radius)

Depict an objective function in the specified dimensions over the Manhattan
neighbourhood with the given centre and radius.
"""
function depict( obj::Objective, dims=nothing, centre=nothing, radius=nothing)
	if dims === nothing
		# Depict up to the first two dimensions:
		dims = 1:min(2,length(obj.domain))
	end

	if isempty(dims)
		# If dims==[], plot a slice through all dimensions of Objective
		# (Mainly for plotting mepi()):
		ndims = length(obj.domain)
		args = [
			[first.(obj.domain)];
			[[last.(obj.domain[1:d]);first.(obj.domain[d+1:ndims])] for d in 1:ndims-1];
			[last.(obj.domain)]
		]
		return lines( 0:ndims, evaluate( obj, args))
	end

	# dims is from this point on correctly defined.
	if centre === nothing
		if any((rng->any(abs.(rng).==Inf)).(obj.domain))
			error( "Cannot display an infinite domain")
			return
		end
		centre = sum.(obj.domain)/2
		radius = (last.(obj.domain)-first.(obj.domain))/2
	else
		if radius === nothing
			radius = ones(length(centre))
		end
	end

	# All parameters are now present and correct:
	nGrad = 50			# Number of gradations in depiction
	lwb = max(first.(obj.domain[dims]),(centre[dims]-radius[dims]))
	upb = min( last.(obj.domain[dims]),(centre[dims]+radius[dims]))

	ndims = length(dims)
	if ndims == 1
		# Use plot for single dimension:
		xs = range( lwb[1], upb[1], length = nGrad)
		args = [copy(centre) for _ in 1:nGrad]
		for n in 1:nGrad
			args[n][dims] .= xs[n]
		end
		lines!(xs,evaluate( obj, args))
	elseif ndims == 2
		# Construct arguments for a contour plot:
		xs = range( lwb[1], upb[1], length=nGrad)
		ys = range( lwb[2], upb[2], length=nGrad)
		args = [copy(centre) for _ in 1:nGrad, _ in 1:nGrad]
		for i in 1:nGrad, j in 1:nGrad
			args[i,j][dims[1]] = xs[i]
			args[i,j][dims[2]] = ys[j]
		end
		zs = evaluate(obj,args)
		contourf!(  xs, ys, zs)
	else
		error( "Cannot depict more than 2 dimensions")
	end
end

#-----------------------------------------------------------------------------------------
"""
	mepi( x)

Watson's maximally epistatic objective function.
"""
function mepi( x)
	dim = length(x)

	if dim == 1
		1
	else
		# Form product of the first and second halves of x separately:
		halflen = div(dim,2)
		dim*(1 - prod(x) - prod(1 .- x)) + mepi(x[1:halflen]) + mepi(x[halflen+1:end])
	end
end

#-----------------------------------------------------------------------------------------
"""
	demo1()

A very simple demonstration of the Objective type.
"""
function demo1()
	fig = Figure(resolution=(300, 300))
	Axis( fig[1,1], xlabel="x", ylabel="y", title="Suite function 7")
	depict( Objective(7))
	fig
end

#-----------------------------------------------------------------------------------------
"""
	demo2()

Demo the Objective type.
"""
function demo2()
	fig = Figure(resolution=(600, 300))

	f(x) =
		3*(1-x[1])^2 * exp(-(x[1]^2) - (x[2]+1)^2) -
		10*(x[1]/5 - x[1]^3 - x[2]^5) * exp(-x[1]^2-x[2]^2) -
		1/3*exp(-(x[1]+1)^2 - x[2]^2)
	obj1 = Objective(f,[-3,3],2)
	Axis( fig[1,1], xlabel="x", ylabel="y", title="Valleys")
	depict( obj1)

	Axis( fig[1,2], xlabel="x", ylabel="y", title="Suite function 7")
	depict( Objective(7))

	fig
end

#-----------------------------------------------------------------------------------------
# Module data:

"""
	testSuite

A vector of standard test functions together with appropriate domains. Functions 1-17 are
roughly drawn from De Jong. Function 18 is Watson's (2004) maximally epistatic function.
"""
testSuite = [
	# 1. Minimum: f(0) = 1:
	(x -> abs(x[1]) + cos(x[1]), [-20,20], 1),
	# 2. Minimum: f(0) = 0:
	(x -> abs(x[1]) + sin(x[1]), [-20,20], 1),
	# 3. Minimum: f(0,0) = 1:
	(x -> sum(x.^2), [-5,5], 2),
	# 4. Minimum: f(1,1) = 0:
	(x -> sum(100*(x[2:end]-x[1:end-1].^2)^2 + (1-x[1:end-1])^2), [-1,1], 2),
	# 5. Minimum: f(0^dim) = 1:
	(x -> sum(abs(x) - 10*cos(sqrt(abs(10*x))),1), [-10,10], 2),
	# 6. Minimum: f(9.6204) = -100.22:
	(x -> (x^2+x)*cos(x), [-10,10], 1),
	# 7. Minimum: f(0.9039,0.8668) = -18.5547:
	(x -> x[1]*sin(4*x[1]) + 1.1*x[2]*sin(2*x[2]), [0,10], 2),
	# 8. Minimum: f(0.9039,0.8668) = -18.5547:
	(x -> x[2]*sin(4*x[1]) + 1.1*x[1]*sin(2*x[2]), [0,10], 2),
	# 9. Minimum: varies:
	(x -> (1:length(x))'*(x.^4) + randn(), [-2,2], 2),
	# 10. Minimum: f(0,0) = 0:
	(x -> 10*length(x) + sum(x.^2-10*cos.(2*pi*x)), [-4,4], 2),
	# 11. Minimum: f(0,0) = 0:
	(x -> 1 + sum(x.^2/4000) - prod(cos.(x)), [-10,10], 2),
	# 12. Minimum: f(1.897,1.006) = -0.5231:
	(x -> 0.5 + (sin(sqrt(sum(x.^2))).^2 - 0.5) ./ (1 + 0.1*sum(x.^2)), [-5,5], 2),
	# 13. Minimum: f(0,0) = 0:
	(x -> sum(abs(x),1) + sum(x.^2).^0.25 .* sin(30*(((x[1]+0.5)^2+x[2]^2).^0.1)), [-10,10], 2),
	# 14. Minimum: f(1,16606) = -0.3356:
	(x -> besselj(0,sum(x.^2)) + 0.1*sum(abs(1-x)), [-5,5], 2),
	# 15. Minimum: f(-14.58,-20) = -23.806:
	(	x -> x[1]*sin(sqrt(abs(x[1]-(x[2]+9)))) - (x[2]+9).*sin(sqrt(abs(x[2]+0.5*x[1]+9))),
		[-20,20], 2
	),
	# 16. Watson's maximally epistatic function:
	(mepi, [0,1], 128),
]

end		# of Objectives

ext: jl


text:
# This file contains special cases that help make the library more useful

# Matrix-matrix multiplication
function *{T}(lhs::VectorizedExpression{T,2}, rhs::VectorizedExpression{T,2})
    return ETContainer( (ETContainer(lhs).data)*(ETContainer(rhs).data) )
end

# Matrix-vector multiplication
function *{T}(lhs::VectorizedExpression{T,2}, rhs::VectorizedExpression{T,1})
    return ETContainer( (ETContainer(lhs).data)*(ETContainer(rhs).data) )
end

export *


text:
using Pkg

## add necessary Julia Packets
Pkg.add("JLD")
Pkg.add("Plots")
Pkg.add("Gadfly")
Pkg.add("Debugger")
Pkg.add("PyPlot")
Pkg.add("GR")
Pkg.add("ProgressMeter")
Pkg.add("Combinatorics")
Pkg.add("ArgParse")
Pkg.add("DelimitedFiles")
Pkg.add("Compose")
Pkg.add("Random")
Pkg.add("Cairo")
Pkg.add("Gurobi")
Pkg.add("HDF5")
Pkg.add("LightGraphs")
Pkg.add("HypothesisTests")
Pkg.add("GraphPlot")
Pkg.add("Dates")
Pkg.add("JuMP")
Pkg.add("Distributions")

Pkg.build()

## initialize the results folder
if !isdir("./data/experiment")
    println("Initializing results folder ./data/experiment.")
    mkdir("./data/experiment")
else
    println("Results folder ./data/experiment exists.")
end

if !isdir("./data/parallel_experiment")
    println("Initializing results folder ./data/parallel_experiment")
    mkdir("./data/parallel_experiment")
else
    println("Results folder ./data/parallel_experiment exists.")
end

max_stars_repo_name: hipersys-team/arrow

max_issues_repo_name: arrow-bow/arrow

max_forks_repo_name: arrow-bow/arrow


text:
abstract type ObservableQuote end

"""
    ZeroCouponQuote(price, maturity)

Quote for a set of zero coupon bonds with given `price` and `maturity`.

# Examples

```julia-repl
julia> prices = [1.3, 0.1, 4.5]
julia> maturities = [1.2, 2.5, 3.6]
julia> swq = Yields.ZeroCouponQuote.(prices, maturities)
```
"""
struct ZeroCouponQuote <: ObservableQuote
    price
    maturity
end

"""
    SwapQuote(yield, maturity, frequency)

Quote for a set of interest rate swaps with the given `yield` and `maturity` and a given payment `frequency`.

# Examples

```julia-repl
julia> maturities = [1.2, 2.5, 3.6]
julia> interests = [-0.02, 0.3, 0.04]
julia> prices = [1.3, 0.1, 4.5]
julia> frequencies = [2,1,2]
julia> swq = Yields.SwapQuote.(interests, maturities, frequencies)
```
"""
struct SwapQuote <: ObservableQuote
    yield
    maturity
    frequency
    function SwapQuote(yield, maturity, frequency)
        frequency <= 0 && throw(DomainError("Payment frequency must be positive"))
        return new(yield, maturity, frequency)
    end
end

"""
    BulletBondQuote(yield, price, maturity, frequency)

Quote for a set of fixed interest bullet bonds with given `yield`, `price`, `maturity` and a given payment frequency `frequency`.

Construct a vector of quotes for use with SmithWilson methods, e.g. by broadcasting over an array of inputs.

# Examples

```julia-repl
julia> maturities = [1.2, 2.5, 3.6]
julia> interests = [-0.02, 0.3, 0.04]
julia> prices = [1.3, 0.1, 4.5]
julia> frequencies = [2,1,2]
julia> bbq = Yields.BulletBondQuote.(interests, maturities, prices, frequencies)
```
"""
struct BulletBondQuote <: ObservableQuote
    yield
    price
    maturity
    frequency

    function BulletBondQuote(yield, maturity, price, frequency)
        frequency <= 0 && throw(DomainError("Payment frequency must be positive"))
        return new(yield, maturity, price, frequency)
    end
end

"""
    SmithWilson(zcq::Vector{ZeroCouponQuote}; ufr, α)
    SmithWilson(swq::Vector{SwapQuote}; ufr, α)
    SmithWilson(bbq::Vector{BulletBondQuote}; ufr, α)
    SmithWilson(times<:AbstractVector, cashflows<:AbstractMatrix, prices<:AbstractVector; ufr, α)
    SmithWilson(u, qb; ufr, α)

Create a yield curve object that implements the Smith-Wilson interpolation/extrapolation scheme.

Positional arguments to construct a curve:

- Quoted instrument as the first argument: either a `Vector` of `ZeroCouponQuote`s, `SwapQuote`s, or `BulletBondQuote`s, or
- A set of `times`, `cashflows`, and `prices`, or
- A curve can be with `u` is the timepoints coming from the calibration, and `qb` is the internal parameterization of the curve that ensures that the calibration is correct. Users may prefer the other constructors but this mathematical constructor is also available.

Required keyword arguments:

- `ufr` is the Ultimate Forward Rate, the forward interest rate to which the yield curve tends, in continuous compounding convention.
- `α` is the parameter that governs the speed of convergence towards the Ultimate Forward Rate. It can be typed with `\\alpha[TAB]`
"""
struct SmithWilson{TU<:AbstractVector,TQb<:AbstractVector} <: AbstractYieldCurve
    u::TU
    qb::TQb
    ufr
    α

    # Inner constructor ensures that vector lengths match
    function SmithWilson{TU,TQb}(u, qb; ufr, α) where {TU<:AbstractVector,TQb<:AbstractVector}
        if length(u) != length(qb)
            throw(DomainError("Vectors u and qb in SmithWilson must have equal length"))
        end
        return new(u, qb, ufr, α)
    end
end

SmithWilson(u::TU, qb::TQb; ufr, α) where {TU<:AbstractVector,TQb<:AbstractVector} = SmithWilson{TU,TQb}(u, qb; ufr = ufr, α = α)

__ratetype(::Type{SmithWilson{TU,TQb}}) where {TU,TQb}= Yields.Rate{Float64, Yields.Continuous}

"""
    H_ordered(α, t_min, t_max)

The Smith-Wilson H function with ordered arguments (for better performance than using min and max).
"""
function H_ordered(α, t_min, t_max)
    return α * t_min + exp(-α * t_max) * sinh(-α * t_min)
end

"""
    H(α, t1, t2)

The Smith-Wilson H function implemented in a faster way.
"""
function H(α, t1::T, t2::T) where {T}
    return t1 < t2 ? H_ordered(α, t1, t2) : H_ordered(α, t2, t1)
end

H(α, t1, t2) = H(α, promote(t1, t2)...)

H(α, t1vec::AbstractVector, t2) = [H(α, t1, t2) for t1 in t1vec]
H(α, t1vec::AbstractVector, t2vec::AbstractVector) = [H(α, t1, t2) for t1 in t1vec, t2 in t2vec]
# This can be optimized by going to H_ordered directly, but it might be a bit cumbersome
H(α, tvec::AbstractVector) = H(α, tvec, tvec)

discount(sw::SmithWilson, t) = exp(-sw.ufr * t) * (1.0 + H(sw.α, sw.u, t) ⋅ sw.qb)
Base.zero(sw::SmithWilson, t) = Continuous(sw.ufr - log(1.0 + H(sw.α, sw.u, t) ⋅ sw.qb) / t)
Base.zero(sw::SmithWilson, t, cf::CompoundingFrequency) = convert(cf, zero(sw, t))

function SmithWilson(times::AbstractVector, cashflows::AbstractMatrix, prices::AbstractVector; ufr, α)
    Q = Diagonal(exp.(-ufr * times)) * cashflows
    q = vec(sum(Q, dims = 1))  # We want q to be a column vector
    QHQ = Q' * H(α, times) * Q
    b = QHQ \ (prices - q)
    Qb = Q * b
    return SmithWilson(times, Qb; ufr = ufr, α = α)
end

"""
    timepoints(zcq::Vector{ZeroCouponQuote})
    timepoints(bbq::Vector{BulletBondQuote})

Return the times associated with the `cashflows` of the instruments.
"""
function timepoints(qs::Vector{Q}) where {Q<:ObservableQuote}
    frequency = maximum(q.frequency for q in qs)
    timestep = 1 / frequency
    maturity = maximum(q.maturity for q in qs)
    return [timestep:timestep:maturity...]
end

"""
    cashflows(interests, maturities, frequency)
    timepoints(zcq::Vector{ZeroCouponQuote})
    timepoints(bbq::Vector{BulletBondQuote})

Produce a cash flow matrix for a set of instruments with given `interests` and `maturities`
and a given payment frequency `frequency`. All instruments are assumed to have their first payment at time 1/`frequency`
and have their last payment at the largest multiple of 1/`frequency` less than or equal to the input maturity.
"""
function cashflows(interests, maturities, frequencies)
    frequency = lcm(frequencies)
    fq = inv.(frequencies)
    timestep = 1 / frequency
    floored_mats = floor.(maturities ./ timestep) .* timestep
    times = timestep:timestep:maximum(floored_mats)
    # we need to determine the coupons in relation to the payment date, not time zero
    time_adj = floored_mats .% fq

    cashflows = [
        # if on a coupon date and less than maturity, pay coupon
        ((((t + time_adj[instrument]) % fq[instrument] ≈ 0) && t <= floored_mats[instrument]) ? interests[instrument] / frequencies[instrument] : 0.0) +
        (t ≈ floored_mats[instrument] ? 1.0 : 0.0) # add maturity payment
        for t in times, instrument = 1:length(interests)
    ]

    return cashflows
end

function cashflows(qs::Vector{Q}) where {Q<:ObservableQuote}
    yield = [q.yield for q in qs]
    maturity = [q.maturity for q in qs]
    frequency = [q.frequency for q in qs]
    return cashflows(yield, maturity, frequency)
end

# Utility methods for calibrating Smith-Wilson directly from quotes
function SmithWilson(zcq::Vector{ZeroCouponQuote}; ufr, α)
    n = length(zcq)
    maturities = [q.maturity for q in zcq]
    prices = [q.price for q in zcq]
    return SmithWilson(maturities, Matrix{Float64}(I, n, n), prices; ufr = ufr, α = α)
end

function SmithWilson(swq::Vector{SwapQuote}; ufr, α)
    times = timepoints(swq)
    cfs = cashflows(swq)
    ones(length(swq))
    return SmithWilson(times, cfs, ones(length(swq)), ufr = ufr, α = α)
end

function SmithWilson(bbq::Vector{BulletBondQuote}; ufr, α)
    times = timepoints(bbq)
    cfs = cashflows(bbq)
    prices = [q.price for q in bbq]
    return SmithWilson(times, cfs, prices, ufr = ufr, α = α)
end


text:
"""
    away_frank_wolfe

Frank-Wolfe with away steps.
The algorithm maintains the current iterate as a convex combination of vertices in the
[`FrankWolfe.ActiveSet`](@ref) data structure.
See the [paper](https://arxiv.org/abs/2104.06675) for illustrations of away steps.
"""
function away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0;
    line_search::LineSearchMethod=Adaptive(),
    L=Inf,
    gamma0=0,
    K=2.0,
    step_lim=20,
    epsilon=1e-7,
    away_steps=true,
    lazy=false,
    momentum=nothing,
    max_iteration=10000,
    print_iter=1000,
    trajectory=false,
    verbose=false,
    linesearch_tol=1e-7,
    emphasis::Emphasis=memory,
    gradient=nothing,
    renorm_interval=1000,
    callback=nothing,
    timeout=Inf,
    print_callback=print_callback,
)
    # add the first vertex to active set from initialization
    active_set = ActiveSet([(1.0, x0)])

    # Call the method using an ActiveSet as input
    return away_frank_wolfe(
        f,
        grad!,
        lmo,
        active_set,
        line_search = line_search,
        L=L,
        gamma0=gamma0,
        K=K,
        step_lim=step_lim,
        epsilon=epsilon,
        away_steps=away_steps,
        lazy=lazy,
        momentum=momentum,
        max_iteration=max_iteration,
        print_iter=print_iter,
        trajectory=trajectory,
        verbose=verbose,
        linesearch_tol=linesearch_tol,
        emphasis=emphasis,
        gradient=gradient,
        renorm_interval=renorm_interval,
        callback=callback,
        timeout= timeout,
        print_callback=print_callback,
    )
end

# step away FrankWolfe with the active set given as parameter
# note: in this case I don't need x0 as it is given by the active set and might otherwise lead to confusion
function away_frank_wolfe(
    f,
    grad!,
    lmo,
    active_set::ActiveSet;
    line_search::LineSearchMethod=Adaptive(),
    L=Inf,
    gamma0=0,
    K=2.0,
    step_lim=20,
    epsilon=1e-7,
    away_steps=true,
    lazy=false,
    momentum=nothing,
    max_iteration=10000,
    print_iter=1000,
    trajectory=false,
    verbose=false,
    linesearch_tol=1e-7,
    emphasis::Emphasis=memory,
    gradient=nothing,
    renorm_interval=1000,
    callback=nothing,
    timeout=Inf,
    print_callback=print_callback,
)
# format string for output of the algorithm
    format_string = "%6s %13s %14e %14e %14e %14e %14e %14i\n"

    if isempty(active_set)
        throw(ArgumentError("Empty active set"))
    end

    t = 0
    dual_gap = Inf
    primal = Inf
    x = compute_active_set_iterate(active_set)
    #  not need anymore active_set = ActiveSet([(1.0, x0)]) # add the first vertex to active set from initialization
    tt = regular
    traj_data = []
    if trajectory && callback === nothing
        callback = trajectory_callback(traj_data)
    end
    time_start = time_ns()

    d = similar(x)

    if line_search isa Shortstep && L == Inf
        println("WARNING: Lipschitz constant not set. Prepare to blow up spectacularly.")
    end

    if line_search isa FixedStep && gamma0 == 0
        println("WARNING: gamma0 not set. We are not going to move a single bit.")
    end

    if verbose
        println("\nAway-step Frank-Wolfe Algorithm.")
        numType = eltype(x)
        println(
            "EMPHASIS: $emphasis STEPSIZE: $line_search EPSILON: $epsilon MAXITERATION: $max_iteration TYPE: $numType",
        )
        grad_type = typeof(gradient)
        println(
            "GRADIENTTYPE: $grad_type LAZY: $lazy K: $K MOMENTUM: $momentum AWAYSTEPS: $away_steps",
        )
        if emphasis == memory
            println("WARNING: In memory emphasis mode iterates are written back into x0!")
        end
        headers =
            ("Type", "Iteration", "Primal", "Dual", "Dual Gap", "Time", "It/sec", "#ActiveSet")
        print_callback(headers, format_string, print_header=true)
    end

    # likely not needed anymore as now the iterates are provided directly via the active set
    if gradient === nothing
        gradient = similar(x)
    end
    gtemp = if momentum !== nothing
        similar(gradient)
    else
        nothing
    end

    x = compute_active_set_iterate(active_set)
    grad!(gradient, x)
    v = compute_extreme_point(lmo, gradient)
    phi_value = max(0, fast_dot(x, gradient) - fast_dot(v, gradient))
    gamma = 1.0

    while t <= max_iteration && phi_value >= max(epsilon, eps())

        #####################
        # managing time and Ctrl-C
        #####################
        time_at_loop = time_ns()
        if t == 0
            time_start = time_at_loop
        end
        # time is measured at beginning of loop for consistency throughout all algorithms
        tot_time = (time_at_loop - time_start) / 1e9

        if timeout < Inf
            if tot_time ≥ timeout
                if verbose
                    @info "Time limit reached"
                end
                break
            end
        end

        #####################

        # compute current iterate from active set
        x = compute_active_set_iterate(active_set)
        if isnothing(momentum)
            grad!(gradient, x)
        else
            grad!(gtemp, x)
            @emphasis(emphasis, gradient = (momentum * gradient) + (1 - momentum) * gtemp)
        end

        if away_steps
            if lazy
                d, vertex, index, gamma_max, phi_value, away_step_taken, fw_step_taken, tt =
                    lazy_afw_step(x, gradient, lmo, active_set, phi_value, epsilon; K=K)
            else
                d, vertex, index, gamma_max, phi_value, away_step_taken, fw_step_taken, tt =
                    afw_step(x, gradient, lmo, active_set, epsilon)
            end
        else
            d, vertex, index, gamma_max, phi_value, away_step_taken, fw_step_taken, tt =
                fw_step(x, gradient, lmo)
        end

        if fw_step_taken || away_step_taken
            gamma, L = line_search_wrapper(
                line_search,
                t,
                f,
                grad!,
                x,
                d,
                gradient,
                dual_gap,
                L,
                gamma0,
                linesearch_tol,
                step_lim,
                gamma_max,
            )
            # cleanup and renormalize every x iterations. Only for the fw steps.
            renorm = mod(t, renorm_interval) == 0
            if away_step_taken
                active_set_update!(active_set, -gamma, vertex, true, index)
            else
                active_set_update!(active_set, gamma, vertex, renorm, index)
            end
        end

        if (
            (mod(t, print_iter) == 0 && verbose) ||
            callback !== nothing ||
            !(line_search isa Agnostic || line_search isa Nonconvex || line_search isa FixedStep)
        )
            primal = f(x)
            dual_gap = phi_value
        end

        if callback !== nothing
            state = (
                t=t,
                primal=primal,
                dual=primal - dual_gap,
                dual_gap=phi_value,
                time=tot_time,
                x=x,
                v=vertex,
                gamma=gamma,
                active_set=active_set,
                gradient=gradient,
            )
            callback(state)
        end

        if verbose && (mod(t, print_iter) == 0 || tt == dualstep)
            if t == 0
                tt = initial
            end
            rep = (
                st[Symbol(tt)],
                string(t),
                Float64(primal),
                Float64(primal - dual_gap),
                Float64(dual_gap),
                tot_time,
                t / tot_time,
                length(active_set),
            )
            print_callback(rep, format_string)
            flush(stdout)
        end
        t += 1
    end

    # recompute everything once more for final verfication / do not record to trajectory though for now!
    # this is important as some variants do not recompute f(x) and the dual_gap regularly but only when reporting
    # hence the final computation.
    # do also cleanup of active_set due to many operations on the same set

    if verbose
        x = compute_active_set_iterate(active_set)
        grad!(gradient, x)
        v = compute_extreme_point(lmo, gradient)
        primal = f(x)
        dual_gap = fast_dot(x, gradient) - fast_dot(v, gradient)
        tt = last
        rep = (
            st[Symbol(tt)],
            string(t - 1),
            Float64(primal),
            Float64(primal - dual_gap),
            Float64(dual_gap),
            (time_ns() - time_start) / 1.0e9,
            t / ((time_ns() - time_start) / 1.0e9),
            length(active_set),
        )
        print_callback(rep, format_string)
        flush(stdout)
    end

    active_set_renormalize!(active_set)
    active_set_cleanup!(active_set)
    x = compute_active_set_iterate(active_set)
    grad!(gradient, x)
    v = compute_extreme_point(lmo, gradient)
    primal = f(x)
    dual_gap = fast_dot(x, gradient) - fast_dot(v, gradient)
    if verbose
        tt = pp
        rep = (
            st[Symbol(tt)],
            string(t - 1),
            Float64(primal),
            Float64(primal - dual_gap),
            Float64(dual_gap),
            (time_ns() - time_start) / 1.0e9,
            t / ((time_ns() - time_start) / 1.0e9),
            length(active_set),
        )
        print_callback(rep, format_string)
        print_callback(nothing, format_string, print_footer=true)
        flush(stdout)
    end

    return x, v, primal, dual_gap, traj_data, active_set
end

function lazy_afw_step(x, gradient, lmo, active_set, phi, epsilon; K=2.0)
    v_lambda, v, v_loc, a_lambda, a, a_loc = active_set_argminmax(active_set, gradient)
    #Do lazy FW step
    grad_dot_lazy_fw_vertex = fast_dot(v, gradient)
    grad_dot_x = fast_dot(x, gradient)
    grad_dot_a = fast_dot(a, gradient)
    if grad_dot_x - grad_dot_lazy_fw_vertex >= grad_dot_a - grad_dot_x &&
            grad_dot_x - grad_dot_lazy_fw_vertex >= phi / K &&
            grad_dot_x - grad_dot_lazy_fw_vertex >= epsilon
        tt = lazy
        gamma_max = one(a_lambda)
        d = x - v
        vertex = v
        away_step_taken = false
        fw_step_taken = true
        index = v_loc
    else
        #Do away step, as it promises enough progress.
        if grad_dot_a - grad_dot_x > grad_dot_x - grad_dot_lazy_fw_vertex &&
           grad_dot_a - grad_dot_x >= phi / K
            tt = away
            gamma_max = a_lambda / (1 - a_lambda)
            d = a - x
            vertex = a
            away_step_taken = true
            fw_step_taken = false
            index = a_loc
            #Resort to calling the LMO
        else
            v = compute_extreme_point(lmo, gradient)
            # Real dual gap promises enough progress.
            grad_dot_fw_vertex = fast_dot(v, gradient)
            dual_gap = grad_dot_x - grad_dot_fw_vertex
            if dual_gap >= phi / K
                tt = regular
                gamma_max = one(a_lambda)
                d = x - v
                vertex = v
                away_step_taken = false
                fw_step_taken = true
                index = -1
                #Lower our expectation for progress.
            else
                tt = dualstep
                phi = min(dual_gap, phi / 2.0)
                gamma_max = zero(a_lambda)
                d = zero(x)
                vertex = v
                away_step_taken = false
                fw_step_taken = false
                index = -1
            end
        end
    end
    return d, vertex, index, gamma_max, phi, away_step_taken, fw_step_taken, tt
end

function afw_step(x, gradient, lmo, active_set, epsilon)
    local_v_lambda, local_v, local_v_loc, a_lambda, a, a_loc =
        active_set_argminmax(active_set, gradient)
    v = compute_extreme_point(lmo, gradient)
    grad_dot_x = fast_dot(x, gradient)
    away_gap = fast_dot(a, gradient) - grad_dot_x
    dual_gap = grad_dot_x - fast_dot(v, gradient)
    if dual_gap >= away_gap && dual_gap >= epsilon
        tt = regular
        gamma_max = one(a_lambda)
        d = x - v
        vertex = v
        away_step_taken = false
        fw_step_taken = true
        index = -1
    elseif away_gap >= epsilon
        tt = away
        gamma_max = a_lambda / (1 - a_lambda)
        d = a - x
        vertex = a
        away_step_taken = true
        fw_step_taken = false
        index = a_loc
    else
        tt = away
        gamma_max = zero(a_lambda)
        d = zero(x)
        vertex = a
        away_step_taken = false
        fw_step_taken = false
        index = a_loc
    end
    return d, vertex, index, gamma_max, dual_gap, away_step_taken, fw_step_taken, tt
end

function fw_step(x, gradient, lmo)
    vertex = compute_extreme_point(lmo, gradient)
    return (
        x - vertex,
        vertex,
        nothing,
        1,
        fast_dot(x, gradient) - fast_dot(vertex, gradient),
        false,
        true,
        regular,
    )
end


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

module StructureAnalysisSymmValue

    using Sparso: Sym, Symexpr, TypedExprNode
    using ..SymbolicAnalysis

    function symbolize(e :: Symexpr, tp :: Type)
        if isa(e, Sym)
            if tp <: Vector
                s = MiddleSymbol(:true)
            elseif tp <: SparseMatrixCSC
                s = MiddleSymbol(:true)
            elseif tp <:Int || tp <: Float64 || tp <: Complex || tp <: Bool
                s = MiddleSymbol(:true)
            else
                dump(tp)
                assert(0)
            end
            return s
        else
            assert(0)
        end
    end

    function preprocess(property_proxies, symbol_info)
        pproxies = property_proxies
        predefined = Dict{Sym, AbstractSymbol}()
        for (s, v) in property_proxies
            if !isa(s, Symbol)
                continue
            end
            if isa(v.symmetric_valued, MiddleSymbol)
                assert(v.symmetric_valued.value == :true)
                predefined[s] = v.symmetric_valued
            end
        end
        predefined
    end

    function postprocess(res, property_proxies, symbol_info)
        for (s, v) in res
            assert(isa(s, Sym))
            if isa(v, MiddleSymbol)
                assert(v.value == :true)
                property_proxies[s].symmetric_valued = v
                property_proxies[s].symmetric_structured = v
            end
        end
    end

    function add_sub_action(e)
        if any(a -> a.svalue == BOTTOM_SYMBOL, e.args)
            e.svalue = BOTTOM_SYMBOL
        elseif all(a -> (a.svalue == MiddleSymbol(:true)), e.args)
            e.svalue = MiddleSymbol(:true)
        end
    end

    # A * B
    function A_mul_B_action(e)
        A = e.args[1].raw_expr
        B = e.args[2].raw_expr
        if isa(A, SymbolNode) && typeof(B) <: Expr && B.head == :call
            m, func_name = resolve_call_names(B.args)
            if func_name == "ctranspose" && B.args[2].name == A.name
                e.svalue = MiddleSymbol(:true)
            end
        end
    end

    function A_mul_Bc_action(e)
        A = e.args[1].raw_expr
        B = e.args[2].raw_expr
        if isa(A, SymbolNode) && isa(B, SymbolNode) && A.name == B.name
            e.svalue = MiddleSymbol(:true)
        end
    end

    function speye_action(e)
        e.svalue = MiddleSymbol(:true)
    end

    function pass_a1_action(e)
        #dump(e.args[1])
        e.svalue = e.args[1].svalue
    end

    function pass_a2_action(e)
        e.svalue = e.args[2].svalue
    end

    function assign_action(e)
        e.args[1].svalue = e.args[2].svalue
    end

    const transfer_rules = (
        ((:(=), Any, SparseMatrixCSC), assign_action),
        ((:(=), Factorization, Any), assign_action),

        ((:call, GlobalRef(Main, :*), SparseMatrixCSC, SparseMatrixCSC), A_mul_B_action),
        ((:call, GlobalRef(Main, :spmatmul_witheps), SparseMatrixCSC, SparseMatrixCSC, Any), A_mul_B_action),

        ((:call, GlobalRef(Main, :scale), SparseMatrixCSC, Any), pass_a1_action),

        ((:call, GlobalRef(Main, :A_mul_Bc), SparseMatrixCSC, SparseMatrixCSC), A_mul_Bc_action),
#        ((:call, GlobalRef(Main, :*), SparseMatrixCSC, SparseMatrixCSC, SparseMatrixCSC), mul3_action),
#        ((:call, GlobalRef(Main, :*), SparseMatrixCSC, Vector), M_mul_v_action),
#        ((:call, GlobalRef(Main, :*), Vector, SparseMatrixCSC), v_mul_M_action),
        ((:call, GlobalRef(Main, :+), SparseMatrixCSC, SparseMatrixCSC), add_sub_action),
        ((:call, GlobalRef(Main, :+), SparseMatrixCSC, SparseMatrixCSC, SparseMatrixCSC), add_sub_action),
        ((:call, GlobalRef(Main, :-), Any, SparseMatrixCSC), add_sub_action),
#        ((:call, GlobalRef(Main, :cholfact_int32), SparseMatrixCSC), pass_a1_action),
        ((:call, GlobalRef(Main, :speye), Int), speye_action),
        ((:call, GlobalRef(Main, :speye_int32), Int), speye_action),
        ((:call, GlobalRef(Main, :generate_symmetric_sparse_matrix), Int), speye_action),

        ((:call, GlobalRef(Main, :*), SparseMatrixCSC, Union{Float64, Int64, Int32}), pass_a1_action),
        ((:call, GlobalRef(Main, :/), SparseMatrixCSC, Union{Float64, Int64, Int32}), pass_a1_action),
        ((:call, GlobalRef(Main, :*), Union{Float64, Int64, Int32}, SparseMatrixCSC), pass_a2_action),

        ((:call, TypedExprNode(Function, :call, TopNode(:apply_type), GlobalRef(Main, :SparseMatrixCSC), GlobalRef(Main, :Float64), GlobalRef(Main, :Int32)), Any), pass_a1_action),
    )

    const pass_info = ("SymmetricValue", transfer_rules, preprocess, postprocess, symbolize)
end

max_stars_repo_name: IntelLabs/Sparso

max_issues_repo_name: IntelLabs/Sparso

max_forks_repo_name: IntelLabs/Sparso


text:
using Stheno: GPC, optimal_q, PseudoObs

# Test Titsias implementation by checking that it (approximately) recovers exact inference
# when M = N and Z = X.
@timedtestset "Titsias" begin
    @timedtestset "optimal_q (single conditioning)" begin
        @timedtestset "σ²" begin
            rng, N, σ², gpc = MersenneTwister(123456), 10, 1e-1, GPC()
            x = collect(range(-3.0, 3.0, length=N))
            f = GP(sin, EQ(), gpc)

            for σ² in [1e-2, 1e-1, 1e0, 1e1]
                @timedtestset "σ² = $σ²" begin
                    y = rand(rng, f(x, σ²))

                    # Compute approximate posterior suff. stats.
                    m_ε, Λ_ε, U = optimal_q(f(x, σ²)←y, f(x))
                    f′ = f | (f(x, σ²) ← y)

                    # Check that exact and approx. posteriors are close in this case.
                    @test m_ε ≈ cholesky(cov(f(x))).U' \ (mean(f′(x)) - mean(f(x)))
                    @test U ≈ cholesky(cov(f(x))).U
                    B = U' \ cov(f(x))
                    @test Λ_ε.U ≈ cholesky(B * B' ./ σ² + I).U
                end
            end
        end
        @timedtestset "Diagonal" begin
            rng, N, gpc = MersenneTwister(123456), 11, GPC()
            x = collect(range(-3.0, 3.0, length=N))
            f = GP(sin, EQ(), gpc)
            Σ = Diagonal(exp.(0.1 * randn(rng, N)) .+ 1)
            y = rand(rng, f(x, Σ))

            # Compute approximate posterior suff. stats.
            m_ε, Λ_ε, U = optimal_q(f(x, Σ)←y, f(x))
            f′ = f | (f(x, Σ) ← y)

            # Check that exact and approx. posteriors are close in this case.
            @test m_ε ≈ cholesky(cov(f(x))).U' \ (mean(f′(x)) - mean(f(x)))
            @test U ≈ cholesky(cov(f(x))).U
            B = U' \ cov(f(x))
            @test Λ_ε.U ≈ cholesky(Symmetric(B * (Σ \ B') + I)).U
        end
        @timedtestset "Dense" begin
            rng, N, gpc = MersenneTwister(123456), 10, GPC()
            x = collect(range(-3.0, 3.0, length=N))
            f = GP(sin, EQ(), gpc)
            A = 0.1 * randn(rng, N, N)
            Σ = Symmetric(A * A' + I)
            y = rand(rng, f(x, Σ))

            # Compute approximate posterior suff. stats.
            m_ε, Λ_ε, U = optimal_q(f(x, Σ)←y, f(x))
            f′ = f | (f(x, Σ) ← y)

            # Check that exact and approx. posteriors are close in this case.
            @test m_ε ≈ cholesky(cov(f(x))).U' \ (mean(f′(x)) - mean(f(x)))
            @test U ≈ cholesky(cov(f(x))).U
            B = U' \ cov(f(x))
            @test Λ_ε.U ≈ cholesky(Symmetric(B * (cholesky(Σ) \ B') + I)).U
        end
    end
    @timedtestset "optimal_q (multiple conditioning)" begin
        rng, N, N′, σ², gpc = MersenneTwister(123456), 5, 7, 1e-1, GPC()
        xx′ = collect(range(-3.0, stop=3.0; length=N + N′))
        idx = randperm(rng, length(xx′))[1:N]
        idx_1, idx_2 = idx, setdiff(1:length(xx′), idx)
        x, x′ = xx′[idx_1], xx′[idx_2]

        f = GP(sin, EQ(), gpc)
        y, y′ = rand(rng, [f(x, σ²), f(x′, σ²)])

        # Compute approximate posterior suff. stats.
        obs = PseudoObs(Obs(f(x, σ²) ← y, f(x′, σ²)←y′), f(xx′))
        m_ε, Λ_ε, U = obs.û.m, obs.û.Λ, obs.U
        f′ = f | (f(x, σ²) ← y, f(x′, σ²)←y′)

        # Check that exact and approx. posteriors are close in this case.
        @test m_ε ≈ cholesky(cov(f(xx′))).U' \ (mean(f′(xx′)) - mean(f(xx′)))
        @test U ≈ cholesky(cov(f(xx′))).U
        @test Λ_ε.U ≈ cholesky((U' \ cov(f(xx′))) * (U' \ cov(f(xx′)))' ./ σ² + I).U

        @timedtestset "multiple cond, multiple pseudo points" begin
            z, z′ = randn(rng, 3), randn(rng, 2)
            zz′ = vcat(z, z′)
            u = PseudoObs(Obs(f(x, σ²)←y, f(x′, σ²)←y′), f(zz′))
            u′ = PseudoObs(Obs(f(x, σ²)←y, f(x′, σ²)←y′), f(z), f(z′))
            @test u.û.m ≈ u′.û.m
            @test u.û.Λ.U ≈ u′.û.Λ.U
            @test u.U ≈ u′.U
        end
        @timedtestset "single cond, multiple pseudo points" begin
            z, z′ = randn(rng, 3), randn(rng, 2)
            zz′ = vcat(z, z′)
            u = PseudoObs(f(x, σ²)←y, f(zz′))
            u′ = PseudoObs(f(x, σ²)←y, f(z), f(z′))
            @test u.û.m ≈ u′.û.m
            @test u.û.Λ.U ≈ u′.û.Λ.U
            @test u.U ≈ u′.U
        end
    end
    @timedtestset "Consistency Tests" begin
        rng, N, M, σ², gpc = MersenneTwister(123456), 5, 3, 1e-1, GPC()
        x = collect(range(-3.0, 3.0, length=N))
        z = randn(rng, M)

        # Generate toy problem.
        f = GP(sin, EQ(), gpc)
        y = rand(f(x, σ²))

        # Generate approximate posterior
        m_ε, Λ_ε, U = optimal_q(f(x, σ²)←y, f(z))
        ỹ = PseudoObs(f, z, m_ε, Λ_ε, U)
        f′_approx = f | ỹ

        P, Q = 7, 4
        x0, x1, x2, x3 = randn(rng, P), randn(rng, Q), randn(rng, Q), randn(rng, P)
        abstractgp_interface_tests(f′_approx, f, x0, x1, x2, x3)

        @test_throws ArgumentError Stheno.mean_vector(ỹ.û, randn(rng, P))
        @test_throws ArgumentError cov(ỹ.û, PseudoObs(f, z, m_ε, Λ_ε, U).û, x0, x1)
    end
    @timedtestset "accuracy tests" begin
        rng, N, N′, Nz, σ², gpc = MersenneTwister(123456), 5, 3, 2, 1e-1, GPC()
        x = collect(range(-3.0, 3.0, length=N))
        x′ = collect(range(-3.0, 3.0, length=N′))
        z = x

        # Generate toy problem.
        f = GP(sin, EQ(), gpc)
        y = rand(f(x, σ²))

        # Exact conditioning.
        f′ = f | (f(x, σ²)←y)
        g′ = f | (f(x, σ²)←y)

        # Approximate conditioning that should yield (almost) exact results.
        m_ε, Λ_ε, U = optimal_q(f(x, σ²)←y, f(z))
        ỹ = PseudoObs(f, z, m_ε, Λ_ε, U)
        f′_approx = f | ỹ
        g′_approx = f | ỹ

        @test mean(f′(x′)) ≈ mean(f′_approx(x′))
        @test cov(f′(x′)) ≈ cov(f′_approx(x′))
        @test cov(f′(x′), f′(x)) ≈ cov(f′_approx(x′), f′_approx(x))
        @test cov(f′(x′), g′(x)) ≈ cov(f′_approx(x′), g′_approx(x))
        @test cov(g′(x′), f′(x)) ≈ cov(g′_approx(x′), f′_approx(x))

        @timedtestset "Standardised Tests" begin
            @timedtestset "Dense Obs. Noise" begin
                # N_obs, M_obs = 11, 13
                # x_obs = collect(range(-5.0, 5.0; length=N_obs))
                # z_obs = collect(range(-5.0, 5.0; length=M_obs))
                # A, B = randn(rng, N_obs, N_obs), randn(rng, M_obs, M_obs)
                # y = rand(rng, (2.3 * GP(cos, EQ(l=0.5), GPC()))(x_obs, _to_psd(A)))
                # standard_1D_tests(
                #     MersenneTwister(123456),
                #     Dict(:l=>0.5, :σ=>2.3, :z=>z_obs, :x=>x_obs, :A=>A, :B=>B, :y=>y),
                #     θ->begin
                #         f = θ[:σ] * GP(cos, EQ(l=θ[:l]), GPC())
                #         u = f(θ[:z], _to_psd(θ[:B]))
                #         f′ = f | PseudoPoints(f(θ[:x], _to_psd(θ[:A]))←θ[:y], u)
                #         return f′, f′
                #     end,
                #     17, 19,
                # )
            end
            @timedtestset "Diagonal Obs. Noise" begin

            end
            @timedtestset "Isotropic Obs. Noise" begin

            end
            @timedtestset "BlockDiagonal Obs. Noise" begin

            end
        end
    end
    # @testset "multiple approximate conditioning" begin
    #     rng, N, N′, Nz, σ², gpc = MersenneTwister(123456), 11, 10, 11, 1e-1, GPC()
    #     x = collect(range(-3.0, 3.0, length=N))
    #     x′ = collect(range(-3.0, 3.0, length=N′))
    #     x̂ = collect(range(-4.0, 4.0, length=N))
    #     z, z′ = copy(x), copy(x′)

    #     xx′, zz′ = vcat(x, x′), vcat(z, z′)

    #     # Construct toy problem.
    #     f = GP(sin, EQ(), gpc)
    #     y, y′ = rand(rng, [f(x, σ²), f(x′, σ²)])
    #     yy′ = vcat(y, y′)

    #     # Perform approximate inference with concatenated inputs.
    #     f′_ss = f | PseudoPoints(f(xx′, σ²)←yy′, f(zz′, 1e-9))

    #     # Perform approximate inference with multiple observations.
    #     f′_ms = f | PseudoPoints((f(x, σ²)←y, f(x′, σ²)←y′), f(zz′, 1e-9))

    #     # Perform approximate inference with multiple sets of pseudo points.
    #     f′_sm = f | PseudoPoints(f(xx′, σ²)←yy′, (f(z, 1e-9), f(z′, 1e-9)))

    #     # Perform approximate inference with multiple sets of pseudo points.
    #     f′_mm = f | PseudoPoints((f(x, σ²)←y, f(x′, σ²)←y′), (f(z, 1e-9), f(z′, 1e-9)))

    #     # Check that the above agree.
    #     @test mean(f′_ss(x̂)) ≈ mean(f′_ms(x̂))
    #     @test cov(f′_ss(x̂)) ≈ cov(f′_ms(x̂))

    #     @test mean(f′_ss(x̂)) ≈ mean(f′_sm(x̂))
    #     @test cov(f′_ss(x̂)) ≈ cov(f′_sm(x̂))

    #     @test mean(f′_ss(x̂)) ≈ mean(f′_mm(x̂))
    #     @test cov(f′_ss(x̂)) ≈ cov(f′_mm(x̂))
    # end
end


text:
using .TurbulenceClosures: ▶z_aaf

using Oceananigans.Operators

const Tx = 16 # CUDA threads per x-block
const Ty = 16 # CUDA threads per y-block

"""
    time_step!(model, Nt, Δt; init_with_euler=true)

Step forward `model` `Nt` time steps with step size `Δt`.

If `init_with_euler` is set to true, then the first step will be taken using a first-order
forward Euler method.
"""
function time_step!(model, Nt, Δt; init_with_euler=true)

    if model.clock.iteration == 0
        [ run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ write_output(model, out)    for out  in values(model.output_writers) ]
    end

    FT = eltype(model.grid)
    RHS = model.poisson_solver.storage
    U, Φ, Gⁿ, G⁻, K, p = datatuples(model.velocities, model.tracers, model.timestepper.Gⁿ,
                                    model.timestepper.G⁻, model.diffusivities, model.pressures)

    for n in 1:Nt
        χ = ifelse(init_with_euler && n==1, FT(-0.5), model.timestepper.χ)

        adams_bashforth_time_step!(model, model.architecture, model.grid, model.buoyancy, model.coriolis,
                                   model.closure, model.forcing, model.boundary_conditions,
                                   U, Φ, p, K, RHS, Gⁿ,  G⁻, Δt, χ)

        [ time_to_run(model.clock, diag) && run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ time_to_run(model.clock, out) && write_output(model, out) for out in values(model.output_writers) ]
    end

    return nothing
end

time_step!(model; Nt, Δt, kwargs...) = time_step!(model, Nt, Δt; kwargs...)

"""
Step forward one time step with a 2nd-order Adams-Bashforth method and pressure-correction
substep.
"""
function adams_bashforth_time_step!(model, arch, grid, buoyancy, coriolis, closure, forcing, bcs,
                                    U, Φ, p, K, RHS, Gⁿ, G⁻, Δt, χ)

    # Arguments for user-defined boundary condition functions:
    boundary_condition_args = (model.clock.time, model.clock.iteration, U, Φ, model.parameters)

    # Pre-computations:
    @launch device(arch) config=launch_config(grid, 3) store_previous_source_terms!(grid, Gⁿ, G⁻)
    fill_halo_regions!(merge(U, Φ), bcs.solution, arch, grid, boundary_condition_args...)

    @launch device(arch) config=launch_config(grid, 3) calc_diffusivities!(K, grid, closure, buoyancy, U, Φ)
    fill_halo_regions!(K, bcs.pressure, arch, grid) # diffusivities share bcs with pressure.
    @launch device(arch) config=launch_config(grid, 2) update_hydrostatic_pressure!(p.pHY′, grid, buoyancy, Φ)
    fill_halo_regions!(p.pHY′, bcs.pressure, arch, grid)

    # Calculate tendency terms (minus non-hydrostatic pressure, which is updated in a pressure correction step):
    calculate_interior_source_terms!(Gⁿ, arch, grid, coriolis, closure, U, Φ, p.pHY′, K, forcing,
                                     model.parameters, model.clock.time)
    calculate_boundary_source_terms!(Gⁿ, arch, grid, bcs.solution, boundary_condition_args...)

    # Complete explicit substep:
    @launch device(arch) config=launch_config(grid, 3) adams_bashforth_update_source_terms!(grid, Gⁿ, G⁻, χ)

    # Start pressure correction substep with a pressure solve:
    fill_halo_regions!(Gⁿ[1:3], bcs.tendency[1:3], arch, grid)
    @launch device(arch) config=launch_config(grid, 3) calculate_poisson_right_hand_side!(arch, grid,
                                                                                model.poisson_solver.bcs,
                                                                                Δt, U, Gⁿ, RHS)
    solve_for_pressure!(arch, model)
    fill_halo_regions!(p.pNHS, bcs.pressure, arch, grid)

    # Complete pressure correction step:
    @launch device(arch) config=launch_config(grid, 3) update_velocities_and_tracers!(grid, U, Φ, p.pNHS, Gⁿ, Δt)

    # Recompute vertical velocity w from continuity equation to ensure incompressibility
    fill_halo_regions!(U, bcs.solution[1:3], arch, grid, boundary_condition_args...)
    @launch device(arch) config=launch_config(grid, 2) compute_w_from_continuity!(grid, U)

    model.clock.time += Δt
    model.clock.iteration += 1

    return nothing
end

function solve_for_pressure!(::CPU, model::Model)
    ϕ = model.poisson_solver.storage

    solve_poisson_3d!(model.poisson_solver, model.grid)
    data(model.pressures.pNHS) .= real.(ϕ)
end

function solve_for_pressure!(::GPU, model::Model)
    ϕ = model.poisson_solver.storage

    solve_poisson_3d!(model.poisson_solver, model.grid)
    @launch device(GPU()) config=launch_config(model.grid, 3) idct_permute!(model.grid, model.poisson_solver.bcs, ϕ,
                                                                            model.pressures.pNHS.data)
end

""" Store previous source terms before updating them. """
function store_previous_source_terms!(grid::AbstractGrid, Gⁿ, G⁻)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds G⁻.Gu[i, j, k] = Gⁿ.Gu[i, j, k]
                @inbounds G⁻.Gv[i, j, k] = Gⁿ.Gv[i, j, k]
                @inbounds G⁻.Gw[i, j, k] = Gⁿ.Gw[i, j, k]
                @inbounds G⁻.GT[i, j, k] = Gⁿ.GT[i, j, k]
                @inbounds G⁻.GS[i, j, k] = Gⁿ.GS[i, j, k]
            end
        end
    end
end

"""
Update the hydrostatic pressure perturbation pHY′. This is done by integrating
the `buoyancy_perturbation` downwards:

    `pHY′ = ∫ buoyancy_perturbation dz` from `z=0` down to `z=-Lz`
"""
function update_hydrostatic_pressure!(pHY′, grid, buoyancy, Φ)
    @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
        @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
            @inbounds pHY′[i, j, 1] = - ▶z_aaf(i, j, 1, grid, buoyancy_perturbation, buoyancy, Φ) * grid.Δz
            @unroll for k in 2:grid.Nz
                @inbounds pHY′[i, j, k] =
                    pHY′[i, j, k-1] - ▶z_aaf(i, j, k, grid, buoyancy_perturbation, buoyancy, Φ) * grid.Δz
            end
        end
    end
end

""" Calculate the right-hand-side of the u-momentum equation. """
function calculate_Gu!(Gu, grid, coriolis, closure, U, Φ, pHY′, K, F, parameters, time)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds Gu[i, j, k] = (-u∇u(grid, U.u, U.v, U.w, i, j, k)
                                            - x_f_cross_U(i, j, k, grid, coriolis, U)
                                            - ∂x_p(i, j, k, grid, pHY′)
                                            + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, U.u, U.v, U.w, K)
                                            + F.u(i, j, k, grid, time, U, Φ, parameters))
            end
        end
    end
end

""" Calculate the right-hand-side of the v-momentum equation. """
function calculate_Gv!(Gv, grid, coriolis, closure, U, Φ, pHY′, K, F, parameters, time)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds Gv[i, j, k] = (-u∇v(grid, U.u, U.v, U.w, i, j, k)
                                            - y_f_cross_U(i, j, k, grid, coriolis, U)
                                            - ∂y_p(i, j, k, grid, pHY′)
                                            + ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure, U.u, U.v, U.w, K)
                                            + F.v(i, j, k, grid, time, U, Φ, parameters))
            end
        end
    end
end

""" Calculate the right-hand-side of the w-momentum equation. """
function calculate_Gw!(Gw, grid, coriolis, closure, U, Φ, pHY′, K, F, parameters, time)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds Gw[i, j, k] = (-u∇w(grid, U.u, U.v, U.w, i, j, k)
                                            - z_f_cross_U(i, j, k, grid, coriolis, U)
                                            + ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure, U.u, U.v, U.w, K)
                                            + F.w(i, j, k, grid, time, U, Φ, parameters))
            end
        end
    end
end

""" Calculate the right-hand-side of the temperature advection-diffusion equation. """
function calculate_GT!(GT, grid, closure, U, Φ, pHY′, K, F, parameters, time)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds GT[i, j, k] = (-div_flux(grid, U.u, U.v, U.w, Φ.T, i, j, k)
                                            + ∇_κ_∇T(i, j, k, grid, Φ.T, closure, K)
                                            + F.T(i, j, k, grid, time, U, Φ, parameters))
            end
        end
    end
end

""" Calculate the right-hand-side of the salinity advection-diffusion equation. """
function calculate_GS!(GS, grid, closure, U, Φ, pHY′, K, F, parameters, time)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds GS[i, j, k] = (-div_flux(grid, U.u, U.v, U.w, Φ.S, i, j, k)
                                            + ∇_κ_∇S(i, j, k, grid, Φ.S, closure, K)
                                            + F.S(i, j, k, grid, time, U, Φ, parameters))
            end
        end
    end
end

""" Store previous value of the source term and calculate current source term. """
function calculate_interior_source_terms!(G, arch, grid, coriolis, closure, U, Φ, pHY′, K, F, parameters, time)

    Bx, By, Bz = floor(Int, grid.Nx/Tx), floor(Int, grid.Ny/Ty), grid.Nz  # Blocks in grid
    @launch device(arch) threads=(Tx, Ty) blocks=(Bx, By, Bz) calculate_Gu!(G.Gu, grid, coriolis, closure, U, Φ, pHY′,
                                                                            K, F, parameters, time)

    @launch device(arch) threads=(Tx, Ty) blocks=(Bx, By, Bz) calculate_Gv!(G.Gv, grid, coriolis, closure, U, Φ, pHY′,
                                                                            K, F, parameters, time)

    @launch device(arch) threads=(Tx, Ty) blocks=(Bx, By, Bz) calculate_Gw!(G.Gw, grid, coriolis, closure, U, Φ, pHY′,
                                                                            K, F, parameters, time)

    @launch device(arch) threads=(Tx, Ty) blocks=(Bx, By, Bz) calculate_GT!(G.GT, grid, closure, U, Φ, pHY′,
                                                                            K, F, parameters, time)

    @launch device(arch) threads=(Tx, Ty) blocks=(Bx, By, Bz) calculate_GS!(G.GS, grid, closure, U, Φ, pHY′,
                                                                            K, F, parameters, time)
end

"""
Evaluate the right-hand-side terms at time step n+½ using a weighted 2nd-order
Adams-Bashforth method

    `G^{n+½} = (3/2 + χ)G^{n} - (1/2 + χ)G^{n-1}`
"""
function adams_bashforth_update_source_terms!(grid::AbstractGrid{FT}, Gⁿ, G⁻, χ) where FT
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds Gⁿ.Gu[i, j, k] = (FT(1.5) + χ) * Gⁿ.Gu[i, j, k] - (FT(0.5) + χ) * G⁻.Gu[i, j, k]
                @inbounds Gⁿ.Gv[i, j, k] = (FT(1.5) + χ) * Gⁿ.Gv[i, j, k] - (FT(0.5) + χ) * G⁻.Gv[i, j, k]
                @inbounds Gⁿ.Gw[i, j, k] = (FT(1.5) + χ) * Gⁿ.Gw[i, j, k] - (FT(0.5) + χ) * G⁻.Gw[i, j, k]
                @inbounds Gⁿ.GT[i, j, k] = (FT(1.5) + χ) * Gⁿ.GT[i, j, k] - (FT(0.5) + χ) * G⁻.GT[i, j, k]
                @inbounds Gⁿ.GS[i, j, k] = (FT(1.5) + χ) * Gⁿ.GS[i, j, k] - (FT(0.5) + χ) * G⁻.GS[i, j, k]
            end
        end
    end
end

"""
Calculate the right-hand-side of the elliptic Poisson equation for the non-hydrostatic
pressure

    `∇²ϕ_{NH}^{n+1} = (∇·u^n)/Δt + ∇·(Gu, Gv, Gw)`
"""
function calculate_poisson_right_hand_side!(::CPU, grid::AbstractGrid, ::PoissonBCs, Δt, U, G, RHS)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                # Calculate divergence of the RHS source terms (Gu, Gv, Gw).
                @inbounds RHS[i, j, k] = div_f2c(grid, U.u, U.v, U.w, i, j, k) / Δt +
                                         div_f2c(grid, G.Gu, G.Gv, G.Gw, i, j, k)
            end
        end
    end
end

"""
Calculate the right-hand-side of the elliptic Poisson equation for the non-hydrostatic
pressure and in the process apply the permutation

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

in the z-direction which is required by the GPU fast cosine transform algorithm for
horizontally periodic model configurations.
"""
function calculate_poisson_right_hand_side!(::GPU, grid::AbstractGrid, ::PPN, Δt, U, G, RHS)
    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
    @loop for k in (1:Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                if (k & 1) == 1  # isodd(k)
                    k′ = convert(UInt32, CUDAnative.floor(k/2) + 1)
                else
                    k′ = convert(UInt32, Nz - CUDAnative.floor((k-1)/2))
                end
                @inbounds RHS[i, j, k′] = div_f2c(grid, U.u, U.v, U.w, i, j, k) / Δt +
                                          div_f2c(grid, G.Gu, G.Gv, G.Gw, i, j, k)
            end
        end
    end
end

"""
Calculate the right-hand-side of the elliptic Poisson equation for the non-hydrostatic
pressure and in the process apply the permutation

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

in the y- and z-directions which is required by the GPU fast cosine transform algorithm for
reentrant channel model configurations.
"""
function calculate_poisson_right_hand_side!(::GPU, grid::AbstractGrid, ::PNN, Δt, U, G, RHS)
    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                if (k & 1) == 1  # isodd(k)
                    k′ = convert(UInt32, CUDAnative.floor(k/2) + 1)
                else
                    k′ = convert(UInt32, Nz - CUDAnative.floor((k-1)/2))
                end

                if (j & 1) == 1  # isodd(j)
                    j′ = convert(UInt32, CUDAnative.floor(j/2) + 1)
                else
                    j′ = convert(UInt32, Ny - CUDAnative.floor((j-1)/2))
                end

                @inbounds RHS[i, j′, k′] = div_f2c(grid, U.u, U.v, U.w, i, j, k) / Δt +
                                           div_f2c(grid, G.Gu, G.Gv, G.Gw, i, j, k)
            end
        end
    end
end

"""
Copy the non-hydrostatic pressure into `pNHS` and undo the permutation

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

along the z-direction.
"""
function idct_permute!(grid::AbstractGrid, ::PPN, ϕ, pNHS)
    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
    @loop for k in (1:Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                if k <= Nz/2
                    @inbounds pNHS[i, j, 2k-1] = real(ϕ[i, j, k])
                else
                    @inbounds pNHS[i, j, 2(Nz-k+1)] = real(ϕ[i, j, k])
                end
            end
        end
    end
end

"""
Copy the non-hydrostatic pressure into `pNHS` and undo the permutation

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

along the y- and z-direction.
"""
function idct_permute!(grid::AbstractGrid, ::PNN, ϕ, pNHS)
    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                if k <= Nz/2
                    k′ = 2k-1
                else
                    k′ = 2(Nz-k+1)
                end

                if j <= Ny/2
                    j′ = 2j-1
                else
                    j′ = 2(Ny-j+1)
                end

                @inbounds pNHS[i, j′, k′] = real(ϕ[i, j, k])
            end
        end
    end
end

"""
Update the horizontal velocities u and v via

    `u^{n+1} = u^n + (Gu^{n+½} - δₓp_{NH} / Δx) Δt`

and the tracers via

    `c^{n+1} = c^n + Gc^{n+½} Δt`

Note that the vertical velocity is not explicitly time stepped.
"""
function update_velocities_and_tracers!(grid::AbstractGrid, U, Φ, pNHS, Gⁿ, Δt)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)

                @inbounds U.u[i, j, k] = U.u[i, j, k] + (Gⁿ.Gu[i, j, k] - ∂x_p(i, j, k, grid, pNHS)) * Δt
                @inbounds U.v[i, j, k] = U.v[i, j, k] + (Gⁿ.Gv[i, j, k] - ∂y_p(i, j, k, grid, pNHS)) * Δt
                @inbounds Φ.T[i, j, k] = Φ.T[i, j, k] + (Gⁿ.GT[i, j, k] * Δt)
                @inbounds Φ.S[i, j, k] = Φ.S[i, j, k] + (Gⁿ.GS[i, j, k] * Δt)
            end
        end
    end
end

"""
Compute the vertical velocity w by integrating the continuity equation downwards

    `w^{n+1} = -∫ [∂/∂x (u^{n+1}) + ∂/∂y (v^{n+1})] dz`
"""
function compute_w_from_continuity!(grid::AbstractGrid, U)
    @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
        @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
            @inbounds U.w[i, j, 1] = 0
            @unroll for k in 2:grid.Nz
                @inbounds U.w[i, j, k] = U.w[i, j, k-1] + grid.Δz * ∇h_u(i, j, k-1, grid, U.u, U.v)
            end
        end
    end
end

""" Apply boundary conditions by adding flux divergences to the right-hand-side. """
function calculate_boundary_source_terms!(Gⁿ, arch, grid, bcs, args...)

    # Velocity fields
    for (i, ubcs) in enumerate(bcs[1:3])
        apply_z_bcs!(Gⁿ[i], arch, grid, ubcs.z.left, ubcs.z.right, args...)
    end

    # Tracer fields
    for (i, cbcs) in enumerate(bcs[4:end])
        apply_z_bcs!(Gⁿ[i+3], arch, grid, cbcs.z.left, cbcs.z.right, args...)
    end

    return nothing
end


text:
"""
    MyArray{T,N}

    My super awesome array wrapper!

    # Fields
    # - `data::AbstractArray{T,N}`: stores the array being wrapped
    # - `metadata::Dict`: stores metadata about the array
"""
struct MyArray{T,N} <: AbstractArray{T,N}
    data::AbstractArray{T,N}
    metadata::Dict
end

"""
    mysearch(array::MyArray{T}, val::T; verbose=true) where {T} -> Int

    Searches the `array` for the `val`. For some reason we don't want to use Julia's
    builtin search :)

    # Arguments
    # - `array::MyArray{T}`: the array to search
    # - `val::T`: the value to search for
    #
    # # Keywords
    # - `verbose::Bool=true`: print out progress details
    #
    # # Returns
    # - `Int`: the index where `val` is located in the `array`
    #
    # # Throws
    # - `NotFoundError`: I guess we could throw an error if `val` isn't found.
"""
function mysearch(array::AbstractArray{T}, val::T) where T
    nothing
end

"""
    Manager(args...; kwargs...) -> Manager

    A cluster manager which spawns workers.

    # Arguments
    #
    # - `min_workers::Integer`: The minimum number of workers to spawn or an exception is thrown
    # - `max_workers::Integer`: The requested number of workers to spawn
    #
    # # Keywords
    #
    # - `definition::AbstractString`: Name of the job definition to use. Defaults to the
    #     definition used within the current instance.
    #     - `name::AbstractString`: ...
    #     - `queue::AbstractString`: ...
"""
function Manager(...)
    nothing
end


text:
using SymbolicUtils, Test
using SymbolicUtils: Term, symtype
using BenchmarkTools

function goldstein_price()
    @syms x1 x2
    f1 = x1+x2+1
    f2 = 19-14*x1+3*x1^2-14*x2+6*x1*x2+3*x2^2
    f3 = 2*x1-3*x2
    f4 = 18-32*x1+12*x1^2+48*x2-36*x1*x2+27*x2^2

    # f(x) is the Goldstein-Price function
    f = (1+f1^2*f2)*(30+f3^2*f4)
    @btime expand($f)
    f2 = f^2
    @btime expand($f2)
    f3 = f^3
    @btime expand($f3)
end


text:
module Observations

using DocStringExtensions
using LinearAlgebra
using Statistics

export Observation

"""
    Observation{FT <: AbstractFloat}

Structure that contains the observations

# Fields

$(TYPEDFIELDS)
"""
struct Observation{FT <: AbstractFloat}
    "vector of observational samples, each of length sample_dim"
    samples::Vector{Vector{FT}}
    "covariance of the observational noise (assumed to be normally
    distributed); `sample_dim x sample_dim` (where `sample_dim` is the number of
    elements in each sample), or a scalar if the sample dim is 1. If not
    supplied, `obs_noise_cov` is set to a diagonal matrix whose non-zero elements
    are the variances of the samples, or to a scalar variance in the case of
    1d samples. `obs_noise_cov` is set to nothing if only a single sample is
    provided."
    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, FT, Nothing}
    "sample mean"
    mean::Union{AbstractVector{FT}, FT}
    "names of the data"
    data_names::Union{AbstractVector{String}, String}
end

# Constructors
function Observation(samples::Vector{Vector{FT}}, data_names::Union{Vector{String}, String}) where {FT <: AbstractFloat}

    N_samples = length(samples)
    # convert to sample_dim x N_samples to determine sample covariance

    if N_samples == 1
        # only one sample - this requires a bit more data massaging
        temp1 = convert(Array, reshape(hcat(samples...)', N_samples, :))
        temp = dropdims(temp1, dims = 1)
        samplemean = vec(mean(temp, dims = 2))
        obs_noise_cov = nothing
    else
        temp = convert(Array, reshape(hcat(samples...)', N_samples, :))
        if size(temp, 2) == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the covariance) are scalars
            samplemean = mean(temp)
            obs_noise_cov = var(temp)
        else
            samplemean = vec(mean(temp, dims = 1))
            obs_noise_cov = cov(temp)
        end
    end
    Observation(samples, obs_noise_cov, samplemean, data_names)
end

function Observation(
    samples::AbstractMatrix{FT},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    # samples is of size sample_dim x N_samples
    sample_dim, N_samples = size(samples)

    if N_samples == 1
        # Only one sample, so there is no covariance to be computed and the
        # sample mean equals the sample itself
        obs_noise_cov = nothing
        samplemean = vec(samples)
        samples_vec = vec([vec(samples)])

    else
        # convert matrix of samples to a vector of vectors
        samples_vec = [samples[:, i] for i in 1:N_samples]
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the variance) are scalars
            samplemean = mean(samples)
            obs_noise_cov = var(samples)
        else
            samplemean = vec(mean(samples, dims = 2))
            obs_noise_cov = cov(samples, dims = 2)
        end
    end
    Observation(samples_vec, obs_noise_cov, samplemean, data_names)
end

function Observation(
    samples::Vector{Vector{FT}},
    obs_noise_cov::AbstractMatrix{FT},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    N_samples = length(samples)
    sample_dim = length(samples[1])
    # convert to sample_dim x N_samples to determine sample covariance

    if N_samples == 1
        # only one sample - this requires a bit more data massaging
        temp1 = convert(Array, reshape(hcat(samples...)', N_samples, :))
        temp = dropdims(temp1, dims = 1)
        samplemean = vec(mean(temp, dims = 2))
        obsnoisecov = obs_noise_cov
    else
        temp = convert(Array, reshape(hcat(samples...)', N_samples, :))
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the covariance) are scalars
            samplemean = mean(temp)
            err = ("When sample_dim is 1, obs_cov_noise must be a 1x1 matrix")
            @assert(size(obs_noise_cov) == (1, 1), err)
            obsnoisecov = obs_noise_cov[1]
        else
            samplemean = vec(mean(temp, dims = 1))
            err = ("obs_cov_noise must be of size sample_dim x sample_dim.
                   \tsample_dim: number of elements per observation sample")
            @assert(size(obs_noise_cov) == (sample_dim, sample_dim), err)
            obsnoisecov = obs_noise_cov
        end
    end

    Observation(samples, obsnoisecov, samplemean, data_names)
end

function Observation(
    samples::AbstractMatrix{FT},
    obs_noise_cov::Union{AbstractMatrix{FT}, Nothing},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    # samples is of size sample_dim x N_samples
    sample_dim, N_samples = size(samples)
    if N_samples == 1
        # sample mean equals the sample itself
        obsnoisecov = obs_noise_cov
        samplemean = vec(samples)
        samples_vec = vec([vec(samples)])
    else
        # convert matrix of samples to a vector of vectors
        samples_vec = [samples[:, i] for i in 1:N_samples]
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the variance) are scalars
            samplemean = mean(samples)
            err = ("When sample_dim is 1, obs_cov_noise must be a 1x1 matrix")
            @assert(size(obs_noise_cov) == (1, 1), err)
            obsnoisecov = obs_noise_cov[1]
        else
            samplemean = vec(mean(samples, dims = 2))
            err = ("obs_cov_noise must be of size sample_dim x sample_dim.
                   \tsample_dim: number of elements per observation sample")
            @assert(size(obs_noise_cov) == (sample_dim, sample_dim), err)
            obsnoisecov = obs_noise_cov
        end
    end

    Observation(samples_vec, obsnoisecov, samplemean, data_names)
end

end # module Observations


text:
export ket,
       bra,
       ketbra,
       partialtrace,
       projector,
       trnormalize,
       trnormalize!,
       purify,
       fidelity,
       concurrence,
       avgfidelity,
       ispossemidef,
       pauli_decomp

trnorm(M::AbstractMatrix) = sum(svdvals(M))

"""
`basis_vector(T,i,d)`

Returns the `(i+1)`th basis element in a `d` dimensional Euclidean space
as a vector of type `t`.
"""
function basis_vector( t::Type{T}, i::Int, d::Int ) where T
  result = zeros(T, d)
  result[i+1] = 1
  result
end

"""
`ket`

Returns a vector corresponding to a particular basis unit vector in some
Hilbert space.

The vector can be specified in a number of different ways:
   + by a string and a base
   + by a integer and a dimension

If the vector is specified by a string and a base, the dimension is
inferred from the length of the string. The base is assumed to be 2
unless explicitly specified.

If the basis element is specified by an integer, the dimension is
assumed to be 2 unless explicitly specified. The integer values
allowed are 0 through the dimension minus 1.

Optionally, the type of the elements of the vector can be specified
in the first argument.
"""
function ket(t::Type, label::AbstractString, base::Int)
  basis_vector( t, parseint(label, base), base^length(label) )
end
ket( label::AbstractString, base=2 ) = ket(Float64, label, base)
ket( t::Type, label::AbstractString ) = ket(t, label, 2)

ket( i::Int, d=2 ) = basis_vector( Float64, i, d )

# the duals of the above
"""
`bra`

Returns a dual vector (row vector) corresponding to a particular basis
unit vector in some Hilbert space.

See `ket` for details on how the corresponding vector may be specified.
"""
function bra( t::Type, label::AbstractString, base::Int )
  return ket( t, label, base )'
end

bra( label::AbstractString, base::Int ) = bra( Float64, label, base )
bra( label::AbstractString ) = bra( Float64, label, 2 )
bra( t::Type, label::AbstractString ) = bra( t, label, 2 )
bra( i::Int, d::Int ) = ket( i, d )'

"""
`ketbra`

Returns the outer product between a dual vector (row vector) and a vector,
each corresponding to a particular basis unit vector in some Hilbert space.
This can also be interpreted as a matrix unite element specified by two
unit vectors for a row and a column.

See `ket` for details on how the corresponding vectors may be specified.
"""
ketbra( t::Type, label_left::AbstractString, label_right::AbstractString, base::Int ) = ket( t, label_left, base ) * bra( t, label_right, base )
ketbra( label_left::AbstractString, label_right::AbstractString, base::Int ) = ket( label_left, base ) * bra( label_right, base )
ketbra( label_left::AbstractString, label_right::AbstractString ) = ket( label_left ) * bra( label_right )
ketbra( t::Type, label_left::AbstractString, label_right::AbstractString ) = ket( t, label_left ) * bra( t, label_right )
ketbra( i::Int, j::Int, d::Int ) = ket( i, d ) * bra( j, d )

"""
`projector(v)`

Computes the rank-1 projector operator corresponding to a given vector. If more than one
vector is given, the projector into the space spanned by the vectors is computed.
"""
function projector( v::AbstractVector )
    v*v'/norm(v,2)^2
end

"""
`trnormalize(m)`

Normalizes a matrix with respect to the trace norm (Schatten 1 norm).
"""
trnormalize( m::AbstractMatrix ) = m/trnorm(m)

"""
`trnormalize!(v)`

Normalizes a matrix with respect to the trace norm (Schatten 1
norm) in place.
"""
function trnormalize!( m::AbstractMatrix{T} ) where T <: ComplexF64
    n = convert(ComplexF64, trnorm(m))
    lmul!(1/n,m)
end

function trnormalize!( m::AbstractMatrix{T} ) where T <: Float64
    n = trnorm(m)
    lmul!(1/n,m)
end

function trnormalize!( m::AbstractMatrix{T} ) where T <: Int64
    n = convert(Float64, trnorm(m))
    m = convert(Array{Float64}, m)
    lmul!(1/n,m)
end

"""
Computes the partial trace of a matrix `m`.
"""
function partialtrace( m::AbstractMatrix{T}, ds::AbstractVector, dt::Int ) where T
  s = size(m)
  l = length(ds);
  if s[1] != s[2]
    error("Partial trace only defined for square matrices.")
  elseif prod(ds) != s[1] || prod(ds) != s[2]
    error("Subsystem decomposition does not match matrix dimensions.")
  end
  # the way we order subsystems is the opposite of the way we order
  # significant digits, so we must changes the index of the subsystem
  ds = reverse(ds)
  dt = l - dt + 1
  m = reshape( m, tuple(ds..., ds...))
  # get a list of the indices not traced over
  indices = [1:l;]
  filter!(x -> x != dt, indices)
  # reshape the matrix so that subsystem to be traced over is
  # the "last" subsystem
  m = permutedims(m, tuple(indices..., map(x->x+l,indices)..., dt, dt+l))
  m = reshape( m, (prod(ds[indices])^2, ds[dt], ds[dt]))
  # now trace over the last subsystem
  new_m = zeros( T, prod(ds[indices])^2 )
  for ii in 1:prod(ds[indices])^2
    new_m[ii] = tr(reshape(m[ii,:,:],(ds[dt],ds[dt])))
  end
  reshape(new_m, (prod(ds[indices]),prod(ds[indices])) )
end

# TODO: specialized, fast versions of partial trace

function _max_entangled_state( d::Int )
  me = zeros(Float64,d^2)
  for ii=0:d-1
    me += kron(ket(ii,d),ket(ii,d))/sqrt(d)
  end
  return me
end

"""
Computes a purification of `rho`
"""
function purify( rho::AbstractMatrix )
  d = size(rho,1)
  (vals,vecs) = eig(rho)
  psi = zeros(eltypeof(rho), d^2)
  for ii=1:d
    psi += sqrt(vals[ii]) * kron(vecs[:,ii],ket(ii-1,d))
  end
  psi
end

"""
Computes the concurrence of a bipartite qubit state.
"""
function concurrence(m)
    ms = conj(m)
    σy = [0 -1im; 1im 0];
    mt = kron(σy,σy)*ms*kron(σy,σy)
    R = sqrt(sqrt(m)*mt*sqrt(m))
    res = dot(sort(real(eigvals(R)),rev=true),[1,-1,-1,-1])

    return res
end

"""
Computes the fidelity between two quantum states. By default, the Josza convention
is used. The Uhlmann fidelity may be calculated by using the `kind` keyword argument
with the value `:uhlmann`.
"""
function fidelity(a::AbstractMatrix,b::AbstractMatrix;kind=:josza)
    fu = trnorm(sqrt(a)*sqrt(b))
    if kind==:josza
        return fu^2
    elseif kind==:uhlmann
        return fu
    end
end

function fidelity(a::AbstractVector,b::AbstractVector;kind=:josza)
    return fidelity(projector(a),projector(b),kind=kind)
end

function fidelity(a::AbstractVector,b::AbstractMatrix;kind=:josza)
    return fidelity(projector(a),b,kind=kind)
end

function fidelity(a::AbstractMatrix,b::AbstractVector;kind=:josza)
    return fidelity(b,a,kind=kind)
end

"""
Computes the average fidelity of the outputs of two quantum operations in the Liouville
representation. One of the operations is assumed to be a unitary quantum operation. The
fidelity measure used is the one defined by Josza.
"""
function avgfidelity(liou,liou_uni;kind=:josza)
    d = sqrt(size(liou,2))
    f = (real(tr(liou*liou_uni'))+d)/(d^2+d)
    return f
end

# TODO: entfidelity(liou,liou_uni)
#       avg2entfidelity()
#       ent2avgfidelity()

# fidelity AbstractVector AbstractVector
# fidelity AbstractVector AbstractMatrix
# fidelity AbstractMatrix AbstractMatrix
# superfidelity AbstractMatrix AbstractMatrix
# subfidelity AbstractMatrix AbstractMatrix

"""
Tests if a matrix is positive semidefinite within a given tolerance.
"""
function ispossemidef(m;tol=0.0)
    evs = eigvals(m)
    tol = tol==0.0 ? 1e2*eps(abs.(one(eltype(m)))) : tol
    all(real(evs) .> -tol) && all(abs.(imag(evs)) .< tol)
end

"""
Decompose a density matrix into Pauli operators with a given cutoff.
"""
function pauli_decomp(ρ::AbstractMatrix, cutoff=1e-3)
    r = Dict{Pauli,Float64}()
    d2 = size(ρ, 1)
    d = round(Int, sqrt(d2))
    for p in allpaulis(d)
        val = real(tr(ρ * p)) / d2
        if val > cutoff
            r[p] = val
        end
    end
    return r
end


text:
# functions related to studentized range distribution

import .RFunctions:
    srdistcdf,
    srdistccdf,
    srdistlogcdf,
    srdistlogccdf,
    srdistinvcdf,
    srdistinvccdf,
    srdistinvlogcdf,
    srdistinvlogccdf


text:
using BenchmarkTools, NumericalIntegration
using Perfusion

t = collect(1:600) ./ 60
cp = aif_georgiou(t .- 1)

test_params = (kt = 0.25, kep = 0.5)
ct = model_tofts(t = t, cp = cp, parameters = test_params)
ref_params = (kt = 0.14, kep = 1.0)
crr = model_tofts(t = t, cp = cp, parameters = ref_params)

rel_kt = test_params.kt / ref_params.kt
rel_ve = (test_params.kt / test_params.kep) / (ref_params.kt / ref_params.kep)
rel_kt = round(rel_kt, digits = 3)
rel_ve = round(rel_ve, digits = 3)

num_noisy_replications = 100
σ = 0.1
noisy_ct = zeros(num_noisy_replications, length(t))
for idx = 1:num_noisy_replications
    noisy_ct[idx, :] .= ct .+ σ .* randn(length(ct))
end

estimates = fit_model(:referenceregion, :lls, t = t, crr = crr, ct = noisy_ct).estimates
std(estimates.rel_kt)
std(estimates.rel_ve)
std(estimates.kep)

estimates =
    fit_model(:constrained_referenceregion, :lls, t = t, crr = crr, ct = noisy_ct).estimates
mean(estimates.rel_kt)
mean(estimates.rel_ve)
mean(estimates.kep)


text:
__precompile__(false)

module Model

using Flux
using Images
using Plots
using Juno
using Distributions
using ScikitLearn
using PyCall
using StatsBase
using Interpolations
using SegmentConnection
using DSB2018.Errors
using HypothesisTests
using Distributions
using ImageTransformations
using CoordinateTransformations

import DSB2018.Errors: metric

@pyimport skimage.filters as skimage_filters
@pyimport skimage.restoration as skimage_restoration
@pyimport skimage.morphology as skimage_morphology
@pyimport scipy.ndimage as ndi

export split_connected_regions, metric

using DSB2018.Data

function test_metric(id::String)
    img = Data.loadTrainImage(id)
    this = convert.(Float64, convert.(Gray{Normed{UInt8,8}}, img.image))

    if mean(this) > 0.5
        this = 1.0 - this
    end

    tfm  = recenter(RotMatrix(pi/600), ImageTransformations.center(this));
    this = warp(this, tfm)
    this = collect(this)
    this[isnan.(this)] = 0.0

    if size(this, 1) > size(img.mask, 1)
        this = this[1:size(img.mask, 1), :]
    end
    while size(this, 1) < size(img.mask, 1)
        this = [this; zeros(1, size(this, 2))]
    end

    if size(this, 2) > size(img.mask, 2)
        this = this[:, 1:size(img.mask, 2)]
    end
    while size(this, 2) < size(img.mask, 2)
        this = [this zeros(size(this, 1), 1)]
    end

    @assert all(size(this) .== size(img.mask))

    th = skimage_filters.threshold_otsu(this)

    labeled_prediction = Images.label_components(this .> th)
    labeled_prediction[labeled_prediction .== 1] = 0

    prediction_vector  = split_labeled_matrix(labeled_prediction)
    truth_images = getTrainImageMasks(img)

    mask = img.mask
    Images.save("/home/korelin/mask.png", (mask .> 0)*1.0)
    Images.save("/home/korelin/image.png", (this .> th)*1.0)

    pt = Images.label_components((this .> th)*1.0)
    mt = Images.label_components(mask .> 0)

    truth_mt  = split_labeled_matrix(mt)
    pred_pt   = split_labeled_matrix(pt)

    return metric(pred_pt, truth_mt)
end

function get_predictions_by_density(images::Vector)::Vector{Prediction}
    predictions = Vector{Prediction}(length(images))
    t = Vector{Float64}(0)
    num_labels = Vector{Int64}(0)

    Juno.@progress "Scoring" for k in eachindex(images)
        tic()
        X = typeof(images[k])
        raw_image = X == String ? loadTrainImage(images[k]) : images[k]

        res = to_density(raw_image.image, 4)

        if any(isinf.(res) .| isnan.(res))
            throw("At image $k: non-finite or nan res")
        end

        adaptive_threshold = skimage_filters.threshold_local(res, 41, offset=0.0)
        global_threshold   = skimage_filters.threshold_otsu(res)

        w_adaptve = 0.5
        threshold = adaptive_threshold * w_adaptve .+ global_threshold * (1.0 - w_adaptve)

        img_t = res .> threshold
        img_t = skimage_morphology.remove_small_objects(img_t, min_size=4)
        img_t = ndi.binary_fill_holes(img_t)

        labeled_prediction = Images.label_components(img_t)
        prediction_vector  = split_labeled_matrix(labeled_prediction)

        truth_images = getTrainImageMasks(raw_image)
        if !isempty(truth_images)
            score = metric(prediction_vector, truth_images)
        else
            score = NaN
        end

        predictions[k] = Prediction(raw_image.id, score, labeled_prediction)

        t0 = toq()

        push!(t, t0)
        push!(num_labels, length(prediction_vector))
        n = sum(t)

        scores = [x.score for x in predictions[1:k]]
        m = mean(scores)

        p1 = scatter(scores, title="Mean = $m", legend = false)
        p2 = plot(t, fill = (0, .9, :grey), title="Seconds. Total: $n", legend = false)

        if length(num_labels) < 2
            #some sort of a bug: bar([1]) errors
            display(plot(p1, p2, layout = grid(2, 1)))
            continue
        end

        p3 = bar(num_labels,  fillcolor=:grey, title="Num labels", legend = false)
        display(plot(p1, p2, p3, layout = grid(3, 1)))
    end

    return predictions
end

function split_labeled_matrix(this::Matrix{Int64})::Vector{Matrix{Bool}}
    labels = unique(this)
    labels = labels[labels .> 0]

    out = Vector{Matrix{Bool}}(length(labels))
    for k in eachindex(labels)
        new_labels = falses(size(this))
        new_labels[this .== labels[k]] = true
        out[k] = new_labels
    end

    return out
end

function get_density(this::Vector{Float64})::Vector{Float64}
    st = sum(this)
    if st .< 1e-3
        return zeros(size(this))
    end

    return this / sum(this) / length(this) * 100
end

function distance(x1::CartesianIndex, x2::CartesianIndex)::Float64
    return sqrt((x2[1] - x1[1])^2 + (x2[2] - x1[2])^2)
end

function box_mean(this::Vector{Float64}, n::Int64)::Vector{Float64}
    out = similar(this)
    for k in eachindex(this)
        j1 = max(1, k - n)
        j2 = min(length(this), k + n)
        v = 0.0
        nn = 0.0
        for j in j1 : j2
            v += this[j]
            nn += 1
        end
        out[k] = v / nn
    end
    return out
end

function to_density(colored_input::Matrix{X}, nlag::Int64) where {X <: Any}

    if X <: ColorTypes.AbstractGray
        colored_input = convert.(Gray{Normed{UInt8,8}}, colored_input)
    end

    this = convert.(Float64, colored_input)
    this = skimage_restoration.denoise_bilateral(this, multichannel=false)

    out = zeros(size(this))
    num_visits = zeros(size(this))

    nr = size(this, 1)
    nc = size(this, 2)

    qs = mean(this[:])
    if qs > 0.5
        this = 1.0 - this
    end

    for k = 1 : size(this, 1)
        mvalues = this[k, :]
        densities = get_density(mvalues)
        if (nlag != 1) && (sum(densities) > 0.0)
            densities *= StatsBase.autocor(densities, [nlag])[1]
        end

        out[k, :] .+= densities
        num_visits[k, :] .+= 1
    end

    for k = 1 : size(this, 2)
        mvalues = this[:, k]
        densities = get_density(mvalues)
        if (nlag != 1) && (sum(densities) > 0.0)
            densities *= StatsBase.autocor(densities, [nlag])[1]
        end
        out[:, k] .+= densities
        num_visits[:, k] .+= 1
    end

    res = out ./ num_visits# .* uniform_weights / 2.
    return res
end

function make_batches(img_size::Tuple{Int64, Int64}, sz::Tuple{Int64, Int64}, δ::Int64)
    idx = 1 : sz[1], 1 : sz[2]

    ind = Vector{Tuple{UnitRange{Int64}, UnitRange{Int64}}}(0)

    kx_last = img_size[1] - sz[1]
    ky_last = img_size[2] - sz[2]

    kx_offset = collect(0 : δ : kx_last)
    ky_offset = collect(0 : δ : ky_last)

    if last(kx_offset) != kx_last
        push!(kx_offset, kx_last)
    end

    if last(ky_offset) != ky_last
        push!(ky_offset, ky_last)
    end

    for kx = kx_offset, ky in ky_offset
        ix = idx[1] + kx
        iy = idx[2] + ky
        push!(ind, (ix, iy))
    end

    return ind
end

function make_batches(img::Matrix{T}, idx::Vector{Tuple{UnitRange{Int64}, UnitRange{Int64}}}) where {T <: Any}

    out = Vector{Matrix{T}}(0)
    for j in idx
        push!(out, img[j[1], j[2]])
    end

    return out
end

function apply_batched_model(model::Flux.Chain,
                             sz::Tuple{Int64, Int64},
                             img::T) where {T <: AbstractArray}
    num_visits = zeros(size(img))
    idx = 1 : sz[1], 1 : sz[2]

    @assert size(img, 4) == 1

    out = Matrix{Float64}(size(img)[1], size(img)[2])
    Juno.@progress "Applying model" for kx = 0 :  (size(img, 1) - sz[1])
        for ky = 0 : (size(img, 2) - sz[2])
            res = model(img[idx[1] + kx, idx[2] + ky, :, :])
            num_visits[idx[1] + kx, idx[2] + ky] .+= 1.
            out[idx[1] + kx, idx[2] + ky] .+= res.data
        end
    end

    out_normed = out ./ num_visits
    out_normed[num_visits .== 0.0] = 0.0
    return out_normed
end

function max_distance_to_neighbor(orig::Matrix{Float64}, mask::Matrix{Float64}, sz::Int64)
    x = Vector{Float64}(0)
    y = Vector{Float64}(0)
    z = zeros(size(orig))

    idx = 1 : sz, 1 : sz
    c0 = Int64(round(sz/2)), Int64(round(sz/2))
    Juno.@progress for kx = 0 :  (size(orig, 1) - sz)
        for ky = 0 : (size(orig, 2) - sz)
            v = minimum(orig[idx[1] + kx, idx[2] + ky])
            v0 = orig[c0[1] + kx, c0[2] + ky]

            z[c0[1] + kx, c0[1] + ky] = v/v0
            push!(x, v0 / v * v0)
            push!(y, mask[c0[1] + kx, c0[2] + ky])
        end
    end

    (x, y, z)
end

function metric(prediction::BitArray{2}, labeled_truth::Matrix{Int64})
    labeled_prediction = Images.label_components(prediction)
    prediction_vector  = split_labeled_matrix(labeled_prediction)
    truth_vector = split_labeled_matrix(labeled_truth)
    metric(prediction_vector, truth_vector)
end

end


text:
# helper function used in some definitons, copied from ColorTypes.jl
#Base.@pure basetype(T::Type) = Base.typename(T).wrapper
#Base.@pure basetype(A) = basetype(typeof(A))

macro concrete(expr)
    @assert expr.head == :struct
    S = expr.args[2]
    return quote
        $(esc(expr))

        for n in fieldnames($S)
            if !isconcretetype(fieldtype($S, n))
                error("field $n is not concrete")
            end
        end
    end
end

simplify_number(x::Number) = x
simplify_number(x::AbstractFloat) = isinteger(x) ? Int(x) : x
simplify_number(x::Rational) = isone(denominator(x)) ? numerator(x) : x
simplify_number(x::Complex) = begin
    r, i = reim(x)
    iszero(i) ? simplify_number(r) : complex(simplify_number(r),simplify_number(i))
end


text:
abstract type AbstractOptimizationFunction end
abstract type AbstractADType end

struct AutoForwardDiff <: AbstractADType end
struct AutoReverseDiff <: AbstractADType end
struct AutoTracker <: AbstractADType end
struct AutoZygote <: AbstractADType end
struct AutoFiniteDiff <: AbstractADType end
struct AutoModelingToolkit <: AbstractADType end

struct OptimizationFunction{F,G,H,HV,K} <: AbstractOptimizationFunction
    f::F
    grad::G
    hess::H
    hv::HV
    adtype::AbstractADType
    kwargs::K
end

function OptimizationFunction(f, x, ::AutoForwardDiff; grad=nothing,hess=nothing, p=DiffEqBase.NullParameters(), chunksize = 1, hv = nothing, kwargs...)
    _f = θ -> f(θ,p)[1]
    if grad === nothing
        gradcfg = ForwardDiff.GradientConfig(_f, x, ForwardDiff.Chunk{chunksize}())
        grad = (res,θ) -> ForwardDiff.gradient!(res, _f, θ, gradcfg)
    end

    if hess === nothing
        hesscfg = ForwardDiff.HessianConfig(_f, x, ForwardDiff.Chunk{chunksize}())
        hess = (res,θ) -> ForwardDiff.hessian!(res, _f, θ, hesscfg)
    end

    if hv === nothing
        hv = function (H,θ,v)
            res = Array{typeof(x[1])}(undef, length(θ), length(θ)) #DiffResults.HessianResult(θ)
            hess(res, θ)
            H .= res*v
        end
    end

    return OptimizationFunction{typeof(f),typeof(grad),typeof(hess),typeof(hv),typeof(kwargs)}(f,grad,hess,hv,AutoForwardDiff(),kwargs)
end

function OptimizationFunction(f, x, ::AutoZygote; grad=nothing, hess=nothing, p=DiffEqBase.NullParameters(), hv = nothing, kwargs...)
    _f = θ -> f(θ,p)[1]
    if grad === nothing
        grad = (res,θ) -> res isa DiffResults.DiffResult ? DiffResults.gradient!(res, Zygote.gradient(_f, θ)[1]) : res .= Zygote.gradient(_f, θ)[1]
    end

    if hess === nothing
        hess = function (res,θ)
            if res isa DiffResults.DiffResult
                DiffResults.hessian!(res, ForwardDiff.jacobian(θ) do θ
                                                Zygote.gradient(_f,θ)[1]
                                            end)
            else
                res .=  ForwardDiff.jacobian(θ) do θ
                    Zygote.gradient(_f,θ)[1]
                  end
            end
        end
    end

    if hv === nothing
        hv = function (H,θ,v)
            _θ = ForwardDiff.Dual.(θ,v)
            res = DiffResults.GradientResult(_θ)
            grad(res,_θ)
            H .= getindex.(ForwardDiff.partials.(DiffResults.gradient(res)),1)
        end
    end
    return OptimizationFunction{typeof(f),typeof(grad),typeof(hess),typeof(hv),typeof(kwargs)}(f,grad,hess,hv,AutoZygote(),kwargs)
end