### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ d8a922a0-e5b0-11eb-3290-75b867a2f730
begin
    using Pkg
    Pkg.activate(".")
    Pkg.instantiate()
end

# ╔═╡ 5642cd29-c2ab-4943-985e-64c5999f7bae
using
    DiffEqFlux, DifferentialEquations, GalacticOptim,
    Flux, Optim, BlackBoxOptim,
    Zygote, ForwardDiff,
    Distributions, Statistics,
    Distances,
    KernelDensity,
    CUDA,
    Logging, TerminalLoggers,
    StatsPlots,
    DataFrames

# ╔═╡ 0085e83f-a3df-42f8-8a0e-658ec2afca50
begin
    global_logger(TerminalLogger())
    plotly()
end

# ╔═╡ 3f8520e1-89f4-4b21-8e52-7bf48c760db5
data_dist = Beta(2, 4)

# ╔═╡ 59e766f3-c011-4aa3-a42d-f0cf03d1c00a
plot(data_dist, label=string(data_dist))

# ╔═╡ 7fb8118e-26a2-4d23-adb0-d20d3d024ff6
train_data = rand(data_dist, 1, 100)

# ╔═╡ 5cf5412d-7eb0-410d-897a-a8d0e273b21f
test_data = rand(data_dist, 1, 1000)

# ╔═╡ 74801cce-a47b-4f60-a8c8-f56bcdc2534d
mdl_dist_f_kde = KernelDensity.kernel_dist(Normal, KernelDensity.default_bandwidth(reshape(train_data, size(train_data, 2))))

# ╔═╡ 1cc754a6-385d-4f87-9f06-9d3cb8bb1a3c
mdl_dist_f_data = Normal(mean(train_data), std(train_data))

# ╔═╡ d8869f7a-988c-4dc1-9f8f-6467ff1657c5
mdl_dist_f_dist = Normal(mean(data_dist), std(data_dist))

# ╔═╡ f7563269-4749-41c7-b4a8-1fd8f9b73c14
let
    p = plot(data_dist, label=string(data_dist))
    p = plot!(p, mdl_dist_f_dist, label=string(mdl_dist_f_dist))
    p = plot!(p, mdl_dist_f_data, label=string(mdl_dist_f_data))
    p = plot!(p, mdl_dist_f_kde, label=string(mdl_dist_f_kde))
end

# ╔═╡ d7175898-e654-4f0f-885b-d0e104e62507
n_hidden = 3

# ╔═╡ 929715b5-3e57-4403-a081-3fd6499d916d
actv = tanh

# ╔═╡ fa1edbc5-a05c-463f-83d9-8a4cc492f731
nn = Chain(
    Dense(1, n_hidden, actv),
    # Dense(n_hidden, n_hidden, actv),
    Dense(n_hidden, 1, actv),
) |> f64

# ╔═╡ 35769a13-0537-4413-b5f9-4a7a126bda3c
tspan = (0.0, 1.0)

# ╔═╡ a65312c6-a9c3-4a0f-ba88-afe122e30395
ffjord_mdl = FFJORD(nn, tspan, Tsit5(), basedist=mdl_dist_f_dist)

# ╔═╡ 595b19d9-de04-4038-a87e-316ada76e070
ffjord_mdl.p

# ╔═╡ daab25dd-6a54-46fa-a26c-d2f05183659e
ffjord_mdl.basedist

# ╔═╡ 4bb079bb-fc0a-476e-bb42-af1caa1805f2
plot(ffjord_mdl.basedist)

# ╔═╡ 763966b9-a42c-4787-8e5f-c925a4b44f15
regularize = false

# ╔═╡ ac7cccd8-e564-4b4a-bf10-bcea49d2ff18
monte_carlo = false

# ╔═╡ 30a83fcc-cb24-4033-a665-92869215da78
function loss(θ)
    logpx, λ₁, λ₂ = ffjord_mdl(train_data, θ; regularize, monte_carlo)
    -mean(logpx)
end

# begin
# 	actual_pdf = pdf.(data_dist, train_data)
# 	function loss(p)
# 		learned_pdf = exp.(first(ffjord_mdl(train_data, p; regularize, monte_carlo)))
# 		dis = totalvariation(learned_pdf, actual_pdf) / size(train_data, 2)
# 		-dis
# 	end
# end

# ╔═╡ 842fc843-2a89-445c-8319-50a5fafec73b
adtype = GalacticOptim.AutoZygote()
# adtype = GalacticOptim.AutoForwardDiff()

# ╔═╡ 4b77653f-1cbb-4d7c-8dc8-6ec0e91b4b0d
optm1 = Flux.ADAM(0.1)
# optm1 = Optim.GradientDescent()

# ╔═╡ 9686dc0b-1f5e-4ce0-958c-7059caf93e70
optm2 = Optim.BFGS()

# ╔═╡ 32fc54a0-3bd7-4ab6-a5de-93cb6d1c3e2e
optm3 = BBO(:adaptive_de_rand_1_bin)

# ╔═╡ f56bfb88-8ac1-4227-a0b1-60eb613aec67
maxiters_m1 = 100

# ╔═╡ 114ce7be-ceaa-4efd-8890-95cc95cf3219
maxiters_m2 = 25

# ╔═╡ 4e42aa35-dfae-422e-ac8c-13c6b80efba4
maxiters_m3 = 10

# ╔═╡ bb213e1d-3ece-40de-8007-4fd12a814823
progress = true

# ╔═╡ 7a6d80ad-2d15-42f0-8012-6316f9905ecc
res1 = DiffEqFlux.sciml_train(loss, ffjord_mdl.p, optm1, adtype; progress, maxiters=maxiters_m1)

# ╔═╡ 83018b3a-27e6-42ae-9dc3-0f684022c386
res1.u

# ╔═╡ 987e0cee-36ba-4130-8b7d-65a8bcfd9bb3
res2 = DiffEqFlux.sciml_train(loss, res1.u, optm2, adtype; progress, maxiters=maxiters_m2)
# res2 = DiffEqFlux.sciml_train(loss, res1.u, optm2, adtype; progress)

# ╔═╡ fd432df9-d795-4732-806f-e0e910c3599d
res2.u

# ╔═╡ 7b40ebee-3559-4cfc-a823-d3772177acc3
# res3 = DiffEqFlux.sciml_train(loss, res2.u, optm3, adtype; progress, maxiters=maxiters_m3)

# ╔═╡ 5ed15a42-3ace-4850-99ac-f4ce705034cf
final_res = res2

# ╔═╡ 60e40d80-eee0-4a08-a679-6727f9a91852
loss(final_res.u)

# ╔═╡ 4f4b8781-5eca-442b-8950-b94f866351db
train_dis = let data = train_data
    actual_pdf = pdf.(data_dist, data)
    learned_pdf = exp.(first(ffjord_mdl(data, final_res.u; regularize, monte_carlo)))
    dis = totalvariation(learned_pdf, actual_pdf) / size(data, 2)
end

# ╔═╡ f97042ce-980d-433b-a293-63b84b01c3e9
test_dis = let data = test_data
    actual_pdf = pdf.(data_dist, data)
    learned_pdf = exp.(first(ffjord_mdl(data, final_res.u; regularize, monte_carlo)))
    dis = totalvariation(learned_pdf, actual_pdf) / size(data, 2)
end

# ╔═╡ d9225bf7-d4b6-4462-bd51-0639ff90d0df
begin
    hist_range = range(0, 1, length=1000 + 1)
    hist_data = reshape(collect(hist_range), (1, size(hist_range, 1)))
end

# ╔═╡ ebf17afb-9e1d-4465-a337-603309ef345d
begin
    hist_pdf_0 = exp.(first(ffjord_mdl(hist_data; regularize, monte_carlo)))
    hist_pdf_1 = exp.(first(ffjord_mdl(hist_data, res1.u; regularize, monte_carlo)))
    hist_pdf_2 = exp.(first(ffjord_mdl(hist_data, res2.u; regularize, monte_carlo)))
end

# ╔═╡ faf10042-af69-46d2-a891-7e774fbcea63
kde_mdl = kde(reshape(train_data, size(train_data, 2)))

# ╔═╡ afbc4978-2eee-41fb-a91a-f89fc9989b4a
    kde_train_dis = let data = train_data
    actual_pdf = pdf.(data_dist, data)
    learned_pdf = reshape(pdf(kde_mdl, reshape(data, size(data, 2))), size(data))
    dis = totalvariation(learned_pdf, actual_pdf) / size(data, 2)
end

# ╔═╡ d8c2056b-6af3-42f5-8ce6-dfb3cd8b638f
kde_test_dis = let data = test_data
    actual_pdf = pdf.(data_dist, data)
    learned_pdf = reshape(pdf(kde_mdl, reshape(data, size(data, 2))), size(data))
    dis = totalvariation(learned_pdf, actual_pdf) / size(data, 2)
end

# ╔═╡ d9afd296-82a9-4989-bcf8-e08ec01e25c2
kde_hist_pdf = reshape(pdf(kde_mdl, reshape(hist_data, size(hist_data, 2))), size(hist_data))

# ╔═╡ 38961acf-e62d-4abf-9e21-0f2e5646506a
let data = hist_data
    p = plot(data_dist, label=string(data_dist))
    p = plot!(p, reshape(data, size(data, 2)), reshape(hist_pdf_0, size(data, 2)), label="p(x) - init")
    p = plot!(p, reshape(data, size(data, 2)), reshape(hist_pdf_1, size(data, 2)), label="p(x) - after adam")
    p = plot!(p, reshape(data, size(data, 2)), reshape(hist_pdf_2, size(data, 2)), label="p(x) - after bfgs")
    p = plot!(p, reshape(data, size(data, 2)), reshape(kde_hist_pdf, size(data, 2)), label="p(x) - kde")
end

# ╔═╡ 3e187938-77e0-4fef-b0ee-b04dd0fee97b
result = DataFrame(
    Method=["FFJORD", "FFJORD", "KDE", "KDE"],
    Data=["train_data", "test_data", "train_data", "test_data"],
    TV_Error=[train_dis, test_dis, kde_train_dis, kde_test_dis],
)

# ╔═╡ 537a31bc-360c-4231-9d35-8e899d642126
let data = train_data
    p = plot(data_dist, label=string(data_dist))
    p = scatter!(p, reshape(data, size(data, 2)), reshape(pdf.(data_dist, data), size(data, 2)), label="sample for training")
    p = plot!(p, reshape(hist_data, size(hist_data, 2)), reshape(hist_pdf_2, size(hist_data, 2)), label="p(x) - final")
end

# ╔═╡ Cell order:
# ╠═d8a922a0-e5b0-11eb-3290-75b867a2f730
# ╠═5642cd29-c2ab-4943-985e-64c5999f7bae
# ╠═0085e83f-a3df-42f8-8a0e-658ec2afca50
# ╠═3f8520e1-89f4-4b21-8e52-7bf48c760db5
# ╠═59e766f3-c011-4aa3-a42d-f0cf03d1c00a
# ╠═7fb8118e-26a2-4d23-adb0-d20d3d024ff6
# ╠═5cf5412d-7eb0-410d-897a-a8d0e273b21f
# ╠═74801cce-a47b-4f60-a8c8-f56bcdc2534d
# ╠═1cc754a6-385d-4f87-9f06-9d3cb8bb1a3c
# ╠═d8869f7a-988c-4dc1-9f8f-6467ff1657c5
# ╠═f7563269-4749-41c7-b4a8-1fd8f9b73c14
# ╠═d7175898-e654-4f0f-885b-d0e104e62507
# ╠═929715b5-3e57-4403-a081-3fd6499d916d
# ╠═fa1edbc5-a05c-463f-83d9-8a4cc492f731
# ╠═35769a13-0537-4413-b5f9-4a7a126bda3c
# ╠═a65312c6-a9c3-4a0f-ba88-afe122e30395
# ╠═595b19d9-de04-4038-a87e-316ada76e070
# ╠═daab25dd-6a54-46fa-a26c-d2f05183659e
# ╠═4bb079bb-fc0a-476e-bb42-af1caa1805f2
# ╠═763966b9-a42c-4787-8e5f-c925a4b44f15
# ╠═ac7cccd8-e564-4b4a-bf10-bcea49d2ff18
# ╠═30a83fcc-cb24-4033-a665-92869215da78
# ╠═842fc843-2a89-445c-8319-50a5fafec73b
# ╠═4b77653f-1cbb-4d7c-8dc8-6ec0e91b4b0d
# ╠═9686dc0b-1f5e-4ce0-958c-7059caf93e70
# ╠═32fc54a0-3bd7-4ab6-a5de-93cb6d1c3e2e
# ╠═f56bfb88-8ac1-4227-a0b1-60eb613aec67
# ╠═114ce7be-ceaa-4efd-8890-95cc95cf3219
# ╠═4e42aa35-dfae-422e-ac8c-13c6b80efba4
# ╠═bb213e1d-3ece-40de-8007-4fd12a814823
# ╠═7a6d80ad-2d15-42f0-8012-6316f9905ecc
# ╠═83018b3a-27e6-42ae-9dc3-0f684022c386
# ╠═987e0cee-36ba-4130-8b7d-65a8bcfd9bb3
# ╠═fd432df9-d795-4732-806f-e0e910c3599d
# ╠═7b40ebee-3559-4cfc-a823-d3772177acc3
# ╠═5ed15a42-3ace-4850-99ac-f4ce705034cf
# ╠═60e40d80-eee0-4a08-a679-6727f9a91852
# ╠═4f4b8781-5eca-442b-8950-b94f866351db
# ╠═f97042ce-980d-433b-a293-63b84b01c3e9
# ╠═d9225bf7-d4b6-4462-bd51-0639ff90d0df
# ╠═ebf17afb-9e1d-4465-a337-603309ef345d
# ╠═faf10042-af69-46d2-a891-7e774fbcea63
# ╠═afbc4978-2eee-41fb-a91a-f89fc9989b4a
# ╠═d8c2056b-6af3-42f5-8ce6-dfb3cd8b638f
# ╠═d9afd296-82a9-4989-bcf8-e08ec01e25c2
# ╠═38961acf-e62d-4abf-9e21-0f2e5646506a
# ╠═3e187938-77e0-4fef-b0ee-b04dd0fee97b
# ╠═537a31bc-360c-4231-9d35-8e899d642126


# Global parameters
Nt = 2^17
tmax = 3000
tmin = -3000
dt = (tmax - tmin)/Nt
t = collect(tmin:dt:tmax)[2:end]
α, linf, rsup, δx = 0.5, -500, 500,0.1
xd = linf:δx:rsup

dω= 2π/(Nt*dt)
ω = collect(-Nt/2:1:Nt/2)*dω
ω = ω[1:end-1];

ΔEw = 0.2/ 27.211396
Eα_r = LinRange(0.0,1,600)
δE = Eα_r[2] - Eα_r[1]
ncoef = 150
# Main functions

"""
tFELlimits(FELpulse, Nt, cut = 1e-4)

Computes a cut off of oscillations from the left and the right in the pulse.
Returns the left and right limits.

"""
function tFELlimits(FELpulse, Nt, cut = 1e-4)
    n_tmin = 0
    for i in 1:Nt
        if abs.(FELpulse[i] >= cut)
            n_tmin = i
            break
        end
    end
    n_tmax = 0
    for i in 1:Nt
        if abs.(FELpulse[Nt - i + 1] >= cut)
            n_tmax = Nt - i + 1
            break
        end
    end
    n_tmin, n_tmax
end

"""
G(t, T)

Pulse Envelope. Where T is the intensity FWHM time duration.
"""
G(t, T) = exp.(-2 .*log(2) .*(t ./T ).^2)

"""
g(t, τ)

Pulse. Where τ is the pulse duration (Pulse HW1/e).
In the PCM context is called coherence time.
"""
g(t, τ ) = exp.(-(t./ τ).^2)

"""
FELpulse(t, Nt; T = 3, τ = 0.5, ω₀ = 0.77, intensity = 1.0, fluence = 1,
        gauss_ref = false, asymm = true)

Generates pulses, via the Partial Coherence method (PCM):

\\begin{equation}
f\\_{l}(t) = N\\_{l} G(t) \\mathcal{F}^{-1}\\left[ e^{i\\phi\\_{l}(\\omega)}
\\mathcal{F}[g(t)\\cos(\\omega\\_{0}t)] \\right]
\\end{equation}

with

\\begin{equation}
G(t) = e^{-2\\ln 2\\left(\\frac{t}{T}\\right)^2},\\quad\\mbox{and}\\quad
g(t) = e^{-\\left(\\frac{t}{\\tau}\\right)^2},
\\end{equation}

where \$\\omega_{0}\$ is the carrier frequency and \$\\mathcal{F}\$ \$\\left(\\mathcal{F}^{-1}\\right)\$
the Fourier (inverse Fourier) transform. Also, \$\\tau\$ is the coherence time and \$T\$ the pulse duration.
\$ N\\_{l} \$ is a constant to normalize all pulses to unit pulse energy(default)
otherwise a different intensity(fluence) can be specified.

# Examples
```
julia> Nt = 2^17
julia> tmax = 3000
julia> tmin = -3000
julia> dt = (tmax - tmin)/Nt
julia> t = collect(tmin:dt:tmax)[2:end]
julia> tlims, pulse = FELpulse(t, Nt;) #gaussian pulse

```
"""
function FELpulse(t, Nt; T = 3, τ = 0.5, ω₀ = 0.77, intensity = 1.0,
    fluence = 1.0, gauss_ref = false, asymm = true)

    Tf = 41.34*T
    τf = 41.34*τ
    ϕp = gauss_ref == true ? zeros(Nt) : 2 .*π .*rand(Nt) .- π
    if asymm == true
        ϕp[2:Int64(Nt/2)] .= - reverse(ϕp[Int64(Nt/2 + 2):end]) # antisymmetric
    end
    flr = real(G(t, Tf) .* ifft(ifftshift(exp.(1im*ϕp).*fftshift(fft(g(t, τf) .* cos.(ω₀ .* t))))))
    pulse = flr  .* sqrt(intensity/(3.51ω₀^2))
    nnorm = sum(flr .^2) .* dt
    pulse = gauss_ref == true ? pulse : sqrt(fluence/nnorm) .* flr
    tmin, tmax = tFELlimits(pulse, Nt)
    t[tmin:tmax], pulse[tmin:tmax]
end

function pulsedouble(t, T, τ, ωs, gaFlux)
    Tf = 41.34*T
    τf = 41.34*τ
    dawsonVals = [dawson(i/τf) for i in t]
    flr = (2/√π)*exp.(-2*log(2)*(t/Tf).^2).* dawsonVals .* sin.(ωs .* t)
    nnorm = sum(flr .^2) .* dt
    pulse = sqrt(gaFlux/nnorm) .* flr
    tmin, tmax = tFELlimits(pulse, Nt)
    t[tmin:tmax], pulse[tmin:tmax]
end

function pulseGauss(t, T, τ, ω₀, intensity)
    Tf = 41.34*T
    τf = 41.34*τ
    flr = exp.(-2*log(2)*(t/Tf).^2 - (t/τf).^2) .* cos.(ω₀ .* t)
    pulse = flr  .* sqrt(intensity/(3.51ω₀^2))
    tmin, tmax = tFELlimits(pulse, Nt)
    t[tmin:tmax], pulse[tmin:tmax]
end

function pulsesSep(t,T1, τ1, ω₀, gaFlux)
    τ = 41.34*τ1
    T = 41.34*T1
    flr = (exp.(-2log(2).*(t .+ τ/2).^2/T^2) .+ exp.(-2log(2).*(t .- τ/2).^2/T^2)).*cos.(ω₀*t)
    nnorm = sum(flr .^2) .* dt
    pulse = sqrt(gaFlux/nnorm) .* flr
    tmin, tmax = tFELlimits(pulse, Nt)
    t[tmin:tmax], pulse[tmin:tmax]
end

function pulsesSep(t,T1, τ1, ω₀, frac, gaFlux)
    τ = 41.34*τ1
    T = 41.34*T1
    flr = (frac*exp.(-2log(2).*(t .+ τ/2).^2/T^2) .+ (1-frac)*exp.(-2log(2).*(t .- τ/2).^2/T^2)).*cos.(ω₀*t)
    nnorm = sum(flr .^2) .* dt
    pulse = sqrt(gaFlux/nnorm) .* flr
    tmin, tmax = tFELlimits(pulse, Nt)
    t[tmin:tmax], pulse[tmin:tmax]
end


"""
    dist_extrema(values::T...) where {T <: Number}
Computes the distance between extrema in an ordered set.
"""
function o_dist_extrema(values)
    m, M = extrema(values)
    return Float64(M - m)
end


# discrete MDP 1
ns = 4
na = 6
R = zeros(ns, na)
T = zeros(ns, na, ns) # sp, a, s

# from state 1
T[2,1, 1] = 1.0  # if go, move to 2
T[1,2:na,1] .= 1.0 # else stay

# from state 2
T[1,2,2] = 0.7 # if (s2, safe) end up in 1
T[3,2,2] = 0.3 # if (s2, safe) end up in 3

T[3,3,2] = 0.5
T[4,3,2] = 0.5
T[2,1,2] = 1.
T[2,4:end,2] .= 1.

# from state 3
T[3,:,3] .= 1

# from state 4
T[4,:,4] .= 1.
T[:,4,4] .= 0.
T[1,4,4] = 1.

mdp = TabularMDP(T, R, 0.95)

MEC = maximal_end_components(mdp, verbose=true)
@test MEC == [[3], [4]]

gw = SimpleGridWorld(size=(20,20))
gw_mec = maximal_end_components(gw, verbose=false)

@test length(gw_mec) == 2*gw.size_y + 2*gw.size_x - 3


module SimplePosetRepresentations

# Nothing here yet. We'll do IntervalOrder, SemiOrder, CircleOrder

using SimplePosets, ClosedIntervals

import Base.show
export IntervalOrder, SemiOrder
export Circle, inside, CircleOrder

export poset_builder

### GENERAL MACHINERY

function poset_builder(Xdict::Dict{S,T}, comparator::Function) where {S,T}
  P = SimplePoset{S}()
  for v in keys(Xdict)
    add!(P,v)
  end
  elts = collect(keys(Xdict))
  n = length(elts)
  for i=1:n
    a = elts[i]
    u = Xdict[a]
    for j=1:n
      b = elts[j]
      v = Xdict[b]
      if i!=j && comparator(u,v)
        add!(P,a,b)
      end
    end
  end
  return P
end

function poset_builder(Xlist::Vector{S}, comparator::Function) where S
  d = Dict{Int,S}()
  n = length(Xlist)
  for i=1:n
    d[i] = Xlist[i]
  end
  return poset_builder(d,comparator)
end

function poset_builder(Xset::Set{S}, comparator::Function) where S
  d = Dict{S,S}()
  for s in Xset
    d[s] = s
  end
  return poset_builder(d,comparator)
end

### INTERVAL ORDERS

_interval_compare(I::ClosedInterval{S}, J::ClosedInterval{S}) where S = I << J

"""
`IntervalOrder(Jmap)` creates an interval order from a dictionary
mapping elements to closed intervals.
"""
function IntervalOrder(Jmap::Dict{S,ClosedInterval{T}}) where {S,T}
  return poset_builder(Jmap,_interval_compare)
end

"""
`IntervalOrder(Jlist)` creates an interval order from a list of
intervals. The elements of the poset are named `1:n` where `n`
is the length of the `Jlist`.
"""
function IntervalOrder(Jlist::Vector{ClosedInterval{T}}) where T
  return poset_builder(Jlist,_interval_compare)
end

"""
`IntervalOrder(Jset)` creates an interval order from a set of
closed intervals. The elements of the poset are the intervals.
"""
function IntervalOrder(Jset::Set{ClosedInterval{T}}) where T
  return poset_builder(Jset,_interval_compare)
end

### SEMIORDERS

_unit_compare(x::Real,y::Real) = x+1 < y

"""
`SemiOrder(Xmap)` creates a semiorder whose elements are the
keys in the dictionary `Xmap`. We have `a<b` provided
`Xmap[a]+1 < Xmap[b]`.
"""
function SemiOrder(Xmap::Dict{S,T}) where {S,T<:Real}
  return poset_builder(Xmap, _unit_compare)
end

"""
`SemiOrder(Xlist)` creates a semiorder from a list of real numbers.
The elements are `1:n` and we have `a<b` provided `Xlist[a]+1 < Xlist[b]`.
"""
function SemiOrder(Xlist::Vector{T}) where {T<:Real}
  return poset_builder(Xlist, _unit_compare)
end

"""
`SemiOrder(Xset)` creates a semiorder from a set of real numbers.
The elements of the poset are the elements of `Xset`. For two elements
`a` and `b` we have `a<b` in the poset iff `a+1 < b` as real numbers.
"""
function SemiOrder(Xset::Set{T}) where {T<:Real}
  return poset_builder(Xset,_unit_compare)
end

SemiOrder(Xset::BitSet) = SemiOrder(Set(Xset))

### CIRCLE ORDERS
"""
The `Circle` datatype represents a circle in the
plane specified by a center and a radius like this:
`Circle(x,y,r)`.
"""
struct Circle{X<:Real, Y<:Real, R<:Real}
  x::X     # x-coordinate of center
  y::Y     # y-coordinate of center
  r::R     # radius
end

function show(io::IO, C::Circle)
  print(io, "Circle($(C.x), $(C.y), $(C.r))")
end

"""
`inside(C,D)` returns `true` provided circle `C`
is entirely contained inside circle `D`.
"""
function inside(C::Circle, D::Circle)
  dx = C.x - D.x
  dy = C.y - D.y
  dr = C.r - D.r
  return (dr<=0)&&(dx*dx + dy*dy <= dr*dr)
end

function CircleOrder(Cmap::Dict{S,Circle}) where S
  return poset_builder(Cmap,inside)
end

function CircleOrder(Clist::Vector{Circle})
  return poset_builder(Clist,inside)
end

function CircleOrder(Cset::Set{Circle})
  return poset_builder(Cset,inside)
end

end  # end of module


using GridInterpolations
using Test
using LinearAlgebra
using Random
using DelimitedFiles

# use import otherwise Interpolations.interpolate conflicts with GridInterpolations.interpolate
import Interpolations: BSpline, Linear, OnGrid, Flat
import Interpolations

function compareToGrid(testType::Symbol=:random, numDims::Int=3, pointsPerDim::Int=5, testPointsDim::Int=5; numRandomTests::Int=1000, eps::Float64=1e-10)
    # Compare interpolation results from GridInterpolations with the Grid Module:
    # create a multi-dimensional test Grid with data, cut points, etc.
    # (dataM, gridM, gridI) = compareToGrid(2, 3, 5);
    cutsDim = pointsPerDim*ones(Int,numDims)
    gridData = rand(cutsDim...)

    # Set up the data structures:
    gridI = Interpolations.interpolate(gridData, BSpline(Linear()))      # This is for the Interpolations package
    gridI = Interpolations.extrapolate(gridI, Flat())
    cutPointMat = 0:(1/(pointsPerDim-1)):1
    cutPoints = [cutPointMat for i=1:numDims]
    gridM = RectangleGrid(tuple(cutPoints...)...)         # Mykel's interpolation package

    # Package the data in a way that can be read by the interpolation package:
    dataM = Array{Float64}(undef, length(gridM))
    for i=1:length(gridM)
        # x = ind2x(gridM, i)
        x = ind2x(gridM, i)
        dataM[i] = gridData[getFractionalIndexes(gridM,x)...]
    end

    # Check a finer grid of coordinates, for comprehensiveness?

    if testType==:random
        # Interpolate in GridInterpolations with the same points, report any discrepancy
        # Select a number of random points and compare the interpolation results.
        testPoint = zeros(numDims)
        for i=1:numRandomTests
            rand!(testPoint)
            testI = gridI(((pointsPerDim.-1)*testPoint.+1)...)
            testM = GridInterpolations.interpolate(gridM,dataM,testPoint)

            if (abs(testI-testM)>eps)
                display("Failed Random Point Interpolation Test")
                display("Non-matching interpolation points:")
                display("Test point = $testPoint")
                display("Grid Package interpolation result = $testI")
                display("Interpolation Module result = $testM")
                return false
            end

        end
        return true
    end

    if testType==:extrapNeg
        # Test behavior outside the grid, i.e. extrapolation:
        testPoint = zeros(numDims)
        for i=1:numRandomTests
            testPoint = -rand!(testPoint)
            testI = gridI(((pointsPerDim-1)*testPoint+1)...)
            testM = interpolate(gridM,dataM,testPoint)

            if (abs(testI-testM)>eps)
                display("Failed Random Negative Point Extrapolation Test")
                display("Non-matching interpolation points:")
                display("Test point = $testPoint")
                display("Grid Package interpolation result = $testI")
                display("Interpolation Module result = $testM")
                return false
            end

        end
        return true
    end

    if testType==:extrapPos
        for i=1:numRandomTests
            testPoint = rand(numDims).+1
            testI = gridI(((pointsPerDim-1)*testPoint.+1)...)
            testM = interpolate(gridM,dataM,testPoint)

            if (abs(testI-testM)>eps)
                display("Failed Random Positive Point Extrapolation Test")
                display("Non-matching interpolation points:")
                display("Test point = $testPoint")
                display("Grid Package interpolation result = $testI")
                display("Interpolation Module result = $testM")
                return false
            end

        end
        return true
    end

    return true

end

function getFractionalIndexes(g::AbstractGrid, s::Array)
    # Returns the fractional index of sprime within the grid-defined discretization.

    fracInd = Array{Int64}(undef, length(g.cutPoints))

    for i=1:length(g.cutPoints)
        gridDisc = g.cutPoints[i]
        fracInd[i] = getFracIndex(gridDisc, s[i])
    end

    return fracInd

end

function getFracIndex(vararray::Array, value::Float64)
    # Searches through the array vararray and returns the fractional index position
    # that value lies in.  It CAPassumes vararray increases monotonically, and returns
    # the first or last element if the value is outside the bounds of the array.

    if value <= vararray[1]
        return 1
    end

    if value >= vararray[end]
        return argmax(vararray)
    end

    i=1
    while value > vararray[i+1]
        i+=1
    end

    return i+(value-vararray[i])÷(vararray[i+1] - vararray[i])

end

function simplexMagic(NDISC::Int=20, NPOINTS::Int=3, checkFileName::AbstractString="", eps::Float64=1e-10)
    if isempty(checkFileName)
        checkFileName = joinpath(@__DIR__, "simplexMagicTest20.txt")
    end

    val = transpose([ [8.,1,6] [3,5,7] [4,9,2] ]) # transposed magic(3) from matlab

    sGrid = SimplexGrid(1.:NPOINTS, 1.:NPOINTS)
    sInterpVal = zeros(NDISC, NDISC)

    for i = 1:NDISC, j = 1:NDISC
        valX = mapPt(i, NPOINTS, NDISC)
        valY = mapPt(j, NPOINTS, NDISC)
        sInterpVal[i, j] = interpolate(sGrid, vcat(val...), [valX, valY])
    end

    sInterpValTest = readdlm(checkFileName)

    testErr = sum(abs, sInterpVal-sInterpValTest)

    if (testErr > eps)
        display("Failed Simplex Comparison Test")
        display("Test error = $testErr")
    end

    return testErr<eps

end

# constructs a new grid from repr output and tests to see if it's the same
function reprConstruct()
    g1 = SimplexGrid([1.0, 2.0, 3.0], [1.0, 2.0, 3.0])
    g2 = eval(Meta.parse(repr(g1)))
    return g1.cutPoints == g2.cutPoints
end

function mapPt(pt::Int, NPOINTS::Int, NDISC::Int)
    return pt * (NPOINTS-1) / NDISC + 1
end

function checkCounters()
    dims = 2
    nPts = 100
    grid = RectangleGrid(1.:10, 1.:10)
    if length(grid) != nPts || dimensions(grid) != dims
        return false
    end
    data = rand(10,10)
    temp = interpolate(grid, data, [-1.5,1.5])
    temp = interpolate(grid, data, [1.,1.])
    grid = SimplexGrid(1.:10, 1.:10)
    if length(grid) != nPts || dimensions(grid) != dims
        return false
    end
    temp = interpolate(grid, data, [-1.5,1.5])
    temp = interpolate(grid, data, [1.,1.])
    temp = interpolate(grid, data, [1.5,1.5])
    #    showcompact(grid)
    #    show(grid)
    return true
end

function testMask()
    x = [1.5, 1.5]
    vals = collect(1.:9)
    mask = falses(9)
    grid = RectangleGrid(1.:3, 1.:3)
    interped = interpolate(grid, vals, x)
    masked   = maskedInterpolate(grid, vals, x, mask)
    @test interped ≈ masked atol = 1e-3
    mask[1] = true # mask the first entry
    truth  = 3.6666
    masked = maskedInterpolate(grid, vals, x, mask)
    @test truth ≈ masked atol = 1e-3
    return true
end

@testset "Comparisons" begin
    @test compareToGrid(:random) == true
    @test compareToGrid(:extrapPos) == true
    #@test compareToGrid(:extrapNeg)==true
end

@testset "Magic" begin
    @test simplexMagic() == true
end

# TODO
@testset "Repr" begin
    @test reprConstruct() == true
end

@testset "Counter" begin
    @test checkCounters() == true
end

@testset "Mask" begin
    @test testMask() == true
end

@testset "Duplicates" begin
    # Check that duplicates are disallowed
    @test_throws ErrorException RectangleGrid([1,1])
    @test_throws ErrorException SimplexGrid([1,1])
end

# Sanity check overall implementations
function test_rect_implemented()
    rgrid = RectangleGrid([2,5],[2,5])

    @test length(rgrid) == 4
    @test GridInterpolations.dimensions(rgrid) == 2

    @test ind2x(rgrid,1) == [2,2]
    @test ind2x(rgrid,2) == [5,2]
    @test ind2x(rgrid,3) == [2,5]
    @test ind2x(rgrid,4) == [5,5]
    x = [0,0]
    ind2x!(rgrid, 4, x)
    @test x == [5,5]

    @test interpolants(rgrid, [1,1]) == ([1],[1.0])
    @test interpolants(rgrid, [2,5]) == ([3],[1.0])

    indices, weights = interpolants(rgrid, [1.5,3])
    @test indices == [1,3]
    @test isapprox(weights, [0.666667, 0.333333], rtol=1e-5)

    @test interpolate(rgrid, [1,2,3,4], [1,1]) == 1.0
    @test maskedInterpolate(rgrid, [1,2,3,4], [1,1], BitArray([false, false, false, false])) == 1.0
    @test isnan(maskedInterpolate(rgrid, [1,2,3,4], [1,1], BitArray([true, false, false, false])))

    @test isapprox(interpolate(rgrid, [1,2,3,4], [1.5,3]), 1.66666666, rtol=1e-5)
    @test maskedInterpolate(rgrid, [1,2,3,4], [1.5,3], BitArray([true, false, false, false])) == 3
    @test maskedInterpolate(rgrid, [1,2,3,4], [1.5,3], BitArray([false, false, true, false])) == 1
end

@testset "Rect Impl" begin
    test_rect_implemented()
end

# Sanity check overall implementations
function test_simplex_implemented()
    grid = SimplexGrid([2,5],[2,5])

    @test length(grid) == 4
    @test GridInterpolations.dimensions(grid) == 2

    @test ind2x(grid,1) == [2,2]
    @test ind2x(grid,2) == [5,2]
    @test ind2x(grid,3) == [2,5]
    @test ind2x(grid,4) == [5,5]
    x = [0,0]
    ind2x!(grid, 4, x)
    @test x == [5,5]

    indices, weights = interpolants(grid, [1,1])
    full_weight_indices = findall(x -> x == 1, weights)
    @test  length(full_weight_indices) == 1
    @test  weights[full_weight_indices[1]] == 1.0

    indices, weights = interpolants(grid, [1.5,3])
    full_weight_indices = findall(x -> isapprox(x, 0.666667, rtol=1e-5), weights)
    @test  length(full_weight_indices) == 1
    @test  isapprox(weights[full_weight_indices[1]], 0.666667, rtol=1e-5) == true
    full_weight_indices = findall(x -> isapprox(x, 0.333333, rtol=1e-5), weights)
    @test  length(full_weight_indices) == 1
    @test  isapprox(weights[full_weight_indices[1]], 0.333333, rtol=1e-5)  == true

    @test interpolate(grid, [1,2,3,4], [1,1]) == 1.0
    @test maskedInterpolate(grid, [1,2,3,4], [1,1], BitArray([false, false, false, false])) == 1.0
    @test isnan(maskedInterpolate(grid, [1,2,3,4], [1,1], BitArray([true, false, false, false])))

    @test isapprox(interpolate(grid, [1,2,3,4], [1.5,3]), 1.66666666, rtol=1e-5)
    @test maskedInterpolate(grid, [1,2,3,4], [1.5,3], BitArray([true, false, false, false])) == 3
    @test maskedInterpolate(grid, [1,2,3,4], [1.5,3], BitArray([false, false, true, false])) == 1
end
@testset "Simplex Impl" begin
    test_simplex_implemented()
end

# check whether rectangle & simplex expect sorted order of cutpoints on dimensions
function test_ordering(grid)
    @test ind2x(grid,1) == [2,18] # 1
    @test ind2x(grid,2) == [5,18] # 2
    @test ind2x(grid,3) == [2,15] # 3
    @test ind2x(grid,4) == [5,15] # 4
    @test ind2x(grid,5) == [2,12] # 5
    @test ind2x(grid,6) == [5,12] # 6

    return interpolate(grid, [1,2,3,4,5,6], [6, 20]) == 2.0
end

# tests the order of vertices returned by vertices()
# by comparing against ind2x for each unrolled index
function test_vertices_ordering(grid)

    grid_verts = @inferred vertices(grid)

    @test length(grid_verts) == length(grid)

    for i = 1 : length(grid)
        @test grid_verts[i] == ind2x(grid,i)
    end

    return true
end
@testset "Ordering" begin
    @test_throws ErrorException test_ordering( RectangleGrid([2,5], [18,15,12]) )
    @test_throws ErrorException test_ordering( SimplexGrid([2,5], [18,15,12]) )

    @test test_vertices_ordering( RectangleGrid([2,5], [12,15,18]) ) == true
    @test test_vertices_ordering( SimplexGrid([1,4,6,10], [8,12,17]) ) == true
end

@testset "NaN" begin
    r = GridInterpolations.RectangleGrid(1:4, 1:4)
    @test_throws DomainError interpolants(r, [2.5, NaN])
end

@testset "Iteration and Indexing" begin
    @testset "2d" begin
        g2 = GridInterpolations.RectangleGrid(1:3,4:5)
        pts = ([1.,4.],[2.,4.],[3.,4.],[1.,5.],[2.,5.],[3.,5.])
        @test length(g2) == length(pts)
        @test all(pts .== [x for x in g2])
        @test iterate(g2,7) === nothing
        @test all(pts .== [g2[ci] for ci in reshape(CartesianIndices((3,2)), length(g2))])
        @test all(pts .== reshape([g2[i,j] for i=1:3, j=1:2], length(g2)))
        @test_throws BoundsError g2[CartesianIndex(0,0)]
        @test_throws BoundsError g2[CartesianIndex(4,2)]
        @test_throws BoundsError g2[CartesianIndex(3,3)]
        @test_throws BoundsError g2[0,0]
        @test_throws BoundsError g2[4,2]
        @test_throws BoundsError g2[3,3]
        @test_throws BoundsError g2[1,1,2]
    end
    @testset "3d" begin
        g3 = GridInterpolations.RectangleGrid(1:3,4:5,6:10)
        pts = reshape([[x,y,z] for x=1.:3., y=4.:5., z=6.:10.], length(g3))
        @test length(g3) == length(pts)
        @test all(pts .== [x for x in g3])
        @test iterate(g3,31) === nothing
        @test all(pts .== [g3[ci] for ci in reshape(CartesianIndices((3,2,5)), length(g3))])
        @test all(pts .== reshape([g3[i,j,k] for i=1:3, j=1:2, k=1:5], length(g3)))
        @test_throws BoundsError g3[CartesianIndex(0,0,0)]
        @test_throws BoundsError g3[CartesianIndex(4,2,5)]
        @test_throws BoundsError g3[CartesianIndex(3,3,5)]
        @test_throws BoundsError g3[CartesianIndex(3,2,6)]
        @test_throws BoundsError g3[0,0,0]
        @test_throws BoundsError g3[4,2,5]
        @test_throws BoundsError g3[3,3,5]
        @test_throws BoundsError g3[3,2,6]
        if VERSION.major > 0  # this test is known to fail for julia v0.7
            @test_throws BoundsError g3[1,1]
        end
        @test_throws BoundsError g3[1,1,1,2]
    end
end

include(joinpath(@__DIR__, "..", "bench", "interpBenchmarks.jl"))
compareBenchmarks(4, 10, 100, quiet=false)

println("All tests complete")


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

function mat_vec(a,x)
    mat_vec_ = load_op_and_grad("./build/libMatVec","mat_vec")
    a,x = convert_to_tensor(Any[a,x], [Float64,Float64])
    mat_vec_(a,x)
end

# TODO: specify your input parameters
a0 = rand(10,10)
x0 = rand(10)
a = constant(a0)
u = mat_vec(a0,x0)
sess = Session(); init(sess)
@show run(sess, u)≈a0 * x0


using ADCME
using PyPlot
using AdFem
using MAT

m = 20
n = 20
h = 1/n
bdnode = bcnode("all", m, n, h)
xy = gauss_nodes(m, n, h)
xy_fem = fem_nodes(m, n, h)

not_bdnode = ones(Bool, (m+1)*(n+1))
not_bdnode[bdnode] .= false
not_bdnode = findall(not_bdnode)

σ0 = 0.0
prior = 0.01
if length(ARGS)==2
    global σ0 = parse(Float64, ARGS[1])
    global prior = parse(Float64, ARGS[2])
end

DIR = "uq_nn2-$σ0-$prior"
rm(DIR, force=true, recursive=true)
mkdir(DIR)

function poisson1d(κ)
    κ = compute_space_varying_tangent_elasticity_matrix(κ, m, n, h)
    K = compute_fem_stiffness_matrix1(κ, m, n, h)
    K, _ = fem_impose_Dirichlet_boundary_condition1(K, bdnode, m, n, h)
    rhs = compute_fem_source_term1(ones(4m*n), m, n, h)
    rhs[bdnode] .= 0.0
    sol = K\rhs
end

function κfc(x, θ)
    xc = fem_to_fvm(x, m, n, h)
    xc = reshape(repeat(xc, 1, 4), (-1,))
    @assert length(xc)==4*m*n
    0.1 + abs(squeeze(fc(reshape(xc, (-1,1)), [20,20,20,1], θ)))
end

function param(x, θ)
    xc = fem_to_fvm(x, m, n, h)
    xc = reshape(repeat(xc, 1, 4), (-1,))
    return 0.1 + 1/(1+xc^2) + 100xc^2
end

function poisson1d_nn(fn)
    θ = missing
    function f(θ, u)
        κ = fn(u, θ)
        K = compute_space_varying_tangent_elasticity_matrix(κ, m, n, h)
        K = compute_fem_stiffness_matrix1(K, m, n, h)
        K, _ = fem_impose_Dirichlet_boundary_condition1(K, bdnode, m, n, h)
        rhs = K * u
        rhs = rhs[not_bdnode] - compute_fem_source_term1(ones(4m*n), m, n, h)[not_bdnode]
        rhs, gradients(rhs, u)
    end
    if fn==κfc
        θ = Variable(ae_init([1, 20, 20, 20, 1]))
    end
    u = newton_raphson_with_grad(f, constant(zeros((m+1)*(n+1))) ,θ)
    κnn = fn(u, θ)
    return u, κnn, θ
end

function jacobian(u, θ)
    κ = κfc(u, θ)
    K = compute_space_varying_tangent_elasticity_matrix(κ, m, n, h)
    K = compute_fem_stiffness_matrix1(K, m, n, h)
    K, _ = fem_impose_Dirichlet_boundary_condition1(K, bdnode, m, n, h)
    rhs = K * u
    rhs = rhs[not_bdnode] - compute_fem_source_term1(ones(4m*n), m, n, h)[not_bdnode]
    J = gradients(rhs, u)
    F = gradients(rhs, κ)
    -J\F
end

ADCME.options.newton_raphson.verbose=true
ADCME.options.newton_raphson.rtol = 1e-4
ADCME.options.newton_raphson.tol = 1e-4
y, κ, _ = poisson1d_nn(param)

sess = Session(); init(sess)
SOL = run(sess, y)

umax = maximum(SOL)
xc = LinRange(0, umax, 100)|>Array
u1 = @. 0.1 + 1/(1+xc^2) + 100xc^2

y, κnn, θ  = poisson1d_nn(κfc)

using Random; Random.seed!(233)
idx = fem_randidx(50, m, n, h)
obs = y[idx]
OBS = SOL[idx] .* ( 1. .+ σ0*randn(length(idx)))
loss = sum((obs-OBS)^2)

init(sess)
BFGS!(sess, loss, 200)

close("all")
figure(figsize=(10,4))
subplot(121)
visualize_scalar_on_fem_points(SOL, m, n, h)
subplot(122)
Yvalue = run(sess, y)
visualize_scalar_on_fem_points(Yvalue, m, n, h)
plot(xy_fem[idx,1], xy_fem[idx,2], "o", c="red", label="Observation")
legend()
savefig("$DIR/sol_compare.png")

close("all")
figure(figsize=(10,4))
subplot(121)
κvalue = run(sess, κ)
visualize_scalar_on_gauss_points(κvalue, m, n, h)
title("Exact \$K(x, y)\$")
subplot(122)
κnnvalue = run(sess, κnn)
visualize_scalar_on_gauss_points(κnnvalue, m, n, h)
title("Estimated \$K(x, y)\$")
savefig("$DIR/K_compare.png")

u0 = constant(LinRange(0,umax,100)|>Array )
z = 0.1 + abs(squeeze(fc(reshape(u0, (-1,1)), [20,20,20,1], θ)))
close("all")
plot(run(sess, u0), run(sess, z), "--")
plot(xc, u1, "-")
savefig("$DIR/ucurve.png")

u_est = run(sess, y)
H = jacobian(y, θ)
H = run(sess, H)[idx,:]
y_ = OBS
hs = run(sess, obs)
s = run(sess, κnn)

R = (prior)^2*diagm(0=>ones(length(obs)))
Q = (1e-2)^2*diagm(0=>ones(length(κnn)))
μ, Σ = uqnlin(y_, hs, H, R, s, Q)
# error()
σ = diag(Σ)

close("all")
u1 = fem_to_fvm(u_est, m, n, h)
z1 = μ[1:4:end]
σ1 = 2σ[1:4:end]
Is = sortperm(u1)
u1 = u1[Is]
z1 = z1[Is]
σ1 = σ1[Is]

xc = LinRange(0, umax, 100)|>Array
u2 = @. 0.1 + 1/(1+xc^2) + 100xc^2
plot(xc, u2, "-", linewidth=2, label="Referencee")

plot(run(sess, u0), run(sess, z), "--", linewidth=2, label="Estimated")
plot(u1, z1, "--", linewidth=2, label="Posterior Mean")

fill_between(u1, z1-σ1, z1+σ1, alpha=0.3, color="orange", label="Uncertainty Region")
xlim(minimum(u1), maximum(u1))
legend()
xlabel("\$u\$")
ylabel("\$K(u)\$")
savefig("$DIR/errorbar.png")

matwrite("$DIR/data.mat",
    Dict(
        "SOL"=>SOL,
        "Yvalue"=>Yvalue,
        "kappavalue"=>κvalue,
        "u_est"=>u_est,
        "H"=>H,
        "yobs"=>y_,
        "hs"=>hs,
        "s"=>s,
        "sigma"=>σ,
        "mu"=>μ
    ))


facts("next prime") do

  @fact next_prime(1) => 2
  @fact next_prime(2) => 3
  @fact next_prime(3) => 5
  @fact next_prime(4) => 5
  @fact next_prime(5) => 7
  @fact next_prime(6) => 7
  @fact next_prime(7) => 11
  @fact next_prime(8) => 11
  @fact next_prime(9) => 11
  @fact next_prime(10) => 11
  @fact next_prime(11) => 13
  @fact next_prime(13) => 17
  @fact next_prime(16) => 17
  @fact next_prime(17) => 19
  @fact next_prime(19) => 23
  @fact next_prime(49) => 53

  @fact next_prime(BigInt(49)) => 53

end


using Catalyst
using Test
using DelaySSAToolkit
rn = @reaction_network begin
    ρ, S+I --> E+I
    1/I, I --> R
    r, I --> R
end ρ r

m = get_reaction_idx(rn)
eqs = equations(rn)
@test m[:,2] == [1,3,2]
@test m[:,1] == eqs[[1,3,2]]


#
# A-B Call Center problem
#
using DiscreteEvents, Distributions, Random

mutable struct Caller
    id::Int
    t₁::Float64  # arrival time
    t₂::Float64  # beginning of service time
    t₃::Float64  # end of servive time
end

mutable struct Server
    id::Int
    S::Distribution  # service time distribution
    tbusy::Float64   # cumulative service time
end

function serve(c::Clock, s::Server, input::Channel, output::Vector{Caller}, limit::Int)
    call = take!(input)           # take a call
    call.t₂ = c.time              # record the beginning of service time
    @delay c s.S                  # delay for service time
    call.t₃ = c.time              # record the end of service time
    s.tbusy += call.t₃ - call.t₂  # log service time
    push!(output, call)           # hang up
    call.id ≥ limit && stop!(c)
end

function arrive(c::Clock, input::Channel, count::Vector{Int})
    count[1] += 1
    put!(input, Caller(count[1], c.time, 0.0, 0.0))
end

Random.seed!(123)
const N = 1000
const M_arr = Exponential(2.5)
const M_a   = Exponential(3)
const M_b   = Exponential(4)
count = [0]

clock = Clock()
input = Channel{Caller}(Inf)
output = Caller[]
s1 = Server(1, M_a, 0.0)
s2 = Server(2, M_b, 0.0)
@process serve(clock, s1, input, output, N)
@process serve(clock, s2, input, output, N)
@event arrive(clock,input,count) every M_arr
@run! clock 5000

# using Plots
# wt = [c.t₂ - c.t₁ for c in output]
# plot(wt, title="A-B-Call center ", xlabel="callers", ylabel="waiting time [min]", legend=false)
# savefig("ccenter.png")


export bezout_number

"""
    totaldegree_solutions(F::Vector{<:MP.AbstractPolynomialLike})

Returns an iterator of the solutions of the total degree startsystem of `F`.
"""
function totaldegree_solutions(F::Vector{AP}) where {AP<:MP.AbstractPolynomialLike}
    TotalDegreeSolutionIterator(MP.maxdegree.(F))
end
totaldegree_solutions(degrees::Vector{Int}) = TotalDegreeSolutionIterator(degrees)
totaldegree_solutions(F::TotalDegreeSystem, _) = TotalDegreeSolutionIterator(F.degrees)
function totaldegree_solutions(F::MultiHomTotalDegreeSystem, vargroups)
    MultiBezoutSolutionsIterator(F.D, F.C, vargroups)
end

"""
    TotalDegreeSolutionIterator(degrees)

Given the `Vector`s `degrees` `TotalDegreeSolutionIterator` enumerates all solutions
of the system
```math
\\begin{align*}
    z_1^{d_1} - 1 &= 0 \\\\
    z_1^{d_2} - 1 &= 0 \\\\
    &\\vdots \\\\
    z_n^{d_n} - 1 &= 0 \\\\
\\end{align*}
```
where ``d_i`` is `degrees[i]`.
"""
struct TotalDegreeSolutionIterator{Iter}
    degrees::Vector{Int}
    iterator::Iter
end
function TotalDegreeSolutionIterator(degrees::Vector{Int})
    iterator = Base.Iterators.product(map(d -> 0:d-1, degrees)...)
    TotalDegreeSolutionIterator(degrees, iterator)
end
function Base.show(io::IO, iter::TotalDegreeSolutionIterator)
    print(io, "TotalDegreeSolutionIterator(degrees=$(iter.degrees))")
end

function Base.iterate(iter::TotalDegreeSolutionIterator)
    indices, state = iterate(iter.iterator)
    _value(iter, indices), state
end
function Base.iterate(iter::TotalDegreeSolutionIterator, state)
    it = iterate(iter.iterator, state)
    it === nothing && return nothing
    _value(iter, it[1]), it[2]
end

function _value(iter::TotalDegreeSolutionIterator, indices)
    value = Vector{Complex{Float64}}()
    for (i, k) in enumerate(indices)
        push!(value, cis(2π * k / iter.degrees[i]))
    end
    value
end

Base.length(iter::TotalDegreeSolutionIterator) = length(iter.iterator)
Base.eltype(iter::Type{<:TotalDegreeSolutionIterator}) = Vector{Complex{Float64}}

##################
# Multihomogeneous
##################

struct MultiBezoutIndicesIterator{N}
    D::Matrix{Int} # Degree matrix
    k::NTuple{N, Int} # (projective) dimensions

    function MultiBezoutIndicesIterator(D::Matrix{Int}, k::NTuple{M, Int}) where {M}
        @assert size(D, 1) == M
        new{M}(D, k)
    end
end
MultiBezoutIndicesIterator(D::Matrix, k) = MultiBezoutIndicesIterator(D, tuple(k...))
MultiBezoutIndicesIterator(D::Matrix, groups::VariableGroups) = MultiBezoutIndicesIterator(D, projective_dims(groups))

Base.IteratorSize(::Type{MultiBezoutIndicesIterator}) = Base.SizeUnknown()
Base.eltype(::Type{MultiBezoutIndicesIterator}) = Tuple{Int, Vector{Int}}
function Base.iterate(iter::MultiBezoutIndicesIterator, state=[i for i=1:length(iter.k) for j=1:iter.k[i]])
    d = 0
    m, n = size(iter.D)
    p = state
    while d == 0
        state[1] > n && return nothing
        p, state = nextpermutation(1:m, state)
        let p=p, D=iter.D
            d = prod(i -> D[p[i], i], 1:n)::Int
        end
    end

    (d, p), state
end

# Adopted from Combinatorics.jl:
# https://github.com/JuliaMath/Combinatorics.jl/blob/8d9571402319799b29da2005a65b627e8771c1e4/src/permutations.jl#L47
function nextpermutation(m, state)
    n = length(state)
    perm = [m[state[i]] for i in 1:n]
    s = copy(state)
    i = n - 1
    while i>=1 && s[i] >= s[i+1]
        i -= 1
    end
    if i > 0
        j = n
        while j > i && s[i] >= s[j]
            j -= 1
        end
        s[i], s[j] = s[j], s[i]
        reverse!(s, i+1)
    else
        s[1] = n+1
    end
    return (perm, s)
end

"""
    bezout_number(F::MPPolys; variable_groups=[variables(F)], homvars=nothing, parameters=nothing)
    bezout_number(multidegrees, groups::VariableGroups)

Compute the multi-homogeneous bezout number associated to the given system and variable groups.
"""
bezout_number(D::Matrix, groups::VariableGroups) = bezout_number(D, projective_dims(groups))
bezout_number(D::Matrix, k) = sum(first, MultiBezoutIndicesIterator(D, k))
function bezout_number(F::MPPolys; parameters=nothing, variable_groups=nothing, homvars=nothing)
    if variable_groups === nothing
        return prod(maxdegrees(F; parameters=parameters))
    end

    vars = variables(F; parameters=parameters)
    hominfo = HomogenizationInformation(variable_groups=variable_groups, homvars=homvars)
    D = multidegrees(F, variable_groups)
    if ishomogeneous(F, hominfo; parameters=parameters)
        bezout_number(D, length.(variable_groups) .- 1)
    else
        bezout_number(D, length.(variable_groups))
    end
end

"""
    multi_bezout_coefficients(D, projective_dims)

Compute coefficients for a multi-bezout start system.
`projective_dims` is the projective dimensions of the projective spaces.
"""
function multi_bezout_coefficients(D::Matrix, k::NTuple{M, Int}) where M
    @assert size(D, 1) == M
    @assert size(D, 2) == sum(k)

    n = size(D, 2)
    C = zeros(n+M, n)
    for i in 1:n
        j = 1
        for kⱼ in k
            for l in 1:kⱼ
                C[j, i] = begin
                    if (i == l && i ≤ kⱼ)
                        1.0
                    elseif i ≠ l && i ≤ kⱼ
                        0.0
                    else
                        randn()
                    end
                end
                j += 1
            end
            C[j, i] = 1.0
            j += 1
        end
    end
    C
end

"""
    totaldegree_polysystem(multidegrees::Matrix, variable_groups::VariableGroups, coeffs)

The multi-homogeneous totaldegree start system described in [Wampler, 93]. Returns a tuple, the system
and the coefficients``c_{i,j,l}`` described in [Wampler, 93] as a `Matrix{Vector{Float64}}`.

[Wampler, 93]: An efficient start system for multi-homogeneous polynomial continuation (https://link.springer.com/article/10.1007/BF01385710).
"""
function totaldegree_polysystem(D::Matrix, vargroups::VariableGroups, C::Matrix)
    Z = variable_groups(vargroups)
    m, n = size(D)
    G = map(1:n) do i
        s = Ref(1)
        prod(1:m) do j
            dᵢⱼ = D[j,i]
            kⱼ = length(Z[j]) - 1
            if dᵢⱼ == 0
                s[] += kⱼ + 1
                return 1
            end
            bᵢⱼ = sum(1:kⱼ) do l
                cz = C[s[], i] * Z[j][l]
                s[] += 1
                cz
            end
            s[] += 1
            bᵢⱼ^dᵢⱼ - Z[j][end]^dᵢⱼ
        end
    end
    G
end

struct MultiBezoutSolutionsIterator{N}
    indices::MultiBezoutIndicesIterator{N}
    coeffs::Matrix{Float64}
    # precomputed stuff and cache
    roots_of_unity::Matrix{ComplexF64} # Lower Triangular matrix storing the precomputed roots
    A::NTuple{N, Matrix{Float64}}
    b::NTuple{N, Vector{ComplexF64}}
    ranges::NTuple{N, UnitRange{Int}}
end

function MultiBezoutSolutionsIterator(indices::MultiBezoutIndicesIterator, coeffs::Matrix)
    d_max = maximum(indices.D)
    roots_of_unity = zeros(ComplexF64, d_max, d_max)
    for i in 1:d_max, j in 0:i-1
        roots_of_unity[i, j+1] = cis(2π * j / i)
    end
    A = map(kⱼ -> zeros(kⱼ, kⱼ), indices.k)
    b = map(kⱼ -> zeros(ComplexF64, kⱼ), indices.k)
    r_stop = Ref(1)
    ranges = map(indices.k) do kⱼ
        range = r_stop[]:(kⱼ + r_stop[] - 1)
        r_stop[] += kⱼ + 1
        range
    end
    MultiBezoutSolutionsIterator(indices, coeffs, roots_of_unity, A, b, ranges)
end
function MultiBezoutSolutionsIterator(D::Matrix, C::Matrix, vargroups::VariableGroups)
    MultiBezoutSolutionsIterator(MultiBezoutIndicesIterator(D, vargroups), C)
end

function Base.show(io::IO, iter::MultiBezoutSolutionsIterator)
    print(io, "Solutions iterator for a multi-homogeneous start system")
end
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::MultiBezoutSolutionsIterator) = x

Base.length(iter::MultiBezoutSolutionsIterator) = bezout_number(iter.indices.D, iter.indices.k)
Base.IteratorSize(::Type{<:MultiBezoutSolutionsIterator}) = Base.SizeUnknown()
Base.eltype(::Type{MultiBezoutSolutionsIterator{N}}) where N = ProjectiveVectors.PVector{ComplexF64, N}

function Base.iterate(iter::MultiBezoutSolutionsIterator)
    (_, perm), indices_state = iterate(iter.indices)
    dᵢ = ntuple(i -> iter.indices.D[perm[i],i], length(perm))
    d_iter = Iterators.product(map(dᵢ -> 1:dᵢ, dᵢ)...)
    q, d_state = iterate(d_iter)

    x = compute_solution(iter, perm, q, dᵢ)
    x, (perm, indices_state, dᵢ, d_iter, d_state)
end

function Base.iterate(iter::MultiBezoutSolutionsIterator, state)
    perm, indices_state, dᵢ, d_iter, d_state = state

    next_d_state = iterate(d_iter, d_state)
    if next_d_state !== nothing
        q, d_state = next_d_state
        x = compute_solution(iter, perm, q, dᵢ)
        return x, (perm, indices_state, dᵢ, d_iter, d_state)
    end

    # We exhausted our current batch. move indices forward
    indices_newstate = iterate(iter.indices, indices_state)
    # check whether we are completely done
    indices_newstate === nothing && return nothing

    (_, perm), indices_state = indices_newstate
    # Create new d_iter
    new_dᵢ::typeof(dᵢ) = let D=iter.indices.D, perm=perm
        ntuple(i -> D[perm[i],i], length(perm))
    end
    new_d_iter::typeof(d_iter) = Iterators.product(map(dᵢⱼ -> 1:dᵢⱼ, new_dᵢ)...)
    q, d_state = iterate(new_d_iter)
    x = compute_solution(iter, perm, q, new_dᵢ)
    return x, (perm, indices_state, new_dᵢ, new_d_iter, d_state)
end

function compute_solution(iter::MultiBezoutSolutionsIterator, perm, q, dᵢ)
    m, n = size(iter.indices.D)
    # for each variable group we have to solve a linear system
    t = 1
    data = zeros(ComplexF64, n + m)
    for j=1:m
        kⱼ = iter.indices.k[j]
        Aⱼ = iter.A[j]
        bⱼ = iter.b[j]
        s = 1
        for i=1:n
            perm[i] == j || continue
            range = iter.ranges[j]
            for (t, l)  in enumerate(range)
                Aⱼ[s, t] = iter.coeffs[l, i]
            end
            bⱼ[s] = iter.roots_of_unity[dᵢ[i], q[i]]
            s += 1
        end

        LinearAlgebra.ldiv!(LinearAlgebra.lu!(Aⱼ), bⱼ)

        data[t:(t+kⱼ-1)] .= bⱼ
        data[t+kⱼ] = 1
        t += kⱼ + 1
    end

    x = ProjectiveVectors.PVector(data, iter.indices.k)
end


function part1(input, preamble_length_str = "25")
    preamble_length = parse(Int, preamble_length_str)
    numbers = parse.(Int, readlines(input))
    return find_invalid(numbers, preamble_length)
end

function part2(input, preamble_length_str = "25")
    preamble_length = parse(Int, preamble_length_str)
    numbers = parse.(Int, readlines(input))
    invalid = find_invalid(numbers, preamble_length)
    for i = 1:length(numbers)
        partial_sum = numbers[i]
        j = i + 1
        while j <= length(numbers)
            partial_sum += numbers[j]
            if partial_sum == invalid
                return sum(extrema(numbers[i:j]))
            elseif partial_sum > invalid
                break
            end
            j += 1
        end
    end
end

function find_invalid(numbers, preamble_length)
    for i = (preamble_length + 1):length(numbers)
        if !is_sum_of_two(numbers[i], numbers[(i - preamble_length):(i - 1)])
            return numbers[i]
        end
    end
    @assert false
end

function is_sum_of_two(n, v)
    for i = 1:length(v)
        for j = (i + 1):length(v)
            v[i] + v[j] == n && return true
        end
    end
    return false
end


# BL solution
# plotting the results of water-flooding oil recovery
# optimizing the rel-perm parametrs using the recovery data
using Roots, Dierckx, Plots, JLD
# include("rel_perms.jl")

# plot the fractional flow functions for a water-wet and an oil-wet system
include("rel_perms_real.jl")
# Water-wet Mixed-wet Oil-wet
# Swi 0.2 0.2 0.2
# Sor 0.1 0.2 0.3
# Ew=Krw(Sor) 0.6 0.8 0.8
# Eo=Kro(Swi) 0.8 0.8 0.5
# Nw 3 2 3
# No 2 2 3
krw0_ww = 0.3
krw0_ow = 1.0
kro0_ww = 0.6
kro0_ow = 0.76
nw_ww = 2.4
nw_ow= 2.4
no_ww = 2.0
no_ow= 2.0
sor_ww=0.15
sor_ow=0.2
sor_mw=0.08
swc_ww=0.09
swc_ow=0.09
teta_ow=deg2rad(180)
teta_ww=deg2rad(0)
teta_mw=deg2rad(90)
teta_prime=teta_mw-π/2 # see my report
a=(sor_mw-0.5*(sor_ww-sor_ow)*cos(teta_mw)-0.5*(sor_ww+sor_ow))/(cos(teta_mw)^2-1)
b=0.5*(sor_ww-sor_ow)
c=0.5*(sor_ww+sor_ow)-a

SF=0.0 # 1.0 is water-wet, 0.0 is oil-wet
sw_plot_ow=linspace(swc_ow,1-sor_ow,100)
sw_plot_ww=linspace(swc_ww,1-sor_ww,100)

plot(size=(500,400), xtickfont = font(12, "Courier"), ytickfont=font(12, "Courier"),
  ylabel="krw, kro", xlabel="Water saturation", legendfont=font(12, "Courier"),
  guidefont=font(12, "Courier"))
plot!(sw_plot_ow, krw.(sw_plot_ow, krw0_ow, sor_ow, swc_ow, nw_ow),
  xlims=(0.0,1.0), linewidth=3, linestyle=:dash, linecolor=:red,
  label="krw, Oil-wet")
plot!(sw_plot_ow, kro.(sw_plot_ow, kro0_ow, sor_ow, swc_ow, no_ow),
  xlims=(0.0,1.0), linewidth=3, linestyle=:dash, linecolor=:red,
  label="kro, Oil-wet")
plot!(sw_plot_ww, krw.(sw_plot_ww, krw0_ww, sor_ww, swc_ww, nw_ww),
  xlims=(0.0,1.0), linewidth=1, linestyle=:solid, linecolor=:blue,
  label="krw, Water-wet")
plot!(sw_plot_ww, kro.(sw_plot_ww, kro0_ww, sor_ww, swc_ww, no_ww),
  xlims=(0.0,1.0), linewidth=1, linestyle=:solid, linecolor=:blue,
  label="kro, Water-wet")
savefig("two_rel_perms.png")

# plot the recovery factors for a water-wet, a mixed-wet, and an oil-wet system
include("frac_flow_funcs.jl")
L_res=1000 # m
perm_res=0.001e-12 # Darcy
inj_pv=5.0 # pore volumes
poros_res=0.35 # reservoir porosity
inj_vel=1.0/(3600*24) # m/s
mu_oil=2e-3 # oil viscosity
mu_water=1e-3 # water viscosity

# water-wet reservoir
(xt_shock_ww, sw_shock_ww, xt_prf_ww, sw_prf_ww, t_ww, p_inj_ww, R_oil_ww)=
frac_flow_wf(muw=mu_water, muo=mu_oil, ut=inj_vel, phi=poros_res,
  k=perm_res, swc=swc_ww, sor=sor_ww, kro0=kro0_ww, no=no_ww, krw0=krw0_ww,
  nw=nw_ww, sw0=swc_ww, sw_inj=1.0, L=L_res, pv_inj=inj_pv)
# oil-wet reservoir
(xt_shock_ow, sw_shock_ow, xt_prf_ow, sw_prf_ow, t_ow, p_inj_ow, R_oil_ow)=
frac_flow_wf(muw=mu_water, muo=mu_oil, ut=inj_vel, phi=poros_res,
  k=perm_res, swc=swc_ow, sor=sor_ow, kro0=kro0_ow, no=no_ow, krw0=krw0_ow,
  nw=nw_ow, sw0=swc_ow, sw_inj=1.0, L=L_res, pv_inj=inj_pv)

plot(t_ow/(24*3600), R_oil_ow, size=(400,300), linecolor=:black, linewidth=2,
  xtickfont = font(10, "Courier"), ytickfont=font(10, "Courier"),
  xlabel="time [day]", ylabel="Recovery factor", label="Oil-wet",
  legendfont=font(10, "Courier"), guidefont=font(12, "Courier"), ylims=(0.0,1.0))
# savefig("recovery_time_ow.png")
plot!(t_ww/(24*3600), R_oil_ww, linecolor=:blue, label="Water-wet", linewidth=2,
  linestyle=:dash)
# savefig("recovery_time_ww_ow.png")

# BL for different wettability
plot(size=(500,400), xtickfont = font(10, "Courier"), ytickfont=font(10, "Courier"),
  xlabel="t [day]", ylabel="Recovery factor", legendfont=font(10, "Courier"),
  guidefont=font(12, "Courier"))
for SF in 0.0:0.25:1.0
  krw0=krw0_ww*SF+krw0_ow*(1-SF)
  kro0=kro0_ww*SF+kro0_ow*(1-SF)
  #sor=sor_ww*SF+sor_ow*(1-SF)
  swc=swc_ww*SF+swc_ow*(1-SF)
  no= no_ww*SF+no_ow*(1-SF)
  nw= nw_ww*SF+nw_ow*(1-SF)
  teta=teta_ww*SF+teta_ow*(1-SF)
  # sor=(max(sor_ow, sor_ww)-sor_mw)*cos(teta-teta_prime)^2+sor_mw
  sor=a*cos(teta).^2+b*cos(teta)+c
  (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj, R_oil)=
  frac_flow_wf(muw=mu_water, muo=mu_oil, ut=inj_vel, phi=poros_res,
    k=perm_res, swc=swc, sor=sor, kro0=kro0, no=no, krw0=krw0,
    nw=nw, sw0=swc, sw_inj=1.0, L=L_res, pv_inj=inj_pv)
    plot!(t/(24*3600), R_oil, size=(500,400), linewidth=2,
      xtickfont = font(10, "Courier"), ytickfont=font(10, "Courier"),
      xlabel="time [day]", ylabel="Recovery factor", label="$SF Sal. Factor",
      legendfont=font(10, "Courier"), guidefont=font(12, "Courier"), ylims=(0.0,1.0))
    savefig("recovery_forced_imb_$SF-WW.png")
end
plot!()

# plot(t, R_oil)
# plot the imbibition Pc-Sw curve for different SF
plot(size=(500,400), xtickfont = font(10, "Courier"), ytickfont=font(10, "Courier"),
  xlabel="Sw", ylabel="Pc [Pa]", legendfont=font(10, "Courier"),
  guidefont=font(12, "Courier"))
for SF in 0.1:0.2:0.9
  krw0=krw0_ww*SF+krw0_ow*(1-SF)
  kro0=kro0_ww*SF+kro0_ow*(1-SF)
  #sor=sor_ww*SF+sor_ow*(1-SF)
  swc=swc_ww*SF+swc_ow*(1-SF)
  no= no_ww*SF+no_ow*(1-SF)
  nw= nw_ww*SF+nw_ow*(1-SF)
  teta=teta_ww*SF+teta_ow*(1-SF)
  # sor=(max(sor_ow, sor_ww)-sor_mw)*cos(teta-teta_prime)^2+sor_mw
  sor=a*cos(teta).^2+b*cos(teta)+c
  pce= 5e2
  sw_plot=linspace(0.0,1.0,100)

  pc=pc_imb.(sw_plot, pce, swc, sor, teta=teta)
  plot!(sw_plot, pc, xlims=(0.0,1.0), ylims=(-2000, 2000), linewidth=2,
  label="$SF Sal. Factor")
  savefig("pc_sw_$SF-WW.png")
end
plot!()


using ValidatedNumerics
#
# HornerInt initially stores a list of coefficients of a polynomial, a evaluation point,
# and optionally a corresponding list of absolute errors of the coefficients.
# On demand it modifies the values in order to obtain the values of the selected derivatives
# at the given point and the corresponding error value.
#

if ! isdefined(:HornerInt)
type HornerInt{S<:AbstractFloat, T<:AbstractFloat}
    x::Interval{S}						# evaluation spot
    coeff::AbstractVector{Interval{T}}	# polynomial coefficients
    valid::Int						# evaluation status
end
end

# Create Horner table for evaluation at point x.
# Provide polynomial coefficients and optionally absolute errors for x and the coefficients.
function HornerInt{S<:Number,T<:Number,U<:AbstractFloat,V<:AbstractFloat}(xdown::S, xup::T, coeffdown::AbstractVector{U}, coeffup::AbstractVector{V})

    if coeffup == []
        coeffup = copy(coeffdown)
    end
    xdown, xup = min(xdown, xup), max(xdown, xup)
    coeffdown, coeffup = min(coeffdown, coeffup), max(coeffdown, coeffup)
    if length(coeffdown) != length(coeffup) throw(ArgumentError("coeffup length wrong")) end
    R = promote_type(S,T,U,V)
	HornerInt(R(xdown), R(xup), map(R,coeffdown), map(R,coeffup), -1)
end

# Evaluate Horner table up to and including m deviations.
# return the m'th deviation of the polynomial and the error estimation
# If m >= degree, all deviations ot the polynomial have been calculated at x.
# Further calls just return the cached results.
function call{S}(h::HornerInt{S}, m::Int)
    if m < 0 throw(ArgumentError("index $m is not >= 0")) end
	n = length(h.coeffdown) - 1
    if m > h.valid
        if h.xdown != 0 || h.xup != 0
            r = rounding(S)
            setronding(S, RoundDown)
            b, bu = h.coeffdown, h.coeffup
            x, xu = h.xdown, h.xup
            for k = max(h.valid+1,0):min(m, n-1)
                for j = 1:n-k
                    b[j+1], bu[j+1] = intmin((b[j+1], bu[j+1), (b[j], bu[j]), (x, xu))
                end
            end
            setRounding(S, r)
        end
        h.valid = m
    end
    m <= n ? (h.coeffdown[n+1-m], h.coeffup[n+1-m]) : (zero(S), zero(S))
end

#
# evaluate polynomial f and all derivatives at given approximate root-values z
#
function horner_analysis{S<:Number,T<:Number}(z::AbstractVector{S}, f::AbstractVector{T})
    thresh = 1000.0
    m = length(z)
    n = length(z)-1
    if m > length(f)-1
        throw(ArgumentError("polynomial of degree $n with $m zeros"))
    end
	h = Array{Horner}(m)
    for i = 1:m
        h[i] = Horner(z[i], f, eps(abs(z[i])) * 100, map(eps, f) * 100)
    end
	for i = 1:m
        println("root $i: $(z[i])")
        mult = 0
        search = true
		for k = 0:n
            hv, he = h[i](k)
            deviation = abs(hv / (he + eps(zero(real(S)))))
            @printf("%d %d %6.2g %8.8g %8.8g\n", i, k, deviation, abs(hv), he)
            if search && deviation < thresh
                mult = k + 1
            else
                search = false
            end
        end
        println("root $i: $(z[i]) estimated multiplicity: $mult")
    end
end


using CuArrays
using NNlib: logsoftmax, logσ

# Cost functions

mse(ŷ, y) = sum((ŷ .- y).^2) * 1 // length(y)

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Nothing)
  return -sum(y .* log.(ŷ)) * 1 // size(y, 2)
end

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Number)
  return -sum(y .* log.(ŷ)) .* weight * 1 // size(y, 2)
end

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::AbstractVector)
  return -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2)
end

crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight=nothing) = _crossentropy(ŷ, y, weight)

function logitcrossentropy(logŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1)
  return -sum(y .* logsoftmax(logŷ) .* weight) * 1 // size(y, 2)
end

"""
    binarycrossentropy(ŷ, y; ϵ=eps(ŷ))

Return `-y*log(ŷ + ϵ) - (1-y)*log(1-ŷ + ϵ)`. The ϵ term provides numerical stability.

    julia> binarycrossentropy.(σ.([-1.1491, 0.8619, 0.3127]), [1, 1, 0.])
    3-element Array{Float64,1}:
    1.4244
    0.352317
    0.86167
"""
binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)

# Re-definition to fix interaction with CuArrays.
CuArrays.@cufunc binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)

"""
    logitbinarycrossentropy(logŷ, y)

`logitbinarycrossentropy(logŷ, y)` is mathematically equivalent to `binarycrossentropy(σ(logŷ), y)`
but it is more numerically stable.

    julia> logitbinarycrossentropy.([-1.1491, 0.8619, 0.3127], [1, 1, 0.])
    3-element Array{Float64,1}:
     1.4244
     0.352317
     0.86167
"""
logitbinarycrossentropy(logŷ, y) = (1 - y)*logŷ - logσ(logŷ)

# Re-definition to fix interaction with CuArrays.
CuArrays.@cufunc logitbinarycrossentropy(logŷ, y) = (1 - y)*logŷ - logσ(logŷ)

"""
    normalise(x::AbstractArray; dims=1)

Normalises `x` to mean 0 and standard deviation 1, across the dimensions given by `dims`. Defaults to normalising over columns.

    julia> a = reshape(collect(1:9), 3, 3)
    3×3 Array{Int64,2}:
     1  4  7
     2  5  8
     3  6  9

    julia> normalise(a)
    3×3 Array{Float64,2}:
     -1.22474  -1.22474  -1.22474
      0.0       0.0       0.0
      1.22474   1.22474   1.22474

    julia> normalise(a, dims=2)
    3×3 Array{Float64,2}:
     -1.22474  0.0  1.22474
     -1.22474  0.0  1.22474
     -1.22474  0.0  1.22474
"""
function normalise(x::AbstractArray; dims=1)
  μ′ = mean(x, dims = dims)
  σ′ = std(x, dims = dims, mean = μ′, corrected=false)
  return (x .- μ′) ./ σ′
end


using COESA
using Test

M0 = COESA.M0

# altitude conversions
@test isapprox(round(COESA.geopotential_altitude(0)), 0)
@test isapprox(round(COESA.geopotential_altitude(-5000)), -5004)
@test isapprox(round(COESA.geopotential_altitude(85_500)), 84_365)
@test isapprox(round(COESA.geopotential_altitude(1_000_000)), 864_071)

# out of bounds errors
@test_throws ErrorException atmosphere(prevfloat(-5_000.0))
@test_throws ErrorException atmosphere(nextfloat(1_000_000.0))

# b = 0, Z < 0
Z = -5_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 320.676)
@test isapprox(round(pressure(atmos), sigdigits=5), 1.7776e5)
@test isapprox(round(density(atmos), sigdigits=5), 1.9311)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 358.99)

# b = 0, Z = 0
Z = 0

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 288.150)
@test isapprox(round(pressure(atmos), sigdigits=6), 101325)
@test isapprox(round(density(atmos), sigdigits=5), 1.2250)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 340.29)

# b = 0, Z > 0
Z = 5_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 255.676)
@test isapprox(round(pressure(atmos), sigdigits=5), 5.4048e4)
@test isapprox(round(density(atmos), sigdigits=5), 7.3643e-1)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 320.55)

# b = 1
Z = 15_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 216.650)
@test isapprox(round(pressure(atmos), sigdigits=4), round(1.2111e4, sigdigits=4))
@test isapprox(round(density(atmos), sigdigits=5), 1.9476e-1)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 295.07)

# b = 2
Z = 25_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 221.552)
@test isapprox(round(pressure(atmos), sigdigits=5), 2.5492e3)
@test isapprox(round(density(atmos), sigdigits=5), 4.0084e-2)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 298.39)

# b = 3
Z = 40_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 250.350)
@test isapprox(round(pressure(atmos), sigdigits=5), 2.8714e2)
@test isapprox(round(density(atmos), sigdigits=5), 3.9957e-3)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 317.19)

# b = 4
Z = 50_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 270.650)
@test isapprox(round(pressure(atmos), sigdigits=5), 7.9779e1)
@test isapprox(round(density(atmos), sigdigits=5), 1.0269e-3)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 329.80)

# b = 5
Z = 60_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 247.021)
@test isapprox(round(pressure(atmos), sigdigits=4), round(2.1958e1, sigdigits=4))
@test isapprox(round(density(atmos), sigdigits=5), 3.0968e-4)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 315.07)

# b = 6
Z = 75_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=5), 28.964)
@test isapprox(round(temperature(atmos), sigdigits=6), 208.399)
@test isapprox(round(pressure(atmos), sigdigits=5), 2.3881)
@test isapprox(round(density(atmos), sigdigits=5), 3.9921e-5)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 289.40)

# b = 7
# The tables show values of M that have a discontinuity at Z = 86km, but our
# routine applies the suggested blending
Z = 85_000

atmos = atmosphere(Z)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=6), round(M0 * 0.999694, sigdigits=6))
@test isapprox(round(temperature(atmos), sigdigits=6), round(188.893 / M0 * mean_molecular_weight(atmos), sigdigits=6))
@test isapprox(round(pressure(atmos), sigdigits=5), 4.4568e-1)
@test isapprox(round(density(atmos), sigdigits=4), round(8.2196e-6, sigdigits=4))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 275.52)

# b = 8, on a breakpoint
Z = 86_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 186.87)
@test isapprox(round(pressure(atmos), sigdigits=5), 3.7338e-1)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=4), 28.95)
@test isapprox(round(density(atmos), sigdigits=3), round(6.958e-6, sigdigits=3))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 8, in between breakpoints
Z = 86_500

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 186.87)
@test isapprox(round(pressure(atmos), sigdigits=5), 3.4163e-1)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=4), 28.95)
@test isapprox(round(density(atmos), sigdigits=4), 6.366e-6)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 8, in between breakpoints
Z = 100_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 195.08)
@test isapprox(round(pressure(atmos), sigdigits=2), round(3.2011e-2, sigdigits=2))
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=4), 28.40)
@test isapprox(round(density(atmos), sigdigits=2), round(5.604e-7, sigdigits=2))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 9
Z = 115_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 300.00)
@test isapprox(round(pressure(atmos), sigdigits=5), 4.0096e-3)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=4), 26.68)
@test isapprox(round(density(atmos), sigdigits=4), 4.289e-8)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 10
Z = 200_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 854.56)
@test isapprox(round(pressure(atmos), sigdigits=5), 8.4736e-5)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=4), 21.30)
@test isapprox(round(density(atmos), sigdigits=3), round(2.541e-10, sigdigits=3))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 11
Z = 750_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=5), 999.99)
@test isapprox(round(pressure(atmos), sigdigits=5), 2.2599e-8)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=3), 6.58)
@test isapprox(round(density(atmos), sigdigits=3), round(1.788e-14, sigdigits=3))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 11, in between breakpoints
Z = 985_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=6), 1000.00)
@test isapprox(round(pressure(atmos), sigdigits=3), round(7.9185e-9, sigdigits=3))
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=3), 3.99)
@test isapprox(round(density(atmos), sigdigits=3), round(3.797e-15, sigdigits=3))
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)

# b = 12
Z = 1_000_000

atmos = atmosphere(Z)
@test isapprox(round(temperature(atmos), sigdigits=6), 1000.00)
@test isapprox(round(pressure(atmos), sigdigits=5), 7.5138e-9)
@test isapprox(round(mean_molecular_weight(atmos), sigdigits=3), 3.94)
@test isapprox(round(density(atmos), sigdigits=4), 3.561e-15)
@test isapprox(round(speed_of_sound(atmos), sigdigits=5), 274.10)


using InteractionNetworkModels, BenchmarkTools, Distances

d_lcs = LCS()

x,y = (rand(1:3,10), rand(1:3, 10))

@btime d(x,y)

d_fast = FastLCS(100)
@btime d_fast(x,y)

d = MatchingDist(LCS())
d_fast = FastMatchingDist(LCS(), 50)
x,y = (

    [rand(1:3, rand(4:6)) for i in 1:rand(4:5)],
    [rand(1:3, rand(4:6)) for i in 1:rand(4:5)]
)
C = zeros(100,100)
@btime d(x,y)
@btime d_fast(x,y)
@btime matching_dist_with_memory!(x,y,d_lcs,C)

@btime pairwise_inbounds(d_lcs, x,y)
@btime Distances.pairwise(d_lcs, x,y)

d = MatchingDist(LCS())
d_norm = Normalised(d)

d_norm([[1,2,1]], [[1,2,3,4,5,4],[1]])

typeof(Normalised(d)) <: InteractionSetDistance

supertype(typeof(Normalised(d)))
supertype(typeof(Euclidean()))


using Distributed
using SharedArrays

@everywhere using LatentOrderBookModel

###########
# boltz_const      [0, 20]
# sigma            [0, 0.1]
# D                [0, 20]
# eta              [0, 0.1]
# lambda           [0, 20]
# mu               [0, 1]
# alpha            [0, 2.0]
###########

# param1_name::String
# param1_values::Array{Any,1}
# param2_name::String
# param2_values::Array{Any,1}
# surface_points::Int64
# replications::Int64
# params::Dict{String, Any}
# log_prices::Array{Float64,1}

cpus_allocated = size(Sys.cpu_info(),1)
println("$cpus_allocated cores allocated for job")

replications = 10
surface_points = 32
println("Number of price paths generated: $(surface_points^2*replications)")
param1_name = "α"
param1_values = repeat(repeat(collect(range(1.0, 10.0, length=surface_points)),inner=surface_points),replications)

param2_name = "ν"
param2_values = repeat(repeat(collect(range(0.0001, 0.0003, length=surface_points)),outer=surface_points),replications)
println("Creating surface plot for parameter $param1_name [1.0, 10.0] and parameter $param2_name [0.0001, 0.0003]")

fixed_params=Dict()
fixed_params["T"]=2300
fixed_params["τ"]=20
fixed_params["initial_mid_price"]=238.745
fixed_params["n_spatial_points"]=4501
fixed_params["boltz_const"]=1.0
fixed_params["sample_std"]=7.415
fixed_params["σ"]=0.001
fixed_params["D"]=5.0
fixed_params["ν"]=0.001
fixed_params["α"]=1.0
fixed_params["λ"]=1.0
fixed_params["μ"]=0.5

objective_surface = ObjectiveSurface(
    param1_name, param1_values,
    param2_name, param2_values,
    surface_points, replications,
    fixed_params)

random_seed = Int(rand(UInt32))
price_paths = @time objective_surface(random_seed, true)
println("Random seed is $random_seed")
println(size(price_paths))
io = open("alpha_vs_nu_price_paths.csv", "w")
for i in 1:size(price_paths,1)
    write(io, join(Array(price_paths[i,:]), ", "), "\n")
end
close(io)

#
# include("preprocess_bar_data.jl")
# include("mom_basic.jl")
#
# bar_data = process_bar_data("Original_Price_Bars_2300.csv")
# log_prices = bar_data.log_price
#
# num_bootstap = 20000
# b = 20
# W = weight_matrix(4547, log_prices, b, num_bootstap)
#
# function redim_pp(price_paths, T, replications, surface_points)
#     price_paths_new_dim = Array{Float32, 3}(undef, T,replications, surface_points^2)
#     for i in 1:surface_points^2
#         for k in 1:replications
#             price_paths_new_dim[:,k,i] = price_paths[:,i + (k-1)*surface_points^2]
#         end
#
#     end
#     return price_paths_new_dim
# end
#
# price_paths_new_dim = redim_pp(price_paths, fixed_params["T"], replications, surface_points)
#
# function objective_value_surface_plot(price_paths, W, log_prices)
#     T, replications, surface_points_sq = size(price_paths)
#     objective_values = zeros(Float32, surface_points_sq)
#     objective_value_fn(price_path) = objective_value(price_path, log_prices, W)
#     for i in 1:surface_points_sq
#         objective_values[i] = objective_value_fn(price_paths[:,:,i])
#     end
#     return objective_values
# end
#
# objective_values = objective_value_surface_plot(price_paths_new_dim, W,
#     log_prices)
#
# x_s = param1_values[1:surface_points^2]
# y_s = param2_values[1:surface_points^2]
# z_s = objective_values[1:end]
# z_s[.!(isnan.(z_s))]
# z_max = log(maximum(z_s[.!(isnan.(z_s))]))
# z_min = log(minimum(z_s[.!(isnan.(z_s))]))
#
#
# using Plots
# pyplot()
# using PyPlot
# pygui(true)
#
# plt = Plots.plot(x_s, y_s, z_s, st=:surface, size=[1000,800],
#     xlabel=param1_name, ylabel=param2_name,
#     zlabel="objective value", camera=(0,30), show=true)
#
# using JLD
# save("D:/Documents/UCT 2019/dissertation/scripts/masters-dissertation/data/alpha_nu_price_paths_001.jld", "price_paths",
#     price_paths, "objective_surface", objective_surface, "objective_values", objective_values, compress=true)


# Replicate Huggett (1993)

# --------------------------------------------------------------------------- #
# Set parameters:
β = 0.9932
BoundSet = (-2.0, -4.0, -6.0, -8.0)

# Exogenous state variable grid
eH = 1.0
eL = 0.1
E = [eL; eH]

# Transition matrix: Pr(j|i) = Π(i,j) (columns sum to 1)
πHH = 0.925 # e_H | e_H
πHL = 0.5   # e_H | e_L
Π = [1 - πHL  πHL
    1 - πHH πHH]

# Guess initial price
q0 = 1.1

# --------------------------------------------------------------------------- #
# Solve for σ = 1.5
σ = 1.5

# Start loop over lower bound
for a_lb in BoundSet

    # Generate grid points for endogenous state variable
    a_max = 5
    a_size = 350
    A = range(a_lb, a_max, length = a_size)
    A = [A; ];

    (q, policy, Ψ) = step3(A, E, Π, q0, β, σ)

    println("- Solving the model for σ = $σ and lower bound of a = $a_lb.")
    println("  The optimal price is q = $q.")
    println("  ")

end

# --------------------------------------------------------------------------- #
# Solve for σ = 3
σ = 3

# Start loop over lower bound
for a_lb in BoundSet

    # Generate grid points for endogenous state variable
    a_max = 5
    a_size = 350
    A = range(a_lb, a_max, length = a_size)
    A = [A; ];

    (q, policy, Ψ) = step3(A, E, Π, q0, β, σ)

    println("- Solving the model for σ = $σ and lower bound of a = $a_lb.")
    println("  The optimal price is q = $q.")
    println("  ")

end

# --------------------------------------------------------------------------- #
# Graphs
q0 = 1.1
σ = 1.5
a_lb = -2.0
a_max = 5
a_size = 350
A = range(a_lb, a_max, length = a_size)
A = [A; ];

(q, policy, Ψ) = step3(A, E, Π, q0, β, σ)

# Find endogenous a upper bar (use aUB+20 in graphs to see the intersection)
aUB = findfirst(policy[:,2] - A[:] .<= 0.001)

# Unconstrained plot for Policy Function
p1 = plot(A, [policy A], linewidth = 2, title = "Policy function",
    label = ["eL" "eH" "45°"], legend =:topleft)

# Constrained  plot for Policy Function
p1B = plot(A[1:aUB+20], [policy[1:aUB+20,:] A[1:aUB+20]], linewidth = 2,
    title = "Policy function", label = ["eL" "eH" "45°"], legend =:topleft)

# Unconstrained plot for Stationary Distribution Ψ
p2 = plot(A, Ψ, linewidth = 2, title = "Stationary distribution",
    label = ["eL" "eH"], legend =:topleft)

# Constrained plot for Stationary Distribution Ψ
p2B = plot(A[1:aUB+20], Ψ[1:aUB+20,:], linewidth = 2,
    title = "Stationary distribution", label = ["eL" "eH"], legend =:topleft)

# Find rescaling value
scale = sum(Ψ[end, :])
F = Ψ ./ scale

# Unconstrained plot for Stationary Distribution Ψ
p3 = plot(A, F, linewidth = 2, title = "Stationary distribution",
    label = ["eL" "eH"], legend =:topleft)

# Constrained plot for Stationary Distribution Ψ
p3B = plot(A[1:aUB+20], F[1:aUB+20], linewidth = 2,
    title = "Stationary distribution", label = ["eL" "eH"], legend =:topleft)

include("Step3.jl")
include("Step2.jl")
include("Step1Grid.jl")


"""
Created on Mon Apr 13 2020

@author: Yoann Pradat

Code supporting presentation on NMF.
"""

include("../../../src/VNMF.jl")
using .VNMF
using Distributions
using Random
include("../util/plot_cost.jl")

gr() # set the backend to GR

#### # 1. SIMULATE DATA
#### # ################################################################################################################

#### simulate matrix for testing NMF
F = 10
N = 25
K = 5
rng = MersenneTwister(123)

#### truncated-normal factors with noise
μ = 2
σ = √2
dist_W = truncated(Normal(μ, σ), 0, Inf)
dist_H = truncated(Normal(μ, σ), 0, Inf)

dist_E = Normal(μ, 0.2*σ)
dist_E_type = "Normal"

# α = μ
# θ = 1
# dist_E = Gamma(α, θ)
# dist_E_type = "Gamma"

W = rand(rng, dist_W, F, K)
H = rand(rng, dist_H, K, N)
E = rand(rng, dist_E, F, N)

V = W * H + E

fo = "../../plot/simulation/nmf_noise"
mkpath(fo)

#### global params
global_params = NMFParams(
    init          = :random,
    dist          = Uniform(0,1),
    rank          = K,
    max_iter      = 20_000,
    stopping_crit = :none,
    verbose       = true,
)

#### # 2. FEVOTTE-IDIER NMF ALGORITHMS WITH KNOWN RANK
#### # ################################################################################################################

list_res = Array{NMFResults{Float64}, 1}()
list_lab = Array{String, 1}()
list_ls  = Array{Symbol, 1}()
list_col = Array{String, 1}()

local_params = FIParams(
    scale_W = true,
    α_H     = 0,
    α_W     = 0,
)

####
#### 2.1 β divergences, FI algorithms
####

for (β, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params.β = β

    run_mm = true
    run_me = true
    run_h  = true

    if β == 2
        run_h = false
    elseif β == 1
        run_h  = false
        run_me = false
    end

    if run_mm
        local_params.alg = :mm

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{MM} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :dash)
        push!(list_col, col)
    end

    if run_me
        local_params.alg = :me

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{ME} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :solid)
        push!(list_col, col)
    end

    if run_h
        local_params.alg = :h

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{H} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :dot)
        push!(list_col, col)
    end

end

####
#### 2.2 PLOT
####

fn = joinpath(fo, "nmf_FI_noise_$(dist_E_type).pdf")
plot_cost(list_res, list_lab, list_ls, list_col, fn, relative=true)

#### # 3. MU NMF ALGORITHMS WITH KNOWN RANK
#### # ################################################################################################################

list_res = Array{NMFResults{Float64}, 1}()
list_lab = Array{String, 1}()
list_ls  = Array{Symbol, 1}()
list_col = Array{String, 1}()

####
#### 2.1 β divergences
####

for (β, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params = MUParams(
        β       = β,
        scale_W = true,
        div     = :β,
        alg     = :mu
    )

    push!(list_res, nmf_MU(V, global_params, local_params))
    push!(list_lab, raw"$\mathrm{MU} \quad \beta =" * "$(β)" * raw"$")
    push!(list_ls, :solid)
    push!(list_col, col)
end

####
#### 2.2 α divergences
####

for (α, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params = MUParams(
        α       = α,
        scale_W = true,
        div     = :α,
        alg     = :mu
    )

    push!(list_res, nmf_MU(V, global_params, local_params))
    push!(list_lab, raw"$\mathrm{MU} \quad \alpha =" * "$(α)" * raw"$")
    push!(list_ls, :dash)
    push!(list_col, col)
end

####
#### 2.3 PLOT
####

fn = joinpath(fo, "nmf_MU_noise_$(dist_E_type).pdf")
plot_cost(list_res, list_lab, list_ls, list_col, fn, relative=true)

#### # Cichocki NMF algorithms
#### #################################################################################################################

#params = CParams(
#    α             = 2,
#    max_iter      = convert(Int, 1e4),
#    stopping_crit = :none,
#    stopping_tol  = 1e-6,
#    α_H           = 0,
#    α_W           = 0,
#    l₁ratio_H     = 1,
#    l₁ratio_W     = 1,
#    μ_H           = 0,
#    μ_W           = 0,
#    α_sparse_H    = 0,
#    α_sparse_W    = 0,
#    ω             = 1,
#    verbose       = true,
#)
#
#
#res_α = nmf_α(
#    V,
#    W_init,
#    H_init,
#    params
#)


@testset "$(rpad("Symmetry",80))" begin

    for s in 1:10
        @test check_symmetry(TableauGauss(s)) == Array{Bool}(ones(s,s))
        @test issymmetric(TableauGauss(s))
    end

    for s in 2:10
        @test check_symmetry(TableauLobattoIIIA(s)) == Array{Bool}(ones(s,s))
        @test check_symmetry(TableauLobattoIIIB(s)) == Array{Bool}(ones(s,s))
        @test check_symmetry(TableauLobattoIIIC(s)) != Array{Bool}(ones(s,s))
        @test check_symmetry(TableauLobattoIIIC̄(s)) != Array{Bool}(ones(s,s))
        @test check_symmetry(TableauLobattoIIID(s)) == Array{Bool}(ones(s,s))
        @test check_symmetry(TableauLobattoIIIE(s)) == Array{Bool}(ones(s,s))
        @test check_symmetry(TableauRadauIA(s))     != Array{Bool}(ones(s,s))
        @test check_symmetry(TableauRadauIIA(s))    != Array{Bool}(ones(s,s))

        @test  issymmetric(TableauLobattoIIIA(s))
        @test  issymmetric(TableauLobattoIIIB(s))
        @test !issymmetric(TableauLobattoIIIC(s))
        @test !issymmetric(TableauLobattoIIIC̄(s))
        @test  issymmetric(TableauLobattoIIID(s))
        @test  issymmetric(TableauLobattoIIIE(s))
        @test !issymmetric(TableauRadauIA(s))
        @test !issymmetric(TableauRadauIIA(s))
    end

end


using PolyChaosODE,
    Distributions, DifferentialEquations, LaTeXStrings, LinearAlgebra, Parameters, Plots

### Linear ODE example ###

μ = 0.0
σ = sqrt(2.0) / 2

prob_dim = 4

struct ODEComp{T1,T2,T3}
    Λ::T1
    Q::T2
    error_vec::T3
end

function ode_func!(du, u, Λ, Q, error_vec, x, t)
    du .= Q' * diagm(Λ .+ error_vec .* x) * Q * u
end

function ode_func_closure!(du, u, comp::ODEComp, x, t)
    @unpack Λ, Q, error_vec = comp
    ode_func!(du, u, Λ, Q, error_vec, x, t)
end

P = randn(prob_dim, prob_dim)
Q = qr(P).Q
Λ = rand(Uniform(-0.200, -0.100), prob_dim)
error_vec = rand(Uniform(0.050, 0.100), prob_dim)

u0 = 10.0 .* randn(prob_dim)
tspan = (0.0, 5.0)
compartment = ODEComp(Λ, Q, error_vec)

function closed_form_variance(t)
    # z = Q * u = diag(e^(X_1), e^(X_2),..., e^(X_n)) * z0, with z0 = Q * u0 and X_i ~ N(Λ[i] * t, (σ * error_vec[i] * t)^2)
    return Q' .^ 2 * [
        (expm1((σ * error_vec[i] * t)^2)) *
        exp(2 * Λ[i] * t + (σ * error_vec[i] * t)^2) *
        (Q*u0)[i]^2 for i = 1:prob_dim
    ]
end

vars = Tuple(Normal(μ, σ) for i = 1:prob_dim)

stoch_galerkin_ode = StochGalerkinODE(ode_func_closure!, prob_dim, vars)
sol = stoch_galerkin_ode(u0, tspan, compartment; alg = VCABM())

step_t = 0.01
interval_t = tspan[1]:step_t:tspan[2]
sol_u = [sol(t) for t in interval_t]

pce_variance =
    compute_expectation_and_diag_variance(stoch_galerkin_ode, interval_t, sol_u)[2]
closed_form = [closed_form_variance(t) for t in interval_t]

error_arr =
    ((x, y) -> (sqrt.(x) .- sqrt.(y)) ./ sqrt.(y)).(pce_variance[2:end], closed_form[2:end])
abs_error_arr = (x -> abs.(x)).(error_arr)

plot_arr = mapreduce(transpose, vcat, abs_error_arr)

plot(
    interval_t[2:end],
    plot_arr;
    label = [L"err[\sigma_1]" L"err[\sigma_2]" L"err[\sigma_3]" L"err[\sigma_4]"],
    legend = :outertopright,
    size = (500, 400),
)
xlabel!(L"t")
ylabel!(L"\left\vert \frac{\sigma_{pce} - \sigma_{real}}{\sigma_{real}} \right\vert")

# savefig("linear_ode_example.png")


import Distributions

export Dists
const Dists = Distributions

function Measure(dist::Dists.Distribution{F,S}) where {F,S}
    X = Base.eltype(dist)
    DistributionMeasure{F,S,X}(dist)
end

struct DistributionMeasure{F, S, X} <: AbstractMeasure
    dist :: Dists.Distribution{F, S}
end

function basemeasure(μ::DistributionMeasure{F,S,X}) where {F, S <: Dists.Continuous, X}
    return Lebesgue(X)
end

export logdensity

function logdensity(μ::DistributionMeasure{F,S,X}, x::X) where {F, S, X}
    return Dists.logpdf(μ.dist, x)
end

function Base.rand(rng::AbstractRNG, μ::DistributionMeasure{F,S,X}) where {F, S, X}
    return rand(rng, μ.dist)
end

basemeasure(d::Dists.AbstractMvNormal) = Lebesgue(ℝ)^size(d)

function basemeasure(μ::Dists.Distribution{Dists.Univariate,Dists.Continuous})
    return Lebesgue(ℝ)
end

function basemeasure(μ::Dists.Distribution{Dists.Univariate,Dists.Discrete})
    return CountingMeasure(ℤ)
end

∫(::typeof(identity), ::Dists.Distribution) = 1.0

logdensity(μ::Dists.Distribution, x) = Dists.logpdf(μ,x)

density(μ::Dists.Distribution, x) = Dists.pdf(μ,x)

testvalue(d::Dists.Distribution) = rand(d)