using OrthogonalPolynomialsQuasi, FillArrays, BandedMatrices, ContinuumArrays, QuasiArrays, LazyArrays, FastGaussQuadrature, Test
import OrthogonalPolynomialsQuasi: recurrencecoefficients, basis

@testset "Jacobi" begin
    @testset "JacobiWeight" begin
        a,b = 0.1,0.2
        w = JacobiWeight(a,b)
        @test w.^2 == JacobiWeight(2a,2b)
        @test sqrt.(w) == JacobiWeight(a/2,b/2)
        @test JacobiWeight(0.2,0.3) .* w == JacobiWeight(a+0.2,b+0.3)
        @test LegendreWeight() .* w == w .* LegendreWeight() == w
        @test ChebyshevWeight() .* w == w .* ChebyshevWeight() == JacobiWeight(a-1/2,b-1/2)
    end

    @testset "basics" begin
        @test Legendre() == Jacobi(0,0)
        @test Jacobi(0,0) == Legendre()
    end

    @testset "basis" begin
        b,a = 0.1,0.2
        P = Jacobi(a,b)
        @test P[0.1,2] ≈ 0.16499999999999998
        P = Jacobi(b,a)
        @test P[-0.1,2] ≈ -0.16499999999999998

        A,B,C = recurrencecoefficients(P)
        X = jacobimatrix(P)
        @testset "recurrence coefficient and jacobimatrix" begin
            @test 1/A[1] ≈ X[2,1]
            @test -B[1]/A[1] ≈ X[1,1]
            @test C[2]/A[2] ≈ X[1,2]
            @test 1/A[2] ≈ X[3,2]
            @test -B[2]/A[2] ≈ X[2,2]
            @test C[3]/A[3] ≈ X[2,3]

            @test A[1] ≈ 1/X[2,1]
            @test B[1] ≈ -X[1,1]/X[2,1]
            @test C[2] ≈ X[1,2]/X[3,2]
            @test A[2] ≈ 1/X[3,2]
            @test B[2] ≈ -X[2,2]/X[3,2]
            @test C[3] ≈ X[2,3]/X[4,3]
        end
    end

    @testset "orthogonality" begin
        a,b = 0.1,0.2
        x,w = gaussjacobi(3,a,b)
        P = Jacobi(a,b)

        M = P[x,1:3]'Diagonal(w)*P[x,1:3]
        @test M ≈ Diagonal(M)
        x,w = gaussradau(3,a,b)
        M = P[x,1:3]'Diagonal(w)*P[x,1:3]
        @test M ≈ Diagonal(M)

        w = JacobiWeight(a,b)
        w_2 = sqrt.(w)
        @test w_2 .^ 2 == w
        @test (P' * (w .* P))[1:3,1:3] ≈ (P' * (w .* P))[1:3,1:3] ≈ ((w_2 .* P)'*(w_2 .* P))[1:3,1:3] ≈ M

        @test (Jacobi(0,0)'Jacobi(0,0))[1:10,1:10] ≈ (Legendre()'Legendre())[1:10,1:10]
    end

    @testset "operators" begin
        a,b = 0.1,0.2
        S = Jacobi(a,b)

        @testset "Jacobi" begin
            x = 0.1
            @test S[x,1] === 1.0
            X = jacobimatrix(S)
            @test X[1,1] ≈ (b^2-a^2)/((a+b)*(a+b+2))
            @test X[2,1] ≈ 2/(a+b+2)
            @test S[x,2] ≈ 0.065
            @test S[x,10] ≈ 0.22071099583604945

            w = JacobiWeight(a,b)
            @test w[x] ≈ (1-x)^a * (1+x)^b
            @test OrthogonalPolynomial(w) == S
            wS = w.*S
            @test wS == WeightedJacobi(a,b) == WeightedJacobi{Float64}(a,b)
            @test wS[0.1,1] ≈ w[0.1]
            @test wS[0.1,1:2] ≈ w[0.1] .* S[0.1,1:2]

            w_A = WeightedJacobi(-1/2,0)
            w_B =  WeightedJacobi(1/2,0)

            u = w_A * [1 ; 2; zeros(∞)]
            v = w_B * [1 ; 2; zeros(∞)]
            @test basis(u + v) == w_A
            @test (u+v)[0.1] == u[0.1] + v[0.1]
        end

        @testset "Clenshaw" begin
            x = axes(S,1)
            a = S * (S \ exp.(x))
            A = S \ (a .* S)
            @test (S * (A * (S \ a)))[0.1] ≈ exp(0.2)
            a = S * [1; Zeros(∞)]
            A = S \ (a .* S)
            @test A[1:10,1:10] == I
        end

        @testset "Derivative" begin
            a,b,c = 0.1,0.2,0.3
            S = Jacobi(a,b)
            x = axes(S,1)
            D = Derivative(x)
            u = S \ exp.(x)
            x̃ = 0.1
            @test ((D*S) * u)[x̃] ≈ exp(x̃)
            @test (D * (JacobiWeight(0,b) .* S) * u)[x̃] ≈ exp(x̃) * (1+x̃)^(b-1) * (1+b+x̃)
            @test (D * (JacobiWeight(a,0) .* S) * u)[x̃] ≈ exp(x̃) * (1-x̃)^(a-1) * (1-a-x̃)
            @test (D * (JacobiWeight(a,b) .* S) * u)[x̃] ≈ -exp(x̃)*(1-x̃)^(-1+a)*(1+x̃)^(-1+b)*(a*(1+x̃)+(-1+x̃)*(1+b+x̃))
            @test (D * (JacobiWeight(0,c) .* S) * u)[x̃] ≈ exp(x̃) * (1+x̃)^(c-1) * (1+c+x̃)
            @test (D * (JacobiWeight(c,0) .* S) * u)[x̃] ≈ exp(x̃) * (1-x̃)^(c-1) * (1-c-x̃)
            @test (D * (JacobiWeight(0,0) .* S) * u)[x̃] ≈ exp(x̃)
        end
    end

    @testset "functions" begin
        b,a = 0.1,0.2
        P = Jacobi(a,b)
        D = Derivative(axes(P,1))

        f = P*Vcat(randn(10), Zeros(∞))
        @test (Jacobi(a,b+1) * (Jacobi(a,b+1)\f))[0.1] ≈ f[0.1]
        h = 0.0000001
        @test (D*f)[0.1] ≈ (f[0.1+h]-f[0.1])/h atol=100h

        (D*(JacobiWeight(a,b) .* f))
    end

    @testset "expansions" begin
        P = Jacobi(1/2,0.)
        x = axes(P,1)
        @test (P * (P \ exp.(x)))[0.1] ≈ exp(0.1)

        wP = WeightedJacobi(1/2,0.)
        f = @.(sqrt(1 - x) * exp(x))
        @test wP[0.1,1:100]'*(wP[:,1:100] \ f) ≈ sqrt(1-0.1) * exp(0.1)
        @test (wP * (wP \ f))[0.1] ≈ sqrt(1-0.1) * exp(0.1)

        P̃ = P[affine(Inclusion(0..1), x), :]
        x̃ = axes(P̃, 1)
        @test (P̃ * (P̃ \ exp.(x̃)))[0.1] ≈ exp(0.1)
        wP̃ = wP[affine(Inclusion(0..1), x), :]
        f̃ = @.(sqrt(1 - x̃) * exp(x̃))
        @test wP̃[0.1,1:100]'*(wP̃[:,1:100] \ f̃) ≈ sqrt(1-0.1) * exp(0.1)
        @test (wP̃ * (wP̃ \ f̃))[0.1] ≈ sqrt(1-0.1) * exp(0.1)
    end

    @testset "trivial weight" begin
        S = JacobiWeight(0.0,0.0) .* Jacobi(0.0,0.0)
        @test S == S
        @test Legendre() == S
        @test Legendre()\S isa Eye
    end

    @testset "Jacobi integer" begin
        S = Jacobi(true,true)
        D = Derivative(axes(S,1))
        P = Legendre()

        @test pinv(pinv(S)) === S
        @test P\P === pinv(P)*P === Eye(∞)

        Bi = pinv(Jacobi(2,2))
        @test Bi isa QuasiArrays.PInvQuasiMatrix

        A = Jacobi(2,2) \ (D*S)
        @test typeof(A) == typeof(pinv(Jacobi(2,2))*(D*S))
        @test A isa BandedMatrix
        @test bandwidths(A) == (-1,1)
        @test size(A) == (∞,∞)
        @test A[1:10,1:10] == diagm(1 => 2:0.5:6)

        M = @inferred(D*S)
        @test M isa MulQuasiMatrix
        @test M.args[1] == Jacobi(2,2)
        @test M.args[2][1:10,1:10] == A[1:10,1:10]
    end

    @testset "Weighted Jacobi integer" begin
        S = Jacobi(true,true)
        w̃ = JacobiWeight(false,true)
        A = Jacobi(true,false)\(w̃ .* S)
        @test A isa BandedMatrix
        @test size(A) == (∞,∞)
        @test A[1:10,1:10] ≈ (Jacobi(1.0,0.0) \ (JacobiWeight(0.0,1.0) .* Jacobi(1.0,1.0)))[1:10,1:10]

        w̄ = JacobiWeight(true,false)
        A = Jacobi(false,true)\(w̄.*S)
        @test A isa BandedMatrix
        @test size(A) == (∞,∞)
        @test A[1:10,1:10] ≈ (Jacobi(0.0,1.0) \ (JacobiWeight(1.0,0.0) .* Jacobi(1.0,1.0)))[1:10,1:10]

        P = Legendre()
        w̄ = JacobiWeight(true,false)
        @test_broken P \ (w̃ .* Jacobi(true,false))
        w̄ = JacobiWeight(false,true)
        @test (P \ (w̃ .* Jacobi(false,true)))[1:10,1:10] == diagm(0 => ones(10), -1 => ones(9))

        w = JacobiWeight(true,true)
        A,B = (P'P),P\(w.*S)

        M = Mul(A,B)
        @test M[1,1] == 4/3

        M = ApplyMatrix{Float64}(*,A,B)
        M̃ = M[1:10,1:10]
        @test M̃ isa BandedMatrix
        @test bandwidths(M̃) == (2,0)

        @test A*B isa MulMatrix
        @test bandwidths(A*B) == bandwidths(B)

        A,B,C = (P\(w.*S))',(P'P),P\(w.*S)
        M = ApplyArray(*,A,B,C)
        @test bandwidths(M) == (2,2)
        @test M[1,1] ≈  1+1/15
        M = A*B*C
        @test bandwidths(M) == (2,2)
        @test M[1,1] ≈  1+1/15

        S = Jacobi(1.0,1.0)
        w = JacobiWeight(1.0,1.0)
        wS = w .* S

        W = Diagonal(w)
        @test W[0.1,0.2] ≈ 0.0
    end

    @testset "Jacobi and Chebyshev" begin
        T = ChebyshevT()
        U = ChebyshevU()
        JT = Jacobi(T)
        JU = Jacobi(U)

        @testset "recurrence degenerecies" begin
            A,B,C = recurrencecoefficients(JT)
            @test A[1] == 0.5
            @test B[1] == 0.0
        end

        @test JT[0.1,1:4] ≈ [1.0,0.05,-0.3675,-0.0925]

        @test ((T \ JT) * (JT \ T))[1:10,1:10] ≈ Eye(10)
        @test ((U \ JU) * (JU \ U))[1:10,1:10] ≈ Eye(10)

        @test T[0.1,1:10]' ≈ JT[0.1,1:10]' * (JT \ T)[1:10,1:10]
        @test U[0.1,1:10]' ≈ JU[0.1,1:10]' * (JU \ U)[1:10,1:10]

        w = JacobiWeight(1,1)
        @test (w .* U)[0.1,1:10] ≈ (T * (T \ (w .* U)))[0.1,1:10]
        @test (w .* U)[0.1,1:10] ≈ (JT * (JT \ (w .* U)))[0.1,1:10]
        @test (w .* JU)[0.1,1:10] ≈ (T * (T \ (w .* JU)))[0.1,1:10]
        @test (w .* JU)[0.1,1:10] ≈ (JT * (JT \ (w .* JU)))[0.1,1:10]
    end

    @testset "Jacobi-Chebyshev-Ultraspherical transforms" begin
        @test Jacobi(0.0,0.0) \ Legendre() == Eye(∞)
        @test ((Ultraspherical(3/2) \ Jacobi(1,1))*(Jacobi(1,1) \ Ultraspherical(3/2)))[1:10,1:10] ≈ Eye(10)
        f = Jacobi(0.0,0.0)*[[1,2,3]; zeros(∞)]
        g = (Legendre() \ f) - f.args[2]
        @test_skip norm(g) ≤ 1E-15
        @test_broken (Legendre() \ f) == f.args[2]
        @test (Legendre() \ f)[1:10] ≈ f.args[2][1:10]
        f = Jacobi(1.0,1.0)*[[1,2,3]; zeros(∞)]
        g = Ultraspherical(3/2)*(Ultraspherical(3/2)\f)
        @test f[0.1] ≈ g[0.1]

        @testset "Chebyshev-Legendre" begin
            T = Chebyshev()
            P = Legendre()
            @test T[:,Base.OneTo(5)] \ P[:,Base.OneTo(5)] == (T\P)[1:5,1:5]

            x = axes(P,1)
            u = P * (P \ exp.(x))
            @test u[0.1] ≈ exp(0.1)

            P = Legendre{BigFloat}()
            x = axes(P,1)
            u = P * (P \ exp.(x))
            @test u[BigFloat(1)/10] ≈ exp(BigFloat(1)/10)
        end
    end

    @testset "hcat" begin
        L = LinearSpline(range(-1,1;length=2))
        S = JacobiWeight(1.0,1.0) .* Jacobi(1.0,1.0)
        P = apply(hcat,L,S)
        @test P isa ApplyQuasiArray
        @test axes(P) == axes(S)
        V = view(P,0.1,1:10)
        @test all(arguments(V) .≈ [L[0.1,:], S[0.1,1:8]])
        @test P[0.1,1:10] == [L[0.1,:]; S[0.1,1:8]]
        D = Derivative(axes(P,1))
        # applied(*,D,P) |> typeof
        # MemoryLayout(typeof(D))
    end

    @testset "Jacobi Clenshaw" begin
        P = Jacobi(0.1,0.2)
        x = axes(P,1)
        a = P * (P \ exp.(x))
        @test a[0.1] ≈ exp(0.1)
        M = P \ (a .* P);
        u = [randn(1000); zeros(∞)];
        @test (P * (M*u))[0.1] ≈ (P*u)[0.1]*exp(0.1)
    end

    @testset "mapped" begin
        R = jacobi(0,1/2,0..1) \ jacobi(0,-1/2,0..1)
        R̃ = Jacobi(0,1/2) \ Jacobi(0,-1/2)
        @test R[1:10,1:10] == R̃[1:10,1:10]
    end

    @testset "Christoffel–Darboux" begin
        a,b = 0.1,0.2
        P = Jacobi(a,b)
        X = P\ (axes(P,1) .* P)
        Mi = inv(P'*(JacobiWeight(a,b) .* P))
        x,y = 0.1,0.2
        n = 10
        Pn = Diagonal([Ones(n); Zeros(∞)])
        Min = Pn * Mi
        @test norm((X*Min - Min*X')[1:n,1:n]) ≤ 1E-13
        β = X[n,n+1]*Mi[n+1,n+1]
        @test (x-y) * P[x,1:n]'Mi[1:n,1:n]*P[y,1:n] ≈ P[x,n:n+1]' * (X*Min - Min*X')[n:n+1,n:n+1] * P[y,n:n+1] ≈ P[x,n:n+1]' * [0 -β; β 0] * P[y,n:n+1]

        @testset "extrapolation" begin
            x,y = 0.1,3.4
            @test (x-y) * P[x,1:n]'Mi[1:n,1:n]*Base.unsafe_getindex(P,y,1:n) ≈ P[x,n:n+1]' * [0 -β; β 0] * Base.unsafe_getindex(P,y,n:n+1)
        end
    end
end


"""
Compute the vertical integrated volume flux from the bottom to z=0 (i.e. linear free-surface)

```
U★ = ∫ᶻ Ax * u★ dz`
V★ = ∫ᶻ Ay * v★ dz`
"""
### Note - what we really want is RHS = divergence of the vertically integrated volume flux
###        we can optimize this a bit later to do this all in one go to save using intermediate variables.
function compute_vertically_integrated_volume_flux!(∫ᶻ_U, model)

    # Fill halo regions for predictor velocity.
    fill_halo_regions!(model.velocities, model.architecture, model.clock, fields(model))

    # Compute the vertically integrated volume flux
    event = launch!(model.architecture, model.grid, :xy,
                    _compute_vertically_integrated_volume_flux!,
                    ∫ᶻ_U, model.grid, model.velocities,
                    dependencies = device_event(model.architecture))
                    # include_right_boundaries = true) # <- get this to work with layout = :xy

    wait(device(model.architecture), event)

    # We didn't include right boundaries, so...
    fill_halo_regions!(∫ᶻ_U, model.architecture, model.clock, fields(model))

    return nothing
end

@kernel function _compute_vertically_integrated_volume_flux!(∫ᶻ_U, grid, U)
    i, j = @index(Global, NTuple)

    @inbounds begin

        ∫ᶻ_U.u[i, j, 1] = 0
        ∫ᶻ_U.v[i, j, 1] = 0

        @unroll for k in 1:grid.Nz
            ∫ᶻ_U.u[i, j, 1] += Axᶠᶜᶜ(i, j, k, grid) * U.u[i, j, k]
            ∫ᶻ_U.v[i, j, 1] += Ayᶜᶠᶜ(i, j, k, grid) * U.v[i, j, k]
        end

     end
end


using CompScienceMeshes, BEAST, Base.Test

fn = joinpath(dirname(@__FILE__),"assets","sphere35.in")

sol = 1.0
D, Δx = 1.0, 0.35
#Γ = meshsphere(D, Δx)
Γ = readmesh(fn)
X = raviartthomas(Γ)
Xm = subset(X,[1])
Xn = subset(X,[100])
Δt, Nt = 0.13/sol, 40
T = timebasisshiftedlagrange(Δt, Nt, 3)
U = timebasisdelta(Δt, Nt)

W = Xm ⊗ U
V = Xn ⊗ T

t = MWSingleLayerTDIO(sol,-1/sol,-sol,2,0)
Z1 = assemble(t, W, V)

## Run again, this time scales with another speedOfLight

sol = 3.0
D, Δx = 1.0, 0.35
#Γ = meshsphere(D, Δx)
Γ = readmesh(fn)
X = raviartthomas(Γ)
Xm = subset(X,[1])
Xn = subset(X,[100])
Δt, Nt = 0.13/sol, 40
T = timebasisshiftedlagrange(Δt, Nt, 3)
U = timebasisdelta(Δt, Nt)

W = Xm ⊗ U
V = Xn ⊗ T

t = MWSingleLayerTDIO(sol,-1/sol,-sol,2,0)
Z3 = assemble(t, W, V)

m = n = 1
@test find(Z1[m,n,:]) == find(Z3[m,n,:])

I = find(Z1[m,n,:])
@test all(sol*Z1[m,n,I] .≈ Z3[m,n,I])


function fibonnaci(n)
	t1 = 0
	t2 = 1
	nextTerm = 0
	for i in 1:n
		if i == 1
			nextTerm = t1
		elseif i == 2
			nextTerm = t2
		end
		nextTerm = t1 + t2
		t1 = t2
		t2 = nextTerm
	#println(nextTerm)
	end
	println(nextTerm)
end

fibonnaci(1500)


# This file is a part of Julia. License is MIT: https://julialang.org/license

# THE SIX LOGICAL FUNCTIONS
#
# Bit shifting and rotation (used by the six SHA-XYZ logical functions:
#
#   NOTE:  The naming of R and S appears backwards here (R is a SHIFT and
#   S is a ROTATION) because the SHA2-256/384/512 description document
#   (see http://csrc.nist.gov/cryptval/shs/sha256-384-512.pdf) uses this
#   same "backwards" definition.

# 32-bit Rotate-right (equivalent to S32 in SHA-256) and rotate-left
rrot(b,x,width) = ((x >> b) | (x << (width - b)))
lrot(b,x,width) = ((x << b) | (x >> (width - b)))

# Shift-right (used in SHA-256, SHA-384, and SHA-512):
R(b,x)   = (x >> b)
# 32-bit Rotate-right (used in SHA-256):
S32(b,x) = rrot(b,x,32)
# 64-bit Rotate-right (used in SHA-384 and SHA-512):
S64(b,x) = rrot(b,x,64)
# 64-bit Rotate-left (used in SHA3)
L64(b,x) = lrot(b,x,64)

# Two of six logical functions used in SHA-256, SHA-384, and SHA-512:
Ch(x,y,z)  = ((x & y) ⊻ (~x & z))
Maj(x,y,z) = ((x & y) ⊻ (x & z) ⊻ (y & z))

# Four of six logical functions used in SHA-256:
Sigma0_256(x) = (S32(2,  UInt32(x)) ⊻ S32(13, UInt32(x)) ⊻ S32(22, UInt32(x)))
Sigma1_256(x) = (S32(6,  UInt32(x)) ⊻ S32(11, UInt32(x)) ⊻ S32(25, UInt32(x)))
sigma0_256(x) = (S32(7,  UInt32(x)) ⊻ S32(18, UInt32(x)) ⊻ R(3 ,   UInt32(x)))
sigma1_256(x) = (S32(17, UInt32(x)) ⊻ S32(19, UInt32(x)) ⊻ R(10,   UInt32(x)))

# Four of six logical functions used in SHA-384 and SHA-512:
Sigma0_512(x) = (S64(28, UInt64(x)) ⊻ S64(34, UInt64(x)) ⊻ S64(39, UInt64(x)))
Sigma1_512(x) = (S64(14, UInt64(x)) ⊻ S64(18, UInt64(x)) ⊻ S64(41, UInt64(x)))
sigma0_512(x) = (S64( 1, UInt64(x)) ⊻ S64( 8, UInt64(x)) ⊻ R( 7,   UInt64(x)))
sigma1_512(x) = (S64(19, UInt64(x)) ⊻ S64(61, UInt64(x)) ⊻ R( 6,   UInt64(x)))

# Let's be able to bswap arrays of these types as well
bswap!(x::Vector{<:Integer}) = map!(bswap, x, x)


module SymbolicML

using CategoricalArrays
using CSV
using DataFrames
using ExprRules
using ExprOptimization
# using MultivariateTimeSeries
using Plots
using StatsPlots
using Random
# using GBDTs

# include("dataset.jl")
include("ds.jl")
include("siemens.jl")
include("functions.jl")

export discrete_derivative, discrete_derivatives
export siemensDataInfo, siemensData

export select_attributes
export abs_energy
export absolute_maximum
export absolute_sum_of_changes
export agg_autocorrelation
export autocorrelation
export benford_correlation
export binned_entropy
export c3
export change_quantiles
export cid_ce
export count_above
export count_above_mean
export count_below
export count_below_mean
export first_location_of_maximum
export first_location_of_minimum
export has_duplicate
export has_duplicate_max
export has_duplicate_min
export large_standard_deviation
export last_location_of_maximum
export last_location_of_minimum
export mean_abs_change
export mean_change
export mean_n_absolute_max
export longest_strike_above_mean
export longest_strike_below_mean
export mean_second_derivative_central
export number_crossing_m
export percentage_of_reoccurring_datapoints_to_all_datapoints
export percentage_of_reoccurring_values_to_all_values
export range_count
export ratio_beyond_r_sigma
export ratio_value_number_to_time_series_length
export root_mean_square
export sample_entropy
export sum_of_reoccurring_data_points

export ClassificationDataset
export ModalFrame
export ModalInstance

export attributes

end


using ExactD, Test, LinearAlgebra

@testset "parity operator" begin
    @test ExactD.parity(10, 1, 2) == 1.
    @test ExactD.parity(10, 2, 0) == -1.
    @test ExactD.parity(10, 1, 3) == 1.
    @test ExactD.parity(14, 1, 3) == -1.
end

@testset "hard-core boson many body operators" begin
@testset "dense many-body operator" begin
    L = 4
    N = 2
    J::Array{Float64, 2} = reshape(collect(1:L^2), (L, L))
    V::Array{Float64, 2} = reshape(collect(1:L^2), (L, L))
    Op = build_many_body_op(L, N, J, V)

    # Diagonal terms.
    @test Op[1, 1] ≈ 7. + 7.
    @test Op[4, 4] ≈ 17. + 17.
    @test Op[6, 6] ≈ 27. + 27.

    # Off-diagonal terms.
    @test Op[1, 2] ≈ 10.
    @test Op[2, 5] ≈ 0.
    @test Op[4, 2] ≈ 12.
    @test Op[6, 1] ≈ 0.
    @test Op[3, 5] ≈ 15.

    C = 1.3
    Op2 = build_many_body_op(L, N, J, V, C)

    # Diagonal terms.
    @test Op2[1, 1] ≈ 8.3 + 7.
    @test Op2[4, 4] ≈ 18.3 + 17.
    @test Op2[6, 6] ≈ 28.3 + 27.
end

@testset "sparse many-body operator" begin
    L = 6
    N = 3
    basis = make_LN_basis(L, N)
    J::Array{Float64, 2} = reshape(sin.(1:L^2), (L, L))
    V::Array{Float64, 2} = reshape(tan.(1:L^2), (L, L))
    C = 1.3
    sp_Op = build_sparse_many_body_op(L, basis, J, V, C)
    de_Op = build_many_body_op(L, basis, J, V, C)
    @test norm(Array(sp_Op) - de_Op) ≈ 0.0 atol=1e-12
end
end  # testset "hard-core boson many body operators"

@testset "dense many-body operator with L spins s=1" begin
    L = 4
    Sz = 2
    J = reshape(collect(1.:L^2), (L, L))
    W = reshape(collect(1.:L^2), (L, L))
    Jz = collect(1.:L)/2
    Jx = zeros(L)
    C = 1.33
    Op = build_spin1_many_body_op(L, Sz, J, W, Jz, Jx, C)

    # Diagonal terms: S^z + (S^z_i)^2 + S^z_i*S^z_j.
    @test Op[1, 1] ≈ 1.5 + 7 + 5 + C
    @test Op[4, 4] ≈ 2.5 + 17 + 13 + C
    @test Op[7, 7] ≈ 4. + 34 + 12 + C
    @test Op[10, 10] ≈ 1. + 34 - 18 + C

    # Off-diagonal terms.
    @test Op[1, 2] ≈ 20.
    @test Op[1, 4] ≈ 28.
    @test Op[1, 6] ≈ 0.
    @test Op[7, 6] ≈ 4.
    @test Op[6, 7] ≈ 10.
    @test Op[1, 9] ≈ 30.
    @test Op[9, 1] ≈ 24.
    @test Op[1, 10] ≈ 24.

    L = 3
    Sz = 0
    J = reshape(collect(1.:L^2), (L, L))
    W = reshape(collect(1.:L^2), (L, L))
    Jz = collect(1.:L)/4
    Jx = zeros(L)
    Op = build_spin1_many_body_op(L, Sz, J, W, Jz, Jx)
    # Off-diagonal terms.
    @test Op[4, 2] ≈ 12.
    @test Op[4, 5] ≈ 4.

    # Test only Sx terms.
    L = 3
    basis = make_spin1_L_basis(L)
    J = zeros(L, L)
    W = zeros(L, L)
    Jz = zeros(L)
    Jx = collect(1.:L)
    Op = build_spin1_many_body_op(L, basis, J, W, Jz, Jx)

    @test all(Op .≈ transpose(Op))
    @test Op[1, 3] ≈ 1/sqrt(2)
    @test Op[1, 2] ≈ 1/sqrt(2)
    @test Op[12, 15] ≈ 2/sqrt(2)
    @test Op[12, 18] ≈ 2/sqrt(2)
    @test Op[27, 9] ≈ 3/sqrt(2)
end


m=rand(2:10)
n=rand(1:10)
Z = Frobenius(m, randn(n)) #size m+n

#Special properties
@test_approx_eq full(inv(Z)) inv(full(Z))

#Matvec product
b = randn(m+n)
@test_approx_eq Z*b full(Z)*b

#Matmul product
Y1 = Frobenius(m, randn(n)) #Another one of the same column
@test_approx_eq full(Y1*Z) full(Y1)*full(Z)

n2=rand(1:(m+n-1))
Y2 = Frobenius(m+n-n2, randn(n2)) #Probably not the same column
@test_approx_eq full(Y2*Z) full(Y2)*full(Z)


@testset "Trends" begin
  rng = MersenneTwister(42)

  # constant trend
  d = georef((z=rand(rng,100),), CartesianGrid(100))
  z̄ = trend(d, :z)[:z]
  @test all(abs.(diff(z̄)) .< 0.01)

  # linear trend
  μ = range(0,stop=1,length=100)
  ϵ = 0.1rand(rng,100)
  d = georef((z=μ+ϵ,), CartesianGrid(100))
  z̄ = trend(d, :z)[:z]
  @test all([abs(z̄[i] - μ[i]) < 0.1 for i in 1:length(z̄)])

  # quadratic trend
  r = range(-1,stop=1,length=100)
  μ = [x^2 + y^2 for x in r, y in r]
  ϵ = 0.1rand(rng,100,100)
  d = georef((z=μ+ϵ,))
  d̄ = trend(d, :z, degree=2)
  z̄ = reshape(d̄[:z], 100, 100)
  @test all([abs(z̄[i] - μ[i]) < 0.1 for i in 1:length(z̄)])

  if visualtests
    p₁ = heatmap(μ+ϵ, title="z")
    p₂ = heatmap(z̄, title="z trend")
    plt = plot(p₁, p₂, size=(900,300))
    @test_reference "data/trends.png" plt
  end

  d = georef((x=rand(rng,10), y=rand(rng,10)), rand(rng,2,10))
  𝒯 = d |> trend |> values
  s = Tables.schema(𝒯)
  @test s.names == (:x, :y)
  @test s.types == (Float64, Float64)
end


"""
return maximum distance across each dimensions in the fist particles ranks
"""
function maxdist_bandwidth(distances::Array{G, 1}, ordinals::Array{Int64, 1}, nparticles::Integer) where G <: Real
    return distances[ordinals[nparticles]]
end

function maxdist_bandwidth(distances::Array{G, 2}, ordinals::Array{Int64, 1}, nparticles::Integer) where G <: Real
    return vec(maximum(subset(distances, ordinals[1:nparticles]), 1))
end

"""
Determine bandwidth based on 0.8 (default) acceptance rate for largest distance of accepted particle using a guassian acceptance kernel
"""
function gaussianprob_bandwidth(distances::Array{G, 1}, ordinals::Array{Int64, 1}, nparticles::Integer; paccept::Float64 = 0.5) where G <: Real
    maxdist = maxdist_bandwidth(distances, ordinals, Int(nparticles / 2))
    return maxdist / sqrt(-2.0 * log(paccept))
end

function gaussianprob_bandwidth(distances::Array{G, 2}, ordinals::Array{Int64, 1}, nparticles::Integer; paccept::Float64 = 0.5) where G <: Real
    maxdist = maxdist_bandwidth(distances, ordinals, Int(nparticles / 2))
    return maxdist ./ sqrt(-2.0 * log(paccept) / length(maxdist))
end


using UpdatableCholeskyFactorizations
using LinearAlgebra
using BenchmarkTools

elty = Float64
n = 128
m = 16
A_full = randn(elty, 2n, 2n)
A_full = A_full'A_full
A = A_full[1:n, 1:n]
a = A_full[1:n+1, n+1]
B = A_full[1:2n, n+1:2n]
Bm = A_full[1:n+m, n+1:n+m] # adding m columns
C = updatable_cholesky(A)

set = @benchmarkset "updatable cholesky" begin # for n in 1:16 IDEA: could have loop over different sizes
    @case "$n by $n cholesky" cholesky($A)
    @case "$(2n) by $(2n) cholesky" cholesky($A_full)
    @case "$n by $n updatable_cholesky" updatable_cholesky($A)
    @case "adding vector to $n by $n UpdatableCholesky" add_column!(D, $a) setup=(D=copy($C))
    @case "adding $m vectors to $n by $n UpdatableCholesky" add_column!(D, $Bm) setup=(D=copy($C))
    @case "adding $n x $n matrix to $n by $n UpdatableCholesky" add_column!(D, $B) setup=(D=copy($C))
end

result = run(set)

r = minimum(result["updatable cholesky"])
# times = Vector{Float64}(undef, length(r))
# titles = Vector{String}(undef, length(r))
# i = 1
for (title, trial) in r.data
    println()
    println(title)
    display(trial)
    # times[i] = trial.time
    # titles[i] = title
    # i += 1
end


import AlphaZero.GI

using StaticArrays
using Crayons

const Cell = UInt8
const EMPTY = 0x00
const ARROW = 0x03

tmp = [0 1 0 0 1 0;
       1 0 0 0 0 1;
       0 0 0 0 0 0;
       0 0 0 0 0 0;
       2 0 0 0 0 2;
       0 2 0 0 2 0]

const NUM_COLS = size(tmp,1)
const NUM_ROWS = size(tmp,2)
const NUM_CELLS = length(tmp)
const Board = SMatrix{NUM_COLS, NUM_ROWS, Cell, NUM_CELLS}

const Player = UInt8
const WHITE = 0x01
const BLACK = 0x02

other(p::Player) = 0x03 - p

const INITIAL_BOARD = SMatrix{NUM_COLS,NUM_ROWS}(Player.(tmp))
const INITIAL_STATE = (board=INITIAL_BOARD, curplayer=WHITE)

struct GameSpec <: GI.AbstractGameSpec end

mutable struct GameEnv <: GI.AbstractGameEnv
  board :: Board
  curplayer :: Player
  finished :: Bool
  winner :: Player
  #selected works as long as size of board is not too big
end

GI.spec(::GameEnv) = GameSpec()

function GI.init(::GameSpec)
  board = INITIAL_STATE.board
  curplayer = INITIAL_STATE.curplayer
  finished = false
  winner = 0x00
  return GameEnv(board, curplayer, finished, winner)
end

GI.two_players(::GameSpec) = true

GI.actions(::GameSpec) = collect(1:9412)
#=
15,15,15,15,15,15
15,17,17,17,17,15
15,17,19,19,17,15
15,17,19,19,17,15
15,17,17,17,17,15
15,15,15,15,15,15

Above is queen moves available at each square.
Because every move can be uniquely identified by a→b→c,
so if b is fixed, it can come from x places and go to x places,
contributing to a total of x^2 moves
So the total number of moves is
20*15^2+12*17^2+4*19^2=9412
=#

flip_cell_color(c::Cell) = (0 < c < 3) ? other(c) : c

flip_colors(board::Board) = flip_cell_color.(board)

function GI.vectorize_state(::GameSpec, state)
  tmp = zeros(Float32,NUM_ROWS,NUM_COLS,3)
  board = state.curplayer == WHITE ? state.board : flip_colors(state.board)
  for i = 1:3
    tmp[:,:,i] = board .== i
  end
  return tmp
end

function GI.set_state!(g::GameEnv, state)
  g.board = state.board
  g.curplayer = state.curplayer
  if any(GI.actions_mask(g))
    g.winner = 0x00
    g.finished = false
  else
    g.winner = other(g.curplayer)
    g.finished = true
  end
  return
end

GI.current_state(g::GameEnv) = (board=g.board, curplayer=g.curplayer)

# TODO: maybe MCTS should make the copy itself. The performance cost should not be great
# and it would probably avoid people a lot of pain.

GI.game_terminated(g::GameEnv) = g.finished

GI.white_playing(g::GameEnv) = g.curplayer == WHITE

const DIRECTIONS = CartesianIndex.([(-1,-1),(-1,0),(-1,1),(0,1),(1,1),(1,0),(1,-1),(0,-1)])

pos_of_xy((x, y)) = (y - 1) * NUM_COLS + (x - 1) + 1
xy_of_pos(pos) = ((pos - 1) % NUM_COLS + 1, (pos - 1) ÷ NUM_COLS + 1)

function get_queen_moves!(b,loc_linear,container)
  loc = CartesianIndex(xy_of_pos(loc_linear))
  valid_moves = reshape(container,NUM_ROWS,NUM_COLS)
  for dir∈DIRECTIONS
    for i=1:max(size(b)...)
      tmploc = loc + i * dir
      if checkbounds(Bool,b,tmploc) && b[tmploc]==0
        valid_moves[tmploc] = true
      else
        break
      end
    end
  end
end

function get_queen_moves(b::Board,loc_linear)
  valid_moves = falses(length(b))
  get_queen_moves!(b,loc_linear,valid_moves)
  return valid_moves
end

function precompute_mask()
  mask = falses(36*36*36)
  tmp = reshape(mask,36,36,36)
  tmp_board = zeros(eltype(INITIAL_BOARD),NUM_ROWS,NUM_COLS)
  tmp_board = SMatrix{NUM_COLS,NUM_ROWS}(Player.(tmp_board))
  moves = falses(NUM_ROWS,NUM_COLS)
  shoots = falses(NUM_ROWS,NUM_COLS)
  for i=1:36
    fill!(moves,false)
    get_queen_moves!(tmp_board,i,moves)
    for j=1:36
      !moves[j] && continue
      fill!(shoots,false)
      get_queen_moves!(tmp_board,j,shoots)
      for k=1:36
        !shoots[k] && continue
        tmp[i,j,k]=true
      end
    end
  end
  return mask
end

const PRECOMPUTED_MASK = precompute_mask()
const NON_ZERO_LOCS = findall(!iszero,PRECOMPUTED_MASK)

#=
Because of the representation, the matrix is 36^3=46656, but there are only 9412 valid moves
So the strategy is, given a state, find the valid moves, fill them into the 36^3 matrix, which has better representation
but lower efficiency. Most of the entries will always be zero, so filter them out using a precomputed mask.

Inversely, given a number x between 1:9412, find the location of the xth non-zero entry in the precomputed mask
=#

function GI.actions_mask(g::GameEnv)
  moves = falses(NUM_ROWS,NUM_COLS)
  shoots = falses(NUM_ROWS,NUM_COLS)
  board = zeros(Cell,NUM_ROWS,NUM_COLS)
  mask = falses(NUM_CELLS^3)
  tmp = reshape(mask,NUM_CELLS,NUM_CELLS,NUM_CELLS)
  for i=1:NUM_CELLS
    board[i] = g.board[i]
  end
  for i=1:NUM_CELLS
    board[i]!=g.curplayer && continue
    fill!(moves,false)
    get_queen_moves!(board,i,moves)
    for j=1:NUM_CELLS
      !moves[j] && continue
      fill!(shoots,false)
      #don't forget in this case, the amazon could shoot
      #to where it came from. So temporarily delete the current
      #piece, and then add it back
      board[i]=0x00
      get_queen_moves!(board,j,shoots)
      board[i]=g.curplayer
      for k=1:NUM_CELLS
        !shoots[k] && continue
        tmp[i,j,k]=true
      end
    end
  end
  return mask[PRECOMPUTED_MASK]
end

function GI.play!(g::GameEnv, a)
  #first map from the compressed representation to the full representation,
  #then map from 1 based index to 0 based index
  a = NON_ZERO_LOCS[a]
  a = a - 1
  source = a % NUM_CELLS
  tmp = (a-source)÷NUM_CELLS
  move = tmp % NUM_CELLS
  shoot = (tmp - move) ÷ NUM_CELLS

  shoot += 1
  move += 1
  source += 1

  g.board = Base.setindex(g.board,0,source)
  g.board = Base.setindex(g.board,g.curplayer,move)
  g.board = Base.setindex(g.board,3,shoot)
  g.curplayer = other(g.curplayer)
  new_actions = GI.actions_mask(g)
  if !any(new_actions)
    g.finished=true
    g.winner = other(g.curplayer)
  end
end

function GI.white_reward(g::GameEnv)
  g.winner == WHITE && (return  1.)
  g.winner == BLACK && (return -1.)
  return 0.
end

function generate_dihedral_symmetries()
  N = NUM_COLS
  rot((x, y)) = (y, N - x + 1) # 90° rotation
  flip((x, y)) = (x, N - y + 1) # flip along vertical axis
  ap(f) = p -> pos_of_xy(f(xy_of_pos(p)))
  sym(f) = map(ap(f), collect(1:NUM_CELLS))
  rot2 = rot ∘ rot
  rot3 = rot2 ∘ rot
  return [
    sym(rot), sym(rot2), sym(rot3),
    sym(flip), sym(flip ∘ rot), sym(flip ∘ rot2), sym(flip ∘ rot3)]
end

const SYMMETRIES = generate_dihedral_symmetries()

function GI.symmetries(::GameSpec, s)
  return [
    ((board=Board(s.board[sym]), curplayer=s.curplayer,phase=s.phase,selected=s.selected), sym)
    for sym in SYMMETRIES]
end

#####
##### Interface for interactive exploratory tools
#####

player_color(p) = p == WHITE ? crayon"light_red" : crayon"light_blue"
player_name(p)  = p == WHITE ? "Red" : "Blue"
function cell_mark(c)
  if c == EMPTY
    return "."
  elseif c == ARROW
    return "x"
  else
    return "o"
  end
end

function cell_color(c)
  if c == ARROW || c == EMPTY
    return crayon""
  else
    return player_color(c)
  end
end

function GI.render(g::GameEnv; with_position_names=true, botmargin=true)
  pname = player_name(g.curplayer)
  pcol = player_color(g.curplayer)
  print(pcol, pname," next")
  print(crayon"reset", "\n\n")
  # Print legend
  print("  ")
  for col in 1:NUM_COLS
    print(col, " ")
  end
  print("\n")
  # Print board
  for row in 1:NUM_ROWS
    print(row," ")
    for col in 1:NUM_COLS
      x = g.board[row,col]
      print(cell_color(x), cell_mark(x), crayon"reset", " ")
    end
    print("\n")
  end
  botmargin && print("\n")
end

function GI.action_string(::GameSpec, a)
  a = NON_ZERO_LOCS[a]
  a = a - 1
  source = a % NUM_CELLS
  tmp = (a-source)÷NUM_CELLS
  move = tmp % NUM_CELLS
  shoot = (tmp - move) ÷ NUM_CELLS

  shoot += 1
  move += 1
  source += 1
  return "move from " * string(xy_of_pos(source)) *
         " to "       * string(xy_of_pos(move)) *
         " shoot to " * string(xy_of_pos(shoot))
end

function GI.parse_action(::GameSpec, str)
  try
    x,y = parse(Int, str[1]),parse(Int,str[3])
    ((1 <= x,y <= NUM_COLS)|>all) ? CartesianIndex(x,y) : nothing
  catch
    nothing
  end
end

#=
function GI.read_state(::GameSpec)
  board = Array(INITIAL_BOARD)
  try
    for col in 1:NUM_COLS
      input = readline()
      for (row, c) in enumerate(input)
        c = lowercase(c)
        if c ∈ ['o', 'w', '1']
          board[col, row] = WHITE
        elseif c ∈ ['x', 'b', '2']
          board[col, row] = BLACK
        end
      end
    end
    nw = count(==(WHITE), board)
    nb = count(==(BLACK), board)
    if nw == nb
      curplayer = WHITE
    elseif nw == nb + 1
      curplayer = BLACK
    else
      return nothing
    end
    return (board=board, curplayer=curplayer)
  catch e
    return nothing
  end
end

=#


# Dependencies
using Distributions # not required per se
using Plots
using Logging
import Base: show, zero
# Constants
const states = ['U','N','A','F']
const cs = Dict('N'=>1/2, 'A'=>1/4, 'F'=>1/8)
# helper function for indices
inds(i,j) = CartesianIndex.([(i-1, j-1), (i,j-1), (i+1, j-1),
            (i-1, j),             (i+1, j),
            (i-1, j+1), (i,j+1), (i+1, j+1)])

"""
    Cell

Cell in our world. State should be one of 'U', 'N', 'A','F','B' for
respectively Urban, Nature, Agriculture, Forest and Border.
"""
struct Cell
    state::Char
    age::Int
end
zero(::Type{Cell}) = Cell('b',0)
show(io::IO, c::Cell) = print(io, "($(c.state), $(c.age))")

"""
    getneigbors(world::Array{T,2}, i::Int,j::Int)

return the eight neighbors of the world in position (i,j)
"""
function getneigbors(world::Array{T,2}, i::Int,j::Int) where T
    return view(world, inds(i,j))
end

"""
    trfprob(c::Cell)

Return transformation probability for a non-urban to urban cell.
"""
function trfprob(c::Cell)
    return tanh((c.age+1/2)/40)
end

"""
    urbanisation(c::Cell, n)

return the next states for a cell and its neighbors

see also: [getneigbors](@ref)
"""
function urbanisation(c::Cell, n)
    # urban stays urban
    if c.state == 'U'
        @debug "Town gets older"
        return Cell('U', c.age + 1)
    # nature evolves to urban
    else
        # within mutation treshold
        if rand() < sum(map(trfprob, filter(x-> x.state =='U',n))) / 8 * cs[c.state]
            return Cell('U', 0)
            @info "Added new city!!!"
        # otherwise just age
        else
            @debug "$(c.state) gets older"
            return Cell(c.state, c.age + 1)
        end
    end
end

"""
    ruralisation(c::Cell, n)

return the next states for a cell and its neighbors

see also: [getneigbors](@ref)
"""
function ruralisation(c::Cell,n)
    # urban stays urban
    if c.state == 'U'
        @debug "Town gets older"
        return Cell(c.state, c.age + 1)
    # forest stays forest
    elseif c.state =='F'
        @debug "Forest grows older"
        return Cell(c.state, c.age + 1)
    # nature to forest
    elseif c.state == 'N'
        # within mutation treshold
        if rand() < 1/9 * tanh(c.age/20) * (1 + length(filter(x-> x.state =='F',n)) )
            @debug "NEW FOREST!!"
            return Cell('F',0)
        # otherwise just age
        else
            return Cell(c.state, c.age +1)
        end
    # agriculture to nature
    elseif c.state == 'A'
        if rand() < 1/8 * sum( map(x -> tanh(x.age / 10), filter(x-> x.state =='N',n) ) )
            @debug "NEW NATURE!!"
            return Cell('N',0)
        else
            return Cell(c.state, c.age +1)
        end
    end
end

"""
    update(c::Cell, n)

Given a cell and its neighbors, determine the next state
"""
function update(c::Cell, n)
    if rand() < 5/10
        return urbanisation(c, n)
    else
        return ruralisation(c,n)
    end
end

"""
    genworld(N::Int)

Generate an initial world (random composition) with a border.

statedist: default to uniform, but you can also use a custom distribution
"""
function genworld(N::Int;
                 statedist=Categorical(1/length(states).* ones(length(states))))
    @debug statedist
    world = zeros(Cell, N+2, N+2)
    world[2:end-1, 2:end-1] .= reshape([Cell(states[rand(statedist)], 7) for _ in 1:N^2], N, :)
    return world
end

"""
    newworld!(w::Array{Cell,2}, nw::Array{Cell,2})

From a given world, determine the next state and return it.
"""
function newworld!(w::Array{Cell,2}, nw::Array{Cell,2})
    for i in 2:size(w,1)-1
        for j in 2:size(w,2)-1
            nw[i,j] = update(w[i,j], getneigbors(w, i, j))
        end
    end
    w .= nw
    return w
end

# actual program
function main(  N::Int=10,
                d=Categorical(1/length(states).* ones(length(states)));
                nsteps::Int=100)
    w = genworld(N, statedist=d)
    nw = zeros(Cell, size(w))
    record = zeros(nsteps,4)
    for i in 1:nsteps
        newworld!(w, nw)
        record[i,1] = sum(map(x-> x.state == 'U', w)) / N^2
        record[i,2] = sum(map(x-> x.state == 'F', w)) / N^2
        record[i,3] = sum(map(x-> x.state == 'N', w)) / N^2
        record[i,4] = sum(map(x-> x.state == 'A', w)) / N^2
    end
    return record
end

# logger settings
Logging.global_logger(ConsoleLogger(stdout,Logging.Debug))
Logging.disable_logging(Logging.LogLevel(-1000))

# do it
@info "running main"
record = main(100)

# plotting (not was not required)
const cdict = Dict(65=>0, 70=>1, 78=>2, 85=>3, 98=>4)
const cmap = cgrad([    RGBA(222/255,166/255,53/255),       # agriculture brown
                        RGBA(80/255,118/255,40/255),        # forest green (dark)
                        RGBA(105/255,224/255,62/255),       # nature green (light)
                        RGBA(181/255,181/255,181/255),      # urban gray
                        RGBA(240/255,240/255,240/255)]);    # border off-white
const heatmapsettings = Dict(:yflip=>true, :color=>cmap, :axis=>false, :size=>(500,350), :clim=>(0,4));

"""
    plotworld(w::Array{Cell,2}; title::String="")

Make a heatmap of the world
"""
function plotworld(w::Array{Cell,2}; title::String="")
    return heatmap(map(x->cdict[Int(x.state)], w); title=title, heatmapsettings...)
end

"""
    testworld()

Generate a test world and show it
"""
function testworld()
    N = 7
    world = zeros(Cell, N+2, N+2)
    world[2:4, 2:4] = reshape([Cell('A', 0) for _ in 1:3^2], 3, :) # agricultre
    world[2:4, 6:8] = reshape([Cell('N', 0) for _ in 1:3^2], 3, :) # nature
    world[6:8, 2:4] = reshape([Cell('F', 0) for _ in 1:3^2], 3, :) # forest
    world[6:8, 6:8] = reshape([Cell('U', 0) for _ in 1:3^2], 3, :) # urban
    p = plotworld(world)
    title!("Testworld: [agriculture, nature; forest, urban]")
    return p
end

@info "plotting testworld"
display(testworld())
@info "plotting evolution"
plot(1:size(record,1), record, marker=:circle, ylims=(-0.05,1.05),
    label=["urban" "forest" "nature" "agriculture"],
    xlabel="time [years]", ylabel="% of land usage")


using Pkg
Pkg.activate("c:/git/JLBoost")
@time using DataFrames
@time using JDF
@time using JLBoost, LossFunctions
using JLBoost: LogitLogloss

###############################################################################
# fitting GiveMeSomeCredit
###############################################################################
if !isdir("c:/data/GiveMeSomeCredit/cs-training.jdf")
	using CSV
	a = CSV.read("c:/data/GiveMeSomeCredit/cs-training.csv", missingstrings=["NA"])
	rename!(a, Symbol("")=>:column1)
	create_missing!(a, :MonthlyIncome)
	create_missing!(a, :NumberOfDependents)
	type_compress!(a, compress_float=true)
	savejdf("c:/data/GiveMeSomeCredit/cs-training.jdf", a)
end

###############################################################################
# testing find_best_split
###############################################################################
@time a = loadjdf("c:/data/GiveMeSomeCredit/cs-training.jdf")

loss = LogitLogloss()

features = setdiff(names(a), [:SeriousDlqin2yrs])
target = :SeriousDlqin2yrs
warmstart = fill(0, length(a.age))
verbose = false
lambda = 0
gamma = 0

using ForwardDiff: derivative

f(warmstart, target) = begin
	p = 1/(1+exp(-(warmstart)))
	-(target*log(p) + (1-target)*log(1-p))
end

f(target) = warmstart -> f(warmstart, target)

x = derivative.(f.(a.SeriousDlqin2yrs), fill(0.0, nrow(a)))

x1 = deriv.(loss, a.SeriousDlqin2yrs, fill(0.0, nrow(a)))

x == x1


module GeometryPrimitives
using StaticArrays, LinearAlgebra
using Statistics: mean
using Makie

export Shape, Shape1, Shape2, Shape3
export level, surfpt_nearby, normal, bounds, translate

abstract type Shape{N,N²} end # a solid geometric shape in N dimensions (N² = N*N is needed in some shapes, e.g., Cuboid)
const Shape1 = Shape{1,1}
const Shape2 = Shape{2,4}
const Shape3 = Shape{3,9}

Base.ndims(o::Shape{N}) where {N} = N

# The following functions return Any due to the limitations of Julia's dispatch system.
# Therefore, always call them with return type assertions.  See
# https://discourse.julialang.org/t/extending-base-in-type-stably/5341/12
# https://github.com/JuliaLang/julia/issues/23210
level(x::AbstractVector{<:Real}, o::Shape{N}) where {N} = level(SVector{N}(x), o)
Base.in(x::AbstractVector{<:Real}, o::Shape{N}) where {N} = level(x,o) ≤ 0
surfpt_nearby(x::AbstractVector{<:Real}, o::Shape{N}) where {N} = surfpt_nearby(SVector{N}(x), o)
normal(x::AbstractVector{<:Real}, o::Shape) = surfpt_nearby(x, o)[2]  # outward direction even for x inside o
translate(s::Shape{N}, ∆::AbstractVector{<:Real}) where {N} = translate(s, SVector{N}(∆))
translate(s::Shape{N}, ∆::SVector{N,<:Real}) where {N} = (s2 = deepcopy(s); s2.c += ∆; s2)  # default implementation

function orthoaxes(n::SVector{3,<:Real})
    u_temp = abs(n[3]) < abs(n[1]) ? SVector(0,0,1) : SVector(1,0,0)
    v = normalize(n × u_temp)
    u = v × n

    return u, v
end

function orthoaxes(n::SVector{N,<:Real}) where {N}
    u_temp = abs(n[3]) < abs(n[1]) ? SVector(0,0,1) : SVector(1,0,0)
    v = normalize(n × u_temp)
    u = v × n

    return u, v
end

include("hyper/hyper.jl")
include("planar/planar.jl")
include("prism/prism.jl")
include("util/util.jl")

end # module


"""
    codas(decisionMat, weights, fs)
Apply CODAS (COmbinative Distance-based ASsessment) method for a given matrix, weights and, type of criteria.
# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n alternatives and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fs::Array{Function,1}`: m-vector of type of criteria. The benefit criteria shown with "maximum", and the cost criteria shown with "minimum".
 - `tau::Float64`: tau parameter for the algorithm. The default is 0.02.
# Description
codas() applies the CODAS method to rank n alternatives subject to m criteria and criteria type vector.
# Output
- `::CODASResult`: CODASResult object that holds multiple outputs including scores and best index.
# Examples
```julia-repl
julia> decmat
7×5 Array{Float64,2}:
 60.0   0.4   2540.0   500.0   990.0
  6.35  0.15  1016.0  3000.0  1041.0
  6.8   0.1   1727.2  1500.0  1676.0
 10.0   0.2   1000.0  2000.0   965.0
  2.5   0.1    560.0   500.0   915.0
  4.5   0.08  1016.0   350.0   508.0
  3.0   0.1   1778.0  1000.0   920.0

julia> df = DataFrame(decmat)
7×5 DataFrame
 Row │ x1       x2       x3       x4       x5
     │ Float64  Float64  Float64  Float64  Float64
─────┼─────────────────────────────────────────────
   1 │   60.0      0.4    2540.0    500.0    990.0
   2 │    6.35     0.15   1016.0   3000.0   1041.0
   3 │    6.8      0.1    1727.2   1500.0   1676.0
   4 │   10.0      0.2    1000.0   2000.0    965.0
   5 │    2.5      0.1     560.0    500.0    915.0
   6 │    4.5      0.08   1016.0    350.0    508.0
   7 │    3.0      0.1    1778.0   1000.0    920.0
julia> result = codas(df, w, fs);
julia> result.bestIndex
2
julia> result.scores
7-element Array{Float64,1}:
  0.5121764914884954
  1.463300034504913
  1.0715325899642418
 -0.21246799780012637
 -1.8515205523193041
 -1.1716767695713806
  0.18865620373316055

```
# References
Keshavarz Ghorabaee, M., Zavadskas, E. K., Turskis, Z., & Antucheviciene, J. (2016). A new combinative distance-based assessment (CODAS) method for multi-criteria decision-making. Economic Computation & Economic Cybernetics Studies & Research, 50(3), 25-44.
"""
function codas(decisionMat::DataFrame, weight::Array{Float64,1}, fns::Array{Function,1}, tau::Float64=0.02)::CODASResult

    #mat = convert(Matrix, decisionMat)
    mat = Matrix(decisionMat)

    nrows, ncols = size(decisionMat)
    w = unitize(weight)

    colMax = colmaxs(decisionMat)
    colMin = colmins(decisionMat)

    A = similar(decisionMat)

    for i in 1:ncols
        if fns[i] == maximum
            @inbounds A[:, i] = decisionMat[:, i] ./ colMax[i]
        elseif fns[i] == minimum
            @inbounds A[:, i] = colMin[i] ./ decisionMat[:, i]
        end
    end

    wA = similar(A)

    for i in 1:ncols
        @inbounds wA[:, i] = A[:, i] .* w[i]
    end

    wAmin = colmins(wA)

    E = similar(A)
    for i in 1:nrows
        for j in 1:ncols
            E[i,j] = (wA[i,j] .- wAmin[j])^2
        end
    end

    Euc = zeros(Float64, nrows)
    for i in 1:nrows
        Euc[i] = sqrt(sum(E[i,:]))
    end

    T = similar(A)
    for i in 1:nrows
        for j in 1:ncols
            T[i,j] = abs(wA[i,j] .- wAmin[j])
        end
    end

    Tax = zeros(Float64, nrows)
    for i in 1:nrows
        Tax[i] = sum(T[i,:])
    end

    EA = zeros(Float64, nrows, nrows)
    TA = zeros(Float64, nrows, nrows)
    RA = zeros(Float64, nrows, nrows)

    for i in 1:nrows
        for j in 1:nrows
            EA[i,j] = Euc[i] - Euc[j]
            TA[i,j] = Tax[i] - Tax[j]
            if abs(Euc[i] - Euc[j]) < tau
                RA[i,j] = 0
            else
                RA[i,j] = 1
            end
        end
    end

    scores = zeros(Float64, nrows)
    for i in 1:nrows
        scores[i] = sum(EA[i,:] .+ (RA[i,:] .* TA[i,:]))
    end

    rankings = sortperm(scores)
    bestIndex = rankings |> last

    result = CODASResult(
        decisionMat,
        w,
        scores,
        rankings,
        bestIndex
    )

    return result
end


@reexport module Threshold

using Adapt
using Random
using CUDA

import ..SpikingNN: excite!, evaluate!, reset!, isactive

export AbstractThreshold, isactive

abstract type AbstractThreshold end

"""
    Ideal(vth::Real)

An ideal threshold spikes when `v > vth`.
"""
struct Ideal{T<:Real} <: AbstractThreshold
    vth::T
end

isactive(threshold::Ideal, t::Integer; dt::Real = 1.0) = false

"""
    evaluate!(threshold::Ideal, t::Real, v::Real; dt::Real = 1.0)
    (threshold::Ideal)(t::Real, v::Real; dt::Real = 1.0)
    evaluate!(thresholds::AbstractArray{<:Ideal}, t::Integer, v; dt::Real = 1.0)

Return `t` when `v > threshold.vth`.
"""
evaluate!(threshold::Ideal, t::Real, v::Real; dt::Real = 1.0) = (v >= threshold.vth) ? t : zero(t)
(threshold::Ideal)(t::Real, v::Real; dt::Real = 1.0) = evaluate!(threshold, t, v; dt = dt)
evaluate!(thresholds::T, t::Integer, v; dt::Real = 1.0) where T<:AbstractArray{<:Ideal} =
    Int.((v .>= thresholds.vth) .* adapt(typeof(v), fill(t, size(v))))

"""
    Poisson(ρ₀::Real, Θ::Real, Δᵤ::Real, rng::AbstractRNG)

Choose to output a spike based on a inhomogenous Poisson process given by

``X < \\mathrm{d}t \\: \\rho_0 \\exp\\left(\\frac{v - \\Theta}{\\Delta_u}\\right)``

where ``X \\sim \\mathrm{Unif}([0, 1])``.
Accordingly, `dt` must be set correctly so that the neuron does not always spike.

Fields:
- `ρ₀::Real`: baseline firing rate at threshold
- `Θ::Real`: firing threshold
- `Δᵤ::Real`: voltage resolution
- `rng`: random number generation
"""
struct Poisson{T<:Real,  RT<:AbstractRNG} <: AbstractThreshold
    ρ₀::T
    Θ::T
    Δᵤ::T
    rng::RT
end

function Poisson{T}(ρ₀::Real, Θ::Real, Δᵤ::Real; rng::RT = Random.GLOBAL_RNG) where {T<:Real, RT}
    Poisson{T,RT}(ρ₀, Θ, Δᵤ, rng)
end
Poisson(ρ₀::Real, Θ::Real, Δᵤ::Real; kwargs...) = Poisson{Real}(ρ₀, Θ, Δᵤ; kwargs...)

isactive(threshold::Poisson, t::Integer; dt::Real = 1.0) = true

"""
    poisson(baserate, theta, deltav, v; dt::Real, rng::AbstractRNG)
    poisson(baserate::AbstractArray{<:Real}, theta::AbstractArray{<:Real}, deltav::AbstractArray{<:Real}, v::AbstractArray{<:Real}; dt::Real, rng::AbstractRNG)
    poisson(baserate::CuVector{<:Real}, theta::CuVector{<:Real}, deltav::CuVector{<:Real}, v::CuVector{<:Real}; dt::Real, rng::AbstractRNG)

Evaluate inhomogeneous Poisson process threshold functions.
Modeled as

``X < \\mathrm{d}t \\rho_0 \\exp\\left(\\frac{v - \\Theta}{\\Delta_u}\\right)``

where ``X \\sim \\mathrm{Unif}([0, 1])``.

Use `CuVector` instead of `Vector` to evaluate on GPU.

# Fields
- `baserate`: base line firing rate
- `theta`: threshold potential
- `deltav`: potential resolution
- `v`: current membrane potential
- `dt`: simulation timestep
- `rng`: random number generation
"""
function poisson(baserate, theta, deltav, v; dt::Real, rng::AbstractRNG = Random.GLOBAL_RNG)
    rho = baserate * exp((v - theta) / deltav)

    return rand(rng) < rho * dt
end
function poisson(baserate::AbstractArray{<:Real}, theta::AbstractArray{<:Real}, deltav::AbstractArray{<:Real}, v::AbstractArray{<:Real}; dt::Real, rng::AbstractRNG = Random.GLOBAL_RNG)
    rho = baserate .* exp.((v .- theta) ./ deltav)

    return rand(rng, length(rho)) .< rho .* dt
end
function poisson(baserate::CuVector{<:Real}, theta::CuVector{<:Real}, deltav::CuVector{<:Real}, v::CuVector{<:Real}; dt::Real, rng::AbstractRNG = Random.GLOBAL_RNG)
    rho = baserate .* exp.((v .- theta) ./ deltav)

    return CuArrays.rand(rng, length(rho)) .< rho .* dt
end

"""
    evaluate!(threshold::Poisson, t::Integer, v::Real; dt::Real = 1.0)
    (::Poisson)(t::Integer, v::Real; dt::Real = 1.0)
    evaluate!(thresholds::AbstractArray{<:Poisson}, t::Integer, v; dt::Real = 1.0)

Evaluate Poisson threshold function. See [`Threshold.Poisson`](@ref).
"""
evaluate!(threshold::Poisson, t::Integer, v::Real; dt::Real = 1.0) =
    poisson(threshold.ρ₀, threshold.Θ, threshold.Δᵤ, v; dt = dt, rng = threshold.rng) ? t : zero(t)
(threshold::Poisson)(t::Integer, v::Real; dt::Real = 1.0) = evaluate!(threshold, t, v; dt = dt)
evaluate!(thresholds::T, t::Integer, v; dt::Real = 1.0) where T<:AbstractArray{<:Poisson} =
    Int.(poisson(thresholds.ρ₀, thresholds.Θ, thresholds.Δᵤ, v; dt = dt, rng = first(thresholds.rng)) .* adapt(typeof(v), fill(t, size(v))))

end


module FDA

using BSplines

struct basisobj
    basis_type
    basis
end

struct fdobj
    basis_type
    basis
    coefs
end

function second_order_diff(f, x, h)
    n = f(x+h) - 2*f(x) + f(x-h)
    return n / (h^2)
end

function estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
    dpsi = zeros(length(basis), length(argvals))
    if basis_type == "BSpline"
        for i in 1:length(argvals)
            tpsi = Psi_matrix[i,:]
            tid = [j for j in 1:length(tpsi) if tpsi[j]!=0]
            t_time = argvals[i]
            temp_deriv = bsplines(basis, t_time, BSplines.Derivative(nderiv))
            for k in tid
                dpsi[k, i] = temp_deriv[:,1][k]
            end
        end
    elseif basis_type == "Fourier"
        h = (argvals[2] - argvals[1])/10
        for i in 1:length(basis)
            tf = basis[i]
            temp_deriv = [second_order_diff(tf, t, h) for t in argvals]
            dpsi[i,:] = temp_deriv
        end
    end

    R_matrix = dpsi * transpose(dpsi)
    R_matrix = R_matrix ./ length(argvals)
    return R_matrix
end

function gen_bspline_basis(range_eval, nbasis, order=4)
    nbreaks = nbasis - order + 2
    break_points = range(range_eval[1], range_eval[2]; length = nbreaks)
    basis = BSplineBasis(order, break_points)
    r = basisobj("BSpline", basis)
    return r
end

function gen_fourier_basis(range_eval, nbasis)
    normalizer = (range_eval[2] - range_eval[1])
    basis = []

    function tf1(t)
        return 1/normalizer
    end

    push!(basis, tf1)

    n = Int(ceil((nbasis-1)/2))
    for i in 1:n
        function tfa(t)
            return sin(2*pi*i*t/normalizer)
        end
        push!(basis, tfa)
        function tfb(t)
            return cos(2*pi*i*t/normalizer)
        end
        push!(basis, tfb)
    end
    basis = basis[1:nbasis]

    r = basisobj("Fourier", basis)
end

function get_coef_mat(argvals, basisobj, lambda, nderiv)
    basis = basisobj.basis
    basis_type = basisobj.basis_type
    if basis_type == "BSpline"
        Psi_matrix = basismatrix(basis, argvals)
        R_matrix = estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
    elseif basis_type == "Fourier"
        Psi_matrix = zeros(length(argvals), length(basis))
        for i in 1:length(basis)
            tf = basis[i]
            Psi_matrix[:,i] = tf.(argvals)
            R_matrix = estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
        end
    end

    coef_mat = inv(transpose(Psi_matrix) * Psi_matrix .+ R_matrix .* lambda) *
            transpose(Psi_matrix)
    return coef_mat

end

function Data2fd(argvals, y, basisobj;
    lambda=3e-8/(maximum(argvals) - minimum(argvals)),
    nderiv=2)

    basis_type = basisobj.basis_type
    basis = basisobj.basis

    coefs = zeros(length(basis), ncol(y))
    coef_mat = get_coef_mat(argvals, basisobj,lambda, nderiv)
    for i in 1:ncol(y)
        coefs[:,i] = coef_mat * y[:,i]
    end

    res = fdobj(basis_type, basis, coefs)
    return res
end

function eval_fd(eval_grid, fdobj)
    basis = fdobj.basis
    coefs = fdobj.coefs
    if fdobj.basis_type == "BSpline"
        Psi_matrix = basismatrix(basis, eval_grid)
    elseif fdobj.basis_type == "Fourier"
        Psi_matrix = zeros(length(eval_grid), length(basis))
        for i in 1:length(basis)
            tf = basis[i]
            Psi_matrix[:,i] = tf.(eval_grid)
        end
    end
    obs = Psi_matrix * coefs
    return obs
end

end


# This file includes the simulation of four node network simulation

import Cairo, Fontconfig
using GraphPlot
using Graphs
using SimpleWeightedGraphs
using LinearAlgebra
using DynamicalNetworks
using DifferentialEquations
using PyPlot
pygui(true)
close("all")

# Helper functions
getwave(x, i, d) = getindex.(x, (i - 1) * d + 1)

# Node settings
nodetype = Lorenz           # Node type
P = diagm([1.0, 0.0, 0])    # State coupling matrix
ηt = 10.0                   # Threshold for maximum Lyapuno exponents

# Topology
n = 5
graph = SimpleWeightedGraph(cycle_graph(n))
add_edge!(graph, 1, 3, 1)
add_edge!(graph, 1, 4, 1)
add_edge!(graph, 2, 5, 5)
gplot(graph, nodelabel = 1:nv(graph), layout = circular_layout)

# Coupling strength for the first symmetry
W1 = laplacian_matrix(graph)
Q1 = [
    1 0 0 0 0
    0 0 0 0 1
    0 0 0 1 0
    0 0 1 0 0
    0 1 0 0 0
]
M1 = eigvecs(Q1)
G1 = inv(M1) * (W1 ./ diag(W1)) * M1
B1 = G1[1:2, 1:2]
D1 = G1[3:5, 3:5]
λt1 = eigvals(B1)
λs1 = eigvals(D1)
λt1[1] > λs1[2] || error("Synchronization is not possible")
ϵr1 = ηt ./ [λt1[1], λs1[2]]

# Coupling strength for the second symmetry
graph = SimpleWeightedGraph(cycle_graph(n))
add_edge!(graph, 1, 3, 1)
add_edge!(graph, 1, 4, 1)
add_edge!(graph, 2, 5, 1)
W2 = laplacian_matrix(graph)
Q2 = [
    1 0 0 0 0
    0 0 0 0 1
    0 0 0 1 0
    0 0 1 0 0
    0 1 0 0 0
]
M2 = eigvecs(Q2)
G2 = inv(M2) * (W2 ./ diag(W2)) * M2
B2 = G2[1:2, 1:2]
D2 = G2[2:5, 2:5]
λt2 = eigvals(B2)
λs2 = eigvals(D2)
# λt2[1] > λs2[2] || error("Synchronization is not possible")
ϵr2 = ηt ./ [λt2[1], λs2[2]]

ϵl = maximum([ϵr1[1], ϵr2[1]])
ϵh = minimum([ϵr2[2], ϵr2[2]])
ϵ = ϵl + abs(ϵh - ϵl) * 0.5

# Callback
nbits = 10
tbit = 100.0
bits = rand(Bool, nbits)
ti = 0.0
tf = nbits * tbit
tswitch = ti:tbit:tf-tbit
srcnode1, dstnode1 = 2, 5     # Control link nodes
high = 5.0          # Control coupling strength
low = 1.0           # Control coupling strength
condition(u, t, integrator) = integrator.t ∈ tswitch
function action!(integrator)
    net, bits = integrator.sol.prob.p
    if popfirst!(bits)
        add_edge!(net.graph, srcnode1, dstnode1, high)
    else
        add_edge!(net.graph, srcnode1, dstnode1, low)
    end
end
callback = DiscreteCallback(condition, action!)

# Network
σ = 0.0  # Noise strength
H = getnoisemat(graph) * σ
net = SDENetwork(nodetype, ϵ, graph, H, P)
sol = solvenet!(net, (ti, tf), (copy(bits),), solverkwargs = (saveat = 0.001,), probkwargs = (callback = callback, tstops = tswitch))
t = sol.t
x = sol.u

# Plots
d = dimension(nodetype)
subplot(221)
plot(t, abs.(getwave(x, 1, d) - getwave(x, 2, d)))
axvline.(tswitch, linestyle = "dashed", color = "red")
xlabel("Time")
ylabel("e12")
title(join(string.(Int.(bits))))
subplot(222)
plot(t, abs.(getwave(x, 2, d) - getwave(x, 5, d)))
axvline.(tswitch, linestyle = "dashed", color = "red")
xlabel("Time")
ylabel("e25")
title(join(string.(Int.(bits))))
subplot(223)
plot(t, getwave(x, 1, d))
xlabel("Time")
ylabel("x1")
subplot(224)
plot(getindex.(x, 1), getindex.(x, 2))
xlabel("x1")
ylabel("x2")
tight_layout()

figure()
plot(λt1, zeros(length(λt1)), marker = "*", label = L"\lambda_t")
plot(λs1, ones(length(λs1)), marker = "o", label = L"\lambda_s")
axvline.([λt1[1], λs1[2]], linestyle = "dashed", color = "red")
axvline(ηt / ϵ, linestyle = "dashed", color = "green")
xlabel(L"\lambda")
title("Eigenvalue distribution")
legend()

figure()
plot(λt2, zeros(length(λt2)), marker = "*", label = L"\lambda_t")
plot(λs2, ones(length(λs2)), marker = "o", label = L"\lambda_s")
axvline.([λt2[1], λs2[2]], linestyle = "dashed", color = "red")
axvline(ηt / ϵ, linestyle = "dashed", color = "green")
xlabel(L"\lambda")
title("Eigenvalue distribution")
legend()


using StochasticProcesses

bm = BrownianMotion()
t = linspace(0, 1, 10000)
cumsim(bm, t, 100)[end,:]

println("cumsim(BrownianMotion) 10000x10000:")
for i in 1:5
  @time cumsim(bm, t, 10000)[end,:]
end


"""
    result = pslq(x::Vector; criteria)

Perform PSLQ iterations for the vector x until termination criteria are fulfilled.
`result` always contains the fields: `code`, `iter`, `A`, `B`, `H`.
The returned result is a named tuple and may contain different result variants:

`result.code == 1`:

 A solution `m` was found `|m' * x| <= cr.atol + cr.rtol*|x|`
 with `m = result.B[:,result.col]`

`result.code == 2`:
    No solution `m` possible with `|m| <= cr.solutionmax`

`result.code == 3`:
    No decision after number of iterations exceeds `cr.itermax`

`result.code == 4`:
    Inaccuracy in calculation - retry with higher precision input data

The input criteria are transformed to object `cr` by method `make_criteria(criteria, x)`.
"""
function pslq(x::Vector{T}; criteria=nothing) where T<:Real
    n = length(x)
    H = make_H(x)
    # criteria = make_criteria(criteria, x)
    A = Matrix(I*integertype(T)(1), n, n)
    B = copy(A)
    D = UnitLowerTriangular(similar(B)) # working area
    iter = 0
    result = (code=0, result=similar(B, n, 0))
    while iter < criteria.itermax && result.code == 0
        iter += 1
        pslq_step!(H, A, B, D)
        result = convergence(x, H, A, B, iter, criteria)
    end
    result
end

"""
    result = convergence(x, H, A, B, iter, criteria)

Termination check and provide result for `pslq`.
The arguments `H`, `A`, `B` are as described in the algoritm paper.
`result.code == 0` means that iterations should continue.
"""
function convergence(x::Vector{T}, H::AbstractMatrix, A::Matrix, B::Matrix,
                     iter::Integer, criteria::NamedTuple) where T
    n = size(x, 1)
    function xtb(j)
        s1 = s2 = T(0)
        for k = 1:n
            bkjx = B[k,j] * x[k]
            s1 += bkjx
            s2 += abs(bkjx)
        end
        s1, s2
    end
    resvec = []
    for j = 1:n
        s1, s2 = xtb(j)
        if abs(s1) <= s2 * criteria.rtol + criteria.atol
            push!(resvec, (code=1, col=j, xb=s1, xc=s2))
        end
    end
    if length(resvec) >= 1
        r = resvec[1]
        return (code=r.code, iter=iter, col=r.col, xb=r.xb, xc=r.xc, A=A, B=B, H=H)
    end
    (code=0, iter=iter, A=A, B=B, H=H)
end

function make_criteria(criteria, x)
    return (rtol=10*eps(eltype(x)), atol=10*eps(norm(x)), itermax=10000)
end


@testset "isolateroots" begin
    prec = 128
    RR = RealField(prec)

    problems = [(cos, 0, 10, [RR(π)/2, 3RR(π)/2, 5RR(π)/2]),
                (x -> (x - 1)*(x - 2)*(x - 3), 0, 5, [1, 2, 3]),
                (x -> exp(x) - 1, -eps(), 5, [0])]

    for (f, a, b, roots) in problems
        a = RR(a)
        b = RR(b)
        roots = RR.(roots)

        for evaltype in (:ball, :taylor)
            found, flags = isolateroots(f, a, b, evaltype = evaltype)

            @test length(found) == length(roots) && length(flags) == length(roots)
            if evaltype == :ball
                @test !any(flags)
            else
                @test all(flags)
            end
            @test all(contains(ArbTools.setinterval(found[i]...), roots[i]) for i in 1:length(roots))
        end
    end
end


# Reactive message-passing in linear Gaussian dynamical system
# Experiment to pass messages without a schedule
#
# Wouter Kouw, BIASlab
# 12-08-2019

using Random
using Revise
using CPUTime
using JLD
using ForneyLab
using Plots
pyplot()

# Functions for generating data
include(joinpath(@__DIR__, "../gen_data.jl"))

"""
Experiment parameters
"""

# Signal time horizon
T = 500

# Reaction-time clock
TT = 10

# Known transition and observation matrices
gain = 0.7
emission = 1.0

# Noise parameters
process_noise = 0.6
measurement_noise = 0.3

# Clamped parameters
params_x0 = [0.0, 1.0]
params_γ0 = [1.0, 1.0]

# Generate data
observed, hidden = gendata_LGDS(gain,
                                emission,
                                process_noise,
                                measurement_noise,
                                params_x0[1],
                                params_x0[2],
                                time_horizon=T)

# # Write data set to file
# save(pwd()*"/experiment_infer-processnoise/data/LGDS_T"*string(T)*".jld", "observed", observed, "hidden", hidden)

# # Load previously generated data
# X = load(pwd()*"/experiment_infer-processnoise/data/LGDS_T"*string(T)*".jld")
# observed = X["observed"]
# hidden = X["hidden"]

"""
Model specification
"""

# Start timer
CPUtic()

g = FactorGraph()

@RV γ ~ Gamma(params_γ0[1], params_γ0[2])
@RV x_0 ~ GaussianMeanVariance(params_x0[1], params_x0[2])

global x = Vector{Variable}(undef, T)
global y = Vector{Variable}(undef, T)
global x_tmin = x_0
for t = 1:T

    @RV x[t] ~ GaussianMeanPrecision(x_tmin, γ)
    @RV y[t] ~ GaussianMeanVariance(x[t], measurement_noise)

    global x_tmin = x[t]

    placeholder(y[t], :y, index=t)
end

"""
Mean-field algorithm generation
"""

# Define recognition factorization
RecognitionFactorization()

q_γ = RecognitionFactor(γ)
q_x_0 = RecognitionFactor(x_0)

q_x = Vector{RecognitionFactor}(undef, T)
for t=1:T
    q_x[t] = RecognitionFactor(x[t])
end

# Compile algorithm
algo_γ_mf = variationalAlgorithm(q_γ, name="PMF")
algo_x_mf = variationalAlgorithm([q_x_0; q_x], name="XMF")
algo_F_mf = freeEnergyAlgorithm(name="MF");

"""
Run inference algorithm
"""

# Load algorithm
eval(Meta.parse(algo_γ_mf))
eval(Meta.parse(algo_x_mf))
eval(Meta.parse(algo_F_mf))

# Initialize data
global data = Dict(:y => observed)

# Initial recognition distributions
global marginals_mf = Dict{Symbol, ProbabilityDistribution}()
global marginals_mf[:γ] = vague(Gamma)
for t = 0:T
    global marginals_mf[:x_*t] = vague(GaussianMeanPrecision)
end

# Preallocate
global estimated_states = zeros(T,2,TT)
global estimated_noises = zeros(T,2,TT)

# Run algorithm
global F_mf = Vector{Float64}(undef, TT)
for tt = 1:TT
    stepPMF!(data, marginals_mf)
    stepXMF!(data, marginals_mf)

    for t = 1:T
        estimated_states[t,1,tt] = mean(marginals_mf[:x_*t])
        estimated_states[t,2,tt] = sqrt(var(marginals_mf[:x_*t]))
        estimated_noises[t, 1, tt] = mean(marginals_mf[:γ])
        estimated_noises[t, 2, tt] = sqrt(var(marginals_mf[:γ]))
    end

    global F_mf[tt] = freeEnergyMF(data, marginals_mf)
end

CPUtoc()

"""
Visualize experimental results
"""

# Visualize final estimates over time
plot(hidden[2:end], color="red", label="states", ylims=[-5, 2.5])
plot!(estimated_states[:,1,end], color="blue", label="estimates")
plot!(estimated_states[:,1,end],
      ribbon=[estimated_states[:,2,end], estimated_states[:,2,end]],
      linewidth=2,
      color="blue",
      fillalpha=0.2,
      fillcolor="blue",
      label="")
scatter!(observed, color="black", label="observations")
xlabel!("time (t)")
title!("State estimates, q(x_t)")
savefig(pwd()*"/experiment_infer-processnoise/viz/baseline_state_estimates.png")

# Visualize final noise estimates over time
plot(inv(process_noise)*ones(T,1), color="black", label="true process noise", linewidth=2, ylims=[0., 2.0])
plot!(estimated_noises[:,1,end], color="blue", label="estimates")
plot!(estimated_noises[:,1,end],
      ribbon=[estimated_noises[:,2,end], estimated_noises[:,2,end]],
      linewidth=2,
      color="blue",
      fillalpha=0.2,
      fillcolor="blue",
      label="")
xlabel!("time (t)")
title!("Noise estimates, q(γ_x)")
savefig(pwd()*"/experiment_infer-processnoise/viz/baseline_noise_estimates.png")


@namedargs mutable struct DICparam
  p::Vector{Matrix{Float64}}
  mu::Vector{Matrix{Float64}}
  sig::Vector{Vector{Float64}}
  y::Vector{Matrix{Float64}}
end

function computeLoglikeDIC(d::Data, param::DICparam)
  ll = 0.0

  for i in 1:d.I
    for j in 1:d.J
      for n in 1:d.N[i]
        y_inj_is_missing = (d.m[i][n, j] == 1)

        # NOTE: Refer to `../compute_loglike.jl` for reasoning.
        if y_inj_is_missing

          # Compute p(m_inj | y_inj, theta) term.
          ll += log(param.p[i][n, j])
        end

        # Compute p(y_inj | theta) term.
        ll += logpdf(Normal(param.mu[i][n, j], param.sig[i][n]),
                     param.y[i][n, j])
      end
    end
  end

  return ll
end

function updateParams(d::MCMC.DICstream{DICparam}, param::DICparam)
  d.paramSum.p += param.p
  d.paramSum.mu += param.mu
  d.paramSum.sig += param.sig
  d.paramSum.y += param.y

  return
end

function paramMeanCompute(d::MCMC.DICstream{DICparam})::DICparam
  return DICparam(d.paramSum.p / d.counter,
                  d.paramSum.mu / d.counter,
                  d.paramSum.sig / d.counter,
                  d.paramSum.y / d.counter)
end


include("HexGrid.jl")

function dist(p1::Array{Float64},p2::Array{Float64})
    d = sqrt((p1[1]-p2[1])^2+(p1[2]-p2[2])^2)
    return d
end

                                # for the intensity to index coding (never equals one)
gauss(d::Float64, node_size::Float64) = 0.999999 * exp(-2*(d/node_size)) # considered 2sigma as d in the gaussian

function input_spikes(obs::Array{Float64}, grid_cells::Array{HexGrid})
    pos = [obs[1],obs[3]%pi]     # obs pos and theta
    vel = [obs[2],obs[4]]     # obs pos_v and theta_v
    bound = [-15.0, 15.0]
    pos_scaled = (pos .- bound[1]) / (bound[2]-bound[1])
    vel_scaled = (vel .- bound[1]) / (bound[2]-bound[1])
    n_grid = length(grid_cells)
    spikes = Tuple{Int64,Float64}[]
    for i in 1:n_grid
        d = grid_cells[i].node_size
        close_pos = filter(p -> dist(p,pos_scaled) < 1.5*d, grid_cells[i].grid) # 1.5 is to reach up to 3sigmas
        close_vel =  filter(p -> dist(p,vel_scaled) < 1.5*d, grid_cells[i].grid)
        if length(close_pos) != 0
            neuron = i
            d_pos = dist(pos_scaled,close_pos[1])
            intensity = gauss(d_pos,d)
            push!(spikes,(neuron,intensity))
        end
        if length(close_vel) != 0
            neuron = n_grid + i
            d_vel = dist(vel_scaled,close_vel[1])
            intensity = gauss(d_vel,d)
            push!(spikes,(neuron,intensity))
        end
    end
    return spikes
end

function latency_spikes(obs::Array{Float64}, grid_cells::Array{HexGrid}, win_size::Int64)
    pos = [obs[1],obs[3]%pi]     # obs pos and theta
    vel = [obs[2],obs[4]]     # obs pos_v and theta_v
    bound = [-15.0, 15.0]
    pos_scaled = (pos .- bound[1]) / (bound[2]-bound[1])
    vel_scaled = (vel .- bound[1]) / (bound[2]-bound[1])
    n_grid = length(grid_cells)
    spikes = [Int64[] for _ in 1:win_size]
    for i in 1:n_grid
        d = grid_cells[i].node_size
        close_pos = filter(p -> dist(p,pos_scaled) < 1.25*d, grid_cells[i].grid) # 1.5 is to reach up to 3sigmas
        close_vel =  filter(p -> dist(p,vel_scaled) < 1.25*d, grid_cells[i].grid)
        if length(close_pos) != 0
            neuron = i
            d_pos = dist(pos_scaled,close_pos[1])
            intensity = gauss(d_pos,d)
            ind = win_size - floor(Int64, win_size * intensity)
            push!(spikes[ind],neuron)
        end
        if length(close_vel) != 0
            neuron = n_grid + i
            d_vel = dist(vel_scaled,close_vel[1])
            intensity = gauss(d_vel,d)
            ind = win_size - floor(Int64, win_size * intensity)
            push!(spikes[ind],neuron)
        end
    end
    return spikes
end

# # %% test input_spikes
#
# include("CartpoleGridGenerator.jl")
# grids = grid_gen()
#
# # %%
#
# obs = [0.523324424,0.323242424243,0.4242424225,0.224253525253]
#
#
# @time spikes = latency_spikes(obs,grids,7)


# Following order of operations on of solid geometries:
#   1. bottomArea(shape): returns the bottom area of a solid shape
#   2. topArea(shape): returns the top area of a solid shape
#   3. volume(shape): computes the volume of a solid shape
#   3. longestEdge(shape): returns is the maximum of all directions (especially needed for FileMesh)
#   5. lengthGeo(shape): returns length of geometries in z direction
#   6. centroid(shape): returns position vector from solid reference to centroid
#   7. inertiaMatrix(shape, mass): returns inertia matrix of solid with respect to reference frame
#      Useful web pages:
#         https://en.wikipedia.org/wiki/List_of_moments_of_inertia
#         https://en.wikipedia.org/wiki/List_of_second_moments_of_area

using LinearAlgebra
EYE3() = Matrix(1.0I,3,3)

const InertiaMatrix  = SMatrix{3,3,Float64,9}

### ------------------------ bottom and top area -------------------------------
bottomArea(shape::Modia3D.AbstractGeometry)  = error(typeof(shape), ": has no bottom or top area!")
bottomArea(shape::Box)       = shape.lengthX*shape.lengthY
bottomArea(shape::Cylinder)  = pi*((shape.diameter/2)^2 - (shape.innerDiameter/2)^2)
bottomArea(shape::Cone)      = pi*(shape.diameter/2)^2
bottomArea(shape::Beam)      = shape.length*shape.width + pi*(shape.width/2)^2

## top area: Cone for other geometrie take bottom area (if defined)
topArea(shape::Modia3D.AbstractGeometry) = bottomArea(shape)
topArea(shape::Cone) = pi*(shape.topDiameter/2)^2

### ------------------------- volume -------------------------------------------
## volume: Sphere, Ellipsoid, Capsule, Cone, FileMesh
## others are computed via bottomArea * lengthGeo
"""
    V = volume(shape)

Return the volume of the solid shape `shape::Modia3D.AbstractGeometry` in [m^3].
"""
volume(shape::Modia3D.AbstractGeometry) = bottomArea(shape)*lengthGeo(shape)
volume(shape::Sphere)    = 4/3*pi*(shape.diameter/2)^3
volume(shape::Ellipsoid) = 4/3*pi*shape.lengthX/2*shape.lengthY/2*shape.lengthZ/2
volume(shape::Cone)      = pi/12*shape.length*(shape.diameter^2 + shape.diameter*shape.topDiameter + shape.topDiameter^2)

function volume(shape::Capsule)
    r = shape.diameter/2
    h = shape.length
    return pi*r^2*h + 4/3*pi*r^3
end

function volume(shape::FileMesh)
  if !isempty(shape.facesIndizes)
    return shape.volume
  else
    println("FileMesh: ", shape.filename, ". The surface areas must be triangular, and each triangle should be specified in right-handed/counter-clockwise order. Otherwise it is not possible to compute a volume.")
    return nothing
end; end

### --------------------------- longest edge  ----------------------------------
#   is the maximum of all directions (especially needed for FileMesh)
longestEdge(shape::Sphere)    = shape.diameter
longestEdge(shape::Ellipsoid) = max(shape.lengthX, shape.lengthY, shape.lengthZ)
longestEdge(shape::Box)       = max(shape.lengthX, shape.lengthY, shape.lengthZ)
longestEdge(shape::Cylinder)  = max(shape.diameter, shape.length)
longestEdge(shape::Cone)      = max(shape.diameter, shape.length)
longestEdge(shape::Capsule)   = shape.length + shape.diameter
longestEdge(shape::Beam)      = max(shape.length + shape.width, shape.thickness)
longestEdge(shape::FileMesh)  = shape.longestEdge
longestEdge(shape::Modia3D.AbstractGeometry) = max(shape.Dx,shape.Dy,shape.Lz)

### -------------------------- length of shape ------------------------------
lengthGeo(shape::Modia3D.AbstractGeometry) = shape.Lz
lengthGeo(shape::Sphere)                   = shape.diameter
lengthGeo(shape::Ellipsoid)                = shape.lengthZ
lengthGeo(shape::Box)                      = shape.lengthZ
lengthGeo(shape::Cylinder)                 = shape.length
lengthGeo(shape::Cone)                     = shape.length
lengthGeo(shape::Beam)                     = shape.thickness
lengthGeo(shape::FileMesh)                 = error("lengthGeo(FileMesh) is not implemented yet!")

### ------------------------ centroid of geometries ----------------------------
#   it is the zero vector, only for Cone and FileMesh it is different
"""
    r = centroid(shape)

Return position vector from solid reference frame to [centroid](https://en.wikipedia.org/wiki/Centroid)
of solid `shape::Modia3D.AbstractGeometry` in [m]. If the solid has a uniform density,
the centroid is identical to the *center of mass*.
"""
@inline centroid(shape::Sphere) = Modia3D.ZeroVector3D
@inline centroid(shape::Ellipsoid) = Modia3D.ZeroVector3D
@inline centroid(shape::Box) = Modia3D.ZeroVector3D
@inline centroid(shape::Cylinder) = Modia3D.ZeroVector3D
@inline centroid(shape::Capsule) = Modia3D.ZeroVector3D
@inline centroid(shape::Beam) = Modia3D.ZeroVector3D

@inline centHeight(shape::Cone) = shape.length/4*(shape.diameter^2 + 2*shape.diameter*shape.topDiameter + 3*shape.topDiameter^2)/(shape.diameter^2 + shape.diameter*shape.topDiameter + shape.topDiameter^2)  # https://mathworld.wolfram.com/ConicalFrustum.html
@inline function centroid(shape::Cone)
    if shape.axis == 1
        return SVector{3,Float64}([centHeight(shape), 0.0, 0.0])
    elseif shape.axis == 2
        return SVector{3,Float64}([0.0, centHeight(shape), 0.0])
    else
        return SVector{3,Float64}([0.0, 0.0, centHeight(shape)])
    end
end

@inline function centroid(shape::FileMesh)
  if !isempty(shape.facesIndizes)
    return shape.centroidAlgo
  else
    return shape.centroid
end; end

### ----------------------- inertia matrix of geometries -----------------------
"""
   I = inertiaMatrix(shape, mass)

Return [inertia matrix] (https://en.wikipedia.org/wiki/Moment_of_inertia) `I` of solid
`shape::Modia3D.AbstractGeometry` w.r.t. center of mass resolved in object3d frame in
[kg*m^2] as `SMatrix{3,3,Float64,9}`. Hereby it is assumed that `shape` has uniform
density and `mass` is the mass of `shape` in [kg].
"""
inertiaMatrix(shape::Sphere, mass::Number) = InertiaMatrix(mass/10*shape.diameter^2*EYE3())

inertiaMatrix(shape::Ellipsoid, mass::Number) =
                InertiaMatrix(mass/20 * Diagonal([shape.lengthY^2 + shape.lengthZ^2, shape.lengthX^2 + shape.lengthZ^2, shape.lengthX^2 + shape.lengthY^2]))

inertiaMatrix(shape::Box, mass::Number) =
                InertiaMatrix(1/12*mass * Diagonal([shape.lengthY^2 + shape.lengthZ^2, shape.lengthX^2 + shape.lengthZ^2, shape.lengthX^2 + shape.lengthY^2]))

function inertiaMatrix(shape::Cylinder, mass::Number)  # https://en.wikipedia.org/wiki/List_of_moments_of_inertia
    Iax  = mass/2*((shape.diameter/2)^2 + (shape.innerDiameter/2)^2)
    Irad = mass/12*(3*((shape.diameter/2)^2 + (shape.innerDiameter/2)^2) + shape.length^2)
    if shape.axis == 1
        return InertiaMatrix(Diagonal([Iax, Irad, Irad]))
    elseif shape.axis == 2
        return InertiaMatrix(Diagonal([Irad, Iax, Irad]))
    else
        return InertiaMatrix(Diagonal([Irad, Irad, Iax]))
    end
end

function inertiaMatrix(shape::Cone, mass::Number)  # https://en.wikipedia.org/wiki/List_of_moments_of_inertia
    if shape.topDiameter == 0.0
        Iax  = 3/40*mass*shape.diameter^2
        Irad = 3/80*mass*(shape.diameter^2 + shape.length^2)
    else
        fullLength = shape.length*shape.diameter/(shape.diameter - shape.topDiameter)  # length from base to apex
        pruneLength = fullLength - shape.length                                  # length from top to apex
        fullVolume = pi/12*fullLength*shape.diameter^2                           # volume from base to apex
        pruneVolume = pi/12*pruneLength*shape.topDiameter^2                      # volume from top to apex
        fullMass = mass*fullVolume/(fullVolume - pruneVolume)                  # mass from base to apex
        pruneMass = mass*pruneVolume/(fullVolume - pruneVolume)                # mass from top to apex
        fullIrad = 3/80*fullMass*(shape.diameter^2 + 16*fullLength^2)            # moment of inertia from base to apex w.r.t. apex
        pruneIrad = 3/80*pruneMass*(shape.topDiameter^2 + 16*pruneLength^2)      # moment of inertia from top to apex w.r.t. apex
        dist = fullLength - centHeight(shape)                                    # distance between apex and centroid
        Iax  = 3/40*(fullMass*shape.diameter^2 - pruneMass*shape.topDiameter^2)
        Irad = (fullIrad - pruneIrad) - mass*dist^2                            # radial moment of intertia w.r.t. centroid
    end
    if shape.axis == 1
        return InertiaMatrix(Diagonal([Iax, Irad, Irad]))
    elseif shape.axis == 2
        return InertiaMatrix(Diagonal([Irad, Iax, Irad]))
    else
        return InertiaMatrix(Diagonal([Irad, Irad, Iax]))
    end
end

function inertiaMatrix(shape::Capsule, mass::Number)
    rho = mass/volume(shape)
    rad = shape.diameter/2

    lenCyl = shape.length
    volCyl = pi*rad^2*lenCyl
    massCyl = volCyl*rho

    volHalfSph = 4/3*pi*rad^3/2
    massHalfSph = volHalfSph*rho

    # moments of inertia of cylinder w.r.t. center of mass of capsule
    IaxCyl = massCyl/2*rad^2
    IradCyl = massCyl/12*(3*rad^2 + lenCyl^2)

    # moment of inertia of half sphere w.r.t. center of full sphere (same formula as for full sphere because of half mass)
    IaxHalfSph = 2/5*massHalfSph*rad^2

    # distance between centers of mass of full sphere and of half sphere https://en.wikipedia.org/wiki/List_of_centroids
    distCoM = 3/8*rad

    # radial moment of inertia of half sphere w.r.t. center of mass of half sphere
    IradHalfSph = IaxHalfSph - massHalfSph*distCoM^2

    # radial moment of inertia of half cylinder w.r.t. center of mass of capsule
    IradHalfSph = IradHalfSph + massHalfSph*(distCoM + lenCyl/2)^2

    Iax  = IaxCyl  + 2*IaxHalfSph
    Irad = IradCyl + 2*IradHalfSph
    if shape.axis == 1
        return InertiaMatrix(Diagonal([Iax, Irad, Irad]))
    elseif shape.axis == 2
        return InertiaMatrix(Diagonal([Irad, Iax, Irad]))
    else
        return InertiaMatrix(Diagonal([Irad, Irad, Iax]))
    end
end

function inertiaMatrix(shape::Beam, massGeo::Number)
    rho = massGeo/volume(shape)

    volBox = shape.length*shape.thickness*shape.width
    massBox = volBox*rho

    radCyl = shape.width/2
    volHalfCyl = pi*radCyl^2*shape.thickness/2
    massHalfCyl = volHalfCyl*rho

    # moments of inertia of box w.r.t. center of mass of beam
    IlBox = massBox/12*(shape.width^2 + shape.thickness^2)
    IwBox = massBox/12*(shape.thickness^2 + shape.length^2)
    ItBox = massBox/12*(shape.length^2 + shape.width^2)

    # moments of inertia of half cylinder w.r.t. center of full cylinder (same formulas as for full cylinder because of half mass)
    IlCylHalf = massHalfCyl/12*(3*radCyl^2 + shape.thickness^2)
    IwCylHalf = massHalfCyl/12*(3*radCyl^2 + shape.thickness^2)
    ItCylHalf = massHalfCyl/2*radCyl^2

    # distance between centers of mass of full cylinder and of half cylinder https://en.wikipedia.org/wiki/List_of_centroids
    distCoM = 4/(3*pi)*radCyl

    # moments of inertia of half cylinder w.r.t. center of mass of half cylinder
    IwCylHalf = IwCylHalf - massHalfCyl*distCoM^2
    ItCylHalf = ItCylHalf - massHalfCyl*distCoM^2

    # moments of inertia of half cylinder w.r.t. center of mass of beam
    IwCylHalf = IwCylHalf + massHalfCyl*(distCoM + shape.length/2)^2
    ItCylHalf = ItCylHalf + massHalfCyl*(distCoM + shape.length/2)^2

    Il = IlBox + 2*IlCylHalf
    Iw = IwBox + 2*IwCylHalf
    It = ItBox + 2*ItCylHalf
    if shape.axis == 1
        return InertiaMatrix(Diagonal([Il, Iw, It]))
    elseif shape.axis == 2
        return InertiaMatrix(Diagonal([It, Il, Iw]))
    else
        return InertiaMatrix(Diagonal([Iw, It, Il]))
    end
end

function inertiaMatrix(shape::FileMesh, mass::Number)
    if !isempty(shape.facesIndizes)
        return shape.inertia.*mass
    else
        println("FileMesh: ", shape.filename, ". The surface areas must be triangular, and each triangle should be specified in right-handed/counter-clockwise order. Otherwise it is not possible to compute an inertia tensor.")
        return nothing
    end
end


# fdf(x, α, d) = f(x + α * d), ∇f(x + α * d)
"""
    DescentMethods.strong_backtracking!(fdf, x₀, d, [y₀, g₀]; [α, αmax, β, σ])

Find `α` such that `f(x₀ + αd) ≤ f(x₀) + β α d⋅∇f(x)` and `|d⋅∇f(x₀ + αd)| ≤ σ|d⋅∇f(x₀)|`
    (strong Wolfe conditions) using the backtracking line search with cubic interpolation
    and Hager-Zhang approximation when steps are very small. Optionally, `y₀ = f(x₀)` and
    `g₀ = ∇f(x₀)` may be provided to avoid recalculation.

# Arguments
- `fdf`: a function `fdf(x, α, d)` returning a tuple `f(x + α * d), ∇f(x + α * d)` where `f`
    is the minimized function
- `x₀::AbstractVector`: the initial point
- `d::AbstractVector`: the search direction. Must be a descent direction, i.e. `d⋅∇f(x₀) < 0`
- `y₀`: (optional) the value of `f(x₀)` if it is known beforehand
- `g₀`: (optional) the value of `∇f(x₀)` if it is known beforehand

# Keywords
- `α=1`: the initial value of `α`
- `αmax=Inf`: the maximum allowed value of α
- `β=1e-4`: the coefficient in Wolfe conditions
- `σ=0.5`: the coefficient in Wolfe conditions
"""
function strong_backtracking!(
    fdf::F, x0, d, y0::T, grad0;
    α = one(y0),
    αmax = convert(T, Inf),
    β = convert(T, 1e-4),
    σ = convert(T, 0.5)
) where {F,T}
    α = min(α, αmax / 2)
    α_prev = zero(α)
    y_prev = y0
    ylo = y_prev
    yhi = convert(T, NaN)
    αlo = αhi = convert(T, NaN)
    g0 = dot(grad0, d)
    @assert g0 < 0 "derivative is non-negative: $g0"

    glo = g_prev = g0
    ghi = convert(typeof(g0), NaN)
    # if g0 > 0
    #     g0 = -g0
    #     rmul!(d, -1)
    # end
    mag = max(abs(y0), -α * g0)
    ϵ = sqrt(eps(one(y0))) * mag

    wolfe1 = β * g0
    wolfe2 = -σ * g0

    # bracketing phase
    min_factor = 17/16 # min. factor to expand bracketing interval

    nbracket_max = 200
    for nbracket in 1:nbracket_max
        y, grad = fdf(x0, α, d)
        g = dot(grad, d)
        Δyp = (g + g0) * α / 2 # parabolic approximation
        if abs(Δyp) < ϵ
            Δy = Δyp
        else
            Δy = y - y0
        end
        if Δy > α * wolfe1 || y >= y_prev + ϵ # x >= NaN is always false
            αlo, αhi, ylo, yhi, glo, ghi = α_prev, α, y_prev, y, g_prev, g
            break
        end
        if abs(g) <= wolfe2
            return α
        elseif g >= 0
            αlo, αhi, ylo, yhi, glo, ghi = α_prev, α, y_prev, y, g_prev, g
            break
        end

        nbracket == nbracket_max && error("Failed to find bracketing")
        # cubic interpolation (Nocedal & Wright 2nd ed., p.59)
        Δα = α - α_prev

        # return 0 signaling that bracketing failed
        if iszero(Δα)
            return Δα
        end
        d1 = g_prev + g - 3 * (y - y_prev) / Δα
        det = d1^2 - g * g_prev
        if det < 0
            αnew = min(2, sqrt(αmax / α)) * α
        else
            d2 = sqrt(det)
            αnew = α - Δα * (g + d2 - d1) / (g - g_prev + 2 * d2)
        end
        if α < αnew < αmax && min_factor * α < αnew
            α_prev, α = α, αnew
        else
            α_prev, α = α, α * min(2, sqrt(αmax / α))
        end
        y_prev, g_prev = y, g
    end

    # zoom phase
    # αlo is lower bound, αhi is higher
    small_α = sqrt(eps(one(αlo)))
    for nzoom in 1:200
        # cubic interpolation (Nocedal & Wright 2nd ed., p.59)
        Δα = αhi - αlo
        if Δα < eps(αhi) * 32
            @warn "Step too small; returning"
            return αlo + Δα / 2
        end
        d1 = ghi + glo - 3 * (yhi - ylo) / Δα
        det = d1^2 - ghi * glo
        if det < 0
            α = αlo + Δα / 2
        else
            d2 = sqrt(det)
            α = αlo - Δα * (ghi + d2 - d1) / (ghi - glo + 2 * d2)
            # ensure that new α is in (αlo; αhi) and not too close to bounds
            if !(α - αlo > small_α * αlo && αhi - α > small_α * αhi)
                α = αlo + Δα / 2
            end
        end
        y, grad = fdf(x0, α, d)
        g = grad ⋅ d
        Δyp = (g + g0) * α / 2 # parabolic approximation
        if abs(Δyp) < ϵ
            Δy = Δyp
        else
            Δy = y - y0
        end
        if Δy > α * wolfe1 || y >= ylo + ϵ
            αhi, yhi, ghi = α, y, g
        else
            abs(g) <= wolfe2 && return α
            if g > 0
                αhi, yhi, ghi = α, y, g
            else
                αlo, ylo, glo = α, y, g
            end
        end
    end
end

function strong_backtracking!(
    fdf::F, x0::AbstractVector, d::AbstractVector{T};
    α = one(T),
    αmax = convert(T, Inf),
    β = convert(T, 1e-4),
    σ = convert(T, 0.5)
) where {F,T}
    y0, grad0 = fdf(x0, zero(T), d)
    return strong_backtracking!(fdf, x0, d, y0, grad0, α = α, αmax = αmax, β = β, σ = σ)
end


# Test other equivalences of prox operations which are not covered by calculus rules

using Random
using SparseArrays

Random.seed!(0)

################################################################################
### testing consistency of simplex/L1 ball
################################################################################
# Inspired by Condat, "Fast projection onto the simplex and the l1 ball", Mathematical Programming, 158:575–585, 2016.
# See Prop. 2.1 there and following remarks.

n = 20
N = 10

# projecting onto the L1 ball
for i = 1:N
  x = randn(n)
  r = 5*rand()
  f = IndSimplex(r)
  g = IndBallL1(r)

  y1, fy1 = prox(f, abs.(x))
  y1 = sign.(x).*y1
  y2, gy2 = prox(g, x)

  @test y1 ≈ y2
end

# projecting onto the simplex
for i = 1:N
  x = randn(n)
  r = 5*rand()
  f = IndSimplex(r)
  g = IndBallL1(r)

  y1, fy1 = prox(f, x)
  y2, gy2 = prox(g, x .- minimum(x) .+ r./n)

  @test y1 ≈ y2
end

################################################################################
### testing consistency of hinge loss/box indicator
################################################################################

n = 20
N = 10

# test using Moreau identity: prox(f, x, gamma) = x - gamma*prox(f*, x/gamma, 1/gamma)
# and a couple of other calculus rules in the case b = ±ones(n)
#
# f(x) = max(0, x) is conjugate to h = IndBox(0,1)
# g(x) = HingeLoss(b, mu)(x) = mu*f(1-b.*x)
# prox(g, x, gamma) = (prox(f, 1-b.*x, mu*gamma) - 1)./(-b)
#   = x + mu*gamma*prox(h, (1-b.*x)/(mu*gamma), 1/(mu*gamma))./b

b = sign.(randn(n))
mu = 0.1+rand()
g = HingeLoss(b, mu)
h = IndBox(0, 1)

for i = 1:N
  x = randn(n)
  gamma = 0.1 + rand()
  y1, ~ = prox(g, x, gamma)
  z, ~ = prox(h, (1 .- b.*x)./(mu*gamma), 1/(mu*gamma))
  y2 = mu*gamma*(z./b) + x
  @test y1 ≈ y2
end

################################################################################
### testing regularize
################################################################################

lambda = rand()
rho = rand()
g = Regularize(NormL1(lambda),rho)

x = randn(10)
y,f = prox(g,x)
y2,f2 = prox(ElasticNet(lambda,rho),x)

@test f ≈ f2
@test y ≈ y2

################################################################################
### testing IndAffine
################################################################################

A = sprand(50,100, 0.1)
b = randn(50)

g1 = IndAffine(A, b)
g2 = IndAffine(Matrix(A), b)

x = randn(100)
y1, f1 = prox(g1, x)
y2, f2 = prox(g2, x)

@test f1 ≈ f2
@test y1 ≈ y2


## PulseDesigner
module PulseDesigner
using ..Koma
using ..Koma: γ, Scanner, get_bvalue, get_max_grad

###############
## RF Pulses ##
###############
RF_hard(B1, T, sys::Scanner; G=[0,0,0]) = begin
	ζ = sum(G) / sys.Smax
	EX = Sequence([	Grad(G[1],T,ζ);	 #Gx
					Grad(G[2],T,ζ);  #Gy
					Grad(G[3],T,ζ);;], #Gz
					 [RF(B1,T,0,ζ);;]	 #RF
					)
	EX
end
##################
## Gradient SEQ ##
##################
# # Without moment nulling
# DIF_base(G0,Δ,δ;verbose=false) = begin
# 	DIF = Sequence([Grad(G0,δ) Delay(Δ-δ) -Grad(G0,δ); #Gx
# 	         		Grad(0, δ) Delay(Δ-δ)  Grad(0 ,δ)]) #Gy
# 	if verbose
# 	b_value = get_bvalue(DIF)[1]
# 	GDmax = get_max_grad(DIF)
# 	println("Sequence with diffusion gradient: "*string(round(GDmax*1e3,digits=2))*" mT/m") # s/mm2
# 	println("Sequence with b-value: "*string(round(b_value*1e-6,digits=2))*" s/mm2") # s/mm2
# 	end
# 	DIF
# end
# # With moment nulling: M0,M1,M2
# DIF_null(G0,Δ,δ,verbose=false) = begin
# 	if abs(Δ-4δ/3)<1e-6 #Necessary condition
# 	#Stoecker
# 	DIF = Sequence([Grad(G0,δ/3) -Grad(G0,2δ/3) Delay(Δ-δ) Grad(G0,2δ/3) -Grad(G0,δ/3); #Gx
#              		Delay(δ/3) Delay(2δ/3) Delay(Δ-δ) Delay(2δ/3) Delay(δ/3)]) #Gy
# 	else
# 	#Castillo
# 	DIF = Sequence([ Grad(G0,δ/4) -Grad(G0,δ/2) Grad(G0,δ/4) Delay(Δ-δ)
# 					-Grad(G0,δ/4)  Grad(G0,δ/2)-Grad(G0,δ/4); #Gx
#              		 Delay(δ/4) Delay(δ/2) Delay(δ/4) Delay(Δ-δ)
# 					 Delay(δ/4) Delay(δ/2) Delay(δ/4)])
# 	end

# 	if verbose
# 	b_value = get_bvalue(DIF)[1]
# 	GDmax = get_max_grad(DIF)
# 	println("Sequence with diffusion gradient: "*string(round(GDmax*1e3,digits=2))*" mT/m") # s/mm2
# 	println("Sequence with b-value: "*string(round(b_value*1e-6,digits=2))*" s/mm2") # s/mm2
# 	end
# 	DIF
# end
# EPI
EPI(FOV::Float64, N::Int, sys::Scanner) = begin
    #TODO: consider when N is even
	Δt = sys.ADC_Δt
	Gmax = sys.Gmax
	Nx = Ny = N #Square acquisition
	Δx = FOV/(Nx-1)
	Ta = Δt*(Nx-1) #4-8 us
    Δτ = Ta/(Ny-1)
	Ga = 1/(γ*Δt*FOV)
	ζ = Ga / sys.Smax
	Ga ≥ sys.Gmax ? error("Ga=$(Ga*1e3) mT/m exceeds Gmax=$(Gmax*1e3) mT/m, increase Δt to at least Δt_min="
	*string(round(1/(γ*Gmax*FOV),digits=2))*" us.") : 0
	ϵ1 = Δτ/(Δτ+ζ)
	#EPI base
	EPI = Sequence(vcat(
	    [mod(i,2)==0 ? Grad(Ga*(-1)^(i/2),Ta,ζ) : Grad(0.,Δτ,ζ) for i=0:2*Ny-2],  #Gx
	 	[mod(i,2)==1 ? ϵ1*Grad(Ga,Δτ,ζ) :         Grad(0.,Ta,ζ) for i=0:2*Ny-2])) #Gy
	EPI.ADC = [mod(i,2)==1 ? ADC(0,Δτ,ζ) : ADC(N,Ta,ζ) for i=0:2*Ny-2]
	# Relevant parameters
	Δfx_pix = 1/Ta
	Δt_phase = (Ta+2ζ)*Ny + (Δτ+2ζ)*Ny
	Δfx_pix_phase = 1/Δt_phase
	println("## EPI parameters ##")
	println("Δx = $(round(Δx*1e3,digits=2)) mm")
	println("Pixel Δf in freq. direction $(round(Δfx_pix,digits=2)) Hz")
	println("Pixel Δf in phase direction $(round(Δfx_pix_phase,digits=2)) Hz")
	#Pre-wind and wind gradients
	ϵ2 = Ta/(Ta+ζ)
    PHASE =   Sequence(1/2*[Grad(-Ga, Ta, ζ)      ; ϵ2*Grad(-Ga, Ta, ζ);;]) #This needs to be calculated differently
	DEPHASE = Sequence(1/2*[Grad((-1)^N*Ga, Ta, ζ); ϵ2*Grad(-Ga, Ta, ζ);;]) #for even N
	seq = PHASE+EPI+DEPHASE
	#Saving parameters
	seq.DEF = Dict("Nx"=>Nx,"Ny"=>Ny,"Nz"=>1,"Name"=>"epi")
	seq
end

# Radial
radial_base(FOV::Float64, Nr::Int, sys::Scanner) = begin
	Δt = sys.ADC_Δt
	Gmax = sys.Gmax
	Δx = FOV/(Nr-1)
	Ta = Δt*(Nr-1)
	Ga = 1/(γ*Δt*FOV)
	ζ = Ga / sys.Smax
	Ga ≥ sys.Gmax ? error("Ga=$(Ga*1e3) mT/m exceeds Gmax=$(Gmax*1e3) mT/m, increase Δt to at least Δt_min="
	*string(round(1/(γ*Gmax*FOV),digits=2))*" us.") : 0
	#Radial base
	rad = Sequence([Grad(Ga,Ta,ζ)]) #Gx
	rad.ADC = [ADC(Nr,Ta,ζ)]
	# Acq/Recon parameters
	Nspokes = ceil(Int64, π/2 * Nr ) #Nyquist in the radial direction
	Δθ = π / Nspokes
	println("## Radial parameters ##")
	println("FOVr = $(round(FOV*1e2,digits=2)) cm; Δr = $(round(Δx*1e3,digits=2)) mm")
	println("Nspokes = $Nspokes, to satisfy the Nyquist criterion")
    PHASE = Sequence([Grad(-Ga/2, Ta, ζ)])
	seq = PHASE+rad+PHASE
	#Saving parameters
	seq.DEF = Dict("Nx"=>Nr,"Ny"=>Nspokes,"Nz"=>1,"Δθ"=>Δθ,"Name"=>"radial","FOV"=>[FOV, FOV, 0])

	seq
end

export EPI, radial_base
end


using DiffEqFlux, Flux, GalacticOptim, OrdinaryDiffEq, Test

n = 2  # number of ODEs
tspan = (0.0, 1.0)

d = 5  # number of data pairs
x = rand(n, 5)
y = rand(n, 5)

cb = function (p,l)
  @show l
  false
end

using Random
Random.seed!(100)

NN = Chain(Dense(n, 10n, tanh),
           Dense(10n, n))

@info "ROCK4"
nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ) = Flux.mse(y, nODE(x, θ))
l1 = loss_function(nODE.p)

res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), GalacticOptim.AutoZygote(), maxiters = 200, cb=cb)
@test 5loss_function(res.minimizer) < l1
res = DiffEqFlux.sciml_train(loss_function, nODE.p, Optim.KrylovTrustRegion(), GalacticOptim.AutoZygote(), maxiters = 200, cb=cb)
@test 5loss_function(res.minimizer) < l1

NN = FastChain(FastDense(n, 10n, tanh),
               FastDense(10n, n))

@info "ROCK2"
nODE = NeuralODE(NN, tspan, ROCK2(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ) = Flux.mse(y, nODE(x, θ))
l1 = loss_function(nODE.p)
optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_function(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, nODE.p,)

res = GalacticOptim.solve(optprob, NewtonTrustRegion(), maxiters = 200, cb=cb)
@test 5loss_function(res.minimizer) < l1
res = GalacticOptim.solve(optprob, Optim.KrylovTrustRegion(), maxiters = 200, cb=cb, allow_f_increases = true)
@test 5loss_function(res.minimizer) < l1


# distance to center of the simulation cell
function gen_dr_center( r, LatVecs )
    #
    cx = 0.5*sum(LatVecs[:,1])
    cy = 0.5*sum(LatVecs[:,2])
    cz = 0.5*sum(LatVecs[:,3])
    center = [cx,cy,cz]
    #
    Npoints = size(r)[2]
    dr = Array{Float64}(undef,Npoints)
    for ip=1:Npoints
        dr[ip] = norm( r[:,ip] - center )
    end
    return dr
end


using Test

#=
    91.2    Basic Unit Tests
=#

foo(x) = length(x)^2

@test foo("hey") == 9
@test foo("there") > 20

@test_throws BoundsError [1,2,3][4]

@test_throws DimensionMismatch [1,2,3] + [1,2]

# @test_throws BoundsError [1,2,3][3]
@test π ≈ 3.14 atol=0.01

result = @test ≈(π, 3.14, atol=0.01)

@show result.test_type
@show result
@show result.value
@show result.orig_expr
@show typeof(result.orig_expr)
@show result.data

# This throws exception MethodError
# @test foo(:cat) == 1
@test_throws MethodError foo(:cat)

#=
    91.3    Working With Test Sets
=#

result = @testset "My Description of this test set" begin
    @test foo("hello") == 25
    @test foo("bar") == 9
end

@show result
@show result.description
@show result.results
@show result.n_passed
@show result.anynonpass

# r = @test_throws TestSetException begin
#         @testset "My Description of this test set" begin
#             @test foo("hello") == 25
#             @test foo("bar") == 9
#             @test foo("bars") == 9
#         end
# end;
# @show r

try
    @testset "My new testset description" begin
        @test foo("hello") == 25
        @test foo("bar") == 9
        @test foo("bars") == 16
    end
catch e
    println(e);
    #rethrow(e)
    nothing
finally
    println("Finally!")
    nothing
end

# Nested testsets

result = @testset "Foo Tests" begin
    @testset "Animals" begin
        @test foo("cat") == 9
        @test foo("dog") == foo("cat")
    end
    @testset "Arrays $i" for i in 1:3
        @test foo(zeros(i)) == i^2
        @test foo(fill(1.0, i)) == i^2
    end
end;

@show result;

@show result.description

for (i,r) in enumerate(result.results)
    println(i, ": ", r)
end

@show result.anynonpass

@testset "Foo Tests" begin
    @testset "Animals" begin
        @testset "Felines" begin
            @test foo("cat") == 9
        end
        @testset "Canines" begin
            @test foo("dog") == 9
        end
    end
    @testset "Arrays" begin
        @test foo(zeros(2)) == 4
        @test foo(fill(1.0, 4)) == 16
    end
end


# This is part of Calendars.jl. See the copyright note there.
# ======================== DayNumbers =======================

# The fix day number of CE 1582-10-15.
const GregorysBreak = DPart(577736)
const JDN_DN  = DPart(1721425)
const DN_MJDN = DPart(678576)

# Days considered here start at midnight 00:00.

# Convert a fix day number to a Julian day number.
# If mod = true return the modified Julian day number.
# (which has epoch CE 1858-11-17, 00:00 UT).
function FixNumToJulianNum(rd::DPart, mod=false)
    if mod
        return rd - DN_MJDN
    else
        return rd + JDN_DN
    end
end

# Convert a Julian day number to a fix day number.
# If mod = true assume a modified Julian day number.
function JulianNumToFixNum(jdn::DPart, mod=false)
    if mod
        return jdn + DN_MJDN
    else
        return jdn - JDN_DN
    end
end

# Convert a fix day number to an European day number.
FixNumToEuroNum(fn::DPart) = fn + 2

# Convert an European day number to a fix day number.
EuroNumToFixNum(en::DPart) = en - 2

"""
Convert a Julian day number to an European day number.
```julia
julia> JulianNumToEuroNum(2440422)
```
The European day number of the first crewed moon landing (Apollo 11) is EN#719000.
"""
JulianNumToEuroNum(jn::DPart) = jn - JDN_DN + 2

"""
Convert an European day number to an Julian day number.
```julia
julia> EuroNumToJulianNum(719000)
```
The Julian day number of the first crewed moon landing (Apollo 11) is JN#2440422.
"""
EuroNumToJulianNum(en::DPart) = en + JDN_DN - 2