using StochasticDiffEq, DiffEqBase, Random

function f(du,u,p,t)
  for i in 1:length(u)
    du[i] = (0.2/length(u))*u[i]
  end
end

function g(du,u,p,t)
  for i in 1:length(u)
    du[i] = (0.2/length(u))*u[i]
  end
end

function condition(u,t,integrator)
  1-maximum(u)
end

function callback_affect!(integrator)
  u = integrator.u
  resize!(integrator,length(u)+1)
  maxidx = findmax(u)[2]
  Θ = rand()
  u[maxidx] = Θ
  u[end] = 1-Θ
  nothing
end

callback = ContinuousCallback(condition,callback_affect!)

u0 = [0.2]
tspan = (0.0,100.0)
prob = SDEProblem(f,g,u0,tspan)

Random.seed!(3)
sol = solve(prob,SRIW1(),callback=callback)

#=
using Plots; pyplot()
p1 = plot(sol,vars=(0,1),plotdensity=10000,title="Amount of X in Cell 1")
scatter!(sol,denseplot=false)
p2 = plot(sol.t,map((x)->length(x),sol[:]),lw=3,
     ylabel="Number of Cells",xlabel="Time")
plot(p1,p2,layout=(2,1),size=(600,1000))
=#

Random.seed!(3)
sol = solve(prob,EM(),callback=callback,dt=1/4)
sol = solve(prob,RKMil(),callback=callback,dt=1/4)

function g(du,u,p,t)
  for i in 1:length(u)
    du[i] = (0.3/length(u))
  end
end
prob = SDEProblem(f,g,u0,tspan)

sol = solve(prob,SRA1(),callback=callback)


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 36d178f5-95fd-455f-954e-c85a5c92eebc
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="Plots", version="1"),
    ])
	using PlutoUI
	using Plots
	using Random
	PlutoUI.TableOfContents()
end

# ╔═╡ 8a1e5a8b-7fbd-4818-af91-a52df41e0e51
md"""
# Line plots

## Line plots (I)
"""

# ╔═╡ f27bb3ad-ba52-4334-9aac-a9fa48b07f47
x = 0:0.1:2pi

# ╔═╡ 7a692a81-a0fe-4d67-b6b1-697d976f89af
y1 = cos.(x)

# ╔═╡ 03e58a6b-02b0-4133-91f4-775a6aee6a39
y2 = sin.(x)

# ╔═╡ 9b2ea567-bdec-4211-a5d2-670a50b82c7f
begin
	p1 = plot(x, y1, c="blue", linewidth=3)
	plot!(p1, x, y2, c="red", line=:dash)
	title!(p1, "Trigonometric functions")
	xlabel!(p1, "angle")
	ylabel!(p1, "sin(x) and cos(x)")
	plot!(p1, xlims=(0,2pi), ylims=(-2, 2), size=(600, 600))
	p1
end

# ╔═╡ a8aa4f51-3e82-47d2-aa09-89ece6c7bbd4
begin
	p2 = plot(x, y1,
    		  c="blue",
    		  linewidth=3,
    		  title="Trigonometric functions",
    		  xlabel="angle",
    		  ylabel="sin(x) and cos(x)")
	plot!(p2, x, y2, c="red", line=:dash)

	plot!(p2, xlims=(0,2pi), ylims=(-2, 2), size=(600, 600))
end

# ╔═╡ e649e45c-7e3c-4b59-b0a4-7ca650077fac
md"""
## Line plots (II)
"""

# ╔═╡ 15d78933-4ac8-4c60-b080-e8d32a5022d2
begin
	Random.seed!(2018)
	time = 30
	walker1 = cumsum(randn(time))
	walker2 = cumsum(randn(time))
	walker3 = cumsum(randn(time))
	walker4 = cumsum(randn(time))
	walker5 = cumsum(randn(time))
	plot(1:time, [walker1 walker2 walker3 walker4 walker5],
    xlabel="time",
    ylabel="position",
    label=["walker1" "walker2" "walker3" "walker4" "walker5"],
    leg=:bottomleft)
end

# ╔═╡ be3813ec-be35-11eb-0f75-956c2874118b
md"""
# Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ Cell order:
# ╠═8a1e5a8b-7fbd-4818-af91-a52df41e0e51
# ╠═f27bb3ad-ba52-4334-9aac-a9fa48b07f47
# ╠═7a692a81-a0fe-4d67-b6b1-697d976f89af
# ╠═03e58a6b-02b0-4133-91f4-775a6aee6a39
# ╠═9b2ea567-bdec-4211-a5d2-670a50b82c7f
# ╠═a8aa4f51-3e82-47d2-aa09-89ece6c7bbd4
# ╠═e649e45c-7e3c-4b59-b0a4-7ca650077fac
# ╠═15d78933-4ac8-4c60-b080-e8d32a5022d2
# ╠═be3813ec-be35-11eb-0f75-956c2874118b
# ╠═36d178f5-95fd-455f-954e-c85a5c92eebc


using QXGraphDecompositions
using Test

import LightGraphs; lg = LightGraphs

include("LabeledGraph_tests.jl")
include("elimination_order_tests.jl")
include("tree_decomposition_tests.jl")
include("treewidth_deletion_tests.jl")


using DiffEqBase, OrdinaryDiffEq
using Plots; plotly()
############### Parameters

#=
cₕ = 15.0
cₖ = 10.0
κₕ = .2
κₖ = 7.0
dₕ = 1.0
dₖ = 1.0
nₕ = 2
nₖ = 2
m = 2
=#

cₕ = 3
cₖ = 5
κₕ = 1
κₖ = 1
dₕ = 0.4
dₖ = 0.4
nₕ = 2
nₖ = 2
m = 2
aₕ = 0.05
aₖ = 0.05

#=
const cₕ = 3
const cₖ = 5
const κₕ = 1
const κₖ = 1
const dₕ = 0.4
const dₖ = 0.4
const nₕ = 2
const nₖ = 2
const m = 2
const aₕ = 0.05
const aₖ = 0.05
=#

################# Setup the problem

type HoxKrox
  RAin::Float64
end
function (p::HoxKrox)(du,u,p,t)
  gₕ,gₖ = u
  du[1] .= (cₕ.*gₕ.^nₕ .+ (κₕ*p.RAin).^m) ./ (1 .+ cₕ.*gₕ.^nₕ .+ cₖ.*gₖ.^nₖ .+ (κₕ*p.RAin).^m) - dₕ.*gₕ
  du[2] .= (cₖ.*gₖ.^nₖ .+ (κₖ*p.RAin).^m) ./ (1 .+ cₕ.*gₕ.^nₕ .+ cₖ.*gₖ.^nₖ .+ (κₖ*p.RAin).^m) - dₖ.*gₖ
end
u0 = [0.5,2.0]
hk = HoxKrox(2.0)
tspan = (0.0,500.0)
prob = ODEProblem(hk,u0,tspan)

####################### Setup Bif values

SShox = Vector{Float64}()
SSkrox = Vector{Float64}()
SShox2 = Vector{Float64}()
SSkrox2 = Vector{Float64}()
####################### Get steady states at different RAin

RAins = 0:0.01:2.5

for RAin_scalar in RAins
  prob.f.RAin = RAin_scalar
  sol = solve(prob,Tsit5(),save_everystep=false)
  push!(SShox,sol[end][1])
  push!(SSkrox,sol[end][2])
end

prob.u0 = [2.0,0.5]

for RAin_scalar in RAins
  prob.f.RAin = RAin_scalar
  sol = solve(prob,Tsit5(),save_everystep=false)
  push!(SShox2,sol[end][1])
  push!(SSkrox2,sol[end][2])
end

p1 = plot(RAins,[SShox,SSkrox],title="Steady State Krox Start",xaxis="RAin",
     yaxis="Steady State Value",label=["Hox" "Krox"],
     lw = 3);

p2 = plot(RAins,[SShox2,SSkrox2],title="Steady State Hox Start",xaxis="RAin",
     yaxis="Steady State Value",label=["Hox" "Krox"],
     lw = 3);
display(plot(p1,p2))

#=
####################### Get steady states at different RAin

RAin = ra_sol[end][:,:,2][:,21:61]
x = X[5,21:61]
RAins = RAin[5,:]
SShox = Vector{Float64}()
SSkrox = Vector{Float64}()
SShox2 = Vector{Float64}()
SSkrox2 = Vector{Float64}()

prob.u0 = [0.5,2.0]

for RAin_scalar in RAins
  prob.f.RAin = RAin_scalar
  sol = solve(prob,Tsit5(),save_everystep=false)
  push!(SShox,sol[end][1])
  push!(SSkrox,sol[end][2])
end

p21 = plot(x,[SShox,SSkrox],title="Steady State Hox/Krox vs RAin",xaxis="x (Anterior -> Posterior)",
     yaxis="Steady State Value",label=["Hox" "Krox"],
     lw = 3)

prob.u0 = [2.0,0.5]

for RAin_scalar in RAins
  prob.f.RAin = RAin_scalar
  sol = solve(prob,Tsit5(),save_everystep=false)
  push!(SShox2,sol[end][1])
  push!(SSkrox2,sol[end][2])
end

p22= plot(x,[SShox2,SSkrox2],title="Steady State Hox/Krox vs RAin",xaxis="x (Anterior -> Posterior)",
     yaxis="Steady State Value",label=["Hox" "Krox"],
     lw = 3)

display(plot(p21,p22))
=#


"""
    VariableLinearAdvectionNonperiodicSemidiscretization(D, Di, a, split_form,
                                                         left_bc, right_bc)

A semidiscretization of the linear advection equation
    \$\\partial_t u(t,x) + \\partial_x ( a(x) u(t,x) ) = 0\$
with boundary conditions `left_bc(t)`, `right_bc(t)`.

`D` is an SBP derivative operator, `Di` an associated dissipation operator
or `nothing`, `a(x)` the variable coefficient, and `split_form::Union{Val(false), Val(true)}`
determines whether the canonical split form or the conservative form should be
used.
"""
struct VariableLinearAdvectionNonperiodicSemidiscretization{T,Derivative<:AbstractDerivativeOperator{T},
                                                            Dissipation,
                                                            SplitForm<:Union{Val{false}, Val{true}},
                                                            LeftBC, RightBC} <: AbstractSemidiscretization
    derivative::Derivative
    dissipation::Dissipation
    a::Vector{T}
    tmp1::Vector{T}
    tmp2::Vector{T}
    split_form::SplitForm
    left_bc::LeftBC
    right_bc::RightBC

    function VariableLinearAdvectionNonperiodicSemidiscretization(derivative::Derivative, dissipation::Dissipation, afunc, split_form::SplitForm, left_bc::LeftBC, right_bc::RightBC) where {T, Derivative<:AbstractDerivativeOperator{T}, Dissipation, SplitForm<:Union{Val{false}, Val{true}}, LeftBC, RightBC}
        if dissipation !== nothing
            @argcheck size(derivative) == size(dissipation) DimensionMismatch
            @argcheck grid(derivative) == grid(dissipation) ArgumentError
        end
        N = size(derivative, 2)
        a = compute_coefficients(afunc, derivative)
        tmp1 = Array{T}(undef, N)
        tmp2 = Array{T}(undef, N)
        new{T,Derivative,Dissipation,SplitForm,LeftBC,RightBC}(derivative, dissipation, a, tmp1, tmp2, split_form, left_bc, right_bc)
    end
end

function Base.show(io::IO, semi::VariableLinearAdvectionNonperiodicSemidiscretization)
    if get(io, :compact, false)
        print(io, "Semidiscretization of the linear advection equation (periodic)")
    else
        print(io, "Semidiscretization of the linear advection equation\n")
        print(io, "  \$ \\partial_t u(t,x) + \\partial_x ( a(x) u(t,x) ) = 0 \$ \n")
        print(io, "with nonperiodic boundaries using")
        if semi.split_form == Val{true}()
            print(io, " a split form and: \n")
        else
            print(io, " no split form and: \n")
        end
        print(io, semi.derivative)
        print(io, semi.dissipation)
    end
end

function godunov_flux_variablelinearadvection(uₗ::T, uᵣ::T, a::T) where {T<:Real}
    ifelse(a > 0, a*uₗ, a*uᵣ)
end

function (disc::VariableLinearAdvectionNonperiodicSemidiscretization)(du, u, p, t)
    @unpack a, tmp1, tmp2, derivative, dissipation, split_form, left_bc, right_bc = disc
    @boundscheck begin
        @argcheck length(u) == length(a)
        @argcheck length(u) == length(tmp1)
        @argcheck length(u) == length(tmp2)
        @argcheck length(u) == length(du)
    end

    # volume terms
    if typeof(split_form) <: Val{true}
        m1_2 = -one(eltype(u)) / 2

        ## a * D * u
        mul!(tmp1, derivative, u)
        @. du = m1_2 * a * tmp1
        ## u * D * a
        mul!(tmp1, derivative, a)
        @. du += m1_2 * u * tmp1
        ## D * a*u
        @. tmp2 = a*u
        mul!(tmp1, derivative, tmp2)
        @. du += m1_2 * tmp1
    else
        @. tmp2 = a * u
        mul!(tmp1, derivative, tmp2)
        @. du = -tmp1
    end

    # dissipation
    if dissipation !== nothing
        mul!(tmp1, dissipation, u)
        @. du += tmp1
    end

    # boundary conditions via Godunov's flux
    @inbounds fnum_left = godunov_flux_variablelinearadvection(left_bc(t), u[1], a[1])
    @inbounds du[1] += (fnum_left - a[1]*u[1]) / left_boundary_weight(derivative)
    @inbounds fnum_right = godunov_flux_variablelinearadvection(u[end], right_bc(t), a[end])
    @inbounds du[end] -= (fnum_right - a[end]*u[end]) / right_boundary_weight(derivative)

    nothing
end


function δ(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E}
    for s in species(N)
        @assert haskey(L, s)
    end

    this_l = filter(p -> p.first in species(N), L)
    tl = [this_l[s] for s in species(N; dims=1)]
    bl = [this_l[s] for s in species(N; dims=2)]
    D = tl .== bl'
    return D
end

"""
    Q(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E}

Modularity of a network and its partition. The second argument is a dictionary
where every species of `N` is associated to an `Int64` value representing the
identity of the module. This function returns the same value of bipartite
networks and their unipartite projection.

#### References

Barber, M.J., 2007. Modularity and community detection in bipartite networks.
Phys. Rev. E 76, 066102. https://doi.org/10.1103/PhysRevE.76.066102

Newman, M.E., 2006. Modularity and community structure in networks. Proceedings
of the national academy of sciences 103, 8577–8582.
"""
function Q(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E}
    for s in species(N)
        @assert haskey(L, s)
    end

    # Degrees
    dkin, dkout = degree_in(N), degree_out(N)
    kin = map(x -> dkin[x], species(N; dims=2))
    kout = map(x -> dkout[x], species(N; dims=1))

    # Value of m -- sum of weights, total number of int, ...
    m = links(N)

    # Null model
    kikj = (kout .* kin')
    Pij = kikj ./ m

    # Difference
    diff = N.A .- Pij

    # Diff × delta
    dd = diff .* δ(N,L)

    return sum(dd)/m
end

"""
    Qr(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E}

Realized modularity -- this function returns a value giving the proportion of
all links that are within the same module. Higher values reflect a more strongly
modular partition (whereas `Q` represents the deviation of modularity from the
random expectation).

#### References

Poisot, T., 2013. An a posteriori measure of network modularity. F1000Research
2. https://doi.org/10.12688/f1000research.2-130.v3
"""
function Qr(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E}
    for s in species(N)
        @assert haskey(L, s)
    end
    W = sum(N.A .* δ(N, L))
    B = links(N)
    return 2.0 * (W/B) - 1.0
end

function tidy_modules!(L::Dict{E,Int64}) where {E}
  l_values = sort(unique(collect(values(L))))
  for i in 1:length(l_values)
    for (k,v) in L
      if v == l_values[i]
        L[k] = i
      end
    end
  end
end


using NCDatasets
using Formatting

fn1 = "clim_LENS_B1850C5CN_005_gx3v7_TEMP.nc"
fn2 = "clim_LENS_B1850C5CN_005_gx3v7_zMLMML_TEMP.nc"
varname = "TEMP"

idx_i = 45   # longitude idx
idx_j = 10   # latitude  idx

Dataset(fn1, "r") do ds

    global Nx = ds.dim["Nx"]
    global Ny = ds.dim["Nx"]
    global Nz_old = ds.dim["Nz"]

    global lat = ds["lat"][:]
    global lon = ds["lon"][:]
    global zs_old = (ds["z_w_top"][:] + ds["z_w_bot"][:] ) / 2 / 100.0

    global data_old = replace(ds[varname][idx_i, idx_j, :, 1], missing=>NaN)

end

Dataset(fn2, "r") do ds

    global Nz_new = ds.dim["Nz"]

    global zs_new = ( ds["zs"][1:end-1] +  ds["zs"][2:end] ) / 2

    global data_new = replace(ds[varname][idx_i, idx_j, :, 1], missing=>NaN)

end

print("Loading PyPlot...")
using PyPlot
println("done.")

plt[:figure]()
plt[:plot](zs_old, data_old, "b-",  label="old data")
plt[:plot](zs_new, data_new, "r--", label="new data")
plt[:legend]()

plt[:title](format("(lat, lon) = ({:.2f}, {:.2f})", lat[idx_i, idx_j], lon[idx_i, idx_j]))
plt[:show]()


include("utilities.jl")

using SEAL
using Printf

"""
    example_rotation_bfv()

Perform some rotation on data encryped with the BFV scheme. This routine is based on the file
`native/examples/5_rotation.cpp` of the original SEAL library and should yield the exact same
output.

* [SEAL](https://github.com/microsoft/SEAL)
* [native/examples/5_rotation.cpp](https://github.com/microsoft/SEAL/blob/master/native/examples/5_rotation.cpp)
"""
function example_rotation_bfv()
  print_example_banner("Example: Rotation / Rotation in BFV")

  parms = EncryptionParameters(SchemeType.bfv)

  poly_modulus_degree = 8192
  set_poly_modulus_degree!(parms, poly_modulus_degree)

  set_coeff_modulus!(parms, coeff_modulus_bfv_default(poly_modulus_degree))
  set_plain_modulus!(parms, plain_modulus_batching(poly_modulus_degree, 20))

  context = SEALContext(parms)
  print_parameters(context)
  println()

  keygen = KeyGenerator(context)
  secret_key_ = secret_key(keygen)
  public_key_ = PublicKey()
  create_public_key!(public_key_, keygen)
  relin_keys_ = RelinKeys()
  create_relin_keys!(relin_keys_, keygen)
  encryptor = Encryptor(context, public_key_)
  evaluator = Evaluator(context)
  decryptor = Decryptor(context, secret_key_)

  batch_encoder = BatchEncoder(context)
  slot_count_ = slot_count(batch_encoder)
  row_size = div(slot_count_, 2)
  println("Plaintext matrix row size: ", row_size)

  pod_matrix = zeros(UInt64, slot_count_)
  pod_matrix[1] = 0
  pod_matrix[2] = 1
  pod_matrix[3] = 2
  pod_matrix[4] = 3
  pod_matrix[row_size + 1] = 4
  pod_matrix[row_size + 2] = 5
  pod_matrix[row_size + 3] = 6
  pod_matrix[row_size + 4] = 7

  println("Input plaintext matrix:")
  print_matrix(pod_matrix, row_size)

  plain_matrix = Plaintext()
  print_line(@__LINE__)
  println("Encode and encrypt.")
  encode!(plain_matrix, pod_matrix, batch_encoder)
  encrypted_matrix = Ciphertext()
  encrypt!(encrypted_matrix, plain_matrix, encryptor)
  println("    + Noise budget in fresh encryption: ",
          invariant_noise_budget(encrypted_matrix, decryptor),
          " bits")
  println()

  galois_keys = GaloisKeys()
  create_galois_keys!(galois_keys, keygen)

  print_line(@__LINE__)
  println("Rotate rows 3 steps left.")
  rotate_rows_inplace!(encrypted_matrix, 3, galois_keys, evaluator)
  plain_result = Plaintext()
  println("    + Noise budget after rotation: ",
          invariant_noise_budget(encrypted_matrix, decryptor),
          " bits")
  println("    + Decrypt and decode ...... Correct.")
  decrypt!(plain_result, encrypted_matrix, decryptor)
  decode!(pod_matrix, plain_result, batch_encoder)
  print_matrix(pod_matrix, row_size)

  print_line(@__LINE__)
  println("Rotate columns.")
  rotate_columns_inplace!(encrypted_matrix, galois_keys, evaluator)
  println("    + Noise budget after rotation: ",
          invariant_noise_budget(encrypted_matrix, decryptor),
          " bits")
  println("    + Decrypt and decode ...... Correct.")
  decrypt!(plain_result, encrypted_matrix, decryptor)
  decode!(pod_matrix, plain_result, batch_encoder)
  print_matrix(pod_matrix, row_size)

  print_line(@__LINE__)
  println("Rotate rows 4 steps right.")
  rotate_rows_inplace!(encrypted_matrix, -4, galois_keys, evaluator)
  plain_result = Plaintext()
  println("    + Noise budget after rotation: ",
          invariant_noise_budget(encrypted_matrix, decryptor),
          " bits")
  println("    + Decrypt and decode ...... Correct.")
  decrypt!(plain_result, encrypted_matrix, decryptor)
  decode!(pod_matrix, plain_result, batch_encoder)
  print_matrix(pod_matrix, row_size)
end

"""
    example_rotation_ckks()

Perform some rotation on data encryped with the CKKS scheme. This routine is based on the file
`native/examples/5_rotation.cpp` of the original SEAL library and should yield the exact same
output.

* [SEAL](https://github.com/microsoft/SEAL)
* [native/examples/5_rotation.cpp](https://github.com/microsoft/SEAL/blob/master/native/examples/5_rotation.cpp)
"""
function example_rotation_ckks()
  print_example_banner("Example: Rotation / Rotation in CKKS")

  parms = EncryptionParameters(SchemeType.ckks)

  poly_modulus_degree = 8192
  set_poly_modulus_degree!(parms, poly_modulus_degree)
  set_coeff_modulus!(parms, coeff_modulus_create(poly_modulus_degree, [40, 40, 40, 40, 40]))

  context = SEALContext(parms)
  print_parameters(context)
  println()

  keygen = KeyGenerator(context)
  public_key_ = PublicKey()
  create_public_key!(public_key_, keygen)
  secret_key_ = secret_key(keygen)
  relin_keys_ = RelinKeys()
  create_relin_keys!(relin_keys_, keygen)
  galois_keys_ = GaloisKeys()
  create_galois_keys!(galois_keys_, keygen)
  encryptor = Encryptor(context, public_key_)
  evaluator = Evaluator(context)
  decryptor = Decryptor(context, secret_key_)

  encoder = CKKSEncoder(context)

  slot_count_ = slot_count(encoder)
  println("Number of slots: ", slot_count_)
  input = collect(range(0.0, 1.0, length=slot_count_))
  println("Input vector:")
  print_vector(input, 3, 7)

  initial_scale = 2.0^50

  print_line(@__LINE__)
  println("Encode and encrypt.")
  plain = Plaintext()
  encode!(plain, input, initial_scale, encoder)
  encrypted = Ciphertext()
  encrypt!(encrypted, plain, encryptor)

  rotated = Ciphertext()
  print_line(@__LINE__)
  println("Rotate 2 steps left.")
  rotate_vector!(rotated, encrypted, 2, galois_keys_, evaluator)
  println("    + Decrypt and decode ...... Correct.")
  decrypt!(plain, rotated, decryptor)
  result = similar(input)
  decode!(result, plain, encoder)
  print_vector(result, 3, 7)

  return
end

"""
    example_rotation()

Perform some rotation on data encryped with the BFV and CKKS schemes. This routine is based on the
file `native/examples/5_rotation.cpp` of the original SEAL library and should yield the exact same
output.

* [SEAL](https://github.com/microsoft/SEAL)
* [native/examples/5_rotation.cpp](https://github.com/microsoft/SEAL/blob/master/native/examples/5_rotation.cpp)
"""
function example_rotation()
  print_example_banner("Example: Rotation")

  example_rotation_bfv()
  example_rotation_ckks()
end


module EmpiricalLikelihood

using DiffResults
using ForwardDiff
using JuMP
using Optim
using PyCall
using Statistics

include("HigherOrderKernels.jl")

const scipy_spatial = PyNULL()

function __init__()
    copy!(scipy_spatial, pyimport_conda("scipy.spatial", "scipy"))
end

export optimize_el, setup_jump_problem

const Kernel = EpanechnikovKernel{2}

function optimize_el(n_moments::Integer, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T})::Tuple{Vector{T}, T} where T <: AbstractFloat
    # Automatic bandwith selection according to Silverman's rule of thumb.
    # This is not the best we can do, e.g. in the paper it is suggested to use
    # cross-validation a la Newey (1993). However, since according to the paper
    # the bandwith is not very important, we just use something simple.
    # So we take the average bandwith over the bandwith's suggested
    b = mean(density_bandwidth(Kernel, X[:, i]) for i in 1:size(X, 2))
    return optimize_el(n_moments, moment_func, X, Y, start_param, b)
end

function optimize_el(n_moments::Int, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T}, bandwidth::T, trimming::T=.99999)::Tuple{Vector{T}, T} where T <: AbstractFloat
    return optimize_el(convert(UInt, n_moments), moment_func, X, Y, start_param, bandwidth, trimming)
end

function optimize_el(n_moments::UInt, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T}, bandwidth::T, trimming::T=.99999)::Tuple{Vector{T}, T} where T <: AbstractFloat
    weights, not_trimmed = weights_and_trims(n_moments, moment_func, X, Y, start_param, bandwidth, trimming)

    res = Optim.optimize(
        Optim.only_fg!((f, g, x) -> el_objective!(n_moments, moment_func, X, Y, weights, not_trimmed, x, f, g)),
        start_param,
        LBFGS()
    )

    return Optim.minimizer(res), -1 * Optim.minimum(res)
end

function setup_jump_problem(n_moments::Int, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T}, bandwidth::T, trimming::T=.99999) where T <: AbstractFloat
    return setup_jump_problem(convert(UInt, n_moments), moment_func, X, Y, start_param, bandwidth, trimming)
end

function setup_jump_problem(n_moments::UInt, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T}, bandwidth::T, trimming::T=.99999) where T <: AbstractFloat
    weights, not_trimmed = weights_and_trims(n_moments, moment_func, X, Y, start_param, bandwidth, trimming)

    model = Model()
    n_params = length(start_param)
    @variable(model, theta[1:n_params])
    for i in 1:n_params
        set_start_value(theta[i], start_param[i])
    end
    neg_ll(x...) = begin
        x_vec = vcat(x...)
        el_objective!(n_moments, moment_func, X, Y, weights, not_trimmed, x_vec, 1, nothing)
    end
    ∇neg_ll(grad, x...) = begin
        x_vec = vcat(x...)
        el_objective!(n_moments, moment_func, X, Y, weights, not_trimmed, x_vec, nothing, grad)
    end
    register(model, :neg_empirical_likelihood, n_params, neg_ll, ∇neg_ll)
    @NLobjective(model, Min, neg_empirical_likelihood(theta...))

    return model, neg_ll, ∇neg_ll
end

function weights_and_trims(n_moments::UInt, moment_func, X::Matrix{T}, Y::Matrix{T}, start_param::Vector{T}, bandwidth::T, trimming::T)::Tuple{Matrix{T}, Vector{Bool}} where T <: AbstractFloat
    # Calcualte weights
    n_obs, n_conditioning_vars = size(X)
    weights = Matrix{T}(undef, n_obs, n_obs)
    not_trimmed = Vector{Bool}(undef, n_obs)  # false = we will trim
    for i in 1:n_obs
        for j in 1:n_obs
            weights[i, j] = prod(
                density_kernel(Kernel, (X[i, k] - X[j, k]) / bandwidth)
                for k in 1:n_conditioning_vars
            )
        end

        # Determine if the observation will be trimmed.
        # Contrary to what is claimed in the paper, I find it impossible to get the
        # estimator even working without trimming. (Just consider the case where w_{ij} =1,
        # i.e. an observation without "neighbours". In these cases the inner loop is
        # always ill-defined.)
        # Nadaray-Watson estimator of pdf of X at x_i
        hhat = sum(weights[i, :]) / (n_obs * bandwidth^n_conditioning_vars)
        not_trimmed[i] = hhat >= bandwidth^trimming

        # Normalize weights
        weights[i, :] ./= sum(weights[i, :])
    end

    if !any(not_trimmed)
        throw("All observations are trimmed. Decrease the bandwith or increase the trimming parameter.")
    end

    # Check if starting paramater is valid
    moment_vals, _ = moment_vals_and_jac(n_moments, moment_func, X, Y, start_param)
    if !zero_in_hull_for_all(moment_vals, weights, not_trimmed)
        throw("Zero is not in convex hull of moments at starting value. Try a different starting value or increase the bandwidth.")
    end

    return weights, not_trimmed
end

function el_objective!(n_moments::UInt, moment_func, X::Matrix{S}, Y::Matrix{S}, weights::Matrix{S}, not_trimmed::Vector{Bool}, theta::Vector{T}, out, grad::Union{Nothing, AbstractVector{S}})::Union{Nothing, T} where {T <: Real, S <: AbstractFloat}
    moment_vals, moment_jac = moment_vals_and_jac(n_moments, moment_func, X, Y, theta)

    if !zero_in_hull_for_all(moment_vals, weights, not_trimmed)
        if !isnothing(grad)
            grad[:] .= NaN
        end
        if !isnothing(out)
            return -Inf
        end
        return
    end

    # Inner loop
    n_obs = size(X, 1)
    n_params = length(theta)
    λ = Matrix{T}(undef, n_obs, n_moments)
    # Obj contains the negative of the objective at the top of p. 1674, as we will minimize and not maximize
    obj = 0.
    if grad != nothing
        grad[:] = zeros(T, n_params)
    end
    for i in 1:n_obs
        if !not_trimmed[i]
            # This adds 0 to the SER, just go on
            continue
        end

        # Find λ maximizing as in (2.7) in the paper
        weight_i = weights[i, :]
        res = Optim.optimize(
            Optim.only_fgh!((f, g, h, γ) -> λ_obj_fgh(f, g, h, γ, moment_vals, weight_i)),
            # γ -> -1 * λ_obj(γ, moment_vals, weight_i),
            # (g, γ) -> -1 * λ_grad!(g, γ, moment_vals, weight_i),
            zeros(n_moments),
            Newton()
        )
        if !Optim.converged(res)
            throw("Inner loop did not converge")
        end
        λ[i, :] = Optim.minimizer(res)
        # As in the paper, we will not trim
        #obj += λ_obj(λ[i, :], moment_vals, weight_i)
        obj += -1 * Optim.minimum(res)
        if grad != nothing
            # Use envelope theorem to obtain gradient
            grad[:] += sum(weight_i[j] * λ[i, :] * moment_jac[j, :, :] / (1 + transpose(λ[i, :]) * moment_vals[j, :]) for j in 1:n_obs, dims=1)[1, :]
        end
    end

    if out != nothing
        return obj
    end
end

function moment_vals_and_jac(n_moments::UInt, moment_func, X::Matrix{T}, Y::Matrix{T}, theta::Vector{T})::Tuple{Matrix{T}, Array{T, 3}} where T <: AbstractFloat
    n_obs = size(X, 1)
    n_params = length(theta)
    moment_vals = Matrix{T}(undef, n_obs, n_moments)
    moment_jac = Array{T, 3}(undef, n_obs, n_moments, n_params)
    out = DiffResults.JacobianResult(zeros(n_moments), theta)
    for i in 1:n_obs
        ForwardDiff.jacobian!(out, x -> moment_func(X[i, :], Y[i, :], x), theta)
        moment_vals[i, :] = DiffResults.value(out)
        moment_jac[i, :, :] = DiffResults.jacobian(out)
    end
    return moment_vals, moment_jac
end

function zero_in_hull_for_all(moments, weights, not_trimmed)
    n_obs = size(moments, 1)
    for i in 1:n_obs
        # For trimmed observations, we don't care as they will drop from the SER
        if not_trimmed[i]
            pos_weight = weights[i, :] .> 0
            if !zero_in_hull(moments[pos_weight, :])
                return false
            end
        end
    end
    return true
end

function zero_in_hull(moments)
    if size(moments, 2) == 1
        return any(moments .>= 0) && any(moments .<= 0)
    end
    # Based on https://stackoverflow.com/a/16898636/7388096
    hull = scipy_spatial.Delaunay(moments)
    return hull.find_simplex(zeros(size(moments, 2))) .>= 0
end

function λ_obj_fgh(f, grad, hess, γ, moment_vals, weights)
    n_obs = size(moment_vals, 1)
    if grad != nothing
        grad[:] .= 0
    end
    if hess != nothing
        hess[:, :] .= 0
    end
    out = 0.
    for j in 1:n_obs
        if weights[j] > 0
            inlog = 1 + transpose(γ) * moment_vals[j, :]
            if inlog <= 0
                if grad != nothing
                    grad[:] .= NaN
                end
                if hess != nothing
                    hess[:, :] .= NaN
                end
                if f != nothing
                    return Inf
                end
            end
            if f != nothing
                out -= weights[j] * log(inlog)
            end
            if grad != nothing
                grad[:] -= weights[j] * moment_vals[j, :] / inlog
            end
            if hess != nothing
                hess[:, :] -= weights[j] * moment_vals[j, :] * transpose(moment_vals[j, :]) / inlog
            end
        end
    end
    if f != nothing
        return out
    end
end

end


function add_velocity_bound!(model::AbstractGameModel, game_con::GameConstraintValues, v_max::AbstractVector, v_min::AbstractVector)
	n = model.n
	p = model.p
	@assert length(v_max) == length(v_min) == p
	for i = 1:p
		if (v_max[i] != Inf) || (v_min[i] != -Inf)
			add_velocity_bound!(model, game_con, i, v_max[i], v_min[i])
		end
	end
	return nothing
end

function add_velocity_bound!(model::AbstractGameModel, game_con::GameConstraintValues, i::Int, v_max::Real, v_min::Real)
	@assert (v_max != Inf) || (v_min != -Inf)
	n = model.n
	p = model.p
	# Compute the bound
	x_max = Inf*ones(n)
	x_min = -Inf*ones(n)
	vel_index_i = velocity_index(model, i)
	x_max[vel_index_i] = v_max
	x_min[vel_index_i] = v_min
	# Add the state bound to all players
	for j = 1:p
		add_state_bound!(game_con, j, x_max, x_min)
	end
	return nothing
end

function velocity_index(model::UnicycleGame, i::Int)
	@assert i ∈ (1:model.p)
	return model.pz[i][4]
end

function velocity_index(model::BicycleGame, i::Int)
	@assert i ∈ (1:model.p)
	return model.pz[i][3]
end

function velocity_index(model::DoubleIntegratorGame, i::Int)
	error("Velocity Index is not implemented for DoubleIntegratorGame.")
	return nothing
end


using MXNet
using Distributions
using Plots

# create training and evaluation data sets
samplesize  = 1000
TrainInput = rand(1,samplesize)*10
TrainOutput = sin.(TrainInput)+rand(1,samplesize)*0.5
ValidationInput = rand(1,samplesize)*10
ValidationOutput = sin.(ValidationInput)+rand(1,samplesize)*0.5

# how to set up data providers using data in memory
function data_source(batchsize = 100)
  train = mx.ArrayDataProvider(
    :data => TrainInput,
    :label => TrainOutput,
    batch_size = batchsize,
    shuffle = true,
    )
  valid = mx.ArrayDataProvider(
    :data => ValidationInput,
    :label => ValidationOutput,
    batch_size = batchsize,
    shuffle = true,
    )

  train, valid
end

# create a two hidden layer MPL: try varying num_hidden, and change tanh to relu,
# or add/remove a layer
data = mx.Variable(:data)
label = mx.Variable(:label)
net = @mx.chain     mx.Variable(:data) =>
                    mx.FullyConnected(num_hidden=6) =>
                    mx.Activation(act_type=:tanh) =>
                    #  mx.FullyConnected(num_hidden=6) =>
                    #  mx.Activation(act_type=:tanh) =>
                    #  mx.FullyConnected(num_hidden=6) =>
                    #  mx.Activation(act_type=:tanh) =>
                    #  mx.FullyConnected(num_hidden=6) =>
                    #  mx.Activation(act_type=:tanh) =>
                    mx.FullyConnected(num_hidden=1) =>
                    mx.LinearRegressionOutput(mx.Variable(:label))

# final model definition, don't change, except if using gpu
model = mx.FeedForward(net, context=mx.cpu())
#  model = mx.FeedForward(net, context=mx.gpu())

# set up the optimizer: select one, explore parameters, if desired
#optimizer = mx.SGD(η=0.01, μ=0.9, λ=0.00001)
optimizer = mx.ADAM()

# more training with the full sample
trainprovider, evalprovider = data_source(#= batchsize =# 200)
mx.fit(model, optimizer, trainprovider,
       initializer = mx.NormalInitializer(0.0, 0.1),
       eval_metric = mx.MSE(),
       eval_data = evalprovider,
       n_epoch = 10,  # previous setting is batchsize = 200, epoch = 20
                       # implies we did (5000 / 200) * 20 times update in previous `fit`
       callbacks = [mx.speedometer()])

# obtain predictions
plotprovider = mx.ArrayDataProvider(:data => ValidationInput, :label => ValidationOutput)
fit = mx.predict(model, plotprovider)
println("correlation between fitted values and true regression line: ", cor(vec(fit), vec(ValidationOutput)))
scatter(ValidationOutput',fit',w = 3, xlabel="true", ylabel="predicted", title="45º line is what we hope for", show=true)


# main.jl

# (c) 2020 Johannes Boehm <johannes.boehm@sciencespo.fr>

using Revise
using DataFrames

using Random
using Distances

using Test

include("src/JacobiAlgorithm.jl")

# Locations
rng = MersenneTwister(777)
nbz = 3
z_df = DataFrame(h = rand(rng, nbz), v = rand(rng, nbz))

# Supply
nbi = 200
rng = MersenneTwister(778)
i_df = DataFrame(h = rand(rng, nbi),
    v = rand(rng, nbi),
    τ = rand(rng, nbi),
    λ = 10.0 * rand(rng, nbi)
    )

# Demand
nbj = 500
rng = MersenneTwister(779)
j_df = DataFrame(h = rand(rng, nbj),
    v = rand(rng, nbj),
    σ = 1.0 ./ rand(rng, nbj),
    ϵ = 20.0 .* rand(rng, nbj),
    η = rand(rng, nbj)
    )

avg_speed_drive = 25.0
avg_speed_walk = 4.0

T_iz = pairwise(Euclidean(), Matrix(i_df[:,[:h,:v]])', Matrix(z_df[:,[:h,:v]])' ) ./ avg_speed_drive
T_jz = pairwise(Euclidean(), Matrix(j_df[:,[:h,:v]])', Matrix(z_df[:,[:h,:v]])' ) ./ avg_speed_walk

function utility_firms(p::R, τ::R, λ::R, T::R) where R<:Real
    return p^(one(R) - τ) - λ * T
end
function utility_consumers(p::R, σ::R, ϵ::R, T::R) where R<:Real
    expo = one(R) - one(R) / σ
    return ( p^expo + (ϵ*T)^expo )^(one(R)/expo)
end

# it would be more "Julian" to have a separate type for firms

function demand_z(prices_z::Vector{R}, consumers_j::DataFrame, T::Matrix{R}) where R<:Real
    number_of_markets = size(prices_z, 1)
    number_of_consumers = size(consumers_j, 1)
    d = zeros(R, number_of_markets)
    # construct matrix of utilities
    u_jz = zeros(R, number_of_consumers, number_of_markets)
    for z = 1:number_of_markets
        for j = 1:number_of_consumers
            u_jz[j,z] = utility_consumers(prices_z[z],consumers_j[j,:σ],consumers_j[j,:ϵ], T[j,z])
        end
    end
    # for each firm i, increment the supply vector at the position where it has the highest value
    for j=1:number_of_consumers
        min_index = -1
        min = consumers_j[j,:η]
        for z = 1:number_of_markets
            if u_jz[j,z] < min
                min = u_jz[j,z]
                min_index = z
            end
        end
        if min_index != -1
            d[min_index] += 1
        else
            # outside option is better
        end
    end
    return d
end
function supply_z(prices_z::Vector{R}, firms_i::DataFrame, T::Matrix{R}) where R<:Real
    number_of_firms = size(firms_i, 1)
    number_of_markets = size(prices_z, 1)
    s = zeros(R, number_of_markets)
    # construct matrix of utilities
    u_iz = zeros(R, number_of_firms, number_of_markets)
    for z = 1:number_of_markets
        for i = 1:number_of_firms
            u_iz[i,z] = utility_firms(prices_z[z],firms_i[i,:τ],firms_i[i,:λ],T[i,z])
        end
    end
    # for each firm i, increment the supply vector at the position where it has the highest value
    for i=1:number_of_firms
        max_index = -1
        max = zero(R)
        for z = 1:number_of_markets
            if u_iz[i,z] > max
                max = u_iz[i,z]
                max_index = z
            end
        end
        s[max_index] += 1
    end
    return s
end

function ssmooth_z(prices_z::Vector{R}, firms_i::DataFrame, T::Matrix{R}, smoothpar::R) where R<:Real
    number_of_firms = size(firms_i, 1)
    number_of_markets = size(prices_z, 1)
    s = zeros(R, number_of_markets)
    # construct matrix of utilities
    u_iz = zeros(R, number_of_firms, number_of_markets)
    for z = 1:number_of_markets
        for i = 1:number_of_firms
            u_iz[i,z] = utility_firms(prices_z[z],firms_i[i,:τ],firms_i[i,:λ],T[i,z])
        end
    end
    # calculate smooth max
    u_i0 = zeros(R, number_of_firms)
    max_i = [max(zero(R), maximum([u_iz[i,z] for z=1:number_of_markets])) for i=1:number_of_firms]
    utilde_i = max_i .- smoothpar.* log.( exp.( (u_i0 .- max_i)./smoothpar ) .+ sum( exp.( (u_iz .- repeat(max_i,1,number_of_markets))./smoothpar ), dims=2 ) )
    return sum( exp.( ( u_iz .- repeat(utilde_i,1,number_of_markets)) ./ smoothpar) , dims=1)
end
function dsmooth_z(prices_z::Vector{R}, consumers_j::DataFrame, T::Matrix{R}, smoothpar::R) where R<:Real
    # here I denote the cost by u_jz
    number_of_markets = size(prices_z, 1)
    number_of_consumers = size(consumers_j, 1)
    d = zeros(R, number_of_markets)
    # construct matrix of utilities
    u_jz = zeros(R, number_of_consumers, number_of_markets)
    for z = 1:number_of_markets
        for j = 1:number_of_consumers
            u_jz[j,z] = utility_consumers(prices_z[z],consumers_j[j,:σ],consumers_j[j,:ϵ], T[j,z])
        end
    end
    # calculate smooth max
    u_j0 = [consumers_j[j,:η] for j=1:number_of_consumers]
    min_j = [min(u_j0[j], minimum([u_jz[j,z] for z=1:number_of_markets])) for j=1:number_of_consumers]
    utilde_j = min_j .- smoothpar.* log.( exp.( (-u_j0 .+ min_j)./smoothpar ) .+ sum( exp.( (-u_jz .+ repeat(min_j,1,number_of_markets))./smoothpar ), dims=2 ) )
    return sum( exp.( ( -u_jz .+ repeat(utilde_j,1,number_of_markets)) ./ smoothpar) , dims=1)
end
# excess supply function, smoothed
function esmooth_z(prices_z::Vector{R}, firms_i::DataFrame, T_firms::Matrix{R}, consumers_j::DataFrame, T_consumers::Matrix{R}, smoothpar::R) where R<:Real
    return ssmooth_z(prices_z, firms_i, T_firms, smoothpar) .- dsmooth_z(prices_z, consumers_j, T_consumers, smoothpar)
end

# show the validity of the approximations

# exact
supply_z(ones(Float64,3), i_df, T_iz)
# approximation
ssmooth_z(ones(Float64,3), i_df, T_iz, 0.0000001)
# test it
@test maximum(abs.(supply_z(ones(Float64,3), i_df, T_iz) .- ssmooth_z(ones(Float64,3), i_df, T_iz, 0.0000001))) ≈ 0.0 atol= 1e04

# exact
demand_z(ones(Float64,3).*0.1, j_df, T_jz)
# approximation
dsmooth_z(ones(Float64,3).*0.1, j_df, T_jz, 0.0000001)
# test it
@test maximum(abs.(demand_z(ones(Float64,3).*0.1, j_df, T_jz) .- dsmooth_z(ones(Float64,3).*0.1, j_df, T_jz, 0.0000001))) ≈ 0.0 atol= 1e04

include("src/JacobiAlgorithm.jl")

# define the problem
e = JacobiAlgorithm.EquilibriumProblem(nbz, p->esmooth_z(p,i_df, T_iz, j_df, T_jz, 1e-7), 0.0, 1e5, zeros(Float64,nbz), zeros(Float64,nbz) )

# try out excess supply fct
e.e_z(zeros(Float64,3)) .- e.q_z'

# try out coordinate update
p_z = ones(Float64, 3).*0.1
JacobiAlgorithm.coordinate_update(e,1,p_z)
JacobiAlgorithm.coordinate_update(e,2,p_z)
JacobiAlgorithm.coordinate_update(e,3,p_z)

# try out the first iterations of Jacobi algorithm
p_z = ones(Float64, 3).*0.1
for i = 1:5
    global p_z
    p_z , _ = JacobiAlgorithm.fJ_z(e,p_z)
    @show p_z
end

# solve the problem using the jacobi algorithm: about 4.2 seconds
JacobiAlgorithm.solve(e, steptol=1e-12, valtol=1e-5, output=2)

# threaded version: < 2 seconds
JacobiAlgorithm.solve_threaded(e, steptol=1e-12, valtol=1e-5, output=2)


function test_location_scale(
    rng::Union{AbstractRNG, Missing},
    μ::Real, σ::Real, ρ::UnivariateDistribution, dref::UnivariateDistribution,
)
    d = Distributions.AffineDistribution(μ, σ, ρ)
    @test params(d) == (μ,σ,ρ)
    @test eltype(d) === eltype(dref)

    # Different ways to construct the AffineDistribution object
    if dref isa DiscreteDistribution
        # floating point division introduces numerical errors
        # Better: multiply with rational numbers
        d_dict = Dict(
            "original" => d,
            "sugar" => σ * ρ + μ,
        )
    else
        d_dict = Dict(
            "original" => d,
            "sugar" => σ * ρ + μ,
            "composed" =>  μ + ((2 * ρ) * σ - 2) / 2 + 1
        )
    end

    @test d == deepcopy(d)

    #### Support / Domain

    @testset "Support" begin
        @testset "$k" for (k,dtest) in d_dict
            @test minimum(dtest) == minimum(dref)
            @test maximum(dtest) == maximum(dref)
            @test extrema(dtest) == (minimum(dref), maximum(dref))
            @test extrema(support(dtest)) == extrema(dref)
            @test support(dtest) == support(dref)
        end
    end

    #### Promotions and conversions

    @testset "Promotions and conversions" begin
        @testset "$k" for (k,dtest) in d_dict
            if dtest isa LocationScale
                @test typeof(dtest.μ) === typeof(dtest.σ)
                @test location(dtest) ≈ μ atol=1e-15
                @test    scale(dtest) ≈ σ atol=1e-15
            end
        end
    end

    #### Statistics

    @testset "Statistics" begin
        @testset "$k" for (k,dtest) in d_dict
            @test mean(dtest) ≈ mean(dref)
            @test median(dtest) ≈ median(dref)
            @test mode(dtest) ≈ mode(dref)
            @test modes(dtest) ≈ modes(dref)

            @test var(dtest) ≈ var(dref)
            @test std(dtest) ≈ std(dref)

            @test skewness(dtest) ≈ skewness(dref)
            @test kurtosis(dtest) ≈ kurtosis(dref)

            @test isplatykurtic(dtest) == isplatykurtic(dref)
            @test isleptokurtic(dtest) == isleptokurtic(dref)
            @test ismesokurtic(dtest) == ismesokurtic(dref)

            @test entropy(dtest) ≈ entropy(dref)
            @test mgf(dtest,-0.1) ≈ mgf(dref,-0.1)
        end
    end

    #### Evaluation & Sampling

    @testset "Evaluation & Sampling" begin
        @testset "$k" for (k,dtest) in d_dict
            xs = rand(dref, 5)
            x = first(xs)
            insupport(dtest, x) == insupport(dref, x)
            # might return `false` for discrete distributions
            insupport(dtest, -x) == insupport(dref, -x)

            @test pdf(dtest, x) ≈ pdf(dref, x)
            @test pdf.(dtest, xs) ≈ pdf.(dref, xs)
            @test logpdf(dtest, x) ≈ logpdf(dref, x)
            @test logpdf.(dtest, xs) ≈ logpdf.(dref, xs)
            @test loglikelihood(dtest, x) ≈ loglikelihood(dref, x)
            @test loglikelihood(dtest, xs) ≈ loglikelihood(dref, xs)

            @test cdf(dtest, x) ≈ cdf(dref, x)
            @test logcdf(dtest, x) ≈ logcdf(dref, x)
            @test ccdf(dtest, x) ≈ ccdf(dref, x) atol=1e-14
            @test logccdf(dtest, x) ≈ logccdf(dref, x) atol=1e-14

            @test quantile(dtest, 0.1) ≈ quantile(dref, 0.1)
            @test quantile(dtest, 0.5) ≈ quantile(dref, 0.5)
            @test quantile(dtest, 0.9) ≈ quantile(dref, 0.9)

            @test cquantile(dtest, 0.1) ≈ cquantile(dref, 0.1)
            @test cquantile(dtest, 0.5) ≈ cquantile(dref, 0.5)
            @test cquantile(dtest, 0.9) ≈ cquantile(dref, 0.9)

            @test invlogcdf(dtest, log(0.2)) ≈ invlogcdf(dref, log(0.2))
            @test invlogcdf(dtest, log(0.5)) ≈ invlogcdf(dref, log(0.5))
            @test invlogcdf(dtest, log(0.8)) ≈ invlogcdf(dref, log(0.8))

            @test invlogccdf(dtest, log(0.2)) ≈ invlogccdf(dref, log(0.2))
            @test invlogccdf(dtest, log(0.5)) ≈ invlogccdf(dref, log(0.5))
            @test invlogccdf(dtest, log(0.8)) ≈ invlogccdf(dref, log(0.8))

            r = Array{float(eltype(dtest))}(undef, 100000)
            if ismissing(rng)
                rand!(dtest, r)
            else
                rand!(rng, dtest, r)
            end
            @test mean(r) ≈ mean(dref) atol=0.02
            @test std(r) ≈ std(dref) atol=0.01
            @test cf(dtest, -0.1) ≈ cf(dref,-0.1)

            if dref isa ContinuousDistribution
                @test gradlogpdf(dtest, 0.1) ≈ gradlogpdf(dref, 0.1)
            end
        end
    end
end

function test_location_scale_normal(
    rng::Union{AbstractRNG, Missing}, μ::Real, σ::Real, μD::Real, σD::Real,
)
    ρ = Normal(μD, σD)
    dref = Normal(μ + σ * μD, σ * σD)
    @test dref === μ + σ * ρ
    return test_location_scale(rng, μ, σ, ρ, dref)
end

function test_location_scale_discretenonparametric(
    rng::Union{AbstractRNG, Missing}, μ::Real, σ::Real, support, probs,
)
    ρ = DiscreteNonParametric(support, probs)
    dref = DiscreteNonParametric(μ .+ σ .* support, probs)
    return test_location_scale(rng, μ, σ, ρ, dref)
end

@testset "AffineDistribution" begin
    rng = MersenneTwister(123)

    for _rng in (missing, rng)
        test_location_scale_normal(_rng, 0.3, 0.2, 0.1, 0.2)
        test_location_scale_normal(_rng, -0.3, 0.1, -0.1, 0.3)
        test_location_scale_normal(_rng, 1.3, 0.4, -0.1, 0.5)
    end
    test_location_scale_normal(rng, ForwardDiff.Dual(0.3), 0.2, 0.1, 0.2)

    probs = normalize!(rand(10), 1)
    for _rng in (missing, rng)
        test_location_scale_discretenonparametric(_rng, 1//3, 1//2, 1:10, probs)
        test_location_scale_discretenonparametric(_rng, -1//4, 1//3, (-10):(-1), probs)
        test_location_scale_discretenonparametric(_rng, 6//5, 3//2, 15:24, probs)
    end

    @test_logs Distributions.AffineDistribution(1.0, 1, Normal())

    @test_deprecated ls_norm = LocationScale(1.0, 1, Normal())
    @test ls_norm isa LocationScale{Float64, Continuous, Normal{Float64}}
    @test ls_norm isa Distributions.AffineDistribution{Float64, Continuous, Normal{Float64}}
end


using JuMP, DataStructures

mutable struct nodo
  nivel::Int64
  Model::JuMP.Model
  inf::Float64
  sup::Float64
  x_relax::Vector{Float64}
  frac::Vector{Float64}
  x_int::Vector{Float64}
  status::Symbol
end

function fractional(x,var_bin)
  s = length(var_bin)
  frac = Array{Float64,1}(s)

  for i=1:s
    frac[i] = modf(x[var_bin[i]])[1]
  end
  return frac
end

function branch(x,var_bin)
  f,i = findmin(abs(x-0.5))
  return var_bin[i]
end

function poda(lista_,sense,liminf,limsup, P)
  L = size(lista_)[1]
  erase = []
  P_ = 0
  for l=1:L
    if lista_[l].status == :Infeasible
      erase = copy([erase;[l]])
      P_ = P_ + 1
    elseif lista_[l].frac == 0
      erase = copy([erase;[l]])
      P_ = P_ + 1
    elseif (sense == :Max && lista_[l].sup <= liminf) ||(sense == :Min && lista_[l].inf >= limsup)
       erase = copy([erase;[l]])
       P_ = P_ + 1
    end
  end
  P =copy(P + P_)
  deleteat!(lista_,erase)
  return  P
end

function deep(m::JuMP.Model,var_bin,frac, solver::MathProgBase.AbstractMathProgSolver = JuMP.UnsetSolver())
    m_ = deepcopy(m)
    N = 1
    liminf = []
    limsup = []
    obj = []
    x_best = []
    x = []
    k = 1
    sense = m.objSense
    status1 = []
    status2 = []
   time = 0

    while frac != zeros(length(var_bin))
      tic()
      m1 = deepcopy(m_)
      m2 = deepcopy(m_)
      i_ = branch(frac,var_bin)
      if mod(k,2) == 0

         m1.colUpper[i_] = 0
         status1 = solve(m1, relaxation=true)
         x = copy(m1.colVal)
         obj = copy(m1.objVal)
         frac = fractional(x,var_bin)
         m_ = deepcopy(m1)
      elseif  mod(k,2) == 1

        m2.colLower[i_] = 1
        status2 = solve(m2, relaxation=true)
        x = copy(m2.colVal)
        obj = copy(m2.objVal)
        frac = fractional(x,var_bin)
        m_ = deepcopy(m2)
      end
     time_aux = toc()
     time = time + time_aux
      if status1 == :Infeasible && status2 == :Optimal
        m_ = deepcopy(m2)
      elseif status2 == :Infeasible && status1 == :Optimal
        m_ = deepcopy(m1)
      elseif  status1 == :Infeasible && status2 == :Infeasible
      break
      end
      if time >= 120
        break
      end
      k = k + 1
      N = N + 1
    end

  if sense == :Max
    if frac == zeros(length(var_bin))
     liminf = copy(obj)
     x_best = copy(x)
    else
     liminf = -Inf
     N = Inf
   end
  else
    if frac == zeros(length(var_bin))
       limsup = copy(obj)
       x_best = copy(x)
    else
       limsup = Inf
       N = Inf
   end
  end

return liminf, limsup, x_best,N
end

function largura(lista,sense,var_bin,liminf,limsup)
  liminf_ = copy(liminf)
  limsup_ = copy(limsup)
  Podas_aux = 0
  AUX = []
  int_sol = []
  x_best = []
  IntSol = 0
  lista_aux = copy(lista)
  Nivel = 5000
  k = 1
  time = 0
  sol = []
  time_aux = 0
  Visit_ = 0
  Visit_aux = 0
   while k <= Nivel && IntSol == 0
     tic()
     L = length(lista_aux)
     Visit_aux = copy(L)
     for l=1:L
        filhos(lista_aux[l],lista_aux,var_bin,sense,liminf_,limsup_)
     end
     L = length(lista_aux)
     erase = []
     for l=1:L
       if lista_aux[l].nivel == k - 1
         erase = copy([erase;[l]])
       end
     end
     deleteat!(lista_aux,erase)
     liminf_, limsup_, AUX, int_sol, IntSol  = gap(lista_aux,sense,liminf_,limsup_,AUX,int_sol,IntSol)
     Podas_aux = poda(lista_aux,sense,liminf,limsup,Podas_aux)
     k = k + 1
     time_aux = toc()
     time = time + time_aux
     Visit_ = Visit_ + Visit_aux
     if time >= 350
       break
     end
   end
   if AUX != []
     if sense == :Max
         t,r = size(AUX)
         for i=1:t
           if liminf_ == AUX[i,1]
             x_best = AUX[i,2:end]
             sol = liminf_
           end
         end
       else
         t,r = size(AUX)
         for i=1:t
           if limsup_ == AUX[i,1]
             x_best = AUX[i,2:end]
             sol = limsup_
           end
         end
     end
   end
   if sense == :Max
     if IntSol == 0
        liminf_ = -Inf
    end
   else
     if IntSol == 0
        limsup_ = Inf
    end
   end
   return liminf_, limsup_, x_best, time, Podas_aux, k, sol, Visit_, IntSol
end

function filhos(pai,lista_,var_bin,sense,liminf,limsup)
  #############################################################################
  # Nodo filho 1
  #############################################################################
    m1 = deepcopy(pai.Model)
    frac = copy(pai.frac)
    i_ = branch(frac,var_bin)
    m1.colUpper[i_] = 0
    status = solve(m1, relaxation=true)
    x = copy(m1.colVal)
    obj = copy(m1.objVal)
    frac = fractional(x,var_bin)

    if sense == :Max
      sup = copy(obj)
      inf = liminf
    else
      inf = copy(obj)
      sup = limsup
    end

    if frac == zeros(length(var_bin))
      x_int = copy(x)
    else
      x_int = []
    end
  nivel = pai.nivel + 1
  filho1 = nodo(nivel,m1,inf,sup,x,frac,x_int,status)
  push!(lista_,filho1)

  #############################################################################
  # Nodo filho 2
  #############################################################################
  m2 = deepcopy(pai.Model)
  m2.colLower[i_] = 1
  status = solve(m2, relaxation=true)
  x = copy(m2.colVal)
  obj = copy(m2.objVal)
  frac = fractional(x,var_bin)

  if sense == :Max
    sup = copy(obj)
    inf = copy(liminf)
  else
    inf = copy(obj)
    sup = copy(limsup)
  end

  if frac  == zeros(length(var_bin))
    x_int = copy(x)
  else
    x_int = []
  end

  filho2 = nodo(nivel,m2,inf,sup,x,frac,x_int,status)
  push!(lista_,filho2)
return nothing
end

function selection(lista,sense,liminf,limsup)
  pai = []
  l = []
  L = length(lista)
  for i=1:L
    if sense == :Max && lista[i].sup == limsup
       pai = lista[i]
       l = i
    elseif sense == :Min && lista[i].inf == liminf
      pai = lista[i]
      l = i
    end
  end
  return pai, l
end

function gap(lista,sense,liminf,limsup,AUX,int_sol,IntSol)
  ##############################################################################
  # Gap
  ##############################################################################

  SUP = []
  INFI = []

  L = length(lista)
    if sense ==:Max
      for l=1:L
        if  lista[l].status != :Infeasible && lista[l].status != :Unbounded
          SUP = copy([SUP;[lista[l].sup]])
        end
        if lista[l].x_int != []
            int_sol = copy([int_sol;[lista[l].sup]])
            IntSol = IntSol + 1
            AUX = copy([AUX;Matrix([lista[l].sup lista[l].x_relax'])])
          else
            int_sol = copy([int_sol;[liminf]])
        end
      end
      limsup = maximum(SUP)
      liminf = maximum(int_sol)
    else
      for l=1:L
        if lista[l].status != :Infeasible && lista[l].status != :Unbounded
         INFI = copy([INFI;[lista[l].inf]])
       end
        if lista[l].x_int != []
          int_sol = copy([int_sol;[lista[l].inf]])
          IntSol = IntSol + 1
          AUX = copy([AUX;Matrix([lista[l].inf lista[l].x_relax'])])
        else
          int_sol = copy([int_sol;[limsup]])
        end
      end
      liminf = minimum(INFI)
      limsup = minimum(int_sol)
    end
    return liminf, limsup, AUX, int_sol,IntSol
end

function BNB(m::JuMP.Model, solver::MathProgBase.AbstractMathProgSolver = JuMP.UnsetSolver())

  #-----------------------------------------------------------------------------
  # Nodo raiz
  #-----------------------------------------------------------------------------
  # inicialização da Lista
  m.ext[:time] = []
  m.ext[:status] = []
  m.ext[:nodes] = []
  m.ext[:Podas] = []
  m.ext[:Iter] = []
  m.ext[:intsols] = []
  Tempo = 0
  lista = []
  sol = []
  x_best = []
  IntSol = 0
  AUX = []
  # Relaxação linear

  m_ = deepcopy(m)
  status = solve(m_, relaxation=true)
  x = copy(m_.colVal)
  obj = copy(m_.objVal)

  sense = m_.objSense

  index = m_.colCat
  var_bin = []
  for  i =1:size(index)[1]
    if index[i] == :Bin
      var_bin = copy([var_bin;[i]])
    end
  end

  frac = fractional(x,var_bin)

  if frac == zeros(length(var_bin))
    m.objVal = copy(obj)
    m.colVal = copy(x)
    m.objBound = copy(obj)
    m.ext[:time] = 0
    m.ext[:status] = status
    m.ext[:nodes] = 0
    m.ext[:Podas] = 0
    m.ext[:Iter] = 0
    m.ext[:intsols] = 1
    return nothing
  end

  if status == :Infeasible
    m.objVal = NaN
    m.colVal = []
    m.objBound = NaN
    m.ext[:time] = 0
    m.ext[:status] = status
    m.ext[:nodes] = 0
    m.ext[:Podas] = 0
    m.ext[:Iter] = 0
    m.ext[:intsols] = 0
    return nothing
  end

 liminf, limsup, x_best, Max_nivel = deep(m_,var_bin,frac)

 if x_best != []
   IntSol = 1
 end

  if sense == :Max
    sup = copy(obj)
    inf = copy(liminf)
    limsup = copy(obj)
  else
    inf = copy(obj)
    sup = copy(limsup)
    liminf = copy(obj)
  end

raiz = nodo(0,m_,inf,sup,x,frac,[],status)

# Dados iniciais
nivel = 1
ϵ = limsup - liminf
iter = 1

push!(lista, raiz)

Visit_ = 0
if x_best ==[]
  liminf_, limsup_, x_best, time, Podas_aux, iter_aux, sol, Visit_, IntSol = largura(lista,sense,var_bin,liminf,limsup)
  ϵ_aux = limsup_ - liminf_
  if x_best !=[] && ϵ_aux <= 1.0e-1
      m.objVal = sol
      m.colVal = x_best
      if sense == :Max
        m.objBound = limsup_
      else
        m.objBound = liminf_
      end
      m.ext[:time] = time
      m.ext[:status] = status
      m.ext[:nodes] = Visit_
      m.ext[:Podas] = Podas_aux
      m.ext[:Iter] = iter_aux
      m.ext[:intsols] = IntSol
      return nothing
  elseif x_best !=[] && ϵ_aux > 1.0e-1
      if sense == :Max
        liminf = copy(liminf_)
      else
        limsup = copy(limsup_)
      end
  end
end

L = 1
int_sol = []

ϵ = limsup - liminf
iter = 1
Visit = 0
Podas = 0

tempo_aux = 0

Max_iter =5000

while ϵ >= 1.0e-1 && iter <= Max_iter
  tic()
  #############################################################################
  # Seleção do problema da lista para fazer branch
  #############################################################################
  pai, l = selection(lista,sense,liminf,limsup)
  deleteat!(lista,l)
  Visit = Visit + 1 # Número do nós onde se fez branch
  #############################################################################
  # Adicionar os nós filhos do problema selecionado da lista
  #############################################################################
  filhos(pai,lista,var_bin,sense,liminf,limsup)
  #############################################################################
  # Atualização do lowere (upper bound) para o problema de Min (Max)
  #############################################################################
  liminf, limsup, AUX, int_sol, IntSol = gap(lista,sense,liminf,limsup,AUX,int_sol,IntSol)
  #############################################################################
  # Podas da árvore
  #############################################################################
  Podas = poda(lista,sense,liminf,limsup,Podas)
  #############################################################################
  # Atualização do upper (lower) bound pra o problema de Min (Max)
  #############################################################################
   L = size(lista)[1]
   ϵ = limsup - liminf # Atualização do ϵ e do x_best (melhor solução inteira)

    if AUX != []
      if sense == :Max
          t,r = size(AUX)
          for i=1:t
            if liminf == AUX[i,1]
              x_best = AUX[i,2:end]
              sol = liminf
            end
          end
        else
          t,r = size(AUX)
          for i=1:t
            if limsup == AUX[i,1]
              x_best = AUX[i,2:end]
              sol = limsup
            end
          end
      end
    end
    tempo_aux = toc()
    Tempo = Tempo + tempo_aux
    iter = copy(iter) + 1
    if Tempo >= 300
      break
    end
end

if iter - 1 == Max_iter &&  ϵ > 1.0e-1 && x_best != []
  status = :Subotimal
  if sense == :Max
    sol = copy(liminf)
  else
    sol = copy(limsup)
  end
elseif iter - 1 == Max_iter && x_best == []
  status = :unsolved
  sol = NaN
elseif lista == [] && x_best == []
  status = :Infeasible
  sol = NaN
else
  status = :Optimal
end
#############################################################################
# Prenchimento dos outputs
#############################################################################
m.objVal = sol
m.colVal = x_best
if sense == :Max
  m.objBound = limsup
else
  m.objBound = liminf
end
m.ext[:time] = Tempo
m.ext[:status] = status
m.ext[:nodes] = Visit
m.ext[:Podas] = Podas
m.ext[:Iter] = iter
m.ext[:intsols] = IntSol

return nothing

end


struct Network

    lsizes::Tuple{Vararg{Int}};
    b::Tuple{Vararg{Vector{Float64}}};
    w::Tuple{Vararg{Matrix{Float64}}};
    z::Tuple{Vararg{Vector{Float64}}};
    a::Tuple{Vararg{Vector{Float64}}};
    depth::Int;
    inlength::Int;
    outlength::Int;

    function Network(lsizes::Tuple{Vararg{Int}})
        @assert length(lsizes) > 0 && all(lsizes .> 0)

        #Initializing the bias randomly
        #For the first layer the bias is not usable.
        b = Vector{Vector{Float64}}();
        push!(b,zeros(1));
        for l in 2:length(lsizes)
            push!(b,randn(lsizes[l]))
        end

        #Initializing the ws
        #For the first layer the weight is not usable.
        w = Vector{Matrix{Float64}}();
        push!(w, zeros(1,1));
        for l in 2:length(lsizes)
            push!(w,randn(lsizes[l],lsizes[l - 1]))
        end

        #The network store the values z and a of the last feedforward
        z = Vector{Vector{Float64}}();
        a = Vector{Vector{Float64}}();
        for l in 1:length(lsizes)
            push!(z, zeros(lsizes[l]));
            push!(a, zeros(lsizes[l]));
        end

        #Constructing
        new(lsizes, tuple(b...), tuple(w...), tuple(z...), tuple(a...),
            length(lsizes), first(lsizes), last(lsizes))
    end
end
Network(lsizes::Int...) = Network(lsizes);
Network(lsizes::Vector{Int}) = Network(lsizes...);
Base.size(net::Network) = net.lsizes;
Base.size(net::Network, l::Int) = net.lsizes[l];
output(net::Network) = last(net.a);
input(net::Network) = first(net.a);
outlength(net::Network) = length(output(net));
inlength(net::Network) = length(input(net));

function feedforward!(net::Network, input::Vector{Float64})::Vector{Float64}
    @assert net.inlength == length(input)

    #first layer
    net.z[1] .= input;
    net.a[1] .= input;

    #Other layers
    for l in 2:net.depth
        net.z[l] .= net.w[l]*net.a[l-1] + net.b[l];
        net.a[l] .= act_function.(net.z[l]);
    end

    #The output
    return output(net);
end

feedforward!(net::Network, testdata) = feedforward!(net, input(testdata));


##
## rand for Flint-Finite fields
##
## fq_nmod has no base ring
function rand(R::FqNmodFiniteField)
  #gen is not a generator for the group!
  r = R(0)
  for i=0:degree(R)
    r *= gen(R)
    r += rand(1:characteristic(R))
  end

  return r
end

function rand(R::FqFiniteField)
  #gen is not a generator for the group!
  r = R(0)
  for i=0:degree(R)
    r *= gen(R)
    r += rand(1:characteristic(R))
  end

  return r
end


using XDiff

ex = quote
    y[i] = W[i,k] * x[k]
    v[i] = y[i] - x[i]
    z = v[i] * I[i]
end
inputs = [:W => rand(2, 2), :x => rand(2)]
ctx = Dict()
@time dexs = rdiff(ex; ctx=ctx, inputs...)


"""
    predict(y, x, xb, xt, hyp, ∇)

    predictive distribution for SPGP given hyperparameters and pseudo-inputs

    y -- training targets (N x 1)
    x -- training inputs (N x dim)
    xb -- pseudo inputs (n x dim)
    xt -- test inputs (Nt x dim)
    hyp -- hyperparameters (including noise)
           for Gaussian covariance: (dim+2 x 1)

          hyp(1:dim) = log( b )
          hyp(dim+1) = log( c )
          hyp(dim+2) = log( sig )

         where cov = c * exp[-0.5 * sum_d b_d * (x_d - x'_d)^2]
                         + sig * delta(x,x')

    ∇ -- OPTIONAL jitter (default 1e-6)
    μ -- predictive mean
    σ² -- predictive variance of latent function
          (add noise if full variance is required)
"""
function predict(y, x, xb, xt, hyp, ∇=1e-6)
    n_observ, n_dim = size(x, 1), size(x, 2)
    n_pseudo = size(xb, 1)
    n_train = size(xt, 1)

    σ = exp(hyp[end])

    K = kernel(xb, xb, hyp) + ∇ * I
    L = cholesky(Hermitian(K)).L
    K = kernel(xb, x, hyp)
    V = L \ K

    @inline function kdiag(x, hyp)
        c = exp(hyp[end - 1])
        Kd = fill(c, size(x, 1), 1)

        return Kd
    end

    ɛ = 1 .+ (kdiag(x, hyp) - transpose(sum(V.^2, dims=1))) / σ

    V = V ./ repeat(transpose(sqrt.(ɛ)), n_pseudo, 1);
    y = y ./ sqrt.(ɛ);
    Lm = transpose(cholesky(σ * I + V * transpose(V)).U);
    bet = Lm \ (V * y);

    K = kernel(xb, xt, hyp);
    lst = L \ K;
    lmst = Lm \ lst;
    μ = transpose((transpose(bet) * lmst));

    σ² = kdiag(xt, hyp) .- transpose(sum(lst.^2, dims=1)) .+ σ .* transpose(sum(lmst.^2, dims=1));

    return μ, σ²
end


"""
    gradient(f, p::AbstractParticles)

Calculate the gradient of `f` in `p`. This corresponds to a smoothed finite-difference approximation where the smoothing kernel is given by the distribution of `p`.
Return mean and std.
"""
function gradient(f,p::MonteCarloMeasurements.AbstractParticles)
    r = 2(p\f(p))
    pmean(r), pstd(r)
end
function gradient(f,p::Union{Integer, AbstractFloat})
    p = p  ± 0.000001
    r = 2(p\f(p))
    pmean(r)
end

function gradient(f::Function,p::MonteCarloMeasurements.MvParticles)
    r = (p-pmean.(p))\(f(p) - f(pmean.(p)))
end

"""
    jacobian(f::Function, p::Vector{<:AbstractParticles})

Calculate the Jacobian of `f` in `p`. This corresponds to a smoothed finite-difference approximation where the smoothing kernel is given by the distribution of `p`.
"""
function jacobian(f::Function,p::MonteCarloMeasurements.MvParticles)
    r = (p-pmean.(p))\(f(p) - f(pmean.(p)))
end


using Pkg
#Pkg.add("StatsBase")
using StatsBase

include("ant.jl")
include("graph.jl")

function algorithmAntColonyOptimization(graph::Graph, maxIteration:: Int64,antsCount:: Int64, pheromoneExponent::Int64, desirableExponent::Float64, pheromoneDepositFactor::Int64, evaporationRate:: Float64)

    bestFitness = Inf
    bestPath = []

    for i in 1:maxIteration
        ants = generateAntPopulation(antsCount,1)
        for j in 1:antsCount
            antReachDestination = false
            while (!antReachDestination)
                # trazimo mogucu putanju
                possiblePaths = findPossiblePaths(ants[j],graph)
                # racunamo verovatnocu
                probabilities = calculateProbabilities(graph, ants[j], possiblePaths, pheromoneExponent, desirableExponent)
                # biramo putanju
                chosenNode = sample(possiblePaths,Weights(probabilities))
                # pomeramo mrava
                push!(ants[j].route,ants[j].currentIndexNode)
                ants[j].currentIndexNode = chosenNode
                if chosenNode == graph.foodIndex
                    antReachDestination = true
                    push!(ants[j].route, ants[j].currentIndexNode)
                end
            end
            # trazenje najboljeg resenja za iteraciju
            fitness = calculateFitness(ants[j],graph)
            ants[j].fitness = fitness
            if fitness < bestFitness
                bestFitness = fitness
                bestPath = deepcopy(ants[j].route)
            end
        end
        # update feromona
        updatePheromones!(ants,graph, pheromoneDepositFactor, evaporationRate)
    end
    return bestPath, bestFitness, graph
end

function findPossiblePaths(ant:: Ant,graph:: Graph)
    possiblePaths = findall(graph.adjMatrix[ant.currentIndexNode,:].!= -1)
    if !isempty(ant.route)
        possiblePaths = deleteat!(possiblePaths,findall(possiblePaths.== ant.route[end]))
    end
    return possiblePaths
end

function calculateProbabilities(graph::Graph, ant::Ant, possiblePaths::Array{Int64,1}, pheromoneExponent, desirableExponent)
    pathSum = 0
    for possibleNode in possiblePaths
        pathSum += ((graph.pheromoneMatrix[ant.currentIndexNode, possibleNode])^pheromoneExponent) *
                    ((graph.adjMatrix[ant.currentIndexNode,possibleNode])^desirableExponent)
    end

    probabilities = Float64[]

    for possibleNode in possiblePaths
        probability = ((graph.pheromoneMatrix[ant.currentIndexNode, possibleNode]) ^ pheromoneExponent) *
        ((graph.adjMatrix[ant.currentIndexNode,possibleNode]) ^ desirableExponent) / pathSum

        push!(probabilities,probability)
    end
    return probabilities
end

function calculateFitness(ant::Ant,graph::Graph)
    pathLength = 0
    for i in 1:(length(ant.route)-1)
        pathLength += graph.adjMatrix[ant.route[i],ant.route[i+1]]
    end
    return pathLength
end

function updatePheromones!(ants, graph, pheromoneDepositFactor, evaporationRate)
    pheromoneDepositMatrix = zeros(Float64, graph.nodesCount, graph.nodesCount)

    for currentAnt in ants
        route = currentAnt.route
        for j in length(route)-1
            pheromoneDepositMatrix[route[j],route[j+1]] += pheromoneDepositFactor / currentAnt.fitness
        end
    end

    for i in 1:graph.nodesCount
        for j in 1:graph.nodesCount
           pheromone = graph.pheromoneMatrix[i,j]
           graph.pheromoneMatrix[i,j] = (1 - evaporationRate) * pheromone + pheromoneDepositMatrix[i,j]
           pheromoneDepositMatrix[j,i] = graph.pheromoneMatrix[i,j]
        end
    end
end


@testset "Augmented Sigma points" begin
    @testset "Create sigma points" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        A = [4.0 0.0; 0.0 4.0]
        Q = [4.0 0.0; 0.0 4.0]
        x = randn(2)
        χ = @inferred KalmanFilters.calc_sigma_points(x, KalmanFilters.Augmented(A, Q), weight_params)
        weight = KalmanFilters.calc_cholesky_weight(weight_params, 4)
        @test χ.x0 == x
        @test χ.P_chol == cholesky(A .* weight).L

        P_chol_temp = zeros(2,2)
        Q_chol_temp = zeros(2,2)
        χ = @inferred KalmanFilters.calc_sigma_points!(
            KalmanFilters.Augmented(P_chol_temp, Q_chol_temp),
            x,
            KalmanFilters.Augmented(A, Q),
            weight_params
        )
        @test χ.x0 == x
        @test χ.P_chol == cholesky(A .* weight).L
    end

    @testset "Covariance" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        𝓨 = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2), weight_params)
        noise = Diagonal(ones(5))
        P = @inferred KalmanFilters.cov(𝓨, Augment(noise))
        @test P == -16.25 .* ones(5) * ones(5)' .+ 2 .* ones(5,5) * ones(5,5)' .* 16 .+
            2 .* ones(5,5) * ones(5,5)' .* 4

        P = zeros(5,5)
        P = @inferred KalmanFilters.cov!(P, 𝓨, Augment(noise))
        @test P == -16.25 .* ones(5) * ones(5)' .+ 2 .* ones(5,5) * ones(5,5)' .* 16 .+
            2 .* ones(5,5) * ones(5,5)' .* 4

        P_chol = LowerTriangular(ones(5,5))
        x = ones(5) .* 4
        𝓨 = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2, ones(5,5) .* 3, ones(5,5) .* 1), weight_params)
        χ = KalmanFilters.AugmentedSigmaPoints(x, P_chol, LowerTriangular(collect(noise)), weight_params)
        P = @inferred KalmanFilters.cov(χ, 𝓨)
        weight_0, weight_i = KalmanFilters.calc_cov_weights(𝓨)
        @test P == weight_i .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 4 .-
#            weight_i .* noise * ones(5,5)' .* 2 .-
            weight_i .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 3
#            weight_i .* noise * ones(5,5)' .* 1

        P = zeros(5,5)
        P = @inferred KalmanFilters.cov!(P, χ, 𝓨)
        @test P == weight_i .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 4 .-
#            weight_i .* noise * ones(5,5)' .* 2 .-
            weight_i .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 3
#            weight_i .* noise * ones(5,5)' .* 1
    end

    @testset "Transform sigma points" begin
        P = [2.0 0.0; 0.0 2.0]
        noise = collect(Diagonal(ones(2) .* 3))
        x = [1.0, 1.0]
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        χ = KalmanFilters.AugmentedSigmaPoints(x, LowerTriangular(P), LowerTriangular(noise), weight_params)
        f(x) = x .* 2
        f(x, n) = x .* 2 .+ n
        𝓨 = @inferred KalmanFilters.transform(f, χ)
        @test 𝓨 == [ones(2) .* 2 [6 2; 2 6] x .* 2 .+ noise [-2 2; 2 -2] x .* 2 .- noise]
        @test 𝓨.weight_params == weight_params

        other_weight_params = MeanSetWeightingParameters(0.5)
        𝓨_temp = KalmanFilters.TransformedSigmaPoints(zeros(2), zeros(2,8), other_weight_params)
        F!(y, x) = y .= x .* 2
        F!(y, x, n) = y .= x .* 2 .+ n
        xi_temp = KalmanFilters.Augmented(zeros(length(x)), zeros(length(x)))
        𝓨 = @inferred KalmanFilters.transform!(𝓨_temp, xi_temp, F!, χ)
        @test 𝓨 == [ones(2) .* 2 [6 2; 2 6] x .* 2 .+ noise [-2 2; 2 -2] x .* 2 .- noise]
        @test 𝓨.weight_params == weight_params
    end
end


quad_params = (m=0.5,
             J=SMatrix{3,3}(Diagonal([0.0023, 0.0023, 0.004])),
             Jinv=SMatrix{3,3}(Diagonal(1.0./[0.0023, 0.0023, 0.004])),
             gravity=SVector(0,0,-9.81),
             motor_dist=0.1750,
             kf=1.0,
             km=0.0245)

include("quaternions.jl")
function quadrotor_dynamics!(ẋ::AbstractVector,x::AbstractVector,u::AbstractVector,params) where T
      #TODO change concatentations to make faster!
      # Quaternion representation
      # Modified from D. Mellinger, N. Michael, and V. Kumar,
      # "Trajectory generation and control for precise aggressive maneuvers with quadrotors",
      # In Proceedings of the 12th International Symposium on Experimental Robotics (ISER 2010), 2010.

      ## States: X ∈ R^13; q = [s;v]
      # x
      # y
      # z
      # q0
      # q1
      # q2
      # q3
      # xdot
      # ydot
      # zdot
      # omega1
      # omega2
      # omega3

      # x = X[1:3]
      q = normalize(Quaternion(view(x,4:7)))
      # q = view(x,4:7)
      # normalize!(q)
      v = view(x,8:10)
      omega = view(x,11:13)

      # Parameters
      m = params[:m] # mass
      J = params[:J] # inertia matrix
      Jinv = params[:Jinv] # inverted inertia matrix
      g = params[:gravity] # gravity
      L = params[:motor_dist] # distance between motors

      w1 = u[1]
      w2 = u[2]
      w3 = u[3]
      w4 = u[4]

      kf = params[:kf]; # 6.11*10^-8;
      F1 = kf*w1;
      F2 = kf*w2;
      F3 = kf*w3;
      F4 = kf*w4;
      F = @SVector [0., 0., F1+F2+F3+F4] #total rotor force in body frame

      km = params[:km]
      M1 = km*w1;
      M2 = km*w2;
      M3 = km*w3;
      M4 = km*w4;
      tau = @SVector [L*(F2-F4), L*(F3-F1), (M1-M2+M3-M4)] #total rotor torque in body frame

      ẋ[1:3] = v # velocity in world frame
      # ẋ[4:7] = 0.5*qmult(q,[0;omega]) #quaternion derivative
      ẋ[4:7] = SVector(0.5*q*Quaternion(zero(x[1]), omega...))
      ẋ[8:10] = g + (1/m)*(q*F) #acceleration in world frame
      ẋ[11:13] = Jinv*(tau - cross(omega,J*omega)) #Euler's equation: I*ω + ω x I*ω = constraint_decrease_ratio
      return tau, omega, J, Jinv
end

quadrotor = Model(quadrotor_dynamics!, 13, 4, quad_params)


using Gadfly

n = 10
plot(x=rand(n), y=rand(n), color=8.^rand(n),
     Scale.color_log2(minvalue=1, maxvalue=8))


using JuMP, EAGO, Gurobi, CSV, DataFrames, Ipopt

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[1, 2, 3, 4, 5, 6, 7, 8]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[1], 100.0)
JuMP.set_upper_bound(x[1], 10000.0)
JuMP.set_lower_bound(x[2], 1000.0)
JuMP.set_upper_bound(x[2], 10000.0)
JuMP.set_lower_bound(x[3], 1000.0)
JuMP.set_upper_bound(x[3], 10000.0)
JuMP.set_lower_bound(x[4], 10.0)
JuMP.set_upper_bound(x[4], 1000.0)
JuMP.set_lower_bound(x[5], 10.0)
JuMP.set_upper_bound(x[5], 1000.0)
JuMP.set_lower_bound(x[6], 10.0)
JuMP.set_upper_bound(x[6], 1000.0)
JuMP.set_lower_bound(x[7], 10.0)
JuMP.set_upper_bound(x[7], 1000.0)
JuMP.set_lower_bound(x[8], 10.0)
JuMP.set_upper_bound(x[8], 1000.0)

# ----- Constraints ----- #
@NLconstraint(m, e2, 833.33252*x[4]/x[1]/x[6]+100/x[6]-83333.333/(x[1]*x[6]) <= 1.0)
@NLconstraint(m, e3, 1250*x[5]/x[2]/x[7]+x[4]/x[7]-1250*x[4]/x[2]/x[7] <= 1.0)
@NLconstraint(m, e4, 1250000/(x[3]*x[8])+x[5]/x[8]-2500*x[5]/x[3]/x[8] <= 1.0)
@constraint(m, e5, 0.0025*x[4]+0.0025*x[6] <= 1.0)
@constraint(m, e6, -0.0025*x[4]+0.0025*x[5]+0.0025*x[7] <= 1.0)
@constraint(m, e7, -0.01*x[5]+0.01*x[8] <= 1.0)

# ----- Objective ----- #
@objective(m, Min, x[1] + x[2] + x[3])

JuMP.optimize!(m)

backend_opt = backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


using Test

@testset "./src/LocalPolytope/utils.jl" begin

using BellScenario

@testset "dimension()" begin

    @testset "vectors" begin
        vertices = [[1,0,0,0],[0,1,0,0],[0,0,1,0],[1,1,0,0]]
        @test LocalPolytope.dimension(vertices) == 3
    end

    @testset "matrices" begin
        vertices = [[1 0;0 0],[0 0;0 1],[1 1;0 0],[1 1;0 1]]
        @test LocalPolytope.dimension(vertices) == 3
    end

    @testset "Deterministic Strategies" begin
        vertices = DeterministicStrategy.([[1 1;0 0],[0 0;1 1],[1 0;0 1]])
        @test LocalPolytope.dimension(vertices) == 2
    end
end

end


#using HiOscSolver
using Test
include("testfftbig.jl")
include("testexpmatrices.jl")
include("testpolyexp.jl")
# include("testcoefexp_ab.jl")
# include("testpreparephi.jl")
# include("testtwoscales_pure_ab.jl")
include("testinterface.jl")
include("testcompfloat.jl")


# Tudor Berariu <tudor.berariu@gmail.com>, 2015

module Bullseye

export bullseye

function bullseye{T<:FloatingPoint}(N::Integer, D::Integer,
                                    radii::Array{T,1}, widths::Array{T,1})
    if length(radii) != length(widths)
        error("Parameters radii and widths must have same length");
    end
    const M = length(radii);
    data = rand(T, D, N);
    labels = rand(1:M, N);
    for i = 1:N
        len = radii[labels[i]] + randn() * widths[labels[i]] * 0.5;
        angle = rand(D-1)*2*pi;
        sins = one(T);
        for j = 1:D-1
            data[j,i] = len * sins * cos(angle[j]);
            sins *= sin(angle[j]);
        end
        data[D,i] = len * sins;
    end
    return data, labels;
end

function bullseye2{T<:FloatingPoint}(N::Integer, D::Integer,
                                    radii::Array{T,1}, widths::Array{T,1})
    if length(radii) != length(widths)
        error("Parameters radii and widths must have same length");
    end
    const M = length(radii);
    data = rand(T, D, N);
    labels = rand(1:M, N);
    len = radii[labels] + randn(N) .* widths[labels] * 0.5;
    angle = rand(N, D-1) * 2 * pi;
    for i = 1:N
        len = radii[labels[i]] + randn() * widths[labels[i]] * 0.5;
        angle = rand(D-1)*2*pi;
        sins = one(T);
        for j = 1:D-1
            data[j,i] = len * sins * cos(angle[j]);
            sins *= sin(angle[j]);
        end
        data[D,i] = len * sins;
    end
    return data, labels;
end

end


# Ex 1.1
println("Hello World")

# Ex 1.2
for i in 1:10
    println("Hello World")
end

# Ex 1.3
for i in 1:10
    println(i)
end

# Ex 1.4
println("What is your name?")
x = chomp(readline())
if x == "Frederik"
    println("That is a good name, ", x, ". Welcome!")
else
    println("Welcome ", x, "!")
end

# Ex 1.5
println("Enter a number")
x = parse(Float64, chomp(readline()))
println("Enter another number")
y = parse(Float64, chomp(readline()))
println(x, " + ", y, " = ", x + y)

# Ex 1.6
println("Enter a number")
x = parse(Float64, chomp(readline()))
println("Enter another number")
y = parse(Float64, chomp(readline()))

println("Enter an operator (+, -, *, /)")
a = chomp(readline())
f = getfield(Main, Symbol(a))
println(x, " ", a, " ", y, " = ", f(x,y))

# Ex 1.7
println("Enter a number")
x = parse(Int64, chomp(readline()))
println("Enter another number")
y = parse(Int64, chomp(readline()))

for i in x:y
    println(i)
end

# Ex 1.8
println("Enter a number")
x = parse(Int64, chomp(readline()))
println("Enter another number")
y = parse(Int64, chomp(readline()))

if x > y
    x, y = y, x
end

for i in x:y
    println(i)
end

# Ex 1.9
function gt_previous(val, list)
    for elem in list
        if val < elem
            return false
        end
    end
    true
end

println("Enter a number")
x = parse(Int64, chomp(readline()))
num_list = []

while gt_previous(x, num_list)
    push!(num_list, x)
    println("Enter a number")
    global x = parse(Int64, chomp(readline()))
end

# Ex 1.10
println("Enter a non-negative number")
x = parse(Int64, chomp(readline()))

function fact(n::Int64)::Int64
    if n < 0
        error("Negative values are not allowed for factorial.")
    end
    res = 1
    for i in 1:n
        res *= i
    end
    return res
end
println(fact(x))

println("Enter a non-negative number")
x = parse(Int64, chomp(readline()))

fact(n) = (n > 0) ? *(1:n...) :
            (n < 0) ? error("Only non-negative vals allowed.") : 1

println(fact(x))

# Ex 1.11
println("Enter a number")
x = parse(Int64, chomp(readline()))

#if x > 0
#    r = 1:x
#else
#    r = x:-1
#end

#sum = +(r...)

sum = +(1:abs(x)...) * sign(x)

println(sum)


# ------------------------------------------------------------------------------------------
# # Julia is fast
#
# Very often, benchmarks are used to compare languages.  These benchmarks can lead to long
# discussions, first as to exactly what is being benchmarked and secondly what explains the
# differences.  These simple questions can sometimes get more complicated than you at first
# might imagine.
#
# The purpose of this notebook is for you to see a simple benchmark for yourself.  One can
# read the notebook and see what happened on the author's Macbook Pro with a 4-core Intel
# Core I7, or run the notebook yourself.
#
# (This material began life as a wonderful lecture by Steven Johnson at MIT:
# https://github.com/stevengj/18S096-iap17/blob/master/lecture1/Boxes-and-registers.ipynb.)
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # `sum`: An easy enough function to understand
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Consider the  **sum** function `sum(a)`, which computes
# $$
# \mathrm{sum}(a) = \sum_{i=1}^n a_i,
# $$
# where $n$ is the length of `a`.
# ------------------------------------------------------------------------------------------

a = rand(10^7) # 1D vector of random numbers, uniform on [0,1)

sum(a)

# ------------------------------------------------------------------------------------------
# The expected result is 0.5 * 10^7, since the mean of each entry is 0.5
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Benchmarking a few ways in a few languages
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Julia has a `BenchmarkTools.jl` package for easy and accurate benchmarking:
# ------------------------------------------------------------------------------------------

#Pkg.add("BenchmarkTools")

using BenchmarkTools

# ------------------------------------------------------------------------------------------
# #  1. The C language
#
# C is often considered the gold standard: difficult on the human, nice for the machine.
# Getting within a factor of 2 of C is often satisfying. Nonetheless, even within C, there
# are many kinds of optimizations possible that a naive C writer may or may not get the
# advantage of.
#
# The current author does not speak C, so he does not read the cell below, but is happy to
# know that you can put C code in a Julia session, compile it, and run it. Note that the
# `"""` wrap a multi-line string.
# ------------------------------------------------------------------------------------------

C_code = """
#include <stddef.h>
double c_sum(size_t n, double *X) {
    double s = 0.0;
    for (size_t i = 0; i < n; ++i) {
        s += X[i];
    }
    return s;
}
"""

const Clib = tempname()   # make a temporary file

# compile to a shared library by piping C_code to gcc
# (works only if you have gcc installed):

open(`gcc  -fPIC -O3 -ffast-math -msse3 -xc -shared -o $(Clib * "." * Libdl.dlext) -`, "w") do f
    print(f, C_code)
end

# define a Julia function that calls the C function:
c_sum(X::Array{Float64}) = ccall(("c_sum", Clib), Float64, (Csize_t, Ptr{Float64}), length(X), X)

c_sum(a)

c_sum(a) ≈ sum(a) # type \approx and then <TAB> to get the ≈ symbolb

≈  # alias for the `isapprox` function

?isapprox

# ------------------------------------------------------------------------------------------
# We can now benchmark the C code directly from Julia:
# ------------------------------------------------------------------------------------------

c_bench = @benchmark c_sum($a)

println("C: Fastest time was $(minimum(c_bench.times) / 1e6) msec")

#d = Dict()  # a "dictionary", i.e. an associative array
d["C"] = minimum(c_bench.times) / 1e6  # in milliseconds
d

using Plots
gr()

t = c_bench.times / 1e6 # times in milliseconds
m, σ = minimum(t), std(t)

histogram(t, bins=500,
    xlim=(m - 0.01, m + σ),
    xlabel="milliseconds", ylabel="count", label="")

# ------------------------------------------------------------------------------------------
# # 2. Python's built in `sum`
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# The `PyCall` package provides a Julia interface to Python:
# ------------------------------------------------------------------------------------------

#Pkg.add("PyCall")

using PyCall

# Call a low-level PyCall function to get a Python list, because
# by default PyCall will convert to a NumPy array instead (we benchmark NumPy below):

apy_list = PyCall.array2py(a, 1, 1)

# get the Python built-in "sum" function:
pysum = pybuiltin("sum")

pysum(a)

pysum(a) ≈ sum(a)

py_list_bench = @benchmark $pysum($apy_list)

d["Python built-in"] = minimum(py_list_bench.times) / 1e6
d

# ------------------------------------------------------------------------------------------
# # 3. Python: `numpy`
#
# ## Takes advantage of hardware "SIMD", but only works when it works.
#
# `numpy` is an optimized C library, callable from Python.
# It may be installed within Julia as follows:
# ------------------------------------------------------------------------------------------

using Conda
#Conda.add("numpy")

numpy_sum = pyimport("numpy")["sum"]
apy_numpy = PyObject(a) # converts to a numpy array by default

py_numpy_bench = @benchmark $numpy_sum($apy_numpy)

numpy_sum(apy_list) # python thing

numpy_sum(apy_list) ≈ sum(a)

d["Python numpy"] = minimum(py_numpy_bench.times) / 1e6
d

# ------------------------------------------------------------------------------------------
# # 4. Python, hand-written
# ------------------------------------------------------------------------------------------

py"""
def py_sum(a):
    s = 0.0
    for x in a:
        s = s + x
    return s
"""

sum_py = py"py_sum"

py_hand = @benchmark $sum_py($apy_list)

sum_py(apy_list)

sum_py(apy_list) ≈ sum(a)

d["Python hand-written"] = minimum(py_hand.times) / 1e6
d

# ------------------------------------------------------------------------------------------
# # 5. Julia (built-in)
#
# ## Written directly in Julia, not in C!
# ------------------------------------------------------------------------------------------

@which sum(a)

j_bench = @benchmark sum($a)

d["Julia built-in"] = minimum(j_bench.times) / 1e6
d

# ------------------------------------------------------------------------------------------
# # 6. Julia (hand-written)
# ------------------------------------------------------------------------------------------

function mysum(A)
    s = 0.0  # s = zero(eltype(A))
    for a in A
        s += a
    end
    s
end

j_bench_hand = @benchmark mysum($a)

d["Julia hand-written"] = minimum(j_bench_hand.times) / 1e6
d

# ------------------------------------------------------------------------------------------
# # Summary
# ------------------------------------------------------------------------------------------

for (key, value) in sort(collect(d))
    println(rpad(key, 20, "."), lpad(round(value, 1), 8, "."))
end

for (key, value) in sort(collect(d), by=x->x[2])
    println(rpad(key, 20, "."), lpad(round(value, 2), 10, "."))
end

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/22. Julia is fast.jl


### Without renewables for validation of IEEE test cases
###################
# Run DCOPF #######
###################

DCOPF_mod, P_opt, Θ_opt, P_inj_opt, P_flow_opt =
	DCOPF(time, nodes, lines, pp)

# Optimisation status
JuMP.termination_status(DCOPF_mod)
JuMP.primal_status(DCOPF_mod)
JuMP.dual_status(DCOPF_mod)

# Results
JuMP.objective_value(DCOPF_mod)
P_opt
Θ_opt
P_inj_opt
P_flow_opt

###################
# Run ED ##########
###################

ED_mod, P_opt, price = EconomicDispatch(time, nodes, pp)

# Optimisation status
JuMP.termination_status(ED_mod)
JuMP.primal_status(ED_mod)
JuMP.dual_status(ED_mod)

# Results
JuMP.objective_value(ED_mod)
P_opt
price

# Dispatch: Unit commitment and power generation
output_EconomicDispatch = DataFrame(time=Int64[],unit=Int64[],output=Float64[])
for t in time
	for g in pp.unit
		val = P_opt[t, g]
		push!(output_EconomicDispatch, [t g val])
	end
end

CSV.write("output/EconomicDispatch.csv", output_EconomicDispatch)
electrondisplay(output_EconomicDispatch)
electrondisplay(powerplants_df)

###################
# Run CM ##########
###################

CM_mod, ΔP_up_opt, ΔP_dn_opt, Θ_opt, P_inj_opt, P_flow_opt =
	CongestionManagement(time, nodes, lines, price, P_opt, pp)

# Optimisation status
JuMP.termination_status(CM_mod)
JuMP.primal_status(CM_mod)
JuMP.dual_status(CM_mod)

# Results
JuMP.objective_value(CM_mod)
ΔP_up_opt
ΔP_dn_opt
Θ_opt
P_inj_opt
P_flow_opt

# Testing
P_opt.data .+ ΔP_up_opt.data .- ΔP_dn_opt.data
P_opt.data


using Test, LinearAlgebra
using DiffEqOperators, OrdinaryDiffEq

@testset "Matrix Free Operator" begin
  f = (du, u, p) -> begin
    cumsum!(du, u)
    @. du = p*du
  end
  p = 1.
  A = MatrixFreeOperator(f, (p,), size=(5,5), opnorm=(p)->5)
  b = rand(5)
  @test is_constant(A)
  prob = ODEProblem(A, b, (0,1.), p)
  @test_nowarn solve(prob, Tsit5())

  f = (du, u, p, t) -> cumsum!(du, u)
  args = (1, 1)
  A = MatrixFreeOperator(f,args, size=(5,5), opnorm=(p)->5)
  @test is_constant(A) == false
  b = rand(5)
  @test_nowarn solve(prob, LawsonEuler(krylov=true, m=5), dt=0.1)

  A1 = [[4 -6]; [1 -1]]
  A2 = [[1 2]; [3 2]]
  THRESHOLD = 1
  # analytic solution for the DE
  u1(t) = t < THRESHOLD ? exp(t)*(-1+3exp(t)) : exp(-3-t)*(exp(5)*(-2+7exp(1))+exp(5t)*(-3+8exp(1)))/5
  u2(t) = t < THRESHOLD ? -exp(t)/2+exp(2t)   : (2-7exp(1))*exp(2-t)/5+exp(-3+4t)*(-3+8exp(1))*(3/10)
  sol_analytic(t) = [u1(t), u2(t)]
  # setups for DE solvers
  function Q!(df, f, p, t)
    i = t >= THRESHOLD ? 2 : 1
    A = p[i]
    mul!(df, A, f)
  end
  p = (A1,A2)
  O = MatrixFreeOperator(Q!, (p, 0.), size=(2,2), opnorm=(p)->10)
  # solve DE numerically
  T = 2
  f_0 = [2.0; 1/2]
  prob = ODEProblem(O,f_0,(0.0,T),p)
  sol = solve(prob, LinearExponential(krylov=:simple, m=3), tstops=THRESHOLD)
  @test sol(THRESHOLD) ≈ sol_analytic(THRESHOLD) atol = 1e-3
  @test sol(T) ≈ sol_analytic(T) atol = 1e-3
end


using BoltzmannMachines
import BoltzmannMachinesPlots # only needed for plotting capabilities

# ==============================================================================
# Binary data: BernoulliRBM and BasicDBM
# ------------------------------------------------------------------------------

# Generate data and split it in training and test data set
nsamples = 100;
nvariables = 36;
using Random
Random.seed!(1);
x = barsandstripes(nsamples, nvariables);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);

# Fit a BernoulliRBM and collect monitoring data during training.
# For small models, the model's loglikelihood can be computed exactly.
monitor = Monitor();
rbm = fitrbm(x; nhidden = 12,
      epochs = 80, learningrate = 0.006,
      monitoring = (rbm, epoch) -> begin
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict)
            monitorreconstructionerror!(monitor, rbm, epoch, datadict)
      end);

# Plot the collected data (requires Gadfly)
BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)
BoltzmannMachinesPlots.plotevaluation(monitor, monitorreconstructionerror)

Random.seed!(12);

# DBM-Fitting: Pretraining and Fine-Tuning combined in one function
dbm = fitdbm(x, nhiddens = [6;2], epochs = 20, learningrate = 0.05);

# .. with extensive monitoring
monitor = Monitor(); monitor1 = Monitor(); monitor2 = Monitor();
dbm = fitdbm(x, epochs = 20, learningrate = 0.05,
      monitoringdatapretraining = datadict,
      pretraining = [
            TrainLayer(nhidden = 6, monitoring = (rbm, epoch, datadict) ->
                  monitorreconstructionerror!(monitor1, rbm, epoch, datadict));
            TrainLayer(nhidden = 2, monitoring = (rbm, epoch, datadict) ->
                  monitorreconstructionerror!(monitor2, rbm, epoch, datadict))
            ],
      monitoring = (dbm, epoch) ->
            monitorexactloglikelihood!(monitor, dbm, epoch, datadict));

# Monitoring plots
BoltzmannMachinesPlots.plotevaluation(monitor1, monitorreconstructionerror)
BoltzmannMachinesPlots.plotevaluation(monitor2, monitorreconstructionerror)
BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)

# Evaluate final result with exact computation of likelihood
exactloglikelihood(dbm, xtest)

# If the model has more parameters, in this case more hidden nodes,
# the exact calclation is not feasible any more:
# We need to calculate the likelihood using AIS.

Random.seed!(12);
monitor = Monitor();
rbm = fitrbm(x, nhidden = 36, epochs = 100,
      learningrates = [0.008*ones(80); 0.001*ones(20)],
      monitoring = (rbm, epoch) ->
            monitorloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorloglikelihood; sdrange = 3.0)

# In the DBM, the estimation of the log likelihood is much slower than in the
# RBM because it needs a AIS run for each sample.
# The lower bound of the log probability is faster to estimate:
# (But on the other hand, in some cases it might be too conservative to be
# useful as it is only a lower bound.)

Random.seed!(2);
monitor = Monitor();
dbm = stackrbms(x; nhiddens = [36;10;5], predbm = true, learningrate = 0.05);
traindbm!(dbm, x; epochs = 50, learningrate = 0.008,
      monitoring = (rbm, epoch) ->
            monitorlogproblowerbound!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorlogproblowerbound; sdrange = 3.0)

# Evaluate final result with AIS-estimated likelihood
loglikelihood(dbm, xtest)

# ==============================================================================
# Real valued data: GaussianBernoulliRBM
# ------------------------------------------------------------------------------

# Use "iris" dataset as example data to train a GaussianBernoulliRBM
using DelimitedFiles
x = convert(Matrix{Float64}, readdlm(
            joinpath(dirname(pathof(BoltzmannMachines)), "..", "test/data/iris.csv"), ',',
            header = true)[1][:,1:4]);
Random.seed!(12);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);
monitor = Monitor();
rbm = fitrbm(x, rbmtype = GaussianBernoulliRBM,
      nhidden = 3, epochs = 70, learningrate = 0.001,
      monitoring = (rbm, epoch) ->
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)

# ==============================================================================
# Data with binomially distributed values in {0,1,2}: Binomial2BernoulliRBM
# ------------------------------------------------------------------------------

# Generate data and split it in training and test data set
nsamples = 100;
nvariables = 25;
Random.seed!(12);
x = barsandstripes(nsamples, nvariables) + barsandstripes(nsamples, nvariables);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);

# Fit Binomial2BernoulliRBM.
monitor = Monitor();
rbm = fitrbm(x, rbmtype = Binomial2BernoulliRBM,
      nhidden = 10, epochs = 100, learningrates = [0.003*ones(60); 0.001*ones(40)],
      monitoring = (rbm, epoch) ->
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)

# ==============================================================================
# Examples for cross-validation
# ------------------------------------------------------------------------------

nsamples = 500;
nvariables = 16;
x = barsandstripes(nsamples, nvariables);

# Determine the optimal number of training epochs for a RBM
monitor = crossvalidation(x,
      (x, datadict) ->
            begin
               monitor = BoltzmannMachines.Monitor()
               BoltzmannMachines.fitrbm(x, epochs = 50,
                     learningrate = 0.005,
                     monitoring = (rbm, epoch) ->
                           BoltzmannMachines.monitorexactloglikelihood!(
                                 monitor, rbm, epoch, datadict))
               monitor
            end);
BoltzmannMachinesPlots.crossvalidationcurve(monitor)

# Determine the optimal number of pretraining epochs for a DBM
# given the other parameters
using Distributed
@everywhere using BoltzmannMachines
@everywhere function my_pretraining_monitoring(
      x::Matrix{Float64}, datadict::DataDict, epoch::Int)

   monitor = Monitor()
   monitorlogproblowerbound!(monitor,
         stackrbms(x,
               nhiddens = [16; 12],
               learningrate = 0.001,
               predbm = true,
               epochs = epoch),
         epoch,
         datadict,
         parallelized = false)
end

monitor = crossvalidation(x, my_pretraining_monitoring, 10:10:200);
BoltzmannMachinesPlots.crossvalidationcurve(monitor, monitorlogproblowerbound)

# ==============================================================================
# Examples for custom optimization algorithm
# ------------------------------------------------------------------------------

x = barsandstripes(100, 16);

struct MyRegularizedOptimizer{R<:AbstractRBM} <: AbstractOptimizer{R}
   llopt::LoglikelihoodOptimizer{R}
end

function MyRegularizedOptimizer()
   MyRegularizedOptimizer(loglikelihoodoptimizer(learningrate = 0.01))
end

import BoltzmannMachines: initialized, computegradient!, updateparameters!

function initialized(myoptimizer::MyRegularizedOptimizer,
      rbm::R) where {R<: AbstractRBM}
   MyRegularizedOptimizer(initialized(myoptimizer.llopt, rbm))
end

function computegradient!(myoptimizer::MyRegularizedOptimizer,
      v, vmodel, h, hmodel, rbm)
   computegradient!(myoptimizer.llopt, v, vmodel, h, hmodel, rbm)
   # add L2 regularization term
   myoptimizer.llopt.gradient.weights .-= 0.1 .* rbm.weights
   nothing
end

function updateparameters!(rbm::R, myoptimizer::MyRegularizedOptimizer{R}
      ) where {R<: AbstractRBM}
   updateparameters!(rbm, myoptimizer.llopt)
end

# The optimizer can be used for fitting RBMs ...
rbm = fitrbm(x; optimizer = MyRegularizedOptimizer())

# and also for DBMs
dbm = fitdbm(x; optimizer = MyRegularizedOptimizer())