using ADCME
using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

SMALL_SIZE = 22
MEDIUM_SIZE = 24
BIGGER_SIZE = 24

plt.rc("font", size=SMALL_SIZE)          # controls default text sizes
plt.rc("axes", titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc("axes", labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc("xtick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc("ytick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc("legend", fontsize=SMALL_SIZE)    # legend fontsize
plt.rc("figure", titlesize=BIGGER_SIZE)  # fontsize of the figure title

# Example 1:

function u_exact(x,y)
    if y == 0.0
        return 1.0
    else
        return 0.0
    end
end

function v_exact(x,y)
    0.0
end

function p_exact(x,y)
    0.0
end

function ffunc(x, y)
    0.0
end
function gfunc(x, y)
    0.0
end

function hfunc(x,y)
    0.0
end

function nu_exact(x, y)
    1 + 6 * x^2 + x / (1 + 2 * y^2)
end

function nu_nn(x, y)
    1.5 + fc([x y], [20,20,20,1])|>squeeze
end

m = 20
n = 20
h = 1/n
# nu = 0.01

F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)
# nu_gauss = Variable(ones(4*m*n))
# nu_fvm = Variable(ones(m*n))
nu_fvm = matread("sess_figures/ex1/ones50.mat")["Variablecolon0"]
nu_fvm = constant(nu_fvm)
nu_gauss = reshape(repeat(nu_fvm, 1, 4), (-1,))

Laplace = compute_fem_laplace_matrix1(nu_gauss, m, n, h)

# Laplace = constant(nu*compute_fem_laplace_matrix1(m, n, h))
function compute_residual(S)
    u, v, p = S[1:(m+1)*(n+1)], S[(m+1)*(n+1)+1:2(m+1)*(n+1)], S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    R = [F;G;H0]
    return R
end

function compute_jacobian(S)
    u, v, p = S[1:(m+1)*(n+1)], S[(m+1)*(n+1)+1:2(m+1)*(n+1)], S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    J0 = [Fu Fv
          Gu Gv]
    J = [J0 -B'
        -B spdiag(zeros(size(B,1)))]
end

NT = 5
# S = zeros(m*n+2(m+1)*(n+1), NT+1)
# S[1:m+1, 1] = ones(m+1,1)

bd = bcnode("all", m, n, h)
# bd = [bd; bd .+ (m+1)*(n+1); ((1:m) .+ 2(m+1)*(n+1))]
bd = [bd; bd .+ (m+1)*(n+1); 2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2] # only apply Dirichlet to velocity

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    # residual[bd] .= 0.0
    residual = scatter_update(residual, bd, zeros(length(bd)))

    d = J\residual
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, [u0; v0; p0])

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n))

S_true = matread("sess_figures/ex1/steady_cavity_data.mat")["V"]
loss = mean((S[end,1:2*(m+1)*(n+1)] - S_true[end,1:2*(m+1)*(n+1)])^2)
loss = loss * 1e10

# vmax = maximum(nu_exact.(x,y))
# vmin = minimum(nu_exact.(x,y))

sess = Session(); init(sess)
# ADCME.load(sess, "sess_figures/ex1/ones50.mat")
nu_fvm = matread("sess_figures/ex1/ones50.mat")["Variablecolon0"]

vmax = 6.8
vmin = -0.5

figure();visualize_scalar_on_fvm_points(nu_exact.(x,y), m, n, h, vmax=vmax, vmin=vmin); tight_layout()
savefig("sess_figures/ex1/ones/viscosity_exact.png")

# figure();visualize_scalar_on_gauss_points(run(sess, nu_gauss), m, n, h); #gca().invert_yaxis()
# savefig("sess_figures/ex1/ones/viscosity_prediction.png")

figure();visualize_scalar_on_gauss_points(run(sess, nu_gauss), m, n, h, vmax=vmax, vmin=vmin); tight_layout() ;#gca().invert_yaxis()
savefig("sess_figures/ex1/ones/viscosity_prediction.png")

# figure();visualize_scalar_on_fvm_points(nu_exact.(x,y).-run(sess, nu_gauss)[1:4:end], m, n, h);
# savefig("sess_figures/ex1/ones/viscosity_error.png")

figure();visualize_scalar_on_fvm_points(nu_exact.(x,y).-nu_fvm, m, n, h); tight_layout()
savefig("sess_figures/ex1/ones/viscosity_error.png")

vmax = 41
vmin = -39
figure();visualize_scalar_on_fvm_points(S_true[end, 2*(m+1)*(n+1)+1:end], m, n, h, vmax=vmax, vmin=vmin); tight_layout()
savefig("sess_figures/ex1/ones/pressure_data.png")

figure();visualize_scalar_on_fvm_points(run(sess, S[end,2*(m+1)*(n+1)+1:end]), m, n, h, vmax=vmax, vmin=vmin); tight_layout()
savefig("sess_figures/ex1/ones/pressure_prediction.png")

figure();visualize_scalar_on_fvm_points(S_true[end, 2*(m+1)*(n+1)+1:end] - run(sess, S[end,2*(m+1)*(n+1)+1:end]), m, n, h); tight_layout()
savefig("sess_figures/ex1/ones/pressure_error.png")

figure();visualize_scalar_on_fvm_points_nocontour(S_true[end, 2*(m+1)*(n+1)+1:end] - run(sess, S[end,2*(m+1)*(n+1)+1:end]), m, n, h); tight_layout()
savefig("sess_figures/ex1/ones/pressure_error_nocontour.png")

function visualize_scalar_on_fvm_points_nocontour(φ::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)
    φ = reshape(φ, m, n)'|>Array
    m_ = mean(φ)
    s = std(φ)
    vmin, vmax = m_ - 2s, m_ + 2s
    x = (1:m)*h .- 0.5h
    y = (1:n)*h .- 0.5h
    ln = pcolormesh(x, y, φ, vmin= vmin, vmax=vmax)
    colorbar()
    # c = contour(φ[1,:,:], 10, cmap="jet", vmin=vmin,vmax=vmax)
    axis("scaled")
    xlabel("x")
    ylabel("y")
    ln = gca().pcolormesh(x, y, φ, vmin= vmin, vmax=vmax)
    # c = gca().contour(x, y, φ, 10, cmap="jet", vmin=vmin,vmax=vmax)
    gca().invert_yaxis()
end


using InvertibleNetworks, LinearAlgebra, Test, Random

# Data
nx = 16
ny = 16
n_in = 3
batchsize = 4
X = randn(Float32, nx, ny, n_in, batchsize)

# Network
nscales = 2
steps_per_scale = 4
nc = 2
hf = 32
α0 = .2f0
H = NetworkHyperbolic(nx, ny, n_in, batchsize, nscales, steps_per_scale; α=α0, hidden_factor=hf, ncenter=nc)

###################################################################################################
# Test invertibility

X_ = H.inverse(H.forward(X)[1])
@show norm(X - X_), norm(X)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1e-2)

X_ = H.forward(H.inverse(X))[1]
@show norm(X - X_), norm(X)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1e-2)

####################################################################################################
# Training

# Loss
function loss(H, X, ind)
    Y, logdet = H.forward(X)
    f = -log_likelihood(Y) - logdet
    ΔY = -∇log_likelihood(Y)
    ΔX, X = H.backward(ΔY, Y)
    return f, ΔX, H.HL[ind].W.grad, H.AL.s.grad
end

# Data
X = randn(Float32, nx, ny, n_in, batchsize)
X0 = randn(Float32, nx, ny, n_in, batchsize)
dX = X - X0

# Gradient test w.r.t. input X0
Y = H.forward(X)
f0, ΔX = loss(H, X0, 1)[1:2]
h = 0.1f0
maxiter = 6
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test hyperbolic network\n")
for j=1:maxiter
    f = loss(H, X0 + h*dX, 1)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

@show rate_1 = sum(err1[1:end-1]./err1[2:end])/(maxiter - 1)
@show rate_2 = sum(err2[1:end-1]./err2[2:end])/(maxiter - 1)

# Gradient test w.r.t. weights of hyperbolic network
H = NetworkHyperbolic(nx, ny, n_in, batchsize, nscales, steps_per_scale; α=α0, hidden_factor=hf, ncenter=nc)
H0 = NetworkHyperbolic(nx, ny, n_in, batchsize, nscales, steps_per_scale; α=α0, hidden_factor=hf, ncenter=nc)
H.forward(X)
H0.forward(X)   # evaluate to initialize actnorm layer

s0 = H0.AL.s.data
ds = H.AL.s.data - H0.AL.s.data

maxiter = 6
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

for i=1:length(H0.HL)
    h1 = 0.1f0
    print("\nGradient test invertible layer for layer $(i)\n")
    W0s = H0.HL[i].W.data
    dWg = H.HL[i].W.data - W0s
    f01, ΔX1, ΔWg, _ = loss(H0, X, i)
    for j=1:maxiter
        H0.HL[i].W.data = W0s + h1*dWg
        f = loss(H0, X, i)[1]
        err3[j] = abs(f - f01)
        err4[j] = abs(f - f01 - h1*dot(dWg, ΔWg))
        print(err3[j], "; ", err4[j], "\n")
        h1 = h1/2f0
    end

    @show local rate_1 = sum(err3[1:end-1]./err3[2:end])/(maxiter - 1)
    @show local rate_2 = sum(err4[1:end-1]./err4[2:end])/(maxiter - 1)
    H0.HL[i].W.data = W0s
    @test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=2f1)
    @test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=2f1)
end

print("\nGradient test invertible layer for affine layer\n")
f0, ΔX, ΔW, Δs = loss(H0, X, 1)
for j=1:maxiter
    H0.AL.s.data = s0 + h*ds
    f = loss(H0, X, 1)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(ds, Δs))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@show rate_1 = sum(err3[1:end-1]./err3[2:end])/(maxiter - 1)
@show rate_2 = sum(err4[1:end-1]./err4[2:end])/(maxiter - 1)

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=2f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=2f1)


function sumsqrtn_fixed(n)
r = 0.0
for i = 1:n
r = r + sqrt(i)
end
return r
end


using BenchmarkTools
using LinearAlgebra
using StaticArrays

add!(C, A, B) = (C .= A .+ B)

function simple_bench(N, T=Float64)
    A = rand(T,N,N)
    A = A'*A
    B = copy(A)
    SA = SMatrix{N,N}(A)
    MA = MMatrix{N,N}(A)
    MB = copy(MA)

    print("""
============================================
    Benchmarks for $N×$N $T matrices
============================================
""")
    ops = [
           ("Matrix multiplication              ", *,        (A, A),     (SA, SA)),
           ("Matrix multiplication (mutating)   ", mul!,     (B, A, A),  (MB, MA, MA)),
           ("Matrix addition                    ", +,        (A, A),     (SA, SA)),
           ("Matrix addition (mutating)         ", add!,     (B, A, A),  (MB, MA, MA)),
           ("Matrix determinant                 ", det,      (A,),       (SA,)),
           ("Matrix inverse                     ", inv,      (A,),       (SA,)),
           ("Matrix symmetric eigendecomposition", eigen,    (A,),       (SA,)),
           ("Matrix Cholesky decomposition      ", cholesky, (A,),       (SA,)),
           ("Matrix LU decomposition            ", lu,       (A,),       (SA,)),
           ("Matrix QR decomposition            ", qr,       (A,),       (SA,)),
    ]
    for (name, op, Aargs, SAargs) in ops
        # We load from Ref's here to avoid the compiler completely removing the
        # benchmark in some cases.
        #
        # Like any microbenchmark, the speedups you see here should only be
        # taken as roughly indicative of the speedup you may see in real code.
        if length(Aargs) == 1
            A1  = Ref(Aargs[1])
            SA1 = Ref(SAargs[1])
            speedup = @belapsed($op($A1[])) / @belapsed($op($SA1[]))
        elseif length(Aargs) == 2
            A1  = Ref(Aargs[1])
            A2  = Ref(Aargs[2])
            SA1 = Ref(SAargs[1])
            SA2 = Ref(SAargs[2])
            speedup = @belapsed($op($A1[], $A2[])) / @belapsed($op($SA1[], $SA2[]))
        elseif length(Aargs) == 3
            A1  = Ref(Aargs[1])
            A2  = Ref(Aargs[2])
            A3  = Ref(Aargs[3])
            SA1 = Ref(SAargs[1])
            SA2 = Ref(SAargs[2])
            SA3 = Ref(SAargs[3])
            speedup = @belapsed($op($A1[], $A2[], $A3[])) / @belapsed($op($SA1[], $SA2[], $SA3[]))
        else
        end
        println(name*" -> $(round(speedup, digits=1))x speedup")
    end
end

simple_bench(3)

simple_bench(256)

N=50
T=Float64

A = rand(T,N,N)
A = A'*A
B = copy(A)
SA = SMatrix{N,N}(A)
MA = MMatrix{N,N}(A)
MB = copy(MA)

@benchmark SMatrix{N,N}(A)
@benchmark rand(T,N,N)

function random_access(data::Vector{UInt}, N::Integer)
    n = rand(UInt)
    mask = length(data) - 1
    @inbounds for i in 1:N
        n = (n >>> 7) ⊻ data[n & mask + 1]
    end
    return n
end

function linear_access(data::Vector{UInt}, N::Integer)
    n = rand(UInt)
    mask = length(data) - 1
    for i in 1:N
        n = (n >>> 7) ⊻ data[i & mask + 1]
    end
    return n
end


import DoubleFloats: Double64, ComplexDF64

@testset "LinearAlgebra" begin
    @testset "Constructor + basics" begin
        A = rand(ComplexF32, 12, 12)
        WS = HC.MatrixWorkspace(A)
        @test WS isa HC.MatrixWorkspace{Float64}
        @test_throws ArgumentError HC.update!(WS, rand(12, 13))

        WS.factorized[] = true
        B = randn(ComplexF64, 12, 12)
        HC.update!(WS, B)
        @test WS.A == B
        @test WS.A !== B
        @test WS.factorized[] == false

        HC.factorization!(WS, HC.LU_FACT)
        @test WS.fact[] == HC.LU_FACT
        WS.factorized[] = true
        HC.factorization!(WS, HC.QR_FACT)
        @test WS.fact[] == HC.QR_FACT
        @test WS.factorized[] == false

        HC.factorization!(WS, HC.LU_FACT)
        @test (@allocated HC.factorize!(WS)) == 0
        HC.factorization!(WS, HC.QR_FACT)
        @test WS.factorized[] == false
        HC.factorize!(WS) # first factorize resizes the work buffers
        WS.factorized[] = false
        @test (@allocated HC.factorize!(WS)) == 0

# test indexing etc
        @test WS[2, 1] == B[2, 1]
        @test WS[2] == B[2]
        WS[2, 3] = 5.0
        @test WS[2, 3] == 5.0
        WS[5] = 2.32
        @test WS[5] == 2.32
        @test size(WS) == (12, 12)
        C = randn(ComplexF64, 12, 12)
        WS .= C
        @test WS.A == C
    end

    @testset "ldiv" begin
        A = randn(ComplexF64, 13, 13)
        b = randn(ComplexF64, 13)
        x = zeros(ComplexF64, 13)
        WS = HC.MatrixWorkspace(A)
        HC.update!(WS, A)
        ldiv!(x, WS, b)
        WS.factorized[] = false
        @test (@allocated ldiv!(x, WS, b)) == 0
        @test (lu(A) \ b) ≈ x rtol = 1e-12

        HC.factorization!(WS, HC.QR_FACT)
        x .= 0
        ldiv!(x, WS, b)
        @test (qr(A, Val(true)) \ b) ≈ x rtol = 1e-12
        WS.factorized[] = false
        @test (@allocated ldiv!(x, WS, b)) == 0

# overdetermined
        A = randn(ComplexF64, 13, 8)
        x = randn(ComplexF64, 8)
        b = A * x
        WS = HC.MatrixWorkspace(A)
        HC.update!(WS, A)
        @test_throws ArgumentError HC.factorization!(WS, HC.LU_FACT)
        ldiv!(x, WS, b)
        WS.factorized[] = false
        x .= 0
        @test (@allocated ldiv!(x, WS, b)) == 0
        @test (qr(A, Val(true)) \ b) ≈ x rtol = 1e-12
    end

    @testset "(Mixed precision) residual" begin
        A = randn(ComplexF64, 13, 13)
        x = randn(ComplexF64, 13)
        b = A * x
        û = convert.(ComplexF64, ComplexDF64.(A) * ComplexDF64.(x) - ComplexDF64.(b))

        u = zeros(ComplexF64, 13)
# fixed precision residual should be of order unit roundoff
        HC.residual!(u, A, x, b)
        @test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) < 1e-14
        @test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) > 1e-18

        HC.residual!(u, A, x, b, Complex{Double64})
        @test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) < 1e-26
    end

    @testset "(Mixed precision) iterative refinement" begin
        A = randn(ComplexF64, 4, 4)
        b = randn(ComplexF64, 4)
        x̂ = zeros(ComplexF64, 4)
        x = copy(x̂)
        WS = HC.MatrixWorkspace(A)
        U, S, VT = svd(A)
        S[end] *= 1e-4
        A2 = U * diagm(0 => S) * VT'
        HC.update!(WS, A2)
        ldiv!(x̂, WS, b)
        δx1 = HC.iterative_refinement_step!(x, WS, x̂, b, HC.InfNorm(), ComplexDF64)
        @test δx1 / norm(x, Inf) > 1e-14
        δx2 = HC.iterative_refinement_step!(WS, x, b, HC.InfNorm(), ComplexDF64)
        @test δx2 / norm(x, Inf) < eps()

        δx = HC.iterative_refinement_step!(x, WS, x̂, b)
        @test δx1 / norm(x, Inf) > 1e-14

# overdetermined
        A = randn(ComplexF64, 4, 3)
        x̂ = zeros(ComplexF64, 3)
        x = copy(x̂)
        WS = HC.MatrixWorkspace(A)
        U, S, VT = svd(A)
        S[end] *= 1e-4
        A2 = U * diagm(0 => S) * VT'
        HC.update!(WS, A2)
        b = A2 * randn(ComplexF64, 3)
        ldiv!(x̂, WS, b)
        δx1 = HC.iterative_refinement_step!(x, WS, x̂, b, HC.InfNorm(), ComplexDF64)
        @test δx1 / norm(x, Inf) > 1e-14
        δx2 = HC.iterative_refinement_step!(WS, x, b, HC.InfNorm(), ComplexDF64)
        @test δx2 / norm(x, Inf) < eps()

        δx = HC.iterative_refinement_step!(x, WS, x̂, b)
        @test δx1 / norm(x, Inf) > 1e-14
    end

    @testset "rcond" begin
        A = randn(ComplexF64, 13, 13)
        WS = HC.MatrixWorkspace(A)
        U, S, VT = svd(A)
        S[end] *= 1e-5
        A2 = U * diagm(0 => S) * VT'
        HC.update!(WS, A2)
        HC.factorization!(WS, HC.LU_FACT)
        HC.factorize!(WS)
        @test cond(A2, Inf) ≈ inv(HC.rcond(WS))
        HC.factorization!(WS, HC.QR_FACT)
        @test 0.5 ≤ cond(A2, Inf) / inv(HC.rcond(WS)) ≤ 1.5
# @test (@allocated rcond(WS)) == 0

        A = randn(ComplexF64, 13, 10)
        WS = HC.MatrixWorkspace(A)
        U, S, VT = svd(A)
        S[end] *= 1e-8
        A2 = U * diagm(0 => S) * VT'
        HC.update!(WS, A2)
        HC.factorize!(WS)
        @test 1 / HC.rcond(WS) ≈ cond(UpperTriangular(WS.qr.R), Inf) rtol = 1e-14
# @test (@allocated rcond(WS)) == 0

# meddle aroudn with the internals to check dead branch
        WS.fact[] = HC.LU_FACT
        @test isnan(HC.rcond(WS))
    end

    @testset "inf norm est" begin
        A = randn(ComplexF64, 5, 5)
        U, s, Vt = svd(A)
        s[end] *= 1e-6
        B = U * diagm(0 => s) * Vt'

        d = rand(5) * 1000
        D = diagm(0 => d)
        g = rand(5) .* 1000
        G = diagm(0 => g)

        @test 0.1 ≤ opnorm(inv(B), Inf) / HC.inf_norm_est(lu(B)) ≤ 10
        @test 0.1 ≤ opnorm(inv(D) * inv(B), Inf) / HC.inf_norm_est(lu(B), nothing, d) ≤ 10
        @test 0.1 ≤ opnorm(inv(B) * G, Inf) / HC.inf_norm_est(lu(B), g, nothing) ≤ 10
        @test 0.1 ≤ opnorm(inv(D) * inv(B) * G, Inf) / HC.inf_norm_est(lu(B), g, d) ≤ 10

        WS = HC.MatrixWorkspace(B)
        @test 0.1 ≤ opnorm(inv(B), Inf) / HC.inf_norm_est(WS) ≤ 10
        @test 0.1 ≤ opnorm(inv(D) * inv(B), Inf) / HC.inf_norm_est(WS, nothing, d) ≤ 10
        @test 0.1 ≤ opnorm(inv(B) * G, Inf) / HC.inf_norm_est(WS, g, nothing) ≤ 10
        @test 0.1 ≤ opnorm(inv(D) * inv(B) * G, Inf) / HC.inf_norm_est(WS, g, d) ≤ 10
    end

    @testset "JacobianMonitor" begin
        A = randn(ComplexF64, 6, 6)
        b = randn(ComplexF64, 6)
        x̂ = similar(b)
        x = similar(b)
        JM = HC.JacobianMonitor(zeros(ComplexF64, 6, 6))
        @test JM isa HC.JacobianMonitor
        HC.jacobian(JM) .= A
        HC.updated!(JM)
        ldiv!(x̂, JM, b)
        ldiv!(x, jacobian(JM), b)
        @test x == x̂
    end
    @testset "Forward error" begin
        A = [
            1.81451 + 0.46473im -0.473651 + 0.813117im 0.130822 - 0.254704im -0.203277 +
                                                                             0.397106im
            -2.04468 - 1.59228im 0.441157 - 1.39994im -0.0984804 + 1.42619im -0.818292 -
                                                                             0.487898im
            0.35501 + 0.0235574im 1.32618 + 0.577264im 1.12561 + 0.43383im 0.579533 +
                                                                           1.49484im
            0.144381 + 0.104175im 0.665384 - 1.43389im 0.22097 + 0.258016im 0.374103 -
                                                                            0.635066im
        ]
        b = [
            -0.14581 - 0.868847im,
            0.0162454 + 0.91156im,
            -0.133769 + 0.510834im,
            -0.867022 + 0.752196im,
        ]
        x̂ = similar(b)
        x = similar(b)
        JM = HC.JacobianMonitor(A)
        HC.updated!(JM)
        ldiv!(x̂, JM, b)

        ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())
        @test HC.forward_err(JM) == ferr
        @test ferr < eps() * cond(A)
        @test ferr > 1e-12
        ferrD64 = HC.forward_err!(x, JM, x̂, b, HC.InfNorm(), ComplexDF64)
        @test 0.5 ≤ ferrD64 / ferr ≤ 2

        ferr = HC.forward_err!(JM, x̂, b, HC.InfNorm(), ComplexF64)
        @test ferr < eps() * cond(A)
        @test ferr > 1e-12

        D = diagm(0 => [10.0^(2i) for i = 1:4])
        JM = HC.JacobianMonitor(D * A)
        HC.updated!(JM)
        ldiv!(x̂, JM, D * b)
        conderr = HC.cond(D * A) * eps()
        sferr = HC.strong_forward_err!(x, JM, x̂, D * b, HC.InfNorm())
        @test HC.forward_err(JM) == sferr
        ferr = HC.forward_err!(x, JM, x̂, D * b, HC.InfNorm())
        @test ferr < sferr < conderr * 1e-4 # check that error estimate is actually substantially less

        sferr2 = HC.strong_forward_err!(JM, x̂, D * b, HC.InfNorm())
        @test sferr == sferr2
        @test x̂ ≈ x

        DR = diagm(0 => [10.0^(2i - 4) for i = 1:4])
        JM = HC.JacobianMonitor(A * DR)
        ldiv!(x̂, JM, b)

        WN = HC.WeightedNorm(HC.InfNorm(), x̂)
        WN .= abs.(x̂)

        wferr = HC.forward_err!(x, JM, x̂, b, WN)
        ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())

        strong_wferr = HC.strong_forward_err!(x, JM, x̂, b, WN)
        strong_ferr = HC.strong_forward_err!(x, JM, x̂, b, HC.InfNorm())
        @test ferr ≤ strong_ferr ≤ eps() * inv(HC.rcond(jacobian(JM)))
        @test wferr ≤ strong_wferr ≤ eps() * inv(HC.rcond(jacobian(JM)))

# overdetermined
        A = randn(ComplexF64, 3, 2)
        x = randn(ComplexF64, 2)
        b = A * x
        x̂ = similar(x)

        JM = HC.JacobianMonitor(A)
        ldiv!(x̂, JM, b)

        ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())
        strong_ferr = HC.strong_forward_err!(x, JM, x̂, b, HC.InfNorm())
        @test ferr < 1e-13
        @test ferr == strong_ferr # test fallback to ferr for qr
    end

    @testset "cond" begin
        A = randn(ComplexF64, 13, 13)
        D = diagm(0 => [2.0^(i - 5) for i = 1:13])
        JM = HC.JacobianMonitor(D * A)
        comp_cond = HC.cond!(JM)
        @test comp_cond === cond(JM)
        B = diagm(0 => inv.(JM.J.row_scaling)) * D * A
        @test abs(cond(JM) - cond(B, Inf)) / cond(JM) < 2

        A = randn(ComplexF64, 13, 12)
        JM = HC.JacobianMonitor(D * A)
        comp_cond = HC.cond!(JM)
        @test comp_cond === cond(JM)
    end

    @testset "JacobianMonitor ldiv updates" begin
        A = randn(ComplexF64, 6, 6)
        b = randn(ComplexF64, 6)
        x̂ = similar(b)
        x = similar(b)
        JM = HC.JacobianMonitor(zeros(ComplexF64, 6, 6))
        @test JM isa HC.JacobianMonitor
        ldiv!(x, JM, b, HC.InfNorm(), HC.JAC_MONITOR_UPDATE_FERR)
        @test HC.forward_err(JM) != 0.0
        ldiv!(x, JM, b, HC.InfNorm(), HC.JAC_MONITOR_UPDATE_COND)
        @test cond(JM) != 1.0
        JM = HC.JacobianMonitor(zeros(ComplexF64, 6, 6))
        ldiv!(x, JM, b, HC.InfNorm(), HC.JAC_MONITOR_UPDATE_ALL)
        @test cond(JM) != 1.0
        @test HC.forward_err(JM) != 0.0
    end

    @testset "Hermite Normal Form" begin
        A = [0 3 1 0; -2 2 -1 2; 1 -1 2 3; -3 3 3 2]
        H, U = HC.hnf(A)
        @test A * U == H

        D = zeros(Int, 1, 1)
        D[1, 1] = -2
        H, U = HC.hnf(D)
        @test H[1, 1] == 2
        @test U[1, 1] == -1

# overflow A
        A = [
            2 4 4 2 -4 -3 -4 3 0 3
            -4 -2 -5 -5 1 0 -3 0 -1 -2
            -3 5 -3 1 4 3 -1 -4 -1 0
            1 3 3 4 -2 -3 -2 -5 -4 -3
            2 2 0 2 -4 4 3 -4 -2 -4
            -3 -3 1 1 -4 4 0 4 4 -4
            -3 -1 -1 -1 2 -4 -3 -4 4 -4
            0 1 -2 4 5 4 3 1 -5 2
            -5 -4 -5 3 1 5 0 -3 -3 -1
            2 -5 -3 -1 -1 5 -4 5 -3 1
        ]
        @test_throws OverflowError HC.hnf(A)
        @test_throws OverflowError HC.hnf(A, Int128)
        H, U = HC.hnf(A, BigInt)
        @test A * U == H
    end
end


using FaultTolerantControl
using Test
using LinearAlgebra
using Transducers

@testset "costs" begin
    Q_p = Matrix(I, 3, 3)
    Q_ω = Matrix(I, 3, 3)
    F_p = Matrix(I, 3, 3)
    F_ω = Matrix(I, 3, 3)
    cf = PositionAngularVelocityCostFunctional(Q_p, Q_ω, F_p, F_ω)
    # data
    t0, tf = 0.0, 1.0
    Δt = tf - t0
    ts = t0:0.01:tf
    e_p_nom = ones(3)
    e_ω_nom = ones(3)
    e_ps = ts |> Map(t -> e_p_nom) |> collect
    e_ωs = ts |> Map(t -> e_ω_nom) |> collect
    J = cost(cf, ts, e_ps, e_ωs)
    @test J ≈ e_p_nom'*Q_p*e_p_nom*Δt + e_ω_nom'*Q_ω*e_ω_nom + e_ps[end]'*F_p*e_ps[end] + e_ωs[end]'*F_ω*e_ωs[end]
end


function test_hinge_loss_layer(backend::Backend, T, eps)
  println("-- Testing HingeLossLayer on $(typeof(backend)){$T}...")

  ############################################################
  # Prepare Data for Testing
  ############################################################
  tensor_dim = abs(rand(Int)) % 4 + 2
  dims = tuple((abs(rand(Int,tensor_dim)) % 6 + 6)...)
  println("    > $dims")
  preds = rand(T, dims)*4-2
  labels = round(rand(T, dims))*2-1
  errs_mask = preds.*labels .< one(T)

  ############################################################
  # Setup
  ############################################################
  layer  = HingeLossLayer(; bottoms=[:predictions, :labels])
  pred_blob  = make_blob(backend, T, dims)
  label_blob = make_blob(backend, T, dims)
  diff_blob  = make_blob(backend, T, dims)
#   diff_blob2 = make_blob(backend, T, dims)
  copy!(pred_blob, preds)
  copy!(label_blob, labels)
  inputs = Blob[pred_blob, label_blob]
  diffs  = Blob[diff_blob, NullBlob()]
#   diffs2 = Blob[diff_blob, diff_blob2]

  state  = setup(backend, layer, inputs, diffs)

  forward(backend, state, inputs)

  loss = sum(max(one(T) .- preds.*labels, zero(T))) / dims[end]
  @test -eps < loss-state.loss < eps

  backward(backend, state, inputs, diffs)
  grad  = -labels .* errs_mask / dims[end]
#   grad2 = -preds  .* errs_mask / dims[end]
  diff  = similar(grad)
  copy!(diff, diffs[1])

  @test all(-eps .< grad - diff .< eps)

  # check accumulated/reset statistics
  forward(backend, state, Blob[label_blob, label_blob])
  @test 0 == state.loss
  @test -eps < 0.5loss-state.loss_accum < eps
  reset_statistics(state)
  forward(backend, state, inputs)
  @test -eps < loss-state.loss_accum < eps

  shutdown(backend, state)
end

function test_hinge_loss_layer(backend::Backend)
  test_hinge_loss_layer(backend, Float32, 1e-2)
  test_hinge_loss_layer(backend, Float64, 1e-8)
end

if test_cpu
  test_hinge_loss_layer(backend_cpu)
end
if test_gpu
  test_hinge_loss_layer(backend_gpu)
end


"""
    ngcd(p,q, [k]; kwargs...)

Find numerical GCD of polynomials `p` and `q`. Refer to [`NGCD.ngcd`](@ref) for details.

In the case `degree(p) ≫ degree(q)`,  a heuristic is employed to first call one step of the Euclidean gcd approach, and then call `ngcd` with relaxed tolerances.

"""
function ngcd(p::P, q::Q,
              args...; kwargs...) where {T,X,P<:StandardBasisPolynomial{T,X},
                                         S,Y,Q<:StandardBasisPolynomial{S,Y}}

    if (degree(q) > degree(p))
        u,w,v,Θ,κ =  ngcd(q,p,args...;kwargs...)
        return (u=u,v=v,w=w, Θ=Θ, κ=κ)
    end
    if degree(p) > 5*(1+degree(q))
        a,b = divrem(p,q)
        return ngcd(q,b, args...; λ=100, kwargs...)
    end

    # easy cases
    degree(p) < 0  && return (u=q,      v=p, w=one(q),  θ=NaN, κ=NaN)
    degree(p) == 0 && return (u=one(q), v=p, w=q,       θ=NaN, κ=NaN)
    degree(q) < 0  && return (u=one(q), v=p, w=zero(q), θ=NaN, κ=NaN)
    degree(q) == 0 && return (u=one(p), v=p, w=q,       θ=NaN, κ=NaN)
    p ≈ q          && return (u=p,v=one(p),  w=one(p),  θ=NaN, κ=NaN)
    Polynomials.assert_same_variable(p,q)

    R = promote_type(float(T), float(S))
    𝑷 = Polynomials.constructorof(promote_type(P,Q)){R,X}

    ps = R[pᵢ for pᵢ ∈ coeffs(p)]
    qs = R[qᵢ for qᵢ ∈ coeffs(q)]

    # cancel zeros
    nz = min(findfirst(!iszero, ps), findfirst(!iszero, qs))
    if nz == length(qs)
        x = variable(p)
        u = x^(nz-1)
        v,w = 𝑷(ps[nz:end]), 𝑷(qs[nz:end])
        return (u=u, v=v, w=w, Θ=NaN, κ=NaN)
    end

    ## call ngcd
    p′ = PnPolynomial{R,X}(ps[nz:end])
    q′ = PnPolynomial{R,X}(qs[nz:end])
    out = NGCD.ngcd(p′, q′, args...; kwargs...)

    𝑷 = Polynomials.constructorof(promote_type(P,Q)){R,X}
    u,v,w = convert.(𝑷, (out.u,out.v,out.w))
    if nz > 1
        u *= variable(u)^(nz-1)
    end
    (u=u,v=v,w=w, Θ=out.Θ, κ = out.κ)

end

## ---- the work is done in this module

module NGCD
using Polynomials, LinearAlgebra
import Polynomials: PnPolynomial, constructorof

"""
    ngcd(ps::PnPolynomial{T,X}, qs::PnPolynomial{T,X}, [k::Int]; scale::Bool=false, atol=eps(T), rtol=eps(T), satol=atol, srtol=rtol)

Return `u, v, w, Θ, κ` where ``u⋅v ≈ ps`` and ``u⋅w ≈ qs`` (polynomial multiplication); `Θ` (`\\Theta[tab]`) is the residual error (``‖ [u⋅v,u⋅w] - [ps,qs] ‖``); and `κ` (`\\kappa[tab]`) is the numerical gcd condition number estimate. When `scale=true`, ``u⋅v ≈ ps/‖ps‖`` and ``u⋅w ≈ qs/‖qs‖``

The numerical GCD problem is defined in [1] (5.4). Let ``(p,q)`` be a
polynomial pair with degree m,n. Let Ρmn be set of all such pairs. Any
given pair of polynomials has an exact greatest common divisor, ``u``, of
degree ``k``, defined up to constant factors. Let ``Ρᵏmn`` be the manifold of
all such ``(p,q)`` pairs with exact gcd of degree ``k``. A given pair ``(p,q)`` with exact gcd of degree ``j`` will
have some distance ``Θᵏ`` from ``Pᵏ``.  For a given threshold ``ϵ > 0`` a numerical GCD
of ``(p,q)`` within ``ϵ`` is an exact GCD of a pair ``(p̂,q̂)`` in ``Ρᵏ`` with

``‖ (p,q) - (p̂,q̂) ‖ <= Θᵏ``, where ``k`` is the largest value for which ``Θᵏ < ϵ``.

(In the ``ϵ → 0`` limit, this would be the exact GCD.)

Suppose ``(p,q)`` is an ``ϵ`` pertubation from ``(p̂,q̂)`` where ``(p̂,q̂)`` has an exact gcd of degree ``k``, then ``degree(gcdₑ(p,q)) = k``; as ``ϵ → 0``, ``gcdₑ(p,q) → gcd(p̂, q̂)``, and

``limsup_{(p,q)→(p̂,q̂)} inf{ ‖ (u,v,w) - (û,v̂,ŵ) ‖} / ‖ (p,q) - (p̂,q̂) ‖ < κₑ(p,q)``.

``κ`` is called the numerical GCD condition number.

The Zeng algorithm proposes a degree for ``u`` and *if* a triple ``(u,v,w)`` with ``u`` of degree ``k`` and ``(u⋅v, u⋅w)`` in ``Ρᵏmn`` can be found satisfying ``‖ (u⋅v, u⋅w) - (p,q) ‖ < ϵ`` then ``(u,v,w)`` is returned; otherwise the proposed degree is reduced and the process repeats. If not terminated, at degree ``0`` a constant gcd is returned.

The initial proposed degree is the first ``j``,  ``j=min(m,n):-1:1``, where ``Sⱼ`` is believed to have a singular value of ``0`` (``Sⱼ`` being related to the Sylvester matrix of `ps` and `qs`). The verification of the proposed degree is done using a Gauss-Newton iteration scheme holding the degree of ``u`` constant.

## Scaling:

If `scale=true` the gcd of ``p/‖p‖`` and ``q/‖q‖`` is identified. Scaling can reduce the condition numbers significantly.

## Tolerances:

There are two places where tolerances are utilized:

* in the identification of the rank of ``Sⱼ`` a value ``σ₁ = ‖Rx‖`` is identified. To test if this is zero a tolerance of `max(satol, ‖R‖ₒₒ ⋅ srtol)` is used.

* to test if ``(u ⋅ v, u ⋅ w) ≈ (p,q)`` a tolerance of `max(atol, λ⋅rtol)` is used with `λ` chosen to be  ``(‖(p,q)‖⋅n⋅m)⋅κ′⋅‖A‖ₒₒ`` to reflect the scale of ``p`` and ``q`` and an estimate for the condition number of ``A`` (a Jacobian involved in the solution).

This seems to work well for a reasaonable range of polynomials, however there can be issues: when the degree of ``p`` is much larger than the degree of ``q``, these choices can fail; when a higher rank is proposed, then too large a tolerance for `rtol` or `atol` can lead to a false verification; when a tolerance for `atol` or `rtol` is too strict, then a degree may not be verified.

!!! note:
    These tolerances are adjusted from those proposed in [1].

## Specified degree:

When `k` is specified, a value for ``(u,v,w)`` is identified with ``degree(u)=k``. No tolerances are utilized in computing ``Θᵏ``.

Output:

The function outputs a named tuple with names (`u`, `v`, `w`, `Θ`, `κ`). The components `u`,`v`,`w` estimate the gcd and give the divisors. The value `Θ` estimates ``Θᵏ`` and `κ` estimates the numerical condition number.

Example:

```
using Polynomials
x = variable(Polynomial{Float64})
p = (x+10)*(x^9 + x^8/3 + 1)
q = (x+10)*(x^9 + x^8/7 - 6/7)
gcd(p,q) # u a constant
gcd(p,q, method=:numerical)  # u a degree 1 polynomial
Polynomials.NGCD.ngcd(coeffs(p), coeffs(q), verbose=true) # to see some computations
```

Reference:

[1] The Numerical Greatest Common Divisor of Univariate Polynomials
by Zhonggang Zeng;
[url](http://homepages.neiu.edu/~zzeng/uvgcd.pdf);
[doi](https://doi.org/10.1090/conm/556/11014)

Note: Based on work by Andreas Varga; Requires `VERSION >= v"1.2"`.

"""
function ngcd(p::PnPolynomial{T,X},
              q::PnPolynomial{T,X};
              scale::Bool=false,
              atol = eps(real(T)),
              rtol = Base.rtoldefault(real(T)),
              satol = eps(real(T))^(5/6),
              srtol = eps(real(T)),
              verbose=false,
              minⱼ = -1,
              λ = 1
              ) where {T, X}

    m,n = length(p)-1, length(q)-1
    @assert m >= n

    ## --- begin
    if scale
        p ./= norm(p)
        q ./= norm(q)
    end
    atol *= λ
    rtol *= λ

    # storage
    A0 = zeros(T, m+1, 2)
    A0[:,1] = coeffs(p)
    A0[end-n:end,2] = coeffs(q)

    # pre-allocate storage for Sylvester Matrices, S₁, S₂...
    Q = zeros(T, m + n, m + n)
    R = zeros(T, m + n, m + n)
    Sₓ = hcat(convmtx(p,1),  convmtx(q, m-n+1))

    uv = copy(p)
    uw = copy(q)

    local x::Vector{T}

    F = qr(Sₓ)
    nr, nc = size(Sₓ) # m+1, m-n+2
    Q[1:nr, 1:nr] .= F.Q
    R[1:nc, 1:nc] .= F.R

    j = n  # We count down Sn, S_{n-1}, ..., S₂, S₁

    while true

        V = view(R, 1:nc, 1:nc)
        flag, σ, x = smallest_singular_value(V, satol *  sqrt(1 + m - j), srtol)
        verbose && println("------ degree $j ----- σ₁: $σ  --- $flag")

        if (flag == :iszero || flag == :ispossible)
            u, v, w = initial_uvw(Val(flag), j, p, q, x)
            flag, ρ₁, σ₂, ρ = refine_uvw!(u,v,w, p, q, uv, uw, atol, rtol)

            verbose && println("   --- Θᵏ: $ρ₁ --- $flag (ρ=$(ρ))")

            if flag == :convergence
                return (u=u, v=v, w=w, Θ=ρ₁, κ=σ₂) # (u,v,w) verified
            end
        end

        # reduce possible degree of u and try again with Sⱼ₋₁
        # unless we hit specified minimum, in which case return it
        if j == minⱼ
            u, v, w = initial_uvw(Val(:ispossible), j, p, q, x)
            flag, ρ₁, σ₂, ρ = refine_uvw!(u,v,w, p, q, uv, uw, atol, rtol)
            return (u=u, v=v, w=w, Θ=ρ₁, κ=σ₂)
        end

        j -= 1
        nr += 1
        nc += 2
        nc > nr && break
        extend_QR!(Q,R, nr, nc, A0) # before Q⋅R = Sⱼ, now Q⋅R = Sⱼ₋₁

    end

    # u is a constant
    verbose && println("------ GCD is constant ------")

    u, v, w = initial_uvw(Val(:constant), j, p, q, x)
    flag, ρ₁, κ, ρ = refine_uvw!(u,v,w, p, q, uv, uw, atol, rtol)
    return (u=u, v=v, w=w, Θ=ρ₁, κ=κ)

end

# fix the degree, k
function ngcd(p::P,
              q::P,
              k::Int;
              kwargs...
              ) where {T <: AbstractFloat,X, P <: PnPolynomial{T,X}}

    m::Int, n::Int = length(p)-1, length(q)-1

    #    u,v,w = initial_uvw(Val(:iszero), k, ps, qs, nothing)
    Sⱼ = [convmtx(p, n-k+1) convmtx(q, m-k+1)]
    F = qr(Sⱼ)
    flag, σ, x = smallest_singular_value(F.R, eps(T) *  sqrt(1 + m - k), eps(T))
    u,v,w = initial_uvw(Val(:k), flag, k, p, q, x)
    flag, ρ₁, κ, ρ = refine_uvw!(u,v,w, copy(p), copy(q), copy(p), copy(q),
                                 T(Inf), T(Inf)) # no tolerances
    return (u=u, v=v, w=w, Θ=ρ₁, κ=κ)

end

## -----

# return guess at smallest singular value and right sinuglar value, x
# for an upper triangular matrix, V
function smallest_singular_value(V::AbstractArray{T,2},
                                 atol=eps(real(T)),
                                 rtol=zero(real(T))) where {T}

    R = UpperTriangular(V)
    k = size(R)[1]/2
    if iszero(det(R))
        return (:iszero, zero(T), T[])
    end

    m,n = size(R)

    # we are testing if ‖Ax‖ ≈ 0
    # If x is a perfect 0, but x is approximated by x' due to round off
    # then ‖A(x-x')‖ <= ‖A‖⋅‖x - x'‖ so we use ‖A‖ as scale factor
    δ = max(atol,  norm(R,Inf) * rtol)

    x = ones(T, n)
    y = zeros(T, m)
    σ₀ = σ₁ = Inf*one(real(T))
    steps, minᵢ = 1, 5

    while true
        y .= R' \ x # use iteration to converge on singular value
        x .= R  \ y
        x ./= norm(x,2)
        σ₁ = norm(R * x, 2)

        if (steps <= 50) && (steps <= minᵢ || σ₁ < 0.05*σ₀) # decreasing, keep going
            σ₀ = σ₁
        else
            break
        end
        steps += 1
    end

    if σ₁ < δ
        return (:ispossible, σ₁, x)
    else
        return (:constant, σ₁, x)
    end

end

## --------------------------------------------------
## Refine u,v,w

## Find u₀,v₀,w₀ from right singular vector
function initial_uvw(::Val{:ispossible}, j, p::P, q::Q, x) where {T,X,
                                                              P<:PnPolynomial{T,X},
                                                              Q<:PnPolynomial{T,X}}

    # Sk*[w;-v] = 0, so pick out v,w after applying permutation
    m,n = length(p)-1, length(q)-1
    vᵢ = vcat(2:m-n+2, m-n+4:2:length(x))
    wᵢ = m-n+3 > length(x) ? [1] : vcat(1, (m-n+3):2:length(x))
    #    v = 𝑷{m-j}(-x[vᵢ])
    v = P(-x[vᵢ])
    w = P(x[wᵢ])
    # p194 3.9 C_k(v) u = p or Ck(w) u = q; this uses 10.2
    u = solve_u(v,w,p,q,j)
    return u,v,w

end

function initial_uvw(::Val{:iszero}, j, p::P, q::Q, x) where {T,X,
                                                              P<:PnPolynomial{T,X},
                                                              Q<:PnPolynomial{T,X}}

    m,n = length(p)-1, length(q)-1
    S = [convmtx(p, n-j+1) convmtx(q, m-j+1)]

    F = qr(S)
    R = UpperTriangular(F.R)

    if iszero(det(R))
        x = eigvals(R)[:,1]
    else
        x = ones(T, size(R,2))
        x .= R \ (R' \ (x/norm(x)))
        x ./= norm(x)
    end

    w = P(x[1:n-j+1])
    v = P(-x[(n-j+2):end])

    u = solve_u(v,w,p,q,j)
    return u,v,w
end

function initial_uvw(::Val{:constant}, j, p::P, q, x) where {T,X,P<:PnPolynomial{T,X}}
    u = one(P)
    w = q
    v = p
    u,v,w
end

function initial_uvw(::Val{:k}, flag, k, p::P, q, x) where {T,X,P<:PnPolynomial{T,X}}
    flag == :iszero && return initial_uvw(Val(flag), k, p, q, nothing)
    n = length(q)-1
    w, v = P(x[1:n-k+1]), P(-x[n-k+2:end])
    u = solve_u(v,w,p,q,k)
    return (u,v,w)
end

# find estimate for σ₂, used in a condition number (κ = 1/σ)
function σ₂(J)
    F = qr(J)
    flag, σ, x = smallest_singular_value(F.R)
    σ
end

## attempt to refine u,v,w
## check that [u * v; u * w] ≈ [p; q]
function refine_uvw!(u::U, v::V, w::W, p, q, uv, uw, atol, rtol) where {T,X,
                                                                        U<:PnPolynomial{T,X},
                                                                        V<:PnPolynomial{T,X},
                                                                        W<:PnPolynomial{T,X}}

    m,n,l = length(u)-1, length(v)-1, length(w)-1

    mul!(uv, u, v)
    mul!(uw, u, w)

    ρ₀, ρ₁ = one(T), residual_error(p,q,uv,uw)

    # storage
    b = zeros(T, (m+n) + (m+l) + 3) # degree(p) + degree(q) + 3 = 1 + length(p) + length(q))
    Δf = zeros(T, m + n + l + 3)
    steps = 0

    h, β =  u, norm(u)^2
    minᵢ, Maxᵢ = 5, 20
    κ = NaN
    A=zeros(T, JF_size(u, v, w)...)
    JF!(A, h, u, v, w)
    Fmp!(b,  dot(h,u) - β, p, q, uv, uw)

    Δvᵢ = 1:(n+1)
    Δwᵢ = (n+1+1):(n+1+l+1)
    Δuᵢ = (n+1+l+1+1):length(Δf)

    while ρ₁ > 0.0

        # update A, b, then solve A\b
        qrsolve!(Δf, A, b)

        # m + n = degree(p)
        # m + l = degree(q)
        # b has length degree(p)+degree(q) + 3
        Δv = view(Δf, Δvᵢ)
        Δw = view(Δf, Δwᵢ)
        Δu = view(Δf, Δuᵢ)

        u .-= Δu
        v .-= Δv
        w .-= Δw

        mul!(uv, u, v)
        mul!(uv, u, w)

        ρ₀, ρ′ = ρ₁, residual_error(p, q, uv, uw)

        # don't worry about first few, but aftewards each step must be productive
        # though we can have really bad first steps, which we cap
        if  (steps <= Maxᵢ) && (steps <= minᵢ || ρ′ < 0.95 * ρ₀) && (  ρ′ < 100*ρ₁ )
            ρ₁ = ρ′
            steps += 1
        else
            break
        end

        # update A,b for next iteration
        JF!(A, h, u, v, w)
        Fmp!(b,  dot(h,u) - β, p, q, uv, uw)

    end

    # this is a heuristic
    # sensitivity is Δu / Δp <= ‖ A+ ‖ = κ
    # we use an estimate for ‖(p,q)‖ error times ‖A⁺‖⋅‖A‖ₒₒ
    κ = 1/σ₂(A) # ≈ ‖A⁺‖
    λ = norm((norm(p), norm(q))) * (m * n) * min(1, κ) * norm(A, Inf)
    ρ = max(atol, rtol * λ)

    if ρ₁ <= ρ
        return :convergence, ρ₁, κ, ρ
    else
        return :no_convergence, ρ₁, κ, ρ
    end

end

## ---- QR factorization

function qrsolve!(y::Vector{T}, A, b) where {T}
    y .= qr(A) \ b
end

# # Fast least-squares solver for full column rank Hessenberg-like matrices
# # By Andreas Varga
function qrsolve!(y::Vector{T}, A, b) where {T <: Float64}
    Base.require_one_based_indexing(A)
    m, n = size(A)
    m < n && error("Column dimension exceeds row dimension")
    _, τ = LinearAlgebra.LAPACK.geqrf!(A)
    T <: Complex ? tran = 'C' : tran = 'T'
    LinearAlgebra.LAPACK.ormqr!('L', tran, A, τ, view(b,:,1:1))
    y .= UpperTriangular(triu(A[1:n,:]))\b[1:n]
end

# extend QR to next size
# Q gets a 1 in nc,nc, 0s should be elswhere
function extend_QR!(Q,R, nr, nc, A0)

    #old Q is m x m, old R is n x n; we add to these
    n = nc-2
    m = nr - 1
    k,l = size(A0)

    # add two columns to R
    # need to apply Q to top part of new columns
    R[nr-k+1:nr, (nc-1):nc] = A0
    R[1:nr-1, (nc-1):nc] = (view(Q, 1:nr-1, 1:nr-1))' *  R[1:nr-1, (nc-1):nc]

    # extend Q with row and column with identity
    Q[nr,nr] = 1

    # Make R upper triangular using Givens rotations
    for j in nr-1:-1:nc-1
        gj,_ = givens(R[j,nc-1], R[j+1,nc-1], j, j+1)
        rmul!(Q, gj')
        lmul!(gj, R)
    end

    for j in nr-1:-1:nc
        gj,_ = givens(R[j,nc], R[j+1,nc], j, j+1)
        rmul!(Q, gj')
        lmul!(gj, R)
    end

    return nothing

end

## Jacobian F(u,v,w) = [p,p'] is J(u,v,w)
function JF_size(u, v, w)

    m, k, j = length(u)-1, length(v)-1, length(w)-1
    n, l = m + k, m + j

    ai, aj = convmtx_size(v, m + 1)
    bi, bj = convmtx_size(u, k + 1)
    di, dj = convmtx_size(w, m + 1)
    fi, fj = convmtx_size(u, j + 1)
    ci, cj = ai, fj
    ei, ej = di, bj

    (1 + ai + di, aj + bj + cj)
end

# Jacobian matrix
function JF(u::Vector{U}, v::Vector{V}, w::Vector{W}) where {U,V,W}
    R = promote_type(U,V, W)
    M = zeros(R, JF_size(u, v, w)...)
    JF!(M, u, v, w)
    M
end

function JF!(M, h,  u::P, v, w) where {T,X,P<:AbstractPolynomial{T,X}}

    du, dv, dw = length(u)-1, length(v)-1, length(w)-1
    m, n = du + dv, du + dw

    # JF_size should return these
    r11,c11 = convmtx_size(u, dv+1)
    r13,c13 = convmtx_size(v, du+1)
    r22,c22 = convmtx_size(u, dw+1)
    r23,c23 = convmtx_size(w, du+1)

    J11 = view(M, 1:r11, 1:c11)
    J13 = view(M, 1:r13, c11 + c22 .+ (1:c23))
    J22 = view(M, r11 .+ (1:r22), c11 .+ (1:c22))
    J23 = view(M, r13 .+ (1:r23), (c11 + c22) .+ (1:c23))
    convmtx!(J11, u, dv+1)
    convmtx!(J13, v, du+1)
    convmtx!(J22, u, dw+1)
    convmtx!(J23, w, du+1)
    M[end, end-du:end] = coeffs(h)'

    return nothing
end

## compute F(u,v,w) - [p, p'] = [u*v, u*w] - [p, p']
function Fmp!(b, γ, p, q, uv, uw)
    b[end] = γ
    for i in 1:1+length(p)-1
        j = i
        b[i] = uv[j] - p[j]
    end
    for i in 1+length(p):length(b)-1
        j = i - length(p)
        b[i] = uw[j] - q[j]
    end
    return nothing
end

function residual_error(p::P,q,uv,uw) where {T,X,P<:AbstractPolynomial{T,X}}
    tot = zero(real(T))
    for (pᵢ, uvᵢ) in zip(p,uv)
        tot += norm(pᵢ-uvᵢ)^2
    end
    for (qᵢ, uwᵢ) in zip(q, uw)
        tot += norm(qᵢ-uwᵢ)^2
    end
    sqrt(tot)
end

## ---- utils

"""
    convmtx(v, n::Int)
    convmtx!(M,v,n)
    convmtx_size(v,n)

Convolution matrix.
C = convmtx(v,n) returns the convolution matrix C for a vector v.
If q is a column vector of length n, then C*q is the same as conv(v,q).

"""
function convmtx!(C, v::AbstractVector{T}, n::Int) where T

    #   Form C as the Toeplitz matrix
    #   C = Toeplitz([v; zeros(n-1)],[v[1]; zeros(n-1));  put Toeplitz code inline

    nv = length(v)-1

    @inbounds for j = 1:n
        C[j:j+nv,j] = v
    end

    nothing

end
convmtx_size(v::AbstractVector, n) = (n + length(v) - 1, n)
function convmtx(v::AbstractVector{T}, n::Int) where {T}
    C = zeros(T, convmtx_size(v, n)...)
    convmtx!(C, v, n)
    C
end

# multroot uses vector/matrix interface.
convmtx!(C, v::AbstractPolynomial, n::Int) = convmtx!(C, coeffs(v), n)
convmtx_size(v::AbstractPolynomial, n) = (n + length(v)-1, n)
function convmtx(v::AbstractPolynomial{T}, n::Int) where {T}
    d = length(v)-1
    C = zeros(T, (n + d, n))
    convmtx!(C, v, n)
    C
end

# solve for u from [v,w] \ [p,q]
function solve_u(v::P,w,p,q, k) where {T,X,P<:PnPolynomial{T,X}}
    A = [convmtx(v,k+1); convmtx(w, k+1)]
    b = vcat(coeffs(p), coeffs(q))
    u = P(A \ b)
    return u
end

end


# Sample postprocessing script for plotting a tensor from one snapshot.
#
# Hongyang Zhou, hyzhou@umich.edu

using Vlasiator, PyPlot, LaTeXStrings, Printf

file = "/home/hongyang/Vlasiator/test/equatorial/timevarying/bulk.0001582.vlsv"

axisunit = EARTH
colorscale = Linear
cmap = matplotlib.cm.turbo

vars = ["proton/vg_ptensor_diagonal", "proton/vg_ptensor_offdiagonal"]

#################

meta = load(file)

pArgs1 = Vlasiator.set_args(meta, vars[1], axisunit; normal=:none)
pArgs2 = Vlasiator.set_args(meta, vars[2], axisunit; normal=:none)

x, y = Vlasiator.get_axis(pArgs1)
# [nPa]
pxx = Vlasiator.prep2d(meta, vars[1], :x)' .* 1e9
pyy = Vlasiator.prep2d(meta, vars[1], :y)' .* 1e9
pzz = Vlasiator.prep2d(meta, vars[1], :z)' .* 1e9
pyz = Vlasiator.prep2d(meta, vars[2], :1)' .* 1e9
pxz = Vlasiator.prep2d(meta, vars[2], :2)' .* 1e9
pxy = Vlasiator.prep2d(meta, vars[2], :3)' .* 1e9

P = (pxx, pyy, pzz, pyz, pxz, pxy)
P_str = ("Pxx", "Pyy", "Pzz", "Pyz", "Pxz", "Pxy")

vmin = minimum(minimum.(P))
vmax = maximum(maximum.(P))

cnorm1, cticks1 = Vlasiator.set_colorbar(colorscale, vmin, vmax)

fig, axs = subplots(3,2,
   figsize=(6, 12), sharex=true, sharey=true, constrained_layout=true)

c1 = axs[1].pcolormesh(x, y, pxx; norm=cnorm1, cmap)

for i in eachindex(axs)[2:end]
   axs[i].pcolormesh(x, y, P[i]; norm=cnorm1, cmap)
end

for (i, ax) in enumerate(axs)
   ax.set_aspect("equal")
   ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
   ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
   ax.set_title(P_str[i])
end

for ax in axs[3,:]
   ax.set_xlabel(L"x [$R_E$]")
end

for ax in axs[:,1]
   ax.set_ylabel(L"y [$R_E$]")
end

# One colorbar for all subplots
cb1 = colorbar(c1; ax=axs, ticks=cticks1, fraction=0.046, pad=0.04)
cb1.ax.set_ylabel("[nPa]")
cb1.outline.set_linewidth(1.0)

str_title = @sprintf "Pressure Tensor at t = %4.1fs" meta.time

fig.suptitle(str_title, fontsize="xx-large")

savefig("ptensor.png", bbox_inches="tight")


using LinearAlgebra

include("byrow.jl")
include("dare.jl")
include("dlyap.jl")
include("simula01.jl")


# Exports
export RK4, istep

#######
# RK4 #
#######

"""
RK4 is a numerical integrator object which stores a function of signature
f(epc, x,...; kwargs...) along with its associated settings parameters.

Arguments:
- `f::Function` Function for integration.
- `auxParams::Dict`
"""
struct RK4
    f::Function
    auxParams
end

function RK4(f::Function; kwargs...)
    return RK4(f, kwargs)
end

"""
Internal funciton used to compute the variational matrix.

Arguments:
- `rk4::RK4` 4-th order Runge-Kutta numerical integrator object
- `epc::Union{Real, Epoch}` Absolute time of start of integration step
- `x::Array{<:Real, 1}` State vector for linearizing perturbations around
- `apert::Union{Real, Array{<:Real,1}}` Absolute perturbations to state elements
used to numerically calculate the partial derivatives through a forward difference
method.

Returns:
- `V:Array{Float64, 2}` Variational equation drivative matrix. This is the
matrix of partial derivatives of the dynamics function _f_ with respect to the
input state _x_.
"""
function varmat(rk4::RK4, epc::Union{Real,Epoch}, x::Array{<:Real, 1}; apert=1.0::Union{Real, Array{<:Real,1}})
    # Get state length
    n = length(x)

    # Expand perturbation step-size as vector
    if typeof(apert) <: Real
        apert = ones(typeof(apert), n)
    end

    # Initialize variational matrix
    V = zeros(Float64, n, n)

    # Evaluate unperturbed state derivative function
    fx = rk4.f(epc, x; rk4.auxParams...)

    # Compute perturbed acceleration for each field
    for i in 1:n
        # Perturbed state vector
        px     = copy(x)
        px[i] += apert[i]

        # Perturbed state derivative
        pfx = rk4.f(epc, px; rk4.auxParams...)

        # Compute partial of state derivative with respect to state using
        # forward difference method
        dfdx = (pfx - fx)/apert[i]

        # Set i-th column as state partial
        V[:, i] = dfdx
    end

    return V
end

function istep(rk4::RK4, epc::Union{Real,Epoch}, dt::Real, x::Array{<:Real, 1})
    # Compute State coefficients

    # Perform internal steps
    k1 = rk4.f(epc,                      x; rk4.auxParams...)
    k2 = rk4.f(epc + dt/2.0, x+(dt/2.0)*k1; rk4.auxParams...)
    k3 = rk4.f(epc + dt/2.0, x+(dt/2.0)*k2; rk4.auxParams...)
    k4 = rk4.f(epc + dt,           x+dt*k3; rk4.auxParams...)

    # Compute updated state
    xu = x + (dt/6.0)*(k1 + 2.0*k2 + 2.0*k3 + k4)

    return xu
end

"""
Performs single integration step of numerical integrator.

Arguments:
- `rk4::RK4` 4-th order Runge-Kutta numerical integrator object
- `epc::Union{Real, Epoch}` Absolute time of start of integration step
- `dt::Real` Integration step size
- `x::Array{<:Real, 1}` State vector
- `phi::Array{<:Real, 2}` State transition matrix at start of integration step
"""
function istep(rk4::RK4, epc::Union{Real,Epoch}, dt::Real, x::Array{<:Real, 1}, phi::Array{<:Real, 2}; apert=1.0::Union{Real, Array{<:Real,1}})
    # Compute State coefficients

    # Perform state integration
    xk1 = rk4.f(epc,                         x; rk4.auxParams...)
    xk2 = rk4.f(epc + dt/2.0, x + (dt/2.0)*xk1; rk4.auxParams...)
    xk3 = rk4.f(epc + dt/2.0, x + (dt/2.0)*xk2; rk4.auxParams...)
    xk4 = rk4.f(epc + dt,           x + dt*xk3; rk4.auxParams...)

    # Perform Integration of variational equations
    phik1 = varmat(rk4, epc, x, apert=apert)*phi
    phik2 = varmat(rk4, epc + dt/2.0, x + dt/2.0*xk1, apert=apert)*(phi + (dt/2.0)*phik1)
    phik3 = varmat(rk4, epc + dt/2.0, x + dt/2.0*xk2, apert=apert)*(phi + (dt/2.0)*phik2)
    phik4 = varmat(rk4,     epc + dt,     x + dt*xk3, apert=apert)*(phi + dt*phik3)

    # Compute updated state
    xu   = x   + (dt/6.0)*(xk1   + 2.0*xk2   + 2.0*xk3   + xk4)
    phiu = phi + (dt/6.0)*(phik1 + 2.0*phik2 + 2.0*phik3 + phik4)

    return xu, phiu
end


using MathPhysicalConstants, Measurements, Unitful

using Test

using MathPhysicalConstants

@testset begin
    @test MathPhysicalConstants.MKS.PlancksConstantH == 6.62606896e-34
    @test iszero(measurement(MathPhysicalConstants.IS.SanchezElectrConstant) - measurement(MathPhysicalConstants.IS.SanchezElectrConstant))
    # \\\ @test setprecision(BigFloat, 768) do; precision(ustrip(big(c))) end == 768
    #\\\@test measurement(h) === measurement(h)
    #\\\@test iszero(measurement(α) - measurement(α))
    #\\\@test isone(measurement(BigFloat, atm) / measurement(BigFloat, atm))
    #\\\@test iszero(measurement(BigFloat, ħ) - (measurement(BigFloat, h) / 2big(pi)))
    #\\\@test isone(measurement(BigFloat, ħ) / (measurement(BigFloat, h) / 2big(pi)))
end


using SparsityOperators
using Test
using Random
using LinearAlgebra
using FFTW

include("testOperators.jl")


#export impute_means, impute_zeros,
#      #Closed-form factorizations with regularization
#      matrixRegFact, quadgraphRegFact

# Subtracts each column's mean (if center=true),
# divides by each column's standard deviation (if scale=true).
# Returns (scaledData, mean, std), where mean or std may be
#  empty matrices if center or scale are false
function standardize!{T}(X::Matrix{T} ; center=true, scale=true)
    n = size(X,1)

    local m
    if center
        m = mean(X,1)
        X = X .- m
    else
        m = similar(X, 0, 0)
    end

    local s
    if scale
        s = std(X,1)
        #Make sure there are no NaN elements or 0 elements of the standard deviation
        for i in 1:length(s)
          if s[i] == 0 || s[i] === NaN
            s[i] = 1
          end
        end
        X = X ./ s
    else
        s = similar(X, 0, 0)
    end

    X, m, s
end

##Subtracts each column mean (if center=true)
##Divides each column standard deviation (if scale=true)
##Skips NA entries in DataArrays
##Returns (scaledData, mean, std)
function standardize!{T}(X::DataMatrix{T}; center=true, scale=true)
  n = size(X,1)

  local m
  if center
    m = mean(X,1,skipna=true)
    X = X .- m
  else
    m = similar(X,0,0)
  end

  local s
  if scale
    s = sqrt(var(X,1,skipna=true))
    #Make sure there are no NaN elements or 0 elements of the standard deviation
    for i in 1:length(s)
      if s[i] == 0 || s[i] === NaN
        s[i] = 1
      end
    end
    X = X ./ s
  else
    s = similar(X,0,0)
  end

  X, m, s
end

standardize(X; center=true, scale=true) = standardize!(copy(X), center=center, scale=scale)

##Imputes the column means to a DataMatrix and returns an ordinary Matrix
##of the same size.
function impute_means{T}(X::DataMatrix{T})
  m, n = size(X)
  colmeans = mean(X,1,skipna=true)
  Xmat = zeros(m,n)
  for j in 1:n
    #For entirely NA columns, just use zero
    if all(isna(X[:,j]))
      Xmat[:,j] = 0
    else
      #Impute either the mean or the existing value
      for i in 1:m
        if X[i,j] === NA
          Xmat[i,j] = colmeans[j]
        else
          Xmat[i,j] = X[i,j]
        end
      end
    end
  end
  Xmat
end

##Impute zeros on every NA value of a DataMatrix
function impute_zeros{T}(X::DataMatrix{T})
  m, n = size(X)
  Xmat = zeros(m,n)
  for j in 1:n
    for i in 1:m
      if !(X[i,j] === NA)
        Xmat[i,j] = X[i,j]
      end
    end
  end
  Xmat
end

##Closed-form solution to a quadratic loss function with graph+quadratic
##regularization on both X and Y. Performance may be poor for large number of rows
function matrixRegFact(A::AbstractMatrix, k::Int,
  α::Number, K::AbstractMatrix,
  β::Number, L::AbstractMatrix)

  n, d = size(A)
  @assert (n,n) == size(K)
  @assert (d,d) == size(L)
  iKA = (I + α*K) \ A # (I + αK)⁻¹ * A
  iKAt = iKA' # A'*(I + αK)⁻¹
  fst = -2*A'*iKA
  snd = iKAt*iKA
  thd = α*iKAt*K*iKA
  res = β*L
  for i in 1:length(res)
    res[i] += fst[i] + snd[i] + thd[i] #Add the results in a loop to save space
  end
  #=
  invIαK = inv(I + α*K)
  #Three terms to find the eigenvectors of, save time by adding in a loop
  fst = -2*A'*invIαK*A
  snd = A'*invIαK^2*A
  thd = α*A'*invIαK*K*invIαK*A
  res = β*L
  for i in 1:length(res)
    res[i] += fst[i] + snd[i] + thd[i]
  end
  =#
  Yt = eigvecs(res)[:,1:k]
  X = invIαK*A*Yt
  X,Yt'
end

##Closed-form solution to quadratic loss function with quadratic regularization
##on X and a graph+quadratic regularizer on Y.
function quadgraphRegFact(A::Matrix{Float64}, k::Int, α::Number,
  β::Number, L::AbstractMatrix)

  n, d = size(A)
  @assert (d,d) == size(L)

  ALmat = -1/(1+α)*A'*A + β*L
  Yt = eigvecs(ALmat)[:,1:k]
  X = 1/(1 + α)*A*Yt
  X,Yt'
end

"""
Initialize the GLRM with an SVD decomposition for faster convergence.
This also solves the problem where identical quadratic regularization is on
both X and Y
"""
function init_qqreg!(gg::GGLRM)
  A = impute_zeros(gg.A)
  F = svdfact(A)
  A_mul_B!(gg.X, F[:U][:,1:gg.k], Diagonal(F[:S][1:gg.k]))
  A_mul_Bt!(gg.Y, Diagonal(F[:S][1:gg.k]), F[:Vt][:,1:gg.k])
end


function vec_isless(x::AbstractArray, y::AbstractArray)
    for (xi,yi) in zip(x,y)
        if xi >= yi return false end
    end
    return true
end

mutable struct SurvivorApproximator
    sample::MatF64
    α::Float64
end

function (surv::SurvivorApproximator)(x::AbstractVector{T}) where T<:Real
    count=0
    N = size(surv.sample, 2)
    # count = @parallel (+) for i in 1:N
    #     convert(Int, vec_isless(x, view(surv.sample, :, i)))
    # end
    for i in 1:N
        count += convert(Int, vec_isless(x, view(surv.sample, :, i)))
    end
    prob=count/N
    se=sqrt((prob-prob^2)/N)
    err=norminvcdf((1 + surv.α)/2)*se
    return prob,err
end

function below_nonrisk_frontier(x::AbstractVector{T}, frontier::LinkedList{VecF64}) where T <: Real
    for p in frontier
        if vec_isless(x,p) return true end
    end
    return false
end

function above_risk_frontier(x::AbstractVector{T}, frontier::LinkedList{VecF64}) where T <: Real
    for p in frontier
        if vec_isless(p, x) return true end
    end
    return false
end

function decapitate!(l::Cons)
    l.head=l.tail.head
    l.tail=l.tail.tail
    return l
end


function test_function_differentiation(T=Float64)
    grid = RegularRectilinearGrid(T; size=(3, 3, 3), extent=(3, 3, 3))
    ϕ = rand(T, 3, 3, 3)
    ϕ² = ϕ.^2

    ∂x_ϕ_f = ϕ²[2, 2, 2] - ϕ²[1, 2, 2]
    ∂x_ϕ_c = ϕ²[3, 2, 2] - ϕ²[2, 2, 2]

    ∂y_ϕ_f = ϕ²[2, 2, 2] - ϕ²[2, 1, 2]
    ∂y_ϕ_c = ϕ²[2, 3, 2] - ϕ²[2, 2, 2]

    ∂z_ϕ_f = ϕ²[2, 2, 2] - ϕ²[2, 2, 1]
    ∂z_ϕ_c = ϕ²[2, 2, 3] - ϕ²[2, 2, 2]

    f(i, j, k, grid, ϕ) = ϕ[i, j, k]^2

    return (
        ∂xᶜᵃᵃ(2, 2, 2, grid, f, ϕ) == ∂x_ϕ_c &&
        ∂xᶠᵃᵃ(2, 2, 2, grid, f, ϕ) == ∂x_ϕ_f &&

        ∂yᵃᶜᵃ(2, 2, 2, grid, f, ϕ) == ∂y_ϕ_c &&
        ∂yᵃᶠᵃ(2, 2, 2, grid, f, ϕ) == ∂y_ϕ_f &&

        ∂zᵃᵃᶜ(2, 2, 2, grid, f, ϕ) == ∂z_ϕ_c &&
        ∂zᵃᵃᶠ(2, 2, 2, grid, f, ϕ) == ∂z_ϕ_f
    )
end

function test_function_interpolation(T=Float64)
    grid = RegularRectilinearGrid(T; size=(3, 3, 3), extent=(3, 3, 3))
    ϕ = rand(T, 3, 3, 3)
    ϕ² = ϕ.^2

    ℑx_ϕ_f = (ϕ²[2, 2, 2] + ϕ²[1, 2, 2]) / 2
    ℑx_ϕ_c = (ϕ²[3, 2, 2] + ϕ²[2, 2, 2]) / 2

    ℑy_ϕ_f = (ϕ²[2, 2, 2] + ϕ²[2, 1, 2]) / 2
    ℑy_ϕ_c = (ϕ²[2, 3, 2] + ϕ²[2, 2, 2]) / 2

    ℑz_ϕ_f = (ϕ²[2, 2, 2] + ϕ²[2, 2, 1]) / 2
    ℑz_ϕ_c = (ϕ²[2, 2, 3] + ϕ²[2, 2, 2]) / 2

    f(i, j, k, grid, ϕ) = ϕ[i, j, k]^2

    return (
        ℑxᶜᵃᵃ(2, 2, 2, grid, f, ϕ) == ℑx_ϕ_c &&
        ℑxᶠᵃᵃ(2, 2, 2, grid, f, ϕ) == ℑx_ϕ_f &&

        ℑyᵃᶜᵃ(2, 2, 2, grid, f, ϕ) == ℑy_ϕ_c &&
        ℑyᵃᶠᵃ(2, 2, 2, grid, f, ϕ) == ℑy_ϕ_f &&

        ℑzᵃᵃᶜ(2, 2, 2, grid, f, ϕ) == ℑz_ϕ_c &&
        ℑzᵃᵃᶠ(2, 2, 2, grid, f, ϕ) == ℑz_ϕ_f
    )
end

@testset "Operators" begin
    @info "Testing operators..."

    @testset "Grid lengths, areas, and volume operators" begin
        @info "  Testing grid lengths, areas, and volume operators..."

        FT = Float64
        grid = RegularRectilinearGrid(FT, size=(1, 1, 1), extent=(π, 2π, 3π))

        @testset "Easterly lengths" begin
            @info "    Testing easterly lengths..."
            for δ in (Δx, Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ)
                @test δ(1, 1, 1, grid) == FT(π)
            end
        end

        @testset "Westerly lengths" begin
            @info "    Testing westerly lengths..."
            for δ in (Δy, Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ)
                @test δ(1, 1, 1, grid) == FT(2π)
            end
        end

        @testset "Vertical lengths" begin
            @info "    Testing vertical lengths..."
            for δ in (ΔzF, ΔzC)
                @test δ(1, 1, 1, grid) == FT(3π)
            end
        end

        @testset "East-normal areas in the yz-plane" begin
            @info "    Testing areas with easterly normal in the yz-plane..."
            for A in (Axᵃᵃᶜ, Axᵃᵃᶠ, Axᶠᶜᶜ)
                @test A(1, 1, 1, grid) == FT(6 * π^2)
            end
        end

        @testset "West-normal areas in the xz-plane" begin
            @info "    Testing areas with westerly normal in the xz-plane..."
            for A in (Ayᵃᵃᶜ, Ayᵃᵃᶠ, Ayᶜᶠᶜ)
                @test A(1, 1, 1, grid) == FT(3 * π^2)
            end
        end

        @testset "Horizontal areas in the xy-plane" begin
            @info "    Testing horizontal areas in the xy-plane..."
            for A in (Azᵃᵃᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ)
                @test A(1, 1, 1, grid) == FT(2 * π^2)
            end
        end

        @testset "Volumes" begin
            @info "    Testing volumes..."
            for V in (Vᵃᵃᶜ, Vᵃᵃᶠ, Vᶜᶜᶜ)
                @test V(1, 1, 1, grid) == FT(6 * π^3)
            end
        end

    end

    @testset "Function differentiation" begin
        @info "  Testing function differentiation..."
        @test test_function_differentiation()
    end

    @testset "Function interpolation" begin
        @info "  Testing function interpolation..."
        @test test_function_interpolation()
    end

    @testset "2D operators" begin
        @info "  Testing 2D operators..."

        Nx, Ny, Nz = 32, 16, 8
        Lx, Ly, Lz = 100, 100, 100

        arch = CPU()
        grid = RegularRectilinearGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
        bcs = FieldBoundaryConditions(grid, (Center, Center, Center))

        Hx, Hy, Hz = grid.Hx, grid.Hy, grid.Hz
        Tx, Ty, Tz = Nx+2Hx, Ny+2Hy, Nz+2Hz

        A3 = OffsetArray(zeros(Tx, Ty, Tz), 1-Hx:Nx+Hx, 1-Hy:Ny+Hy, 1-Hz:Nz+Hz)
        @. @views A3[1:Nx, 1:Ny, 1:Nz] = rand()
        fill_halo_regions!(A3, bcs, arch, grid)

        # A yz-slice with Nx==1.
        A2yz = OffsetArray(zeros(1+2Hx, Ty, Tz), 1-Hx:1+Hx, 1-Hy:Ny+Hy, 1-Hz:Nz+Hz)
        grid_yz = RegularRectilinearGrid(size=(1, Ny, Nz), extent=(Lx, Ly, Lz))

        # Manually fill in halos for the slice.
        A2yz[0:2, 0:Ny+1, 1:Nz] .= A3[1:1, 0:Ny+1, 1:Nz]
        A2yz[:, :, 0] .= A2yz[:, :, 1]
        A2yz[:, :, Nz+1] .= A2yz[:, :, Nz]

        # An xz-slice with Ny==1.
        A2xz = OffsetArray(zeros(Tx, 1+2Hy, Tz), 1-Hx:Nx+Hx, 1-Hy:1+Hy, 1-Hz:Nz+Hz)
        grid_xz = RegularRectilinearGrid(size=(Nx, 1, Nz), extent=(Lx, Ly, Lz))

        # Manually fill in halos for the slice.
        A2xz[0:Nx+1, 0:2, 1:Nz] .= A3[0:Nx+1, 1:1, 1:Nz]
        A2xz[:, :, 0] .= A2xz[:, :, 1]
        A2xz[:, :, Nz+1] .= A2xz[:, :, Nz]

        test_indices_3d = [(4, 5, 5), (21, 11, 4), (16, 8, 4),  (30, 12, 3), (11, 3, 6), # Interior
                           (2, 10, 4), (31, 5, 6), (10, 2, 4), (17, 15, 5), (17, 10, 2), (23, 5, 7),  # Borderlands
                           (1, 5, 5), (32, 10, 3), (16, 1, 4), (16, 16, 4), (16, 8, 1), (16, 8, 8),  # Edges
                           (1, 1, 1), (32, 16, 8)] # Corners

        test_indices_2d_yz = [(1, 1, 1), (1, 1, 2), (1, 2, 1), (1, 2, 2),
                              (1, 1, 5), (1, 5, 1), (1, 5, 5), (1, 11, 4),
                              (1, 15, 7), (1, 15, 8), (1, 16, 7), (1, 16, 8)]

        test_indices_2d_xz = [(1, 1, 1), (1, 1, 2), (2, 1, 1), (2, 1, 2),
                              (1, 1, 5), (5, 1, 1), (5, 1, 5), (17, 1, 4),
                              (31, 1, 7), (31, 1, 8), (32, 1, 7), (32, 1, 8)]

        for idx in test_indices_2d_yz
            @test δxᶜᵃᵃ(idx..., grid_yz, A2yz) ≈ 0
            @test δxᶠᵃᵃ(idx..., grid_yz, A2yz) ≈ 0
            @test δyᵃᶜᵃ(idx..., grid_yz, A2yz) ≈ δyᵃᶜᵃ(idx..., grid_yz, A3)
            @test δyᵃᶠᵃ(idx..., grid_yz, A2yz) ≈ δyᵃᶠᵃ(idx..., grid_yz, A3)
            @test δzᵃᵃᶜ(idx..., grid_yz, A2yz) ≈ δzᵃᵃᶜ(idx..., grid_yz, A3)
            @test δzᵃᵃᶠ(idx..., grid_yz, A2yz) ≈ δzᵃᵃᶠ(idx..., grid_yz, A3)
        end

        for idx in test_indices_2d_xz
            @test δxᶜᵃᵃ(idx..., grid_xz, A2xz) ≈ δxᶜᵃᵃ(idx..., grid_xz, A3)
            @test δxᶠᵃᵃ(idx..., grid_xz, A2xz) ≈ δxᶠᵃᵃ(idx..., grid_xz, A3)
            @test δyᵃᶜᵃ(idx..., grid_xz, A2xz) ≈ 0
            @test δyᵃᶠᵃ(idx..., grid_xz, A2xz) ≈ 0
            @test δzᵃᵃᶜ(idx..., grid_xz, A2xz) ≈ δzᵃᵃᶜ(idx..., grid_xz, A3)
            @test δzᵃᵃᶠ(idx..., grid_xz, A2xz) ≈ δzᵃᵃᶠ(idx..., grid_xz, A3)
        end
    end
end


c = one(PermCycs)
p = one(PermList)
s = one(PermSparse)
m = one(PermMat)

@test null(c) == null(p) == null(s)

for el in ( :c, :p, :s, :m)
    @eval begin
        @test isid($el)
        @test issym($el)
        @test isposdef($el)
        @test iseven($el)
    end
end

##

p = PermList([3,1,2])
c = cycles(p)
s = psparse(p)
m = matrix(p)

for el in ( :c, :p, :s, :m)
    @eval begin
        @test ! isid($el)
        @test ! issym($el)
        @test ! isposdef($el)
        @test inv($el) == transpose($el)
        @test rank($el) == plength($el)
        @test det($el) == sign($el)
        @test iseven($el)
    end
end

p = PermList([2,7,8,3,1,9,6,4,5,10])
@test pivtopermlist(topiv(p)) == p

q = PermList([5,6,9,1,4,7,8,3,2,10])
q1 = pivtopermlist([5,6,9,5,5,7,8,9,9,10])
@test q1 == q


import Chemfiles
using CellListMap
using StaticArrays
using Test

const ε = 0.0441795
const σ = 2*1.64009

function lj_NE(d2,u)
    d = sqrt(d2)
    u += ε*( (σ/d)^12 - 2*(σ/d)^6 )
end

function getcoor(file)
    traj = redirect_stdout(() -> Chemfiles.Trajectory(file), devnull)
    frame = Chemfiles.read_step(traj,0)
    Chemfiles.close(traj)
    return reinterpret(reshape,SVector{3,Float64},Chemfiles.positions(frame))
end

function test_newcl(file,unit_cell,correct,lcell)
    coordinates = getcoor(file)
    box = Box(unit_cell, 10., lcell=lcell)
    cl = CellList(coordinates,box)
    u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
    return u ≈ correct
end

@testset "NAMD" begin

  dir="./namd"

  lcell=1

  # Some orthorhombic cells

  unit_cell = [ 50.,   50.,  50. ]
  correct = 32230.01699504111
  @test test_newcl("$dir/o1.dcd", unit_cell, correct, lcell)

  unit_cell = [ 80.,  70.,  50. ]
  correct = 1093.7225407797744
  @test test_newcl("$dir/o2.dcd", unit_cell, correct, lcell)

  # Orthorhombic but rotated

  unit_cell = [ 50.      0.     50.
                50.     50.      0.
                 0.     50.     50. ]
  correct = 1724.3195067566828
  @test test_newcl("$dir/o3.dcd", unit_cell, correct, lcell)

  unit_cell = transpose([ 70.7107   0.0      0.0
                         35.3553  61.2372   0.0
                         35.3553  20.4124  57.735 ])
  correct = 1754.0802503953591
  @test test_newcl("$dir/o4.dcd", unit_cell, correct, lcell)

  unit_cell= transpose([ 70.7107   0.0      0.0
                         35.3553  61.2372   0.0
                         35.3553  20.4124  57.735 ])
  correct = 1765.1389457850137
  @test test_newcl("$dir/o5.dcd", unit_cell, correct, lcell)

  unit_cell = [ 80.,  80., 80. ]
  correct = -158.04751357760088
  @test test_newcl("$dir/o6.dcd", unit_cell, correct, lcell)

  # Some triclinic cells

  unit_cell = [ 80.      0.     30.
                30.     80.      0.
                0.     40.     80. ]
  correct = -116.53213607052128
  @test test_newcl("$dir/t1.dcd", unit_cell, correct, lcell)

  unit_cell = [ 50.      0.      0.
                50.     50.      0.
                0.     50.     50. ]
  correct = 32096.48839031735
  @test test_newcl("$dir/t2.dcd", unit_cell, correct, lcell)

  #
  # Check cell list updating routine
  #

  coordinates = getcoor("$dir/o1.dcd")
  unit_cell = [ 50.      0.      0.
                 0.     50.      0.
                 0.      0.     50. ]
  correct = 32230.01699504111
  box = Box(unit_cell, 10., lcell=lcell)
  cl = CellList(coordinates,box)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

  coordinates = getcoor("$dir/o2.dcd")
  unit_cell = [ 80.,  70., 50. ]
  correct = 1093.7225407797744
  box = Box(unit_cell, 10., lcell=lcell)
  cl = UpdateCellList!(coordinates,box,cl)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

  # Test preallocated AuxThreaded struct
  aux = CellListMap.AuxThreaded(cl)
  coordinates = getcoor("$dir/t1.dcd")
  unit_cell = [ 80.      0.     30.
                30.     80.      0.
                0.     40.     80. ]
  correct = -116.53213607052128
  box = Box(unit_cell, 10., lcell=lcell)
  cl = UpdateCellList!(coordinates,box,cl,aux)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

  # Same thing with a different lcell

  lcell=3

  # Some orthorhombic cells

  unit_cell = [ 50.,   50.,  50. ]
  correct = 32230.01699504111
  @test test_newcl("$dir/o1.dcd", unit_cell, correct, lcell)

  unit_cell = [ 80.,  70.,  50. ]
  correct = 1093.7225407797744
  @test test_newcl("$dir/o2.dcd", unit_cell, correct, lcell)

  # Orthorhombic but rotated

  unit_cell = [ 50.      0.     50.
                50.     50.      0.
                 0.     50.     50. ]
  correct = 1724.3195067566828
  @test test_newcl("$dir/o3.dcd", unit_cell, correct, lcell)

  unit_cell= transpose([ 70.7107   0.0      0.0
                         35.3553  61.2372   0.0
                         35.3553  20.4124  57.735 ])
  correct = 1754.0802503953591
  @test test_newcl("$dir/o4.dcd", unit_cell, correct, lcell)

  unit_cell= transpose([ 70.7107   0.0      0.0
                         35.3553  61.2372   0.0
                         35.3553  20.4124  57.735 ])
  correct = 1765.1389457850137
  @test test_newcl("$dir/o5.dcd", unit_cell, correct, lcell)

  unit_cell = [ 80.,  80., 80. ]
  correct = -158.04751357760088
  @test test_newcl("$dir/o6.dcd", unit_cell, correct, lcell)

  # Some triclinic cells

  unit_cell = [ 80.      0.     30.
                30.     80.      0.
                0.     40.     80. ]
  correct = -116.53213607052128
  @test test_newcl("$dir/t1.dcd", unit_cell, correct, lcell)

  unit_cell = [ 50.      0.      0.
                50.     50.      0.
                0.     50.     50. ]
  correct = 32096.48839031735
  @test test_newcl("$dir/t2.dcd", unit_cell, correct, lcell)

  #
  # Check cell list updating routine
  #

  coordinates = getcoor("$dir/o1.dcd")
  unit_cell = [ 50.      0.      0.
                 0.     50.      0.
                 0.      0.     50. ]
  correct = 32230.01699504111
  box = Box(unit_cell, 10., lcell=lcell)
  cl = CellList(coordinates,box)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

  coordinates = getcoor("$dir/o2.dcd")
  unit_cell = [ 80.,  70., 50. ]
  correct = 1093.7225407797744
  box = Box(unit_cell, 10., lcell=lcell)
  cl = UpdateCellList!(coordinates,box,cl)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

  # Test preallocated AuxThreaded struct
  aux = CellListMap.AuxThreaded(cl)
  coordinates = getcoor("$dir/t1.dcd")
  unit_cell = [ 80.      0.     30.
                30.     80.      0.
                0.     40.     80. ]
  correct = -116.53213607052128
  box = Box(unit_cell, 10., lcell=lcell)
  cl = UpdateCellList!(coordinates,box,cl,aux)
  u = map_pairwise!((x,y,i,j,d2,u) -> lj_NE(d2,u),0.0,box,cl)
  @test u ≈ correct

end


immutable Spiral
    m::Int
    n::Int
    cmax::Int
    dir::Array{Array{Int,1},1}
    bdelta::Array{Array{Int,1},1}
end

function Spiral(m::Int, n::Int)
    cmax = m*n
    dir = Array{Int,1}[[0,1], [1,0], [0,-1], [-1,0]]
    bdelta = Array{Int,1}[[0,0,0,1], [-1,0,0,0],
                          [0,-1,0,0], [0,0,1,0]]
    Spiral(m, n, cmax, dir, bdelta)
end

function spiral(m::Int, n::Int)
    0<m&&0<n || error("The matrix dimensions must be positive.")
    Spiral(m, n)
end
spiral(n::Int) = spiral(n, n)

type SpState
    cnt::Int
    dirdex::Int
    cell::Array{Int,1}
    bounds::Array{Int,1}
end

Base.length(sp::Spiral) = sp.cmax
Base.start(sp::Spiral) = SpState(1, 1, [1,1], [sp.n,sp.m,1,1])
Base.done(sp::Spiral, sps::SpState) = sps.cnt > sp.cmax

function Base.next(sp::Spiral, sps::SpState)
    s = sub2ind((sp.m, sp.n), sps.cell[1], sps.cell[2])
    if sps.cell[rem1(sps.dirdex+1, 2)] == sps.bounds[sps.dirdex]
        sps.bounds += sp.bdelta[sps.dirdex]
        sps.dirdex = rem1(sps.dirdex+1, 4)
    end
    sps.cell += sp.dir[sps.dirdex]
    sps.cnt += 1
    return (s, sps)
end


# grids.jl

using Adapt

mutable struct Grid{T}
    # grid_type::String
    nx::Int64
    ny::Int64
    nz::Int64
    Δx::Float64
    Δy::Float64
    Δz::Float64
    x::T
    y::T
    z::T
end

function grid_coords(n1, n2, n3, Δ1, Δ2, Δ3, useCUDA)
    # Create 3D structured grid
    # For now cell centered
    coord1 = Array{Float64,1}(undef, n1)
    for n = 1:n1
        coord1[n] = (n - 0.5) * Δ1
    end

    coord2 = Array{Float64,1}(undef, n2)
    for n = 1:n2
        coord2[n] = (n - 0.5) * Δ2
    end

    coord3 = Array{Float64,1}(undef, n3)
    for n = 1:n3
        coord3[n] = (n - 0.5) * Δ3
    end

    if useCUDA
        # Convert to CUDA Arrays
        println(typeof(coord1))
        coord1 = adapt(CuArray, coord1)
        coord2 = adapt(CuArray, coord2)
        coord3 = adapt(CuArray, coord3)
        println(typeof(coord1))
    end
    return coord1, coord2, coord3
end

# function sigma_grid!(grid, state)
#     nx  = grid[1]
#     nz  = grid[2]
#     Δx  = grid.Δx
#     Δz  = grid.Δz
#     xc  = grid.xc
#     zc  = grid.zc
#     xu  = grid.xu
#     zu  = grid.zu
#     xw  = grid.xw
#     zw  = grid.zw
#
#     dzs = state.Δzₛ
#     dzu = state.Δzᵤ
#
#     # Co-ordinates of the cell centres
#     zc[:,1] = 0.5*dzs[:,1]
#     for i = 1:nx
#         #for j = 1:ny
#         for k = 2:nz
#             xc[i,k] = (i-0.5)*Δx
#             zc[i,k] = zc[i,k-1] + 0.5*(dzs[i,k-1]+dzs[i,k])
#         end
#     end
#
#     # Co-ordinates of the horizontal velocity points
#     zu[:,1] = 0.5*dzu[:,1]
#     for i = 1:nx+1
#         #for j = 1:ny
#         for k = 2:nz
#             xu[i,k] = (i-1)*Δx
#             zu[i,k] = zu[i,k-1] + 0.5*(dzu[i,k-1]+dzu[i,k])
#         end
#     end
#
#     # Co-ordinates of the vertical velocity points
#     zw[:,1] .= 0.0
#     for i = 1:nx
#         #for j = 1:ny
#         for k = 2:nz+1
#             xw[i,k] = (i-0.5)*Δx
#             zw[i,k] = zw[i,k-1] + dzs[i,k-1]
#         end
#     end
#
# end


#=
day17:
- Julia version: 1.7.0
- Author: Paul.Mealor
- Date: 2021-12-17
=#

using Test, AoC

function launch(v, target)
    p = [0,0]
    maxHeight = p[2]
    while (p[1] <= target[1,2]) && (p[2] >= target[2,1])
        if p[2] > maxHeight
            maxHeight = p[2]
        end
        if (target[1,1] <= p[1] <= target[1,2]) && (target[2,1] <= p[2] <= target[2,2])
            return maxHeight
        end
        p = p + v
        v[2] -= 1
        v[1] -= sign(v[1])
    end
    return false
end
@test launch([17,-4], [[20, -10] [30, -5]]) == false
@test launch([7,2], [[20, -10] [30, -5]]) == 3
@test launch([6,3], [[20, -10] [30, -5]]) == 6
@test launch([6,9], [[20, -10] [30, -5]]) == 45
@test launch([23,-10], [[20, -10] [30, -5]]) == 0

exampleLines(17,1) .|> line -> @test launch(parse.(Int, split(line, ",")), [[20, -10] [30, -5]]) !== false

function part1(target)
    maxHeight = 0
    for v_x ∈ 1:target[1,2]
        for v_y ∈ 1:abs(target[2,1])+1
            height = launch([v_x,v_y], target)
            if height > maxHeight
                maxHeight = height
            end
        end
    end
    return maxHeight
end

function part2(target)
    hits = 0
    for v_x ∈ 1:target[1,2]
        for v_y ∈ -abs(target[2,1])-1:abs(target[2,1])+1
            height = launch([v_x,v_y], target)
            if height !== false
                hits += 1
            end
        end
    end
    return hits
end

@test part1([[20, -10] [30, -5]]) == 45
@test part2([[20, -10] [30, -5]]) == 112

@show @time part1([[96, -144] [125,-98]])
@show @time part2([[96, -144] [125,-98]])


using Surrogates

#=
#1D MOE
n = 30
lb = 0.0
ub = 5.0
x = Surrogates.sample(n,lb,ub,UniformSample())
f = x-> 2*x
y = f.(x)
#Standard definition
my_moe = MOE(x,y,lb,ub)
val = my_moe(3.0)

#Local surrogates redefinition
my_local_kind = [InverseDistanceStructure(p = 1.0),
                 KrigingStructure(p=1.0, theta=1.0)]
my_moe = MOE(x,y,lb,ub,k = 2,local_kind = my_local_kind)

#ND MOE
n = 30
lb = [0.0,0.0]
ub = [5.0,5.0]
x = sample(n,lb,ub,LatinHypercubeSample())
f = x -> x[1]*x[2]
y = f.(x)
my_moe_ND = MOE(x,y,lb,ub)
val = my_moe_ND((1.0,1.0))

#Local surr redefinition
my_locals = [InverseDistanceStructure(p = 1.0),
             RandomForestStructure(num_round=1)]
my_moe_redef = MOE(x,y,lb,ub,k = 2,local_kind = my_locals)
=#


include("generar-vecino.jl")

"""
PrimeroMejor(numero_maximo_vecinos_sin_mejora::Int,
            evaluaciones_maximas_funcion_objetivo::Int,
            numero_atributos::Int,
            umbral_atributo::Real,
            F)::Vector{<:Real}

Devuelve mejor vector de peso encontrando de acorde a función
de evaluación `F`
"""
function PrimeroMejor(numero_maximo_vecinos_sin_mejora::Int,
    evaluaciones_maximas_funcion_objetivo::Int,
    numero_atributos::Int,
    umbral_atributo::Real,
    F,
    w = 0)::Vector{<:Real}
    if (w == 0)
        w = rand(numero_atributos) # vector de tamaño atributo uniformement inicializado
    end
    # Ponemos a cero atributos menores que umbral atributo
    w = map(x-> (x<umbral_atributo) ? 0 : x , w)

    iteraciones = 0
    evaluaciones = 0
    F_w = F(w)
    indice::Int = 1
    while iteraciones < numero_maximo_vecinos_sin_mejora &&
        evaluaciones < evaluaciones_maximas_funcion_objetivo
        w_n = GenNeighbourhoodIesimo(w,0, 0.3, indice)
        F_w_vecino = F(w_n)

        if(F_w < F_w_vecino)
            w = w_n
            F_w = F_w_vecino
            iteraciones = 0
        else
            # Si no hay mejora se cambia la exploración del atributo
            indice = indice % (numero_atributos)+1
        end
        iteraciones += 1
        evaluaciones += 1
    end
    return w
end


"""
    module DielectricFunctions

for defining dielectric and pump functions either as piecewise constant
or with user-supplied functions
"""
module DielectricFunctions

export DielectricFunction,
PumpFunction,
piecewise_constant_ε,
piecewise_constant_F

"""
    struct DielectricFunction

---

    DielectricFunction(n1=1,n2=0) -> dielectric
    DielectricFunction(n::Complex) -> dielectric

piecewise constant dielectric function with index given by `n1+1im*n2`

---

    DielectricFunction(ε::Function, parameters::Dict) -> dielectric

ε is any function with the signature ε(x,y,parameters::Dict{Symbol,Float64})
which evaluates to the dielectric at position `x,y`. For example:
    ε(x,y,params) = 2+sin(2πx/params[:period])

---

    (::DielectricFunction)(x,y) -> ε
"""
struct DielectricFunction
    ε::Function
    parameters::Dict{Symbol,Float64}

    DielectricFunction(ε::Function=piecewise_constant_ε,
        parameters::Dict{Symbol,T}=Dict(:n₁=>1,:n₂=>0)
        ) where T<:Number = new(ε,parameters)

    DielectricFunction(n::Complex) = DielectricFunction(real(n),imag(n))

    function DielectricFunction(n1::Real,n2::Real=0.0)
        parameters = Dict(:n₁=>n1,:n₂=>n2)
        return DielectricFunction(piecewise_constant_ε,parameters)
    end

    (df::DielectricFunction)(x::Number,y::Number) = df.ε(x,y,df.parameters)

    function Base.show(io::IO,d::DielectricFunction)
        print(io, "DielectricFunction: ")
        if d.ε==piecewise_constant_ε
            print(io, "piecewise_constant_ε: ",d.parameters[:n₁]+1im*d.parameters[:n₂])
        else
            print(io, d.ε,": ", d.parameters)
        end
    end
end

"""
    struct PumpFunction

---

    PumpFunction(F=0) -> pump

piecewise constant pump function with pump parameter given by `F`

---

    PumpFunction(F::Function, parameters::Dict) -> pump

F is any real function with the signature F(x,y,parameters::Dict{Symbol,Float64})
which evaluates to the pump profile at position `x,y`. For example:
    F(x,y,params) = sin(2πx/params[:period])

---

    (::PumpFunction)(x,y) -> F
"""
struct PumpFunction
    F::Function
    parameters::Dict{Symbol,Float64}

    PumpFunction(F::Function=piecewise_constant_F,
        parameters::Dict{Symbol,T}=Dict(:F=>0)
        ) where T<:Number = new(F,parameters)

    function PumpFunction(F::Real)
        parameters = Dict(:F=>F)
        return PumpFunction(piecewise_constant_F,parameters)
    end

    (pf::PumpFunction)(x::Number,y::Number) = pf.F(x,y,pf.parameters)

    function Base.show(io::IO,d::PumpFunction)
        print(io, "PumpFunction: ")
        if d.F==piecewise_constant_F
            print(io, "piecewise_constant_F: ",d.parameters[:F])
        else
            print(io, d.F,": ", d.parameters)
        end
    end
end

"""
    piecewise_constant_ε(x,y,parameters) -> ε
"""
function piecewise_constant_ε(x::Number,y::Number,parameters::Dict)
    n1 = get!(parameters,:n₁,nothing)
    n2 = get!(parameters,:n₂,nothing)
    return complex(n1,n2)^2
end

"""
    piecewise_constant_F(x,y,parameters) -> F
"""
piecewise_constant_F(x::Number,y::Number,parameters::Dict) = get!(parameters,:F,nothing)

end # module


function centroid!{T <: Real}(o::Vector{T},
	                          a::Vector{T},
	                          b::Vector{T})
	for i in 1:length(o)
		o[i] = (a[i] + b[i]) / 2
	end
	return
end


using BSON
using Flux
using Flux: onehotbatch, onecold
using MLDatasets

Core.eval(Main, :(using Flux)) # hide

function reshape_data(X::AbstractArray{T, 3}, y::AbstractVector) where T
    s = size(X)
    return reshape(X, s[1], s[2], 1, s[3]), reshape(y, 1, :)
end

function train_or_load!(file_name, m, X, y; force=false, kwargs...)

    !isdir(dirname(file_name)) && mkpath(dirname(file_name))

    if force || !isfile(file_name)
        train_model!(m, X, y; file_name=file_name, kwargs...)
    else
        m_loaded = BSON.load(file_name)[:m]
        Flux.loadparams!(m, params(m_loaded))
    end
end

function load_data(dataset; onehot=false, T=Float32)
    classes = 0:9
    X_train, y_train = reshape_data(dataset.traindata(T)...)
    X_test, y_test = reshape_data(dataset.testdata(T)...)
    y_train = T.(y_train)
    y_test = T.(y_test)

    if onehot
        y_train = onehotbatch(y_train[:], classes)
        y_test = onehotbatch(y_test[:], classes)
    end

    return X_train, y_train, X_test, y_test
end

using Plots

plot_image(x::AbstractArray{T, 2}) where T = plot(Gray.(x'), axis=nothing)

function plot_image(x::AbstractArray{T, 4}) where T
    @assert size(x,4) == 1
    plot_image(x[:,:,:,1])
end

function plot_image(x::AbstractArray{T, 3}) where T
    @assert size(x,3) == 1
    plot_image(x[:,:,1])
end

T = Float32
dataset = MLDatasets.MNIST

X_train, y_train, X_test, y_test = load_data(dataset; T=T, onehot=true)

m = Chain(
    Conv((2,2), 1=>16, sigmoid),
    MaxPool((2,2)),
    Conv((2,2), 16=>8, sigmoid),
    MaxPool((2,2)),
    flatten,
    Dense(288, size(y_train,1)), softmax)

file_name = joinpath("data", "mnist_sigmoid.bson")
train_or_load!(file_name, m, X_train, y_train)

ii1 = findall(onecold(y_train, 0:9) .== 1)[1:5]
ii2 = findall(onecold(y_train, 0:9) .== 9)[1:5]

for qwe = 0:9
    ii0 = findall(onecold(y_train, 0:9) .== qwe)[1:5]

    p0 = [plot_image(X_train[:,:,:,i:i][:,:,1,1]) for i in ii0]
    p1 = [plot_image((m[1:2](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii0]
    p2 = [plot_image((m[1:4](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii0]

    p = plot(p0..., p1..., p2...; layout=(3,5))
    display(p)
end

p0 = [plot_image(X_train[:,:,:,i:i][:,:,1,1]) for i in ii1]
p1 = [plot_image((m[1:2](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii1]
p2 = [plot_image((m[1:4](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii1]

plot(p0..., p1..., p2...; layout=(3,5))

p0 = [plot_image(X_train[:,:,:,i:i][:,:,1,1]) for i in ii2]
p1 = [plot_image((m[1:2](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii2]
p2 = [plot_image((m[1:4](X_train[:,:,:,i:i]))[:,:,1,1]) for i in ii2]

plot(p0..., p1..., p2...; layout=(3,5))

for i in 1:length(m)
    println(size(m[1:i](X_train[:,:,:,1:1])))
end


# Find some reasonable values for tick marks.
#
# This is basically Wilkinson's ad-hoc scoring method that tries to balance
# tight fit around the data, optimal number of ticks, and simple numbers.
#
# Args:
#   x_min: minimum value occuring in the data.
#   x_max: maximum value occuring in the data.
#
# Returns:
#   A Float64 vector containing tick marks.
#
function optimize_ticks(x_min::Float64, x_max::Float64)
    if x_min == x_max
        return [x_min]
    end

    # TODO: these should perhaps be part of the theme
    const Q = {(1,1), (5, 0.9), (2, 0.7), (25, 0.5), (3, 0.2)}
    const n = length(Q)
    const k_min   = 2
    const k_max   = 10
    const k_ideal = 5

    xspan = x_max - x_min
    z = ceil(log10(xspan))

    high_score = -Inf
    z_best = 0.0
    k_best = 0.0
    r_best = 0.0
    q_best = 0.0

    while k_max * 10^(z+1) > xspan
        for k in k_min:k_max
            for (q, qscore) in Q
                span = (k - 1) * q * 10^z
                if span < xspan
                    continue
                end

                r = ceil((x_max - span) / (q*10^z))
                while r*q*10^z <= x_min
                    has_zero = r <= 0 && abs(r) < k

                    # simplicity
                    s = has_zero ? 1.0 : 0.0

                    # granularity
                    g = 0 < k < 2k_ideal ? 1 - abs(k - k_ideal) / k_ideal : 0.0

                    # coverage
                    c = 1.5 * xspan/span

                    score = (1/4)g + (1/6)s + (1/3)c + (1/4)qscore

                    # strict limits on coverage
                    if span >= 2*xspan || span < xspan
                        score -= 1000
                    end

                    if score > high_score
                        (q_best, r_best, k_best, z_best) = (q, r, k, z)
                        high_score = score
                    end
                    r += 1
                end
            end
        end
        z -= 1
    end

    S = Array(Float64, int(k_best))
    for i in 0:(k_best - 1)
        S[i+1] = (r_best + i) * q_best * 10^z_best
    end

    S
end


struct BinaryRelation{O, S, T}
    x::S
    y::T
end

BinaryRelation{O}(x::X, y::Y) where {O,X,Y} = BinaryRelation{O,X,Y}(x, y)

x = DiscreteVariable(:x, 1:5)
y = DiscreteVariable(:y, 2:6)

# BinaryRelation(==, x, y)
# BinaryRelation(==, x, 1)

# Use NodeVariables to break up binary relations like x + y ~ 3 and ternary such as x + y <= z
struct NodeVariable <: Var
    name
    op
    x
    y
    domain
    booleans
    varmap
end

Base.show(io::IO, var::NodeVariable) = print(io, "$(var.name) = $(var.op)($(var.x), $(var.y))")

find_domain(op, x, y) = sort(collect(Set([op(i, j) for i in domain(x) for j in domain(y)])))

domain(v::Var) = v.domain
booleans(v::Var) = v.booleans

domain(x::Real) = x:x

"All booleans inside an expression that must be added to the encoding"
recursive_booleans(v::NodeVariable) = v.booleans ∪ booleans(v.x) ∪ booleans(v.y)

function NodeVariable(op, x, y, name=gensym())
    domain = find_domain(op, x, y)
    booleans = [Variable(name, i) for i ∈ indices(domain)]
    varmap = Dict(i => v for (i, v) in zip(domain, booleans))
    NodeVariable(name, op, x, y, domain, booleans, varmap)
end

Base.:+(x::Var, y) = NodeVariable(+, x, y)
Base.:+(x, y::Var) = NodeVariable(+, x, y)

Base.:*(x::Var, y) = NodeVariable(*, x, y)
Base.:*(x, y::Var) = NodeVariable(*, x, y)

Base.:^(x::Var, y) = NodeVariable(^, x, y)
Base.:^(x, y::Var) = NodeVariable(^, x, y)

function clauses(var::NodeVariable)

    op = var.op
    x = var.x
    y = var.y

    clauses = exactly_one(var.booleans)

    # deal with 2 + X
    if x isa Real
        for j in domain(var.y)
            value = op(x, j)
            # y == j => var == op(x, j)
            push!(clauses, ¬(y.varmap[j]) ∨ var.varmap[value])

        end

    # deal with x + 2
    elseif y isa Real
        for i in domain(var.x)
            value = op(i, y)
            # x == i => var == op(i, y)
            push!(clauses, ¬(x.varmap[i]) ∨ var.varmap[value])

        end

    else

        # neither x nor y is real
        for i in domain(var.x)
            for j in domain(var.y)
                value = op(i, j)
                # x == i && y == j => var == op(i, j)
                push!(clauses, ¬(x.varmap[i]) ∨ ¬(y.varmap[j]) ∨ var.varmap[value])
            end
        end

    end

    return clauses

end

Base.:(==)(v::Var, w) = BinaryRelation{==}(v, w)
Base.:(<=)(v::Var, w) = BinaryRelation{<=}(v, w)

"Encode relation like x == 1"
function encode(rel::BinaryRelation{==, <:Var, <:Real})

    x = rel.x
    y = rel.y

    if y ∉ domain(x)
        error("$y is not in the domain of $x")
    end

    boolean = x.varmap[y]

    return [boolean]
end

"Encode relation like x != 1"
function encode(rel::BinaryRelation{!=, <:Var, <:Real})

    x = rel.x
    y = rel.y

    if y ∈ domain(x)
        boolean = ¬(x.varmap[y])
    end

    return [boolean]
end

"Encode relation like x == y"
function encode(rel::BinaryRelation{==, <:Var, <:Var})
    x = rel.x
    y = rel.y

    clauses = []

    for i in domain(x)
        if i in domain(y)
        # (x == i) => (y == i)
            push!(clauses, ¬(x.varmap[i]) ∨ (y.varmap[i]))

        else
            push!(clauses, ¬(x.varmap[i]))
        end
    end

    for i in domain(y)
        if i in domain(x)
        # (y == i) => (x == i)
            push!(clauses, ¬(y.varmap[i]) ∨ (x.varmap[i]))

        else
            push!(clauses, ¬(y.varmap[i]))
        end
    end

    return clauses
end

function encode(rel::BinaryRelation{!=, <:Var, <:Var})
    x = rel.x
    y = rel.y

    clauses = []

    for i in domain(x)
        if i in domain(y)
        # (x == i) => (y != i)
            push!(clauses, ¬(x.varmap[i]) ∨ ¬(y.varmap[i]))
        end
    end

    return clauses
end

function parse_expression(varmap, ex)
    op = operation(ex)
    args = arguments(ex)

    new_args = parse_expression.(Ref(varmap), args)

    return op, new_args
end

parse_expression(varmap, ex::Term) = varmap[ex]
parse_expression(varmap, ex::Sym) = varmap[ex]
parse_expression(varmap, ex::Num) = parse_expression(varmap, value(ex))
parse_expression(varmap, ex::Real) = ex

"Parse a symbolic expression into a relation"
function parse_relation(varmap, ex)
    op = operation(ex)
    args = arguments(ex)

    lhs = args[1]
    rhs = args[2]

    return BinaryRelation{op}(parse_expression(varmap, lhs), parse_expression(varmap, rhs))

end

# function parse_relation!(varmap, ex::Equation)

#     varmap[ex.lhs] = parse_expression(varmap, ex.rhs)

# end

is_variable(ex) = !istree(ex) || (operation(ex) == getindex)

function process(constraint)

    new_constraints = []

    constraint2 = value(constraint)

    op = operation(constraint2)
    args = arguments(constraint2)

    lhs = args[1]
    rhs = args[2]

    intermediates_generated = false

    if !is_variable(lhs)
        lhs = ReversePropagation.cse_equations(lhs)

        append!(new_constraints, lhs[1])
        lhs = lhs[2]

    end

    if !is_variable(rhs)
        rhs = ReversePropagation.cse_equations(rhs)

        append!(new_constraints, rhs[1])
        rhs = rhs[2]

    end

    # @show new_constraints
    # @show op, lhs, rhs
    push!(new_constraints, op(lhs, rhs))

    return new_constraints

end

process(constraints::Vector) = reduce(vcat, process(constraint) for constraint in constraints)

function parse_constraint!(domains, ex)

    expr = value(ex)
    op = operation(expr)

    new_constraints = []

    if op == ∈   # assumes right-hand side is an explicit set specifying the doomain

        args = arguments(expr)

        var = args[1]
        domain = args[2]

        domains[var] = domain

    else
        new_constraints = process(expr)

        println()
        # @show expr
        # @show new_constraints

    end

    return domains, new_constraints

end

function parse_constraints(constraints)
    # domains = Dict(var => -Inf..Inf for var in vars)

    additional_vars = []
    domains = Dict()
    all_new_constraints = []  # excluding domain constraints

    for constraint in constraints

        # binarize constraints:
        domains, new_constraints = parse_constraint!(domains, value(constraint))

        for statement in new_constraints
            if statement isa Assignment
                push!(additional_vars, statement.lhs)
            end
        end

        append!(all_new_constraints, new_constraints)
    end

    # append!(additional_vars, keys(domains))

    return domains, all_new_constraints, additional_vars
end

Base.isless(x::Sym, y::Sym) = isless(x.name, y.name)
Base.isless(x::Term, y::Term) = isless(string(x), string(y))
Base.isless(x::Term, y::Sym) = isless(string(x), string(y))
Base.isless(x::Sym, y::Term) = isless(string(x), string(y))

struct ConstraintSatisfactionProblem
    original_vars
    additional_vars   # from binarizing
    domains
    constraints
end

function ConstraintSatisfactionProblem(constraints)

    domains, new_constraints, additional_vars = parse_constraints(constraints)
    @show keys(domains)
    vars = sort(identity.(keys(domains)))
    additional_vars = sort(identity.(additional_vars))

    return ConstraintSatisfactionProblem(vars, additional_vars, domains, new_constraints)
end

struct DiscreteCSP
    original_vars
    additional_vars
    constraints
    varmap   # maps symbolic variables to the corresponding DiscreteVariable
end

DiscreteCSP(constraints) = DiscreteCSP(ConstraintSatisfactionProblem(constraints))

function DiscreteCSP(prob::ConstraintSatisfactionProblem)
    variables = []
    varmap = Dict()

    new_constraints = []

    for (var, domain) in prob.domains
        variable = DiscreteVariable(var, domain)
        push!(variables, variable)
        push!(varmap, var => variable)
    end

    for constraint in prob.constraints
        constraint = value(constraint)
        if constraint isa Assignment

            lhs = constraint.lhs

            op, new_args = parse_expression(varmap, constraint.rhs)   # makes a NodeVariable

            # @show op, new_args

            variable = NodeVariable(op, new_args[1], new_args[2], lhs)

            push!(variables, variable)
            push!(varmap, lhs => variable)

        else
            # @show constraint
            push!(new_constraints, parse_relation(varmap, constraint))
        end
    end

    original_vars = [varmap[var] for var in prob.original_vars]
    additional_vars = [varmap[var] for var in prob.additional_vars]

    return DiscreteCSP(original_vars, additional_vars, new_constraints, varmap)
end

function encode(prob::DiscreteCSP)
    all_variables = []
    all_clauses = []

    variables = Any[prob.original_vars; prob.additional_vars]

    domains = Dict(v.name => v.domain for v in variables)

    for var in variables
        append!(all_variables, var.booleans)
        append!(all_clauses, clauses(var))
    end

    # @show prob.constraints

    for constraint in prob.constraints
        # @show constraint
        append!(all_clauses, encode(constraint))
    end

    # @show all_variables
    # @show all_clauses

    return SymbolicSATProblem(identity.(all_variables), identity.(all_clauses))
end

function solve(prob::DiscreteCSP)
    symbolic_sat_problem = encode(prob)

    status, result_dict = solve(symbolic_sat_problem)

    (status == :unsat) && return status, missing

    # variables = Any[prob.original_vars; prob.additional_vars]
    variables = prob.original_vars

    return status, decode(prob, result_dict)
end

decode(prob::DiscreteCSP, result_dict) = Dict(v.name => decode(result_dict, v) for v in prob.original_vars)

function all_solutions(prob::DiscreteCSP)
    sat_problem = encode(prob)
    # variables = Any[prob.original_vars; prob.additional_vars]
    variables = prob.original_vars

    prob2 = encode(prob)
    prob3 = prob2.p   # SATProblem

    sat_solutions = all_solutions(prob3)

    if isempty(sat_solutions)
        return sat_solutions
    end

    solutions = []

    for solution in sat_solutions
        push!(solutions, decode(prob, decode(prob2, solution)))
    end

    return solutions

end