export cross_validation

struct CrossValidation
    abs_errors::Matrix{Float64}
    mae::Vector{Float64}
    crps_scores::Matrix{Float64}
    mean_crps::Vector{Float64}
    function CrossValidation(n::Int, steps_ahead::Int)
        abs_errors = Matrix{Float64}(undef, steps_ahead, n)
        crps_scores = Matrix{Float64}(undef, steps_ahead, n)
        mae = Vector{Float64}(undef, steps_ahead)
        mean_crps = Vector{Float64}(undef, steps_ahead)
        return new(abs_errors, mae, crps_scores, mean_crps)
    end
end

discrete_crps_indicator_function(val::T, z::T) where {T} = val < z
function crps(val::T, scenarios::Vector{T}) where {T}
    sorted_scenarios = sort(scenarios)
    m = length(scenarios)
    crps_score = zero(T)
    for i = 1:m
        crps_score +=
            (sorted_scenarios[i] - val) *
            (m * discrete_crps_indicator_function(val, sorted_scenarios[i]) - i + 0.5)
    end
    return (2 / m^2) * crps_score
end
evaluate_abs_error(y::Vector{T}, forecast::Vector{T}) where T = abs.(y - forecast)
function evaluate_crps(y::Vector{T}, scenarios::Matrix{T}) where {T}
    crps_scores = Vector{T}(undef, length(y))
    for k = 1:length(y)
        crps_scores[k] = crps(y[k], scenarios[k, :])
    end
    return crps_scores
end

"""
TODO
"""
function cross_validation(gas::Model{<:Distribution, T}, y::Vector{T}, steps_ahead::Int, start_idx::Int;
                  S::Int = 10_000,
                  initial_params::Matrix{T} = DEFAULT_INITIAL_PARAM,
                  opt_method = NelderMead(gas, DEFAULT_NUM_SEEDS)) where T
    num_mle = length(y) - start_idx - steps_ahead
    b = CrossValidation(num_mle, steps_ahead)
    for i in 1:num_mle
        println("CrossValidation: step $i of $num_mle")
        gas_to_fit = deepcopy(gas)
        y_to_fit = y[1:start_idx - 1 + i]
        y_to_verify = y[start_idx + i:start_idx - 1 + i + steps_ahead]
        fit!(gas_to_fit, y_to_fit; initial_params=initial_params, opt_method=opt_method)
        forec = initial_params !== DEFAULT_INITIAL_PARAM ?
                    forecast(y_to_fit, gas_to_fit, steps_ahead; S=S, initial_params=initial_params) :
                    forecast(y_to_fit, gas_to_fit, steps_ahead; S=S)
        abs_errors = evaluate_abs_error(y_to_verify, forec.observation_forecast)
        crps_scores = evaluate_crps(y_to_verify, forec.observation_scenarios)
        b.abs_errors[:, i] = abs_errors
        b.crps_scores[:, i] = crps_scores
    end
    for i in 1:steps_ahead
        b.mae[i] = mean(b.abs_errors[i, :])
        b.mean_crps[i] = mean(b.crps_scores[i, :])
    end
    return b
end


#!/usr/bin/env julia

using Gadfly
using RDatasets
iris = dataset("datasets", "iris")

p = plot(iris, x=:SepalLength, y=:SepalWidth, Geom.point);

img = SVG("iris_plot.svg", 14cm, 8cm)
draw(img, p)

# Look at this, coloring by species.
#p2 = plot(iris, x=:SepalLength, y=:SepalWidth, color=:Species, Geom.point);
#img2 = SVG("iris_plot2.svg", 14cm, 8cm)
#draw(img2, p2)

# Egads.
# ./hello.jl  42.91s

# Same time for 2 plots, though.
# ./hello.jl  49.89s

# Load iris data alone.
# ./hello.jl  15.18s
# With using Gadfly
# hello.jl  20.15s
# entire first plot.
# 48.74
# So plotting is about 28 seconds.


function RK4!(du, u, p, update_du!, Δt::T) where T<:AbstractFloat

  # initialize derivative ponderated sum
  du_sum = zeros(eltype(u),size(u))

  # increment derivative ponderated sum
  update_du!(du,u,p,0.0)
  du_sum .+= du

  update_du!(du,u.+((Δt/2.0).*du),p,0.0)
  du_sum .+= 2.0 .* du

  update_du!(du,u.+(Δt/2.0).*du,p,0.0)
  du_sum .+= 2.0 .* du

  update_du!(du,u.+Δt.*du,p,0.0)
  du_sum .+= du

  # update u with the mean of the ponderated sum over Δt
  return (Δt/6).*du_sum .+ u
end

function RK4(u, p, update_du, Δt::T) where T<:AbstractFloat

  # increment derivative ponderated sum
  du1 = update_du(u,p,0.0)
  du2 = update_du(u.+((Δt/2.0).*du1),p,0.0)
  du3 = update_du(u.+((Δt/2.0).*du2),p,0.0)
  du4 = update_du(u.+Δt.*du3,p,0.0)
  du_sum = du1 + 2*du2 + 2*du3 + du4

  # update u with the mean of the ponderated sum over Δt
  return (Δt/6)*du_sum + u
end

function timestep_RK4_adaptative!(du, u, p, update_du!, Δt::T, e₀::T) where T<: AbstractFloat
  u1 = RK4!(du, u, p, update_du!, Δt)
  #println("RK4 u1",u1)
  uint = RK4!(du, u, p, update_du!, Δt/2)
  #println("RK4 uint",uint)
  u2 = RK4!(du, uint, p, update_du!, Δt/2)
  #println("RK4 u2",u2)

  e = abs(u2[1]-u1[1]) # error calculated on damage only
  if e < e₀
      return u1, Δt, min(Δt*abs(e₀/e),Δt*2)
  else
      timestep_RK4_adaptative(du, u, p, update_du!, float(Δt*abs(e₀/e)^2), e₀)
  end
end

function timestep_RK4_adaptative_2points(u, p, update_du, Δt::T, e₀::T) where T<: AbstractFloat
  u1 = RK4(u, p, update_du, Δt)
  #println("RK4 u1",u1)
  uint = RK4(u, p, update_du, Δt/2)
  #println("RK4 uint",uint)
  u2 = RK4(uint, p, update_du, Δt/2)
  #println("RK4 u2",u2)
  Ṡ1, σ̇ⁱξξ1, σ̇ⁱoop1, ϵ̇ⁱηη1, Ḋᵒ1, Ḋⁱ1 = u1
  Ṡ2, σ̇ⁱξξ2, σ̇ⁱoop2, ϵ̇ⁱηη2, Ḋᵒ2, Ḋⁱ2 = u2

  eD = max(abs((Ḋⁱ2-Ḋⁱ1)/Ḋⁱ2),abs((Ḋᵒ2-Ḋᵒ1)/Ḋᵒ2))
  eS = abs((Ṡ2-Ṡ1)/Ṡ2)
  eσ = max(abs((σ̇ⁱξξ2-σ̇ⁱξξ1)/σ̇ⁱξξ2),abs((σ̇ⁱoop2-σ̇ⁱoop1)/σ̇ⁱoop2))
  eϵ = abs((ϵ̇ⁱηη2-ϵ̇ⁱηη1)/ϵ̇ⁱηη2)
  e_vec = Vec(eD,eS,eσ,eϵ)
  e_normalized = e_vec./e₀
  e_norm_max, id_e_norm_max = findmax(e_normalized)

  if e_norm_max < e₀
      return u2, Δt, min(Δt*abs(e₀/e_norm_max),Δt*2)
  else
      timestep_RK4_adaptative_2points(u, p, update_du, float(Δt*abs(e₀/e_norm_max)^2), e₀)
  end
end

function compute_RK4_adaptative(u0,p,update_du,tspan,Δt_ini,e₀)
  u_vec = Vector{SVector{6,Float64}}()
  t_vec = Float64[tspan[1]]
  #u_next = similar(u0)
  cohesion_loss_flag = false
  push!(u_vec,u0)
  next_Δt = Δt_ini
  while t_vec[end] < tspan[2]
      (t_vec[end]+next_Δt > tspan[2]) && (next_Δt = tspan[2] - t_vec[end]) # prevents time overshoot
      u_next, used_Δt, next_Δt = timestep_RK4_adaptative_2points(u_vec[end], p, update_du,next_Δt,e₀)
      push!(u_vec,u_next)
      push!(t_vec,t_vec[end]+used_Δt)

      # break if D >= 1
      if any(@view(u_next[5:6]) .>= 0.98)
          cohesion_loss_flag = true
          break
      end
  end
  return t_vec, u_vec, cohesion_loss_flag
end


function cnsactivations!(system::CVSystem,n::Int64)
    # activation based on Gompertz's sigmoidal function
    push!(system.cns.np,system.cns.a*exp(
        -system.cns.b*exp(-system.cns.c*
        (system.cns.Paverage[system.solverparams.numbeats]/
        system.cns.Ptarget - 1))));
    push!(system.cns.ns,system.cns.a*exp(
        -system.cns.b*exp(system.cns.c*
        (system.cns.Paverage[system.solverparams.numbeats]/
        system.cns.Ptarget - 1))));
    # # hard-code activations
    # nst = 0.25;
    # r = 1/system.cns.c*log(log(nst^(-1/system.cns.b)))+1;
    # ns = exp(-system.cns.b*exp(system.cns.c*(r-1)));
    # np = exp(-system.cns.b*exp(-system.cns.c*(r-1)));
    # push!(system.cns.ns,ns);
    # push!(system.cns.np,np);
end


using AdvancedTopics
using DataFrames
using JLD2

include("util.jl")

function perform_tests(training_samples::T₁,
                       training_targets::T₂,
                       testing_samples::T₃,
                       testing_targets::T₄,
                       n_neurons::Int;
                       boundary_offset::Int = 1,
                       batch_size::Int = 100) where {T₁ <: AbstractMatrix,
                                                     T₂ <: AbstractVector,
                                                     T₃ <: AbstractMatrix,
                                                     T₄ <: AbstractVector}
    X₀ = training_samples
    t₀ = training_targets
    X₁ = testing_samples
    t₁ = testing_targets
    L = n_neurons

    W, fs = train_kde_comparators(Float64, X₀, t₀, L, boundary_offset = boundary_offset)

    elm_sig = ELM{Float64}(X₀, t₀, W, [Sigmoid() for f in 1:length(fs)], batch_size = batch_size)
    elm_kde = ELM{Float64}(X₀, t₀, W, fs, batch_size = batch_size)

    y₀_sig = predict(elm_sig, X₀)
    y₀_kde = predict(elm_kde, X₀)
    y₁_sig = predict(elm_sig, X₁)
    y₁_kde = predict(elm_kde, X₁)

    training_error_sig = calculate_error(y₀_sig, t₀)
    training_error_kde = calculate_error(y₀_kde, t₀)
    testing_error_sig = calculate_error(y₁_sig, t₁)
    testing_error_kde = calculate_error(y₁_kde, t₁)

    (training_error_sig, training_error_kde,
     testing_error_sig, testing_error_kde)
end

function run_experiments()
    X₀, t₀ = mnist_training_data()
    X₁, t₁ = mnist_testing_data()
    n_neurons = [2500]
    trials = 1:30

    for L in n_neurons
        results = DataFrame(:trial => collect(trials),
                            :sigmoid_training_error => zeros(Float64, length(trials)),
                            :kde_training_error => zeros(Float64, length(trials)),
                            :sigmoid_testing_error => zeros(Float64, length(trials)),
                            :kde_testing_error => zeros(Float64, length(trials)))

        for trial in trials
            (training_error_sig, training_error_kde,
             testing_error_sig, testing_error_kde) = perform_tests(X₀, t₀, X₁, t₁, L, batch_size = 100)

            results[trial,:sigmoid_training_error] = training_error_sig
            results[trial,:kde_training_error] = training_error_kde
            results[trial,:sigmoid_testing_error] = testing_error_sig
            results[trial,:kde_testing_error] = testing_error_kde

            println("$(L) neurons, trial $(trial)")
        end

        @save "results/mnist_$(L)_neurons.jld2" results
    end
end

run_experiments()

@load "results/mnist_2500_neurons.jld2" results
@show results[:,:sigmoid_training_error]
@show results[:,:kde_training_error]


using Scientific_Programming_in_Julia
using Base: tail

import Scientific_Programming_in_Julia.ReverseDiff

using Base: tail
struct Dense{M<:TrackedMatrix, V<:TrackedVector, F}
    W::M
    b::V
    f::F
end

Dense(W::Matrix, b::Vector, f) = Dense(track(W),track(b),f)
Dense(in,out,f=identity) = Dense(rand(out,in) .-0.5, rand(out) .-0.5, f)
(m::Dense)(x) = m.f(m.W*x + m.b)

struct Chain{T<:Tuple}
    layers::T
end
Chain(ls...) = Chain(ls)

(m::Chain)(x) = applychain(m.layers, x)
applychain(ls::Tuple{}, x) = x
applychain(ls::Tuple, x) = applychain(tail(ls), first(ls)(x))

params(m::Dense) = (W=m.W, b=m.b)
params(m::Chain) = params(m.layers)
params(ls::Tuple) = (params(first(ls))..., params(tail(ls))...)
params(ls::Tuple{}) = ls

function ReverseDiff.reset!(m::Chain)
    map(ReverseDiff.reset!, m.layers)
    nothing
end

function ReverseDiff.reset!(m::Dense)
    ReverseDiff.reset!(m.W)
    ReverseDiff.reset!(m.b)
    nothing
end

function train_step!(loss, network, xs, ys; λ=0.01)
    l = loss(network, xs, ys)
    l.grad = 1.0
    for w in params(network)
        w.data .-= λ .* accum!(w)
    end
    l
end

# task
f(x,y) = y^2 + sin(x)
f(xs) = [f(xs[1],xs[2])]

# data
xs = map(1:30) do i
    x = rand(-4:0.1:4)
    y = rand(-2:0.1:2)
    [x,y]
end
ys = track.(f.(xs))
xs = track.(xs)

hdim = 15
network = Chain(
    Dense(2,hdim,σ),
    Dense(hdim,hdim,σ),
    Dense(hdim,1)
)
function loss(network,xs,ys)
    errs = hcat([abs2(network(x)-y) for (x,y) in zip(xs,ys)]...)
    sum(errs)
end

using Plots
color_scheme = cgrad(:RdYlBu_5, rev=true)

forward(network,x,y) = network(track([x,y])).data[1]

# training
anim = @animate for i in 1:2000
    ReverseDiff.reset!.(xs)
    ReverseDiff.reset!.(ys)
    ReverseDiff.reset!(network)
    l = train_step!(loss, network, xs, ys, λ=0.003)
    if mod1(i,50) == 1
        @info i l
        p1 = contour(-4:0.3:4, -2:0.3:2, f, fill=true, c=color_scheme, xlabel="x", ylabel="y", title="Truth")
        p2 = contour(-4:0.3:4, -2:0.3:2, (x,y)->forward(network,x,y), fill=true, c=color_scheme, xlabel="x",title="Iteration: $i")
        p = plot(p1,p2,size=(1200,400)) |> display
    end
end every 50

gif(anim, "anim.gif", fps=15)
error()

for i in 1:20000
    ReverseDiff.reset!.(xs)
    ReverseDiff.reset!.(ys)
    ReverseDiff.reset!(network)
    l = train_step!(loss, network, xs, ys, λ=0.002)
    if mod1(i,1000) == 1
        @info i l
        p1 = contour(-4:0.3:4, -2:0.3:2, f, fill=true, c=color_scheme, xlabel="x", ylabel="y", title="Truth")
        p2 = contour(-4:0.3:4, -2:0.3:2, (x,y)->forward(network,x,y), fill=true, c=color_scheme, xlabel="x",title="Iteration: $i")
        plot(p1,p2) |> display
    end
end


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 89ab29f2-1ad2-11eb-2aa1-a7fae56d7fd8
begin
	using Plots
	using PlutoUI
end

# ╔═╡ 32e91c48-1ad2-11eb-262e-e97ac9fa9a2d
function grid_with_boundary(k)

    n = 2^k

    # Grid:
    f = zeros(n+1, n+1)

    # Boundary Conditions:

	# Bottom
	f[1, 1:(n÷2 +1)] .= range(13,5;length= n÷2 +1)
	f[1, (n÷2+1):(3*n÷4)] .= 5
	f[1, (3*n÷4+1):(n+1)] .= range(5,13;length=n÷4+1)

    # Top
    f[n+1, 1:n+1] .= 21

    # Left
    f[1:(3*n÷8+1), 1] .= range(13,40;length=3*n÷8+1)
    f[(n÷2+1):n+1, 1] .= range(40,21,length=n÷2+1)

    # Right
    f[1:(n÷2+1), n+1] .= range(13,40;length=n÷2+1)
    f[(5*n÷8+1):n+1, n+1] .= range(40,21;length=3*n÷8+1)

    # Heaters:
    f[(3*n÷8+1):(n÷2+1), 1:(n÷8+1)] .= 40
    f[(n÷2+1):(5*n÷8+1), (7*n÷8+1):n+1] .= 40

    return f

end

# ╔═╡ 79741a14-1ad2-11eb-34e1-7feadae4168e
function jacobi_step(T::AbstractMatrix)

    m, n = size(T)

    T′ = deepcopy(T)

    # iterate over interior:
    for i in 2:m-1, j in 2:n-1

		# Jacobi update
        T′[i, j] = (T[i+1, j] + T[i-1, j] + T[i, j-1] + T[i, j+1]) / 4

    end

    return T′

end

# ╔═╡ 664637ba-1ad2-11eb-0382-b99596159809
function add_heaters!(T)

    n = size(T, 1)-1

    T[(3*n÷8+1):(n÷2+1), 1:(n÷8+1)] .= 40
    T[(n÷2+1):(5*n÷8+1), (7*n÷8+1):n+1] .= 40

end

# ╔═╡ 5805d042-1ad5-11eb-0f93-110cd0c6d41e
function simulation(T, ϵ=1e-3, num_steps=200)

    results = [T]

    for i in 1:num_steps

        T′ = jacobi_step(T)
        add_heaters!(T′)

        push!(results, copy(T′))

		if maximum(abs.(T′ - T)) < ϵ
			return results
		end

        T = copy(T′)

    end

    return results
end

# ╔═╡ 9a57746c-1ad0-11eb-3b48-c1d7877ef61a
T = grid_with_boundary(5)

# ╔═╡ bf076518-1ad0-11eb-3b74-5bb405031a2b
@elapsed results = simulation(T,1e-10, 5000)

# ╔═╡ 9997a2de-1ad2-11eb-2eaa-6dcd337db44d
@bind i Slider(1:length(results), show_value=true)

# ╔═╡ ab4b6a1a-1ad2-11eb-0536-873dab9733ee
i

# ╔═╡ 5d594710-1b80-11eb-0c14-8b824654c196
showp(h) = PlutoUI.Show(MIME"image/svg+xml"(), repr(MIME"image/svg+xml"(), h))

# ╔═╡ 98fb5f6a-1c7f-11eb-0b2e-a5185e8d7649
myheatmap(T) = heatmap(T, yflip=true, ratio=1)

# ╔═╡ af8777ae-1ad2-11eb-14b4-9d557cb5efa2
showp(myheatmap(results[i]))

# ╔═╡ 113b8616-1ad3-11eb-1099-ebfe49db3815
extrema((results[end] - results[end-1]))

# ╔═╡ 449547f4-1ad3-11eb-23e4-39b0d4e8f0f7
md"""
## Better initial condition
"""

# ╔═╡ 95da5910-1ad3-11eb-1872-155aad1d8ae2
md"""
We've seen that Jacobi is a *local* method.

Putting zeros everywhere is a terrible initial condition.
"""

# ╔═╡ d0f8cda6-1ad3-11eb-0717-57ce76bfd3c8
function initial_condition(T)

	m, n = size(T)

	T′ = copy(T)

	for i in 2:m-1
		for j in 2:n-1

		# weighted average of boundaries:
		T′[i, j] = ( j * T[i, n] + (n + 1 - j) * T[i, 1] +
					(m + 1 - i) * T[1, j] + i * T[m, j]) / (m + n + 2)

		end
	end

	return T′

end

# ╔═╡ 6a7a1656-1ad4-11eb-2fd0-3fe69d6c6c4d
TT = grid_with_boundary(10);

# ╔═╡ 74e6e18a-1ad4-11eb-389e-6b95f00eac8c
TT′ = initial_condition(TT);

# ╔═╡ 923bd7c6-1ad4-11eb-35ce-55d99debaa5e
add_heaters!(TT′);

# ╔═╡ 857aebe2-1ad4-11eb-373c-bb0039279a74
heatmap(TT′, yflip=true, ratio=1)

# ╔═╡ 9877d516-1ad4-11eb-096b-75d3e8811259
@elapsed results2 = simulation(TT′, 1e-2,500)

# ╔═╡ aa80f116-1ad4-11eb-0742-a992b8278c4b
@bind i2 Slider(1:length(results2), show_value=true)

# ╔═╡ b5af2d1e-1ad4-11eb-01bb-3ffe82f8bbb0
showp(heatmap(results2[i2], yflip=true, ratio=1))

# ╔═╡ 861efa52-1c82-11eb-3e5f-9f45aadd3054
extrema((results2[end] - results2[end-1]))

# ╔═╡ 3047db8e-1ad5-11eb-0b3c-6da3fc2251b0
md"""
## Thinking Hierarchically: Coarse to fine grid

Solve the problem almost exactly on a coarse grid
"""

# ╔═╡ 8d946d20-1c7e-11eb-1d87-3d7228c4945e
coarse = initial_condition(grid_with_boundary(4));

# ╔═╡ 4dc4ec56-1ad5-11eb-0969-c7cd1728bb68
add_heaters!(coarse);

# ╔═╡ debadd2e-1ad5-11eb-0bb6-99037ade1f6f
heatmap(coarse, yflip=true, ratio=1)

# ╔═╡ 96032a3c-1ad5-11eb-242b-b7d2b6c54568
coarse_results = simulation(coarse, 1e-10, 5000);

# ╔═╡ b01fbeee-1ad5-11eb-0c27-359d42df7775
length(coarse_results)

# ╔═╡ 1a04f8ce-1ad6-11eb-3e8b-ad10bbc2a217
myheatmap(coarse_results[end])

# ╔═╡ 27b3c57c-1ad6-11eb-2546-0fe68a40bf06
begin
	fine_size = 2^10
	coarse_size = 2^4

	fine_size^2 / coarse_size^2
end

# ╔═╡ 4325b748-1ad6-11eb-20cc-d75cdbb1fbe5
md"""
So the whole Jacobi calculation on the coarse grid corresponds to <1 step on the fine grid.
"""

# ╔═╡ 4fada552-1ad6-11eb-0e25-77b981c007c2
function uplevel(T)

	n = size(T, 1)
	k = floor(Int, log2(n - 1))

	T′ = grid_with_boundary(k + 1)

	add_heaters!(T′)

	new_n = size(T′, 1)

	for i in 2:new_n-1
		for j in 2:new_n-1

			# T'[i,j] = T[i÷2,j÷2]

			T′[i, j] = (T[Int(floor((i+1)/2)), Int(floor((j+1)/2))]
						+ T[Int(ceil((i+1)/2)), Int(floor((j+1)/2))]
						+ T[Int(floor((i+1)/2)), Int(ceil((j+1)/2))]
						+ T[Int(ceil((i+1)/2)), Int(ceil((j+1)/2))])/4

		end
	end

	return T′
end

# ╔═╡ bdf66322-1ad7-11eb-1fc4-05c6877a1782
next = uplevel(coarse_results[end]);

# ╔═╡ 521c2858-1c80-11eb-0b03-c90d72040f5c
add_heaters!(next);

# ╔═╡ c45f7ec2-1ad7-11eb-0ff1-890a3dac8e95
myheatmap(next)

# ╔═╡ 1ab1b858-1ad8-11eb-09e8-47133b186c2b
function multilevel(T, k_final)

	n = size(T, 1)
	k_start = floor(Int, log2(n - 1))

	ks = Int[]
	iterations = Int[]

	for k in k_start: k_final
		push!(ks, k)

		results = simulation(T, 1e-2, 5000)
		push!(iterations, length(results))

		T = results[end]

		if k < k_final
			T = uplevel(T)
			add_heaters!(T)
		end
	end

	return T, ks, iterations
end

# ╔═╡ 882bf17a-1ad8-11eb-254e-1bafa2b2f599
starting_point = initial_condition(grid_with_boundary(4));

# ╔═╡ fcfebd94-1c80-11eb-35a5-4b6c26df5913
add_heaters!(starting_point);

# ╔═╡ 4e3e2892-1ad9-11eb-1cf8-61058302f3cb
@elapsed last, ks, iterations = multilevel(starting_point, 10)

# ╔═╡ 77d779a2-1c81-11eb-14fc-bd9a86cb7090
ks

# ╔═╡ 302dc78c-1c81-11eb-09c5-355d90caf574
iterations

# ╔═╡ b6802886-1c81-11eb-0f02-db9fd5c65a4d
cost = iterations[7] + iterations[6]/4 + iterations[5]/16 + iterations[4]/64 + iterations[3]/256 + iterations[2]/1024 + iterations[1]/2048

# ╔═╡ 7bb49244-1c81-11eb-1c7e-8bd4b88d9214
myheatmap(last)

# ╔═╡ Cell order:
# ╠═89ab29f2-1ad2-11eb-2aa1-a7fae56d7fd8
# ╠═32e91c48-1ad2-11eb-262e-e97ac9fa9a2d
# ╠═5805d042-1ad5-11eb-0f93-110cd0c6d41e
# ╠═79741a14-1ad2-11eb-34e1-7feadae4168e
# ╟─664637ba-1ad2-11eb-0382-b99596159809
# ╠═9a57746c-1ad0-11eb-3b48-c1d7877ef61a
# ╠═bf076518-1ad0-11eb-3b74-5bb405031a2b
# ╠═9997a2de-1ad2-11eb-2eaa-6dcd337db44d
# ╠═ab4b6a1a-1ad2-11eb-0536-873dab9733ee
# ╟─5d594710-1b80-11eb-0c14-8b824654c196
# ╠═98fb5f6a-1c7f-11eb-0b2e-a5185e8d7649
# ╠═af8777ae-1ad2-11eb-14b4-9d557cb5efa2
# ╠═113b8616-1ad3-11eb-1099-ebfe49db3815
# ╟─449547f4-1ad3-11eb-23e4-39b0d4e8f0f7
# ╟─95da5910-1ad3-11eb-1872-155aad1d8ae2
# ╠═d0f8cda6-1ad3-11eb-0717-57ce76bfd3c8
# ╠═6a7a1656-1ad4-11eb-2fd0-3fe69d6c6c4d
# ╠═74e6e18a-1ad4-11eb-389e-6b95f00eac8c
# ╠═923bd7c6-1ad4-11eb-35ce-55d99debaa5e
# ╠═857aebe2-1ad4-11eb-373c-bb0039279a74
# ╠═9877d516-1ad4-11eb-096b-75d3e8811259
# ╠═aa80f116-1ad4-11eb-0742-a992b8278c4b
# ╠═b5af2d1e-1ad4-11eb-01bb-3ffe82f8bbb0
# ╠═861efa52-1c82-11eb-3e5f-9f45aadd3054
# ╠═3047db8e-1ad5-11eb-0b3c-6da3fc2251b0
# ╠═8d946d20-1c7e-11eb-1d87-3d7228c4945e
# ╠═4dc4ec56-1ad5-11eb-0969-c7cd1728bb68
# ╠═debadd2e-1ad5-11eb-0bb6-99037ade1f6f
# ╠═96032a3c-1ad5-11eb-242b-b7d2b6c54568
# ╠═b01fbeee-1ad5-11eb-0c27-359d42df7775
# ╠═1a04f8ce-1ad6-11eb-3e8b-ad10bbc2a217
# ╠═27b3c57c-1ad6-11eb-2546-0fe68a40bf06
# ╟─4325b748-1ad6-11eb-20cc-d75cdbb1fbe5
# ╠═4fada552-1ad6-11eb-0e25-77b981c007c2
# ╠═bdf66322-1ad7-11eb-1fc4-05c6877a1782
# ╠═521c2858-1c80-11eb-0b03-c90d72040f5c
# ╠═c45f7ec2-1ad7-11eb-0ff1-890a3dac8e95
# ╠═1ab1b858-1ad8-11eb-09e8-47133b186c2b
# ╠═882bf17a-1ad8-11eb-254e-1bafa2b2f599
# ╠═fcfebd94-1c80-11eb-35a5-4b6c26df5913
# ╠═4e3e2892-1ad9-11eb-1cf8-61058302f3cb
# ╠═77d779a2-1c81-11eb-14fc-bd9a86cb7090
# ╠═302dc78c-1c81-11eb-09c5-355d90caf574
# ╠═b6802886-1c81-11eb-0f02-db9fd5c65a4d
# ╠═7bb49244-1c81-11eb-1c7e-8bd4b88d9214


import StatsBase

"""
    median(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the median of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the median is computed for each of those length-`L`
realisations, yielding a distribution of median estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the median for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the median of `x`, which is returned as
    a vector.
"""
function StatsBase.median(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.median, x, n)
end

"""
    mean(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the mean of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the mean is computed for each of those length-`L`
realisations, yielding a distribution of mean estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the mean for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the mean of `x`, which is returned as
    a vector.
"""
function StatsBase.mean(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.mean, x, n)
end

"""
    mode(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the mode of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the mode is computed for each of those length-`L`
realisations, yielding a distribution of mode estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the mode for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the mode of `x`, which is returned as
    a vector.
"""
function StatsBase.mode(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.mode, x, n)
end

"""
    quantile(x::UVAL_COLLECTION_TYPES, q, n::Int)

Obtain a distribution for the quantile(s) `q` of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the quantile is computed for each of those length-`L`
realisations, yielding a distribution of quantile estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the quantile for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the quantile of `x`, which is returned as
    a vector.
"""
function StatsBase.quantile(x::UVAL_COLLECTION_TYPES, q, n::Int)
    resample(StatsBase.quantile, x, n, q)
end

"""
    iqr(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the interquartile range (IQR), i.e. the 75th
percentile minus the 25th percentile, of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the IQR is computed for each of those length-`L`
realisations, yielding a distribution of IQR estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the IQR for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the IQR of `x`, which is returned as
    a vector.
"""
function StatsBase.iqr(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.iqr, x, n)
end

"""
    middle(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the middle of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the middle is computed for each of those length-`L`
realisations, yielding a distribution of middle estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the middle for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the middle of `x`, which is returned as
    a vector.
"""
function StatsBase.middle(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.middle, x, n)
end

"""
    std(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the standard deviation of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the standard deviation is computed for each of those length-`L`
realisations, yielding a distribution of standard deviation estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the std for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the standard deviation of `x`, which is returned as
    a vector.
"""
function StatsBase.std(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.std, x, n)
end

"""
    var(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the variance of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the variance is computed for each of those length-`L`
realisations, yielding a distribution of variance estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the variance for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the variance of `x`, which is returned as
    a vector.
"""
function StatsBase.var(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.var, x, n)
end

"""
    genmean(x::UVAL_COLLECTION_TYPES, p, n::Int)

Obtain a distribution for the generalized/power mean with exponent `p` of a
collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the generalized mean is computed for each of those length-`L`
realisations, yielding a distribution of generalized mean estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the generalized mean for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the generalized mean of `x`, which is returned as
    a vector.
"""
function StatsBase.genmean(x::UVAL_COLLECTION_TYPES, p, n::Int)
    resample(StatsBase.genmean, x, n, p)
end

"""
    genvar(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the generalized sample variance of a collection of
uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the generalized sample variance is computed for each of
those length-`L`
realisations, yielding a distribution of generalized sample variance estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the generalized sample variance for the realisation, which is a
    vector of length `L`.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the generalized sample variance of `x`,
    which is returned as a vector.
"""
function StatsBase.genvar(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.genvar, x, n)
end

"""
    harmmean(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the harmonic mean of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the harmonic mean is computed for each of those length-`L`
realisations, yielding a distribution of harmonic mean estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the harmonic mean for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the harmonic mean of `x`, which is returned as
    a vector.
"""
function StatsBase.harmmean(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.harmmean, x, n)
end

"""
    geomean(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the geometric mean of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the geometric mean is computed for each of those length-`L`
realisations, yielding a distribution of geometric mean estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the geometric mean for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the geometric mean of `x`, which is returned as
    a vector.
"""
function StatsBase.geomean(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.geomean, x, n)
end

"""
    kurtosis(x::UVAL_COLLECTION_TYPES, n::Int, f = StatsBase.mean)

Obtain a distribution for the kurtosis of a collection of uncertain values.

This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the kurtosis is computed for each of those length-`L`
realisations, yielding a distribution of kurtosis estimates.

Optionally, a center function `f` can be specified. This function is used
to compute the center of each draw, i.e. for the i-th draw, call
`StatsBase.kurtosis(draw_i, f(draw_i))`.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the kurtosis for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the kurtosis of `x`, which is returned as a vector.
"""
function StatsBase.kurtosis(x::UVAL_COLLECTION_TYPES, n::Int, f = StatsBase.mean)

    kurtosis_estimates = zeros(Float64, n)

    for i = 1:n
        draw = resample(x)
        kurtosis_estimates[i] = StatsBase.kurtosis(draw, f(draw))
    end

    return kurtosis_estimates
end

"""
    moment(x::UVAL_COLLECTION_TYPES, k, n::Int)

Obtain a distribution for the `k`-th order central moment of a collection
of uncertain values.

This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the `k`-th order central moment is computed for each
of those length-`L`realisations, yielding a distribution of `k`-th
order central moment estimates.

The procedure is as follows.

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the `k`-th order central moment for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the `k`-th order central moment of `x`,
    which is returned as a vector.
"""
function StatsBase.moment(x::UVAL_COLLECTION_TYPES, k, n::Int)
    resample(StatsBase.moment, x, n, k)
end

"""
    percentile(x::UVAL_COLLECTION_TYPES, p, n::Int)

Obtain a distribution for the percentile(s) `p` of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the percentile is computed for each of those length-`L`
realisations, yielding a distribution of percentile estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the percentile for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the percentile of `x`, which is returned as
    a vector.
"""
function StatsBase.percentile(x::UVAL_COLLECTION_TYPES, p, n::Int)
    resample(StatsBase.percentile, x, n, p)
end

"""
    renyientropy(x::UVAL_COLLECTION_TYPES, α, n::Int)

Obtain a distribution for the Rényi (generalized) entropy of
order `α` of a collection of uncertain values.

This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the generalized entropy is computed for each
of those length-`L`realisations, yielding a distribution of
generalized entropy estimates.

The procedure is as follows.

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the Rényi (generalized) entropy of order `α` for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the Rényi (generalized) entropy of
    order `α` of `x`, which is returned as a vector.
"""
function StatsBase.renyientropy(x::UVAL_COLLECTION_TYPES, α, n::Int)
    resample(StatsBase.renyientropy, x, n, α)
end

"""
    rle(x::UVAL_COLLECTION_TYPES, α, n::Int)

Obtain a distribution for the run-length encoding of a
collection of uncertain values.

This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the run-length encoding is computed for each
of those length-`L`realisations, yielding a distribution of
run-length encoding estimates.

Returns a vector of tuples of run-length encodings.

The procedure is as follows.

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the run-length encoding for the realisation. This gives a
    tuple, where the first element of the tuple is a vector of
    values of the input and the second is the number of consecutive occurrences of
    each element.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the run-length encoding of `x`,
    which is returned as a vector of the run-length encoding tuples.
"""
function StatsBase.rle(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.rle, x, n)
end

"""
    sem(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the standard error of the mean of a
collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the standard error of the mean is computed for
each of those length-`L`
realisations, yielding a distribution of standard error of the
mean estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the standard error of the mean for the realisation, which is
    a vector of length `L`.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the standard error of the mean of `x`,
    which is returned as a vector.
"""
function StatsBase.sem(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.sem, x, n)
end

"""
    skewness(x::UVAL_COLLECTION_TYPES, n::Int, f = StatsBase.mean)

Obtain a distribution for the skewness of a collection of uncertain values.

This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the skewness is computed for each of those length-`L`
realisations, yielding a distribution of skewness estimates.

Optionally, a center function `f` can be specified. This function is used
to compute the center of each draw, i.e. for the i-th draw, call
`StatsBase.skewness(draw_i, f(draw_i))`.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the skewness for the realisation.
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the skewness of `x`,
    which is returned as a vector.
"""
function StatsBase.skewness(x::UVAL_COLLECTION_TYPES, n::Int, f = StatsBase.mean)
    skewness_estimates = zeros(Float64, n)

    for i = 1:n
        draw = resample(x)
        skewness_estimates[i] = StatsBase.skewness(draw, f(draw))
    end

    return skewness_estimates
end

"""
    span(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the span of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the span is computed for each of those length-`L`
realisations, yielding a distribution of span estimates.

Returns a length-`L` vector of `span`s, where the i-th span is the range
`minimum(draw_x_i):maximum(draw_x_i)`.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the span for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the span of `x`, which is returned as
    a vector.
"""
function StatsBase.span(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.span, x, n)
end

"""
    totalvar(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the total variance of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the total variance is computed for each of those length-`L`
realisations, yielding a distribution of total variance estimates.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the total variance for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the total variance of `x`, which is returned as
    a vector.
"""
function StatsBase.totalvar(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.totalvar, x, n)
end

"""
    summarystats(x::UVAL_COLLECTION_TYPES, n::Int)

Obtain a distribution for the summary statistics of a collection of uncertain values.
This is done by first drawing `n` length-`L` realisations of `x`, where
`L = length(x)`. Then, the summary statistics is computed for each of those length-`L`
realisations, yielding a distribution of summary statistics estimates.

Returns a length-`L` vector of `SummaryStats` objects containing the mean, minimum,
25th percentile, median, 75th percentile, and maximum for each draw of `x`.

Detailed steps:

1. First, draw a length-`L` realisation of `x` by drawing one random
    number from  each uncertain value furnishing the dataset. The draws are
    independent, so that no element-wise dependencies (e.g. sequential
    correlations) that are not already present in the data are introduced in
    the realisation.
2. Compute the summary statistics for the realisation, which is a vector of length `L`
3. Repeat the procedure `n` times, drawing `n` independent realisations of `x`.
    This yields `n` estimates of the summary statistics of `x`, which is returned as
    a vector.
"""
function StatsBase.summarystats(x::UVAL_COLLECTION_TYPES, n::Int)
    resample(StatsBase.summarystats, x, n)
end


# # PTDF with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl)

# **Originally Contributed by**: Sourabh Dalvi

# ## Introduction

# PowerSimulations.jl supports linear PTDF optimal power flow formulation. This example shows a
# single multi-period optimization of economic dispatch with a linearilzed DC-OPF representation of
# using PTDF power flow and how to extract duals values or locational marginal prices for energy.

# ## Dependencies
# We can use the same RTS data and some of the initialization as in
# [OperationsProblem example](../../notebook/3_PowerSimulations_examples/1_operations_problems.ipynb)
# by sourcing it as a dependency.
using SIIPExamples
pkgpath = dirname(dirname(pathof(SIIPExamples)))
include(joinpath(
    pkgpath,
    "test",
    "3_PowerSimulations_examples",
    "01_operations_problems.jl",
));

# Since we'll be retreving duals, we need a solver that returns duals values
# here we use Ipopt.
using Ipopt
solver = optimizer_with_attributes(Ipopt.Optimizer)

# In the [OperationsProblem example](../../notebook/3_PowerSimulations_examples/1_operations_problems.ipynb)
# we defined a unit-commitment problem with a copper plate representation of the network.
# Here, we want do define an economic dispatch (linear generation decisions) with
# linear DC-OPF using PTDF network representation.
# So, starting with the network, we can select from _almost_ any of the endpoints on this
# tree:
#nb TypeTree(PSI.PM.AbstractPowerModel,  init_expand = 10, scopesep="\n")

# For now, let's just choose a standard PTDF formulation.
ed_template = template_economic_dispatch(network = StandardPTDFModel)

# Currently  energy budget data isn't stored in the RTS-GMLC dataset.
ed_template.devices[:Hydro] = DeviceModel(HydroEnergyReservoir, HydroDispatchRunOfRiver)

# Calculate the PTDF matrix.
PTDF_matrix = PTDF(sys)

# Now we can build a 4-hour economic dispatch / PTDF problem with the RTS data.
# Here, we have to pass the keyword argument `constraint_duals` to OperationsProblem
# with the name of the constraint for which duals are required for them to be returned in the results.
problem = OperationsProblem(
    EconomicDispatchProblem,
    ed_template,
    sys,
    horizon = 4,
    optimizer = solver,
    balance_slack_variables = true,
    constraint_duals = [:CopperPlateBalance, :network_flow],
    PTDF = PTDF_matrix,
)

# And solve the problem and collect the results
res = solve!(problem)

# Here we collect the dual values from the results for the `CopperPlateBalance` and `network_flow`
# constraints. In the case of PTDF network formulation we need to compute the final LMP for each bus in the system by
# subtracting the duals (μ) of `network_flow` constraint multipled by the PTDF matrix
# from the  dual (λ) of `CopperPlateBalance` constraint.
# Note:we convert the results from DataFrame to Array for ease of use.
λ = convert(Array, res.dual_values[:CopperPlateBalance])
μ = convert(Array, res.dual_values[:network_flow])

# Here we create Dict to store the calculate congestion component of the LMP which is a product of μ and the PTDF matrix.
buses = get_components(Bus, sys)
congestion_lmp = Dict()
for bus in buses
    congestion_lmp[get_name(bus)] = μ * PTDF_matrix[:, get_number(bus)]
end
congestion_lmp = DataFrame(congestion_lmp)

# Finally here we get the LMP for each node in a lossless DC-OPF using the PTDF formulation.
LMP = λ .- congestion_lmp


########################################################################################
# Sampling from a solution
########################################################################################
"""Helper function to sample from our covariances, which often have a "cross" of zeros
For the 0-cov entries the outcome of the sampling is deterministic!"""
function _rand(x::Gaussian, n::Int=1)
    m, C = x.μ, x.Σ
    @assert C isa SRMatrix

    sample = m .+ C.squareroot * randn(length(m), n)
    return sample
end

function sample_states(sol::ProbODESolution, n::Int=1)
    @assert sol.interp.smooth "sampling not implemented for non-smoothed posteriors"
    return sample_states(sol.t, sol.x_filt, sol.diffusions, sol.t, sol.interp, n)
end
function sample(sol::ProbODESolution, n::Int=1)
    d, q = sol.interp.d, sol.interp.q
    sample_path = sample_states(sol, n)
    return sample_path[:, 1:q+1:d*(q+1), :]
end
function sample_states(ts, xs, diffusions, difftimes, posterior, n::Int=1)
    @assert length(diffusions) + 1 == length(difftimes)

    @unpack A, Q, d, q = posterior
    D = d * (q + 1)

    x = xs[end]
    sample = _rand(x, n)
    @assert size(sample) == (D, n)

    sample_path = zeros(length(ts), D, n)
    sample_path[end, :, :] .= sample
    # @info "final value and samples" x.μ sample sample_path[end, :]

    for i in length(xs)-1:-1:1
        dt = ts[i+1] - ts[i]

        i_diffusion = sum(difftimes .<= ts[i])
        diffusion = diffusions[i_diffusion]

        Qh = apply_diffusion(Q, diffusion)
        make_preconditioners!(posterior, dt)
        P, PI = posterior.P, posterior.PI

        for j in 1:n
            sample_p = P * sample_path[i+1, :, j]
            x_prev_p = P * xs[i]

            prev_sample_p, _ =
                smooth(x_prev_p, Gaussian(sample_p, SRMatrix(zeros(D, D))), A, Qh)

            # sample_path[i, :, j] .= PI*prev_sample_p.μ
            sample_path[i, :, j] .= PI * _rand(prev_sample_p)[:]
        end
    end

    return sample_path
end
function dense_sample_states(sol::ProbODESolution, n::Int=1; density=1000)
    @assert sol.interp.smooth "sampling not implemented for non-smoothed posteriors"
    times = range(sol.t[1], sol.t[end], length=density)
    states = StructArray([
        sol.interp(t, sol.t, sol.x_filt, sol.x_smooth, sol.diffusions; smoothed=false)
        for t in times
    ])

    return sample_states(times, states, sol.diffusions, sol.t, sol.interp, n), times
end
function dense_sample(sol::ProbODESolution, n::Int=1; density=1000)
    samples, times = dense_sample_states(sol, n; density=density)
    d, q = sol.interp.d, sol.interp.q
    return samples[:, 1:q+1:d*(q+1), :], times
end


let
    # fitting noisy data to an exponential model
    # TODO: Change to `.-x` when 0.5 support is dropped
    @. model(x, p) = p[1] * exp(-x * p[2])
    model_inplace(F, x, p) = (@. F = p[1] * exp(-x * p[2]))

    # some example data
    Random.seed!(12345)
    xdata = range(0, stop=10, length=500000)
    ydata = model(xdata, [1.0, 2.0]) + 0.01*randn(length(xdata))
    p0 = [0.5, 0.5]

    # if your model is differentiable, it can be faster and/or more accurate
    # to supply your own jacobian instead of using the finite difference
    function jacobian_model(x,p)
        J = Array{Float64}(undef, length(x), length(p))
        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
        @. @views J[:,2] = -x*p[1]*J[:,1]
        J
    end

    function jacobian_model_inplace(J::Array{Float64,2},x,p)
        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
        @. @views J[:,2] = -x*p[1]*J[:,1]
    end

    f(p) = model(xdata, p) - ydata
    g(p) = jacobian_model(xdata, p)
    df = OnceDifferentiable(f, g, p0, similar(ydata); inplace = false);
    evalf(x) = NLSolversBase.value!!(df, x)
    evalg(x) = NLSolversBase.jacobian!!(df, x)
    r = evalf(p0);
    j = evalg(p0);

    f_inplace = (F,p) -> (model_inplace(F, xdata, p); @. F = F - ydata)
    g_inplace = (G,p) -> jacobian_model_inplace(G, xdata, p)
    df_inplace = OnceDifferentiable(f_inplace, g_inplace, p0, similar(ydata); inplace = true);
    evalf_inplace(x) = NLSolversBase.value!!(df_inplace,x)
    evalg_inplace(x) = NLSolversBase.jacobian!!(df_inplace,x)
    r_inplace = evalf_inplace(p0);
    j_inplace = evalg_inplace(p0);

    @test r == r_inplace
    @test j == j_inplace

    println("--------------\nPerformance of non-inplace")
    println("\t Evaluation function")

    stop = 8 #8 because the tests afterwards will call the eval function 8 or 9 times, so it makes it easy to compare
    step = 1

    @time for i=range(0,stop=stop,step=step)
        evalf(p0);
    end

    println("\t Jacobian function")
    @time for i=range(0,stop=stop,step=step)
        evalg(p0);
    end

    println("--------------\nPerformance of inplace")
    println("\t Evaluation function")
    @time for i=range(0,stop=stop,step=step)
        evalf_inplace(p0);
    end

    println("\t Jacobian function")
    @time for i=range(0,stop=stop,step=step)
        evalg_inplace(p0);
    end

    curve_fit(model, xdata, ydata, p0; maxIter=100); #warmup
    curve_fit(model_inplace, xdata, ydata, p0; inplace = true, maxIter=100);

    #explicit jac
    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=100);
    curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, p0; inplace = true, maxIter=100);

    println("--------------\nPerformance of curve_fit")

    println("\t Non-inplace")
    fit = @time curve_fit(model, xdata, ydata, p0; maxIter=100)
    @test fit.converged

    println("\t Inplace")
    fit_inplace = @time curve_fit(model_inplace, xdata, ydata, p0; inplace = true, maxIter=100)
    @test fit_inplace.converged

    @test fit_inplace.param == fit.param

    println("\t Non-inplace with jacobian")
    fit_jac = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=100)
    @test fit_jac.converged

    println("\t Inplace with jacobian")
    fit_inplace_jac = @time curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, p0; inplace = true, maxIter=100)
    @test fit_inplace_jac.converged

    @test fit_jac.param == fit_inplace_jac.param

    # some example data
    yvars = 1e-6*rand(length(xdata))
    ydata = model(xdata, [1.0, 2.0]) + sqrt.(yvars) .* randn(length(xdata))

    println("--------------\nPerformance of curve_fit with weights")

    curve_fit(model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]);
    curve_fit(model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100);

    curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]);
    curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100);

    println("\t Non-inplace with weights")
    fit_wt = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; maxIter=100)
    @test fit_wt.converged

    println("\t Inplace with weights")
    fit_inplace_wt = @time curve_fit(model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100)
    @test fit_inplace_wt.converged

    @test maximum(abs.(fit_wt.param - fit_inplace_wt.param)) < 1e-15

    println("\t Non-inplace with jacobian with weights")
    fit_wt_jac = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; maxIter=100)
    @test fit_wt_jac.converged

    println("\t Inplace with jacobian with weights")
    fit_inplace_wt_jac = @time curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100)
    @test fit_inplace_wt_jac.converged

    @test maximum(abs.(fit_wt_jac.param - fit_inplace_wt_jac.param)) < 1e-15

end


using DSP, FFTW

"""
    melspectrogram(signal, sr, window=ones, windowsize=2048, overlap=512, nmels=26)

Compute the melspectrogram of the given signal.

**Parameters:**

`signal`: The signal on which the coefficients will be computed.

`sr`: Sample rate of the signal.

`window`: The window function to be applied on the signal by the STFT.

`windowsize`: The size of the STFT window.

`overlap`: The overlap of two adjacent windows of the STFT in samples.

`nmels`: The amount of mel bands to be extracted from the signal.
"""
function melspectrogram(
	signal, sr, window=ones, windowsize=2048, overlap=512, nmels=128
)

	# Transform the input to mel scale
	mel(frequency) = 1127*log(1 + frequency/700)

	# Transform the input back from mel scale
	imel(mel) = 700*(exp(mel/1127) - 1)

	# Return an array of mel spaced points
	melspace(start, length, stop) = imel.(range(mel(start), mel(stop), length=length))

	# Each column is a frame of stft
	sigstft = stft(signal, windowsize, overlap, window=window)

	# If the signal is real then the absolute value of it's spectrum is symmetric
	eltype(signal)<:Real && (windowsize = windowsize ÷ 2 + 1)

	# Evaluate the periodogram of the signal
	periodgrm = [abs2.(frame)/windowsize for frame in sigstft]

	# Create a set of linearly spaced points in the mel scale
	fmax = sr/2
	melpoints = ceil.(Integer, melspace(1, nmels + 2, fmax)*windowsize/fmax)

	# Create the filterbank...
	windows = map(1:nmels) do m
		# ...by creating a triangular window for each set of three consecutive points
		start, apex, stop = melpoints[m:m + 2]
		w = zeros(stop - start)

		for k in 1:(apex - start)
			w[k] = k/(apex - start)
		end

		for k in (apex + 1):stop
			w[k - start] = (stop - k)/(stop - apex)
		end

		return 2w/(stop - start)	# Normalized so that the triangle has area 1
	end

	windowstarts = melpoints[1:nmels]

	# Apply the filterbank to each frame of the periodogram...
	return mapslices(periodgrm, dims = 1) do frame
		# ...by multiplying the non zero part of each window in the filterbank
		return map(windows, windowstarts) do melwindow, winstart
			return melwindow'*frame[winstart:winstart + length(melwindow) - 1]
		end
	end

end


function solve()
    n = parse(Int, readline())
    as = [parse(Int, x) for x in split(readline())]

    ret = 0
    for l=1:n
        x = as[l]
        for r=l:n
            x = min(x, as[r])
            ret = max(ret, x*(r-l+1))
        end
    end
    ret
end

println(solve())


"""
This can contain a maximum of `length(par)` parenting-relations
par is an array of `Int`, which is the index of the parent node.
"""
mutable struct DisjointSet
    par::Vector{Int}
    function DisjointSet(size)
        x=[i for i in 1:size]
        return new(x)
    end
end

"""
Find the ancestor of node `x`.
"""
function find(set::DisjointSet,x::Int)::Int
    if set.par[x]==x return x
    else return set.par[x]=find(set,set.par[x])
    end
end

function Base.merge!(set::DisjointSet,x::Int,y::Int)
    x=find(set,x)
    y=find(set,y)
    set.par[x]=y
end


"""
    convolve(inp, kernel)

Convolves an inpL x inpL array with the FFT of a centered kernel
of size kerL x kerL to produce an output of size (kerL - inpL) x (kerL - inpL).

TODO: maybe this should return a real array
"""
function convolve(inp, kernel)
    inpL = size(inp)[1]
    kerL = size(kernel)[1]
    outL = kerL - inpL

    arr_pad = [inp zeros(inpL, outL); zeros(outL, inpL) zeros(outL, outL)]
    #out_pad = P \ ((P * arr_pad) .* kernel)
    out_pad = planned_ifft(planned_fft(arr_pad) .* kernel)
    out = out_pad[inpL+1:kerL, inpL+1:kerL]

    out
end

function convolveT(out, kernel)
    outL = size(out)[1]
    kerL = size(kernel)[1]
    inpL = kerL - outL

    out_pad = [zeros(inpL, inpL) zeros(inpL, outL); zeros(outL, inpL) out]
    #arr_pad = P * ((P \ out_pad) .* kernel)
    arr_pad = planned_fft(planned_ifft(out_pad) .* kernel)
    arr = arr_pad[1:inpL, 1:inpL]

    arr
end

# padded stores the intermediate padded result
function convolve!(inp, kernel, padded)
    inpL = size(inp)[1]
    kerL = size(kernel)[1]
    outL = kerL - inpL

    padded .= 0
    padded[1:inpL, 1:inpL] .= inp
    planned_fft!(padded)
    padded .*= kernel
    planned_ifft!(padded)
    out = @view padded[inpL+1:kerL, inpL+1:kerL]

    out
end

function convolveT!(out, kernel, padded)
    outL = size(out)[1]
    kerL = size(kernel)[1]
    inpL = kerL - outL

    padded .= 0
    padded[inpL+1:inpL+outL, inpL+1:inpL+outL] .= out
    planned_ifft!(padded)
    padded .*= kernel
    planned_fft!(padded)
    inp = @view padded[1:inpL, 1:inpL]

    inp
end


using ShrinkageEstimators
using Base.Test

d = 1
n = 10000
b = sqrt(n) # getfrequencies function uses sqrt(n) bins

arr = rand(n, d)
arrX = arr
arrY = rand(10000, 1)
arrZ = rand(10000, 1)

# Since the estimator is only an estimator, the tests are only approximate.
# In general a tolerance of 0.5 is used, which is relatively large. This
# should pass even if the estimate is slightly off, but fail if something
# is going drastically wrong.

# Check entropy is roughly log_base(b) for uniform distribution with b bins
@test_approx_eq_eps get_shrinkage_entropy(arr) log2(b) 0.5
println("Entropy passed")

# Check change of base works
@test_approx_eq_eps get_shrinkage_entropy(arr, base=e) log(b) 0.5
println("Entropy with change of base passed")

# Check specifying lambda works
@test_approx_eq_eps get_shrinkage_entropy(arr, lambda=0) log2(b) 0.5
println("Entropy with with lambda 0 passed")

# Check uniform width discretization works
@test_approx_eq_eps get_shrinkage_entropy(arr, mode="uniformcount") log2(b) 0.5
println("Entropy with with uniform count passed")

# TODO: Work out how to test "bayesianblocks" mode

# Check joint entropy between multiple identical distributions is roughly
# the same as entropy for one of them
@test_approx_eq_eps get_shrinkage_entropy(arr) get_shrinkage_entropy(arr, arr) 0.5
println("Joint entropy between 2 variables passed")
@test_approx_eq_eps get_shrinkage_entropy(arr) get_shrinkage_entropy(arr, arr, arr) 0.5
println("Joint entropy between 3 variables passed")

# Check I(X; Y) is H(X) + H(Y) - H(X, Y)
a = get_shrinkage_entropy(arrX) + get_shrinkage_entropy(arrY) - get_shrinkage_entropy(arrX, arrY)
@test a == get_shrinkage_mi(arrX, arrY)
println("Mutual information passed")

# Check I(X; Y | Z) is H(X, Z) + H(Y, Z) - H(X, Y, Z) - H(Z)
a = get_shrinkage_entropy(arrX, arrZ) + get_shrinkage_entropy(arrY, arrZ) - get_shrinkage_entropy(arrX, arrY, arrZ) - get_shrinkage_entropy(arrZ)
@test a == get_shrinkage_cmi(arrX, arrY, arrZ)
println("Conditional mutual information passed")


using Surrogates
using LinearAlgebra
using Zygote
using Test
#using Zygote: @nograd
#=
#FORWARD
###### 1D ######
lb = 0.0
ub = 10.0
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> x^2
y = f.(x)

#Radials
my_rad = RadialBasis(x,y,lb,ub,x->norm(x),2)
g = x -> ForwardDiff.derivative(my_rad,x)
g(5.0)

#Kriging
p = 1.5
my_krig = Kriging(x,y,p)
g = x -> ForwardDiff.derivative(my_krig,x)
g(5.0)

#Linear Surrogate
my_linear = LinearSurrogate(x,y,lb,ub)
g = x -> ForwardDiff.derivative(my_linear,x)
g(5.0)

#Inverse distance
p = 1.4
my_inverse = InverseDistanceSurrogate(x,y,p,lb,ub)
g = x -> ForwardDiff.derivative(my_inverse,x)
g(5.0)

#Lobachevsky
n = 4
α = 2.4
my_loba = LobachevskySurrogate(x,y,α,n,lb,ub)
g = x -> ForwardDiff.derivative(my_loba,x)
g(5.0)

#Second order polynomial
my_second = SecondOrderPolynomialSurrogate(x,y,lb,ub)
g = x -> ForwardDiff.derivative(my_second,x)
g(5.0)

###### ND ######
lb = [0.0,0.0]
ub = [10.0,10.0]
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> x[1]*x[2]
y = f.(x)

#Radials
my_rad = RadialBasis(x,y,[lb,ub],z->norm(z),2)
g = x -> ForwardDiff.gradient(my_rad,x)
g([2.0,5.0])

#Kriging
theta = [2.0,2.0]
p = [1.9,1.9]
my_krig = Kriging(x,y,p,theta)
g = x -> ForwardDiff.gradient(my_krig,x)
g([2.0,5.0])

#Linear Surrogate
my_linear = LinearSurrogate(x,y,lb,ub)
g = x -> ForwardDiff.gradient(my_linear,x)
g([2.0,5.0])

#Inverse Distance
p = 1.4
my_inverse = InverseDistanceSurrogate(x,y,p,lb,ub)
g = x -> ForwardDiff.gradient(my_inverse,x)
g([2.0,5.0])

#Lobachevsky
alpha = [1.4,1.4]
n = 4
my_loba_ND = LobachevskySurrogate(x,y,alpha,n,lb,ub)
g = x -> ForwardDiff.gradient(my_loba_ND,x)
g([2.0,5.0])

#Second order polynomial
my_second = SecondOrderPolynomialSurrogate(x,y,lb,ub)
g = x -> ForwardDiff.gradient(my_second,x)
g([2.0,5.0])
=#

##############
### ZYGOTE ###
##############

############
#### 1D ####
############
lb = 0.0
ub = 3.0
n = 10
x = sample(n,lb,ub,SobolSample())
f = x -> x^2
y = f.(x)

#Radials
@testset "Radials 1D" begin
    my_rad = RadialBasis(x,y,lb,ub,rad = linearRadial())
    g = x -> my_rad'(x)
    g(5.0)
end

# #Kriging
@testset "Kriging 1D" begin
    my_p = 1.5
    my_krig = Kriging(x,y,lb,ub,p=my_p)
    g = x -> my_krig'(x)
    g(5.0)
end

# #Linear Surrogate
@testset "Linear Surrogate" begin
    my_linear = LinearSurrogate(x,y,lb,ub)
    g = x -> my_linear'(x)
    g(5.0)
end

#Inverse distance
@testset "Inverse Distance" begin
    my_p = 1.4
    my_inverse = InverseDistanceSurrogate(x,y,lb,ub,p=my_p)
    g = x -> my_inverse'(x)
    g(5.0)
end

#Second order polynomial
@testset "Second Order Polynomial" begin
    my_second = SecondOrderPolynomialSurrogate(x,y,lb,ub)
    g = x -> my_second'(x)
    g(5.0)
end

#Lobachevsky
@testset "Lobachevsky" begin
    n = 4
    α = 2.4
    my_loba = LobachevskySurrogate(x,y,lb,ub, alpha = α, n = 4)
    g = x -> my_loba'(x)
    g(0.0)
end

#Wendland
@testset "Wendland" begin
    my_wend = Wendland(x,y,lb,ub)
    g = x -> my_wend'(x)
    g(3.0)
end

# MOE and VariableFidelity for free because they are Linear combinations
# of differentiable surrogates

# #Gek
@testset "GEK" begin
    n = 10
    lb = 0.0
    ub = 5.0
    x = sample(n,lb,ub,SobolSample())
    f = x-> x^2
    y1 = f.(x)
    der = x->2*x
    y2 = der.(x)
    y = vcat(y1,y2)
    my_gek = GEK(x,y,lb,ub)
    g = x-> my_gek'(x)
    g(3.0)
end

################
###### ND ######
################

lb = [0.0,0.0]
ub = [10.0,10.0]
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> x[1]*x[2]
y = f.(x)

#Radials
@testset "Radials ND" begin
    my_rad = RadialBasis(x,y,lb,ub,rad = linearRadial(), scale_factor = 2.1)
    g = x -> Zygote.gradient(my_rad,x)
    g((2.0,5.0))
end

# Kriging
@testset "Kriging ND" begin
    my_theta = [2.0,2.0]
    my_p = [1.9,1.9]
    my_krig = Kriging(x,y,lb,ub,p=my_p,theta=my_theta)
    g = x -> Zygote.gradient(my_krig,x)
    g((2.0,5.0))
end

# Linear Surrogate
@testset "Linear Surrogate ND" begin
    my_linear = LinearSurrogate(x,y,lb,ub)
    g = x -> Zygote.gradient(my_linear,x)
    g((2.0,5.0))
end

#Inverse Distance
@testset "Inverse Distance ND" begin
    my_p = 1.4
    my_inverse = InverseDistanceSurrogate(x,y,lb,ub,p=my_p)
    g = x -> Zygote.gradient(my_inverse,x)
    g((2.0,5.0))
end

#Lobachevsky not working yet weird issue with Zygote @nograd
#=
Zygote.refresh()
alpha = [1.4,1.4]
n = 4
my_loba_ND = LobachevskySurrogate(x,y,alpha,n,lb,ub)
g = x -> Zygote.gradient(my_loba_ND,x)
g((2.0,5.0))
=#

# Second order polynomial mutating arrays
@testset "SecondOrderPolynomialSurrogate ND" begin
    my_second = SecondOrderPolynomialSurrogate(x,y,lb,ub)
    g = x -> Zygote.gradient(my_second,x)
    g((2.0,5.0))
end

#wendland
@testset "Wendland ND" begin
    my_wend_ND = Wendland(x,y,lb,ub)
    g = x -> Zygote.gradient(my_wend_ND,x)
    g((2.0,5.0))
end
# #MOE and VariableFidelity for free because they are Linear combinations
# #of differentiable surrogates


# SPDX-License-Identifier: MIT

"""
    AbstractNormalDiag{D} <: Distribution

Abstract type for Normal distribution (with diagonal covariance matrix)
implementations.
"""
abstract type AbstractNormalDiag{D} <: Distribution end

function DefaultNormalDiagParameter(μ::AbstractVector{T},
                                    v::AbstractVector{T}) where T
    λ = 1 ./ v
    ξ = vcat(λ .* μ, -T(.5)*λ)
    DefaultParameter(ξ)
end

"""
    mutable struct NormalDiag{D} <: AbstractNormalDiag{D}
        param::P where P <: AbstractParameter
    end

Normal distribution with a diagonal covariance matrix.

# Constructors

    NormalDiag(μ, v)

where `μ` is the mean `v` is the diagonal of the covariance matrix.

# Examples
```jldoctest
julia> NormalDiag([1.0, 1.0], [2.0, 1.0])
NormalDiag{2}:
  μ = [1.0, 1.0]
  Σ = [2.0 0.0; 0.0 1.0]
```
"""
mutable struct NormalDiag{D} <: AbstractNormalDiag{D}
    param::P where P <: AbstractParameter
end

function NormalDiag(μ::AbstractVector{T}, v::AbstractVector{T}) where T
    param = DefaultNormalDiagParameter(μ, v)
    P = typeof(param)
    D = length(μ)
    NormalDiag{D}(param)
end
NormalDiag(μ::AbstractVector, Σ::Diagonal) = NormalDiag(μ, diag(Σ))

#######################################################################
# Distribution interface.

function basemeasure(::AbstractNormalDiag, x::AbstractVector{T}) where T
    -T(.5) * length(x) * log(T(2π))
end

function lognorm(n::AbstractNormalDiag{D},
                 η::AbstractVector{T} = naturalform(n.param)) where {T,D}
    η₁, η₂ = η[1:D], η[D+1:end]
    -T(.5) * sum(log.(-T(2)*η₂)) -T(.25) * dot(η₁, (1 ./ η₂) .* η₁)
end

function gradlognorm(n::AbstractNormalDiag{D}, η = naturalform(n.param)) where D
    T = eltype(η)
    Λμ, nhλ = η[1:D], η[D+1:end]
    v = 1 ./ (-T(2) * nhλ)
    μ = v .* Λμ
    m = v .+ μ.^2
    vcat(μ, m)
end

function sample(n::AbstractNormalDiag{D}, size) where D
    μ, Σ = stdparam(n)
    σ = sqrt.(diag(Σ))
    ϵ = randn!(similar(μ, D, size))
    μ .+ σ .* ϵ
end

stats(::AbstractNormalDiag, x) = vcat(x, x.^2)

_splitgrad_normaldiag(D, m) = m[1:D], m[D+1:end]
splitgrad(n::AbstractNormalDiag{D}, m) where D = m[1:D], m[D+1:end]

function stdparam(n::AbstractNormalDiag{D},
                  η::AbstractVector{T} = naturalform(n.param)) where {T,D}
    Λμ, nhλ = η[1:D], η[D+1:end]
    v = 1 ./ (-T(2) * nhλ)
    μ = v .* Λμ
    (μ = μ, Σ = Diagonal(v))
end


function compute(file::String)::Int
    chars = Dict(Char(Int('A') + i) => i + 1 for i ∈ 0:25)
    return sum(i * chars[c] for (i, name) ∈ enumerate(sort(split(replace(file, "\"" => ""), ","))) for c ∈ name)
end


############################################################################################
## This file is part of RLinearAlgebra.jl
##
## Overview: abstractions and methods for performing various deploying Random Linear System
## solvers
##
## Contents
## - `RLSSolver` struct
## - `rsolve` functions
## - Randomized linear Solvers
## - Export Statements
##
############################################################################################

# Dependencies:

#############################################
# `RLSSolver` stucture
#############################################

"""
    RLSSolver{S<:LinSysSampler, R<:LinSysSolveRoutine,
        L<:LinSysSolverLog, C<:LinSysStopCriterion}

An encapsulation for a (R)andomized (L)inear (S)ysterm Solver as a mutable structure.

# Fields
- `sampler::S`, a procedure for sampling, sketching or selecting with respect to a linear
    system.
- `routine::R`, a routine for randomly generating the subproblem generated by the sampling,
    sketching or selecting operation
- `log::L`, a logger for the progress and behavior of the randomized linear system solver.
- `stop::C`, a stopping criterion for the randomized linear system solver.
- `x::Union{Vector{Float64},Nothing}`, the result produced by the randomized linear system
    solver.

# Constructions
- `RLSSolver(iter::Int64)` specifies a solver with: random cyclic row sampling;  partially
    orthogonalized row projections (memory of five) subproblem solver; a full residual
    logger; and a maximum iteration stopping criterion specified by `iter`.
"""
mutable struct RLSSolver{S<:LinSysSampler, R<:LinSysSolveRoutine,
    L<:LinSysSolverLog, C<:LinSysStopCriterion}

    sampler::S
    routine::R
    log::L
    stop::C
    x::Union{Vector{Float64},Nothing}
end

RLSSolver(iter::Int64) = RLSSolver(
    LinSysVecRowRandCyclic(),   # Random Cyclic Sampling
    LinSysVecRowProjPO(),       # Partially Orthogonalized Row Projection, 5 vector memory
    LSLogFull(),                # Full Logger: maintains residual history
    LSStopMaxIterations(iter),  # Maximum iterations stopping criterion
    nothing                     # System solution
)

#############################################
# `rsolve` functions
#############################################
"""
    rsolve([solver::RLSSolver,] A, b)

Deploys an iterative randomized solver for the linear system whose coefficient matrix is
    encapsulated in `A` and whose constant vector is encapsulated in `b`. If the `solver`
    is not specified, then the default `RLSSolver` is constructed and used with a maximum
    number of iterations of `10 * length(b)`. See `RLSSolver` for more details on the
    default solver.

The function returns the solution.
"""
function rsolve(A, b)
    solver = RLSSolver(10 * length(b))
    x = zeros(size(A, 2))
    rsolve!(solver, A, b, x)

    return x
end

function rsolve(solver::RLSSolver, A, b)
    x = zeros(size(A, 2))
    rsolve!(solver, A, b, x)

    return x
end

"""
    rsolve!([solver::RLSSolver,] A, b, x::AbstractVector)

Identical to `rsolve` except an initial iterate is supplied by the argument `x`. The
    function overwrites the solution into the vector `x`.
"""
function rsolve!(A, b, x::AbstractVector)
    solver = RLSSolver(10 * length(b))
    rsolve!(solver, A, b, x)

    return nothing
end

function rsolve!(solver::RLSSolver, A, b, x::AbstractVector)

    iter = 0
    log_update!(solver.log, solver.sampler, x, (), iter, A, b)

    while check_stop_criterion(solver.log, solver.stop) == false
        iter += 1
        samp = sample(solver.sampler, A, b, x, iter)
        rsubsolve!(solver.routine, x, samp, iter)
        log_update!(solver.log, solver.sampler, x, samp, iter, A, b)
    end

    solver.x = x

    return nothing
end

#############################################
# Randomized Linear Solvers
#############################################

#############################################
# Export Statements
#############################################
# See RLinearAlgebra.jl
#export RLSSolver, rsolve, rsolve!


module TestCore

using AnisotropicDiffusion
using AnisotropicDiffusion: generic_step!, hand_tuned_step!
using ImageFiltering
using LinearAlgebra
using Test

@testset "1d edge" begin
    img = zeros(20)
    img[1:10] .= 100
    img_noisy = img + randn(size(img)...)
    img_clean = denoise(img_noisy, PeronaMalik(niter=3, lambda=0.25));
    ker = centered(normalize(ones(3),1))
    img_baseline = imfilter(img_noisy, ker)
    @test sum(abs2, img - img_clean) < sum(abs2, img - img_baseline)
end

@testset "generic vs hand tuned $(dim)d" for dim in [1,2]
    alg1 = PeronaMalik(step = generic_step!, niter=1)
    alg2 = PeronaMalik(step = hand_tuned_step!, niter=1)
    img = randn([10,11,12][1:dim]...)
    ret1 = denoise(img, alg1)
    ret2 = denoise(img, alg2)
    @test ret1 ≈ ret2
    @test ret1 !== ret2
end

end#module


@add_problem mip function mip_lp_fallback_interface(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable()
    p = minimize(x, x >= 4.3; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4.3 atol = atol rtol = rtol
    end

    x = Variable(2)
    p = minimize(norm(x, 1), x[1] >= 4.3; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 4.3 atol = atol rtol = rtol
    end
end

@add_problem mip function mip_integer_variables(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(:Int)
    p = minimize(x, x >= 4.3; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol = atol rtol = rtol
    end

    x = Variable(2, :Int)
    p = minimize(sum(x), x >= 4.3; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 10 atol = atol rtol = rtol
    end

    x = Variable(:Int)
    y = Variable()
    p = minimize(sum(x + y), x >= 4.3, y >= 7; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 12 atol = atol rtol = rtol
    end

    x = Variable(2, :Int)
    p = minimize(norm(x, 1), x[1] >= 4.3; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol = atol rtol = rtol
    end

    x = Variable(2, :Int)
    p = minimize(sum(x), x[1] >= 4.3, x >= 0; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol = atol rtol = rtol
    end

    x = Variable(2, :Int)
    p = minimize(sum(x), x >= 0.5; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol = atol rtol = rtol
    end
end

@add_problem mip function mip_binary_variables(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2, :Bin)
    p = minimize(sum(x), x >= 0.5; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol = atol rtol = rtol
    end

    x = Variable(2, :Bin)
    p = minimize(sum(x), x[1] >= 0.5, x >= 0; numeric_type = T)

    if test
        @test vexity(p) == AffineVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 1 atol = atol rtol = rtol
    end
end


using Random
using NeuralVerification;

"""
    make_random_network(layer_sizes::Vector{Int}, [min_weight = -1.0], [max_weight = 1.0], [min_bias = -1.0], [max_bias = 1.0], [rng = 1.0])
    read_layer(output_dim::Int, f::IOStream, [act = ReLU()])

Generate a network with random weights and bias. The first layer is treated as the input.
The values for the weights and bias will be uniformly drawn from the range between min_weight
and max_weight and min_bias and max_bias respectively. The last layer will have an ID()
activation function and the rest will have ReLU() activation functions. Allow a random number
generator(rng) to be passed in. This allows for seeded random network generation.
"""
function make_random_network(layer_sizes::Vector{Int}, min_weight = -1.0, max_weight = 1.0, min_bias = -1.0, max_bias = 1.0, rng=MersenneTwister(0))
    # Create each layer based on the layer_size
    layers = []
    for index in 1:(length(layer_sizes)-1)
        cur_size = layer_sizes[index]
        next_size = layer_sizes[index+1]
        # Use Id activation for the last layer - otherwise use ReLU activation
        if index == (length(layer_sizes)-1)
            cur_activation = NeuralVerification.Id()
        else
            cur_activation = NeuralVerification.ReLU()
        end

        # Dimension: num_out x num_in
        cur_weights = min_weight .+ (max_weight - min_weight) * rand(rng, Float64, (next_size, cur_size))
        cur_weights = reshape(cur_weights, (next_size, cur_size)) # for edge case where 1 dimension is equal to 1 this keeps it from being a 1-d vector

        # Dimension: num_out x 1
        cur_bias = min_bias .+ (max_bias - min_bias) * rand(rng, Float64, (next_size))
        push!(layers, NeuralVerification.Layer(cur_weights, cur_bias, cur_activation))
    end

    return Network(layers)
end


module PkgTest

using BenchmarkTools, Random, Statistics, Test, OTSM
using COSMO, SCS #, MosekTools
using LinearAlgebra

@testset "cat" begin
A = Matrix{Matrix{Int}}(undef, 2, 2)
A[1, 1] = [1 1 1; 1 1 1]
A[1, 2] = [2 2; 2 2]
A[2, 1] = [3 3 3]
A[2, 2] = [3 3]
@test cat(A) == [1 1 1 2 2; 1 1 1 2 2; 3 3 3 3 3]
end

@testset "otsm_pba (port wine, MAXDIFF)" begin
A, maxdiff_optim, _ = portwine_data()
Smaxdiff = [A[i]'A[j] for i in 1:4, j in 1:4]
OTSM.verify_input_data!(Smaxdiff, 3, true) # set S[i, i] to 0
for r in 1:3, init_fun in [init_eye, init_sb, init_tb]
    @info "rank = $r, init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(Smaxdiff, r;
    verbose = true, O = init_fun(Smaxdiff, r))
    @show ts_ba
    @test ts_ba ≈ maxdiff_optim[r]
    @show test_optimality(Ô_ba, Smaxdiff, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, Smaxdiff, method = :lww)[[1, end]]
end
# bm = @benchmark tsm_ba(S, 1)
# display(bm); println()
end

@testset "otsm_pba (port wine, MAXBET)" begin
A, _, maxbet_optim = portwine_data()
Smaxbet = [A[i]'A[j] for i in 1:4, j in 1:4]
for r in 1:3, init_fun in [init_eye, init_sb, init_tb, init_lww1]
    @info "rank = $r, init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(Smaxbet, r;
    verbose = true, O = init_fun(Smaxbet, r))
    @show ts_ba
    if r == 3 && init_fun ∈ [init_eye, init_lww1]
        # r=3, init=init_eye or init_lww1 leads to an inferior local maximum
        @test ts_ba ≈ 818.7063749651282
    else
        @test ts_ba ≈ maxbet_optim[r]
    end
    # check global optimality certificate
    @show test_optimality(Ô_ba, Smaxbet, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, Smaxbet, method = :lww)[[1, end]]
end
# bm = @benchmark tsm_ba(S, 1)
# display(bm); println()
end

@testset "SDP relaxation (port wine, MAXDIFF)" begin
A, maxdiff_optim, _ = portwine_data()
Smaxdiff = [A[i]'A[j] for i in 1:4, j in 1:4]
OTSM.verify_input_data!(Smaxdiff, 3, true) # set S[i, i] to 0
for r in 1:3
    for solver in [
        #Mosek.Optimizer(LOG=0),
        COSMO.Optimizer(max_iter=5000, verbose=false), # takes long at r = 1
        SCS.Optimizer(verbose=0)
        ]
        @info "rank = $r, solver = $solver"
        @time Ô_sdp, ts_sdp, = otsm_sdp(Smaxdiff, r, solver = solver)
        @show ts_sdp
        @test ts_sdp ≈ maxdiff_optim[r]
    end
end
end

@testset "Procrustes (random data)" begin
Random.seed!(123)
n, d, m = 10, 3, 50
A, S, Ā, A_manopt = generate_procrustes_data(m, n, d)
@info "Proximal block ascent algorithm:"
for init_fun in [init_eye, init_sb, init_tb]
    @info "init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(S, d; verbose = true, O = init_fun(S, d))
    @show ts_ba
    @show test_optimality(Ô_ba, S, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, S, method = :lww)[[1, end]]
end
@info "SDP relaxation (failed to find global solution):"
for solver in [
    #Mosek.Optimizer(LOG=0),
    COSMO.Optimizer(max_iter=5000, verbose=false),
    SCS.Optimizer(verbose=0)
    ]
    @info "solver = $solver"
    @time Ô_sdp, ts_sdp, = otsm_sdp(S, d, solver = solver)
    @show ts_sdp
end
end

@testset "MaxBet (Liu-Wang-Wang Example 5.1)" begin
A = [4.3299 2.3230 -1.3711 -0.0084 -0.7414;
2.3230 3.1181 1.0959 0.1285 0.0727;
-1.3711 1.0959 6.4920 -1.9883 -0.1878;
-0.0084 0.1285 -1.9883 2.4591 1.8463;
-0.7414 0.0727 -0.1878 1.8463 5.8875]
ns    = [2, 3]
m     = length(ns)
r     = 1
start = [1; cumsum(ns)[1:end-1] .+ 1]
stop  = cumsum(ns)
S     = [A[start[i]:stop[i], start[j]:stop[j]] for i in 1:m, j in 1:m]
@info "Proximal block ascent algorithm:"
for init_fun in [init_eye, init_sb, init_tb, init_lww1]
    @info "init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(S, r; verbose = true, O = init_fun(S, r))
    @show test_optimality(Ô_ba, S, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, S, method = :lww)[[1, end]]
    if init_fun == init_eye
        # init_eye leads to a inferior local maximum
        @test abs(ts_ba - 14.10123) < 1e-3
    elseif init_fun == init_sb
        # the test matrix by Shapiro-Botha (init_sb) is singular with largest eigenvalue
        # 0 (with multiplicity 3). Different machines give different init_sb. Some of
        # them lead to inferior local maximum.
        @test abs(ts_ba - 14.10123) < 1e-3 || abs(ts_ba - 14.73022) < 1e-3
    else
        @test abs(ts_ba - 14.73022) < 1e-3
    end
end
end

@testset "MaxBet (Liu-Wang-Wang Example 5.2)" begin
A = Float64.([45 -20 5 6 16 3;
-20 77 -20 -25 -8 -21;
5 -20 74 47 18 -32;
6 -25 47 54 7 -11;
16 -8 18 7 21 -7;
3 -21 -32 -11 -7 70])
ns    = [2, 2, 2]
m     = length(ns)
r     = 1
start = [1; cumsum(ns)[1:end-1] .+ 1]
stop  = cumsum(ns)
S     = [A[start[i]:stop[i], start[j]:stop[j]] for i in 1:m, j in 1:m]
@info "Proximal block ascent algorithm:"
for init_fun in [init_eye, init_sb, init_tb, init_lww1]
    @info "init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(S, r; verbose = true, O = init_fun(S, r))
    @test abs(ts_ba - 378.9624) < 1e-3
    @show test_optimality(Ô_ba, S, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, S, method = :lww)[[1, end]]
end
end

@testset "Ten Berge 1977 p270" begin
m    = 3
d    = 3
A    = Vector{Matrix{Float64}}(undef, d)
A[1] = Matrix([Matrix(I, d, d); Matrix(I, d, d); zeros(d, d)])
A[2] = Matrix([-Matrix(I, d, d); zeros(d, d); Matrix(I, d, d)])
A[3] = Matrix([zeros(d, d); Matrix(I, d, d); Matrix(I, d, d)])
Stb  = [A[i]'A[j] for i in 1:m, j in 1:m]
OTSM.verify_input_data!(Stb, 1, true) # make it MAXDIFF
for r in 1:3, init in [init_eye, init_tb, init_sb]
    Ô_ba, ts_ba, obj_ba, history = otsm_pba(Stb, r; verbose = false, log = true, O = init(Stb, r))
    println("r = $r, init = $init, obj = $ts_ba, niters = $(length(history[:tracesum]))")
    @show test_optimality(Ô_ba, Stb)[1]
end
end

@testset "Procrustes (restrict to SO(d))" begin
# under this setting, the solution from `otsm_pba(soconstr=false)` has some Oi
# with determinant -1;
# but `otsm_pba(soconstr=true)` will enforce the solution in SO(d)
Random.seed!(123)
n, d, m = 10, 3, 50
A, S, Ā, A_manopt = generate_procrustes_data(m, n, d, 1.0)
@info "Proximal block ascent algorithm:"
for init_fun in [init_eye, init_sb, init_tb]
    @info "init = $init_fun"
    @time Ô_ba, ts_ba, = otsm_pba(S, d; verbose = true, O = init_fun(S, d))
    @test any(det.(Ô_ba) .< 0)
    @time Ô_ba, ts_ba, = otsm_pba(S, d; verbose = true, O = init_fun(S, d), soconstr=true)
    # @show det.(Ô_ba)
    @test all(det.(Ô_ba) .≥ 0)
    @show test_optimality(Ô_ba, S, method = :wzl)[[1, end]]
    @show test_optimality(Ô_ba, S, method = :lww)[[1, end]]
end
end

end # PkgTest module


module CoverRanges #src
export CompInf, min_covering, solve_min_covering #src

include("access_dict.jl") #src
export AccessDict, total_gets, total_sets #src

# The idea of the algorithm is that we can simply start from the left-most element
# of the range `r` that we wish to cover, and decide whether or not to include it
# in our partial cover. If we do want to include it, the optimal range to use to
# cover it is the one which extends the furthest to the right (out of all the
# ranges which include the element). Then it remains to solve a smaller instance
# of the problem. Likewise, if we don't include the element, we again are left
# with a smaller instance of the problem. In other words, we aim to solve the
# dynamic programming equation
#
# ```julia
# min_number_to_cover(a:b, K) = min(1 + min_number_to_cover(f:b, K - (f-a+1)), min_number_to_cover((a+1):b, K))
# ```
#
# where
# `min_number_to_cover(r, K)` is the minimum number of ranges needed to cover at
# least `K` elements of `r` and `f` is the endpoint furthest to the right of any
# range in `C` that includes `a`.
#
# Why is this the right recursive equation for this problem? Let us say we reach
# node `a`. We can either
#
# * include node `a` in our set of nodes covered, in which case we choose the
#   range that includes `a` and stretches as far as possible to the right; in that
#   case, we have to solve the problem on the remaining set of nodes (this
#   corresponds to `min_number_to_cover(f:b, K - (f-a+1))`).
# * Otherwise, we skip node `a` and need to solve the problem with the remaining
# nodes `a+1:b`,
#   still needing to cover `K` nodes (this corresponds to `min_number_to_cover(
#   (a+1):b, K))`).
#

# Let's begin the implementation.
# First, we need an object to represent an impossible task, which is never chosen in a
# minimum and which is invariant under adding `+1`. Julia's `Inf` does this for
# us, but `Inf` is a floating point number, and we'd prefer to stick to integers
# (with the exception of this infinity).
struct ComparisonInf end
const CompInf = ComparisonInf()
Base.isless(::ComparisonInf, ::ComparisonInf) = false

for T in (Int8, Int16, Int32, Int64, Int128)
    @eval Base.isless(::ComparisonInf, ::$T) = false
    @eval Base.isless(::$T, ::ComparisonInf) = true
    @eval Base.:+(::$T, c::ComparisonInf) = c
    @eval Base.:+(c::ComparisonInf, ::$T) = c
end

# In the above list, we could add other types if needed, such as
# [SaferIntegers.jl](https://github.com/JeffreySarnoff/SaferIntegers.jl)
# numeric types to ensure our algorithm does not overflow.

# Next, we define an struct to hold an instance or subinstance of our problem.
struct LeftToCover{T}
    "Range to partially cover"
    r::UnitRange{T}
    "Number of nodes we are required to cover"
    K::T
end

# And likewise, one to hold a solution (or part of one).
struct PartialCover{T}
    "Minimum number needed to partially cover the range"
    m::Union{T, ComparisonInf}
    "Index of latest range added"
    i::Int
    "Key to point towards next range in the partial cover"
    key::LeftToCover{T}
end

# Next, to find the range with the furthest right endpoint
# including a given node, we just write a simple loop.
"""
    find_furthest_range(a, C) -> index, right_endpoint

Given a node `a`, finds the range `r` in `C` such that `r` contains `a` and has
the largest right endpoint.
"""
function find_furthest_range(a, C)
    ## This could be optimized more by sorting `C` first and terminating early
    best_f_so_far = typemin(a) # right endpoint
    best_i_so_far = nothing # index
    for (i, r) in enumerate(C)
        a ∈ r || continue
        if last(r) > best_f_so_far
            best_f_so_far = last(r)
            best_i_so_far = i
        end
    end
    return best_i_so_far, best_f_so_far
end

# Now comes the heart of the algorithm. We compute
#
# ```julia
# min_number_to_cover(a:b, K) = min(1 + min_number_to_cover(f:b, K - (f-a+1)), min_number_to_cover((a+1):b, K))
# ```
#
# recursively with the help of a memoization cache, which maps from
# (sub)instances of the problem to solutions.
function min_covering(key::LeftToCover{T}, furthest_ranges, cache::AbstractDict{LeftToCover{T}, PartialCover{T}}) where {T}
    get!(cache, key) do
        r = key.r;  K = key.K
        K <= 0 && return PartialCover{T}(zero(T), 0, key) # Nothing left to cover
        K > length(r) && return PartialCover{T}(CompInf, 0, key) # Impossible task
        a, b = first(r), last(r)
        i_a, f_a = furthest_ranges[a] # index and right endpoint of optimal range including `a`

        ## Solve the subproblem in which we skip `a`
        skip_cover = min_covering(LeftToCover{T}(a+1:b, K), furthest_ranges, cache)
        i_a === nothing && return skip_cover # cannot cover `a` means we must skip it

        ## Solve the subproblem in which we include `a` in our partial cover
        add_key = LeftToCover{T}( (f_a+1) : b, K - (f_a - a + 1))
        add_cover = min_covering(add_key, furthest_ranges, cache)

        ## Decide whether or not to include `a` and return the better result
        if add_cover.m + 1 <= skip_cover.m
            return PartialCover{T}(add_cover.m + T(1), i_a, add_key)
        else
            return skip_cover
        end
    end
end

# Note that in the case that we choose to skip `a`, we return `skip_cover`, which
# is the solution for the subproblem that skips `a`. But when we choose to include
# `a` in the partial cover, we make a new `PartialCover` object instead of
# returning `add_cover`. That's because if we include `a` by using the range that
# ends at `f_a`, we really are adding a range to the partial cover represented by
# `add_cover`, so we make a new `PartialCover` object, one whose value of `m` is one
# larger (corresponding to the extra range we added), which holds the index `i_a`
# of the range that we added, and which holds `add_key`, which we can use to
# recover `add_cover` from `cache` later. In all, this lets us hold the sequence
# of ranges used in the final optimal cover efficiently (as a linked list).

# Lastly, we wrap `min_covering` in a helper function that constructs `furthest_ranges`,
# the memoization cache, and unwinds the linked list that holds the solution in order
# to recover the indices of the ranges included in the minimal partial cover we have found.
"""
    solve_min_covering(r::UnitRange{T}, K, C; [cache]) -> Union{Nothing, Vector{Int}}

Returns either `nothing` or a vector of indices of `C` corresponding to a
minimal set of ranges which covers at least `K` elements of `r`. Returns
`nothing` if and only if the task is impossible.

* `r`: range to partially cover
* `K`: number of elements of `r` that must be
covered
* `C`: vector of ranges to use in constructing the partial cover
* `cache`: optionally provide an `AbstractDict{LeftToCover{T}, PartialCover{T}}`
which is used as a memoization cache

!!! warning
    Overflow is not checked.
"""
function solve_min_covering(r::UnitRange{T}, K, C; cache =  Dict{LeftToCover{T},
PartialCover{T}}()) where {T}
    ## `furthest_ranges` is all we need from `C`
    furthest_ranges = Dict(a => find_furthest_range(a, C) for a in r)

    ## Recursive function to compute the covering
    result = min_covering(LeftToCover{T}(r, K), furthest_ranges, cache)

    k = result.m # minimum number of ranges needed
    k == CompInf && return nothing

    ## We store the list of ranges essentially as a linked list.
    ## To recover the optimal list of ranges, we traverse through
    ## the cache, at each step getting the `key` for the next element in the cache.
    range_inds = Int[]
    sizehint!(range_inds, k)
    while k > 1
        push!(range_inds, result.i)
        k = result.m
        result = cache[result.key]
    end
    return range_inds
end

end #src


using Ipopt, ForwardDiff
import Vecchia: ipopt_hessian

# For reference, compute the exact MLE. This file also loads example_setup.jl,
# which creates the simulated data and stuff.
println("Exact estimation for reference:")
include("example_exact_reference.jl")
println("\nExact MLE:\n")
display(exact_mle)
println("\n\nExact Hessian:\n")
display(exact_obs_inf_mat)
print("\n\n\n")

# Very important for multi-threading and small linear algebra. I don't set this
# in the package source, but you really should if you're going to use this.
BLAS.set_num_threads(1)

# Load in the functions for the objective (_nll), gradient (grad!), and Hessian
# (hess), which are made slightly more complicated to squeeze out redundant AD
# computations.
include("ad_nll_derivatives.jl")

# Set up and solve the problem in Ipopt.
prob = CreateIpoptProblem(
         2,ones(2).*1.0e-3, ones(2).*100.0,
         0,Float64[], Float64[], 0,3,
         _nll,  (args...)->nothing,
         grad!, (args...)->nothing,
         (x,r,c,o,l,v)->ipopt_hessian(x,r,c,o,l,v,hess,Function[],0))
AddIpoptNumOption(prob, "tol",      1.0e-5)
AddIpoptIntOption(prob, "max_iter", 1000)
prob.x = ones(2)
println("Solving problem with Vecchia approximation:")
@time status = IpoptSolve(prob)

# Here are your MLE and observed information matrix, with pretty display:
mle = copy(prob.x)
obs_inf_mat = hess(mle)
println("\nMLE:\n")
display(mle)
println("\n\nHessian:\n")
display(obs_inf_mat)
print("\n\n\n")


# using Revise
using Test
using Logging
using LarSurf
# Logging.configure(level==Logging.Debug)

# include("../src/LarSurf.jl")
# include("../src/block.jl")

@testset "Boundary matrix shape test" begin
    grid_size = [2,3,4]
    b3, larmodel = LarSurf.calculate_boundary3(grid_size)
#     println(size(b3))
#     @test size(b3) =
#     slides0 = LarSurf.random_image([7, 7, 7], [1,2,2], [3, 4, 5], 2)
#     Matrix(slides0)
#     @test all(slides0 .== [1, 5, 6, 10, 11, 15])

    # slides0 = LarSurf.data_sub_from_block_sub([5, 5, 5], 1, [1,2,3])
    # @test all(slides0 .== [0, 6, 5, 11, 10, 16])
    Cnum = prod(grid_size)
    Fnum = LarSurf.grid_number_of_faces(grid_size)
    @test size(b3, 1) == Cnum
    @test size(b3, 2) == Fnum

    V, (VV, EV, FV, CV) = larmodel
    @test size(b3, 1) == length(CV)
    @test size(b3, 2) == length(FV)

    @test size(b3) == (24, 98)
    FV = larmodel[2][3]
    # display(FV)
    # println(Fnum)

end

@testset "reagular brick boundary matrix to image file" begin
    # using ImageView
    sz = 4
    grid_size = [sz, sz, sz]
    b3, larmodel = LarSurf.calculate_boundary3(grid_size)
    barr = Matrix(b3 .> 0)
    nvoxels = sz^3
    nfaces = 3 * (sz^2 * (sz + 1) )
    @test size(barr, 1) == nvoxels
    @test size(barr, 2) == nfaces
    # imshow(barr)
end

@testset "Tests get boundary" begin
    grid_size = [2,3,4]
    LarSurf.set_param(
        boundary_allow_memory=false,
        boundary_allow_read_files=false,
        boundary_allow_write_files=false
        )
    b3a, larmodel_a = LarSurf.get_boundary3(grid_size)

    LarSurf.set_param(boundary_allow_memory=true)
    b3b, larmodel_b = LarSurf.get_boundary3(grid_size)

    @test size(b3a) == (24, 98)
    @test size(b3b) == (24, 98)

end