text:
export
    tensor_assignment,
    tensor_species_map!,
    reduced_site_tensor,
    tensor_size,
    tensor,
    right_env,
    left_env,
    dressed_mps,
    mpo, mps

function assign_tensors!(network::PEPSNetwork{false})
    _types = (:site, :central_h, :central_v, :gauge_h)
    for type ∈ _types
        push!(network.tensor_spiecies, tensor_assignment(network, type)...)
    end
end

function assign_tensors!(network::PEPSNetwork{true})
    _types = (:site, :central_h, :central_v, :virtual, :central_d, :gauge_h)
    for type ∈ _types
        push!(network.tensor_spiecies, tensor_assignment(network, type)...)
    end
end

tensor_assignment(
    network::AbstractGibbsNetwork{S, T},
    s::Symbol
) where {S, T} = tensor_assignment(network, Val(s))

tensor_assignment(
    network::AbstractGibbsNetwork{S, T},
    ::Val{:site}
) where {S, T} = Dict(
    (i, j) => :site for i ∈ 1:network.nrows, j ∈ 1:network.ncols
)

tensor_assignment(
    network::AbstractGibbsNetwork{S, T},
    ::Val{:central_h}
) where {S, T} = Dict(
    (i, j + 1//2) => :central_h for i ∈ 1:network.nrows, j ∈ 1:network.ncols
)

tensor_assignment(
    network::AbstractGibbsNetwork{S, T},
    ::Val{:central_v}
) where {S, T} = Dict(
    (i + 1//2, j) => :central_v for i ∈ 1:network.nrows-1, j ∈ 1:network.ncols
)

tensor_assignment(
    network::PEPSNetwork{false},
    ::Val{:gauge_h}
) = Dict((i + δ, j) => :gauge_h
    for i ∈ 1:network.nrows-1, j ∈ 1:network.ncols, δ ∈ (1//6, 2//6, 4//6, 5//6)
)

tensor_assignment(
    network::PEPSNetwork{true},
    ::Val{:gauge_h}
) = Dict((i + δ, r) => :gauge_h
    for i ∈ 1:network.nrows-1, r ∈ 1:1//2:network.ncols, δ ∈ (1//6, 2//6, 4//6, 5//6)
)

tensor_assignment(
    network::PEPSNetwork{true},
    ::Val{:virtual}
) = Dict(
    (i, j + 1//2) => :virtual for i ∈ 1:network.nrows, j ∈ 1:network.ncols-1
)

tensor_assignment(
    network::PEPSNetwork{true},
    ::Val{:central_d}
) = Dict(
    (i + 1//2, j + 1//2) => :central_d for i ∈ 1:network.nrows-1, j ∈ 1:network.ncols-1
)

function tensor(
    network::AbstractGibbsNetwork,
    v
)
    if v ∈ keys(network.tensor_spiecies)
        tensor(network, v, Val(network.tensor_spiecies[v]))
    else
        ones(1, 1, 1, 1)
    end
end

function tensor_size(
    network::AbstractGibbsNetwork, v
)
    if v ∈ keys(network.tensor_spiecies)
        tensor_size(network, v, Val(network.tensor_spiecies[v]))
    else
        (1, 1, 1, 1)
    end
end

function tensor(
    network::AbstractGibbsNetwork{S}, v::S, ::Val{:site}
) where {S}
    loc_exp = exp.(-network.β .* local_energy(network, v))
    projs = projectors(network, v)
    @cast A[σ, _] := loc_exp[σ]
    for pv ∈ projs @cast A[σ, (c, γ)] |= A[σ, c] * pv[σ, γ] end
    B = dropdims(sum(A, dims=1), dims=1)
    reshape(B, size.(projs, 2))
end

function tensor_size(
    network::AbstractGibbsNetwork, v, ::Val{:site}
)
    dims = size.(projectors(network, v))
    pdims = first.(dims)
    @assert all(σ -> σ == first(pdims), pdims)
    last.(dims) # TODO: CZEMU?
end

function tensor(
    network::AbstractGibbsNetwork,
    v::Tuple{Rational{Int}, Int},
    ::Val{:central_v}
)
    r, j = v
    i = floor(Int, r)
    h = connecting_tensor(network, (i, j), (i+1, j))
    @cast A[_, u, _, d] := h[u, d]
    A
end

function tensor_size(
    network::AbstractGibbsNetwork,
    v::Tuple{Rational{Int}, Int},
    ::Val{:central_v}
)
    r, j = v
    i = floor(Int, r)
    u, d = size(interaction_energy(network, (i, j), (i+1, j)))
    (1, u, 1, d)
end

function tensor(
    network::AbstractGibbsNetwork,
    w::Tuple{Int, Rational{Int}},
    ::Val{:central_h}
)
    i, r = w
    j = floor(Int, r)
    v = connecting_tensor(network, (i, j), (i, j+1))
    @cast A[l, _, r, _] := v[l, r]
    A
end

function tensor_size(
    network::AbstractGibbsNetwork,
    w::Tuple{Int, Rational{Int}},
    ::Val{:central_h}
)
    i, r = w
    j = floor(Int, r)
    l, r = size(interaction_energy(network, (i, j), (i, j+1)))
    (l, 1, r, 1)
end

function tensor(
    network::AbstractGibbsNetwork, v::Tuple{Rational{Int}, Rational{Int}}, ::Val{:central_d}
)
    r, s = v
    i = floor(Int, r)
    j = floor(Int, s)
    NW = connecting_tensor(network, (i, j), (i + 1, j + 1))
    NE = connecting_tensor(network, (i, j + 1), (i + 1, j))
    @cast A[_, (u, ũ), _, (d, d̃)] := NW[u, d] * NE[ũ, d̃]
    A
end

function tensor_size(
    network::AbstractGibbsNetwork, v::Tuple{Rational{Int}, Rational{Int}}, ::Val{:central_d}
)
    r, s = v
    i = floor(Int, r)
    j = floor(Int, s)
    u, d = size(interaction_energy(network, (i, j), (i + 1, j + 1)))
    ũ, d̃ = size(interaction_energy(network, (i, j + 1), (i + 1, j)))
    (1, u*ũ, 1, d*d̃)
end

function _all_fused_projectors(
    network::AbstractGibbsNetwork, v::Tuple{Int, Rational{Int}},
)
    i, s = v
    j = floor(Int, s)

    left_nbrs = ((i+1, j+1), (i, j+1), (i-1, j+1))
    prl = projector.(Ref(network), Ref((i, j)), left_nbrs)
    p_lb, p_l, p_lt = last(fuse_projectors(prl))

    right_nbrs = ((i+1, j), (i, j), (i-1, j))
    prr = projector.(Ref(network), Ref((i, j+1)), right_nbrs)
    p_rb, p_r, p_rt = last(fuse_projectors(prr))

    p_lb, p_l, p_lt, p_rb, p_r, p_rt
end

function tensor(
    network::AbstractGibbsNetwork, v::Tuple{Int, Rational{Int}}, ::Val{:virtual}
)
    p_lb, p_l, p_lt,
    p_rb, p_r, p_rt = _all_fused_projectors(network, v)

    h = connecting_tensor(network, floor.(Int, v), ceil.(Int, v))

    @tensor B[l, r] := p_l[l, x] * h[x, y] * p_r[r, y]
    @cast A[l, (ũ, u), r, (d̃, d)] |= B[l, r] * p_lt[l, u] * p_rb[r, d] *
                                     p_rt[r, ũ] * p_lb[l, d̃]
    A
end

function tensor_size(
    network::AbstractGibbsNetwork, v::Tuple{Int, Rational{Int}}, ::Val{:virtual}
)
    p_lb, p_l, p_lt,
    p_rb, p_r, p_rt = _all_fused_projectors(network, v)
    (size(p_l, 1), size(p_lt, 2) * size(p_rt, 2),
     size(p_r, 1), size(p_rb, 2) * size(p_lb, 2))
end

function tensor(
    network::AbstractGibbsNetwork, v, ::Val{:gauge_h}
)
    X = network.gauges[v]
    @cast A[_, u, _, d] := Diagonal(X)[u, d]
    A
end

function tensor_size(
    network::AbstractGibbsNetwork, v, ::Val{:gauge_h}
)
    X = network.gauges[v]
    u = size(X, 1)
    (1, u, 1, u)
end

function connecting_tensor(
    network::AbstractGibbsNetwork{S}, v::S, w::S
) where {S}
    en = interaction_energy(network, v, w)
    exp.(-network.β .* (en .- minimum(en)))
end

function reduced_site_tensor(
    network::PEPSNetwork{false}, v::Tuple{Int, Int}, l::Int, u::Int
)
    i, j = v
    eng_local = local_energy(network, v)
    pl = projector(network, v, (i, j-1))
    eng_pl = interaction_energy(network, v, (i, j-1))
    @matmul eng_left[x] := sum(y) pl[x, y] * eng_pl[y, $l]

    pu = projector(network, v, (i-1, j))
    eng_pu = interaction_energy(network, v, (i-1, j))
    @matmul eng_up[x] := sum(y) pu[x, y] * eng_pu[y, $u]

    en = eng_local .+ eng_left .+ eng_up
    loc_exp = exp.(-network.β .* (en .- minimum(en)))

    pr = projector(network, v, (i, j+1))
    pd = projector(network, v, (i+1, j))
    @cast A[r, d, σ] := pr[σ, r] * pd[σ, d] * loc_exp[σ]
    A
end

function reduced_site_tensor(
    network::PEPSNetwork{true}, v::Tuple{Int, Int}, ld::Int, l::Int, d::Int, u::Int
)
    i, j = v
    eng_local = local_energy(network, v)

    pl = projector(network, v, (i, j-1))
    eng_pl = interaction_energy(network, v, (i, j-1))
    @matmul eng_left[x] := sum(y) pl[x, y] * eng_pl[y, $l]

    pd = projector(network, v, (i-1, j-1))
    eng_pd = interaction_energy(network, v, (i-1, j-1))
    @matmul eng_diag[x] := sum(y) pd[x, y] * eng_pd[y, $d]

    pu = projector(network, v, (i-1, j))
    eng_pu = interaction_energy(network, v, (i-1, j))
    @matmul eng_up[x] := sum(y) pu[x, y] * eng_pu[y, $u]

    en = eng_local .+ eng_left .+ eng_diag .+ eng_up
    loc_exp = exp.(-network.β .* (en .- minimum(en)))

    p_lb = projector(network, (i, j-1), (i+1, j))
    p_rb = projector(network, (i, j), (i+1, j-1))
    pr = projector(network, v, ((i+1, j+1), (i, j+1), (i-1, j+1)))
    pd = projector(network, v, (i+1, j))

    @cast A[r, d, (k̃, k), σ] := p_rb[σ, k] * p_lb[$ld, k̃] * pr[σ, r] *
                                pd[σ, d] * loc_exp[σ]
    A
end

function tensor_size(
    network::PEPSNetwork{false}, v::Tuple{Int, Int}, ::Val{:reduced}
)
    i, j = v
    pr = projector(network, v, (i, j+1))
    pd = projector(network, v, (i+1, j))
    @assert size(pr, 1) == size(pr, 1)
    (size(pr, 2), size(pd, 2), size(pd, 1))
end

function mpo(::Type{T},
    peps::AbstractGibbsNetwork, r::Union{Rational{Int}, Int}
) where {T <: Number}
    W = MPO(T, length(peps.columns_MPO) * peps.ncols)
    layers = Iterators.product(peps.columns_MPO, 1:peps.ncols)
    Threads.@threads for (k, (d, j)) ∈ collect(enumerate(layers))
        W[k] = tensor(peps, (r, j + d))
    end
    W
end

@memoize Dict mpo(
    peps::AbstractGibbsNetwork, r::Union{Rational{Int}, Int}
) = mpo(Float64, peps, r)

function mps(::Type{T}, peps::AbstractGibbsNetwork, i::Int) where {T <: Number}
    if i > peps.nrows return IdentityMPS() end
    ψ = mps(peps, i+1)
    for r ∈ peps.layers_MPS ψ = mpo(peps, i+r) * ψ end
    compress(ψ, peps)
end

@memoize Dict mps(peps::AbstractGibbsNetwork, i::Int) = mps(Float64, peps, i)

@memoize Dict function dressed_mps(peps::AbstractGibbsNetwork, i::Int)
    ψ = mps(peps, i+1)
    for r ∈ peps.layers_left_env ψ = mpo(peps, i+r) * ψ end
    ψ
end

function compress(ψ::AbstractMPS, peps::AbstractGibbsNetwork)
    if bond_dimension(ψ) < peps.bond_dim return ψ end
    SpinGlassTensors.compress(ψ, peps.bond_dim, peps.var_tol, peps.sweeps)
end

@memoize Dict function _mpo(peps::AbstractGibbsNetwork, i::Int)
    prod(mpo.(Ref(peps), i .+ reverse(peps.layers_right_env)))
end

@memoize Dict function right_env(peps::AbstractGibbsNetwork, i::Int, ∂v::Vector{Int})
    l = length(∂v)
    if l == 0 return ones(1, 1) end
    R̃ = right_env(peps, i, ∂v[2:l])
    ϕ = dressed_mps(peps, i)
    W = _mpo(peps, i)
    k = length(W)
    M = ϕ[k-l+1]
    M̃ = W[k-l+1]
    K = @view M̃[:, ∂v[1], :, :]
    @tensor R[x, y] := K[y, β, γ] * M[x, γ, α] * R̃[α, β] order = (β, γ, α)
    R
end

@memoize Dict function left_env(peps::AbstractGibbsNetwork, i::Int, ∂v::Vector{Int})
    l = length(∂v)
    if l == 0 return ones(1) end
    L̃ = left_env(peps, i, ∂v[1:l-1])
    ϕ = dressed_mps(peps, i)
    m = ∂v[l]
    M = ϕ[l]
    @matmul L[x] := sum(α) L̃[α] * M[α, $m, x]
    L
end


text:
using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
using DynamicalSystems, PyPlot, Statistics
include(srcdir("style.jl"))

using OrdinaryDiffEq, DiffEqBase

# This function over-writes the currently unexported version of ChaosTools, which
# has nested `εs` but doesn't work because of diffeq error, see here:

function ChaosTools.transit_time_statistics(ds::ContinuousDynamicalSystem, u0, εs, T;
        alg = Tsit5(), diffeq...
    )

    eT = eltype(ds.t0)
    ChaosTools.check_εs_sorting(εs, length(u0))
    exits = [eT[] for _ in 1:length(εs)]
    entries = [eT[] for _ in 1:length(εs)]

    # Make the magic callback:
    crossing(u, t, integ) = ChaosTools.εdistance(u, u0, εs[1])
    negative_affect!(integ) = push!(entries[1], integ.t)
    positive_affect!(integ) = push!(exits[1], integ.t)
    cb = ContinuousCallback(crossing, positive_affect!, negative_affect!;
        save_positions = (false, false)
    )

    prob = ODEProblem(ds, (eT(0), eT(T)); u0 = u0)
    sol = solve(prob, alg;
        callback=cb, save_everystep = false, dense = false,
        save_start=false, save_end = false, diffeq...
    )
    return exits, entries
end

ro = Systems.roessler()
u0 = SVector(
    4.705494942754781,
    -10.221120945130545,
    0.06186563933318555,
)

exponents = (-4:0.5:-1)
n = Float64[]
εs = Float64[]

for e in exponents
    x = [ℯ^e]
    exits, entries = transit_time_statistics(ro, u0, x, 100.0^(-e))
    transits, returns = transit_return(exits, entries)
    m = mean.(returns)
    push!(εs, x[1])
    push!(n, m[1])
end

# plot
fig, ax = subplots()

o, d = linreg(log.(εs), log.(n))
ax.plot(log.(εs), log.(n), marker = "o", ls="-", lw = 1.0, ms = 5, color = "C0", label = rdspl(abs(d)),)

ax.legend()
ax.set_ylabel("log(⟨t⟩)")
ax.set_xlabel("log(ε)")
ax.set_title("Towel map")
fig.tight_layout()

ext: jl

max_stars_repo_name: Datseris/FractalDimension

max_issues_repo_name: Datseris/FractalDimension

max_forks_repo_name: Datseris/FractalDimension


text:
__precompile__()
"""
ProcGA

Process scheduling program with GA written in Julia.
"""
module ProcGA

export DVector

# type definition for easy reading
DVector{T} = Vector{Vector{T}} where T

include("jobassign.jl") # assignment job table
include("jobsequencial.jl") # sequencial job table
include("procga_base.jl") # basics for population

end

max_stars_repo_name: Yoshinobu-Ishizaki/ProcGA

max_issues_repo_name: Yoshinobu-Ishizaki/ProcGA

max_forks_repo_name: Yoshinobu-Ishizaki/ProcGA


text:
using PySyntax
using Base.Test

# write your own tests here
@test @py(:(df)) == :(df)
@test @py(:(df.loc)) == :(df[:loc])
@test @py(:(df.loc[:])) == :(get(df[:loc], (pybuiltin("slice"))(nothing, nothing, nothing)))
@test @py(:(df.loc[0:5, "salary1":"salary2"].plot.bar(linestyle="--"))) == :((((get(df[:loc], (pybuiltin("slice"))(0, 5, nothing), (pybuiltin("slice"))("salary1", "salary2", nothing)))[:plot])[:bar])(linestyle="--"))


text:
function largest_independence_number(graph::AbstractGraph)
    return largest_clique_number(complement(graph))
end

function largest_independents_set(graph::AbstractGraph)
    return largest_cliques_set(complement(graph))
end

function largest_clique_number(graph::AbstractGraph)
    return length(largest_cliques_set(graph)[1])
end

function largest_cliques_set(graph::AbstractGraph)
    max_cliques = maximal_cliques(graph)
    largest_cliques = Vector{Vector{Integer}}()
    number = 0
    for clique = max_cliques
        lc=length(clique)
        number = lc>number ? lc : number
    end
    for clique = max_cliques
        if length(clique)==number; push!(largest_cliques, clique) end
    end

    return largest_cliques
end


text:
module MaxLFSR

export LFSR, FastLFSR

# Coefficients are taken from https://users.ece.cmu.edu/~koopman/lfsr/index.html
#
# Index by the number of bits needed. Using fewer than 4 bits is apparently not supported,
# which is fine because usually one would want longer sequences anyways.
#
# Shorter sequences are simply generated using a longer LFSR and iterating multiple times.
include("terms.jl")

function nbits(x::Integer)
    if x <= zero(x) || x > typemax(Int64)
        throw(ArgumentError("Expected `length` to be between 1 and 2^64 - 1: Got: $x"))
    end

    # Add 1 if `x` is a power of 2 because a LFSR with `n` bits generates `2^n - 1` unique
    # numbers (it cannot generate zero)
    bits = ceil(Int, log2(x)) + convert(Int, ispow2(x))
    return max(bits, 4)
end

abstract type AbstractLFSR end

#####
##### Generic Length LFSR
#####
struct LFSR <: AbstractLFSR
    len::Int
    mask::Int
    seed::Int
end

function Base.show(io::IO, A::T) where {T <: AbstractLFSR}
    print(io, "$(string(T)) for 1:$(length(A)) starting at $(seed(A))")
end

"""
    LFSR(length; seed = 1)

Construct a maximum length shift register that will randomly cycle once through the numbers
`1` to `length`.

Keyword argument `seed` defines the starting point for the LFSR.

Example
-------
```julia
# Standard LFSR
julia> A = LFSR(10)
LFSR for 1:10 starting at 1

julia> for i in A
           println(i)
       end
1
9
7
10
5
6
3
8
4
2

# Change the starting location
julia> A = LFSR(10; seed = 7)
LFSR for 1:10 starting at 7

julia> for i in A
           println(i)
       end
7
10
5
6
3
8
4
2
1
9
```
"""
function LFSR(length::Integer; seed = 1)
    # Bounds checking on `seed`
    if seed < 1 || seed > length
        throw(ArgumentError("Expected `seed` to be between 1 and `length`"))
    end

    # Find the number of bits needed to represent length
    bits = nbits(length)
    mask = FEEDBACK[bits]
    return LFSR(length, mask, seed)
end

Base.eltype(L::AbstractLFSR) = Int
Base.length(L::AbstractLFSR) = L.len
mask(L::AbstractLFSR) = L.mask
seed(L::AbstractLFSR) = L.seed

# Iterator Interface
@inline Base.iterate(A::AbstractLFSR) = (seed(A), seed(A))
@inline function Base.iterate(A::AbstractLFSR, x)
    # Iterate until we reach a result that is within the correct range.
    while true
        x = step(A, x)

        # If we arrive back at out seed, we're done
        (x == seed(A)) && return nothing

        # Otherwise, perform a length check and exit.
        (x <= length(A)) && return (x, x)
    end
end

@inline function step(A::AbstractLFSR, x)
    m = isodd(x) ? mask(A) : 0
    return xor(x >> 1, m)
end

#####
##### Fast LFSR
#####
struct FastLFSR <: AbstractLFSR
    len::Int
    mask::Int
    seed::Int
end

"""
    FastLFSR(length; seed = 1)

Construct a **fast** maximum length shift register that will randomly cycle once through the numbers
`1` to `length`.
**NOTE:** This type requires `ispow2(length + 1) == true`.

Keyword argument `seed` defines the starting point for the LFSR.

Example
-------
```julia
# Standard FastLFSR
julia> A = FastLFSR(15)
LFSR for 1:10 starting at 1

julia> println.(A);
1
9
13
15
14
7
10
5
11
12
6
3
8
4
2

# Change the starting location
julia> A = LFSR(15; seed = 7)
FastLFSR for 1:15 starting at 7

julia> println.(A);
7
10
5
11
12
6
3
8
4
2
1
9
13
15
14

# Errors if length is incorrect.
julia> FastLFSR(10)
ERROR: ArgumentError: FastLFSR expects `ispow2(length + 1) == true`
```
"""
function FastLFSR(length::Integer; seed = 1)
    # Check that `length` is a power of 2
    if !ispow2(length + 1)
        throw(ArgumentError("FastLFSR expects `ispow2(length + 1) == true`"))
    end

    # Because we're restricted to a minimum of 4 bits, the minimum length we can provide
    # is 15
    if length < 15
        throw(ArgumentError("Length must be 15 or greater!"))
    end

    # Bounds checking on `seed`
    if seed < 1 || seed > length
        throw(ArgumentError("Expected `seed` to be between 1 and `length`"))
    end

    # Find the number of bits needed to represent length
    bits = nbits(length)
    mask = FEEDBACK[bits]
    return FastLFSR(length, mask, seed)
end

@inline function Base.iterate(A::FastLFSR, x)
    x = step(A, x)
    return (x == seed(A)) ? nothing : (x, x)
end

end # module


text:
### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 538f05be-5053-11eb-19a0-959e34e1c2a1
using Pkg, DrWatson

# ╔═╡ 5d84f90c-5053-11eb-076b-5f30fc9685e3
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanQuap
	using StatisticalRethinking
end

# ╔═╡ 05d45c8f-0500-4142-8dae-b6208127fe6f
#include(joinpath(sr_path(), "require", "stan", "optimize.jl"))

# ╔═╡ 235c5298-5053-11eb-0608-435d1aa4716c
md" ## Clip-07-11s.jl"

# ╔═╡ 861accd4-5053-11eb-0432-81db212f4f38
begin
	sppnames = [:afarensis, :africanus, :hapilis, :boisei,
		:rudolfensis, :ergaster, :sapiens]
	brainvol = [438, 452, 612, 521, 752, 871, 1350]
	masskg = [37, 35.5, 34.5, 41.5, 55.5, 61, 53.5]
	df = DataFrame(species = sppnames, brain = brainvol, mass = masskg)
end;

# ╔═╡ ee275e7a-5067-11eb-325b-7760b758e85e
stan7_2 = "
data {
 int < lower = 1 > N; 			// Sample size
 int < lower = 1 > K;			// Degree of polynomial
 vector[N] brain; 				// Outcome
 matrix[N, K] mass; 			// Predictor
}

parameters {
 real a;                        // Intercept
 vector[K] bA;                  // K slope(s)
 real < lower = 0 > sigma;    	// Error SD
}

transformed parameters {
    vector[N] mu;
    mu = a + mass * bA;
}

model {
  a ~ normal(0.5, 1);         	//Priors
  bA ~ normal(0, 10);
  sigma ~ lognormal(0, 1);
  brain ~ normal(mu , sigma);   // Likelihood
}
";

# ╔═╡ bc44cdc8-5422-11eb-0034-93d8626c23f7
fig7_4 = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2);

# ╔═╡ 919ade48-5a83-11eb-2fd5-75d756d4f4d6
md"
!!! tip
	In this case `fnc` creates a polynomial mass matrix.

	This could also be a more complicated function.
"

# ╔═╡ b3fc5ef8-2309-465d-873f-05690bcad3cf
methods(quap)

# ╔═╡ a9980819-c0eb-49dc-9b07-aef79409b87c
methods(quap)

# ╔═╡ 95cbb456-520e-11eb-2e0a-1d173ddc9dec
for (findx, K) in enumerate([1, 6])
	fig7_4[findx] = plot(;ylims=(200,1500), leg=false)
	for i in 1:6
		df1 = sample(df, 6, replace=false)
		scale!(df1, [:mass, :brain])
		df1.brain_std = df1.brain/maximum(df1.brain)
		mass = create_observation_matrix(df1.mass_s, K)
		N = size(df1, 1)
		data = (N = N, K = K, brain = df1.brain_s, mass = mass)
		init = (a = 0.0, bA = ones(K), sigma = 2)
		q7_2s, m7_2s, o7_2s = stan_quap("m7.2s", stan7_2; data, init)
		linkvars = [:a, :bA, :sigma]
		fnc=create_observation_matrix

		if !isnothing(q7_2s)
			nt7_2s = read_samples(m7_2s, :namedtuple)
			lab="$(size(nt7_2s[Symbol(linkvars[2])], 1))-th degree polynomial"
			title="R^2 = $(r2_is_bad(nt7_2s, df1))"
			fig7_4[findx] = plot!(;lab, title)
			fig7_4[findx] = plotlines(df1, :mass, :brain, nt7_2s,
				linkvars, fig7_4[findx];
			stepsize=0.05, fnc, lab, title, leg=false)
		end
	end
	scatter!(fig7_4[findx], df.mass, df.brain;
		xlab="body mass [kg]", ylab="brain vol [cc]", lab="Observations")
	for (ind, species) in pairs(df.species)
		annotate!(fig7_4[findx], [(df[ind, :mass] + 1, df[ind, :brain] + 90,
			Plots.text(df[ind, :species], 6, :red, :right))])
	end
end;

# ╔═╡ 2389915c-51c7-11eb-15c6-4fe1fa67fbfe
md" ### Snippet 7.11"

# ╔═╡ eae3bf44-51c6-11eb-19b7-a79473f2bf36
plot(fig7_4[1], fig7_4[2])

# ╔═╡ 91e9af0a-5065-11eb-212c-f751fd114263
md" ## End of clip-07-11s.jl"

# ╔═╡ Cell order:
# ╟─235c5298-5053-11eb-0608-435d1aa4716c
# ╠═538f05be-5053-11eb-19a0-959e34e1c2a1
# ╠═5d84f90c-5053-11eb-076b-5f30fc9685e3
# ╠═861accd4-5053-11eb-0432-81db212f4f38
# ╠═ee275e7a-5067-11eb-325b-7760b758e85e
# ╠═bc44cdc8-5422-11eb-0034-93d8626c23f7
# ╟─919ade48-5a83-11eb-2fd5-75d756d4f4d6
# ╠═b3fc5ef8-2309-465d-873f-05690bcad3cf
# ╠═05d45c8f-0500-4142-8dae-b6208127fe6f
# ╠═a9980819-c0eb-49dc-9b07-aef79409b87c
# ╠═95cbb456-520e-11eb-2e0a-1d173ddc9dec
# ╟─2389915c-51c7-11eb-15c6-4fe1fa67fbfe
# ╠═eae3bf44-51c6-11eb-19b7-a79473f2bf36
# ╟─91e9af0a-5065-11eb-212c-f751fd114263


text:
# status: WIP
# still does not allow to come back from bankruptcy punishment
# this means that final period bankrtuptcy punishment is very important
function bk!(m::BModel,p::Param)

    cmat = fill(-Inf,p.nD,p.na)
    cmat2 = fill(-Inf,p.nD,p.na)
    vmat = fill(-Inf,p.nD,p.na)
    ctmp = fill(-Inf,p.nD,p.ny,p.na)
    # vtmp = fill(-Inf,p.nD,p.ny,p.na)

    for it in p.nT:-1:1
        # println(it)
        for iy in 1:p.ny  # current state
            for id in 1:p.nD  # current dchoice
                # println("id = $id")
                filer = id==2  # filing is id=2

                if it==p.nT
                    # final period: consume everyting.
                    if id == 1
                        m.c[id,iy,it] = Envelope(MLine(vcat(p.a_lowT,p.a_high),vcat( 0.0 , p.a_high)) )
                        m.v[id,iy,it] = Envelope(MLine(vcat(p.a_lowT,p.a_high),vcat( 0.0 , NaN)) )
                    else
                        m.c[id,iy,it] = Envelope(MLine(vcat(p.a_lowT,p.a_high),vcat( 0.0 , p.a_high)) )
                        m.v[id,iy,it] = Envelope(MLine(vcat(p.a_lowT,p.a_high),vcat( 0.0 + p.alphaT  , NaN)) )
                        m.v[id,iy,it].vbound = p.alphaT
                    end
                    # initialize value function with vf(1) = 0
                    # m.v[id,iy,it] = Envelope(MLine(vcat(p.a_lowT,p.a_high),vcat(id == 1 ? 0.0 : p.alphaT,p.alphaT)) )

                    # with bk flag on:
                    m.cbk[iy,it] = Envelope(MLine(vcat(0.0,p.a_high),vcat(0.0,p.a_high)) )
                    # initiz,iae value function with vf(1) = 0
                    m.vbk[iy,it] = Envelope(MLine(vcat(0.0,p.a_high),vcat(p.alphaT,NaN)) )
                    m.vbk[iy,it].vbound = p.alphaT
                else
                    # next period consumption and values y-coords
                    # for each d-choice
                    # reset all value matrices
                    fill!(vmat,0.0)
                    fill!(ctmp,-Inf)
                    fill!(cmat,0.0)
                    fill!(cmat2,0.0)

                    # fill!(vtmp,0.0)#

                    # you are the bank
                    # you know how pairs (ia,iy) map into bankruptcies in any given period
                    # like you observe people in t+1 and you can now know that a certain asset state ia together with iy
                    # corresponds to "bankruptcy".
                    # in period t, you can compute an expectation over this.
                    prnobk_at_ia = zeros(p.na)
                    for jy in 1:p.ny # future state: owner, renter, income, etc
                        # E[ m.v[1,jy,it+1] vs m.v[2,jy,it+1] | iy ]
                        # m.v is defined on cash on hand M = a + y, so I can evaluate it at any M
                        # say M = iy + ia
                        pr = m.ywgt[iy,jy]  # transprob
                        # println("v1 = $(gety(interp(m.v[1,jy,it+1].env, income(it+1,p,m.yvec[jy]) .+ m.avec )))")
                        if it+1 == p.nT
                            I = trues(p.na)
                        else
                            may = income(it+1,p,m.yvec[jy]) .+ m.avec
                            b = [gety(interp(m.v[1,jy,it+1].env, may )) gety(interp(m.v[2,jy,it+1].env, may))]
                            b[isnan.(b)] .= -Inf
                            I = b[:,1] .> b[:,2]
                        end
                        prnobk_at_ia[:] += pr * I
                    end

                    # risk adjusted return times asset grid
                    m.rgrid[id,iy,it] = clamp!([(p.R / prnobk_at_ia[ia]) for ia in 1:p.na], p.R, p.maxinterest)
                    # rgrid = p.R
                    # println("rgrid = $rgrid")

                    ragrid = m.rgrid[id,iy,it] .* m.avec

                    if !filer
                        fill!(cmat,0.0)

                        for jy in 1:p.ny # future state: owner, renter, income, etc
                            pr = m.ywgt[iy,jy]  # transprob

                            for iid in 1:p.nD  # future dchoice
                                # only feasible choices at this state
                                # if renter, cannot sell etc
                                if iid == 1
                                    # cash on hand : R (a[ia])*a[ia] + y(t+1)
                                    # R[ia] = (1+r) / pr_no_bk_iy[ia]
                                    # cash on hand:
                                    m1 = ragrid .+ income(it+1,p,m.yvec[jy])
                                    # m1 = m.m1[it+1][iid][jy,:]  # state specific mvec
                                    c1 = interp(m.c[iid,jy,it+1].env, m1) # C(d',y',m')
                                    floory!(c1,p.cfloor)   # floor negative consumption
                                    ctmp[iid,jy,:] = gety(c1)
                                    # vtmp[iid,jy,:] = vfun(iid,it+1,ctmp[iid,jy,:],m1,m.v[iid,jy,it+1],p)
                                    vmat[iid,:] += pr * vfun(x->u(x,filer,p),it+1,ctmp[iid,jy,:],m1,m.v[iid,jy,it+1],p)
                                    cmat[iid,:] += pr * up(ctmp[iid,jy,:],p)
                                else
                                    # you are a filer: no savings and analytic value function
                                    ctmp[iid,jy,:] .= income(it+1,p,repeat([m.yvec[jy]],p.na))
                                    vmat[iid,:] += pr * (u(ctmp[iid,jy,:],true,p) .+ p.beta * bound(m.vbk[jy,it+1]))
                                    cmat[iid,:] += pr * up(ctmp[iid,jy,:],p)
                                    vmat[iid,m.iazero:end] .= -Inf   # filing only with negative assets!
                                    cmat[iid,m.iazero:end] .= 0.0   # filing only with negative assets!
                                end
                            end
                        end # end future state

                        # Discrete choice setup next period
                        # =================================

                        # get ccp of filing next period at m1 : P(d'|iy), pfile
                        pnofile = ccp(vmat,p)   # prob of filing TOMORROW from today's perspective
                        # m.pnofile[iy,it,:] = pnofile   # prob of filing at state ia next period, given iy in period it.

                        # compute tomorrow's marginal utility
                        mu1 = pnofile .* cmat[1,:] .+ (1.0 .- pnofile) .* cmat[2,:] # 1,na

                        #RHS
                        # TODO R here will be computed above as the expected return of the bank
                        # in each discrete choice
                        RHS = p.beta * p.R * mu1

                        #optimal cons
                        c0 = iup(RHS,p)
                        # set optimal consumption function today. endo grid m and cons c0
                        cline = MLine(m.avec .+ c0, c0)
                        # consumption function done.

                        # compute value function
                        # ----------------------
                        ev =  logsum(vmat,p)  # E[ max(V(1,y'),V(2,y')) | y]
                        vline = MLine(m.avec .+ c0, u(c0,id==2,p) .+ p.beta * ev[:])
                        minv = gety(vline)[1]

                        # store the expected value at the lower boundary
                        # in a separate object
                        minv = ev[1] < minv ? ev[1] : minv

                        # SECONDARY ENVELOPE COMPUTATION
                        # ==============================
                        m.v[id,iy,it], m.c[id,iy,it] = do_secondary(vline,cline,filer,minv,p)
                        m.v[id,iy,it].vbound = minv

                        # this creates the credit constrained region
                        prepend!(m.c[id,iy,it].env,[Point(m.avec[1],0.0)])
                        prepend!(m.v[id,iy,it].env,[Point(m.avec[1],m.v[id,iy,it].vbound)])
                        sortx!(m.c[id,iy,it].env)  # sort cons by default

                    else  # if you are filing today (or you have filed previously)

                        iid = 2 # hard code to 2 - we only will use the second row of vmat here
                        # there is no discrete choice for you next period.

                        # we compute the value and savings function of being in bk state first.
                        for jy in 1:p.ny # future state:
                            pr = m.ywgt[iy,jy]  # transprob
                            m1 = m.m1[it+1][jy,:]  # cash on hand in bk state
                            c1 = interp(m.cbk[jy,it+1].env, m1) # C(d',y',m')
                            floory!(c1,p.cfloor)   # floor negative consumption
                            ctmp[iid,jy,:] = gety(c1)
                            # vtmp[iid,jy,:] = vfun(iid,it+1,ctmp[iid,jy,:],m1,m.v[iid,jy,it+1],p)
                            vmat[iid,:] += (pr * vfun(x->u(x,filer,p),it+1,ctmp[iid,jy,:],m1,m.vbk[jy,it+1],p))
                            cmat[iid,:] += (pr * up(ctmp[iid,jy,:],p))  # E[ c(id',y') | y]

                            # # now get the same for non-bk state
                            # c1 = interp(m.c[1,jy,it+1].env, m1) # C(d',y',m')
                            # floory!(c1,p.cfloor)   # floor negative consumption
                            # ctmp[iid,jy,:] = gety(c1)
                            # # vtmp[iid,jy,:] = vfun(iid,it+1,ctmp[iid,jy,:],m1,m.v[iid,jy,it+1],p)
                            # vmat[iid,:] += (pr * vfun(x->u(x,filer,p),it+1,ctmp[iid,jy,:],m1,m.vbk[jy,it+1],p))
                            # cmat[iid,:] += (pr * up(ctmp[iid,jy,:],p))  # E[ c(id',y') | y]

                        end  # future states

                        mu1 = cmat[iid,:]
                        #RHS
                        RHS = p.beta * p.R * mu1
                        #optimal cons
                        c0 = iup(RHS,p)
                        # set optimal consumption function today. endo grid m and cons c0
                        cline = MLine(m.aposvec .+ c0, c0)
                        # consumption function done.
                        ev = vmat[iid,:]
                        vline = MLine(m.aposvec .+ c0, u(c0,id==2,p) .+ p.beta * ev[:])

                        m.cbk[iy,it] = Envelope(cline)
                        m.vbk[iy,it] = Envelope(vline)
                        m.vbk[iy,it].vbound = ev[1]

                        # this creates the credit constrained region
                        prepend!(m.cbk[iy,it].env,[Point(m.aposvec[1],0.0)])
                        prepend!(m.vbk[iy,it].env,[Point(m.aposvec[1],ev[1])])
                        sortx!(m.cbk[iy,it].env)  # sort cons by default
                        # end value in bk state

                        # value and cons function if filing today
                        c0 = income(it,p,repeat([m.yvec[iy]],p.na))
                        c01 = copy(c0)
                        c01[m.iazero:end] .= NaN
                        cline = MLine(m.avec .+ c0, c01)
                        vline = MLine(m.avec .+ c0, u(c01, id == 2, p) .+ p.beta * ev[1])
                        m.c[id,iy,it] = Envelope(cline)
                        m.v[id,iy,it] = Envelope(vline)
                        prepend!(m.c[id,iy,it].env,[Point(m.avec[1],c01[1])])
                        prepend!(m.v[id,iy,it].env,[Point(m.avec[1],gety(vline)[1])])

                        m.v[id,iy,it].vbound = NaN   # does not exist - you don't save when filing.
                    end
                end # if last period
            end  # current id
        end  # iy
    end # it
end

function runbk(; par = Dict(:nT => 50, :a_low => -5.0,:a_lowT => -5.0,:na =>101, :alpha => 0.0,  :alphaT => 0.0, :lambda => 0.5))
    p = Param(par = par)
    m = BModel(p)
    bk!(m,p)
    m,p
end


text:
@testset "background noise mask" begin
    for i in 1:10
        x = rand(10:30)
        y = rand(10:30)
        z = rand(13:30)
        rad = rand(20:100.0, x,y,z)
        clouds = rand(1:30, x,y,z)
        @test size(background_noise_mask(rad,clouds,rand(0:0.9))) == (x,y)
    end
end


text:
"""
A list of physical constants.

# References

## CODATA 2018 constants

* Mohr, P. J., Newell, D. B., & Taylor, B. N. (2016a). CODATA recommended
  values of the fundamental physical constants: 2014. _Reviews of Modern
  Physics_, _88_(3), 035009. <https://doi.org/10.1103/RevModPhys.88.035009>
* Mohr, P. J., Newell, D. B., & Taylor, B. N. (2016b). CODATA Recommended
  Values of the Fundamental Physical Constants: 2014. _Journal of Physical and
  Chemical Reference Data_, _45_(4), 043102.
  <https://doi.org/10.1063/1.4954402>
* Mohr, P. J., Newell, D. B., Taylor, B. N., & Tiesinga, E. (2018). Data and
  analysis for the CODATA 2017 special fundamental constants adjustment.
  _Metrologia_, _55_(1), 125--146. <https://doi.org/10.1088/1681-7575/aa99bc>

## Molar masses

* Rumble, J. R. (Ed.). (2017). _CRC handbook of chemistry and physics_ (98th
  ed.). CRC Press, Taylor & Francis Group: Boca Raton, FL, USA. ISBN:
  9781498784542.
  Retrieved from <http://hbcponline.com/faces/contents/ContentsSearch.xhtml> on
  2018-06-12.
"""
module Constants

# Notes
#
# 1. Use lowercase variable names, except when there is already a commonly used
#    notation otherwise.
# 2. Keep the long names self-explanatory.

export c,
    h,
    N_A,
    k_B,
    sigma_SB,
    c_1,
    c_1L,
    c_2,
    R,
    g,
    atm,
    T_0,
    kappa,
    M_w,
    M_c,
    M_d,
    cp_d

"""Speed of light in the vacuum [m s^-1]."""
const c = 2.997_924_58e8

"""Planck's constant [J s]."""
const h = 6.626_070_15e-34

"""First radiation constant \$2\\pi hc^2\$ [W m^2]."""
const c_1 = 3.741_771_852e-16

"""First radiation constant for spectral radiance \$2hc^2\$ [W m^-2 sr^-1]."""
const c_1L = 1.191_042_972e-16

"""Second radiation constant \$ hc / k\$ [m K]."""
const c_2 = 1.438_776_877e-2

"""Stefan--Boltzmann constant [W m^-2 K^-4]."""
const sigma_SB = 5.670_374_419e-8

"""Boltzmann constant [J K^-1]."""
const k_B = 1.380_649e-23

"""Avogadro constant [mol^-1]."""
const N_A = 6.022_140_76e23

"""Molar gas constant [J mol^-1 K^-1]."""
const R = 8.314_462_618

"""Standard atmospheric pressure [Pa]."""
const atm = 1.013_25e5

"""Standard acceleration of gravity [m s^-2]."""
const g = 9.806_65

"""Zero Celsius [K]."""
const T_0 = 273.15

"""Von Kármán's constant [-]."""
const kappa = 0.40

"""Molar mass of water [kg mol^-1]."""
const M_w = 18.015_28e-3

"""Molar mass of CO2 [kg mol^-1]."""
const M_c = 44.010e-3

"""Molar mass of dry air [kg mol^-1]."""
const M_d = 28.9645e-3

"""Molar isobaric specific heat capacity of dry air [J K^-1 mol^-1]."""
const cp_d = 1.004e3 * M_d
# Note: using a constant value has less than 2% error between -40 and 60°C

"""A dictionary of physical constants."""
const constants = Dict{String,Float64}(
    "speed of light" => c,
    "Planck constant" => h,
    "Avogadro constant" => N_A,
    "Boltzmann constant" => k_B,
    "Stefan-Boltzmann constant" => sigma_SB,
    "first radiation constant" => c_1,
    "first radiation constant for spectral radiance" => c_1L,
    "second radiation constant" => c_2,
    "molar gas constant" => R,
    "standard acceleration of gravity" => g,
    "standard atmosphere" => atm,
    "zero Celsius" => T_0,
    "von Karman constant" => kappa,
    "molar mass of water" => M_w,
    "molar mass of CO2" => M_c,
    "molar mass of dry air" => M_d,
    "molar isobaric specific heat capacity of dry air" => cp_d,
)

end  # module


text:
using ClimateMachine

ClimateMachine.init()

using ClimateMachine.CartesianDomains

@testset "Cartesian domains" begin

    for FT in (Float64, Float32)
        domain = RectangularDomain(
            FT,
            Ne = (16, 24, 1),
            Np = 4,
            x = (0, π),
            y = (0, 1.1),
            z = (-1, 0),
            periodicity = (false, false, false),
        )

        @test eltype(domain) == FT
        @test domain.Ne == (x = 16, y = 24, z = 1)
        @test domain.Np == 4
        @test domain.L.x == FT(π)
        @test domain.L.y == FT(1.1)
        @test domain.L.z == FT(1)
    end
end


text:
using Harlequin
using Test

@testset "Quaternions" begin
    include("test_quaternions.jl")
end

@testset "Pointing generation" begin
    include("test_genpointings.jl")
end

@testset "Segmented time spans" begin
    include("test_sts.jl")
end

@testset "Dipole analysis" begin
    include("test_dipole.jl")
end

@testset "Beam functions" begin
    include("test_beams.jl")
end

@testset "Map-making functions" begin
    include("test_mapmaking.jl")
end

@testset "MPI" begin
    import MPI: mpiexec

    nprocs = 2  # test_mapmaking.jl assumes that only two processes are being used!
    juliaexec = joinpath(Sys.BINDIR, Base.julia_exename())
    @test success(`$mpiexec -n $nprocs $juliaexec test_mpi.jl`)
end


text:
###############################################################################
#
#   arb_mat.jl : Arb matrices over arb
#
###############################################################################

export rows, cols, zero, one, deepcopy, -, transpose, +, *, &, ==, !=,
       strongequal, overlaps, contains, inv, divexact, charpoly, det,
       lufact, lufact!, solve, solve!, solve_lu_precomp, solve_lu_precomp!,
       swap_rows, swap_rows!, bound_inf_norm

###############################################################################
#
#   Similar
#
###############################################################################

function similar(x::arb_mat)
   z = arb_mat(rows(x), cols(x))
   z.base_ring = x.base_ring
   return z
end

function similar(x::arb_mat, r::Int, c::Int)
   z = arb_mat(r, c)
   z.base_ring = x.base_ring
   return z
end

###############################################################################
#
#   Basic manipulation
#
###############################################################################

parent_type(::Type{arb_mat}) = ArbMatSpace

elem_type(::Type{ArbMatSpace}) = arb_mat

base_ring(a::ArbMatSpace) = a.base_ring

base_ring(a::arb_mat) = a.base_ring

parent(x::arb_mat, cached::Bool = true) =
      MatrixSpace(base_ring(x), rows(x), cols(x))

prec(x::ArbMatSpace) = prec(x.base_ring)

function getindex!(z::arb, x::arb_mat, r::Int, c::Int)
  v = ccall((:arb_mat_entry_ptr, :libarb), Ref{arb},
              (Ref{arb_mat}, Int, Int), x, r - 1, c - 1)
  ccall((:arb_set, :libarb), Void, (Ref{arb}, Ref{arb}), z, v)
  return z
end

@inline function getindex(x::arb_mat, r::Int, c::Int)
  @boundscheck Generic._checkbounds(x, r, c)

  z = base_ring(x)()
  v = ccall((:arb_mat_entry_ptr, :libarb), Ref{arb},
              (Ref{arb_mat}, Int, Int), x, r - 1, c - 1)
  ccall((:arb_set, :libarb), Void, (Ref{arb}, Ref{arb}), z, v)
  return z
end

for T in [Int, UInt, fmpz, fmpq, Float64, BigFloat, arb, AbstractString]
   @eval begin
      @inline function setindex!(x::arb_mat, y::$T, r::Int, c::Int)
         @boundscheck Generic._checkbounds(x, r, c)

         z = ccall((:arb_mat_entry_ptr, :libarb), Ref{arb},
                   (Ref{arb_mat}, Int, Int), x, r - 1, c - 1)
         Nemo._arb_set(z, y, prec(base_ring(x)))
      end
   end
end

Base.@propagate_inbounds setindex!(x::arb_mat, y::Integer,
                                 r::Int, c::Int) =
         setindex!(x, fmpz(y), r, c)

Base.@propagate_inbounds setindex!(x::arb_mat, y::Rational{T},
                                 r::Int, c::Int) where {T <: Integer} =
         setindex!(x, fmpz(y), r, c)

zero(a::ArbMatSpace) = a()

function one(x::ArbMatSpace)
  z = x()
  ccall((:arb_mat_one, :libarb), Void, (Ref{arb_mat}, ), z)
  return z
end

rows(a::arb_mat) = a.r

cols(a::arb_mat) = a.c

function deepcopy_internal(x::arb_mat, dict::ObjectIdDict)
  z = arb_mat(rows(x), cols(x))
  ccall((:arb_mat_set, :libarb), Void, (Ref{arb_mat}, Ref{arb_mat}), z, x)
  z.base_ring = x.base_ring
  return z
end

###############################################################################
#
#   String I/O
#
###############################################################################

function show(io::IO, a::ArbMatSpace)
   print(io, "Matrix Space of ")
   print(io, a.rows, " rows and ", a.cols, " columns over ")
   print(io, base_ring(a))
end

function show(io::IO, a::arb_mat)
   r = rows(a)
   c = cols(a)
   for i = 1:r
      print(io, "[")
      for j = 1:c
         print(io, a[i, j])
         if j != c
            print(io, " ")
         end
      end
      print(io, "]")
      if i != r
         println(io, "")
      end
   end
end

################################################################################
#
#  Unary operations
#
################################################################################

function -(x::arb_mat)
  z = similar(x)
  ccall((:arb_mat_neg, :libarb), Void, (Ref{arb_mat}, Ref{arb_mat}), z, x)
  return z
end

################################################################################
#
#  Transpose
#
################################################################################

function transpose(x::arb_mat)
  z = similar(x, cols(x), rows(x))
  ccall((:arb_mat_transpose, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}), z, x)
  return z
end

################################################################################
#
#  Binary operations
#
################################################################################

function +(x::arb_mat, y::arb_mat)
  check_parent(x, y)
  z = similar(x)
  ccall((:arb_mat_add, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb_mat}, Int),
              z, x, y, prec(parent(x)))
  return z
end

function -(x::arb_mat, y::arb_mat)
  check_parent(x, y)
  z = similar(x)
  ccall((:arb_mat_sub, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb_mat}, Int),
              z, x, y, prec(parent(x)))
  return z
end

function *(x::arb_mat, y::arb_mat)
  cols(x) != rows(y) && error("Matrices have wrong dimensions")
  z = similar(x, rows(x), cols(y))
  ccall((:arb_mat_mul, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb_mat}, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

################################################################################
#
#   Ad hoc binary operators
#
################################################################################

function ^(x::arb_mat, y::UInt)
  rows(x) != cols(x) && error("Matrix must be square")
  z = similar(x)
  ccall((:arb_mat_pow_ui, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, UInt, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

function *(x::arb_mat, y::Int)
  z = similar(x)
  ccall((:arb_mat_scalar_mul_si, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Int, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

*(x::Int, y::arb_mat) = y*x

function *(x::arb_mat, y::fmpz)
  z = similar(x)
  ccall((:arb_mat_scalar_mul_fmpz, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{fmpz}, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

*(x::fmpz, y::arb_mat) = y*x

function *(x::arb_mat, y::arb)
  z = similar(x)
  ccall((:arb_mat_scalar_mul_arb, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb}, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

*(x::arb, y::arb_mat) = y*x

for T in [Integer, fmpz, fmpq, arb]
   @eval begin
      function +(x::arb_mat, y::$T)
         z = deepcopy(x)
         for i = 1:min(rows(x), cols(x))
            z[i, i] += y
         end
         return z
      end

      +(x::$T, y::arb_mat) = y + x

      function -(x::arb_mat, y::$T)
         z = deepcopy(x)
         for i = 1:min(rows(x), cols(x))
            z[i, i] -= y
         end
         return z
      end

      function -(x::$T, y::arb_mat)
         z = -y
         for i = 1:min(rows(y), cols(y))
            z[i, i] += x
         end
         return z
      end
   end
end

function +(x::arb_mat, y::Rational{T}) where T <: Union{Int, BigInt}
   z = deepcopy(x)
   for i = 1:min(rows(x), cols(x))
      z[i, i] += y
   end
   return z
end

+(x::Rational{T}, y::arb_mat) where T <: Union{Int, BigInt} = y + x

function -(x::arb_mat, y::Rational{T}) where T <: Union{Int, BigInt}
   z = deepcopy(x)
   for i = 1:min(rows(x), cols(x))
      z[i, i] -= y
   end
   return z
end

function -(x::Rational{T}, y::arb_mat) where T <: Union{Int, BigInt}
   z = -y
   for i = 1:min(rows(y), cols(y))
      z[i, i] += x
   end
   return z
end

###############################################################################
#
#   Shifting
#
###############################################################################

doc"""
    ldexp(x::acb_mat, y::Int)
> Return $2^yx$. Note that $y$ can be positive, zero or negative.
"""
function ldexp(x::arb_mat, y::Int)
  z = similar(x)
  ccall((:arb_mat_scalar_mul_2exp_si, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Int), z, x, y)
  return z
end

###############################################################################
#
#   Comparisons
#
###############################################################################

doc"""
    isequal(x::arb_mat, y::arb_mat)
> Return `true` if the matrices of balls $x$ and $y$ are precisely equal,
> i.e. if all matrix entries have the same midpoints and radii.
"""
function isequal(x::arb_mat, y::arb_mat)
  r = ccall((:arb_mat_equal, :libarb), Cint,
              (Ref{arb_mat}, Ref{arb_mat}), x, y)
  return Bool(r)
end

function ==(x::arb_mat, y::arb_mat)
  check_parent(x, y)
  r = ccall((:arb_mat_eq, :libarb), Cint, (Ref{arb_mat}, Ref{arb_mat}), x, y)
  return Bool(r)
end

function !=(x::arb_mat, y::arb_mat)
  r = ccall((:arb_mat_ne, :libarb), Cint, (Ref{arb_mat}, Ref{arb_mat}), x, y)
  return Bool(r)
end

doc"""
    overlaps(x::arb_mat, y::arb_mat)
> Returns `true` if all entries of $x$ overlap with the corresponding entry of
> $y$, otherwise return `false`.
"""
function overlaps(x::arb_mat, y::arb_mat)
  r = ccall((:arb_mat_overlaps, :libarb), Cint,
              (Ref{arb_mat}, Ref{arb_mat}), x, y)
  return Bool(r)
end

doc"""
    contains(x::arb_mat, y::arb_mat)
> Returns `true` if all entries of $x$ contain the corresponding entry of
> $y$, otherwise return `false`.
"""
function contains(x::arb_mat, y::arb_mat)
  r = ccall((:arb_mat_contains, :libarb), Cint,
              (Ref{arb_mat}, Ref{arb_mat}), x, y)
  return Bool(r)
end

###############################################################################
#
#   Ad hoc comparisons
#
###############################################################################

doc"""
    contains(x::arb_mat, y::fmpz_mat)
> Returns `true` if all entries of $x$ contain the corresponding entry of
> $y$, otherwise return `false`.
"""
function contains(x::arb_mat, y::fmpz_mat)
  r = ccall((:arb_mat_contains_fmpz_mat, :libarb), Cint,
              (Ref{arb_mat}, Ref{fmpz_mat}), x, y)
  return Bool(r)
end

doc"""
    contains(x::arb_mat, y::fmpq_mat)
> Returns `true` if all entries of $x$ contain the corresponding entry of
> $y$, otherwise return `false`.
"""
function contains(x::arb_mat, y::fmpq_mat)
  r = ccall((:arb_mat_contains_fmpq_mat, :libarb), Cint,
              (Ref{arb_mat}, Ref{fmpq_mat}), x, y)
  return Bool(r)
end

==(x::arb_mat, y::Integer) = x == parent(x)(y)

==(x::Integer, y::arb_mat) = y == x

==(x::arb_mat, y::fmpz) = x == parent(x)(y)

==(x::fmpz, y::arb_mat) = y == x

==(x::arb_mat, y::fmpz_mat) = x == parent(x)(y)

==(x::fmpz_mat, y::arb_mat) = y == x

###############################################################################
#
#   Inversion
#
###############################################################################

doc"""
    inv(M::arb_mat)
> Given a  $n\times n$ matrix of type `arb_mat`, return an
> $n\times n$ matrix $X$ such that $AX$ contains the
> identity matrix. If $A$ cannot be inverted numerically an exception is raised.
"""
function inv(x::arb_mat)
  cols(x) != rows(x) && error("Matrix must be square")
  z = similar(x)
  r = ccall((:arb_mat_inv, :libarb), Cint,
              (Ref{arb_mat}, Ref{arb_mat}, Int), z, x, prec(base_ring(x)))
  Bool(r) ? (return z) : error("Matrix cannot be inverted numerically")
end

###############################################################################
#
#   Exact division
#
###############################################################################

function divexact(x::arb_mat, y::arb_mat)
   cols(x) != cols(y) && error("Incompatible matrix dimensions")
   x*inv(y)
end

###############################################################################
#
#   Ad hoc exact division
#
###############################################################################

function divexact(x::arb_mat, y::Int)
  y == 0 && throw(DivideError())
  z = similar(x)
  ccall((:arb_mat_scalar_div_si, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Int, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

function divexact(x::arb_mat, y::fmpz)
  z = similar(x)
  ccall((:arb_mat_scalar_div_fmpz, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{fmpz}, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

function divexact(x::arb_mat, y::arb)
  z = similar(x)
  ccall((:arb_mat_scalar_div_arb, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb}, Int),
              z, x, y, prec(base_ring(x)))
  return z
end

################################################################################
#
#  Characteristic polynomial
#
################################################################################

function charpoly(x::ArbPolyRing, y::arb_mat)
  base_ring(y) != base_ring(x) && error("Base rings must coincide")
  z = x()
  ccall((:arb_mat_charpoly, :libarb), Void,
              (Ref{arb_poly}, Ref{arb_mat}, Int), z, y, prec(base_ring(y)))
  return z
end

###############################################################################
#
#   Determinant
#
###############################################################################

function det(x::arb_mat)
  cols(x) != rows(x) && error("Matrix must be square")
  z = base_ring(x)()
  ccall((:arb_mat_det, :libarb), Void,
              (Ref{arb}, Ref{arb_mat}, Int), z, x, prec(base_ring(x)))
  return z
end

################################################################################
#
#   Exponential function
#
################################################################################

doc"""
    exp(x::arb_mat)
> Returns the exponential of the matrix $x$.
"""
function Base.exp(x::arb_mat)
  cols(x) != rows(x) && error("Matrix must be square")
  z = similar(x)
  ccall((:arb_mat_exp, :libarb), Void,
              (Ref{arb_mat}, Ref{arb_mat}, Int), z, x, prec(base_ring(x)))
  return z
end

###############################################################################
#
#   Linear solving
#
###############################################################################

function lufact!(P::Generic.perm, x::arb_mat)
  cols(x) != rows(x) && error("Matrix must be square")
  parent(P).n != rows(x) && error("Permutation does not match matrix")
  P.d .-= 1
  r = ccall((:arb_mat_lu, :libarb), Cint,
              (Ref{Int}, Ref{arb_mat}, Ref{arb_mat}, Int),
              P.d, x, x, prec(base_ring(x)))
  r == 0 && error("Could not find $(rows(x)) invertible pivot elements")
  P.d .+= 1
  inv!(P)
  return rows(x)
end

function lufact(x::arb_mat, P = PermGroup(rows(x)))
  p = P()
  R = base_ring(x)
  L = similar(x)
  U = deepcopy(x)
  n = cols(x)
  r = lufact!(p, U)
  for i = 1:n
    for j = 1:n
      if i > j
        L[i, j] = U[i, j]
        U[i, j] = R()
      elseif i == j
        L[i, j] = one(R)
      else
        L[i, j] = R()
      end
    end
  end
  return r, p, L, U
end

function solve!(z::arb_mat, x::arb_mat, y::arb_mat)
  r = ccall((:arb_mat_solve, :libarb), Cint,
              (Ref{arb_mat}, Ref{arb_mat}, Ref{arb_mat}, Int),
              z, x, y, prec(base_ring(x)))
  r == 0 && error("Matrix cannot be inverted numerically")
  nothing
end

function solve(x::arb_mat, y::arb_mat)
  cols(x) != rows(x) && error("First argument must be square")
  cols(x) != rows(y) && error("Matrix dimensions are wrong")
  z = similar(y)
  solve!(z, x, y)
  return z
end

function solve_lu_precomp!(z::arb_mat, P::Generic.perm, LU::arb_mat, y::arb_mat)
  Q = inv(P)
  ccall((:arb_mat_solve_lu_precomp, :libarb), Void,
              (Ref{arb_mat}, Ref{Int}, Ref{arb_mat}, Ref{arb_mat}, Int),
              z, Q.d .- 1, LU, y, prec(base_ring(LU)))
  nothing
end

function solve_lu_precomp(P::Generic.perm, LU::arb_mat, y::arb_mat)
  cols(LU) != rows(y) && error("Matrix dimensions are wrong")
  z = similar(y)
  solve_lu_precomp!(z, P, LU, y)
  return z
end

################################################################################
#
#   Row swapping
#
################################################################################

function swap_rows(x::arb_mat, i::Int, j::Int)
  Generic._checkbounds(rows(x), i) || throw(BoundsError())
  Generic._checkbounds(rows(x), j) || throw(BoundsError())
  z = deepcopy(x)
  swap_rows!(z, i, j)
  return z
end

function swap_rows!(x::arb_mat, i::Int, j::Int)
  ccall((:arb_mat_swap_rows, :libarb), Void,
              (Ref{arb_mat}, Ref{Void}, Int, Int),
              x, C_NULL, i - 1, j - 1)
end

################################################################################
#
#   Norm
#
################################################################################

doc"""
    bound_inf_norm(x::arb_mat)
> Returns a nonnegative element $z$ of type `arb`, such that $z$ is an upper
> bound for the infinity norm for every matrix in $x$
"""
function bound_inf_norm(x::arb_mat)
  z = arb()
  t = ccall((:arb_rad_ptr, :libarb), Ref{mag_struct}, (Ref{arb}, ), z)
  ccall((:arb_mat_bound_inf_norm, :libarb), Void,
              (Ref{mag_struct}, Ref{arb_mat}), t, x)
  s = ccall((:arb_mid_ptr, :libarb), Ref{arf_struct}, (Ref{arb}, ), z)
  ccall((:arf_set_mag, :libarb), Void,
              (Ref{arf_struct}, Ref{mag_struct}), s, t)
  ccall((:mag_zero, :libarb), Void,
              (Ref{mag_struct},), t)
  return base_ring(x)(z)
end

################################################################################
#
#   Unsafe functions
#
################################################################################

for (s,f) in (("add!","arb_mat_add"), ("mul!","arb_mat_mul"),
              ("sub!","arb_mat_sub"))
  @eval begin
    function ($(Symbol(s)))(z::arb_mat, x::arb_mat, y::arb_mat)
      ccall(($f, :libarb), Void,
                  (Ref{arb_mat}, Ref{arb_mat}, Ref{arb_mat}, Int),
                  z, x, y, prec(base_ring(x)))
      return z
    end
  end
end

###############################################################################
#
#   Parent object call overloads
#
###############################################################################

function (x::ArbMatSpace)()
  z = arb_mat(x.rows, x.cols)
  z.base_ring = x.base_ring
  return z
end

function (x::ArbMatSpace)(y::fmpz_mat)
  (x.cols != cols(y) || x.rows != rows(y)) &&
      error("Dimensions are wrong")
  z = arb_mat(y, prec(x))
  z.base_ring = x.base_ring
  return z
end

function (x::ArbMatSpace)(y::Array{T, 2}) where {T <: Union{Int, UInt, fmpz, fmpq, Float64, BigFloat, arb, AbstractString}}
  _check_dim(x.rows, x.cols, y)
  z = arb_mat(x.rows, x.cols, y, prec(x))
  z.base_ring = x.base_ring
  return z
end

function (x::ArbMatSpace)(y::Array{T, 1}) where {T <: Union{Int, UInt, fmpz, fmpq, Float64, BigFloat, arb, AbstractString}}
  _check_dim(x.rows, x.cols, y)
  z = arb_mat(x.rows, x.cols, y, prec(x))
  z.base_ring = x.base_ring
  return z
end

function (x::ArbMatSpace)(y::Union{Int, UInt, fmpz, fmpq, Float64,
                          BigFloat, arb, AbstractString})
  z = x()
  for i in 1:rows(z)
      for j = 1:cols(z)
         if i != j
            z[i, j] = zero(base_ring(x))
         else
            z[i, j] = y
         end
      end
   end
   return z
end

(x::ArbMatSpace)(y::arb_mat) = y

###############################################################################
#
#   Matrix constructor
#
###############################################################################

function matrix(R::ArbField, arr::Array{T, 2}) where {T <: Union{Int, UInt, fmpz, fmpq, Float64, BigFloat, arb, AbstractString}}
   z = arb_mat(size(arr, 1), size(arr, 2), arr, prec(R))
   z.base_ring = R
   return z
end

function matrix(R::ArbField, r::Int, c::Int, arr::Array{T, 1}) where {T <: Union{Int, UInt, fmpz, fmpq, Float64, BigFloat, arb, AbstractString}}
   _check_dim(r, c, arr)
   z = arb_mat(r, c, arr, prec(R))
   z.base_ring = R
   return z
end

function matrix(R::ArbField, arr::Array{<: Integer, 2})
   arr_fmpz = map(fmpz, arr)
   return matrix(R, arr_fmpz)
end

function matrix(R::ArbField, r::Int, c::Int, arr::Array{<: Integer, 1})
   arr_fmpz = map(fmpz, arr)
   return matrix(R, r, c, arr_fmpz)
end

function matrix(R::ArbField, arr::Array{Rational{T}, 2}) where {T <: Integer}
   arr_fmpz = map(fmpq, arr)
   return matrix(R, arr_fmpz)
end

function matrix(R::ArbField, r::Int, c::Int, arr::Array{Rational{T}, 1}) where {T <: Integer}
   arr_fmpz = map(fmpq, arr)
   return matrix(R, r, c, arr_fmpz)
end

###############################################################################
#
#  Zero matrix
#
###############################################################################

function zero_matrix(R::ArbField, r::Int, c::Int)
   z = arb_mat(r, c)
   z.base_ring = R
   return z
end

###############################################################################
#
#  Identity matrix
#
###############################################################################

function identity_matrix(R::ArbField, n::Int)
   z = arb_mat(n, n)
   ccall((:arb_mat_one, :libarb), Void, (Ref{arb_mat}, ), z)
   z.base_ring = R
   return z
end

###############################################################################
#
#   Promotions
#
###############################################################################

promote_rule(::Type{arb_mat}, ::Type{T}) where {T <: Integer} = arb_mat

promote_rule(::Type{arb_mat}, ::Type{Rational{T}}) where T <: Union{Int, BigInt} = arb_mat

promote_rule(::Type{arb_mat}, ::Type{fmpz}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{fmpq}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{arb}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{Float64}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{BigFloat}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{fmpz_mat}) = arb_mat

promote_rule(::Type{arb_mat}, ::Type{fmpq_mat}) = arb_mat

###############################################################################
#
#   MatrixSpace constructor
#
###############################################################################

function MatrixSpace(R::ArbField, r::Int, c::Int; cached = true)
  (r <= 0 || c <= 0) && error("Dimensions must be positive")
  return ArbMatSpace(R, r, c, cached)
end


text:
using MBTR
using Test, LinearAlgebra, BenchmarkTools
using JuLIP, JuLIP.Testing
using MBTR: mbtr

@testset "MBTR.jl" begin
    include("test.jl")
end


text:
# The following variables need to be defined before `include`-ing this file:
# `testset_name_suffix`
# `n_values_1`
# `n_values_2`
# `m_values_1`
# `m_values_2`
# `sz_values`

@time @testset "Float64 Matrix-Vector $(testset_name_suffix)" begin
    for n ∈ n_values_1
        for m ∈ m_values_1
            u1 = zeros(n)
            u2 = zeros(n)
            A = rand(n, m)
            v = rand(m)
            @test Gaius.mul!(u1, A, v) ≈ A * v
            @test Gaius.mul_serial!(u2, A, v) ≈ A * v
            @test u1 ≈ Gaius.mul(A, v)
            @test u2 ≈ Gaius.mul(A, v)
            @test u1 ≈ Gaius.mul_serial(A, v)
            @test u2 ≈ Gaius.mul_serial(A, v)
        end
    end
end

@time @testset "ComplexF64 Matrix-Vector $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(ComplexF64, n)   |> StructArray
            u2 = zeros(ComplexF64, n)   |> StructArray
            A = rand(ComplexF64, n, m) |> StructArray
            v = rand(ComplexF64, m)    |> StructArray
            @test Gaius.mul!(u1, A, v) ≈ collect(A) * collect(v)
            @test Gaius.mul_serial!(u2, A, v) ≈ collect(A) * collect(v)
            @test u1 ≈ Gaius.mul(A, v)
            @test u2 ≈ Gaius.mul(A, v)
            @test u1 ≈ Gaius.mul_serial(A, v)
            @test u2 ≈ Gaius.mul_serial(A, v)
        end
    end
end

@time @testset "Float32 Matrix-Vector $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(Float32, n)
            u2 = zeros(Float32, n)
            A = rand(Float32, n, m)
            v = rand(Float32, m)
            @test Gaius.mul!(u1, A, v) ≈ A * v
            @test Gaius.mul_serial!(u2, A, v) ≈ A * v
            @test u1 ≈ Gaius.mul(A, v)
            @test u2 ≈ Gaius.mul(A, v)
            @test u1 ≈ Gaius.mul_serial(A, v)
            @test u2 ≈ Gaius.mul_serial(A, v)
        end
    end
end

@time @testset "Int64 Matrix-Vector $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(Int, n)
            u2 = zeros(Int, n)
            A = rand(-20:20, n, m)
            v = rand(-20:20, m)
            @test Gaius.mul!(u1, A, v) == A * v
            @test Gaius.mul_serial!(u2, A, v) == A * v
            @test u1 == Gaius.mul(A, v)
            @test u2 == Gaius.mul(A, v)
            @test u1 == Gaius.mul_serial(A, v)
            @test u2 == Gaius.mul_serial(A, v)
        end
    end
end

@time @testset "ComplexInt32 Matrix-Vector $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(Complex{Int32}, n) |> StructArray
            u2 = zeros(Complex{Int32}, n) |> StructArray
            A = StructArray{Complex{Int32}}((rand(Int32.(-10:10), n, m), rand(Int32.(-10:10), n, m)))
            v = StructArray{Complex{Int32}}((rand(Int32.(-10:10),    m), rand(Int32.(-10:10),    m)))
            @test Gaius.mul!(u1, A, v) == collect(A) * collect(v)
            @test Gaius.mul_serial!(u2, A, v) == collect(A) * collect(v)
            @test u1 == Gaius.mul(A, v)
            @test u2 == Gaius.mul(A, v)
            @test u1 == Gaius.mul_serial(A, v)
            @test u2 == Gaius.mul_serial(A, v)
        end
    end
end

@time @testset "Float64 CoVector-Matrix $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(m)
            u2 = zeros(m)
            A = rand(n, m)
            v = rand(n)
            @test Gaius.mul!(u1', v', A) ≈ v' * A
            @test Gaius.mul_serial!(u2', v', A) ≈ v' * A
            @test u1' ≈ Gaius.mul(v', A)
            @test u2' ≈ Gaius.mul(v', A)
            @test u1' ≈ Gaius.mul_serial(v', A)
            @test u2' ≈ Gaius.mul_serial(v', A)

            @test Gaius.mul!(transpose(u1), transpose(v), A) ≈ transpose(v) * A
            @test Gaius.mul_serial!(transpose(u2), transpose(v), A) ≈ transpose(v) * A
            @test transpose(u1) ≈ Gaius.mul(transpose(v), A)
            @test transpose(u2) ≈ Gaius.mul(transpose(v), A)
            @test transpose(u1) ≈ Gaius.mul_serial(transpose(v), A)
            @test transpose(u2) ≈ Gaius.mul_serial(transpose(v), A)
        end
    end
end

@time @testset "ComplexF32 CoVector-Matrix $(testset_name_suffix)" begin
    for n ∈ n_values_2
        for m ∈ m_values_2
            u1 = zeros(Complex{Float32}, m)   |> StructArray
            u2 = zeros(Complex{Float32}, m)   |> StructArray
            A = rand(Complex{Float32}, n, m) |> StructArray
            v = rand(Complex{Float32}, n)    |> StructArray
            @test Gaius.mul!(u1', v', A) ≈ v' * A
            @test Gaius.mul_serial!(u2', v', A) ≈ v' * A
            @test u1' ≈ Gaius.mul(v', A)
            @test u2' ≈ Gaius.mul(v', A)
            @test u1' ≈ Gaius.mul_serial(v', A)
            @test u2' ≈ Gaius.mul_serial(v', A)

            @test Gaius.mul!(transpose(u1), transpose(v), A) ≈ transpose(v) * A
            @test Gaius.mul_serial!(transpose(u2), transpose(v), A) ≈ transpose(v) * A
            @test transpose(u1) ≈ Gaius.mul(transpose(v), A)
            @test transpose(u2) ≈ Gaius.mul(transpose(v), A)
            @test transpose(u1) ≈ Gaius.mul_serial(transpose(v), A)
            @test transpose(u2) ≈ Gaius.mul_serial(transpose(v), A)
        end
    end
end

@time @testset "ComplexFloat64 Matrix-Matrix $(testset_name_suffix)" begin
    for sz ∈ sz_values
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C1 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C2 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            A  = StructArray{ComplexF64}((randn(n, k), randn(n, k)))
            B  = StructArray{ComplexF64}((randn(k, m), randn(k, m)))

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 ≈ C1
            @test C0 ≈ C2
            @test C1 ≈ C2

            @test Gaius.mul_serial(A, B) ≈ C0
        end
    end
    for sz ∈ sz_values
        n, k, m = shuffle([sz + rand(-5:5), sz + rand(-5:5), 10])
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C1 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C2 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            A  = StructArray{ComplexF64}((randn(n, k), randn(n, k)))
            B  = StructArray{ComplexF64}((randn(k, m), randn(k, m)))

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 ≈ C1
            @test C0 ≈ C2
            @test C1 ≈ C2

            @test Gaius.mul_serial(A, B) ≈ C0

            @test Gaius.mul(A, B) ≈ C0
        end
    end
end

@time @testset "Float64 Matrix-Matrix $(testset_name_suffix)" begin
    for sz ∈ sz_values
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = zeros(n, m)
            C1 = zeros(n, m)
            C2 = zeros(n, m)
            A  = randn(n, k)
            B  = randn(k, m)
            At = copy(A')
            Bt = copy(B')

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 ≈ C1
            @test C0 ≈ C2
            @test C1 ≈ C2

            @test Gaius.mul_serial(A, B) ≈ C0

            fill!(C1, NaN); @test Gaius.mul!(C1, At', B)   ≈ C0
            fill!(C1, NaN); @test Gaius.mul!(C1, A,   Bt') ≈ C0
            fill!(C1, NaN); @test Gaius.mul!(C1, At', Bt') ≈ C0

            fill!(C2, NaN); @test Gaius.mul_serial!(C2, At', B)   ≈ C0
            fill!(C2, NaN); @test Gaius.mul_serial!(C2, A,   Bt') ≈ C0
            fill!(C2, NaN); @test Gaius.mul_serial!(C2, At', Bt') ≈ C0
        end
    end
    for sz ∈ sz_values
        n, k, m = shuffle([sz + rand(-5:5), sz + rand(-5:5), 10])
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = zeros(n, m)
            C1 = zeros(n, m)
            C2 = zeros(n, m)
            A  = randn(n, k)
            B  = randn(k, m)
            At = copy(A')
            Bt = copy(B')

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 ≈ C1
            @test C0 ≈ C2
            @test C1 ≈ C2

            @test Gaius.mul_serial(A, B) ≈ C0

            fill!(C1, NaN); @test Gaius.mul!(C1, At', B)   ≈ C0
            fill!(C1, NaN); @test Gaius.mul!(C1, A,   Bt') ≈ C0
            fill!(C1, NaN); @test Gaius.mul!(C1, At', Bt') ≈ C0

            fill!(C2, NaN); @test Gaius.mul_serial!(C2, At', B)   ≈ C0
            fill!(C2, NaN); @test Gaius.mul_serial!(C2, A,   Bt') ≈ C0
            fill!(C2, NaN); @test Gaius.mul_serial!(C2, At', Bt') ≈ C0
        end
    end
end

@time @testset "Float32 Matrix-Matrix $(testset_name_suffix)" begin
    for sz ∈ sz_values
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = zeros(Float32, n, m)
            C1 = zeros(Float32, n, m)
            C2 = zeros(Float32, n, m)
            A  = randn(Float32, n, k)
            B  = randn(Float32, k, m)
            At = copy(A')
            Bt = copy(B')

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 ≈ C1
            @test C0 ≈ C2
            @test C1 ≈ C2

            @test Gaius.mul_serial(A, B) ≈ C0

            fill!(C1, NaN32); @test Gaius.mul!(C1, At', B)   ≈ C0
            fill!(C1, NaN32); @test Gaius.mul!(C1, A,   Bt') ≈ C0
            fill!(C1, NaN32); @test Gaius.mul!(C1, At', Bt') ≈ C0

            fill!(C2, NaN32); @test Gaius.mul_serial!(C2, At', B)   ≈ C0
            fill!(C2, NaN32); @test Gaius.mul_serial!(C2, A,   Bt') ≈ C0
            fill!(C2, NaN32); @test Gaius.mul_serial!(C2, At', Bt') ≈ C0
        end
    end
end

@time @testset "Int64 Matrix-Matrix $(testset_name_suffix)" begin
    for sz ∈ sz_values
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = zeros(Int64, n, m)
            C1 = zeros(Int64, n, m)
            C2 = zeros(Int64, n, m)
            A  = rand(Int64.(-100:100), n, k)
            B  = rand(Int64.(-100:100), k, m)
            At = copy(A')
            Bt = copy(B')

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 == C1
            @test C0 == C2
            @test C1 == C2

            @test Gaius.mul_serial(A, B) == C0

            fill!(C1, typemax(Int64)); @test Gaius.mul!(C1, At', B)   == C0
            fill!(C1, typemax(Int64)); @test Gaius.mul!(C1, A,   Bt') == C0
            fill!(C1, typemax(Int64)); @test Gaius.mul!(C1, At', Bt') == C0

            fill!(C2, typemax(Int64)); @test Gaius.mul_serial!(C2, At', B)   == C0
            fill!(C2, typemax(Int64)); @test Gaius.mul_serial!(C2, A,   Bt') == C0
            fill!(C2, typemax(Int64)); @test Gaius.mul_serial!(C2, At', Bt') == C0

            @test Gaius.mul(A,   B) == C0
            @test Gaius.mul(At', B) == C0
            @test Gaius.mul(A,   Bt') == C0
            @test Gaius.mul(At', Bt') == C0
        end
    end
end

@time @testset "Int32 Matrix-Matrix $(testset_name_suffix)" begin
    for sz ∈ sz_values
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m) $(testset_name_suffix)" begin
            C0 = zeros(Int32, n, m)
            C1 = zeros(Int32, n, m)
            C2 = zeros(Int32, n, m)
            A  = rand(Int32.(-100:100), n, k)
            B  = rand(Int32.(-100:100), k, m)
            At = copy(A')
            Bt = copy(B')

            LinearAlgebra.mul!(C0, A, B)
            Gaius.mul!(C1, A, B)
            Gaius.mul_serial!(C2, A, B)

            @test C0 == C1
            @test C0 == C2
            @test C1 == C2

            @test Gaius.mul_serial(A, B) == C0

            fill!(C1, typemax(Int32)); @test Gaius.mul!(C1, At', B)   == C0
            fill!(C1, typemax(Int32)); @test Gaius.mul!(C1, A,   Bt') == C0
            fill!(C1, typemax(Int32)); @test Gaius.mul!(C1, At', Bt') == C0

            fill!(C2, typemax(Int32)); @test Gaius.mul_serial!(C2, At', B)   == C0
            fill!(C2, typemax(Int32)); @test Gaius.mul_serial!(C2, A,   Bt') == C0
            fill!(C2, typemax(Int32)); @test Gaius.mul_serial!(C2, At', Bt') == C0
        end
    end
end


text:
"""
    house_fits_pattern(i::Int, p::Int, d::StudentHousingData)

Returns a Boolean denoting whether someone with preference pattern `p` would
be willing to live in the `i`^th house.
"""
function house_fits_pattern(i::Int, p::Int, d::StudentHousingData)
    h = d.houses[i]
    for c in d.patterns_allow[p]
        if house_fits_characteristic(h, d.all_characteristics[c], d.market_data)
            return true
        end
    end
    false
end

function house_allowedby(i::Int, d::StudentHousingData)
    ret = Int[]
    for p = 1:length(d.patterns_allow)
      if house_fits_pattern(i, p, d)
          push!(ret, p)
      end
    end
    ret
end

# ==============================================================================
# One stage

function onestagemodel(d::StudentHousingData)

    houses = d.houses
    nhouses = length(d.houses)
    # One stage and deterministic
    demand = d.demands[:, 1, 1]

    # The number of only legal patterns
    npatterns = length(d.patterns_allow)

    m = Model(solver = GurobiSolver(OutputFlag=0))
    # m = Model(solver = CplexSolver(CPX_PARAM_SCRIND = 0, CPX_PARAM_MIPDISPLAY=0))
    @variables(m, begin
        shortage[1:npatterns] >= 0
        investment[1:nhouses], Bin
        assignment[i=1:nhouses, p=house_allowedby(i,d)] >= 0, Int
    end)

    @constraints(m, begin
        [p = 1:npatterns], shortage[p] == demand[p] - sum(assignment[i, p] for i = 1:nhouses if house_fits_pattern(i, p, d))
        [i = 1:nhouses], sum(beds_needed(p, d) * assignment[i, p] for p = 1:npatterns if house_fits_pattern(i, p, d)) <= investment[i] * beds_avail(i, houses)
        sum(investment[i] * maintenance(i, houses) for i = 1:nhouses) <= d.budget
    end)

    @objective(m, Min, sum(shortage))
    m
end

# ==============================================================================
# Multistage
function multistagemodel(d::StudentHousingData)

    houses = d.houses
    nhouses = length(d.houses)
    demands = d.demands
    upperbound = sum(sum(sum(demands)))

    # The number of only legal patterns
    npatterns = length(d.patterns_allow)

    m = SDDPModel(stages = d.nstages,
            objective_bound = 0.0,
            sense = :Min,
            solver = GurobiSolver(OutputFlag=0)) do sp, stage

        @binarystate(sp, 0 <= investment[1:nhouses] <= 1,
                investment0 == 0, Bin)

        @variables(sp, begin
            shortage[1:npatterns] >= 0
            assignment[1:nhouses, 1:npatterns] >= 0, Int
        end)

        if stage == 1
            @constraint(sp, [p = 1:npatterns], shortage[p] + sum(house_fits_pattern(i, p, d) * assignment[i, p] for i = 1:nhouses) == demands[p, 1, 1])
        else
            for p = 1:npatterns
                @rhsnoise(sp, D = demands[p, :, stage],
                    shortage[p] + sum(house_fits_pattern(i, p, d) * assignment[i, p] for i = 1:nhouses) == D)
            end
        end

        @constraints(sp, begin
            # Can't uninvest
            investment0 .<= investment
            # Only legal assignments
            [i = 1:nhouses], sum(house_fits_pattern(i, p, d) * assignment[i, p] * beds_needed(p, d) for  p = 1:npatterns) <= investment[i] * beds_avail(i, houses)
            # Budget constraint
            sum(investment[i] * maintenance(i, houses) for i = 1:nhouses) <= d.budget
        end)

        # Minimize all unmet demand
        @stageobjective(sp, sum(shortage))

        setSDDiPsolver!(sp, method=KelleyMethod(), pattern = Pattern(benders = 0, lagrangian=1))
    end
    m

end


text:
# Domain

export Domain
"""
Bounding box for a domain. The domain does not need to be recangular.
"""
struct Domain{S,D}
    min::SVector{D,S}
    max::SVector{D,S}
end

Base.ndims(::Domain{S,D}) where {S,D} = D

Base.eltype(::Domain{S}) where {S} = S

Base.min(dom::Domain) = dom.min
Base.max(dom::Domain) = dom.max


text:
const CONFIG = MOIT.Config(Int)

const OPTIMIZER = JaCoP.Optimizer()
const BRIDGED_OPTIMIZER = MOI.Bridges.full_bridge_optimizer(OPTIMIZER, Int32)

COIT.runtests(
    # BRIDGED_OPTIMIZER,
    OPTIMIZER,
    CONFIG,
)


text:
"""
$(SIGNATURES)
    Determine if solving a minimum time problem
"""
function is_min_time(solver::Solver)
    if solver.dt == 0 && solver.N > 0
        return true
    end
    return false
end

"""
$(SIGNATURES)
    Determine an initial dt for minimum time solver
"""
function get_initial_dt(solver::Solver)
    if is_min_time(solver)
        if solver.opts.minimum_time_dt_estimate > 0.0
            dt = solver.opts.minimum_time_dt_estimate
        elseif solver.opts.minimum_time_tf_estimate > 0.0
            dt = solver.opts.minimum_time_tf_estimate / (solver.N - 1)
            if dt > solver.opts.max_dt
                dt = solver.opts.max_dt
                @warn "Specified minimum_time_tf_estimate is greater than max_dt. Capping at max_dt"
            end
        else
            dt  = solver.opts.max_dt / 2
        end
    else
        dt = solver.dt
    end
    return dt
end

function total_time(solver::Solver, results::SolverVectorResults)
    if is_min_time(solver)
        m̄,mm = get_num_controls(solver)
        T = sum([u[m̄]^2 for u in results.U[1:solver.N-1]])
    else
        T = solver.dt*(solver.N-1)
    end
    return T::Float64
end

function total_time(solver::Solver, results::DircolVars)
    N = solver.N
    if is_min_time(solver)
        m̄, = get_num_controls(solver)
        T = sum(results.U[m̄,1:N-1])
    else
        T = solver.dt*(solver.N-1)
    end
end

function get_time(solver::Solver)
    range(0,stop=solver.obj.tf,length=solver.N)
end

function get_time(solver::Solver, results)
    if solver.state.minimum_time
        cumsum([u[end]^2 for u in results.U])
    else
        get_time(solver)
    end
end


text:
using Test, MLToolkit

@testset "Activations" begin
    dims = (10, 5)

    @testset "softplus" begin
        for _ = 1:10
            x = randn(dims...)
            y = softplus.(x)
            @test all(y .> 0)
        end
    end

    @testset "leaky_relu" begin
        for _ = 1:10
            x = randn(dims...)
            y = leaky_relu.(x; alpha=0)
            @test all(y .>= 0)
        end
    end
end


text:
# The following block allows reading Union types written prior to v0.2
const LEGACY_H5TYPE_UNION = VariableLengthDatatype(H5TYPE_DATATYPE)

function jlconvert(::ReadRepresentation{Union, Vlen{DataTypeODR()}}, f::JLDFile,
                   ptr::Ptr, header_offset::RelOffset)
    v = Union{jlconvert(ReadRepresentation{DataType,Vlen{DataTypeODR()}}(), f, ptr, NULL_REFERENCE)...}
    track_weakref!(f, header_offset, v)
    v
end

constructrr(::JLDFile, ::Type{T}, dt::VariableLengthDatatype, ::Vector{ReadAttribute}) where {T<:Union} =
    dt == LEGACY_H5TYPE_UNION ? (ReadRepresentation{Union,Vlen{DataTypeODR()}}(), true) :
                         throw(UnsupportedFeatureException())


text:
using TestImages: testimage
using CSV, FileIO, Images, Colors, Augmentor, Distributions

"""
    spaghetti_fun(x, y, a, b, τ, len)

Gives the intensity for a single "spaghetti" at (`x`, `y`) with the params `a`, `b`, `τ` and `len`.

The intensity is equal to:
```math
I(\\mathrm{x,y})=\\left(1-\\frac{f(\\mathrm{x, y})}{τ^{2}}\\right)\\left(1+\\mathrm{e}
^{\\left(\\left(\\mathrm{x}^2 +\\mathrm{y}^2 \\right) -\\mathrm{len}^{2}\\right) /\\left(50 τ^{2}
\\right)}\\right)^{-1} \\mathrm{e}^{-f(\\mathrm{x, y}) /\\left(2 τ^{2}\\right)}
```
where
```math
f(\\mathrm{x, y})=\\left|y-a x^{3}-b x\\right|^{2} /\\left(1+\\left|3 a x^{2}-b\\right|^{2}\\right)
```

See also: [`gen_spaghetti`](@ref), [`gen_one_spaghetti`](@ref), [`gen_spaghetti`](@ref)
"""
function spaghetti_fun(x, y, a, b, τ, len)
    spag(x, y, a, b) = (abs(y - a*x^3 - b*x)^2)/(1 + abs(3a*x^2 - b)^2)
    args = (x, y, a, b)
    return (1 - spag(args...)/τ^2) * (1 + exp(((x^2 + y^2) - len)/(50*τ^2)))^(-1) * exp((-spag(args...))/(2*τ^2))
end

# TODO edit wording
"""
    gen_one_spaghetti(img_size, τ, len, spread)

Generate a single spaghetti image with the params `τ` (thickness), `len` (length) and `spread`.
The coordinates `x` and `y` in the equation are scaled from `(-1, 1)` and the `spread` gives the length of interval
in which the `(-1, 1)` is randomly located. The genereated spaghetti is also randomly rotated.

(Example for spread 10, the interval for `x` and `y` is randomly anywhere between `(-9:1)` and `(-1:9)`)

The intensity is equal to:
```math
I(\\mathrm{x,y})=\\left(1-\\frac{f(\\mathrm{x, y})}{τ^{2}}\\right)\\left(1+\\mathrm{e}
^{\\left(\\left(\\mathrm{x}^2 +\\mathrm{y}^2 \\right) -\\mathrm{len}^{2}\\right) /\\left(50 τ^{2}
\\right)}\\right)^{-1} \\mathrm{e}^{-f(\\mathrm{x, y}) /\\left(2 τ^{2}\\right)}
```
where
```math
f(\\mathrm{x, y})=\\left|y-a x^{3}-b x\\right|^{2} /\\left(1+\\left|3 a x^{2}-b\\right|^{2}\\right)
```
`a` and `b` are chosen randomly from a normal distribution.

See also: [`gen_spaghetti`](@ref), [`spaghetti_fun`](@ref)
"""
function gen_one_spaghetti(img_size, τ, len, spread)
    (a, b) = randn(2)

    c = rand()*(spread-2)+1
    d = rand()*(spread-2)+1
    enlarged_img_size = ceil.(Int, img_size .* sqrt(2)) # so that the rotation doesnt make have to extrapolate

    X = ones(enlarged_img_size[1]) * collect(range(c-spread,c,length=enlarged_img_size[2]))'
    Y = collect(range(d-spread,d,length=enlarged_img_size[1])) * ones(enlarged_img_size[2])'

    rand_rot = Rotate(0:360)
    crop = CropSize(img_size...)
    pl = rand_rot |>
         # ShearX(-5:5) * ShearY(-5:5) |>
         crop

    X_rot, Y_rot = augment((X, Y), pl);

    img = map((x, y) -> spaghetti_fun(x, y, a, b, τ, len), X_rot, Y_rot)
end

"""
    gen_spaghetti(img_size = (200, 200),
                  τ = 0.2,
                  len = 30;
                  count=25,
                  spread=10)

Generate an image with `count` spaghetti images with the params `τ` (thickness), `len` (length) and `spread`.
The coordinates `x` and `y` in the equation are scaled from `(-1, 1)` and the `spread` gives the length of interval
in which the `(-1, 1)` is randomly located. The genereated spaghetti is also randomly rotated.

(Example for spread 10, the interval for `x` and `y` is randomly anywhere between `(-9:1)` and `(-1:9)`)

# Example
```@example
# feature-full spaghetti image
img = gen_spaghetti((256, 256), 0.15, 70, spread=20)
# feature-less spaghetti image
img = gen_spaghetti((256, 256), 1, 50, spread=12)
```

Note: For the exact intensity function see [`gen_one_spaghetti`](@ref)

See also:  [`spaghetti_fun`](@ref)
"""
function gen_spaghetti(img_size = (256, 256),
                       τ = 0.2,
                       len = 30;
                       count=25,
                       spread=10)

    img = zeros(img_size)
    for _ in 1:count
        adding_img = gen_one_spaghetti(img_size, τ, len, spread)
        img = img + adding_img
    end
    img = LAP_julia.normalize_to_zero_one(img)
    return img
end

rand_uniform(a, b) = rand() * (b-a) + a

"""
    rand_H(flow_size)

Generate a semi-random homography transfomration matrix. See the implementation for used constants and distributions.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref), [`gen_homo_flow`](@ref)
"""
function rand_H(flow_size)
    a = flow_size[1] * 0.4
    H = [rand(Normal(1, 0.4)) rand(Normal(0, 0.2)) rand_uniform(-a, a);
         rand(Normal(0, 0.2)) rand(Normal(1, 0.4)) rand_uniform(-a, a);
         rand(Normal(0, 0.0002)) rand(Normal(0, 0.0002)) 1.]
    @info H
    return H
end

"""
    gen_homo_flow(flow_size, max_magnitude=10)

Generate a smoothly varying locally constant random flow of size `img_size` and with a maximal
displacement of `max_magnitude` using a randomly generated homography matrix.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref), [`rand_H`](@ref)
"""
function gen_homo_flow(flow_size, max_magnitude=10)
    H = rand_H(flow_size)
    flow = LAP_julia.make_flow_from_H(H, flow_size)
    max_len = maximum(LAP_julia.vec_len.(flow))
    flow = flow .* (max_magnitude/max_len)
    @assert median(LAP_julia.vec_len.(flow)) > (0.1*max_magnitude) H
    return flow
end

"""
    gen_anhir(base_path = "/Users/MrTrololord/Documents/anhir/";
              mutate=true,
              diag_pixels=500)

Generate a random ANHIR image pair `target` and `source` images (in this order) from the ANHIR dataset at `base_path`.

See also: [`load_anhir_image_pair`](@ref)
"""
function gen_anhir(base_path = "/Users/MrTrololord/Documents/anhir/";
                   mutate=true,
                   diag_pixels=500)

    loc_table = CSV.read(joinpath(base_path, "location_table.csv"))
    train_rows = loc_table[loc_table[:status] .== "training", :]

    random_row = train_rows[rand(1:size(train_rows, 1)), :]
    target_path = random_row[Symbol("Target image")]
    source_path = random_row[Symbol("Source image")]

    target, source = load_anhir_image_pair(target_path, source_path, base_path=joinpath(base_path, "dataset"))

    if any(size(target) .!= size(source))
        println("sizes: ", size(target), ", ", size(source))
    end
    target, source = pad_images(target, source)

    if mutate
        target = resize_to_diag_size(target, diag_pixels)
        source = resize_to_diag_size(source, diag_pixels)
    end

    return target, source
end

"""
    resize_to_diag_size(img, desired_diag_size)

Resize `img` so that it has `desired_diag_size` pixels on its diagonal.

See also: [`gen_anhir`](@ref), [`load_anhir_image_pair`](@ref)
"""
function resize_to_diag_size(img, desired_diag_size)
    actual_diag_pixels = sqrt(sum((size(img).^2)))
    resize_ratio = desired_diag_size/actual_diag_pixels
    img = imresize(img, ratio = resize_ratio)
    return img, resize_ratio
end

"""
    load_anhir_image_pair(target_path,
                          source_path;
                          base_path = "/Users/MrTrololord/Documents/anhir/dataset")

Load an image pair in the locations `target_path` and `source_path` from the location `base_path`.

See also: [`gen_anhir`](@ref)
"""
function load_anhir_image_pair(target_path,
                               source_path;
                               base_path = "/Users/MrTrololord/Documents/anhir/dataset")

    println(joinpath(base_path, target_path))
    target = Float64.(Gray.(load(joinpath(base_path, target_path))))
    source = Float64.(Gray.(load(joinpath(base_path, source_path))))

    return target, source
end

function load_image_gray(img_path)::Image
    img = Float64.(Gray.(load(img_path)))
end

"""
    gen_lena()

Get a `256x256` grayscale "lena" image.
"""
function gen_lena()
    img = testimage("lena_gray")
    img = Float64.(img)
    return img
end

"""
    gen_quad_flow(img_size, max_magnitude=10)

Generate a smoothly varying locally constant random flow of size `img_size` and with a maximal
displacement of `max_magnitude` using a quadratic function:

```math
f(z) = a + b*z + c*z^2,
```
where:
```math
z = x + y*i, x, y ∈ (0,1)
```

The constants `a`, `b` and `c` are random numbers from the normal distribution with mean 0 and standard deviation 1.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref)
"""
function gen_quad_flow(img_size, max_magnitude=10)

    a, b, c = randn(), randn(), randn()
    f(z) = a + b*z + c*z^2

    X = ones(img_size[1]) * collect(range(0,1,length=img_size[2]))'
    Y = collect(range(0,1,length=img_size[1])) * ones(img_size[2])'
    A = Y .* im + X

    B = f.(A)
    max_len = maximum(vec_len.(B))
    B = B .* (max_magnitude/max_len)

    return B
end

"""
    gen_uniform_flow(flow_size=(200, 200), vector=1 + 1im, max_magnitude=vec_len(vector))

Generate a uniform flow of size `flow_size`, where every displacement vector is `vector`, scaled by `max_magnitude`.

Note: `max_magnitude` = `vec_len(vector)` by default,
so if nothing is enetered as `max_magnitude` the flow is made up of `vector` and is not scaled.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref)
"""
function gen_uniform_flow(flow_size=(200, 200),
                          vector=1 + 1im,
                          max_magnitude=vec_len(vector))

    flow = ones(flow_size) .* vector

    max_len = vec_len(vector)
    flow = flow .* (max_magnitude/max_len)
    return flow
end

"""
    gen_tiled_flow(flow_size::Tuple{T, T}=(200, 200),
    max_magnitude::Real=20, tile_size=nothing; filter_amp=nothing)::Flow where {T <: Integer}

Generate a smoothly varying random flow. The flow parameters are set by the function arguments.

It works by generating a tiled flow where each tile has a random uniform flow. Then it smooths these tiles by filtering with a gaussian.

# Arguments
- `flow_size::Tuple{T, T}=(200, 200)`: dimensions of the flow.
- `max_magnitude::Real=20`: maximum allowed amplitude of the displacement.
- `tile_size=nothing`: size of random uniform flow tiles that make up the start matrix. (The larger the slower the flow.) Note: If set to `flow_size` or larger it will generate a uniform pixel shift in a random direction.
- `filter_amp=nothing`: size of the gaussian filter which is used to smooth the random start matrix.

See also: [`showflow`](@ref), [`Flow`](@ref), [`gen_chess`](@ref), [`gen_quad_flow`](@ref)
"""
function gen_tiled_flow(flow_size::Tuple{T, T}=(200, 200),
                        max_magnitude::Real=10,
                        tile_size=nothing,
                        filter_amp=nothing)::Flow where {T <: Integer}

    # set default values
    if tile_size == nothing
        tile_size = ceil(Int64, flow_size[1]/6)
    end
    if filter_amp == nothing
        filter_amp = ceil(Int64, tile_size/2)
        if isodd(filter_amp)
            filter_amp += 1
        end
    end

    tile_count = ceil.(Int64, flow_size ./ tile_size)

    uv_base = Array{Float64}(undef, tile_count..., 2)
    uv_base[:,:,1] .= -max_magnitude .+ (2 .* max_magnitude .* rand(tile_count...))
    uv_base[:,:,2] .= -max_magnitude .+ (2 .* max_magnitude .* rand(tile_count...))

    uv_flow = Array{Float64}(undef, tile_size*tile_count[1], tile_size*tile_count[2], 2)
    uv_flow[:,:,1] .= repeat(uv_base[:,:,1], inner=(tile_size, tile_size))
    uv_flow[:,:,2] .= repeat(uv_base[:,:,2], inner=(tile_size, tile_size))

    # cut off what we dont want
    #rand_flow = zeros(Complex{Float64}, flow_size)
    rand_flow = uv_flow[1:flow_size[1], 1:flow_size[2], 1] .+ (im * uv_flow[1:flow_size[1], 1:flow_size[2], 2]);

    # blur to make it continuous
    rand_flow = smooth_with_gaussian!(rand_flow, [filter_amp, filter_amp])

    # set maximum magnitude to actually be the max_magnitude
    max_len = maximum(vec_len.(rand_flow))
    rand_flow = rand_flow .* (max_magnitude/max_len)

    return rand_flow
end

"""
    gen_chess(tile_size::Integer=50, board_size::Integer=4)

Create a chessboard image with `board_size` tiles in each dimension, where each tile is `tile_size` pixels in each dimension.
Note that `board_size` has to be even.

See also: [`imgshow`](@ref), [`gen_tiled_flow`](@ref)
"""
function gen_chess(tile_size::Integer=50, board_size::Integer=4)
    mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
                  ones(tile_size, tile_size) zeros(tile_size, tile_size)]

    chessboard = repeat(mini_board, outer=(convert(Integer, (board_size/2)), convert(Integer, (board_size/2))))
    return chessboard
end

"""
    gen_init(type::Symbol=:lena; flow_args=[])

Create the usual testing data; `img`, `imgw`, `flow` using the given parameters.

# Arguments
- `img_type::Symbol=:lena`: what base image is used. [Options: `:lena`, `:chess`]
- `flow_type::Symbol=:quad`: what flow generation function is used. [Options: `:tiled`, `:quad`, `:uniform`, `:homo`]

# Keyword Arguments
- `flow_args=[]`: arguments passed to the flow generation function besides the flow size.
- `img_args=[]`: arguments passed to the img generation function if `:chess` is chosen.
- `flipped::Bool=false`: the image will be flipped along the horizontal axis, so that the origin is at the bottom.

# Example
```@example
# chess image, warped chess image, flow with maximal displacement 20 generated by the `gen_quad_flow` function
img, imgw, flow = gen_init(:chess, :quad, flow_args=[20])

See also: [`gen_anhir`](@ref), [`gen_lena`](@ref), [`gen_spaghetti`](@ref), [`gen_tiled_flow`](@ref), [`gen_quad_flow`](@ref), [`gen_uniform_flow`](@ref)
```
"""
function gen_init(img_type::Symbol=:lena, flow_type::Symbol=:quad; flow_args=[], img_args=[], flipped::Bool=false)
    if img_type == :lena
        img = gen_lena()
    elseif img_type == :chess
        img = gen_chess((Int64.(img_args))...)
    elseif img_type == :spaghetti
        img = gen_spaghetti(img_args...)
    end

    if flipped
        img = reverse(img, dims=1)
    end

    if flow_type == :quad
        flow = gen_quad_flow(size(img), flow_args...)
    elseif flow_type == :tiled
        flow = gen_tiled_flow(size(img), flow_args...)
    elseif flow_type == :uniform
        flow = gen_uniform_flow(size(img), flow_args...)
    elseif flow_type == :homo
        flow = gen_homo_flow(size(img), flow_args...)
    end

    imgw = warp_img(img, -real(flow), -imag(flow), border_strat=:zeros)

    return img, imgw, flow
end


text:
function find_direction(s::AbstractString)
  isnothing(findfirst("<", s)) ? :rightarrow : :leftarrow
end

function replace_dag_term(s::AbstractString, from::Int, to::Int)
  term_vars = replace(s[from+1:to-1], " " => ",")
  "$(s[1:from-1])[$(term_vars)]$(s[to+1:end])"
end

function handle_term(s::AbstractString)
  ind = findfirst("[", s)
  if isnothing(ind)
    # Just a single Synbol
    return Symbol(s)
  else
    from = ind.start
    ind = findfirst("]", s)
    if isnothing(ind)
      @error "$s not well formed."
    end
    to = ind.start
    t = split(s[from+1:to-1], ",")
    return Symbol.(t)
  end
end

function split_and_insert!(f, t, s)
  new_t = [t[1]]
  for e in t[2:end-1]
    push!(new_t, e)
    push!(new_t, e)
  end
  push!(new_t, t[end])
  indx = sort(vcat(findall(">", s), findall("<", s)))
  @assert length(new_t) == 2length(indx)
  j = 1
  for i in 1:2:length(new_t)
    push!(f, "$(new_t[i]) $(s[indx[j].start]) $(new_t[i+1])")
    j += 1
  end
end

function from_dagitty(str::AbstractString; tolowercase=false)
  str = replace(str, "<-" => "<")
  str = replace(str, "->" => ">")
  startp = findall("{", str)
  endp = findall("}", str)
  f = split(str[startp[1].start+1:endp[end].stop-1], ";")

  dct = OrderedDict{SymbolList, SymbolList}()

  while length(f) > 0
    #display(f)
    s = pop!(f)
    dir = find_direction(s)
    #println(dir)

    ind = findfirst("{", s)
    while !isnothing(ind)
      from = ind.start
      to = findfirst("}", s)
      to = to.start
      s = replace_dag_term(s, from, to)
      ind = findfirst("{", s)
    end

    t = String.(replace(strip(s), " " => ""))
    t = split(t, x -> x in "><", keepempty=false)
    if length(t) > 2
      split_and_insert!(f, t, s)
      continue
    end

    term1 = handle_term(t[1])
    term2 = handle_term(t[2])

    if dir == :leftarrow
      #println([term1, dir, term2, term1 in keys((dct))])
      if term1 in keys(dct)
        #println(dct[term1])
        if typeof(dct[term1]) == Symbol
          #println(typeof(dct[term1]) == Symbol)
          dct[term1] = push!([dct[term1]], term2)
        else
          push!(dct[term1], term2)
        end
      else
        dct[term1] = term2
      end
    else
      if term2 in keys(dct)
        if typeof(dct[term2]) == Symbol
          dct[term2] = push!([dct[term2]], term1)
        else
          push!(dct[term2], term1)
        end
      else
        dct[term2] = term1
      end
    end
    #println(dct)
  end

  dct
end

function to_dagitty(d::OrderedDict; touppercase=false)
  str = "dag { "
  for (ind, key) in enumerate(keys(d))
    str = ind > 1 ? "$(str); " : "$(str)"
    rhs = d[key]
    if typeof(rhs) == Symbol
      rhs_str = String(rhs)
    elseif length(rhs) == 0
      # Not used yet
    elseif length(rhs) == 1
      rhs_str = String(rhs[1])
    elseif length(rhs) > 1
      rhs_str = "{ $(String(rhs[1]))"
      for s in rhs[2:end]
        rhs_str = "$(rhs_str) $(String(s))"
      end
      rhs_str = "$(rhs_str) }"
    end
    # Split multiple lhs symbols (if present)
    if typeof(key) == Symbol
      str = "$(str)$(String(key)) <- $rhs_str"
    elseif length(key) == 1
      str = "$(str)$(String(key[1])) <- $rhs_str"
    elseif length(key) > 1
      str = "$(str)$(String(key[1])) <- $(rhs_str); "
      for s in key[2:end-1]
        str = "$(str)$(String(s)) <- $(rhs_str)"
      end
      str = "$(str)$(String(key[end])) <- $(rhs_str)"
    end
  end
  str = "$(str) }"
  touppercase ? uppercase(str) : str
end

function to_dagitty(d::DAG; touppercase=false)
  to_dagitty(d.d; touppercase=touppercase)
end

export
  from_dagitty,
  to_dagitty


text:
## PARTE 1
cd(@__DIR__)
rx = r"([0-9]+)x([0-9]+)x([0-9]+)"
input = map(eachline("input.txt")) do line
    m = match(rx, line)
    sort(parse.(Int, m.captures))
end;

##
total_area(x, y, z) = 3x*y + 2y*z + 2x*z
sum(x->total_area(x...), input)

## PARTE 2
total_ribbon(xs) = 2*sum(xs[1:2]) + prod(xs)
sum(total_ribbon, input)