text:
module EllZeroTrendFiltering

export QuadraticPolynomial, PiecewiseQuadratic, QuadraticForm
export fit_pwl_constrained, fit_pwl_regularized
export pwq_dp_regularized, pwq_dp_constrained
#Should we export the following?
export compute_transition_costs, compute_discrete_transition_costs
export recover_optimal_index_set, recover_solution, brute_force_search
export snp500_data

import Base.-
import Base.+
import Base.show
import Base: iterate, length, zero, getindex, ==

import IterTools

using StaticArrays
using QuadGK

#StdLibrary
using LinearAlgebra
import DelimitedFiles: readdlm
import Printf: @printf
#Pkg, DelimitedFiles
global const DEBUG = false
global const DEBUG2 = false
global const COUNTER_TEST = false
global const OPTIMIZE = true

# TODO Tweak this accuracy, sqrt(eps()) was not good enough for snp 500, M = 12 test
global const ACCURACY = sqrt(eps())/100000

include(joinpath("types","QuadraticPolynomial.jl"))
include(joinpath("types","PiecewiseQuadratic.jl"))
include(joinpath("types","QuadraticForm.jl"))
include(joinpath("types","transition_costs.jl"))

AbstractTransitionCost{T} = Union{Array{QuadraticForm{T},2}, TransitionCostDiscrete{T}, TransitionCostContinuous{T}}

include("transition_cost_computation.jl")
include("brute_force_search.jl")
include("solve.jl")

snp500_data() = readdlm(joinpath(dirname(@__FILE__),"..","examples","data","snp500.txt"))

end # module


text:
@testset "Transforms" begin
  # compositions
  for c in [Composition(1,2,3),
            Composition(0,2,3),
            Composition(1,2,0)]
    @test alrinv(alr(c)) == c
    @test clrinv(clr(c)) == c
    @test ilrinv(ilr(c)) == c
  end

  # TableTransform.jl API
  t = (a=Float64[1,0,1], b=Float64[2,2,2], c=Float64[3,3,0])
  n, c = apply(ALR(), t)
  talr = revert(ALR(), n, c)
  n, c = apply(CLR(), t)
  tclr = revert(CLR(), n, c)
  n, c = apply(ILR(), t)
  tilr = revert(ILR(), n, c)
  @test Tables.matrix(talr) ≈ Tables.matrix(tclr)
  @test Tables.matrix(tclr) ≈ Tables.matrix(tilr)
  @test Tables.matrix(talr) ≈ Tables.matrix(tilr)

  n, c = apply(ALR(:b), t)
  talr = revert(ALR(:b), n, c)
  @test n    |> Tables.columnnames == (:a, :c)
  @test talr |> Tables.columnnames == (:a, :b, :c)

  n, c = apply(ILR(:b), t)
  tilr = revert(ILR(:b), n, c)
  @test n    |> Tables.columnnames == (:a, :c)
  @test tilr |> Tables.columnnames == (:a, :b, :c)

  # Tests for Closure
  t = (a=Float64[2,66,0], b=Float64[4,22,2], c=Float64[4,12,98])
  n, c = apply(Closure(), t)
  tcls = revert(Closure(), n, c)
  @test Tables.matrix(n) ≈ [0.2 0.4 0.4; 0.66 0.22 0.12; 0.00 0.02 0.98;]
  @test Tables.matrix(tcls) ≈ Tables.matrix(t)
end

ext: jl


text:
# Minimal implementation. Sort by columns

function _col_lt(icols...)
    function (x, y)
        @inbounds for c in icols
            x[c] == y[c] && continue
            return isless(x[c], y[c])
        end
        return false
    end
end

"""
    sort!(rt::RowTable; cols=[], kws...)

Sort `rt` by columns `cols`.
"""
function Base.sort!(rt::RowTable, cols=[]; kws...)
    if isempty(cols)
        icols = collect(1:size(rt, 2))
    else
        icols = cindex(rt)[cols]
    end
    nc = length(icols)
    if nc == 1
        @inbounds c = icols[1]
        @inbounds Base.sort!(rows(rt); lt = (x, y)-> isless(x[c], y[c]), kws...)
    else
        Base.sort!(rows(rt); lt=_col_lt(icols...), kws...)
    end
    return rt
end

Base.sort(rt::RowTable, cols=[]; kws...) = sort!(copy(rt), cols; kws...)


text:
using Documenter, BenchSweeps

makedocs(;
    modules=[BenchSweeps],
    format=:html,
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/tkf/BenchSweeps.jl/blob/{commit}{path}#L{line}",
    sitename="BenchSweeps.jl",
    authors="Takafumi Arakaki",
    assets=[],
)

deploydocs(;
    repo="github.com/tkf/BenchSweeps.jl",
    target="build",
    julia="1.0",
    deps=nothing,
    make=nothing,
)


text:
module ServerHandler

using CSV
using DataFrames
using LibPQ

#include( *( @__DIR__, "\\Code\\data.sql" ) )

params = Dict(
    :dbname => "data",
    :host => "localhost",
    :port => "5432",
    :user => "postgress",
    :password => "password",
    :connect_timeout => "10"
)

#   con = LibPQ.Connection( "dbname=$(params[:dbname]) host=$(params[:host]) port=$(params[:port]) user=$(params[:user]) password=$(params[:password])" )

"""
    updateProductsVal( key::AbstractString, table::AbstractString, attribute::AbstractString, newval::Any )

For the row in table `table` of the database that has `key` as its pk, update "attribute" to have value `newval`
"""
function updateProductsVal( key::AbstractString, table::AbstractString, attribute::AbstractString, newval )
    # Connettiti al db
    # Effettua l'update dell'attributo nella tabella del db indicata alla riga indicata da key

    println("Valore modificato")
end

"""
    updateProductsVal( keys::AbstractVector{AbstractString}, table::AbstractString, attribute::AbstractString, newval::Any )

For all the rows in table `table` of the database that have one of `keys` as their pk, update `attribute` to have value `newval`
"""
function updateProductsVal( keys::AbstractVector{AbstractString}, table::AbstractString, attribute::AbstractString, newval )
    # Connettiti al db
    # Effettua l'update dell'attributo nella tabella del db indicata alla riga indicata da key

    println("Valore modificato")
end

"""
    updateProductsVal( condition::Function, table::AbstractString, attribute::AbstractString, newval::Any )

For all the rows in table `table` of the database that verify `condition`, update `attribute` to have value `newval`
"""
function updateProductsVal( condition::Function, table::AbstractString, attribute::AbstractString, newval )
    # Connettiti al db
    # Effettua l'update dell'attributo nella tabella del db indicata per tutte le righe che rispettano la condizione

    println("Valore modificato")
end

end # module

max_stars_repo_path: src/Data Gathering/ServerHandler.jl

max_stars_repo_name: DavideFavaro/Tirocinio

max_issues_repo_path: src/Data Gathering/ServerHandler.jl

max_issues_repo_name: DavideFavaro/Tirocinio

max_forks_repo_path: src/Data Gathering/ServerHandler.jl

max_forks_repo_name: DavideFavaro/Tirocinio


text:
@static if VERSION >= v"0.7.0-DEV.2915"
using Distributed
end
addprocs(5)

using HTTP
using HTTP.Dates
using HTTP.Test

@testset "HTTP" begin
    include("utils.jl");
    include("fifobuffer.jl");
    include("sniff.jl");
    include("uri.jl");
    include("url.jl");
    include("cookies.jl");
    include("parser.jl");

    include("loopback.jl");
    include("WebSockets.jl");
    include("messages.jl");
    include("client.jl");

    include("handlers.jl")
    include("server.jl")

    include("async.jl");
end;


text:
using Measurements
r_ini = [14.0,0.1,
        0.14,0.4,
        5.0,0.2,
        20.0,4.0]

model = bimolecular_diss_bell_activated_state_w_cons!
# u₀ = [1.0,0.0]
u₀ = [] # don't provide initial states
u₀_type = "equilibrium" # calculate initial states based on equilibrium
# u₀ = "detailed_balance"
# u₀ = "force"
# u₀ = []
tspan = (0.0,maximum(γϵ.t))
# k₋ᵒf_γϵ = 14.1
k₋ᵒf_γϵ = 12.2
γϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[k₋ᵒf_γϵ],0.0,Rodas5)
# γϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[0.5],0.0,Rodas5)

pinit = [   0.01,
        0.1,0.6,
        3.,0.2,
        150.0,4.0]
opt =  [        (0.,10.),
        (0.,1.),(0.,10.),
        (0.,10.),(0.,10.),
        (0.,500.),(0.,10.)]

@time mle_loss(pinit,γϵ,γϵ_ft_M)
# loss_in = (p) -> mle_loss(p,γϵ,γϵ_ft_M)
# hes = ForwardDiff.hessian(loss_in,rates_γϵ)
# σ_γϵ = hessian2σ(hes)

@time rates_γϵ, l_γϵ, hes_γϵ = optimization(γϵ_ft_M,γϵ,opt)
ks_γϵ = rates_γϵ
σ_γϵ = hessian2σ(hes_γϵ)
k_γϵ = rates_γϵ .± σ_γϵ

plotly()
plot_ft_fit(γϵ_ft_M,ks_γϵ,γϵ)
# plot_ft_fit(γϵ_ft_M,ks_γϵ,γϵ,dF=30)
plot(); Color = :blue
plot_mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,γϵ_binned;species="γϵ")
# sols_γϵ = ftsolve(γϵ_ft_M,k_γϵ) #change solver to Tsit5
# plot(Fs,mean_lifetime.(sols_γϵ),ylabel="<t>",xlabel="F",title="γϵ")
# plot_F_vs_t(clean_dissociation_data(γϵ_binned))
u₀γϵ = initial_u(ks_γϵ,γϵ_ft_M)
________________________________________________________________________________________________
## δϵ fitting with 2 species model
model = bimolecular_diss_bell_activated_state_w_cons!
u₀ = []
u₀_type = "equilibrium"
tspan = (0.0,maximum(δϵ.t))
# k₋ᵒf_δϵ = 11.1
k₋ᵒf_δϵ = 13.8
δϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[k₋ᵒf_δϵ],0.0,Rodas5)
pinit = [   0.1,
        0.1,0.5,
        0.1,0.5,
        10.0,5.0]
opt =  [        (0.,10.),
        (0.,1.),(0.,10.),
        (0.,10.),(0.,10.),
        (0.,100.),(0.,10.)]

@time mle_loss(pinit,δϵ,δϵ_ft_M)

@time rates_δϵ, l_δϵ, hes_δϵ = optimization(δϵ_ft_M,δϵ,opt)
ks_δϵ = rates_δϵ
σ_δϵ = hessian2σ(hes_δϵ)
k_δϵ = rates_δϵ .± σ_δϵ
u₀δϵ = initial_u(ks_δϵ,δϵ_ft_M)

plot_ft_fit(δϵ_ft_M,ks_δϵ,δϵ)
plot();
Color = :red
plot_mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,δϵ_binned;species="δϵ")
# plot(Fs,mean_lifetime.(sols_δϵ),ylabel="<t>",xlabel="F",title="δϵ")
# plot_F_vs_t(clean_dissociation_data(δϵ_binned))

## mix fitting with 2 species model

model = cd3_tri_dissociation_bell!
# u₀_type = "equilibrium"
u₀_type = []
# u₀ = vcat(u₀γϵ./2, u₀δϵ./2, 0.0,0.0)
u₀ = vcat(0.454*u₀γϵ,0.290*u₀δϵ, 0.256*[1.,0.])
# u₀ = vcat(0.454*u₀γϵ,0.290*u₀δϵ, 0.256*[0.,0.])/(1-0.256)
sum(u₀)
# u₀ = vcat(u₀γϵ/3, u₀δϵ/3)
tspan = (0.0,maximum(mix.t))
k₋ᵒf_mix = 0.6
cons = [k₋ᵒf_γϵ;
        ks_γϵ;
        k₋ᵒf_δϵ;
        ks_δϵ;
        k₋ᵒf_mix ]
mix_ft_M = M(model,u₀,u₀_type,tspan,[],cons,0.0,Rodas5)

pinit = [   0.1,
        0.1,0.5,
        0.1,0.5,
        10.0,5.0]
# pinit = zeros(7)
opt =  [        (0.,10.),
        (0.,1.),(0.,1.),
        (0.,10.),(0.,1.),
        (0.,100.),(0.,10.)]

@time mle_loss(pinit,mix,mix_ft_M)
loss_in = (p) -> mle_loss(p,mix,mix_ft_M)
hes_mix = ForwardDiff.hessian(loss_in,pinit)
# hes = Zygote.hessian(loss_in,pinit)

@time rates_mix, l_mix, hes_mix = optimization(mix_ft_M,mix,opt)
ks_mix = rates_mix
σ_mix = hessian2σ(hes_mix)
k_mix = rates_mix .± σ_mix

ts = 0:dt:maximum(mix.t)
Fs = 0:1:maximum(mix.F)
plot_ft_fit(mix_ft_M,ks_mix,mix)
# sols_mix = plot_ft_fit(mix_ft_M,pinit,mix)

plot_ft_fit(mix_ft_M,ks_mix,mix)
plot(); Color = :green
plot_mean_ft_fit(mix_ft_M,k_mix,mix,mix_binned;species="mix")
# plot(Fs,mean_lifetime.(sols_mix),ylabel="<t>",xlabel="F",title="mix")
# plot_F_vs_t(clean_dissociation_data(mix_binned))

## plot rates

plt = plot_bell_equations(vcat(k₋ᵒf_γϵ, ks_γϵ))
title!("γϵ")
plt = plot_bell_equations(vcat(k₋ᵒf_δϵ, ks_δϵ))
title!("δϵ")
plt = plot_bell_equations(vcat(k₋ᵒf_mix, ks_mix))
title!("mix")

pyplot()
plot(); Color = :blue
plot_mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,γϵ_binned;species="γϵ")
Color = :red
plot_mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,δϵ_binned;species="δϵ")
Color = :green
plot_mean_ft_fit(mix_ft_M,k_mix,mix,mix_binned;species="mix")

plot(); Color = :blue
extrapolate_fit(γϵ_ft_M,ks_γϵ,γϵ,150.0;species="γϵ",dF=10.0)
Color = :red
extrapolate_fit(δϵ_ft_M,ks_δϵ,δϵ,150.0;species="δϵ",dF=10.0)
Color = :green
extrapolate_fit(mix_ft_M,ks_mix,mix,150.0;species="mix",dF=10.0)

mean_t_γϵ = mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,35;species=[],dt=0.1,dF=1.0)
mean_t_δϵ = mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,35;species=[],dt=0.1,dF=1.0)
mean_t_mix = mean_ft_fit(mix_ft_M,k_mix,mix,35;species=[],dt=0.1,dF=1.0)

mean_t_sum_bi = 0.610*mean_t_γϵ + 0.390*mean_t_δϵ
plot!(mean_t_sum_bi)

pyplot()
lftm_coop = (mean_t_mix -mean_t_sum_bi) ./ mean_t_sum_bi
plot(0:35, 100 .* lftm_coop,errorbar=:ribbon, xlabel="F",ylabel="lifetime cooperativity (%)")

CSV.write("results/γϵ_fit.csv", measurements2Df(0:35,mean_t_γϵ))
CSV.write("results/δϵ_fit.csv", measurements2Df(0:35,mean_t_δϵ))
CSV.write("results/mix_fit.csv", measurements2Df(0:35,mean_t_mix))
CSV.write("results/lftm_coop.csv", measurements2Df(0:35,lftm_coop))


text:
using DimArrays
using Test

a = DimVector([1,2,3,4], :aa, 100)

@test DimArrays.dname(a,1) == :aa
@test DimArrays.ifunc(a,1)(1) == 100
@test DimArrays.haslabel(a) == false

@test typeof(a') <: DimMatrix
@test ndims( (a')' )==1
@test DimArrays.dname( (a')' ,1) == :aa

@test typeof( map(sqrt, a) ) <: DimVector
@test typeof( [sqrt(i) for i in a] ) <: DimVector

b = DimArray([1.0 2; 3 4]; label=:c)

@test string(b) == "DimArray([1.0 2.0; 3.0 4.0]; label = :c)"

@test string(name!(selectdim(b, :col, 2), "one", "cont")) == "DimArray([2.0, 4.0], :one; label = :cont)"

@test permutedims(b, (:col, :row))[1,2] ≈ 3

using Statistics

@test sum(b; dims=[:row,:col])[1] ≈ 10
@test std(b; dims=:col)[1] ≈ 0.7071067811865476
@test typeof( mean(b; dims=2) ) <: DimMatrix

@test selectdim(b, :col, 1)[2] ≈ 3
@test typeof( selectdim(b, 2, 1) ) <: DimVector

@test (b .+ 99)[1,1] ≈ 100
# @test typeof(b .+ 99) <: DimMatrix

@test sum( hcat(DimArray([1,3]), [2,4]) ./ b ) ≈ 4

@test maximum( nest([DimArray([1,3]), [2,4]]) .- b ) ≈ 0

# aabb = vcat(hcat(a,a),b,b)

# @test size(aabb) == (8,2)
# @test DimArrays.dnames(aabb) == [:aa, :col]

push!(a, 55,66)
@test DimArrays.ifuncs(a)[1](length(a)) == 600

c = dictvector(ones(2), [:α, "β"], :ω, :N)

push!(c, 77, 88)
append!(c, [0,0])
push!(c, 99, :γ)

@test c isa DimVector
# @test sum(c, :ω)[1] ≈ 266 # todo: dims=: etc


text:
module FieldInterfaceTests

using Test
using InplaceArrays.TensorValues
using InplaceArrays.Fields
using InplaceArrays.Fields: MockField

v = 3
w = 4.0
d = 2
xi = Point(3,4)
np = 4
x = fill(xi,np)
f = MockField{d}(v)
g = MockField{d}(w)

cf = field_cache(f,x)
cg = field_cache(g,x)
c = field_caches((f,g),x)
@test (cf,cg) == c

fx = evaluate_field!(cf,f,x)
gx = evaluate_field!(cg,g,x)
fgx = evaluate_fields!(c,(f,g),x)
@test (fx,gx) == fgx

Tf = field_return_type(f,x)
Tg = field_return_type(g,x)
Tfg = field_return_types((f,g),x)
@test (Tf,Tg) == Tfg

∇f = field_gradient(f)
∇g = field_gradient(g)
∇fg = field_gradients(f,g)
@test (∇f,∇g) == ∇fg

x = Point(1,2)
@test gradient_type(Float64,x) == VectorValue{2,Float64}
@test gradient_type(VectorValue{2,Float64},x) == TensorValue{2,Float64,4}
@test gradient_type(VectorValue{3,Float64},x) == MultiValue{Tuple{2,3},Float64,2,6}

end # module


text:
im_mean = reshape([0.485, 0.456, 0.406], 1, 1, 3)
global original_size = (224, 224)

im2arr_rgb(img) = permutedims(float.(channelview(imresize(img, (224, 224)))), (3, 2, 1))

"""
    set_training_false()

Function to set the model into `testmode`, so that the output result do not vary while evaluating gradients(or other results) for the models which contain layers like `BatchNormalization`, `Dropout` etc.
"""
function set_training_false()
    @eval Flux.istraining() = false
    @eval Zygote.@adjoint Flux.istraining() = false, _ -> nothing
end

"""
    set_training_true()

Function to revert back the model into `trainmode`, once the gradients(or other results) are evaluated.
"""
function set_training_true()
    @eval Flux.istraining() = true
    @eval Zygote.@adjoint Flux.istraining() = true, _ -> nothing
end

# one_hot_encode: Utility to return one_hot array for provided index with size same as that of `preds`
function one_hot_encode(preds, idx)
  one_hot = zeros(eltype(preds), size(preds)[1], 1)
  one_hot[idx.I[1] ,1] = 1
  one_hot |> gpu
end

# get_topk: Utility to return top_k predicted classes with their probabilities
function get_topk(probs; k = 5)
  T = eltype(probs)
  prob = Array{T, 1}()
  idx = []
  while(k!=0)
    push!(idx, argmax(probs))
    push!(prob, probs[idx[end]])
    probs[idx[end]] = 0.0
    k -= 1
  end
  (prob, idx)
end

"""
    save_gradient_images(gradient, file_name; gray=false)

Saves the gradient as an RGB image by default. If `gray` is set to true, gradients are sasved as a gray image with the provided `filename`.
"""
function save_gradient_images(gradient, file_name; gray=false)
  gradient = gradient .- minimum(gradient)
  gradient ./= maximum(gradient)
  gradient = permutedims(dropdims(gradient, dims=4), (3, 2, 1))
  img = colorview(RGB{eltype(gradient)}, gradient)
  if (gray)
      img = Gray.(img)
  end
  display(img)
  @info("Saving Gradient Image......")
  save(file_name, img)
end

"""
    positive_negative_saliency(gradient)

Returns a tuple of positive and negative saliency maps based on the provided gradients.
"""
function positive_negative_saliency(gradient)
  pos_saliency = max.(zero(gradient), gradient) ./ maximum(gradient)
  neg_saliency = max.(zero(gradient), -gradient) ./ maximum(-gradient)
  (pos_saliency, neg_saliency)
end

"""
    image_to_arr(img; preprocess = true)

Convert an RGB type image to an array. If `preprocess` is set as `true`, subtract the mean from the image array.
"""
function image_to_arr(img; preprocess = true)
  local x = img
  x = Float32.(channelview(img))
  x = permutedims(x, [3,2,1])
  if(preprocess)
    x = x .- im_mean
  end
  x = reshape(x, size(x,1), size(x,2), size(x,3), 1) * 255 |> gpu
end

"""
    load_image(path, resize = false; size_save = true)

Function to load the images from the specified `path`. Returns as an array of image. If `resize` is set as `true`, resize the image to shape (224, 224).
"""
function load_image(path, resize = false; size_save = true)
  img = load(path)
  if(size_save)
    global original_size = size(img)
  end
  if(resize)
    image_to_arr(imresize(img, (224, 224)))
  else
    image_to_arr(img)
  end
end

"""
     generate_image(x; resize_original=false, custom_size=Nothing)

Save the image as an RGB image defined the by the array `x`. If `resize_original` is set to be true, then the image can be resized to the `custom_size`. If not provided explicitly, then the image is resized to the `original_size` of the image saved while loading the image or the default one(ie. (224, 224))
"""
function generate_image(x; resize_original = false, custom_size=Nothing)
  x = reshape(x, size(x)[1:3]...)/255 |> cpu
  x = x .+ im_mean
  x = permutedims(x, [3,2,1])
  x .-= minimum(x)
  x ./= maximum(x)
  if resize_original
    if custom_size!=Nothing
        return imresize(colorview(RGB, x), custom_size)
    else
        global original_size
        return imresize(colorview(RGB, x), original_size)
    end
  else
    return colorview(RGB, x)
  end
end

rgb2arr(a) = permutedims(Float32.(channelview(a)), [3,2,1])
arr2rgb(a) = colorview(RGBA{eltype(a)}, permutedims(a, [3,2,1]))

function arr2rgb_normalize(a)
    a .-= minimum(a)
    a ./= maximum(a)
    a = permutedims(a, [3,2,1])
    colorview(RGB{eltype(a)},a)
end

"""
    save_image(path, x)

Save the image defined by `x` to the specified `path`.
"""
save_image(path, x) = save(path, generate_image(x; resize_original = true))

"""
    apple_colormap(org_img, activation_map; map_alpha=0.4, process_img=true)

Creates a `rainbow` based colormap for the image, which is then applied onto `org_img`(original image) with different transparency to generate the combined heatmap with the `org_img` (original image). Function returns a tuple of `combined colormap` and `activation map alone based heatmap`.
"""
function apply_colormap(org_img, activation_map; map_alpha=0.4, process_img=true)
    cm = cgrad(:rainbow)
    no_trans_heat_map = RGB.(cm[activation_map])
    heat_map = copy(rgb2arr(no_trans_heat_map))
    if (process_img==true)
        org_img = rgb2arr(org_img)
    end
    x = org_img*0.7 + heat_map*0.3
    return (arr2rgb_normalize(x), heatmap)
end


text:
using Test

using Statistics: mean
using CUDA

using Oceananigans
using Oceananigans.CubedSpheres
using Oceananigans.Models.HydrostaticFreeSurfaceModels
using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField

include("data_dependencies.jl")

@testset "Cubed spheres" begin

    @testset "Conformal cubed sphere grid" begin
        @info "  Testing conformal cubed sphere grid..."

        grid = ConformalCubedSphereGrid(face_size=(10, 10, 1), z=(-1, 0))
        @test try show(grid); println(); true; catch; false; end
    end

    for arch in archs

        @info "Constructing a ConformalCubedSphereGrid from file [$(typeof(arch))]..."

        # Prototype grid and model for subsequent tests
        cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2"
        grid = ConformalCubedSphereGrid(cs32_filepath, architecture=arch, Nz=1, z=(-1, 0))

        @info "Constructing a HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..."

        model = HydrostaticFreeSurfaceModel(
                  architecture = arch,
                          grid = grid,
            momentum_advection = VectorInvariant(),
                  free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1),
                      coriolis = nothing,
                       closure = nothing,
                       tracers = :c,
                      buoyancy = nothing
        )

        @testset "Constructing a grid from file [$(typeof(arch))]" begin
            @test grid isa ConformalCubedSphereGrid
        end

        @testset "CubedSphereData and CubedSphereFields [$(typeof(arch))]" begin
            @info "Testing CubedSphereData and CubedSphereFields [$(typeof(arch))]..."
            c = model.tracers.c
            η = model.free_surface.η

            set!(c, 0)
            set!(η, 0)

            CUDA.allowscalar(true)
            @test all(all(face_c .== 0) for face_c in faces(c))
            @test all(all(face_η .== 0) for face_η in faces(η))
            CUDA.allowscalar(false)

            @test maximum(abs, c) == 0
            @test minimum(abs, c) == 0
            @test mean(c) == 0

            @test maximum(abs, η) == 0
            @test minimum(abs, η) == 0
            @test mean(η) == 0
        end

        @testset "Constructing a HydrostaticFreeSurfaceModel [$(typeof(arch))]" begin
            @test model isa HydrostaticFreeSurfaceModel
        end

        @testset "Time stepping a HydrostaticFreeSurfaceModel [$(typeof(arch))]" begin
            @info "Time-stepping HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..."
            time_step!(model, 1)
            @test try time_step!(model, 1); true; catch; false; end
        end

        @testset "KernelComputedField on ConformalCubedSphereGrid [$(typeof(arch))]" begin
            @info "Testing KernelComputedField on a ConformalCubedSphereGrid [$(typeof(arch))]..."
            ζ = VerticalVorticityField(model)

            @test ζ isa KernelComputedField

            set!(model, u = (x, y, z) -> rand())

            @test try compute!(ζ); true; catch; false; end
            @test maximum(abs, ζ) > 0 # fingers crossed
        end
    end
end


text:
macro publish(mod,name)
  quote
    using Gridap.$mod: $name; export $name
  end
end

# Reexport from LinearAlgebra (just for convenience)
using LinearAlgebra:  det, inv, tr, cross, dot, norm, ×, ⋅
export det, inv, tr, cross, dot, norm, ×, ⋅

@publish Helpers GridapType

@publish Algebra solve
@publish Algebra solve!
@publish Algebra symbolic_setup
@publish Algebra numerical_setup
@publish Algebra numerical_setup!
@publish Algebra LUSolver
@publish Algebra BackslashSolver
@publish Algebra zero_initial_guess
@publish Algebra NLSolver
@publish Algebra get_matrix
@publish Algebra get_vector
@publish Algebra SparseMatrixCSR
@publish Algebra SymSparseMatrixCSR
@publish Algebra jacobian
@publish Algebra hessian

@publish Arrays array_cache
@publish Arrays getindex!
@publish Arrays get_array
@publish Arrays lazy_map
@publish Arrays Reindex
@publish Arrays Broadcasting
@publish Arrays Operation
@publish Arrays print_op_tree
using Gridap.Arrays: ∑; export ∑

@publish TensorValues VectorValue
@publish TensorValues TensorValue
@publish TensorValues inner
@publish TensorValues outer
@publish TensorValues diagonal_tensor
@publish TensorValues num_components
using Gridap.TensorValues: ⊙; export ⊙
using Gridap.TensorValues: ⊗; export ⊗

@publish Fields gradient
@publish Fields ∇
@publish Fields ∇∇
@publish Fields integrate
@publish Fields Point
@publish Fields evaluate
@publish Fields evaluate!
@publish Fields curl
@publish Fields laplacian
@publish Fields divergence
@publish Fields Δ
@publish Fields ε
@publish Fields symmetric_gradient

@publish Integration get_coordinates
@publish Integration get_weights

@publish ReferenceFEs is_simplex
@publish ReferenceFEs is_n_cube
@publish ReferenceFEs simplexify
@publish ReferenceFEs num_dims
@publish ReferenceFEs num_cell_dims
@publish ReferenceFEs num_point_dims
@publish ReferenceFEs VERTEX
@publish ReferenceFEs SEGMENT
@publish ReferenceFEs TRI
@publish ReferenceFEs QUAD
@publish ReferenceFEs TET
@publish ReferenceFEs HEX
@publish ReferenceFEs WEDGE
@publish ReferenceFEs PYRAMID
@publish ReferenceFEs is_first_order
@publish ReferenceFEs is_Q
@publish ReferenceFEs is_P
@publish ReferenceFEs is_S
@publish ReferenceFEs VERTEX1
@publish ReferenceFEs SEG2
@publish ReferenceFEs TRI3
@publish ReferenceFEs QUAD4
@publish ReferenceFEs TET4
@publish ReferenceFEs HEX8
@publish ReferenceFEs Polytope
@publish ReferenceFEs ReferenceFE
@publish ReferenceFEs Lagrangian
@publish ReferenceFEs RaviartThomas
@publish ReferenceFEs Nedelec
@publish ReferenceFEs lagrangian
@publish ReferenceFEs raviart_thomas
@publish ReferenceFEs nedelec

@publish Geometry get_triangulation
@publish Geometry num_cells
@publish Geometry num_facets
@publish Geometry num_vertices
@publish Geometry num_edges
@publish Geometry num_faces
@publish Geometry Triangulation
@publish Geometry get_cell_coordinates
@publish Geometry get_cell_ref_coordinates
@publish Geometry get_cell_map
@publish Geometry get_cell_to_bgcell
@publish Geometry CartesianGrid
@publish Geometry CartesianDiscreteModel
@publish Geometry DiscreteModel
@publish Geometry DiscreteModelFromFile
@publish Geometry num_tags
@publish Geometry num_entities
@publish Geometry get_grid
@publish Geometry get_face_labeling
@publish Geometry add_tag!
@publish Geometry add_tag_from_tags!
@publish Geometry BoundaryTriangulation
@publish Geometry SkeletonTriangulation
@publish Geometry RestrictedTriangulation
@publish Geometry InterfaceTriangulation

@publish CellData CellQuadrature
@publish CellData Measure
@publish CellData DomainStyle
@publish CellData ReferenceDomain
@publish CellData PhysicalDomain
@publish CellData get_cell_points
@publish CellData CellField
@publish CellData CellState
@publish CellData jump
@publish CellData mean
@publish CellData update_state!
@publish CellData get_normal_vector
using Gridap.CellData: ∫; export ∫
@publish CellData get_cell_measure
@publish CellData get_physical_coordinate
@publish CellData DiracDelta

@publish FESpaces FESpace
@publish FESpaces TrialFESpace
@publish FESpaces TestFESpace
#@publish FESpaces FETerm
#@publish FESpaces FEEnergy
#@publish FESpaces AffineFETerm
#@publish FESpaces LinearFETerm
#@publish FESpaces FESource
@publish FESpaces AffineFEOperator
@publish FESpaces LinearFESolver
@publish FESpaces get_free_values
@publish FESpaces get_dirichlet_values
@publish FESpaces num_dirichlet_dofs
@publish FESpaces num_free_dofs
@publish FESpaces num_dirichlet_tags
@publish FESpaces get_cell_dof_ids
@publish FESpaces get_cell_dof_values
@publish FESpaces get_cell_shapefuns
@publish FESpaces get_cell_shapefuns_trial
@publish FESpaces FEFunction
@publish FESpaces interpolate
@publish FESpaces interpolate_everywhere
@publish FESpaces interpolate_dirichlet
@publish FESpaces assemble_vector
@publish FESpaces assemble_matrix
@publish FESpaces assemble_matrix_and_vector
@publish FESpaces FEOperator
@publish FESpaces FESolver
@publish FESpaces SparseMatrixAssembler
@publish FESpaces FiniteElements

@publish MultiField MultiFieldFESpace
@publish MultiField num_fields

@publish Visualization writevtk
@publish Visualization createvtk

export apply
function apply(args...)
  Helpers.@unreachable """\n
  Function apply has been removed and replaced by lazy_map.
  This error message will be deleted in future versions.
  """
end

export cell_measure
function cell_measure(args...)
  Helpers.@unreachable """\n
  Function cell_measure(a,b) has been removed and replaced by get_cell_measure(a).
  This error message will be deleted in future versions.
  """
end

export restrict
function restrict(args...)
  Helpers.@unreachable """\n
  Function restrict has been removed. The user does not need to explicitly
  restrict to a given Triangulation any more. The code does it undere the hood.
  This error message will be deleted in future versions.
  """
end

export FETerm
function FETerm(args...)
  Helpers.@unreachable """\n
  Function FETerm has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end

export FEEnergy
function FEEnergy(args...)
  Helpers.@unreachable """\n
  Function FEEnergy has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end

export AffineFETerm
function AffineFETerm(args...)
  Helpers.@unreachable """\n
  Function AffineFETerm has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end

export LinearFETerm
function LinearFETerm(args...)
  Helpers.@unreachable """\n
  Function LinearFETerm has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end

export FESource
function FESource(args...)
  Helpers.@unreachable """\n
  Function FESource has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end


text:
"""
AbstractControlLogicCore.jl defines the AbstractControlLogicCore type and
methods

------------------------------------------------------------------------------
COPYRIGHT/LICENSE. This file is part of the XYZ package. It is subject to
the license terms in the LICENSE file found in the top-level directory of
this distribution. No part of the XYZ package, including this file, may be
copied, modified, propagated, or distributed except according to the terms
contained in the LICENSE file.
------------------------------------------------------------------------------
"""
# --- Exports

# ------ Types

export AbstractControlLogicCore

# ------ Functions

export wait_for_input_ready
export decode_control_signal, process_control_signal!

# --- Type definitions

"""
    AbstractControlLogicCore

Supertype for control logic core types.

Interface
=========

    wait_for_input_ready(logic_core::AbstractControlLogicCore,
                         input_channels::Vector{InputChannel})

    decode_control_signal(::Type{<:AbstractControlLogicCore},
                          bytes::Vector{UInt8})

    process_control_signal!(signal, logic_core::AbstractControlLogicCore)

    get_state(logic_core::AbstractControlLogicCore)

    set_state!(logic_core::AbstractControlLogicCore, state::Dict)

Common Signals to Support
-------------------------
* Request for the output type. To support this signal, include the Node or
  OutputChannel as a field in the concrete subtype.
"""
abstract type AbstractControlLogicCore end

# --- Method definitions
#
# Note: the following method definitions are no-op place holders to provide
#       a central location for docstrings.

"""
    wait_for_input_ready(logic_core::AbstractControlLogicCore,
                         input_channels::Vector{InputChannel})

Wait for input channels to be ready for processing.
"""
wait_for_input_ready(logic_core::AbstractControlLogicCore,
                     input_channels::Vector{InputChannel}) = nothing

"""
    decode_control_signal(::Type{<:AbstractControlLogicCore},
                          bytes::Vector{UInt8})

Convert `bytes` to a signal that is understood by `process_control_signal!()`.
Throw an ArgumentError if `bytes` does not convert to a valid signal.
"""
decode_control_signal(::Type{<:AbstractControlLogicCore},
                      bytes::Vector{UInt8}) = nothing

"""
    process_control_signal!(logic_core::AbstractControlLogicCore, signal)

Process the control `signal` updating the data fields of `logic_core`, if
necesasry. Return an appropriate status response.
"""
process_control_signal!(logic_core::AbstractControlLogicCore, signal) = nothing

"""
    get_state(logic_core::AbstractControlLogicCore)

Return a Dict containing sufficient information to reconstruct the operational
state of `logic_core`.
"""
get_state(logic_core::AbstractControlLogicCore) = nothing

"""
    set_state!(logic_core::AbstractControlLogicCore, state::Dict)

Set the operational state of `logic_core` using values from `state`.
"""
set_state!(logic_core::AbstractControlLogicCore, state::Dict) = nothing


text:
using Distributions
using SMC
using KernelDensity
using PyPlot
using LaTeXStrings

include("../math.jl")
include("../aide.jl")

# where to write plots to
const PLOT_DIR = "plots"
if !Base.Filesystem.isdir(PLOT_DIR)
    Base.Filesystem.mkdir(PLOT_DIR)
end

"""
An unnormalized target density defined by a Gaussian prior and a mixture-of-Gaussian likelihood.
"""
struct Problem
    prior_mean::Float64
    prior_std::Float64
    proportions::Array{Float64,1}
    likelihood_means::Array{Float64,1}
    likelihood_stds::Array{Float64,1}
end

function log_joint_density(problem::Problem, x::Float64)
    prior = Normal(problem.prior_mean, problem.prior_std)
    ld = logpdf(prior, x)
    log_likelihoods = Array{Float64,1}(length(problem.likelihood_means))
    for i=1:length(problem.likelihood_means)
        likelihood = Normal(problem.likelihood_means[i],
                            problem.likelihood_stds[i])
        log_likelihoods[i] = log(problem.proportions[i]) + logpdf(likelihood, x)
    end
    ld + logsumexp(log_likelihoods)
end

"""
Exact sampler for the given inference problem.
"""
struct ExactSampler <: Sampler
    problem::Problem
end

function simulate(s::ExactSampler)
    error("not implemented!")
end

function regenerate(s::ExactSampler, x)
    log_joint_density(s.problem, x)
end

"""
Sampling importance resampling scheme.

The distribution is a Gaussian distribution with mean `mean` and
standard-deviation `std`.  `SMC.forward` Samples from the proposal and returns
the importance weight with respect to the given inference problem `p`.
`SMC.backward` returns the importance weight given a latent sample.
"""
struct GaussianInitializer
    p::Problem
    mean::Float64
    std::Float64
end

function SMC.forward(init::GaussianInitializer)
    dist = Normal(init.mean, init.std)
    x = rand(dist)
    (x, log_joint_density(init.p, x) - logpdf(dist, x))
end

function SMC.backward(init::GaussianInitializer, x::Float64)
    dist = Normal(init.mean, init.std)
    log_joint_density(init.p, x) - logpdf(dist, x)
end

function make_sir_scheme(problem::Problem,
                   init_mean::Float64,
                   init_std::Float64,
                   num_particles::Int)
    initializer = GaussianInitializer(problem, init_mean, init_std)
    SMCScheme(initializer, [], num_particles)
end

function make_offset_proposal_scheme(problem::Problem, num_particles::Int)
    make_sir_scheme(problem, 2.0, 0.5, num_particles)
end

function make_broad_proposal_scheme(problem::Problem, num_particles::Int)
    make_sir_scheme(problem, 0.0, 1.0, num_particles)
end

function make_sampler(sir_scheme::SMCScheme, problem::Problem)
    SMCSampler(sir_scheme, (x) -> log_joint_density(problem, x))
end

const OFFSET_PROPOSAL = "offset"
const BROAD_PROPOSAL = "broad"

"""
Returns a dictionary mapping key to a vector of tuples (sample, log marginal
likelihood estimate)
"""
function run_sir(problem::Problem, num_particles_list::Vector{Int}, num_replicates::Int)
    sir_outputs = Dict()
    for n in num_particles_list
        key = (OFFSET_PROPOSAL, n)
        scheme = make_offset_proposal_scheme(problem, n)
        sir_outputs[key] = map((i) -> smc(scheme), 1:num_replicates)
        key = (BROAD_PROPOSAL, n)
        scheme = make_broad_proposal_scheme(problem, n)
        sir_outputs[key] = map((i) -> smc(scheme), 1:num_replicates)
    end
    return sir_outputs
end

function get_log_marginal_likelihood_estimates(sir_outputs::Dict, proposal_name::String, num_particles::Int)
    map((output) -> output[2],  sir_outputs[(proposal_name, num_particles)])
end

function get_samples(sir_outputs::Dict, proposal_name::String, num_particles::Int)
    map((output) -> output[1],  sir_outputs[(proposal_name, num_particles)])
end

"""
Returns a dictionary mapping key to kernel density estimators
"""
function kernel_density_estimators(sir_outputs::Dict)
    kd_estimators = Dict()
    for (key, output) in sir_outputs
        samples = map((o) -> o[1], output)
        kd_estimators[key] = InterpKDE(kde(samples))
    end
    return kd_estimators
end

function plot_sir_sampling_densities(problem::Problem, kd_estimators::Dict, fname::String)

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width * 1.3, height))
    ax = plt[:gca]()

    # number of particles to plot sampling densities for
    all_num_particles_to_plot_samples = [1, 10, 100, 1000]

    # location of sampling densities on the y-axis
    yticks = [4, 3, 2, 1]

    # location of floor of plot on the y-axis
    density_y_base = 0

    # height of the target density on the y-axis
    density_actual_span = 0.7

    # x-axis limits
    data_lim = (-3, 4)

    # x-coordinates at which to evaluate densities
    density_xs = linspace(-3, 4, 100)

    # plot the unnormalized target density
    density_ys = map((x) -> exp(log_joint_density(problem, x)), density_xs)
    density_ys_span = maximum(density_ys) - minimum(density_ys)
    density_ys *= density_actual_span / density_ys_span
    plt[:plot](density_xs, density_y_base + density_ys, color="black")

    # plot the sampling densities (kernel density estimates)
    for (i, num_particles) in enumerate(all_num_particles_to_plot_samples)
        ys = pdf(kd_estimators[(OFFSET_PROPOSAL, num_particles)], density_xs)
        ys = (ys / maximum(ys)) * density_actual_span + yticks[i]
        plt[:plot](density_xs, ys, color="magenta", zorder=100, alpha=0.7)
        ys = pdf(kd_estimators[(BROAD_PROPOSAL, num_particles)], density_xs)
        ys = (ys / maximum(ys)) * density_actual_span + yticks[i]
        plt[:plot](density_xs, ys, color="blue", alpha=0.7)
    end

    # labels
    plt[:text](1.5, density_y_base + density_actual_span/9.0, "posterior\ndensity", fontsize=10)
    plt[:text](-1.7, density_y_base + density_actual_span/6.0, "L", fontsize=12)
    plt[:text](0.3, density_y_base + density_actual_span/6.0, "R", fontsize=12)
    ax[:set_xlim](data_lim)
    exponents = [0, 1, 2, 3]
    exponent_strings = map((i) -> @sprintf("%d", i), exponents)
    ax[:set_yticks](yticks)
    ax[:set_yticklabels](map((exponent_string) -> latexstring("\$10^{$exponent_string}\$"),
                             exponent_strings))
    ax[:tick_params](axis="both", which="major", labelsize=9)
    plt[:ylabel]("Number of particles\n")

    # arrow for the number of particles on y-axis
    ax[:annotate]("", xy=(-0.2, 0.0), xycoords="axes fraction", xytext=(-0.2, 1.0),
                arrowprops=Dict([("arrowstyle", "->"), ("color", "black")]))

    plt[:tight_layout]()
    plt[:savefig](fname)
end

function draw_arrow(x, y, xlen, ylen, ax)
    ax[:arrow](x, y, xlen, ylen, fc="black", ec="black",
	    transform=ax[:transAxes], length_includes_head=true,
	    head_width=0.05, head_length=0.05)
end

function plot_log_marginal_likelihoods(sir_outputs::Dict, all_num_particles::Vector{Int}, fname::String)

    function plot_single_series(proposal_name::String, label, color)
        # a vector of a vector of IID estimates
        # there is one vector of IID estimates for each number of particles
        individual_estimates = map((num_particles) -> get_log_marginal_likelihood_estimates(sir_outputs, proposal_name, num_particles), all_num_particles)
        estimates = map(mean, individual_estimates)
        stderrs = map((data::Vector{Float64}) -> std(data)/sqrt(length(data)), individual_estimates)
        plt[:plot](all_num_particles, estimates, color=color, label=label)
        lower = estimates - stderrs
        upper = estimates + stderrs
        plt[:fill_between](all_num_particles, lower, estimates, color=color, alpha=0.3)
        plt[:fill_between](all_num_particles, estimates, upper, color=color, alpha=0.3)
    end

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width, height))
    ax = plt[:gca]()

    # plot log marginal likelihood estimates
    plot_single_series(OFFSET_PROPOSAL, "Offset proposal", "magenta")
    plot_single_series(BROAD_PROPOSAL, "Broad proposal", "blue")
    gold_standard_lml = mean(get_log_marginal_likelihood_estimates(sir_outputs, "broad", maximum(all_num_particles)))
    xlim = ax[:get_xlim]()
    plt[:plot](xlim, gold_standard_lml * ones(2), color="orange", label="Gold-standard")

    # arrow and label for small penalty
    plt[:text](12, -26, "Small penalty")
    draw_arrow(0.8, 0.73, 0.0, 0.19, ax)

    # plot formatting and labels
    ax[:set_xscale]("log")
    plt[:minorticks_off]()
    ax[:tick_params](axis="both", which="major", labelsize=8)
    ax[:set_xlim](xlim)
    plt[:legend](loc="lower right", fontsize=8.7)
    plt[:xlabel]("Number of particles")
    plt[:title]("Log marginal\nlikelihood (nats)")
    plt[:tight_layout]()
    plt[:savefig](fname)
end

function aide_estimate(gold_standard_sampler::Sampler, target_sampler::Sampler, num_replicates::Int)
    lowers = map((i::Int) -> estimate_elbo(target_sampler, gold_standard_sampler, 1, 1),
                 1:num_replicates)
    uppers = map((i::Int) -> -estimate_elbo(gold_standard_sampler, target_sampler, 1, 1),
                 1:num_replicates)
    return (uppers - lowers)
end

function generate_aide_estimates(problem::Problem, all_num_particles::Vector{Int}, num_replicates::Int)
    gold_standard_scheme = make_broad_proposal_scheme(problem, maximum(all_num_particles))
    gold_standard_sampler = make_sampler(gold_standard_scheme, problem)
    aide_estimates = Dict()
    for num_particles in all_num_particles
        offset_proposal_scheme = make_offset_proposal_scheme(problem, num_particles)
        offset_proposal_sampler = make_sampler(offset_proposal_scheme, problem)
        broad_proposal_scheme = make_broad_proposal_scheme(problem, num_particles)
        broad_proposal_sampler = make_sampler(broad_proposal_scheme, problem)
        aide_estimates[(OFFSET_PROPOSAL, num_particles)] = aide_estimate(gold_standard_sampler, offset_proposal_sampler, num_replicates)
        aide_estimates[(BROAD_PROPOSAL, num_particles)] = aide_estimate(gold_standard_sampler, broad_proposal_sampler, num_replicates)
    end
    return aide_estimates
end

function plot_aide_estimates(aide_estimates::Dict, all_num_particles::Vector{Int}, fname::String)

    function plot_single_series(proposal_name::String, label, color)
        # a vector of a vector of IID estimates
        # there is one vector of IID estimates for each number of particles
        individual_estimates = map((num_particles) -> aide_estimates[(proposal_name, num_particles)],
                                   all_num_particles)
        estimates = map(mean, individual_estimates)
        stderrs = map((data::Vector{Float64}) -> std(data)/sqrt(length(data)), individual_estimates)
        plt[:plot](all_num_particles, estimates, color=color, label=label)
        lower = estimates - stderrs
        upper = estimates + stderrs
        plt[:fill_between](all_num_particles, lower, estimates, color=color, alpha=0.3)
        plt[:fill_between](all_num_particles, estimates, upper, color=color, alpha=0.3)
    end

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width, height))
    ax = plt[:gca]()

    # plot log marginal likelihood estimates
    plot_single_series(OFFSET_PROPOSAL, "Offset proposal", "magenta")
    plot_single_series(BROAD_PROPOSAL, "Broad proposal", "blue")

    # arrow and label for large penalty
    plt[:text](10, 29, "Large penalty for\nmissing mode")
    draw_arrow(0.8, 0.36, 0.0, -0.22, ax)

    # plot formatting and labels
    ymax = ax[:get_ylim]()[2]
    ax[:set_ylim](0, ymax)
    ax[:set_xscale]("log")
    plt[:minorticks_off]()
    ax[:tick_params](axis="both", which="major", labelsize=8)
    plt[:legend](loc="upper right", fontsize=8.7)
    plt[:xlabel]("Number of particles")
    plt[:title]("AIDE estimate\n(nats)")
    plt[:tight_layout]()
    plt[:savefig](fname)
end

# generate plots
srand(1)
problem = Problem(0.0, 10.0, [0.7, 0.3], [-1, 1], [0.1, 0.1])
exponents = linspace(0, 3, 10)
num_particles_list = Vector{Int}(floor.(10.^exponents))

# generate subfigure (a)
println("Generating density plot...")
sir_outputs_for_kde = run_sir(problem, num_particles_list, 10000)
kd_estimators = kernel_density_estimators(sir_outputs_for_kde)
plot_sir_sampling_densities(problem, kd_estimators, "$PLOT_DIR/bimodal_densities.pdf")

# generate subfigure (b)
println("Generating AIDE estimate plot...")
aide_estimates = generate_aide_estimates(problem, num_particles_list, 100)
plot_aide_estimates(aide_estimates, num_particles_list, "$PLOT_DIR/bimodal_symkl.pdf")

# generate subfigure (c)
println("Generating log marginal likelihood plot...")
sir_outputs_for_lml = run_sir(problem, num_particles_list, 100)
plot_log_marginal_likelihoods(sir_outputs_for_lml, num_particles_list, "$PLOT_DIR/bimodal_lml.pdf")


text:
using GR

const hello = Dict(
    "Chinese" => "你好世界",
    "Dutch" => "Hallo wereld",
    "English" => "Hello world",
    "French" => "Bonjour monde",
    "German" => "Hallo Welt",
    "Greek" => "γειά σου κόσμος",
    "Italian" => "Ciao mondo",
    "Japanese" => "こんにちは世界",
    "Korean" => "여보세요 세계",
    "Portuguese" => "Olá mundo",
    "Russian" => "Здравствуй, мир",
    "Spanish" => "Hola mundo"
)

function say_hello()
    y = 0.9
    for (lang, trans) in hello
        text(0.1, y, lang)
        text(0.4, y, trans)
        y -= 0.072
    end
    updatews()
end

say_hello()


text:
#=
Simulated Annealing placement.
=#

################################################################################
# Type to keep track of previous moves.
################################################################################
abstract type AbstractCookie end
mutable struct MoveCookie{T} <: AbstractCookie
    cost_of_move::Float64
    index_of_moved_node::Int64
    move_was_swap::Bool
    index_of_other_node::Int64
    old_location::T
end

function MoveCookie(sa_struct::SAStruct)
    return MoveCookie(0.0, 0, false, 0, zero(eltype(sa_struct)))
end

################################################################################
# Types that control various parameters of the placement.
################################################################################

##########
# SAWarm #
##########
"""
Abstract type controlling the warming cycle of Simulated Annealing placement.

API
---
* [`warm!`](@ref)

Implementations
---------------
* [`DefaultSAWarm`](@ref)
"""
abstract type SAWarm end

#-----#
# API #
#-----#
"""
    warm!(warmer :: SAWarm, state :: SAState)

Increase `state.temperature` according to `warmer`. If `warmer` decides that
the system is done warming up, it must set `state.warming = false`.

Method List
-----------
$(METHODLIST)
"""
function warm! end

#----------------#
# Implementation #
#----------------#

"""
Default imeplementation of [`SAWarm`](@ref).

On each invocation, will multiply the temperature of
the anneal by `multiplier`. When the acceptance ratio rises above `ratio`,
warming routine will end.

To prevent unbounded warming, the `ratio` field is multiplied by the `decay`
field on each invocation.
"""
mutable struct DefaultSAWarm <: SAWarm
    ratio::Float64
    multiplier::Float64
    decay::Float64
end

@inline function warm!(w::DefaultSAWarm, state::SAState)
    # Compute acceptance ratio from the state
    acceptance_ratio = state.recent_accepted_moves / state.recent_successful_moves

    # Update temperature
    state.temperature *= w.multiplier
    # Check if acceptance ratio has achieved the desired ratio
    if acceptance_ratio > w.ratio
        state.warming = false
    else
        # Decay acceptance ratio
        w.ratio *= w.decay
    end

    return nothing
end

##########
# SACool #
##########

"""
Cooling schedule for Simulated Annealing placement.

API
---
* [`cool!`](@ref)

Implementations
---------------
* [`DefaultSACool`](@ref)

"""
abstract type SACool end

#-----#
# API #
#-----#

"""
    cool!(cooler :: SACool, state :: SAState)

When called, may decrease `state.temperature`.

Method List
-----------
$(METHODLIST)
"""
function cool! end

#----------------#
# Implementation #
#----------------#

"""
Default cooling schedule. Each invocation, will multiply the temperature
of the anneal by the `alpha` parameter.
"""
struct DefaultSACool <: SACool
    alpha::Float64
end

@inline cool!(c::DefaultSACool, state::SAState) = (state.temperature *= c.alpha)

####################
# Distance Limiter #
####################

"""
Type that controls the length contraction mechanism of Simulated Annealing
placement.

API
---
* [`limit!`](@ref)

Implmementations
----------------
* [`DefaultSALimit`](@ref)

"""
abstract type SALimit end

#-----#
# API #
#-----#
"""
    limit!(limiter :: SALimit, state :: SAState)

Mutate `state.distance_limit` when called.
"""
function limit! end

#----------------#
# Implementation #
#----------------#

"""
Default distance limit updater. Will adjust the distance limit so approximate
`ratio` of moves are accepted. Will not set `state.limit` lower than `minimum`.
"""
struct DefaultSALimit <: SALimit
    ratio::Float64
    # Minimum allows distance limit.
    minimum::Int
end
DefaultSALimit(ratio::Float64) = DefaultSALimit(ratio, 1)

function limit!(l::DefaultSALimit, state::SAState)
    # Compute acceptance ratio from state
    acceptance_ratio = state.recent_accepted_moves / state.recent_successful_moves
    # Update distance limit
    temporary = (1 - l.ratio + acceptance_ratio)
    state.distance_limit = clamp(
        state.distance_limit * temporary, l.minimum, state.max_distance_limit
    )
    return nothing
end

##########
# SADone #
##########

"""
Control when to terminate Simulated Annealing placement.

API
---
* [`sa_done`](@ref)

Implementations
---------------
* [`DefaultSADone`]

"""
abstract type SADone end

#-----#
# API #
#-----#

"""
    sa_done(doner :: SADone, state :: SAState)

Return `true` to finish placement.
"""
function sa_done end

#---------------#
# Implmentation #
#---------------#

"""
Default end detection. Will return `true` when objective deviation is
less than `atol`.
"""
struct DefaultSADone <: SADone
    atol::Float64
end

@inline function sa_done(d::DefaultSADone, state::SAState)
    return !state.warming && (state.deviation < d.atol)
end

################################################################################
# Placement Routine!
################################################################################

function place!(
    sa_struct::SAStruct{A,U,D};
    # Number of moves before doing a parameter update.
    move_attempts = 20000,
    initial_temperature = 1.0,
    supplied_state = nothing,
    movegen::MoveGenerator = CachedMoveGenerator(sa_struct),
    # Parameters for high-level control
    warmer::SAWarm = DefaultSAWarm(0.96, 2.0, 0.97),
    cooler::SACool = DefaultSACool(0.999),
    doner::SADone = DefaultSADone(10.0^-5),
    limiter::SALimit = DefaultSALimit(0.44, 2),
    kwargs...,
) where {A,U,D}
    @info "Running Simulated Annealing Placement."
    # Set the random number generator for repeatable results.

    # Unpack SA
    num_nodes = length(sa_struct.nodes)
    num_channels = length(sa_struct.channels)

    # Get the largest address
    max_addresses = size(sa_struct.pathtable)
    largest_address = maximum(max_addresses)

    initial_move_limit = distancelimit(movegen, sa_struct)

    # Initialize structure to help during placement.
    cost = map_cost(sa_struct)
    cookie = MoveCookie(sa_struct)

    # Main Simulated Annealing Loop
    loop = true

    # Initialize the main state variable. State variable's timer begins when
    # the structure is created.
    if supplied_state === nothing
        state = SAState(initial_temperature, Float64(initial_move_limit), cost)
    else
        state = supplied_state
        reset!(state)
    end

    # Initialize move generator
    initialize!(movegen, sa_struct, state.distance_limit)

    # Print out the header for the statistic columns.
    show_stats(state, true)
    first_display = true
    while loop
        ##################
        # Pre-loop setup #
        ##################

        # Invert temperature to perform floating point multiplication rather
        # than division. Set local counters for this iteration.
        one_over_T = 1 / state.temperature
        accepted_moves = 0
        successful_moves = 0
        objective = state.objective
        sum_cost_difference = zero(typeof(cost))

        ##############
        # Inner Loop #
        ##############
        while successful_moves < move_attempts
            # Try to generate a move. If it failed, try again.
            #@inbounds success = generate_move!(sa_struct, movegen, cookie)
            success = generate_move!(sa_struct, movegen, cookie)
            if !success
                continue
            end

            successful_moves += 1
            # Get the cost of the move
            cost_of_move = cookie.cost_of_move
            if cost_of_move <= zero(typeof(cost_of_move)) ||
               rand() < exp(-cost_of_move * one_over_T)
                accepted_moves += 1

                if cost_of_move > 0
                    sum_cost_difference += cost_of_move
                end

                objective += cost_of_move
            else
                undo_move(sa_struct, cookie)
            end
        end

        ###########################
        # Post iteration routines #
        ###########################
        # Update cost for numerical stability reasons
        state.objective = map_cost(sa_struct)
        # Sanity Check

        if !isapprox(objective, state.objective)
            @warn """
            Objective mismatch.
            actual: $(state.objective).
            calculated: $objective.
            """
        end

        # Update some statistics in the state variable
        state.recent_move_attempts = move_attempts
        state.recent_successful_moves = successful_moves
        state.recent_accepted_moves = accepted_moves
        # Quick check to avoid NaN's showing up in the case of zero accepted
        # moves.
        if iszero(accepted_moves)
            state.recent_deviation = 0
        else
            state.recent_deviation = sum_cost_difference / accepted_moves
        end
        state.aux_cost = aux_cost(sa_struct)

        # Adjust temperature
        state.warming ? warm!(warmer, state) : cool!(cooler, state)
        # Adjust distance limit - only if we're not warming up.
        state.warming || limit!(limiter, state)

        # State updates - check if we potentially need to do some recomputation
        # for the move generator
        update_movegen = update!(state)

        if update_movegen
            update!(movegen, sa_struct, state.distance_limit_int)
        end

        # Exit Condition
        loop = !sa_done(doner, state)
    end
    # Show final statistics
    show_stats(state)

    return state
end

################################################################################
# Movement related functions
################################################################################
@propagate_inbounds function move_with_undo(
    sa_struct::SAStruct, cookie, index::Int64, new_location
)
    node = sa_struct.nodes[index]

    # Store the old information
    old_location = location(node)
    cookie.index_of_moved_node = index
    cookie.old_location = old_location
    # Check to see if there is a node already mapped to this location
    occupying_node_index = sa_struct.grid[new_location]
    if occupying_node_index == 0
        cookie.move_was_swap = false

        base_cost = node_cost(sa_struct, index)
        move(sa_struct, index, new_location)
        moved_cost = node_cost(sa_struct, index)

        cookie.cost_of_move = moved_cost - base_cost
    else
        occupying_node = sa_struct.nodes[occupying_node_index]

        # Check if the node at the new location can be moved to the old location
        isvalid(sa_struct.maptable, getclass(occupying_node), old_location) || return false
        cookie.move_was_swap = true
        # Save the index of the occupying node
        cookie.index_of_other_node = occupying_node_index
        # Compute the cost before the move
        base_cost = node_pair_cost(sa_struct, index, occupying_node_index)
        # Swap nodes
        swap(sa_struct, index, occupying_node_index)
        # Get cost after move and store
        moved_cost = node_pair_cost(sa_struct, index, occupying_node_index)
        cookie.cost_of_move = moved_cost - base_cost
    end
    return true
end

@propagate_inbounds function undo_move(sa_struct::SAStruct, cookie)
    # If the last move was a swap, need a swap in order to undo it.
    if cookie.move_was_swap
        swap(sa_struct, cookie.index_of_moved_node, cookie.index_of_other_node)
        # Otherwise, a simple move is just fine.
    else
        move(sa_struct, cookie.index_of_moved_node, cookie.old_location)
    end
    return nothing
end


text:
#--------------------General Utilities---------------------

function zoom_image(x, scale_x, scale_y)
  img = generate_image(x)
  inv_scalex = 1 / scale_x
  inv_scaley = 1 / scale_y
  local s = size(img)
  local r1 = clamp(Int(ceil(s[1] * (1 - inv_scalex) / 2)), 1, s[1]) : clamp(Int(ceil(s[1] * (1 + inv_scalex) / 2)), 1, s[1])
  local r2 = clamp(Int(ceil(s[2] * (1 - inv_scaley) / 2)), 1, s[2]) : clamp(Int(ceil(s[2] * (1 + inv_scaley) / 2)), 1, s[2])
  img = imresize(img[r1, r2], s)
  image_to_arr(img)
end

function load_model(layer, m = VGG19)
  model = m()
  global model = Chain(model.layers[1:layer]...) |> gpu
end

function argmax(A, dims)
  z = findmax(A, dims)[2] .% size(A, dims)
  z[z.==0] .= size(A,dims)
  z
end

#---------------------Core Functions-----------------------

function make_step(img, iterations, η, solo_call = false; path = "")
  if(solo_call && path == "")
    error("Image Save Path must be specified for solo calls")
  end
  input = param(img)
  for i in 1:iterations
    out = model(input)
    Flux.back!(out, out.data)
    input.data .= input.data + η * input.grad / mean(abs.(input.grad))
    info("$i iterations complete")
  end
  if(solo_call)
    save_image(path, input.data)
  end
  return input.data
end

function deepdream(base_img, iterations, η, octave_scale, num_octaves, path, guide = 1.0)
  octaves = [copy(base_img)]
  for i in 1:(num_octaves-1)
    push!(octaves, zoom_image(octaves[end], octave_scale, octave_scale))
  end
  detail = zeros(octaves[end])
  out = base_img
  for (i, oct) in enumerate(octaves[length(octaves):-1:1])
    info("OCTAVE NUMBER = $i")
    w, h = size(oct)[1:2]
    if(i > 1)
      w1, h1 = size(detail)[1:2]
      detail = zoom_image(detail, w1 / w, h1 / h)
    end
    input_oct = (oct + detail) |> gpu
    if(guide != 1.0)
      out = guided_step(input_oct, guide, iterations, η)
    else
      out = make_step(input_oct, iterations, η)
    end
    detail = out - oct
  end
  save_image(path, out)
end


text:
"""
    EpisodeTurnBuffer{names, types, Tbs} <: AbstractTurnBuffer{names, types}
    EpisodeTurnBuffer{names, types}() where {names, types}

Using a Vector to store each element specified by `names` and `types`.

See also: [`EpisodeSARDBuffer`](@ref), [`EpisodeSARDSBuffer`](@ref), [`EpisodeSARDSABuffer`](@ref)
"""
struct EpisodeTurnBuffer{names, types, Tbs} <: AbstractTurnBuffer{names, types}
    buffers::Tbs
    function EpisodeTurnBuffer{names, types}() where {names, types}
        length(names) != length(types.parameters) && throw(DimensionMismatch("length of names and types must match!"))
        buffers = merge(NamedTuple(),
                        (names[i], Vector{types.parameters[i]}()) for i in 1:length(names))
        new{names, types, typeof(buffers)}(buffers)
    end
end

capacity(b::EpisodeTurnBuffer) = isfull(b) ? length(b) : typemax(Int)

##############################
# EpisodeSARDBuffer
##############################
"`EpisodeSARDBuffer` is just an alias for `EpisodeTurnBuffer{(:state, :action, :reward, :isdone)}`"
const EpisodeSARDBuffer = EpisodeTurnBuffer{SARD}

EpisodeSARDBuffer(;state_type::Type=Int, action_type::Type=Int) = EpisodeSARDBuffer{Tuple{state_type, action_type, Float64, Bool}}()

"only valid when buffer is empty"
function push!(b::EpisodeSARDBuffer{Tuple{Ts, Ta, Float64, Bool}}, s::Ts, a::Ta) where {Ts, Ta}
    !isempty(b) && empty!(b)
    push!(b.state, s)
    push!(b.action, a)
end

function push!(b::EpisodeSARDBuffer{Tuple{Ts, Ta, Float64, Bool}}, r::Float64, d::Bool, ns::Ts, na::Ta) where {Ts, Ta}
    push!(b.reward, r)
    push!(b.isdone, d)
    push!(b.state, ns)
    push!(b.action, na)
end

length(b::EpisodeSARDBuffer) = length(b.isdone)
isfull(b::EpisodeSARDBuffer) = length(b.isdone) > 0 && b.isdone[end]

##############################
# EpisodeSARDSBuffer
##############################
"`EpisodeSARDSBuffer` is just an alias for `EpisodeTurnBuffer{(:state, :action, :reward, :isdone, :nextstate)}`"
const EpisodeSARDSBuffer = EpisodeTurnBuffer{SARDS}

EpisodeSARDSBuffer(;state_type::Type=Int, action_type::Type=Int) = EpisodeSARDSBuffer{Tuple{state_type, action_type, Float64, Bool, state_type}}()

function push!(b::EpisodeSARDSBuffer{Tuple{Ts, Ta, Float64, Bool, Ts}}, s::Ts, a::Ta, r::Float64, d::Bool, ns::Ts) where {Ts, Ta}
    isfull(b) && empty!(b)
    push!(b.state, s)
    push!(b.action, a)
    push!(b.reward, r)
    push!(b.isdone, d)
    push!(b.nextstate, ns)
end

length(b::EpisodeSARDSBuffer) = length(b.isdone)
isfull(b::EpisodeSARDSBuffer) = length(b.isdone) > 0 && b.isdone[end]

##############################
# EpisodeSARDSABuffer
##############################
"`EpisodeSARDSABuffer` is just an alias for `EpisodeTurnBuffer{(:state, :action, :reward, :isdone, :nextstate, :nextaction)}`"
const EpisodeSARDSABuffer = EpisodeTurnBuffer{SARDSA}

EpisodeSARDSABuffer(;state_type::Type=Int, action_type::Type=Int) = EpisodeSARDSABuffer{Tuple{state_type, action_type, Float64, Bool, state_type, action_type}}()

function push!(b::EpisodeSARDSABuffer{Tuple{Ts, Ta, Float64, Bool, Ts, Ta}}, s::Ts, a::Ta, r::Float64, d::Bool, ns::Ts, na::Ta) where {Ts, Ta}
    isfull(b) && empty!(b)
    push!(b.state, s)
    push!(b.action, a)
    push!(b.reward, r)
    push!(b.isdone, d)
    push!(b.nextstate, ns)
    push!(b.nextaction, na)
end

length(b::EpisodeSARDSABuffer) = length(b.isdone)
isfull(b::EpisodeSARDSABuffer) = length(b.isdone) > 0 && b.isdone[end]


text:
"""
Type: NLPStopping

Methods: start!, stop!, update\\_and\\_start!, update\\_and\\_stop!, fill\\_in!, reinit!, status
KKT, unconstrained\\_check, unconstrained2nd\\_check, optim\\_check\\_bounded

Specialization of GenericStopping. Stopping structure for non-linear programming problems using NLPModels.

Attributes:
- pb         : an AbstractNLPModel
- state      : The information relative to the problem, see GenericState
- (opt) meta : Metadata relative to stopping criterion, see `StoppingMeta`.
- (opt) main_stp : Stopping of the main loop in case we consider a Stopping
                          of a subproblem.
                          If not a subproblem, then nothing.
- (opt) listofstates : ListofStates designed to store the history of States.
- (opt) stopping_user_struct : Contains any structure designed by the user.

`NLPStopping(:: AbstractNLPModel, :: AbstractState; meta :: AbstractStoppingMeta = StoppingMeta(), max_cntrs :: Dict = init_max_counters(), main_stp :: Union{AbstractStopping, Nothing} = nothing, list :: Union{ListofStates, Nothing} = nothing, stopping_user_struct :: Any = nothing, kwargs...)`

 Note:
- designed for `NLPAtX` State. Constructor checks that the State has the
 required entries.

 There is an additional default constructor creating a Stopping where the State is by default and the
 optimality function is the function `KKT()``.

 `NLPStopping(pb :: AbstractNLPModel; kwargs...)`

 Note: Kwargs are forwarded to the classical constructor.
 """
mutable struct NLPStopping{
  Pb, M, SRC, T, MStp, LoS
}  <: AbstractStopping{Pb, M, SRC, T, MStp, LoS}

  # problem
  pb                   :: Pb

  # Common parameters
  meta                 :: M
  stop_remote          :: SRC

  # current state of the problem
  current_state        :: T

  # Stopping of the main problem, or nothing
  main_stp             :: MStp

  # History of states
  listofstates         :: LoS

  # User-specific structure
  stopping_user_struct :: AbstractDict

end

function NLPStopping(pb             :: Pb,
                     meta           :: M,
                     stop_remote    :: SRC,
                     current_state  :: T;
                     main_stp       :: AbstractStopping = VoidStopping(),
                     list           :: AbstractListofStates = VoidListofStates(),
                     user_struct    :: AbstractDict = Dict(),
                     kwargs...
                     ) where {Pb  <: AbstractNLPModel,
                              M   <: AbstractStoppingMeta,
                              SRC <: AbstractStopRemoteControl,
                              T   <: AbstractState}

    return NLPStopping(pb, meta, stop_remote, current_state,
                       main_stp, list, user_struct)
end

function NLPStopping(pb             :: Pb,
                     meta           :: M,
                     current_state  :: T;
                     main_stp       :: AbstractStopping = VoidStopping(),
                     list           :: AbstractListofStates = VoidListofStates(),
                     user_struct    :: AbstractDict = Dict(),
                     kwargs...
                     ) where {Pb <: AbstractNLPModel,
                              M  <: AbstractStoppingMeta,
                              T  <: AbstractState}

  stop_remote = StopRemoteControl() #main_stp == VoidStopping() ? StopRemoteControl() : cheap_stop_remote_control()

  return NLPStopping(pb, meta, stop_remote, current_state,
                     main_stp, list, user_struct)
end

function NLPStopping(pb             :: Pb,
                     current_state  :: T;
                     main_stp       :: AbstractStopping = VoidStopping(),
                     list           :: AbstractListofStates = VoidListofStates(),
                     user_struct    :: AbstractDict = Dict(),
                     kwargs...
                     ) where {Pb <: AbstractNLPModel, T <: AbstractState}

  if :max_cntrs in keys(kwargs)
    mcntrs = kwargs[:max_cntrs]
  else
    mcntrs = init_max_counters()
  end

  if :optimality_check in keys(kwargs)
    oc = kwargs[:optimality_check]
  else
    oc = KKT
  end

    meta = StoppingMeta(;max_cntrs = mcntrs, optimality_check = oc, kwargs...)
    stop_remote = StopRemoteControl()

  return NLPStopping(pb, meta, stop_remote, current_state,
                     main_stp, list, user_struct)
end

function NLPStopping(pb :: AbstractNLPModel;
                     n_listofstates :: Integer = 0,
                     kwargs...)
  #Create a default NLPAtX
  nlp_at_x = NLPAtX(pb.meta.x0)

  if n_listofstates > 0 && :list ∉ keys(kwargs)
    list = ListofStates(n_listofstates, Val{typeof(nlp_at_x)}())
    return NLPStopping(pb, nlp_at_x, list = list, optimality_check = KKT; kwargs...)
  end

  return NLPStopping(pb, nlp_at_x, optimality_check = KKT; kwargs...)
end

"""
init\\_max\\_counters:
initialize the maximum number of evaluations on each of
the functions present in the NLPModels.Counters, e.g.

`init_max_counters(; allevals :: T = typemax(T), obj = allevals, grad = allevals, cons = allevals, jcon = allevals, jgrad = allevals, jac = allevals, jprod = allevals, jtprod = allevals, hess = allevals, hprod = allevals, jhprod = allevals, sum = 11 * allevals, kwargs...)`

`:neval_sum` is by default limited to `|Counters| * allevals`.
"""
function init_max_counters(; allevals :: T = typemax(Int), kwargs...) where {T <: Integer}

  entries = [Meta.parse(split("$(f)", '_')[2]) for f in fieldnames(Counters)]
  lim_fields = keys(kwargs)
  cntrs = Dict{Symbol,T}([
    (Meta.parse("neval_$(t)"), t in lim_fields ? kwargs[t] : allevals) for t in entries
  ])
  push!(cntrs,
    (:neval_sum => :sum in lim_fields ? kwargs[:sum] : typemax(T) )
  )

  return cntrs
end

function max_evals!(stp :: NLPStopping, allevals :: Integer)
  stp.meta.max_cntrs = init_max_counters(allevals = allevals)
  return stp
end

function max_evals!(stp :: NLPStopping; allevals :: T = typemax(Int), kwargs...) where {T <: Integer}
  stp.meta.max_cntrs = init_max_counters(allevals = allevals; kwargs...)
  return stp
end

"""
init\\_max\\_counters\\_NLS:
initialize the maximum number of evaluations on each of
the functions present in the `NLPModels.NLSCounters`, e.g.

`init_max_counters_NLS(; allevals = typemax(T), residual = allevals, jac_residual = allevals, jprod_residual = allevals, jtprod_residual = allevals, hess_residual = allevals, jhess_residual = allevals, hprod_residual = allevals, kwargs...)`
"""
function init_max_counters_NLS(; allevals :: T = typemax(Int), kwargs...) where {T <: Integer}

  cntrs_nlp = init_max_counters(; allevals = allevals, kwargs...)

  entries = [Meta.parse(split("$(f)", '_')[2]) for f in setdiff(fieldnames(NLSCounters),[:counters])]
  lim_fields = keys(kwargs)
  cntrs = Dict{Symbol,T}([
    (Meta.parse("neval_$(t)"), t in lim_fields ? kwargs[t] : allevals) for t in entries
  ])

  return merge(cntrs_nlp, cntrs)
end

"""
fill_in!: (NLPStopping version) a function that fill in the required values in the `NLPAtX`.

`fill_in!( :: NLPStopping, :: Union{AbstractVector, Nothing}; fx :: Union{AbstractVector, Nothing} = nothing, gx :: Union{AbstractVector, Nothing} = nothing, Hx :: Union{MatrixType, Nothing} = nothing, cx :: Union{AbstractVector, Nothing} = nothing, Jx :: Union{MatrixType, Nothing} = nothing, lambda :: Union{AbstractVector, Nothing} = nothing, mu :: Union{AbstractVector, Nothing} = nothing, matrix_info :: Bool = true, kwargs...)`
"""
function fill_in!(stp         :: NLPStopping{Pb, M, SRC, NLPAtX{S, T, HT, JT}, MStp, LoS},
                  x           :: T;
                  fx          :: Union{eltype(T), Nothing} = nothing,
                  gx          :: Union{T, Nothing}         = nothing,
                  Hx          :: Union{HT, Nothing}        = nothing,
                  cx          :: Union{T, Nothing}         = nothing,
                  Jx          :: Union{JT, Nothing}        = nothing,
                  lambda      :: Union{T, Nothing}         = nothing,
                  mu          :: Union{T, Nothing}         = nothing,
                  matrix_info :: Bool                      = true,
                  kwargs...) where {Pb, M, SRC, MStp, LoS, S, T, HT, JT}

  gfx = isnothing(fx)  ? obj(stp.pb, x)   : fx
  ggx = isnothing(gx)  ? grad(stp.pb, x)  : gx

  if isnothing(Hx) && matrix_info
    gHx = hess(stp.pb, x).data
  else
    gHx = Hx
  end

  if stp.pb.meta.ncon > 0
    gJx = isnothing(Jx) ? jac(stp.pb, x)  : Jx
    gcx = isnothing(cx) ? cons(stp.pb, x) : cx
  else
    gJx = stp.current_state.Jx
    gcx = stp.current_state.cx
  end

  #update the Lagrange multiplier if one of the 2 is asked
  if (stp.pb.meta.ncon > 0 || has_bounds(stp.pb)) && (isnothing(lambda) || isnothing(mu))
    lb, lc = _compute_mutliplier(stp.pb, x, ggx, gcx, gJx; kwargs...)
  elseif  stp.pb.meta.ncon == 0 && !has_bounds(stp.pb) && isnothing(lambda)
    lb, lc = mu, stp.current_state.lambda
  else
    lb, lc = mu, lambda
  end

  return update!(stp, x = x, fx = gfx, gx = ggx, Hx = gHx,
                             cx = gcx, Jx = gJx, mu = lb,
                             lambda = lc)
end

function fill_in!(stp         :: NLPStopping{Pb, M, SRC, OneDAtX{S,T}, MStp, LoS},
                  x           :: T;
                  fx          :: Union{T, Nothing} = nothing,
                  gx          :: Union{T, Nothing} = nothing,
                  f₀          :: Union{T, Nothing} = nothing,
                  g₀          :: Union{T, Nothing} = nothing,
                  kwargs...) where {Pb, M, SRC, S, T, MStp, LoS}

 gfx = isnothing(fx) ? obj(stp.pb, x)    : fx
 ggx = isnothing(gx) ? grad(stp.pb, x)   : gx
 gf₀ = isnothing(f₀) ? obj(stp.pb, 0.0)  : f₀
 gg₀ = isnothing(g₀) ? grad(stp.pb, 0.0) : g₀

 return update!(stp.current_state, x=x, fx = gfx, gx = ggx, f₀ = gf₀, g₀ = gg₀)
end

"""
For NLPStopping, `rcounters` set as true also reinitialize the counters.
"""
function reinit!(stp       :: NLPStopping;
                 rstate    :: Bool = false,
                 rlist     :: Bool = true,
                 rcounters :: Bool = false,
                 kwargs...)

  stp.meta.start_time  = NaN
  stp.meta.optimality0 = 1.0

  #reinitialize the boolean status
  reinit!(stp.meta)

  #reinitialize the counter of stop
  stp.meta.nb_of_stop = 0

  #reinitialize the list of states
  if rlist && (typeof(stp.listofstates) != VoidListofStates)
    #TODO: Warning we cannot change the type of ListofStates
    stp.listofstates = rstate ? VoidListofStates() : ListofStates(stp.current_state)
  end

  #reinitialize the state
  if rstate
    reinit!(stp.current_state; kwargs...)
  end

  #reinitialize the NLPModel Counters
  if rcounters && typeof(stp.pb) <: AbstractNLPModel
    NLPModels.reset!(stp.pb)
  end

  return stp
end

"""
\\_resources\\_check!: check if the optimization algorithm has exhausted the resources.
                   This is the NLP specialized version that takes into account
                   the evaluation of the functions following the sum_counters
                   structure from NLPModels.

`_resources_check!(:: NLPStopping, :: T)`

Note:
- function uses counters in `stp.pb`, and update the counters in the state.
- function is compatible with `Counters`, `NLSCounters`, and any type whose entries match the entries of `max_cntrs`.
- all the NLPModels have an attribute `counters` and a function `sum_counters(nlp)`.
"""
function _resources_check!(stp    :: NLPStopping,
                           x      :: T
                           ) where T <: Union{AbstractVector, Number}
  max_cntrs = stp.meta.max_cntrs

  if max_cntrs == Dict{Symbol,Int64}()
    return stp.meta.resources
  end

  # check all the entries in the counter
  max_f = false
  if typeof(stp.pb) <: AbstractNLSModel || (:counters in fieldnames(typeof(stp.pb)) && typeof(stp.pb.counters) == NLSCounters)
    for f in intersect(fieldnames(NLSCounters), keys(max_cntrs))
      max_f = f != :counters ? (max_f || (eval(f)(stp.pb) > max_cntrs[f])) : max_f
    end
    for f in intersect(fieldnames(Counters), keys(max_cntrs))
      max_f = max_f || (eval(f)(stp.pb) > max_cntrs[f])
    end
  elseif typeof(stp.pb) <: AbstractNLPModel || (:counters in fieldnames(typeof(stp.pb)) && typeof(stp.pb.counters) == Counters)
    for f in intersect(fieldnames(Counters), keys(max_cntrs))
      max_f = max_f || (eval(f)(stp.pb) > max_cntrs[f])
    end
  end

  # Maximum number of function and derivative(s) computation
  if :neval_sum in keys(stp.meta.max_cntrs)
    max_evals = sum_counters(stp.pb) > max_cntrs[:neval_sum]
  end

  # global user limit diagnostic
  if (max_evals || max_f) stp.meta.resources = true end

  return stp.meta.resources
end

"""
`_unbounded_problem_check!`: This is the NLP specialized version that takes into account
                   that the problem might be unbounded if the objective or the
                   constraint function are unbounded.

`_unbounded_problem_check!(:: NLPStopping, :: AbstractVector)`

Note:
- evaluate the objective function if `state.fx` for NLPAtX or `state.fx` for OneDAtX is `_init_field` and store in `state`.
- do NOT evaluate the constraint function if `state.cx` is `_init_field` and store in `state`.
- if minimize problem (i.e. nlp.meta.minimize is true) check if
`state.fx <= - meta.unbounded_threshold`,
otherwise check `state.fx ≥ meta.unbounded_threshold`.
- `state.cx` is unbounded if larger than `|meta.unbounded_threshold|`.
"""
function _unbounded_problem_check!(stp  :: NLPStopping{Pb, M, SRC, NLPAtX{S, T, HT, JT}, MStp, LoS},
                                   x    :: AbstractVector
                                  ) where {Pb, M, SRC, MStp, LoS, S, T, HT, JT}

  if isnan(stp.current_state.fx)
	  stp.current_state.fx = obj(stp.pb, x)
  end

  if stp.pb.meta.minimize
    f_too_large = stp.current_state.fx <= - stp.meta.unbounded_threshold
  else
    f_too_large = stp.current_state.fx >=   stp.meta.unbounded_threshold
  end

  c_too_large = false
  #we do not evaluate the constraint if not in the state.
  if stp.pb.meta.ncon != 0 && (stp.current_state.cx != _init_field(typeof(stp.current_state.cx)))
    c_too_large = norm(stp.current_state.cx) >= abs(stp.meta.unbounded_threshold)
  end

  if (f_too_large || c_too_large) stp.meta.unbounded_pb = true end

  return stp.meta.unbounded_pb
end

function _unbounded_problem_check!(stp  :: NLPStopping{Pb, M, SRC, OneDAtX{S,T}, MStp, LoS},
                                   x    :: Union{AbstractVector, Number}
                                  ) where {Pb, M, SRC, MStp, LoS, S, T}
  if isnan(stp.current_state.fx)
	  stp.current_state.fx = obj(stp.pb, x)
  end

  if stp.pb.meta.minimize
    f_too_large = stp.current_state.fx <= - stp.meta.unbounded_threshold
  else
    f_too_large = stp.current_state.fx >=   stp.meta.unbounded_threshold
  end

  return stp.meta.unbounded_pb
end
"""
\\_infeasibility\\_check!: This is the NLP specialized version.

Note:
  - check wether the `current_score` contains Inf.
  - check the feasibility of an optimization problem in the spirit of a convex
  indicator function.
"""
function _infeasibility_check!(stp  :: NLPStopping,
                               x    :: T) where T
#=
#- evaluate the constraint function if `state.cx` is `nothing` and store in `state`.
#- check the Inf-norm of the violation ≤ stp.meta.atol
 if stp.pb.meta.ncon != 0 #if the problems has constraints, check |c(x)|
  cx = stp.current_state.cx
  if cx == _init_field(typeof(stp.current_state.cx))
   cx = cons(stp.pb, x)
  end
  vio = max.(max.(cx - stp.pb.meta.ucon, 0.), max.(stp.pb.meta.lcon - cx, 0.))
  tol = Inf #stp.meta.atol
  stp.meta.infeasible = _inequality_check(vio, stp.meta.atol, 0.) ? true : stp.meta.infeasible
 end
 =#

  if stp.pb.meta.minimize
    vio = any(z-> z ==  Inf, stp.current_state.current_score)
    if vio stp.meta.infeasible = true end
  else
    vio = any(z-> z == -Inf, stp.current_state.current_score)
    if vio stp.meta.infeasible = true end
  end

  return stp.meta.infeasible
end

################################################################################
# Nonlinear problems admissibility functions
# Available: unconstrained_check(...), optim_check_bounded(...), KKT
################################################################################
include("nlp_admissible_functions.jl")

################################################################################
# line search admissibility functions
#
# TODO: change the ls_admissible_functions and use tol_check et tol_check_neg to
# handle the inequality instead of a max.
################################################################################
include("ls_admissible_functions.jl")

#=
"""
"""
function feasibility_optim_check(pb, state; kwargs...)
     vio = _feasibility(pb, state)
     tol = Inf #stp.meta.atol
     return _inequality_check(vio, tol, 0.)
end
=#

################################################################################
# Functions computing Lagrange multipliers of a nonlinear problem
# Available: _compute_mutliplier(...)
################################################################################
include("nlp_compute_multiplier.jl")


text:
platforms_and_sources = Dict(
    Platform("x86_64", "linux") => [
        ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.2/11.4_07062021/cudnn-11.4-linux-x64-v8.2.2.26.tgz",
                      "fbc631ce19688e87d7d2420403b20db97885b17f718f0f51d7e9fc0905d86e07")],
    Platform("powerpc64le", "linux") => [
        ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.2/11.4_07062021/cudnn-11.4-linux-ppc64le-v8.2.2.26.tgz",
                      "b11b9e515a86978dc21ab50a7d2320bfb505cbce9dffa25480225c597c682b43")],
    Platform("aarch64", "linux") => [
        ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.2/11.4_07062021/cudnn-11.4-linux-aarch64sbsa-v8.2.2.26.tgz",
                      "e240d45d79eecb2257fcb8a219324f19d8e2d6e145fbd035a38d267580d65e9a")],
    Platform("x86_64", "windows") => [
        ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.2/11.4_07062021/cudnn-11.4-windows-x64-v8.2.2.26.zip",
                      "fce18f1f515a480e33fa3df8bfd8616463ee44b101e6e149d46f29a66c1c306e")],
)

max_stars_repo_name: jameskermode/Yggdrasil

max_issues_repo_name: jameskermode/Yggdrasil

max_forks_repo_name: jameskermode/Yggdrasil


text:
using AbInitioSoftwareBase: parentdir
using AbInitioSoftwareBase.Inputs: Setter, Input
using AbInitioSoftwareBase.Commands: MpiexecConfig
using Dates: format, now
using Express: Calculation, Scf
using Express.EquationOfStateWorkflow: VcOptim
using Express.PhononWorkflow: Dfpt, RealSpaceForceConstants, PhononDispersion, VDos
# using QuantumESPRESSO.Inputs: QuantumESPRESSOInput
using QuantumESPRESSO.Inputs.PWscf:
    PWInput,
    CellParametersCard,
    AtomicPositionsCard,
    CellParametersCardSetter,
    AtomicPositionsCardSetter
using QuantumESPRESSO.Inputs.PHonon:
    PhInput, Q2rInput, MatdynInput, VerbositySetter, relayinfo
using QuantumESPRESSO.Commands: pw, ph, q2r, matdyn
using QuantumESPRESSO.Outputs.PWscf: tryparsefinal
using Setfield: @set!
using UnifiedPseudopotentialFormat  # To work with `download_potential`

import Express.PhononWorkflow: MakeInput, RunCmd, parsecell, inputtype, buildjob
import Express.Shell: distprocs

inputtype(x::Calculation) = inputtype(typeof(x))
inputtype(::Type{Scf}) = PWInput
inputtype(::Type{Dfpt}) = PhInput
inputtype(::Type{RealSpaceForceConstants}) = Q2rInput
inputtype(::Type{<:Union{PhononDispersion,VDos}}) = MatdynInput

parsecell(str) =
    tryparsefinal(AtomicPositionsCard, str), tryparsefinal(CellParametersCard, str)

(::MakeInput{Scf})(template::PWInput, args...) =
    (customizer(args...) ∘ normalizer(Scf(), template))(template)
(::MakeInput{Dfpt})(template::PhInput, previnp::PWInput) =
    normalizer(Dfpt(), previnp)(template)
(::MakeInput{RealSpaceForceConstants})(template::Q2rInput, previnp::PhInput) =
    normalizer(RealSpaceForceConstants(), previnp)(template)
(::MakeInput{T})(
    template::MatdynInput,
    a::Q2rInput,
    b::PhInput,
) where {T<:Union{PhononDispersion,VDos}} = normalizer(T(), (a, b))(template)

struct CalculationSetter <: Setter
    calc::Union{Scf,Dfpt}
end
function (::CalculationSetter)(template::PWInput)
    @set! template.control.calculation = "scf"
    return template
end

struct RelayArgumentsSetter <: Setter
    input::Union{Input,Tuple}
end
(x::RelayArgumentsSetter)(template) = relayinfo(x.input, template)
function (x::RelayArgumentsSetter)(template::MatdynInput)
    template = relayinfo(x.input[1], template)
    template = relayinfo(x.input[2], template)
    return template
end

struct DosSetter <: Setter
    dos::Bool
end
function (x::DosSetter)(template::MatdynInput)
    @set! template.input.dos = x.dos
    return template
end

struct RecoverySetter <: Setter end
function (::RecoverySetter)(template::PhInput)
    @set! template.inputph.recover = true
    return template
end

normalizer(::Scf, args...) = VerbositySetter("high") ∘ CalculationSetter(Scf())
normalizer(::Dfpt, input::PWInput) =
    RelayArgumentsSetter(input) ∘ VerbositySetter("high") ∘ RecoverySetter()
normalizer(::RealSpaceForceConstants, input::PhInput) = RelayArgumentsSetter(input)
normalizer(
    ::PhononDispersion,
    inputs::Union{Tuple{Q2rInput,PhInput},Tuple{PhInput,Q2rInput}},
) = RelayArgumentsSetter(inputs) ∘ DosSetter(false)
normalizer(::VDos, inputs::Union{Tuple{Q2rInput,PhInput},Tuple{PhInput,Q2rInput}}) =
    RelayArgumentsSetter(inputs) ∘ DosSetter(true)

struct OutdirSetter <: Setter
    timefmt::String
end
function (x::OutdirSetter)(template::PWInput)
    @set! template.control.outdir = abspath(
        joinpath(
            template.control.outdir,
            join((template.control.prefix, format(now(), x.timefmt), rand(UInt)), '_'),
        ),
    )
    mkpath(template.control.outdir)
    return template
end

customizer(
    ap::AtomicPositionsCard,
    cp::CellParametersCard,
    timefmt::AbstractString = "Y-m-d_H:M:S",
) = OutdirSetter(timefmt) ∘ CellParametersCardSetter(cp) ∘ AtomicPositionsCardSetter(ap)

(x::RunCmd{Scf})(input, output = mktemp(parentdir(input))[1]; kwargs...) =
    pw(input, output; kwargs...)
(x::RunCmd{Dfpt})(input, output = mktemp(parentdir(input))[1]; kwargs...) =
    ph(input, output; kwargs...)
(x::RunCmd{RealSpaceForceConstants})(
    input,
    output = mktemp(parentdir(input))[1];
    kwargs...,
) = q2r(input, output; kwargs...)
(x::RunCmd{<:Union{VDos,PhononDispersion}})(
    input,
    output = mktemp(parentdir(input))[1];
    kwargs...,
) = matdyn(input, output; kwargs...)


text:
module Bridges

using Reexport
using HyperGraphs

pkgmodules = ["Graphs"]

for pkgmodule in pkgmodules
    include("$pkgmodule.jl")
	eval(Meta.parse("@reexport using .$(Symbol(pkgmodule))"))
end

end # module


text:
#-----------------------------------------------------------
# RBF basis calculation
#-----------------------------------------------------------

function calculaterbfbasis!(Φ, x, gridpoints, r=1.0)

    M = length(gridpoints)

    N = length(x)

    for m in 1:M

        for n in 1:N

            @inbounds Φ[n,m] = exp(-norm(gridpoints[m] - x[n])^2 / r^2)

        end

    end

    nothing

end

function calculaterbfbasis(x, gridpoints, r=1.0)

    M = length(gridpoints)

    N = length(x)

    Φ = zeros(N, M)

    calculaterbfbasis!(Φ, x, gridpoints, r)

    return Φ

end


text:
#=
#
# Writes a .msh file from a MeshData struct
=#
export write_mesh

# Writes to the file stream using the mesh
function write_mesh_MSHv2(file, mesh::MeshData)
    # Write the format info
    println(file, "\$MeshFormat");
    println(file, "2.2 0 8");
    println(file, "\$EndMeshFormat");

    # Write the nodes
    println(file, "\$Nodes");
    println(file, string(mesh.nx));
    tmp = [0.0,0.0,0.0];
    for i=1:mesh.nx
        if config.dimension == 1
            tmp[1] = mesh.nodes[1,i];
        elseif config.dimension == 2
            tmp[1:2] = mesh.nodes[:,i];
        else
            tmp = mesh.nodes[:,i];
        end
        println(file, string(mesh.indices[i])*" "*string(tmp[1])*" "*string(tmp[2])*" "*string(tmp[3]));
    end
    println(file, "\$EndNodes");

    # Print the elements
    println(file, "\$Elements");
    println(file, string(mesh.nel));
    for i=1:mesh.nel
        line = string(i)*" "*string(mesh.etypes[i])*" 0 ";
        for j=1:mesh.nv[i]
            line = line*" "*string(mesh.elements[j,i]);
        end
        println(file, line);
    end
    println(file, "\$EndElements");
end

# Writes to the file stream using the mesh
function write_mesh_MSHv4(file, mesh::MeshData)
    # Write the format info
    println(file, "\$MeshFormat");
    println(file, "4.1 0 8");
    println(file, "\$EndMeshFormat");

    # Write the nodes
    minn = mesh.indices[1];
    maxn = minn;
    for i=2:mesh.nx
        minn = min(minn,mesh.indices[i]);
        maxn = max(maxn,mesh.indices[i]);
    end
    dim = 1;
    if mesh.etypes[1] > 1
        dim = 2;
    end
    if mesh.etypes[1] > 3
        dim = 3;
    end
    println(file, "\$Nodes");
    println(file, "1 "*string(mesh.nx)*" "*string(minn)*" "*string(maxn));
    println(file, string(dim)*" 1 0 "*string(mesh.nx));
    for i=1:mesh.nx
        println(file, string(mesh.indices[i]));
    end
    for i=1:mesh.nx
        println(file, string(mesh.nodes[i,1])*" "*string(mesh.nodes[i,2])*" "*string(mesh.nodes[i,3]))
    end
    println(file, "\$EndNodes");

    # Print the elements
    println(file, "\$Elements");
    println(file, "1 "*string(mesh.nel)*" 1 "*string(mesh.nel));
    println(file, string(dim)*" 1 "*string(mesh.etypes[1])*" "*string(mesh.nel));
    for i=1:mesh.nel
        line = string(i);
        for j=1:mesh.nv[1]
            line = line*" "*string(mesh.elements[i,j]);
        end
        println(file, line);
    end
    println(file, "\$EndElements");
end

function write_mesh(file, format, mesh::MeshData)
    # don't proceed unless the mesh contains things
    if mesh.nx == 0 || mesh.nel == 0
        println("Tried to write a MSH file with an empty mesh.");
        return;
    end
    # select your file format
    if format == MSH_V2
        write_mesh_MSHv2(file, mesh);
    elseif format == MSH_V4
        write_mesh_MSHv4(file, mesh);
    end
end


text:
using RuntimeGeneratedFunctions
const RGF = RuntimeGeneratedFunctions

# see https://github.com/SciML/RuntimeGeneratedFunctions.jl/issues/28
function closure_generator(mod::Module, expr)
    (mod != @__MODULE__) && !isdefined(mod, RGF._tagname) && RGF.init(mod)
    RuntimeGeneratedFunction(mod, mod, expr)
end

# TODO ugly initialization hack?
init(mod) = closure_generator(mod, :(x -> x))

macro metatheory_init()
    quote Metatheory.init($__module__) end
end


text:
Base.@kwdef struct Spec{T} <: AbstractElement
    analyses::Tuple=()
    pkeys::NamedTuple=NamedTuple()
    mapping::Mapping=Mapping()
    options::NamedTuple=NamedTuple()
end

Spec(ctx::AbstractContext) = Spec{Any}(mapping=Mapping(ctx))
Spec(mapping::Mapping) = Spec{Any}(mapping=mapping)
Spec(s::Spec) = s

visual(args...; kwargs...) = Spec{Any}(options=namedtuple(args...; kwargs...))
visual(T::Union{Type, Symbol}, args...; kwargs...) = Spec{T}(options=namedtuple(args...; kwargs...))

@deprecate spec(args...; kwargs...) visual(args...; kwargs...)

plottype(::Spec{T}) where {T} = T

function Base.merge(t1::Spec{T1}, t2::Spec{T2}) where {T1, T2}
    T = T2 === Any ? T1 : T2
    analyses = (t1.analyses..., t2.analyses...)
    pkeys = merge(t1.pkeys, t2.pkeys)
    mapping = merge(t1.mapping, t2.mapping)
    options = merge(t1.options, t2.options)
    return Spec{T}(analyses, pkeys, mapping, options)
end

Base.:*(a1::AbstractElement, a2::AbstractElement) = merge(Spec(a1), Spec(a2))

layers(g::AbstractElement) = AlgebraicList([Spec(g)])
layers(g::AlgebraicList) = g

Base.:*(s1::AbstractElement, s2::AlgebraicList) = layers(s1) * s2
Base.:*(s1::AlgebraicList, s2::AbstractElement) = s1 * layers(s2)

const ElementOrList = Union{AbstractElement, AlgebraicList}

Base.:+(s1::ElementOrList, s2::ElementOrList) = layers(s1) + layers(s2)

# pipeline

# Expand pairs and run the analyses
function expand(sp::Spec{T}) where {T}
    analyses, pkeys, mapping, options = sp.analyses, sp.pkeys, sp.mapping, sp.options
    @assert isempty(pkeys)
    v = [Spec{T}(mapping=val, pkeys=key, options=options) for (key, val) in pairs(mapping)]
    list = AlgebraicList(v)
    return foldl((ls, an) -> apply(an, ls), analyses, init=list)
end

global_options(f, d::AlgebraicList) = NamedTuple()

function compute(f, layer, precomputed_attributes)
    args, kwargs = split(layer.mapping.value)
    return f(args...; precomputed_attributes..., kwargs...)
end

# default fallback to apply a callable to a vector of key => value pairs
# if customized, it must return a vector of key => value pairs
function apply(f, d::AlgebraicList)
    precomputed_attributes = global_options(f, d)
    v = map(parent(d)) do layer
        analyses, pkeys, options = layer.analyses, layer.pkeys, layer.options
        T = plottype(layer)
        res = compute(f, layer, precomputed_attributes) * Spec{T}(analyses=analyses, options=options, pkeys=pkeys)
        return parent(layers(res))
    end
    return AlgebraicList(reduce(vcat, v))
end

# Expand mappings, apply analyses, compute scales, and return vector of traces
function run_pipeline(s::ElementOrList)
    nested = [parent(expand(layer)) for layer in layers(s)]
    computescales(reduce(vcat, nested))
end


text:
"""
	BoxCox{T}(λ1, λ2=0)

This struct is used to store the lambda values for boxcox transform, since the value is only search-able on train data.
"""
struct BoxCox
	λ1
	λ2
	BoxCox(λ1, λ2=0) = new(λ1, λ2)
end

function BoxCox(X::AbstractArray)
	λ1, λ2 = boxcox_gridsearch(X)
	BoxCox(λ1, λ2)
end

function (f::BoxCox)(X; inv=false)
	if inv
		inv_boxcox(X, f.λ1, f.λ2)
	else
		boxcox(X, f.λ1, f.λ2)
	end
end

"""
	boxcox(x, λ₁, λ₂ = 0)

Return the boxcox transformed `x`.
"""
function boxcox(x, λ1, λ2 = 0)
	if λ1 == 0
		log(x + λ2)
	else
		((x + λ2)^λ1 - 1) / λ1
	end
end

function inv_boxcox(x, λ1, λ2 = 0)
	y = if λ1 == 0
		exp(x)
	else
		(x * λ1 + 1)^(1/λ1)
	end
	y - λ2
end

"""
	auto_boxcox(X, λ₂ = 0; r = -5:0.01:5)

Return tuple `(λ₁, λ₂, Y, pv)` where:
- `λ₁`, `λ₂` are boxcox λ-s value
- `Y` are boxcox transformed values of `X`
- `pv` is `pvalue` of the `OneSampleADTest`

Lambda values are automatically calculated using grid search by finding the max of the `pvalue` of `OneSampleADTest`. One can pass a range to `λ₂` to grid search for `λ₂`, too.
"""
function auto_boxcox(X, λ2 = 0; r = -5:0.01:5, inv=false)
	λ1, λ2 = boxcox_gridsearch(X, r, λ2)
	if inv
		boxcox.(X, λ1, λ2)
	else
		inv_boxcox.(X, λ1, λ2)
	end
end

"""
	boxcox_gridsearch(X, l1 = range(start=-5, stop=5, step=0.01), l2 = 0)

Return (λ1, λ2) of BoxCox transformation, grid search evaluated by finding max of `pvalue(OneSampleADTest(...))`.
"""
function boxcox_gridsearch(X, l1 = range(start=-5, stop=5, step=0.01), l2 = 0)
	iter = Iterators.product(l1, l2)
	dist = Normal(0, 1)
	grid = map(iter) do (l1, l2)
		@chain begin
			boxcox.(X, l1, l2)
			standardize(ZScoreTransform, _)
			OneSampleADTest(_, dist)
			pvalue(_)
			(l1 = l1, l2 = l2, pv = _)
		end
	end
	@chain begin
		grid
		filter(x -> !isnan(x.pv), _)
		sort(_, by = x -> x.pv, rev=true)
		first(_)
	end
end

function create_boxcox_transform(df, numeric; exclude=String[], use_log1p=false)
	t = map(numeric) do name
		col = Vector(df[!, name])
		if skewness(col) ≥ 0.75
			f = if use_log1p
				(x; inv=false) -> inv ? exp1m(x) : log1p(x)
			else
				l1, l2 = boxcox_gridsearch(col)
				BoxCox(l1, l2)
			end
			name => ByRow(f) => name
		else
			nothing
		end
	end
	filter(!isnothing, t)
end


text:
function run_update_manifests(argv=ARGS)

    argset = ArgParse.ArgParseSettings()
    ArgParse.@add_arg_table! argset begin
        "rootdir"
            help="Root directory"
            arg_type = String
            default = pwd()
        "--precompile", "-c"
            help = "Precompile the projects after update"
            action = :store_true
        "--pkgs", "-p"
            help = "Specific packages to update"
            arg_type = String
            default = ""
    end

    parsed_args = ArgParse.parse_args(argv, argset)
    rootdir = parsed_args["rootdir"]
    precompile = parsed_args["precompile"]
    packages = _split_arglist(parsed_args["pkgs"])

    _walk_pkgs(;rootdir) do path, proj
        Pkg.activate(path)
        println()
        if isempty(packages)
            try;
                @info("Updating All")
                Pkg.update()
            catch err
                (err isa InterruptException) && rethrow(err)
                @error("Update Fails", err)
            end
        else
            for name in packages
                isempty(name) && continue
                try;
                    @info("Updating", name)
                    spec = PackageSpec(;name)
                    Pkg.update(spec; level=Pkg.UPLEVEL_MAJOR, mode=Pkg.PKGMODE_MANIFEST)
                catch err
                    (err isa InterruptException) && rethrow(err)
                    @error("Update Fails", err)
                end
            end
        end

        precompile && Pkg.precompile()
    end
end


text:
using EconPDEs

# PDE with 1 state variable
include("./AssetPricing/CampbellCochrane.jl")
m = CampbellCochraneModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time y, result, distance = pdesolve(m, stategrid, y0)
#   0.011112 seconds (286.45 k allocations: 12.268 MiB)

# PDE with 2 state variables
include("./AssetPricing/BansalYaron.jl")
m = BansalYaronModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.081095 seconds (485.84 k allocations: 78.910 MiB, 11.41% gc time)

# System of 4 PDEs with 1 state variable
include("./AssetPricing/GarleanuPanageas.jl")
m = GarleanuPanageasModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.047346 seconds (167.95 k allocations: 65.049 MiB, 12.90% gc time)

# System of 3 PDEs with 2 state variables
include("./AssetPricing/DiTella.jl")
m = DiTellaModel()
stategrid = initialize_stategrid(m; xn = 80, νn = 10)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.912719 seconds (2.55 M allocations: 1.303 GiB, 9.01% gc time)


text:
function juliaset(z, z0, nmax::Int64)
    for n = 1:nmax
        if abs(z) > 2 (return n-1) end
        z = z^2 + z0
    end
    return nmax
end