using Test
using Miletus
using Dates: today, Day

using CSV, DataFrames

s = SingleStock()
d1 = today()

bs = CSV.read("bs.csv")

# 1. Test with floats
calls = EuropeanCall.(d1 .+ Day.(bs[:days]), [s], float.(bs[:strikeprice]))
puts = EuropeanPut.(d1 .+ Day.(bs[:days]), [s], float.(bs[:strikeprice]))
m = GeomBMModel.([d1], float.(bs[:startprice]), bs[:interestrate], bs[:carryrate], bs[:sigma])

atol=1e-6; rtol=1e-3

@test all(isapprox.(value.(m,calls), bs[:c], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,puts), bs[:p],  rtol=rtol,atol=atol))

# @test all(isapprox.(ivol.(m,calls,value.(m,calls)), bs[:sigma]))

crr = CSV.read("crr.csv")

eucalls = EuropeanCall.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice]))
euputs  = EuropeanPut.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] ))
amcalls = AmericanCall.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice]))
amputs  = AmericanPut.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] ))

m = CRRModel.(d1, d1 .+ Dates.Day.(crr[:days]), crr[:nsteps], float.(crr[:startprice]), crr[:interestrate], crr[:carryrate], crr[:sigma])

@test all(isapprox.(value.(m,eucalls) , crr[:ce], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,euputs)  , crr[:pe], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amcalls) , crr[:ca], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amputs)  , crr[:pa], rtol=rtol,atol=atol))

tian = CSV.read("tian.csv")

eucalls = EuropeanCall.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice]))
euputs  = EuropeanPut.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] ))
amcalls = AmericanCall.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice]))
amputs  = AmericanPut.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] ))

m = TianModel.([d1],d1 .+ Dates.Day.(tian[:days]), tian[:nsteps], float.(tian[:startprice]), tian[:interestrate], tian[:carryrate], tian[:sigma])

@test all(isapprox.(value.(m,eucalls) , tian[:ce], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,euputs)  , tian[:pe], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amcalls) , tian[:ca], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amputs)  , tian[:pa], rtol=rtol,atol=atol))

# 2. Test with currencies
calls = EuropeanCall.(d1 .+ Day.(bs[:days]), [s], float.(bs[:strikeprice].*[USD]))
puts = EuropeanPut.(d1 .+ Day.(bs[:days]), [s], float.(bs[:strikeprice].*[USD]))
m = GeomBMModel.([d1], float.(bs[:startprice].*[USD]), bs[:interestrate], bs[:carryrate], bs[:sigma])

atol=1e-6; rtol=1e-3

@test all(isapprox.(value.(m,calls) ./ [USD], bs[:c], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,puts) ./ [USD], bs[:p],  rtol=rtol,atol=atol))

crr = CSV.read("crr.csv")

eucalls = EuropeanCall.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] .* [USD]))
euputs  = EuropeanPut.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] .* [USD]))
amcalls = AmericanCall.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] .* [USD]))
amputs  = AmericanPut.(d1 .+ Dates.Day.(crr[:days]), [s], float.(crr[:strikeprice] .* [USD]))

m = CRRModel.([d1],d1 .+ Dates.Day.(crr[:days]), crr[:nsteps], float.(crr[:startprice].*[USD]), crr[:interestrate], crr[:carryrate], crr[:sigma])

@test all(isapprox.(value.(m,eucalls) ./ [USD], crr[:ce], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,euputs)  ./ [USD], crr[:pe], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amcalls) ./ [USD], crr[:ca], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amputs)  ./ [USD], crr[:pa], rtol=rtol,atol=atol))

tian = CSV.read("tian.csv")

eucalls = EuropeanCall.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] .* [USD]))
euputs  = EuropeanPut.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] .* [USD]))
amcalls = AmericanCall.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] .* [USD]))
amputs  = AmericanPut.(d1 .+ Dates.Day.(tian[:days]), [s], float.(tian[:strikeprice] .* [USD]))

m = TianModel.([d1],d1 .+ Dates.Day.(tian[:days]), tian[:nsteps], float.(tian[:startprice].*[USD]), tian[:interestrate], tian[:carryrate], tian[:sigma])

@test all(isapprox.(value.(m,eucalls) ./ [USD], tian[:ce], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,euputs)  ./ [USD], tian[:pe], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amcalls) ./ [USD], tian[:ca], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amputs)  ./ [USD], tian[:pa], rtol=rtol,atol=atol))

jr = CSV.read("jr.csv")

eucalls = EuropeanCall.(d1 .+ Dates.Day.(jr[:days]), [s], float.(jr[:strikeprice] .* [USD]))
euputs  = EuropeanPut.(d1 .+ Dates.Day.(jr[:days]), [s], float.(jr[:strikeprice] .* [USD]))
amcalls = AmericanCall.(d1 .+ Dates.Day.(jr[:days]), [s], float.(jr[:strikeprice] .* [USD]))
amputs  = AmericanPut.(d1 .+ Dates.Day.(jr[:days]), [s], float.(jr[:strikeprice] .* [USD]))

m = JRModel.([d1],d1 .+ Dates.Day.(jr[:days]), jr[:nsteps], float.(jr[:startprice].*[USD]), jr[:interestrate], jr[:carryrate], jr[:sigma])

@test all(isapprox.(value.(m,eucalls) ./ [USD], jr[:ce], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,euputs)  ./ [USD], jr[:pe], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amcalls) ./ [USD], jr[:ca], rtol=rtol,atol=atol))
@test all(isapprox.(value.(m,amputs)  ./ [USD], jr[:pa], rtol=rtol,atol=atol))

# Test Black-76 model
d2 = d1 + Day(120)
cm = CoreForwardModel(d1, d2, 110.00USD, 0.02)
m = GeomBMModel(cm, 0.3)
o = EuropeanCall( d2, s, 100.00USD)
@test isapprox(value(m,o), 13.1833USD, rtol=1e-5)


using Base.Iterators

export all_functions, sqrt, has_sqrt, all_sqrts, random_square

# This code is for finding compositional square roots of functions.
# It is extremely inefficient!!!

"""
`all_functions(n)` returns a generator that produces all
`DiscreteFunctions` on `{1,2,...,n}`
"""
function all_functions(n::Int)
    @assert n>0 "Argument must be a positive integer"
    arg = (1:n for t=1:n)
    iter = product(arg...)
    (DiscreteFunction(collect(t)) for t in iter)
end
#
# OLD VERSION OF sqrt
#
# """
# `sqrt(f::DiscreteFunction)` returns a `DiscreteFunction`
# `g` such that `g*g==f` or throws an error if no such `g` exists.
# """
# function sqrt(f::DiscreteFunction)::DiscreteFunction
#     n = length(f)
#     gen = all_functions(n)
#     for g in gen
#         if g*g == f
#             return g
#         end
#     end
#     error("This function does not have a square root")
# end

"""
`all_sqrts(f::DiscreteFunction)` returns an array consisting
of all `g` such that `g*g==f`.
"""
function all_sqrts(f::DiscreteFunction)
    n = length(f)
    gen = all_functions(n)
    [ g for g in gen if g*g==f ]
end

"""
`all_squares(n)` returns a list of all function on `[n]` that are
perfect compositional squares.
"""
function all_squares(n::Int)
    G = all_functions(n)
    S = Set(f*f for f in G)
    collect(S)
end

"""
`random_square(n)` returns a `DiscreteFunction` on `n` elements
that has a compositional square root.
"""
function random_square(n::Int)::DiscreteFunction
    @assert n>0 "Require `n` to be positive"
    while true
        f = RandomFunction(n)
        if has_sqrt(f)
            return f
        end
    end
end


@testset "MonteCarlo" begin
    mc = MonteCarlo(1000)

    @test isa(mc, MonteCarlo)
    @test mc.n == 1000
end

@testset "QuasiMonteCarlo" begin

    # Make sure that all QMC Types implement qmc_samples
    for qmc in subtypes(AbstractQuasiMonteCarlo)
        @test hasmethod(qmc_samples, (qmc, Integer))
    end

    @testset "SobolSampling" begin
        sobol = SobolSampling(1000)

        @test isa(sobol, SobolSampling)
        @test sobol.n == 1000

        @testset "sample" begin
            inputs = [RandomVariable.(Uniform(), [:a, :b]); Parameter(1, :c)]

            samples = sample(inputs, SobolSampling(4))

            @test isapprox(samples.a, [0.375, 0.875, 0.625, 0.125])
            @test isapprox(samples.b, [0.375, 0.875, 0.125, 0.625])
            @test samples.c == [1.0, 1.0, 1.0, 1.0]
        end
    end

    @testset "HaltonSampling" begin
        halton = HaltonSampling(1000)

        @test isa(halton, HaltonSampling)
        @test halton.n == 1000

        @testset "sample" begin
            inputs = [RandomVariable.(Uniform(), [:a, :b]); Parameter(1, :c)]

            samples = sample(inputs, HaltonSampling(4))

            @test isapprox(samples.a, [0.5, 0.25, 0.75, 0.125])
            @test isapprox(samples.b, [1 / 3, 2 / 3, 1 / 9, 4 / 9])
            @test samples.c == [1.0, 1.0, 1.0, 1.0]
        end
    end
end

@testset "LatticeRuleSampling" begin
    lattice = LatticeRuleSampling(1000)
    @test isa(lattice, LatticeRuleSampling)
    @test lattice.n == 1000
    @testset "sample" begin
        inputs = [RandomVariable.(Uniform(), [:a, :b]); Parameter(1, :c)]

        Random.seed!(1)
        s = QuasiMonteCarlo.sample(4, [0, 0], [1, 1], LatticeRuleSample())
        Random.seed!(1)
        samples = sample(inputs, LatticeRuleSampling(4))
        Random.seed!()

        @test isapprox(samples.a, s[1, :])
        @test isapprox(samples.b, s[2, :])
        @test samples.c == [1.0, 1.0, 1.0, 1.0]
    end
end


n = 10
using Test

pairwise_funcs = [
    StatsBase.cov,
    StatsBase.cor,
    StatsBase.countne,
    StatsBase.counteq,
    StatsBase.corkendall,
    StatsBase.corspearman,
    StatsBase.maxad,
    StatsBase.meanad,
    StatsBase.msd,
    StatsBase.psnr,
    StatsBase.rmsd,
    StatsBase.sqL2dist,
    StatsBase.crosscor,
    StatsBase.crosscov
]

@testset "Pairwise statistics on uncertain values" begin
    @testset "$(pairwise_funcs[i])" for i = 1:length(pairwise_funcs)
        f = pairwise_funcs[i]
        @testset for (i, uval) in enumerate(example_uvals)
            if f == StatsBase.psnr
                maxv = 100
                @test f(uval, uval, maxv, n) isa T where T <: Real

            elseif f ∈ [StatsBase.crosscor, StatsBase.crosscov]
                @test f(uval, uval, n) isa AbstractVector{T} where T <: Real
                @test f(uval, uval, 1:5, n) isa AbstractVector{T} where T <: Real
            else
                @test f(uval, uval, n) isa T where T <: Real
            end
        end
    end;
end

# Functions that under the hood use functions with strictly positive domains
# special_pairwise_funcs = [
#     StatsBase.gkldiv,
#     StatsBase.kldivergence,
# ]

# @testset "$(pairwise_funcs[i])" for i = 1:length(pairwise_funcs)
#     f = pairwise_funcs[i]
#     @testset for (i, uval) in enumerate(example_uvals)
#         @test f(uval, uval, n) isa T where T <: Real
#     end
# end;


using DataFrames
using RDatasets
using Distributions
using Gadfly

srand(619)

names(Normal)

params(dist1)

dist1.μ

dist1.σ

x = rand(dist1, 1000)

dist2 = Truncated(dist1, -4.0, 6.0)

plot(x = rand(dist1, 1000), Geom.histogram(bincount = 30))


using SyntheticDatasets
using DataFrames
using Test

@testset "SkLearn Generators" begin
    samples = 20000
    features = 20

    data = SyntheticDatasets.make_blobs(centers = [-1 1;-0.5 0.75],
                                        cluster_std = 0.225,
                                        n_samples = 20000,
                                        center_box = (-1.5, 1.5))

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    samples = 20000
    data = SyntheticDatasets.make_moons(n_samples = 20000)

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    data = SyntheticDatasets.make_s_curve(  n_samples = samples,
                                            noise = 2.2,
                                            random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 4

    data = SyntheticDatasets.make_circles(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                n_targets = 2,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 2

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                n_targets = 3,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 3

    data = SyntheticDatasets.make_classification(   n_samples = samples,
                                                    n_features = features,
                                                    n_classes = 1)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_friedman1(n_samples = samples,
                                            n_features = features)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_friedman2(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 5

    data = SyntheticDatasets.make_friedman3(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 5

    data = SyntheticDatasets.make_low_rank_matrix(  n_samples = samples,
                                                    n_features = features,
                                                    effective_rank = 10,
                                                    tail_strength = 0.5,
                                                    random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features

    data = SyntheticDatasets.make_swiss_roll(   n_samples = samples,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 4

    data = SyntheticDatasets.make_hastie_10_2(  n_samples = samples,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 11

    data = SyntheticDatasets.make_gaussian_quantiles(   n_samples = samples,
                                                        n_features = features,
                                                        random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    @test_throws DimensionMismatch SyntheticDatasets.make_gaussian_quantiles(   n_samples = 300,
                                                                                n_features = 4,
                                                                                random_state = 5,
                                                                                mean = [1, 2, 3])
end

@testset "Matlab Generators" begin
    samples = 20000

    data = SyntheticDatasets.make_twospirals(n_samples = samples,
                                                 noise = 2.2)
    @test size(data)[1] == samples

    data = SyntheticDatasets.make_halfkernel(n_samples = samples)
    @test size(data)[1] == samples

    data = SyntheticDatasets.make_outlier(n_samples = samples)
    @test size(data)[1] == samples
    @test size(data)[2] == 3
end


# Adders

function HalfAdder(; name)
    @named x1 = DigitalPin()
    @named x2 = DigitalPin()
    @named y1 = DigitalPin()
    @named y2 = DigitalPin()
    @variables sum(t), carry(t)

    eqs = [
        y1.val ~ _xor(x1.val, x2.val)
        y2.val ~ _and(x1.val, x2.val)
        sum ~ y1.val
        carry ~ y2.val
    ]
    ODESystem(eqs, t, [sum, carry], [], systems=[x1, x2, y1, y2], name=name)
end

function FullAdder(; name)
    @named x1 = DigitalPin()
    @named x2 = DigitalPin()
    @named x3 = DigitalPin()
    @named y1 = DigitalPin()
    @named y2 = DigitalPin()
    @variables sum(t), carry(t)

    eqs = [
        y1.val ~ _xor(x1.val, x2.val, x3.val)
        y2.val ~ _or(_and(x3.val, _xor(x1.val, x2.val)), _and(x1.val, x2.val))
        sum ~ y1.val
        carry ~ y2.val
    ]
    ODESystem(eqs, t, [sum, carry], [], systems=[x1, x2, x3, y1, y2], name=name)
end

# Multiplexers

# This selects data from the `N` input ports (`d₀` to `dₙ₋₁`)
# using values of `n` select lines, where `N = 2^n`
function MUX(; name, N=4)
    n = log2(N)
    try n = Int(n) catch(e) throw("`N` must be a power of 2") end
    s = map(0:n-1) do i
        DigitalPin(; name=Symbol(:s, i))
    end
    d = map(0:N-1) do i
        DigitalPin(; name=Symbol(:d, i))
    end
    @named y = DigitalPin()

    nodes = Num[]
    for i in 1:N
        bin = digits!(zeros(Int64,n), i-1, base=2)
        statelist = Term{Real, Nothing}[]
        for j in 1:n
            varstate = bin[j] == 0 ? _not(s[j].val) : s[j].val
            push!(statelist, varstate)
        end
        push!(nodes, _and(statelist..., d[i].val))
    end

    eqs = Equation[
        y.val ~ _or(nodes...)
    ]

    ODESystem(eqs, t, [], [], systems=[d..., s..., y], name=name)
end

# This selects one of the `N` output ports (`y₀` to `yₙ₋₁`)
# to transmit data `d` using values of `n` select lines, where `N = 2^n`
function DEMUX(; name, N=4)
    n = log2(N)
    try n = Int(n) catch(e) throw("`N` must be a power of 2") end
    @named d = DigitalPin()
    s = map(0:n-1) do i
        DigitalPin(; name=Symbol(:s, i))
    end
    y = map(0:N-1) do i
        DigitalPin(; name=Symbol(:y, i))
    end

    eqs = Equation[]
    for i in 1:N
        bin = digits!(zeros(Int64,n), i-1, base=2)
        statelist = Term{Real, Nothing}[]
        for j in 1:n
            varstate = bin[j] == 0 ? _not(s[j].val) : s[j].val
            push!(statelist, varstate)
        end
        push!(eqs, y[i].val ~ _and(statelist..., d.val))
    end

    ODESystem(eqs, t, [], [], systems=[d, s..., y...], name=name)
end

# Encoder-Decoder

# Encodes `N` inputs to `n` outputs, where `N = 2^n`
function Encoder(; name, N=4)
    n = log2(N)
    try n = Int(n) catch(e) throw("`N` must be a power of 2") end
    d = map(0:N-1) do i
        DigitalPin(; name=Symbol(:d, i))
    end
    y = map(0:n-1) do i
        DigitalPin(; name=Symbol(:y, i))
    end

    nodes = Vector{Term{Real, Nothing}}[]
    i = 0
    for j in 1:n
        counter = 1
        statelist = Term{Real, Nothing}[]
        while i < N
            while counter <= 2^j
                counter > 2^(j-1) && push!(statelist, d[i+1].val)
                counter += 1
                i = i+1
            end
            counter = 1
        end
        i = 0
        push!(nodes, statelist)
    end

    eqs = Equation[]
    for i in n:-1:1
        push!(eqs, y[i].val ~ _or(nodes[i]...))
    end

    ODESystem(eqs, t, [], [], systems=[d..., y...], name=name)
end

# Decodes `n` inputs to `N` outputs, where `N = 2^n`
function Decoder(; name, n=2)
    N = 2^n
    d = map(0:n-1) do i
        DigitalPin(; name=Symbol(:d, i))
    end
    y = map(0:N-1) do i
        DigitalPin(; name=Symbol(:y, i))
    end

    nodes = Vector{Term{Real, Nothing}}[]
    for i in 1:N
        bin = digits!(zeros(Int64, n), i-1, base=2)
        statelist = Term{Real, Nothing}[]
        for j in 1:n
            varst = bin[j] == 0 ? _not(d[j].val) : d[j].val
            push!(statelist, varst)
        end
        push!(nodes, statelist)
    end

    eqs = Equation[]
    for i in N:-1:1
        push!(eqs, y[i].val ~ _and(nodes[i]...))
    end

    ODESystem(eqs, t, [], [], systems=[d..., y...], name=name)
end


# In-place update of the initial state distribution.
function update_a!(a::AbstractVector, α::AbstractMatrix, β::AbstractMatrix)
    @argcheck size(α, 1) == size(β, 1)
    @argcheck size(α, 2) == size(β, 2) == size(a, 1)

    K = length(a)
    c = 0.0

    for i in OneTo(K)
        a[i] = α[1,i] * β[1,i]
        c += a[i]
    end

    for i in OneTo(K)
        a[i] /= c
    end
end

# In-place update of the transition matrix.
function update_A!(A::AbstractMatrix, ξ::AbstractArray, α::AbstractMatrix, β::AbstractMatrix, LL::AbstractMatrix)
    @argcheck size(α, 1) == size(β, 1) == size(LL, 1) == size(ξ, 1)
    @argcheck size(α, 2) == size(β, 2) == size(LL, 2) == size(A, 1) == size(A, 2) == size(ξ, 2) == size(ξ, 3)

    T, K = size(LL)

    @inbounds for t in OneTo(T - 1)
        m = vec_maximum(view(LL, t+1, :))
        c = 0.0

        for i in OneTo(K), j in OneTo(K)
            ξ[t,i,j] = α[t,i] * A[i,j] * exp(LL[t + 1,j] - m) * β[t + 1,j]
            c += ξ[t,i,j]
        end

        for i in OneTo(K), j in OneTo(K)
            ξ[t,i,j] /= c
        end
    end

    fill!(A, 0.0)

    @inbounds for i in OneTo(K)
        c = 0.0

        for j in OneTo(K)
            for t in OneTo(T - 1)
                A[i,j] += ξ[t,i,j]
            end
            c += A[i,j]
        end

        for j in OneTo(K)
            A[i,j] /= c
        end
    end
end

# In-place update of the observations distributions.
function update_B!(B::AbstractVector, γ::AbstractMatrix, observations, estimator)
    @argcheck size(γ, 1) == size(observations, 1)
    @argcheck size(γ, 2) == size(B, 1)
    K = length(B)
    for i in OneTo(K)
        if sum(γ[:,i]) > 0
            B[i] = estimator(typeof(B[i]), permutedims(observations), γ[:,i])
        end
    end
end

function fit_mle!(hmm::AbstractHMM, observations; display = :none, maxiter = 100, tol=1e-3, robust = false, estimator=fit_mle)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    T, K = size(observations, 1), size(hmm, 1)
    history = EMHistory(false, 0, [])

    # Allocate memory for in-place updates
    c = zeros(T)
    α = zeros(T, K)
    β = zeros(T, K)
    γ = zeros(T, K)
    ξ = zeros(T, K, K)
    LL = zeros(T, K)

    loglikelihoods!(LL, hmm, observations)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    forwardlog!(α, c, hmm.a, hmm.A, LL)
    backwardlog!(β, c, hmm.a, hmm.A, LL)
    posteriors!(γ, α, β)

    logtot = sum(log.(c))
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it in 1:maxiter
        update_a!(hmm.a, α, β)
        update_A!(hmm.A, ξ, α, β, LL)
        update_B!(hmm.B, γ, observations, estimator)

        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check isprobvec(hmm.a)
        @check istransmat(hmm.A)

        loglikelihoods!(LL, hmm, observations)
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        forwardlog!(α, c, hmm.a, hmm.A, LL)
        backwardlog!(β, c, hmm.a, hmm.A, LL)
        posteriors!(γ, α, β)

        logtotp = sum(log.(c))
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) && println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    history
end

mutable struct EMHistory
    converged::Bool
    iterations::Int
    logtots::Vector{Float64}
end


export l1Reg

"""
l1Reg(m,mref,M)

l1-regularizer: R = sum(abs(m-mref))
"""
function l1Reg(m::Vector{Float64},mref::Vector{Float64},M::AbstractMesh)
	dm   = m .- mref
    dR   = ones(size(m))
    Rc   = sum(dm)
    return Rc,dR,1e-4*speye(length(m))
end


import Random
using Random: MersenneTwister
using Distributions
using Test
using StaticArrays
using LinearAlgebra
using Unitful

μ = rand()
x = rand()
σ = rand()
Σ = σ*σ'

# Check type conversions
GFloat = Gaussian{Vector{Float64}, Matrix{Float64}}
v = GFloat[Gaussian([1.0], Matrix(1.0I, 2, 2)),
           Gaussian(SVector(1.0), @SMatrix [1.0 0.0; 0.0 1.0])]
@test mean.(v) == [[1.0], [1.0]]

G = Gaussian([1.0], Matrix(1.0I, 2, 2))
m, K = G
@test m === mean(G)
@test K === cov(G)

p = pdf(Normal(μ, √Σ), x)
@test pdf(Gaussian(μ, Σ), x) ≈ p
@test pdf(Gaussian(μ, Σ*I), x) ≈ p
@test pdf(Gaussian([μ], [σ]*[σ]'), [x]) ≈ p

@test pdf(Gaussian((@SVector [μ]), @SMatrix [Σ]), @SVector [x]) ≈ p

for d in 1:3
    μ = rand(d)
    x = rand(d)
    σ = tril(rand(d,d))
    Σ = σ*σ'
    p = pdf(MvNormal(μ, Σ), x)

    @test pdf(Gaussian(μ, Σ), x) ≈ p
    @test pdf(Gaussian(μ, PSD(σ)), x) ≈ p
    @test pdf(Gaussian(SVector{d}(μ), SMatrix{d,d}(Σ)), x) ≈ p
    @test pdf(Gaussian(SVector{d}(μ), PSD(SMatrix{d,d}(σ))), x) ≈ p
end

for d in 1: 3
    local μ = rand(d)
    local x = rand(d)
    local σ = rand()
    local Σ = Matrix(1.0I, d, d).*σ^2
    local p = pdf(MvNormal(μ, Σ), x)

    @test pdf(Gaussian(μ, σ^2*I), x) ≈ p
    @test pdf(Gaussian(SVector{d}(μ), SDiagonal(σ^2*ones(SVector{d}))), x) ≈ p
    @test pdf(Gaussian(SVector{d}(μ), SMatrix{d,d}(Σ)), x) ≈ p
end

@test rand(Random.GLOBAL_RNG, Gaussian(1.0, 0.0)) == 1.0
@test mean(rand(MersenneTwister(1), Gaussian([1., 2], Matrix(1.0I, 2, 2)), 100000)) ≈ 1.5 atol=0.02
@test mean(rand(MersenneTwister(1), Gaussian(1.5u"m", 2.0u"m^2"), 100000)) ≈ 1.5u"m" atol=0.02u"m"
@test std(Gaussian(1.0u"m", 4.0u"m^2")) == 2.0u"m"

@test rand(Gaussian(1.0,0.0)) == 1.0
Random.seed!(5)
x = randn()
Random.seed!(5)
@test rand(Gaussian(1.0,2.0)) == 1.0 .+ sqrt(2)*x

Random.seed!(5)
@test rand(Gaussian(1.0,2.0), (1,)) == [1.0 + sqrt(2)*x]

g = Gaussian([1,2], Matrix(1.0I, 2, 2))
@test mean(g + [10,10]) == [11,12]
@test (g - [10,10]) + [10,10] == g
@test mean(g + [10, 20]) == [11,22]
m = [1.0 2.0; 0.0 2.0]
@test cov(m * g) == m * m'

using GaussianDistributions: ⊕
@test cov(m*g ⊕ g) == cov(m*g) + cov(g)
@test mean(m*g ⊕ g) == mean(m*g) + mean(g)

@test cov(g ⊕ m*mean(g)) == cov(g)
@test mean(m*mean(g) ⊕ g) == mean(m*g) + mean(g)

g2 = Gaussian(SVector(1.000001, 2.0), [1.00001 0.0; 0.0 0.999999])
@test !(g ≈ g2)
@test g ≈ g2 rtol=0.0001


# Load X and y variable
using JLD
data = load("multiData.jld")
(X,y,Xtest,ytest) = (data["X"],data["y"],data["Xtest"],data["ytest"])

# Data is already roughly standardized, but let's add bias
n = size(X,1)
X = [ones(n,1) X]

# Do the same transformation to the test data
t = size(Xtest,1)
Xtest = [ones(t,1) Xtest]

# Fit one-vs-all logistic regression model
#include("logReg.jl")
#model = logRegOnevsAll(X,y)
include("softmaxClassifier.jl")
model = softmaxClassifier(X,y)

# Compute training and validation error
yhat = model.predict(X)
trainError = mean(yhat .!= y)
@show(trainError)
yhat = model.predict(Xtest)
validError = mean(yhat .!= ytest)
@show(validError)

# Plot results
k = maximum(y)
include("plot2Dclassifier.jl")
plot2Dclassifier(X,y,model,Xtest=Xtest,ytest=ytest,biasIncluded=true,k=5)
show()


using TracyWidom
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

# See https://arxiv.org/abs/0904.1581
@test TWcdf(0,beta=1) ≈ 0.83190806620295 atol=1e-14
@test TWcdf(0,beta=2) ≈ 0.96937282835526 atol=1e-14


export compute_dt

"""
    compute_dt( sqg, ogird, w)

CFL of the (large-scale) advection

"""
function compute_dt(sqg)

    update_velocities!(sqg)
    dx = sqg.grid.dx
    dy = sqg.grid.dy

    bound1 = maximum(abs.(sqg.u_x .* π ./ dx) .+ abs.(sqg.u_y .* π ./ dy))
    bound1 = 1 / bound1 / 4

    # CFL of the hyperviscosity

    dx2 = ((dx, dy) ./ pi) .^ 2
    bound2 = 1 / sqg.hv_val * (prod(dx2) / sum(dx2))^(sqg.hv_order ÷ 2)

    # Minimum of the CFL
    min(minimum(bound1), bound2)

end


module ConvexHull
using LinearAlgebra
using StaticArrays

import Polyhedra # for convex hull
import GLPK

module LibTriangle
	using Triangle
end
include("Projectors.jl")
using .Projectors
# using MiniQhull

# tools««1
# findextrema««2
"""
    findextrema(itr; lt=isless)

Like `findmin`, except that
 - it returns both extrema, as a `NamedTuple`;
 - it accepts an `lt` parameter, like `sort`.
"""
function findextrema(itr; lt=isless)
  p = pairs(itr); y = iterate(p)
  if y == nothing
    throw(ArgumentError("collection must be non-empty"))
  end
  (mi, m), s = y; (Mi, M) = (mi, m); i = mi
  while true
    y = iterate(p, s)
    y == nothing && break
    (ai, a), s = y
    if lt(a, m) m = a; mi = ai; end
    if lt(M, a) M = a; Mi = ai; end
  end
  return (min=(m, mi), max=(M, Mi))
end
# determinants««2
@inline det2(x,y) = x[1]*y[2]-x[2]*y[1]
@inline det2(x,y,z) = det2(y-x, z-x)
# triangulations««2
"""
    triangulate_face(points; direction, map)

Returns a triangulation of the face (assumed convex; points in any order)
Optional keyword arguments:
 - `direction` is a normal vector (used for projecting to 2d).
 - `map` is a labeling of points (default is identity map).
 - `convex` is a Val(true) or Val(false).

The triangulation is returned as a vector of StaticVector{3,Int},
containing the labels of three points of each triangle.
"""
function triangulate_face(points::AbstractVector{<:StaticVector{3}};
		direction::AbstractVector = face_normal(points),
		map::AbstractVector{<:Integer} = [1:length(points)...],
		)
	axis = main_axis(direction)
	N = length(points)
	# this common case deserves a shortcut:
	N == 3 && return [(map[1], map[2], map[3])]

	points2d = Matrix{Float64}(undef, N, 2)
	for (i, p) in pairs(points)
		points2d[i,:] .= project2d(axis, p)
	end
	tri = LibTriangle.basic_triangulation(points2d, map)
	return [(t[1], t[2], t[3]) for t in tri]
end
# Polyhedra interface««1
@inline polyhedra_lib(T::Type{<:Real}) =
	Polyhedra.DefaultLibrary{T}(GLPK.Optimizer)

# fixing an oversight in Polyhedra.jl: it has multiplications but no
# divisions
@inline Base.:(/)(h::Polyhedra.HyperPlane, α::Real) =
	Polyhedra.HyperPlane(h.a / α, h.β / α)

# HRepElement is the supertype of HalfSpace and HyperPlane
@inline direction(h::Polyhedra.HRepElement) = h.a
@inline function normalize(h::Polyhedra.HRepElement)
	n = norm(direction(h))
	(n ≠ 0) ? (h / n) : h
end
@inline (h::Polyhedra.HRepElement)(p) = dot(h.a, p) - h.β
@inline ∈(p, h::Polyhedra.HyperPlane) = iszero(h(p))
@inline Base.convert(T::Type{<:Polyhedra.HRepElement},
		h::Polyhedra.HRepElement) = T(h.a, h.β)

# 2d convex hull ««1

"""
    convex_hull_list(points)

Returns the convex hull of the points, as a list of indexes (in direct
order, starting at a reproducible index in the list of points).
"""
function convex_hull_list(points)
  # Uses Graham scan
  # In practice this is faster than using `Polyhedra.jl`.
#   println("points=$points, length=$(length(points))")
  i0 = findextrema(points;
    lt=(p,q)->(p[2]<q[2])|| (p[2]==q[2] && p[1]>q[1])).min[2]
  @inline detp2(i,j,k) = det2(points[[i,j,k]]...)
	# 1024 is an arbitrary value for detecting “aligned” points (i.e. up to
	# representation errors), which should be fast for both Float and Fixed
	# division
  @inline function are_aligned(i,j,k)
    v1 = points[j]-points[i]
    v2 = points[k]-points[j]
    d = det2(v1, v2)
    c = dot(v1, v2)
    return abs(d) < abs(c)/1024
  end
  scan = sort(filter(!isequal(i0), eachindex(points)),
    lt=(p,q)->detp2(i0,p,q) > 0)
#   println("i0=$i0, scan=$scan")
  stack = [i0, scan[1]]
  for h in scan[2:end]
#     println("scanning: $stack + $h")
    v1 = points[stack[end]] - points[stack[end-1]]
    v2 = points[h] - points[stack[end]]
    s = det2(v1, v2)
    c = dot(v1, v2)
    if abs(s) < abs(c)/1024 && c < 0 # points are aligned and backwards
			# here we know that we can insert at (end)
			# look for an insertion point i:
			i = length(stack)
			while i > 2
# 				println(" try to insert at $i")
				v1 = points[stack[i]] - points[stack[i-1]]
				v2 = points[h] - points[stack[i]]
				s = det2(v1, v2)
				c = dot(v1, v2)
				if s < -1e-3*abs(c)
# 					println(" break at $i")
					break
				end
				i -= 1
			end
# 			println("  inserting at $i")
			insert!(stack, i, h)
			continue
# 			println("  now stack=$stack")
    end
    while detp2(last(stack), h, stack[end-1]) < 0
      pop!(stack)
    end
    push!(stack, h)
  end
  return stack
end

"""
    convex_hull([vector of 2d points])

Returns the convex hull (as a vector of 2d points, ordered in direct
order).
"""
function convex_hull(points::AbstractVector{<:StaticVector{2}})
	upoints = unique(sort(points))
	return upoints[convex_hull_list(upoints)]
end

# """
#     convex_hull(x::Geometry{3}...)
#
# Returns the convex hull of the union of all the given solids, as a
# pair `(points, faces)`. `faces` is a list of triangles.
# """
# @inline convex_hull(x::AbstractGeometry{3}) =
# 	convex_hull(vcat([vertices(y) for y in x]...))

# this is the version using MiniQhull: #««
# convex_hull(points::AbstractVector{<:AnyVec(2)}) =
#		# Delaunay triangulation:
#		let T = MiniQhull.delaunay([p[i] for i in 1:2, p in points]),
#				N = length(points),
#				M = zeros(Bool, N, N) # ⚠ memory O(N²)
#		# mark all edges:
#		for (a,b,c) in eachcol(T)
#			b1 = points[b] - points[a]
#			c1 = points[c] - points[a]
#			d = b1[1]*c1[2] - b1[2]*c1[1] # determinant === orientation of triangle
#			if d < 0
#				M[a,b] = M[b,c] = M[c,a] = true
#			else
#				M[b,a] = M[c,b] = M[a,c] = true
#			end
#		end
#		# list of remaining edges (retrograde oriented)
#		L= sort([(i,j) for i in 1:N, j in 1:N if M[i,j] && ! M[j,i]], by=v->v[1])
#		next(i) = L[searchsorted(L, i, by=y->y[1])][1][2]
#		R = zeros(Int, length(L))
#		R[1:2] .= L[1] # initialize with first edge
#		for i in 3:length(R)
#			R[i] = next(R[i-1])
#		end
#		# returns in retrograde ordering (OpenSCAD convention):
#		points[R]
# end#»»
# this is the version using Polyhedra:««
# """
#     convex_hull([vector of 2d points])
#
# Returns the convex hull (as a vector of 2d points).
# """
# function convex_hull(points::Union{AbstractVector{<:Vec{2}},
# 	AbstractMatrix{<:Number}})
# # M is a matrix with the points as *columns*, hence the transpose
# # below:
# 	PH = Polyhedra
# 	poly = vpoly(points)
# 	PH.removevredundancy!(poly)
# 	return Vec{2}.(collect(PH.points(poly)))
# end
# »»

# 3d convex hull««1
"""
    convex_hull(vector of 3d points)

Returns the convex hull of these points, as a pair `(points, faces)`.
All the faces are triangles.
"""
function convex_hull(p::AbstractVector{<:StaticVector{3,T}}) where{T}
	M = reduce(hcat, p)
	PH = Polyhedra
	poly = PH.polyhedron(PH.vrep(transpose(M)), polyhedra_lib(T))
	R = PH.removevredundancy!(poly)
	V = SVector{3,T}.(collect(PH.points(poly)))

	triangles = SVector{3,Int}[]
	for i in PH.eachindex(PH.halfspaces(poly)) # index of halfspace
		h = PH.get(poly, i)
		pts = PH.incidentpointindices(poly, i) # vector of indices of points
		vlist = [SVector{3,T}(PH.get(poly, j)) for j in pts]
		for t in triangulate_face( vlist;
				direction = h.a,
				map = [j.value for j in pts])
			(a,b,c) = (V[j] for j in t)
			k = det([b-a c-a h.a])
			push!(triangles, (k > 0) ? t : (t[1], t[3], t[2]))
		end
	end
	return (points=V, faces=triangles)
end

#»»1
export convex_hull
end # module


using Makie, ZigZagBoomerang, SparseArrays, LinearAlgebra
# using CairoMakie
include("faberschauder.jl")

# Drift
const α = 1.0
const β = 0.1
b(x) = -β*x + α*sin(2pi*x)
# First derivative
b′(x) = -β + 2*α*2pi*cos(2pi*x)
# Second derivative
b″(x) = -2*α*(2pi)^2*sin(2pi*x)

L = 7
n = (2 << L) + 1
T = 2.0 # length diffusion bridge
ξ0 = 0randn(n)
u, v = 1.5, -0.5  # initial and fianl point
ξ0[1] = u/sqrt(T)
ξ0[end] = v/sqrt(T)
c = ones(n)
c[end] = c[1] = 0.0
θ0 = rand((-1.0, 1.0), n)
θ0[end] = θ0[1] =  0.0 # fix final point
T′ = 2000.0 # final clock of the pdmp

Γ = sparse(1.0I, n, n)
#trace, (t, ξ, θ), (acc, num) = @time pdmp(∇ϕ!, 0.0, ξ0, θ0, T, 10.0, Boomerang(Γ, ξ0*0, 0.1; ρ=0.9), 1, L, adapt=false);
#trace, (t, ξ, θ), (acc, num) = @time pdmp(∇ϕ, 0.0, ξ0, rand((-1.0, 1.0), n), T, 40.0*ones(n), ZigZag(Γ, ξ0*0), 5, L, adapt=false);
trace, (t, ξ, θ), (acc, num), c = @time spdmp(∇ϕmoving, 0.0, ξ0,
    θ0, T′, c, ZigZag(Γ, ξ0*0), SelfMoving(), L, T, adapt=true);
#trace, (t, ξ, θ), (acc, num) = @time pdmp(∇ϕ, 0.0, ξ0, rand((-1.0,1.0), n), T, 100.0*ones(n), FactBoomerang(Γ, ξ0*0, 0.1), 5, L, adapt=false);
ts, ξs = splitpairs(discretize(trace, T′/n))
S = T*(0:n)/(n+1)

#using CairoMakie
p1 = lines(S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
for ξ in ξs[1:5:end]
    lines!(p1, S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
end
display(p1)

p3 = hbox([lines(ts, getindex.(ξs, i)) for i in [1,2,4,8,16,(n+1)÷2]]...)

save("figures/diffbridges.png", p1)
p0 = vbox(p1, p3)


using AssociativeWindowAggregation
using BenchmarkTools
using RollingFunctions

function run_example(window::Integer, data::Vector{Int})
    state = FixedWindowAssociativeOp{Int,+}(window)
    result = Vector{Int}(undef, length(data) - (window - 1))
    i = 1
    for x in data
        update_state!(state, x)  # Do work
        window_full(state) || continue  # Only record numbers when the window is full
        @inbounds result[i] = window_value(state)  # Extract mean
        i += 1
    end
    return result
end

run_rolling(window::Integer, data::Vector{Int}) = rolling(sum, data, window)

data = rand(-100:100, 100000)
window = 200

@assert run_example(window, data) == run_rolling(window, data)

@benchmark run_example(window, data)
@benchmark run_rolling(window, data)

# Observations on an M1 MacBook Air, running Julia 1.6.2:
#
# Currently this package is faster for windows >= 20, but is slower for smaller windows.
# This is because we have better complexity [O(1) vs O(n) for RollingFunctions], however we
# incur some additional overhead that results in worse constant factors.


using FastAI
using Flux: mse, @functor, Dense, Descent, train!, Params, gradient, update!
using Flux.Data: DataLoader

"A simple DataBunch where `x` is random and `y = a*x + b` plus some noise."
function test_dbunch(a=2, b=3, bs=16, n_train=10, n_valid=2)

    function get_data(n)
        xy = []
        for i in 1:bs*n
            x = rand()
            y = a*x .+ b + 0.1*rand()
            push!(xy,([x],[y]))
        end
        return DataLoader(xy, batchsize=bs, shuffle=true)
    end

    train_dl = get_data(n_train)
    valid_dl = get_data(n_valid)
    return DataBunch(train_dl, valid_dl)
end

struct TestCallback <: AbstractCallback end

function test_learner(n_train=10, n_valid=2, cuda=false, lr=0.01)
    data = test_dbunch() #n_train=n_train,n_valid=n_valid)
    return Learner(data, Dense(1,1), loss=mse, opt=Descent(0.001))
end

function one_batch(dl::DataLoader)
    for d in dl
        return d
    end
end

function is_close(a,b;eps=1e-5)
    #@show a
    #@show b
    return (a-b)^2 < eps^2
end

function run_learner(learn)

    xys = learn |> data_bunch |> train |> one_batch
    lf = loss(learn)
    mf = model(learn)

    function lff(xy)
        x = xy[1]
        p = mf(x)[1]
        y = xy[2]
        return lf(p,y)
    end

    init_loss = sum(lff.(xys))
    fit!(learn,16)
    final_loss = sum(lff.(xys))
    @test final_loss < init_loss
end


__precompile__()
module ComplexSystems
    using DataFrames
    using Distributions
    using GLM
    using Statistics
    using LibGEOS

    export distance,freeze,look_around,random_walk,walker_set,dla
    export neighboors,initialize_population,friends_opinion,avg_friends_opinion
    export convince,convince_2_sites,convince_3_sites,sznajd_one_site,sznajd_two_site,sznajd_three_site
    export stupid_bots, smart_bots, cambridge_analytica, gen_box, make_grid
    export inside_box, euler, heun, lv, lot3, my_integrator
    export evolve, paste_string, lsystem,mult_evol_lsystem,turtle_step,update_angle

    # An AbstractType with an empty property
    mutable struct box_2d
        squares::Polygon
        empty::Bool
    end

    # walker type, with a position and if it is a wanderer
    mutable struct walker
        x::Float64
        y::Float64
        loner::Bool
    end

    include("diffusion_limited_aggregation.jl")
    include("dynamical_system.jl")
    include("l_system.jl")
    include("box_counting.jl")
    include("CA.jl")
end # module


using ADCME
using PyPlot
matplotlib.use("agg")

n = 80
pml = 30
N = n + 2pml + 1
NT = 1500
Δt = 1.5/NT
Δx = 1/n
x0 = LinRange(0, 1, n+1)
xE = Array((0:n+2pml)*Δx .- pml*Δx)
xH = (xE[2:end]+xE[1:end-1])/2
C = 1000.0

## PML
σE = zeros(N)
for i = 1:pml
    d = i*Δx
    σE[pml + n + 1 + i] = C* (d/(pml*Δx))^3
    σE[pml+1-i] = C* (d/(pml*Δx))^3
end

σH = zeros(N-1)
for i = 1:pml
    d = (i-1/2)*Δx
    σH[pml + n + i] = C* (d/(pml*Δx))^3
    σH[pml+1-i] = C* (d/(pml*Δx))^3
end

## Source
function ricker(dt = 0.002, f0 = 3.0)
    nw = 2/(f0*dt)
    nc = floor(Int, nw/2)
    t = dt*collect(-nc:1:nc)
    b = (π*f0*t).^2
    w = @. (1 - 2b)*exp(-b)
    # w = @. exp(-b)
end
R = ricker() * 1e7
if length(R)<NT+1
    R = [R;zeros(NT+1-length(R))]
else
    R = R[1:NT+1]
end
R_ = constant(R)

idx_src = pml + Int(round(0.0/Δx))+1
Z = zeros(N)
Z[idx_src] = 1.0
Z = Z[2:end-1]

idx_rcv = [pml + Int(round(0.0/Δx))+2, pml + Int(round(1.0/Δx))+1]

## Anomaly
model = "step"
cH = ones(length(xH)) * 3.0
cH_init = copy(cH)
cH[pml + Int(round(0.6/Δx)):end] .= 1.0
# model = "slop"
# cH = ones(length(xH)) * 2.0
# cH_init = copy(cH)
# cH[pml + Int(round(0.3/Δx)):end] = range(2.0, 3.0, length=length(cH[pml + Int(round(0.3/Δx)):end]))

cH_true = copy(cH)
cE = (cH[1:end-1]+cH[2:end])/2

## Similation
function condition(i, E_arr, H_arr)
    i<=NT+1
end

function body(i, E_arr, H_arr)
    E = read(E_arr, i-1)
    H = read(H_arr, i-1)

    # second order
    # ΔH = cH * (E[2:end]-E[1:end-1])/Δx - σH*H
    # H += ΔH * Δt
    # ΔE = cE * (H[2:end]-H[1:end-1])/Δx - σE[2:end-1]*E[2:end-1] + R_[i] * Z
    # E = scatter_add(E, 2:N-1, ΔE * Δt)

    ## forth order
    ΔH = cH * (E[2:end]-E[1:end-1])/Δx - σH*H
    H += ΔH * Δt
    ΔH = 1/(24Δx) * cH[2:N-2] * (-E[4:end] + 3E[3:end-1] - 3E[2:end-2] + E[1:end-3])
    H = scatter_add(H, 2:N-2, ΔH * Δt)
    ΔE = cE * (H[2:end]-H[1:end-1])/Δx - σE[2:end-1]*E[2:end-1] + R_[i] * Z
    E = scatter_add(E, 2:N-1, ΔE * Δt)
    ΔE = 1/(24Δx) * cE[2:end-1] * (-H[4:end] + 3H[3:end-1] - 3H[2:end-2] + H[1:end-3])
    E = scatter_add(E, 3:N-2, ΔE*Δt)

    i+1, write(E_arr, i, E), write(H_arr, i, H)
end

E_arr = TensorArray(NT+1)
H_arr = TensorArray(NT+1)

E_arr = write(E_arr, 1, zeros(N))
H_arr = write(H_arr, 1, zeros(N-1))

i = constant(2, dtype = Int32)

_, E, H = while_loop(condition, body, [i, E_arr, H_arr])

E = stack(E)
H = stack(H)

sess = Session(); init(sess)
E_true, H_true = run(sess, [E, H])

## data
figure(figsize = (14,4))
subplot(132)
pcolormesh(xE, (0:NT)*Δt, E_true, cmap="seismic", vmax=maximum(E_true), vmin=-maximum(E_true), rasterized=true)
plot(xE[idx_src]*ones(NT+1), (0:NT)*Δt, "r", label="Source", alpha=0.8)
plot(xE[idx_rcv[1]]*ones(NT+1), (0:NT)*Δt, "--", label="Receiver #1", alpha=0.8)
plot(xE[idx_rcv[2]]*ones(NT+1), (0:NT)*Δt, "--", label="Receiver #2", alpha=0.8)
autoscale(enable=true, axis="x", tight=true)
legend(loc="upper center")
xlabel("Localtion (km)")
ylabel("Time (s)")
subplot(133)
plot((0:NT)*Δt, E_true[:,idx_rcv[1]], label="Receiver #1", alpha=0.8)
plot((0:NT)*Δt, E_true[:,idx_rcv[2]], label="Receiver #2", alpha=0.8)
legend()
gca().ticklabel_format(axis="y", style="sci", scilimits=(0,0))
autoscale(enable=true, axis="x", tight=true)
xlabel("Time (s)")
ylabel("Amplitude")
subplot(131)
plot(xH, cH_true)
autoscale(enable=true, axis="x", tight=true)
xlabel("Location (km)")
ylabel("Velocity (km/s)")
tight_layout()
savefig("data.png", bbox_inch="tight")
savefig("data.pdf", bbox_inch="tight", dpi=300)


using DataStructures

#mutable struct Node
#    ip::Int
#    neighbours::Array{Tuple{Node,Int}}
#    risk::Union{Float64, Nothing}
#    predecessor::Union{Node, Nothing}
#    probability::Float64
#    Node() = (obj = new(0);)
#end

function create_network()
    nodes = []
    for i = 1:4
        push!(nodes, Node() )
        nodes[i].ip = i
    end

    nodes[1].probability = 1.0
    nodes[1].risk = 0
    nodes[1].neighbours = [(nodes[2], 6), (nodes[3], 3)]

    nodes[2].probability = 0.5
    nodes[2].risk = 12
    nodes[2].neighbours = [(nodes[4], 3)]

    nodes[3].probability = 0.75
    nodes[3].risk = 4
    nodes[3].neighbours = [(nodes[4], 4)]

    nodes[4].probability = 0.8
    nodes[4].risk = 9

    return nodes
end

function initialize_single_source!(graph, start)
    for v in graph
        v.risk = typemax(Float64)
        v.predecessor = nothing
    end
    start.risk = 0
end

function relax!(from_node,to_node,cost)
    tot_risk = from_node.risk + (cost/to_node.probability)
    if (tot_risk < to_node.risk)
        to_node.risk = tot_risk
        to_node.predecessor = from_node
    end
end

function extract_min(list)
    min = typemax(Float64)
    ret_node = list[1]
    delete_idx = 1
    for i = 1:length(list)
        if list[i].risk < min
            ret_node = list[i]
            delete_idx = i
        end
    end
    deleteat!(list, delete_idx)
    return ret_node
end

function dijkstra!(graph,start)
    initialize_single_source!(graph, start)
    pq = PriorityQueue()                  # create a priority que
    enqueue!(pq, start, start.risk)
    while !isempty(pq)
        from_node = dequeue!(pq)
        for to_node in from_node.neighbours
            relax!(from_node, to_node[1], to_node[2])
            for node in pq # check if node already exists in que
                if node[1] == to_node[1] # if so change it
                    pq[to_node[1]] = to_node[1].risk # priority changed to new risk
                end
                println(node[to_node[1]])
            end
            enqueue!(pq, to_node[1], to_node[1].risk)
        end
    end
end

test = PriorityQueue()
network = create_network()
enqueue!(test, network[1].neighbours[1], network[1].neighbours[1][1].risk)
enqueue!(test, network[1].neighbours[2], network[1].neighbours[2][1].risk)

for i in test
    println(i[1])
    if i[1] == network[1].neighbours[2]
        #println(i[1])
    end
end


__precompile__()

module BilevelTests

export test_bilevel_learn

using Printf
using SparseArrays
using FileIO
using ColorTypes: Gray
import ColorVectorSpace
using ImageContrastAdjustment

using AlgTools.Util
using AlgTools.LinOps
using AlgTools.LinkedLists
using ImageTools.Visualise

using VariationalImaging

using VariationalImaging.Util
using VariationalImaging.GradientOps
using VariationalImaging.OpDenoise

using VariationalImaging.TestDatasets
using VariationalImaging.Bilevel

# Parameters
const default_save_prefix="bilevel_result_"

const default_params = (
    verbose_iter = 1,
    maxiter = 20,
    save_results = true,
    dataset_name = "cameraman128_5",
    save_iterations = false
)

const bilevel_params = (
    η₁ = 0.25,
    η₂ = 0.75,
    β₁ = 0.25,
    β₂ = 2.0,
    Δ₀ = 1.0,
    x₀ = 0.001
)

function save_results(params, b, b_data, x, opt_img,st)
    if params.save_results
        perffile = params.save_prefix * ".txt"
        println("Saving " * perffile)
        write_log(perffile, st.log, "# params = $(params), x = $x\n")
        fn = (t, ext) -> "$(params.save_prefix)_$(t).$(ext)"
        save(File(format"PNG", fn("true", "png")), grayimg(b))
        save(File(format"PNG", fn("data", "png")), grayimg(b_data))
        save(File(format"PNG", fn("reco", "png")), grayimg(opt_img))
    end
end

###############
# Learning function
###############
function learning_function(x,data)
    op = FwdGradientOp()
    u = denoise(data[2],x,op)
    cost = 0.5*norm₂²(u-data[1])
    grad = gradient(x,u,data[1],op)
    return u,cost,grad
end

function denoise(data,x,op::LinOp)
    denoise_params = (
        ρ = 0,
        α = x,
        op=op,
        # PDPS
        τ₀ = 5,
        σ₀ = 0.99/5,
        accel = true,
        save_results = false,
        maxiter = 1000,
        verbose_iter = 1001,
        save_iterations = false
    )
    st_opt, iterate_opt = initialise_visualisation(false)
    opt_img = op_denoise_pdps(data; iterate=iterate_opt, params=denoise_params)
    finalise_visualisation(st_opt)
    return opt_img
end

function gradient(α,u::AbstractArray{T,3},ū::AbstractArray{T,3},op::LinOp) where T

	M,N,O = size(u)
	grad = 0
	for i = 1:O
		u1 = @view u[:,:,i]
		u2 = @view ū[:,:,i]
		g = gradient(α,u1,u2,op)
		grad += g
	end
	return grad
end

function gradient(α,u::AbstractArray{T,2},ū::AbstractArray{T,2},op::LinOp) where T
	u = Float64.(Gray{Float64}.(u))
	ū = Float64.(Gray{Float64}.(ū))
	# Obtain Active and inactive sets
	n = size(u,1)

	# Generate centered gradient matrix
    G = matrix(op,n)
	#G = createDivMatrix(n)
	Gu = G*u[:]
	nGu = xi(Gu)
	act = nGu .< 1e-12
	inact = 1 .- act
	Act = spdiagm(0=>act)
	Inact = spdiagm(0=>inact)

	# Vector with grad norm in inactive components and one in the active
	den = Inact*nGu+act
	Den = spdiagm(0=>1 ./den)

	# prod KuKuᵗ/norm³
	prodKuKu = prodesc(Gu ./den.^3,Gu)

	Adj = [spdiagm(0=>ones(n^2)) α*G';
			Act*G+Inact*(prodKuKu-Den)*G Inact+eps()*Act]

	Track=[(u[:]-ū[:]);zeros(2*n^2)]
	mult = Adj\Track
	p = @view mult[1:n^2]
	return -p'*(G'*Inact*Den*Gu)
end

###############
# Bilevel learn test
###############

function test_bilevel_learn(;
    visualise=true,
    save_prefix=default_save_prefix,
    kwargs...)

    # Parameters for this experiment
    params = default_params ⬿ bilevel_params ⬿ kwargs
    params = params ⬿ (save_prefix = save_prefix * "denoise_" * params.dataset_name,)

    # Load dataset
    b,b_noisy = TestDatasets.testdataset(params.dataset_name)
    b = Float64.(Gray{Float64}.(b))
    b_noisy = Float64.(Gray{Float64}.(b_noisy))

    # Launch (background) visualiser
    st, iterate = initialise_visualisation(visualise)

    # Run algorithm
    x, u, st = bilevel_learn((b,b_noisy),learning_function; xinit=params.x₀,iterate=iterate, params=params)

    adjust_histogram!(u,LinearStretching())

    save_results(params, b, b_noisy, x, u, st)

    # Exit background visualiser
    finalise_visualisation(st)
end

end


function recursive_backtracker!(grid, start_pos=(1, 1))
    start_at = grid[start_pos...]
    stack = [start_at]
    while !isempty(stack)
        current = stack[end]
        nbs = filter(x -> isempty(x.connections), neighbors(grid, current))
        if isempty(nbs)
            pop!(stack)
        else
            neighbor = rand(nbs)
            connect!(current, neighbor)
            push!(stack, neighbor)
        end
    end
    return grid
end


# ===== VSB Verifications: Pressure Distribution =====

println("=== BEGINNING Verifications1.jl ===")

# == Imports ==
println("   IMPORTING PACKAGES...")
using VSB
using SimpleVPM
using LinearAlgebra
using Plots
pyplot()
using LaTeXStrings
using Printf

cd("/Users/Damyn/Documents/BYU/FLOW Lab/VSB/verifications")

# === Boundary Geometry and Plotting Grids ===
println("   GENERATING GEOMETRY...")
NPTS = 150                              # Number of Points on Boundary
del_theta = (2*pi)/(NPTS - 1)           # Step in theta around body
R = 1.0                                 # Cylinder Radius
X_coor = zeros(NPTS)                    # X-coordinates
Y_coor = zeros(NPTS)                    # Y-coordinates
X_tHat = zeros(NPTS)                    # X Tangent Vector
Y_tHat = zeros(NPTS)                    # X Tangent Vector
X_nHat = zeros(NPTS)                    # Y Normal Vector
Y_nHat = zeros(NPTS)                    # Y Normal Vector
theta = zeros(NPTS)                     # Angle from positive X-axis
for i=1:NPTS
    X_coor[i] = R*cos((i-1)*del_theta)
    Y_coor[i] = R*sin((i-1)*del_theta)

    X_tHat[i] = sin((i-1)*del_theta)
    Y_tHat[i] = -cos((i-1)*del_theta)

    X_nHat[i] = cos((i-1)*del_theta)
    Y_nHat[i] = sin((i-1)*del_theta)

    theta[i] = (i-1)*del_theta
end
body_pts = [[X_coor[i], Y_coor[i], 0.0] for i in 1:NPTS]
t_hats = [[X_tHat[i], Y_tHat[i], 0.0] for i in 1:NPTS]
n_hats = [[X_nHat[i], Y_nHat[i], 0.0] for i in 1:NPTS]
body = VSB.Boundary(body_pts,t_hats,n_hats)

# === CP: Analytical ===
println("   CALCULATING CP_ANALYTICAL...")
CP_analytical = [1 - 4*sin(theta[i])^2 for i in 1:NPTS]

# === CP: VSB ===
println("   CALCULATING CP_VSB...")
println("   CALCULATING PROBLEM PARAMETERS...")

U_inf = [1.0, 0.0, 0.0]
U_inf_func(X) = U_inf

params = VSB.Parameters(body, U_inf_func, verbose = true)
alphas = params.alphas

U_gamma(X) = VSB.CalcVSVelocityCirlce(body, alphas, X)

CP_VSB = zeros(NPTS)
for i in 1:NPTS
    if i % 10 == 0
        println("     - ON POINT ",i," OF ",NPTS)
    end
    scale_factor = 0.000
    X_eval = body.bodyPTS[i] + scale_factor .* body.nHats[i]
    CP_VSB[i] = 1 - (norm( U_gamma(X_eval)) / norm( U_inf_func(X_eval) ))^2
end

# === CP: VPM ===
println("   CALCULATING CP_VPM...")
CP_VPM = zeros(NPTS)

nu = 1e1
Uinf(t) = U_inf
sigma = 1.0

print("     - CREATING PARTICLE FIELD...")
pfield = SimpleVPM.ParticleField(nu,                 # Kinematic viscosity
                                 zeta_gauserf,       # Basis function
                                 g_gauserf;          # Regularizing function of zeta
                                 transposed=false,    # Transposed scheme
                                 relax=!true,        # Pedrizzeti relaxation scheme
                                 rlxf=0.3,           # Relaxation factor
                                 integration="euler",# Time integration scheme
                                 Uinf=Uinf)          # Free stream velocity
print("   DONE!")
println()

print("     - ADDING PARTICLES...")
for i=1:NPTS
    loc = body.bodyPTS[i]
    Gamma = [0.0, 0.0, VSB.CalcVS(body,alphas,loc)]
    vol = 1.0
    particle = Particle(loc, Gamma, sigma, vol)
    addparticle(pfield, particle)
end
print("   DONE!")
println()

for i=1:NPTS
    if i % 10 == 0
        println("     - ON POINT ",i," OF ",NPTS)
    end
    CP_VPM[i] = -(1 - ( norm( U(pfield, body.bodyPTS[i]) ) / norm( U_inf ))^2)
end

println()
println("   RESULTS...")
println("CP_ANALYTICAL    CP_VSB    CP_VPM")
for i in 1:NPTS
    if i % 10 == 0
        println(@sprintf("%13.3f",CP_analytical[i]),"  ",
                @sprintf("%8.3f",CP_VSB[i]),"  ",
                @sprintf("%8.3f",CP_VPM[i]))
    end
end

# === Plotting ===
println("   PLOTTING...")
plt1 = plot(theta, CP_analytical, line = (:red), label = "Analytical")
plot!(theta, CP_VSB, line = (:blue, :dash), label = "VSB")
plot!(title = "VSB Verification - Pressure Distribution",
      xlabel = L"\theta",
      xlims = (0.0, 2*pi),
      xticks = ([0.0,pi/2,pi,3*pi/2,2*pi],["0",L"\frac{\pi}{2}",L"\pi",L"\frac{3\pi}{2}",L"2\pi"]),
      ylabel = L"C_P",
      ylims = (-4,2),
      yflip = true)

plt2 = plot(theta, alphas, line=(:red), title = "Alpha Coefficients")
savefig(plt1, "VSB_Verification_CP.pdf")
savefig(plt2, "VSB_Verification_Alphas.pdf")

println("=== END OF Verifications1.jl ===")


########################################################################################
# Algorithm
########################################################################################
abstract type GaussianODEFilter <: OrdinaryDiffEq.OrdinaryDiffEqAdaptiveAlgorithm end
abstract type AbstractEK <: GaussianODEFilter end

"""
    EK0(; prior=:ibm, order=3, diffusionmodel=DynamicDiffusion(), smooth=true)

**Gaussian ODE filtering with zeroth order extended Kalman filtering.**

Currently, only the integrated Brownian motion prior `:ibm` is supported.
For the diffusionmodel, chose one of
`[DynamicDiffusion(), DynamicMVDiffusion(), FixedDiffusion(), FixedMVDiffusion()]`.

See also: [`EK1`](@ref)

# References:
- N. Bosch, P. Hennig, F. Tronarp: **Calibrated Adaptive Probabilistic ODE Solvers** (2021)
- F. Tronarp, H. Kersting, S. Särkkä, and P. Hennig: **Probabilistic Solutions To Ordinary Differential Equations As Non-Linear Bayesian Filtering: A New Perspective** (2019)
- M. Schober, S. Särkkä, and P. Hennig: **A Probabilistic Model for the Numerical Solution of Initial Value Problems** (2018)
"""
Base.@kwdef struct EK0{DT,IT} <: AbstractEK
    prior::Symbol = :ibm
    order::Int = 3
    diffusionmodel::DT = DynamicDiffusion()
    smooth::Bool = true
    initialization::IT = TaylorModeInit()
end

"""
    EK1(; prior=:ibm, order=3, diffusionmodel=DynamicDiffusion(), smooth=true)

**Gaussian ODE filtering with first order extended Kalman filtering.**

Currently, only the integrated Brownian motion prior `:ibm` is supported.
For the diffusionmodel, chose one of
`[DynamicDiffusion(), DynamicMVDiffusion(), FixedDiffusion(), FixedMVDiffusion()]`.

See also: [`EK0`](@ref)

# References:
- N. Bosch, P. Hennig, F. Tronarp: **Calibrated Adaptive Probabilistic ODE Solvers** (2021)
- F. Tronarp, H. Kersting, S. Särkkä, and P. Hennig: **Probabilistic Solutions To Ordinary Differential Equations As Non-Linear Bayesian Filtering: A New Perspective** (2019)
"""
struct EK1{CS,AD,DiffType,DT,IT} <: AbstractEK
    prior::Symbol
    order::Int
    diffusionmodel::DT
    smooth::Bool
    initialization::IT
end
EK1(;
    prior=:ibm,
    order=3,
    diffusionmodel=DynamicDiffusion(),
    smooth=true,
    initialization=TaylorModeInit(),
    chunk_size=0,
    autodiff=true,
    diff_type=Val{:forward},
) = EK1{chunk_size,autodiff,diff_type,typeof(diffusionmodel),typeof(initialization)}(
    prior,
    order,
    diffusionmodel,
    smooth,
    initialization,
)

Base.@kwdef struct EK1FDB{DT,IT} <: AbstractEK
    prior::Symbol = :ibm
    order::Int = 3
    diffusionmodel::DT = DynamicDiffusion()
    smooth::Bool = true
    initialization::IT = TaylorModeInit()
    jac_quality::Int = 1
end


module Dual

using LinearAlgebra
using LogExpFunctions
import OptimalTransport.add_singleton
import OptimalTransport.dot_vecwise

export ot_entropic_semidual
export ot_entropic_semidual_grad
export getprimal_ot_entropic_semidual
export ot_entropic_dual
export ot_entropic_dual_grad
export getprimal_ot_entropic_dual

function logKexp(logK::AbstractMatrix, x::AbstractVector)
    return logsumexp(logK .+ add_singleton(x, Val(1)); dims=2)
end

function logKexp(logK::AbstractMatrix, x::AbstractMatrix; dims)
    return logsumexp(logK + x; dims=dims)
end

"""
    ot_entropic_semidual(μ, v, eps, K)

Computes the semidual (in the second argument) of the entropic optimal transport loss, with source marginal `μ`, regularization parameter `ε`, and Gibbs kernel `K`.

That is,
```math
    \\operatorname{OT}_{\\varepsilon}(\\mu, \\nu) = \\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle + \\varepsilon \\Omega(\\gamma)
```
with ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{ij} \\log \\gamma_{ij}``, then the semidual in the second argument ν is [^Z21]
```math
\\begin{aligned}
    \\operatorname{OT}_{\\varepsilon}^*(\\mu, v) &= \\sup_{\\nu} \\langle v, \\nu \\rangle - \\operatorname{OT}_{\\varepsilon}(\\mu, \\nu) \\ \\
    &= -\\varepsilon \\left\\langle \\mu, \\log\\left( \\dfrac{\\mu}{K e^{v/\\varepsilon}} \\right) - 1\\right\\rangle.
\\end{aligned}
```
Notably, the semidual is computationally advantageous for solving variational problems since it is a smooth and unconstrained function of `v` since it admits a closed form gradient. See [^CP16] for a detailed discussion of dual methods for variational problems in optimal transport.

[^CP16]: Cuturi, Marco, and Gabriel Peyré. "A smoothed dual approach for variational Wasserstein problems." SIAM Journal on Imaging Sciences 9.1 (2016): 320-343.
[^Z21]: Zhang, Stephen Y. “A Unified Framework for Non-Negative Matrix and Tensor Factorisations with a Smoothed Wasserstein Loss.” ArXiv: Machine Learning, 2021.
"""
function ot_entropic_semidual(μ, v, eps, K)
    return eps *
           (-dot_vecwise(xlogx.(μ) .- μ, one.(μ)) + dot_vecwise(μ, log.(K * exp.(v / eps))))
end

"""
    ot_entropic_semidual_grad(μ, v, eps, K)

Computes the gradient with respect to `v` of the semidual of the entropic optimal transport loss. That is,
```math
\\nabla_v \\operatorname{OT}^*_{\\varepsilon}(\\mu, v) = K^\\top \\left( \\dfrac{\\mu}{K e^{v/\\varepsilon}} \\right) \\odot e^{v/\\varepsilon}.
```

See also: [`ot_entropic_semidual`](@ref)
"""
function ot_entropic_semidual_grad(μ, v, eps, K)
    return K' * (μ ./ (K * exp.(v / eps))) .* exp.(v / eps)
end

"""
    getprimal_ot_entropic_semidual(μ, v, eps, K)

Computes the the primal variable `ν` corresponding to the dual variable `v` at optimality. That is,
```math
\\nu^\\star = e^{v^\\star/\\varepsilon} \\odot K^\\top \\dfrac{\\mu}{K e^{v^\\star/\\varepsilon}}.
```

See also: [`ot_entropic_semidual`](@ref)
"""
function getprimal_ot_entropic_semidual(μ, v, eps, K)
    return exp.(v / eps) .* (K' * (μ ./ (K * exp.(v / eps))))
end

"""
    ot_entropic_dual(u, v, eps, K)

Computes the dual in both arguments of entropic optimal transport loss, where `u` and `v` are the dual variables associated with the source and target marginals respectively.

That is,
```math
    \\begin{aligned}
    \\operatorname{OT}_{\\varepsilon}^*(u, v) &= \\sup_{\\mu, \\nu} \\langle u, \\mu \\rangle + \\langle v, \\nu \\rangle - \\operatorname{OT}_\\varepsilon(\\mu, \\nu) \\ \\
    &= \\varepsilon \\log \\langle e^{u/\\varepsilon}, K e^{v/\\varepsilon} \\rangle.
    \\end{aligned}
```
"""
function ot_entropic_dual(u, v, eps, K)
    # (μ, ν) → min_{γ ∈ Π(μ, ν)} ε H(γ | K)
    # has Legendre transform
    # (u, v) → ε log < exp(u/ε), K exp(v/ε) >
    return eps * log.(dot_vecwise(exp.(u / eps), K * exp.(v / eps)))
end

"""
    ot_entropic_dual_grad(u, v, eps, K)

Computes the gradient with respect to `u` and `v` of the dual of the entropic optimal transport loss. That is,
```math
\\begin{aligned}
\\nabla_u \\operatorname{OT}^*_{\\varepsilon}(u, v) &= \\dfrac{e^{u/\\varepsilon} \\odot K e^{v/\\varepsilon}}{\\langle e^{u/\\varepsilon}, K e^{v/\\varepsilon} \\rangle} \\ \\
\\nabla_v \\operatorname{OT}^*_{\\varepsilon}(u, v) &= \\dfrac{e^{v/\\varepsilon} \\odot K^\\top e^{u/\\varepsilon}}{\\langle e^{v/\\varepsilon}, K^\\top e^{u/\\varepsilon} \\rangle}.
\\end{aligned}
```

See also: [`ot_entropic_dual`](@ref)
"""
function ot_entropic_dual_grad(u, v, eps, K)
    U = exp.(u / eps)
    V = exp.(v / eps)
    grad_u = (U .* (K * V)) ./ dot_vecwise(U, K * V)'
    grad_v = (V .* (K' * U)) ./ dot_vecwise(V, K' * U)'
    return grad_u, grad_v
end

"""
    getprimal_ot_entropic_dual(u, v, eps, K)

Computes the the primal variable `γ` corresponding to the dual variable `u, v` at optimality. That is,
```math
    \\gamma = \\operatorname{softmax}(\\mathrm{diag}(e^{u/\\varepsilon}) K \\mathrm{diag}(e^{v/\\varepsilon}))
```

See also: [`ot_entropic_dual`](@ref)
"""
function getprimal_ot_entropic_dual(u, v, eps, K)
    γ = K .* add_singleton(exp.(-u / eps), Val(2)) .* add_singleton(exp.(-v / eps), Val(1))
    return γ
end

end


module ContComplianceDemo1

using Makie, TopOpt, LinearAlgebra, StatsFuns
using TopOptMakie

E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0 # downward force
rmin = 4.0 # filter radius
xmin = 0.0001 # minimum density
problem_size = (60, 20)
V = 0.5 # maximum volume fraction
p = 4.0 # penalty
x0 = fill(V, prod(problem_size)) # initial design

problem = HalfMBB(Val{:Linear}, problem_size, (1.0, 1.0), E, v, f)

solver = FEASolver(Displacement, Direct, problem, xmin = xmin)
cheqfilter = DensityFilter(solver, rmin = rmin)
comp = TopOpt.Compliance(problem, solver)

function obj(x)
    # minimize compliance
    return comp(cheqfilter(x))
end
function constr(x)
    # volume fraction constraint
    return sum(cheqfilter(x)) / length(x) - V
end

m = Model(obj)
addvar!(m, zeros(length(x0)), ones(length(x0)))
Nonconvex.add_ineq_constraint!(m, constr)

options = MMAOptions(
    maxiter=1000, tol = Tolerance(kkt = 1e-4, x = 1e-4, f = 1e-4),
)
TopOpt.setpenalty!(solver, p)
# Method of Moving Asymptotes
@time r = Nonconvex.optimize(m, MMA87(), x0, options = options);

@show obj(r.minimizer)
@show constr(r.minimizer)
topology = cheqfilter(r.minimizer);
fig = visualize(problem, solver.u;
    topology = topology, default_exagg_scale=0.0, scale_range=10.0)
Makie.display(fig)

end


saa_model = @stochastic_model begin
    @stage 1 begin
        @variable(model, x >= 0)
    end
    @stage 2 begin
        @decision x
        @uncertain ξ
        @variable(model, y)
        @objective(model, Min, y)
        @constraint(model, y == x)
        @constraint(model, y >= ξ)
    end
end


using SumOfSquares
using DynamicPolynomials
using MosekTools

# Create symbolic variables (not JuMP decision variables)
@polyvar x1 x2

# Create a Sum of Squares JuMP model with the Mosek solver
model = SOSModel(Mosek.Optimizer)

# Create a JuMP decision variable for the lower bound
@variable(model, γ)

# f(x) is the Goldstein-Price function
f1 = x1+x2+1
f2 = 19-14*x1+3*x1^2-14*x2+6*x1*x2+3*x2^2
f3 = 2*x1-3*x2
f4 = 18-32*x1+12*x1^2+48*x2-36*x1*x2+27*x2^2

f = (1+f1^2*f2)*(30+f3^2*f4)

# Constraints f(x) - γ to be sum of squares
@constraint(model, f >= γ)

@objective(model, Max, γ)

optimize!(model)

# The lower bound found is 3
println(objective_value(model))


module _Machin

using Brainstorm._Math._Series._Arctan: arctanhsum

# http://numbers.computation.free.fr/Constants/Log2/log2.html
ln2_N2a(digits::Integer) = arctanhsum(digits, [5,  7], [2, 2])
ln2_N2b(digits::Integer) = arctanhsum(digits, [6, 99], [4, 2])
ln2_N2c(digits::Integer) = arctanhsum(digits, [7, 17], [4, 2])

ln2_N3a(digits::Integer) = arctanhsum(digits, [11,  111, 19601], [ 8, -4, -2])
ln2_N3b(digits::Integer) = arctanhsum(digits, [17,   79,  1351], [10,  8,  4])
ln2_N3c(digits::Integer) = arctanhsum(digits, [23,   65,   485], [14,  6, -4])
ln2_N3d(digits::Integer) = arctanhsum(digits, [26, 4801,  8749], [18, -2,  8])
ln2_N3e(digits::Integer) = arctanhsum(digits, [31,   49,   161], [14, 10,  6])

ln2_N4a(digits::Integer) = arctanhsum(digits, Int64[ 52, 4801, 8749, 70226], [ 36, -2,   8,  18])
ln2_N4b(digits::Integer) = arctanhsum(digits,      [127,  449, 4801,  8749], [ 72, 54,  34, -10])
ln2_N4c(digits::Integer) = arctanhsum(digits,      [251,  449, 4801,  8749], [144, 54, -38,  62])

ln2_N5a(digits::Integer) = arctanhsum(digits, Int64[575, 3361, 8749, 13121, 56251], [342, 198, 222, 160, 106])

end


using SIAN

println("Setting up the problem")

ode = @ODEmodel(
    x1'(t) = a * x1(t) - b * x1(t) * x2(t),
    x2'(t) = -c * x2(t) + d * x1(t) * x2(t),
    y1(t) = x1(t)
)

res = res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=0, nthrds=1)

println(res)


using OptimTestProblems
using OptimTestProblems.MultivariateProblems

import LinearAlgebra: norm

using Test

verbose = false

@testset "Bounded univariate problems" begin
    uvp = OptimTestProblems.UnivariateProblems.examples
    for (name, p) in uvp
        verbose && print_with_color(:green, "Problem: $name \n")
        for (miz, mia) in zip(p.minimizers, p.minima)
            @test p.f(miz) ≈ mia
        end
    end

end

@testset "Unconstrained multivariate problems" begin
    muvp = MultivariateProblems.UnconstrainedProblems.examples
    for (name, p) in muvp
        verbose && print_with_color(:green, "Problem: $name \n")
        soltest = !any(isnan, p.solutions)

        if startswith(name, "Penalty Function I")
            # The provided solutions are not exact
            tol = 1e-16
        else
            tol = 1e-32
        end

        f = objective(p)
        soltest && @test f(p.solutions) ≈ p.minimum

        gs = similar(p.initial_x)
        g! = gradient(p)
        g!(gs, p.solutions)
        soltest && @test norm(gs, Inf) < tol

        fg! = objective_gradient(p)
        fgs = similar(gs)
        g!(gs, p.initial_x)

        @test fg!(fgs, p.initial_x) ≈ f(p.initial_x)
        @test norm(fgs.-gs, Inf)  < eps(eltype(gs))
    end
end

@testset "Constrained multivariate problems" begin
    mcvp = MultivariateProblems.ConstrainedProblems.examples
    for (name, p) in mcvp
        verbose && print_with_color(:green, "Problem: $name \n")
        soltest = all(isfinite, p.solutions)

        f = objective(p)
        soltest && @test f(p.solutions) ≈ p.minimum

        if !isempty(p.constraintdata.lx)
            @test all(p.solutions .> p.constraintdata.lx)
        end
        if !isempty(p.constraintdata.ux)
            @test all(p.solutions .< p.constraintdata.ux)
        end

        lc = p.constraintdata.lc
        if !isempty(lc)
            c = fill(eltype(lc)(0.0), size(lc))
            p.constraintdata.c!(c, p.solutions)
            @test all(c .>= p.constraintdata.lc)
            @test all(c .<= p.constraintdata.uc)
        end

        # TODO: How do we test for optimality here?
        gs = similar(p.initial_x)
        g! = gradient(p)
        g!(gs, p.solutions)
        #soltest && @test norm(gs, Inf) < tol

        fg! = objective_gradient(p)
        fgs = similar(gs)
        g!(gs, p.initial_x)

        @test fg!(fgs, p.initial_x) ≈ f(p.initial_x)
        @test norm(fgs.-gs, Inf)  < eps(eltype(gs))
    end
end


const data = readlines("day18.txt")

function evalinorder(text::AbstractString)
    withendingparens = replace(text, r"([+*] \d+)" => s"\1)")
    numparens = count(r"[+*]", text)
    withparens = '('^numparens * withendingparens
    return eval(Meta.parse(withparens))
end

function custommath(text::AbstractString)
    parensblock = r"\([^()]+?\)"
    while true
        m = match(parensblock, text)
        if m === nothing
            return evalinorder(text)
        end
        text = replace(text, parensblock => evalinorder(m.match), count=1)
    end
end

function part1(lines)
    return sum(custommath(line) for line in lines)
end
(9 + (9 * 5 * 8 * 2) + 7 + 9) + 5
println("Part1: $(part1(data))")

function evalplusfirst(text::AbstractString)
    allpluses = r"(\d+) [+] (\d+)"
    while occursin("+", text)
        m = match(allpluses, text)
        text = replace(text, m.match => sum(parse.(Int, m.captures)), count=1)
    end
    return evalinorder(text)
end

function custommath_p2(text::AbstractString)
    # println("line:\t ", text)
    parensblock = r"\([^()]+?\)"
    while true
        m = match(parensblock, text)
        if m === nothing
            res = evalplusfirst(text)
            # println("result:\t ", res)
            return res
        end
        text = replace(text, parensblock => evalplusfirst(m.match), count=1)
        # println("current: ", text)
    end
end

function part2(lines)
    return sum(custommath_p2(line) for line in lines)
end

println("Part2: $(part2(data))")

# Bonus: AST manipulation!

# The idea is to use the parser. Change + to ^ so the normal rules apply, but change it
# back to + before evaluating, this will give the good result!

# Seen on stackoverflow:
# https://stackoverflow.com/questions/51928959/how-to-find-and-replace-subexpression-of-ast-in-julia
# Here we do not care not changing the input, so no deepcopy below
function expr_replace(expr, old, new)
    expr == old && return new
    if expr isa Expr
        for i in eachindex(expr.args)
            expr.args[i] = expr_replace(expr.args[i], old, new)
        end
    end
    expr
end

function mathfromast(text::AbstractString)
    new_text = replace(text, '+' => '^')
    expr = Meta.parse(new_text)
    expr_replace(expr, :^, :+)
    return eval(expr)
end

function part2bis(lines)
    return sum(mathfromast(line) for line in lines)
end

println("Part2 with ast manipulation: $(part2bis(data))")


function pad!(v::Vector{Int},n::Int64)
    push!(v,repeat([0],n-length(v))...)
end

function pad_left_right(v::Vector{Float64},l::Int64, r::Int64; p::T = 0) where {T}
    vcat(repeat([p],l),v,repeat([p],r))
end

moving_average(time_series::Union{Vector{Float64},Vector{Int64}}, n₋::Int64, n₊::Int64) = [mean(time_series[(i-n₋):(i+n₊)]) for i in (1+n₋):(length(time_series)-n₊)]
moving_average(time_series::Union{Vector{Float64},Vector{Int64}},window::Int64) = moving_average(time_series,0,window-1)


## 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


#=
Tools for working with Markov Chains

@author : Spencer Lyon, Zac Cranko

@date: 07/10/2014

References
----------

http://quant-econ.net/jl/finite_markov.html
=#
import LightGraphs: DiGraph, period, attracting_components,
                    strongly_connected_components, is_strongly_connected

____solver_docs = """
solve x(P-I)=0 using either an eigendecomposition, lu factorization, or an
algorithm presented by Grassmann-Taksar-Heyman (GTH)

##### Arguments

- `p::Matrix` : valid stochastic matrix

##### Returns

- `x::Matrix`: A matrix whose columns contain stationary vectors of `p`

##### References

The following references were consulted for the GTH algorithm

- W. K. Grassmann, M. I. Taksar and D. P. Heyman, "Regenerative Analysis and
Steady State Distributions for Markov Chains, " Operations Research (1985),
1107-1116.
- W. J. Stewart, Probability, Markov Chains, Queues, and Simulation, Princeton
University Press, 2009.

"""

"""
Finite-state discrete-time Markov chain.

It stores useful information such as the stationary distributions, and
communication, recurrent, and cyclic classes, and allows simulation of state
transitions.

##### Fields

- `p::Matrix` The transition matrix. Must be square, all elements must be
positive, and all rows must sum to unity
"""
type MarkovChain{T<:Real}
    p::Matrix{T} # valid stochastic matrix

    function MarkovChain(p::Matrix)
        n, m = size(p)

        if n != m
            throw(ArgumentError("stochastic matrix must be square"))
        elseif any(p.<0)
            throw(ArgumentError("stochastic matrix must have nonnegative elements"))
        elseif any(x->!isapprox(sum(x), 1), [p[i, :] for i = 1:n])
            throw(ArgumentError("stochastic matrix rows must sum to 1"))
        end

        return new(p)
    end
end

# Provide constructor that infers T from eltype of matrix
MarkovChain{T<:Real}(p::Matrix{T}) = MarkovChain{T}(p)

"Number of states in the markov chain `mc`"
n_states(mc::MarkovChain) = size(mc.p, 1)

function Base.show(io::IO, mc::MarkovChain)
    println(io, "Discrete Markov Chain")
    println(io, "stochastic matrix:")
    println(io, mc.p)
end

# solve x(P-I)=0 by eigen decomposition
"$(____solver_docs)"
function eigen_solve{T}(p::Matrix{T})
    ef = eigfact(p')
    isunit = map(x->isapprox(x, 1), ef.values)
    x = real(ef.vectors[:, isunit])
    x ./= sum(x, 1) # normalisation
    for i = 1:length(x)
        x[i] = isapprox(x[i], zero(T)) ? zero(T) : x[i]
    end
    any(x .< 0) && warn("something has gone wrong with the eigen solve")
    x
end

# solve x(P-I)=0 by lu decomposition
"$(____solver_docs)"
function lu_solve{T}(p::Matrix{T})
    n, m = size(p)
    x   = vcat(Array(T, n-1), one(T))
    u   = lufact(p' - one(p))[:U]
    for i = n-1:-1:1 # backsubstitution
        x[i] = -sum([x[j]*u[i, j] for j=i:n])/u[i, i]
    end
    x ./= norm(x, 1) # normalisation
    for i = 1:length(x)
        x[i] = isapprox(x[i], zero(T)) ? zero(T) : x[i]
    end
    any(x .< 0) && warn("something has gone wrong with the lu solve")
    x
end

"$(____solver_docs)"
gth_solve{T<:Integer}(a::Matrix{T}) = gth_solve(convert(Array{Rational, 2}, a))

# solve x(P-I)=0 by the GTH method
function gth_solve{T<:Real}(original::Matrix{T})
    a = copy(original)
    n = size(a, 1)
    x = zeros(T, n)

    @inbounds for k in 1:n-1
        scale = sum(a[k, k+1:n])
        if scale <= zero(T)
            # There is one (and only one) recurrent class contained in
            # {1, ..., k};
            # compute the solution associated with that recurrent class.
            n = k
            break
        end
        a[k+1:n, k] /= scale

        for j in k+1:n, i in k+1:n
            a[i, j] += a[i, k] * a[k, j]
        end
    end

    # backsubstitution
    x[n] = 1
    @inbounds for k in n-1:-1:1, i in k+1:n
        x[k] += x[i] * a[i, k]
    end

    # normalisation
    x / sum(x)
end

"""
Find the recurrent classes of the `MarkovChain`

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix

##### Returns

- `x::Vector{Vector}`: A `Vector` containing `Vector{Int}`s that describe the
recurrent classes of the transition matrix for p

"""
recurrent_classes(mc::MarkovChain) = attracting_components(DiGraph(mc.p))

communication_classes(mc::MarkovChain) = strongly_connected_components(DiGraph(mc.p))

is_irreducible(mc::MarkovChain) =  is_strongly_connected(DiGraph(mc.p))

is_aperiodic(mc::MarkovChain) = period(mc) == 1

function period(mc::MarkovChain)
    g = DiGraph(mc.p)
    recurrent = attracting_components(g)
    periods   = Int[]

    for r in recurrent
        push!(periods, period(g[r]))
    end
    return lcm(periods)
end
"""
calculate the stationary distributions associated with a N-state markov chain

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix
- `;method::Symbol(:gth)`: One of `gth`, `lu`, and `eigen`; specifying which
of the three `_solve` methods to use.

##### Returns

- `dists::Matrix{Float64}`: N x M matrix where each column is a stationary
distribution of `mc.p`

"""
function mc_compute_stationary{T}(mc::MarkovChain{T}; method::Symbol=:gth)
    solvers = Dict(:gth => gth_solve, :lu => lu_solve, :eigen => eigen_solve)
    solve   = solvers[method]

    recurrent = recurrent_classes(mc)

    # unique recurrent class
    length(recurrent) == 1 && return solve(mc.p)

    # more than one recurrent classes
    stationary_dists = Array(T, n_states(mc), length(recurrent))
    for (i, class) in enumerate(recurrent)
        dist        = zeros(T, n_states(mc))
        dist[class] = solve(mc.p[class, class]) # recast to correct dimension
        stationary_dists[:, i] = dist
    end
    return stationary_dists
end
"""
Simulate a Markov chain starting from an initial state

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix
- `init::Int(rand(1:n_states(mc)))` : The index of the initial state. This should
be an integer between 1 and `n_states(mc)`
- `sample_size::Int(1000)`: The number of samples to collect
- `;burn::Int(0)`: The burn in length. Routine drops first `burn` of the
`sample_size` total samples collected

##### Returns

- `samples::Vector{Int}`: Vector of simulated states

"""
function mc_sample_path(mc::MarkovChain,
                        init::Int=rand(1:n_states(mc)),
                        sample_size::Int=1000;
                        burn::Int=0)
    samples = Array(Int, sample_size)
    samples[1] = init
    mc_sample_path!(mc, samples)
    samples[burn+1:end]
end

"""
Simulate a Markov chain starting from an initial distribution

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix
- `init::Vector` : A vector of length `n_state(mc)` specifying the number
probability of being in seach state in the initial period
- `sample_size::Int(1000)`: The number of samples to collect
- `;burn::Int(0)`: The burn in length. Routine drops first `burn` of the
`sample_size` total samples collected

##### Returns

- `samples::Vector{Int}`: Vector of simulated states

"""
function mc_sample_path(mc::MarkovChain,
                        init::Vector = vcat(1, zeros(n_states(mc)-1)),
                        sample_size::Int=1000; burn::Int=0)
    # ensure floating point input for Categorical()
    init = convert(Vector{Float64}, init)
    return mc_sample_path(mc, rand(Categorical(init)), sample_size, burn=burn)
end

# duplicate python functionality
function mc_sample_path(mc::MarkovChain,
                        state::Real,
                        sample_size::Int=1000; burn::Int=0)
    init = zeros(n_states(mc))
    init[state] = 1
    return mc_sample_path(mc, init, sample_size, burn=burn)
end

"""
Fill `samples` with samples from the Markov chain `mc`

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix
- `samples::Array{Int}` : Pre-allocated vector of integers to be filled with
samples from the markov chain `mc`. The first element will be used as the
initial state and all other elements will be over-written.

##### Returns

None modifies `samples` in place
"""
function mc_sample_path!(mc::MarkovChain, samples::Array)
    # ensure floating point input for Categorical()
    p       = convert(Matrix{Float64}, mc.p)
    dist    = [Categorical(vec(p[i, :])) for i=1:n_states(mc)]
    for t=2:length(samples)
        samples[t] = rand(dist[samples[t-1]])
    end
    Void
end

function simulate(mc::MarkovChain,
                  init::Vector = vcat(1, zeros(n_states(mc)-1)),
                  sample_size=1000)
    return mc_sample_path(mc, init, sample_size)
end


using Primes

function randfact(N; min_len=1, max_len=10)
    l = rand(min_len:max_len)
    base_arr = factor(Array, N)
    ind = rand(1:l, length(base_arr))
    return prod.(getindex.(Ref(base_arr),　(i -> ind .== i).(1:l)))
end

function generate_shape(n_elemtnt; min_dim=1, max_dim=5)
    return Tuple(randfact(n_elemtnt, max_len=max_dim)), Tuple(randfact(n_elemtnt, min_len=min_dim, max_len=max_dim))
end

function generate_expr(N)
    randop() = sample(["+", "-", "/", "*", "^"])
    randnum(scale=100) = @sprintf("%.2f", scale * (rand() - 0.5))
    randarg() = rand() > 0.5 ? "x" : randnum()
    S = "$(randarg()) $(randop()) $(randarg())"
    N = 10
    for _ in 1:N
        S = "($S) $(randop()) $(randarg())"
    end
    return S
end

function numerical_diff(f::Function, x::Real; e=1e-7)
    return (f(x + e) - f(x - e)) / 2e
end

function numerical_diff(f::Function, xs::AbstractArray; e=1e-7)
    grads = zeros(length(xs))
    for idx in 1:length(xs)
        tmp_val = xs[idx]
        xs[idx] = tmp_val + e
        fxh1 = f(xs...)
        xs[idx] = tmp_val - e
        fxh2 = f(xs...)
        grads[idx] = (fxh1 - fxh2) / 2e
        xs[idx] = tmp_val
    end
    return grads
end

function backward_diff(f, xs::AbstractArray)
    inputs = Variable.(xs)
    outs = f(inputs...)
    backward!(outs)
    return (input -> input.grad.values).(inputs)
end

function isAbout(x, y; e=1e-4)
    return ((x - y) < 1e-15) || (-e < ((x - y) / y) < e)
end

function isAbout(X::AbstractArray, Y::AbstractArray; e=1e-4)
    for (x, y) in zip(X, Y)
        if !(((x - y) < 1e-15) || (-e < ((x - y) / y) < e))
            return false
        end
    end
    return true
end

function random_arr(;min_dim=1, max_dim=4, min_r=1, max_r=10, collection=Float64)
    dim = rand(min_dim:max_dim)
    shape = Tuple(rand(min_r:max_r, dim))
    return rand(collection, shape)
end


using IncrementalInference
using Test

##

@testset "test _evalFactorTemporary" begin
## test utility to build a temporary graph

fct = EuclidDistance(Normal(10,1))
varTypes = (ContinuousScalar,ContinuousScalar);
varPts = ([0;],[9.5;])

##

dfg, _dfgfct = IIF._buildGraphByFactorAndTypes!(fct, varTypes, varPts)

@test length(intersect(ls(dfg), [:x1; :x2])) == 2
@test lsf(dfg) == [:x1x2f1;]

## test  the evaluation of factor without

B = IIF._evalFactorTemporary!(EuclidDistance(Normal(10,1)), varTypes, 2, [[10;]], varPts );

@test_broken B isa Vector{Vector{Float64}}
@test isapprox( B[1], [10.0;], atol=1e-6)

##
end

@testset "test residual slack prerequisite for numerical factor gradients, Euclidean(1)" begin
##

fct = EuclidDistance(Normal(10,1))
measurement = [[10;]]
varTypes = (ContinuousScalar,ContinuousScalar)
pts = ([0;],[9.5;])

##

slack_resid = calcFactorResidualTemporary(fct, varTypes, measurement[1], pts)

##

coord_1 = IIF._evalFactorTemporary!(fct, varTypes, 1, measurement, pts, _slack=slack_resid )
@test length(coord_1) == 1
@test isapprox( coord_1[1], [0.0], atol=1e-6)

coord_2 = IIF._evalFactorTemporary!(fct, varTypes, 2, measurement, pts, _slack=slack_resid )
@test length(coord_2) == 1
@test isapprox( coord_2[1], [9.5], atol=1e-6)

##

coord_1 = IIF._evalFactorTemporary!(fct, varTypes, 1, measurement, pts )
@test length(coord_1) == 1
@test isapprox( coord_1[1], [-0.5], atol=1e-6)

coord_2 = IIF._evalFactorTemporary!(fct, varTypes, 2, measurement, pts )
@test length(coord_2) == 1
@test isapprox( coord_2[1], [10.0], atol=1e-6)

##
end

@testset "test residual slack prerequisite for numerical factor gradients, Euclidean(2)" begin
##

fct = LinearRelative(MvNormal([10;0],[1 0; 0 1]))
measurement = [[10.0;0.0]]
varTypes = (ContinuousEuclid{2},ContinuousEuclid{2})
pts = ([0;0.0] ,[9.5;0])

## test the building of factor graph to be correct

_fg,_ = IIF._buildGraphByFactorAndTypes!(fct, varTypes, pts);

@test length(getVal(_fg[:x1])) == 1
@test length(getVal(_fg[:x1])[1]) == 2

@test length(getVal(_fg[:x2])) == 1
@test length(getVal(_fg[:x2])[1]) == 2

##

_fg = initfg()

slack_resid = calcFactorResidualTemporary(fct, varTypes, measurement[1], pts, tfg=_fg)

@test length(getVal(_fg[:x1])) == 1
@test length(getVal(_fg[:x1])[1]) == 2

@test length(getVal(_fg[:x2])) == 1
@test length(getVal(_fg[:x2])[1]) == 2

## Manually provide a common temp graph and force no factor and same variables via keywords

tfg,_ = IIF._buildGraphByFactorAndTypes!(fct, varTypes, pts);
coord_1 = IIF._evalFactorTemporary!(fct, varTypes, 1, measurement, pts, _slack=slack_resid, tfg=tfg, newFactor=false, currNumber=0 )

##

@test length(coord_1) == 1
@test length(coord_1[1]) == 2

@test isapprox( coord_1[1], [0;0.0], atol=1e-6)

coord_2 = IIF._evalFactorTemporary!(fct, varTypes, 2, measurement, pts, _slack=slack_resid )
@test length(coord_2) == 1
@test length(coord_2[1]) == 2

@test isapprox( coord_2[1], [9.5; 0], atol=1e-6)

## Repeat the same test but allow _evalFactorTemporary to self construct internal temporary graph

coord_1 = IIF._evalFactorTemporary!(fct, varTypes, 1, measurement, pts )

##
@test length(coord_1) == 1
@test length(coord_1[1]) == 2

@test isapprox( coord_1[1], [-0.5;0], atol=1e-6)

coord_2 = IIF._evalFactorTemporary!(fct, varTypes, 2, measurement, pts )
@test length(coord_2) == 1
@test length(coord_2[1]) == 2

@test isapprox( coord_2[1], [10.0;0], atol=1e-6)

##
end

#


using Revise
using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using Test
using SparseArrays
Random.seed!(233)

sparse_indexing = load_op_and_grad("build/libSparseIndexing", "sparse_indexing", multiple=true)
# function Base.:getindex(s::SparseTensor, i1::Union{PyObject,Array{S,1}},
#          i2::Union{PyObject,Array{T,1}}) where {S<:Real,T<:Real}
#     m_, n_ = length(i1), length(i2)
#     i1 = convert_to_tensor(i1, dtype=Int64)
#     i2 = convert_to_tensor(i2, dtype=Int64)
#     ii1, jj1, vv1 = find(s)
#     m = tf.convert_to_tensor(s.o.shape[1],dtype=tf.int64)
#     n = tf.convert_to_tensor(s.o.shape[2],dtype=tf.int64)
#     ii2, jj2, vv2 = sparse_indexing(ii1,jj1,vv1,m,n,i1,i2)
#     SparseTensor(ii2, jj2, vv2, m_, n_)
# end
################## End Load Operator ##################

i1 = unique(rand(1:20,3))
j1 = unique(rand(1:30,3))
A = sprand(20,30,0.3)
@show i1, j1
Ad = Array(A[i1, j1])
B = SparseTensor(A)
Bd = Array(B[i1, j1])
sess = Session()
init(sess)
Bd_ = run(sess, Bd)
@test Ad≈Bd_

# error()
# TODO: change your test parameter to `m`
# gradient check -- v
function scalar_function(m_)
    A = SparseTensor(ii1, jj1, m_, 20, 30)
    B = A[i1, j1]
    return sum(B)^2
end

# TODO: change `m_` and `v_` to appropriate values
m = 20
n = 30
ii1,jj1,vv1 = find(B)
m_ = constant(rand(length(B.o.values)))
v_ = rand(length(B.o.values))
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session()
init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


using AlgebraicDynamics.UWDDynam
using AlgebraicDynamics.DWDDynam

using StaticArrays
using DynamicalSystems
import DynamicalSystems: trajectory

using Base.Iterators

"""    trajectory(r::DiscreteResourceSharer, u0::AbstractVector, p, nsteps::Int; dt::Int = 1)

Evolves the resouce sharer `r` for `nsteps` times with step size `dt`, initial condition `u0`, and parameters `p`.
"""
trajectory(r::DiscreteResourceSharer, u0::AbstractVector, p, T::Int; dt::Int= 1) =
  trajectory((u,p,t) -> eval_dynamics(r, u, p, t), u0, p, T; dt = dt)

"""    trajectory(m::DiscreteMachine, u0::AbstractVector, xs::AbstractVector, p, nsteps::Int; dt::Int = 1)

Evolves the machine `m` for `nsteps` times with step size `dt`, initial condition `u0`, and parameters `p`. Any inputs to `m` are determied by `xs`. If `m` has no inputs then you can omit `xs`.
"""
trajectory(m::DiscreteMachine, u0::AbstractVector, xs::AbstractVector,  p, T::Int; dt::Int= 1) =
  trajectory((u,p,t) -> eval_dynamics(m, u, xs, p, t), u0, p, T; dt = dt)

trajectory(m::DiscreteMachine, u0::AbstractVector, x,  p, T::Int; dt::Int = 1) =
  trajectory(m, u0, collect(repeated(x, ninputs(m))), p, T; dt = dt)

trajectory(m::DiscreteMachine, u0::AbstractVector, p, T::Int; dt::Int = 1) =
  trajectory(m, u0, [], p, T; dt = dt)

function trajectory1d(update::Function, u0::Vector{S}, p, T::Int; dt::Int) where S
  dt = round(Int, dt)
  tvec = 0:dt:T
  L = length(tvec)
  data = Vector{S}(undef, L)
  data[1] = u0[1]
  u = copy(u0)
  for i in 2:L
    for j in 1:dt
      u = update(u, p, (i - 1)*dt + j)
    end
    data[i] = u[1]
  end
  return data
end

function trajectory(update::Function, u0::Vector{S}, p, T::Int; dt::Int) where S
  if length(u0) == 1
    return trajectory1d(update, u0, p, T; dt = dt)
  end
  dt = round(Int, dt)
  tvec = 0:dt:T
  L = length(tvec)
  data = Vector{SVector{length(u0), S}}(undef, L)
  data[1] = u0
  u = copy(u0)
  for i in 2:L
    for j in 1:dt
      u = update(u, p, (i - 1)*dt + j)
    end
    data[i] = u
  end
  return Dataset(data)
end


using SymmetricPolynomials
using Test
import SymmetricPolynomials: semi_elementary_polynomial, dim, is_elementary, to_string
import SymmetricPolynomials: push!, to_polynomial

@testset "Monomials" begin
    include("monomialTests.jl")
end

@testset "Polynomials" begin
    include("polynomialTests.jl")
end

@testset "Decompose" begin
    include("decomposeTests.jl")
end

@testset "Evaluate" begin
    include("evaluateTests.jl")
end


struct Solution
    ret_val::Int
    primal::Vector{Float64}
    dual::Vector{Float64}
    slack::Vector{Float64}
    objval::Float64
end
Solution() = Solution(0, # SCS_UNFINISHED
                      Float64[], Float64[], Float64[], NaN)

mutable struct Data
    m::Int
    n::Int
    I::Vector{Int}
    J::Vector{Int}
    V::Vector{Float64}
    b::Vector{Float64}
    objconstant::Float64
    c::Vector{Float64}
end

mutable struct Cone
    f::Int # number of linear equality constraints
    l::Int # length of LP cone
    q::Int # length of SOC cone
    qa::Vector{Int} # array of second-order cone constraints
    s::Int # length of SD cone
    sa::Vector{Int} # array of semi-definite constraints
    ep::Int # number of primal exponential cone triples
    ed::Int # number of dual exponential cone triples
    p::Vector{Float64} # array of power cone params
    setconstant::Dict{Int, Float64} # For the constant of EqualTo, LessThan and GreaterThan
    nrows::Dict{Int, Int} # The number of rows of each vector sets
    function Cone()
        new(0, 0,
            0, Int[],
            0, Int[],
            0, 0, Float64[],
            Dict{Int, Float64}(),
            Dict{Int, UnitRange{Int}}())
    end
end


# Out-of-place C = A*B
@inline function gemm(A::AbstractMatrix, B::AbstractMatrix)
    m, n, p = size(A,1), size(B,2), size(A,2)

    C = Matrix{eltype(A)}(undef, m, n)

    @turbo for i in 1:m, k in 1:n
        aux = zero(eltype(A[1]))
        for j in 1:p
            aux += A[i,j]*B[j,k]
        end
        C[i,k] = aux
    end
    return C

end

# threaded version of vectorized gemm
@inline function gemmt(A::AbstractMatrix, B::AbstractMatrix)
    m, n, p = size(A,1), size(B,2), size(A,2)

    C = Matrix{eltype(A)}(undef, m, n)

    @tturbo for i in 1:m, k in 1:n
        aux = zero(eltype(A[1]))
        for j in 1:p
            aux += A[i,j]*B[j,k]
        end
        C[i,k] = aux
    end
    return C

end

# In-place C = A*B
@inline function gemm!(C::AbstractArray{T},A::AbstractArray{T},B::AbstractArray{T}) where T

    m, n, p = size(A,1), size(B,2), size(A,2)
    @turbo for i in 1:m, k in 1:n
        aux = zero(eltype(A[1]))
        for j in 1:p
            aux += A[i,j]*B[j,k]
        end
        C[i,k] = aux
    end
end

@inline function gemmt!(C::AbstractArray{T},A::AbstractArray{T},B::AbstractArray{T}) where T

    m, n, p = size(A,1), size(B,2), size(A,2)
    @tturbo for i in 1:m, k in 1:n
        aux = zero(eltype(A[1]))
        for j in 1:p
            aux += A[i,j]*B[j,k]
        end
        C[i,k] = aux
    end
end

# x + p*y -> y (vec+scalar*vec)
function xpy!(y::Vector{T}, x::Vector{T}, p::T) where T
    @tturbo for i in eachindex(y)
        y[i] = x[i] + p * y[i]
    end
end

# x + p*y -> y (vec+scalar*vec)
# xpy(y::CuArray, x::CuArray, p::Number) = x .+ p .* y

# out-of-place Sparse CSC matrix times dense vector
function spmv(A::AbstractSparseMatrix, x::DenseVector)
    out = zeros(eltype(x), A.m)
    Aj = rowvals(A)
    nzA = nonzeros(A)
    @inbounds for col in 1:A.n
        xj = x[col]
        @fastmath for j in nzrange(A, col)
            out[Aj[j]] += nzA[j]*xj
        end
    end
    out
end

function _spmvt!(out, A, Aj, xj, nzA, col)
    @fastmath for j in nzrange(A, col)
        out[Aj[j]] += nzA[j]*xj
    end
end

# in-of-place Sparse CSC matrix times dense vector
function spmv!(out::DenseVector, A::AbstractSparseMatrix, x::DenseVector)
    fill!(out, zero(eltype(out)))
    Aj = rowvals(A)
    nzA = nonzeros(A)
    @inbounds for col in 1:A.n
        xj = x[col]
        @fastmath for j in nzrange(A, col)
            out[Aj[j]] += nzA[j]*xj
        end
    end
    out
end

mynorm(r::Vector) = sqrt(mydot(r,r))

# Dot product
function mydot(a::AbstractArray{T}, b::AbstractArray{T}) where {T}
    n = zero(T)
    @turbo for i in eachindex(a)
        n += a[i]*b[i]
    end
    n
end

# Threaded dot product
function mydott(a::AbstractArray{T}, b::AbstractArray{T}) where {T}
    n = zero(T)
    @tturbo for i in eachindex(a)
        n += a[i]*b[i]
    end
    n
end


module QCurlGradMonomialBasesTests

using Test
using TensorPolynomialBases
using TensorValues

P = VectorValue{3}
V = VectorValue{3,Float64}
order = 1
basis = QCurlGradMonomialBasis{P,V}(order)
G = gradient_type(basis)

x = VectorValue(2,3,5)

v = V[
  (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0),
  (2.0, 0.0, 0.0), (0.0, 3.0, 0.0), (0.0, 0.0, 5.0)]

g = G[
  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
  (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
  (0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0),
  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0)]

test_polynomial_basis(basis,x,v,g)

P = VectorValue{2}
V = VectorValue{2,Float64}
order = 2
basis = QCurlGradMonomialBasis{P,V}(order)
G = gradient_type(basis)

x = VectorValue(2,3)

v = V[
  (1.0, 0.0), (0.0, 1.0), (2.0, 0.0), (0.0, 3.0),
  (4.0, 0.0), (0.0, 9.0), (3.0, 0.0), (0.0, 2.0),
  (6.0, 0.0), (0.0, 6.0), (12.0, 0.0), (0.0, 18.0)]

g = G[
  (0.0, 0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 0.0),
  (1.0, 0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0),
  (4.0, 0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 6.0),
  (0.0, 1.0, 0.0, 0.0), (0.0, 0.0, 1.0, 0.0),
  (3.0, 2.0, 0.0, 0.0), (0.0, 0.0, 3.0, 2.0),
  (12.0, 4.0, 0.0, 0.0),(0.0, 0.0, 9.0, 12.0)]

test_polynomial_basis(basis,x,v,g)

end # module