using Tables
using MLJBase
using StatsBase: sample, mean, quantile
using Statistics
using CategoricalArrays
using Distributions
using EvoTrees
using EvoTrees: logit, sigmoid
# import EvoTrees: EvoTreeRegressor, EvoTreeClassifier, EvoTreeCount, EvoTreeGaussian

##################################################
### Regression - small data
##################################################
features = rand(10_000) .* 5 .- 2
X = reshape(features, (size(features)[1], 1))
Y = sin.(features) .* 0.5 .+ 0.5
Y = logit(Y) + randn(size(Y))
Y = sigmoid(Y)
y = Y
X = Tables.table(X)

# @load EvoTreeRegressor
# linear regression
tree_model = EvoTreeRegressor(max_depth=5, η=0.05, nrounds=10)
# logistic regression
tree_model = EvoTreeRegressor(loss=:logistic, max_depth=5, η=0.05, nrounds=10)
# quantile regression
tree_model = EvoTreeRegressor(loss=:quantile, α=0.75, max_depth=5, η=0.05, nrounds=10)

tree = machine(tree_model, X, y)
train, test = partition(eachindex(y), 0.7, shuffle=true); # 70:30 split
fit!(tree, rows=train, verbosity=1)

tree.model.nrounds += 10
fit!(tree, rows=train, verbosity=1)

# predict on train data
pred_train = predict(tree, selectrows(X, train))
println(mean(abs.(pred_train - selectrows(Y, train))))

# predict on test data
pred_test = predict(tree, selectrows(X, test))
println(mean(abs.(pred_test - selectrows(Y, test))))

##################################################
### classif
##################################################
# features = rand(10_000) .* 5 .- 2
# X = reshape(features, (size(features)[1], 1))
# Y = sin.(features) .* 0.5 .+ 0.5
# Y = logit(Y) + randn(size(Y))
# Y = sigmoid(Y)
# y = Int.(round.(Y)) .+ 1
# y = _levels[y]
# # y = string.(y)
# y = CategoricalArray(y, ordered=false)
# X = Tables.table(X)
# X_matrix = MLJBase.matrix(X)

X, y = @load_crabs
X_matrix = matrix(X)

# define hyperparameters
tree_model = EvoTreeClassifier(max_depth=4, η=0.05, λ=0.0, γ=0.0, nbins=32, nrounds=20, metric=:none)
model = fit_evotree(tree_model, X_matrix, y);
pred = predict(model, X_matrix)
pred_cat = pred .> 0.5
sum((y .== "B") .== pred_cat[:,1]) / length(y)

# @load EvoTreeRegressor
tree = machine(tree_model, X, y)
train, test = partition(eachindex(y), 0.7, shuffle=true); # 70:30 split
fit!(tree, rows=train, verbosity=1)
report(tree)[:feature_importances]

tree.model.nrounds += 10
fit!(tree, rows=train, verbosity=1)
report(tree)[:feature_importances]

pred_train = predict(tree, selectrows(X, train))
pred_train_mode = predict_mode(tree, selectrows(X, train))
println(cross_entropy(pred_train, selectrows(y, train)) |> mean)
println(sum(pred_train_mode .== y[train]))

pred_test = predict(tree, selectrows(X, test))
pred_test_mode = predict_mode(tree, selectrows(X, test))
println(cross_entropy(pred_test, selectrows(y, test)) |> mean)
println(sum(pred_test_mode .== y[test]))
pred_test_mode = predict_mode(tree, selectrows(X, test))

# using LossFunctions, Plots
# evo_model = EvoTreeClassifier(max_depth=6, η=0.05, λ=1.0, γ=0.0, nrounds=10, nbins=64)
# evo = machine(evo_model, X, y)
# r = range(evo_model, :nrounds, lower=1, upper=500)
# @time curve = learning_curve!(evo, range=r, resolution=10, measure=HingeLoss())
# plot(curve.parameter_values, curve.measurements)

##################################################
### regression - Larger data
##################################################
features = rand(1_000_000, 100)
# features = rand(100, 10)
X = features
Y = rand(size(X, 1))
𝑖 = collect(1:size(X, 1))

# train-eval split
𝑖_sample = sample(𝑖, size(𝑖, 1), replace=false)
train_size = 0.8
𝑖_train = 𝑖_sample[1:floor(Int, train_size * size(𝑖, 1))]
𝑖_eval = 𝑖_sample[floor(Int, train_size * size(𝑖, 1)) + 1:end]

X_train, X_eval = X[𝑖_train, :], X[𝑖_eval, :]
Y_train, Y_eval = Y[𝑖_train], Y[𝑖_eval]

# @load EvoTreeRegressor
tree_model = EvoTreeRegressor(
    loss=:linear, metric=:mae,
    nrounds=10,
    λ=0.0, γ=0.0, η=0.1,
    max_depth=6, min_weight=1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

X = Tables.table(X);
X = Tables.rowtable(X);
# X = Tables.columntable(X)
# X_matrix = MLJBase.matrix(X)

# typeof(X)
@time tree = machine(tree_model, X, Y);
train, test = partition(eachindex(Y), 0.8, shuffle=true); # 70:30 split
@time fit!(tree, rows=train, verbosity=1, force=true)

using LossFunctions
using MLJ
r = range(tree_model, :nrounds, lower=1, upper=100)
m = rms
@time curve = learning_curve!(evo, range=r, resolution=100, measure=m)

tree.model.nrounds += 1
@time update(tree.model, 0, tree.fitresult, tree.cache, X, Y);

tree.model.nrounds += 1
@time fit!(tree, rows=train, verbosity=1)
# @time MLJBase.fit!(tree, rows=train, verbosity=1)

# yhat = MLJBase.predict(tree.model, tree.fitresult, MLJ.selectrows(X,test))
pred_train = predict(tree, selectrows(X, train))
mean(abs.(pred_train - selectrows(Y, train)))

##################################################
### count - Larger data
##################################################
features = rand(100_000, 100)
# features = rand(100, 10)
X = features
Y = rand(UInt8, size(X, 1))
𝑖 = collect(1:size(X, 1))

# train-eval split
𝑖_sample = sample(𝑖, size(𝑖, 1), replace=false)
train_size = 0.8
𝑖_train = 𝑖_sample[1:floor(Int, train_size * size(𝑖, 1))]
𝑖_eval = 𝑖_sample[floor(Int, train_size * size(𝑖, 1)) + 1:end]

X_train, X_eval = X[𝑖_train, :], X[𝑖_eval, :]
Y_train, Y_eval = Y[𝑖_train], Y[𝑖_eval]

# @load EvoTreeRegressor
tree_model = EvoTreeCount(
    loss=:poisson, metric=:poisson,
    nrounds=10,
    λ=0.0, γ=0.0, η=0.1,
    max_depth=6, min_weight=1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

X = Tables.table(X)
X = Tables.rowtable(X)
X = Tables.columntable(X)
X_matrix = MLJBase.matrix(X)

# typeof(X)
@time tree = machine(tree_model, X, Y)
train, test = partition(eachindex(Y), 0.8, shuffle=true); # 70:30 split
@time fit!(tree, rows=train, verbosity=1, force=true)

tree.model.nrounds += 10
@time MLJBase.update(tree.model, 0, tree.fitresult, tree.cache, X, Y)

tree.model.nrounds += 10
@time fit!(tree, rows=train, verbosity=1)
# @time MLJBase.fit!(tree, rows=train, verbosity=1)

# yhat = MLJBase.predict(tree.model, tree.fitresult, MLJ.selectrows(X,test))
pred = predict(tree, selectrows(X, train))
pred_mean = predict_mean(tree, selectrows(X, train))
pred_mode = predict_mode(tree, selectrows(X, train))

##################################################
### Gaussian - Larger data
##################################################
features = rand(100_000, 100)
# features = rand(100, 10)
X = features
Y = rand(size(X, 1))
𝑖 = collect(1:size(X, 1))

# train-eval split
𝑖_sample = sample(𝑖, size(𝑖, 1), replace=false)
train_size = 0.8
𝑖_train = 𝑖_sample[1:floor(Int, train_size * size(𝑖, 1))]
𝑖_eval = 𝑖_sample[floor(Int, train_size * size(𝑖, 1)) + 1:end]

X_train, X_eval = X[𝑖_train, :], X[𝑖_eval, :]
Y_train, Y_eval = Y[𝑖_train], Y[𝑖_eval]

# @load EvoTreeRegressor
tree_model = EvoTreeGaussian(
    loss=:gaussian, metric=:gaussian,
    nrounds=10,
    λ=0.0, γ=0.0, η=0.1,
    max_depth=6, min_weight=1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

X = Tables.table(X)
X_matrix = MLJBase.matrix(X)

# typeof(X)
@time tree = machine(tree_model, X, Y)
train, test = partition(eachindex(Y), 0.8, shuffle=true); # 70:30 split
@time fit!(tree, rows=train, verbosity=1, force=true)

tree.model.nrounds += 10
@time MLJBase.update(tree.model, 0, tree.fitresult, tree.cache, X, Y)

tree.model.nrounds += 10
@time fit!(tree, rows=train, verbosity=1)
# @time MLJBase.fit!(tree, rows=train, verbosity=1)

# yhat = MLJBase.predict(tree.model, tree.fitresult, MLJ.selectrows(X,test))
pred = predict(tree, selectrows(X, train))
pred_mean = predict_mean(tree, selectrows(X, train))
pred_mode = predict_mode(tree, selectrows(X, train))
mean(abs.(pred_mean - selectrows(Y, train)))

q_20 = quantile.(pred, 0.20)
q_20 = quantile.(pred, 0.80)

#########################################
# MLJ2 test
#########################################
using EvoTrees
using MLJModelInterface
using MLJBase
using StatsBase: sample, mean, quantile
using Tables

X = rand(1_000_000, 100);
Y = rand(size(X, 1))

# @load EvoTreeRegressor
tree_model = EvoTreeRegressor(
    loss=:linear, metric=:mae,
    nrounds=10,
    λ=0.0, γ=0.0, η=0.1,
    max_depth=6, min_weight=1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

X = Tables.table(X);
# X = Tables.rowtable(X);
# X = Tables.columntable(X);
# X_matrix = MLJBase.matrix(X);

# typeof(X)
@time tree = machine(tree_model, X, Y);
train, test = partition(eachindex(Y), 0.8, shuffle=true); # 70:30 split
@time fit!(tree, rows=train, verbosity=1, force=false)

tree.model.nrounds += 1
@time fit!(tree, rows=train, verbosity=1)


# # Mauna Loa time series example
#
# In this notebook, we apply Gaussian process regression to the Mauna Loa CO₂
# dataset. This showcases a rich combination of kernels, and how to handle and
# optimize all their parameters.

# ## Setup
#
# We make use of the following packages:

using CSV, DataFrames  # data loading
using AbstractGPs  # exact GP regression
using ParameterHandling  # for nested and constrained parameters
using Optim  # optimization
using Zygote  # auto-diff gradient computation
using Plots  # visualisation

# Let's load and visualize the dataset.

# !!! tip
#     The `let` block [creates a new
#     scope](https://docs.julialang.org/en/v1/manual/variables-and-scoping/#scope-of-variables),
#     so any utility variables we define in here won't leak outside. This is
#     particularly helpful to keep notebooks tidy! The return value of the
#     block is given by its last expression.

(xtrain, ytrain), (xtest, ytest) = let
    data = CSV.read(joinpath(@__DIR__, "CO2_data.csv"), Tables.matrix; header=0)
    year = data[:, 1]
    co2 = data[:, 2]

    ## We split the data into training and testing set:
    idx_train = year .< 2004
    idx_test = .!idx_train

    (year[idx_train], co2[idx_train]), (year[idx_test], co2[idx_test])  # block's return value
end
## The utility variables such as `idx_train` and `idx_test` are not available outside the `let` scope

function plotdata()
    plot(; xlabel="year", ylabel="CO₂ [ppm]", legend=:bottomright)
    scatter!(xtrain, ytrain; label="training data", ms=2, markerstrokewidth=0)
    return scatter!(xtest, ytest; label="test data", ms=2, markerstrokewidth=0)
end

plotdata()

# ## Prior
#
# We will model this dataset using a sum of several kernels which describe
#
# - smooth trend: squared exponential kernel with long lengthscale;
# - seasonal component: periodic covariance function with period of one year,
#   multiplied with a squared exponential kernel to allow decay away from exact
#   periodicity;
# - medium-term irregularities: rational quadratic kernel;
# - noise terms: squared exponential kernel with short lengthscale
#   and uncorrelated observation noise.
#
# For more details, see [Rasmussen & Williams (2005), chapter 5](http://www.gaussianprocess.org/gpml/chapters/RW5.pdf).

# We will use
# [ParameterHandling.jl](https://invenia.github.io/ParameterHandling.jl/) for
# handling the (hyper)parameters of our model. It provides functions such as
# `positive` with which we can put constraints on the hyperparameters, and
# allows us to represent all required parameters as a nested NamedTuple:

#! format: off
## initial values to match http://stor-i.github.io/GaussianProcesses.jl/latest/mauna_loa/
θ_init = (;
    se_long = (;
        σ = positive(exp(4.0)),
        ℓ = positive(exp(4.0)),
    ),
    seasonal = (;
        ## product kernels only need a single overall signal variance
        per = (;
            ℓ = positive(exp(0.0)),  # relative to period!
            p = fixed(1.0),  # 1 year, do not optimize over
        ),
        se = (;
            σ = positive(exp(1.0)),
            ℓ = positive(exp(4.0)),
        ),
    ),
    rq = (;
        σ = positive(exp(0.0)),
        ℓ = positive(exp(0.0)),
        α = positive(exp(-1.0)),
    ),
    se_short = (;
        σ = positive(exp(-2.0)),
        ℓ = positive(exp(-2.0)),
    ),
    noise_scale = positive(exp(-2.0)),
)
#! format: on
#md nothing #hide

# We define a couple of helper functions to simplify the kernel construction:

SE(θ) = θ.σ^2 * with_lengthscale(SqExponentialKernel(), θ.ℓ)
## PeriodicKernel is broken, see https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/issues/389
##Per(θ) = with_lengthscale(PeriodicKernel(; r=[θ.ℓ/2]), θ.p)  # NOTE- discrepancy with GaussianProcesses.jl
Per(θ) = with_lengthscale(SqExponentialKernel(), θ.ℓ) ∘ PeriodicTransform(1 / θ.p)
RQ(θ) = θ.σ^2 * with_lengthscale(RationalQuadraticKernel(; α=θ.α), θ.ℓ)
#md nothing #hide

# This allows us to write a function that, given the nested tuple of parameter values, constructs the GP prior:

function build_gp_prior(θ)
    k_smooth_trend = SE(θ.se_long)
    k_seasonality = Per(θ.seasonal.per) * SE(θ.seasonal.se)
    k_medium_term_irregularities = RQ(θ.rq)
    k_noise_terms = SE(θ.se_short) + θ.noise_scale^2 * WhiteKernel()
    kernel = k_smooth_trend + k_seasonality + k_medium_term_irregularities + k_noise_terms
    return GP(kernel)  # [`ZeroMean`](@ref) mean function by default
end
#md nothing #hide

# ## Posterior
#
# To construct the posterior, we need to first build a [`FiniteGP`](@ref),
# which represents the infinite-dimensional GP at a finite number of input
# features:

function build_finite_gp(θ)
    f = build_gp_prior(θ)
    return f(xtrain)
end
#md nothing #hide

# !!! info "`WhiteKernel` vs `FiniteGP` observation noise"
#     In this notebook, we already included observation noise through the
#     `WhiteKernel` as part of the GP prior covariance in `build_gp_prior`. We
#     therefore call `f(xtrain)` which implies zero (additional) observation
#     noise.
#
#     Alternatively, we could have omitted the `θ.noise_scale^2 *
#     WhiteKernel()` term and instead passed the noise variance as a second
#     argument to the GP call in `build_finite_gp`, `f(xtrain,
#     θ.noise_scale^2)`.
#
#     These two approaches have slightly different semantics: In the first one,
#     the `WhiteKernel` contributes non-zero variance to the `[i, j]` element
#     of the covariance matrix of the `FiniteGP` if `xtrain[i] == xtrain[j]`
#     (based on the values of the features). In the second one, the observation
#     noise variance passed to `FiniteGP` only contributes to the diagonal
#     elements of the covariance matrix, i.e. for `i == j`.
#
#     Moreover, the variance (uncertainty) of the posterior predictions
#     includes the variance from the `WhiteKernel`, but does not include the
#     variance of the observation noise passed to the `FiniteGP`. To include
#     the observation noise in posterior predictions from the second approach,
#     call `fpost_opt(xtest, noise_scale^2)`.
#
# !!! tip
#     For most use-cases and if in any doubt, we recommend that you pass in
#     observation noise to the `FiniteGP`, and omit the explicit `WhiteKernel`.
#     This is slightly faster (no need to check `xtrain[i] == xtrain[j]` for
#     all pairs `i`, `j`), and `WhiteKernel` will not give stable gradients if
#     you wish to compute the gradient of the log marginal likelihood
#     `logpdf(f(x), y)` w.r.t. `x`.

# We obtain the posterior, conditioned on the (finite) observations, by calling
# [`posterior`](@ref):

function build_posterior_gp(θ)
    fx = build_finite_gp(θ)
    return posterior(fx, ytrain)
end
#md nothing #hide

# Now we can put it all together to obtain a [`PosteriorGP`](@ref).
# The call to `ParameterHandling.value` is required to replace the constraints
# (such as `positive` in our case) with concrete numbers:

fpost_init = build_posterior_gp(ParameterHandling.value(θ_init))

# Let's visualize what the GP fitted to the data looks like, for the initial choice of kernel hyperparameters.
#
# We use the following function to plot a GP `f` on a specific range, using the
# AbstractGPs [plotting
# recipes](https://juliagaussianprocesses.github.io/AbstractGPs.jl/dev/concrete_features/#Plotting).
# By setting `ribbon_scale=2` we visualize the uncertainty band with ``\pm 2``
# (instead of the default ``\pm 1``) standard deviations.

plot_gp!(f; label) = plot!(f(1920:0.2:2030); ribbon_scale=2, linewidth=1, label)
#md nothing #hide

plotdata()
plot_gp!(fpost_init; label="posterior f(⋅)")

# A reasonable fit to the data, but poor extrapolation away from the observations!

# ## Hyperparameter Optimization
#
# To improve the fit, we want to maximize the (log) marginal likelihood with
# respect to the hyperparameters.
# [Optim.jl](https://julianlsolvers.github.io/Optim.jl/stable/) expects to
# minimize a loss, so we define it as the negative log marginal likelihood:

function loss(θ)
    fx = build_finite_gp(θ)
    lml = logpdf(fx, ytrain)  # this computes the log marginal likelihood
    return -lml
end
#md nothing #hide

# !!! note "Work-in-progress"
#     In the future, we are planning to provide the `optimize_loss` utility
#     function as part of JuliaGaussianProcesses -- for now, we just define it
#     inline.
#
# The L-BFGS parameters were chosen because they seem to work well empirically.
# You could also try with the defaults.

default_optimizer = LBFGS(;
    alphaguess=Optim.LineSearches.InitialStatic(; scaled=true),
    linesearch=Optim.LineSearches.BackTracking(),
)

function optimize_loss(loss, θ_init; optimizer=default_optimizer, maxiter=1_000)
    options = Optim.Options(; iterations=maxiter, show_trace=true)

    θ_flat_init, unflatten = ParameterHandling.value_flatten(θ_init)
    loss_packed = loss ∘ unflatten

    ## https://julianlsolvers.github.io/Optim.jl/stable/#user/tipsandtricks/#avoid-repeating-computations
    function fg!(F, G, x)
        if F !== nothing && G !== nothing
            val, grad = Zygote.withgradient(loss_packed, x)
            G .= only(grad)
            return val
        elseif G !== nothing
            grad = Zygote.gradient(loss_packed, x)
            G .= only(grad)
            return nothing
        elseif F !== nothing
            return loss_packed(x)
        end
    end

    result = optimize(Optim.only_fg!(fg!), θ_flat_init, optimizer, options; inplace=false)

    return unflatten(result.minimizer), result
end
#md nothing #hide

# We now run the optimization:

θ_opt, opt_result = optimize_loss(loss, θ_init)
opt_result

# The final value of the log marginal likelihood is:

-opt_result.minimum

# !!! warning
#     To avoid bad local optima, we could (and should) have carried out several
#     random restarts with different initial values for the hyperparameters,
#     and then picked the result with the highest marginal likelihood. We omit
#     this for simplicity. For more details on how to fit GPs in practice,
#     check out [A Practical Guide to Gaussian
#     Processes](https://tinyurl.com/guide2gp).
#
# Let's construct the posterior GP with the optimized hyperparameters:

fpost_opt = build_posterior_gp(ParameterHandling.value(θ_opt))
#md nothing #hide

# This is the kernel with the point-estimated hyperparameters:

fpost_opt.prior.kernel

# Let's print the optimized values of the hyperparameters in a more helpful format:

# !!! note "Work-in-progress"
#     This is another utility function we would eventually like to move out of this notebook:

using Printf

show_params(nt::Union{Dict,NamedTuple}) = String(take!(show_params(IOBuffer(), nt)))
function show_params(io, nt::Union{Dict,NamedTuple}, indent::Int=0)
    for (s, v) in pairs(nt)
        if v isa Union{Dict,NamedTuple}
            println(io, " "^indent, s, ":")
            show_params(io, v, indent + 4)
        else
            println(io, " "^indent, s, " = ", @sprintf("%.3f", v))
        end
    end
    return io
end

print(show_params(ParameterHandling.value(θ_opt)))

# And, finally, we can visualize our optimized posterior GP:

plotdata()
plot_gp!(fpost_opt; label="optimized posterior f(⋅)")


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMCoupling.jl/blob/master/LICENSE

using FEMCoupling: get_C
using Base.Test

@testset "Plain strain kinematic Coupling" begin
# plane strain problem
# Square shaped plane strain element(8x8) + one single node(2x2)
K=zeros(10,10)
K[1:8,1:8]=
  [500 225 -350 75 -250 -225 100 -75;
   225 500 -75 100 -225 -250 75 -350;
  -350 -75 500 -225  100 75 -250 225;
   75 100 -225  500 -75 -350 225 -250;
  -250 -225 100 -75   500 225 -350 75;
  -225 -250 75  -350  225 500 -75 100;
   100 75 -250 225 -350 -75 500 -225;
  -75 -350 225 -250 75 100 -225 500]

# Removing fixed DOFs (1 and 2) from nodes 1 and 4
K=K[1:end .!=1, 1:end .!=1]
K=K[1:end .!=1, 1:end .!=1]
K=K[1:end .!=5, 1:end .!=5]
K=K[1:end .!=5, 1:end .!=5]

# Force vector
f=zeros(10,1)
f[9,1]=1200

# Removing fixed DOFs
f=f[1:end .!= 1]
f=f[1:end .!= 1]
f=f[1:end .!= 5]
f=f[1:end .!= 5]

############### Calculations by hand

# Making the C matrix by hand
A=zeros(4,10)
A[1,3]-=1;  A[1,9]=1; A[2,4]=-1; A[2,10]=1;
A[3,5]-=1;  A[3,9]=1; A[4,6]=-1; A[4,10]=1

# Eliminating fixed dofs (1,2) from nodes 1 and 4
A=A[1:end , 1:end .!=1 ]
A=A[1:end , 1:end .!=1 ]
A=A[1:end , 1:end .!=5 ]
A=A[1:end , 1:end .!=5 ]

# Renaming variables to match with get_C.jl
C_expected=A
g_expected=zeros(4,1)
D_expected=zeros(4,4)

# Assembly for solving
K_expected= [K            C_expected';
             C_expected   D_expected]

f_expected= [f;
             g_expected]

u_expected = K_expected\f_expected

############### Calculating C,D and g with get_C.jl
#                  get_C(refnode,slaves,dofs,ndofs,K_size)
K_size=size(K,1)
C,D,g= FEMCoupling.get_C(5,[2,3],[1,2],2,10)

# Removing fixed DOFs
C=C[1:end , 1:end .!=1 ]
C=C[1:end , 1:end .!=1 ]
C=C[1:end , 1:end .!=5 ]
C=C[1:end , 1:end .!=5 ]

KK=[K C';
    C D]
ff=[f;
    g]
u = lufact(KK) \ full(ff)

@test isapprox(C,C_expected,rtol=0.0001)
@test isapprox(D,D_expected,rtol=0.0001)
@test isapprox(g,g_expected,rtol=0.0001)
@test isapprox(ff,f_expected,rtol=0.0001)
@test isapprox(KK,K_expected,rtol=0.0001)
@test isapprox(u,u_expected,rtol=0.0001)
# If the last test passes, all other tests will pass too.
# Other tests are made to help tracing why the last test doesn't pass.
end


function gmres(a::AbstractArray{T,2},b::AbstractArray{T,1},x::AbstractArray{T,1}) where {T <: AbstractFloat} # Exact GMRES

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b)

    e1    = zeros(c)
    e1[1] = norm(b)

    g = zeros(T,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres(a::AbstractArray{T,2},b::AbstractArray{T,1},x::AbstractArray{T,1},forcing_term::T) where {T <: AbstractFloat} # Inexact GMRES

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b,forcing_term*norm(b))

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres(a::AbstractArray{T,2},b::AbstractArray{T,1},x::AbstractArray{T,1},forcing_term::T,m::S) where {T <: AbstractFloat, S <: Integer} # Restarted inexact GMRES

    x = copy(x)
    n = length(x)

    m = min(m,n)

    step = Array{T,1}(undef,n)

    g = Array{T,2}(undef,2,2)

    loop_max = ceil(Int,n/m)
    loop_count = 1
    while loop_count <= loop_max

        # Create Hessenberg form using Arnoldi iteration
        q, h, c = arnoldi(a,b,forcing_term*norm(b),m)

        e1    = zeros(c)
        e1[1] = norm(b)

        # Use Givens rotations to transform to upper-triangular form
        @views for i = 2:c
            givens!(g,h[i-1:i,i-1])
            h[i-1:i,:] = g'h[i-1:i,:]
            e1[i-1:i] = g'e1[i-1:i]
        end

        step .= q*(h\e1)

        if maximum(abs,a*step-b) < forcing_term*norm(b)
            return step, true
        else
            x .= x + step
        end
        loop_count += 1
    end

    return step, false

end

function gmres(f::Function,x::Array{T,1}) where {T <: AbstractFloat} # Exact GMRES

    n = length(x)
    a = ForwardDiff.jacobian(f,x)
    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b)

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres_sparse(f::Function,x::Array{T,1}) where {T <: AbstractFloat} # Exact GMRES

    n = length(x)
    a = sparse(ForwardDiff.jacobian(f,x))
    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b)

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres(f::Function,x::Array{T,1},forcing_term::T) where {T <: AbstractFloat} # Inexact GMRES

    n = length(x)
    a = ForwardDiff.jacobian(f,x)
    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b,forcing_term*norm(b))

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres_sparse(f::Function,x::Array{T,1},forcing_term::T) where {T <: AbstractFloat} # Inexact GMRES

    n = length(x)
    a = sparse(ForwardDiff.jacobian(f,x))
    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(a,b,forcing_term*norm(b))

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function gmres(f::Function,x::Array{T,1},forcing_term::T,m::S) where {T <: AbstractFloat, S <: Integer} # Restarted inexact GMRES

    x = copy(x)
    n = length(x)

    step = similar(x)

    m = min(m,n)

    j = Array{T,2}(undef,n,n)
    b = Array{T,1}(undef,n)

    g = Array{T,2}(undef,2,2)

    loop_max = ceil(Int,n/m)
    loop_count = 1
    while loop_count <= loop_max

        j .= ForwardDiff.jacobian(f,x)
        b .= -f(x)

        # Create Hessenberg form using Arnoldi iteration
        q, h, c = arnoldi(j,b,forcing_term*norm(b),m)

        e1    = zeros(c)
        e1[1] = norm(b)

        # Use Givens rotations to transform to upper-triangular form
        @views for i = 2:c
            givens!(g,h[i-1:i,i-1])
            h[i-1:i,:] = g'h[i-1:i,:]
            e1[i-1:i] = g'e1[i-1:i]
        end

        step .= q*(h\e1)

        if maximum(abs,j*step-b) < forcing_term*norm(b)
            return step, true
        else
            x .= x + step
        end
        loop_count +=1
    end

    return step, false

end

function gmres_sparse(f::Function,x::Array{T,1},forcing_term::T,m::S) where {T <: AbstractFloat, S <: Integer} # Restarted inexact GMRES

    x = copy(x)
    n = length(x)

    step = similar(x)

    m = min(m,n)

    j = sparse(Array{T,2}(undef,n,n))
    b = Array{T,1}(undef,n)

    g = Array{T,2}(undef,2,2)

    loop_max = ceil(Int,n/m)
    loop_count = 1
    while loop_count <= loop_max

        j .= sparse(ForwardDiff.jacobian(f,x))
        b .= -f(x)

        # Create Hessenberg form using Arnoldi iteration
        q, h, c = arnoldi(j,b,forcing_term*norm(b),m)

        e1    = zeros(c)
        e1[1] = norm(b)

        # Use Givens rotations to transform to upper-triangular form
        @views for i = 2:c
            givens!(g,h[i-1:i,i-1])
            h[i-1:i,:] = g'h[i-1:i,:]
            e1[i-1:i] = g'e1[i-1:i]
        end

        step .= q*(h\e1)

        if maximum(abs,j*step-b) < forcing_term*norm(b)
            return step, true
        else
            x .= x + step
        end
        loop_count +=1
    end

    return step, false

end

function jacobian_free_gmres(f::Function,x::Array{T,1}) where {T <: AbstractFloat} # Jacobian-free exact GMRES

    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(f,x,b)

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function jacobian_free_gmres(f::Function,x::Array{T,1},forcing_term::T) where {T <: AbstractFloat} # Jacobian-free inexact GMRES

    b = -f(x)

    # Create Hessenberg form using Arnoldi iteration
    q, h, c = arnoldi(f,x,b,forcing_term*norm(b))

    e1    = zeros(c)
    e1[1] = norm(b)

    g = Array{T,2}(undef,2,2)
    # Use Givens rotations to transform to upper-triangular form
    @views for i = 2:c
        givens!(g,h[i-1:i,i-1])
        h[i-1:i,:] = g'h[i-1:i,:]
        e1[i-1:i] = g'e1[i-1:i]
    end

    step = q*(h\e1)

    return step

end

function jacobian_free_gmres(f::Function,x::Array{T,1},forcing_term::T,m::S) where {T <: AbstractFloat, S <: Integer} # Jacobian-free restarted inexact GMRES

    x = copy(x)
    n = length(x)
    jvec = zeros(n)
    step = similar(x)

    m = min(m,n)

    b = Array{T,1}(undef,n)

    g = Array{T,2}(undef,2,2)

    loop_max = ceil(Int,n/m)
    loop_count = 1
    while loop_count <= loop_max

        b .= -f(x)

        # Create Hessenberg form using Arnoldi iteration
        q, h, c = arnoldi(f,x,b,forcing_term*norm(b),m)

        e1    = zeros(c)
        e1[1] = norm(b)

        # Use Givens rotations to transform to upper-triangular form
        @views for i = 2:c
            givens!(g,h[i-1:i,i-1])
            h[i-1:i,:] = g'h[i-1:i,:]
            e1[i-1:i] = g'*e1[i-1:i]
        end

        step .= q*(h\e1)

        jacvec!(jvec,f,x,step)

        if maximum(abs,jvec-b) < forcing_term*norm(b)
            return step, true
        else
            x .= x + step
        end
        loop_count += 1
    end

    return step, false

end

function jacobian_free_gmres_inplace(f::Function,x::Array{T,1},forcing_term::T,m::S) where {T <: AbstractFloat, S <: Integer} # Jacobian-free restarted inexact GMRES

    x = copy(x)
    n = length(x)
    jvec = zeros(n)
    step = similar(x)

    m = min(m,n)

    b   = Array{T,1}(undef,n)
    ffk = Array{T,1}(undef,n)

    g = Array{T,2}(undef,2,2)

    loop_max = ceil(Int,n/m)
    loop_count = 1
    while loop_count <= loop_max

        f(ffk,x)
        b .= -ffk

        # Create Hessenberg form using Arnoldi iteration
        q, h, c = arnoldi_inplace(f,x,b,forcing_term*norm(b),m)

        e1    = zeros(c)
        e1[1] = norm(b)

        # Use Givens rotations to transform to upper-triangular form
        @views for i = 2:c
            givens!(g,h[i-1:i,i-1])
            h[i-1:i,:] = g'h[i-1:i,:]
            e1[i-1:i] = g'*e1[i-1:i]
        end

        step .= q*(h\e1)

        jacvec_inplace!(jvec,f,x,step)

        if maximum(abs,jvec-b) < forcing_term*norm(b)
            return step, true
        else
            x .= x + step
        end
        loop_count += 1
    end

    return step, false

end


"""
Functions to deal with various types of norms and seminorms
"""

using FastRounding
using IntervalOptimisation
using TaylorSeries: Taylor1
using SparseArrays: getcolptr
using .DynamicDefinition

using .DynamicDefinition: derivative
using .PwDynamicDefinition: PwMap

"""
'Absolute value' definition that returns mag(I) for an interval and abs(x) for a real
"""
abs_or_mag(x::Number) = abs(x)
abs_or_mag(x::Interval) = mag(x)

"""
Certified upper bound to ||A|| (of specified NormKind)
"""
function opnormbound(::Type{L1}, A::AbstractVecOrMat{T}) where {T}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for j = 1:size(A, 2)
            nrmj::Tsum = 0
            for i = 1:size(A, 1)
                nrmj = nrmj ⊕₊ abs_or_mag(A[i,j])
            end
            nrm = max(nrm,nrmj)
        end
    end
    return convert(Tnorm, nrm)
end

function opnormbound(::Type{Linf}, A::AbstractVecOrMat{T}) where {T}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for i = 1:size(A, 1)
            nrmi::Tsum = 0
            for j = 1:size(A, 2)
                nrmi = nrmi ⊕₊ abs_or_mag(A[i,j])
            end
            nrm = max(nrm,nrmi)
        end
    end
    return convert(Tnorm, nrm)
end

function opnormbound(::Type{L1}, A::SparseMatrixCSC) where {T}
    # partly taken from JuliaLang's Sparsearray/src/linalg.jl
    m, n = size(A)
    Tnorm = typeof(abs_or_mag(float(real(zero(eltype(A))))))
    Tsum = promote_type(Float64,Tnorm)
    nA::Tsum = 0
    @inbounds begin
        for j=1:n
            colSum::Tsum = 0
            for i = getcolptr(A)[j]:getcolptr(A)[j+1]-1
                colSum = colSum ⊕₊ abs_or_mag(nonzeros(A)[i])
            end
            nA = max(nA, colSum)
        end
    end
    return convert(Tnorm, nA)
end

function opnormbound(::Type{Linf}, A::SparseMatrixCSC) where {T}
    # partly taken from JuliaLang's Sparsearray/src/linalg.jl
    m, n = size(A)
    Tnorm = typeof(abs_or_mag(float(real(zero(eltype(A))))))
    Tsum = promote_type(Float64,Tnorm)
    rowSum = zeros(Tsum,m)
    @inbounds begin
        for i=1:length(nonzeros(A))
            rowSum[rowvals(A)[i]] = rowSum[rowvals(A)[i]] ⊕₊ abs_or_mag(nonzeros(A)[i])
        end
    end
    return convert(Tnorm, maximum(rowSum))
end

"""
Rigorous upper bound on a vector norm. Note that Linf, L1 are the "analyst's" norms
"""
normbound(N::Type{L1}, v::AbstractVector) = opnormbound(L1, v) ⊘₊ Float64(length(v), RoundDown)
normbound(N::Type{Linf}, v::AbstractVector) = opnormbound(Linf, v)

"""
Types to compute norms iteratively by "adding a column at a time".
"""
abstract type NormCacher{T} end

mutable struct NormCacherL1 <: NormCacher{L1}
    C::Float64
    function NormCacherL1(n)
        new(0)
    end
end
"""
Create a new NormCacher to compute the normbound of the empty matrix with n rows
"""
NormCacher{L1}(n) = NormCacherL1(n)

mutable struct NormCacherLinf <: NormCacher{Linf}
    C::Vector{Float64}
    function NormCacherLinf(n)
        new(zeros(n))
    end
end
NormCacher{Linf}(n) = NormCacherLinf(n)

"""
Update a NormCacher to add one column to the matrix it is computing a norm of.
This column may be affected by an error ε (in the same norm).
"""
function add_column!(Cacher::NormCacherL1, v::AbstractVector, ε::Float64)
    Cacher.C = max(Cacher.C, opnormbound(L1, v) ⊕₊ ε)
end

function add_column!(Cacher::NormCacherLinf, v::AbstractVector, ε::Float64)
    Cacher.C = Cacher.C .⊕₊ abs.(v) .⊕₊ ε
end

"""
Return the norm of the matrix the NormCacher is working on.
"""
function get_norm(Cacher::NormCacherL1)
    return Cacher.C
end

function get_norm(Cacher::NormCacherLinf)
    return maximum(Cacher.C)
end

"""
(A, B) = dfly(strongnorm, auxnorm, dynamic)

Constants (A, B) such that ||Lf||_s ≦ A||f||_s + B||f||_aux
"""
dfly(::Type{<:NormKind}, ::Type{<:NormKind}, ::Dynamic) = @error "Not implemented"

function dfly(::Type{TotalVariation}, ::Type{L1}, D::Dynamic)
    dist = maximise(x -> distorsion(D, x), domain(D))[1]
    lam = maximise(x-> abs(1/derivative(D, x)), domain(D))[1]

    if !(abs(lam) < 1) # these are intervals, so this is *not* equal to abs(lam) >= 1.
        @error "The function is not expanding"
    end

    if is_full_branch(D)
        return lam.hi, dist.hi
    else
        if !(abs(lam) < 0.5)
            @error "Expansivity is insufficient to prove a DFLY. Try with an iterate."
        end
        # We need a way to estimate the branch widths
        @error "Not implemented"
    end
end

function dfly(::Type{TotalVariation}, ::Type{L1}, D::PwMap)
    dist = 0
    lam = 0
    disc = 0
    for i in 1:nbranches(D)
        f(x) = D.Ts[i](x)
        domain = hull(D.endpoints[i], D.endpoints[i+1])
        fprime(x) = f(Taylor1([x, 1], 1))[1]
        fsecond(x) = f(Taylor1([x, 1], 2))[2]*2
        distorsion(x)=abs(fsecond(x)/(fprime(x)^2))
        lambda(x) = abs(1/fprime(x))
        dist = max(dist, maximise(distorsion, domain)[1].hi)
        lam = max(lam, maximise(lambda, domain)[1].hi)
        low_rad = (abs(D.endpoints[i]-D.endpoints[i+1])/2).lo
        disc = max(disc, ((1/Interval(low_rad)).hi))
    end

    if is_full_branch(D)
        if !(abs(lam) < 1) # these are intervals, so this is *not* equal to abs(lam) >= 1.
            @error "The function is not expanding"
        end
        return lam, dist
    else
        if !(abs(2⊗₊lam) < 1)
            @error "Expansivity is insufficient to prove a DFLY. Try with an iterate."
        end
        return 2⊗₊lam, dist ⊕₊ disc
    end
end

function dfly(::Type{Lipschitz}, ::Type{L1}, D::Dynamic)
    # TODO: should assert that D is globally C2 instead, but we don't have that kind of infrastructure yet.
    @assert is_full_branch(D)

    dist = maximise(x -> distorsion(D, x), domain(D))[1]
    lam = maximise(x-> abs(1/derivative(D, x)), domain(D))[1]

    lam = lam.hi
    dist = dist.hi

    return lam ⊗₊ (2. ⊗₊  dist ⊕₊ 1.), dist ⊗₊ (2. ⊗₊ dist ⊕₊ 1.)
end


export projectGradient, line_search, projected_newton

# Write projected Newton method, according to
# Bertsekas, Dimitri P. "Projected Newton methods for optimization problems with simple constraints." SIAM Journal on control and Optimization 20.2 (1982): 221-246.

# Code inspired from Ricardo Baptista, Youssef Marzouk MIT
function projectGradient(x::Array{Float64,1}, G::Array{Float64,1})
    PG = copy(G);
    for i=1:size(PG,1)
        @inbounds begin
        if x[i]==0.0 && G[i] >=0
        PG[i] = 0.0
        end
        end
    end
    return PG
end

function linesearch(x::Array{Float64,1}, G::Array{Float64,1},
         p::Array{Float64,1}, J::Float64,Idx::Array{Int64,1}, Lhd::LHD)

    # Fixed parameters
    itmax = 15
    σ = 1e-4
    β = 2

    it = 0
    Jx_α_lin = copy(J)+1 # To enter the while loop
    α = copy(β)# The first iteration inside the while loop will be for α = 1
    α_p = copy(p)
    x_α = zero(x)
    while (J<Jx_α_lin) && (it<itmax)
        α /= β
        α_p .= copy(p)
        rmul!(α_p, copy(α))
        x_α .= max.(0.0, x - α_p)
        Jx_α = copy(Lhd(x_α, false, noutput = 1))
        α_p[Idx] .= x[Idx] - x_α[Idx]
        Jx_α_lin = Jx_α + σ*dot(G, α_p)
        it += 1
    end

    if (J<Jx_α_lin)
        print("Line search method reached maximum number of iterations")
        xh = copy(x)
        xh[Idx] .= 0.0
        index = Int64[]

        for i=1:size(x,1)
            @inbounds begin
            if (xh[i]>0.0) && (p[i]>0.0)
                push!(index, copy(i))
            end
            end
        end
        if isempty(index)
            α = 1.0
        else
            α = min(x[index]./p[index])
        end
    end
return α
end

function projected_newton(x0::Array{Float64,1}, Lhd::LHD, type::String)
# Default parameters
rtol_J = 1e-6;
rtol_G = 1e-6;
itmax = 30;
ϵ = 0.01;

n = size(x0,1)

for i=1:n
    @assert sign(x0[i])==1.0 "Initial conditions are not feasible"
end

x = copy(x0)
Lhd(x, true)
@get Lhd (G,H)

Jold =copy(Lhd.J[1])
J = copy(Lhd.J[1])

norm_PG0 = norm(projectGradient(x, G))
tol_G = norm_PG0*rtol_G

norm_PG = copy((norm_PG0))
rδJ = rtol_J + 1
it = 0

# Pre-allocate space
p =zero(x)
# h = zero(x)
z = zero(x)

## Iteration
while (rδJ > rtol_J) && (norm_PG > tol_G) && (it<itmax)
    # Define search direction
    wk = norm(x - max.(0.0, x-G))
    ϵk = min(ϵ, wk)
    Idx= Int64[]
    @inbounds for i=1:n
        if x[i] <= ϵk && G[i] >0.0
        push!(Idx, copy(i))
        end
    end

    if isempty(Idx)==false
        z .= zeros(n)
        @inbounds for i in Idx
            z[i] = copy(H[i,i])
        end
        # h .= deepcopy(diag(H))
        # z .= zeros(n)
        # z[Idx] .= h[Idx]
        H[Idx,:] .= 0.0
        H[:,Idx] .= 0.0
        H .+= diagm(z)
    end

    # Switch according to search direction
    if type =="TrueHessian"
    p .= copy(G)
    # Perform in-place Cholesky decomposition and store H\G into p
    LAPACK.posv!('U', H, p)
    elseif type =="ModifHessian"
        itmax = 100
        # Try if the TrueHessian technique works
        if isposdef(H)
            p .= copy(G)
            # Perform in-place Cholesky decomposition and store H\G into p
            LAPACK.posv!('U', H, p)
        else
        # else Strategy 1- flip negative eigenvalues
            if !issymmetric(H)
                H .+= H'
                rmul!(H, 0.5)
            end

            print("Hessian H is not positive-definite")
            # EH = eigen(Symmetric(H))
            EH = eigen(H)
            for i=1:size(EH.values,1)
                @inbounds EH.values[i] = copy( 1/(max(1e-8, abs(EH.values[i]))) )
            end

            # p .= deepcopy(G)
            mul!(p, EH.vectors', G)
            p .*= EH.values
            mul!(p, EH.vectors, copy(p))

            # p .= deepcopy(EH.vectors*( EH.values .* (EH.vectors'*G)))
        end
    elseif type =="Gradient"
        itmax = 1000
        p = copy(G)

    else
        error("Method not implemented yet")
    end

    # Perform line search
    α = linesearch(x,G,p,J,Idx,Lhd)
    # @show α

    # Update the estimate
    x .= copy(max.(0.0, x - α*p))
    # @show x
    Lhd(x, true)
    @get Lhd (G,H)
    J = copy(Lhd.J[1])

    # Convergence criteria
    rδJ = abs(J-Jold)/abs(Jold)
    # @show rδJ
    Jold = copy(J)
    norm_PG = copy(norm(projectGradient(x,G)))
    it +=1
end

if it>=itmax
    print("Max number of iterations is reached during the optimization")
end

return x, rδJ, norm_PG, it

end


srcdir = "/Users/anray/Desktop/TDGP/src"
any(srcdir .== LOAD_PATH) || push!(LOAD_PATH, srcdir)
any(pwd() .== LOAD_PATH) || push!(LOAD_PATH, pwd())
using PyPlot, Random, Statistics, Revise, LinearAlgebra, MCMC_Driver
import GP
##
Random.seed!(13)
xtest = collect(0:0.01:1)'
λ = [0.05]
δ = 0.02
δtry = 5δ
p = 2
demean = true

y = zeros(size(xtest,2))
change2 = round(Int, 3/4*length(y))
change1 = round(Int, 1/4*length(y))
y[change2:end] .= 2
y[change1:change2-1] .= 1
ynoisy = y + δ*randn(length(y))
dec = 4
ytrain = ynoisy[1:dec:end]
xtrain = xtest[1:dec:end]'

# f1 = figure()
# plot(xtest[:][:],y, "-k",label="true")
# plot(xtrain[:],ytrain, "+m", label="train",markersize=10)
#
ytest, σ2, σ_prior  = GP.GPfit(ytrain, xtrain, xtest, λ, δtry, p=p, demean=demean)
# plot(xtest[:],ytest,label="testing")
# plot(xtest[:],ytest+2sqrt.(diag(σ2)),"--r",alpha=0.2)
# plot(xtest[:],ytest-2sqrt.(diag(σ2)),"--b",alpha=0.2)

## non stationary
λ[1] = 0.05
λtest = λ[1]*ones(size(xtest))
λtrain = λ[1]*ones(Float64, size(xtrain))
ytest_ns, σ2_ns, σ_prior_ns  = GP.GPfit(ytrain, xtrain, xtest, λtest, λtrain, δtry, p=p, demean=demean)

f2, ax2 = plt.subplots(2,1, sharex=true, figsize=(7,7))
ax2[1].plot(xtest[:],y, "-k",label="true", linewidth=2)
ax2[1].plot(xtrain[:],ytrain, "xm", label="train",markersize=10, alpha=0.8)
ax2[1].plot(xtest[:],ytest_ns,label="fixed λ", linewidth=3)
# plot(xtest[:],ytest_ns+2sqrt.(diag(σ2_ns)),"--r",alpha=0.2)
# plot(xtest[:],ytest_ns-2sqrt.(diag(σ2_ns)),"--b",alpha=0.2)
s = gca()
##
lscale_max = 0.05
dropoff = 0.0025
λtest = lscale_max*(1 .+eps() .- (exp.(-1/dropoff*abs.(xtest.-xtest[change1+1]).^2) +
                    exp.(-1/dropoff*abs.(xtest.-xtest[change2-2]).^2)))
λtrain[:] = λtest[1:dec:end]
ytest_ns, σ2_ns, σ_prior_ns  = GP.GPfit(ytrain, xtrain,
    xtest, λtest, λtrain, δtry, p=2, demean=demean, nogetvars=false)

ax2[1].plot(xtest[:],ytest_ns,label="variable λ", linewidth=3, alpha=0.7, "-r")
ax2[2].plot(xtest[:],λtest[:], linewidth=3, color="r", alpha=0.7)
ax2[1].set_ylabel(L"f(y)")
ax2[2].set_xlabel(L"y")
ax2[2].set_ylabel(L"\lambda")
ax2[1].grid()
ax2[2].grid()
MCMC_Driver.nicenup(gcf(), fsize=14)
savefig("variable.png", dpi=300)
## timing
@info "meshkernel with broadcast"
@time for i = 1:1000
    GP.meshkernel(xtrain, xtest, λtest, λtrain, p)
end
@info "mapkernel with pairwise"
@time for i = 1:1000
    GP.mapkernel(xtrain, xtest, λtest, λtrain)
end
@info "loops and full expression"
@time for i = 1:1000
    GP.makekernel(xtrain, xtest, λtest, λtrain, p)
end


module Simpson
    using Base.Threads

    function simpson(func, x1::BigFloat, x2::BigFloat, n)
        dh = (x2 - x1) / BigFloat(n)

        # Simpsonの公式によって数値積分する
        m = div(n, 2)
        results = Array{BigFloat}(undef, m)
        @inbounds @threads for i = 1:m
            fx1 = func(x1 + dh * BigFloat(i * 2 - 1))
            fx2 = func(x1 + dh * BigFloat(i * 2))

            results[i] = BigFloat(4) * fx1 + BigFloat(2) * fx2
        end
        sum = func(x1) + reduce(+, results) - func(x2)

        return dh / BigFloat(3) * sum
    end
end


module Metrics

"""
    creategrid!(x1, elemtocoord, r)

Create a 1-D grid using `elemtocoord` (see [`brickmesh`](@ref)) using the 1-D
`(-1, 1)` reference coordinates `r`. The element grids are filled using linear
interpolation of the element coordinates.

If `Nq = length(r)` and `nelem = size(elemtocoord, 3)` then the preallocated
array `x1` should be `Nq * nelem == length(x1)`.
"""
function creategrid!(x1, e2c, r)
    (d, nvert, nelem) = size(e2c)
    @assert d == 1
    Nq = length(r)
    x1 = reshape(x1, (Nq, nelem))

    # linear blend
    @inbounds for e in 1:nelem
        for i in 1:Nq
            x1[i, e] =
                ((1 - r[i]) * e2c[1, 1, e] + (1 + r[i]) * e2c[1, 2, e]) / 2
        end
    end
    nothing
end

"""
    creategrid!(x1, x2, elemtocoord, r)

Create a 2-D tensor product grid using `elemtocoord` (see [`brickmesh`](@ref))
using the 1-D `(-1, 1)` reference coordinates `r`. The element grids are filled
using bilinear interpolation of the element coordinates.

If `Nq = length(r)` and `nelem = size(elemtocoord, 3)` then the preallocated
arrays `x1` and `x2` should be `Nq^2 * nelem == size(x1) == size(x2)`.
"""
function creategrid!(x1, x2, e2c, r)
    (d, nvert, nelem) = size(e2c)
    @assert d == 2
    Nq = length(r)
    x1 = reshape(x1, (Nq, Nq, nelem))
    x2 = reshape(x2, (Nq, Nq, nelem))

    # bilinear blend of corners
    @inbounds for (f, n) in zip((x1, x2), 1:d)
        for e in 1:nelem
            for j in 1:Nq
                for i in 1:Nq
                    f[i, j, e] =
                        (
                            (1 - r[i]) * (1 - r[j]) * e2c[n, 1, e] +
                            (1 + r[i]) * (1 - r[j]) * e2c[n, 2, e] +
                            (1 - r[i]) * (1 + r[j]) * e2c[n, 3, e] +
                            (1 + r[i]) * (1 + r[j]) * e2c[n, 4, e]
                        ) / 4
                end
            end
        end
    end
    nothing
end

"""
    creategrid!(x1, x2, x3, elemtocoord, r)

Create a 3-D tensor product grid using `elemtocoord` (see [`brickmesh`](@ref))
using the 1-D `(-1, 1)` reference coordinates `r`. The element grids are filled
using trilinear interpolation of the element coordinates.

If `Nq = length(r)` and `nelem = size(elemtocoord, 3)` then the preallocated
arrays `x1`, `x2`, and `x3` should be `Nq^3 * nelem == size(x1) == size(x2) ==
size(x3)`.
"""
function creategrid!(x1, x2, x3, e2c, r)
    (d, nvert, nelem) = size(e2c)
    @assert d == 3
    # TODO: Add asserts?
    Nq = length(r)
    x1 = reshape(x1, (Nq, Nq, Nq, nelem))
    x2 = reshape(x2, (Nq, Nq, Nq, nelem))
    x3 = reshape(x3, (Nq, Nq, Nq, nelem))

    # trilinear blend of corners
    @inbounds for (f, n) in zip((x1, x2, x3), 1:d)
        for e in 1:nelem
            for k in 1:Nq
                for j in 1:Nq
                    for i in 1:Nq
                        f[i, j, k, e] =
                            (
                                (1 - r[i]) *
                                (1 - r[j]) *
                                (1 - r[k]) *
                                e2c[n, 1, e] +
                                (1 + r[i]) *
                                (1 - r[j]) *
                                (1 - r[k]) *
                                e2c[n, 2, e] +
                                (1 - r[i]) *
                                (1 + r[j]) *
                                (1 - r[k]) *
                                e2c[n, 3, e] +
                                (1 + r[i]) *
                                (1 + r[j]) *
                                (1 - r[k]) *
                                e2c[n, 4, e] +
                                (1 - r[i]) *
                                (1 - r[j]) *
                                (1 + r[k]) *
                                e2c[n, 5, e] +
                                (1 + r[i]) *
                                (1 - r[j]) *
                                (1 + r[k]) *
                                e2c[n, 6, e] +
                                (1 - r[i]) *
                                (1 + r[j]) *
                                (1 + r[k]) *
                                e2c[n, 7, e] +
                                (1 + r[i]) *
                                (1 + r[j]) *
                                (1 + r[k]) *
                                e2c[n, 8, e]
                            ) / 8
                    end
                end
            end
        end
    end
    nothing
end

"""
    computemetric!(x1, J, ξ1x1, sJ, n1, D)

Compute the 1-D metric terms from the element grid arrays `x1`. All the arrays
are preallocated by the user and the (square) derivative matrix `D` should be
consistent with the reference grid `r` used in [`creategrid!`](@ref).

If `Nq = size(D, 1)` and `nelem = div(length(x1), Nq)` then the volume arrays
`x1`, `J`, and `ξ1x1` should all have length `Nq * nelem`.  Similarly, the face
arrays `sJ` and `n1` should be of length `nface * nelem` with `nface = 2`.
"""
function computemetric!(x1, J, ξ1x1, sJ, n1, D)
    Nq = size(D, 1)
    nelem = div(length(J), Nq)
    x1 = reshape(x1, (Nq, nelem))
    J = reshape(J, (Nq, nelem))
    ξ1x1 = reshape(ξ1x1, (Nq, nelem))
    nface = 2
    n1 = reshape(n1, (1, nface, nelem))
    sJ = reshape(sJ, (1, nface, nelem))

    d = 1

    @inbounds for e in 1:nelem
        J[:, e] = D * x1[:, e]
    end
    ξ1x1 .= 1 ./ J

    n1[1, 1, :] .= -sign.(J[1, :])
    n1[1, 2, :] .= sign.(J[Nq, :])
    sJ .= 1
    nothing
end

"""
    computemetric!(x1, x2, J, ξ1x1, ξ2x1, ξ1x2, ξ2x2, sJ, n1, n2, D)

Compute the 2-D metric terms from the element grid arrays `x1` and `x2`. All the
arrays are preallocated by the user and the (square) derivative matrix `D`
should be consistent with the reference grid `r` used in [`creategrid!`](@ref).

If `Nq = size(D, 1)` and `nelem = div(length(x1), Nq^2)` then the volume arrays
`x1`, `x2`, `J`, `ξ1x1`, `ξ2x1`, `ξ1x2`, and `ξ2x2` should all be of size `(Nq, Nq,
nelem)`.  Similarly, the face arrays `sJ`, `n1`, and `n2` should be of size
`(Nq, nface, nelem)` with `nface = 4`.
"""
function computemetric!(x1, x2, J, ξ1x1, ξ2x1, ξ1x2, ξ2x2, sJ, n1, n2, D)
    T = eltype(x1)
    Nq = size(D, 1)
    nelem = div(length(J), Nq^2)
    d = 2
    x1 = reshape(x1, (Nq, Nq, nelem))
    x2 = reshape(x2, (Nq, Nq, nelem))
    J = reshape(J, (Nq, Nq, nelem))
    ξ1x1 = reshape(ξ1x1, (Nq, Nq, nelem))
    ξ2x1 = reshape(ξ2x1, (Nq, Nq, nelem))
    ξ1x2 = reshape(ξ1x2, (Nq, Nq, nelem))
    ξ2x2 = reshape(ξ2x2, (Nq, Nq, nelem))
    nface = 4
    n1 = reshape(n1, (Nq, nface, nelem))
    n2 = reshape(n2, (Nq, nface, nelem))
    sJ = reshape(sJ, (Nq, nface, nelem))

    @inbounds for e in 1:nelem
        for j in 1:Nq, i in 1:Nq
            x1ξ1 = x1ξ2 = zero(T)
            x2ξ1 = x2ξ2 = zero(T)
            for n in 1:Nq
                x1ξ1 += D[i, n] * x1[n, j, e]
                x1ξ2 += D[j, n] * x1[i, n, e]
                x2ξ1 += D[i, n] * x2[n, j, e]
                x2ξ2 += D[j, n] * x2[i, n, e]
            end
            J[i, j, e] = x1ξ1 * x2ξ2 - x2ξ1 * x1ξ2
            ξ1x1[i, j, e] = x2ξ2 / J[i, j, e]
            ξ2x1[i, j, e] = -x2ξ1 / J[i, j, e]
            ξ1x2[i, j, e] = -x1ξ2 / J[i, j, e]
            ξ2x2[i, j, e] = x1ξ1 / J[i, j, e]
        end

        for i in 1:Nq
            n1[i, 1, e] = -J[1, i, e] * ξ1x1[1, i, e]
            n2[i, 1, e] = -J[1, i, e] * ξ1x2[1, i, e]
            n1[i, 2, e] = J[Nq, i, e] * ξ1x1[Nq, i, e]
            n2[i, 2, e] = J[Nq, i, e] * ξ1x2[Nq, i, e]
            n1[i, 3, e] = -J[i, 1, e] * ξ2x1[i, 1, e]
            n2[i, 3, e] = -J[i, 1, e] * ξ2x2[i, 1, e]
            n1[i, 4, e] = J[i, Nq, e] * ξ2x1[i, Nq, e]
            n2[i, 4, e] = J[i, Nq, e] * ξ2x2[i, Nq, e]

            for n in 1:4
                sJ[i, n, e] = hypot(n1[i, n, e], n2[i, n, e])
                n1[i, n, e] /= sJ[i, n, e]
                n2[i, n, e] /= sJ[i, n, e]
            end
        end
    end

    nothing
end

"""
    computemetric!(x1, x2, x3, J, ξ1x1, ξ2x1, ξ3x1, ξ1x2, ξ2x2, ξ3x2, ξ1x3,
                   ξ2x3, ξ3x3, sJ, n1, n2, n3, D)

Compute the 3-D metric terms from the element grid arrays `x1`, `x2`, and `x3`.
All the arrays are preallocated by the user and the (square) derivative matrix
`D` should be consistent with the reference grid `r` used in
[`creategrid!`](@ref).

If `Nq = size(D, 1)` and `nelem = div(length(x1), Nq^3)` then the volume arrays
`x1`, `x2`, `x3`, `J`, `ξ1x1`, `ξ2x1`, `ξ3x1`, `ξ1x2`, `ξ2x2`, `ξ3x2`, `ξ1x3`,
`ξ2x3`, and `ξ3x3` should all be of length `Nq^3 * nelem`.  Similarly, the face
arrays `sJ`, `n1`, `n2`, and `n3` should be of size `Nq^2 * nface * nelem` with
`nface = 6`.

The curl invariant formulation of Kopriva (2006), equation 37, is used.

Reference:
  Kopriva, David A. "Metric identities and the discontinuous spectral element
  method on curvilinear meshes." Journal of Scientific Computing 26.3 (2006):
  301-327. <https://doi.org/10.1007/s10915-005-9070-8>
"""
function computemetric!(
    x1,
    x2,
    x3,
    J,
    ξ1x1,
    ξ2x1,
    ξ3x1,
    ξ1x2,
    ξ2x2,
    ξ3x2,
    ξ1x3,
    ξ2x3,
    ξ3x3,
    sJ,
    n1,
    n2,
    n3,
    D,
)
    T = eltype(x1)

    Nq = size(D, 1)
    nelem = div(length(J), Nq^3)

    x1 = reshape(x1, (Nq, Nq, Nq, nelem))
    x2 = reshape(x2, (Nq, Nq, Nq, nelem))
    x3 = reshape(x3, (Nq, Nq, Nq, nelem))
    J = reshape(J, (Nq, Nq, Nq, nelem))
    ξ1x1 = reshape(ξ1x1, (Nq, Nq, Nq, nelem))
    ξ2x1 = reshape(ξ2x1, (Nq, Nq, Nq, nelem))
    ξ3x1 = reshape(ξ3x1, (Nq, Nq, Nq, nelem))
    ξ1x2 = reshape(ξ1x2, (Nq, Nq, Nq, nelem))
    ξ2x2 = reshape(ξ2x2, (Nq, Nq, Nq, nelem))
    ξ3x2 = reshape(ξ3x2, (Nq, Nq, Nq, nelem))
    ξ1x3 = reshape(ξ1x3, (Nq, Nq, Nq, nelem))
    ξ2x3 = reshape(ξ2x3, (Nq, Nq, Nq, nelem))
    ξ3x3 = reshape(ξ3x3, (Nq, Nq, Nq, nelem))

    nface = 6
    #= This code is broken when views are used
    n1 = reshape(n1, Nq, Nq, nface, nelem)
    n2 = reshape(n2, Nq, Nq, nface, nelem)
    n3 = reshape(n3, Nq, Nq, nface, nelem)
    sJ = reshape(sJ, Nq, Nq, nface, nelem)
    =#

    JI2 = similar(J, (Nq, Nq, Nq))
    (yzr, yzs, yzt) = (similar(JI2), similar(JI2), similar(JI2))
    (zxr, zxs, zxt) = (similar(JI2), similar(JI2), similar(JI2))
    (xyr, xys, xyt) = (similar(JI2), similar(JI2), similar(JI2))

    ξ1x1 .= zero(T)
    ξ2x1 .= zero(T)
    ξ3x1 .= zero(T)
    ξ1x2 .= zero(T)
    ξ2x2 .= zero(T)
    ξ3x2 .= zero(T)
    ξ1x3 .= zero(T)
    ξ2x3 .= zero(T)
    ξ3x3 .= zero(T)

    @inbounds for e in 1:nelem
        for k in 1:Nq, j in 1:Nq, i in 1:Nq
            x1ξ1 = x1ξ2 = x1ξ3 = zero(T)
            x2ξ1 = x2ξ2 = x2ξ3 = zero(T)
            x3ξ1 = x3ξ2 = x3ξ3 = zero(T)
            for n in 1:Nq
                x1ξ1 += D[i, n] * x1[n, j, k, e]
                x1ξ2 += D[j, n] * x1[i, n, k, e]
                x1ξ3 += D[k, n] * x1[i, j, n, e]
                x2ξ1 += D[i, n] * x2[n, j, k, e]
                x2ξ2 += D[j, n] * x2[i, n, k, e]
                x2ξ3 += D[k, n] * x2[i, j, n, e]
                x3ξ1 += D[i, n] * x3[n, j, k, e]
                x3ξ2 += D[j, n] * x3[i, n, k, e]
                x3ξ3 += D[k, n] * x3[i, j, n, e]
            end
            J[i, j, k, e] = (
                x1ξ1 * (x2ξ2 * x3ξ3 - x3ξ2 * x2ξ3) +
                x2ξ1 * (x3ξ2 * x1ξ3 - x1ξ2 * x3ξ3) +
                x3ξ1 * (x1ξ2 * x2ξ3 - x2ξ2 * x1ξ3)
            )

            JI2[i, j, k] = 1 / (2 * J[i, j, k, e])

            yzr[i, j, k] = x2[i, j, k, e] * x3ξ1 - x3[i, j, k, e] * x2ξ1
            yzs[i, j, k] = x2[i, j, k, e] * x3ξ2 - x3[i, j, k, e] * x2ξ2
            yzt[i, j, k] = x2[i, j, k, e] * x3ξ3 - x3[i, j, k, e] * x2ξ3
            zxr[i, j, k] = x3[i, j, k, e] * x1ξ1 - x1[i, j, k, e] * x3ξ1
            zxs[i, j, k] = x3[i, j, k, e] * x1ξ2 - x1[i, j, k, e] * x3ξ2
            zxt[i, j, k] = x3[i, j, k, e] * x1ξ3 - x1[i, j, k, e] * x3ξ3
            xyr[i, j, k] = x1[i, j, k, e] * x2ξ1 - x2[i, j, k, e] * x1ξ1
            xys[i, j, k] = x1[i, j, k, e] * x2ξ2 - x2[i, j, k, e] * x1ξ2
            xyt[i, j, k] = x1[i, j, k, e] * x2ξ3 - x2[i, j, k, e] * x1ξ3
        end

        for k in 1:Nq, j in 1:Nq, i in 1:Nq
            for n in 1:Nq
                ξ1x1[i, j, k, e] += D[j, n] * yzt[i, n, k]
                ξ1x1[i, j, k, e] -= D[k, n] * yzs[i, j, n]
                ξ2x1[i, j, k, e] += D[k, n] * yzr[i, j, n]
                ξ2x1[i, j, k, e] -= D[i, n] * yzt[n, j, k]
                ξ3x1[i, j, k, e] += D[i, n] * yzs[n, j, k]
                ξ3x1[i, j, k, e] -= D[j, n] * yzr[i, n, k]
                ξ1x2[i, j, k, e] += D[j, n] * zxt[i, n, k]
                ξ1x2[i, j, k, e] -= D[k, n] * zxs[i, j, n]
                ξ2x2[i, j, k, e] += D[k, n] * zxr[i, j, n]
                ξ2x2[i, j, k, e] -= D[i, n] * zxt[n, j, k]
                ξ3x2[i, j, k, e] += D[i, n] * zxs[n, j, k]
                ξ3x2[i, j, k, e] -= D[j, n] * zxr[i, n, k]
                ξ1x3[i, j, k, e] += D[j, n] * xyt[i, n, k]
                ξ1x3[i, j, k, e] -= D[k, n] * xys[i, j, n]
                ξ2x3[i, j, k, e] += D[k, n] * xyr[i, j, n]
                ξ2x3[i, j, k, e] -= D[i, n] * xyt[n, j, k]
                ξ3x3[i, j, k, e] += D[i, n] * xys[n, j, k]
                ξ3x3[i, j, k, e] -= D[j, n] * xyr[i, n, k]
            end
            ξ1x1[i, j, k, e] *= JI2[i, j, k]
            ξ2x1[i, j, k, e] *= JI2[i, j, k]
            ξ3x1[i, j, k, e] *= JI2[i, j, k]
            ξ1x2[i, j, k, e] *= JI2[i, j, k]
            ξ2x2[i, j, k, e] *= JI2[i, j, k]
            ξ3x2[i, j, k, e] *= JI2[i, j, k]
            ξ1x3[i, j, k, e] *= JI2[i, j, k]
            ξ2x3[i, j, k, e] *= JI2[i, j, k]
            ξ3x3[i, j, k, e] *= JI2[i, j, k]
        end

        for j in 1:Nq, i in 1:Nq
            #= This code is broken when views are used
            n1[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x1[ 1, i, j, e]
            n1[i, j, 2, e] =  J[Nq, i, j, e] * ξ1x1[Nq, i, j, e]
            n1[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x1[ i, 1, j, e]
            n1[i, j, 4, e] =  J[ i,Nq, j, e] * ξ2x1[ i,Nq, j, e]
            n1[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x1[ i, j, 1, e]
            n1[i, j, 6, e] =  J[ i, j,Nq, e] * ξ3x1[ i, j,Nq, e]
            n2[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x2[ 1, i, j, e]
            n2[i, j, 2, e] =  J[Nq, i, j, e] * ξ1x2[Nq, i, j, e]
            n2[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x2[ i, 1, j, e]
            n2[i, j, 4, e] =  J[ i,Nq, j, e] * ξ2x2[ i,Nq, j, e]
            n2[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x2[ i, j, 1, e]
            n2[i, j, 6, e] =  J[ i, j,Nq, e] * ξ3x2[ i, j,Nq, e]
            n3[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x3[ 1, i, j, e]
            n3[i, j, 2, e] =  J[Nq, i, j, e] * ξ1x3[Nq, i, j, e]
            n3[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x3[ i, 1, j, e]
            n3[i, j, 4, e] =  J[ i,Nq, j, e] * ξ2x3[ i,Nq, j, e]
            n3[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x3[ i, j, 1, e]
            n3[i, j, 6, e] =  J[ i, j,Nq, e] * ξ3x3[ i, j,Nq, e]

            for n = 1:6
              sJ[i, j, n, e] = hypot(n1[i, j, n, e], n2[i, j, n, e], n3[i, j, n, e])
              n1[i, j, n, e] /= sJ[i, j, n, e]
              n2[i, j, n, e] /= sJ[i, j, n, e]
              n3[i, j, n, e] /= sJ[i, j, n, e]
            end
            =#

            ije = i + (j - 1) * Nq + (e - 1) * nface * Nq^2
            n1[ije + (1 - 1) * Nq^2] = -J[1, i, j, e] * ξ1x1[1, i, j, e]
            n1[ije + (2 - 1) * Nq^2] = J[Nq, i, j, e] * ξ1x1[Nq, i, j, e]
            n1[ije + (3 - 1) * Nq^2] = -J[i, 1, j, e] * ξ2x1[i, 1, j, e]
            n1[ije + (4 - 1) * Nq^2] = J[i, Nq, j, e] * ξ2x1[i, Nq, j, e]
            n1[ije + (5 - 1) * Nq^2] = -J[i, j, 1, e] * ξ3x1[i, j, 1, e]
            n1[ije + (6 - 1) * Nq^2] = J[i, j, Nq, e] * ξ3x1[i, j, Nq, e]
            n2[ije + (1 - 1) * Nq^2] = -J[1, i, j, e] * ξ1x2[1, i, j, e]
            n2[ije + (2 - 1) * Nq^2] = J[Nq, i, j, e] * ξ1x2[Nq, i, j, e]
            n2[ije + (3 - 1) * Nq^2] = -J[i, 1, j, e] * ξ2x2[i, 1, j, e]
            n2[ije + (4 - 1) * Nq^2] = J[i, Nq, j, e] * ξ2x2[i, Nq, j, e]
            n2[ije + (5 - 1) * Nq^2] = -J[i, j, 1, e] * ξ3x2[i, j, 1, e]
            n2[ije + (6 - 1) * Nq^2] = J[i, j, Nq, e] * ξ3x2[i, j, Nq, e]
            n3[ije + (1 - 1) * Nq^2] = -J[1, i, j, e] * ξ1x3[1, i, j, e]
            n3[ije + (2 - 1) * Nq^2] = J[Nq, i, j, e] * ξ1x3[Nq, i, j, e]
            n3[ije + (3 - 1) * Nq^2] = -J[i, 1, j, e] * ξ2x3[i, 1, j, e]
            n3[ije + (4 - 1) * Nq^2] = J[i, Nq, j, e] * ξ2x3[i, Nq, j, e]
            n3[ije + (5 - 1) * Nq^2] = -J[i, j, 1, e] * ξ3x3[i, j, 1, e]
            n3[ije + (6 - 1) * Nq^2] = J[i, j, Nq, e] * ξ3x3[i, j, Nq, e]

            for n in 1:6
                sJ[ije + (n - 1) * Nq^2] = hypot(
                    n1[ije + (n - 1) * Nq^2],
                    n2[ije + (n - 1) * Nq^2],
                    n3[ije + (n - 1) * Nq^2],
                )
                n1[ije + (n - 1) * Nq^2] /= sJ[ije + (n - 1) * Nq^2]
                n2[ije + (n - 1) * Nq^2] /= sJ[ije + (n - 1) * Nq^2]
                n3[ije + (n - 1) * Nq^2] /= sJ[ije + (n - 1) * Nq^2]
            end
        end
    end

    nothing
end

creategrid1d(elemtocoord, r) = creategrid(Val(1), elemtocoord, r)
creategrid2d(elemtocoord, r) = creategrid(Val(2), elemtocoord, r)
creategrid3d(elemtocoord, r) = creategrid(Val(3), elemtocoord, r)

"""
    creategrid(::Val{1}, elemtocoord::AbstractArray{S, 3},
               r::AbstractVector{T}) where {S, T}

Create a grid using `elemtocoord` (see [`brickmesh`](@ref)) using the 1-D `(-1,
1)` reference coordinates `r`. The element grids are filled using bilinear
interpolation of the element coordinates.

The grid is returned as a tuple of with `x1` array
"""
function creategrid(
    ::Val{1},
    e2c::AbstractArray{S, 3},
    r::AbstractVector{T},
) where {S, T}
    (d, nvert, nelem) = size(e2c)
    @assert d == 1
    Nq = length(r)
    x1 = Array{T, 2}(undef, Nq, nelem)
    creategrid!(x1, e2c, r)
    (x1 = x1,)
end

"""
  creategrid(::Val{2}, elemtocoord::AbstractArray{S, 3},
             r::AbstractVector{T}) where {S, T}

Create a 2-D tensor product grid using `elemtocoord` (see [`brickmesh`](@ref))
using the 1-D `(-1, 1)` reference coordinates `r`. The element grids are filled
using bilinear interpolation of the element coordinates.

The grid is returned as a tuple of the `x1` and `x2` arrays
"""
function creategrid(
    ::Val{2},
    e2c::AbstractArray{S, 3},
    r::AbstractVector{T},
) where {S, T}
    (d, nvert, nelem) = size(e2c)
    @assert d == 2
    Nq = length(r)
    x1 = Array{T, 3}(undef, Nq, Nq, nelem)
    x2 = Array{T, 3}(undef, Nq, Nq, nelem)
    creategrid!(x1, x2, e2c, r)
    (x1 = x1, x2 = x2)
end

"""
  creategrid(::Val{3}, elemtocoord::AbstractArray{S, 3},
             r::AbstractVector{T}) where {S, T}

Create a 3-D tensor product grid using `elemtocoord` (see [`brickmesh`](@ref))
using the 1-D `(-1, 1)` reference coordinates `r`. The element grids are filled
using bilinear interpolation of the element coordinates.

The grid is returned as a tuple of the `x1`, `x2`, `x3` arrays
"""
function creategrid(
    ::Val{3},
    e2c::AbstractArray{S, 3},
    r::AbstractVector{T},
) where {S, T}
    (d, nvert, nelem) = size(e2c)
    @assert d == 3
    Nq = length(r)
    x1 = Array{T, 4}(undef, Nq, Nq, Nq, nelem)
    x2 = Array{T, 4}(undef, Nq, Nq, Nq, nelem)
    x3 = Array{T, 4}(undef, Nq, Nq, Nq, nelem)
    creategrid!(x1, x2, x3, e2c, r)
    (x1 = x1, x2 = x2, x3 = x3)
end

"""
  computemetric(x1::AbstractArray{T, 2}, D::AbstractMatrix{T}) where T

Compute the 1-D metric terms from the element grid array `x1` using the
derivative matrix `D`. The derivative matrix `D` should be consistent with the
reference grid `r` used in [`creategrid!`](@ref).

The metric terms are returned as a 'NamedTuple` of the following arrays:

- `J` the Jacobian determinant
- `ξ1x1` derivative ∂r / ∂x1'
- `sJ` the surface Jacobian
- 'n1` outward pointing unit normal in \$x1\$-direction
"""
function computemetric(x1::AbstractArray{T, 2}, D::AbstractMatrix{T}) where {T}

    Nq = size(D, 1)
    nelem = size(x1, 2)
    nface = 2

    J = similar(x1)
    ξ1x1 = similar(x1)

    sJ = Array{T, 3}(undef, 1, nface, nelem)
    n1 = Array{T, 3}(undef, 1, nface, nelem)

    computemetric!(x1, J, ξ1x1, sJ, n1, D)

    (J = J, ξ1x1 = ξ1x1, sJ = sJ, n1 = n1)
end

"""
  computemetric(x1::AbstractArray{T, 3}, x2::AbstractArray{T, 3},
                D::AbstractMatrix{T}) where T

Compute the 2-D metric terms from the element grid arrays `x1` and `x2` using
the derivative matrix `D`. The derivative matrix `D` should be consistent with
the reference grid `r` used in [`creategrid!`](@ref).

The metric terms are returned as a 'NamedTuple` of the following arrays:

- `J` the Jacobian determinant
- `ξ1x1` derivative ∂ξ1 / ∂x1'
- `ξ2x1` derivative ∂ξ2 / ∂x1'
 - `ξ1x2` derivative ∂ξ1 / ∂x2'
 - `ξ2x2` derivative ∂ξ2 / ∂x2'
 - `sJ` the surface Jacobian
 - 'n1` outward pointing unit normal in \$x1\$-direction
 - 'n2` outward pointing unit normal in \$x2\$-direction
"""
function computemetric(
    x1::AbstractArray{T, 3},
    x2::AbstractArray{T, 3},
    D::AbstractMatrix{T},
) where {T}
    @assert size(x1) == size(x2)
    Nq = size(D, 1)
    nelem = size(x1, 3)
    nface = 4

    J = similar(x1)
    ξ1x1 = similar(x1)
    ξ2x1 = similar(x1)
    ξ1x2 = similar(x1)
    ξ2x2 = similar(x1)

    sJ = Array{T, 3}(undef, Nq, nface, nelem)
    n1 = Array{T, 3}(undef, Nq, nface, nelem)
    n2 = Array{T, 3}(undef, Nq, nface, nelem)

    computemetric!(x1, x2, J, ξ1x1, ξ2x1, ξ1x2, ξ2x2, sJ, n1, n2, D)

    (
        J = J,
        ξ1x1 = ξ1x1,
        ξ2x1 = ξ2x1,
        ξ1x2 = ξ1x2,
        ξ2x2 = ξ2x2,
        sJ = sJ,
        n1 = n1,
        n2 = n2,
    )
end

"""
    computemetric(x1::AbstractArray{T, 3}, x2::AbstractArray{T, 3},
                  D::AbstractMatrix{T}) where T

Compute the 3-D metric terms from the element grid arrays `x1`, `x2`, and `x3`
using the derivative matrix `D`. The derivative matrix `D` should be consistent
with the reference grid `r` used in [`creategrid!`](@ref).

The metric terms are returned as a 'NamedTuple` of the following arrays:

 - `J` the Jacobian determinant
 - `ξ1x1` derivative ∂ξ1 / ∂x1'
 - `ξ2x1` derivative ∂ξ2 / ∂x1'
 - `ξ3x1` derivative ∂ξ3 / ∂x1'
 - `ξ1x2` derivative ∂ξ1 / ∂x2'
 - `ξ2x2` derivative ∂ξ2 / ∂x2'
 - `ξ3x2` derivative ∂ξ3 / ∂x2'
 - `ξ1x3` derivative ∂ξ1 / ∂x3'
 - `ξ2x3` derivative ∂ξ2 / ∂x3'
 - `ξ3x3` derivative ∂ξ3 / ∂x3'
 - `sJ` the surface Jacobian
 - 'n1` outward pointing unit normal in \$x1\$-direction
 - 'n2` outward pointing unit normal in \$x2\$-direction
 - 'n3` outward pointing unit normal in \$x3\$-direction

!!! note

   The storage of the volume terms and surface terms from this function are
   slightly different. The volume terms used Cartesian indexing whereas the
   surface terms use linear indexing.
"""
function computemetric(
    x1::AbstractArray{T, 4},
    x2::AbstractArray{T, 4},
    x3::AbstractArray{T, 4},
    D::AbstractMatrix{T},
) where {T}

    @assert size(x1) == size(x2) == size(x3)
    Nq = size(D, 1)
    nelem = size(x1, 4)
    nface = 6

    J = similar(x1)
    ξ1x1 = similar(x1)
    ξ2x1 = similar(x1)
    ξ3x1 = similar(x1)
    ξ1x2 = similar(x1)
    ξ2x2 = similar(x1)
    ξ3x2 = similar(x1)
    ξ1x3 = similar(x1)
    ξ2x3 = similar(x1)
    ξ3x3 = similar(x1)

    sJ = Array{T, 3}(undef, Nq^2, nface, nelem)
    n1 = Array{T, 3}(undef, Nq^2, nface, nelem)
    n2 = Array{T, 3}(undef, Nq^2, nface, nelem)
    n3 = Array{T, 3}(undef, Nq^2, nface, nelem)

    computemetric!(
        x1,
        x2,
        x3,
        J,
        ξ1x1,
        ξ2x1,
        ξ3x1,
        ξ1x2,
        ξ2x2,
        ξ3x2,
        ξ1x3,
        ξ2x3,
        ξ3x3,
        sJ,
        n1,
        n2,
        n3,
        D,
    )

    (
        J = J,
        ξ1x1 = ξ1x1,
        ξ2x1 = ξ2x1,
        ξ3x1 = ξ3x1,
        ξ1x2 = ξ1x2,
        ξ2x2 = ξ2x2,
        ξ3x2 = ξ3x2,
        ξ1x3 = ξ1x3,
        ξ2x3 = ξ2x3,
        ξ3x3 = ξ3x3,
        sJ = sJ,
        n1 = n1,
        n2 = n2,
        n3 = n3,
    )
end

end # module


csv_dir = Pkg.dir("D4M") * "/examples/1Intro/2EdgeArt/Edge.csv"
print(csv_dir)

E = ReadCSV(csv_dir);                       # Read CSV file into associative array.
Ev = logical( E[:, StartsWith("V,")] );        # Get vertices and convert to numbers.

EvO = Ev[StartsWith("O,"),:];                  # Get orange edges.
EvG = Ev[StartsWith("G,"),:];                  # Get green edges.

AvOG = transpose(EvO) * EvG;                   # Compute vertex adjacency graph.
       print(AvOG);                             # Show empty structure.

       AeOG = EvO * transpose(EvG);                   # Compute edge adjacency graph.
       print(AeOG);                             # Show in tabular form.

       AeOG = CatKeyMul(EvO,transpose(EvG));          # Compute edge adjacency graph preserving pedigree.
       print(AeOG);


using CVChannel
using Test

"""
This script verifies the multiplicativity of the depolarizing channel
with itself for small dimensions.
"""

@testset "Verifying Depolarizing Channel Multiplicativity for d=3" begin
    println("Here we verify multiplicativity of the depolarizing channel with itself for d=3.")
    d = 3
    scan_range = [0:0.05:1;]
    is_mult = true
    max_val = 0
    for q in scan_range
        println("Now checking q = ", q)
        depolqChan(X) = depolarizingChannel(X,q)
        depolq_chan = Choi(depolqChan,d,d)

        par_cv, = pptCV(parChoi(depolq_chan, depolq_chan), :dual)

        target = (d * (1-q) + q)^2
        !isapprox(par_cv,target,atol=3e-6) ? is_mult = false : nothing
        max_val < abs(par_cv-target) ? max_val = abs(par_cv-target) : nothing
    end
    @test is_mult
end

"""
This verifies multiplicativity for d=4 but it takes forever
and doesn't matter for running in parallel with other channels
because it is too big.

@testset "Verifying Multiplicativity for d=4" begin
    d = 4
    scan_range = [0.75:0.05:1;]
    is_mult = true
    max_val = 2.264111881089e-6
    for q in scan_range
        println("Now checking q =", q, ".")
        depolqChan(X) = depolarizingChannel(X,q)
        depolq_chan = Choi(depolqChan,d,d)

        par_dims = d*[1,1,1,1]
        par_JN = permuteSubsystems(kron(depolq_chan.JN, depolq_chan.JN), [1,3,2,4], par_dims)
        par_in_dim = d^2
        par_out_dim = d^2
        par_cv, = pptCVDual(par_JN, d^2, d^2)

        target = (d*(1-q)+q)^2
        !isapprox(par_cv,target,atol=3e-6) ? is_mult = false : nothing
        max_val < abs(par_cv-target) ? max_val = abs(par_cv-target) : nothing
    end
    @test is_mult
end

"""


@testset "$(rpad("Order Conditions",80))" begin

    for s in 1:5
        g = TableauGauss(s)

        for σ in 1:s
            @test check_order_conditions_b(g,σ) == true
            @test check_order_conditions_c(g,σ) == Array{Bool}(ones(s))
            @test check_order_conditions_d(g,σ) == Array{Bool}(ones(s))
        end

        @test satisfies_simplifying_assumption_b(g)
        @test satisfies_simplifying_assumption_c(g)
        @test satisfies_simplifying_assumption_d(g)
    end

    for s in 2:5
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIIA(s))
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIIB(s))
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIIC(s))
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIIC̄(s))
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIID(s))
        @test  satisfies_simplifying_assumption_b(TableauLobattoIIIE(s))
        @test  satisfies_simplifying_assumption_b(TableauRadauIA(s))
        @test  satisfies_simplifying_assumption_b(TableauRadauIIA(s))

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

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

end


"""
Vshale(args...)

Calculate the Vshale from logs

`` Vshale=\\frac { { GR }_{ log }-{ GR }_{ Sand } }{ { GR }_{ Sh }-{ GR }_{ Sand } } ``

## Input Example

```julia
Vshale(Gr,GrSand,GrShale)
```
"""
function Vshale(Gr,GrSand,GrShale)
    Vsh=(Gr.-GrSand)./(GrShale.-GrSand)
end


#
# cg-tests.jl -
#
# Tests for linear conjugate gradients methods.
#

#isdefined(:LazyAlgebra) || include("../src/LazyAlgebra.jl")

module TestingLazyAlgebraConjGrad

using LazyAlgebra
using Test

const DEBUG = true

@testset "ConjGrad" begin
    types = (Float32, Float64)
    rows = (7,8)
    cols = (2,3,4)
    @testset "least-square fit ($T)" for T in types
        H = GeneralMatrix(randn(T, rows..., cols...))
        x = randn(T, cols)
        y = H*x + 0.01*randn(T, rows)

        # Least-squares solution using Julia linear algebra.
        yj = reshape(y, length(y))
        Hj = reshape(coefficients(H), length(y), length(x))
        Aj = Hj'*Hj
        bj = Hj'*yj
        x0 = reshape(Aj\bj, cols)

        # Tolerances.
        rtol = 1e-4
        atol = rtol*vnorm2(x0)

        # LHS matrix and RHS vector of the normal equations.
        A = H'*H
        b = H'*y
        n = length(b)
        x1 = conjgrad(A, b; maxiter=10n, restart=n,
                      quiet=false, verb=true,
                      gtol=(0,0), ftol=0, xtol=rtol/10)
        DEBUG && println("x1: ", vnorm2(x1 - x0)/vnorm2(x0))
        @test vnorm2(x1 - x0) ≤ atol

        # Exercise random starting point.
        x2 = conjgrad(A, b, randn(T, cols); maxiter=2n, restart=n,
                      quiet=true, verb=false,
                      gtol=(0,0), ftol=0, xtol=rtol/10)
        DEBUG && println("x2: ", vnorm2(x2 - x0)/vnorm2(x0))
        @test vnorm2(x2 - x0) ≤ atol

        # Exercise other convergence tests.
        x3 = conjgrad(A, b; maxiter=n, restart=n,
                      quiet=false, verb=true,
                      gtol=(0,0), ftol=0, xtol=0)
        DEBUG && println("x3: ", vnorm2(x3 - x0)/vnorm2(x0))
        @test vnorm2(x3 - x0) ≤ atol
        x4 = conjgrad(A, b; maxiter=10n, restart=n,
                      quiet=false, verb=true,
                      gtol=(0,0), ftol=1e-9, xtol=0)
        DEBUG && println("x4: ", vnorm2(x4 - x0)/vnorm2(x0))
        @test vnorm2(x4 - x0) ≤ atol
        x5 = conjgrad(A, b; maxiter=10n, restart=n,
                      quiet=false, verb=true,
                      gtol=(0,1e-6), ftol=0, xtol=0)
        DEBUG && println("x5: ", vnorm2(x5 - x0)/vnorm2(x0))
        @test vnorm2(x5 - x0) ≤ atol

        # Directly use an array.
        x6 = conjgrad(reshape(Aj, cols..., cols...), b;
                      maxiter=2n, restart=n,
                      quiet=true, verb=false,
                      gtol=(0,0), ftol=0, xtol=rtol/10)
        DEBUG && println("x6: ", vnorm2(x6 - x0)/vnorm2(x0))
        @test vnorm2(x6 - x0) ≤ atol

        # Force no iterations (should yield exactly the initial solution).
        x7 = conjgrad(A, b, x0;
                      maxiter=0, restart=n,
                      quiet=true, verb=false,
                      gtol=(0,0), ftol=0, xtol=0)
        DEBUG && println("x7: ", vnorm2(x7 - x0)/vnorm2(x0))
        @test vnorm2(x7 - x0) ≤ 0

        # Force restarts.
        x8 = conjgrad(A, b;
                      maxiter=4n, restart=(n>>1),
                      quiet=true, verb=false,
                      gtol=(0,0), ftol=0, xtol=0)
        DEBUG && println("x8: ", vnorm2(x8 - x0)/vnorm2(x0))
        @test vnorm2(x8 - x0) ≤ atol

        # Non-positive definite.
        @test vnorm2(conjgrad(-A, b; strict=false)) ≤ 0
        @test_throws LazyAlgebra.NonPositiveDefinite conjgrad(-A, b; verb=true,
                                                              quiet=false)

    end
end

nothing

end # module


module DirichletFESpacesTests

using Test
using Gridap.Algebra
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.CellData

n = 2
mesh = (n,n)
domain = (0,1,0,1)
order = 1
dirichlet_tags = [1,3,4,6,7]
model = CartesianDiscreteModel(domain, mesh)
labels = get_face_labeling(model)
add_tag_from_tags!(labels,"dirichlet",dirichlet_tags)

Vf = TestFESpace(
  reffe=:Lagrangian, order=order, valuetype=Float64,
  conformity=:H1, model=model, labels=labels, dirichlet_tags="dirichlet")

Vd = DirichletFESpace(Vf)
matvecdata = ([],[],[])
matdata = ([],[],[])
vecdata = ([],[])
test_single_field_fe_space(Vd,matvecdata,matdata,vecdata)

Uf = TrialFESpace(Vf,1)

Ud = TrialFESpace(Vd)
test_single_field_fe_space(Ud,matvecdata,matdata,vecdata)

trian = Triangulation(model)
degree = 2*order
quad = CellQuadrature(trian,degree)

a(u,v) = u*v
l(v) = v*0
t_Ω = AffineFETerm(a,l,trian,quad)

op_ff = AffineFEOperator(Uf,Vf,t_Ω)
op_fd = AffineFEOperator(Ud,Vf,t_Ω)
op_df = AffineFEOperator(Uf,Vd,t_Ω)

xd = get_dirichlet_values(Uf)

@test maximum(abs.(get_vector(op_ff) + get_matrix(op_fd)*xd)) < 1.0e-10

end # module


#this example is modified from https://numericalenvironmental.wordpress.com/2018/05/26/a-steady-state-variably-saturated-flow-model-in-vertical-cross-section-a-finite-difference-approach-using-julia/
import NLsolve
import NonlinearEquations
import PyPlot
import Random
import Zygote

include("utilities.jl")
include("inputdeck.jl")

@NonlinearEquations.equations exclude=(coords, dirichletnodes, neighbors, areasoverlengths) function reqnumenv(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	NonlinearEquations.setnumequations(length(psi))
	for i = 1:length(psi)
		if i in dirichletnodes
			NonlinearEquations.addterm(i, psi[i] - dirichletpsis[i])
		else
			NonlinearEquations.addterm(i, Qs[i])
		end
	end
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
				relperm = hm(kr(psi[node1], alphas[i], Ns[i]), kr(psi[node2], alphas[i], Ns[i]))
				NonlinearEquations.addterm(node1, relperm * Ks[i] * (psi[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i])
			elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
				relperm = hm(kr(psi[node1], alphas[i], Ns[i]), kr(dirichletpsis[node2], alphas[i], Ns[i]))
				NonlinearEquations.addterm(node1, relperm * Ks[i] * (dirichletpsis[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i])
			end
		end
	end
end

@NonlinearEquations.equations exclude=(coords, dirichletnodes, neighbors, areasoverlengths) function req(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	NonlinearEquations.setnumequations(length(psi))
	for i = 1:length(psi)
		if i in dirichletnodes
			NonlinearEquations.addterm(i, psi[i] - dirichletpsis[i])
		else
			NonlinearEquations.addterm(i, Qs[i])
		end
	end
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(psi[node2], alphas[i], Ns[i])) * Ks[i] * (psi[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i])
			elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(dirichletpsis[node2], alphas[i], Ns[i])) * Ks[i] * (dirichletpsis[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i])
			end
		end
	end
end

psi0 = -ones(size(coords, 2))
psi0[dirichletnodes] = zeros(length(dirichletnodes))

function solveforpsi(Ks, psi0=psi0; doplot=false, donewtonish=true)
	resnumenv(psi) = reqnumenv_residuals(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	jacnumenv(psi) = reqnumenv_psi(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	res(psi) = req_residuals(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	jac(psi) = req_psi(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	j!(J, psi) = req_psi!(J, psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
	f!(residuals, psi) = copy!(residuals, res(psi))
	if donewtonish
		psi1 = NonlinearEquations.newtonish(resnumenv, jacnumenv, psi0; numiters=10, rate=0.05)
	else
		psi1 = psi0
	end
	psi1 = soln_dbe[:, end]
	#callback(psi1, res(psi1), jac(psi1), 0)
	df = NLsolve.OnceDifferentiable(f!, j!, psi1, res(psi0), jac(psi0))
	nls = NLsolve.nlsolve(df, psi1; show_trace=false, iterations=200, ftol=1e-12)
	if doplot
		callback(nls.zero, res(nls.zero), jac(nls.zero), 0)
	end
	return nls.zero
end

psi = solveforpsi(Ks; doplot=true)

#compute the gradient
Random.seed!(1)
obsnode = rand(1:length(psi))
g(psi, Ks) = psi[obsnode]
function g_psi(psi, p)
	retval = zeros(length(psi))
	retval[obsnode] = 1.0
	return retval
end
g_Ks(psi, Ks) = zeros(length(Ks))
f_psi(psi, Ks) = req_psi(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
f_Ks(psi, Ks) = req_Ks(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs)
grad = NonlinearEquations.gradient(psi, Ks, g_psi, g_Ks, f_psi, f_Ks)
plotgradient(grad, obsnode)

#=
fdgrad = similar(grad)
sortedgradindices = sort(1:length(grad); by=i->abs(grad[i]), rev=true)
@time for i = 1:length(fdgrad)
	global fdgrad
	global Ks
	global psi
	dk = 1e-8
	psi0 = psi
	Ks0 = copy(Ks)
	thisKs = copy(Ks)
	thisKs[i] += dk
	thispsi = solveforpsi(thisKs, psi0; donewtonish=false)
	fdgrad[i] = (thispsi[obsnode] - psi0[obsnode]) / dk
end
plotgradient(fdgrad, obsnode)
=#


function save_radau_table_to_file(n_rule::Int64)
    @eval begin  # https://stackoverflow.com/questions/26071317/declaring-top-level-variables-in-julia-using-metaprogramming/26071597#26071597
        n_rule = $n_rule
        A, b, c, bi, b_hat = radau_butcher_table_plus(2 * n_rule - 1)
        inv_A = inv(A)
        lambda, T = big_eigen(inv_A)
        inv_T = inv(T)
        f_64_type(v::VecOrMat{BigFloat}) = Float64
        f_64_type(v::VecOrMat{Complex{BigFloat}}) = Complex{Float64}
        symbols_save = (:A, :inv_A, :b, :c, :lambda, :T, :inv_T, :bi, :b_hat)
        symbols_type = f_64_type.(eval.(symbols_save))

        ### Table Data
        radau_table_path = get_radau_table_path()
        save_path = joinpath(radau_table_path, string(n_rule) * "_rule")
        for (k, var) = enumerate(symbols_save)
            file_path = joinpath(save_path, string(var) * ".txt")
            open(file_path, "w") do io
                d_type_k = symbols_type[k]
                writedlm(io, d_type_k.(eval(var)))
            end
        end
    end

    ### Table Types
    save_path_types = joinpath(radau_table_path, "type.txt")
    open(save_path_types, "w") do io
        writedlm(io, string.(symbols_type))
    end

    ### Table Symbols
    save_path_names = joinpath(radau_table_path, "symbols.txt")
    open(save_path_names, "w") do io
        writedlm(io, string.(symbols_save))
    end
end


using SparseArrays

@i function i_mul!(C::StridedVecOrMat, A::AbstractSparseMatrix, B::StridedVector{T}, α::Number, β::Number) where T
    @safe size(A, 2) == size(B, 1) || throw(DimensionMismatch())
    @safe size(A, 1) == size(C, 1) || throw(DimensionMismatch())
    @safe size(B, 2) == size(C, 2) || throw(DimensionMismatch())
    nzv ← nonzeros(A)
    rv ← rowvals(A)
    if (β != 1, ~)
        @safe error("only β = 1 is supported, got β = $(β).")
    end
    # Here, we close the reversibility check inside the loop to increase performance
    @invcheckoff for k = 1:size(C, 2)
        @inbounds for col = 1:size(A, 2)
            αxj ← zero(T)
            αxj += B[col,k] * α
            for j = SparseArrays.getcolptr(A)[col]:(SparseArrays.getcolptr(A)[col + 1] - 1)
                C[rv[j], k] += nzv[j]*αxj
            end
            αxj -= B[col,k] * α
        end
    end
end

@i function i_dot(r::T, A::SparseMatrixCSC{T},B::SparseMatrixCSC{T}) where {T}
    m ← size(A, 1)
    n ← size(A, 2)
    @invcheckoff branch_keeper ← zeros(Bool, 2*m)
    @safe size(B) == (m,n) || throw(DimensionMismatch("matrices must have the same dimensions"))
    @invcheckoff @inbounds for j = 1:n
        ia1 ← A.colptr[j]
        ib1 ← B.colptr[j]
        ia2 ← A.colptr[j+1]
        ib2 ← B.colptr[j+1]
        ia ← ia1
        ib ← ib1
        @inbounds for i=1:ia2-ia1+ib2-ib1-1
            ra ← A.rowval[ia]
            rb ← B.rowval[ib]
            if (ra == rb, ~)
                r += A.nzval[ia]' * B.nzval[ib]
            end
            ## b move -> true, a move -> false
            branch_keeper[i] ⊻= ia == ia2-1 || (ib != ib2-1 && ra > rb)
            ra → A.rowval[ia]
            rb → B.rowval[ib]
            if (branch_keeper[i], ~)
                INC(ib)
            else
                INC(ia)
            end
        end
        ~@inbounds for i=1:ia2-ia1+ib2-ib1-1
            ## b move -> true, a move -> false
            branch_keeper[i] ⊻= ia == ia2-1 || (ib != ib2-1 && A.rowval[ia] > B.rowval[ib])
            if (branch_keeper[i], ~)
                INC(ib)
            else
                INC(ia)
            end
        end
    end
    @invcheckoff branch_keeper → zeros(Bool, 2*m)
end


struct PieceWiseNormal <: Gen.Distribution{Float64} end
const piecewise_normal = PieceWiseNormal()

# piece 1 has interval (-Inf, boundaries[1])
# piece 2 has interval (boundaries[1], boundaries[2])
# ..
# piece n has interval (boundaries[n-1], boundaries[n])
# piece n+1 has interval (boundaries[n], Inf)

function Gen.random(::PieceWiseNormal, probabilities::Vector{Float64},
                    mus::Vector{Float64}, stds::Vector{Float64},
                    boundaries::Vector{Float64})
    n_segments = length(probabilities)
    @assert n_segments == length(mus)
    @assert n_segments == length(stds)
    @assert n_segments == length(boundaries)+1
    idx = categorical(probabilities)
    dist = Distributions.TruncatedNormal(
        mus[idx],
        stds[idx],
        idx == 1 ? -Inf : boundaries[idx-1], idx == n_segments ? Inf : boundaries[idx])
    rand(dist)
end

function Gen.logpdf(::PieceWiseNormal, x::Float64, probabilities::Vector{Float64},
                    mus::Vector{Float64}, stds::Vector{Float64},
                    boundaries::Vector{Float64})
    n_segments = length(probabilities)
    @assert n_segments == length(mus)
    @assert n_segments == length(stds)
    @assert n_segments == length(boundaries)+1
    log_probs = Vector{Float64}(undef, n_segments)
    for idx=1:length(probabilities)
        dist = Distributions.TruncatedNormal(
            mus[idx],
            stds[idx],
            idx == 1 ? -Inf : boundaries[idx-1], idx == n_segments ? Inf : boundaries[idx])
        log_probs[idx] = log(probabilities[idx]) + Distributions.logpdf(dist, x)
    end
    logsumexp(log_probs)
end

Gen.has_output_grad(::PieceWiseNormal) = false
Gen.has_argument_grads(::PieceWiseNormal) = (false, false, false, false)


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -11.08168115564476 <= q <= 12.629895565298366)

                     add_NL_constraint(m, :(swish(-0.8013086439641905 + -0.577380120181807*swish(-0.15161094843442724 + 0.3901524114721284*$(x[1]) + -0.7620628362671806*$(x[2]) + 0.5531251925343885*$(x[3]) + 0.8877751701397409*$(x[4])) + -0.9073468716494122*swish(-0.9517329762307463 + 0.7324043290359472*$(x[1]) + -0.6011940599882695*$(x[2]) + 0.8132214951552479*$(x[3]) + 0.2016959253369457*$(x[4])) + -0.34890000179172853*swish(0.5700564024677015 + -0.6508434167160906*$(x[1]) + 0.3531203702734036*$(x[2]) + 0.9592602519851474*$(x[3]) + -0.2457804299770685*$(x[4]))) + swish(-0.6400914085711236 + 0.5999574270201746*swish(-0.15161094843442724 + 0.3901524114721284*$(x[1]) + -0.7620628362671806*$(x[2]) + 0.5531251925343885*$(x[3]) + 0.8877751701397409*$(x[4])) + -0.3302091884595031*swish(-0.9517329762307463 + 0.7324043290359472*$(x[1]) + -0.6011940599882695*$(x[2]) + 0.8132214951552479*$(x[3]) + 0.2016959253369457*$(x[4])) + 0.9262651374928534*swish(0.5700564024677015 + -0.6508434167160906*$(x[1]) + 0.3531203702734036*$(x[2]) + 0.9592602519851474*$(x[3]) + -0.2457804299770685*$(x[4]))) + swish(0.7171327203675486 + 0.9873737246227612*swish(-0.15161094843442724 + 0.3901524114721284*$(x[1]) + -0.7620628362671806*$(x[2]) + 0.5531251925343885*$(x[3]) + 0.8877751701397409*$(x[4])) + -0.03055409136192777*swish(-0.9517329762307463 + 0.7324043290359472*$(x[1]) + -0.6011940599882695*$(x[2]) + 0.8132214951552479*$(x[3]) + 0.2016959253369457*$(x[4])) + 0.20254351925716962*swish(0.5700564024677015 + -0.6508434167160906*$(x[1]) + 0.3531203702734036*$(x[2]) + 0.9592602519851474*$(x[3]) + -0.2457804299770685*$(x[4]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using MathProgBase, GLPKMathProgInterface

include("misc.jl")

function leastMax(X, y)
  (n, d) = size(X)
  X = [ones(n, 1) X]
  d += 1
  y = y[:,1]
  c = [ones(n); zeros(d)]
  A = zeros(0, d + n)
  b = Float64[]

  for j = 1:d
    colj = zeros(d, d)
    colj[j,j] = 1
    A = [A; eye(n) X*colj; eye(n) -X*colj]
    b = [b; -y; y]
  end

  display(size(y))
  display(size(c))
  display(size(A))
  display(size(b))

  l = [zeros(n); ones(d) * -Inf]

  sol = linprog(c, A, '>', b, l, Inf, GLPKSolverLP())
  w = sol.sol[n+1:n+d]
  if sol.status == :Optimal
      println("Optimal objective value is $(sol.objval)")
      println("Optimal solution vector is: $(w)")
  else
      println("Error: solution status $(sol.status)")
  end
  predict(Xtilde) = [ones(size(Xtilde,1),1) Xtilde]*w

  # Return model
  return LinearModel(predict, w)
end


using NearestNeighbors
using Distances: chebyshev
using SpecialFunctions: digamma
using StatsBase: autocor
using KernelDensity

export mutinfo, mutinfo_delaycurve
#####################################################################################
#                                Mutual Information                                 #
#####################################################################################
"""
    mutinfo(k, X1, X2[, ..., Xm]) -> MI

Calculate the mutual information `MI` of the given vectors
`X1, X2, ...`, using `k` nearest-neighbors.

The method follows the second algorithm ``I^{(2)}`` outlined by Kraskov in [1].

## References
[1] : A. Kraskov *et al.*, [Phys. Rev. E **69**, pp 066138 (2004)](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.69.066138)

## Performance Notes
This functin gets very slow for large `k`.

See also [`estimate_delay`](@ref) and [`mutinfo_delaycurve`](@ref).
"""
function mutinfo(k, Xm::Vararg{<:AbstractVector,M}) where M
    @assert M > 1
    @assert (size.(Xm,1) .== size(Xm[1],1)) |> prod
    k += 1
    N = size(Xm[1],1)
    invN = 1/N

    d = Dataset(Xm...)
    tree = KDTree(d.data, Chebyshev())

    n_x_m = zeros(M)

    Xm_sp = zeros(Int, N, M)
    Xm_revsp = zeros(Int, N, M)
    for m in 1:M
        Xm_sp[:,m] .= sortperm(Xm[m]; alg=QuickSort)
        Xm_revsp[:,m] .= sortperm(Xm_sp[:,m]; alg=QuickSort)
    end

    I = digamma(k) - (M-1)/k + (M-1)*digamma(N)

    nns = (x = knn(tree, d.data, k)[1]; [ind[1] for ind in x])

    I_itr = zeros(M)
    # Makes more sense computationally to loop over N rather than M
    for i in 1:N
        ϵ = abs.(d[nns[i]] - d[i])./2

        for m in 1:M # this loop takes 8% of time
            hb = lb = Xm_revsp[i,m]
            while abs(Xm[m][Xm_sp[hb,m]] - Xm[m][i]) <= ϵ[m] && hb < N
                hb += 1
            end
            while abs(Xm[m][Xm_sp[lb,m]] - Xm[m][i]) <= ϵ[m] && lb > 1
                lb -= 1
            end
            n_x_m[m] = hb - lb
        end

        I_itr .+= digamma.(n_x_m)
    end

    I_itr .*= invN

    I -= sum(I_itr)

    return max(0, I)
end

"""
    mutinfo_delaycurve(x; maxtau=100, k=1)

Return the [`mutinfo`](@ref) between `x` and itself for delays of `1:maxtau`.
"""
function mutinfo_delaycurve(X::AbstractVector; maxtau=100, k=1)
    I = zeros(maxtau)

    @views for τ in 1:maxtau
        I[τ] = mutinfo(k, X[1:end-τ],X[τ+1:end])
    end

    return I
end

### Fraser's and Swinney's algorithm ###

# Minimum number of points in the margin of a node
const minnodesize = 6

"""
    FSNode

Tree structure for 4^m partitions of a plane:
`low`, `high` are the corners of a node of the partition.
`middle` is the middle point that defines the subpartition.
`children` contains the four sub-partitions.
"""
struct FSNode{T}
    low::Tuple{T,T}
    high::Tuple{T,T}
    middle::Tuple{T,T}
    children::Array{FSNode,1}
end

# Recursive creation of partitions
FSNode(low::Int, high::Int) = FSNode((low,low), (high,high))
function FSNode(low::Tuple{T,T}, high::Tuple{T,T}) where {T}
    half = [div(high[1]-low[1], 2), div(high[2]-low[2],2)]
    middle = (low[1]+half[1], low[2]+half[2])
    node = FSNode(low, high, middle, FSNode[])
    if maximum(half) > minnodesize # or another condition
        append!(node.children,
                [FSNode((low[1],     low[2]),      (middle[1],middle[2])),
                 FSNode((middle[1]+1,low[2]),      (high[1],  middle[2])),
                 FSNode((low[1],     middle[2]+1), (middle[1],high[2])),
                 FSNode((middle[1]+1,middle[2]+1), (high[1],  high[2]))])
    end
    return node
end

"""
    mi_fs(s,τs)

Compute the mutual information between the series `s` and various
images of the same series delayed by `τ` ∈ `τs`, according to
Fraser's and Swinney's algorithm [1].

## References:
[1]: Fraser A.M. & Swinney H.L. "Independent coordinates for strange attractors
from mutual information" *Phys. Rev. A 33*(2), 1986, 1134:1140.
"""
function mi_fs(s::AbstractVector,τs)
    n = length(s)
    nτ = n-maximum(τs)       # trim the series for the maximum delay
    tree = FSNode(1, nτ)      # tree structure of the subpartitions
    perm = sortperm(s[1:nτ])  # order of the original series
    mi = zeros(length(τs))
    for (i,τ) in enumerate(τs)
        indices = sortperm(sortperm(s[τ+1:n])) # rank values of the delayed image of `s`
        mi[i] = _recursive_mi(tree, indices[perm])/nτ - log2(nτ) # eq. (19) of [1]
    end
    return mi
end

# Recursive implementation of equations (20a) and (20b) in [1]
function _recursive_mi(node::FSNode, s)
    # Get view of the `s` within the first dimension of the node
    v = @view s[node.low[1]:node.high[1]]
    if isempty(node.children) # terminal node, no possible substructure
        # Count points in the range of the second dimension
        n = _count(v, node.low[2], node.high[2])
        mi_value = (n==0) ? 0 : n*log2(n)    # eq. (20a) in [1]
    else
        # Split the view in two halves of the first dimension
        # and count point in the two subranges of the second
        n1 = node.middle[1] - node.low[1] + 1
        vhalf = @view v[1:n1]
        n11, n21 = _count(vhalf, node.low[2], node.middle[2], node.high[2])
        vhalf = @view v[n1+1:end]
        n12, n22 = _count(vhalf, node.low[2], node.middle[2], node.high[2])
        n = n11+n12+n21+n22
        if _homogeneous(n11,n12,n21,n22) # no substructure (homogeneous table)
            mi_value = n*log2(n) # eq. (20a) in [1]
        else
            mi_value = 2.0*n     # eq. (20b) in [1]
            for (i, ni) in enumerate((n11,n12,n21,n22))
                (ni > 0) && (mi_value += _recursive_mi(node.children[i], s))
            end
        end
    end
    return mi_value
end

# count values of `s` inside `low:high`
function _count(s, low, high)
    n = 0
    @inbounds for i in s
        (low <= i <= high) && (n+= 1)
    end
    return n
end
# count values of `s` inside `low:middle` and `low:high`
function _count(s, low, middle, high)
    n1 = 0
    n2 = 0
    @inbounds for i in s
        if low <= i
            if i <= middle
                n1 += 1
            elseif i <= high
                n2 += 1
            end
        end
    end
    return n1, n2
end

# homogeneity test - equation (22) in [1]
function _homogeneous(n11,n12,n21,n22)
    n = n11+n12+n21+n22
    nh = n/4
    χ2 = 4/3/n*((n11-nh)^2+(n12-nh)^2+(n21-nh)^2+(n22-nh)^2)
    return χ2 < 1.547
end

function ami_kde(x::AbstractVector{T}, τs::AbstractVector{Int}) where {T}
    n = length(x)
    np = min(2048, prevpow(2, n))
    dx = kde(x, npoints=np)
    #
    tails_width = 4.0*KernelDensity.default_bandwidth(x)
    tp = round(Int, np*tails_width/(dx.x[end]-dx.x[1]))
    #
    # dx.density .*= step(dx.x)
    px = @view dx.density[tp+1:end-tp]
    hx = -sum(px.*log2.(px))*step(dx.x)
    ami = zeros(T, length(τs))
    #
    np2 = min(256, 2^floor(Int, log2(n)))
    tp = div(tp*np2,np)
    for (i,τ) in enumerate(τs)
        hxy = zero(T)
        dxy = kde((x[1:n-τ], x[τ+1:n]), npoints=(np2,np2))
        # dxy.density .*= (step(dxy.x) * step(dxy.y))
        pxy = @view dxy.density[tp+1:end-tp, tp+1:end-tp]
        for p in pxy
            if p ≉ 0
                hxy -= p*log2(p)
            end
        end
        hxy *= (step(dxy.x)*step(dxy.y))
        mi[i] = 2hx - hxy
    end
    return ami
end