println("------------------------------------")
println("|     Fields and vector spaces     |")
println("------------------------------------")
ti = time()
@timedtestset "Fields" begin
    @test isa(ℝ, Field)
    @test isa(ℂ, Field)
    @test eval(Meta.parse(sprint(show, ℝ))) == ℝ
    @test eval(Meta.parse(sprint(show, ℂ))) == ℂ
    @test ℝ ⊆ ℝ
    @test ℝ ⊆ ℂ
    @test ℂ ⊆ ℂ
    @test !(ℂ ⊆ ℝ)

    for T in (Int8, Int16, Int32, Int64, BigInt)
        @test one(T) ∈ ℝ
        @test one(Rational{T}) ∈ ℝ
        @test !(one(Complex{T}) ∈ ℝ)
        @test !(one(Complex{Rational{T}}) ∈ ℝ)
        @test one(T) ∈ ℂ
        @test one(Rational{T}) ∈ ℂ
        @test one(Complex{T}) ∈ ℂ
        @test one(Complex{Rational{T}} ∈ ℂ)

        @test T ⊆ ℝ
        @test Rational{T} ⊆ ℝ
        @test !(Complex{T} ⊆ ℝ)
        @test !(Complex{Rational{T}} ⊆ ℝ)
        @test T ⊆ ℂ
        @test Rational{T} ⊆ ℂ
        @test Complex{T} ⊆ ℂ
        @test Complex{Rational{T}} ⊆ ℂ
    end
    for T in (Float32, Float64, BigFloat)
        @test one(T) ∈ ℝ
        @test !(one(Complex{T}) ∈ ℝ)
        @test one(T) ∈ ℂ
        @test one(Complex{T} ∈ ℂ)

        @test T ⊆ ℝ
        @test !(Complex{T} ⊆ ℝ)
        @test T ⊆ ℂ
        @test Complex{T} ⊆ ℂ
    end
end

@timedtestset "ElementarySpace: CartesianSpace" begin
    d = 2
    V = ℝ^d
    @test eval(Meta.parse(sprint(show, V))) == V
    @test eval(Meta.parse(sprint(show, typeof(V)))) == typeof(V)
    @test isa(V, VectorSpace)
    @test isa(V, ElementarySpace)
    @test isa(V, InnerProductSpace)
    @test isa(V, EuclideanSpace)
    @test isa(V, CartesianSpace)
    @test !isdual(V)
    @test !isdual(V')
    @test V == CartesianSpace(Trivial() => d) == CartesianSpace(Dict(Trivial() => d))
    @test @constinferred(hash(V)) == hash(deepcopy(V))
    @test V == @constinferred(dual(V)) == @constinferred(conj(V)) == @constinferred(adjoint(V))
    @test field(V) == ℝ
    @test @constinferred(sectortype(V)) == Trivial
    @test ((@constinferred sectors(V))...,) == (Trivial(),)
    @test length(sectors(V)) == 1
    @test @constinferred(TensorLabXD.hassector(V, Trivial()))
    @test @constinferred(dim(V)) == d == @constinferred(dim(V, Trivial()))
    @test dim(@constinferred(typeof(V)())) == 0
    @test (sectors(typeof(V)())...,) == ()
    @test @constinferred(TensorLabXD.axes(V)) == Base.OneTo(d)
    @test ℝ^d == ℝ[](d) == CartesianSpace(d) == typeof(V)(d)
    W = @constinferred ℝ^1
    @test @constinferred(oneunit(V)) == W == oneunit(typeof(V))
    @test @constinferred(⊕(V,V)) == ℝ^(2d)
    @test @constinferred(⊕(V,oneunit(V))) == ℝ^(d+1)
    @test @constinferred(⊕(V,V,V,V)) == ℝ^(4d)
    @test @constinferred(fuse(V,V)) == ℝ^(d^2)
    @test @constinferred(fuse(V,V',V,V')) == ℝ^(d^4)
    @test @constinferred(flip(V)) == V'
    @test flip(V) ≅ V
    @test flip(V) ≾ V
    @test flip(V) ≿ V
    @test V ≺ ⊕(V,V)
    @test !(V ≻ ⊕(V,V))
    @test @constinferred(infimum(V, ℝ^3)) == V
    @test @constinferred(supremum(V', ℝ^3)) == ℝ^3
end

@timedtestset "ElementarySpace: ComplexSpace" begin
    d = 2
    V = ℂ^d
    @test eval(Meta.parse(sprint(show, V))) == V
    @test eval(Meta.parse(sprint(show, V'))) == V'
    @test eval(Meta.parse(sprint(show, typeof(V)))) == typeof(V)
    @test isa(V, VectorSpace)
    @test isa(V, ElementarySpace)
    @test isa(V, InnerProductSpace)
    @test isa(V, EuclideanSpace)
    @test isa(V, ComplexSpace)
    @test !isdual(V)
    @test isdual(V')
    @test V == ComplexSpace(Trivial() => d) == ComplexSpace(Dict(Trivial() => d))
    @test @constinferred(hash(V)) == hash(deepcopy(V)) != hash(V')
    @test @constinferred(dual(V)) == @constinferred(conj(V)) == @constinferred(adjoint(V)) != V
    @test @constinferred(field(V)) == ℂ
    @test @constinferred(sectortype(V)) == Trivial
    @test @constinferred(sectortype(V)) == Trivial
    @test ((@constinferred sectors(V))...,) == (Trivial(),)
    @test length(sectors(V)) == 1
    @test @constinferred(TensorLabXD.hassector(V, Trivial()))
    @test @constinferred(dim(V)) == d == @constinferred(dim(V, Trivial()))
    @test dim(@constinferred(typeof(V)())) == 0
    @test (sectors(typeof(V)())...,) == ()
    @test @constinferred(TensorLabXD.axes(V)) == Base.OneTo(d)
    @test ℂ^d == Vect[Trivial](d) == Vect[](Trivial()=>d) == ℂ[](d) == typeof(V)(d)
    W = @constinferred ℂ^1
    @test @constinferred(oneunit(V)) == W == oneunit(typeof(V))
    @test @constinferred(⊕(V, V)) == ℂ^(2d)
    @test_throws SpaceMismatch (⊕(V, V'))
    @test_throws MethodError (⊕(ℝ^d, ℂ^d))
    @test_throws MethodError (⊗(ℝ^d, ℂ^d))
    @test @constinferred(⊕(V,V)) == ℂ^(2d)
    @test @constinferred(⊕(V,oneunit(V))) == ℂ^(d+1)
    @test @constinferred(⊕(V,V,V,V)) == ℂ^(4d)
    @test @constinferred(fuse(V,V)) == ℂ^(d^2)
    @test @constinferred(fuse(V,V',V,V')) == ℂ^(d^4)
    @test @constinferred(flip(V)) == V'
    @test flip(V) ≅ V
    @test flip(V) ≾ V
    @test flip(V) ≿ V
    @test V ≺ ⊕(V,V)
    @test !(V ≻ ⊕(V,V))
    @test @constinferred(infimum(V, ℂ^3)) == V
    @test @constinferred(supremum(V', (ℂ^3)')) == dual(ℂ^3) == conj(ℂ^3)
end

@timedtestset "ElementarySpace: GeneralSpace" begin
    d = 2
    V = GeneralSpace{ℂ}(d)
    @test eval(Meta.parse(sprint(show, V))) == V
    @test eval(Meta.parse(sprint(show, dual(V)))) == dual(V)
    @test eval(Meta.parse(sprint(show, conj(V)))) == conj(V)
    @test eval(Meta.parse(sprint(show, typeof(V)))) == typeof(V)
    @test !isdual(V)
    @test isdual(V')
    @test !isdual(conj(V))
    @test isdual(conj(V'))
    @test !TensorLabXD.isconj(V)
    @test !TensorLabXD.isconj(V')
    @test TensorLabXD.isconj(conj(V))
    @test TensorLabXD.isconj(conj(V'))
    @test isa(V, VectorSpace)
    @test isa(V, ElementarySpace)
    @test !isa(V, InnerProductSpace)
    @test !isa(V, EuclideanSpace)
    @test @constinferred(hash(V)) == hash(deepcopy(V)) != hash(V')
    @test @constinferred(dual(V)) != @constinferred(conj(V)) != V
    @test @constinferred(field(V)) == ℂ
    @test @constinferred(sectortype(V)) == Trivial
    @test @constinferred(TensorLabXD.hassector(V, Trivial()))
    @test @constinferred(dim(V)) == d == @constinferred(dim(V, Trivial()))
    @test @constinferred(TensorLabXD.axes(V)) == Base.OneTo(d)
end

@timedtestset "ElementarySpace: $(TensorLabXD.type_repr(Vect[I]))" for I in sectorlist
    if Base.IteratorSize(values(I)) === Base.IsInfinite()
        set = unique(vcat(one(I), [randsector(I) for k = 1:10]))
        gen = (c=>2 for c in set)
    else
        gen = (values(I)[k]=>(k+1) for k in 1:length(values(I)))
    end
    V = GradedSpace(gen)
    @test eval(Meta.parse(TensorLabXD.type_repr(typeof(V)))) == typeof(V)
    @test eval(Meta.parse(sprint(show, V))) == V
    @test eval(Meta.parse(sprint(show, V'))) == V'
    @test V' == GradedSpace(gen; dual = true)
    @test V == @constinferred GradedSpace(gen...)
    @test V' == @constinferred GradedSpace(gen...; dual = true)
    @test V == @constinferred GradedSpace(tuple(gen...))
    @test V' == @constinferred GradedSpace(tuple(gen...); dual = true)
    @test V == @constinferred GradedSpace(Dict(gen))
    @test V' == @constinferred GradedSpace(Dict(gen); dual = true)
    @test V == @inferred Vect[I](gen)
    @test V' == @constinferred Vect[I](gen; dual = true)
    @test V == @constinferred Vect[I](gen...)
    @test V' == @constinferred Vect[I](gen...; dual = true)
    @test V == @constinferred Vect[I](Dict(gen))
    @test V' == @constinferred Vect[I](Dict(gen); dual = true)
    @test V == @constinferred typeof(V)(c=>dim(V,c) for c in sectors(V))
    if I isa ZNIrrep
        @test V == @constinferred typeof(V)(V.dims)
        @test V' == @constinferred typeof(V)(V.dims; dual = true)
    end
    @test @constinferred(hash(V)) == hash(deepcopy(V)) != hash(V')
    @test V == GradedSpace(reverse(collect(gen))...)
    @test eval(Meta.parse(sprint(show, V))) == V
    @test eval(Meta.parse(sprint(show, typeof(V)))) == typeof(V)
    # space with no sectors
    @test dim(@constinferred(typeof(V)())) == 0
    # space with a single sector
    W = @constinferred GradedSpace(one(I)=>1)
    @test W == GradedSpace(one(I)=>1, randsector(I) => 0)
    @test @constinferred(oneunit(V)) == W == oneunit(typeof(V))
    # randsector never returns trivial sector, so this cannot error
    @test_throws ArgumentError GradedSpace(one(I)=>1, randsector(I) => 0, one(I)=>3)
    @test eval(Meta.parse(sprint(show, W))) == W
    @test isa(V, VectorSpace)
    @test isa(V, ElementarySpace)
    @test isa(V, InnerProductSpace)
    @test isa(V, EuclideanSpace)
    @test isa(V, GradedSpace)
    @test isa(V, GradedSpace{I})
    @test @constinferred(dual(V)) == @constinferred(conj(V)) == @constinferred(adjoint(V)) != V
    @test @constinferred(field(V)) == ℂ
    @test @constinferred(sectortype(V)) == I
    slist = @constinferred sectors(V)
    @test @constinferred(TensorLabXD.hassector(V, first(slist)))
    @test @constinferred(dim(V)) == sum(dim(s)*dim(V, s) for s in slist)
    @constinferred dim(V, first(slist))
    if hasfusiontensor(I)
        @test @constinferred(TensorLabXD.axes(V)) == Base.OneTo(dim(V))
    end
    @test @constinferred(⊕(V,V)) == Vect[I](c=>2dim(V,c) for c in sectors(V))
    @test @constinferred(⊕(V,V,V,V)) == Vect[I](c=>4dim(V,c) for c in sectors(V))
    @test @constinferred(⊕(V,oneunit(V))) ==
            Vect[I](c=>isone(c)+dim(V,c) for c in sectors(V))
    @test @constinferred(fuse(V,oneunit(V))) == V
    d = Dict{I,Int}()
    for a in sectors(V), b in sectors(V)
        for c in a ⊗ b
            d[c] = get(d, c, 0) + dim(V, a)*dim(V, b)*Nsymbol(a,b,c)
        end
    end
    @test @constinferred(fuse(V,V)) == GradedSpace(d)
    @test @constinferred(flip(V)) ==
            Vect[I](conj(c)=>dim(V,c) for c in sectors(V))'
    @test flip(V) ≅ V
    @test flip(V) ≾ V
    @test flip(V) ≿ V
    @test @constinferred(⊕(V, V)) == @constinferred supremum(V, ⊕(V, V))
    @test V == @constinferred infimum(V, ⊕(V, V))
    @test V ≺ ⊕(V,V)
    @test !(V ≻ ⊕(V,V))
    @test infimum(V, GradedSpace(one(I)=>3)) == GradedSpace(one(I)=>2)
    @test_throws SpaceMismatch (⊕(V, V'))
end

@timedtestset "ProductSpace{ℂ}" begin
    V1, V2, V3, V4 = ℂ^1, ℂ^2, ℂ^3, ℂ^4
    P = @constinferred ProductSpace(V1, V2, V3, V4)
    @test eval(Meta.parse(sprint(show, P))) == P
    @test eval(Meta.parse(sprint(show, typeof(P)))) == typeof(P)
    @test isa(P, VectorSpace)
    @test isa(P, CompositeSpace)
    @test spacetype(P) == ComplexSpace
    @test sectortype(P) == Trivial
    @test @constinferred(hash(P)) == hash(deepcopy(P)) != hash(P')
    @test P == deepcopy(P)
    @test P == typeof(P)(P...)
    @constinferred (x->tuple(x...))(P)
    @test @constinferred(dual(P)) == P'
    @test @constinferred(field(P)) == ℂ
    @test @constinferred(*(V1, V2, V3, V4)) == P
    @test @constinferred(⊗(V1, V2, V3, V4)) == P
    @test @constinferred(⊗(V1, V2⊗V3⊗V4)) == P
    @test @constinferred(⊗(V1⊗V2, V3⊗V4)) == P
    @test @constinferred(⊗(V1, V2, V3⊗V4)) == P
    @test @constinferred(⊗(V1, V2⊗V3, V4)) == P
    @test @constinferred(insertunit(P, 3)) == V1 * V2 * oneunit(V1) * V3 * V4
    @test fuse(V1, V2', V3) ≅ V1 ⊗ V2' ⊗ V3
    @test fuse(V1, V2', V3) ≾ V1 ⊗ V2' ⊗ V3
    @test fuse(V1, V2', V3) ≿ V1 ⊗ V2' ⊗ V3
    @test fuse(flip(V1), V2, flip(V3)) ≅ V1 ⊗ V2 ⊗ V3
    @test @constinferred(⊗(P)) == P
    @test @constinferred(⊗(V1)) == ProductSpace(V1)
    @test eval(Meta.parse(sprint(show, ⊗(V1)))) == ⊗(V1)
    @test @constinferred(one(V1)) == @constinferred(one(typeof(V1))) ==
                @constinferred(one(P)) == @constinferred(one(typeof(P))) ==
                ProductSpace{ComplexSpace}(())
    @test eval(Meta.parse(sprint(show, one(P)))) == one(P)
    @test @constinferred(⊗(one(P), P)) == P
    @test @constinferred(⊗(P, one(P))) == P
    @test @constinferred(⊗(one(P), one(P))) == one(P)
    @test @constinferred(adjoint(P)) == dual(P) == V4' ⊗ V3' ⊗ V2' ⊗ V1'
    @test @constinferred(dims(P)) == map(dim, (V1, V2, V3, V4))
    @test @constinferred(dim(P)) == prod(dim, (V1, V2, V3, V4))
    @test @constinferred(dim(P, 2)) == dim(V2)
    @test @constinferred(sectors(P)) ==
            (mapreduce(sectors, (a, b)->tuple(a..., b...), (V1, V2, V3, V4)),)
    cube(x) = x^3
    @test @constinferred(cube(V1)) == V1 ⊗ V1 ⊗ V1
    N = 3
    @test V1^N == V1 ⊗ V1 ⊗ V1
    @test P^2 == P ⊗ P
    @test @constinferred(dims(P, first(sectors(P)))) == dims(P)
    @test ((@constinferred blocksectors(P))...,) == (Trivial(),)
    @test (blocksectors(P ⊗ ℂ^0)...,) == ()
    @test @constinferred(blockdim(P, first(blocksectors(P)))) == dim(P)
    @test Base.IteratorEltype(P) == Base.IteratorEltype(typeof(P)) ==
                                    Base.IteratorEltype(P.spaces)
    @test Base.IteratorSize(P) == Base.IteratorSize(typeof(P)) ==
                                    Base.IteratorSize(P.spaces)
    @test Base.eltype(P) == Base.eltype(typeof(P)) == typeof(V1)
    @test eltype(collect(P)) == typeof(V1)
    @test collect(P) == [V1, V2, V3, V4]
end

@timedtestset "ProductSpace{SU₂Space}" begin
    V1, V2, V3 = SU₂Space(0=>3, 1//2=>1), SU₂Space(0=>2, 1=>1), SU₂Space(1//2=>1, 1=>1)'
    P = @constinferred ProductSpace(V1, V2, V3)
    @test eval(Meta.parse(sprint(show, P))) == P
    @test eval(Meta.parse(sprint(show, typeof(P)))) == typeof(P)
    @test isa(P, VectorSpace)
    @test isa(P, CompositeSpace)
    @test spacetype(P) == SU₂Space
    @test sectortype(P) == Irrep[SU₂] == SU2Irrep
    @test @constinferred(hash(P)) == hash(deepcopy(P)) != hash(P')
    @test @constinferred(dual(P)) == P'
    @test @constinferred(field(P)) == ℂ
    @test @constinferred(*(V1, V2, V3)) == P
    @test @constinferred(⊗(V1, V2, V3)) == P
    @test @constinferred(adjoint(P)) == dual(P) == V3' ⊗ V2' ⊗ V1'
    @test @constinferred(insertunit(P, 3; conj = true)) == V1 * V2 * oneunit(V1)' * V3
    @test fuse(V1, V2', V3) ≅ V1 ⊗ V2' ⊗ V3
    @test fuse(V1, V2', V3) ≾ V1 ⊗ V2' ⊗ V3 ≾ fuse(V1 ⊗ V2' ⊗ V3)
    @test fuse(V1, V2') ⊗ V3 ≾ V1 ⊗ V2' ⊗ V3
    @test fuse(V1, V2', V3) ≿ V1 ⊗ V2' ⊗ V3 ≿ fuse(V1 ⊗ V2' ⊗ V3)
    @test V1 ⊗ fuse(V2', V3) ≿ V1 ⊗ V2' ⊗ V3
    @test fuse(flip(V1) ⊗ V2) ⊗ flip(V3) ≅ V1 ⊗ V2 ⊗ V3
    @test @constinferred(⊗(V1)) == ProductSpace(V1)
    @test @constinferred(one(V1)) == @constinferred(one(typeof(V1))) ==
                @constinferred(one(P)) == @constinferred(one(typeof(P))) ==
                ProductSpace{ComplexSpace}(())
    @test @constinferred(dims(P)) == map(dim, (V1, V2, V3))
    @test @constinferred(dim(P)) == prod(dim, (V1, V2, V3))
    for s in @constinferred(sectors(P))
        @test hassector(P, s)
        @test @constinferred(dims(P, s)) == dim.((V1, V2, V3), s)
    end
    @test sum(dim(c)*blockdim(P, c) for c in @constinferred(blocksectors(P))) == dim(P)
end

@timedtestset "Deligne tensor product of spaces" begin
    V1 = SU₂Space(0=>3, 1//2=>1)
    V2 = SU₂Space(0=>2, 1=>1)'
    V3 = ℤ₃Space(0=>3, 1=>2, 2=>1)
    V4 = ℂ^3

    for W1 in (V1, V2, V3, V4)
        for W2 in (V1, V2, V3, V4)
            for W3 in (V1, V2, V3, V4)
                for W4 in (V1, V2, V3, V4)
                    Ws = @constinferred(W1 ⊠ W2 ⊠ W3 ⊠ W4)
                    @test Ws == @constinferred((W1 ⊠ W2) ⊠ (W3 ⊠ W4)) ==
                                @constinferred(((W1 ⊠ W2) ⊠ W3) ⊠ W4) ==
                                @constinferred((W1 ⊠ (W2 ⊠ W3)) ⊠ W4) ==
                                @constinferred(W1 ⊠ ((W2 ⊠ W3)) ⊠ W4) ==
                                @constinferred(W1 ⊠ (W2 ⊠ (W3 ⊠ W4)))
                    I1, I2, I3, I4 = map(sectortype, (W1, W2, W3, W4))
                    I = sectortype(Ws)
                    @test I == @constinferred((I1 ⊠ I2) ⊠ (I3 ⊠ I4)) ==
                                @constinferred(((I1 ⊠ I2) ⊠ I3) ⊠ I4) ==
                                @constinferred((I1 ⊠ (I2 ⊠ I3)) ⊠ I4) ==
                                @constinferred(I1 ⊠ ((I2 ⊠ I3)) ⊠ I4) ==
                                @constinferred(I1 ⊠ (I2 ⊠ (I3 ⊠ I4)))
                    @test dim(Ws) == dim(W1) * dim(W2) * dim(W3) * dim(W4)
                end
            end
        end
    end
    @test sectortype(@constinferred((V1 ⊗ V2) ⊠ V3)) == @constinferred(Irrep[SU₂ × ℤ₃])
    @test dim((V1 ⊗ V2) ⊠ V3) == dim(V1) * dim(V2) * dim(V3)
    @test sectortype((V1 ⊗ V2) ⊠ V3 ⊠ V4) == Irrep[SU₂ × ℤ₃]
    @test dim((V1 ⊗ V2) ⊠ V3 ⊠ V4) == dim(V1) * dim(V2) * dim(V3) * dim(V4)
    @test fuse(V2 ⊠ V4) == fuse(V4 ⊠ V2) == SU₂Space(0=>6, 1=>3)
    @test fuse(V3 ⊠ V4) == fuse(V4 ⊠ V3) == ℤ₃Space(0=>9, 1=>6, 2=>3)
end

@timedtestset "HomSpace" begin
    V1, V2, V3, V4, V5 = SU₂Space(0=>3, 1//2=>1), SU₂Space(0=>2, 1=>1),
                            SU₂Space(1//2=>1, 1=>1)', SU₂Space(0=>2, 1//2=>2),
                            SU₂Space(0=>1, 1//2=>1, 3//2=>1)'
    W = TensorLabXD.HomSpace(V1 ⊗ V2, V3 ⊗ V4 ⊗ V5)
    @test W == (V3 ⊗ V4 ⊗ V5 → V1 ⊗ V2)
    @test W == (V1 ⊗ V2 ← V3 ⊗ V4 ⊗ V5)
    @test W' == (V1 ⊗ V2 → V3 ⊗ V4 ⊗ V5)
    @test eval(Meta.parse(sprint(show, W))) == W
    @test eval(Meta.parse(sprint(show, typeof(W)))) == typeof(W)
    @test spacetype(W) == SU₂Space
    @test sectortype(W) == Irrep[SU₂]
    @test W[1] == V1
    @test W[2] == V2
    @test W[3] == V3'
    @test W[4] == V4'
    @test W[5] == V5'
    @test @constinferred(hash(W)) == hash(deepcopy(W)) != hash(W')
    @test W == deepcopy(W)
end

tf = time()
printstyled("Finished vector space tests in ",
            string(round(tf-ti; sigdigits=3)),
            " seconds."; bold = true, color = Base.info_color())
println()


using Printf
using SparseArrays

import PyPlot
const plt = PyPlot

include("../LF_common/m_LF1d.jl")
include("../LF_common/m_LF3d.jl")

function test_main( NN::Array{Int64} )
  AA = [0.0, 0.0, 0.0]
  BB = [16.0, 16.0, 16.0]
  LF = init_LF3d_c( NN, AA, BB )

  Npoints = prod(NN)
  println("Step 1: ")

  NNZ = ccall( (:get_nnz_laplacian3d_, "./Laplace3d_sparse.so"), Int64,
         (Ptr{Int64}, Ref{Int64}, Ptr{Int64},
          Ptr{Float64}, Ptr{Float64}, Ptr{Float64} ),
         NN, Npoints, LF.lin2xyz, LF.LFx.D2jl, LF.LFy.D2jl, LF.LFz.D2jl )

  #
  row    = zeros(Int64,NNZ)
  column = zeros(Int64,NNZ)
  values = zeros(Float64,NNZ)
  #
  println("Step 2: ")

  ccall( (:init_laplacian3d_, "./Laplace3d_sparse.so"), Nothing,
         ( Ref{Int64}, Ptr{Int64}, Ref{Int64}, Ptr{Int64},
           Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Int64}, Ptr{Int64}, Ptr{Float64} ),
         NNZ, NN, Npoints, LF.lin2xyz, LF.LFx.D2jl, LF.LFy.D2jl, LF.LFz.D2jl,
         row, column, values )

  #
  #for ii = 1:NNZ
  #  @printf("%5d %5d %18.10f\n", row[ii], column[ii], values[ii])
  #end
  println("Step 3: ")

  Laplacian = sparse( row, column, values )

  plt.clf()
  plt.spy(Matrix(Laplacian))
  filename = "Laplacian_"*string(NN[1])*"_"*string(NN[2])*"_"*string(NN[3])*".pdf"
  plt.savefig(filename)

end

test_main([3,5,5])
#test_main([20,20,20])
#test_main([30,30,30])
#test_main([40,40,40])
#test_main([50,50,50])
#for ii = 10:10:60
#  @time test_main( ii*ones(Int64,3) )
#end


"""
    ConjugateModel <: ProbabilisticModel

`ConjugateModel` is a `ProbabilisticModel` with a conjugate prior of the
corresponding likelihood function.
"""
abstract type ConjugateModel <: ProbabilisticModel end

function Base.show(io::IO, model::ConjugateModel)
    modelinfo = replace(string(model.dist), "Distributions."=>"", count=1)
    message = "$(typeof(model)) : $(modelinfo)"
    print(io, message)
end

"""
    ConjugateBernoulli <: ConjugateModel

Bernoulli likelihood with Beta distribution as the conjugate prior.

```julia
ConjugateBernoulli(α, β)              # construct a ConjugateBernoulli

update!(model, stats)             # update model with statistics from data
samplepost(model, numsamples)    # sampling from the posterior distribution
samplestats(model, numsamples)   # sampling statistics from the data generating distribution
```

"""
mutable struct ConjugateBernoulli <: ConjugateModel
    dist::Beta
    function ConjugateBernoulli(α, β)
        return new(Beta(α, β))
    end
end

defaultparams(::ConjugateBernoulli) = [:θ]

function update!(model::ConjugateBernoulli, stats::BetaStatistics)
    numsuccesses = stats.s
    numtrials = stats.n
    α = model.dist.α + numsuccesses
    β = model.dist.β + numtrials - numsuccesses
    model.dist = Beta(α, β)
    return nothing
end

"""
    ConjugateExponential <: ConjugateModel

Exponential likelihood with Gamma distribution as the conjugate prior.

```julia
ConjugateExponential(α, β)            # construct a ConjugateExponential

update!(model, stats)             # update model with statistics from data
samplepost(model, numsamples)    # sampling from the posterior distribution
samplestats(model, numsamples)   # sampling statistics from the data generating distribution
```
"""
mutable struct ConjugateExponential <: ConjugateModel
    dist::Gamma
    function ConjugateExponential(α, θ)
        return new(Gamma(α, θ))
    end
end

defaultparams(::ConjugateExponential) = [:θ]

mutable struct ConjugatePoisson <: ConjugateModel
    dist::Gamma
    function ConjugatePoisson(α, θ)
        return new(Gamma(α, θ))
    end
end

defaultparams(::ConjugatePoisson) = [:λ]

function update!(model::T, stats::S) where
        {T<:Union{ConjugateExponential, ConjugatePoisson},
         S<:Union{GammaStatistics}}
    n = stats.n
    x̄ = stats.x̄
    α = model.dist.α + n
    θ = model.dist.θ / (1 + model.dist.θ * n * x̄)
    model.dist = Gamma(α, θ)
    return nothing
end

"""
    ConjugateNormal <: ConjugateModel

Normal likelihood and Normal Inverse Gamma distribution as the
conjugate prior.

## Parameters

- μ: mean of normal distribution
- v:  scale variance of Normal
- α:  shape of Gamma distribution
- θ:  scale of Gamma distribution

```julia
ConjugateNormal(μ, v, α, θ)           # construct a ConjugateNormal

update!(model, stats)             # update model with statistics from data
samplepost(model, numsamples)    # sampling from the posterior distribution
samplestats(model, numsamples)   # sampling statistics from the data generating distribution
```

## References

- The update rule for Normal distribution is based on this
  [lecture notes](https://people.eecs.berkeley.edu/~jordan/courses/260-spring10/other-readings/chapter9.pdf).
"""
mutable struct ConjugateNormal{S} <: ConjugateModel
    dist::S
    function ConjugateNormal{Normal}(μ::Real, σ::Real)
        T = promote_type(typeof(μ), typeof(σ))
        return new(Normal{T}(μ, σ))
    end
    function ConjugateNormal{NormalInverseGamma}(μ::Real, v::Real, α::Real, θ::Real)
        return new(NormalInverseGamma(μ, v, α, θ))
    end
end

ConjugateNormal(μ::T, v::T, α::T, θ::T) where T <: Real = ConjugateNormal{NormalInverseGamma}(μ, v, α, θ)
ConjugateNormal(μ::T, σ::T) where T <: Real = ConjugateNormal{Normal}(μ, σ)

defaultparams(::ConjugateNormal{NormalInverseGamma}) = [:μ]

"""
    ConjugateLogNormal(μ, v, α, θ)

A model with Normal likelihood and Normal Inverse distribution with log transformed data.
Notice `LogNormal` in `Distributions.jl` takes mean and standard deviation of \$\\log(x)\$
instead of \$x\$ as the input parameters.

```julia
ConjugateLogNormal(μ, v, α, θ) # construct a ConjugateLogNormal

lognormalparams(μ_logx, σ²_logx) # convert normal parameters to log-normal parameters
update!(model, stats)              # update model with statistics from data
samplepost(model, numsamples)     # sampling from the posterior distribution
samplestats(model, numsamples)    # sampling statistics from the data generating distributio
```

"""
mutable struct ConjugateLogNormal{S} <: ConjugateModel
    dist::S
    function ConjugateLogNormal{NormalInverseGamma}(μ::Real, v::Real, α::Real, θ::Real)
        return new(NormalInverseGamma(μ, v, α, θ))
    end
end

ConjugateLogNormal(μ::T, v::T, α::T, θ::T) where T <: Real = ConjugateLogNormal{NormalInverseGamma}(μ, v, α, θ)

function lognormalparams(μ_logx, σ²_logx)
    μ_x  = @. exp(μ_logx + σ²_logx / 2)
    σ²_x = @. (exp(σ²_logx) - 1) * exp(2 * μ_logx + σ²_logx)
    return (μ_x, σ²_x)
end

defaultparams(::ConjugateLogNormal{NormalInverseGamma}) = [:μ_x]

function update!(model::Union{ConjugateNormal{NormalInverseGamma},ConjugateLogNormal{NormalInverseGamma}},
        stats::Union{NormalStatistics,LogNormalStatistics})
    n, x̄, sdx = getall(stats)
    ΣΔx² = sdx^2 * (n - 1)

    μ0 = model.dist.μ
    v0 = model.dist.v
    α0 = model.dist.α
    θ0 = model.dist.θ
    inv_v0 = 1.0 / v0
    inv_v = inv_v0 + n

    μ = (inv_v0 * μ0 + n * x̄) / inv_v
    v = 1 / inv_v
    α = α0 + n / 2
    θ = θ0 + 0.5 * (ΣΔx² + (n * inv_v0) * (x̄ - μ0)^2 * v)
    model.dist = NormalInverseGamma(μ, v, α, θ)
    return nothing
end

abstract type ChainOperator end

struct MultiplyOperator <: ChainOperator end
(m::MultiplyOperator)(a, b) = .*(a, b)

"""
    ChainedModel <: ProbabilisticModel

`ChainedModel` is a combination of `ConjugateModel`s chained by the specified operator.
It can be used to model a multiple step process.
"""
struct ChainedModel <: ProbabilisticModel
    models::Vector{ConjugateModel}
    operators::Vector{ChainOperator}
    function ChainedModel(models::Vector{ConjugateModel}, operators)
        length(models) - 1 == length(operators) || error("need to specify (number of model - 1) chaining operators")
        return new(models, operators)
    end
end

function ChainedModel(models::Vector{ConjugateModel})
    operators = [MultiplyOperator() for _ in 1:(length(models) -1)]
    return ChainedModel(models, operators)
end

function update!(chainedmodel::ChainedModel, listofstats::Vector{T}) where T <: ModelStatistics
    @assert(length(chainedmodel.models) == length(listofstats),
            "Number of model should be equal to number of statistics.")
    for (model, stats) = zip(chainedmodel.models, listofstats)
        update!(model, stats)
    end
    return nothing
end

function defaultparams(chainedmodel::ChainedModel)
    params = Symbol[]
    for model in chainedmodel.models
        push!(params, defaultparams(model)[1])
    end
    return params
end

"""
    samplepost(model, numsamples)

Sample from the posterior distribution of the model.
"""
function samplepost(model::ConjugateBernoulli, numsamples::Int)
    return BernoulliPosteriorSample(rand(model.dist, numsamples))
end

function samplepost(model::ConjugateExponential, numsamples::Int)
    # Gamma distribution generates the rate (λ) of Exponential distribution
    # convert it to scale θ to make it consistent with Distributions.jl
    λs = rand(model.dist, numsamples)
    θs = 1 ./ λs
    return ExponentialPosteriorSample(θs)
end

function samplepost(model::ConjugatePoisson, numsamples::Int)
    inv_rates = rand(model.dist, numsamples)
    λs = 1 ./ inv_rates
    return PoissonPosteriorSample(λs)
end

function samplepost(model::ConjugateNormal{NormalInverseGamma}, numsamples::Int)
    μ, σ² = rand(model.dist, numsamples)
    return NormalPosteriorSample(μ, σ²)
end

function samplepost(model::ConjugateLogNormal{NormalInverseGamma}, numsamples::Int)
    # sample for normal means and variances
    μ_logx, σ²_logx = rand(model.dist, numsamples)
    μ_x, σ²_x = lognormalparams(μ_logx, σ²_logx)
    return LogNormalPosteriorSample(μ_logx, σ²_logx, μ_x, σ²_x)
end

function samplepost(model::T, parameter::Symbol, numsamples::Int) where T <: ConjugateModel
    samples = samplepost(model, numsamples)
    return getfield(samples, parameter)
end

function samplepost(model::T, parameters::Vector{Symbol}, numsamples::Int)  where T <: ConjugateModel
    parameter = toparameter(parameters)
    return samplepost(model, parameter, numsamples)
end

function samplepost(models::Vector{T}, parameters::Vector{Symbol}, numsamples::Int) where T <: ProbabilisticModel
    samples = Array{Float64,2}(undef, numsamples, length(models))
    for (modelindex, model) in enumerate(models)
        sample = samplepost(model, parameters, numsamples)
        samples[:, modelindex] = sample
    end
    return samples
end

function samplepost(chainedmodel::ChainedModel, parameters::Vector{Symbol}, numsamples::Int)
    length(parameters) == length(chainedmodel.models) ||
        throw(ArgumentError("Number of parameters must be equal to number of chained models"))
    models = chainedmodel.models
    operators = chainedmodel.operators
    post_samples = [samplepost(model, parameter, numsamples) for (model, parameter) in zip(models, parameters)]
    chainedsample = post_samples[1]
    for i = 2:length(post_samples)
        chainedsample = operators[i - 1](chainedsample, post_samples[i])
    end
    return chainedsample
end

"""
    samplestats(model, dist, numsamples)

Sample from the distribution of the data generating process, and
calculate the corresponding statistics for the model.
"""
function samplestats(::ConjugateBernoulli, dist::Bernoulli, numsamples::Integer)
    return BetaStatistics(rand(dist, numsamples))
end

function samplestats(::ConjugateExponential, dist::Exponential, numsamples::Integer)
    return GammaStatistics(rand(dist, numsamples))
end

function samplestats(::ConjugatePoisson, dist::Poisson, numsamples::Integer)
    return GammaStatistics(rand(dist, numsamples))
end

function samplestats(::ConjugateNormal{NormalInverseGamma}, dist::Normal, numsamples::Integer)
    return NormalStatistics(rand(dist, numsamples))
end

function samplestats(::ConjugateLogNormal{NormalInverseGamma}, dist::LogNormal, numsamples::Integer)
    return LogNormalStatistics(rand(dist, numsamples))
end

function samplestats(chainedmodel::ChainedModel, dists::Vector{T}, listofnumsamples::Vector{Int}) where T
    listofstats = Vector{ModelStatistics}()
    for (model, dist, numsamples) in zip(chainedmodel.models, dists, listofnumsamples)
        push!(listofstats, samplestats(model, dist, numsamples))
    end
    return listofstats
end


"""
    get_p_min(ν::Real, B::Real)

Calculate the smallest synchrotron bright moment following Donnert+16, Eq. 22
"""
function smallest_synch_bright_p(ν::Real, B::Real)
    sqrt( ν / (C_crit * B) )
end


include("SLPDE.jl")
include("etdrk4.jl")
include("evaluate_lambda.jl")
include("sphcesaro.jl")

using PyPlot

srand(0)

n = 128
U0 = zeros(4n, 8n-1);
U0[1:n, 1:2n-1] = sphrandn(Float64, n, n)/n
n = size(U0, 1)

θ = [0;(0.5:n-0.5)/n;1]
φ = [(0:2n-2)*2/(2n-1);2]
x = [cospi(φ)*sinpi(θ) for θ in θ, φ in φ]
y = [sinpi(φ)*sinpi(θ) for θ in θ, φ in φ]
z = [cospi(θ) for θ in θ, φ in φ]

V = sph2fourier(U0)
F = zero(V)
A_mul_B!(F, FastTransforms.plan_synthesis(V), V)
F = [mean(F[1,:])*ones(1, size(F, 2)); F; mean(F[end,:])*ones(1, size(F, 2))]
F = [F F[:,1]]

function ETDRK4_free_energy(SLPDE::SemiLinearPDE{1, S}, T::Number, n::Int) where S
    U0 = SLPDE.U0[1]
    U = deepcopy(U0) # solution spherical harmonic coefficients
    V = zero(U) # solution Fourier coefficients
    F = zero(U) # solution function values
    NU = zero(U) # nonlinearity spherical harmonic coefficients
    NV = zero(U) # nonlinearity Fourier coefficients
    NF = zero(U) # nonlinearity function values

    # ETDRK4 stages
    A = zero(U)
    B = zero(U)
    C = zero(U)
    NA = zero(U)
    NB = zero(U)
    NC = zero(U)

    P = plan_sph2fourier(U; sketch = :none)
    Ps = FastTransforms.plan_synthesis(U)
    Pa = FastTransforms.plan_analysis(U)

    L = create_linear_operator(SLPDE.L, size(U, 1), size(U, 2)÷2 + 1)
    h = S(T)/n

    NE = NonlinearOperator(u -> 0.5*(u^2-1))

    ENERGY = zeros(S, n+1)
    ENERGY[1] = -0.5*sum(U.*L.*U) + sum(evaluate_nonlinear_operator!(NE, P, Ps, Pa, U, V, F, NF, NV, NU).^2)

    # Create all the matrices of diagonal scalings used in ETDRK4. These operate
    # component-wise on the spherical harmonic coefficients.

    Z = L*h
    ez = exp.(Z)
    ez2 = exp.(Z./2)
    h2ez2m1 = (h/2.0).*φ1.(Z./2)
    hezα = h.*expα.(Z)
    h2ezβ = 2h.*expβ.(Z)
    hezγ = h.*expγ.(Z)
    ez2u = zero(Z)

    for k=1:n
        ez2u .= ez2.*U
        evaluate_nonlinear_operator!(SLPDE.N, P, Ps, Pa, U, V, F, NF, NV, NU)
        A .= ez2u .+ h2ez2m1.*NU
        evaluate_nonlinear_operator!(SLPDE.N, P, Ps, Pa, A, V, F, NF, NV, NA)
        B .= ez2u .+ h2ez2m1.*NA
        evaluate_nonlinear_operator!(SLPDE.N, P, Ps, Pa, B, V, F, NF, NV, NB)
        C .= ez2.*A .+ h2ez2m1.*(2.0.*NB .- NU)
        evaluate_nonlinear_operator!(SLPDE.N, P, Ps, Pa, C, V, F, NF, NV, NC)
        U .= ez.*U .+ hezα.*NU .+ h2ezβ.*(NA .+ NB) .+ hezγ.*NC
        ENERGY[k+1] = -0.5*sum(U.*L.*U) + sum(evaluate_nonlinear_operator!(NE, P, Ps, Pa, U, V, F, NF, NV, NU).^2)
    end

    ENERGY
end

T = 64.0
n = round(Int, 10*T)
t = linspace(0, T, n+1)

SLPDE = SemiLinearPDE(LaplaceBeltrami(0.01), NonlinearOperator(u->u-u^3), U0)
local_energy = ETDRK4_free_energy(SLPDE, T, n)

SLPDE = SemiLinearPDE(NonlocalLaplaceBeltrami(0.01, -0.5, 1.0), NonlinearOperator(u->u-u^3), U0)
nonlocal_energy = ETDRK4_free_energy(SLPDE, T, n)

clf()
plot(t, local_energy, "-k", t, nonlocal_energy, "--k")
xlabel("\$t\$"); ylabel("\${\\cal E}(u)\$"); grid()
legend(["Local","Nonlocal"])
savefig("GL_free_energy.pdf")


@inline unorm(U) = sum(map((x->x.^2),U))

@inline function unpackfields(eqn::Euler{d},U) where {d}
    return first(U),ntuple(i->U[i+1],d),last(U)
end

"""
    prim_to_cons(eqn::Euler{d},rho,U,p) where {d}

convert primitive variables (ρ,U...,p) to conservative vars (ρ,ρU,E).
n-dimensional version where U = tuple(u1,...,u_d)
"""
function prim_to_cons(eqn::Euler{d},Q) where {d}
    rho,U,p = unpackfields(eqn,Q)
    rhoU = map(x->rho.*x,U)
    Unorm = unorm(U)
    E = @. p /(eqn.γ-1) + .5*rho*Unorm
    return SVector{d+2}(rho,rhoU...,E)
end

"""
    cons_to_prim_beta(eqn::Euler{d},U) where {d}

converts conservative variables to `primitive' variables which make evaluating EC
fluxes cheaper.
"""
@inline function cons_to_prim_beta(eqn::Euler{d},U) where {d}
    rho,rhoU,E = unpackfields(eqn,U)
    return SVector{d+2}(rho, map(x->x./rho,rhoU)..., betafun(eqn,U))
end

"""
    pfun(eqn::Euler{d},U) where {d}

Computes pressure (assuming ideal gas law) given array/tuple of conservative variables `U`
"""
@inline function pfun(eqn::Euler{d},U) where {d}
    rho,rhoU,E = unpackfields(eqn,U)
    rhoUnorm = unorm(rhoU)
    return @. (eqn.γ - 1)*(E - .5*rhoUnorm / rho)
end

"""
    betafun(eqn::Euler{d},U) where {d}

Converts "inverse temperature" β = ρ/(2p) given array/tuple of conservative variables `U`
"""
@inline function betafun(eqn::Euler{d},U) where {d}
    rho = first(U)
    p = pfun(eqn,U)
    return @. rho/(2*p)
end

function sfun(eqn::Euler{d},U) where {d}
    rho = first(U)
    p = pfun(eqn,U)
    return @. log(p/(rho^eqn.γ))
end

"""
    Sfun(eqn::Euler{d},U) where {d}

Converts mathematical entropy -ρs(U) given array/tuple of conservative variables `U`
"""
function Sfun(eqn::Euler{d},U) where {d}
    rho = first(U)
    return -rho.*sfun(eqn,U)
end

"""
    v_ufun(eqn::Euler{d}, U) where {d}

Returns entropy variables given tuple/array of conservative variables `U`.
"""
function v_ufun(eqn::Euler{d}, U) where {d}
    @unpack γ = eqn

    s = sfun(eqn,U)
    p = pfun(eqn,U)

    rho,rhoU,E = unpackfields(eqn,U)
    v1 = @. (γ+1-s) - (γ-1)*E/p
    vU = (x->@. x*(γ-1)/p).(rhoU)
    vE = (x->@. x*(γ-1)/p)(-rho)

    return SVector(v1,vU...,vE)
end

function s_vfun(eqn::Euler{d},V) where {d}
    @unpack γ = eqn
    v1,vU,vE = unpackfields(eqn,V)
    vUnorm = unorm(vU)
    return @. γ - v1 + vUnorm/(2*vE)
end

function rhoe_vfun(eqn::Euler{d},V) where {d}
    @unpack γ = eqn
    s = s_vfun(eqn,V)
    v1,vU,vE = unpackfields(eqn,V)
    return @. ((γ-1)/((-vE)^γ))^(1/(γ-1)) * exp(-s/(γ-1))
end

"""
    u_vfun(eqn::Euler{d},V) where {d}

Returns conservative variables given tuple/array of entropy variables `V`.
"""
function u_vfun(eqn::Euler{d},V) where {d}
    v1,vU,vE  = unpackfields(eqn,V)
    rhoeV     = rhoe_vfun(eqn,V)
    vUnorm    = unorm(vU)
    rho       = (@. rhoeV*(-vE))
    rhoU      = (x->rhoeV.*x).(vU)
    E         = (@. rhoeV*(1-vUnorm/(2*vE)))
    return SVector{d+2}(rho,rhoU...,E)
end

"""
    wavespeed(eqn::Euler{1},U)

Returns 1D wavespeed (used for flux penalization terms) given conservative variables U
"""
function wavespeed(eqn::Euler{1},U)
    rho,rhou,_ = unpackfields(eqn,U)
    p = pfun(eqn,U)
    cvel = @. sqrt(eqn.γ*p/rho)
    return @. abs(rhou/rho) + cvel
end

# function dUdV_explicit(v1,vU1,vU2,vE)
#     rho,rhou,rhov,E = u_vfun(v1,vU1,vU2,vE)
#     u,v = (x->x./rho).((rhou,rhov))
#     p = pfun(rho,rhou,rhov,E)
#     a2 = γ*p/rho
#     H = a2/(γ-1) + (u^2+v^2)/2
#
#     dUdV = @SMatrix [rho  rhou        rhov        E;
#                      rhou rhou*u + p  rhou*v      rhou*H;
#                      rhov rhov*u      rhov*v + p  rhov*H;
#                      E    rhou*H      rhov*H      rho*H^2-a2*p/(γ-1)]
#
#     return dUdV*(1/(γ-1))
# end
#
# function dVdU_explicit(rho,rhou,rhov,E)
#     rhoe = rhoe_ufun(rho,rhou,rhov,E)
#     V = v_ufun(rho,rhou,rhov,E)
#     k = .5*(V[2]^2+V[3]^2)/V[4]
#
#     dVdU = @SMatrix [γ+k^2      k*V[2]          k*V[3]         V[4]*(k+1);
#                     k*V[2]      V[2]^2-V[4]     V[2]*V[3]      V[2]*V[4];
#                     k*V[3]      V[2]*V[3]       V[3]^2-V[4]    V[3]*V[4]
#                     V[4]*(k+1)  V[2]*V[4]       V[3]*V[4]      V[4]^2]
#     return -dVdU/(rhoe*V[4])
# end


#!/usr/local/bin/julia
# last tested Julia version: 1.6.1
# Hodgkin-Huxley model on a 2D lattice
# FvW 03/2018

using NPZ
using PyCall
using PyPlot
using Statistics
using VideoIO
@pyimport matplotlib.animation as anim

function hh2d(N, T, t0, dt, s, D, C, gNa, gK, gL, VNa, VK, VL, I0, stim, blocks)
    # initialize Hodgkin-Huxley system
	C1 = 1.0/C
    v = -10*ones(Float64,N,N)
    m = zeros(Float64,N,N)
    n = zeros(Float64,N,N)
    h = zeros(Float64,N,N)
    dvdt = zeros(Float64,N,N)
    dmdt = zeros(Float64,N,N)
    dndt = zeros(Float64,N,N)
    dhdt = zeros(Float64,N,N)
    s_sqrt_dt = s*sqrt(dt)
    X = zeros(Float64,T,N,N)
    # stimulation protocol
    I = zeros(Float64,t0+T,N,N)
    for st in stim
        t_on, t_off = st[1]
        x0, x1 = st[2]
        y0, y1 = st[3]
        I[t0+t_on:t0+t_off, x0:x1, y0:y1] .= I0
    end
    # iterate
    for t in range(2, stop=t0+T, step=1)
        (t%100 == 0) && print("    t = ", t, "/", t0+T, "\r")
        # HH equations
        # dV/dt (membrane potential)
        dvdt = C1.*(-gNa.*(m.^3).*h.*(v.-VNa)
                    .- gK.*(n.^4).*(v.-VK)
                    .- gL*(v.-VL) .+ I[t,:,:]) .+ D.*L(v)
        # dm/dt (Na+ channel activation)
        alpha_m = 0.1.*(-v.+25)./(exp.((-v.+25)./10).-1)
        beta_m = 4 * exp.(-v./18)
        dmdt = (alpha_m .* (1 .- m) .- beta_m .* m)
        # dh/dt (Na+ channel inactivation)
        alpha_h = 0.07 * exp.(-v./20)
        beta_h = 1.0 ./ (exp.((-v .+ 30) ./ 10) .+ 1)
        dhdt = (alpha_h .* (1.0 .- h) .- beta_h.*h)
        # dn/dt (K+ channel gating)
        alpha_n = 0.01 .* (-v .+ 10)./(exp.((-v .+ 10)./10) .- 1)
        beta_n = 0.125 .* exp.(-v./80)
        dndt = (alpha_n .* (1.0 .- n) .- beta_n.*n)
        # stochastic integration
        v += (dvdt*dt + s_sqrt_dt*randn(N,N))
        m += (dmdt*dt)
        n += (dndt*dt)
        h += (dhdt*dt)
        # dead block(s):
        for bl in blocks
            v[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
            m[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
            n[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
            h[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
        end
        # store
        (t > t0) && (X[t-t0,:,:] = v)
    end
    println("\n")
    return X
end

function animate_pyplot(fname, data)
    """
    Animate 3D array as .mp4 using PyPlot, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    nt, nx, ny = size(data)
    vmin = minimum(data)
    vmax = maximum(data)
    # setup animation image
    fig = figure(figsize=(6,6))
    axis("off")
    t = imshow(data[1,:,:], origin="lower", cmap=ColorMap("gray"),
			   vmin=vmin, vmax=vmax)
    tight_layout()
    # frame generator
    println("[+] animate")
    function animate(i)
        (i%100 == 0) && print("    t = ", i, "/", nt, "\r")
        t.set_data(data[i+1,:,:])
    end
    # create animation
    ani = anim.FuncAnimation(fig, animate, frames=nt, interval=10)
    println("\n")
    # save animation
    ani[:save](fname, bitrate=-1,
               extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])
    show()
end

function animate_video(fname, data)
    """
    Animate 3D array as .mp4 using VideoIO, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    # BW
    y = UInt8.(round.(255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    # BW inverted
    #y = UInt8.(round.(255 .- 255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    encoder_options = (color_range=2, crf=0, preset="medium")
    framerate=30
    T = size(data,1)
    open_video_out(fname, y[1,end:-1:1,:], framerate=framerate,
                   encoder_options=encoder_options) do writer
        for i in range(2,stop=T,step=1)
            write(writer, y[i,end:-1:1,:])
        end
    end
end

function L(x)
    # Laplace operator
    # periodic boundary conditions
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
    # non-periodic boundary conditions
    Lx[1,:] .= 0.0
    Lx[end,:] .= 0.0
    Lx[:,1] .= 0.0
    Lx[:,end] .= 0.0
    return Lx
end

function main()
    println("Hodgkin-Huxley (HH) lattice model\n")
    N = 128
    T = 5000
    t0 = 2000
    dt = 0.01
    s = 1.0
    D = 1.0
    # Hodgkin-Huxley parameters
    C = 1
    gNa = 120
    gK = 36
    gL = 0.3
    VNa = 115
    VK = -12
    VL = 10.6
    I = 45.0
    println("[+] Lattice size N: ", N)
    println("[+] Time steps T: ", T)
    println("[+] Warm-up steps t0: ", t0)
    println("[+] Integration time step dt: ", dt)
    println("[+] Noise intensity s: ", s)
    println("[+] Diffusion coefficient D: ", D)
    println("[+] HH parameter gNa: ", gNa)
    println("[+] HH parameter gK: ", gK)
    println("[+] HH parameter gL: ", gL)
    println("[+] HH parameter VNa: ", VNa)
    println("[+] HH parameter VK: ", VK)
    println("[+] HH parameter VL: ", VL)
    println("[+] HH parameter C: ", C)
    println("[+] Stimulation current I: ", I)

    # stimulation protocol, array of elements [[t0,t1], [x0,x1], [y0,y1]]
    stim = [ [[50,350], [1,5], [1,5]],
             [[1900,2200], [25,30], [1,15]] ]
    #stim = [] # empty, no current

    # dead blocks, array of elementy [[x0,x1], [y0,y1]]
    blocks = [ [[1,15], [5,10]] ]
    #blocks = [] # empty, no dead areas

    # run simulation
    data = hh2d(N, T, t0, dt, s, D, C, gNa, gK, gL, VNa, VK, VL, I, stim, blocks)
    println("[+] Data dimensions: ", size(data))

    # plot mean voltage
    m = mean(reshape(data, (T,N*N)), dims=2)
    plot(m, "-k"); show()

    # save data
    I_str = rpad(I, 4, '0') # stim. current amplitude as 4-char string
    s_str = rpad(s, 4, '0') # noise as 4-char string
    D_str = rpad(D, 4, '0') # diffusion coefficient as 4-char string
    #fname1 = string("ml2d_I_", I_str, "_s_", s_str, "_D_", D_str, ".npy")
    #npzwrite(fname1, data)
    #println("[+] Data saved as: ", data_filename)

    # video
    fname2 = string("hh2d_I_", I_str, "_s_", s_str, "_D_", D_str, ".mp4")
    #animate_pyplot(fname2, data) # slow
    animate_video(fname2, data) # fast
    println("[+] Animation saved as: ", fname2)
end

main()


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using BenchmarkTools
using CAGD

function generateCubeGrids(n=1, m=1, p=1)
    V,(VV,EV,FV,CV) = Lar.cuboidGrid([n,m,p],true)
    mybox = (V,CV,FV,EV)

    twocubs = Lar.Struct([mybox, Lar.t(.3,.4,.5), Lar.r(pi/5,0,0), Lar.r(0,0,pi/12), mybox])

    V,CV,FV,EV = Lar.struct2lar(twocubs)

    cop_EV = convert(Lar.ChainOp, Lar.coboundary_0(EV::Lar.Cells))
    cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
    model = CAGD.Model(V)
    CAGD.addModelCells!(model, 1, cop_EV)
    CAGD.addModelCells!(model, 2, cop_FE)

    return model
end

m = generateCubeGrids(1, 1, 1)

@btime CAGD.spatial_arrangement(m)


### Kinetic matrix elements

function kinetic(a::PGBF,b::PGBF)
    return a.norm*b.norm*kinetic(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
                                b.expn,b.x,b.y,b.z,b.I,b.J,b.K)
end

function kinetic(aexpn::Float64,ax::Float64,ay::Float64,az::Float64,aI::Int64,aJ::Int64,aK::Int64,
            bexpn::Float64,bx::Float64,by::Float64,bz::Float64,bI::Int64,bJ::Int64,bK::Int64)
    overlap0 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)
    overlapx1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI+2,bJ,bK)
    overlapy1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ+2,bK)
    overlapz1 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK+2)
    overlapx2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI-2,bJ,bK)
    overlapy2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ-2,bK)
    overlapz2 = overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK-2)
    term0 = bexpn*(2*(bI+bJ+bK)+3)*overlap0
    term1 = -2*(bexpn^2)*(overlapx1+overlapy1+overlapz1)
    term2 = -0.5*(bI*(bI-1)*overlapx2+bJ*(bJ-1)*overlapy2+bK*(bK-1)*overlapz2)
    return term0+term1+term2
end

kinetic(a::CGBF,b::CGBF) = contract(kinetic,a,b)

function test_kinetic()
    @testset "test_kinetic" begin
    s = pgbf(1.0)
    c = cgbf(0.0,0.0,0.0)
    push!(c,1,1)
    @test amplitude(c,0,0,0) ≈ 0.71270547
    @test kinetic(1.,0.,0.,0.,0,0,0,1.,0.,0.,0.,0,0,0) ≈ 2.9530518648229536
    @test kinetic(s,s) ≈ 1.5
    @test kinetic(c,c) ≈ 1.5
    end #testset
end


module ForwardOptimalControl

import ..DynamicsModelInterface
import ..JuMPUtils
import Ipopt
import JuMP

using JuMP: @variable, @constraint, @objective
using UnPack: @unpack

export forward_quadratic_objective, linear_dynamics_constraints, solve_lqr, solve_optimal_control

#============================================= LQ case =============================================#

"The performance index for the forward optimal control problem."
function forward_quadratic_objective(x, u; Q, R)
    T = last(size(x))
    sum(x[:, t]' * Q * x[:, t] + u[:, t]' * R * u[:, t] for t in 1:T)
end

function linear_dynamics_constraints(x, u; A, B)
    reduce(hcat, ((x[:, t + 1] - A[t] * x[:, t] - B[t] * u[:, t]) for t in axes(x)[2][1:(end - 1)]))
end

"Solves a forward LQR problem using JuMP."
function solve_lqr(
    A,
    B,
    Q,
    R,
    x0,
    T;
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
)
    n_states, n_controls = size(only(unique(B)))
    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])
    @constraint(opt_model, linear_dynamics_constraints(x, u; A, B) .== 0)
    @constraint(opt_model, x[:, 1] .== x0)
    @objective(opt_model, Min, forward_quadratic_objective(x, u; Q, R))
    @time JuMP.optimize!(opt_model)
    JuMPUtils.isconverged(opt_model), JuMPUtils.get_values(; x, u), opt_model
end

#=========================================== Non-LQ-Case ===========================================#

"Solves a forward optimal control problem with protentially nonlinear dynamics and nonquadratic
costs using JuMP."
function solve_optimal_control(
    control_system,
    cost_model,
    x0,
    T;
    fixed_inputs = (),
    init = (),
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
    verbose = false,
)
    @unpack n_states, n_controls = control_system

    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    # decision variables
    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])

    # initial guess
    JuMPUtils.init_if_hasproperty!(x, init, :x)
    JuMPUtils.init_if_hasproperty!(u, init, :u)

    # fix certain inputs
    for i in fixed_inputs
        JuMP.fix.(u[i, :], init.u[i, :])
    end

    DynamicsModelInterface.add_dynamics_constraints!(control_system, opt_model, x, u)
    if !isnothing(x0)
        @constraint(opt_model, x[:, 1] .== x0)
    end
    cost_model.add_objective!(opt_model, x, u; cost_model.weights)
    time = @elapsed JuMP.optimize!(opt_model)
    verbose && @info time

    JuMPUtils.isconverged(opt_model), JuMPUtils.get_values(; x, u), opt_model
end

end


### This implementation is roughly in the order defined in
### `Uniform Discrete Curvelet Transform' by Truong T. Nguyen and Hervé Chauris
### Equations will be referenced out of this paper.

### Common Lisp comment style will be used here for clarity
module Curvelet
    export curveletTransform, inverseCurveletTransform
    using FFTW
    using DSP

    Curvelet_verbose = false

    ### Eqn. 13 defines β(t) s.t. it smoothly maps the range -1...1 to 0...1
    ### This function is used to specify the transition regions of the windows
    function _beta(t::AbstractFloat)
        ## While the absolute value is not included in any of the equations, it seems
        ## to be required to properly calculate values with t~ -1.0
        if(abs(t) <= 1)
            sqrt(abs(sum([-5/32 21/32 -35/32 35/32 1/2].*(t.^[7 5 3 1 0]))))
        elseif(1 <= t)
            1.0
        else
            0.0
        end
    end

    ### ηa & ηb are used to define the respective radial and angular falloff rates
    ### Within the paper the value of 0.15 was used for both, though recovery
    ### guarentees hold for more values.
    const ETA_A = 0.15
    const ETA_B = 0.15

    #### Windowing functions
    ###
    _w1p(t) = _beta((pi-abs(t))/(pi*ETA_A))
    _w0p(t) = _w1p(2*t*(1+ETA_A))
    _w0(w::Tuple{Float64,Float64}) = _w0p(w[1]).*_w0p(w[2])

    _w1(w::Tuple{Float64, Float64}) = ((1-_w0(w)^2)^0.5)*_w1p(w[1])*_w1p(w[2])

    _v1p(t::Float64,N) = _beta((2/N-1-t)/(2*ETA_B/N))*_beta((t+1)/(2ETA_B/N))

    _vp(t,N,l) = _v1p(t-(2*(l-1))/N,N)

    ## Eqn. 29 defines T(θ) as a warping function that warps the 1D function ṽ(t,l)
    ## to the angular function v(θ,l)
    function T(theta)
        if(-pi/4 <= theta && theta <= pi/4)
            tan(theta)
        elseif(pi/4 < theta && theta < pi/2)
            2-tan(pi/2-theta)
        elseif(-pi/2 < theta && theta < -pi/4)
            -2-tan(pi/2-theta)
        elseif(theta < -pi/2)
            -2
        else#  theta > pi/2
            2
        end
    end

    ### Warped version of ṽ(t,l)
    function _vl(theta,N,l)
        if(l<N)
            _vp(float(T(theta)), N, l)
        else
            _vp(T(pi/2-theta),N,2*N+1-l)
        end
    end

    ### Intermediate in defining the windowing function for l!=0
    function _ulp(w::Tuple{AbstractFloat, AbstractFloat}, N, l)
        _vl(atan2(w[2],w[1]),N,l)*_w1(w)
    end

    function _ul(w::Tuple{AbstractFloat, AbstractFloat}, N::Int, l::Int)
        if(l==0)
            _w0(w)
        else
            result = 0
            ## These extra loops should not be needed, but something is missed
            ## if these are ommited
            for a=[-2pi 0.0 2pi], b=[-2pi 0.0 2pi]
                result += _ulp((w[1]+a,w[2]+b), N, l)
            end
            result
        end
    end

    grid(x,size) = (-pi+2.0*pi/size*(x-1))

    ### Defines the window (without scaling term) exactly as described in the paper.
    ### This approach is quite expensive though, resulting ~100x the work really
    ### needed to get the output as many points outside the ROS are dealt with
    function direct_window_definition(size::Int, N::Int, l::Int)
        fn = (x,y)->_ul((x,y),N,l)

        result = zeros(size, size)

        ## Calculate the window, allowing it to have its full image expressed by
        ## wraping around the edge
        for x=1:size, y=1:size,
            result[y,x] = fn(grid(x,size), grid(y,size))
        end
        return result
    end

    ###Generate scaled window, see fast_window_definition in the optimization
    ### section
    function window(size::Int, N::Int, l::Int)
        if(l==0)
            fast_window_definition(size, N, l)*2
        else
            fast_window_definition(size, N, l)*2*sqrt(2)
        end
    end

    ### Performs the downsample by a factor of two in both directions
    function downsample(X)
        DSP.fftshift(DSP.fft(DSP.ifft(DSP.ifftshift(X))[1:2:end,1:2:end]))
    end

    function downsample2(X)
        Xx = DSP.fftshift(X)
        y  = X[1:Int(end/2),1:Int(end/2)]
        y += X[1:Int(end/2),Int(end/2)+1:end]
        y += X[Int(end/2)+1:end,Int(end/2)+1:end]
        y += X[Int(end/2)+1:end,1:Int(end/2)]
        DSP.fftshift(y)/4
    end

    function super_downsample(X)
        Xx = DSP.fftshift(X)
        y  = X[1:end/2,1:end/2]
        y += X[1:end/2,end/2+1:end]
        y += X[end/2+1:end,end/2+1:end]
        y += X[end/2+1:end,1:end/2]
        yy = y[1:end/2,:] + y[end/2+1:end,:]
        yyy = yy[1:end/2,:] + yy[end/2+1:end,:]
        DSP.fftshift(yyy)
    end

    function super_upsample(X)
        Xx = DSP.fftshift(X)
        XX = [Xx Xx; Xx Xx]
        XXX = [XX; XX]
        XXXX = [XXX; XXX]
    end

    function upsample(X)
        N = size(X,1)
        xx = zeros(2*N,2*N)+0im
        xx[1:2:end,1:2:end] = DSP.ifft(DSP.ifftshift(X))
        DSP.fftshift(DSP.fft(xx))
    end

    function upsample2(X)
        Xx = DSP.fftshift(X)
        [Xx Xx; Xx Xx]
    end

    C1 = 7
    C2 = 2*C1+1

    function _gen_windows(Wsize,N)
        win = Array{Matrix{Float64}}(2*N+1)
        for i=0:N
            if(Curvelet_verbose)
                print('.')
            end
            win[i+1] = window(Wsize,N,i)
        end
        for i=N+1:2N
            win[i+1] = flipdim(win[i-N+1],1)'
            win[i+1][:,[2:end..., 1]] = win[i+1][:, [1:end...]]
        end
        win
    end

    struct CurveletPlan
        dim::Int
        angles::Int
        windows::Vector{Matrix{Float64}}
        subplan::Union{CurveletPlan,Void}
    end

    ## no need to print out a ton of matricies...
    import Base.show
    show(io::IO, cp::Curvelet.CurveletPlan) = print(io, "CurveletPlan")

    function planCurveletTransform(dim::Int, N::Int)
        if(Curvelet_verbose)
            print("<planning>")
        end
        @assert mod(dim,2)==0
        @assert dim > 0
        @assert N > 0

        if(dim <= 16)
            return nothing
        end

        CurveletPlan(dim,2*N+1,_gen_windows(dim,N),planCurveletTransform(Int(dim/2), N))
    end

    struct CurveletCoeff
        HP::Vector{Matrix{Complex{Float64}}}
        LP::Union{CurveletCoeff,Matrix{Complex{Float64}}}
    end

    curveletTransform(x) = curveletTransform(x, planCurveletTransform(size(x,1), C1))

    function curveletTransform(x::Matrix{Float64}, plan::CurveletPlan)
        if(Curvelet_verbose)
            println("FFT")
        end
        S = DSP.fftshift(DSP.fft(x))

        tf = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("Transforming")
        end
        for i=1:plan.angles
            if(Curvelet_verbose)
                print(".")
            end
            tf[i] = S.*plan.windows[i]
        end
        if(Curvelet_verbose)
            println()
        end

        ds = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("Downsampling")
        end
        for i=1:plan.angles
            ds[i] = downsample2(tf[i])
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end

        coeff = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("creating coeff")
        end
        for i=1:plan.angles
            coeff[i] = DSP.ifft(DSP.fftshift(ds[i]))
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end

        CurveletCoeff(coeff[2:end], curveletTransform(ds[1], plan.subplan))
    end

    function curveletTransform(S::Matrix{Complex{Float64}}, plan::Void)
        DSP.ifft(DSP.fftshift(S))
    end

    function curveletTransform(S::Matrix{Complex{Float64}}, plan::CurveletPlan)
        tf = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("Transforming")
        end
        for i=1:plan.angles
            if(Curvelet_verbose)
                print(".")
            end
            tf[i] = S.*plan.windows[i]
        end
        if(Curvelet_verbose)
            println()
        end

        ds = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("Downsampling")
        end
        for i=1:plan.angles
            ds[i] = downsample2(tf[i])
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end

        coeff = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("creating coeff")
        end
        for i=1:plan.angles
            coeff[i] = DSP.ifft(DSP.fftshift(ds[i]))
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end

        CurveletCoeff(coeff[2:end], curveletTransform(ds[1], plan.subplan))
    end

    inverseCurveletTransform(x) = inverseCurveletTransform(x, planCurveletTransform(size(x.HP[1],1)*2, C1))

    function inverseCurveletTransform(cc::CurveletCoeff, plan::CurveletPlan)
        Wsize = plan.dim

        us = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("upsampling")
        end
        for i=1:plan.angles-1
            us[i+1] = upsample2(DSP.fftshift(DSP.fft(cc.HP[i])))
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end
        us[1] = upsample2(inverseCurveletTransform(cc.LP, plan.subplan, true))
        if(Curvelet_verbose)
            println()
        end

        prev = zeros(Wsize,Wsize)
        if(Curvelet_verbose)
            print("Inverting")
        end
        for i=1:plan.angles
            if(Curvelet_verbose)
                print(".")
            end
            prev += us[i].*plan.windows[i]
        end
        if(Curvelet_verbose)
            println()
        end

        if(Curvelet_verbose)
            println("IFFT")
            println("imaginary : ",norm(imag(DSP.ifft(DSP.ifftshift(prev)))))
        end
        real(DSP.ifft(DSP.ifftshift(prev)))
    end

    function inverseCurveletTransform(cc::CurveletCoeff, plan::CurveletPlan, asdf::Bool)
        Wsize = plan.dim

        us = Array{Matrix{Complex{Float64}}}(plan.angles)
        if(Curvelet_verbose)
            print("upsampling")
        end
        for i=1:plan.angles-1
            us[i+1] = upsample2(DSP.fftshift(DSP.fft(cc.HP[i])))
            if(Curvelet_verbose)
                print(".")
            end
        end
        if(Curvelet_verbose)
            println()
        end
        #println("types are = ", typeof(cc.LP), " ", typeof(plan.subplan))
        us[1] = upsample2(inverseCurveletTransform(cc.LP, plan.subplan, true))
        if(Curvelet_verbose)
            println()
        end

        prev = zeros(Wsize,Wsize)
        if(Curvelet_verbose)
            print("Inverting")
        end
        for i=1:plan.angles
            if(Curvelet_verbose)
                print(".")
            end
            prev += us[i].*plan.windows[i]
        end
        if(Curvelet_verbose)
            println()
        end

        prev
    end

    function inverseCurveletTransform(cc::Matrix{Complex{Float64}}, plan::Void, asdf::Bool)
        DSP.fftshift(DSP.fft(cc))
    end

    #### Optimization methods
    ### While the overall transform as laid out above is quick, the generation of
    ### necessary windowing functions is not very fast if it is done with only first
    ### principle definitions.
    ### The overall windowing functions have a Region Of Support (ROS) which is
    ### mostly convex and slowly changing.
    ### This feature can be exploited to only calculate the points that are non-zero
    ### within the [-π,π]^2 square
    ###
    ### Through symmetry we can limit this optimization to the windows on the
    ### eastern side.
    ### From N&l we can estimate the approximate angle and use that to find a
    ### point within the ROS.
    ### From there the upper and lower bounds can be adjusted as the window is
    ### calculated column-wise.

    ### Fit the window bounds starting from the previous bounds
    function refit(f, low::Int, high::Int)
        while(f(low-1) != 0)
            low -= 1
        end
        while(f(low) == 0 && low != high)
            low += 1
        end
        while(f(high+1) != 0)
            high += 1
        end
        while(f(high) == 0 && high != low)
            high -= 1
        end

        return (low, high)
    end

    ## The oddness of the below code is either the result of some strange bug
    ## that I introduced or some WEIRD compiler bug with reguards to
    ## lambda precision...
    function run_cols(f, x, b, A, dir, N, l)
        #V = x::Float64->-pi+2.0*pi/size(A,1)*(x-1)
        M = size(A,1)
        (low, high) = refit(y->f(grid(x,M),grid(y,M)), b[1], b[2])

        for y=low:high
            x_::Int = mod(x-1, size(A,1))+1
            y_::Int = mod(y-1, size(A,2))+1
            A[y_, x_] = _ul((grid(x_,M),grid(y_,M)), N, l)
        end

        if(low != high)
            run_cols(f, x+dir, (low, high), A, dir, N, l)
        end
    end

    function fast_window_definition(size::Int, N::Int, l::Int)
        if(l>N)
            result = flipdim(fast_window_definition(size,N,N-l),1)'
            result[:,[2:end, 1]] = result[:, [1:end]]
            return result
        end
        fn = (x,y)->_ul((x,y),N,l)

        result = zeros(size, size)

        V = x->-pi+2*pi/size*(x-1)
        #f = (x,y)->fn(idxToVal(x), idxToVal(y))

        #-pi/2..3*pi/2 FIXME see eqn 28. T(theta)
        interp = x->(-1/2+1*x/N)
        slope  = (atan(pi*interp(l-1))+atan(pi*interp(l)))/2

        x::Int = l != 0 ? floor(size*4/5) : size/2 #TODO enhance this
        y::Int = round((size/(2*pi)*(V(x)*slope+pi))) + 1

        if(fn(V(x),V(y)) == 0)
            println("Warning: bad y value(",y,"), correcting...")
            ## Bad guess, which should not happen if the above math is
            ## right
            ## oh well, do a quick search to find a valid point
            for ny = 1:size
                if(fn(V(x),V(ny)) != 0)
                    y = ny
                    break
                end
            end
        end

        ## Real failure happens here
        @assert fn(V(x),V(y)) != 0

        bounds = refit(y->fn(V(x),V(y)), y, y)
        run_cols(fn, x, (y,y), result, -1, N, l)
        run_cols(fn, x+1, (y,y), result, +1, N, l)

        return result
    end
end


"""
    compute_E_M_Ising(σ::Array{Int8, 2})

Computes the Ising model internal energy and magnetization for the input spin configuration `σ`.

Arguments
* `σ`     : A 2D array representing a square lattice in the Ising model
Returns
* `(E, M)`: The system's energy and magnetization
"""
function compute_E_M_Ising(σ::Array{Int8,2})
    E::Int = 0
    M::Int = 0
    L = size(σ, 1)
    for j = 1:L, i = 1:L
        E -= σ[i, j] * (σ[plus(i, L), j] + σ[i, plus(j, L)])
        M += σ[i, j]
    end

    return (E, M)
end

"""
    compute_E_Potts(σ::Array{Int8, 2})

Computes the Potts model internal energy for the input spin configuration `σ`.

Arguments
* `σ` : A 2D array representing a square lattice in the Potts model
Returns
* `E` : The system's energy
"""
function compute_E_Potts(σ::Array{Int8,2})
    E::Int = 0
    L = size(σ, 1)
    for j = 1:L, i = 1:L
        E -= Int(σ[i, j] == σ[plus(i, L), j]) + Int(σ[i, j] == σ[i, plus(j, L)])
    end

    return E
end

"""
    compute_E_M_XY(σ::Array{NTuple{2, Float64}, 2})

Computes the XY model internal energy and magnetization for the input spin configuration `σ`.

Arguments
* `σ`          : A 2D array representing a square lattice in the Ising model
Returns
* `(E, Mx, My)`: The system's energy, x-magnetization, and y-magnetization
"""
function compute_E_M_XY(σ::Array{NTuple{2,Float64},2})
    E = 0.0
    Mx = 0.0
    My = 0.0
    L = size(σ, 1)
    for j = 1:L, i = 1:L
        E -= dot(σ[i, j], σ[plus(i, L), j]) + dot(σ[i, j], σ[i, plus(j, L)])
        Mx += σ[i, j][1]
        My += σ[i, j][2]
    end

    return (E, Mx, My)
end

"""
    update_observables!(model::Ising)

Pushes the current energy and magnetization to their respective arrays.

Arguments
* `model`: Ising type
"""
function update_observables!(model::Ising)
    E, M = compute_E_M_Ising(model.σ)
    push!(model.observables.E, E)
    push!(model.observables.M, M)
end

"""
    update_observables!(model::Potts)

Pushes the current energy and magnetization to their respective arrays.

Arguments
* `model`: Potts type
"""
function update_observables!(model::Potts)
    E = compute_E_Potts(model.σ)
    push!(model.observables.E, E)
end

"""
    update_observables!(model::XY)

Pushes the current energy and magnetization to their respective arrays.

Arguments
* `model`: XY type
"""
function update_observables!(model::XY)
    E, Mx, My = compute_E_M_XY(model.σ)
    push!(model.observables.E, E)
    push!(model.observables.Mx, Mx)
    push!(model.observables.My, My)
end

"""
    compute_observables_statistics!(model::AbstractMagnetismModel)

Computes the average energy per spin, average magnetization per spin,
heat capacity per spin, magnetic susceptibility per spin, Binder parameter,
and the integrated autocorrelation times (computed using the binning method).
"""
function compute_observables_statistics!(model::AbstractMagnetismModel)
    V = model.params.L^2
    model.observables.statistics[:e] = mean(model.observables.E) / V
    model.observables.statistics[:e²] = mean(model.observables.E .^ 2) / V^2
    model.observables.statistics[:c] = model.params.β^2 * var(model.observables.E) / V

    if isa(model, Ising)
        model.observables.statistics[:m] = mean(abs.(model.observables.M)) / V
        model.observables.statistics[:m²] = mean(model.observables.M .^ 2) / V^2
        model.observables.statistics[:χ] =
            model.params.β * var(abs.(model.observables.M)) / V
        model.observables.statistics[:U] =
            1 - mean(model.observables.M .^ 4) / 3 / mean(model.observables.M .^ 2)^2
    elseif isa(model, XY)
        model.observables.statistics[:mx] = mean(model.observables.Mx) / V
        model.observables.statistics[:mx²] = mean(model.observables.Mx .^ 2) / V^2
        model.observables.statistics[:my] = mean(model.observables.My) / V
        model.observables.statistics[:my²] = mean(model.observables.My .^ 2) / V^2
        model.observables.statistics[:m²] =
            mean(model.observables.Mx .^ 2 .+ model.observables.My .^ 2) / V^2
        model.observables.statistics[:χ] = model.observables.statistics[:m²] * V
    end

    model.observables.statistics[:τ_int] =
        compute_autocorrelation_time(model, collect(1:400))
end


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

function sample(RNG::Random.AbstractRNG,w::Vector)
    w /= sum(w)
    p = rand(RNG)
    runningsum = 0
    for i=1:length(w)
        runningsum += w[i]
        if p < runningsum
            return i
        end
    end
end

sample(w::Vector) = sample(Random.globalRNG(),w)

###---- FUNCTIONS FOR GENERATING CONDITIONED RANDOM WALK ----------

lgamma(x) = first(SpecialFunctions.logabsgamma(x))
lbinom(x) = first(SpecialFunctions.logabsbinom(x))
lbinom(n,k) = lgamma(n+1) - lgamma(k+1) - lgamma(n-k+1)
p0(a,c,n::BigInt) = mod(n+c-a,2) ≠ 0 ? 0.0 : exp(lbinom(n,(n+c-a)/2) + n * log(1/2))
p0float(a,c,n) = exp(lbinom(n,(n+c-a)/2) + n * log(1/2))
p0precise(a,c,n) = Float64(binomial(big(n),big(div(n+c-a,2))) * big(1/2)^n)
function p0(a,c,n::Int64)
    if mod(n+c-a,2) ≠ 0 || abs(a-c) > n
        return 0.0
    elseif n < 15_000
        return p0float(a,c,n)
    else
        # Use binomial approxmation
        k = div(n+c-a,2)
        return sqrt(2/(pi*n))*exp(-2*(k-n/2)^2/n)
    end
end
p(a,c,n) = p0(a,c,n) - p0(a,-c,n)
q(a,b,n) = p(b,3,n)*p(a,1,n) - p(b,1,n)*p(a,3,n)

function normalize(v)
    return v / sum(v)
end

function transweights(a,b,n)
    return normalize([q(i,j,n) for i=a-1:2:a+1, j=b-1:2:b+1])
end

function takestep(RNG::Random.AbstractRNG, a, b, n)
    return [(i,j) for j=b-1:2:b+1 for i=a-1:2:a+1][
                            sample(RNG,transweights(a,b,n)[:])]
end

takestep(a,b,n) = takestep(Random.GLOBAL_RNG,a,b,n)

function contour_to_tree(v::Array{Int64,1})
    n = div(length(v)-1,2)
    nbs = [Int64[] for _=1:n+3]
    stack = Int64[n+1]
    heads = Int64[]
    num_ones = 0
    bluetree = Dict{Int64,Int64}()
    for s in diff(v)
        if s == 1
            push!(stack,num_ones+1)
            num_ones += 1
        else
            e = pop!(stack)
            push!(heads,e)
            bluetree[e] = stack[end]
            push!(nbs[e],stack[end])
            push!(nbs[stack[end]],e)
        end
    end
    return PlanarMap(map(reverse,nbs)), bluetree, heads
end

function pair_dyck_paths(RNG::Random.AbstractRNG,
                         n::Integer;
                         verbose=false)
    X = Int64[1]
    Y = Int64[3]
    power_of_ten = 10^floor(Integer,log10(n/2))
    for k=1:n
        (a,b) = takestep(RNG,X[end],Y[end],(n-k))
        push!(X,a)
        push!(Y,b)
    end
    X = X .- 1
    Y = Y .- 3
    return X, Y
end

pair_dyck_paths(n::Integer;kwargs...) =
    pair_dyck_paths(Random.GLOBAL_RNG,n)
#----------------------------------------------------

#--- FUNCTIONS FOR GENERATING WOODS -----------------

import Base: keys, getindex
keys(R::RootedTree) = keys(R.ancestors)
getindex(R::RootedTree,k) = getindex(R.ancestors,k)

function add_red_tree(M::PlanarMap,
                      bluetree::Dict{Int64,Int64},
                      Y::Array{Int64,1},
                      heads::Array{Int64,1})
    nbs = deepcopy(M.nbs)
    headDict = Dict(map(reverse,enumerate(heads)))
    n = (length(Y)-1) ÷ 2
    tails = Int64[]
    num_ones = 0
    redtree = Dict{Int64,Int64}()
    tails = cumsum(diff(Y).==1)[findall(diff(Y) .== -1)]
    tails = [t == n ? n+2 : t+1 for t in tails]
    unmatchedHeadInds = Set(1:length(heads))
    for tail in tails
        matchInd = tail == n+2 ? n : findfirst(x -> x ≥ tail, heads) - 1
        while matchInd ∉ unmatchedHeadInds
            matchInd -= 1
        end
        pop!(unmatchedHeadInds,matchInd)
        head = heads[matchInd]
        redtree[head] = tail
        nbs[head] = insertccw(nbs[head],bluetree[head],tail)
        if tail == n+2
            if length(nbs[tail]) == 0
                nbs[tail] = NeighborCycle([head])
            else
                nbs[tail] = insertcw(nbs[tail],nbs[tail][1],head)
            end
        else
            cwneighbor = bluetree[tail]
            while true
                cwneighbor = cw(M,tail,cwneighbor)
                if (cwneighbor == n+1 ||
                    bluetree[cwneighbor] == tail ||
                        bluetree[tail] == cwneighbor)
                    break
                end
            end
            nbs[tail] = insertccw(nbs[tail],cwneighbor,head)
        end
    end
    PlanarMap(nbs), redtree
end

function find_split(F::Face,
                    bluetree::Dict{Int64,Int64},
                    redtree::Dict{Int64,Int64})
    # finds which vertex in each face has two incident faces outgoing red and blue
    f = [F.elements; F.elements[1:1]]
    for (i,v) in enumerate(f)
        prv = f[i == 1 ? length(f)-1 : i - 1]
        nxt = f[i == length(f)-1 ? 1 : i + 1]
        if v in keys(redtree) && redtree[v] == prv && bluetree[v] == nxt
            return v, prv, nxt
        end
    end
    error("Split not found for face "*string(f))
    (0,0,0)
end

function rotate(a::Array{Int64,1},n::Integer)
    vcat(a[n:end],a[1:n-1])
end

import Base.getindex, Base.findfirst
getindex(Z::Base.Iterators.Zip,k::Integer) = (Z.a[k],Z.b[k])
findfirst(C::Face,t::Tuple) = findfirst(isequal(t),pairs(C))

function add_green_tree(M::PlanarMap,
                        bluetree::Dict{Int64,Int64},
                        redtree::Dict{Int64,Int64})
    n = maximum(keys(bluetree))
    nbs = deepcopy(M.nbs)
    greentree = Dict{Int64,Int64}()
    found = false # whether outer face has been found yet
    FL = faces(M;outeredge=(n+1,1))
    # connect edges on left side of the planar map to
    # the green root:
    F = outerface(FL)
    leftedges = F[1:findfirst(isequal(n+2), F.elements)-1]
    for (i,w) in enumerate(leftedges[2:end])
        greentree[w] = n+3
        nbs[w] = insertcw(nbs[w],leftedges[i],n+3)
        if length(nbs[n+3]) == 0
            nbs[n+3] = NeighborCycle([w])
        else
            nbs[n+3] = insertccw(nbs[n+3],nbs[n+3][end],w)
        end
    end
    # triangulate each interior face with green edges
    for F in interiorfaces(FL)
        if length(F) > 3
            v, prv, nxt = find_split(F,bluetree,redtree)
            # idx stores index in the neighbor list of v
            # where incoming green edges should be inserted
            fc = rotate(F,findfirst(isequal(nxt), F))[2:end-1]
            for (i,w) in enumerate(fc[1:end-1])
                greentree[w] = v
                nbs[v] = insertcw(nbs[v],prv,w)
                nbs[w] = insertccw(nbs[w],fc[i+1],v)
            end
        end
    end
    return PlanarMap(nbs), greentree
end

function UWT(RNG::Random.AbstractRNG,n::Integer)
    X,Y = pair_dyck_paths(RNG,2n)
    M, bluetree, heads = contour_to_tree(X)
    M, redtree = add_red_tree(M,bluetree,Y,heads)
    M, greentree = add_green_tree(M,bluetree,redtree)
    add_edge!(M,n+1,n+2)
    add_edge!(M,n+2,n+3)
    add_edge!(M,n+3,n+1)
    return WoodedTriangulation(M,map(RootedTree,
                    (bluetree,redtree,greentree),
                    (n+1,n+2,n+3))...)
end

UWT(n::Integer) = UWT(Random.GLOBAL_RNG,n)

function descendants(d::Dict{Int64,Int64})
    n = maximum(keys(d))
    desc = zeros(Int64,n)
    for j = 1:n
        k = j
        while true
            k = d[k]
            if k < n+1
                desc[k] += 1
            else
                break
            end
        end
    end
    return desc
end

function flowline(v::Int64,d::Dict{Int64,Int64})
    n = maximum(keys(d))
    branch = Int64[v]
    while branch[end] < n+1
        push!(branch,d[branch[end]])
    end
    return branch
end

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


function drawcircle!(img::Matrix{T}, col::T, x0::Int, y0::Int, radius::Int) where T
    x = radius - 1
    y = 0
    δx = δy = 1
    er = δx - (radius << 1)

    s = x + y
    while x ≥ y
        for opx in (+, -), opy in (+, -), el in (x, y)
            @inbounds img[opx(x0, el) + 1, opy(y0, s - el) + 1] = col
        end
        if er ≤ 0
            y  += 1
            er += δy
            δy += 2
        end
        if er > 0
            x  -= 1
            δx += 2
            er += (-radius << 1) + δx
        end
        s = x + y
    end
    return img
end

# Test
using Images

img = fill(Gray(255.0), 25, 25);
drawcircle!(img, Gray(0.0), 12, 12, 12)


########################################################
#           ONE LAYER NEURONAL NETWORK TYPE
#   and evaluation with forward propagation
########################################################
# Constructores
export RandomWeightsNN
export FromMatrixNN
# Tipo
export AbstractOneLayerNeuralNetwork

"""
    AbstractOneLayerNeuralNetwork
The basic elements that define a one layer neural network
Must have two matrix:
W1: Matrix n x (d+1)
W2: Matris s x n
"""
abstract type AbstractOneLayerNeuralNetwork end

"""
    RandomWeightsNN
Return a random initialized Neuronal Network
"""
mutable struct RandomWeightsNN  <: AbstractOneLayerNeuralNetwork
    entry_dimesion :: Int
    number_of_hide_units :: Int
    output_dimension :: Int
    W1  # pesos de la entrada a la capa oculta A S (sesgo última columna)
    W2  # pesos de la capa oculta a la salida

    function RandomWeightsNN(entry_dimesion,
                                                number_of_hide_units,
                                                output_dimension)

        W1 = rand(Float64, number_of_hide_units, entry_dimesion+1)
        W2 = rand(Float64, output_dimension, number_of_hide_units)
        return new(
            entry_dimesion,
            number_of_hide_units,
            output_dimension,
            W1,
            W2
        )

    end
end

"""
    RandomWeightsNN
Return a  Neuronal Network inizialized  by three matrix
"""
mutable struct FromMatrixNN <: AbstractOneLayerNeuralNetwork
    W1 :: Matrix # pesos de la entrada a la capa oculta
    W2 :: Matrix# pesos de la capa oculta a la salida
    function FromMatrixNN(S,A,B)
        # Comprobación de que los tipos son correctos
        if !( typeof(S) <: Vector && typeof(A) <: Matrix && typeof(B) <: Matrix )
            throw(ArgumentError("El tipo de los argumentos no es el correcto\n
            Debería de ser:\n
             typeof(S) <: Vector && typeof(A) <: Matrix && typeof(B) <: Matrix \n
            pero se ha encontrado: \n
                typeof(S)= $(typeof(S))  typeof(A) $(typeof(A))  typeof(B) $(typeof(B))
            "))
        end
        # Comprobaciones de que los tamaños son coherentes
        (n_a,d_a) = size(A)
        l_s = length(S)
        (s_b, n_b) = size(B)
        ## Coherencia A y S
        if n_a != l_s
            throw(ArgumentError("El número de filas de A (que es $(n_a))debe de ser igual que
            la longitud de S (que es $(l_s))
            Los tamaños encontrados son:
                size(S)=$(size(S)).
                size(A)=$(size(A))
            "))
        end
        # Coherencia A y B
        if n_a != n_b
            throw(ArgumentError("El número de fila de A (que es $(n_a)) no es coherente con el número de columnas de B (que es $(n_b)).
            Ambos debería de ser iguales."))
        end
        return new(hcat(A,S), B)
    end
end

"""
    Base.show(io::IO, h <:AbstractOneLayerNeuralNetwork)
Implementamos el algoritmo de visualización de nuestras matrices
"""
function Base.show(io::IO, h ::AbstractOneLayerNeuralNetwork)
    display(Text("La matrix de pesos de las neuronas, W1, es:\n"))
    display(h.W1)
    display(Text("\nLa matrix de pesos de la salida, W2, es:\n"))
    display(h.W2)
end


using ExpmV
using Base.Test

for i = 1:100, d in 10:10:60, herm in [true, false]

    d = 60

    r = sprandn(d,d,.1)+1im*sprandn(d,d,.1)
    if herm
        r = (r-r')/2
    end

    rv = randn(d)+1im*randn(d)
    rv = rv / norm(rv,2)

    rt = randn()

    x = expmv(rt,r,rv)

    @test_approx_eq_eps norm(x-expm(full(rt*r))*rv,2) 0.0 1e-9
end


# Logistic regression model for the Prima Indian data

# load packages
using Pkg
using PyPlot
using LinearAlgebra
using Random
using Printf
using GLM
using DataFrames

# include files
include(pwd()*"/utilities.jl")

Random.seed!(1234) # set random numbers

# load data
x_train,y_train,x_val,y_val,x_test,y_test = load_data([0,1],0.2)

# create data frame
training_data = zeros(size(x_train,2), size(x_train,1)+1)
training_data[:,1:end-1] = x_train'
training_data[:,end] = y_train
train_df = DataFrame(training_data)
test_df = DataFrame(x_test')

model = @formula(x9 ~ x1+x2+x3+x4+x5+x6+x7+x8)

@time logisticregresson = glm(model, train_df, Binomial(), LogitLink())

print(logisticregresson)

# compute predictions
y_pred_train = class_logistic(predict(logisticregresson))
y_pred_test = class_logistic(predict(logisticregresson,test_df))

# check classification results
class_res_training = classification_results(y_pred_train, y_train, [0,1])
class_res_test = classification_results(y_pred_test, y_test, [0,1])


function sumdigits(number)
    rest = abs(number)
    sum = 0
    while rest > 0
      sum += rest % 10
      rest = rest ÷ 10
    end
    return sum
end


abstract type AbstractMedium{T<:Real} end

"""
    refractive_index(medium :: AbstractMedium{T}, λ :: T) where T

Return the refractive index of medium at the given wavelength λ in μm.
"""
function refractive_index end

"""
    ConstantIndex{T}

Definition of a medium with no dispersion.

# Fields
- `n :: T`: The refractive index of the medium
"""
struct ConstantIndex{T} <: AbstractMedium{T}
    n::T
end

function refractive_index(medium::ConstantIndex{T}, λ::T) where {T}
    medium.n
end

"""
    Sellmeier_1{T}

Definition of a material's dispersion using the Sellmeier equation. This definition is
equivalent to the Zemax Sellmeier_1 dispersion equation.

``n^2(λ) = 1 + \\frac{k_1 λ^2}{λ^2 - l_1} + \\frac{k_2 λ^2}{λ^2 - l_1} + \\frac{k_3 λ^2}{λ^2 - l_3}``
"""
struct Sellmeier_1{T} <: AbstractMedium{T}
    k1::T
    l1::T
    k2::T
    l2::T
    k3::T
    l3::T
end

function refractive_index(medium::Sellmeier_1{T}, λ::T) where {T}
    λ2 = λ^2
    a = medium.k1 * λ2 / (λ2 - medium.l1)
    b = medium.k2 * λ2 / (λ2 - medium.l2)
    c = medium.k3 * λ2 / (λ2 - medium.l3)
    sqrt(a + b + c + 1)
end

#
# Definition of a few materials
#

"""
    Silica

Dispersion of fused silica, defined as Sellmeier equation coefficients.
The coefficients match Zemax.
"""
const Silica = Sellmeier_1(
    6.96166300e-1,
    4.67914800e-3,
    4.07942600e-1,
    1.35120600e-2,
    8.974794e-1,
    9.7934e+1,
)

"""
    Air

The refractive index of air is defined to be 1.0 at all wavelengths. This follows the usual
convention to define refractive indices relative to air.
"""
const Air = ConstantIndex(1.0)


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

# Compute X <- a X + b Y.
function matfun_axpby!(X,a,b,Y)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:m
        @inbounds for j=1:n
            if (b isa ValueOne)
                X[i,j]+=Y[i,j]
            else
                X[i,j]+=b*Y[i,j]
            end
        end
    end
end

@inline function exp_native_jl_rho0_015(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_native_jl_rho0_015!(A_copy)
end

@inline function exp_native_jl_rho0_015!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=4
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: A2 Ua U V Z X P
    # Computing A2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing Ua = x*I+x*A2
    coeff1=60.0
    coeff2=1.0
    memslots3 .= coeff2.*memslots2
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing U with operation: mult
    mul!(memslots4,memslots3,memslots1)
    # Deallocating Ua in slot 3
    # Deallocating A in slot 1
    # Computing V = x*I+x*A2
    coeff1=120.0
    coeff2=12.0
    # Smart lincomb recycle A2
    memslots2 .= coeff2.*memslots2
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Z = x*V+x*U
    coeff1=1.0
    coeff2=-1.0
    memslots1 .= coeff1.*memslots2 .+ coeff2.*memslots4
    # Computing X = x*V+x*U
    coeff1=1.0
    coeff2=1.0
    # Smart lincomb recycle V
    memslots2 .= coeff1.*memslots2 .+ coeff2.*memslots4
    # Deallocating U in slot 4
    # Computing P with operation: ldiv
    memslots3 .=memslots1\memslots2
    # Deallocating Z in slot 1
    # Deallocating X in slot 2
    return memslots3 # Returning P
end


using Unitful
Unitful.preferunits(u"cm,g,s"...)

function main()

    a = 9.81*u"m/s^2"
    v0 = 5u"m/s"

    x(t) = v0*t + 1/2*a*t^2
    v(t) = v0 + a*t

    t1 = 5*u"s"

    # displaying units
    println(a)
    println(v0)
    println(x(t1))
    println(v(t1))

    # checking unit math
    # println(x(t1) + v(t1))

    # converting units
    println(uconvert(u"yr", t1))

    # checking dimensions
    h = 10*u"ħ"
    println(h)
    println(typeof(h))

    # looking at SI units
    println(upreferred(h))

end

main()


"""
    generate_complex_potential_source(; pos = ComplexQuantity((0.0, 0.0)m)::ComplexQuantity, massflowout = 1.0m²/s)
    -> f::Z->R

We find the velocity potential function of a source or sink by realizing that
´´´
    2π·r·u_r = q
       where u_r is radial velocity
             q is flow
        ↕
    u_r = q / (2π·r)
´´´
The velocity potential, ϕ, is the imaginary valued function found by integrating
velocity along the velocity gradient. The sign and endpoint of the integration
can be defined as we want. With the limitation above:
´´´
    u_r  = ∇ϕ = δu / δr
    ⇑
    ϕ = ∫ q / (2πr) dr = q / (2πr) · ∫ 1 / r dr = q / (2πr) · ln(r)
´´´
"""
function generate_complex_potential_source(; pos = complex(0.0, 0.0)m, massflowout = 1.0m²/s)
    # Note: log means natural logarithm, not log10
    #
    # Note: natural logarithm of a quantity gives no meaning. We can drop the unit as long as we
    # use the scalar consistently.
    @assert massflowout != zero(massflowout)
    # Would call this ϕ_s(p)
    (p::ComplexQuantity) -> begin
        Δp = p - pos
        r = hypot(Δp)
        massflowout ∙ log(r / oneunit(r)) / (2π)
    end
end

"""
    generate_complex_potential_vortex(; pos = ComplexQuantity((0.0, 0.0)m)::ComplexQuantity, vorticity= 1.0m²/s)
    -> f::Z->R

We find the velocity potential function of a vortex by realizing that velocity
´´´
    u_θ = K / r
       where u_θ is tangential velocity, i.e. length per time in the tangential direction.
             K is vorticity
             Γ = K / 2π  is circulation
´´´
The velocity potential, ϕ, is the scalar valued function found by integrating
velocity along the velocity gradient. The sign and endpoint of the integration
can be defined as we want.
"""
function generate_complex_potential_vortex(; pos = complex(0.0, 0.0)m, vorticity = 1.0m²/s)
    # Would call this ϕ_v(p)
    (p::ComplexQuantity) -> begin
        Δp = p - pos
        Θ = -angle(Δp)
        vorticity ∙ Θ
    end
end

absolute_scale() = ColorSchemes.linear_grey_10_95_c0_n256
complex_arg0_scale() = ColorSchemes.linear_ternary_red_0_50_c52_n256

"Scale and move real values to fit in the closed interval 0.0 to 1.0, given minimum and maximum values in A"
normalize_datarange_real(mi, ma, A) = map(A) do x
    @assert ma != mi
    if isnan(x) || isinf(x)
        mi
    else
        (x - mi) / (ma - mi)
    end
end
"""
Scale and move the magnitude (modulus, absolute value) of complex values
to fit in the closed interval 0.0 to 1.0, without changing the argument (angle).
The argument of the minimum value(s) is lost, since the angle is undefined for zero.
"""
normalize_datarange_complex(mi, ma, A) = map(A) do x
    @assert ma != mi
    if isnan(x) || isinf(x)
        mi
    else
        θ = angle(x)
        r = (hypot(x) - mi) / (ma - mi)
        r * complex(cos(θ), sin(θ))
    end
end
"""
    normalize_datarange(A)

normalize the elements in matrix A from 0.0 to 1.0, type preserving type.
This includes complex values, in which case the magnitude is normalized.
"""
function normalize_datarange(A::AbstractArray{<:RealQuantity})
    # Minimum and maximum value
    mi, ma = lenient_min_max(A)
    normalize_datarange_real(mi, ma, A)
end
function normalize_datarange(A::AbstractArray{<:Real})
    # Minimum and maximum value
    mi, ma = lenient_min_max(A)
    normalize_datarange_real(mi, ma, A)
end
function normalize_datarange(A)
    # Minimum and maximum magnitude
    mi, ma = lenient_min_max(A)
    normalize_datarange_complex(mi, ma, A)
end


#############################################################################
# geom_mean_hypocone.jl
# Constrains T to
#   A #_t B ⪰ T
# where:
#   * A #_t B is the t-weighted geometric mean of A and B:
#     A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}
#     Parameter t should be in [0,1].
#   * Constraints A ⪰ 0, B ⪰ 0 are added.
#
# Note on parameter fullhyp:
#   In many applications one doesn't need the full hypograph
#     hyp_t = {(A,B,T) : A #_t B ⪰ T}
#   but rather it is enough to work with a convex set C_t that satisfies
#     (A,B,A #_t B) \in C_t
#     (A,B,T) \in C_t  =>  A #_t B ⪰ T
#   In this case one should set fullhyp = false. The SDP description will be
#   (slightly) smaller. (By default fullhyp is set to true).
#
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#
#REFERENCE
#   Ported from CVXQUAD which is based on the paper: "Lieb's concavity
#   theorem, matrix geometric means and semidefinite optimization" by Hamza
#   Fawzi and James Saunderson (arXiv:1512.03401)
#############################################################################

struct GeomMeanHypoCone
    A::AbstractExpr
    B::AbstractExpr
    t::Rational
    size::Tuple{Int,Int}
    fullhyp::Bool

    function GeomMeanHypoCone(
        A::AbstractExpr,
        B::AbstractExpr,
        t::Rational,
        fullhyp::Bool = true,
    )
        if size(A) != size(B)
            throw(DimensionMismatch("A and B must be the same size"))
        end
        n = size(A)[1]
        if size(A) != (n, n)
            throw(DimensionMismatch("A and B must be square"))
        end
        if t < 0 || t > 1
            throw(DomainError(t, "t must be in the range [0, 1]"))
        end
        return new(A, B, t, (n, n), fullhyp)
    end

    function GeomMeanHypoCone(
        A::Value,
        B::AbstractExpr,
        t::Rational,
        fullhyp::Bool = true,
    )
        return GeomMeanHypoCone(Constant(A), B, t, fullhyp)
    end
    function GeomMeanHypoCone(
        A::AbstractExpr,
        B::Value,
        t::Rational,
        fullhyp::Bool = true,
    )
        return GeomMeanHypoCone(A, Constant(B), t, fullhyp)
    end
    function GeomMeanHypoCone(
        A::Value,
        B::Value,
        t::Rational,
        fullhyp::Bool = true,
    )
        return GeomMeanHypoCone(Constant(A), Constant(B), t, fullhyp)
    end

    function GeomMeanHypoCone(
        A::AbstractExprOrValue,
        B::AbstractExprOrValue,
        t::Integer,
        fullhyp::Bool = true,
    )
        return GeomMeanHypoCone(A, B, t // 1, fullhyp)
    end
end

struct GeomMeanHypoConeConstraint <: Constraint
    head::Symbol
    id_hash::UInt64
    T::AbstractExpr
    cone::GeomMeanHypoCone

    function GeomMeanHypoConeConstraint(T::AbstractExpr, cone::GeomMeanHypoCone)
        if size(T) != cone.size
            throw(DimensionMismatch("T must be size $(cone.size)"))
        end
        id_hash = hash((cone.A, cone.B, cone.t, :GeomMeanHypoCone))
        return new(:GeomMeanHypoCone, id_hash, T, cone)
    end

    function GeomMeanHypoConeConstraint(T::Value, cone::GeomMeanHypoCone)
        return GeomMeanHypoConeConstraint(Constant(T), cone)
    end
end

in(T, cone::GeomMeanHypoCone) = GeomMeanHypoConeConstraint(T, cone)

function AbstractTrees.children(constraint::GeomMeanHypoConeConstraint)
    return (
        constraint.T,
        constraint.cone.A,
        constraint.cone.B,
        "t=$(constraint.cone.t)",
    )
end

# For t ∈ [0,1] the t-weighted matrix geometric mean is matrix concave (arxiv:1512.03401).
# So if A and B are convex sets, then T ⪯ A #_t B will be a convex set.
function vexity(constraint::GeomMeanHypoConeConstraint)
    A = vexity(constraint.cone.A)
    B = vexity(constraint.cone.B)
    T = vexity(constraint.T)

    # NOTE: can't say A == NotDcp() because the NotDcp constructor prints a warning message.
    if typeof(A) == ConvexVexity || typeof(A) == NotDcp
        return NotDcp()
    end
    if typeof(B) == ConvexVexity || typeof(B) == NotDcp
        return NotDcp()
    end
    # Copied from vexity(c::GtConstraint)
    vex = -ConcaveVexity() + T
    if vex == ConcaveVexity()
        vex = NotDcp()
    end
    return vex
end

function conic_form!(
    constraint::GeomMeanHypoConeConstraint,
    unique_conic_forms::UniqueConicForms,
)
    if !has_conic_form(unique_conic_forms, constraint)
        A = constraint.cone.A
        B = constraint.cone.B
        t = constraint.cone.t
        T = constraint.T
        fullhyp = constraint.cone.fullhyp

        is_complex =
            sign(A) == ComplexSign() ||
            sign(B) == ComplexSign() ||
            sign(T) == ComplexSign()
        if is_complex
            make_temporary = () -> HermitianSemidefinite(size(A)[1])
        else
            make_temporary = () -> Semidefinite(size(A)[1])
        end

        if fullhyp && t != 0 && t != 1
            W = make_temporary()
            conic_form!(
                W in GeomMeanHypoCone(A, B, t, false),
                unique_conic_forms,
            )
            conic_form!(W ⪰ T, unique_conic_forms)
            cache_conic_form!(
                unique_conic_forms,
                constraint,
                Array{Convex.ConicConstr,1}(),
            )
        else
            p = t.num
            q = t.den

            if t == 0
                #println("geom_mean_hypocone p=$p q=$q t==0")
                conic_form!(A ⪰ 0, unique_conic_forms)
                conic_form!(B ⪰ 0, unique_conic_forms)
                conic_form!(A ⪰ T, unique_conic_forms)
            elseif t == 1
                #println("geom_mean_hypocone p=$p q=$q t==1")
                conic_form!(A ⪰ 0, unique_conic_forms)
                conic_form!(B ⪰ 0, unique_conic_forms)
                conic_form!(B ⪰ T, unique_conic_forms)
            elseif t == 1 // 2
                #println("geom_mean_hypocone p=$p q=$q t==1/2")
                conic_form!([A T; T' B] ⪰ 0, unique_conic_forms)
            elseif ispow2(q)
                #println("geom_mean_hypocone p=$p q=$q ispow2(q)")
                Z = make_temporary()
                if t < 1 / 2
                    conic_form!(
                        Z in GeomMeanHypoCone(A, B, 2 * t, false),
                        unique_conic_forms,
                    )
                    conic_form!([A T; T' Z] ⪰ 0, unique_conic_forms)
                else
                    conic_form!(
                        Z in GeomMeanHypoCone(A, B, 2 * t - 1, false),
                        unique_conic_forms,
                    )
                    conic_form!([B T; T' Z] ⪰ 0, unique_conic_forms)
                end
            elseif ispow2(p) && t > 1 // 2
                #println("geom_mean_hypocone p=$p q=$q ispow2(p) && t>1/2")
                Z = make_temporary()
                conic_form!(
                    Z in GeomMeanHypoCone(A, T, (2 * p - q) // p, false),
                    unique_conic_forms,
                )
                conic_form!([Z T; T B] ⪰ 0, unique_conic_forms)
            elseif t < 1 / 2
                #println("geom_mean_hypocone p=$p q=$q t<1/2")
                X = make_temporary()
                # Decompose t = (p/2^l) * (2^l/q) where l=floor(log2(q))
                l = floor(Int, log2(q))
                conic_form!(
                    X in GeomMeanHypoCone(A, B, p // (2^l), false),
                    unique_conic_forms,
                )
                conic_form!(
                    T in GeomMeanHypoCone(A, X, (2^l) // q, false),
                    unique_conic_forms,
                )
            else
                #println("geom_mean_hypocone p=$p q=$q else")
                conic_form!(
                    T in GeomMeanHypoCone(B, A, 1 - t, false),
                    unique_conic_forms,
                )
            end

            cache_conic_form!(
                unique_conic_forms,
                constraint,
                Array{Convex.ConicConstr,1}(),
            )
            #cache_conic_form!(unique_conic_forms, constraint, ConicConstr([A,B,T], :geom_mean_hypocone, [size(A), size(B), size(T)]))
            #cache_conic_form!(unique_conic_forms, conic_form!(T, unique_conic_forms))
        end
    end
    return get_conic_form(unique_conic_forms, constraint)
end


cd(@__DIR__); include("setups/simpletree.jl")
gr(dpi = 200)

## build rNGWF
distROT = natural_eigdist(𝚽, 𝛌, Q; α = 1.0, input_format = :pmf1, distance = :ROT)
rNGWF, dic_l2x = rngwf_all_vectors(distROT, 𝚽; σ = 0.1 * maximum(distROT), thres = 0.15)
Γ = rngwf_lx(dic_l2x)

just_eigenvec = [2, 16]
important_idx = sortperm((rNGWF' * f).^2; rev = true)
plot(layout = Plots.grid(3, 7), size = (1400, 1500))
for i = 1:21
    l, x = Γ[important_idx[i]]
    w = rNGWF[:, important_idx[i]]
    w ./= norm(w, 2)
    scatter!(X[:, 1], X[:, 2], marker_z = w, ms = 5, c = :viridis, subplot = i, mswidth = 0)
    plot!(frame = :none, cbar = false, grid = false, legend = false, subplot = i)
    if l in just_eigenvec
        plot!(title = latexstring("\\bar{\\psi}_{", l, ",", x, "} \\equiv \\mathbf{\\phi}_{$(l)}"), titlefont = 18, subplot = i)
    else
        plot!(title = latexstring("\\bar{\\psi}_{", l, ",", x, "}"), titlefont = 18, subplot = i)
    end
end
plt = current()
savefig(plt, "../figs/simpletree_f_rngwf_top21.png")