### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ ce85691e-fc77-11ea-0aee-df1ba52cecd7
using Pkg, DrWatson

# ╔═╡ ce85a7b2-fc77-11ea-2bec-0d7f0af3c59b
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanQuap
	using StatisticalRethinking
end

# ╔═╡ 9741f266-fc76-11ea-0689-cd04b7b16135
md"## Clip-05-10-11s.jl"

# ╔═╡ b9c3f450-fcdb-11ea-2ca9-27f5dbb73a98
md"##### The model m5.3s represents a regression of Divorce on both Marriage rate and MedianAgeMarriage and is defined as:"

# ╔═╡ aba487fe-fcdb-11ea-25e1-67c6f9596a16
md"
```
model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);         // Priors
  bA ~ normal(0, 0.5);
  bM ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bA * A + bM * M;
  D ~ normal(mu , sigma);     // Likelihood
}
```
"

# ╔═╡ c62fab76-fcdb-11ea-1b75-e71f5c77eccf
md"##### D (Divorce rate), M (Marriage rate) and A (MediumAgeMarriage) are all standardized."

# ╔═╡ ce861b20-fc77-11ea-0b79-8dd922420260
begin
	df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);
	scale!(df, [:Marriage, :MedianAgeMarriage, :Divorce])
	data = (N=size(df, 1), D=df.Divorce_s, A=df.MedianAgeMarriage_s,
		M=df.Marriage_s)
end

# ╔═╡ ce8df1ec-fc77-11ea-1c7d-ad620c19fbbd
stan5_1 = "
	data {
	 int N; // Sample size
	 vector[N] D; // Outcome
	 vector[N] A; // Predictor
	}

	parameters {
	 real a; // Intercept
	 real bA; // Slope (regression coefficients)
	 real < lower = 0 > sigma;    // Error SD
	}

	model {
	  vector[N] mu;               // mu is a vector
	  a ~ normal(0, 0.2);         // Priors
	  bA ~ normal(0, 0.5);
	  sigma ~ exponential(1);
	  mu = a + bA * A;
	  D ~ normal(mu , sigma);   // Likelihood
	}
";

# ╔═╡ 8a7c3e5c-8116-11eb-1911-b940b352dfc9
md"##### Result rethinking:"

# ╔═╡ 8a7c7662-8116-11eb-3540-1dea64d7eec6
rethinking = "
	   mean   sd  5.5% 94.5%
a      0.00 0.10 -0.16  0.16
bA    -0.57 0.11 -0.74 -0.39
sigma  0.79 0.08  0.66  0.91
";

# ╔═╡ 8a7d461c-8116-11eb-043e-ffb6bf7d0737
begin
	data = (N = size(df, 1), D = df.Divorce_s, A = df.MedianAgeMarriage_s,
		M = df.Marriage_s)
	init = (a=1.0, bA=1.0, bM=1.0, sigma=10.0)
	q5_1s, m5_1s, om5_1s = stan_quap("m5.1s", stan5_1; data, init)
	if !isnothing(q5_1s)
		quap5_1s_df = sample(q5_1s)
		PRECIS(quap5_1s_df)
	end
end

# ╔═╡ b7c70164-8116-11eb-0f7e-754d76c430ec
stan5_2 = "
data {
  int N;
  vector[N] D;
  vector[N] M;
}
parameters {
  real a;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + bM * M;
  a ~ normal( 0 , 0.2 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ normal( mu , sigma );
}
";

# ╔═╡ bdc8090a-8116-11eb-07c8-d343006e2d79
begin
	q5_2s, m5_2s, om = stan_quap("m5.2s", stan5_2; data)
	if !isnothing(q5_2s)
		quap5_2s_df = sample(q5_2s)
	end
	PRECIS(quap5_2s_df)
end

# ╔═╡ ce79d490-8116-11eb-0f1d-8d1719cf1e5b
stan5_3 = "
data {
  int N;
  vector[N] D;
  vector[N] M;
  vector[N] A;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + + bA * A + bM * M;
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ normal( mu , sigma );
}
";

# ╔═╡ e4731414-8116-11eb-1340-0f2a41529348
begin
	q5_3s, m5_3s, om5_3 = stan_quap("m5.3s", stan5_3; data, init)
	if !isnothing(q5_3s)
		quap5_3s_df = sample(q5_3s)
	end
	PRECIS(quap5_3s_df)
end

# ╔═╡ 47566d90-fcda-11ea-0662-8fa0dd0583c4
md"##### Normal estimates:"

# ╔═╡ ce8e68c0-fc77-11ea-30eb-3fcbdd292d2c
if !isnothing(m5_1s) &&!isnothing(m5_2s) && !isnothing(m5_3s)
	(s1, p1) = plot_model_coef([m5_1s, m5_2s, m5_3s], [:bA, :bM];
		title="Particles (Normal) estimates")
	p1
end

# ╔═╡ cea73280-fc77-11ea-348f-8973a4e7a5d3
md"## End of clip-05-10-11s.jl"

# ╔═╡ Cell order:
# ╟─9741f266-fc76-11ea-0689-cd04b7b16135
# ╠═ce85691e-fc77-11ea-0aee-df1ba52cecd7
# ╠═ce85a7b2-fc77-11ea-2bec-0d7f0af3c59b
# ╟─b9c3f450-fcdb-11ea-2ca9-27f5dbb73a98
# ╟─aba487fe-fcdb-11ea-25e1-67c6f9596a16
# ╟─c62fab76-fcdb-11ea-1b75-e71f5c77eccf
# ╟─ce861b20-fc77-11ea-0b79-8dd922420260
# ╠═ce8df1ec-fc77-11ea-1c7d-ad620c19fbbd
# ╟─8a7c3e5c-8116-11eb-1911-b940b352dfc9
# ╠═8a7c7662-8116-11eb-3540-1dea64d7eec6
# ╠═8a7d461c-8116-11eb-043e-ffb6bf7d0737
# ╠═b7c70164-8116-11eb-0f7e-754d76c430ec
# ╠═bdc8090a-8116-11eb-07c8-d343006e2d79
# ╠═ce79d490-8116-11eb-0f1d-8d1719cf1e5b
# ╠═e4731414-8116-11eb-1340-0f2a41529348
# ╟─47566d90-fcda-11ea-0662-8fa0dd0583c4
# ╠═ce8e68c0-fc77-11ea-30eb-3fcbdd292d2c
# ╟─cea73280-fc77-11ea-348f-8973a4e7a5d3


# Belief Propagation                #
# --------------------------------- #
@rule NormalMeanVariance(:v, Marginalisation) (m_out::PointMass, m_μ::UnivariateNormalDistributionsFamily) = begin
    m_out_mean        = mean(m_out)
    m_μ_mean, m_μ_cov = mean_cov(m_μ)

    return ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> -log(m_μ_cov + x) / 2 - 1 / (2x) * (m_out_mean - m_μ_mean) ^ 2)
end

@rule NormalMeanVariance(:v, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_μ::PointMass) = begin
    m_out_mean, m_out_cov = mean_cov(m_out)
    m_μ_mean              = mean(m_μ)

    return ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> -log(m_out_cov + x) / 2 - 1 / (2x) * (m_out_mean - m_μ_mean) ^ 2)
end

@rule NormalMeanVariance(:v, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_μ::UnivariateNormalDistributionsFamily) = begin
    m_out_mean, m_out_cov = mean_cov(m_out)
    m_μ_mean, m_μ_cov     = mean_cov(m_μ)

    return ContinuousUnivariateLogPdf(DomainSets.HalfLine(), (x) -> -log(m_out_cov + m_μ_cov + x) / 2 - 1 / (2x) * (m_out_mean - m_μ_mean) ^ 2)
end


@diffusion_process SineAux{K,R,S} begin
    :parameters
    (a, b, c, σ) --> K

    :auxiliary_info
    t0 --> Float64
    T --> Float64
    v0 --> R
    vT --> S

    :additional
    constdiff --> true
    linear --> true
end

DiffusionDefinition.B(t, P::SineAux) = 0.0 + (t/P.T)/5.0
DiffusionDefinition.β(t, P::SineAux) = (P.vT-P.v0)/P.T
DiffusionDefinition.σ(t, x, P::SineAux) = P.σ

DiffusionDefinition.default_type(::SineAux) = Float64
DiffusionDefinition.default_wiener_type(::SineAux) = Float64


"""
Scaling matrix.
"""
struct ScalMat{T<:Real} <: AbstractPDMat{T}
    dim::Int
    value::T
    inv_value::T
end

ScalMat(d::Int,v::Real) = ScalMat{typeof(inv(v))}(d, v, inv(v))

### Conversion
Base.convert(::Type{ScalMat{T}}, a::ScalMat) where {T<:Real} = ScalMat(a.dim, T(a.value), T(a.inv_value))
Base.convert(::Type{AbstractArray{T}}, a::ScalMat) where {T<:Real} = convert(ScalMat{T}, a)

### Basics

dim(a::ScalMat) = a.dim
Base.Matrix(a::ScalMat) = Matrix(Diagonal(fill(a.value, a.dim)))
LinearAlgebra.diag(a::ScalMat) = fill(a.value, a.dim)
LinearAlgebra.cholesky(a::ScalMat) = cholesky(Diagonal(fill(a.value, a.dim)))

### Arithmetics

function pdadd!(r::Matrix, a::Matrix, b::ScalMat, c)
    @check_argdims size(r) == size(a) == size(b)
    if r === a
        _adddiag!(r, b.value * c)
    else
        _adddiag!(copyto!(r, a), b.value * c)
    end
    return r
end

*(a::ScalMat, c::T) where {T<:Real} = ScalMat(a.dim, a.value * c)
/(a::ScalMat{T}, c::T) where {T<:Real} = ScalMat(a.dim, a.value / c)
*(a::ScalMat, x::StridedVecOrMat) = a.value * x
\(a::ScalMat, x::StridedVecOrMat) = a.inv_value * x
Base.kron(A::ScalMat, B::ScalMat) = ScalMat( dim(A) * dim(B), A.value * B.value )

### Algebra

Base.inv(a::ScalMat) = ScalMat(a.dim, a.inv_value, a.value)
LinearAlgebra.logdet(a::ScalMat) = a.dim * log(a.value)
LinearAlgebra.eigmax(a::ScalMat) = a.value
LinearAlgebra.eigmin(a::ScalMat) = a.value

### whiten and unwhiten

function whiten!(r::StridedVecOrMat, a::ScalMat, x::StridedVecOrMat)
    @check_argdims dim(a) == size(x, 1)
    mul!(r, x, sqrt(a.inv_value))
end

function unwhiten!(r::StridedVecOrMat, a::ScalMat, x::StridedVecOrMat)
    @check_argdims dim(a) == size(x, 1)
    mul!(r, x, sqrt(a.value))
end

### quadratic forms

quad(a::ScalMat, x::AbstractVector) = sum(abs2, x) * a.value
invquad(a::ScalMat, x::AbstractVector) = sum(abs2, x) * a.inv_value

quad!(r::AbstractArray, a::ScalMat, x::StridedMatrix) = colwise_sumsq!(r, x, a.value)
invquad!(r::AbstractArray, a::ScalMat, x::StridedMatrix) = colwise_sumsq!(r, x, a.inv_value)

### tri products

function X_A_Xt(a::ScalMat, x::StridedMatrix)
    @check_argdims dim(a) == size(x, 2)
    lmul!(a.value, x * transpose(x))
end

function Xt_A_X(a::ScalMat, x::StridedMatrix)
    @check_argdims dim(a) == size(x, 1)
    lmul!(a.value, transpose(x) * x)
end

function X_invA_Xt(a::ScalMat, x::StridedMatrix)
    @check_argdims dim(a) == size(x, 2)
    lmul!(a.inv_value, x * transpose(x))
end

function Xt_invA_X(a::ScalMat, x::StridedMatrix)
    @check_argdims dim(a) == size(x, 1)
    lmul!(a.inv_value, transpose(x) * x)
end


# Calibrate convective adjustment closure parameters to LESbrary 2-day "free_convection" simulation

using OceanTurbulenceParameterEstimation, LinearAlgebra, CairoMakie
using Oceananigans.Units

include("utils/lesbrary_paths.jl")
include("utils/one_dimensional_ensemble_model.jl")

# Build an observation from "free convection" LESbrary simulation

LESbrary_directory = "/Users/adelinehillier/Desktop/dev/2DaySuite/"

suite = OrderedDict("2d_free_convection" => (
    filename = joinpath(LESbrary_directory, "free_convection/instantaneous_statistics.jld2"),
    fields = (:b,)))

observations = SyntheticObservationsBatch(suite; first_iteration = 13, last_iteration = nothing, normalize = ZScore, Nz = 32)

closure = ConvectiveAdjustmentVerticalDiffusivity(;
    convective_κz = 1.0,
    background_κz = 1e-4
)

# Build an ensemble simulation based on observation

ensemble_size = 30
ensemble_model = OneDimensionalEnsembleModel(observations;
    architecture = CPU(),
    ensemble_size = ensemble_size,
    closure = closure
)

ensemble_simulation = Simulation(ensemble_model; Δt = 10seconds, stop_time = 4days)

# Specify priors

priors = (
    convective_κz = ConstrainedNormal(0.0, 1.0, 0.1, 1.0),
    background_κz = ConstrainedNormal(0.0, 1.0, 0.0, 10e-4)
)

free_parameters = FreeParameters(priors)

# Specify an output map that tracks 3 uniformly spaced time steps, ignoring the initial condition
track_times = floor.(range(1, stop = length(observations[1].times), length = 3))
popfirst!(track_times)
output_map = ConcatenatedOutputMap(track_times)

# Build `InverseProblem`
calibration = InverseProblem(observations, ensemble_simulation, free_parameters; output_map = output_map)

# Ensemble Kalman Inversion

eki = EnsembleKalmanInversion(calibration; noise_covariance = 0.01)

iterations = 10
iterate!(eki; iterations = iterations)

# Visualize the outputs of EKI calibration. Plots will be stored in `directory`.

directory = "calibrate_convadj_to_lesbrary/$(iterations)_iters_$(ensemble_size)_particles/"
isdir(directory) || mkpath(directory)

### Parameter convergence plot

# Vector of NamedTuples, ensemble mean at each iteration
ensemble_means = getproperty.(eki.iteration_summaries, :ensemble_mean)

# N_param x N_iter matrix, ensemble covariance at each iteration
θθ_std_arr = sqrt.(hcat(diag.(getproperty.(eki.iteration_summaries, :ensemble_cov))...))

N_param, N_iter = size(θθ_std_arr)
iter_range = 0:(N_iter-1)
pnames = calibration.free_parameters.names

n_cols = 3
n_rows = Int(ceil(N_param / n_cols))
ax_coords = [(i, j) for i = 1:n_rows, j = 1:n_cols]

f = Figure(resolution = (500n_cols, 200n_rows))
for (i, pname) in enumerate(pnames)
    coords = ax_coords[i]
    ax = Axis(f[coords...],
        xlabel = "Iteration",
        xticks = iter_range,
        ylabel = string(pname))

    ax.ylabelsize = 20

    mean_values = [getproperty.(ensemble_means, pname)...]
    lines!(ax, iter_range, mean_values)
    band!(ax, iter_range, mean_values .+ θθ_std_arr[i, :], mean_values .- θθ_std_arr[i, :])
end

save(joinpath(directory, "conv_adj_to_LESbrary_parameter_convergence.pdf"), f);

### Pairwise ensemble plots

N_param, N_iter = size(θθ_std_arr)
for pname1 in pnames, pname2 in pnames
    if pname1 != pname2

        f = Figure()
        axtop = Axis(f[1, 1])
        axmain = Axis(f[2, 1],
            xlabel = string(pname1),
            ylabel = string(pname2)
        )
        axright = Axis(f[2, 2])
        scatters = []
        for iteration in [0, 1, 2, N_iter - 1]
            ensemble = eki.iteration_summaries[iteration].parameters
            ensemble = [[particle[pname1], particle[pname2]] for particle in ensemble]
            ensemble = transpose(hcat(ensemble...)) # N_ensemble x 2
            push!(scatters, scatter!(axmain, ensemble))
            density!(axtop, ensemble[:, 1])
            density!(axright, ensemble[:, 2], direction = :y)
        end
        colsize!(f.layout, 1, Fixed(300))
        colsize!(f.layout, 2, Fixed(200))
        rowsize!(f.layout, 1, Fixed(200))
        rowsize!(f.layout, 2, Fixed(300))
        Legend(f[1, 2], scatters,
            ["Initial ensemble", "Iteration 1", "Iteration 2", "Iteration $N_iter"],
            position = :lb)
        hidedecorations!(axtop, grid = false)
        hidedecorations!(axright, grid = false)
        xlims!(axright, 0, 10)
        ylims!(axtop, 0, 10)
        save(joinpath(directory, "conv_adj_to_LESbrary_eki_$(pname1)_$(pname2).pdf"), f)
    end
end

# Compare EKI result to true values

y = observation_map(calibration)
output_distances = [mapslices(norm, (forward_map(calibration, [ensemble_means...])[:, 1:N_iter] .- y), dims = 1)...]

f = Figure()
lines(f[1, 1], iter_range, output_distances, color = :blue, linewidth = 2,
    axis = (title = "Output distance",
        xlabel = "Iteration",
        ylabel = "|G(θ̅ₙ) - y|",
        yscale = log10))
save(joinpath(directory, "conv_adj_to_LESbrary_error_convergence_summary.pdf"), f);

include("./utils/visualize_profile_predictions.jl")
visualize!(calibration, ensemble_means[end];
    field_names = (:b,),
    directory = directory,
    filename = "realizations.pdf"
)

θglobalmin = NamedTuple((:convective_κz => 0.275, :background_κz => 0.000275))
visualize!(calibration, θglobalmin;
    field_names = (:b,),
    directory = directory,
    filename = "realizations_θglobalmin.pdf"
)

## Visualize loss landscape

name = "Loss landscape"

pvalues = Dict(
    :convective_κz => collect(0.075:0.025:1.025),
    :background_κz => collect(0e-4:0.25e-4:10e-4),
)

ni = length(pvalues[:convective_κz])
nj = length(pvalues[:background_κz])

params = hcat([[pvalues[:convective_κz][i], pvalues[:background_κz][j]] for i = 1:ni, j = 1:nj]...)
xc = params[1, :]
yc = params[2, :]

# build an `InverseProblem` that can accommodate `ni*nj` ensemble members
ensemble_model = OneDimensionalEnsembleModel(observations;
    architecture = CPU(),
    ensemble_size = ni * nj,
    closure = closure)
ensemble_simulation = Simulation(ensemble_model; Δt = 10seconds, stop_time = 6days)
calibration = InverseProblem(observations, ensemble_simulation, free_parameters)

y = observation_map(calibration)

using FileIO
# a = forward_map(calibration, params) .- y
# save("calibrate_convadj_to_lesbrary/loss_landscape.jld2", "a", a)

a = load("calibrate_convadj_to_lesbrary/loss_landscape.jld2")["a"]
zc = [mapslices(norm, a, dims = 1)...]

# 2D contour plot with EKI particles superimposed
begin
    f = Figure()
    ax1 = Axis(f[1, 1],
        title = "EKI Particle Traversal Over Loss Landscape",
        xlabel = "convective_κz",
        ylabel = "background_κz")

    co = CairoMakie.contourf!(ax1, xc, yc, zc, levels = 50, colormap = :default)

    cvt(iter) = hcat(collect.(eki.iteration_summaries[iter].parameters)...)
    diffc = cvt(2) .- cvt(1)
    diff_mag = mapslices(norm, diffc, dims = 1)
    # diffc ./= 2
    us = diffc[1, :]
    vs = diffc[2, :]
    xs = cvt(1)[1, :]
    ys = cvt(1)[2, :]

    arrows!(xs, ys, us, vs, arrowsize = 10, lengthscale = 0.3,
        arrowcolor = :yellow, linecolor = :yellow)

    am = argmin(zc)
    minimizing_params = [xc[am] yc[am]]

    scatters = [scatter!(ax1, minimizing_params, marker = :x, markersize = 30)]
    for (i, iteration) in enumerate([1, 2, iterations])
        ensemble = eki.iteration_summaries[iteration].parameters
        ensemble = [[particle[:convective_κz], particle[:background_κz]] for particle in ensemble]
        ensemble = transpose(hcat(ensemble...)) # N_ensemble x 2
        push!(scatters, scatter!(ax1, ensemble))
    end
    Legend(f[1, 2], scatters,
        ["Global minimum", "Initial ensemble", "Iteration 1", "Iteration $(iterations)"],
        position = :lb)

    save(joinpath(directory, "loss_contour.pdf"), f)
end

# 3D loss landscape
begin
    f = Figure()
    ax1 = Axis3(f[1, 1],
        title = "Loss Landscape",
        xlabel = "convective_κz",
        ylabel = "background_κz",
        zlabel = "MSE loss"
    )

    CairoMakie.surface!(ax1, xc, yc, zc, colorscheme = :thermal)

    save(joinpath(directory, "loss_landscape.png"), f)
end


using SpatialFields
using BenchmarkTools
import StaticArrays: SVector

suite = BenchmarkGroup()

let
    suite["hrbf"] = BenchmarkGroup(["hrbf"])
    points = [rand(SVector{3, Float64}) for i in 1:100]
    normals = [rand(SVector{3, Float64}) for i in 1:100];
    suite["hrbf"]["construct"] = @benchmarkable field = HermiteRadialField($points, $normals)

    field = HermiteRadialField(points, normals)
    x = zeros(SVector{3, Float64})
    suite["hrbf"]["evaluate"] = @benchmarkable ($field)($x)
end

let
    suite["interpolating"] = BenchmarkGroup(["interpolating"])
    edge_points = [rand(SVector{3, Float64}) for i in 1:50]
    interior_points = [rand(SVector{3, Float64}) for i in 1:50]
    points = vcat(edge_points, interior_points)
    values = vcat([0.0 for i in edge_points], [-1.0 for i in interior_points])

    suite["interpolating"]["construct"] = @benchmarkable surface = InterpolatingSurface($points, $values)
    surface = InterpolatingSurface(points, values)
    x = zeros(SVector{3, Float64})
    suite["interpolating"]["evaluate_svector"] = @benchmarkable ($surface)($x)

    x = zeros(3)
    suite["interpolating"]["evaluate_vector"] = @benchmarkable ($surface)($x)
end

tune!(suite)
results = run(suite, verbose=false)
showall(results)
println()


#=
using Revise
=#
module HMC

using Flux
using Distributions

"""
Get random momentum from a TrackedArray parameter.
"""
rand_momentum(x::TrackedArray) = randn(size(x))

"""
Compute kinetic energy given momentum
"""
function compute_kinetic_energy(momentum)
  return sum([sum(v .^ 2) for v in momentum]) / 2.0
end

# Given an arbitrary type, get the Flux parameters
function get_params(s::S) where S
  ps = []
  for key in fieldnames(S)
    if isdefined(s, key)
      f = getfield(s, key)
      if Tracker.istracked(f)
        append!(ps, [f])
      end
    else
      ErrorException("Field $key is undefined in $(S.name)")
    end
  end
  return Flux.Params(ps)
end

"""
See p. 14 of
https://arxiv.org/pdf/1206.1901.pdf
"""
function hmc_update(curr_state::S, log_prob::Function,
                    num_leapfrog_steps::Int, eps::Float64;
                    momentum_sd::Real=1) where S
  # TODO: Add an invM argument. See BDA3 p. 301 - 302.

  state = deepcopy(curr_state)

  # Get tracked parameters
  qs = get_params(state)

  # Create a dictionary state -> momentum
  ps = [rand_momentum(q) * momentum_sd for q in qs]

  # Current kinetic energy
  curr_K = compute_kinetic_energy(values(ps))

  # Potential enerdy
  U(s) = -log_prob(s)

  # Current potential energy
  curr_U = Tracker.data(U(state))

  # Compute gradient of U
  grad_U(qs) = Tracker.gradient(() -> U(state), qs)

  # Make a half step for momentum at the beginning
  gs = grad_U(qs)
  for (q, p) in zip(qs, ps)
    # p .-= eps * Tracker.data(gs[q]) / 2.0
    Flux.Tracker.update!(p, -eps * gs[q] / 2)
  end

  # Alternate full steps for position and momentum
  for i in 1:num_leapfrog_steps
    # Make a full step for the position
    for (q, p) in zip(qs, ps)
      Flux.Tracker.update!(q, eps * p)
    end

    # Make a full step for the momentum, except at end of trajectory
    if i < num_leapfrog_steps
      gs = grad_U(qs)
      for (q, p) in zip(qs, ps)
        # p .-= eps * Tracker.data(gs[q])
        Flux.Tracker.update!(p, -eps * gs[q])
      end
    end
  end

  # Make a half step for momentum at the end
  gs = grad_U(qs)
  for (q, p) in zip(qs, ps)
    # p .-= eps * Tracker.data(gs[q]) / 2.0
    Flux.Tracker.update!(p, -eps * gs[q] / 2)
  end

  # Proposed potential energy
  cand_U = Tracker.data(U(state))

  # Proposed kinetic  energy
  cand_K = Tracker.data(compute_kinetic_energy(values(ps)))

  ### Metropolis step ###
  # log acceptance ratio.  NOTE: potential = -log_prob
  log_acceptance_ratio = curr_U + curr_K - cand_U - cand_K

  # Accept or reject
  if log_acceptance_ratio > log(rand())
    return state, -cand_U
  else
    return curr_state, -curr_U
  end
  ### End of Metropolis step ###
end

end # module HMC


"""

```
velocity_distribution(
  velfile::String,l
  vmax::Float64=10.,
  nbins::Int=10,
  T::DataType=Point2D
)
```

Computes the distribution of velocities. If `T=Point3D` the system will be read as three-dimensional.

"""
function velocity_distribution(
  sys::System{T},
  velfile::String;
  vmax::Float64=2.,
  nbins::Int=100
) where T

  # Initialize histogram
  @unpack n = sys
  Δv = vmax/nbins
  vvals = vmax*[ (i-1/2)*Δv for i in 1:nbins ]
  vcount = zeros(nbins)
  vavg = zero(T)
  vavg_norm = 0.
  kavg = 0.

  # Reading velocities, computing distances, and adding to histogram
  nframes = 0
  open(velfile) do file
    while !(eof(file))
      readline(file) # number of atoms, ignore
      readline(file) # title, ignore
      nframes += 1
      # read atomic coordinates
      for i in 1:n
        atom_data = split(readline(file))
        v = T(ntuple(i -> parse(Float64,atom_data[i+1]), dim(T)))
        ibin = floor(Int,(norm(v)/vmax)/Δv) + 1
        if ibin <= nbins
          vcount[ibin] += 1
        end
        vavg += v
        vavg_norm += norm(v)
        kavg += norm2(v)/2
      end
    end
  end
  println("Number of frames read = ", nframes)

  # Normalizing by the number of frames and number of atoms
  for i in 1:nbins
    vcount[i] = vcount[i] / (nframes*n*Δv)
  end
  vavg = vavg / (nframes*n)
  println(" Average velocity = ", vavg)
  println(" Average velocity norm = ", vavg_norm/(nframes*n))
  println(" Average kinetic energy = ", kavg/(nframes*n))

  return vvals, vcount
end


using Flux
using Flux.Data: DataLoader
using BSON: @save, @load # for data saving
using ProgressMeter
using LinearAlgebra
using Dates

pdename = ARGS[1]
@info "Learning Green function of $pdename equation"

cp = "checkpoints_$(pdename)_$(now())"
mkdir(cp)
@info "$cp folder created"

# Load mat file
using MAT
file = matopen("greenlearning/examples/datasets/$(pdename).mat")
varnames = collect(names(file))
for varname in varnames
    val = read(file, varname)
    @eval $(Symbol(varname)) = $val
end
close(file)
X = reshape(X, 100)
Y = reshape(Y, 200)
xmin, xmax = X[begin], X[end]
ymin, ymax = Y[begin], Y[end]
dx = X[2]-X[1]
dy = Y[2]-Y[1]
Nu, Nf = length(X), length(Y)
@info "Train data loaded"

hidden_size = 50
G_network = Chain(
    Dense(2, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, 1)
)
U_hom_network = Chain(
    Dense(1, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, hidden_size, relu),
    Dense(hidden_size, 1)
)
parameters = Flux.params(G_network, U_hom_network)
@info "Model" G_network U_hom_network

# integrator weights
trapezoidal_x = ones(Nu); trapezoidal_x[[begin,end]] .= 0.5
trapezoidal_y = ones(Nf); trapezoidal_y[[begin,end]] .= 0.5

function loss(U_train, F_train)
    l = 0.0
    U_hom_predict = [U_hom_network([x])[1] for x in X]
    # G_predict = [G_network([x,y])[1] for x in X, y in Y]
    G_predict = [[G_network([x,y])[1] for x in X] for y in Y]
    G_predict = hcat(G_predict...)
    num = size(U_train, 2)
    for n in 1:num
        @views u = U_train[:,n] # 100×1
        @views f = F_train[:,n] # 200×1
        Gf = dy*G_predict*(trapezoidal_y.*f)
        integ = (u-U_hom_predict-Gf).^2
        loss_u = dx*dot(integ, trapezoidal_x)
        u_L2 = dx*dot(u.^2, trapezoidal_x)
        l += loss_u/u_L2
    end
    return l/num
end

opt = ADAM()
@info "Optimiser" opt

batch_size = 100
train_loader = DataLoader((U, F), batchsize=batch_size, shuffle=true)
data_num = size(U, 2)
epochs = parse(Int, ARGS[2])
digit_epochs = length(string(epochs))
iteration_num = Int(epochs*batch_size/data_num)
@info "Training start" data_num epochs batch_size iteration_num

start = now()

learning = Dict([
    ("Equation name", pdename),
    ("Total epochs", epochs),
    ("Batch size", batch_size),
    ("Number of data", data_num),
    ("Current epoch number", 0),
    ("Current loss value", []),
    ("Start time", start),
    ("Time", 0)
])

skip = 5
counter = 0
@showprogress for epoch in 1:epochs
    for (U_train, F_train) in train_loader
        gs = gradient(()->loss(U_train, F_train), parameters)
        Flux.Optimise.update!(opt, parameters, gs)
    end
    global counter += 1
    if counter % skip == 0
        d = now()
        learning["Current epoch number"] = epoch
        learning["Time"] = d
        current_loss = loss(U, F)
        push!(learning["Current loss value"], current_loss)
        fname1 = "$cp/$(pdename)_description_$(lpad(epoch, digit_epochs, "0")).bson"
        fname2 = "$cp/$(pdename)_model_$(lpad(epoch, digit_epochs, "0")).bson"
        @save fname1 learning
        @save fname2 G_network U_hom_network opt
    end
end

d = now()
learning["Current epoch number"] = epochs
learning["Time"] = d
fname1 = "$cp/$(pdename)_description_$epochs.bson"
fname2 = "$cp/$(pdename)_model_$epochs.bson"
@save fname1 learning
@save fname2 G_network U_hom_network opt
@info "Training ended\nModel saved to $fname2"


using KissSmoothing
using Test
using Random

@testset "Goodness of Smoothing" begin
    rng = Random.MersenneTwister(1337)
    for N in [100, 200, 400, 800]
        for α in [0.01, 0.05, 0.1]
            x = LinRange(0, 1, N)
            y = identity.(x)
            n = randn(rng, length(x)) .* α
            yr = y .+ n
            ys, yn = denoise(yr)
            c1 = sum(abs2, yr .- y)
            c2 = sum(abs2, ys .- y)
            @test c2 < c1
            s1 = sum(abs2, n)
            s2 = sum(abs2, yn)
            @test abs(log2(s2 / s1)) < 1
        end
    end
end

@testset "Too few points" begin
    rng = Random.MersenneTwister(1337)
    O = [1.0, 2.0, 10.0]
    S, N = denoise(O)
    @test all(O .== S)
    @test all(N .== 0)
    @test length(S) == 3
    @test length(N) == 3
end

@testset "Infinite smoothing" begin
    O = sign.(sin.(1:1000)) .+ 1
    S, N = denoise(O, factor = Inf)
    @test all(S .≈ 1)
    @test all(abs.(N) .≈ 1)
end

@testset "Verbose" begin
    O = sign.(sin.(1:1000)) .+ 1
    S, N = denoise(O, factor = 0.0, verbose = true)
    @test all(abs.(S .- O) .< 1e-10)
end

@testset "Fit 1D RBF" begin
    for μ in LinRange(-100,100,5)
        t = LinRange(0,2pi,150)
        y = sin.(t).+μ
        fn = fit_rbf(t,y,LinRange(0,2pi,50))
        pred_y = fn(t)
        error = sqrt(sum(abs2, pred_y .- y)/length(t))
        @test error < 0.0002
    end
end


#=
# 3D heat eq. Dirichlet bc, CG
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("heat3d");

# Try making an optional log
@useLog("heat3dlog")

n = 6;
ord = 2;

# Set up the configuration (order doesn't matter)
@domain(3, SQUARE, UNIFORM_GRID)    # dimension, geometry, decomposition
@solver(CG)                         # DG, CG, etc.
@functionSpace(LEGENDRE, ord)         # function, order (or use testFunction and trialFunction)
@nodes(LOBATTO)                     # elemental node arrangement
@stepper(EULER_IMPLICIT)            # time stepper (optional second arg is CFL#)

# Specify the problem
@mesh(HEXMESH, n)                   # .msh file or generate our own

@variable(u)                        # same as @variable(u, SCALAR)

@testSymbol(v)                    # sets the symbol for a test function

T = 1;
@timeInterval(T)                    # (start, end) using this sets problem to time dependent
@initial(u, "abs(x-0.5)+abs(y-0.5)+abs(z-0.5) < 0.2 ? 1 : 0")  # initial condition needed if time dependent

@boundary(u, 1, DIRICHLET, 0)

# Write the weak form
@coefficient(f, "-0.1*sin(pi*x)*sin(pi*y)*sin(pi*z)")
@weakForm(u, "Dt(u*v) + 0.01 * dot(grad(u),grad(v)) - f*v")

solve(u);

# solution is stored in the variable's "values"
#using Plots
#pyplot();
#N = n*ord+1;
#half = Int(round(N/2));
#range = (N*N*half+1):(N*N*(half+1));
#display(plot(Finch.grid_data.allnodes[1,range], Finch.grid_data.allnodes[2,range], u.values[range], st=:surface))

# check
log_dump_config(Finch.config);
log_dump_prob(Finch.prob);

@finalize()


module BayesTesting

using Distributions, KernelDensity, LinearAlgebra, StatsBase, Plots, StatsPlots
# using PlotRecipes
export update_mean, marginal_posterior_mu, mcodds, todds, bayespval, pdr_pval
export blinreg, bayesregNIG, bayesreg, gsreg, compare_means, compare_proportions
export beta_posterior, beta_update, equiv_test, Bayesian_ttest, correlation_ttest
# package code goes here
include("bayes_meta_analysis.jl")
include("bayes_odds_pvalue.jl")
include("bayesreg.jl")
include("gsreg.jl")
include("compare_means.jl")
include("bayesian_ttest.jl")
include("correlation_coefficient_v1.jl")
end # module


immutable DifferentiableFunction{F <: Function, G <: Function} <: MathProgBase.AbstractNLPEvaluator
    func::F
    gradient!::G
end

DifferentiableFunction(f::Function) = DifferentiableFunction(f, (result, x) -> ForwardDiff.gradient!(result, f, x))

function MathProgBase.initialize(gd::DifferentiableFunction, requested_features::Vector{Symbol})
end

MathProgBase.features_available(gd::DifferentiableFunction) = [:Grad]

MathProgBase.jac_structure(gd::DifferentiableFunction) = (tuple(), tuple())

MathProgBase.eval_grad_f(gd::DifferentiableFunction, g, x) = gd.gradient!(g, x)

MathProgBase.eval_f(gd::DifferentiableFunction, x) = gd.func(x)

function getmodel(problem::DifferentiableFunction,
                num_vars::Integer,
                solver=MathProgBase.defaultNLPsolver,
                lower_bound=[-Inf for i in 1:num_vars],
                upper_bound=[Inf for i in 1:num_vars])
    model = MathProgBase.NonlinearModel(solver)
    MathProgBase.loadproblem!(model, num_vars, 0, lower_bound, upper_bound, Float64[], Float64[], :Min, problem)
    model
end

getmodel(f::Function, args...) = getmodel(DifferentiableFunction(f), args...)


#
# Functions and types to interpolate ephemeris data with
# a Cubic Hermite interpolator.
#
# This functionality is ~entirely~ enabled by Interpolations.jl!
#

export ContinuousEphemeris

"""
An abstract supertype for all possible interpolations for ephemeris data.
"""
abstract type AbstractEphemerisInterpolation end

"""
A callable structure which holds ephemeris data for a body's positions and velocities between the timepoints specified by `T`.
Call the object with a Julian day input to interpolate the body's positions and velocities using `Interpolations.CubicHermite`.
"""
struct ContinuousEphemeris{D<:Number, TX<:Interpolations.CubicHermite, TY<:Interpolations.CubicHermite, TZ<:Interpolations.CubicHermite} <:AbstractEphemerisInterpolation
    T::Tuple{D,D}
    X::TX
    Y::TY
    Z::TZ
end

"""
Construct a `ContinuousEphemeris` instance from an appropriately formatted
`DataFrame`. This is mostly an internal package method. See `ephemeris(args...; continuous = Val{true}, kwargs...)`
for a more user friendly constructor!
"""
function ContinuousEphemeris(data::DataFrame)

    T = eltype(data.dⱼ)
    timespan = (
        data.dⱼ[1],
        data.dⱼ[end],
    )

    X = Interpolations.CubicHermite(data.dⱼ, data.x, data.ẋ)
    Y = Interpolations.CubicHermite(data.dⱼ, data.y, data.ẏ)
    Z = Interpolations.CubicHermite(data.dⱼ, data.z, data.ż)

    return ContinuousEphemeris{T, typeof(X), typeof(Y), typeof(Z)}(
        timespan, X, Y, Z
    )
end

"""
Interpolate the body's position and velocity by a Julian day input.
"""
function (ephemeris::ContinuousEphemeris)(t::Number; bias = false, type = nothing)
    if bias
        time = t + ephemeris.T[1]
    else
        time = t
    end
    if isnothing(type)
        return vcat(
            ephemeris.X(time),
            ephemeris.Y(time),
            ephemeris.Z(time),
            Interpolations.gradient(ephemeris.X, time),
            Interpolations.gradient(ephemeris.Y, time),
            Interpolations.gradient(ephemeris.Z, time),
        )
    else
        return (
            ephemeris.X(time),
            ephemeris.Y(time),
            ephemeris.Z(time),
            Interpolations.gradient(ephemeris.X, time),
            Interpolations.gradient(ephemeris.Y, time),
            Interpolations.gradient(ephemeris.Z, time),
        ) |> type
    end
end

function Base.show(io::IO, ephemeris::ContinuousEphemeris)
    print(io, "Cubic Hermite interpolation of ephemeris data")
end


"""
Run an intcode program: 1 adds, 2 multiplies, 99 halts
prog: the program to run as an array of integers, will be modified in place
"""
function run_intcode!(prog)
    pc = 1
    while prog[pc] != 99
        # Remember to add 1 for Julia 1-based indices!
        arg1_index = prog[pc+1] + 1
        arg2_index = prog[pc+2] + 1
        res_index = prog[pc+3] + 1
        if prog[pc] == 1
            # 1 means add
            prog[res_index] = prog[arg1_index] + prog[arg2_index]
        elseif prog[pc] == 2
            # 2 means mul
            prog[res_index] = prog[arg1_index] * prog[arg2_index]
        end
        pc += 4
    end
end

# Part 1: load input and patch it, then say what was in first element after running
program = open(f -> parse.(Int32, split(readline(f), ',')), "day02_input.txt")
patched = copy(program)
patched[2] = 12
patched[3] = 2
run_intcode!(patched)
println("First element: $(patched[1])")

# Part 2: do a stupid grid search to find "noun" and "verb"
for noun in 0:99, verb in 0:99
    test = copy(program)
    test[2] = noun
    test[3] = verb
    run_intcode!(test)
    if test[1] == 19690720
        # This is the one!
        println("Noun=$noun, verb=$verb, 100*noun+verb=$(100*noun + verb)")
    end
end


"""
    ZeroDGrid()

Constructs a placeholder grid object for "0D" problems (in other words, systems of ODEs).
"""
struct ZeroDGrid <: AbstractGrid end

"""
    OneDGrid(nx, Lx; x0=-Lx/2, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE)

Constrcut a OneDGrid object with size `Lx`, resolution `nx`, and leftmost
position `x0`. FFT plans are generated for `nthreads` CPUs using
FFTW flag `effort`.
"""
struct OneDGrid{T} <: AbstractOneDGrid
  nx::Int
  nk::Int
  nkr::Int
  Lx::T
  dx::T
  x::Array{T,1}
  k::Array{T,1}
  kr::Array{T,1}
  invksq::Array{T,1}
  invkrsq::Array{T,1}

  fftplan::FFTW.cFFTWPlan{Complex{T},-1,false,1}
  rfftplan::FFTW.rFFTWPlan{T,-1,false,1}

  # Range objects that access the aliased part of the wavenumber range
  kalias::UnitRange{Int}
  kralias::UnitRange{Int}
end

function OneDGrid(nx, Lx; x0=-0.5*Lx, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE)
  T = typeof(Lx)
  dx = Lx/nx
  x = Array{T}(range(x0, step=dx, length=nx))

  nk = nx
  nkr = Int(nx/2+1)

  i1 = 0:Int(nx/2)
  i2 = Int(-nx/2+1):-1
  k = Array{T}(2π/Lx*cat(i1, i2; dims = 1))
  kr = Array{T}(2π/Lx*cat(i1; dims = 1))

  invksq = @. 1/k^2
  invksq[1] = 0
  invkrsq = @. 1/kr^2
  invkrsq[1] = 0

  FFTW.set_num_threads(nthreads)
  fftplan   = plan_fft(Array{Complex{T},1}(undef, nx); flags=effort)
  rfftplan  = plan_rfft(Array{T,1}(undef, nx); flags=effort)

  # Index endpoints for aliased i, j wavenumbers
  iaL, iaR = Int(floor(nk/3))+1, 2*Int(ceil(nk/3))-1
  kalias  = iaL:iaR
  kralias = iaL:nkr

  OneDGrid(nx, nk, nkr, Lx, dx, x, k, kr, invksq, invkrsq,
           fftplan, rfftplan, kalias, kralias)
end

"""
    TwoDGrid(nx, Lx)
    TwoDGrid(nx, Lx, ny, Ly; x0=-Lx/2, y0=-Ly/2, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE)

Constrcut a TwoDGrid object. The two-dimensional domain has size (Lx, Ly),
resolution (nx, ny) and bottom left corner at (x0, y0). FFT plans are generated
which use nthreads threads with the specified planning effort.
"""
struct TwoDGrid{T} <: AbstractTwoDGrid
  nx::Int
  ny::Int
  nk::Int
  nl::Int
  nkr::Int
  Lx::T
  Ly::T
  dx::T
  dy::T
  x::Array{T,2}
  y::Array{T,2}
  X::Array{T,2}
  Y::Array{T,2}
  k::Array{T,2}
  l::Array{T,2}
  kr::Array{T,2}
  K::Array{T,2}
  L::Array{T,2}
  Kr::Array{T,2}
  Lr::Array{T,2}
  KKsq::Array{T,2}      # K^2 + L^2
  invKKsq::Array{T,2}   # 1/KKsq, invKKsq[1, 1]=0
  KKrsq::Array{T,2}     # Kr^2 + Lr^2
  invKKrsq::Array{T,2}  # 1/KKrsq, invKKrsq[1, 1]=0

  fftplan::FFTW.cFFTWPlan{Complex{T},-1,false,2}
  rfftplan::FFTW.rFFTWPlan{T,-1,false,2}

  # Range objects that access the aliased part of the wavenumber range
  kalias::UnitRange{Int}
  kralias::UnitRange{Int}
  lalias::UnitRange{Int}
end

function TwoDGrid(nx, Lx, ny=nx, Ly=Lx; x0=-0.5*Lx, y0=-0.5*Ly, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE,
                  T=Float64)
  dx = Lx/nx
  dy = Ly/ny
  nk = nx
  nl = ny
  nkr = Int(nx/2+1)

  # Physical grid
  x = Array{T}(reshape(range(x0, step=dx, length=nx), (nx, 1)))
  y = Array{T}(reshape(range(y0, step=dy, length=ny), (1, ny)))
  X = [ x[i] for i = 1:nx, j = 1:ny]
  Y = [ y[j] for i = 1:nx, j = 1:ny]

  # Wavenubmer grid
  i1 = 0:Int(nx/2)
  i2 = Int(-nx/2+1):-1
  j1 = 0:Int(ny/2)
  j2 = Int(-ny/2+1):-1

  k  = reshape(2π/Lx*cat(i1, i2, dims=1), (nk, 1))
  kr = reshape(2π/Lx*cat(i1, dims=1), (nkr, 1))
  l  = reshape(2π/Ly*cat(j1, j2, dims=1), (1, nl))

  K = [ k[i] for i = 1:nk, j = 1:nl]
  L = [ l[j] for i = 1:nk, j = 1:nl]
  Kr = [ kr[i] for i = 1:nkr, j = 1:nl]
  Lr = [ l[j]  for i = 1:nkr, j = 1:nl]

  KKsq  = @. k^2 + l^2
  invKKsq = 1 ./ KKsq
  invKKsq[1, 1] = 0

  KKrsq = @. kr^2 + l^2
  invKKrsq = 1 ./ KKrsq
  invKKrsq[1, 1] = 0

  # FFT plans
  FFTW.set_num_threads(nthreads)
  fftplan   = plan_fft(Array{Complex{T},2}(undef, nx, ny); flags=effort)
  rfftplan  = plan_rfft(Array{T,2}(undef, nx, ny); flags=effort)

  # Index endpoints for aliased i, j wavenumbers
  iaL, iaR = Int(floor(nk/3))+1, 2*Int(ceil(nk/3))-1
  jaL, jaR = Int(floor(nl/3))+1, 2*Int(ceil(nl/3))-1

  kalias  = iaL:iaR
  kralias = iaL:nkr
  lalias  = jaL:jaR

  TwoDGrid(nx, ny, nk, nl, nkr, Lx, Ly, dx, dy, x, y, X, Y,
           k, l, kr, K, L, Kr, Lr, KKsq, invKKsq, KKrsq, invKKrsq,
           fftplan, rfftplan, kalias, kralias, lalias)
end

"""
    gridpoints(g)

Returns the collocation points of the grid `g` in 2D arrays `X, Y`.
"""
function gridpoints(g)
  X = [ g.x[i] for i=1:g.nx, j=1:g.ny]
  Y = [ g.y[j] for i=1:g.nx, j=1:g.ny]
  X, Y
end

"""
    dealias!(a, g, kalias)

Dealias `a` on the grid `g` with aliased x-wavenumbers `kalias`.
"""
function dealias!(a, g::OneDGrid)
  kalias = size(a, 1) == g.nkr ? g.kralias : g.kalias
  dealias!(a, g, kalias)
  nothing
end

function dealias!(a, g::OneDGrid, kalias)
  @views @. a[kalias, :] = 0
  nothing
end

function dealias!(a, g::TwoDGrid)
  kalias = size(a, 1) == g.nkr ? g.kralias : g.kalias
  dealias!(a, g, kalias)
  nothing
end

function dealias!(a, g::TwoDGrid, kalias)
  @views @. a[kalias, g.lalias, :] = 0
  nothing
end

"""
    makefilter(K; order=4, innerK=0.65, outerK=1)

Returns a filter acting on the non-dimensional wavenumber K that decays exponentially
for K>innerK, thus removing high-wavenumber content from a spectrum it is multiplied with.
The decay rate is determined by order and outerK determines the outer wavenumber at which
the filter is smaller than machine precision.
"""
function makefilter(K; order=4, innerK=0.65, outerK=1)
  TK = typeof(K)
  K = Array(K)
  decay = 15*log(10) / (outerK-innerK)^order # decay rate for filtering function
  filt = @. exp( -decay*(K-innerK)^order )
  filt[real.(K) .< innerK] .= 1
  TK(filt)
end

function makefilter(g::AbstractTwoDGrid; realvars=true, kwargs...)
  K = realvars ?
      @.(sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2)) : @.(sqrt((g.k*g.dx/π)^2 + (g.l*g.dy/π)^2))
  makefilter(K; kwargs...)
end

function makefilter(g::AbstractOneDGrid; realvars=true, kwargs...)
  K = realvars ? g.kr*g.dx/π : @.(abs(g.k*g.dx/π))
  makefilter(K; kwargs...)
end

makefilter(g, T, sz; kwargs...) = ones(T, sz).*makefilter(g; realvars=sz[1]==g.nkr, kwargs...)


x, y = generators(MPoly{Float64}, :x, :y)

@test vars(x) == [:x, :y]
@test nvars(x) == 2
@test eltype(x) == Float64


function searchsortedstart(sortedArray, x)
    left = 1
    right = length(sortedArray)
    mid = 0
    index = -1
    while(left<=right)
        mid = div(left+right, 2)
        if(sortedArray[mid] == x)
            index = mid
            break
        end
        if(sortedArray[mid] > x)
            right = mid-1
        else
            left = mid+1
        end
    end
    if(index == -1)
        if(x>sortedArray[mid]) index = mid+1
        else index = mid end
    else
        if((index-1 == 0) || sortedArray[index-1]!=x)
            return index
        else
            left = 1
            right = index
            mid = 0
            index = -1
            while(left<=right)
                mid = div(left+right, 2)
                if((mid==1 && sortedArray[mid]==x) || (sortedArray[mid]==x && sortedArray[mid-1]!=x))
                    index = mid
                    break
                end
                if(sortedArray[mid] == x)
                    right = mid-1
                else
                    left = mid+1
                end
            end
        end
    end
    return index
end

function searchsortedlast(sortedArray, x)
    left = 1
    right = length(sortedArray)
    mid = 0
    index = -1
    while(left<=right)
        mid = div(left+right, 2)
        if(sortedArray[mid] == x)
            index = mid
            break
        end
        if(sortedArray[mid] > x)
            right = mid-1
        else
            left = mid+1
        end
    end
    if(index == -1)
        if(x>sortedArray[mid]) index = mid
        else index = mid-1 end
    else
        if((index+1 > length(sortedArray)) || sortedArray[index+1]!=x)
            return index
        else
            left = index
            right = length(sortedArray)
            mid = 0
            index = -1
            while(left<=right)
                mid = div(left+right, 2)
                if((mid==length(sortedArray) && sortedArray[mid]==x) || (sortedArray[mid]==x && sortedArray[mid+1]!=x))
                    index = mid
                    break
                end
                if(sortedArray[mid] == x)
                    left = mid+1
                else
                    right = mid-1
                end
            end
        end
    end
    return index
end

function searchSorted(sortedArray, x)
    return string(searchsortedstart(sortedArray, x),":",searchsortedlast(sortedArray, x))
end


# ---
# title: 1658. Minimum Operations to Reduce X to Zero
# id: problem1658
# author: AquaIndigo
# date: 2020-11-23
# difficulty: Medium
# categories: Two Pointers, Binary Search, Greedy
# link: <https://leetcode.com/problems/minimum-operations-to-reduce-x-to-zero/description/>
# hidden: true
# ---
#
# You are given an integer array `nums` and an integer `x`. In one operation,
# you can either remove the leftmost or the rightmost element from the array
# `nums` and subtract its value from `x`. Note that this **modifies** the array
# for future operations.
#
# Return _the **minimum number** of operations to reduce _`x` _to **exactly**_
# `0` _if it 's possible_ _, otherwise, return_`-1`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,1,4,2,3], x = 5
#     Output: 2
#     Explanation: The optimal solution is to remove the last two elements to reduce x to zero.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [5,6,7,8,9], x = 4
#     Output: -1
#
#
# **Example 3:**
#
#
#
#     Input: nums = [3,2,20,1,1,3], x = 10
#     Output: 5
#     Explanation: The optimal solution is to remove the last three elements and the first two elements (5 operations in total) to reduce x to zero.
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 105`
#   * `1 <= nums[i] <= 104`
#   * `1 <= x <= 109`
#
#
## @lc code=start
using LeetCode

function min_operations(nums::Vector{Int}, x::Int)
    target = sum(nums) - x
    target < 0 && return -1
    println(target)
    left, right, maxn = 0, 0, -1
    len = length(nums)
    while right < len
        target -= nums[right += 1]
        while target < 0
            target += nums[left += 1]
        end
        (target == 0) && (maxn = max(maxn, right - left))
    end
    maxn == -1 ? maxn : len - maxn
end
## @lc code=end


# Bogumił Kamiński, 2019-03-25

function uppersum(M::AbstractMatrix{T}) where {T<:Number}
    n = size(M, 1)
    @assert size(M, 1) == size(M, 1)
    s = zero(T)
    for i in 2:n
        @simd for j in 1:i
            @inbounds s += M[j, i]
        end
    end
    s
end

@time m = randn(10_000, 10_000);
@time m = randn(10_000, 10_000);

m

@time sum(m)
@time sum(m)

@time uppersum(M)
@time uppersum(M)


# ---
# title: 1300. Sum of Mutated Array Closest to Target
# id: problem1300
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/sum-of-mutated-array-closest-to-target/description/>
# hidden: true
# ---
#
# Given an integer array `arr` and a target value `target`, return the integer
# `value` such that when we change all the integers larger than `value` in the
# given array to be equal to `value`, the sum of the array gets as close as
# possible (in absolute difference) to `target`.
#
# In case of a tie, return the minimum such integer.
#
# Notice that the answer is not neccesarilly a number from `arr`.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [4,9,3], target = 10
#     Output: 3
#     Explanation: When using 3 arr converts to [3, 3, 3] which sums 9 and that's the optimal answer.
#
#
# **Example 2:**
#
#
#
#     Input: arr = [2,3,5], target = 10
#     Output: 5
#
#
# **Example 3:**
#
#
#
#     Input: arr = [60864,25176,27249,21296,20204], target = 56803
#     Output: 11361
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 10^4`
#   * `1 <= arr[i], target <= 10^5`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Test
using Laplacians
using LinearAlgebra
using Random
using SparseArrays
using LinearAlgebra
using Statistics
using RandomV06

L = Laplacians
include("testIncludes.jl")
include("test_graph_generators.jl")

include("testByFile.jl")

include("testByExport.jl")

include("testDrawing.jl")

include("testSolvers.jl")

include("testPCG.jl")


using OrdinaryDiffEq, StochasticDiffEq, DelayDiffEq, Flux,
      Zygote, Test
println("Starting tests")
using DiffEqFlux, DiffEqSensitivity
mp = Float32[0.1,0.1]
x = Float32[2.; 0.]
xs = Float32.(hcat([0.; 0.], [1.; 0.], [2.; 0.]))
tspan = (0.0f0,1.0f0)
dudt = Chain(Dense(2,50,tanh),Dense(50,2))
fastdudt = FastChain(FastDense(2,50,tanh),FastDense(50,2))
staticdudt = FastChain(StaticDense(2,50,tanh),StaticDense(50,2))

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(x)

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(xs)

@info "Test some gradients"

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP()))
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

## Fast

@info "Test some fast layers"

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

goodgrad = grads[node.p]
p = node.p

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])
goodgrad2 = grads[node.p]
@test goodgrad ≈ goodgrad2 # Make sure adjoint overloads are correct

node = NeuralODE(staticdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

@test_throws ErrorException grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))

@info "Test some adjoints"

# Adjoint
@testset "adjoint mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])
end

@info "Test Tracker"

# RD
@testset "Tracker mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])
end

@info "Test non-ODEs"

dudt2 = Chain(Dense(2,50,tanh),Dense(50,2))
fastdudt2 = FastChain(FastDense(2,50,tanh),FastDense(50,2))
NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.1)(x)
sode = NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test ! iszero(grads[xs])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

sode = NeuralDSDE(fastdudt,fastdudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

dudt22 = Chain(Dense(2,50,tanh),Dense(50,4),x->reshape(x,2,2))
fastdudt22 = FastChain(FastDense(2,50,tanh),FastDense(50,4),(x,p)->reshape(x,2,2))
NeuralSDE(dudt,dudt22,(0.0f0,.1f0),2,LambaEM(),saveat=0.01)(x)

sode = NeuralSDE(dudt,dudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

sode = NeuralSDE(fastdudt,fastdudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

ddudt = Chain(Dense(6,50,tanh),Dense(50,2))
NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)

grads = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(grads[x])
@test ! iszero(grads[dode.p])

@test_broken grads = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[dode.p])

fastddudt = FastChain(FastDense(6,50,tanh),FastDense(50,2))
NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)

grads = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(grads[x])
@test ! iszero(grads[dode.p])

@test_broken grads = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[dode.p])


unitrng(; eps = 1e-5) = 0 + eps:0.005:1 - eps

function fx(xrv::RandVar;
            ΩT = defΩ(),
            xdim = 1,
            verbose = false)
  ω = ΩT()
  xrv(ω)
  function f(x)
    ω = Space.update(ω, xdim, x)
    y = xrv(ω)
    verbose && println(x, " -> ", y)
    y
  end
end

function fxy(xrv::RandVar;
                  ΩT = defΩ(),
                  xdim = 1,
                  ydim = 2)
  ω = ΩT()
  xrv(ω)
  function f(x, y)
    ω = Space.update(ω, xdim, x)
    ω = Space.update(ω, ydim, y)
    xrv(ω)
  end
end

"""
Contour Plot of two dimensions of Ω

```julia
using Omega
x = normal(0.0, 1.0)
y = normal(0.0, 1.0)
c1 = logerr(x ==ₛ y)
c2 = logerr(x >ₛ y)
ωcontour(c2)
```

"""
function ωcontour(xrv::RandVar;
                  ΩT = defΩ(),
                  xdim = 1,
                  ydim = 2,
                  xrng = unitrng(),
                  yrng = unitrng(),
                  plt = Plots.plot(),
                  fill = true,
                  aspect_ratio = 1,
                  kwargs...)
  ω = ΩT()
  xrv(ω)
  function f(x, y)
    ω = Space.update(ω, xdim, x)
    ω = Space.update(ω, ydim, y)
    @show xrv(ω)
  end
  Plots.contour!(plt, xrng, yrng, f; fill = fill, aspect_ratio = aspect_ratio, kwargs...)
end

function ωcontourmk(xrv::RandVar;
                  ΩT = defΩ(),
                  xdim = 1,
                  ydim = 2,
                  xrng = unitrng(),
                  yrng = unitrng(),
                  aspect_ratio = 1,
                  kwargs...)
  ω = ΩT()
  xrv(ω)
  function f(x, y)
    ω = Space.update(ω, xdim, x)
    ω = Space.update(ω, ydim, y)
    @show abs(xrv(ω)) / 20000
  end
  AbstractPlotting.surface(xrng, yrng, f; kwargs...)
end

# "x ∈ [0, 1]?"
# isunit(x) = 0.0 <= x <= 1.0

"Plot path in Omega Space from sequence of samples"
function ωtrace(ωs::Vector;
                   xdim = 1,
                   ydim = 2,
                   plt = plot())
  xs, ys = ntranspose(map(ω -> (getdim(ω, xdim), getdim(ω, ydim)), ωs))
  plot!(plt, xs, ys, arrow = :arrow, linealpha = 0.5, legend = false)
end


@testset "Testing the BBox structure" begin
    X = [
        0.512363  0.523803  0.610511;
        0.207648  0.54746   0.21826;
        0.472262  0.919383  0.313127;
        0.812478  0.223824  0.593296;
        0.109389  0.569246  0.625016
    ]
    Y = [
        0.551976   0.105557   0.733297
        0.0419843  0.0205705  0.545029
        0.237319   0.844252   0.217371
        0.759811   0.638246   0.404687
    ]

    bx = BBox(X)
    by = BBox(Y)

    # constructor
    @test_throws DomainError BBox([1 2 3 4])

    # minimum
    @test get_min(bx) == [0.109389, 0.223824, 0.21826]
    @test get_min(by) == [0.0419843, 0.0205705, 0.217371]

    # maximum
    @test get_max(bx) == [0.812478, 0.919383, 0.625016]
    @test get_max(by) == [0.759811, 0.844252, 0.733297]

    # center
    @test get_ctr(bx) ≈ [0.4609335, 0.5716035, 0.421638]
    @test get_ctr(by) ≈ [0.40089765, 0.43241125, 0.475334]

    # edge
    @test get_edg(bx) ≈ [0.703089, 0.695559, 0.406756]
    @test get_edg(by) ≈ [0.7178267, 0.8236815, 0.515926]
end


struct Information
    f_optimum::Float64
    x_optimum::Array{Float64}
end

"""
    Information Structure

`Information` can be used to store the true optimum in order to stop a metaheuristic early.

Properties:

- `f_optimum` known minimum.
- `x_optimum` known minimizer.

If `Options` is provided, then `optimize` will stop when `|f(x) - f(x_optimum)| < Options.f_tol`
or `‖ x - x_optimum ‖ < Options.x_tol` (euclidean distance).

# Example

If you want an approximation to the minimum with accuracy of `1e-3` (|f(x) - f(x*)| < 1e-3),
then you may use `Information`.

```jldoctest
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> bounds = [  -10.0 -10 -10; # lower bounds
                    10.0  10 10 ] # upper bounds
2×3 Array{Float64,2}:
 -10.0  -10.0  -10.0
  10.0   10.0   10.0

julia> information = Information(f_optimum = 0.0)
Information(0.0, Float64[])

julia> options = Options(f_tol = 1e-3)
Options(0.0, 0.001, 0.0, 0.0, 1000.0, 0.0, 0.0, 0, false, true, false, :minimize)

julia> state = optimize(f, bounds, ECA(information=information, options=options))
+=========== RESULT ==========+
| Iter.: 22
| f(x) = 0.000650243
| solution.x = [0.022811671589729583, 0.007052331140376011, -0.008951836265056107]
| f calls: 474
| Total time: 0.0106 s
+============================+
```
"""
function Information(;#
    f_optimum = NaN,
    x_optimum::Array{Float64} = Float64[],
)

    Information(Float64(f_optimum), x_optimum)

end


using DelayEmbeddings, Statistics, TimeseriesSurrogates

"""
    construct_candidate_variables(
        source::Vector{AbstractVector},
        target::Vector{AbstractVector},
        [cond::Vector{AbstractVector}];
        k::Int = 1, include_instantaneous = true,
        τexclude::Union{Int, Nothing} = nothing,
        maxlag::Union{Int, Float64} = 0.05
    ) → ([τs_source, τs_target, τs_cond, ks_targetfuture], [js_source, js_target, js_cond, js_targetfuture])

Construct candidate variables from input time series. `source` is a vector of equal-length time series
assumed to represent the putative source process. `target` and `cond` are the same, but contains time series
of the target process and of the conditional processes, respectively. `k` is the desired prediction lag.

If `include_instantaneous == true`, then the analysis will also consider instantaneous interactions between
the variables.

If `maxlag` is an integer, `maxlag` is taken as the maximum allowed embedding lag. If `maxlag` is a float,
then the maximum embedding lag is taken as `maximum([length.(source); length.(target); length.(cond)])*maxlag`.

If `τexclude` is an integer, all variables whose embedding lag has absolute value equal to `exclude` will be
excluded.
"""
function construct_candidate_variables(source, target, cond;
        k::Int = 1,
        τexclude::Union{Int, Nothing} = nothing,
        include_instantaneous = true,
        method_delay = "ac_min",
        maxlag::Union{Int, Float64} = 0.05)

    # Ensure all time series are of the same length.
    Ls = [length.(source); length.(target); length.(cond)]
    @assert all(Ls .== maximum(Ls))

    if maxlag isa Int
        τs = 1:maxlag
    else
        τs = 1:ceil(Int, maximum(Ls)*maxlag)
    end

    # Find the maximum allowed embedding lag for each of the candidates.
    τsmax_source = [estimate_delay(s, method_delay, τs) for s in source]
    τsmax_target = [estimate_delay(t, method_delay, τs) for t in target]
    τsmax_cond = [estimate_delay(c, method_delay, τs) for c in cond]

    # The set of candidate variables generated from the target
    # time series must not overlap with the prediction variables,
    # so the k-th lag variable is never included in the candidate set.
    τs_target = [[0:-1:-τ...,] for τ in τsmax_target]

    # Source and conditional variables may have instantaneous
    # interactions with target, so include lag `k` (if desired)
    # in the set of candidate variables.
    if include_instantaneous
        τs_source = [[k; [0:-1:-τ...,]] for τ in τsmax_source]
        τs_cond = [[k; [0:-1:-τ...,]] for τ in τsmax_cond]
    else
        τs_source = [[0:-1:-τ...,] for τ in τsmax_source]
        τs_cond = [[0:-1:-τ...,] for τ in τsmax_cond]
    end

    τs = [τs_source..., τs_target..., τs_cond...]

    # Embedding variables
    js = [[i for x in 1:length(τs[i])] for i = 1:length(τs)]
    js_targetfuture = [i for i in length(τs_source)+1:length(τs_source)+length(τs_target)]

    # Prediction variables
    ks_targetfuture = [k for i in 1:length(target)]

    # Variable filtering, if desired
    if τexclude isa Int
        τs = [filtered_τs(τsᵢ, jsᵢ, τexclude) for (τsᵢ, jsᵢ) in zip(τs, js)]
        js = [filtered_js(τsᵢ, jsᵢ, τexclude) for (τsᵢ, jsᵢ) in zip(τs, js)]
    end
    return [τs..., ks_targetfuture], [js..., js_targetfuture]
end

# Usaully, we use all lags from startlag:-\tau_max to construct variables. In some situations,
# we may want to exclude som of those variables.
function filtered_τs(τs::AbstractVector{Int}, js::AbstractVector{Int}, τexclude::Int)
    [τ for τ in τs if abs(τ) != abs.(τexclude)]
end

function filtered_js(τs::AbstractVector{Int}, js::AbstractVector{Int}, τexclude::Int)
    [j for (τ, j) in zip(τs, js) if abs(τ) != abs.(τexclude)]
end

# source & target variant
function construct_candidate_variables(source, target;
        k::Int = 1,
        τexclude::Union{Int, Nothing} = nothing,
        include_instantaneous = true,
        method_delay = "ac_min",
        maxlag::Union{Int, Float64} = 0.05)

    # Ensure all time series are of the same length.
    Ls = [length.(source); length.(target)]
    @assert all(Ls .== maximum(Ls))

    if maxlag isa Int
        τs = 1:maxlag
    else
        τs = 1:ceil(Int, maximum(Ls)*maxlag)
    end

    # Find the maximum allowed embedding lag for each of the candidates.
    τsmax_source = [estimate_delay(s, method_delay, τs) for s in source]
    τsmax_target = [estimate_delay(t, method_delay, τs) for t in target]

    # The set of candidate variables generated from the target
    # time series must not overlap with the prediction variables,
    # so the k-th lag variable is never included in the candidate set.
    τs_target = [[0:-1:-τ...,] for τ in τsmax_target]

    # Source variables may have instantaneous
    # interactions with target, so include lag `k` (if desired)
    # in the set of candidate variables.
    if include_instantaneous
        τs_source = [[k; [0:-1:-τ...,]] for τ in τsmax_source]
    else
        τs_source = [[0:-1:-τ...,] for τ in τsmax_source]
    end

    τs = [τs_source..., τs_target...]

    ks_targetfuture = [k for i in 1:length(target)]
    js_targetfuture = [i for i in length(τs_source)+1:length(τs_source)+length(τs_target)]
    τs = [τs_source..., τs_target...,]
    js = [[i for x in 1:length(τs[i])] for i = 1:length(τs)]

    # Variable filtering, if desired
    if τexclude isa Int
        τs = [filtered_τs(τsᵢ, jsᵢ, τexclude) for (τsᵢ, jsᵢ) in zip(τs, js)]
        js = [filtered_js(τsᵢ, jsᵢ, τexclude) for (τsᵢ, jsᵢ) in zip(τs, js)]
    end

    return [τs..., ks_targetfuture], [js..., js_targetfuture]
end

# source, target & cond variant
function candidate_embedding(source, target, cond;
        η::Int = 1,
        τexclude::Union{Int, Nothing} = nothing,
        include_instantaneous = true,
        method_delay = "mi_min",
        maxlag::Union{Int, Float64} = 0.05)

    τs, js = construct_candidate_variables(
        source, target, cond,
        k = η,
        τexclude = τexclude,
        method_delay = method_delay,
        maxlag = maxlag,
        include_instantaneous = include_instantaneous)

    # TODO: This is more efficient if not using datasets. Re-do manually.
    data = Dataset([source..., target..., cond...,]...,)
    ℰ = genembed(data, ((τs...)...,), ((js...)...,))

    # Get all variables except the target future (which are the last columns of ℰ)
    n_timeseries = size(ℰ, 2)
    n_timeseries_target = length(target)
    Ω = [ℰ[:, i] for i = 1:n_timeseries - n_timeseries_target]
    Y⁺ = ℰ[:, n_timeseries - n_timeseries_target+1:end]

    # We need to keep track of which variables are from the source, because
    # when computing the final TE, we need a marginal which is 𝒮 \ 𝒮_source.
    # Hence, we need to know which indices in `js` correspond to the source.
    idxs_source = 1:length(source)
    idxs_target = length(source)+1:length(source)+length(target)
    idxs_cond = length(source)+length(target)+1:length(source)+length(target)+length(cond)

    return Ω, Y⁺, τs, js, idxs_source, idxs_target, idxs_cond
end

# source & target variant
function candidate_embedding(source, target;
        η::Int = 1,
        τexclude::Union{Int, Nothing} = nothing,
        include_instantaneous = true,
        method_delay = "mi_min",
        maxlag::Union{Int, Float64} = 0.05)

    τs, js = construct_candidate_variables(
        source, target,
        k = η,
        τexclude = τexclude,
        method_delay = method_delay,
        maxlag = maxlag,
        include_instantaneous = include_instantaneous)

    # TODO: This is more efficient if not using datasets. Re-do manually.
    data = Dataset([source..., target...,]...,)
    ℰ = genembed(data, ((τs...)...,), ((js...)...,))

    # Get all variables except the target future (which are the last columns of ℰ)
    n_timeseries = size(ℰ, 2)
    n_timeseries_target = length(target)
    Ω = [ℰ[:, i] for i = 1:n_timeseries - n_timeseries_target]

    Y⁺ = ℰ[:, n_timeseries - n_timeseries_target+1:end]
    idxs_source = 1:length(source)
    idxs_target = length(source)+1:length(source)+length(target)
    idxs_cond = Int[]

    return Ω, Y⁺, τs, js, idxs_source, idxs_target, idxs_cond
end

function optim_te(Ω, Y⁺, τs, js, idxs_source, idxs_target, idxs_cond, est;
        q = 1, base = 2,
        α = 0.05, nsurr = 100, surr::Surrogate = RandomShuffle())

    τs_comb = [(τs...)...,]
    js_comb = [(js...)...,]

    n_candidate_variables = length(Ω)

    𝒮 = Vector{Vector{Float64}}(undef, 0)
    𝒮_τs = Vector{Int}(undef, 0)
    𝒮_js = Vector{Int}(undef, 0)

    k = 1
    while k <= n_candidate_variables
        n_remaining_candidates = length(Ω)
        CMIs_between_Y⁺_and_candidates = zeros(n_remaining_candidates)

        # At first iteration, only loop through source variable. If no source variable is found that
        # yields significant TE, terminate.
        for i = 1:n_remaining_candidates
            if k == 1 || length(𝒮) == 0
                Cᵢ = Ω[i]
                CMI_Y⁺_Cᵢ =
                    genentropy(Dataset(Y⁺, Dataset(Cᵢ)), est, q = q, base = base) -
                    genentropy(Dataset(Cᵢ), est, q = q, base = base)
            else
                Cᵢ = [Ω[i], 𝒮...]
                CMI_Y⁺_Cᵢ =
                    genentropy(Dataset(Y⁺, Dataset(Cᵢ...,)), est, q = q, base = base) -
                    genentropy(Dataset(Cᵢ...,), est, q = q, base = base)
            end
            CMIs_between_Y⁺_and_candidates[i] = CMI_Y⁺_Cᵢ
        end

        idx = findfirst(x -> x == minimum(CMIs_between_Y⁺_and_candidates), CMIs_between_Y⁺_and_candidates)
        cₖ = Ω[idx]

        # Test the significance of this candidate by using a permutation test. The type of surrogate
        # is given by `surr`, and we will use `nsurr` surrogate realizations.
        CMI_permutations = zeros(nsurr)
        s = surrogenerator(cₖ, surr)

        # If k == 1, no candidates have been selected, so CMI reduces to MI
        if k == 1
            cmiₖ = CMIs_between_Y⁺_and_candidates[idx]

            for i = 1:nsurr
                surr_cₖ = s() # Surrogate version of cₖ
                CMI_permutations[i] = mutualinfo(Y⁺, surr_cₖ, est)
            end
        # If k > 1, at least one candidate has been selected, so we compute CMI
        else
            # Precompute terms that do not change during surrogate loop
            H_Y⁺_𝒮 = genentropy(Dataset(Y⁺, Dataset(𝒮...,)), est, q = q, base = base)
            H_𝒮 = genentropy(Dataset(𝒮...), est, q = q, base = base)

            # Original TE
            cmiₖ = H_Y⁺_𝒮 +
                    genentropy(Dataset([cₖ, 𝒮...,]...,), est, q = q, base = base) -
                    genentropy(Dataset(Y⁺, Dataset([cₖ, 𝒮...,]...,)), est, q = q, base = base) -
                    H_𝒮

            for i = 1:nsurr
                surr_cₖ = s() # Surrogate version of cₖ
                CMI_permutations[i] = H_Y⁺_𝒮 +
                    genentropy(Dataset([surr_cₖ, 𝒮...]...,), est, q = q, base = base) -
                    genentropy(Dataset(Y⁺, Dataset([surr_cₖ, 𝒮...]...,)), est, q = q, base = base) -
                    H_𝒮
            end
        end

        # If the candidate passes the significance test, add it to list of selected candidates
        # and remove it from list of remaining candidates.
        if cmiₖ > quantile(CMI_permutations, 1 - α)
            push!(𝒮, cₖ)
            push!(𝒮_τs, τs_comb[idx])
            push!(𝒮_js, js_comb[idx])
            deleteat!(Ω, idx)
            deleteat!(τs_comb, idx)
            deleteat!(js_comb, idx)
            k = k + 1
        # If the candidate does not pass significance test, terminate.
        else
            k = n_candidate_variables + 1
        end
    end

    # No variables were selected at all.
    if length(𝒮) == 0
        return 0.0, Int[], Int[], idxs_source, idxs_target, idxs_cond
    end

    # If no variables were selected from the source process, then TE is not well-defined.
    n_source_vars_picked = count(x -> x ∈ idxs_source, 𝒮_js)
    if n_source_vars_picked == 0
        return 0.0, Int[], Int[], idxs_source, idxs_target, idxs_cond
    end

    # No variables were selected from the target or conditional processes.
    𝒮_nonX = [ts for (ts, j) in zip(𝒮, 𝒮_js) if j ∉ idxs_source]
    if length(𝒮_nonX) == 0
        return 0.0, Int[], Int[], idxs_source, idxs_target, idxs_cond
    end

    CE2 = genentropy(Dataset(Y⁺, Dataset(𝒮...,)), est, base = base, q = q) -
        genentropy(Dataset(𝒮...,), est, base = base, q = q)

    CE1 = genentropy(Dataset(Y⁺, Dataset(𝒮_nonX...,)), est, base = base, q = q) -
        genentropy(Dataset(𝒮_nonX...,), est, base = base, q = q)

    CMI = CE1 - CE2
    return CMI, 𝒮_js, 𝒮_τs, idxs_source, idxs_target, idxs_cond
end

process_input(ts::Vector{T}) where T <: Number = [ts]
process_input(ts::AbstractVector{V}) where V <: Vector{N} where N <: Number = ts
process_input(ts::Dataset) = [columns(ts)...,]