text:
using Test, Statistics, LinearAlgebra
using DynamicAxisWarping, SlidingDistancesBase
using Distances, Plots

@testset "DynamicAxisWarping" begin
    @info "Testing DynamicAxisWarping"

    @testset "Normalizers" begin
        @info "Testing Normalizers"
        a = randn(2,100)
        @test dtwnn(a,a,SqEuclidean(),3,normalizer=IsoZNormalizer).cost < 1e-20
        @test dtwnn(a,a,SqEuclidean(),3,normalizer=ZNormalizer).cost < 1e-20
    end

    @testset "Basic Dynamic Time Warping" begin
        a = Float64[1, 1, 1, 2, 4, 6, 5, 5, 5, 4, 4, 3, 1, 1, 1]
        b = Float64[1, 1, 2, 4, 6, 6, 6, 5, 4, 4, 4, 3, 3, 3, 1]
        cost, match1, match2 = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == 0
        @test match1 ==
              [1, 2, 3, 4, 5, 6, 6, 6, 7, 8, 9, 10, 10, 11, 12, 12, 12, 13, 14, 15]
        @test match2 == [1, 1, 2, 3, 4, 5, 6, 7, 8, 8, 8, 9, 10, 11, 12, 13, 14, 15, 15, 15]
        @test evaluate(DTW(10), a, b) == cost

        @test @inferred(soft_dtw_cost(Float64.(a),Float64.(b), γ=0.001)) > -0.01

        a[end] += 2
        cost, match1, match2 = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == 4
        @test evaluate(DTW(10), a, b) == cost
        cost, match1, match2 = dtw(Float64.(a), Float64.(b), transportcost=1.1)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test evaluate(DTW(radius=10, transportcost=1.1), a, b) == cost

        @test dtw_cost(a, b, SqEuclidean(), 0) ≈ norm(a-b)^2 # Test that radius 0 reduces to SqEuclidean distance
        @test dtw_cost(a, b, Euclidean(), 0) ≈ sum(abs, a-b) # Test that radius 0 reduces to SqEuclidean distance

        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.001) ≈ cost rtol = 1e-2
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.001) == SoftDTW(0.001)(a,b)

        a = collect(1.0:10)
        b = a .+ 1
        cost, match1, match2 = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == 2
        @test evaluate(DTW(10), a, b) == cost
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.01) ≈ cost rtol = 1e-2
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.01) == SoftDTW(0.01)(a,b) ≈ SoftDTW(0.01)(big.(a),big.(b))

        a = zeros(6)
        b = 1 .+ a
        cost, match1, match2 = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == length(a)
        @test evaluate(DTW(10), a, b) == cost
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.01) ≈ cost rtol = 1e-2
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.01) == SoftDTW(0.01)(a,b)

        # Verify that a tie prefers diagonal moves
        a = Float64[1, 1, 1]
        b = Float64[1, 1, 1]
        cost, pa, pb = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == 0
        @test pa == [1, 2, 3]
        @test pb == [1, 2, 3]
        @test evaluate(DTW(10), a, b) == cost
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.001) ≈ cost atol = 1e-2
        @test soft_dtw_cost(Float64.(a),Float64.(b), γ=0.001) == SoftDTW(0.001)(a,b)

        # Verify that trackback ends properly if it reaches an edge before reaching [1,1]
        # Also check that trackback prefers diagonal moves
        a = [0, 1, 1, 1]
        b = [0, 0, 1, 1]
        cost, pa, pb = dtw(a, b)
        @test dtw_cost(a, b, SqEuclidean(), length(a)) == cost
        @test cost == 0
        @test pa == [1, 1, 2, 3, 4]
        @test pb == [1, 2, 3, 3, 4]
        @test evaluate(DTW(10), a, b) == cost

        # test the distance api with different distances
        a, b = randn(10), randn(10)
        cost, = dtw(a, b, Euclidean())
        @test evaluate(DTW(10,Euclidean()), a, b) == cost
        @test dtw_cost(a, b, Euclidean(), length(a)) == cost
        cost, = dtw(a, b, Cityblock())
        @test evaluate(DTW(10,Cityblock()), a, b) == cost
        @test dtw_cost(a, b, Cityblock(), length(a)) == cost
        cost, = dtw(a, b, Chebyshev())
        @test evaluate(DTW(10,Chebyshev()), a, b) == cost
        @test dtw_cost(a, b, Chebyshev(), length(a)) == cost

        @test_nowarn dtwplot(a, b)
        @test_nowarn matchplot(a, b)
    end

    @testset "DTW with windows" begin
        # Verify that a tie prefers diagonal moves
        dist = SqEuclidean()
        a = [1, 1, 1]
        b = [1, 1, 1]
        cost, pa, pb = dtw(a, b, dist, [1, 1, 1], [3, 3, 3])
        @test cost == 0
        @test pa == [1, 2, 3]
        @test pb == [1, 2, 3]

        # Verify that trackback ends properly if it reaches an edge before reaching [1,1]
        # Also check that trackback prefers diagonal moves
        a = [0, 1, 1, 1]
        b = [0, 0, 1, 1]
        cost, pa, pb = dtw(a, b, dist, [1, 1, 1, 1], [4, 4, 4, 4])
        @test cost == 0
        @test pa == [1, 1, 2, 3, 4]
        @test pb == [1, 2, 3, 3, 4]

        # First do the windowed test w/o windows
        a = [0, 1, 2, 3, 4, 4, 4, 4]
        b = [0, 0, 1, 2, 2, 2, 3, 4]
        best_pa = [1, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8]
        best_pb = [1, 2, 3, 4, 5, 6, 7, 8, 8, 8, 8]
        cost, pa, pb = dtw(a, b)
        @test cost == 0
        @test pa == best_pa
        @test pb == best_pb

        # Wide window, not touching optimal path
        rmin = [1, 1, 1, 2, 3, 4, 5, 6]
        rmax = [4, 6, 7, 8, 8, 8, 8, 8]
        cost, pa, pb = dtw(a, b, dist, rmin, rmax)
        @test cost == 0
        @test pa == best_pa
        @test pb == best_pb

        # Bottom of window is optimal path
        rmin = [1, 3, 4, 7, 8, 8, 8, 8]
        rmax = [4, 6, 7, 8, 8, 8, 8, 8]
        cost, pa, pb = dtw(a, b, dist, rmin, rmax)
        @test cost == 0
        @test pa == best_pa
        @test pb == best_pb

        # Top of window is optimal path
        rmin = [1, 1, 1, 2, 3, 4, 5, 6]
        rmax = [2, 3, 6, 7, 8, 8, 8, 8]
        cost, pa, pb = dtw(a, b, dist, rmin, rmax)
        @test cost == 0
        @test pa == best_pa
        @test pb == best_pb

        # Top and bottom of window are optimal path
        rmin = [1, 3, 4, 7, 8, 8, 8, 8]
        rmax = [2, 3, 6, 7, 8, 8, 8, 8]
        cost, pa, pb = dtw(a, b, dist, rmin, rmax)
        @test cost == 0
        @test pa == best_pa
        @test pb == best_pb

        # Now top of window cuts into optimal path
        rmin = [1, 1, 1, 2, 3, 4, 5, 6]
        rmax = [4, 4, 5, 6, 7, 8, 8, 8]
        cost, pa, pb = dtw(a, b, dist, rmin, rmax)
        @test cost == 2
        @test pa == [1, 1, 2, 3, 3, 4, 5, 6, 7, 8]
        @test pb == [1, 2, 3, 4, 5, 6, 7, 8, 8, 8]

        # Compare windowed and regular on non-square geometry:
        seq1 = collect(1:16)
        seq2 = seq1[1:2:end]
        cost, pa, pb = dtw(seq1, seq2)
        n1, n2 = length(seq1), length(seq2)
        cost2, qa, qb = dtw(seq1, seq2, dist, fill(1, n1), fill(n2, n1))
        @test cost == cost2
        @test pa == qa
        @test pb == qb

        seq1 = [
            1.2051,
            1.7890,
            2.1314,
            2.2423,
            2.2101,
            2.0234,
            1.6898,
            1.3422,
            1.1530,
            1.2197,
            1.5435,
            2.07147,
            2.7359,
            3.469,
            4.2033,
            4.8704,
            5.4056,
            5.75243,
            5.8693,
            5.7362,
            5.35974,
            4.7761,
            4.0482,
            3.2548,
            2.4764,
        ]
        seq2 = [
            1.4927,
            2.1859,
            2.1182,
            1.5186,
            1.1858,
            1.8034,
            3.0968,
            4.5317,
            5.5763,
            5.8038,
            5.0724,
            3.6576,
            2.4823,
        ]

        cost, pa, pb = dtw(seq1, seq2)
        n1, n2 = length(seq1), length(seq2)
        cost2, qa, qb = dtw(seq1, seq2, dist, fill(1, n1), fill(n2, n1))
        @test cost == cost2
        @test pa == qa
        @test pb == qb
    end

    @testset "FastDTW compression" begin
        compress2 = DynamicAxisWarping.compress2
        s = collect(0:2:98)
        s1 = compress2(s)
        s2 = compress2(s1)
        @test s1 == float(collect(1:4:97))
        @test s2 == float(vcat(collect(3:8:91), [97]))

        s = [1]
        s1 = compress2(s)
        @test s1 == [1.0]
    end

    @testset "Window Computations" begin
        computewindow = DynamicAxisWarping.computewindow

        # Simplest path (along the diagonal)
        p = collect(1:8)
        rmin, rmax = computewindow(p, p, 1)
        @test rmin == [1, 1, 1, 2, 3, 4, 5, 6]
        @test rmax == [3, 4, 5, 6, 7, 8, 8, 8]

        rmin, rmax = computewindow(p, p, 2)
        @test rmin == [1, 1, 1, 1, 1, 2, 3, 4]
        @test rmax == [5, 6, 7, 8, 8, 8, 8, 8]

        # A warpy path
        pa = [1, 1, 2, 3, 4, 5, 6, 7, 8, 8, 8]
        pb = [1, 2, 3, 3, 3, 4, 4, 5, 6, 7, 8]
        rmin, rmax = computewindow(pa, pb, 1)
        @test pa[end] == length(rmin)
        @test pa[end] == length(rmax)
        @test rmin == [1, 1, 2, 2, 2, 3, 3, 4]
        @test rmax == [4, 4, 4, 5, 5, 6, 8, 8]

        rmin, rmax = computewindow(pa, pb, 2)
        @test pa[end] == length(rmin)
        @test pa[end] == length(rmax)
        @test rmin == [1, 1, 1, 1, 1, 1, 2, 2]
        @test rmax == [5, 5, 6, 6, 7, 8, 8, 8]

        rmin, rmax = computewindow(pa, pb, 20)
        @test pa[end] == length(rmin)
        @test pa[end] == length(rmax)
        @test rmin == fill(1, 8)
        @test rmax == fill(8, 8)

        # Extreme path: follows left then upper edge
        pa = [1, 1, 1, 1, 1, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8]
        pb = [1, 2, 3, 4, 5, 6, 7, 8, 8, 8, 8, 8, 8, 8, 8]
        rmin, rmax = computewindow(pa, pb, 1)
        @test pa[end] == length(rmin)
        @test pa[end] == length(rmax)
        @test rmin == [1, 1, 7, 7, 7, 7, 7, 7]
        @test rmax == [8, 8, 8, 8, 8, 8, 8, 8]

        rmin, rmax = computewindow(pa, pb, 2)
        @test pa[end] == length(rmin)
        @test pa[end] == length(rmax)
        @test rmin == [1, 1, 1, 6, 6, 6, 6, 6]
        @test rmax == [8, 8, 8, 8, 8, 8, 8, 8]

        # More columns than rows
        pa = [1, 2, 3, 4, 5, 6, 7, 8]
        pb = [1, 2, 3, 4, 4, 4, 4, 4]
        rmin, rmax = computewindow(pa, pb, 1)
        @test rmin == [1, 1, 1, 2, 3, 3, 3, 3]
        @test rmax == [3, 4, 4, 4, 4, 4, 4, 4]

        rmin, rmax = computewindow(pa, pb, 2)
        @test rmin == [1, 1, 1, 1, 1, 2, 2, 2]
        @test rmax == fill(4, 8)

        rmin, rmax = computewindow(pa, pb, 3)
        @test rmin == fill(1, 8)
        @test rmax == fill(4, 8)

        rmin, rmax = computewindow(pa, pb, 4)
        @test rmin == fill(1, 8)
        @test rmax == fill(4, 8)

        rmin, rmax = computewindow(pa, pb, 47)
        @test rmin == fill(1, 8)
        @test rmax == fill(4, 8)

        # More rows than columns
        pa, pb = pb, pa
        rmin, rmax = computewindow(pa, pb, 1)
        @test rmin == [1, 1, 1, 2]
        @test rmax == [3, 4, 8, 8]

        rmin, rmax = computewindow(pa, pb, 2)
        @test rmin == fill(1, 4)
        @test rmax == [5, 8, 8, 8]

        rmin, rmax = computewindow(pa, pb, 3)
        @test rmin == fill(1, 4)
        @test rmax == fill(8, 4)

        rmin, rmax = computewindow(pa, pb, 4)
        @test rmin == fill(1, 4)
        @test rmax == fill(8, 4)

        rmin, rmax = computewindow(pa, pb, 47)
        @test rmin == fill(1, 4)
        @test rmax == fill(8, 4)
    end

    @testset "DTW and FastDTW agreement" begin
        t = collect(1:1600)
        pktimes = [100, 300, 1000, 1300]
        x = 1 * exp.(-0.5 * ((t .- pktimes[1]) / 100) .^ 2)
        x += 2 * exp.(-0.5 * ((t .- pktimes[2]) / 150) .^ 2)
        x += 3 * exp.(-0.5 * ((t .- pktimes[3]) / 250) .^ 2)
        x += 4 * exp.(-0.5 * ((t .- pktimes[4]) / 250) .^ 2)
        y = x[1:2:end]
        cost, px, py = dtw(x, y)
        cost1, qx, qy = fastdtw(x, y, SqEuclidean(), 15)
        @test isapprox(cost, cost1)
        @test px == qx
        @test py == qy
    end

    @testset "DBA" begin
        x = [1.0, 2.0, 2.0, 3.0, 3.0, 4.0]
        y = [1.0, 3.0, 4.0]
        z = [1.0, 2.0, 2.0, 4.0]
        avg, _ = dba([x, y, z], DTW(5), init_center = z)
        @test avg == [1.0, 1.75, 2.75, 4.0]
    end

    @testset "DTW NN" begin
        @info "Testing DTW NN"
        a = randn(Float32, 100)
        b = randn(Float32, 10000)

        function naive(a, b, r=7)
            dists = map(1:length(b)-length(a)+1) do i
                dtw_cost(a, @view(b[i:i+length(a)-1]), SqEuclidean(), r)
            end
        end

        @inferred dtwnn(a, b, SqEuclidean(), 7)

        res = dtwnn(a, b, SqEuclidean(), 7)
        m = findmin(naive(a, b))
        @test m[1] ≈ res.cost
        @test m[2] == res.loc

        res = dtwnn(a, b, SqEuclidean(), 7, saveall=true)
        resn = naive(a, b)
        m = findmin(resn)
        @test m[1] ≈ res.cost
        @test m[2] == res.loc
        @test res.dists ≈ resn

        a = randn(Float64, 100)
        b = randn(Float64, 10000)

        function naive_norm(a, b)
            an = normalize(ZNormalizer, a)
            dists = map(1:length(b)-length(a)+1) do i
                bn = normalize(ZNormalizer, b[i:i+length(a)-1])
                @test mean(bn) ≈ 0 atol = 10eps(eltype(a))
                @test std(bn, corrected=false, mean=0) ≈ 1 atol = sqrt(eps(eltype(a)))
                dtw_cost(an, bn, SqEuclidean(), 7)
            end
        end

        @inferred dtwnn(a, b, SqEuclidean(), 7, normalizer=Val(ZNormalizer))
        res = dtwnn(a, b, SqEuclidean(), 7, normalizer=Val(ZNormalizer), saveall=true)
        resn = naive_norm(a, b)
        m = findmin(resn)
        @test m[1] ≈ res.cost
        @test m[2] == res.loc
        @test res.dists ≈ resn

        a = sin.(0.1 .* (1:100))    .+ 0.05 .* randn.()
        b = sin.(0.1 .* (1:10_000)) .+ 0.05 .* randn.()

        res = dtwnn(a, b, SqEuclidean(), 2, prune_endpoints = false, prune_envelope = true)
        @test res.prunestats.prune_env > 0
        resn = naive(a, b, 2)
        m = findmin(resn)
        @test m[2] == res.loc
        @test m[1] ≈ res.cost

        res = dtwnn(a, b, SqEuclidean(), 2, prune_endpoints = true, prune_envelope = true)
        @test res.prunestats.prune_env > 0
        @test res.prunestats.prune_end > 0
        resn = naive(a, b, 2)
        m = findmin(resn)
        @test m[2] == res.loc
        @test m[1] ≈ res.cost

        @test SlidingDistancesBase.value(res) == res.cost
        @test SlidingDistancesBase.location(res) == res.loc
        @test SlidingDistancesBase.payload(res) == res.dists
        @test SlidingDistancesBase.target(res) == a

        res = dtwnn(a, b, SqEuclidean(), 2, prune_endpoints = true, prune_envelope = true, saveall=true)
        rr = [res, res]
        @test_nowarn 2res
        @test minimum(rr) == SlidingDistancesBase.value(res)
        @test findmin(rr) == (res.cost, 1)
        @test maximum(rr) == SlidingDistancesBase.value(res)
        @test findmax(rr) == (res.cost, 2)

    end

    @testset "DBA clust" begin
        allsame(x) = all(==(x[1]), x)
        data = [randn(100) .+ 2(i ÷ 5) for i = 0:19]
        nclust = 4
        result = dbaclust(
            data,
            nclust,
            DTW(10);
            n_init = 20,
            iterations = 10,
        )
        inds = 1:5
        @test all([allsame(result.clustids[inds .+ 5i]) for i in 0:3])

        result = [dbaclust(
            data,
            nclust,
            FastDTW(10);
            n_init = 20,
            iterations = 10,
        ) for _ in 1:2]
        inds = 1:5
        @test any(all([allsame(result.clustids[inds .+ 5i]) for i in 0:3]) for result in result)
    end

    @testset "sparse_distmat" begin
        @info "Testing sparse_distmat"
        using Distances
        y = [sin.(0.1 .* (1:100) .+ 2pi*rand()) .+ 0.1randn(100) for _ in 1:200]
        dists,inds = DynamicAxisWarping.sparse_distmat(y, 4, SqEuclidean(), 5)
        D = [dtw_cost(y1,y2,SqEuclidean(), 5) for y1 in y, y2 in y]
        for i = 1:200
            @test partialsort(D[:,i], 2:5) == dists[i]
            @test partialsortperm(D[:,i], 2:5) == inds[i]
        end

    end

    @testset "matrix_profile" begin
        @info "Testing matrix_profile"
        include("test_matrixprofile.jl")
    end

    @testset "distance profile" begin
        @info "Testing distance profile"
        Q = randn(10)
        T = randn(100)
        d = DTW(5)
        D = distance_profile(d, Q, T)
        @test length(D) == 100-10+1
        foreach(i->@test(D[i] ≈ d(Q, getwindow(T,10,i))), 1:91)
    end

    @testset "datasets" begin
        @info "Testing datasets"

        @testset "gaussian" begin
            @info "Testing gaussian"
            data, labels = DynamicAxisWarping.Datasets.fakedata_gaussian()
        end

        @testset "ucr" begin
            @info "Testing ucr"

        end

    end

end

# using DynamicAxisWarping, Distances
# a = randn(1000)
# b = randn(1000)
# @btime dtw($a, $b)
#
# # 10.526 ms (9 allocations: 7.68 MiB)
# # add @avx
# # 2.245 ms (9 allocations: 7.68 MiB)
#
# imi,ima = radiuslimits(10,1000,1000)
# @btime dtw($a, $b, SqEuclidean(), $imi, $ima)
# # 196.412 μs (18 allocations: 193.91 KiB)
#
#
#
# s1,s2 = zeros(21),zeros(21)
# @btime dtw_cost($a, $b, SqEuclidean(), 10, cost=$s1, cost_prev=$s2)
#
# dtw(a, b, SqEuclidean(), imi, ima)
# dtw_cost(a, b, SqEuclidean(), 10)

# using DynamicAxisWarping, Distances
# a = randn(Float32, 1000)
# b = randn(Float32, 100000)
#
# function naive(a,b)
#     dists = map(1:length(b)-length(a)) do i
#         dtw_cost(a,@view(b[i:i+length(a)-1]),SqEuclidean(), 7)
#     end
# end
#
# w = DTWWorkspace(a,SqEuclidean(),7)
# @inferred dtwnn(w,b)
#
# res = dtwnn(w,b)
# m = findmin(naive(a,b))
# @test m[1] ≈ res.cost
# @test m[2] == res.loc
#
#
# @btime naive($a,$b)
# @btime dtwnn($w,$b)
#
#
# a = sin.(0.1 .* (1:100)) .+ 0.1 .* randn.()
# b = sin.(0.1 .* (1:100000)) .+ 0.1 .* randn.()
# w = DTWWorkspace(a,SqEuclidean(),7)
# res = dtwnn(w,b)
#
# plot([a b[eachindex(a) .+ (res.loc-1)]])
#
#
# res = dtwnn(w,b)
# m = findmin(naive(a,b))
# @test m[1] ≈ res.cost
# @test m[2] == res.loc

# a      = sin.((1:1000))     .+ 0.05 .* randn.()
# b      = sin.((1:1000_000)) .+ 0.05 .* randn.()
# @time dtwnn(a, b, SqEuclidean(), 7, normalizer=ZNormalizer, saveall=false)
# @btime dtwnn($a, $b, SqEuclidean(), 5, normalizer=Nothing, saveall=false)
# @btime naive_norm($a, $b)

# # Bench prune_envelope
# a = sin.(0.1 .* (1:100))    .+ 0.1 .* randn.()
# b = sin.(0.1 .* (1:1000_000)) .+ 0.1 .* randn.()
# @btime dtwnn($a, $b, SqEuclidean(), 5, prune_endpoints = false, prune_envelope = false)
# @btime dtwnn($a, $b, SqEuclidean(), 5, prune_endpoints = false, prune_envelope = true)
# @btime dtwnn($a, $b, SqEuclidean(), 5, prune_endpoints = true, prune_envelope = false)
# @btime dtwnn($a, $b, SqEuclidean(), 5, prune_endpoints = true, prune_envelope = true)
#
#
# a = sin.(0.1f0 .* (1:100))    .+ 0.1f0 .* randn.(Float32)
# b = sin.(0.1f0 .* (1:1000_000)) .+ 0.1f0 .* randn.(Float32)
# @btime dtwnn($a, $b, SqEuclidean(), 5, prune_endpoints = true, prune_envelope = true, normalizer=Val(ZNormalizer))


text:
export
  default_nm_loc_ham,
  nm_loc_ham,
  nm_lind,
  make_vertex_set,
  nm_glob_ham,
  nm_init,
  nm_measurement

"""

    default_nm_loc_ham(size)

Return default local Hamiltonian of size `size`×`size` for the demoralization
procedure. The Hamiltonian is sparse with nonzero elements on the first
upper diagonal (equal to `1im`) and lower diagonal (equal to `-1im`).

*Note:* This function is used to provide the default argument for
`nm_loc_ham` function.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> QSWalk.default_nm_loc_ham(4)
4×4 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 6 stored entries:
  [2, 1]  =  0.0-1.0im
  [1, 2]  =  0.0+1.0im
  [3, 2]  =  0.0-1.0im
  [2, 3]  =  0.0+1.0im
  [4, 3]  =  0.0-1.0im
  [3, 4]  =  0.0+1.0im
```
"""
function default_nm_loc_ham(size::Int)
  @argumentcheck size>0 "Size of default local Hamiltonian needs to be positive"
  if size == 1
    return spzeros(ComplexF64, 1, 1)
  else
    spdiagm(1=>im*ones(size-1), -1=>-im*ones(size-1))
  end
end

"""

    nm_loc_ham(vertexset[, hamiltoniansByDegree])

Return Hamiltonian acting locally on each vertex from `vertexset` linear
subspace. `hamiltoniansByDegree` is a dictionary `Dict{Int,
SparseDenseMatrix}`, which, for a given dimension of vertex linear subspace,
yields a hermitian operator. Only matrices for existing dimensions needs to be
specified.

*Note:* Value of `vertexset` should be generated by `make_vertex_set` in order
to match demoralization procedure. Numerical analysis suggests, that
hamiltonians should be complex valued.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> vset = VertexSet([[1, 2], [3, 4]])
VertexSet(Vertex[Vertex([1, 2]), Vertex([3, 4])])

julia> Matrix(nm_loc_ham(vset))
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  0.0+1.0im  0.0+0.0im  0.0+0.0im
 0.0-1.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+1.0im
 0.0+0.0im  0.0+0.0im  0.0-1.0im  0.0+0.0im

julia> A = [1 1im; -1im 1]
2×2 Array{Complex{Int64},2}:
 1+0im  0+1im
 0-1im  1+0im

julia> nm_loc_ham(vset, Dict{Int,Matrix{ComplexF64}}(2  => A))
4×4 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 8 stored entries:
  [1, 1]  =  1.0+0.0im
  [2, 1]  =  0.0-1.0im
  [1, 2]  =  0.0+1.0im
  [2, 2]  =  1.0+0.0im
  [3, 3]  =  1.0+0.0im
  [4, 3]  =  0.0-1.0im
  [3, 4]  =  0.0+1.0im
  [4, 4]  =  1.0+0.0im
```
"""
function nm_loc_ham(vset::VertexSet,
                    hamiltonians::Dict{Int, <:AbstractMatrix{<:Number}}
                        = Dict(l =>default_nm_loc_ham(l) for l = length.(vlist(vset))))
  @argumentcheck all([size(h, 1) == size(h, 2) for h = values(hamiltonians)])  "Hamiltonians must consists of square matrices"
  verticeslengths = length.(vlist(vset))
  @assert all([l in keys(hamiltonians) for l = verticeslengths]) "Missing degree in the Dictionary: $verticeslengths needed"

  hamiltonianlist = Dict{Vertex,AbstractMatrix{<:Number}}(v =>hamiltonians[length(v)] for v = vlist(vset))
  nm_loc_ham(vset, hamiltonianlist)
end

"""

    nm_loc_ham(vertexset[, hamiltoniansByVertex])

Return Hamiltonian acting locally on each vertex from `vertexset` linear
subspace. `hamiltoniansByVertex` is a dictionary
`Dict{Vertex, SparseDenseMatrix}`, which, for a given vertex, yields a hermitian
operator of the size equal to the dimension of the vertex subspace.

*Note:* Value of `vertexset` should be generated by `make_vertex_set` in order
to match demoralization procedure. Numerical analysis suggests, that
hamiltonians should be complex valued.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> vset = VertexSet([[1, 2], [3, 4]])
VertexSet(Vertex[Vertex([1, 2]), Vertex([3, 4])])

julia> Matrix(nm_loc_ham(vset))
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  0.0+1.0im  0.0+0.0im  0.0+0.0im
 0.0-1.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+1.0im
 0.0+0.0im  0.0+0.0im  0.0-1.0im  0.0+0.0im

julia> A, B = [1 1im; -1im 1], [0 1; 1 0]
(Complex{Int64}[1 + 0im 0 + 1im; 0 - 1im 1 + 0im], [0 1; 1 0])

julia> v1, v2 = vlist(vset)
2-element Array{Vertex,1}:
 Vertex([1, 2])
 Vertex([3, 4])

julia> nm_loc_ham(vset, Dict{Vertex,Matrix{Number}}(v1  => A, v2  => B))
4×4 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 6 stored entries:
  [1, 1]  =  1.0+0.0im
  [2, 1]  =  0.0-1.0im
  [1, 2]  =  0.0+1.0im
  [2, 2]  =  1.0+0.0im
  [4, 3]  =  1.0+0.0im
  [3, 4]  =  1.0+0.0im
```
"""
function nm_loc_ham(vset::VertexSet,
                    hamiltonians::Dict{Vertex,<:AbstractMatrix{<:Number}})
  @argumentcheck all([size(h, 1) == size(h, 2) for h = values(hamiltonians)])  "Hamiltonians must consists of square matrices"
  @assert all([v in keys(hamiltonians) for v = vlist(vset)]) "Missing hamiltonian for some vertex"
  @assert all([length(v) == size(hamiltonians[v], 1) for v = vlist(vset)]) "The vertex length and hamiltonian size do no match"

  result = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for v = vlist(vset)
    result[subspace(v), subspace(v)]  =  hamiltonians[v]
  end
  result
end

"""

    nm_lind(A[, lindbladians][, epsilon])

Return single Lindbladian operator and a vertex set describing how vertices are
bound to subspaces. The operator is constructed according to the
corection scheme presented in [1]. Parameter `A` is a square matrix, describing
the connection between the canonical subspaces in a similar manner as the adjacency matrix.
Parameter `epsilon`, with the default value `eps()`, determines the relevant
values by `abs(A[i, j]) >=  epsilon` formula. List `lindbladians` describes the
elementary matrices used. It can be `Dict{Int, SparseDenseMatrix}`,
which returns the matrix by the indegree, or `Dict{Vertex, SparseDenseMatrix}`
which, for different vertices, may return different
matrix. The matrix should have orthogonal columns and be of the size
outdeg of the vertex. As default the function uses Fourier matrices.

*Note:* It is expected that for all pair of vertices there exists a matrix in the
`lindbladians` list.

*Note:* The orthogonality of matrices in `lindbladians` is not verified.

*Note:* The submatrices of the result matrix are multiplied by corresponding `A`
element.

[1] K. Domino, A. Glos, M. Ostaszewski, Superdiffusive quantum stochastic walk
definable on arbitrary directed graph, Quantum Information & Computation,
Vol.17 No.11&12, pp. 0973-0986, arXiv:1701.04624.

# Examples

```jldoctest; setup = :(using QSWalk, LinearAlgebra)
julia> A = [0 1 0; 1 0 1; 0 1 0]
3×3 Array{Int64,2}:
 0  1  0
 1  0  1
 0  1  0

julia> L, vset = nm_lind(A)
(
  [2, 1]  =  1.0+0.0im
  [3, 1]  =  1.0+0.0im
  [1, 2]  =  1.0+0.0im
  [4, 2]  =  1.0+0.0im
  [1, 3]  =  1.0+0.0im
  [4, 3]  =  1.0+0.0im
  [2, 4]  =  1.0+0.0im
  [3, 4]  =  -1.0+1.22465e-16im, VertexSet(Vertex[Vertex([1]), Vertex([2, 3]), Vertex([4])]))

julia> B1, B2 = 2*diagm(0=>[1.]), 3*ones(2, 2)
([2.0], [3.0 3.0; 3.0 3.0])

julia> nm_lind(A, Dict{Int,Matrix{Float64}}(1=>B1, 2=>B2))
(
  [2, 1]  =  3.0+0.0im
  [3, 1]  =  3.0+0.0im
  [1, 2]  =  2.0+0.0im
  [4, 2]  =  2.0+0.0im
  [1, 3]  =  2.0+0.0im
  [4, 3]  =  2.0+0.0im
  [2, 4]  =  3.0+0.0im
  [3, 4]  =  3.0+0.0im, VertexSet(Vertex[Vertex([1]), Vertex([2, 3]), Vertex([4])]))

julia> v1, v2, v3 = vlist(vset)
3-element Array{Vertex,1}:
 Vertex([1])
 Vertex([2, 3])
 Vertex([4])

julia> nm_lind(A, Dict{Vertex,Matrix{Float64}}(v1 => ones(1, 1), v2 => [2 2; 2 -2], v3 => 3*ones(1, 1)))[1] |> Matrix
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  1.0+0.0im  1.0+0.0im   0.0+0.0im
 2.0+0.0im  0.0+0.0im  0.0+0.0im   2.0+0.0im
 2.0+0.0im  0.0+0.0im  0.0+0.0im  -2.0+0.0im
 0.0+0.0im  3.0+0.0im  3.0+0.0im   0.0+0.0im
```
"""
function nm_lind(A::AbstractMatrix{<:Number},
                 lindbladians::Dict{Int, <:AbstractMatrix{<:Number}};
                 epsilon::Real = eps())
  @argumentcheck epsilon>= 0 "epsilon needs to be nonnegative"

  revincidence_list = reversed_incidence_list(A, epsilon = epsilon)
  vset = revinc_to_vertexset(revincidence_list)
  verticeslengths = length.(vlist(vset))

  @assert all([ l in keys(lindbladians) for l = verticeslengths]) "Missing degrees in lindbladians: $verticeslengths needed"

  L = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for i = 1:size(A, 1), (index, j) = enumerate(revincidence_list[i]), k in subspace(vset[j])
      L[subspace(vset[i]), k]  =  A[i, j]*lindbladians[length(vset[i])][:, index]
  end
  L, vset
end,

function nm_lind(A::AbstractMatrix{<:Number};
                 epsilon::Real = eps())
  vset = make_vertex_set(A, epsilon = epsilon)
  degrees = [length(v) for v = vlist(vset)]

  nm_lind(A, Dict(d =>fourier_matrix(d) for d = degrees), epsilon = epsilon)
end,

function nm_lind(A::AbstractMatrix{<:Number},
                 lindbladians::Dict{Vertex, <:AbstractMatrix{<:Number}};
                 epsilon::Real = eps())
  @argumentcheck all([size(l, 1) == size(l, 2) for l = values(lindbladians)])  "lindbladians should consist of square matrix"
  @argumentcheck epsilon>= 0 "epsilon needs to be nonnegative"

  revincidence_list = reversed_incidence_list(A; epsilon = epsilon)
  vset = revinc_to_vertexset(revincidence_list)

  @argumentcheck all([ v in keys(lindbladians) for v = vlist(vset)]) "Some vertex is missing in lindbladians"
  @argumentcheck all([ length(v) == size(lindbladians[v], 1) for v = vlist(vset)]) "Size of the lindbladians should equal to indegree of the vertex"

  L = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for i = 1:size(A, 1), (index, j) = enumerate(revincidence_list[i]), l = subspace(vset[j])
    L[subspace(vset[i]), l]  =  A[i, j]*lindbladians[vset[i]][:, index]
  end
  L, vset
end

"""

    nm_glob_ham(A[, hamiltonians][, weights, epsilon])

Return global Hamiltonian for the moralization procedure. Matrix `A` should the
adjacency matrix of a directed graph, for which one aims to construct the
nonmoralizing dynamics. Here, `hamiltonians` is an optional argument which is a
Dictionary with keys of type `Tuple{Int, Int}` or `Tuple{Vertex, Vertex}`. The
first collects the submatrices according to their shape, while the second
collects them according to each pair of vertices. As the default all-one
submatrices are chosen. The last argument states that only those elements for
which `abs(A[i, j]) >= epsilon` are considered.

*Note:* The submatrices of the result matrix are scaled by corresponding `weights`
argument, which should be a square matrix of the same dimension as `A`. If `weights`
is not provided, then `weights[i,j]=A[i,j]`, if `A[i,j]` is nonzero and `A[j,i]`
is zero, `weights[i,j]=A[j,i]`, if `A[i,j]` in reverse scenario,
`weights[i,j]=(A[i,j]+A[j,i])/2` if both are nonzero, and zero otherwise.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> A = [ 0 1 0; 1 0 1; 0 1 0]
3×3 Array{Int64,2}:
 0  1  0
 1  0  1
 0  1  0

julia> nm_glob_ham(A) |> Matrix
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  1.0+0.0im  1.0+0.0im  0.0+0.0im
 1.0+0.0im  0.0+0.0im  0.0+0.0im  1.0+0.0im
 1.0+0.0im  0.0+0.0im  0.0+0.0im  1.0+0.0im
 0.0+0.0im  1.0+0.0im  1.0+0.0im  0.0+0.0im

julia> dict_deg = Dict{Tuple{Int,Int},Matrix{ComplexF64}}((1, 2) => (2+1im)*ones(1, 2), (2, 1) =>1im*ones(2, 1));

julia> nm_glob_ham(A, dict_deg) |> Matrix
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  2.0+1.0im  2.0+1.0im  0.0+0.0im
 2.0-1.0im  0.0+0.0im  0.0+0.0im  0.0+1.0im
 2.0-1.0im  0.0+0.0im  0.0+0.0im  0.0+1.0im
 0.0+0.0im  0.0-1.0im  0.0-1.0im  0.0+0.0im

julia> v1, v2, v3 = vlist(make_vertex_set(A))
3-element Array{Vertex,1}:
 Vertex([1])
 Vertex([2, 3])
 Vertex([4])

julia> dict_vec = Dict{Tuple{Vertex,Vertex},Matrix{ComplexF64}}((v1, v2) =>2*ones(1, 2), (v2, v3) =>[1im 2im;]');

julia> nm_glob_ham(A, dict_vec) |> Matrix
4×4 Array{Complex{Float64},2}:
 0.0+0.0im  2.0+0.0im  2.0+0.0im  0.0+0.0im
 2.0+0.0im  0.0+0.0im  0.0+0.0im  0.0-1.0im
 2.0+0.0im  0.0+0.0im  0.0+0.0im  0.0-2.0im
 0.0+0.0im  0.0+1.0im  0.0+2.0im  0.0+0.0im
```
"""
function nm_glob_ham(A::AbstractMatrix{<:Number},
                     hamiltonians::Dict{Tuple{Int, Int}, <:AbstractMatrix{<:Number}};
                     epsilon::Real = eps(),
                     weights::AbstractMatrix{<:Number}=default_weights_glob_ham(A, epsilon))
  @argumentcheck epsilon>= 0 "epsilon needs to be nonnegative"

  revincidence_list = reversed_incidence_list(A, epsilon = epsilon)
  vset = revinc_to_vertexset(revincidence_list)

  H = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for (index, i) = enumerate(revincidence_list), j = i
    ind_min, ind_max = sort([index, j])
    hamiltonianshape = length.((subspace(vset[ind_min]), subspace(vset[ind_max])))
    @argumentcheck hamiltonianshape in keys(hamiltonians) "hamiltonian of size $hamiltonianshape not found"
    @argumentcheck hamiltonianshape == size(hamiltonians[hamiltonianshape]) "hamiltonian for key $(hamiltonianshape) shoud have shape $(hamiltonianshape)"
    H[subspace(vset[ind_min]), subspace(vset[ind_max])]  =  weights[ind_min, ind_max]*hamiltonians[hamiltonianshape]
  end
  H + H'
end,

function nm_glob_ham(A::T;
                     epsilon::Real = eps(),
                     weights::AbstractMatrix{<:Number}=default_weights_glob_ham(A, epsilon)) where T<:AbstractMatrix{<:Number}
  #indlist = incidence_list(A, epsilon = epsilon)
  revindlist = reversed_incidence_list(A, epsilon = epsilon)
  alloneshamiltonians = Dict{Tuple{Int, Int}, T}()
  for v = revindlist, w = v
    w = revindlist[w]
    length_w = max(1, length(w))
    length_v = max(1, length(v))
    alloneshamiltonians[(length_v, length_w)]  =  ones(length_v, length_w)
    alloneshamiltonians[(length_w, length_v)]  =  ones(length_w, length_v)
  end
  nm_glob_ham(A, alloneshamiltonians, epsilon = epsilon)
end,

function nm_glob_ham(A::AbstractMatrix{<:Number},
                     hamiltonians::Dict{Tuple{Vertex, Vertex}, <:AbstractMatrix{<:Number}};
                     epsilon::Real = eps(),
                     weights::AbstractMatrix{<:Number}=default_weights_glob_ham(A, epsilon))
  @argumentcheck epsilon >= 0 "epsilon needs to be nonnegative"
  revincidence_list = reversed_incidence_list(A, epsilon = epsilon)
  vset = revinc_to_vertexset(revincidence_list)

  H = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for (index, i) = enumerate(revincidence_list), j = i
      ind_min, ind_max = sort([index, j])
      key = (vset[ind_min], vset[ind_max])
      @argumentcheck key in keys(hamiltonians) "hamiltonian for $key not found"
      @argumentcheck length.(key) == size(hamiltonians[key]) "hamiltonian for key $key shoud have shape $(length.(key))"
      H[subspace(vset[ind_min]), subspace(vset[ind_max])]  =  weights[ind_min, ind_max]*hamiltonians[key]
  end
  H + H'
end

function default_weights_glob_ham(A::AbstractMatrix{<:Number},
                                  epsilon::Real)
    #@argumentcheck epsilon >= 0 "epsilon needs to be nonnegative"
    result = zeros(ComplexF64, size(A)...)
    for i = 1:size(A, 1), j=(i+1):size(A, 2)
      if abs(A[i,j]) > epsilon
        if abs(A[j,i]) > epsilon
          result[i,j]  =  (A[i,j] + A[j,i])/2
        else
          result[i,j]  =  A[i,j]
        end
      elseif abs(A[j,i]) > epsilon
        result[i,j]  =  A[j,i]
      end
    end
    result
end

function default_weights_glob_ham(A::SparseMatrixCSC{<:Number},
                                  epsilon::Real)
    #@argumentcheck epsilon >= 0 "epsilon needs to be nonnegative"
    result = spzeros(ComplexF64, size(A)...)
    for (i, (row_ind1, row_ind2)) = enumerate(zip(A.colptr[1:end-1], A.colptr[2:end]))
      for j = filter(k -> (i < k), A.rowval[row_ind1:row_ind2-1])
        if abs(A[i,j]) > epsilon
          if abs(A[j,i]) > epsilon
            result[i,j]  =  (A[i,j] + A[j,i])/2
          else
            result[i,j]  =  A[i,j]
          end
        elseif abs(A[j,i]) > epsilon
          result[i,j]  =  A[j,i]
        end
      end
    end
    result
end

"""

    nm_measurement(probability, vertexset)

Return joint probability of `probability`, which is real-valued probability vector
according to `vertexset`.

*Note:* It is up to user to provide proper probability vector.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> probability = [0.05, 0.1, 0.25, 0.3, 0.01, 0.20, 0.04, 0.05]
8-element Array{Float64,1}:
 0.05
 0.1
 0.25
 0.3
 0.01
 0.2
 0.04
 0.05

julia> nm_measurement(probability, VertexSet([[1, 4], [2, 3, 5], [6], [7, 8]]))
4-element Array{Float64,1}:
 0.35
 0.36
 0.2
 0.09
```
"""
function nm_measurement(probability::AbstractVector{<:Number},
                        vset::VertexSet)
  @assert vertexsetsize(vset) == length(probability) "vertexset size and probability vector length do not match"

  [sum(probability[subspace(vertex)]) for vertex = vlist(vset)]
end

"""

    nm_measurement(state, vertexset)

Return joint probability of cannonical measurement of density matrix `state`,
according to `vertexset`.

*Note:* It is up to user to provide proper density state.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> state = [1/6 0 1/6; 0 2/3 0; 1/6 0 1/6]
3×3 Array{Float64,2}:
 0.166667  0.0       0.166667
 0.0       0.666667  0.0
 0.166667  0.0       0.166667

julia> nm_measurement(state, VertexSet([[1, 3], [2]]))
2-element Array{Float64,1}:
 0.3333333333333333
 0.6666666666666666
```
"""
function nm_measurement(state::AbstractMatrix{<:Number},
                        vset::VertexSet)
  @argumentcheck size(state, 1) == size(state, 2) "state should be square matrix"
  @assert vertexsetsize(vset) == size(state, 1) "vertexset size and state size do not match"

  nm_measurement(real.(diag(state)), vset)
end

"""
    nm_init(init_vertices, vertexset)

Create initial state in the case of the nonmoralizing evolution based on
`init_vertices` of type `Vector{Vertex}`. The result is
a block diagonal matrix, where each block corresponds to vertex from `vertexset`.
The final state represent an uniform probability over `nm_measurement`.

*Note:* The function returns sparse matrix with `ComplexF64` field type.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> vset = VertexSet([[1], [2, 3, 4], [5, 6, 7], [8, 9]])
VertexSet(Vertex[Vertex([1]), Vertex([2, 3, 4]), Vertex([5, 6, 7]), Vertex([8, 9])])

julia> nm_init(vset[[1, 3, 4]], vset)
9×9 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 6 stored entries:
  [1, 1]  =  0.333333+0.0im
  [5, 5]  =  0.111111+0.0im
  [6, 6]  =  0.111111+0.0im
  [7, 7]  =  0.111111+0.0im
  [8, 8]  =  0.166667+0.0im
  [9, 9]  =  0.166667+0.0im
```
"""
function nm_init(initialvertices::AbstractVector{Vertex},
                 vset::VertexSet)
  @assert all([v in vlist(vset) for v = initialvertices]) "initialvertices is not a subset of vertexset"

  L = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for v = initialvertices
    normalization = length(v)*length(initialvertices)
    L[subspace(v), subspace(v)]  =  SparseMatrixCSC{ComplexF64}(I, length(v), length(v))/normalization
  end
  L
end

"""
    nm_init(init_states, vertexset)

Create initial state in the case of the nonmoralizing evolution based on `init_states`
of type `Dict{Vertex, <:AbstractMatrix{<:Number}}`. For each given vertex a block
from dictionary is used, otherwise zero matrix is chosen. Each matrix from
dictionary should be nonnegative and sum of all traces should equal one. The
keys of `init_vertices` should be a vertices from `vertexset`.
Note that matrix from `init_states` corresponding to vertex `v` should be of
size `length(v)`×`length(v)`.

*Note:* The function returns sparse matrix with `ComplexF64` field type.

# Examples

```jldoctest; setup = :(using QSWalk)
julia> vset = VertexSet([[1], [2, 3, 4], [5, 6, 7], [8, 9]])
VertexSet(Vertex[Vertex([1]), Vertex([2, 3, 4]), Vertex([5, 6, 7]), Vertex([8, 9])])

julia> A1, A2, A3 = ones(ComplexF64, 1, 1)/4, [ 1/5+0im 0 1/5; 0 1/10 0 ; 1/5 0 1/5 ], [0.125 -0.125+0im; -0.125 0.125]
(Complex{Float64}[0.25 + 0.0im], Complex{Float64}[0.2 + 0.0im 0.0 + 0.0im 0.2 + 0.0im; 0.0 + 0.0im 0.1 + 0.0im 0.0 + 0.0im; 0.2 + 0.0im 0.0 + 0.0im 0.2 + 0.0im], Complex{Float64}[0.125 + 0.0im -0.125 + 0.0im; -0.125 + 0.0im 0.125 + 0.0im])

julia> nm_init(Dict(vset[1] =>A1, vset[3] =>A2, vset[4] =>A3), vset)
9×9 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 10 stored entries:
  [1, 1]  =  0.25+0.0im
  [5, 5]  =  0.2+0.0im
  [7, 5]  =  0.2+0.0im
  [6, 6]  =  0.1+0.0im
  [5, 7]  =  0.2+0.0im
  [7, 7]  =  0.2+0.0im
  [8, 8]  =  0.125+0.0im
  [9, 8]  =  -0.125+0.0im
  [8, 9]  =  -0.125+0.0im
  [9, 9]  =  0.125+0.0im

```
"""
function nm_init(initial_states::Dict{Vertex, <:AbstractMatrix{<:Number}},
                 vset::VertexSet)
  @assert all([size(state, 1) == length(k) for (k, state) = initial_states]) "The size of initial state and the vertex do not match"
  @assert all([k in vlist(vset) for k = keys(initial_states)]) "keys of initial_states is not a subset of vertexset"

  L = spzeros(ComplexF64, vertexsetsize(vset), vertexsetsize(vset))
  for v = keys(initial_states)
    L[subspace(v), subspace(v)]  =  initial_states[v]
  end
  L
end

ext: jl


text:
using DashBootstrapComponents

list_group = dbc_listgroup([
    dbc_listgroupitem("Active item", active = true),
    dbc_listgroupitem("Item 2"),
    dbc_listgroupitem("Item 3"),
]);

max_stars_repo_name: glsdown/dash-bootstrap-components

max_issues_repo_name: glsdown/dash-bootstrap-components

max_forks_repo_name: glsdown/dash-bootstrap-components


text:
# serial part of parallel tests

"""
  Test the parallel communication constructs.
"""
function test_parallel_mpi()
  fname = "input_vals_parallel.jl"
  mesh, sbp, eqn, opts = solvePDE(fname)

  # test the Utils parallel functions
  @testset "----- Testing Parallel Functions -----" begin

    test_tagmanager()

    Tsol = eltype(eqn.q)
    mesh.npeers = 1
    nfaces = length(mesh.bndryfaces)
    shared_data = Array{SharedFaceData{Tsol}}(mesh.npeers)
    resize!(mesh.bndries_local, mesh.npeers)
    resize!(mesh.bndries_remote, mesh.npeers)
    resize!(mesh.shared_interfaces, mesh.npeers)
    resize!(mesh.peer_parts, mesh.npeers)
    resize!(mesh.peer_face_counts, mesh.npeers)
    resize!(mesh.local_element_counts, mesh.npeers)
    resize!(mesh.remote_element_counts, mesh.npeers)
    fill!(mesh.peer_parts, mesh.myrank)
    fill!(mesh.peer_face_counts, nfaces)
    fill!(mesh.local_element_counts, nfaces)
    fill!(mesh.remote_element_counts, nfaces)
    for i=1:mesh.npeers
      println("creating SharedFaceData for peer ", i, ", Tsol = ", Tsol)
      mesh.bndries_local[i] = mesh.bndryfaces
      mesh.bndries_remote[i] = mesh.bndryfaces
      mesh.shared_interfaces[i] = Array{Interface}(0)
      shared_data[i] = SharedFaceData(Tsol, mesh, i, opts["parallel_data"], 100)
    end

    for i=1:mesh.npeers
      @test ( shared_data[i].send_req )== MPI.REQUEST_NULL
      @test ( shared_data[i].recv_req )== MPI.REQUEST_NULL
    end

    data_i = shared_data[1]
    Utils.getSendDataFace(mesh, sbp, eqn, opts, data_i)
    buff = data_i.q_send

    # verify that the inteperpolation was exact
    eqn.params.alpha_x = 1
    eqn.params.alpha_y = 1
    for i=1:length(mesh.bndryfaces)
      for j=1:mesh.numNodesPerFace
        coords = mesh.coords_bndry[:, j, i]
        val_exp = AdvectionEquationMod.calc_p4(eqn.params, coords, 0.0)
        @test isapprox( buff[1, j, i], val_exp) atol=1e-13
      end
    end

    rand!(data_i.q_send)
    data_i.send_req = MPI.Isend(data_i.q_send, data_i.myrank, 1, data_i.comm)
    data_i.recv_req = MPI.Irecv!(data_i.q_recv, data_i.myrank, 1, data_i.comm)

    data_i.recv_status = MPI.Wait!(data_i.recv_req)
    # if the assertiosn do not trigger, then the test passes
    verifyReceiveCommunication(data_i)
    @test isapprox( data_i.q_send, data_i.q_recv) atol=1e-13

    # test calcSharedFaceIntegrals
    mesh.npeers = 1
    mesh.peer_face_counts = [2]

    # do a dummy send
    a = [1]
    ar = [1]
    data_i.send_req = MPI.Isend(a, mesh.myrank, 1, mesh.comm)
    data_i.recv_req = MPI.Irecv!(ar, mesh.myrank, 1, mesh.comm)
    data_i.q_send = Array{Float64}(mesh.numDofPerNode, mesh.numNodesPerFace, 2)
    data_i.q_recv = Array{Float64}(mesh.numDofPerNode, mesh.numNodesPerFace, 2)
    mesh.nrm_sharedface = Array{Array{Float64, 3}}(1)
    mesh.nrm_sharedface[1] = zeros(2, mesh.numNodesPerFace, 2)
    nrm_arr = mesh.nrm_sharedface[1]

    interiorfaceinterpolate!(mesh.sbpface, mesh.interfaces, eqn.q, eqn.q_face)
    for i=1:mesh.peer_face_counts[1]
      iface = mesh.interfaces[i]
      for k=1:mesh.numNodesPerFace
        nrm_arr[:, k, i] = mesh.nrm_face[ :, k, i]
        for j=1:mesh.numDofPerNode
          data_i.q_send[j, k, i] = eqn.q_face[j, 1, k, i]
          data_i.q_recv[j, k, i] = eqn.q_face[j, 2, k, i]
        end
      end
    end

    # get the face flux using calcFaceFlux
    ifaces_orig = mesh.interfaces
    mesh.interfaces = ifaces_orig[1:2]
    fill!(eqn.res, 0.0)
    eqn.flux_face = zeros(mesh.numDofPerNode, mesh.numNodesPerFace, 2)
    AdvectionEquationMod.calcFaceFlux(mesh, sbp, eqn, eqn.flux_func, mesh.interfaces, eqn.flux_face)

    # now get it using Boundary integrate
    eqn.flux_sharedface = Array{Array{Float64, 3}}(1)
    eqn.flux_sharedface[1] = Array{Float64}(mesh.numDofPerNode, mesh.numNodesPerFace, 2)

    data_i.interfaces = ifaces_orig[1:2]
    # create the boundary array
    data_i.bndries_local = Array{Boundary}(2)
    for i=1:2
      data_i.bndries_local[i] = Boundary(ifaces_orig[i].elementL, ifaces_orig[i].faceL)
    end
    # use the identity permutation because the values were already permuted
    for i=1:mesh.numNodesPerFace
      mesh.sbpface.nbrperm[i, 1] = i
    end
    AdvectionEquationMod.calcSharedFaceIntegrals(mesh, sbp, eqn, opts, data_i)

    for i=1:length(eqn.flux_sharedface[1])
      @test isapprox( eqn.flux_sharedface[1][i], eqn.flux_face[i]) atol=1e-13
    end

    # test changing buffers
    @test getParallelData(shared_data) == PARALLEL_DATA_FACE
    @test size(data_i.q_send, 1) == mesh.numDofPerNode
    @test size(data_i.q_send, 2) == mesh.numNodesPerFace
    @test size(data_i.q_send, 3) == mesh.peer_face_counts[1]

    @test size(data_i.q_recv, 1) == mesh.numDofPerNode
    @test size(data_i.q_recv, 2) == mesh.numNodesPerFace
    @test size(data_i.q_recv, 3) == mesh.peer_face_counts[1]

    setParallelData(shared_data, PARALLEL_DATA_ELEMENT)
    @test getParallelData(shared_data) == PARALLEL_DATA_ELEMENT
    @test size(data_i.q_send, 1) == mesh.numDofPerNode
    @test size(data_i.q_send, 2) == mesh.numNodesPerElement
    @test size(data_i.q_send, 3) == mesh.local_element_counts[1]

    @test size(data_i.q_recv, 1) == mesh.numDofPerNode
    @test size(data_i.q_recv, 2) == mesh.numNodesPerElement
    @test size(data_i.q_recv, 3) == mesh.remote_element_counts[1]

    # test copy
    data_i2 = copy(data_i)
    data_i2.q_send[1] = 10
    data_i2.q_recv[1] = 10
    @test data_i2.q_send[1] != data_i.q_send[1]
    @test data_i2.q_recv[1] != data_i.q_recv[1]
    @test pointer(data_i2.q_send) != pointer(data_i.q_send)
    @test pointer(data_i2.q_recv) != pointer(data_i.q_recv)

    copy!(data_i2, data_i)
    @test data_i2.q_send[1] == data_i.q_send[1]
    @test data_i2.q_recv[1] == data_i.q_recv[1]

    old_tag = shared_data[1].tag
    setNewTag(shared_data)
    @test old_tag != shared_data[1].tag

  end  # end facts block

  return nothing
end

function test_tagmanager()
  @testset "MPITagManager" begin
    mgr = MPITagManager()

    for i=1:10
      @test getNextTag(mgr) == i
    end

    # test using tags
    markTagUsed(mgr, 12)
    @test getNextTag(mgr) == 11
    @test getNextTag(mgr) == 13

    freeTag(mgr, 3)
    freeTag(mgr, 4)
    @test getNextTag(mgr) == 3
    @test getNextTag(mgr) == 4
    @test getNextTag(mgr) == 14

    freeTag(mgr, 5)
    freeTag(mgr, 4)
    freeTag(mgr, 6)
    @test getNextTag(mgr) == 4
    @test getNextTag(mgr) == 5
    @test getNextTag(mgr) == 6

    freeTag(mgr, 3)
    @test_throws Exception freeTag(mgr, 3)

    @test_throws Exception markTagUsed(mgr, 4)

    # test starting tag = 5
    mgr = MPITagManager(5)

    for i=1:10
      @test getNextTag(mgr) == i + 5 - 1
    end

    markTagUsed(mgr, 16)
    @test getNextTag(mgr) == 15
    @test getNextTag(mgr) == 17

    freeTag(mgr, 9)
    freeTag(mgr, 7)
    freeTag(mgr, 8)
    @test getNextTag(mgr) == 7
    @test getNextTag(mgr) == 8
    @test getNextTag(mgr) == 9
    @test getNextTag(mgr) == 18

    @test_throws Exception freeTag(mgr, 2)
    @test_throws Exception markTagUsed(mgr, 2)

    freeTag(mgr, 7)
    @test_throws Exception freeTag(mgr, 7)

  end

  return nothing
end

#test_parallel_mpi()
add_func1!(AdvectionTests, test_parallel_mpi, [TAG_SHORTTEST])

function test_parallel_serialpart()
  fname = "input_vals_parallel.jl"
  mesh, sbp, eqn, opts = solvePDE(fname)

  # do a serial rk4 run to compare against later
  start_dir = pwd()
  cd("./rk4/serial")
  opts["order"] = 1
  opts["solve"] = true
  fname = "input_vals_parallel_run"
  make_input(opts, fname)

  fname2 = string(fname, ".jl")
  mesh, sbp, eqn, opts = solvePDE(fname2)
  cd(start_dir)

  # make the parallel version
  opts["smb_name"] = "SRCMESHES/psquare2.smb"
  opts["dmg_name"] = "SRCMESHES/psquare2.dmg"
  make_input( opts, string("./rk4/parallel/", fname, "p"))

  # serial newton
  start_dir = pwd()
  cd("./newton/serial")
  fname = "input_vals_serial.jl"
  mesh, sbp, eqn, opts = solvePDE(fname)
  cd(start_dir)
  opts["smb_name"] = "SRCMESHES/psquare2.smb"
  opts["dmg_name"] = "SRCMESHES/psquare2.dmg"
  make_input(opts, string("./newton/parallel/", "input_vals_parallel"))

  # 3D rk4
  start_dir = pwd()
  cd("./rk4_3d/serial")
  fname = "input_vals_rk4_3d.jl"
  mesh, sbp, eqn, opts = solvePDE(fname)

  # make the parallel version
  cd(start_dir)
  opts["smb_name"] = "SRCMESHES/ptet8cube.smb"
  make_input(opts, string("./rk4_3d/parallel/", "input_vals_parallel"))

  # 3D Newton's method
  opts["run_type"] = 5
  opts["jac_method"] = 2
  opts["jac_type"] = 3
  opts["parallel_type"] = 2
  opts["parallel_data"] = PARALLEL_DATA_ELEMENT
  make_input(opts, string("./newton_3d/parallel/", "input_vals_parallel"))

  opts["smb_name"] = "SRCMESHES/tet8cube.smb"
  make_input(opts, string("./newton_3d/serial/", "input_vals_serial"))

  cd("./newton_3d/serial")
  fname = "input_vals_serial.jl"
  mesh, sbp, eqn, opts = solvePDE(fname)

  cd(start_dir)
end

#test_parallel_serialpart()
add_func1!(AdvectionTests, test_parallel_serialpart, [TAG_SHORTTEST])


text:
#!/usr/bin/env julia

# Copyright 2022 John T. Foster
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
using Test, MPI

@info "Running MPI tests"

testdir = @__DIR__
istest(f) = endswith(f, ".jl") && startswith(f, "test_")
testfile = sort(filter(istest, readdir(testdir)))

for f in testfile
    @testset "$f: Test on $nprocs process(es)" for nprocs in [2, 4]
        mpiexec() do cmd
            run(`$cmd -n $nprocs $(Base.julia_cmd()) $(joinpath(testdir, f))`)
            @test true
        end
    end
end


text:
@testset "Config" begin
    Impl.@polyvar x y z

    f = Polynomial{Float64}(x^2+3y+z*x+y^2+12*x^7*z^5)
    g = Polynomial{Float64}(4*y*x^2-3z+z*x^4+y+y^2-y*x)

    ∇f = [differentiate(f, i) for i=1:3]
    ∇g = [differentiate(g, i) for i=1:3]

    w = rand(3)
    u = zeros(3)

    cfg = config(f, w)
    @test cfg isa GradientConfig
    @test f(w) ≈ evaluate(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient!(u, f, w, cfg)
    @test [p(w) for p in ∇f] ≈ u

    @test_throws BoundsError FixedPolynomials.gradient(f, rand(2), cfg)
    @test_throws BoundsError FixedPolynomials.gradient!(u, f, rand(2), cfg)
    @test_throws BoundsError FixedPolynomials.gradient!(u[1:2], f, w, cfg)

    wc = rand(ComplexF64, 3)
    uc = zeros(ComplexF64, 3)
    cfg = config(f, wc)
    @test cfg isa GradientConfig
    @test f(wc) ≈ evaluate(f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ FixedPolynomials.gradient(f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ FixedPolynomials.gradient!(uc, f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ uc

    cfg = config(f, w)
    @test f(w) ≈ evaluate(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(f, w, cfg, true)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(f, w, cfg)
    @test f(w) ≈ evaluate(f, w, cfg, true)

    u = zeros(3)
    @test f(w) ≈ evaluate(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient!(u, f, w, cfg, true)

    r = GradientDiffResult(cfg)

    FixedPolynomials.gradient!(r, f, w, cfg)
    @test f(w) ≈ value(r)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(r)

    v = zeros(4)
    r = GradientDiffResult((@view v[1:3]))
    FixedPolynomials.gradient!(r, f, w, cfg)
    @test f(w) ≈ value(r)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(r)

    @test_throws MethodError GradientConfig([f, f], w)
end


text:
# Add tests to the quote for functions with  FastMath varients.
function jacobian_via_frule(f,z)
    du_dx, dv_dx = reim(frule((Zero(), 1),f,z)[2])
    du_dy, dv_dy = reim(frule((Zero(),im),f,z)[2])
    return [
        du_dx  du_dy
        dv_dx  dv_dy
    ]
end
function jacobian_via_rrule(f,z)
    _, pullback = rrule(f,z)
    du_dx, du_dy = reim(pullback( 1)[2])
    dv_dx, dv_dy = reim(pullback(im)[2])
    return [
        du_dx  du_dy
        dv_dx  dv_dy
    ]
end

function jacobian_via_fdm(f, z::Union{Real, Complex})
    fR2((x, y)) = (collect ∘ reim ∘ f)(x + im*y)
    v = float([real(z)
               imag(z)])
    j = jacobian(central_fdm(5,1), fR2, v)[1]
    if size(j) == (2,2)
        j
    elseif size(j) == (1, 2)
        [j
         false false]
    else
        error("Invalid Jacobian size $(size(j))")
    end
end

function complex_jacobian_test(f, z)
    @test jacobian_via_fdm(f, z) ≈ jacobian_via_frule(f, z)
    @test jacobian_via_fdm(f, z) ≈ jacobian_via_rrule(f, z)
end

const FASTABLE_AST = quote
    is_fastmath_mode = sin === Base.FastMath.sin_fast

    @testset "Trig" begin
        @testset "Basics" for x = (Float64(π)-0.01, Complex(π, π/2))
            test_scalar(sin, x)
            test_scalar(cos, x)
            test_scalar(tan, x)
        end
        @testset "Hyperbolic" for x = (Float64(π)-0.01, Complex(π-0.01, π/2))
            test_scalar(sinh, x)
            test_scalar(cosh, x)
            test_scalar(tanh, x)
        end
        @testset "Inverses" for x = (0.5, Complex(0.5, 0.25))
            test_scalar(asin, x)
            test_scalar(acos, x)
            test_scalar(atan, x)
        end
        @testset "Multivariate" begin
            @testset "sincos(x::$T)" for T in (Float64, ComplexF64)
                x, Δx, x̄ = randn(T, 3)
                Δz = (randn(T), randn(T))

                frule_test(sincos, (x, Δx))
                rrule_test(sincos, Δz, (x, x̄))
            end
        end
    end

    @testset "exponents" begin
        for x in (-0.1, 7.9, 0.5 + 0.25im)
            test_scalar(inv, x)

            test_scalar(exp, x)
            test_scalar(exp2, x)
            test_scalar(exp10, x)
            test_scalar(expm1, x)

            if x isa Real
                test_scalar(cbrt, x)
            end

            if x isa Complex || x >= 0
                test_scalar(sqrt, x)
                test_scalar(log, x)
                test_scalar(log2, x)
                test_scalar(log10, x)
                test_scalar(log1p, x)
            end
        end
    end

    @testset "rounding" begin
        for x in (-0.6, -0.2, 0.1, 0.6)
            # thanks to RoundNearest
            if 0 > x % 1 > 0.5 || -1 < x % 1 <= 0.5
                test_scalar(round, x; fdm=backward_fdm(5,1))
            else
                test_scalar(round, x; fdm=forward_fdm(5,1))
            end
            test_scalar(floor, x; fdm=backward_fdm(5, 1))
            test_scalar(ceil, x; fdm=forward_fdm(5, 1))
        end
    end

    @testset "Unary complex functions" begin
        for f ∈ (abs, abs2, conj), z ∈ (-4.1-0.02im, 6.4, 3 + im)
            @testset "Unary complex functions f = $f, z = $z" begin
                complex_jacobian_test(f, z)
            end
        end
        # As per PR #196, angle gives a Zero() pullback for Real z and ΔΩ, rather than
        # the one you'd get from considering the reals as embedded in the complex plane
        # so we need to special case it's tests
        for z ∈ (-4.1-0.02im, 6.4 + 0im, 3 + im)
            complex_jacobian_test(angle, z)
        end
        @test frule((Zero(), randn()), angle, randn())[2] === Zero()
        @test rrule(angle, randn())[2](randn())[2]        === Zero()

        # test that real primal with complex tangent gives complex tangent
        ΔΩ = randn(ComplexF64)
        for x in (-0.5, 2.0)
            @test isapprox(
                frule((Zero(), ΔΩ), angle, x)[2],
                frule((Zero(), ΔΩ), angle, complex(x))[2],
            )
        end
    end

    @testset "Unary functions" begin
        for x in (-4.1, 6.4, 0.0, 0.0 + 0.0im, 0.5 + 0.25im)
            test_scalar(+, x)
            test_scalar(-, x)
            test_scalar(atan, x; rtol=(is_fastmath_mode ? 1e-7 : 1e-9))
        end
    end

    @testset "binary functions" begin
        @testset "$f(x, y)" for f in (atan, rem, max, min)
            x, Δx, x̄ = 100rand(3)
            y, Δy, ȳ = 10rand(3)
            Δz = rand()

            frule_test(f, (x, Δx), (y, Δy))
            rrule_test(f, Δz, (x, x̄), (y, ȳ))
        end

        @testset "$f(x::$T, y::$T)" for f in (/, +, -, hypot), T in (Float64, ComplexF64)
            x, Δx, x̄ = 10rand(T, 3)
            y, Δy, ȳ = rand(T, 3)
            Δz = randn(typeof(f(x, y)))

            # some tests struggle in fast_math mode to get accurasy so we relax it some.
            frule_test(f, (x, Δx), (y, Δy); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
            rrule_test(f, Δz, (x, x̄), (y, ȳ); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
        end

        @testset "$f(x::$T, y::$T) type check" for f in (/, +, -,\, hypot, ^), T in (Float32, Float64)
            x, Δx, x̄ = 10rand(T, 3)
            y, Δy, ȳ = rand(T, 3)
            @assert T == typeof(f(x, y))
            Δz = randn(typeof(f(x, y)))

            @test frule((Zero(), Δx, Δy), f, x, y) isa Tuple{T, T}
            _, ∂x, ∂y = rrule(f, x, y)[2](Δz)
            @test extern.((∂x, ∂y)) isa Tuple{T, T}

            if f != hypot
                # Issue #233
                @test frule((Zero(), Δx, Δy), f, x, 2) isa Tuple{T, T}
                _, ∂x, ∂y = rrule(f, x, 2)[2](Δz)
                @test extern.((∂x, ∂y)) isa Tuple{T, T}

                @test frule((Zero(), Δx, Δy), f, 2, y) isa Tuple{T, T}
                _, ∂x, ∂y = rrule(f, 2, y)[2](Δz)
                @test extern.((∂x, ∂y)) isa Tuple{T, T}
            end
        end

        @testset "^(x::$T, n::$T)" for T in (Float64, ComplexF64)
            # for real x and n, x must be >0
            x, Δx, x̄ = rand(T, 3) .+ 3
            y, Δy, ȳ = rand(T, 3) .+ 3
            Δz = rand(T)

            frule_test(^, (x, Δx), (y, Δy))
            rrule_test(^, Δz, (x, x̄), (y, ȳ))

            T <: Real && @testset "discontinuity for ^(x::Real, n::Int) when x ≤ 0" begin
                # finite differences doesn't work for x < 0, so we check manually
                x = -rand(T) .- 3
                y = 3
                Δx = randn(T)
                Δy = randn(T)
                Δz = randn(T)

                @test frule((Zero(), Δx, Δy), ^, x, y)[2] ≈ Δx * y * x^(y - 1)
                @test frule((Zero(), Δx, Δy), ^, zero(x), y)[2] ≈ 0
                _, ∂x, ∂y = rrule(^, x, y)[2](Δz)
                @test ∂x ≈ Δz * y * x^(y - 1)
                @test ∂y ≈ 0
                _, ∂x, ∂y = rrule(^, zero(x), y)[2](Δz)
                @test ∂x ≈ 0
                @test ∂y ≈ 0
            end
        end
    end

    @testset "sign" begin
        @testset "real" begin
            @testset "at $x" for x in (-1.1, -1.1, 0.5, 100.0)
                test_scalar(sign, x)
            end

            @testset "Zero over the point discontinuity" begin
                # Can't do finite differencing because we are lying
                # following the subgradient convention.

                _, pb = rrule(sign, 0.0)
                _, x̄ = pb(10.5)
                @test extern(x̄) == 0

                _, ẏ = frule((Zero(), 10.5), sign, 0.0)
                @test extern(ẏ) == 0
            end
        end
        @testset "complex" begin
            @testset "at $z" for z in (-1.1 + randn() * im, 0.5 + randn() * im)
                test_scalar(sign, z)

                # test that complex (co)tangents with real primal gives same result as
                # complex primal with zero imaginary part

                ż, ΔΩ = randn(ComplexF64, 2)
                Ω, ∂Ω = frule((Zero(), ż), sign, real(z))
                @test Ω == sign(real(z))
                @test ∂Ω ≈ frule((Zero(), ż), sign, real(z) + 0im)[2]

                Ω, pb = rrule(sign, real(z))
                @test Ω == sign(real(z))
                @test pb(ΔΩ)[2] ≈ rrule(sign, real(z) + 0im)[2](ΔΩ)[2]
            end

            @testset "zero over the point discontinuity" begin
                # Can't do finite differencing because we are lying
                # following the subgradient convention.

                _, pb = rrule(sign, 0.0 + 0.0im)
                _, z̄ = pb(randn(ComplexF64))
                @test extern(z̄) == 0.0 + 0.0im

                _, Ω̇ = frule((Zero(), randn(ComplexF64)), sign, 0.0 + 0.0im)
                @test extern(Ω̇) == 0.0 + 0.0im
            end
        end
    end
end

# Now we generate tests for fast and nonfast versions
@eval @testset "fastmath_able Base functions" begin
    $FASTABLE_AST
end

@eval @testset "fastmath_able FastMath functions" begin
    $(Base.FastMath.make_fastmath(FASTABLE_AST))
end


text:
using PlanerTransformer
using Dierckx

in_3f4 =
 (25e3,(82,19),(210,360),
  100e3,(48,19),(170,1070),
  200e3,(35,19),(150,1300),
  400e3,(25,16),(140,1700),
  1e6,(15,20),(78,1900),
  2e6,(5.5,20),(40,1850),
  3e6,(4,40),(26,2000))

# x=log10(B), y=log10(frequency), z=log10(P)

x = log10.([82,   210,  48,    170,   35,    150,   25,    140,   15,  78,   5.5, 20,   4,   26])
y = log10.([25e3, 25e3, 100e3, 100e3, 200e3, 200e3, 400e3, 400e3, 1e6, 1e6,  2e6, 2e6,  3e6, 3e6])
z = log10.([19,   360,  19,    1070,  19,    1300,  16,    1700,  20,  1900, 20,  1850, 40,  2000])
Spline2D(x, y, z; kx=3, ky=3, s=0.0)


text:
using Documenter, WaterLily

makedocs(
    modules = [WaterLily],
    repo="https://github.com/gabrielweymouth/WaterLily.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", nothing) == "true",
        canonical="https:/Zitzeronion.github.io/WaterLily.jl",
        assets=String[],
    ),
    authors = "Gabriel Weymouth",
    sitename = "WaterLily.jl",
    pages = ["index.md"]
    # strict = true,
    # clean = true,
    # checkdocs = :exports,
)

deploydocs(
    repo = "github.com/Zitzeronion/WaterLily.jl.git",
    push_preview = true
)


text:
import Boltzmann: split_evenly

if :unit in TEST_GROUPS
    @test split_evenly(10, 2) == [(1,2), (3,4), (5,6), (7,8), (9,10)]
    @test split_evenly(11, 2) == [(1,2), (3,4), (5,6), (7,8), (9,10), (11,11)]
    @test split_evenly(3, 5) == [(1,3)]
end


text:
using FractionalDiffEq, Plots
h=0.5; t=collect(0:h:1)
f(t) = 0 ≤ t ≤ 1 ? (return 8) : (return 0)
prob=DODEProblem([1, 1, 1], [2, x->6*x*(1-x), 0], (0, 1), f, [0; 0], t)
sol=solve(prob, h, DOMatrixDiscrete())
plot(sol)


text:
module NonlinearEquations

import Calculus
import MacroTools
import SparseArrays

function codegen_addterm_residuals(equationnum, term)
	return quote
		residuals[$equationnum] += $term
	end
end

function codegen_addterm_inplacejacobian(equationnum, term, xsym)
	if MacroTools.inexpr(term, xsym)
		derivatives, refs = NonlinearEquations.differentiatewithrefs(term, xsym)
		for ref in refs
			if length(ref) != 1
				error("must be a reference with a single index")
			end
		end
		refs = map(ref->replaceall(ref[1], :end, :(length($xsym))), refs)
		q = quote end
		for (derivative, ref) in zip(derivatives, refs)
			newcode = quote
				___jacobian_storage___[$equationnum, $ref] += $derivative
			end
			append!(q.args, newcode.args)
		end
		return q
	else
		return :()
	end
end

function codegen_addterm_jacobian(equationnum, term, xsym)
	if MacroTools.inexpr(term, xsym)
		derivatives, refs = NonlinearEquations.differentiatewithrefs(term, xsym)
		for ref in refs
			if length(ref) != 1
				error("must be a reference with a single index")
			end
		end
		refs = map(ref->replaceall(ref[1], :end, :(length($xsym))), refs)
		q = quote end
		for (derivative, ref) in zip(derivatives, refs)
			newcode = quote
				push!(I, $equationnum)
				push!(J, $(ref))
				push!(V, $(derivative))
			end
			append!(q.args, newcode.args)
		end
		return q
	else
		return :()
	end
end

function differentiatewithrefs(exorig, x::Symbol)
	ex, dict = replacerefswithsyms(exorig)
	diffsyms = Symbol[]
	diffrefs = Any[]
	for (junksym, (sym, ref)) in dict
		if sym == x
			push!(diffsyms, junksym)
			push!(diffrefs, ref)
		end
	end
	if typeof(ex) == Symbol#this hack gets around the fact that Calculus.differentiate doesn't have a function differentiate(::Symbol, ::Array{Symbol,1})
		diffs = Any[]
		for i = 1:length(diffsyms)
			if diffsyms[i] == ex
				push!(diffs, :(1))
			else
				push!(diffs, :(0))
			end
		end
	else
		diffs = Calculus.differentiate(ex, diffsyms)
	end
	#diffs = map(diff->Calculus.simplify(replacesymswithrefs(diff, dict)), diffs)
	diffs = map(diff->replacesymswithrefs(diff, dict), diffs)#TODO: get the simplify stuff working again. currently it breaks some precompilation thing. maybe implement a simplify function without using calculus or use sympy.jl or something. hopefully the compiler does some of the optimizations for now
	return diffs, diffrefs
end

macro equations(other, fundef)
	if @MacroTools.capture(other, expandmodule=modulename_)
		return equations(fundef, eval(modulename), Symbol[])
	elseif @MacroTools.capture(other, exclude=(varnames_))
		return equations(fundef, Main, map(x->x, varnames.args))
	end
	return equations(fundef, Main, Symbol[])
end

macro equations(fundef)
	return equations(fundef, Main, Symbol[])
end

function equations(fundef::Expr, macroexpand_module, dont_differentiate_syms::Array{Symbol, 1})
	dict = MacroTools.splitdef(fundef)
	original_body = macroexpand(macroexpand_module, dict[:body])
	#generate the code for computing the residuals
	body_residuals = MacroTools.postwalk(x->replacenumequations(x, :(residuals = zeros(numequations))), original_body)
	body_residuals = MacroTools.postwalk(x->replaceaddterm(x, codegen_addterm_residuals), body_residuals)
	original_name = dict[:name]
	dict[:name] = Symbol(original_name, :_residuals)
	dict[:body] = quote
		$body_residuals
		return residuals
	end
	q_residuals = MacroTools.combinedef(dict)
	q_result = quote
		$(esc(q_residuals))
	end
	#generate the code for the jacobian
	for arg in filter(x->!(x in dont_differentiate_syms), dict[:args])
		arg_name = MacroTools.splitarg(arg)[1]
		body_jacobian = MacroTools.postwalk(x->replacenumequations(x, :()), original_body)
		body_jacobian = MacroTools.postwalk(x->replaceaddterm(x, (eqnum, term)->codegen_addterm_jacobian(eqnum, term, arg_name)), body_jacobian)
		dict[:name] = Symbol(original_name, :_, arg_name)
		dict[:body] = quote
			I = Int[]
			J = Int[]
			V = eltype($arg)[]
			$body_jacobian
			return SparseArrays.sparse(I, J, V, numequations, length($arg_name), +)
		end
		push!(q_result.args, :($(esc(MacroTools.combinedef(dict)))))
	end
	#generate the in-place versions of the jacobian
	pushfirst!(dict[:args], :___jacobian_storage___)
	for arg in filter(x->!(x in dont_differentiate_syms), dict[:args])
		arg_name = MacroTools.splitarg(arg)[1]
		body_inplacejacobian = MacroTools.postwalk(x->replacenumequations(x, :()), original_body)
		body_inplacejacobian = MacroTools.postwalk(x->replaceaddterm(x, (eqnum, term)->codegen_addterm_inplacejacobian(eqnum, term, arg_name)), body_inplacejacobian)
		dict[:name] = Symbol(original_name, :_, arg_name, :!)
		dict[:body] = quote
			fill!(SparseArrays.nonzeros(___jacobian_storage___), zero(eltype($arg)))
			$body_inplacejacobian
			return nothing
		end
		push!(q_result.args, :($(esc(MacroTools.combinedef(dict)))))
	end
	#display(MacroTools.prettify(q_result))
	return q_result
end

function escapesymbols(expr, symbols)
	for symbol in symbols
		expr = NonlinearEquations.replaceall(expr, symbol, Expr(:escape, symbol))
	end
	return expr
end

function gradient(x, p, g_x, g_p, f_x, f_p)
	lambda = transpose(f_x(x, p)) \ g_x(x, p)
	return g_p(x, p) - transpose(f_p(x, p)) * lambda
end

function newtonish(residuals, jacobian, x0; numiters=10, solver=(J, r)->J \ r, rate=0.05, callback=(x, r, J, i)->nothing)
	x = x0
	for i = 1:numiters
		J = jacobian(x)
		r = residuals(x)
		#x = (1 - rate) * x - rate * solver(J, r)
		x = x - rate * solver(J, r)
		callback(x, r, J, i)
	end
	return x
end

function newton(residuals, jacobian, x0; numiters=10, solver=(J, r)->J \ r, callback=(x, r, J, i)->nothing)
	x = x0
	for i = 1:numiters
		J = jacobian(x)
		r = residuals(x)
		x = x - solver(J, r)
		callback(x, r, J, i)
	end
	return x
end

function replace(expr, old, new)
	@MacroTools.capture(expr, $old) || return expr
	return new
end

function replaceall(expr, old, new)
	return MacroTools.postwalk(x->replace(x, old, new), expr)
end

function replacerefswithsyms(expr)
	sym2symandref = Dict()
	function replaceref(expr)
		@MacroTools.capture(expr, x_[y__]) || return expr
		sym = gensym()
		sym2symandref[sym] = (x, y)
		return sym
	end
	newexpr = MacroTools.prewalk(replaceref, expr)
	return newexpr, sym2symandref
end

function replaceaddterm(x, codegen)
	@MacroTools.capture(x, NonlinearEquations.addterm(equationnum_, term_)) || return x
	return codegen(equationnum, term)
end

function replacenumequations(x, postcode)
	@MacroTools.capture(x, NonlinearEquations.setnumequations(numequations_)) || return x
	return quote
		numequations = $numequations
		$postcode
	end
end

function replacesymswithrefs(expr, sym2symandref)
	function replacesym(expr)
		if expr isa Symbol && haskey(sym2symandref, expr)
			sym, ref = sym2symandref[expr]
			return :($sym[$(ref...)])
		else
			return expr
		end
	end
	return MacroTools.postwalk(replacesym, expr)
end

function updateentries!(dest::SparseArrays.SparseMatrixCSC, src::SparseArrays.SparseMatrixCSC)
	if dest.colptr != src.colptr || dest.rowval != src.rowval
		error("Cannot update entries unless the two matrices have the same pattern of nonzeros")
	end
	copy!(dest.nzval, src.nzval)
end

end


text:
function count end

macro count(source, f)
    q = Expr(:quote, f)
    :(count($(esc(source)), $(esc(f)), $(esc(q))))
end

macro count(source)
    :(count($(esc(source))))
end

function default_if_empty end

macro default_if_empty(source)
	:(default_if_empty($(esc(source))))
end

macro default_if_empty(source, default_value)
	:(default_if_empty($(esc(source)), $(esc(default_value))))
end

function filter end

macro filter(source, f)
    q = Expr(:quote, f)
    :(QueryOperators.filter($(esc(source)), $(esc(f)), $(esc(q))))
end

function gather end

macro gather(source, withIndex = False)
	:(groupby($(esc(source)), $(esc(withIndex))))
end

function groupby end

macro groupby(source,elementSelector,resultSelector)
	q_elementSelector = Expr(:quote, elementSelector)
	q_resultSelector = Expr(:quote, resultSelector)

	:(groupby($(esc(source)), $(esc(elementSelector)), $(esc(q_elementSelector)), $(esc(resultSelector)), $(esc(q_resultSelector))))
end

macro groupby_simple(source,elementSelector)
	q_elementSelector = Expr(:quote, elementSelector)

	:(groupby($(esc(source)), $(esc(elementSelector)), $(esc(q_elementSelector))))
end

function groupjoin end

macro groupjoin(outer, inner, outerKeySelector, innerKeySelector, resultSelector)
	q_outerKeySelector = Expr(:quote, outerKeySelector)
	q_innerKeySelector = Expr(:quote, innerKeySelector)
	q_resultSelector = Expr(:quote, resultSelector)

	:(groupjoin($(esc(outer)), $(esc(inner)), $(esc(outerKeySelector)), $(esc(q_outerKeySelector)), $(esc(innerKeySelector)),$(esc(q_innerKeySelector)), $(esc(resultSelector)),$(esc(q_resultSelector))))
end

function join end

macro join(outer, inner, outerKeySelector, innerKeySelector, resultSelector)
	q_outerKeySelector = Expr(:quote, outerKeySelector)
	q_innerKeySelector = Expr(:quote, innerKeySelector)
	q_resultSelector = Expr(:quote, resultSelector)

	:(join($(esc(outer)), $(esc(inner)), $(esc(outerKeySelector)), $(esc(q_outerKeySelector)), $(esc(innerKeySelector)),$(esc(q_innerKeySelector)), $(esc(resultSelector)),$(esc(q_resultSelector))))
end

function map end

macro map(source, f)
    q = Expr(:quote, f)
    :(map($(esc(source)), $(esc(f)), $(esc(q))))
end

function mapmany end

macro mapmany(source,collectionSelector,resultSelector)
	q_collectionSelector = Expr(:quote, collectionSelector)
	q_resultSelector = Expr(:quote, resultSelector)

	:(mapmany($(esc(source)), $(esc(collectionSelector)), $(esc(q_collectionSelector)), $(esc(resultSelector)), $(esc(q_resultSelector))))
end

function orderby end

macro orderby(source, f)
	q = Expr(:quote, f)
    :(orderby($(esc(source)), $(esc(f)), $(esc(q))))
end

function orderby_descending end

macro orderby_descending(source, f)
	q = Expr(:quote, f)
    :(orderby_descending($(esc(source)), $(esc(f)), $(esc(q))))
end

function thenby end

macro thenby(source, f)
	q = Expr(:quote, f)
    :(thenby($(esc(source)), $(esc(f)), $(esc(q))))
end

function thenby_descending end

macro thenby_descending(source, f)
	q = Expr(:quote, f)
    :(thenby_descending($(esc(source)), $(esc(f)), $(esc(q))))
end

function take end

macro take(source, n)
	:(take($(esc(source)), $(esc(n))))
end

function drop end

macro drop(source, n)
	:(drop($(esc(source)), $(esc(n))))
end

function unique end

macro unique(source, f)
    q = Expr(:quote, f)
    :(unique($(esc(source)), $(esc(f)), $(esc(q))))
end


text:
## TODO: Implement support for complex numbers in @p_str macro

remove_spaces(str) = filter(x -> !isspace(x), str)

const operator_lookup = Dict(
    :* => *,
    :/ => /,
    :+ => +,
    :- => -,
    :// => //,
    :^ => ^
)

function parse_term(tree::Expr)
    operator = tree.args[1]
    if operator in (:(*), :(/), :(//))
        @views leftovers = tree.args[2:end]
        ind = findfirst(x -> !(x isa Number), leftovers) + 1
        coeff, exponent, var = parse_term(tree.args[ind])

        deleteat!(tree.args, ind)
        @assert all(x -> x isa Number, tree.args[2:end])

        coeff *= tree.args[2]
        coeff = operator_lookup[operator](1, coeff)

    elseif tree.args[1] == :(+)
        return parse_term(tree.args[2])
    elseif tree.args[1] == :(-)
        coeff, exponent, var = parse_term(tree.args[2])
        coeff = -1*coeff
    else
        coeff = 1
        exponent = tree.args[3]
        @assert exponent > 0
        var = tree.args[2]::Symbol
    end

    return coeff, exponent, var
end

parse_term(term::Symbol) = (1, 1, term)
parse_term(term::Number) = (term, 0, nothing)

split_polynomial(p_str) =
    split(p_str, r"(?=[+,-])") .|> remove_spaces .|> Meta.parse .|> parse_term

function parse_unnested_poly(p_str)
    arr = zip(split_polynomial(p_str)...) |> collect
    symbols = arr[3]
    #@show symbols
    if length(symbols) == 1
        var = symbols[1]
        var = isnothing(var) ? :CONSTANT : var
    else
        ind = findfirst(!isnothing, symbols)
        var = isnothing(ind) ? :CONSTANT : symbols[ind]
    end
    #@show var
    @assert all(x -> isnothing(x) || x == var, symbols)
    var == :im && error("Cannot parse complex polynomial")
    return StaticPolynomial(arr[1], arr[2], var)
end

function tryparse_unnested_poly(p_str)
    try
        return parse_unnested_poly(p_str)
    catch e
        return nothing
    end
end

function Base.parse(::Type{StaticPolynomial}, p_str)
    p = tryparse_unnested_poly(p_str)
    if isnothing(p)
        tree = Meta.parse(p_str)
        new_p_str = string(tree)
        p = tryparse_unnested_poly(new_p_str)
    else
        return p
    end

    if isnothing(p)
        a1 = tree.args[1]
        for (i, arg) in enumerate(tree.args)
            if arg isa Expr || arg == :x
                a = parse(StaticPolynomial, string(arg))
                if !isnothing(a)
                    tree.args[i] = a
                end
            end
        end

        if tree.args[1] isa StaticPolynomial
            pushfirst!(tree.args, :(*))
        end
        op = tree.args[1]
        rest = tree.args[2:end]

        for (i, a) in enumerate(rest)
            if a == :im
                rest[i] = Complex(false, true)
            end
        end
        return StaticPolynomial(operator_lookup[op](rest...))
    else
        return p
    end
end

function Base.tryparse(::Type{StaticPolynomial}, p_str)
    try parse(StaticPolynomial, p_str)
    catch e
        return nothing
    end
end

StaticPolynomial(p_str::String) = parse(StaticPolynomial, p_str)

macro p_str(p)
    StaticPolynomial(p)
end


text:
using Base.LinAlg.BLAS
using HDF5
using PyCall
@pyimport matplotlib.pyplot as plt
@pyimport numpy as np

abstract AbstractMonitor
@runonce type Monitor <: AbstractMonitor
    LastIndex::Int
    UseValidation::Bool
    MonitorEvery::Int
    MonitorVisual::Bool
    MonitorText::Bool
    Epochs::Vector{Float64}
    LearnRate::Vector{Float64}
    Momentum::Vector{Float64}
    PseudoLikelihood::Vector{Float64}
    TAPLikelihood::Vector{Float64}
    ValidationPseudoLikelihood::Vector{Float64}
    ValidationTAPLikelihood::Vector{Float64}
    ReconError::Vector{Float64}
    ValidationReconError::Vector{Float64}
    BatchTime_µs::Vector{Float64}
    FigureHandle
end

function Monitor(n_iter,monitor_every;monitor_vis=false,monitor_txt=true,validation=false)
    len = convert(Int,floor(n_iter/monitor_every))
    blank_vector1 = vec(fill!(Array(Float64,len,1),convert(Float64,NaN)))
    blank_vector2 = copy(blank_vector1)
    blank_vector3 = copy(blank_vector1)
    blank_vector4 = copy(blank_vector1)
    blank_vector5 = copy(blank_vector1)
    blank_vector6 = copy(blank_vector1)
    blank_vector7 = copy(blank_vector1)
    blank_vector8 = copy(blank_vector1)
    blank_vector9 = copy(blank_vector1)
    blank_vector10 = copy(blank_vector1)

    if monitor_vis
        fh = plt.figure(1;figsize=(15,10))
    else
        fh = NaN
    end

    Monitor(0,                   # Last Index
            validation,          # Flag for validation set
            monitor_every,       # When to display
            monitor_vis,         # Monitor visal display flag
            monitor_txt,         # Monitor text display flag
            blank_vector1,       # Epochs (for x-axes)
            blank_vector2,       # Learn Rate
            blank_vector3,       # Momentum
            blank_vector4,       # Pseudo-Likelihood
            blank_vector5,       # Tap-Likelihood
            blank_vector6,       # Validation Pseudo-Likelihood
            blank_vector7,       # Validation  Tap-Likelihood
            blank_vector8,       # ReconError
            blank_vector9,       # ValidationReconError
            blank_vector10,      # BatchTime_µs
            fh)                  # Monitor Figure Handle
end

function UpdateMonitor!(rbm::RBM,mon::Monitor,dataset::Mat{Float64},itr::Int;validation=[],bt=NaN,lr=NaN,mo=NaN)
    nh = size(rbm.W,1)
    nv = size(rbm.W,2)
    N = nh + nv
    nsamps = min(size(dataset,2),5000)      # Set maximum number of samples to test as 5000

    if itr%mon.MonitorEvery==0
        if mon.UseValidation
            vpl = mean(score_samples(rbm, validation))/N
            vtl = mean(score_samples_TAP(rbm, validation))/N
            vre = recon_error(rbm,validation)/N
        else
            vpl = NaN
            vtl = NaN
            vre = NaN
        end
        pl = mean(score_samples(rbm, dataset[:,1:nsamps]))/N
        tl = mean(score_samples_TAP(rbm, dataset[:,1:nsamps]))/N
        re = recon_error(rbm,dataset[:,1:nsamps])/N

        mon.LastIndex+=1
        li = mon.LastIndex

        mon.PseudoLikelihood[li] = pl
        mon.TAPLikelihood[li] = tl
        mon.ReconError[li] = re
        mon.ValidationPseudoLikelihood[li] = vpl
        mon.ValidationTAPLikelihood[li] = vtl
        mon.ValidationReconError[li] = vre
        mon.Epochs[li] = itr
        mon.Momentum[li] = mo
        mon.LearnRate[li] = lr
        mon.BatchTime_µs[li] = bt
    end
end


text:
## Common Subexpression Elimination (CSE)

"""Do common subexpression elimination on the expression `ex`,
by traversing it recursively and reducing to binary operations.

Modifies the `OrderedDict` `dict`.
"""
function cse!(dict, ex)

    if istree(ex)

        args = arguments(ex)
        op = operation(ex)

        if length(args) == 1

            left  = cse!(dict, args[1])

            ex = op(left)

        elseif length(args) == 2

            left  = cse!(dict, args[1])
            right = cse!(dict, args[2])

            ex = op(left, right)

        else
            left  = cse!(dict, args[1])
            right = cse!(dict, op(args[2:end]...))   # use similarterm?

            ex = op(left, right)

        end

        if haskey(dict, ex)
            return dict[ex]

        else
            val = make_variable()
            push!(dict, ex => val)
        end

        return val

    else  # not a tree
        return ex
    end

end

"Do CSE on an expression"
function cse(ex)
    dict = OrderedDict()
    final = cse!(dict, ex)

    return dict, final
end

"Version of CSE returing a vector of equations"
function cse_equations(ex)
    dict, final = cse(ex)

    return [Assignment(rhs, lhs) for (lhs, rhs) in pairs(dict)], final
end

cse(ex::Num) = cse(Symbolics.value(ex))

# @syms x y

# ex = exp(3x^2 + 4x^2 * y)

# dict, final = cse(ex)

# cse_equations(ex)

# dict
# final


text:
# Use at own risk!

#Script para generar series de tiempo conjeturadas de los datos abiertos oficiales.

#Módulos para cargar los datos.
using ZipFile, CSV, DataFrames, Dates

#Definimos una función para clasificar un paciente como vivo (true) o fallecido (false) dependiendo de la fecha de corte:
function clasificación_vivo(fecha_corte, row)

    #Si la fecha de defunción no está dada (missing), devolvemos vcerdadero:
    if ismissing(row[:FECHA_DEF])

        return true
    end

    #Si la fecha de defunción está definida, revisamos la fecha de corte:
    if (row[:FECHA_DEF] != "9999-99-99")

        #Si estamos en la fecha de corte o después, consideramos al paciente como fallecido.
        if Date(row[:FECHA_DEF]) <= fecha_corte

            return false
        end
    end

    #En otro caso, consideramos al paciente como vivo
    return true
end

#Definimos una función que devuelve los dataframes de interés.
#Estos contienen los datos filtrados entre el inicio de la epidemia y una fecha de corte.
function cargar_datos_fecha(datos, fecha_corte)

    #Filtramos los casos ingresados entre el inicio del brote y la fecha de corte.
    datos_corte = filter(row -> row[:FECHA_INGRESO] <= fecha_corte, datos)

    #Clasifiquemos los casos:
    #Positivos -> :CLASIFICACION_FINAL == 1,2,3 (positivo por laboratorio, dictaminación, asociación epidemiológica o prueba de antígeno)
    #Sospechosos -> :CLASIFICACION_FINAL == 6 (resultado pendiente)
    #Negativos -> :CLASIFICACION_FINAL == 7
    positivos = datos_corte[datos_corte.CLASIFICACION_FINAL .∈ [[1, 2, 3]], :]
    sospechosos = datos_corte[datos_corte.CLASIFICACION_FINAL .== 6, :]
    negativos = datos_corte[datos_corte.CLASIFICACION_FINAL .== 7, :]

    #Removemos los fallecidos, si su fallecimiento es anterior o igual a la fecha de corte:
    #Fallecidos -> :FECHA_DEF != "9999-99-99"
    fallecidos = filter(row -> !(clasificación_vivo(fecha_corte, row)), positivos)
    vivos = filter(row -> clasificación_vivo(fecha_corte, row), positivos)

    #Clasificamos a los pacientes vivos en ambulatorios y hospitalizados:
    #Ambulatorios -> :TIPO_PACIENTE == 1
    #Hospitalizados -> :TIPO_PACIENTE == 2
    ambulatorios = vivos[vivos.TIPO_PACIENTE .== 1, :]
    hospitalizados = vivos[vivos.TIPO_PACIENTE .== 2, :]

    #Clasificamos a los pacientes hospitalizados en leves, graves y críticos:
    #Leves -> :UCI == 2 && :INTUBADO == 2
    #Graves -> :UCI == 1 && :INTUBADO == 2 (En UCI, pero sin intubar)
    #Críticos -> :INTUBADO == 1 (Intubados)
    leves = filter(row -> (row[:UCI] == 2) && (row[:INTUBADO] == 2), hospitalizados)
    graves = filter(row -> (row[:UCI] == 1) && (row[:INTUBADO] == 2), hospitalizados)
    críticos = filter(row -> (row[:INTUBADO] == 1), hospitalizados)

    return negativos, sospechosos, positivos, fallecidos, vivos, ambulatorios, hospitalizados, leves, graves, críticos
end

#Definimos una función auxiliar para contar resultados por clave de estado:
#Tomaremos la entidad de la unidad de atención médica para identificar la entidad de la información.
número(clave_entidad, df) = length(filter(row -> row[:ENTIDAD_UM] == clave_entidad, df)[!, :ENTIDAD_UM])
número(clave_entidad) = df -> número(clave_entidad, df)

#Definimos una función auxiliar que junte los datos de un estado:
function cálculo_estado(clave_entidad, negativos, sospechosos, positivos, fallecidos, vivos, ambulatorios, hospitalizados, leves, graves, críticos)

    negativos_edo = número(clave_entidad)(negativos)
    sospechosos_edo = número(clave_entidad)(sospechosos)
    positivos_edo = número(clave_entidad)(positivos)

    fallecidos_edo = número(clave_entidad)(fallecidos)

    ambulatorios_edo = número(clave_entidad)(ambulatorios)
    hospitalizados_edo = número(clave_entidad)(hospitalizados)
    leves_edo = número(clave_entidad)(leves)
    graves_edo = número(clave_entidad)(graves)
    críticos_edo = número(clave_entidad)(críticos)

    return [negativos_edo, sospechosos_edo, positivos_edo, fallecidos_edo, ambulatorios_edo, hospitalizados_edo, leves_edo, graves_edo, críticos_edo]
end

#Definimos una función que realice los cálculos para todos los estados y el país, devolviendo un string de resultados listos para agregar al archivo de registro.
function información_fecha(datos, fecha)

    datos_acumulados = cargar_datos_fecha(datos, fecha)

    resumen_estados = [cálculo_estado(i, datos_acumulados...) for i in 1:32]
    resumen_país = sum(resumen_estados)

    fila = string.(vcat(resumen_estados..., resumen_país))

    #Agrega la fecha a la primer columna:
    fecha_string = string(fecha)
    pushfirst!(fila, fecha_string)

    #Agrega comas para el csv y lo junta todo en un string
    fila_csv = prod(fila[1:(end - 1)].*",")*fila[end]

    return fila_csv
end

#Definamos una función que itere sobre las fechas y genere la tabla resumen:
function tabla_resumen(datos, fecha_análisis; fecha_máxima = fecha_análisis, fecha_mínima = Date("2020-01-01"))

    #Construyamos el header del .csv:
    #Definamos las abreviaturas de las entidades en orden: claves 1, 2, 3,... 32.
    abreviaturas_entidades = ["AGU", "BCN", "BCS", "CAM", "COA", "COL", "CHP", "CHH", "CMX", "DUR", "GUA", "GRO", "HID", "JAL", "MEX", "MIC", "MOR", "NAY", "NLE", "OAX", "PUE", "QUE", "ROO", "SLP", "SIN", "SON", "TAB", "TAM", "TLA", "VER", "YUC", "ZAC", "TOT"] #Añadimos una categoría cumulativa

    #Y los tags de los casos negativos (_N), sospechosos(_S), positivos (_P), positivos fallecidos (_PF), positivos ambulatorios (_PA), positivos hospitalizados (_PH),
    # positivos hospitalizados leves (_PHL), positivos hospitalizados graves (_PHG) y positivos hospitalizados críticos (_PHC)
    tags = ["_N", "_S", "_P", "_PF", "_PA", "_PH", "_PHL", "_PHG", "_PHC"]
    nombres_columnas = [entidad*tag for entidad in abreviaturas_entidades for tag in tags]

    #Generemos el header
    header = "Fecha,"*prod(nombres_columnas[1:(end - 1)].*",")*nombres_columnas[end]

    fecha_archivo = Dates.format(fecha_análisis, "yyyymmdd")

    open("Mexico_COVID19_$(fecha_archivo).csv", "w") do io

        #Escribimos el header:
        write(io, header*"\n")

        #Hacemos el cálculo de cada día y lo agreamos al archivo:
        for fecha in fecha_mínima:Day(1):fecha_máxima

            fila_csv = información_fecha(datos, fecha)

            write(io, fila_csv*"\n")
        end
    end

    return "Done"
end

####
#Construyamos el archivo resumen:

#Cambiamos el directorio de trabajo a dónde se encuentran los datos.
cd("Open_data/COVID-19")

if Time(now()) < Time(19)
    fecha_análisis = today() - Day(1)
else
    fecha_análisis = today()
end

#Los datos se encuentran en formato .csv al interior del archivo .zip:
datos_zip = ZipFile.Reader(Dates.format(fecha_análisis, "yyyymm")*"/"*Dates.format(fecha_análisis, "yyyymmdd")*".zip")
datos = CSV.read(datos_zip.files[1], header = 1, types = Dict(:FECHA_DEF => String), DataFrame)[: , [:FECHA_INGRESO, :ENTIDAD_UM, :CLASIFICACION_FINAL, :FECHA_DEF, :TIPO_PACIENTE, :UCI, :INTUBADO]]
sort!(datos, [:FECHA_INGRESO, :ENTIDAD_UM, :CLASIFICACION_FINAL, :FECHA_DEF, :TIPO_PACIENTE, :UCI, :INTUBADO])

tabla_resumen(datos, fecha_análisis)