module BetaTest

using Test
using ReactiveMP
using Distributions
using Random

@testset "Beta" begin

    # Beta comes from Distributions.jl and most of the things should be covered there
    # Here we test some extra ReactiveMP.jl specific functionality

    @testset "vague" begin
        d = vague(Beta)

        @test typeof(d) <: Beta
        @test mean(d) === 0.5
        @test params(d) === (1.0, 1.0)
    end

    @testset "prod" begin
        @test prod(ProdAnalytical(), Beta(3.0, 2.0), Beta(2.0, 1.0)) ≈ Beta(4.0, 2.0)
        @test prod(ProdAnalytical(), Beta(7.0, 1.0), Beta(0.1, 4.5)) ≈ Beta(6.1, 4.5)
        @test prod(ProdAnalytical(), Beta(1.0, 3.0), Beta(0.2, 0.4)) ≈ Beta(0.19999999999999996, 2.4)
    end

    @testset "logmean" begin
        @test logmean(Beta(1.0, 3.0)) ≈ -1.8333333333333335
        @test logmean(Beta(0.1, 0.3)) ≈ -7.862370395825961
        @test logmean(Beta(4.5, 0.3)) ≈ -0.07197681436958758
    end

    @testset "mirroredlogmean" begin
        @test mirroredlogmean(Beta(1.0, 3.0)) ≈ -0.33333333333333337
        @test mirroredlogmean(Beta(0.1, 0.3)) ≈ -0.9411396776150167
        @test mirroredlogmean(Beta(4.5, 0.3)) ≈ -4.963371962929249
    end

end

end


# Day 16: Packet Decoder
function hex2bits(hex) # "BITS" representation of our data
    # Parse as Int8 as there is not Int4 datatype, but only keep last 4 digits.
    return Vector{Char}(
        join(bitstring(i)[5:8] for i in parse.(Int8, split(hex, ""), base=16))
    )
end
load(::Day{16}, path) = hex2bits(readline(path))

function bin2num(binary_string) # independent of string length
    return sum(c == '1' ? 2^(i - 1) : 0 for (i, c) in enumerate(reverse(binary_string)))
end

# Part 1 & 2 use the same parser
function parse_transmission!(bitstream)
    version_sum = 0

    function parse_packet!(bitstream)
        bits = splice!(bitstream, 1:6)
        version = bin2num(bits[1:3])
        type_id = bin2num(bits[4:6])
        version_sum += version

        # Every parsed packet returns its value and length
        if type_id == 4 # packet is a literal value
            # Iterate over value sub-packets
            is_last_group = false
            val_bits = Vector{Char}()
            subpacket_count = 0

            while !is_last_group
                is_last_group = (popfirst!(bitstream) == '0')
                append!(val_bits, splice!(bitstream, 1:4))
                subpacket_count += 1
            end

            packet_length = 6 + subpacket_count * 5
            value = bin2num(val_bits)
        else
            # Packet is an operator acting on subpackets
            length_type_id = popfirst!(bitstream)
            subpackets = Vector{Tuple{Int,Int}}()
            if length_type_id == '0'
                # Parse subpacket by length of bits
                subpacket_length = bin2num(splice!(bitstream, 1:15))
                subpacket_bits = splice!(bitstream, 1:subpacket_length)
                while !isempty(subpacket_bits)
                    push!(subpackets, parse_packet!(subpacket_bits))
                end
                packet_length = 6 + subpacket_length
            else
                # Parse subpackets by number
                subpacket_number = bin2num(splice!(bitstream, 1:11))
                while length(subpackets) < subpacket_number
                    push!(subpackets, parse_packet!(bitstream))
                end
                packet_length = 6 + sum(last.(subpackets))
            end
            # Apply operators to subpackets
            if type_id == 0
                value = sum(first.(subpackets))
            elseif type_id == 1
                value = prod(first.(subpackets))
            elseif type_id == 2
                value = minimum(first.(subpackets))
            elseif type_id == 3
                value = maximum(first.(subpackets))
            elseif type_id == 5
                value = Int(first.(subpackets[1]) > first.(subpackets[2]))
            elseif type_id == 6
                value = Int(first.(subpackets[1]) < first.(subpackets[2]))
            elseif type_id == 7
                value = Int(first.(subpackets[1]) == first.(subpackets[2]))
            end
        end

        return value, packet_length
    end

    value, _ = parse_packet!(bitstream)
    return version_sum, value
end

# Part 1 – return sum of versions
solve1(::Day{16}, data) = first(parse_transmission!(deepcopy(data)))
# Part 2 – return value
solve2(::Day{16}, data) = last(parse_transmission!(deepcopy(data)))

# No test for for both parts, just pass:
testresult1(::Day{16}) = 0
test1(::Day{16}) = 0
testresult2(::Day{16}) = 0
test2(::Day{16}) = 0


sp = " "
num1 = split("one two three four five six seven eight nine", sp)
num2 = split("eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen nineteen", sp)
num3 = split("ten twenty thirty forty fifty sixty seventy eighty ninety", sp)
num4 = split("hundred thousand", sp)
conj = "and"

function n2w(i)
  words = []
  if i == 1000
    push!(words, num1[1], num4[2])
  else
    if i >= 100
      push!(words, num1[Int(floor(i/100))], num4[1])
      if i % 100 != 0
        push!(words, conj)
      end
    end
    if i % 100 != 0
      if 11 <= i%100 <= 19
        push!(words, num2[Int(floor(i%100%10))])
      else
        if i % 100 >= 10
          push!(words, num3[Int(floor(i%100/10))])
        end
        if i % 10 != 0
          push!(words, num1[Int(floor(i%10))])
        end
      end
    end
  end
  return words
end

total = 0
for i = 1:1000
  global total += sum(map(x -> length(x), n2w(i)))
end
println(total)


immutable Quaternion{T<:Real} <: Number
    s::T
    v1::T
    v2::T
    v3::T
    norm::Bool
end

Quaternion(s::Real, v1::Real, v2::Real, v3::Real, n::Bool = false) =
    Quaternion( promote(s, v1, v2, v3)..., n)
Quaternion(x::Real) = Quaternion(x, zero(x), zero(x), zero(x), abs(x) == one(x))
Quaternion(z::Complex) = Quaternion(z.re, z.im, zero(z.re), zero(z.re), abs(z) == one(z.re))
Quaternion(s::Real, a::Vector) = Quaternion(s, a[1], a[2], a[3])
Quaternion(a::Vector) = Quaternion(0, a[1], a[2], a[3])

convert{T}(::Type{Quaternion{T}}, x::Real) =
    Quaternion(convert(T,x), convert(T,0), convert(T,0), convert(T,0))
convert{T}(::Type{Quaternion{T}}, z::Complex) =
    Quaternion(convert(T,real(z)), convert(T,imag(z)), convert(T,0), convert(T,0))
convert{T<:Real}(::Type{Quaternion{T}}, q::Quaternion{T}) = q
convert{T}(::Type{Quaternion{T}}, q::Quaternion) =
    Quaternion(convert(T,q.s), convert(T,q.v1), convert(T,q.v2), convert(T,q.v3), q.norm)

promote_rule{T<:Real}(::Type{Quaternion{T}}, ::Type{T}) = Quaternion{T}
promote_rule{T<:Real}(::Type{Quaternion}, ::Type{T}) = Quaternion
promote_rule{T<:Real,S<:Real}(::Type{Quaternion{T}}, ::Type{S}) = Quaternion{promote_type(T,S)}
promote_rule{T<:Real,S<:Real}(::Type{Complex{T}}, ::Type{Quaternion{S}}) = Quaternion{promote_type(T,S)}
promote_rule{T<:Real,S<:Real}(::Type{Quaternion{T}}, ::Type{Quaternion{S}}) = Quaternion{promote_type(T,S)}

quat(p, v1, v2, v3, n = false) = Quaternion(p, v1, v2, v3, n)
quat(x) = Quaternion(x)
quat(s, a) = Quaternion(s, a)

function show(io::IO, q::Quaternion)
    pm(x) = x < 0 ? " - $(-x)" : " + $x"
    print(io, q.s, pm(q.v1), "im", pm(q.v2), "jm", pm(q.v3), "km")
end

real{T}(::Type{Quaternion{T}}) = T
real(q::Quaternion) = q.s
imag(q::Quaternion) = [q.v1, q.v2, q.v3]

(/)(q::Quaternion, x::Real) = Quaternion(q.s/x, q.v1/x, q.v2/x, q.v3/x)

conj(q::Quaternion) = Quaternion(q.s, -q.v1, -q.v2, -q.v3, q.norm)
abs(q::Quaternion) = sqrt(q.s*q.s + q.v1*q.v1 + q.v2*q.v2 + q.v3*q.v3)
abs2(q::Quaternion) = q.s*q.s + q.v1*q.v1 + q.v2*q.v2 + q.v3*q.v3
inv(q::Quaternion) = q.norm ? conj(q) : conj(q)/abs2(q)

isfinite(q::Quaternion) = q.norm ? true : (isfinite(q.s) && isfinite(q.v1) && isfinite(q.v2) && isfinite(q.v3))

function normalize(q::Quaternion)
    if (q.norm)
        return q
    end
    q = q / abs(q)
    Quaternion(q.s, q.v1, q.v2, q.v3, true)
end

function normalizea(q::Quaternion)
    if (q.norm)
        return (q,one(q.s))
    end
    a = abs(q)
    q = q / a
    (Quaternion(q.s, q.v1, q.v2, q.v3, true), a)
end

function normalizeq(q::Quaternion)
    a = abs(q)
    if a > 0
        q = q / a
        Quaternion(q.s, q.v1, q.v2, q.v3, true)
    else
        Quaternion(0.0, 1.0, 0.0, 0.0, true)
    end
end

(-)(q::Quaternion) = Quaternion(-q.s, -q.v1, -q.v2, -q.v3, q.norm)

(+)(q::Quaternion, w::Quaternion) =
    Quaternion(q.s + w.s, q.v1 + w.v1, q.v2 + w.v2, q.v3 + w.v3)

(-)(q::Quaternion, w::Quaternion) =
    Quaternion(q.s - w.s, q.v1 - w.v1, q.v2 - w.v2, q.v3 - w.v3)

(*)(q::Quaternion, w::Quaternion) = Quaternion(q.s*w.s - q.v1*w.v1 - q.v2*w.v2 - q.v3*w.v3,
                                               q.s*w.v1 + q.v1*w.s + q.v2*w.v3 - q.v3*w.v2,
                                               q.s*w.v2 - q.v1*w.v3 + q.v2*w.s + q.v3*w.v1,
                                               q.s*w.v3 + q.v1*w.v2 - q.v2*w.v1 + q.v3*w.s,
                                               q.norm && w.norm)
(/)(q::Quaternion, w::Quaternion) = q * inv(w)

angleaxis(q::Quaternion) = angle(q), axis(q)

angle(q::Quaternion) = 2 * acos(real(normalize(q)))

function axis(q::Quaternion)
    q = normalize(q)
    s = sin(angle(q) / 2)
        abs(s) > 0 ?
        [q.v1, q.v2, q.v3] / s :
        [1.0, 0.0, 0.0]
end

function argq(q::Quaternion)
    q = normalize(q)
    q = Quaternion(imag(q))
    normalizeq(q)
end

function exp(q::Quaternion)
    s = q.s
    se = exp(s)
    scale = se
    th = abs(Quaternion(imag(q)))
    if th > 0
        scale *= sin(th) / th
    end
    Quaternion(se*cos(th), scale*q.v1, scale*q.v2, scale*q.v3, abs(s) < eps(typeof(s)))
end

function log(q::Quaternion)
    q, a = normalizea(q)
    s = q.s
    M = abs(Quaternion(imag(q)))
    th = atan2(M, s)
    if M > 0
        M = th / M
        return Quaternion(log(a), q.v1*M, q.v2*M, q.v3*M)
    else
        return Quaternion(log(a), th , 0.0, 0.0)
    end
end

function sin(q::Quaternion)
    L = argq(q)
    return (exp(L*q) - exp(-L*q))/(2*L)
end

function cos(q::Quaternion)
    L = argq(q)
    return (exp(L*q) + exp(-L*q)) / 2
end

(^)(q::Quaternion, w::Quaternion) = exp(w*log(q))

sqrt(q::Quaternion) = exp(0.5*log(q))

function linpol(p::Quaternion, q::Quaternion, t::Real)
    p = normalize(p)
    q = normalize(q)
    qm = -q
    if abs(p - q) > abs(p - qm)
        q = qm
    end
    c = p.s * q.s + p.v1 * q.v1 + p.v2 * q.v2 + p.v3 * q.v3
    if c > - 1.0
        if c < 1.0
            o = acos(c)
            s = sin(o)
            sp = sin((1 - t)*o)/s
            sq = sin(t*o)/s
        else
        sp = 1 - t
        sq = t
    end
    Quaternion(sp*p.s  + sq*q.s,
               sp*p.v1 + sq*q.v1,
               sp*p.v2 + sq*q.v2,
               sp*p.v3 + sq*q.v3, true)
    else
    s  =  p.v3
    v1 = -p.v2
    v2 =  p.v1
    v3 = -p.s
    sp = sin((0.5 - t)*pi)
    sq = sin(t*pi)
    Quaternion(s,
               sp * p.v1 + sq * v1,
               sp * p.v2 + sq * v2,
               sp * p.v3 + sq * v3, true)
    end
end

quatrand()  = quat(randn(), randn(), randn(), randn())
nquatrand() = normalize(quatrand())

## Rotations

function qrotation{T<:Real}(axis::Vector{T}, theta)
    if length(axis) != 3
        error("Must be a 3-vector")
    end
    u = normalize(axis)
    thetaT = convert(eltype(u), theta)
    s = sin(thetaT/2)
    Quaternion(cos(thetaT/2), s*u[1], s*u[2], s*u[3], true)
end

# Variant of the above where norm(rotvec) encodes theta
function qrotation{T<:Real}(rotvec::Vector{T})
    if length(rotvec) != 3
        error("Must be a 3-vector")
    end
    theta = norm(rotvec)
    if theta > 0
        s = sin(theta/2)/theta  # divide by theta to make rotvec a unit vector
        return Quaternion(cos(theta/2), s*rotvec[1], s*rotvec[2], s*rotvec[3], true)
    end
    Quaternion(one(T), zero(T), zero(T), zero(T), true)
end

rotationmatrix(q::Quaternion) = rotationmatrix_normalized(normalize(q))

function rotationmatrix_normalized(q::Quaternion)
    sx, sy, sz = 2q.s*q.v1, 2q.s*q.v2, 2q.s*q.v3
    xx, xy, xz = 2q.v1^2, 2q.v1*q.v2, 2q.v1*q.v3
    yy, yz, zz = 2q.v2^2, 2q.v2*q.v3, 2q.v3^2
    [1-(yy+zz)     xy-sz     xz+sy;
        xy+sz   1-(xx+zz)    yz-sx;
        xz-sy      yz+sx  1-(xx+yy)]
end

function normalize{T}(v::Vector{T})
    nv = norm(v)
    if nv > 0
        return v/nv
    end
    zeros(promote_type(T,typeof(nv)), length(v))
end

function slerp{T}(qa::Quaternion{T}, qb::Quaternion{T}, t::T)
    # http://www.euclideanspace.com/maths/algebra/realNormedAlgebra/quaternions/slerp/
    coshalftheta = qa.s * qb.s + qa.v1 * qb.v1 + qa.v2 * qb.v2 + qa.v3 * qb.v3;

    if coshalftheta < 0
        qm = -qb
        coshalftheta = -coshalftheta
    else
        qm = qb
    end
    abs(coshalftheta) >= 1.0 && return qa

    halftheta    = acos(coshalftheta)
    sinhalftheta = sqrt(one(T) - coshalftheta * coshalftheta)

    if abs(sinhalftheta) < 0.001
        return Quaternion(
            T(0.5) * (qa.s  + qb.s),
            T(0.5) * (qa.v1 + qb.v1),
            T(0.5) * (qa.v2 + qb.v2),
            T(0.5) * (qa.v3 + qb.v3),
        )
    end

    ratio_a = sin((one(T) - t) * halftheta) / sinhalftheta
    ratio_b = sin(t * halftheta) / sinhalftheta

    Quaternion(
        qa.s  * ratio_a + qm.s  * ratio_b,
        qa.v1 * ratio_a + qm.v1 * ratio_b,
        qa.v2 * ratio_a + qm.v2 * ratio_b,
        qa.v3 * ratio_a + qm.v3 * ratio_b,
    )
end


import FinanceLib as Fl

import FinanceLib.FixedIncomes.MoneyMarkets as MM
import FinanceLib.FixedIncomes as FI

@testset "FinanceLib.FixedIncomes                                     " begin

  @testset "MoneyMarkets" begin
    @test MM.tBillR(150,98_000,100_000) ≈ 0.048
    @test MM.tBillD(0.048, 150, 100_000) == 2_000

    @test MM.holdingPerYield(98, 95, 5) == 0.020408163265306145
    @test MM.effAnnYield(150, 98, 95, 5) == 0.05038831660532006
    @test MM.moneyMktYield(150, 98, 95, 5) == 0.04897959183673475

    @test MM.twrr([4,6,5.775,6.72,5.508],[1,-0.5,0.225,-0.6]) ==
    0.06159232319186159
    @test MM.twrr(1.0, [100, 112, 142.64], [0, 20.])==0.21027878787878795

  end

end

include("Bonds/runtests.jl")


@testset "param" begin

# Need to test symmetry of param3dmat for different boundary conditions.
@testset "symmetry of create_paramop, uniform grid" begin
    # Create a grid.
    isbloch = [true, true, true]
    # isbloch = [true, false, false]
    lprim = ([-1.5, -0.5, 0.5, 1.5], [-1.5, -0.5, 0.5, 1.5], [-1.5, -0.5, 0.5, 1.5])
    g3 = Grid(lprim, isbloch)
    N = g3.N

    # Create materials.
    εvac = 1.0
    vac = Material("Vacuum", ε=εvac)

    εdiel = 2.0
    diel = Material("Dielectric", ε=εdiel)

    # Create objects.
    dom_vac = Object(Cuboid(g3.bounds), vac)
    obj_diel = Object(Cuboid([0,0,0], [1,1,1]), diel)
    # obj_diel = Object(Ball([0,0,0], 1), diel)

    # Add objects.
    oind2shp = Shape3[]
    oind2εind = ParamInd[]
    oind2μind = ParamInd[]
    εind2ε = S²ComplexF3[]
    μind2μ = S²ComplexF3[]

    add_obj!(oind2shp, (oind2εind,oind2μind), (εind2ε,μind2μ), dom_vac, obj_diel)

    # Construct arguments and call assign_param!.
    N = g3.N
    ε3d = create_param_array(N)
    εxx_oind3d = create_oind_array(N)
    εyy_oind3d = create_oind_array(N)
    εzz_oind3d = create_oind_array(N)
    εoo_oind3d = create_oind_array(N)

    μ3d = create_param_array(N)
    μxx_oind3d = create_oind_array(N)
    μyy_oind3d = create_oind_array(N)
    μzz_oind3d = create_oind_array(N)
    μoo_oind3d = create_oind_array(N)

    boundft = SVec(EE,EE,EE)
    assign_param!(ε3d, (μxx_oind3d,μyy_oind3d,μzz_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(ε3d, tuple(μoo_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(μ3d, (εxx_oind3d,εyy_oind3d,εzz_oind3d), oind2shp, oind2μind, μind2μ, ft2gt.(HH,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(μ3d, tuple(εoo_oind3d), oind2shp, oind2μind, μind2μ, ft2gt.(HH,boundft), g3.ghosted.τl, g3.isbloch)

    # Perform smoothing.
    smooth_param!(ε3d, (εxx_oind3d,εyy_oind3d,εzz_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.l, g3.ghosted.l, g3.σ, g3.ghosted.∆τ)
    smooth_param!(ε3d, tuple(εoo_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.l, g3.ghosted.l, g3.σ, g3.ghosted.∆τ)

    ∆l = g3.∆l[nDL]
    ∆l′ = g3.∆l[nPR]
    ∆l′⁻¹ = map(x->inv.(x), ∆l′)

    Mε = create_paramop(ε3d, boundft.!=EE, ∆l, ∆l′⁻¹, g3.isbloch, order_cmpfirst=true)

    # @test issymmetric(Mε)
    @test Mε ≈ transpose(Mε)

    # Code I used to reveal the matrix structure while debugging:
    # (i,j) = (nX,nY); full(real.(Mε3))[27(i-1)+1:27i, 27(j-1)+1:27j]
end  # @testset "create_paramop"

@testset "symmetry of create_paramop, nonuniform grid" begin
    # Create a grid.
    isbloch = [true, true, true]
    # isbloch = [true, false, false]
    lprim = ([-1.5, -0.4, 0.3, 1.5], [-1.5, -0.4, 0.3, 1.5], [-1.5, -0.4, 0.3, 1.5])
    g3 = Grid(lprim, isbloch)
    N = g3.N

    # Create materials.
    εvac = 1.0
    vac = Material("Vacuum", ε=εvac)

    εdiel = 2.0
    diel = Material("Dielectric", ε=εdiel)

    # Create objects.
    dom_vac = Object(Cuboid(g3.bounds), vac)
    obj_diel = Object(Cuboid([0,0,0], [1,1,1]), diel)
    # obj_diel = Object(Ball([0,0,0], 1), diel)

    # Add objects.
    oind2shp = Shape3[]
    oind2εind = ParamInd[]
    oind2μind = ParamInd[]
    εind2ε = S²ComplexF3[]
    μind2μ = S²ComplexF3[]

    add_obj!(oind2shp, (oind2εind,oind2μind), (εind2ε,μind2μ), dom_vac, obj_diel)

    # Construct arguments and call assign_param!.
    N = g3.N
    ε3d = create_param_array(N)
    εxx_oind3d = create_oind_array(N)
    εyy_oind3d = create_oind_array(N)
    εzz_oind3d = create_oind_array(N)
    εoo_oind3d = create_oind_array(N)

    μ3d = create_param_array(N)
    μxx_oind3d = create_oind_array(N)
    μyy_oind3d = create_oind_array(N)
    μzz_oind3d = create_oind_array(N)
    μoo_oind3d = create_oind_array(N)

    boundft = SVec(EE,EE,EE)
    assign_param!(ε3d, (μxx_oind3d,μyy_oind3d,μzz_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(ε3d, tuple(μoo_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(μ3d, (εxx_oind3d,εyy_oind3d,εzz_oind3d), oind2shp, oind2μind, μind2μ, ft2gt.(HH,boundft), g3.ghosted.τl, g3.isbloch)
    assign_param!(μ3d, tuple(εoo_oind3d), oind2shp, oind2μind, μind2μ, ft2gt.(HH,boundft), g3.ghosted.τl, g3.isbloch)

    # Perform smoothing.
    smooth_param!(ε3d, (εxx_oind3d,εyy_oind3d,εzz_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.l, g3.ghosted.l, g3.σ, g3.ghosted.∆τ)
    smooth_param!(ε3d, tuple(εoo_oind3d), oind2shp, oind2εind, εind2ε, ft2gt.(EE,boundft), g3.l, g3.ghosted.l, g3.σ, g3.ghosted.∆τ)

    ∆l = g3.∆l[nDL]
    ∆l′ = g3.∆l[nPR]
    ∆l′⁻¹ = map(x->inv.(x), ∆l′)

    Mε = create_paramop(ε3d, boundft.!=EE, ∆l, ∆l′⁻¹, g3.isbloch, order_cmpfirst=true)

    # Ml
    ∆lprim = g3.∆l[nPR]
    ∆ldual = g3.∆l[nDL]

    Nx, Ny, Nz = N
    nX, nY, nZ = 1, 2, 3
    ∆X = repeat(∆ldual[nX], outer=(1,Ny,Nz))
    ∆Y = repeat(reshape(∆ldual[nY], (1,Ny,1)), outer=(Nx,1,Nz))
    ∆Z = repeat(reshape(∆ldual[nZ], (1,1,Nz)), outer=(Nx,Ny,1))

    ∆L = [reshape(∆X,1,:); reshape(∆Y,1,:); reshape(∆Z,1,:)]
    ∆l = ∆L[:]
    Ml⁻¹ = sparse(Diagonal(1 ./ ∆l))

    # Ma
    ∆X = repeat(∆lprim[nX], outer=(1,Ny,Nz))
    ∆Y = repeat(reshape(∆lprim[nY], (1,Ny,1)), outer=(Nx,1,Nz))
    ∆Z = repeat(reshape(∆lprim[nZ], (1,1,Nz)), outer=(Nx,Ny,1))

    ∆YZ = ∆Y .* ∆Z
    ∆ZX = ∆Z .* ∆X
    ∆XY = ∆X .* ∆Y

    ∆A = [reshape(∆YZ,1,:); reshape(∆ZX,1,:); reshape(∆XY,1,:)]
    ∆a = ∆A[:]
    Ma = sparse(Diagonal(∆a))

    Mε_sym = Ma * Mε * Ml⁻¹

    @test Mε_sym ≈ transpose(Mε_sym)

    # Code I used to reveal the matrix structure while debugging:
    # (i,j) = (nX,nY); full(real.(Mε3))[27(i-1)+1:27i, 27(j-1)+1:27j]
end  # @testset "create_paramop"

# Need to test if PML still works after the field-averaging operator applied to the material
# parameters, because effectively it makes the material inhomogeneous along the direction
# normal to the PML interface.

end  # @testset "param"


@testset "Slope Limiter" begin

using ConservationLawsDiffEq

for limiter in (GeneralizedMinmodLimiter(), MinmodLimiter())
    printstyled("Limiter: ", string(typeof(limiter)) ,"\n"; color = :blue)
    @test limiter(1.0,2.0) == 1.0
    @test limiter(-1.0,2.0) == 0.0
    @test limiter(-1.0,-2.0) == -1.0
end

limiter  = OsherLimiter()
    printstyled("Limiter: ", string(typeof(limiter)) ,"\n"; color = :blue)
    @test limiter(1.0,2.0) == 1.0
    @test limiter(-1.0,2.0) == 0.0
    @test limiter(-1.0,-2.0) == 0.0

limiter = SuperbeeLimiter()
    printstyled("Limiter: ", string(typeof(limiter)) ,"\n"; color = :blue)
    @test limiter(1.0,2.0) == 2.0
    @test limiter(-1.0,2.0) == 0.0
    @test limiter(-1.0,-2.0) == 0.0
end


@doc """
	clip(A::AbstractPoint2D, B::AbstractPoint2D) -> C, D

Clip the line segment with endpoints `A` and `B` to the bounding box.
The returned points `C` and `D` are the endpoints of the intersection.

If the line segment is not intersecting the bounding box, both `C` and `D` are `Void`.
"""->
function clip(A::AbstractPoint2D, B::AbstractPoint2D)
	# TODO: Use Line2D from GeometricalPredicates as input?
	if isinside(A) && isinside(B)
		return A, B
	end

	t0 = 0.0
	t1 = 1.0
	D = B - A
	# left, right, bottom, top
	# Parametrization of line segment from A to B
	p = [-getx(D) ; getx(D) ; -gety(D) ; gety(D)]
	q = [getx(A)-LEFT ; RIGHT-getx(A) ; gety(A)-LOWER ; UPPER-gety(A)]
	for k in 1:4
		if p[k] == 0.0
			# Line parallel with k'th edge
			if q[k] < 0.0
				return nothing, nothing
			end
		elseif p[k] < 0.0
			# Outside to inside
			u = q[k] / p[k]
			if u > t1
				return nothing, nothing
			end
			t0 = max( u, t0 )
		else # p[k] > 0.0
			# Inside to outside
			u = q[k] / p[k]
			if u < t0
				return nothing, nothing
			end
			t1 = min( u, t1 )
		end
	end

	return A + t0*D, A + t1*D
end

@doc """
	isinside(p::Point2D) -> Bool

Test if the point `p` is inside the bounding box.
"""->
function isinside(p::AbstractPoint2D)
	LEFT <= getx(p) <= RIGHT && LOWER <= gety(p) <= UPPER
end

@doc """
	isoutside(edge::VoronoiEdge) -> Bool

Test if `edge` intersects the bounding box.
"""->
function isoutside(edge::VoronoiDelaunay.VoronoiEdge{IndexablePoint2D})
	A, B = clip( geta(edge), getb(edge) )

	if isa(A,Void) && isa(B,Void)
		return true
	else
		return false
	end
end


using LinearAlgebra
using Oceananigans.Architectures: array_type
using Oceananigans.TimeSteppers: _compute_w_from_continuity!

function can_solve_single_tridiagonal_system(arch, N)
    ArrayType = array_type(arch)

    a = rand(N-1)
    b = 3 .+ rand(N) # +3 to ensure diagonal dominance.
    c = rand(N-1)
    f = rand(N)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, b, c)
    ϕ_correct = M \ f

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    grid = RegularCartesianGrid(size=(1, 1, N), length=(1, 1, 1))
    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_single_tridiagonal_system_with_functions(arch, N)
    ArrayType = array_type(arch)

    grid = RegularCartesianGrid(size=(1, 1, N), length=(1, 1, 1))

    a = rand(N-1)
    c = rand(N-1)

    @inline b(i, j, k, grid, p) = 3 .+ cos(2π*grid.zC[k])  # +3 to ensure diagonal dominance.
    @inline f(i, j, k, grid, p) = sin(2π*grid.zC[k])

    bₐ = [b(1, 1, k, grid, nothing) for k in 1:N]
    fₐ = [f(1, 1, k, grid, nothing) for k in 1:N]

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, bₐ, c)
    ϕ_correct = M \ fₐ

    # Convert to CuArray if needed.
    a, c = ArrayType.([a, c])

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    a = rand(Nz-1)
    b = 3 .+ rand(Nz) # +3 to ensure diagonal dominance.
    c = rand(Nz-1)
    f = rand(Nx, Ny, Nz)

    M = Tridiagonal(a, b, c)
    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the systems with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        ϕ_correct[i, j, :] .= M \ f[i, j, :]
    end

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(1, 1, 1))
    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(1, 1, 1))

    a = rand(Nz-1)
    c = rand(Nz-1)

    @inline b(i, j, k, grid, p) = 3 + grid.xC[i]*grid.yC[j] * cos(2π*grid.zC[k])
    @inline f(i, j, k, grid, p) = (grid.xC[i] + grid.yC[j]) * sin(2π*grid.zC[k])

    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        bₐ = [b(i, j, k, grid, nothing) for k in 1:Nz]
        M = Tridiagonal(a, bₐ, c)

        fₐ = [f(i, j, k, grid, nothing) for k in 1:Nz]
        ϕ_correct[i, j, :] .= M \ fₐ
    end

    # Convert to CuArray if needed.
    a, c = ArrayType.([a, c])

    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType
    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ) ≈ ϕ_correct
end

function vertically_stretched_poisson_solver_correct_answer(arch, Nx, Ny, zF)
    Lx, Ly, Lz = 1, 1, zF[end]
    Δx, Δy = Lx/Nx, Ly/Ny

    #####
    ##### Vertically stretched operators
    #####

    @inline δx_caa(i, j, k, f) = @inbounds f[i+1, j, k] - f[i, j, k]
    @inline δy_aca(i, j, k, f) = @inbounds f[i, j+1, k] - f[i, j, k]
    @inline δz_aac(i, j, k, f) = @inbounds f[i, j, k+1] - f[i, j, k]

    @inline ∂x_caa(i, j, k, Δx,  f) = δx_caa(i, j, k, f) / Δx
    @inline ∂y_aca(i, j, k, Δy,  f) = δy_aca(i, j, k, f) / Δy
    @inline ∂z_aac(i, j, k, ΔzF, f) = δz_aac(i, j, k, f) / ΔzF[k]

    @inline ∂x²(i, j, k, Δx, f)       = (∂x_caa(i, j, k, Δx, f)  - ∂x_caa(i-1, j, k, Δx, f))  / Δx
    @inline ∂y²(i, j, k, Δy, f)       = (∂y_aca(i, j, k, Δy, f)  - ∂y_aca(i, j-1, k, Δy, f))  / Δy
    @inline ∂z²(i, j, k, ΔzF, ΔzC, f) = (∂z_aac(i, j, k, ΔzF, f) - ∂z_aac(i, j, k-1, ΔzF, f)) / ΔzC[k]

    @inline div_ccc(i, j, k, Δx, Δy, ΔzF, u, v, w) = ∂x_caa(i, j, k, Δx, u) + ∂y_aca(i, j, k, Δy, v) + ∂z_aac(i, j, k, ΔzF, w)

    @inline ∇²(i, j, k, Δx, Δy, ΔzF, ΔzC, f) = ∂x²(i, j, k, Δx, f) + ∂y²(i, j, k, Δy, f) + ∂z²(i, j, k, ΔzF, ΔzC, f)

    #####
    ##### Generate "fake" vertically stretched grid
    #####

    function grid(zF)
        Nz = length(zF) - 1
        ΔzF = [zF[k+1] - zF[k] for k in 1:Nz]
        zC = [(zF[k] + zF[k+1]) / 2 for k in 1:Nz]
        ΔzC = [zC[k+1] - zC[k] for k in 1:Nz-1]
        return zF, zC, ΔzF, ΔzC
    end

    Nz = length(zF) - 1
    zF, zC, ΔzF, ΔzC = grid(zF)

    # Need some halo regions.
    ΔzF = OffsetArray([ΔzF[1], ΔzF...], 0:Nz)
    ΔzC = [ΔzC..., ΔzC[end]]

    # Useful for broadcasting z operations
    ΔzC = reshape(ΔzC, (1, 1, Nz))

    # Temporary hack: Useful for reusing fill_halo_regions! and
    # BatchedTridiagonalSolver which only need Nx, Ny, Nz.
    fake_grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(Lx, Ly, Lz))

    #####
    ##### Generate batched tridiagonal system coefficients and solver
    #####

    function λi(Nx, Δx)
        is = reshape(1:Nx, Nx, 1, 1)
        @. (2sin((is-1)*π/Nx) / Δx)^2
    end

    function λj(Ny, Δy)
        js = reshape(1:Ny, 1, Ny, 1)
        @. (2sin((js-1)*π/Ny) / Δy)^2
    end

    kx² = λi(Nx, Δx)
    ky² = λj(Ny, Δy)

    # Lower and upper diagonals are the same
    ld = [1/ΔzF[k] for k in 1:Nz-1]
    ud = copy(ld)

    # Diagonal (different for each i,j)
    @inline δ(k, ΔzF, ΔzC, kx², ky²) = - (1/ΔzF[k-1] + 1/ΔzF[k]) - ΔzC[k] * (kx² + ky²)

    d = zeros(Nx, Ny, Nz)
    for i in 1:Nx, j in 1:Ny
        d[i, j, 1] = -1/ΔzF[1] - ΔzC[1] * (kx²[i] + ky²[j])
        d[i, j, 2:Nz-1] .= [δ(k, ΔzF, ΔzC, kx²[i], ky²[j]) for k in 2:Nz-1]
        d[i, j, Nz] = -1/ΔzF[Nz-1] - ΔzC[Nz] * (kx²[i] + ky²[j])
    end

    #####
    ##### Random right hand side
    #####

    # Random right hand side
    Ru = CellField(Float64, arch, fake_grid, UVelocityBoundaryConditions(fake_grid))
    Rv = CellField(Float64, arch, fake_grid, VVelocityBoundaryConditions(fake_grid))
    Rw = CellField(Float64, arch, fake_grid, WVelocityBoundaryConditions(fake_grid))

    interior(Ru) .= rand(Nx, Ny, Nz)
    interior(Rv) .= rand(Nx, Ny, Nz)
    interior(Rw) .= zeros(Nx, Ny, Nz)

    U = (u=Ru, v=Rv, w=Rw)
    fill_halo_regions!(U, arch)

    _compute_w_from_continuity!(U, fake_grid)

    fill_halo_regions!(Rw, arch)

    R = zeros(Nx, Ny, Nz)
    for i in 1:Nx, j in 1:Ny, k in 1:Nz
        R[i, j, k] = div_ccc(i, j, k, Δx, Δy, ΔzF, Ru.data, Rv.data, Rw.data)
    end

    # @show sum(R)  # should be zero by construction.

    F = zeros(Nx, Ny, Nz)
    F = ΔzC .* R  # RHS needs to be multiplied by ΔzC

    #####
    ##### Solve system
    #####

    F̃ = fft(F, [1, 2])

    btsolver = BatchedTridiagonalSolver(arch, dl=ld, d=d, du=ud, f=F̃, grid=fake_grid)

    ϕ̃ = zeros(Complex{Float64}, Nx, Ny, Nz)
    solve_batched_tridiagonal_system!(ϕ̃, arch, btsolver)

    ϕ = CellField(Float64, arch, fake_grid, PressureBoundaryConditions(fake_grid))
    interior(ϕ) .= real.(ifft(ϕ̃, [1, 2]))
    ϕ.data .= ϕ.data .- mean(interior(ϕ))

    #####
    ##### Compute Laplacian of solution ϕ to test that it's correct
    #####

    fill_halo_regions!(ϕ, arch)

    ∇²ϕ = CellField(Float64, arch, fake_grid, PressureBoundaryConditions(fake_grid))
    for i in 1:Nx, j in 1:Ny, k in 1:Nz
        ∇²ϕ.data[i, j, k] = ∇²(i, j, k, Δx, Δy, ΔzF, ΔzC, ϕ.data)
    end

    return interior(∇²ϕ) ≈ R
end

@testset "Solvers" begin
    @info "Testing Solvers..."

    for arch in archs
        @testset "Batched tridiagonal solver [$arch]" begin
            for Nz in [8, 11, 18]
                @test can_solve_single_tridiagonal_system(arch, Nz)
                @test can_solve_single_tridiagonal_system_with_functions(arch, Nz)
            end

            for Nx in [3, 8], Ny in [5, 16], Nz in [8, 11]
                @test can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
                @test can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
            end
        end
    end

    for arch in [CPU()]
        @testset "Vertically stretched Poisson solver [FACR, $arch]" begin
            @info "  Testing vertically stretched Poisson solver [FACR, $arch]..."

            Nx = Ny = 8
            zF = [1, 2, 4, 7, 11, 16, 22, 29, 37]
            @test vertically_stretched_poisson_solver_correct_answer(arch, Nx, Ny, zF)
        end
    end
end


import Base: +, -, *, /, \

#--------------------------------------------------
# Vector space algebra

# Unary ops
-(a::StaticArrayLite) = map(-, a)

# Binary ops
# Between arrays
+(a::StaticArrayLite, b::StaticArrayLite) = map(+, a, b)
+(a::AbstractArray, b::StaticArrayLite) = map(+, a, b)
+(a::StaticArrayLite, b::AbstractArray) = map(+, a, b)

-(a::StaticArrayLite, b::StaticArrayLite) = map(-, a, b)
-(a::AbstractArray, b::StaticArrayLite) = map(-, a, b)
-(a::StaticArrayLite, b::AbstractArray) = map(-, a, b)

# Scalar-array
*(a::Number, b::StaticArrayLite) = map(x -> a*x, b)
*(a::StaticArrayLite, b::Number) = map(x -> x*b, a)

/(a::StaticArrayLite, b::Number) = map(x -> x/b, a)
\(a::Number, b::StaticArrayLite) = map(x -> a\x, b)

#--------------------------------------------------
# Matrix algebra

# Transpose, conjugate, etc
conj(a::StaticArrayLite) = map(conj, a)

dot(a::StaticArrayLite, b::LinearAlgebra.Adjoint{<:StaticArrayLite}) = a*adjoint(b)
bilinear_vecdot(a::StaticArrayLite, b::StaticArrayLite) = a*transpose(b)

*(a::LinearAlgebra.Transpose{<:StaticArrayLite{<:Number, 1}}, b::StaticArrayLite{<:Number,1}) = mapreduce(*, +, a, b)
*(a::LinearAlgebra.Adjoint{<:StaticArrayLite{<:Real, 1}}, b::StaticArrayLite{<:Number,1}) = mapreduce(*, +, a, b)
*(a::LinearAlgebra.Adjoint{<:StaticArrayLite{<:Complex, 1}}, b::StaticArrayLite{<:Number,1}) = mapreduce((x, y) -> conj(x)*y, +, a, b)

#--------------------------------------------------
# Norms
LinearAlgebra.norm_sqr(a::StaticArrayLite) = mapreduce(abs2, +, a; init=zero(real(eltype(a))))
LinearAlgebra.norm(a::StaticArrayLite) = sqrt(LinearAlgebra.norm_sqr(a))

LinearAlgebra.normalize(a::StaticArrayLite) = inv(norm(a))*a

function LinearAlgebra.normalize!(a::StaticArrayLite)
    c = inv(norm(a))
    map!(x -> c*x, a, a)
    return a
end


"""
 Holds the tableau of a weak implicit Runge-Kutta method.

   According to Wang, Hong, Xu, "Construction of Symplectic Runge-Kutta Methods for Stochastic Hamiltonian Systems",
   Commun. Comput. Phys. 21(1), 2017
"""
struct TableauWIRK{T} <: AbstractTableauIRK{T}
    name::Symbol
    s::Int
    qdrift0::CoefficientsRK{T}
    qdrift1::CoefficientsRK{T}
    qdiff0::CoefficientsRK{T}
    qdiff1::CoefficientsRK{T}
    qdiff3::CoefficientsRK{T}

    # Order of the tableau is not included, because unlike in the deterministic
    # setting, it depends on the properties of the noise (e.g., the dimension of
    # the Wiener process and the commutativity properties of the diffusion matrix)
    #
    # Orders stored in qdrift and qdiff are understood as the classical orders of these methods.
    #
    # qdrift0, qdrift1 correspond to A0, A1 in the paper
    # qdiff0, qdiff1, qdiff3 correspond to B0, B1, B3
    # qdrift0.b = alpha
    # qdiff0.b  = beta
    # qdrift0.c = c0
    # qdrift1.c = c1

    function TableauWIRK{T}(name, s, qdrift0, qdrift1, qdiff0, qdiff1, qdiff3) where {T}
        # THE COMMENTED OUT PART WAS FOR TableauFIRK. MAY IMPLEMENT SOMETHING
        # SIMILAR FOR TableauSFIRK LATER.

        # if (q.s > 1 && istrilstrict(q.a)) || (q.s==1 && q.a[1,1] == 0)
        #     warn("Initializing TableauFIRK with explicit tableau ", q.name, ".\n",
        #          "You might want to use TableauERK instead.")
        # elseif q.s > 1 && istril(q.a)
        #     warn("Initializing TableauFIRK with diagonally implicit tableau ", q.name, ".\n",
        #          "You might want to use TableauDIRK instead.")
        # end

        @assert s == qdrift0.s == qdrift1.s == qdiff0.s == qdiff1.s == qdiff3.s

        new(name, s, qdrift0, qdrift1, qdiff0, qdiff1, qdiff3)
    end
end

function TableauWIRK(name::Symbol, qdrift0::CoefficientsRK{T}, qdrift1::CoefficientsRK{T}, qdiff0::CoefficientsRK{T}, qdiff1::CoefficientsRK{T}, qdiff3::CoefficientsRK{T}) where {T}
    TableauWIRK{T}(name, qdrift0.s, qdrift0, qdrift1, qdiff0, qdiff1, qdiff3)
end

function TableauWIRK(name::Symbol, A0::Matrix{T}, A1::Matrix{T},
                                   B0::Matrix{T}, B1::Matrix{T}, B3::Matrix{T},
                                    α::Vector{T}, β1::Vector{T},
                                   c0::Vector{T}, c1::Vector{T} ) where {T}
    TableauWIRK(name, CoefficientsRK(name, 0, A0, α, c0), CoefficientsRK(name, 0, A1, α, c1),
                         CoefficientsRK(name, 0, B0, β1, c0), CoefficientsRK(name, 0, B1, β1, c1), CoefficientsRK(name, 0, B3, β1, c1))
end

"Parameters for right-hand side function of weak implicit Runge-Kutta methods."
mutable struct ParametersWIRK{DT, TT, ET <: SDE{DT,TT}, D, M, S} <: Parameters{DT,TT}
    equ::ET
    tab::TableauWIRK{TT}
    Δt::TT
    ΔW::Vector{DT}
    ΔZ::Vector{DT}

    t::TT
    q::Vector{DT}
end

function ParametersWIRK(equ::ET, tab::TableauWIRK{TT}, Δt::TT, ΔW::Vector{DT}, ΔZ::Vector{DT}) where {DT, TT, ET <: SDE{DT,TT}}
    @assert equ.m == length(ΔW) == length(ΔZ)
    ParametersWIRK{DT, TT, ET, equ.d, equ.m, tab.s}(equ, tab, Δt, ΔW, ΔZ, 0, zeros(DT, equ.d))
end

struct IntegratorCacheWIRK{DT}
    # Structure for holding the internal stages Q0, and Q1 the values of the drift vector
    # and the diffusion matrix evaluated at the internal stages V=v(Q0), B=B(Q1),
    # and the increments Y = Δt*a_drift*v(Q) + a_diff*B(Q)*ΔW
    Q0::Vector{Vector{DT}}     # Q0[i][k]   - the k-th component of the internal stage Q^(0)_i
    Q1::Vector{Matrix{DT}}     # Q1[i][k,l] - the k-th component of the internal stage Q^(l)_i
    V ::Vector{Vector{DT}}     # V [i][k]   - the k-th component of the drift vector v(Q0[i][:])
    B ::Vector{Matrix{DT}}     # B [i][:,:] - the diffusion matrix at the stage i such that the l-th column B[i][:,l] is evaluated at Q^(l)_i
    Y0::Vector{Vector{DT}}     # Y0[i][:]   - the increment of the internal stage Q0[i][:]
    Y1::Vector{Matrix{DT}}     # Y1[i][:,l] - the increment of the internal stage Q1[i][:,l]

    v::Vector{DT}
    b::Matrix{DT}
    y::Vector{DT}

    Δy::Vector{DT}

    function IntegratorCacheWIRK{DT}(D, M, S) where {DT}
        # create internal stage vectors
        Q0 = create_internal_stage_vector(DT, D, S)
        Q1 = create_internal_stage_vector(DT, D, M, S)
        V  = create_internal_stage_vector(DT, D, S)
        B  = create_internal_stage_vector(DT, D, M, S)
        Y0 = create_internal_stage_vector(DT, D, S)
        Y1 = create_internal_stage_vector(DT, D, M, S)

        # create velocity and update vector
        v = zeros(DT,D)
        b = zeros(DT,D,M)
        y = zeros(DT,D)

        # create temporary vectors
        Δy  = zeros(DT,M)

        new(Q0, Q1, V, B, Y0, Y1, v, b, y, Δy)
    end
end

"Stochastic implicit Runge-Kutta integrator."
struct IntegratorWIRK{DT, TT, PT <: ParametersWIRK{DT,TT},
                              ST <: NonlinearSolver{DT}, N} <: StochasticIntegrator{DT,TT}
    params::PT
    solver::ST
    cache::IntegratorCacheWIRK{DT}
end

function IntegratorWIRK(equation::SDE{DT,TT,VT,BT,N}, tableau::TableauWIRK{TT}, Δt::TT) where {DT,TT,VT,BT,N}
    D = equation.d
    M = equation.m
    NS= max(equation.ns,equation.ni)
    S = tableau.s

    # create params
    params = ParametersWIRK(equation, tableau, Δt, zeros(DT,M), zeros(DT,M))

    # create solver
    solver = create_nonlinear_solver(DT, D*S*(M+1), params)

    # create cache for internal stage vectors and update vectors
    cache = IntegratorCacheWIRK{DT}(D, M, S)

    # create integrator
    IntegratorWIRK{DT, TT, typeof(params), typeof(solver), N}(params, solver, cache)
end

equation(integrator::IntegratorWIRK) = integrator.params.equ
timestep(integrator::IntegratorWIRK) = integrator.params.Δt
tableau(integrator::IntegratorWIRK) = integrator.params.tab
noisedims(integrator::IntegratorWIRK) = integrator.params.equ.m
Base.eltype(integrator::IntegratorWIRK{DT, TT, PT, ST, N}) where {DT, TT, PT, ST, N} = DT

function update_params!(int::IntegratorWIRK, sol::AtomicSolutionSDE)
    # set time for nonlinear solver and copy previous solution
    int.params.t  = sol.t
    int.params.q .= sol.q
    int.params.ΔW .= sol.ΔW
    int.params.ΔZ .= sol.ΔZ
end

"""
Unpacks the data stored in
x = (Y0[1][1], Y0[1][2]], ... Y0[1][D], ... Y0[S][D], Y1[1][1,1], Y1[1][2,1], ... Y1[1][D,1], Y1[1][1,2], Y1[1][2,2], ... Y1[1][D,2], ... Y1[S][D,M]  )
into Y0::Vector{Vector} and Y1::Vector{Matrix}, calculates the internal stages Q0 and Q1, the values of the RHS
of the SDE ( v(Q0) and B(Q1) ), and assigns them to V and B.
Unlike for FIRK, here Y = Δt a v(Q) + ̃a B(Q) ΔW
"""
@generated function compute_stages!(x::Vector{ST}, Q0::Vector{Vector{ST}}, Q1::Vector{Matrix{ST}}, V::Vector{Vector{ST}},
                                                   B::Vector{Matrix{ST}}, Y0::Vector{Vector{ST}}, Y1::Vector{Matrix{ST}},
                                                   params::ParametersWIRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    tQ::Vector{ST} = zeros(ST,D)
    tV::Vector{ST} = zeros(ST,D)
    tB::Vector{ST} = zeros(ST,D)    # This was tB::Matrix{ST} = zeros(ST,D,M) for SIRK, but here we calculate B column by column.
                                    # tB is actually superfluous, could just use tV; keeping tB for clarity though

    quote
        local tᵢ::TT

        @assert S == length(Q0) == length(Q1) == length(V) == length(B)

        # copy x to Y0 and Y1, and calculate Q0 and Q1
        for i in eachindex(Q0)

            @assert D == length(Q0[i]) == size(Q1[i],1) == length(V[i]) == size(B[i],1)
            @assert M == size(Q1[i],2) == size(B[i],2)

            for k in eachindex(Q0[i])
                Y0[i][k] = x[D*(i-1)+k]
                Q0[i][k] = params.q[k] + Y0[i][k]
            end
        end

        for i in eachindex(Q1)
            for l in 1:size(Q1[i],2)
                for k in 1:size(Q1[i],1)
                    Y1[i][k,l] = x[ D*S + D*M*(i-1) + D*(l-1) + k ]
                    Q1[i][k,l] = params.q[k] + Y1[i][k,l]
                end
            end
        end

        # compute V = v(Q0) and B=B(Q1)
        for i in 1:S
            # time point with the coefficient c0
            tᵢ = params.t + params.Δt * params.tab.qdrift0.c[i]
            # calculates v(t,tQ) and assigns to the i-th column of V
            params.equ.v(tᵢ, Q0[i], V[i])

            # time point with the coefficient c1
            tᵢ = params.t + params.Δt * params.tab.qdrift1.c[i]

            for l=1:M
                # copies Q1[i][:,l] to the vector tQ
                simd_copy_xy_first!($tQ, Q1[i], l)
                # calculates the l-th column of B(t,tQ) and assigns to the vector tB
                params.equ.B(tᵢ, $tQ, $tB, col=l)
                # copies the vector tB to B[i][:,l]
                simd_copy_yx_first!($tB, B[i], l)
            end
        end
    end
end

"Compute stages of weak implicit Runge-Kutta methods."
@generated function function_stages!(x::Vector{ST}, b::Vector{ST}, params::ParametersWIRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    cache = IntegratorCacheWIRK{ST}(D, M, S)

    quote
        compute_stages!(x, $cache.Q0, $cache.Q1, $cache.V, $cache.B, $cache.Y0, $cache.Y1, params)

        local y1::ST
        local y2::ST

        # compute b = - (Y-AV)
        # the terms corresponding to Y0
        for i in 1:S
            for k in 1:D
                y1 = 0
                y2 = 0
                for j in 1:S
                    y1 += params.tab.qdrift0.a[i,j] * $cache.V[j][k] * params.Δt
                    y2 += params.tab.qdrift0.â[i,j] * $cache.V[j][k] * params.Δt
                    for l in 1:M
                        y1 += params.tab.qdiff0.a[i,j] * $cache.B[j][k,l] * params.ΔW[l]
                        y2 += params.tab.qdiff0.â[i,j] * $cache.B[j][k,l] * params.ΔW[l]
                    end
                end
                b[D*(i-1)+k] = - $cache.Y0[i][k] + (y1 + y2)
            end
        end

        # compute b = - (Y-AV)
        # the terms corresponding to Y1
        for i in 1:S
            for l in 1:M
                for k in 1:D
                    y1 = 0
                    y2 = 0
                    for j in 1:S
                        # The drift terms
                        y1 += params.tab.qdrift1.a[i,j] * $cache.V[j][k] * params.Δt
                        y2 += params.tab.qdrift1.â[i,j] * $cache.V[j][k] * params.Δt

                        # The noise terms, calculated either with the B1 or B3 tableau
                        for noise_idx in 1:M
                            if noise_idx==l
                                y1 += params.tab.qdiff1.a[i,j] * $cache.B[j][k,noise_idx] * params.ΔW[noise_idx]
                                y2 += params.tab.qdiff1.â[i,j] * $cache.B[j][k,noise_idx] * params.ΔW[noise_idx]
                            else
                                y1 += params.tab.qdiff3.a[i,j] * $cache.B[j][k,noise_idx] * params.ΔW[noise_idx]
                                y2 += params.tab.qdiff3.â[i,j] * $cache.B[j][k,noise_idx] * params.ΔW[noise_idx]
                            end
                        end
                    end
                    b[D*S + D*M*(i-1) + D*(l-1) + k] = - $cache.Y1[i][k,l] + (y1 + y2)
                end
            end
        end
    end
end

"""
This function computes initial guesses for Y and assigns them to int.solver.x
For WIRK we are NOT IMPLEMENTING an InitialGuess.
"""
function initial_guess!(int::IntegratorWIRK{DT,TT}, sol::AtomicSolutionSDE{DT,TT}) where {DT,TT}
    # The simplest initial guess for Y is 0
    int.solver.x .= 0

    # Using an explicit integrator to predict the next step's value (like in SIRK)
    # does not seem to be a good idea here, because the integrators are convergent
    # in the weak sense only, and there is no guarantee that the explicit integrator
    # will produce anything close to the desired solution...
end

"""
Integrate SDE with a stochastic implicit Runge-Kutta integrator.
  Integrating the m-th sample path
"""
function integrate_step!(int::IntegratorWIRK{DT,TT}, sol::AtomicSolutionSDE{DT,TT}) where {DT,TT}
    # update nonlinear solver parameters from cache
    update_params!(int, sol)

    # compute initial guess and assign to int.solver.x
    initial_guess!(int, sol)

    # reset cache
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute the drift vector field and the diffusion matrix at internal stages
    compute_stages!(int.solver.x, int.cache.Q0, int.cache.Q1, int.cache.V, int.cache.B, int.cache.Y0, int.cache.Y1, int.params)

    # compute final update (same update function as for SIRK)
    update_solution!(sol.q, int.cache.V, int.cache.B, int.params.tab.qdrift0.b, int.params.tab.qdrift0.b̂, int.params.tab.qdiff0.b, int.params.tab.qdiff0.b̂, int.params.Δt, int.params.ΔW, int.cache.Δy)
end


using JLD, GibbsTypePriors, DataFramesMeta, RCall, Optim, OptimTestProblems, StatsBase

include("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/scripts_for_figures/common_functions.jl")

R"
StickBreakingPY <- function(alpha,sigma, N) {
  a<- rep(1-sigma,N-1)
  b<- alpha+ sigma*(1:N-1)
  V <- rbeta(N-1 , a, b)
  p    <- vector('numeric',length=N)
  p[1] <- V[1]
  for(l in 2:(N - 1))p[l] <- prod(1 - V[1:(l - 1)])*V[l]
  p[N] <- prod(1 - V)
  p
  return(p)
}

Prior_on_K_SB_PY<- function(alpha, sigma, N_s=100,N_tr=100, runs=10^4 ){
  array_nc<-c()
  i=1
  while (i<=runs){
    weights_PY<- StickBreakingPY(alpha,sigma,N_tr)
    c <- sample(1:N_tr,size=N_s, replace=TRUE, prob=weights_PY )
    n_c<- length(unique(c))
    array_nc[i]<- n_c
    i=i+1
  }
  p_k = tibble(k=as.numeric(names(table(array_nc))),
               p_k=as.vector(table(array_nc))/sum(table(array_nc)))
  p_zeros= tibble(k=(1:N_s)[!1:N_s%in%p_k$k],
                  p_k=rep(0,length((1:N_s)[!1:N_s%in%p_k$k])))
  pk_padded= rbind(p_k, p_zeros)%>% arrange(k)
  E_k=sum((1:N_s) *(pk_padded$p_k))
  V_k = sum((((1:N_s) - E_k)^2) *(pk_padded$p_k))
  return(list(pk=pk_padded$p_k, Ek= E_k,Vk=V_k ))
}
"

function EV_of_number_of_clusters_PY(n,par_vec)
    #pk_py = Pkn_NGG.(1:n, n, par_vec[1], par_vec[2])
    E =  GibbsTypePriors.E_PY(n, par_vec[1], par_vec[2]) |> Float64
    V  = GibbsTypePriors.V_PY(n, par_vec[1], par_vec[2])  |> Float64
    return E,V
end

sigma = collect(range(0.05,0.99, length=10))
alpha = collect(exp.(range(log(1), log(200), length =10)))
grid = collect(Iterators.product(alpha, sigma))
n=100
grid_vec = vec(grid)
H= 250

save("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_7_5.jld", "grid_vec" ,grid_vec)

## PY exact

EV_PY = EV_of_number_of_clusters_PY.(n, grid_vec)

DF_PY= DataFrame(Exp = first.(EV_PY),Var =last.(EV_PY), beta = first.(grid_vec), sigma = last.(grid_vec))
DF_PY.std = sqrt.(DF_PY.Var)
DF_PY.beta_log = log.(DF_PY.beta)

DF_PY

save("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_7_1.jld", "DF_PY" ,DF_PY)

## PY multinomial process

function EV_of_number_of_clusters_PYM(n,H,par_vec)
    pk_pym = GibbsTypePriors.Pkn_PYM.(1:n, n, H,  par_vec[1], par_vec[2])
    E_pym =  pk_pym|> ar -> map(*, ar, 1:n) |> sum
    x = ((1:n).-E_pym).^2
    V_pym = pk_pym |> ar -> map(*, ar, x) |> sum
    return E_pym, V_pym
end

EV_PYM = EV_of_number_of_clusters_PYM.(n,H,grid_vec)

DF_PYM= DataFrame(Exp = first.(EV_PYM),Var =last.(EV_PYM), beta = first.(grid_vec), sigma = last.(grid_vec))
DF_PYM.std = sqrt.(DF_PYM.Var)
DF_PYM.beta_log = log.(DF_PYM.beta)

DF_PYM

save("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_7_2.jld", "DF_PYM" ,DF_PYM)

## PY predictive approximation

function EV_of_number_of_clusters_PY_approx(n,par_vec)
    pk_approx = GibbsTypePriors.Pkn_PY_pred_approx(n, par_vec[1], par_vec[2])
    E_approx =  pk_approx|> ar -> map(*, ar, 1:n) |> sum
    x = ((1:n).-E_approx).^2
    V_approx = pk_approx |> ar -> map(*, ar, x) |> sum
    return E_approx, V_approx
end

EV_PY_apr= EV_of_number_of_clusters_PY_approx.(n, grid_vec)

DF_PY_apr= DataFrame(Exp = first.(EV_PY_apr),Var =last.(EV_PY_apr), beta = first.(grid_vec), sigma = last.(grid_vec))
DF_PY_apr.std = sqrt.(DF_PY_apr.Var)
DF_PY_apr.beta_log = log.(DF_PY_apr.beta)

DF_PY_apr

save("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_7_3.jld", "DF_PY_apr" ,DF_PY_apr)

##PY stick-breaking approximation

function EV_of_number_of_clusters_PY_SB(n,H, par_vec)
    R"SB_PY = Prior_on_K_SB_PY($(par_vec[1]),$(par_vec[2]), $n, $H, runs= 2*10^2)
    E_py_sb = SB_PY$Ek
    V_py_sb = SB_PY$Vk"
    @rget E_py_sb
    @rget V_py_sb
    return E_py_sb, V_py_sb
end

EV_PY_sb= EV_of_number_of_clusters_PY_SB.(n,H,  grid_vec)

DF_PY_sb= DataFrame(Exp = first.(EV_PY_sb),Var =last.(EV_PY_sb), beta = first.(grid_vec), sigma = last.(grid_vec))
DF_PY_sb.std = sqrt.(DF_PY_sb.Var)
DF_PY_sb.beta_log = log.(DF_PY_sb.beta)

DF_PY_sb

save("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_7_4.jld", "DF_PY_sb" ,DF_PY_sb)


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

# http://ahojukka5.github.io/posts/finite-element-solution-for-one-element-problem/

using JuliaFEM
using JuliaFEM.Testing

@testset "test 2d linear elasticity local matrices" begin
    element = Element(Quad4, [1, 2, 3, 4])
    X = Dict{Int64, Vector{Float64}}(
        1 => [0.0, 0.0],
        2 => [1.0, 0.0],
        3 => [1.0, 1.0],
        4 => [0.0, 1.0])
    u = Dict{Int64, Vector{Float64}}(
        1 => [0.0, 0.0],
        2 => [0.0, 0.0],
        3 => [0.0, 0.0],
        4 => [0.0, 0.0])
    update!(element, "geometry", X)
    update!(element, "displacement", u)
    update!(element, "youngs modulus" =>  288.0, "poissons ratio" => 1/3)
    update!(element, "displacement load", DCTI([4.0, 8.0]))

    problem = Problem(Elasticity, "[0x1] x [0x1] block", 2)
    problem.properties.formulation = :plane_stress
    assemble!(problem, element)
    K = full(problem.assembly.K)
    f = vec(full(problem.assembly.f))

    K_expected = [
        144  54 -90   0 -72 -54  18   0
         54 144   0  18 -54 -72   0 -90
        -90   0 144 -54  18   0 -72  54
          0  18 -54 144   0 -90  54 -72
        -72 -54  18   0 144  54 -90   0
        -54 -72   0 -90  54 144   0  18
         18   0 -72  54 -90   0 144 -54
          0 -90  54 -72   0  18 -54 144]

    f_expected = [1, 2, 1, 2, 1, 2, 1, 2]

    @test isapprox(K, K_expected)
    @test isapprox(f, f_expected)
end


isconvex(tri::Triangle2D) = true
function isconvex(poly::Polygon{N, 2}) where {N}
    # If each of the point triplets in the polygon is CCW oriented, it is convex
    vec₁ = poly[2] - poly[1]
    for i ∈ 1:N-2
        vec₂ = poly[i+2] - poly[i+1]
        vec₁ × vec₂ < 0 && return false
        vec₁ = vec₂
    end
    return 0 ≤ vec₁ × (poly[1] - poly[N])
end


"""
'triangle generalized_preferential_attachment_graph'
===========================================
Generate an instance of a triangle generalized preferential attachment graph which
follows the Eikmeier, Gleich description (https://arxiv.org/pdf/1904.12989.pdf).
This is an undirected graph that is generated as follows:
- Start with an empty graph
- Add n vertices where at each time step one of three events occurs:
A new node is added with probability p, along with two edges to existing nodes;
Three new nodes with two edges between them is added with probability 1 - p
Every m steps, coalesce the nodes added in time step 1

Functions
---------
The following functions are synonyms
- 'triangle_generalized_preferential_attachment_graph_v2'
- 'tgpa_graph_v2'
and
- 'triangle_generalized_preferential_attachment_edges_v2!'
- 'tgpa_edges_v2!'
The computational functions are
- 'tgpa_graph_v2(n,p,k0,m)' Generate a TGPA graph with n total nodes.
    This returns a MatrixNetwork type. Note that this process allows self loops.
The edge functions are
-   'tgpa_edges_v2!(n,p,m,edges,n0)' Add new edges to an existing set, by taking
    n0 time steps. Edges are added in one of two ways: From a new node to
    an existing node with probability p, between three new nodes with probability 1-p

Input
-----
- 'n': the number of nodes in the final graph.
- 'p': The probability of a node event, p must be a constant.
- 'k0': The number of nodes in the starting clique.
- 'm': The frequency of combining nodes
- 'edges': A list of edges to be manipulated in the process of generating
          new edges.
Output
------
- A matrix network type for the triangle generalized preferential attachment graph.
- 'edges': An updated list of edges.
Example:
triangle_generalized_preferential_attachment_graph_v2(100,1/3,1/2)
"""
:triangle_generalized_preferential_attachment_graph_v2, :tgpa_graph_v2,
:triangle_generalized_preferential_attachment_edges_v2!, :tgpa_edges_v2!

using Printf
using MatrixNetworks
using Distributions

function triangle_generalized_preferential_attachment_graph_v2(
    n::Int,p::Float64,k0::Int,m::Int)
    k0 >= 0 || throw(ArgumentError(@sprintf("k0=%i must be >= 0",n)))
    n >= k0 || throw(ArgumentError(@sprintf("n=%i must be >= k0=%i",n,k0)))
    0<=p<=1 || throw(ArgumentError(@sprintf("p=%0.3f must be between 0 and 1",p)))
    m >= 1 || throw(ArgumentError(@sprintf("m=%i must be >= 1",m)))
    m <= max(n,1) || throw(ArgumentError(@sprintf("m=%i must be <= n=%i",m,n)))
    edges = Vector{Tuple{Int,Int}}()
    # add the clique
    for i=1:k0
        for j=1:i-1
            push!(edges, (i,j))
            push!(edges, (j,i))
        end
    end
    return MatrixNetwork(triangle_generalized_preferential_attachment_edges_v2!(n,p,m,edges,k0),n)
end

function _edges_to_neighbors(edges, k0)
    neighbors = Vector{Vector{Int}}(undef, k0)
    for i=1:k0
        neighbors[i] = Vector{Int}()
    end
    for e in edges
        push!(neighbors[e[1]], e[2])
    end
    return neighbors
end

function triangle_generalized_preferential_attachment_edges_v2!(
    n::Int,p::Float64,m::Int,edges::Vector{Tuple{Int,Int}},k0::Int)
    i = k0 #number of nodes added

    # for this code, we need to have random neighbor access, so we'll keep
    neighbors = _edges_to_neighbors(edges, k0)

    while i < n
        #determine how many component events we should have
        #this says for every m steps (every m new nodes) we should have this many new components
        num_component_events = rand(Binomial(m,1-p),1)[1]

        #new component events
        for j = 1:num_component_events
            if i+3 <= n #only allow this step if there is room for three more nodes
                push!(edges, (i+1, i+2)) #add wedge to edge list
                push!(edges, (i+2, i+1))
                push!(edges, (i+1, i+3))
                push!(edges, (i+3, i+1))
                push!(neighbors,Vector{Int}()) ##update neighbor list
                push!(neighbors,Vector{Int}())
                push!(neighbors,Vector{Int}())
                push!(neighbors[i+1],i+2)
                push!(neighbors[i+1],i+3)
                push!(neighbors[i+2],i+1)
                push!(neighbors[i+3],i+1)
                i = i+3
            end
        end
        #new node events
        if i+1 <=n
            push!(neighbors,Vector{Int}())
            for j = 1:m
                if length(edges) > 0
                    v = rand(edges)[1]
                    v_neighbors = neighbors[v]
                else
                    v = i+1
                    v_neighbors = [];
                end
                #edge between new node and v
                push!(edges, (i+1, v))
                push!(edges, (v, i+1))
                #update neighbor vector
                push!(neighbors[i+1],v)
                push!(neighbors[v],i+1)

                #add another edge between new node and a neighbor of v
                if length(v_neighbors) > 0
                    v2 = rand(v_neighbors)
                    push!(edges, (i+1, v2))
                    push!(edges, (v2, i+1))
                    push!(neighbors[i+1],v2)
                    push!(neighbors[v2],i+1)
                end
            end
            i = i+1;
        end

    end
    return edges
end

tgpa_graph_v2 = triangle_generalized_preferential_attachment_graph_v2
tgpa_edges_v2! = triangle_generalized_preferential_attachment_edges_v2!


#start boundary conditions
println("calculating right-hand-side...\n")

F = zeros(dofs,1)

(nbc,ni,ix) = IDs(bc) #nbc has number of conditions
                   #ix has edges nodes ids for each condition
global bc,edge,elem,nodes,i,F
i=1

while true
    global i,ix,elem
    if i>nbc
        break
    end
    rows = [j for k=1:length(ni[i]) for j=1:length(elem[:,1]) if issubset(ix[i][k,1],elem[j,:]) &
                                                                 issubset(ix[i][k,2],elem[j,:]) &
                                                                 issubset(ix[i][k,3],elem[j,:])] #contains superset id element
    xedge = nodes[elem[rows,:],:][:,:,1]; yedge = nodes[elem[rows,:],:][:,:,2] #contains x,y of edges

    superelem = elem[rows,:]
    mapx = [sum(dims=2,superelem[k,:] .== ix[i][k,:]') for k=1:length(superelem[:,1])]#zero the non edge indexes of element
    for k=1:length(mapx)
        g(x) = f(x)[i]
        a = BC(g,xedge[k,:]',yedge[k,:]',mapx[k])
        id = [ix[i][k,:]' (ix[i][k,:]' .+ Int64(0.5*dofs))]
        id = sort(id,dims=2)

        idlocal = [mapx[k]' mapx[k]'] #has the global degrees of freedoms
        idlocal = [idlocal[i] .==1 for i=1:12]
        F[id] .= F[id] .+ a[idlocal]'
    end
    i = i+1
end


using ContactStructures
using Plots

f = (x,y,z) -> cos(2π*x)
g = (x,y,z) -> sin(2π*x)
giroux_form = f * dz + g * dy

plotly()
plot_along_xy_grid(0:0.125:1,0:0.125:1,0,giroux_form;  markersize = 0,ε = 0.05)
plot!(xlims=(-0.1, 1.1), ylims=(-0.1,1.1), zlims=(-0.6,0.6),legend=false)


using Printf

function parse_map(s)
    lines = split(s, "\n")
    map_arr = zeros(Bool, length(lines), length(lines[1]))
    for (row_idx, l) in enumerate(lines)
        for (col_idx, c) in enumerate(l)
            map_arr[row_idx, col_idx] = (c == '#')
        end
    end
    return map_arr
end

function count_trees_on_slope(map_arr, right, down)
    tree_counter = 0
    x, y = 1, 1
    n_rows, n_cols = size(map_arr)
    while x <= n_rows
        if map_arr[x, y]
            tree_counter += 1
        end
        x += down
        y = (y + right) % n_cols
        if y == 0
            y = n_cols
        end
    end
    return tree_counter
end

function run(s)
    # Your code here
    map_arr = parse_map(s)
    slopes = [(1, 1), (3, 1), (5, 1), (7, 1), (1, 2)]
    results = map(x -> count_trees_on_slope(map_arr, x...), slopes)
    return reduce(*, results)
end

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

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


# This code presents an example to illustrate how NLboxsolve.jl can be used

using NLboxsolve

function test_nlboxsolve_outplace()

    function test8(x)

        f = similar(x)

        f[1] = x[1]^2 - 1.0
        for i = 2:length(x)-1
            f[i] = x[i-1] - x[i]^3
        end
        f[end] = x[end-1] - x[end]

        return f

    end

    # Solve test8

    n = 20
    x0 = ones(n)*1.5
    l = ones(n)*0.5
    l[1] = 0.8
    u = ones(n)*2.0

    soln_newton     = constrained_newton(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test8,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(test8,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(test8,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

function test_nlboxsolve_inplace()

    function test8(f,x)

        f[1] = x[1]^2 - 1.0
        for i = 2:length(x)-1
            f[i] = x[i-1] - x[i]^3
        end
        f[end] = x[end-1] - x[end]

    end

    # Solve test8

    n = 20
    x0 = ones(n)*1.5
    l = ones(n)*0.5
    l[1] = 0.8
    u = ones(n)*2.0

    soln_newton     = constrained_newton(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(test8,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test8,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test8,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(test8,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(test8,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

test_nlboxsolve_outplace()
test_nlboxsolve_inplace()


export NormalScore

"""
    struct NormalScore <: Score{Float64}

A score defined by a normal density given the mean and sd of the score.
"""
struct NormalScore <: Score{Float64}
    mean :: Float64
    sd :: Float64
end

@impl begin
    struct NormalScoreGetScore end
    function get_score(sf::NormalScore, i::Float64)::AbstractFloat
        return normal_density(i, sf.mean, sf.sd)
    end
end


using MLJ
import RDatasets: dataset
using PrettyPrinting
using Random

Random.seed!(3203)
X = randn(20, 2)
y = vcat(-ones(10), ones(10))

using PyPlot
figure(figsize=(8,6))

ym1 = y .== -1
ym2 = .!ym1
plot(X[ym1, 1], X[ym1, 2], ls="none", marker="o")
plot(X[ym2, 1], X[ym2, 2], ls="none", marker="x")

X = MLJ.table(X)
y = categorical(y);

SVC = @load SVC pkg=LIBSVM

svc_mdl = SVC()
svc = machine(svc_mdl, X, y)

fit!(svc);

ypred = MLJ.predict(svc, X)
misclassification_rate(ypred, y)

rc = range(svc_mdl, :cost, lower=0.1, upper=5)
tm = TunedModel(model=svc_mdl, ranges=[rc], tuning=Grid(resolution=10),
                resampling=CV(nfolds=3, rng=33), measure=misclassification_rate)
mtm = machine(tm, X, y)

fit!(mtm)

ypred = MLJ.predict(mtm, X)
misclassification_rate(ypred, y)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


using GLMakie, Random
Random.seed!(13)
x = -6:0.5:6
y = -6:0.5:6
z = 6exp.( -(x.^2 .+ y' .^ 2)./4)

box = Rect3(Point3f(-0.5), Vec3f(1))
colors = to_colormap(:linear_worb_100_25_c53_n256)
n = length(colors)
g(x) = x^(1/10)
alphas = [g(x) for x in LinRange(0,1,n)]
cmap_alpha = RGBAf0.(colors, alphas)

set_theme!(theme_dark())
fig, ax, = meshscatter(x, y, z; marker=box, markersize = 0.5,
    color = vec(z), colormap = cmap_alpha, colorrange = (0,6),
    axis = (; type = Axis3, aspect = :data, azimuth = 7.3,
    elevation = 0.189, perspectiveness = 0.5),
    figure = (;resolution =(1200,800)))
meshscatter!(ax, x .+ 7, y, z./2; markersize = 0.25, color = vec(z./2),
    colormap = cmap_alpha, colorrange = (0, 6),
    ambient = Vec3f0(0.85, 0.85, 0.85), backlight = 1.5f0)
xlims!(-5.5,10)
ylims!(-5.5,5.5)
hidedecorations!(ax; grid = false)
hidespines!(ax)
save("meshscatter.png", fig)
display(fig)


mutable struct C2
    Wc1::AbstractArray{T,4} where T
    Wc2::AbstractArray{T,4} where T
    W1::AbstractMatrix
    b::AbstractVector
end

C2() = C2(rand(1,1,1,1), rand(1,1,1,1), rand(1,1), rand(1))

function predict(m, X::AbstractArray)
    c1 = conv2d(X, m.Wc1)
    p1 = pool(c1)
    c2 = conv2d(p1, m.Wc2)
    p2 = pool(c2)
    cc = reshape(p2, 4*4*6, size(p1, 4))
    res = sigmoid.(m.W1 * cc .+ m.b)
    return res
end

function loss(m, X::AbstractArray, y::AbstractMatrix)
    ŷ = predict(m, X)
    # cost = sum((y .- ŷ) .^ 2)
    cost = pytorch_cross_entropy(ŷ, y)
    return cost
end

function partial_fit!(m, X::AbstractArray, y::AbstractMatrix; mem=Dict())
    # opt = ModelOptimizer(m, Momentum(η=0.001, γ=0.9))
    opt = ModelOptimizer(m, Adam(α=0.001))
    cost, dm, dX, dy = xgrad(loss; mem=mem, m=m, X=X, y=y)
    update_params!(opt, m, dm)
    return cost
end

function fit!(m, X::AbstractArray, Y::AbstractMatrix; cuda=false, n_epochs=100)
    mem = Dict()
    if cuda
        m = to_cuda(m)
    end
    @time for epoch=1:n_epochs
        print("Epoch: $epoch: ")
        cost = 0
        for (x, y) in batchview((X, Y))
            x, y = map(copy, (x, y))
            if cuda
                x, y = map(to_cuda, (x, y))
            end
            cost = partial_fit!(m, x, y)
        end
        println(cost)
    end
    return m
end

function predict_label(m, X)
    findmax(predict(m, X))[2]
end


module Lanczos

function lanczos(A,v1,steps)

    size(A)[1] == size(A)[2] || error("matrix must be square")

    ∘ = dot

    v = zeros(size(v1))
    r = v1
    b = 1

    T = zeros(Float64,steps+1,steps+1)      # Lanczos matrix
    K = zeros(Float64,size(A)[1],steps+1)   # vectors in Krylov space

    # First round
    vold = v
    v = r/b
    K[:,1] = v
    Av = A*v
    a = v∘Av
    T[1,1] = a
    r = Av - v*a - vold*b

    # Loop
    for i in 1:steps
        b = norm(r)
        T[i,i+1] = T[i+1,i] = b
        vold = v
        v = r/b                     # next v
        K[:,i+1] = v                # store vector
        Av = A*v
        a = v∘Av
        T[i+1,i+1] = a
        r = Av - v*a - vold*b       # orthogonalize
    end

    return T, K

end

lanczos(A,v1) = lanczos(A,v1,size(A)[1])

end


immutable XCubed <: Function end
@inline (::XCubed)(x) = x^3
derivative(::XCubed) = dXCubed()

immutable dXCubed <: Function end
@inline (::dXCubed)(x) = 3x^2
derivative(::dXCubed) = ddXCubed()

immutable ddXCubed <: Function end
@inline (::ddXCubed)(x) = 6x

immutable XSquaredLogX <: Function end
@inline (::XSquaredLogX)(x) = x^2 * log(x)
derivative(::XSquaredLogX) = dXSquaredLogX()

immutable dXSquaredLogX <: Function end
@inline (::dXSquaredLogX)(x) = 2x*log(x) + x
derivative(::dXSquaredLogX) = ddXSquaredLogX()

immutable ddXSquaredLogX <: Function end
@inline (::ddXSquaredLogX)(x) = 2*log(x) + 3


module Density

# This implements different methods to compute density
#
# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, Unitful, LaTeXStrings
using ..Units
using ..PhaseDiagrams
using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct
using ..MaterialParameters: No_MaterialParam, MaterialParamsInfo
import Base.show, GeoParams.param_info

include("../Utils.jl")
include("../Computations.jl")

abstract type AbstractDensity{T} <: AbstractMaterialParam end

export  compute_density,        # calculation routines
        compute_density!,       # in place calculation
        param_info,             # info about the parameters
        AbstractDensity,
        ConstantDensity,        # constant
        PT_Density,             # P & T dependent density
        Compressible_Density    # Compressible density

# Define "empty" computational routines in case nothing is defined
compute_density!(rho::_T,s::No_MaterialParam{_T}; P::_T=zero(_T),T::_T=zero(_T)) where {_T} = zero(_T)
compute_density(s::No_MaterialParam{_T}; P::_T=zero(_T),T::_T=zero(_T)) where {_T} = zero(_T)

# Constant Density -------------------------------------------------------
"""
    ConstantDensity(ρ=2900kg/m^3)

Set a constant density:
```math
    \\rho  = cst
```
where ``\\rho`` is the density [``kg/m^3``].
"""
@with_kw_noshow struct ConstantDensity{_T,U}   <: AbstractDensity{_T}
    ρ::GeoUnit{_T,U}        =   2900.0kg/m^3                # density
end
ConstantDensity(args...) = ConstantDensity(convert.(GeoUnit,args)...)

function param_info(s::ConstantDensity) # info about the struct
    return MaterialParamsInfo(Equation = L"\rho = cst")
end

# Calculation routines
function compute_density(s::ConstantDensity{_T}, P::Quantity,T::Quantity) where _T
    @unpack_units ρ   = s
    return ρ
end

function compute_density(s::ConstantDensity{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack_val ρ   = s
    return ρ
end

function compute_density!(rho::AbstractArray{_T}, s::ConstantDensity{_T},P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack_val ρ   = s
    rho[:] .= ρ
    return nothing
end

# Print info
function show(io::IO, g::ConstantDensity)
    print(io, "Constant density: ρ=$(UnitValue(g.ρ))")
end
#-------------------------------------------------------------------------

# Pressure & Temperature dependent density -------------------------------
"""
    PT_Density(ρ0=2900kg/m^3, α=3e-5/K, β=1e-9/Pa, T0=0C, P=0MPa)

Set a pressure and temperature-dependent density:
```math
    \\rho  = \\rho_0 (1.0 - \\alpha (T-T_0) + \\beta  (P-P_0) )
```
where ``\\rho_0`` is the density [``kg/m^3``] at reference temperature ``T_0`` and pressure ``P_0``,
``\\alpha`` is the temperature dependence of density and ``\\beta`` the pressure dependence.
"""
@with_kw_noshow struct PT_Density{_T,U1,U2,U3,U4,U5} <: AbstractDensity{_T}
    ρ0 ::GeoUnit{_T,U1}      =   2900.0kg/m^3                # density
    α  ::GeoUnit{_T,U2}      =   3e-5/K                      # T-dependence of density
    β  ::GeoUnit{_T,U3}      =   1e-9/Pa                     # P-dependence of density
    T0 ::GeoUnit{_T,U4}      =   0.0C                        # Reference temperature
    P0 ::GeoUnit{_T,U5}      =   0.0MPa                      # Reference pressure
end
PT_Density(args...) = PT_Density(convert.(GeoUnit,args)...)

function param_info(s::PT_Density)  # info
    return MaterialParamsInfo(Equation = L"\rho = \rho_0(1.0-\alpha (T-T_0) + \beta (P-P_0)" )
end

# Calculation routine in case units are provided
function compute_density(s::PT_Density{_T},P::Quantity,T::Quantity) where _T
    @unpack_units ρ0,α,β,P0, T0   = s

    ρ = ρ0*(1.0 - α*(T - T0) + β*(P - P0) )

    return ρ
end

# All other numbers
function compute_density(s::PT_Density{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack_val ρ0,α,β,P0, T0   = s

    ρ = ρ0*(1.0 - α*(T - T0) + β*(P - P0) )

    return ρ
end

function compute_density!(ρ::AbstractArray, s::PT_Density{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack ρ0,α,β,P0, T0   = s

    ρ .= ρ0*(1.0 - α*(T-T0) + β*(P-P0) )

    return nothing
end

# Print info
function show(io::IO, g::PT_Density)
    print(io, "P/T-dependent density: ρ0=$(UnitValue(g.ρ0)), α=$(UnitValue(g.α)), β=$(UnitValue(g.β)), T0=$(UnitValue(g.T0)), P0=$(UnitValue(g.P0))")
end
#-------------------------------------------------------------------------

# Pressure-dependent density -------------------------------
"""
    Compressible_Density(ρ0=2900kg/m^3, β=1e-9/Pa, P₀=0MPa)

Set a pressure-dependent density:
```math
    \\rho  = \\rho_0 \\exp(β*(P - P\\_0))
```
where ``\\rho_0`` is the density [``kg/m^3``] at reference pressure ``P_0`` and ``\\beta`` the pressure dependence.
"""
@with_kw_noshow struct Compressible_Density{_T,U1,U2,U3} <: AbstractDensity{_T}
    ρ0::GeoUnit{_T,U1}     =   2900.0kg/m^3                # density
    β ::GeoUnit{_T,U2}     =   1e-9/Pa                     # P-dependence of density
    P0::GeoUnit{_T,U3}     =   0.0MPa                      # Reference pressure
end
Compressible_Density(args...) = Compressible_Density(convert.(GeoUnit,args)...)

function param_info(s::Compressible_Density) # info about the struct
    return MaterialParamsInfo(Equation = L"\rho = \rho_0\exp(\beta*(P-P_0))"     )
end

function compute_density(s::Compressible_Density{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack_val ρ0,β,P0   = s

    ρ = ρ0*exp(β*(P - P0) )

    return ρ
end

function compute_density(s::Compressible_Density{_T}, P::Quantity, T::Quantity) where _T
    @unpack ρ0,β,P0   = s

    ρ = ρ0*exp(β*(P - P0) )

    return ρ
end

function compute_density!(ρ::_T, s::Compressible_Density{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack ρ0,β,P0   = s

    return ρ0*exp( β*(P-P0) )
end

function compute_density!(ρ::AbstractArray, s::Compressible_Density{_T}, P::_T=zero(_T),T::_T=zero(_T)) where _T
    @unpack ρ0,β,P0   = s

    ρ .= ρ0*exp( β*(P-P0) )

    return nothing
end

# Print info
function show(io::IO, g::Compressible_Density)
    print(io, "Compressible density: ρ0=$(UnitValue(g.ρ0)), β=$(UnitValue(g.β)), P0=$(UnitValue(g.P0))")
end
#-------------------------------------------------------------------------

#-------------------------------------------------------------------------
# Phase diagrams
function param_info(s::PhaseDiagram_LookupTable) # info about the struct
    return MaterialParamsInfo(Equation = L"\rho = f_{PhaseDiagram}(T,P))" )
end

"""
    compute_density(P,T, s::PhaseDiagram_LookupTable)
Interpolates density as a function of `T,P` from a lookup table
"""
function compute_density(s::PhaseDiagram_LookupTable, P::_T=zero(_T),T::_T=zero(_T)) where _T
    return s.Rho(T,P)
end

"""
    compute_density!(rho::AbstractArray{<:AbstractFloat}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, s::PhaseDiagram_LookupTable)
In-place computation of density as a function of `T,P`, in case we are using a lookup table.
"""
function compute_density!(rho::AbstractArray{_T}, s::PhaseDiagram_LookupTable, P::AbstractArray{_T}=[zero(_T)],T::AbstractArray{_T}=[zero(_T)]) where _T
    rho[:] = s.Rho.(T,P)
    return nothing
end

#------------------------------------------------------------------------------------------------------------------#
# Computational routines needed for computations with the MaterialParams structure

# This assumes that density always has a single parameter. If that is not the case, we will have to extend this (to be done)
function compute_density(s::AbstractMaterialParamsStruct, P::_T=zero(_T),T::_T=zero(_T)) where {_T}
    return compute_density(s.Density[1], P,T)
end

# these routines may come handy when we have >1 field for Density
#function compute_density(s::NTuple{N,AbstractDensity{_T}}, P::_T=zero(_T), T::_T=zero(_T)) where {N,_T}
#    compute_density.(s,P,T)
#end

#now with Tuple of Tuples
#function compute_density!(rho::Vector{NTuple{M,_T}}, P::Number, T::Number, MatParam::NTuple{N, NTuple{M, AbstractMaterialParamsStruct}}) where {N,M,_T}
#    rho .= map(x->compute_density(P,T,x), MatParam)
#end

#-------------------------------------------------------------------------------------------------------------

"""
    compute_density!(rho::AbstractArray{_T, ndim}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, Phases::AbstractArray{_I, ndim}; P=nothing, T=nothing) where {ndim,N,_T,_I<:Integer}

In-place computation of density `rho` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `Phase` of every point is specified as an Integer in the `Phases` array.
# Example
```julia
julia> MatParam = (SetMaterialParams(Name="Mantle", Phase=1,
                        CreepLaws= (PowerlawViscous(), LinearViscous(η=1e23Pa*s)),
                        Density   = PT_Density()
                        ),
                    SetMaterialParams(Name="Crust", Phase=2,
                        CreepLaws= (PowerlawViscous(), LinearViscous(η=1e23Pas)),
                        Density   = ConstantDensity(ρ=2900kg/m^3))
                  );
julia> Phases = ones(Int64,400,400);
julia> Phases[:,20:end] .= 2
julia> rho     = zeros(size(Phases))
julia> T       =  ones(size(Phases))
julia> P       =  ones(size(Phases))*10
julia> compute_density!(rho, MatParam, Phases, P, T)
julia> rho
400×400 Matrix{Float64}:
2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
    ⋮                                            ⋮                                         ⋱     ⋮                                       ⋮
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  …  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
 2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91  2899.91     2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0  2900.0
```
The routine is made to minimize allocations:
```julia
julia> using BenchmarkTools
julia> @btime compute_density!(\$rho, \$MatParam, \$Phases, P=\$P, T=\$T)
    203.468 μs (0 allocations: 0 bytes)
```
_____________________________________________________________________________________________________________________________

    compute_density!(rho::AbstractArray{_T, N}, MatParam::NTuple{K,AbstractMaterialParamsStruct}, PhaseRatios::AbstractArray{_T, M}, P=nothing, T=nothing)

In-place computation of density `rho` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `PhaseRatio` of every point is specified as an Integer in the `PhaseRatios` array, which has one dimension more than the data arrays (and has a phase fraction between 0-1)
"""
compute_density!(args...) = compute_param!(compute_density,args...) #Multiple dispatch to rest of routines found in Computations.jl
compute_density(args...) = compute_param(compute_density,args...)

#=
function compute_density!(rho::AbstractArray{_T, ndim}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, Phases::AbstractArray{_I, ndim}, P=nothing, T=nothing) where {ndim,N,_T,_I<:Integer}

    Phase_tup = ntuple(i->MatParam[i].Phase, Val(N))

    # check if we compute with P/T or use default values
    isnothing(T) ? compute_T = false : compute_T = true
    isnothing(P) ? compute_P = false : compute_P = true

    Tval = zero(_T)
    Pval = zero(_T)

    @inbounds for I in eachindex(Phases)
        phase   = find_ind(Phase_tup, Phases[I])

        # Extract relevant value if requested
        if compute_T; Tval = T[I]; end
        if compute_P; Pval = P[I]; end

        # this computes density for ALL phases, which is a bit of an overkill as we only need density for a single phase
        rho_tup = compute_density(MatParam, Pval, Tval)
        rho[I]  = rho_tup[phase]
    end
end

function compute_density!(rho::AbstractArray{_T, N}, MatParam::NTuple{K,AbstractMaterialParamsStruct}, PhaseRatios::AbstractArray{_T, M}, P=nothing, T=nothing) where {_T<:AbstractFloat, N,M, K}
    if M!=(N+1)
        error("The PhaseRatios array should have one dimension more than the other arrays")
    end

    # check if we compute with P/T or use default values
    isnothing(T) ? compute_T = false : compute_T = true
    isnothing(P) ? compute_P = false : compute_P = true

    Tval = zero(_T)
    Pval = zero(_T)
    if compute_P
        Rindex = CartesianIndices(P)
    else
        Rindex = CartesianIndices(T)
    end

    @inbounds for I in Rindex
        frac    = view(PhaseRatios, Tuple(I)..., 1:K)    # fraction of each phase @ point I

        # Extract relevant value if requested
        if compute_T; Tval = T[I]; end
        if compute_P; Pval = P[I]; end

        # compute point-wise density:
        rho[I] = compute_density_times_frac(frac,MatParam, Pval, Tval)

    end

    return
end

# Multiplies density with the fraction of a phase
function compute_density_times_frac(PhaseRatios::AbstractArray{_T, 1}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, P::_T, T::_T) where {_T, N}

    value_tup = compute_density(MatParam, P, T)

    # Sum & multiply density with fraction
    val = zero(_T)
    @inbounds for j=1:N
        val += PhaseRatios[j]*value_tup[j]
    end

    return val
end
=#

end

# OBSOLETE routines; I leave them in for now (commented), in case we want to reuse pieces of that @ a later stage.
# much of what we have above is more efficient (apart when using the Phase Diagram lookup tables)
#=
# OLD function (which actually works better in case of phase diagrams as it can do computations at omce)
function compute_density!(rho::AbstractArray{<:AbstractFloat, N}, PhaseRatios::AbstractArray{<:AbstractFloat, M}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where {N,M}

    if M!=(N+1)
        error("The PhaseRatios array should have one dimension more than the other arrays")
    end

    rho1        =  zeros(size(rho))
    rho         .=  0.0;
    for i = 1:length(MatParam)
        Fraction    = selectdim(PhaseRatios,M,i);
        if (!isempty(MatParam[i].Density))
            ind         =   Fraction .> 0.0
            rho_local   =   view(rho1, ind )
            P_local     =   view(P   , ind )
            T_local     =   view(T   , ind )
            density = MatParam[i].Density[1]
            compute_density!(rho_local, P_local, T_local, density )
            rho[ind] .+= rho_local.*Fraction[ind]
        end

    end
end
# OBSOLETE?
function compute_density!(rho::AbstractArray{<:AbstractFloat, N}, Phases::AbstractArray{<:Integer, N}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where {N,_T}
    for i = 1:length(MatParam)
        if !isempty(MatParam[i].Density)
            # Create views into arrays (so we don't have to allocate)
            ind = Phases .== MatParam[i].Phase
            rho_local   =   view(rho, ind )
            P_local     =   view(P  , ind )
            T_local     =   view(T  , ind )
           # density::Union{AbstractDensity{_T},PhaseDiagram_LookupTable}     =   MatParam[i].Density[1]
            density    =   MatParam[i].Density[1]

            #@time compute_density!(rho_local, P_local, T_local, dens )
            compute_density!(rho_local, P_local, T_local, density )

        end

    end
end
# Many allocations
function compute_density!(rho::AbstractArray{_T, ndim}, Phases::AbstractArray{<:Integer, ndim}, P::AbstractArray{_T, ndim},T::AbstractArray{_T, ndim}, MatParam::NTuple{N,AbstractMaterialParamsStruct}) where {ndim,N,_T}
    # This is the general way to implement this.
    #  Remarks:
    #       1) It allocates a lot when used with a phase diagram
    #       2) We still have to check if this works with the way ParallelStencil operates

    # The phases in MatParam may be ordered differently (or start with zero)

    Phase_vec = [MatParam[i].Phase for i=1:N]
    Phase_ind = copy(Phases);
    for i=1:N
        Phase_ind[Phases.==Phase_vec[i]] .= i
    end

    rho_local = zeros(_T,N)
    for i in eachindex(Phase_ind)
        phase   =   Phase_ind[i]

        # This does not allocate but computes density for all phases simultaneously.
        #  That is a bit of an overkill, but alternative approaches appear to allocate
        compute_density!(rho_local,P[i],T[i], MatParam)
        rho[i]  =   rho_local[phase]
    end
    return
end
=#


using CompScienceMeshes
using BEAST

using Test

m = meshrectangle(1.0, 1.0, 0.5, 3)
nodes = skeleton(m,0)
srt_bnd_nodes = sort.(skeleton(boundary(m),0))
int_nodes = submesh(nodes) do node
    !(sort(node) in srt_bnd_nodes)
end

@test length(int_nodes) == 1

L0 = BEAST.lagrangec0d1(m, int_nodes)
@test numfunctions(L0) == 1
@test length(L0.fns[1]) == 6
@test length(m) == 8

Lx = BEAST.lagrangecxd0(m)
@test numfunctions(Lx) == 8

Id = BEAST.Identity()
Q1, st1 = BEAST.allocatestorage(Id, L0, Lx, Val{:bandedstorage}, BEAST.LongDelays{:ignore})
BEAST.assemble_local_matched!(Id, L0, Lx, st1)

Q2, st2 = BEAST.allocatestorage(Id, L0, Lx, Val{:bandedstorage}, BEAST.LongDelays{:ignore})
BEAST.assemble_local_mixed!(Id, L0, Lx, st2)

@test isapprox(Q1, Q2, atol=1e-8)

RT = raviartthomas(m)
BC = buffachristiansen(m)

Q1, st1 = BEAST.allocatestorage(Id, BC, RT, Val{:bandedstorage}, BEAST.LongDelays{:ignore})
BEAST.assemble_local_refines!(Id, BC, RT, st1)

Q2, st2 = BEAST.allocatestorage(Id, BC, RT, Val{:bandedstorage}, BEAST.LongDelays{:ignore})
BEAST.assemble_local_mixed!(Id, BC, RT, st2)
@test isapprox(Q1, Q2, atol=1e-8)


using PDMats
using Test

n = 4
m = 7

for T in [Float64, Float32]
    X = randn(T, n, n)
    Y = randn(T, m, m)
    A = X * X'
    B = Y * Y'
    AkB = kron(A, B)
    PDA = PDMat(A)
    PDB = PDMat(B)
    PDAkB1 = PDMat(AkB)
    PDAkB2 = kron(PDA, PDB)
    @test PDAkB1.dim == PDAkB2.dim
    @test PDAkB1.mat ≈ PDAkB2.mat
    @test PDAkB1.chol.L ≈ PDAkB2.chol.L
    @test PDAkB1.chol.U ≈ PDAkB2.chol.U
    @test Matrix(PDAkB2.chol) ≈ PDAkB2.mat
end


"""
STOCHASTIC GRADIENT HAMILTONIAN MONTE CARLO

The following codes explore the use of SGHMC for Bayesian Inference
"""
immutable SGHMC
  dU      ::Function
  dK      ::Function
  dKΣ     ::Array{Float64}
  C       ::Array{Float64}
  V       ::Array{Float64}
  init_est::Array{Float64}
  d       ::Int64
end

function mcmc(parameters::SGHMC;
  leapfrog_params::Dict{Symbol, Real} = Dict([:ɛ => .05, :τ => 20]),
  set_seed::Int64 = 123,
  r::Int64 = 1000)

  dU, dK, dKΣ, C, V, w, d = parameters.dU, parameters.dK, parameters.dKΣ, parameters.C, parameters.V, parameters.init_est, parameters.d
  ɛ, τ = leapfrog_params[:ɛ], leapfrog_params[:τ]

  if typeof(set_seed) == Int64
    srand(set_seed)
  end

  chain = zeros(r, d);
  chain[1, :] = w
  B = .5 * V * ɛ
  D = sqrt(2 * (C - B) * ɛ)
  if size(B) != size(C)
    error("C and V should have the same dimension.")
  else
    if sum(size(B)) > 1
      if det(B) > det(C)
        error("ɛ is too big. Consider decreasing it.")
      end
    else
      if det(B[1]) > det(C[1])
        error("ɛ is too big. Consider decreasing it.")
      end
    end
  end

  for i in 1:(r - 1)
    p = randn(d, 1)

    for j in 1:τ
      p = p - dU(w) * ɛ - C * inv(dKΣ) * p + D * randn(d, 1);
      w = w + dK(p) * ɛ;
    end

    chain[i + 1, :] = w
  end

  return chain
end


using DelimitedFiles
function to_phase(wfc_file, wfci_file, xDim, yDim)
    wfc = reshape(readdlm(wfc_file) + im*readdlm(wfci_file), xDim, yDim)
    hval = maximum(abs2.(wfc))
    return angle.(wfc).*(abs2.(wfc)/hval)
end

function to_density(wfc_file, wfci_file, xDim, yDim)
    wfc = reshape(readdlm(wfc_file) + im*readdlm(wfci_file), xDim, yDim)
    return abs2.(wfc)
end

wfc_file = "data_L10/wfc_L10"
wfci_file = "data_L10/wfci_L10"
xDim = 512
yDim = 512

writedlm("phase_L10.dat", to_phase(wfc_file, wfci_file, xDim, yDim))
writedlm("density_L10.dat", to_density(wfc_file, wfci_file, xDim, yDim))


struct Monomial{N}
    θ::SVector{N,Int}
end
Monomial(p::Int) = Monomial(SVector(p))
Monomial(p::Tuple) = Monomial(SVector(p))

dimension(m::Monomial{N}) where {N} = N

Base.exponent(m::Monomial) = m.θ

function (m::Monomial)(x::SVector)
    @assert length(x) == dimension(m)
    x .^ exponent(m) |> prod
end
(m::Monomial)(x)  = m(SVector(x))


using ReachabilityAnalysis: IntervalExpAlg

@testset "ASB07 algorithm: homogeneous case" begin
    # Example 1 from
    # "Reachability analysis of linear systems with uncertain parameters and inputs"

    # initial set
    X0 = BallInf([1.0, 1.0], 0.1)

    # linear ODE: x' = Ax
    A = IntervalMatrix([-1.0 ± 0.05 -4.0 ± 0.05;
                         4.0 ± 0.05 -1.0 ± 0.05])
    P_lin = @ivp(x' = Ax, x(0) ∈ X0)

    sol1 = solve(P_lin, tspan=(0.0, 1.0), alg=ASB07(δ=0.04));
    @test sol1.alg.recursive == Val(true) # default is recursive TODO add isrecursive ?
    @test dim(sol1) == 2
    @test isa(sol1.alg, ASB07)
    @test sol1.alg.δ == 0.04
    @test sol1.alg.max_order == 5
    @test typeof(sol1.alg.approx_model) == CorrectionHull{IntervalExpAlg}
    @test sol1.alg.approx_model.order == 10
    @test setrep(sol1) <: Zonotope
    @test setrep(sol1) == Zonotope{Float64,Array{Float64,1},Array{Float64,2}}

    # change some of the algorithm-specific parameters
    Calg = CorrectionHull(order=5,               # truncation order in Φ
                          exp=IntervalExpAlg(9)) # truncation order in F
    Ralg = ASB07(δ=0.01,
                 max_order=7, # maximum zonotope order
                 approx_model=Calg);

    sol2 = solve(P_lin, tspan=(0.0, 5.0), alg=Ralg)
    @test dim(sol2) == 2
    @test isa(sol2.alg, ASB07)
    @test sol2.alg.δ == 0.01
    @test sol2.alg.max_order == 7
    @test typeof(sol2.alg.approx_model) == CorrectionHull{IntervalExpAlg}
    @test sol2.alg.approx_model.order == 5
    @test sol2.alg.approx_model.exp.order == 9
    @test setrep(sol2) <: Zonotope
    @test setrep(sol2) == Zonotope{Float64,Array{Float64,1},Array{Float64,2}}

    # compare recursive option: recursive is more precise
    sol1_rec = solve(P_lin, tspan=(0.0, 1.0), alg=ASB07(δ=0.04, recursive=true))
    sol1_nonrec = solve(P_lin, tspan=(0.0, 1.0), alg=ASB07(δ=0.04, recursive=false))
    @test diameter(set(sol1_rec[end])) < diameter(set(sol1_nonrec[end]))

    # FIXME test without IntervalMatrix wrapper
    #A = [-1.0 ± 0.05 -4.0 ± 0.05;
    #     4.0 ± 0.05 -1.0 ± 0.05]
    #P_lin = @ivp(x' = Ax, x(0) ∈ X0)
    #sol1 = solve(P_lin, tspan=(0.0, 1.0), alg=ASB07(δ=0.04));
end

@testset "ASB07 algorithm: inhomogeneous case" begin
    # affine ODE: x' = Ax + Bu
    A = [-1.0 ± 0.05 -4.0 ± 0.05;
          4.0 ± 0.05 -1.0 ± 0.05]
    B = hcat([1.0 ± 0.01; 1.0 ± 0.0])
    U = Interval(-0.05, 0.05)
    X0 = BallInf([1.0, 1.0], 0.1)
    P_aff = @ivp(x' = Ax + Bu, x(0) ∈ X0, u ∈ U, x ∈ Universe(2))
    sol1 = solve(P_aff, tspan=(0.0, 1.0), alg=ASB07(δ=0.04));
end

@testset "ASB07 with time-triggered hybrid system" begin
    prob = embrake_pv_1(ζ=[-1e-8, 1e-7], Tsample=1e-4)
    sol = solve(prob, alg=ASB07(δ=1e-7, max_order=3), max_jumps=2)
    @test dim(sol) == 4
    @test flowpipe(sol) isa HybridFlowpipe

    @test tspan(shift(sol[1],1.0))==tspan(sol[1])+1.0
end