module TestBootsampleNonParametric

using Bootstrap
using Bootstrap.Datasets
using FactCheck

using DataFrames
using StatsBase

function test_bootsample(bs, ref, raw_data, n)

    show(IOBuffer(), bs)
    @fact issubtype(typeof(bs), NonParametricBootstrapSample) --> true
    t0 = original(bs)
    @fact length(t0) --> length(ref)
    [@fact t --> roughly(r) for (t, r) in zip(t0, ref)]

    t1 = straps(bs)
    @fact length(t1) --> length(t0)
    [@fact length(t) --> n for t in t1]
    [@fact (minimum(t) <= tr && maximum(t) >= tr) --> true for (t, tr) in zip(t1, t0)]
    [@fact eltype(t) --> eltype(tr) for (t, tr) in zip(t1, t0)]

    @fact Bootstrap.data(bs) --> raw_data ## TODO: define scope

    @fact nrun(sampling(bs)) --> n
    @fact nrun(bs) --> n

    @fact length(bias(bs)) --> length(ref)
    [@fact (b[1] > -Inf && b[1] < Inf) --> true for b in bias(bs)]
    @fact length(se(bs)) --> length(ref)
    [@fact s --> greater_than(0) for s in se(bs)]

    [@fact original(bs, i) --> original(bs)[i]  for i in 1:nvar(bs)]
    [@fact straps(bs, i) --> straps(bs)[i]  for i in 1:nvar(bs)]
    [@fact bias(bs, i) --> bias(bs)[i]  for i in 1:nvar(bs)]
    [@fact se(bs, i) --> se(bs)[i]  for i in 1:nvar(bs)]

    @fact_throws MethodError model(bs)

    return Void
end

function test_ci(bs)

    cim_all = (BasicConfInt(), PercentileConfInt(), NormalConfInt(), BCaConfInt())
    for cim in cim_all
        c = ci(bs, cim)
        [@fact (x[1] >= x[2]  && x[1] <= x[3]) --> true for x in c]
        [@fact x[1] --> roughly(t0) for (x, t0) in zip(c, original(bs))]
        @fact level(cim) --> 0.95
    end

    return Void
end

n = 250

## 'city' dataset
citya = convert(DataArray, city)

city_ratio(df::DataFrames.DataFrame) = mean(df[:,:X]) ./ mean(df[:,:U])
city_ratio(a::AbstractArray) = mean(a[:,2]) ./ mean(a[:,1])

city_cor(x::AbstractArray) = cor(x[:,1], x[:,2])
city_cor(x::AbstractDataFrame) = cor(x[:,:X], x[:,:U])

facts("Basic resampling") do

    context("city_ratio with DataFrame input") do
        ref = city_ratio(city)
        @fact ref --> roughly(1.5203125)
        bs = bootstrap(city, city_ratio, BasicSampling(n))
        test_bootsample(bs, ref, city, n)
        test_ci(bs)
    end

    context("city_cor with DataFrame input") do
        ref = city_cor(city)
        bs = bootstrap(city, city_cor, BasicSampling(n))
        test_bootsample(bs, ref, city, n)
        test_ci(bs)
    end

    context("city_cor with DataArray input") do
        ref = city_cor(citya)
        bs = bootstrap(citya, city_cor, BasicSampling(n))
        test_bootsample(bs, ref, citya, n)
        test_ci(bs)
    end

    context("mean_and_sd: Vector input, 2 output variables") do
        r = randn(25)
        ref = mean_and_std(r)
        bs = bootstrap(r, mean_and_std, BasicSampling(n))
        test_bootsample(bs, ref, r, n)
        test_ci(bs)
    end

    context("mean_and_sd: Student CI") do
        r = randn(25)
        bs = bootstrap(r, mean_and_std, BasicSampling(n))
        ## Student confint
        cim = StudentConfInt()
        c = ci(bs, straps(bs, 2), cim, 1)
        @fact (c[1] >= c[2]  && c[1] <= c[3]) --> true
        @fact c[1] --> roughly(original(bs, 1))
        @fact level(cim) --> 0.95
    end

end

facts("Antithetic resampling") do

    context("mean_and_sd: Vector input, 2 output variables") do
        r = randn(50)
        ref = mean_and_std(r)
        bs = bootstrap(r, mean_and_std, AntitheticSampling(n))
        test_bootsample(bs, ref, r, n)
        test_ci(bs)
    end

end

facts("Balanced resampling") do

    context("city_ratio with DataFrame input") do
        ref = city_ratio(city)
        @fact ref --> roughly(1.5203125)
        bs = bootstrap(city, city_ratio, BalancedSampling(n))
        test_bootsample(bs, ref, city, n)
        test_ci(bs)
    end

    context("city_cor with DataFrame input") do
        ref = city_cor(city)
        bs = bootstrap(city, city_cor, BalancedSampling(n))
        test_bootsample(bs, ref, city, n)
        test_ci(bs)
    end

    context("city_cor with DataArray input") do
        ref = city_cor(citya)
        bs = bootstrap(citya, city_cor, BalancedSampling(n))
        test_bootsample(bs, ref, citya, n)
        test_ci(bs)
    end

    context("mean_and_sd: Vector input, 2 output variables") do
        r = randn(50)
        ref = mean_and_std(r)
        bs = bootstrap(r, mean_and_std, BalancedSampling(n))
        test_bootsample(bs, ref, r, n)
        test_ci(bs)
        ## mean should be unbiased
        @fact bias(bs)[1] --> roughly(0.0, 1e-8)
    end

end

facts("Exact resampling") do

    nc = Bootstrap.nrun_exact(nrow(city))

    context("city_ratio with DataFrame input") do
        ref = city_ratio(city)
        @fact ref --> roughly(1.5203125)
        bs = bootstrap(city, city_ratio, ExactSampling())
        test_bootsample(bs, ref, city, nc)
        test_ci(bs)
    end

    context("mean: Vector input, 1 output variables") do
        r = randn(10)
        ref = mean(r)
        bs = bootstrap(r, mean, ExactSampling())
        test_bootsample(bs, ref, r, nc)
        test_ci(bs)
    end

end

end


Random.seed!(12345)

@testset "Real Linear Solver" begin
  A = rand(3,2);
  x = rand(2);
  b = A*x;

  solvers = linearSolverListReal()

  for solver in solvers
    solverInfo = SolverInfo(Float64)
    S = createLinearSolver(solver, A, iterations=100, solverInfo=solverInfo, shape=(2,1))
    x_approx = solve(S,b)
    @info "Testing solver $solver ...: $x  == $x_approx"
    @test norm(x - x_approx) / norm(x) ≈ 0 atol=0.1
  end

  for solver in solvers
    solverInfo = SolverInfo(Float32)
    S = createLinearSolver(solver, A, iterations=100, solverInfo=solverInfo, shape=(2,1))
    x_approx = solve(S,b)
    @info "Testing solver $solver ...: $x  == $x_approx"
    @test norm(x - x_approx) / norm(x) ≈ 0 atol=0.1
  end
end

@testset "Complex Linear Solver" begin
  A = rand(3,2)+im*rand(3,2);
  x = rand(2)+im*rand(2);
  b = A*x;

  solvers = linearSolverList()

  for solver in solvers
    solverInfo = SolverInfo(ComplexF64)
    S = createLinearSolver(solver, A, iterations=100, solverInfo=solverInfo)
    x_approx = solve(S,b)
    @info "Testing solver $solver ...: $x  == $x_approx"
    @test norm(x - x_approx) / norm(x) ≈ 0 atol=0.1
  end

  for solver in solvers
    solverInfo = SolverInfo(ComplexF32)
    S = createLinearSolver(solver, A, iterations=100, solverInfo=solverInfo)
    x_approx = solve(S,b)
    @info "Testing solver $solver ...: $x  == $x_approx"
    @test norm(x - x_approx) / norm(x) ≈ 0 atol=0.1
  end
end

@testset "General Convex Solver" begin
  # fully sampled operator, image and data
  N = 256
  numPeaks = 5
  F = [1 / sqrt(N)*exp(-2im * π *j*k/N) for j=0:N-1, k=0:N-1]
  x = zeros(N)
  for i = 1:3
    x[rand(1:N)] = rand()
  end
  b = 1 / sqrt(N)*fft(x)

  # random undersampling
  idx = sort(unique(rand(1:N, div(N,2))))
  b = b[idx]
  F = F[idx,:]

  for solver in ["fista","admm"]
    reg = Regularization("L1",1e-3)
    solverInfo = SolverInfo(ComplexF64)
    S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1

    reg.λ *= length(b)/norm(b,1)
    scale_F = 1e3 # test invariance to the maximum eigenvalue
    S = createLinearSolver(solver, F .* scale_F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=true)
    x_approx = solve(S, b)
    x_approx .*= scale_F
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

  for solver in ["fista","admm"]
    reg = Regularization("L1",1e-3)
    solverInfo = SolverInfo(ComplexF32)
    S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1

    reg.λ *= length(b)/norm(b,1)
    scale_F = 1e3 # test invariance to the maximum eigenvalue
    S = createLinearSolver(solver, F .* scale_F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=true)
    x_approx = solve(S, b)
    x_approx .*= scale_F
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

  # test ADMM with option vary_ρ
  solver = "admm"
  reg = Regularization("L1",1.e-3)
  solverInfo = SolverInfo(ComplexF64)
  S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false, ρ=1e6, vary_ρ=true, verbose=false)
  x_approx = solve(S, b)
  @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
  @test x ≈ x_approx rtol=0.1

  S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false, ρ=1e-6, vary_ρ=true, verbose=false)
  x_approx = solve(S, b)
  @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
  @test x ≈ x_approx rtol=0.1

  # test ADMM with option vary_ρ
  solver = "admm"
  reg = Regularization("L1",1.e-3)
  solverInfo = SolverInfo(ComplexF32)
  S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false, ρ=1e6, vary_ρ=true, verbose=false)
  x_approx = solve(S, b)
  @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
  @test x ≈ x_approx rtol=0.1

  S = createLinearSolver(solver,F; reg=reg, iterations=200, solverInfo=solverInfo, normalizeReg=false, ρ=1e-6, vary_ρ=true, verbose=false)
  x_approx = solve(S, b)
  @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
  @test x ≈ x_approx rtol=0.1

  for solver in ["splitBregman"]
    reg = Regularization("L1",1.e-3)
    solverInfo = SolverInfo(ComplexF64)
    S = createLinearSolver(solver,F; reg=reg,iterations=5,iterationsInner=40,
                                    ρ=1.0,solverInfo=solverInfo, normalizeReg=false)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1

    reg.λ *= length(b)/norm(b,1)
    S = createLinearSolver(solver,F; reg=reg,iterations=5,iterationsInner=40,
                            ρ=1.0,solverInfo=solverInfo, normalizeReg=true)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

  for solver in ["splitBregman"]
    reg = Regularization("L1",1.e-3)
    solverInfo = SolverInfo(ComplexF32)
    S = createLinearSolver(solver,F; reg=reg,iterations=5,iterationsInner=40,
                                    ρ=1.0,solverInfo=solverInfo, normalizeReg=false)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1

    reg.λ *= length(b)/norm(b,1)
    S = createLinearSolver(solver,F; reg=reg,iterations=5,iterationsInner=40,
                            ρ=1.0,solverInfo=solverInfo, normalizeReg=true)
    x_approx = solve(S, b)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

  for solver in ["primaldualsolver"]
    reg = [Regularization("L1",1.e-4), Regularization("TV",1.e-4)]
    solverInfo = SolverInfo(Float64)
    FR = [real.(F./norm(F)); imag.(F./norm(F))]
    bR = [real.(b./norm(F)); imag.(b./norm(F))]
    S = createLinearSolver(solver,FR; reg=reg, regName=["L1","TV"], iterations=1000, solverInfo=solverInfo)
    x_approx = solve(S, bR)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

  for solver in ["primaldualsolver"]
    reg = [Regularization("L1",1.e-4), Regularization("TV",1.e-4)]
    solverInfo = SolverInfo(Float32)
    FR = [real.(F./norm(F)); imag.(F./norm(F))]
    bR = [real.(b./norm(F)); imag.(b./norm(F))]
    S = createLinearSolver(solver,FR; reg=reg, regName=["L1","TV"], iterations=1000, solverInfo=solverInfo)
    x_approx = solve(S, bR)
    @info "Testing solver $solver ...: relative error = $(norm(x - x_approx) / norm(x))"
    @test x ≈ x_approx rtol=0.1
  end

end


using PastaQ
using ITensors
using Random

Random.seed!(1234)

N = 4     # Number of qubits
depth = 4 # Depth of the quantum circuit

# Generate random quantum circuit built out of
# layers of single-qubit random rotations + `CX`
# gates, alternating between even and of odd layers.
println("Random circuit of depth $depth on $N qubits:")
circuit = randomcircuit(N, depth; twoqubitgates="CX", onequbitgates="Rn")
display(circuit)
println()

# 1. Unitary quantum circuit
# Returns the MPS at the output of the quantum circuit:
# `|ψ⟩ = Û|0,0,…,0⟩`
# where `Û` is the unitary circuit.
println("Applying random circuit to compute |ψ⟩ = U|0,0,…,0⟩...")
ψ = runcircuit(circuit)
@show maxlinkdim(ψ)
println()

# A representation of the unitary operation as a MPO
# is obtained using the flag `process=true`:
println("Approximating random circuit as an MPO U...")
U = runcircuit(circuit; process=true)
@show maxlinkdim(U)
println()

# 2. Quantum circuit with noise
# Apply a noise model `noise` (specified by appropriate
# parameters) to each quantum gate in `gates`.
# Returns a mixed density operator as MPO:
# `ρ = ε(|0,0,…⟩⟨0,0,…|)`
# where `ε` is the quantum channel.
# Here, the noise is a single-qubit amplitude damping
# channel with decay rate `γ=0.01`..
println(
  "Running the circuit with amplitude damping to compute the state ρ = ε(|0,0,…⟩⟨0,0,…|)..."
)
ρ = runcircuit(circuit; noise=("amplitude_damping", (γ=0.01,)))
@show maxlinkdim(ρ)
println()

# A representation of the quantum channel as a MPO
# is obtained using the flag `process=true`, which
# returns the Choi matrix `Λ` of the channel:`:
println(
  "Running the circuit with amplitude damping to compute the Choi matrix Λ of the quantum channel...",
)
Λ = runcircuit(circuit; process=true, noise=("amplitude_damping", (γ=0.01,)))
@show maxlinkdim(Λ)


#=
Direct estimation of the probability of for a diffusion with X(0) = x₀ satisfying
X(T) ∈ B for the Muller potential.
=#

using Statistics
using HypothesisTests
using Printf

include("muller_setup.jl");
opts = MDOptions(n_iters=nΔt)

Random.seed!(100);
n_trials = 10^4;
n_in_set = 0;
for j in 1:n_trials
    X = copy(x₀);
    sample_trajectory!(X, sampler, options=opts);
    global n_in_set+= f(X);
end
CI = confint(BinomialTest(n_in_set, n_trials));
@printf("95%% CI = (%g,%g)\n", CI[1], CI[2])


#############################################################################
#Copyright (c) 2016-2017 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# functions.jl contains several mathematic functions
#
#
#############################################################################

"""
	peakmeas(x::Array{Float64},y::Array{Float64};smoothing = "yes", filter = :SavitzkyGolay, M=5,N=2,ese_y=1.,y_smo_out=false)

The peakmeas function allows performing measurements of the position, width, intensity and centroïd of a dominant peak in a provided x-y signal.

It smooths the signal with a Savitzky-Golay filter prior to measuring the peak position, width and intensity. It is advised to check that the M and N values of the Savitzky-Golay filter are adequate for your problem before trusting the results from peakmeas. For that, just use the y_smo_out option.

half-width at half-maximum are calculated as the width of the peak at half its maximum intensity. This calculation is not affected by any asumption of peak symmetry (no fitting is done).

Centroïd is calculated as sum(y./sum(y).*x).

INPUTS:

	x: Array{Float64}, the x values;

	y: Array{Float64}, the y values.

OPTIONS:

	smoothing, String, triggers the smoothing of the spectrum if set to yes (default value);

	filter, Symbol, the filter that will be used. See the smooth function documentation;

	M=5, the M parameter for smoothing y with a Savitzky-Golay filter. See smooth function documentation;

	N=2, the M parameter for smoothing y with a Savitzky-Golay filter. See smooth function documentation;

	y_smo_out=false, the smoothed signal. Signal will be smoothed if set to true, using the SavitzkyGolayFilter function with the M and N values. y_smo output will also be provided.

OUTPUTS:

	intensity: Float64, the intensity of the peak;

	position: Float64, the position of the peak;

	hwhm: Float64, the half-width at half-maximum of the peak;

	centroïd: Float64, the centroïd of the peak;

	if y_smo_out is set to true, then another output is provided:

	y_smo: Array{Float64}, the smoothed y signal.
"""
function peakmeas(x::Array{Float64,1}, y::Array{Float64,1}; smoothing = "yes", method = "savgol", window_length=5, polyorder=2, ese_y=1., y_smo_out=false)
    ### PRELIMINARY CHECK: INCREASING SIGNAL
    if x[end,1] < x[1,1]
        x = flipdim(x,1)
        y = flipdim(y,1)
    end

	if smoothing == "yes"
    	y_smo = vec(smooth(x,y,method = method, window_length=window_length, polyorder=polyorder))
	else
		y_smo = collect(y)
	end

    x_maximum = x[y_smo .== maximum(y_smo)]

	if size(x_maximum,1) >= 2
		error("Something went wrong, is there two peaks with identical intensities in the signal? This function treats signal with one main peak...")
	end

    x_1 = x[x .<x_maximum];
	x_2 = x[x .>=x_maximum];
    y_first_portion = y_smo[x .<x_maximum];
	y_second_portion = y_smo[x .>=x_maximum];

    half_int = maximum(y_smo)/2
    idx_1 = findmin(abs.(y_first_portion.-half_int))
    idx_2 = findmin(abs.(y_second_portion.-half_int))

    hwhm = (x_2[idx_2[2]].-x_1[idx_1[2]])./2
    position = x_maximum[1]
    intensity = maximum(y_smo)
    centroid = sum(y_smo./sum(y_smo).*x)

    if y_smo_out == true
      return intensity, position, hwhm, centroid, y_smo
    elseif y_smo_out ==false
      return intensity, position, hwhm, centroid
    else
      error("Set y_smo_out to true or false.")
    end

end


using MatchCore

# example assuming * operation is always binary

# ENV["JULIA_DEBUG"] = Metatheory

struct NumberFold <: AbstractAnalysis
    egraph::EGraph
    data::Dict{Int64, Any}
end

# Mandatory for AbstractAnalysis
NumberFold(g::EGraph) = NumberFold(g, Dict{Int64, Any}())

# This should be auto-generated by a macro
function EGraphs.make(an::NumberFold, n)
    n isa Number && return n

    if n isa Expr && Meta.isexpr(n, :call)
        if n.args[1] == :*
            l = n.args[2]
            r = n.args[3]

            if get(an, l, nothing) isa Number && get(an, r, nothing) isa Number
                return an[l] * an[r]
            end
        elseif n.args[1] == :+
            l = n.args[2]
            r = n.args[3]

            if get(an, l, nothing) isa Number && get(an, r, nothing) isa Number
                return an[l] + an[r]
            end
        end
    end
    return nothing
end

function EGraphs.join(analysis::NumberFold, from, to)
    if from isa Number
        if to isa Number
            @assert from == to
        else return from
        end
    end
    return to
end

function EGraphs.modify!(an::NumberFold, id::Int64)
    g = an.egraph
    # !haskey(an, id) && return nothing
    if an[id] isa Number
        newclass = EGraphs.add!(g, an[id])
        merge!(g, newclass.id, id)
    end
end

Base.setindex!(an::NumberFold, value, id::Int64) = setindex!(an.data, value, id)
Base.getindex(an::NumberFold, id::Int64) = an.data[id]
Base.haskey(an::NumberFold, id::Int64) = haskey(an.data, id)
Base.delete!(an::NumberFold, id::Int64) = delete!(an.data, id)

comm_monoid = @theory begin
    a * b => b * a
    a * 1 => a
    a * (b * c) => (a * b) * c
end

G = EGraph(:(3 * 4))
addanalysis!(G, NumberFold)
@testset "Basic Constant Folding Example - Commutative Monoid" begin
    # addanalysis!(G, NumberFold())
    @test (true == @areequalg G comm_monoid 3 * 4 12)

    @test (true == @areequalg G comm_monoid 3 * 4 12 4*3  6*2)
end

@testset "Basic Constant Folding Example 2 - Commutative Monoid" begin
    ex = :(a * 3 * b * 4)
    G = EGraph(cleanast(ex))
    addanalysis!(G, NumberFold)
    display(G.M); println()
    display(G.analyses[1].data); println()
    @test (true == @areequalg G comm_monoid (3 * a) * (4 * b) (12*a)*b ((6*2)*b)*a)
end

@testset "Basic Constant Folding Example - Adding analysis after saturation" begin
    G = EGraph(:(3 * 4))
    # addexpr!(G, 12)
    saturate!(G, comm_monoid)
    addexpr!(G, :(a * 2))
    addanalysis!(G, NumberFold)
    saturate!(G, comm_monoid)

    display(G.M); println()
    println(G.root)
    display(G.analyses[1].data); println()

    @test (true == areequal(G, comm_monoid, :(3 * 4), 12, :(4*3), :(6*2)))

    ex = :(a * 3 * b * 4)
    G = EGraph(cleanast(ex))
    addanalysis!(G, NumberFold)
    @test (true == areequal(G, comm_monoid, :((3 * a) * (4 * b)), :((12*a)*b),
        :(((6*2)*b)*a)))
end

@testset "Infinite Loops analysis" begin
    boson = @theory begin
        1 * x => x
    end

    G = EGraph(Util.cleanast( :(1 * x) ))
    saturate!(G,boson, timeout=100)
    extractor = addanalysis!(G, ExtractionAnalysis, astsize)
    ex = extract!(G, extractor)

    display(G.M); println()
    display(G.analyses[1].data); println()

    println(ex)

    using Metatheory.EGraphs
    boson = @theory begin
        (:c * :cdag) => :cdag * :c + 1
        a * (b + c) => (a * b) + (a * c)
        (b + c) * a => (b * a) + (c * a)
        # 1 * x => x
        (a * b) * c => a * (b * c)
        a * (b * c) => (a * b) * c
    end

    G = EGraph(Util.cleanast( :(c * c * cdag * cdag) ))
    saturate!(G,boson)
    extractor = addanalysis!(G, ExtractionAnalysis, astsize_inv)
    ex = extract!(G, extractor)

    println(ex)
end


using ModelingToolkit
using Test

using ModelingToolkit: value
using SymbolicUtils: FnType

const PType = ModelingToolkit.Parameter{Real}

@parameters t
@variables x(t) y(t) # test multi-arg
@variables z(t) # test single-arg

x1 = Num(Sym{FnType{Tuple{Any}, Real}}(:x)(value(t)))
y1 = Num(Sym{FnType{Tuple{Any}, Real}}(:y)(value(t)))
z1 = Num(Sym{FnType{Tuple{Any}, Real}}(:z)(value(t)))

@test isequal(x1, x)
@test isequal(y1, y)
@test isequal(z1, z)

@parameters begin
    t
    s
end
@parameters σ(..)

t1 = Num(Sym{PType}(:t))
s1 = Num(Sym{PType}(:s))
σ1 = Num(Sym{FnType{Tuple, PType}}(:σ))
@test isequal(t1, t)
@test isequal(s1, s)
@test isequal(σ1, σ)

@derivatives D'~t
D1 = Differential(t)
@test D1 == D

@test @macroexpand(@parameters x, y, z(t)) == @macroexpand(@parameters x y z(t))
@test @macroexpand(@variables x, y, z(t)) == @macroexpand(@variables x y z(t))

# Test array expressions
@parameters begin
    t[1:2]
    s[1:2:4,1:2]
end
@parameters σ[1:2](..)

fntype(n, T) = FnType{NTuple{n, Any}, T}
t1 = Num[Variable{PType}(:t, 1), Variable{PType}(:t, 2)]
s1 = Num[Variable{PType}(:s, 1, 1) Variable{PType}(:s, 1, 2);
        Variable{PType}(:s, 3, 1) Variable{PType}(:s, 3, 2)]
σ1 = [Num(Variable{fntype(1, PType)}(:σ, 1)), Num(Variable{fntype(1, PType)}(:σ, 2))]
@test isequal(t1, t)
@test isequal(s1, s)
@test isequal(σ1, σ)

@parameters t
@variables x[1:2](t)
x1 = Num[Variable{FnType{Tuple{Any}, Real}}(:x, 1)(t.val),
      Variable{FnType{Tuple{Any}, Real}}(:x, 2)(t.val)]

@test isequal(x1, x)

@variables a[1:11,1:2]
@variables a()


function LHop!(LH,vbasis,sbasis,quad,linemap::LineMap,NED,scale,nhdofs)

    detjac = determinant_jacobian(linemap)
    for nk in NED
        for (p,w) in quad
            vals = vbasis(linemap(p...))
            svals = sbasis(p)

            M = make_row_matrix(vals,nk)
            N = interpolation_matrix(svals,nhdofs)

            LH .+= M'*N*detjac*scale*w
        end
    end
end

function LHop!(LH,vbasis,sbasis,quad,imap::InterpolatingPolynomial,normals,
    ED,scale,nhdofs)

    @assert length(scale) == length(quad)
    @assert size(normals)[2] == length(quad)

    for (idx,(p,w)) in enumerate(quad)
        n = view(normals,:,idx)
        vals = vbasis(imap(p))
        svals = sbasis(p)
        detjac = determinant_jacobian(imap,p)

        for k in 1:length(ED)
            NED = n[k]*ED[k]

            M = make_row_matrix(vals,NED)
            N = interpolation_matrix(svals,nhdofs)

            LH .+= M'*N*detjac*scale[idx]*w

        end
    end
end

function LHop(vbasis,sbasis,facequads,facemaps,normals,Dhalf,cellmap)

    nfaces = length(facemaps)
    @assert length(normals) == nfaces
    @assert length(facequads) == nfaces

    dim = dimension(vbasis)
    sdim = symmetric_tensor_dimension(dim)
    Ek = vec_to_symm_mat_converter(dim)

    NLF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    isactiveface = active_faces(facequads)

    LH = [zeros(sdim*NLF,dim*NHF) for i in 1:count(isactiveface)]
    facescale = face_determinant_jacobian(cellmap)

    counter = 1
    for faceid in 1:nfaces
        if isactiveface[faceid]
            n = normals[faceid]
            NED = [n[k]*Ek[k]'*Dhalf for k in 1:length(Ek)]
            LHop!(LH[counter],vbasis,sbasis,facequads[faceid],facemaps[faceid],
                NED,facescale[faceid],dim)
            counter += 1
        end
    end
    return LH
end

function LHop_on_interface(vbasis,sbasis,iquad,imap,normals,Dhalf,cellmap)

    dim = dimension(vbasis)
    sdim = symmetric_tensor_dimension(dim)
    nq = length(iquad)

    @assert size(Dhalf) == (sdim,sdim)
    @assert size(normals) == (dim,nq)

    Ek = vec_to_symm_mat_converter(dim)
    NLF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    LH = zeros(sdim*NLF,dim*NHF)
    scale = scale_area(cellmap,normals)

    ED = [E'*Dhalf for E in Ek]
    LHop!(LH,vbasis,sbasis,iquad,imap,normals,ED,scale,dim)

    return LH
end

function UHop!(UH,vbasis,sbasis,quad,linemap::LineMap,
    scale,nudofs)

    detjac = determinant_jacobian(linemap)
    for (p,w) in quad
        vals = vbasis(linemap(p...))
        svals = sbasis(p)

        Nv = interpolation_matrix(vals,nudofs)
        Ns = interpolation_matrix(svals,nudofs)

        UH .+= Nv'*Ns*detjac*scale*w
    end
end

function UHop!(UH,vbasis,sbasis,quad,imap::InterpolatingPolynomial,
    scale,nudofs)

    @assert length(scale) == length(quad)

    for (idx,(p,w)) in enumerate(quad)
        vals = vbasis(imap(p))
        svals = sbasis(p)
        detjac = determinant_jacobian(imap,p)

        Nv = interpolation_matrix(vals,nudofs)
        Ns = interpolation_matrix(svals,nudofs)

        UH .+= Nv'*Ns*detjac*scale[idx]*w
    end
end

function UHop(vbasis,sbasis,facequads,facemaps,stabilization,cellmap)

    nfaces = length(facemaps)
    @assert length(facequads) == nfaces

    dim = dimension(vbasis)
    sdim = dimension(sbasis)
    @assert sdim == dim-1

    NUF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    isactiveface = [length(fq) > 0 ? true : false for fq in facequads]

    facejac = face_determinant_jacobian(cellmap)
    UH = [zeros(dim*NUF,dim*NHF) for i = 1:count(isactiveface)]

    counter = 1
    for faceid in 1:nfaces
        if isactiveface[faceid]
            UHop!(UH[counter],vbasis,sbasis,facequads[faceid],facemaps[faceid],
                stabilization*facejac[faceid],dim)
            counter += 1
        end
    end
    return UH
end

function UHop_on_interface(vbasis,sbasis,squad,imap,normals,
    stabilization,cellmap)

    dim = dimension(vbasis)
    sdim = dimension(sbasis)
    @assert sdim == dim-1

    NUF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    UH = zeros(dim*NUF,dim*NHF)
    scale = stabilization*scale_area(cellmap,normals)

    UHop!(UH,vbasis,sbasis,squad,imap,scale,dim)

    return UH
end

function local_hybrid_operator(vbasis,sbasis,facequads,facemaps,normals,
    Dhalf,stabilization,cellmap)

    LH = LHop(vbasis,sbasis,facequads,facemaps,normals,Dhalf,cellmap)
    UH = UHop(vbasis,sbasis,facequads,facemaps,stabilization,cellmap)

    @assert length(LH) == length(UH)

    return [[LH[i];UH[i]] for i in 1:length(LH)]

end

function local_hybrid_operator_on_interface(vbasis,sbasis,iquad,imap,inormals,
    Dhalf,stabilization,cellmap)

    LH = LHop_on_interface(vbasis,sbasis,iquad,imap,inormals,Dhalf,cellmap)
    UH = UHop_on_interface(vbasis,sbasis,iquad,imap,inormals,
        stabilization,cellmap)

    return [LH;UH]
end


@testset "match wave low frequency" begin

## Low frequency
    medium = Medium(1.0,1.0+0.0im)
    volfracs = [0.1,0.5]
    species = [
        Specie(ρ=5.,r=0.002, c=v, volfrac=v)
    for v in volfracs]

    ω = 0.002
    θin = -0.2
    tol = 1e-8
    hankel_order = 1

    k_eff_lows = map(species) do s
        med_eff = effective_medium(medium, [s])
        ω/med_eff.c
    end
    wave_eff_lows = [
        EffectiveWave(ω, k_eff_lows[i], medium, [species[i]]; tol = tol,
            hankel_order=hankel_order, θin = θin)
    for i in eachindex(species)]

    wave_effs_arr = [
        effective_waves(ω, medium, [s];
            hankel_order=hankel_order,
            num_wavenumbers = 1,
            tol = tol,  θin = θin,
            extinction_rescale = false)
    for s in species]

    @test norm([ws[1].k_eff for ws in wave_effs_arr] - k_eff_lows) < tol

    match_ws = [
        MatchWave(ω, medium, species[i]; hankel_order=hankel_order,
        max_size = 60,
        θin = θin, tol = tol,
        wave_effs = wave_effs_arr[i])
    for i in eachindex(species)]

   @test maximum(match_error.(match_ws)) < tol

    map(eachindex(species)) do i
        x = LinRange(match_ws[i].x_match[end],2pi/real(match_ws[i].effective_waves[1].k_eff),200)
        avg_low = AverageWave(x, wave_eff_lows[i])
        avg = AverageWave(x, match_ws[i].effective_waves)
        @test norm(avg.amplitudes[:] - avg_low.amplitudes[:]) < tol
        @test maximum(abs.(avg.amplitudes[:] - avg_low.amplitudes[:])) < tol
    end
end


using Base: AbstractFloat
########################################################################################################
## Details
# This file will create the type structures for the UQ inference problems.
#
#########################################################################################################
abstract type Potential{T} end

mutable struct SNAP{T<:Base.AbstractFloat} <: Potential{T}
    A :: Matrix{T}
    b :: Vector{T}
    β :: Vector{T}
end

abstract type ArbitraryDistribution{T} end

mutable struct PotentialDistribution{T<:AbstractFloat} <: ArbitraryDistribution{T}
    potential::Potential{T}       # potential defined in Potential.jl
    x                             # Container of learnable parameters
    p                             # Container of nonlearnable parameters
    distribution                  # -log likelihood
end

function PotentialDistribution(potential :: Potential{T}, distribution) where T<: AbstractFloat
    pd = PotentialDistribution(potential, p.x, p.q, distribution)
    return pd
end

#########################################################################################################
## SNAP

# Nonlearnable Q
mutable struct SNAPDistribution{T<:AbstractFloat} <: ArbitraryDistribution{T}
    snap::SNAP{T}               # SNAP potential, assume has A, b, β
    x                        # Container of learnable parameters
    p                        # Container of nonlearnable parameters
    distribution             # -log likelihood
end

function SNAPDistribution(snap :: SNAP{T}, Q) where T<: AbstractFloat
    x = (β = snap.β, )
    p = (A = snap.A, b = snap.b, Q = Q)

    distribution(x, p) = 0.5*tr( transpose(p.A*x.β-p.b)*(p.Q\(p.A*x.β-p.b)) + logdet(p.Q))
    pd = SNAPDistribution{T}(snap, x, p, distribution)
    return pd
end

# Learnable Q
mutable struct SNAPQDistribution{T<:AbstractFloat} <: ArbitraryDistribution{T}
    snap::SNAP{T}               # SNAP potential, assume has A, b, β
    x                           # Container of learnable parameters
    p                           # Container of nonlearnable parameters
    distribution                # SNAP likelihood(p)
end

function SNAPQDistribution(snap :: SNAP{T}, Q) where T<: AbstractFloat
    x = (β = snap.β, Q = Q)
    p = (A = snap.A, b = snap.b)
    distribution(x, p) = 0.5*tr( transpose(p.A*x.β-p.b)*(x.Q\(p.A*x.β-p.b)) + logdet(x.Q))
    pd = SNAPQDistribution{T}(snap, x, p, distribution)
    return pd
end


using LinQu, Test, Random
import LinQu.data # data() is extended

# define customized QGate
struct RandomSingleGate <: QGate
	rng::MersenneTwister
	seed::Int
	maxangle::Number

	RandomSingleGate(rng::MersenneTwister, seed::Int, maxangle::Number) = new(rng, seed, maxangle)
end
RandomSingleGate(seed::Int, maxangle::Number) = RandomSingleGate(MersenneTwister(seed), seed, maxangle)

reset!(gate::RandomSingleGate) = Random.seed!(gate.rng, gate.seed)

function data(gate::RandomSingleGate)
	params = rand(gate.rng, 4)
	params[2:end] ./= sum(params[2:end].^2)
	theta = params[1]* gate.maxangle
	return cos(theta/2)*I_DATA-sin(theta/2)im* (params[2]*X_DATA+params[3]*Y_DATA+params[4]*Z_DATA)
end

# Built-in interface still work for customized gate
state = MPSState(2)
state_copy = copy(state)
customizedgate = RandomSingleGate(5, π)
apply!(state, customizedgate, 1)
apply!(state, customizedgate, 2)
reset!(customizedgate)
apply!(state_copy, customizedgate, 1)
apply!(state_copy, customizedgate, 2)
@test toExactState(state) ≈ toExactState(state_copy)


using DataFrames
using NamedArrays
using SchumakerSpline

bounds = (0.0, 1.0)
x_name = :alpha
evaluations = DataFrame()

function func_inc1(x)
    feval = NamedArray([x, -sqrt(x)])
    setnames!(feval, ["utility", "x_E"], 1)
    return feval
end
const_Array = NamedArray([0.1, sqrt(0.2)])
setnames!(const_Array, ["utility", "x_E"], 1)

# Testing the construction of a dataframe of functions
dd_function = make_function_dataframe([:func_inc1], Array{Any,1}([func_inc1]), [:Increasing])
dd_function2 = make_function_dataframe([:const_Array], Array{Any,1}([const_Array]), [:Constant])

#
grid                = reverse(0.0001:0.01:0.9999)
output_for_envelope = :utility
x_name              = :alpha
env                 = Envelope(dd_function, grid, output_for_envelope,x_name)
isa(env, Envelope)
env2                = Envelope(dd_function2, grid, output_for_envelope,x_name)
isa(env2, Envelope)


#############################################################################
# card.jl
# Handles functions of the form f(A) = g(length(A)), where g is a concave function
#############################################################################

export card, compose, CardBasedAtom
export in
export sign, monotonicity, modularity, evaluate

mutable struct CardBasedAtom{T} <: SubmodFunc
  head::Symbol
  id_hash::UInt64
  children::Tuple{AbstractExpr, SetVariable}
  size::Tuple{Int, Int}
  setvariables::Array{SetVariable}
  func::T
end

function CardBasedAtom(F::Function, S::CombiSet)
  z = Variable(1)
  children = (F(z), S)                           # create an expression
  if F(0) != 0
    error("A function has to return 0 at 0 in order to derive a cardinality-based function.")
  end
  setvariables = get_sv(S)
  return(CardBasedAtom(:card, hash(children), children, (1, 1), setvariables, F))
end

function CardBasedAtom(w::Array{Float64}, S::CombiSet)
  children = (Constant(w), S)                           # create an expression
  setvariables = get_sv(S)
  if length(setvariables) != 1
    error("The expression should have exactly one set variable.")
  elseif size(w) != (setvariables[1].cardinality, )
    error("The sizes of the set variable and the array mismatch.")
  end
  w₁ = copy(w)
  for i = 2:length(w)
    w₁[i] = w[i] - w[i-1]
  end
  return(CardBasedAtom(:card, hash(children), children, (1, 1), setvariables, w₁))
end

card(f::Function, S::CombiSet) = CardBasedAtom(f, S)

card(w::Array{Float64}, S::CombiSet) = CardBasedAtom(w, S)

card(w::Array{Int}, S::CombiSet) = CardBasedAtom(Float64.(w), S)

card(S::CombiSet) = card(x -> x, S)

function sign(F::CardBasedAtom)
  return sign(F.children[1])
end

function sign(F::CardBasedAtom{Array{Float64, 1}})
  if all(x -> x>=0, F.func)
    return Positive()
  elseif all(x -> x<=0, F.func)
    return Negative()
  else
    return NoSign()
  end
end

function monotonicity(F::CardBasedAtom)
  return (monotonicity(F.children[1]), )
end

function monotonicity(F::CardBasedAtom{Array{Float64, 1}})
  if all(x -> x>=0, F.func)
    return Nondecreasing()
  elseif all(x -> x<=0, F.func)
    return Nonincreasing()
  elseif all(x -> x==0, F.func)
    return ConstMonotonicity()
  else
    return NoMonotonicity()
  end
end

function modularity(F::CardBasedAtom)
  if vexity(F.children[1]) == ConcaveVexity()
    return SubModularity()
  elseif vexity(F.children[1]) == ConvexVexity()
    return SuperModularity()
  elseif vexity(F.children[1]) == AffineVexity()
    return Modularity()
  elseif vexity(F.children[1]) == ConstVexity()
    return ConstModularity()
  else
    return NotDetermined()
  end
end

function modularity(F::CardBasedAtom{Array{Float64, 1}})
  w₁ = sort(F.func, rev=true)
  if w₁ == F.func
    return SubModularity()
  else
    w₂ = sort(F.func)
    if w₂ == F.func
      return SuperModularity()
    elseif all(x -> x==0, F.func)
      return ConstModularity()
    else
      return NotDetermined()
    end
  end
end

function evaluate(F::CardBasedAtom)
  var = F.setvariables[1]
  var1 = get_v(F.children[1])
  var1[1].value = length(Set(get_elements(var)))
  evaluate(F.children[1])
end

function evaluate(F::CardBasedAtom{Array{Float64, 1}})
  return sum(F.func[1:length(get_elements(F.children[2]))])
end

function in(w::AbstractArray, p::SubmodPoly{CardBasedAtom})
  n = length(w)
  @assert length(p.V) == n
  checker = true
  S = sort(p.V)
  ordering = sortperm(w, rev = true)
  h = 0
  ind = []
  y = zeros(n + 1)
  for i = 1:length(w)
    y[i + 1] = evaluate(p.F, S[push!(ind, ordering[i])])[1]
    h += - y[i] - w[i] + y[i + 1]
    if h < 0
      checker = false
      break
    end
  end
  return checker
end

function compose(f::Function, g::CardBasedAtom)
  return card(f, g.setvariables[1])
end


"""動力学関連
"""
module Dynamics

using LinearAlgebra

"""慣性行列"""
function M!(q::Vector{Float64}, out::Matrix{Float64})
    ccall(
        (:M, "o/soft_robot/derived/ikko_dake/eqs/c_so/M.so"),
        Cvoid,
        (Cdouble, Cdouble, Cdouble, Ptr{Cdouble}),
        q[1], q[2], q[3], out
    )
end

"""慣性行列"""
function M(q::Vector{Float64})
    Z = Matrix{Float64}(undef, 3, 3)
    M!(q, Z)
    return Z
end

"""慣性行列の逆行列"""
function invM!(q::Vector{Float64}, out::Matrix{Float64})
    ccall(
        (:invM, "o/soft_robot/derived/ikko_dake/eqs/c_so/invM.so"),
        Cvoid,
        (Cdouble, Cdouble, Cdouble, Ptr{Cdouble}),
        q[1], q[2], q[3], out
    )
end

"""慣性行列の逆行列"""
function invM(q::Vector{Float64})
    Z = Matrix{Float64}(undef, 3, 3)
    invM!(q, Z)
    return Z
end

function C!(q::Vector{Float64}, q_dot::Vector{Float64}, out::Matrix{Float64})
    ccall(
        (:C, "o/soft_robot/derived/ikko_dake/eqs/c_so/C.so"),
        Cvoid,
        (Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Ptr{Cdouble}),
        q[1], q[2], q[3], q_dot[1], q_dot[2], q_dot[3], out
    )
end

"""コリオリ＋遠心力行列"""
function C(q::Vector{Float64}, q_dot::Vector{Float64})
    Z = Matrix{Float64}(undef, 3, 3)

    C!(q, q_dot, Z)

    return Z
end

"""Gのラッパー"""
function G!(q::Vector{Float64}, out::Vector{Float64})
    ccall(
        (:G, "o/soft_robot/derived/ikko_dake/eqs/c_so/G.so"),
        Cvoid,
        (Cdouble, Cdouble, Cdouble, Ptr{Cdouble}),
        q[1], q[2], q[3], out
    )
end

"""重力行列"""
function G(q::Vector{Float64})
    Z = Vector{Float64}(undef, 3)

    G!(q, Z)

    return Z
end

const K = diagm([1700.0, 1700.0, 1700.0])  # 剛性行列の真値
const D = diagm([110.0, 110.0, 110.0])  # 減衰行列の真値

const uncertain_K = diagm([1020.0, 1020.0, 1020.0])  # 不確かな剛性行列
const uncertain_D = diagm([77.0, 77.0, 77.0])  # 不確かな減衰行列

const αh = 23.705
const βh = 1.7267
const γh = -42.593

function h_dot(h::T, q::T, q_dot::T) where T
    q * (αh - (βh * sign(q_dot*h) + γh)*abs(h))
end

"""ヒステリシスベクトルの更新式"""
function H_dot(H::Vector{T}, q::Vector{T}, q_dot::Vector{T}) where T
    [
        h_dot(H[1], q[1], q_dot[1]),
        h_dot(H[2], q[2], q_dot[2]),
        h_dot(H[3], q[3], q_dot[3]),
    ]
end

"""加速度"""
function calc_q_dot_dot(
    τ::Vector{T}, q::Vector{T}, q_dot::Vector{T},
    H::Vector{T},
    ) where T
    inv(M(q)) * (τ .- (C(q, q_dot) .+ D)*q_dot .- K*q .- G(q) .- H) |> vec
end

"""加速度

outに加速度ベクトルを書き出す
"""
function q_dot_dot!(
    τ::Vector{Float64}, q::Vector{Float64}, q_dot::Vector{Float64},
    H::Vector{Float64}, out::Vector{Float64}
    )
    ccall(
        (:q_dot_dot, "o/soft_robot/derived/ikko_dake/eqs/c_so/q_dot_dot.so"),
        Cvoid,
        (Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Cdouble, Ptr{Cdouble}),
        H[1], H[2], H[3], q[1], q_dot[1], q[2], q_dot[2], q[3], q_dot[3], τ[1], τ[2], τ[3], out
    )
end

end

# using LinearAlgebra
# using .Dynamics

# m = Dynamics.invM([0.0, 0.01, 0.0])
# m2 = Dynamics.invM([0.0, 0.02, 0.0])
# println((m .- m2) ./ m)
# # e,u=eigen(m)
# # e


using Oceananigans.Advection:
    _advective_momentum_flux_Uu,
    _advective_momentum_flux_Uv,
    _advective_momentum_flux_Vu,
    _advective_momentum_flux_Vv,
    advective_tracer_flux_x,
    advective_tracer_flux_y,
    vertical_vorticity_U,
    vertical_vorticity_V,
    bernoulli_head_U,
    bernoulli_head_V

using Oceananigans.Grids: AbstractGrid
using Oceananigans.Operators: Ax_qᶠᶜᶜ, Ay_qᶜᶠᶜ

#####
##### Momentum flux operators
#####

@inline momentum_flux_huu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uu(i, j, k, grid, advection, solution[1], solution[1]) / solution.h[i, j, k]

@inline momentum_flux_hvu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vu(i, j, k, grid, advection, solution[2], solution[1]) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_huv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uv(i, j, k, grid, advection, solution[1], solution[2]) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_hvv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vv(i, j, k, grid, advection, solution[2], solution[2]) / solution.h[i, j, k]

#####
##### Momentum flux divergence operators
#####

@inline div_mom_u(i, j, k, grid, advection, solution, formulation) =
    1 / Azᶠᶜᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, momentum_flux_huu, advection, solution) +
                                δyᵃᶜᵃ(i, j, k, grid, momentum_flux_hvu, advection, solution))

@inline div_mom_v(i, j, k, grid, advection, solution, formulation) =
    1 / Azᶜᶠᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_huv, advection, solution) +
                                δyᵃᶠᵃ(i, j, k, grid, momentum_flux_hvv, advection, solution))

@inline div_mom_u(i, j, k, grid, advection, solution, ::VectorInvariantFormulation) = (
    + vertical_vorticity_U(i, j, k, grid, advection, solution[1], solution[2])  # Vertical relative vorticity term
    + bernoulli_head_U(i, j, k, grid, advection, solution[1], solution[2]))     # Bernoulli head term

@inline div_mom_v(i, j, k, grid, advection, solution, ::VectorInvariantFormulation) = (
    + vertical_vorticity_V(i, j, k, grid, advection, solution[1], solution[2])  # Vertical relative vorticity term
    + bernoulli_head_V(i, j, k, grid, advection, solution[1], solution[2]))     # Bernoulli head term

# Support for no advection
@inline div_mom_u(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, formulation) where FT = zero(FT)
@inline div_mom_v(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, formulation) where FT = zero(FT)
@inline div_mom_u(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, ::VectorInvariantFormulation) where FT = zero(FT)
@inline div_mom_v(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, ::VectorInvariantFormulation) where FT = zero(FT)

#####
##### Mass transport divergence operator
#####

"""
    div_Uh(i, j, k, grid, advection, solution, formulation)

Calculates the divergence of the mass flux into a cell,

    1/Az * [δxᶜᵃᵃ(Δy * uh) + δyᵃᶜᵃ(Δx * vh)]

which will end up at the location `ccc`.
"""
@inline function div_Uh(i, j, k, grid, advection, solution, formulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, solution[1]) +
                                     δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, solution[2]))
end

@inline div_Uh(i, j, k, grid, advection, solution, formulation::VectorInvariantFormulation) =
        div_Uc(i, j, k, grid, advection, solution, solution.h, formulation)

#####
##### Tracer advection operator
#####

@inline transport_tracer_flux_x(i, j, k, grid, advection, uh, h, c) =
    @inbounds advective_tracer_flux_x(i, j, k, grid, advection, uh, c) / ℑxᶠᵃᵃ(i, j, k, grid, h)

@inline transport_tracer_flux_y(i, j, k, grid, advection, vh, h, c) =
    @inbounds advective_tracer_flux_y(i, j, k, grid, advection, vh, c) / ℑyᵃᶠᵃ(i, j, k, grid, h)

"""
    div_Uc(i, j, k, grid, advection, U, c)

Calculates the divergence of the flux of a tracer quantity c being advected by
a velocity field U = (u, v), ∇·(Uc),

    1/Az * [δxᶜᵃᵃ(Δy * uh * ℑxᶠᵃᵃ(c) / h) + δyᵃᶜᵃ(Δx * vh * ℑyᵃᶠᵃ(c) / h)]

which will end up at the location `ccc`.
"""

@inline function div_Uc(i, j, k, grid, advection, solution, c, formulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, transport_tracer_flux_x, advection, solution[1], solution.h, c) +
                                     δyᵃᶜᵃ(i, j, k, grid, transport_tracer_flux_y, advection, solution[2], solution.h, c))
end

@inline function div_Uc(i, j, k, grid, advection, solution, c, ::VectorInvariantFormulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, advective_tracer_flux_x, advection, solution[1], c) +
                                     δyᵃᶜᵃ(i, j, k, grid, advective_tracer_flux_y, advection, solution[2], c))
end

# Support for no advection
@inline div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c, formulation) where FT = zero(FT)

@inline u(i, j, k, grid, solution) = @inbounds solution.uh[i, j, k] / ℑxᶠᵃᵃ(i, j, k, grid, solution.h)
@inline v(i, j, k, grid, solution) = @inbounds solution.vh[i, j, k] / ℑyᵃᶠᵃ(i, j, k, grid, solution.h)

"""
    c_div_U(i, j, k, grid, advection, U)

Calculates the product of the tracer concentration c with
the horizontal divergence of the velocity field U = (u, v), c ∇·(U),

    1/Az * [δxᶜᵃᵃ(Δy * uh / h) + δyᵃᶜᵃ(Δx * vh / h]

which will end up at the location `ccc`.
"""
@inline c_div_U(i, j, k, grid, solution, c, formulation) =
    @inbounds c[i, j, k] * 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, u, solution) + δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, v, solution))

@inline c_div_U(i, j, k, grid, solution, c, ::VectorInvariantFormulation) =
    @inbounds c[i, j, k] * 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, solution[1]) + δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, solution[2]))

# Support for no advection
@inline c_div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c, formulation) where FT = zero(FT)


"""
The arguments are:

# !    r₁      - Primary object's position vector in ECI coordinates
# !              (1x3 row vector)
# !    v₁      - Primary object's velocity vector in ECI coordinates
# !              (1x3 row vector)
# !    Σ₁      - Primary object's covariance matrix in ECI coordinate frame
# !              (3x3 or 6x6)
# !    r₂      - Secondary object's position vector in ECI coordinates
# !              (1x3 row vector)
# !    v₂      - Secondary object's velocity vector in ECI coordinates
# !              (1x3 row vector)
# !    Σ₂      - Secondary object's covariance matrix in ECI coordinate frame
# !              (3x3 or 6x6)
# !    HBR     - Hard body region

"""
function pc2dfoster_circle(r₁::SVector{3,T}, v₁::SVector{3,T}, Σ₁::SymmetricM3{T},
                            r₂::SVector{3,T}, v₂::SVector{3,T}, Σ₂::SymmetricM3{T}, HBR = T(0.02)) where T
    #
    # ! Construct relative encounter frame
    δr = r₁ - r₂
    δr₀ = norm(δr)
    δv = v₁ - v₂
    z = LinearAlgebra.cross(δr, δv)

    # ! Relative encounter frame
    nδv = normalize(δv)
    nz  = normalize(z)

    Cp = combinecov(Σ₁, Σ₂, nδv, nz)
    Up = revchol(Cp)
    Uv = SVector{3,T}(
        7.071067811865475244008443621048490392848359376884740365883398689953662392310596e-01/Up[1],
        Up[2], Up[3]
    )

    gaussianarea_old(Uv, δr₀, HBR, HBR*HBR)

end

"""
Lv is the Cholesky factor of the 2x2 projection at TCA.
δr₀ is the projected miss distance.

The function will tabulate a series of N * 2 * W misses, where W is the CPU vector width.
W should be a function of both the data type used (eg, it will be twice as large for single
precision as it will for double), and CPU architecture.
For example, with double precision and the sse instruction set (128 bit vectors), W will be 2. With single precision
and the avx512f instruction set (512 bit), W will be 16.

"""
@generated function sample_misses!(table::Array{SVec{W,T},3}, r1, v1, r2, v2, c1, c2, N) where {T,W}
    # W = VectorizationBase.pick_vector_width(T)
    quote
        rng = VectorizedRNG.GLOBAL_vPCG
        # Transpose U; broadcast the individual elements.
        Σ, δr₀ = CovarianceRealism.RIC_to_2D(r1, v1, c1, r2, v2, c2)
        U = chol(Σ)

        L11 = vbroadcast(SVec{$W,$T}, U[1,1])
        L21 = vbroadcast(SVec{$W,$T}, U[1,2])
        L22 = vbroadcast(SVec{$W,$T}, U[2,2])
        vδr₀ = vbroadcast(SVec{$W,$T}, δr₀)
        @inbounds for n ∈ 1:N
            # Sampling random normals in batches of 4W is fastest on the tested architecture.
            z1_and_2 = randn(rng, Vec{$(4W),$T})

            # We want coordinates z1 and z2. Because we sampled 4W, we have 2W of each z1 and z2.
            # In some versions of Julia there is a bug that causes a more typical version of constructing these tuples:
            # z1_1 = ntuple(Val($W)) do i x1_and_2[i] end
            # or equivalently
            # z1_1 = ntuple(i -> x1_and_2[i], Val($W))
            # So we construct the tuple explicitly here. This is fast in all versions of Julia > 1.
            z1_1 = SVec($(Expr(:tuple, [:(z1_and_2[$w]) for w ∈    1: W]...)))
            z1_2 = SVec($(Expr(:tuple, [:(z1_and_2[$w]) for w ∈  W+1:2W]...)))
            z2_1 = SVec($(Expr(:tuple, [:(z1_and_2[$w]) for w ∈ 2W+1:3W]...)))
            z2_2 = SVec($(Expr(:tuple, [:(z1_and_2[$w]) for w ∈ 3W+1:4W]...)))

            # Multiply Calculate L * z = x.
            # Additionally, subtract the distance δr₀, which we treat as [δr₀, 0]
            x1_1 = L11 * z1_1 - vδr₀
            x1_2 = L11 * z1_2 - vδr₀

            x2_1 = L21 * z1_1 + L22 * z2_1
            x2_2 = L21 * z1_2 + L22 * z2_2

            # Calculate distance, multiply 1000 to translate from kilometers to meters.
            δ_1 = sqrt( x1_1*x1_1 + x2_1*x2_1 ) * T(10^3)
            δ_2 = sqrt( x1_2*x1_2 + x2_2*x2_2 ) * T(10^3)
            vi = vbroadcast(SVec{$W,$T}, $T(100))
            l_1_5 = δ_1 < vi
            l_2_5 = δ_2 < vi
            # Every 5 times, we check if we can break early.
            for i ∈ 50:-5:1
                ((l_1_5 === SVec($(Expr(:tuple, [:(Core.VecElement(false)) for w ∈ 1:W]...)))) && (l_2_5 === SVec($(Expr(:tuple, [:(Core.VecElement(false)) for w ∈ 1:W]...))))) && break
                # (any(l_1_5) || any(l_2_5)) || break
                Base.Cartesian.@nexprs 5 j -> begin
                    vi_j = vbroadcast(SVec{$W,$T}, $T(i+1-j))
                    ti_1_j = table[1,j,i+1-j]
                    l_1_j = δ_1 < vi_j
                    ti_1_j = vifelse(l_1_j, ti_1_j + vbroadcast(SVec{W,T},one(T)), ti_1_j)
                    table[1,j,i+1-j] = ti_1_j
                    ti_2_j = table[2,j,i+1-j]
                    l_2_j = δ_2 < vi_j
                    ti_2_j = vifelse(l_2_j, ti_2_j + vbroadcast(SVec{W,T},one(T)), ti_2_j)
                    table[2,j,i+1-j] = ti_2_j
                end
            end
        end

    end
end

function pc2dfoster_circle_old(r₁::SVector{3,T}, v₁::SVector{3,T}, Σ₁::SymmetricM3{T},
                            r₂::SVector{3,T}, v₂::SVector{3,T}, Σ₂::SymmetricM3{T}, HBR = T(0.02)) where T
    #
    # ! Construct relative encounter frame
    δr = r₁ - r₂
    δr₀ = norm(δr)
    δv = v₁ - v₂
    z = LinearAlgebra.cross(δr, δv)

    # ! Relative encounter frame
    nδv = normalize(δv)
    nz  = normalize(z)

    Cp = combinecov(Σ₁, Σ₂, nδv, nz)
    Up = revchol(Cp)
    Uv = SVector{3,T}(
        7.071067811865475244008443621048490392848359376884740365883398689953662392310596e-01/Up[1],
        Up[2], Up[3]
    )

    gaussianarea_old(Uv, δr₀, HBR, HBR*HBR)

end
@inline function gaussianarea_old(U::SVector{3,T}, x₀::T, HBR = T(0.2), HBR²::T = HBR^2) where T

    n = SVector{8,T}(
        0.09226835946330202,0.273662990072083,0.4457383557765383,
        0.6026346363792564,0.7390089172206591,0.8502171357296142,
        0.9324722294043558,0.9829730996839018
    )
    w = SVector{8,T}(
        0.036704932945846216,0.035454990477090234,0.032997670831211,
        0.029416655081518854,0.02483389042441457,0.01940543741387547,
        0.01331615550646369,0.006773407894247478
    )

    s = HBR / U[3]

    @inbounds Pc = w[1] * gdensity(s*n[1], HBR², x₀, U)
    @inbounds for i ∈ 2:8
        pct = w[i] * gdensity(s*n[i], HBR², x₀, U)
        Pc += pct
    end

    Pc * s

end

function pc2dfoster_RIC(r₁::SVector{3,T}, v₁::SVector{3,T}, Σ₁::SymmetricM3{T},
                            r₂::SVector{3,T}, v₂::SVector{3,T}, Σ₂::SymmetricM3{T}, HBR = T(0.02)) where T

    Cp, δr₀ = RIC_to_2D(r₁, v₁, Σ₁, r₂, v₂, Σ₂)
    Up = revchol(Cp)
    Uv = SVector{3,T}(
        7.071067811865475244008443621048490392848359376884740365883398689953662392310596e-01/Up[1],
        Up[2], Up[3]
    )

    gaussianarea_old(Uv, δr₀, HBR, HBR*HBR)

end

#
# @inline function gaussianarea(U::SVector{3,T}, x₀::T, HBR = T(0.2), HBR²::T = HBR^2) where T
#     # ! Relative error within machine error in a test case.
#     # ! This is better than 1e-8 ComputePcUncertainty passed to Pc2dFoster.
#     # ! This algorithm is not adaptive, unlike quadgk.
#     # ! Therefore, it may need validation to ensure that error holds up over a range.
#     # ! Also, currently quadgk is passed AbsTol = 1e-13.
#     # ! With that AbsTol, low probability events guaranteed too much accuracy.
#     # ! quadgk will terminate after its first error estimate
#     # !   - This first error estimate is probably enough for decent relative error
#     # !
#     # !
#     # ! These are Chebyschev nodes and weights of the second kind,
#     # ! with weights adjusted for the implicit weight function.
#     # ! They were generated via the following Julia code:
#     # !
#     # ! #Pkg.add("FastGaussQuadrature") #if it isn't already installed.
#     # ! using FastGaussQuadrature
#     # ! nodes = 32
#     # ! Nh = div(nodes,2)
#     # ! n, w = gausschebyshev(nodes, 2) # Arguments: Number of nodes, Chebyschev-kind (1,2,3,or 4)
#     # ! uw = @. w / sqrt(8pi*(1-abs2(n))) # Counter weight function of Chebyschev's 2nd kind : 1/sqrt(1-x^2)
#     # !
#     # !
#     # ! # Fortran's preprocessor could truncate numbers to single precision if they're not given as double precision.
#     # ! # Fortran also has limits on number of characters per line.
#     # ! # We don't want to have to make these adjustments manually.
#     # ! # The following function converts the vectors into a string that we can then print, copy and paste.
#     # !
#     # ! pa_dp(n[Nh+1:end]) |> print
#     # ! pa_dp(uw[Nh+1:end]) |> print
#     #
#     # ! Note that the rules are symmetric, so I am only copying the positive half.
#     # ! In integrating, the code calculates density for +/- n.
#
#
#     n = SVec{8,T}(
#         0.09226835946330202,0.273662990072083,0.4457383557765383,
#         0.6026346363792564,0.7390089172206591,0.8502171357296142,
#         0.9324722294043558,0.9829730996839018
#     )
#     w = SVec{8,T}(
#         0.036704932945846216,0.035454990477090234,0.032997670831211,
#         0.029416655081518854,0.02483389042441457,0.01940543741387547,
#         0.01331615550646369,0.006773407894247478
#     )
#
#     U₁ = U[1]
#     s = HBR / U[3]
#
#     @fastmath begin
#         z = s * n
#         U₂z = VectorizationBase.extract_data(U[2] * z)
#         U₃z = U[3] * z
#         radical = sqrt( HBR² - U₃z * U₃z )
#         exponentials = exp(-0.5 * z * z)
#         erfs_1 = SLEEFwrap.erfc(erf_args(x₀, U₂z[1].value, U₂z[2].value, radical[1], radical[2], U₁))
#         erfs_2 = SLEEFwrap.erfc(erf_args(x₀, U₂z[3].value, U₂z[4].value, radical[3], radical[4], U₁))
#         erfs_3 = SLEEFwrap.erfc(erf_args(x₀, U₂z[5].value, U₂z[6].value, radical[5], radical[6], U₁))
#         erfs_4 = SLEEFwrap.erfc(erf_args(x₀, U₂z[7].value, U₂z[8].value, radical[7], radical[8], U₁))
#         # erfreductions = SVec{8,Float64}(ntuple(Val(8)) do i
#         #     erfset = erfs[cld(i,2)]
#         #     k = 4*((i-1) % 2)
#         #     VE(erfset[k+1].value + erfset[k+2].value - erfset[k+3].value - erfset[k+4].value)
#         # end)
#         erfreductions = SVec((
#             erfset_reduce(erfs_1, 0), erfset_reduce(erfs_1, 4),
#             erfset_reduce(erfs_2, 0), erfset_reduce(erfs_2, 4),
#             erfset_reduce(erfs_3, 0), erfset_reduce(erfs_3, 4),
#             erfset_reduce(erfs_4, 0), erfset_reduce(erfs_4, 4)
#         ))
#         pc = SVec(VectorizationBase.extract_data(w * exponentials * erfreductions))
#         SIMDPirates.vsum(w * exponentials * erfreductions) * s
#     end
# end
# @inline function erf_args(x₀::T, U₂z₁, U₂z₂, radical₁, radical₂, U₁) where T
#     m₁ = SVec{8,T}(-1, 1,-1, 1,-1, 1,-1, 1)
#     m₂ = SVec{8,T}(-1,-1, 1, 1,-1,-1, 1, 1)
#     U₂z = SVec{8,T}(U₂z₁,U₂z₁,U₂z₁,U₂z₁,U₂z₂,U₂z₂,U₂z₂,U₂z₂)
#     radical = SVec{8,T}(radical₁,radical₁,radical₁,radical₁,radical₂,radical₂,radical₂,radical₂)
#     extract_data(U₁ * muladd(m₂, radical, muladd(m₁, U₂z, x₀)) )
# end
# @inline function erfset_reduce(erfset, k)
#     VE(erfset[k+1].value + erfset[k+2].value - erfset[k+3].value - erfset[k+4].value)
# end

@inline function gdensity(z, HBR², x₀, U)

    @fastmath radical = sqrt( HBR² - (U[3]*z)^2 )

    U₂z = U[2]*z
    # ! Integral over positive and negative half of the ellipse.
    exp(-0.5z^2)*( erfc((x₀-U₂z-radical)*U[1])
                 + erfc((x₀+U₂z-radical)*U[1])
                 - erfc((x₀-U₂z+radical)*U[1])
                 - erfc((x₀+U₂z+radical)*U[1]) )
end

@inline function gdensity(z₁, z₂, HBR², x₀, U)

    @fastmath radical₁ = sqrt( HBR² - (U[3]*z₁)^2 )
    @fastmath radical₂ = sqrt( HBR² - (U[3]*z₂)^2 )

    U₂z₁ = U[2]*z₁
    U₂z₂ = U[2]*z₂

    @fastmath erfcs = SLEEFwrap.erfc(
        VE((x₀-U₂z₁-radical₁)*U[1]),
        VE((x₀+U₂z₁-radical₁)*U[1]),
        VE((x₀-U₂z₁+radical₁)*U[1]),
        VE((x₀+U₂z₁+radical₁)*U[1]),
        VE((x₀-U₂z₂-radical₂)*U[1]),
        VE((x₀+U₂z₂-radical₂)*U[1]),
        VE((x₀-U₂z₂+radical₂)*U[1]),
        VE((x₀+U₂z₂+radical₂)*U[1])
    )

    # ! Integral over positive and negative half of the ellipse.

    nk = SLEEF.exp((VE(-0.5z₁*z₁), VE(-0.5z₂*z₂)))
    erfreduces = (VE( erfcs[1].value
                    + erfcs[2].value
                    - erfcs[3].value
                    - erfcs[4].value ),
                  VE( erfcs[5].value
                    + erfcs[6].value
                    - erfcs[7].value
                    - erfcs[8].value ))
    SIMDPirates.vsum(SIMDPirates.vmul(nk, erfreduces))
end

function pcfoster_core(PriCovECI::SymmetricM3{T}, SecCovECI::SymmetricM3{T},
                        y::SVector{3,T}, z::SVector{3,T}, x₀::T, HBR = T(0.2), HBR² = HBR^2) where T
    Cp = combinecov(PriCovECI, SecCovECI, y, z)
    Up = revchol(Cp)
    Uv = SVector{3,T}(
        7.071067811865475244008443621048490392848359376884740365883398689953662392310596e-01/Up[1],
        Up[2], Up[3]
    )
    gaussianarea(Uv, x₀, HBR, HBR²)
end

function pcfoster_core(Σ::SymmetricM2{T}, x₀::T, HBR = T(0.2), HBR² = HBR^2) where T
    Up = revchol(Σ)
    Uv = SVector{3,T}(
        7.071067811865475244008443621048490392848359376884740365883398689953662392310596e-01/Up[1],
        Up[2], Up[3]
    )
    gaussianarea(Uv, x₀, HBR, HBR²)
end


using OneDimensionalCorrelation
using Test

@testset "Elements" begin
    e1 = Element1D(0, 1.2)
    e2 = Element1D(1.2, 2.4)

    @test (e1 < e2) == false
    @test (e2 > e1) == false
    @test e1 <= e2
    @test e2 >= e1

    p = 64
    eg1 = ElementGrid(0, 1.2, p)
    eg2 = ElementGrid(1.2, 2.4, p)

    @test (eg1 < eg2) == false
    @test (eg2 > eg1) == false
    @test eg1 <= eg2
    @test eg2 >= eg1

    eg = ElementGrid(0, 2π, p)

    # Integration accuracy
    @test isapprox( sum(sin.(eg) .* get_weight(eg)), 0.0; atol=1E-12 )

    # Derivative accuracy
    @test all( derivative_matrix(eg) * sin.(eg) ≈ cos.(eg) )

end

@testset "Basis Gauss-Lagrange" begin
    b = Basis(-π, π, 2, 64)

    # Symmetric range
    @test b[begin] == -b[end]

    w = get_weight(b)
    ∇ = derivative_matrix(b)

    # Integral
    @test isapprox( w' *  sin.(b), 0.0; atol=1E-12 )

    # Derivative
    @test ∇ * sin.(b) ≈ cos.(b)

end

@testset "Basis Gauss-Lobatto" begin
    b = BasisLobatto(-2, 2, 3, 32)

    # Symmetric range
    @test b[begin] == -b[end]

    w = get_weight(b)
    ∇ = derivative_matrix(b)

    # Integral
    @test isapprox( w' *  sin.(b), 0.0; atol=1E-12 )

    # Derivative
    @test  ∇ * sin.(b) ≈ cos.(b)

end

@testset "Integrals" begin
    b = BasisLobatto(-5, 5, 2, 64)

    ψ = exp.(-b.^2)
    I = get_identity(b)

    # Norm of Gaussian
    @test bracket(b, ψ, I, ψ) ≈ sqrt(π/2)
    @test bracket(b, ψ, ψ) ≈ sqrt(π/2)

    # Particle in box states
    ϕ1 = particle_in_box(b, 1)
    ϕ4 = particle_in_box(b, 4)
    @test bracket(b, ϕ1, ϕ1) ≈ 1
    @test bracket(b, ϕ4, ϕ4) ≈ 1
    @test bracket(b, ϕ1, ϕ4) + 1 ≈ 1

    # Particle in box Energy
    E(b, n) = 0.5 * n^2 * π^2 / (get_length(b))^2

    d = derivative_matrix(b)
    e_kin = -0.5 * d^2

    # Kinetic energy accuracy
    abs( bracket(b, ϕ1, e_kin, ϕ1) - E(b, 1) ) < 1E-9
    abs( bracket(b, ϕ4, e_kin, ϕ4) - E(b, 4) ) < 1E-9
end

@testset "Hartree Fock" begin
    b = BasisLobatto(-5, 5, 2, 64)
    orbitals = initial_orbitals(b)
    C = coulomb_matrix(b, orbitals)
    K = exchange_matrix(b, orbitals)
    F = fock_matrix(b)

    # Matrices are Hermitian
    @test C[2,1] ≈ C[2,1]
    @test K[2,1] ≈ K[2,1]
    @test F[2,1] ≈ F[2,1]

    V(x) = -2exp(-0.5x^2)
    E, eo, orbitals = solve_hartree_fock(b, V)
    @test E < -2.2123101
end

@testset "CI" begin
    b = BasisLobatto(-10, 10, 6, 12)
    V(x) = -2exp(-0.5x^2)
    E, eo, orbitals = solve_hartree_fock(b, V)

    eb = block_ci(b, orbitals, V)
    @test E > eb["energy"][1]
end


using Test
using WaveProp
using WaveProp.Geometry
using WaveProp.Integration
using WaveProp.Mesh

@testset "Circle" begin
    geo = Circle()
    bnd = boundary(geo)
    @test length(bnd) == 1
    Ω    = Domain([geo])
    mesh =  Mesh.meshgen(Ω;h=0.1)
    @test ambient_dimension(mesh) == 2
end


# Utility functions
indexOnBorder( index, csize ) = any( index .== 1 ) || any( index .== csize )

neighIndices( idx::II  ) = [ (idx[1]+y,idx[2]+x) for y in -1:1, x in -1:1 ]

neighIndices( idx::III ) = [ (idx[1]+y,idx[2]+x,idx[3]+z) for y in -1:1, x in -1:1, z in -1:1 ]

getNeighbours( index, cmatrix ) =  [ cmatrix[ x... ] for x in neighIndices(index) ]

# Val{true} is passed at the last mp iteration and Val{false} during the previous iterations.

function approxTranslation( cmat::Array{<:AbstractFloat,N}, alg::String, ::Val{false} ) where {N}

    ccenter = div.( size( cmat ) , 2 );
    return firstPeak( cmat )[1][1:N] .- ccenter
end

function approxTranslation( cmat::Array{<:AbstractFloat,N}, alg::String, ::Val{true} ) where {N}

    ccenter = div.( size( cmat ), 2 );
    peak1, maxVal = firstPeak( cmat );
    return subPixel( cmat, alg, peak1[1:N], maxVal ) .- ccenter;
end

function subPixel( cmatrix::Array{T,N}, alg::String, peak1::NTuple{N,I}, maxVal::T ) where {T,N}

    refinement = zeros( T, N );

    if alg == "" || indexOnBorder( peak1, size( cmatrix ) )
        # no sub-pixel refinement
    else
        neighs = getNeighbours( peak1, cmatrix );

        # ensuring values are positive. This does not change subpixel results
		minVal = minimum( neighs )
        for idx in 1:length(neighs)
            neighs[idx] = 1 + neighs[idx] - minVal
        end

        if     alg == "gaussian"
            refinement = gaussian( neighs, 1+maxVal-minVal );
        elseif alg == "centroid"
            refinement = centroid( neighs, maxVal );
        end

        # nan occur when neighbours contains 0's, resulting in 0/0
        refinement = [ isnan(x) ? 0.0 : x for x in refinement ]
    end

    return peak1 .+ refinement
end

""" 3-POINT GAUSSIAN SUBPIXEL """

# 3-point Gaussian subpixel 2D from the pixel neighbourhood (ns) around the max value (mx)
function gaussian( ns::Array{T,2}, mx::T ) where {T}
    return gaussian_2D( log(ns[4]), log(ns[6]), log(ns[2]), log(ns[8]), log(mx) );
end

function gaussian_2D( up::T, down::T, left::T, right::T, mx::T ) where {T<:AbstractFloat}
    return [ ( up  - down )/( 2* up  - 4*mx + 2*down  ),
             (left - right)/( 2*left - 4*mx + 2*right ) ];
end

# 3-point Gaussian subpixel 3D from the voxel neighbourhood (ns) around the max value (mx)
function gaussian( ns::Array{T,3}, mx::T ) where {T}
    return gaussian_3D( log(ns[13]), log(ns[15]), log(ns[11]), log(ns[17]),
                        log(ns[ 5]), log(ns[23]), log(mx) );
end

function gaussian_3D( up::T, down::T, left::T, right::T, front::T, back::T, mx::T ) where {T}
    return [ (  up  - down )/( 2*  up  - 4*mx + 2*down  ),
             ( left - right)/( 2* left - 4*mx + 2*right ),
             (front - back )/( 2*front - 4*mx + 2*back  ) ];
end

""" CENTROID SUBPIXEL """

# Centroid subpixel 2D from the pixel neighbourhood (ns) around the max value (mx)
function centroid( ns::Array{T,2}, mx::T ) where {T}
	return centroid_2D( ns[4], ns[6], ns[2], ns[8], mx )
end

function centroid_2D( up::T, down::T, left::T, right::T, mx::T ) where {T}
    return [ ( up  - down )/(  up  - mx + down  ),
             (left - right)/( left - mx + right ) ];
end

# Centroid subpixel 3D from the voxel neighbourhood (ns) around the max value (mx)
function centroid( ns::Array{T,3}, mx::T ) where {T<:AbstractFloat}
	return centroid_3D( ns[13], ns[15], ns[11], ns[17], ns[5], ns[23], mx )
end

function centroid_3D( up::T, down::T, left::T, right::T, front::T, back::T, mx::T ) where {T}
    return [ (  up  - down )/(   up  - mx + down  ),
             ( left - right)/(  left - mx + right ),
			 (front - back )/( front - mx + back  ) ];
end


module MASW

using FFTW, Plots

export masw

"""
    masw(U::AbstractArray, x::AbstractArray, t::AbstractArray, ct::AbstractArray;
               plotfreqlimits::Tuple=(),
               plotwaterlevel::Real=0.0,
               figdir::String="",
               figname::String="")

Return dispersion curve inferred from Multichannel Analysis of Surface Waves (MASW)

# Arguments
- `U::AbstractArray`: 2D Array for seismci traces: size(U)=(T:length of timeseries, N:number of traces)
- `x::AbstractArray`: Vector for offsets (distance between source and receiver)
- `t::AbstractArray`: Vector for timeseries (ex. 0:1/fs:Tmax)
- `ct::AbstractArray`: Vector for scanning velocity of interest.
- `plotfreqlimits::Tuple=()`: Tuple for minfreq and maxfreq of plotting (ex. (0.0, 10.0))
- `plotwaterlevel::Real=0.0`: Dispersion Image waterlevel expressed as a percent for display purposes.
- `figdir::String=""`: Figure directory name if you want to plot the result figure.
- `figname::String=""`: figname if you want to plot the result figure (please include extention.).

# Usage

`(f, ct, Udisp, DispersionCurve) = masw(U, x, t, ct, plotfreqlimits, plotwaterlevel, figdir, figname)`

# Output

- `f`: Vector for frequency associated with dispersion image.
- `ct`: Vector for scanning velocity associated with dispersion image.
- `Udisp`: Normalized 2D array for dispersion image.
- `DispersionCurve`: DispersionCurve including frequency and automatically picked phase velocity.

# Reference
- Park, C.B., Miller, R.D. and Xia, J. (2001) Offset and resolution of dispersion curve in multichannel analysis of surface waves (MASW). doi:10.4133/1.2922953

Julia version written by kura-okubo (https://github.com/kura-okubo)

This is the translation in Julia based on matlab script written by John Schuh.

Matlab script: Copyright (c) 2018, John Schuh
All rights reserved.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

- John Schuh (2020). MASW Dispersion Curve (https://www.mathworks.com/matlabcentral/fileexchange/65138-masw-dispersion-curve), MATLAB Central File Exchange. Retrieved January 14, 2020.
"""
function masw(U::AbstractArray, x::AbstractArray, t::AbstractArray, ct::AbstractArray;
                plotfreqlimits::Tuple=(),
                plotwaterlevel::Real=0.0,
                figdir::String="",
                figname::String="")

    # precheck of array size
    if size(U) != (length(t), length(x))
        error("size of U is not correct. Please check that size(U) == (length(t), length(x))")
    end

    datanorm = U./maximum(U)

    L=length(t); # number os samples
    T=maximum(t); # max time
    if mod(L,2) == 0
        f = (1/T)*(0:L/2-1); # n is even
    else
        f = (1/T)*(0:(L-1)/2); # n is odd
    end

    Ufft=FFTW.rfft(U, 1) # fft is performed on each column (the time vector)
    w=f*2*pi; # convert the freq vector to rads/sec

    ## 2. Normalize the spectrum
    Rnorm=Ufft./abs.(Ufft); # Normalize the spectrum

    ## 3. Perfrom the summation over N in the frequency domain
    AsSum = zeros(length(w),length(ct))
    @simd for ii=1:length(w) # send in frequencies one at a time
        As=zeros(Complex, length(x),length(ct)); #initialize As or overwite the previous As
        for n=1:length(x) # send in a positon (x) one at a time
            # Iterate As over position
            As[n,:]=exp.(1im*w[ii].*(x[n]./ct)).*Rnorm[ii,n]
        end
         AsSum[ii,:]=abs.(sum(As, dims=1));
    end
    AsSum=AsSum'; # transpose the matrix so velocity is on vertical and freq is on horizontal
    normed=AsSum./maximum(AsSum, dims=1);# normalize the dispersion image so max column values=1

    ############ curve autopicking ###################
    # This section will auto pick the dispersion curve based on the peak
    # value in each column. This will need updating when
    # higher modes are introduced
    remaindat=deepcopy(normed)
    remaindat[remaindat .< 1] .= 0
    ind=findall(x -> x==1, remaindat)
    row = []; col = []
    for i = 1:length(ind)
        push!(row, ind[i][1]); push!(col, ind[i][2])
    end

    autovel=ct[row]; # autopicked velocity indices that contain the peak value
    autofreq=f[col]; # autopicked frequency indices that contain the peak value
    # get rid of all the extra picks at zero frequency
    nz=length(autofreq)-count(x -> x != 0, autofreq); # determine the number of repeated picks at zero
    autofreq=autofreq[nz:end]; # remove the repeated zero picks
    autovel=autovel[nz:end]; # remove the repeated zero picks
    DispersionVelocity=autovel; # need to get the rlowess working again, I no longer have access to the toolbox that allows it
    normed[normed .< plotwaterlevel/100] .= NaN; # throw out data less than plotwaterlevel
    ################ end autopicking ######################

    DispersionCurve=[f, DispersionVelocity]; # this only has meaning if the curve was picked using 'manual' or 'auto'

    if !isempty(figdir)

        if Plots.backend() != Plots.PyPlotBackend()
            @warn ("MASW.jl plotting is available only with pyplot backend. Please type
                    `using Plots;Plots.pyplot()` before call this function. Return result
                     without plotting figure.")
            return (f, ct, normed, DispersionCurve)
        end

        # plot input traces
        l = Plots.@layout [a b]
        p1 = Plots.contourf(x, t, U, levels=200,  yflip = true, xlabel="Distance", ylabel="Time",
                            linewidth=0.0, c=:balance,  title="Traces")
        # Plot the dispersion results
        p2 = Plots.contourf(f,ct,normed, levels=200, xlabel="Freq (Hz)", ylabel="Phase Velocity",
                            clims=(0.0,1.0), linewidth=0.0, c=:thermal, title="Dispersion Image") # plot the normalized dispersion image
        Plots.plot!(DispersionCurve[1],DispersionCurve[2], line=(:black, 2.0), label="dispersion curve") # Plot the picked dispersion curve. This is nothing if 'none' picking was set
        if !isempty(plotfreqlimits)
            Plots.xlims!(plotfreqlimits)
        end
        Plots.scalefontsizes(3)
        p = Plots.plot(p1, p2, size=(1600, 800), layout=l)
        if !ispath(figdir)
            mkpath(figdir)
        end
        Plots.savefig(p, figdir*"/"*figname)
    end

    return (f, ct, normed, DispersionCurve)
end

end # module


@testset "Rosenbrock" begin

    function test_result(result::Evolutionary.EvolutionaryOptimizationResults, N::Int, tol::Float64)
        fitness = minimum(result)
        extremum = Evolutionary.minimizer(result)
        @test sum(abs2,extremum.-ones(N)) ≈ 0.0 atol=tol
        @test fitness ≈ 0.0 atol=tol
    end

    # Objective function
    rosenbrock(x::AbstractVector) = (1.0 - x[1])^2 + 100 * (x[2] - x[1]^2)^2

    # Parameters
    N = 2

    # Testing: (15,3/(+,)100)-ES
    settings = [
        :isotropic=>(gaussian, gaussian, IsotropicStrategy(N)),
        :isotropic=>(cauchy, gaussian, IsotropicStrategy(N)),
        :anisotropic=>(gaussian, gaussian, AnisotropicStrategy(N))
    ]
    selections = [:plus, :comma]
    @testset "ES settings" for (sn,ss) in settings, sel in selections
        result = Evolutionary.optimize( rosenbrock, (() -> rand(N)),
            ES(
                initStrategy = ss[3],
                recombination = average, srecombination = average,
                mutation = ss[1], smutation = ss[2],
                μ = 10, ρ = 3, λ = 100, selection=sel
            ), Evolutionary.Options(iterations=1000, successive_f_tol=25)
        )
        println("(15/3$(sel == :plus ? "+" : ",")100)-ES:$(sn) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
        test_result(result, N, sel == :plus ? 0.1 : 0.5)
    end

    # Testing: CMA-ES
    result = Evolutionary.optimize(rosenbrock, (() -> rand(Float32,N)), CMAES(μ = 5, λ = 100))
    println("(5/5,100)-CMA-ES => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
    test_result(result, N, 1e-2)

    # Testing: GA
    result = Evolutionary.optimize(rosenbrock, (() -> rand(N)),
        GA(
            populationSize = 100,
            ɛ = 0.1,
            selection = rouletteinv,
            crossover = intermediate(0.25),
            mutation = domainrange(fill(0.5,N))
        ))
    println("GA(p=100,x=0.8,μ=0.1,ɛ=0.1) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
    test_result(result, N, 1e-1)

    # Testing: DE
    result = Evolutionary.optimize(rosenbrock, (() -> rand(N)), DE(populationSize = 100))
    println("DE/rand/1/bin(F=1.0,Cr=0.5) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
    test_result(result, N, 1e-1)

end


# https://gridap.github.io/Tutorials/stable/pages/t009_stokes/

using Pkg
Pkg.activate("simulate")

using Gridap

n = 100
domain = (0,1,0,1)
partition = (n,n)
model = CartesianDiscreteModel(domain, partition)

labels = get_face_labeling(model)
add_tag_from_tags!(labels,"diri1",[6,])
add_tag_from_tags!(labels,"diri0",[1,2,3,4,5,7,8])

V = TestFESpace(
  reffe=:QLagrangian, conformity=:H1, valuetype=VectorValue{2,Float64},
  model=model, labels=labels, order=2, dirichlet_tags=["diri0","diri1"])
Q = TestFESpace(
  reffe=:PLagrangian, conformity=:L2, valuetype=Float64,
  model=model, order=1, constraint=:zeromean)
Y = MultiFieldFESpace([V,Q])

u0 = VectorValue(0,0)
u1 = VectorValue(1,0)
U = TrialFESpace(V,[u0,u1])
P = TrialFESpace(Q)
X = MultiFieldFESpace([U,P])

trian = get_triangulation(model); degree = 2
quad = CellQuadrature(trian,degree)

function a(x,y)
  v,q = y
  u,p = x
  ∇(v)⊙∇(u) - (∇⋅v)*p + q*(∇⋅u)
end
t_Ω = LinearFETerm(a,trian,quad)
op = AffineFEOperator(X,Y,t_Ω)
uh, ph = solve(op)

writevtk(trian,"results",order=2,cellfields=["uh"=>uh,"ph"=>ph])


function UCEP_mix(model::EoSModel;v0=nothing)
    if v0 == nothing
        v0 = x0_UCEP_mix(model)
    end
    ts = T_scales(model,[0.5,0.5])
    pmix = p_scale(model,[0.5,0.5])
    f! = (F,x) -> Obj_UCEP_mix(model, F, x[1], x[2], exp10(x[3]), exp10(x[4]), x[5],ts,pmix)
    r  = Solvers.nlsolve(f!,v0[1:end],LineSearch(Newton()))
    sol = Solvers.x_sol(r)
    x = FractionVector(sol[1])
    y = FractionVector(sol[2])
    V_l = exp10(sol[3])
    V_v = exp10(sol[4])
    T = sol[5]
    p = pressure(model, V_l, T, x)
    return (T, p, V_l, V_v, x, y)
end

function Obj_UCEP_mix(model::EoSModel,F,x,y,V_l,V_v,T,ts,ps)
    y    = FractionVector(y)
    x    = FractionVector(x)
    f(z) = eos(model,V_l,T,z)
    H(z) = ForwardDiff.hessian(f,z)/8.134/T
    L(z) = det(H(z))
    dL(z) = ForwardDiff.gradient(L,z)
    M(z) = [H(z)[1:end-1,:];transpose(dL(z))]

    μ_l = VT_chemical_potential(model,V_l,T,x)
    μ_v = VT_chemical_potential(model,V_v,T,y)
    p_l = pressure(model,V_l,T,x)
    p_v = pressure(model,V_v,T,y)

    for i in 1:length(x)
        F[i] = (μ_l[i]-μ_v[i])/(R̄*ts[i])
    end
    F[3] = (p_l-p_v)/ps

    F[4] = L(x)
    F[5] = det(M(x))
    return F
end

function x0_UCEP_mix(model::EoSModel)
    T0 = T_scale(model,[0.5,0.5])*1.5
    x0 = 0.5
    y0 = 0.75
    v0 = x0_bubble_pressure(model,T0,[x0,1-x0])
    return [x0,y0,v0[1],v0[2],T0]
end


"""
    Code for a half-cosine mask shape for use with CCF
Author: Eric Ford
Created: August 2020
"""

"""   CosCCFMask
Cosine mask with one parameter, it's quarter period, i.e., where to truncate it, as a velocity in m/s.
Mask weights are stored separately in a line list.
"""
struct CosCCFMask <: AbstractCCFMaskShape
    half_width::Float64

    """ CosCCFMask( full_width ) """
    function CosCCFMask( w::Real )
        @assert 100 <= w <= 300000   # 0.1 and 300km/s are arbitrary choices for an upper limit
        new(w/2)
    end

end

""" `CosCCFMask( inst  )` """
function CosCCFMask(inst::AbstractInstrument; scale_factor::Real = 1)
    w = scale_factor * default_ccf_mask_v_width(inst)
    CosCCFMask(w)
end

λ_min(m::CosCCFMask,λ::Real) = λ/calc_doppler_factor(m.half_width)
λ_max(m::CosCCFMask,λ::Real) = λ*calc_doppler_factor(m.half_width)

function integrate(m::CosCCFMask, v_lo::Real,v_hi::Real)
    #quadgk(m, v_lo, v_hi)[1]
    0.5*(sin(0.5*π*v_hi/m.half_width)-sin(0.5*π*v_lo/m.half_width))
end

""" Functor for returning PSF for Δv <= half_width.  """
function (m::CosCCFMask)(Δv::Real)
    if abs2(Δv) > abs2(m.half_width)
        return zero(Δv)
    else
        return cos(0.5*π*Δv/m.half_width)*π/(4*m.half_width)
    end
end

function mask_with_increased_fwhm(m::CosCCFMask, Δfwhm::Real )
    width = 2*m.half_width+Δfwhm
    return CosCCFMask(width)
end


# constants
const a = 6.378137e6 # the semi-major axis of the earth
const F = 298.257222101 # the inverse 1st flattening
const n = 1/(2F-1)
const m₀ = 0.9999 # central meridian scale factor
A₀ = 1 + n^2/4 + n^4/64
const A̅ = A₀/(1+n) *m₀*a

## longitude, latitude of origin for zones in Japan
const Origin_LatLon_Japan = [ 33.0  129.5    ; # 1
                              33.0  131.0    ; # 2
                              36.0  132.0+1/6; # 3
                              33.0  133.5    ; # 4
                              36.0  134.0+1/3; # 5
                              36.0  136.0    ; # 6
                              36.0  137.0+1/6; # 7
                              36.0  138.5    ; # 8
                              36.0  139.0+5/6; # 9
                              40.0  140.0+5/6; # 10
                              44.0  140.25   ; # 11
                              44.0  142.25   ; # 12
                              44.0  144.25   ; # 13
                              26.0  142.0    ; # 14
                              26.0  127.5    ; # 15
                              26.0  124.0    ; # 16
                              26.0  131.0    ; # 17
                              20.0  136.0    ; # 18
                              26.0  154.0    ; # 19
                             ]


function getTableauHSPARK(s, σ, o, tsym, g, h, lq, lp, ω, d=Nothing)
    # α_q_1 and α_p_2 need to be conjugate symplectic
    # α_q_1 and α_p_3 need to be conjugate symplectic
    # α_q_2 and α_p_2 need to be conjugate symplectic
    # α_q_2 and α_p_3 need to be conjugate symplectic
    # -> α_q_1 and α_q_2 need to be identical
    # -> α_p_2 and α_p_3 need to be identical

    α_q_2 = lq.a
    α_q_3 = lq.a
    α_p_2 = lp.a
    α_p_3 = lp.a

    b_q_2 = lq.b
    b_q_3 = lq.b
    b_p_2 = lp.b
    b_p_3 = lp.b

    # α_p_1 is free, but determines a_q_1 and a_q_2
    α_q_1 = h.a
    α_p_1 = h.a

    # a_p_1, a_p_2 and a_p_3 are free (they are not used in HSPARKpSecondary integrators anymore)
    # REALLY?
    a_q_1 = g.a
    b_q_1 = g.b

    a_p_1 = g.a
    b_p_1 = g.b

    a_p_2 = g.a
    a_p_3 = g.a

    β_q_1 = lq.b
    β_q_2 = lq.b
    β_q_3 = lq.b

    β_p_1 = lp.b
    β_p_2 = lp.b
    β_p_3 = lp.b

    # a_q_1_ij = b_q_1_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # a_q_1 = zeros(s,σ)
    # for i in 1:s
    #     for j in 1:σ
    #         a_q_1[i,j] = b_q_1[j] / b_p_1[i] * (b_p_1[i] - α_p_1[j,i])
    #     end
    # end

    # a_q_1_ij = b_q_1_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # TODO check this
    a_q_2 = zeros(s,σ)
    for i in 1:s
        for j in 1:σ
            a_q_2[i,j] = b_q_2[j] / b_p_1[i] * (b_p_1[i] - α_p_2[j,i])
        end
    end

    # a_q_2_ij = b_q_2_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # TODO check this
    a_q_3 = zeros(s,σ)
    for i in 1:s
        for j in 1:σ
            a_q_3[i,j] = b_q_3[j] / b_p_1[i] * (b_p_1[i] - α_p_3[j,i])
        end
    end

    a_q = (a_q_1, a_q_2, a_q_3)
    b_q = (b_q_1, b_q_2, b_q_3)
    c_q = g.c

    a_p = (a_p_1, a_p_2, a_p_3)
    b_p = (b_p_1, b_p_2, b_p_3)
    c_p = g.c

    α_q = (α_q_1, α_q_2, α_q_3)
    β_q = (β_q_1, β_q_2, β_q_3)
    γ_q = lq.c

    α_p = (α_p_1, α_p_2, α_p_3)
    β_p = (β_p_1, β_p_2, β_p_3)
    γ_p = lp.c

    coeff_q = CoefficientsSPARK(tsym, o, s, σ, a_q, b_q, c_q)
    coeff_p = CoefficientsSPARK(tsym, o, s, σ, a_p, b_p, c_p)
    coeff_q̃ = CoefficientsSPARK(tsym, o, σ, s, α_q, β_q, γ_q)
    coeff_p̃ = CoefficientsSPARK(tsym, o, σ, s, α_p, β_p, γ_p)

    TableauHSPARKsecondary(tsym, o, s, σ, coeff_q, coeff_p, coeff_q̃, coeff_p̃, ω, d)
end

function TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIII"))
    o = 2s-2
    getTableauHSPARK(s, s, o, name, lq, lp, lq, lp, get_lobatto_ω_matrix(s), get_lobatto_nullvector(s))
end

function TableauHSPARKLobattoIIIAB(s)
    lq = TableauLobattoIIIA(s)
    lp = TableauLobattoIIIB(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIAIIIB"))
end

function TableauHSPARKLobattoIIIBA(s)
    lq = TableauLobattoIIIB(s)
    lp = TableauLobattoIIIA(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIBIIIA"))
end

function TableauHSPARKLobattoIIICC̄(s)
    lq = TableauLobattoIIIC(s)
    lp = TableauLobattoIIIC̄(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIICIIIC̄"))
end

function TableauHSPARKLobattoIIIC̄C(s)
    lq = TableauLobattoIIIC̄(s)
    lp = TableauLobattoIIIC(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIC̄IIIC"))
end

function TableauHSPARKLobattoIIID(s)
    l = TableauLobattoIIID(s)
    TableauHSPARKLobattoIII(s, l, l; name = Symbol("HSPARKLobattoIIID"))
end

function TableauHSPARKLobattoIIIE(s)
    l = TableauLobattoIIIE(s)
    TableauHSPARKLobattoIII(s, l, l; name = Symbol("HSPARKLobattoIIIE"))
end

function TableauHSPARKGLRKLobattoIII(s, σ, lq, lp; name = Symbol("HSPARKGLRKLobattoIII"))
    o = 2s
    g = TableauGauss(s)
    getTableauHSPARK(s, σ, o, name, g, get_lobatto_glrk_coefficients(s, σ), lq, lp, get_GLRK_ω_matrix(σ), get_lobatto_nullvector(σ))
end

function TableauHSPARKGLRKLobattoIIIAB(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIA(σ), TableauLobattoIIIB(σ); name = Symbol("HSPARKGLRKLobattoIIIAIIIB"))
end

function TableauHSPARKGLRKLobattoIIIBA(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIB(σ), TableauLobattoIIIA(σ); name = Symbol("HSPARKGLRKLobattoIIIBIIIA"))
end

function TableauHSPARKGLRKLobattoIIICC̄(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIC(σ), TableauLobattoIIIC̄(σ); name = Symbol("HSPARKGLRKLobattoIIICIIIC̄"))
end

function TableauHSPARKGLRKLobattoIIIC̄C(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIC̄(σ), TableauLobattoIIIC(σ); name = Symbol("HSPARKGLRKLobattoIIIC̄IIIC"))
end

function TableauHSPARKGLRKLobattoIIID(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIID(σ), TableauLobattoIIID(σ); name = Symbol("HSPARKGLRKLobattoIIID"))
end

function TableauHSPARKGLRKLobattoIIIE(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIE(σ), TableauLobattoIIIE(σ); name = Symbol("HSPARKGLRKLobattoIIIE"))
end


# Implements the SpikeSlab likelihood.
# Some constants
u = 0.01
v = 0.1
C = log(1.0/sqrt(2*pi))

@doc """
A function.
""" ->
function logsumexp(values::Vector{Float64})
	biggest = maximum(values)
	x = values - biggest
	return log(sum(exp(x))) + biggest
end

@doc """
Evaluate the log likelihood.
""" ->
function log_likelihood(params::Vector{Float64})
	logl1 = length(params)*(C - log(u))
	logl2 = length(params)*(C - log(v))
	for(i in 1:length(params))
		logl1 += -0.5*((params[i] - 0.5 - 0.031)/u)^2
		logl2 += -0.5*((params[i] - 0.5)/v)^2
	end
	logl1 += log(100.0)
	return logsumexp([logl1, logl2])
end


"""
    ImbensWagerOptRD(B::Float64, solver, num_buckets=2000, variance=Homoskedastic()) <: SharpRD

A `SharpRD` estimator.

`B` is a bound on the second derivative of the mean of
the outcome variable, conditional on the running variable above and below the cutoff.

The `solver` object is a `MathOptInterface` compatible optimizer such as `Mosek.Optimizer` that will be
used to solve the numerical optimization problem described in Imbens & Wager (2019).

The number of bins for the discretization of the running variable that is required to solve
the numerical optimization problem is `num_buckets`.
"""
Base.@kwdef struct ImbensWagerOptRD <: SharpRD
    B::Float64
    solver::Any
    num_buckets = 2000
    variance = Homoskedastic()
end

Base.@kwdef struct FittedOptRD
    tau_est::AbstractFloat
    se_est::AbstractFloat
    maxbias::AbstractFloat
    ci::Array{Float64}
    weights::Array{Float64}
    B::Float64
    coeftable::Any
end

function estimate_var(data::RDData)
    fitted_lm = fit(LinearModel, @formula(Ys ~ Ws * ZsC), data)
    yhat = predict(fitted_lm)
    σ² = mean((data.Ys - yhat) .^ 2) * length(data.Zs) / (length(data.Zs) - 4)
    return σ²
end

function bias_adjusted_gaussian_ci(se; maxbias = 0.0, level = 0.95)
    rel_bias = maxbias / se
    zz = fzero(
        z -> cdf(Normal(), rel_bias - z) + cdf(Normal(), -rel_bias - z) + level - 1,
        0,
        rel_bias - quantile(Normal(), (1 - level) / 2.1),
    )
    zz * se
end

function Statistics.std(::Homoskedastic, γ::Vector, sig2::Float64)
    sqrt(sum(γ .^ 2) * sig2)
end

"""
    fit(method::ImbensWagerOptRD, data::RDData; level=0.95)

Find the min-max optimal RD estimator over a convex class of conditional mean
functions with a bounded second derivative, based on the numerical optimization
procedure described in Imbens & Wager (2019).

Returns a FittedOptRD object with fields `tau_est`, `se_est`, `maxbias`, `ci`,
`weights` the RD estimator weights for each observation, `B` the bound on
the second derivative of the conditional mean functions used in estimation, and
`coeftable` providing an object of type CoefTable for the estimator.
"""
function fit(method::ImbensWagerOptRD, data::RDData; level = 0.95)
    B = method.B
    ZsD = DiscretizedRunningVariable(data.ZsR, method.num_buckets)
    X = ZsD.Zs
    h = ZsD.h
    d = length(X)
    n = ZsD.weights
    ixc = argmin(abs.(ZsD.Zs .- ZsD.cutoff))
    c = ZsD.cutoff
    W = ZsD.Ws
    sig2 = estimate_var(data)
    σ² = zeros(d) .+ sig2
    model = Model(method.solver)

    @variable(model, f[1:d])
    @variable(model, l1 >= 0)
    @variable(model, l2)
    @variable(model, l3)
    @variable(model, l4)
    @variable(model, l5)
    @variable(model, s)

    G = @expression(
        model,
        2 * B .* f .+ l2 .* (1 .- W) +
        l3 .* W +
        l4 .* (X .- c) +
        l5 .* (W .- 0.5) .* (X .- c)
    )

    @constraint(model, f[ixc] == 0)
    @constraint(model, (f[ixc+1] - f[ixc]) / h[ixc] == 0)

    for i = 3:d
        @constraint(model, (f[i] - 2 * f[i-1] + f[i-2]) / (h[i-1] * h[i-2]) - l1 <= 0)
        @constraint(model, (f[i] - 2 * f[i-1] + f[i-2]) / (h[i-1] * h[i-2]) + l1 >= 0)
    end

    qobj = @expression(model, [1 / 2 .* sqrt.(n) .* G ./ sqrt.(σ²); l1])
    @constraint(model, [0.5; s; qobj] in RotatedSecondOrderCone())

    @objective(model, Min, s - l2 + l3)
    optimize!(model)
    γ_xx = -value.(G) ./ (2 .* σ²)
    γ = γ_xx[ZsD.binmap]

    τ = sum(γ .* data.Ys)
    maxbias = value(l1)
    se = Statistics.std(method.variance, γ, sig2)
    ci = bias_adjusted_gaussian_ci(se; maxbias = maxbias, level = level)
    ci = [τ - ci, τ + ci]
    levstr = isinteger(level * 100) ? string(Integer(level * 100)) : string(level * 100)
    res = [τ se maxbias first(ci) last(ci)]
    colnms = ["τ̂"; "se"; "max bias"; "Lower $levstr%"; "Upper $levstr%"]
    rownms = ["Sharp RD estimand"]
    coeftbl = CoefTable(res, colnms, rownms, 0, 0)

    FittedOptRD(
        tau_est = sum(γ .* data.Ys),
        se_est = se,
        maxbias = maxbias,
        ci = ci,
        weights = γ,
        B = B,
        coeftable = coeftbl,
    )
end

function Base.show(io::IO, rdd_fit::RegressionDiscontinuity.FittedOptRD)
    println(io, "Imbens-Wager (2019) optimized regression discontinuity design")
    println(io, "Max Second Derivative Bound: ", string(rdd_fit.B))
    Base.show(io, rdd_fit.coeftable)
end


using CompScienceMeshes
using BEAST
using Test
using LinearAlgebra

o, x, y, z = CompScienceMeshes.euclidianbasis(3)
tet = simplex(x,y,z,o)

rs = BEAST.NDLCDRefSpace{Float64}()
Q = BEAST.restrict(rs, tet, tet)
@test Q ≈ Matrix(1.0LinearAlgebra.I, 4, 4)

rs = BEAST.NDLCCRefSpace{Float64}()
Q = BEAST.restrict(rs, tet, tet)
@test Q ≈ Matrix(1.0LinearAlgebra.I, 6, 6)

c = cartesian(center(tet))
smalltet = simplex(x,y,z,c)
p_smalltet = center(smalltet)
p_tet = neighborhood(tet, carttobary(tet, cartesian(p_smalltet)))
@show cartesian(p_smalltet)
@show cartesian(p_tet)
@assert cartesian(p_tet) ≈ cartesian(p_smalltet)
@assert volume(tet) / volume(smalltet) ≈ 4

Q = BEAST.restrict(rs, tet, smalltet)
Fp = rs(p_tet)
fp = rs(p_smalltet)

for j in axes(Q,1)
    x = Fp[j].value
    y = sum(Q[j,i]*fp[i].value for i in axes(Q,2))
    @show x
    @show y
    @test x ≈ y
end


using Oceananigans.Advection:
    _advective_momentum_flux_Uu,
    _advective_momentum_flux_Uv,
    _advective_momentum_flux_Vu,
    _advective_momentum_flux_Vv

using Oceananigans.Grids: AbstractGrid
using Oceananigans.Operators: Ax_qᶠᶜᶜ, Ay_qᶜᶠᶜ

#####
##### Momentum flux operators
#####

@inline momentum_flux_huu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uu(i, j, k, grid, advection, solution.uh, solution.uh) / solution.h[i, j, k]

@inline momentum_flux_hvu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vu(i, j, k, grid, advection, solution.vh, solution.uh) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_huv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uv(i, j, k, grid, advection, solution.uh, solution.vh) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_hvv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vv(i, j, k, grid, advection, solution.vh, solution.vh) / solution.h[i, j, k]

#####
##### Momentum flux divergence operators
#####

@inline div_hUu(i, j, k, grid, advection, solution) =
    1 / Vᶠᶜᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, momentum_flux_huu, advection, solution) +
                               δyᵃᶜᵃ(i, j, k, grid, momentum_flux_hvu, advection, solution))

@inline div_hUv(i, j, k, grid, advection, solution) =
    1 / Vᶜᶠᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_huv, advection, solution) +
                               δyᵃᶠᵃ(i, j, k, grid, momentum_flux_hvv, advection, solution))

# Support for no advection
@inline div_hUu(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution) where FT = zero(FT)
@inline div_hUv(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution) where FT = zero(FT)

#####
##### Mass transport divergence operator
#####

"""
    div_Uh(i, j, k, grid, solution)

Calculates the divergence of the mass flux into a cell,

    1/V * [δxᶜᵃᵃ(Ax * uh) + δyᵃᶜᵃ(Ay * vh)]

which will end up at the location `ccc`.
"""
@inline function div_Uh(i, j, k, grid, solution)
    1/Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, solution.uh) +
                             δyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, solution.vh))
end

#####
##### Tracer advection operator
#####

@inline transport_tracer_flux_x(i, j, k, grid, advection, uh, h, c) =
    @inbounds advective_tracer_flux_x(i, j, k, grid, advection, uh, c) / h[i, j, k]

@inline transport_tracer_flux_y(i, j, k, grid, advection, vh, h, c) =
    @inbounds advective_tracer_flux_y(i, j, k, grid, advection, vh, c) / h[i, j, k]

"""
    div_Uc(i, j, k, grid, advection, U, c)

Calculates the divergence of the flux of a tracer quantity c being advected by
a velocity field U = (u, v), ∇·(Uc),

    1/V * [δxᶜᵃᵃ(Ax * uh * ℑxᶠᵃᵃ(c) / h) + δyᵃᶜᵃ(Ay * vh * ℑyᵃᶠᵃ(c) / h)]

which will end up at the location `ccc`.
"""
@inline function div_Uc(i, j, k, grid, advection, solution, c)
    1/Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, transport_tracer_flux_x, advection, solution.uh, solution.h, c) +
                             δyᵃᶜᵃ(i, j, k, grid, transport_tracer_flux_y, advection, solution.vh, solution.h, c))
end

# Support for no advection
@inline div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c) where FT = zero(FT)

@inline u(i, j, k, grid, solution) = @inbounds solution.uh[i, j, k] / solution.h[i, j, k]
@inline v(i, j, k, grid, solution) = @inbounds solution.vh[i, j, k] / solution.h[i, j, k]

"""
    c_div_U(i, j, k, grid, advection, U)

Calculates the product of the tracer concentration c with
the horizontal divergence of the velocity field U = (u, v), c ∇·(U),

    1/V * [δxᶜᵃᵃ(Ax * uh / h) + δyᵃᶜᵃ(Ay * vh / h]

which will end up at the location `ccc`.
"""
@inline function c_div_U(i, j, k, grid, solution, c)
    @inbounds c[i, j, k] * 1/Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, u, solution) + δyᵃᶜᵃ(i, j, k, grid, v, solution))
end

# Support for no advection
@inline c_div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c) where FT = zero(FT)


    function PLacExampleUtility(ins)

        # Definition of the inputs for the ODEs
        if length(ins) == 1
            IPTG2,IPTG3 = ins[1];
        else
            IPTG2,IPTG3 = ins;
        end
        inputsMC = [0,IPTG2,IPTG3];

        # Solve ODEs
        solMC = PLacExample_SolveAll(tsMC, pD1MC, spMC, inputsMC, ivss1MC, sampsMC, pre1MC);

        # Extracte wanted vectors (observables) with time reduction
        Obs1_MC = solMC[:,2,:];
    Obs2_MC = 3 .*solMC[:,4,:];

        EntObs1_MC = zeros(size(solMC)[1]);
    EntObs2_MC = zeros(size(solMC)[1]);

        dists = [Beta, Exponential, LogNormal, Normal, Gamma, Laplace, Pareto, Rayleigh, Cauchy, Uniform];
        dists2 = [Exponential, LogNormal, Normal, Gamma, Laplace, Pareto, Rayleigh, Cauchy, Uniform];

        names = ["timePoint"];

        for j in 1:size(solMC)[1]
            fitts1 = Dict();
        bestfit1 = Dict();
        bestfitInd1 = zeros(1,1);
        fitts2 = Dict();
        bestfit2 = Dict();
        bestfitInd2 = zeros(1,1);

            try
                fitts1[names[1]] = fit.(dists, Ref(Obs1_MC[j,:]));
                bestfitInd1[1] = findmax(loglikelihood.(fitts1[names[1]], Ref(Obs1_MC[j,:])))[2];
                bestfit1[names[1]] = fitts1[names[1]][findmax(loglikelihood.(fitts1[names[1]], Ref(Obs1_MC[j,:])))[2]];
            catch
                fitts1[names[1]] = fit.(dists2, Ref(Obs1_MC[j,:]));
                bestfitInd1[1] = findmax(loglikelihood.(fitts1[names[1]], Ref(Obs1_MC[j,:])))[2]+1;
                bestfit1[names[1]] = fitts1[names[1]][findmax(loglikelihood.(fitts1[names[1]], Ref(Obs1_MC[j,:])))[2]];
            end
            try
                fitts2[names[1]] = fit.(dists, Ref(Obs2_MC[j,:]));
                bestfitInd2[1] = findmax(loglikelihood.(fitts2[names[1]], Ref(Obs2_MC[j,:])))[2];
                bestfit2[names[1]] = fitts2[names[1]][findmax(loglikelihood.(fitts2[names[1]], Ref(Obs2_MC[j,:])))[2]];
            catch
                fitts2[names[1]] = fit.(dists2, Ref(Obs2_MC[j,:]));
                bestfitInd2[1] = findmax(loglikelihood.(fitts2[names[1]], Ref(Obs2_MC[j,:])))[2]+1;
                bestfit2[names[1]] = fitts2[names[1]][findmax(loglikelihood.(fitts2[names[1]], Ref(Obs2_MC[j,:])))[2]];
            end

            EntObs1_MC[j] = entropy(bestfit1[names[1]]);
    EntObs2_MC[j] = entropy(bestfit2[names[1]]);

        end

        HES = zeros(1,2);
        HES[1,1] = sum(EntObs1_MC);
    HES[1,2] = sum(EntObs2_MC);

        util = sum(HES)/2

        return(util)
    end